Add support for Atom Indices Import

[bradyajohnston]

Add support for Atom Indices Import

[BradyAJohnston]

If somebody can come up with a parser (one probably exists in a function in MDAnalysis or similar somewhere) that can parse the .ndx files which provide the atom groups for GROMACS etc, then this can become at attribute on import and allow for inclusion of the atom groups inside of Molecular Nodes.

[germanbarcenas]

Hi Brady! I've never contributed before, so I don't know how to do it necessarily. I did however find this other project that seems to accomplish what you need.

https://gist.github.com/jbarnoud/09a22ac41fbd5969af51458eb89317e7

If you'd like, you can give me a little more specific guidance and I can try to add this, or else it looks like it has all the info you might need to do it yourself.

[BradyAJohnston]

oooh thanks for pointing this out! I hadn't looked into it deeply yet, but it seems like this should be straightforward to implement.

Happy to support a PR, if you'd like to have a go at implementing it. I'd like to prioritise encouraging community contributions over doing everything myself.

[BradyAJohnston]

In terms of contributing, I've got some minor details in the README about how to do it in terms of building your own Molecular Nodes. You'll want VSCode and a couple of specific addons for Blender addon development.

The function that does the MDAnalysis importing is here:

MolecularNodes/MolecularNodes/md.py

Lines 65 to 246 in 7979e38

	def load_trajectory(file_top,
	file_traj,
	md_start = 1,
	md_end = 50,
	md_step = 1,
	world_scale = 0.01,
	include_bonds = False,
	del_solvent = False,
	selection = "not (name H* or name OW)",
	name = "default"
	):
	import MDAnalysis as mda
	import MDAnalysis.transformations as trans
	# initially load in the trajectory
	univ = mda.Universe(file_top, file_traj)
	# separate the trajectory, separate to the topology or the subsequence selections
	traj = univ.trajectory[md_start:md_end:md_step]
	# if there is a non-blank selection, apply the selection text to the universe for
	# later use. This also affects the trajectory, even though it has been separated earlier
	if selection != "":
	try:
	univ = univ.select_atoms(selection)
	except:
	warnings.warn(f"Unable to apply selection: '{selection}'. Loading entire topology.")
	# Try and extract the elements from the topology. If the universe doesn't contain
	# the element information, then guess based on the atom names in the toplogy
	try:
	elements = univ.atoms.elements.tolist()
	except:
	elements = [mda.topology.guessers.guess_atom_element(x) for x in univ.atoms.names]
	# determin the bonds for the structure
	if hasattr(univ, 'bonds') and include_bonds:
	bonds = univ.bonds.indices
	else:
	bonds = []
	# create the initial model
	mol_object = create_object(
	name = name,
	collection = coll_mn(),
	locations = univ.atoms.positions * world_scale,
	bonds = bonds
	)
	## add the attributes for the model
	# The attributes for the model are initially defined as single-use functions. This allows
	# for a loop that attempts to add each attibute by calling the function. Only during this
	# loop will the call fail if the attribute isn't accessible, and the warning is reported
	# there rather than setting up a try: except: for each individual attribute which makes
	# some really messy code.
	def att_atomic_number():
	atomic_number = np.array(list(map(
	# if getting the element fails for some reason, return an atomic number of -1
	lambda x: data.elements.get(x, {"atomic_number": -1}).get("atomic_number"),
	np.char.title(elements)
	)))
	return atomic_number
	def att_vdw_radii():
	try:
	vdw_radii = np.array(list(map(
	lambda x: mda.topology.tables.vdwradii.get(x, 1),
	np.char.upper(elements)
	)))
	except:
	# if fail to get radii, just return radii of 1 for everything as a backup
	vdw_radii = np.ones(len(univ.atoms.names))
	warnings.warn("Unable to extract VDW Radii. Defaulting to 1 for all points.")
	return vdw_radii * world_scale
	def att_res_id():
	return univ.atoms.resnums
	def att_res_name():
	res_names = np.array(list(map(lambda x: x[0: 3], univ.atoms.resnames)))
	res_numbers = np.array(list(map(
	lambda x: data.residues.get(x, {'res_name_num': 0}).get('res_name_num'),
	res_names
	)))
	return res_numbers
	def att_b_factor():
	return univ.atoms.tempfactors
	def att_chain_id():
	chain_id = univ.atoms.chainIDs
	chain_id_unique = np.unique(chain_id)
	chain_id_num = np.array(list(map(lambda x: np.where(x == chain_id_unique)[0][0], chain_id)))
	return chain_id_num
	# returns a numpy array of booleans for each atom, whether or not they are in that selection
	def bool_selection(selection):
	return np.isin(univ.atoms.ix, univ.select_atoms(selection).ix).astype(bool)
	def att_is_backbone():
	return bool_selection("backbone or nucleicbackbone")
	def att_is_alpha_carbon():
	return bool_selection('name CA')
	def att_is_solvent():
	return bool_selection('name OW or name HW1 or name HW2')
	def att_is_nucleic():
	return bool_selection('nucleic')
	def att_is_peptide():
	return bool_selection('protein')
	attributes = (
	{'name': 'atomic_number', 'value': att_atomic_number, 'type': 'INT', 'domain': 'POINT'},
	{'name': 'vdw_radii', 'value': att_vdw_radii, 'type': 'FLOAT', 'domain': 'POINT'},
	{'name': 'res_id', 'value': att_res_id, 'type': 'INT', 'domain': 'POINT'},
	{'name': 'res_name', 'value': att_res_name, 'type': 'INT', 'domain': 'POINT'},
	{'name': 'b_factor', 'value': att_b_factor, 'type': 'float', 'domain': 'POINT'},
	{'name': 'chain_id', 'value': att_chain_id, 'type': 'INT', 'domain': 'POINT'},
	{'name': 'is_backbone', 'value': att_is_backbone, 'type': 'BOOLEAN', 'domain': 'POINT'},
	{'name': 'is_alpha_carbon', 'value': att_is_alpha_carbon, 'type': 'BOOLEAN', 'domain': 'POINT'},
	{'name': 'is_solvent', 'value': att_is_solvent, 'type': 'BOOLEAN', 'domain': 'POINT'},
	{'name': 'is_nucleic', 'value': att_is_nucleic, 'type': 'BOOLEAN', 'domain': 'POINT'},
	{'name': 'is_peptide', 'value': att_is_peptide, 'type': 'BOOLEAN', 'domain': 'POINT'},
	)
	for att in attributes:
	# tries to add the attribute to the mesh by calling the 'value' function which returns
	# the required values do be added to the domain.
	try:
	add_attribute(mol_object, att['name'], att['value'](), att['type'], att['domain'])
	except:
	warnings.warn(f"Unable to add attribute: {att['name']}.")
	# add the custom selections if they exist
	custom_selections = bpy.context.scene.trajectory_selection_list
	if custom_selections:
	for sel in custom_selections:
	try:
	add_attribute(
	object=mol_object,
	name=sel.name,
	data=bool_selection(sel.selection),
	type = "BOOLEAN",
	domain = "POINT"
	)
	except:
	warnings.warn("Unable to add custom selection: {}".format(sel.name))
	# create the frames of the trajectory in their own collection to be disabled
	coll_frames = bpy.data.collections.new(name + "_frames")
	coll_mn().children.link(coll_frames)
	add_occupancy = True
	for ts in traj:
	frame = create_object(
	name = name + "_frame_" + str(ts.frame),
	collection = coll_frames,
	locations = univ.atoms.positions * world_scale
	)
	# adds occupancy data to each frame if it exists
	# This is mostly for people who want to store frame-specific information in the
	# b_factor but currently neither biotite nor MDAnalysis give access to frame-specific
	# b_factor information. MDAnalysis gives frame-specific access to the `occupancy`
	# so currently this is the only method to get frame-specific data into MN
	# for more details: https://github.com/BradyAJohnston/MolecularNodes/issues/128
	if add_occupancy:
	try:
	add_attribute(frame, 'occupancy', ts.data['occupancy'])
	except:
	add_occupancy = False
	# disable the frames collection from the viewer
	bpy.context.view_layer.layer_collection.children[coll_mn().name].children[coll_frames.name].exclude = True
	return mol_object, coll_frames

For this to be used, I imagine adding an extra argument to the function maybe file_idx which defaults to None. If a file is given, then use the code above to create a list of boolean selections, each of which can be applied as attributes, similarly to:

MolecularNodes/MolecularNodes/md.py

Lines 206 to 218 in 7979e38

	custom_selections = bpy.context.scene.trajectory_selection_list
	if custom_selections:
	for sel in custom_selections:
	try:
	add_attribute(
	object=mol_object,
	name=sel.name,
	data=bool_selection(sel.selection),
	type = "BOOLEAN",
	domain = "POINT"
	)
	except:
	warnings.warn("Unable to add custom selection: {}".format(sel.name))

The function is called in this operator:

MolecularNodes/MolecularNodes/ui.py

Line 78 in 7979e38

class MOL_OT_Import_Protein_MD(bpy.types.Operator):

so it will need to be updated, along with adding a new bpy.context.scene.* property for the file, and displaying that in the UI.

If it's a bit much, then I'm happy to go about it as well. The code base I think still requires a bit of a refactor, and blender addons can be a bit overwhelming if you haven't worked with them before as well.

[germanbarcenas]

I'll give it a go as well. This isn't top priority for me at the moment, but I really like this project, and it's getting me to learn some more git and Blender, which I'm super happy about!

[BradyAJohnston]

If you have questions about approach to take, please let me know! As a brief outline, you'll want to:

1. Fork the repo into your github account.
2. Download your fork onto you local computer.
3. Tinker around with making the changes.
4. Push your changes back to your fork of the repo.
5. When you are ready to get some feedback on code / ideas, create a pull request to merge your fork back into my repo. That will allow me to do code review & give feedback.
6. You can continue to make changes on your fork, which will be reflected in the pull request.
7. Once it's all good, I'll merge it in to main and it's part of the package!

[germanbarcenas]

Brady thanks for the generous guide! I should have some time next week to start poking at this.