Adds 'Live' MDAnalysis Selections

[BradyAJohnston]

Completely re-works how MDA universes are added and updated under the hood, while adding a panel for MDA selections. These are 'Live' selections which are updated when the selection string is updated and on frame change. Now that the universe linking is working really well - the 'in memory' option is removed.

Fixes #475, #455

Enables functionality to implement #440

The selections appear as boolean attributes on the mesh. It is a True / False for whether the atoms are selected via universe.select_atoms(selection) to the imported universe.

New features:

• Live selections panel
• Options for enable / disable the updating of the selection, as well as being periodic for geometry selections
• Whether to interpolate when using subframes. Non-interpolated means the trajectory skips every n frames before moving onto the next one.

Todo:

• Implement MNSession for tracking imported molecules / universes / ensembles
• change tests to work with new updating system
• Handle deletion of custom selections so they remove the attribute also
• When adding and deleting selections, update and remove the corresponding attribute on the object / mesh
• Have a UUID based linking between blender objects and universes for updating.
• Add tests for updating selections

Selections

The selections are defined by creating an AtomGroup and checking which universe atoms are inside of that group. Only on an update to the selection_str is a new AtomGroup created. For regular frame changes the same group is used, and the selection will updating internally if updating is selected.

• name: Name that will become the attribute name on the object
• selection_str: used inside of mda.Universe.select_atoms(selection_str) for defining the atom group
• updating: whether the selection is updated on change of frame or change of selection_str
• periodic: for geomtric selections whether to look into different periodic images when searching

Proximity selections are a lot faster inside of Geometry Nodes / Molecular Nodes - so combining a selection with the proximity selection node is the fastest way to achieve distance selections. Need to improve proximity selection node related to #524

MNSession

Class for tracking the imported objects in the scene. Each dictionary for the correpsonding items uses a uuid that is generated when imported - which is also stored on the corresponding blender object when created. This allows for cleaner lookup of objects and universes, without relying on the bpy.data.objects['name'] approach which is very easy to break (a frustrating limitation of Blender).

There is a dictionary for each type of item, (Molecule, MNUniverse and Ensemble). They are then available to the session after import is finished. The session won't persist between .blend save and load, so on save we pickle the file and save it next to the .blend file as .blend.MNSession. When loading a .blend file the pickle is loaded again and everything that was in the old session is appended to the new session. You can't pickle Blender's Object or Collection objects, so just before pickling a reference to their names are saved and then on load they are retrieved using the reference.

This session system now handles refreshing for universes, molecules and ensembles, so that it is all united in one location.

CleanShot.2024-06-20.at.22.23.12.mp4

[BradyAJohnston]

@rbdavid this is an example of using the connected MDA Universe that is maintained between sessions. This won't be applicable to other import methods, but a similar approach could be done.

[rbdavid]

This is awesome! From the video, it seems like you can add any number of selection atom groups where custom_selection is the named boolean attribute associated with the selection string. So this should enable a user to make multiple selections and use different materials, representation styles, etc to highlight those. Awesome work! Let me know if you want a second pair of eyes on the code.

[BradyAJohnston]

Yes the idea is that you can add any number of 'custom selections'. Each listed custom selection becomes a named attribute on the mesh, with a boolean selection that is generated through MDAnalysis. This can then be used inside of geometry nodes for styling / selecting etc.

[BradyAJohnston]

Related to #440 - I'm thinking we could manage universe transformations (smoothing, translations etc) from a similar list.

[BradyAJohnston]

@rbdavid if you did want to have a look over the code - I'm pretty happy with how this implementation is now. I explain mostly in the edited PR, but along with adding the 'live' selections - this has now completely overhauled how MDA Universes are imported and tracked in a global MNSession object.

[rbdavid]

Looks great! Sorry for the radio silence from me. I've had many distractions pushing my blender work to lower priority. I'd like to get back to finishing my visualizations as well as helping in development where possible. Hopefully things clear up for me soon.

I likely won't be able to test out the MNSession implementation for at least a week, maybe longer. Also, I'm still stuck on v4.0.9 or so for my visualization work -- whomp whomp. But from the video, the selection strings look intuitive and helpful.

One question, and this might show how out-of-date my knowledge is, the MNSession is only available when the user imports a structure via the MD trajectory import field? Importing from a .pdb or .mmcif file still uses biotite?

[BradyAJohnston]

No worries @rbdavid . I think I'll merge this into the #501 branch and if you wanted to do a more broad review there that would be useful.

Downloading from the PDB, opening .pdb / .mmcif files are still handled by biotite, but they are also store & restored via the MNSession object under MNSession.molecules. Now that I've figured out how to reliably load, save and restore the session, we can track any arbitrary extra python objects. I haven't implemented any extra features for these other Molecule objects yet, but this enables all sorts of functionality that would rely upon the structure file that is currently not possible to implement in geometry nodes. The text-based input for selections is via MDAnalysis, so those other structures imported other ways can't use the selections. We could potentially 'under the hood' create a universe object that corresponds to the structures though, which would enable that kind of text input. This would unify the backend but I don't know if it would be wort it.

class MNSession:
    def __init__(self) -> None:
        self.molecules: Dict[str, Molecule] = {}
        self.universes: Dict[str, MNUniverse] = {}
        self.ensembles: Dict[str, Ensemble] = {}