def select_atoms(mdsys):
    """ Return the selection of atoms that are of interest in the simulation
    """
    # Change this selection in case you have different lipids, and update the
    # onehot encoder dictionary accordingly
    asel = mdsys.select_atoms("resname DPPC or resname DOPC or resname DOPX or \
                               resname DPPX or resname CHOL or resname DIPC or resname DIPX")
    # asel_tails = mdsys.select_atoms("(resname DPPC or resname DOPC or resname DOPX or resname DPPX or resname CHOL or resname DIPC or resname DIPX) and not (name GL1 or name GL2 or name NC3 or name PO4 or name ROH)")
    # asel_heads = mdsys.select_atoms("(resname DPPC or resname DOPC or resname DOPX or resname DPPX or resname CHOL or resname DIPC or resname DIPX) and (name GL1 or name GL2 or name NC3 or name PO4 or name ROH)")

    return asel