diff --git a/gmso/formats/mol2.py b/gmso/formats/mol2.py index 2112062e2..5ed37ce87 100644 --- a/gmso/formats/mol2.py +++ b/gmso/formats/mol2.py @@ -6,6 +6,7 @@ import unyt as u from gmso import Atom, Bond, Box, Topology +from gmso.abc.abstract_site import MoleculeType, ResidueType from gmso.core.element import element_by_name, element_by_symbol from gmso.formats.formats_registry import loads_as @@ -69,6 +70,7 @@ def from_mol2(filename, site_type="atom"): "@BOND": _parse_bond, "@CRYSIN": _parse_box, "@FF_PBC": _parse_box, + "@MOLECULE": _parse_molecule, } for section in sections: if section not in supported_rti: @@ -79,7 +81,6 @@ def from_mol2(filename, site_type="atom"): else: supported_rti[section](topology, sections[section]) - topology.update_topology() # TODO: read in parameters to correct attribute as well. This can be saved in various rti sections. return topology @@ -139,13 +140,14 @@ def parse_ele(*symbols): f"No charge was detected for site {content[1]} with index {content[0]}" ) charge = None - + molecule = top.label if top.__dict__.get("label") else top.name atom = Atom( name=content[1], position=position.to("nm"), element=element, charge=charge, - residue=(content[7], int(content[6])), + residue=ResidueType(content[7], int(content[6])), + molecule=MoleculeType(molecule, 1), ) top.add_site(atom) @@ -178,3 +180,8 @@ def _parse_box(top, section): lengths=[float(x) for x in content[0:3]] * u.Å, angles=[float(x) for x in content[3:6]] * u.degree, ) + + +def _parse_molecule(top, section): + """Parse molecule information from the mol2 file.""" + top.label = str(section[0].strip()) diff --git a/gmso/tests/test_mol2.py b/gmso/tests/test_mol2.py index 1b95b94a7..271512b54 100644 --- a/gmso/tests/test_mol2.py +++ b/gmso/tests/test_mol2.py @@ -133,3 +133,28 @@ def test_neopentane_mol2_elements(self): r"consider manually adding the element to the topology$", ): top = Topology.load(get_fn("neopentane.mol2")) + + def test_mol2_residues(self): + top = Topology.load(get_fn("parmed.mol2")) + assert np.all( + np.array([site.residue.name for site in top.sites]) == "RES" + ) + assert np.all( + np.array([site.residue.number for site in top.sites]) == 1 + ) + + def test_mol2_molecules(self): + top = Topology.load(get_fn("methane.mol2")) + assert np.all( + np.array([site.molecule.name for site in top.sites]) == "MET" + ) + assert np.all( + np.array([site.molecule.number for site in top.sites]) == 1 + ) + + def test_mol2_group(self): + # Is there a place to read from mol2 file? + top = Topology.load(get_fn("ethane.mol2")) + for site in top.sites: + site.group = "ethane" + assert np.all(np.array([site.group for site in top.sites]) == "ethane")