unify mobile,ref argument unpacking

[kain88-de]

I've noticed that I have to write a mobile/reference unpacking and checking everytime I write my own analysis code. Today I wrote a small helper for my own code that is reusable and helps me with that. I currently don't have time myself to integrate it into the analysis module where it would be appropriate, like RMSD and AlignTraj. This shouldn't be hard to add if someone wants to add this early. I also have a collection of unittests for these functions (written for pytest though so they need to be converted to nose).

def select_atoms_wrapper(ag, sel):
    """wrapper for `select_atoms` that expands sel

    Parameters
    ----------
    ag : mda.AtomGroup, mda.Universe
        AtomGroup from which new atoms are selected
    sel : str, list, tuple, mda.AtomGroup
        a selection.

    Returns
    -------
    sel : mda.AtomGroup
        Atomgroup defined by selection
    """
    if isinstance(sel, str):
        return ag.select_atoms(sel)
    elif isinstance(sel, (list, tuple)):
        return ag.select_atoms(*sel)
    else:
        return sel


def parse_common_selection(universe, mobile, ref=None, strict=False):
    """Helper function to unpack mobile and ref selection used in common
    analysis algorithms of MDAnalysis. This also ensures that mobile and
    ref use different universe to ensure that iterating over the universe
    doesn't change the reference.

    Parameters
    ----------
    universe : mda.Universe
        universe from which to select the mobile atomgroup
    mobile : str, list, tuple, mda.AtomGroup
        selection for mobile atomgroup
    ref : str, list, tuple, mda.AtomGroup (optional)
        selection for reference atomgroup. If `None` ref will use the same
        selection as `mobile` from a new universe.
    strict : bool (optional)
        check if `mobile` and `ref` are refering to exactly the same atoms

    Returns
    -------
    mobile : mda.AtomGroup
        Atomgroup that refers to the given universe
    ref : mda.AtomGroup
        reference atomgroup that does not refer to universe

    Raises
    ------
    RuntimeError

    """
    mobile = select_atoms_wrapper(universe, mobile)

    if mobile.universe != universe:
        raise RuntimeError("mobile selection doesn't refer to given universe")

    if ref is None:
        u2 = mda.Universe(universe.filename, universe.trajectory.filename)
        ref = u2.select_atoms(*('bynum {}'.format(i + 1) for i in mobile.ids))
    else:
        ref = select_atoms_wrapper(universe, ref)

    if mobile.n_atoms != ref.n_atoms:
        raise RuntimeError("mobile and ref have different number of atoms")

    if strict:
        resids_ok = np.all(mobile.resids == ref.resids)
        segids_ok = np.all(mobile.segids == ref.segids)
        ids_ok = np.all(mobile.ids == ref.ids)
        names_ok = np.all(mobile.names == ref.names)
        resnames_ok = np.all(mobile.resnames == ref.resnames)
        if not (resids_ok and segids_ok and ids_ok and names_ok and
                resnames_ok):
            raise RuntimeError("mobile and ref aren't strictly the same")

    if universe == ref.universe:
        u2 = mda.Universe(universe.filename, universe.trajectory.filename)
        ref = u2.select_atoms(*('bynum {}'.format(i + 1) for i in ref.ids))

    return mobile, ref

This allows unpacking of mobile ref arguments in a very flexible way

m, r = parse_common_selection(u, ca)
m, r = parse_common_selection(u, 'name CA')
m, r = parse_common_selection(u, ('bynum 1', 'bynum 2'))
m, r = parse_common_selection(u, ['bynum 1', 'bynum 2'])
m, r = parse_common_selection(u, ca, ca)
m, r = parse_common_selection(u, ('bynum 1', 'bynum 2'), ('bynum 2', 'bynum 3'))
m, r = parse_common_selection(u, 'bynum 1-3', 'bynum 2-4')

#combinations that will raise an error
# wrong universe for mobile
u2 = mda.Universe(PSF, DCD)
m, r = consume_args(u, u2.atoms.CA)
# atomgroups are not of equal size
m, r = consume_args(u, 'all', 'name CA') 
# atomgroups are of equal size but refer to different atoms.
m, r = consume_args(u, ('bynum 1', 'bynum 2'), ('bynum 2', 'bynum 3'), strict=True) 

This can unify code in the analysis that takes a mobile and reference argument.

def some_analysis(u, mobile, ref=None):
    u.trajectory[0]
    mobile, ref = parse_common_selection(u, mobile, ref)

# or if the analysis requires complete equality in the atomgroups
def some_analysis_strict(u, mobile, ref=None):
    u.trajectory[0]
    mobile, ref = parse_common_selection(u, mobile, ref, strict=True)

Here the implicit assumption is that if no reference is given we use the mobile selection of the first frame as a reference. I also find it nice that this guarantees that ref won't be changed when I iterate over the u.trajectory.

It would be nice if the behavior of select_atoms_wrapper could be integrated into select_atoms directly if wanted.

[orbeckst]

Overall good idea!

I just wouldn't raise RuntimeError but ValueError (after all, the problem is that the user provided the wrong input). I'd also use ValueError over MDAnalysis.exceptions.SelectionError because the latter can report on actual problems with the selection string.

RuntimeError is a catch-all but I think here we can be more precise.

[kain88-de]

This code is now published and unit-tested in pydiffusion. It worked for me quite well last year. If there is interested I can look into including it into MDAnalysis and see if the analysis classes can be modified to use it.

https://github.com/bio-phys/pydiffusion/blob/master/pydiffusion/util/mda.py