search

openforcefield / openff-interchange

A project (and object) for storing, manipulating, and converting molecular mechanics data.
https://docs.openforcefield.org/projects/interchange
MIT License
69 stars 22 forks source link

JSON roundtrip doesn't work with `from_openmm` generated system #1003

Open IAlibay opened 2 months ago

IAlibay commented 2 months ago

Description

Interchange.parse_raw() yields a KeyError: 'name' when passing a JSON generated from an Interchange object that was created from an OpenMM System & Topology (via Interchange.from_openmm).

Note: the openmm System & Topology were created using SystemGenerator.

Reproduction

from openmm import app, MonteCarloBarostat
from openmm import unit as omm_unit
from openmmforcefields.generators.system_generators import SystemGenerator
from openff.toolkit import Molecule
from openff.units import unit
from openff.units.openmm import to_openmm, from_openmm
from openff.interchange import Interchange
import os

os.environ['INTERCHANGE_EXPERIMENTAL'] = "1"

forcefield_kwargs = {
    'constraints': app.HBonds,
    'rigidWater': True,
    'removeCMMotion': True,
    'hydrogenMass': 4.0 * omm_unit.amu
}

periodic_kwargs = {
    'nonbondedMethod': app.PME,
    'nonbondedCutoff': to_openmm(9 * unit.angstrom)
}

non_periodic_kwargs = {
    'nonbondedMethod': app.NoCutoff,
}

system_generator = SystemGenerator(
    forcefields=['amber/ff14SB.xml', 'amber/tip3p_standard.xml'],
    small_molecule_forcefield='openff-2.1.0',
    forcefield_kwargs=forcefield_kwargs,
    nonperiodic_forcefield_kwargs=non_periodic_kwargs,
    periodic_forcefield_kwargs=periodic_kwargs,
    cache=None,
    barostat=None,
)

off = Molecule.from_smiles('C')
off.generate_conformers()
off.assign_partial_charges(partial_charge_method='am1bcc')

system_generator.create_system(off.to_topology().to_openmm(), molecules=[off])

model = app.Modeller(app.Topology(), [])
model.add(
    off.to_topology().to_openmm(), to_openmm(off.conformers[0])
)

model.addSolvent(
    system_generator.forcefield,
    model='tip3p',
    padding=to_openmm(12 * unit.angstrom),
    positiveIon='Na+', negativeIon='Cl-',
    ionicStrength=to_openmm(0.15 * unit.molar),
    neutralize=True,
)

system = system_generator.create_system(
    model.topology, molecules=[off],
)

inter = Interchange.from_openmm(system, model.topology, model.positions)

inter2 = Interchange.parse_raw(inter.json())

Output

Click me ``` --------------------------------------------------------------------------- KeyError Traceback (most recent call last) Cell In[1], line 65 59 system = system_generator.create_system( 60 model.topology, molecules=[off], 61 ) 63 inter = Interchange.from_openmm(system, model.topology, model.positions) ---> 65 inter2 = Interchange.parse_raw(inter.json()) File ~/software/mambaforge/install/envs/ofe_demo/lib/python3.11/site-packages/pydantic/v1/main.py:539, in BaseModel.parse_raw(cls, b, content_type, encoding, proto, allow_pickle) 528 @classmethod 529 def parse_raw( 530 cls: Type['Model'], (...) 536 allow_pickle: bool = False, 537 ) -> 'Model': 538 try: --> 539 obj = load_str_bytes( 540 b, 541 proto=proto, 542 content_type=content_type, 543 encoding=encoding, 544 allow_pickle=allow_pickle, 545 json_loads=cls.__config__.json_loads, 546 ) 547 except (ValueError, TypeError, UnicodeDecodeError) as e: 548 raise ValidationError([ErrorWrapper(e, loc=ROOT_KEY)], cls) File ~/software/mambaforge/install/envs/ofe_demo/lib/python3.11/site-packages/pydantic/v1/parse.py:37, in load_str_bytes(b, content_type, encoding, proto, allow_pickle, json_loads) 35 if isinstance(b, bytes): 36 b = b.decode(encoding) ---> 37 return json_loads(b) 38 elif proto == Protocol.pickle: 39 if not allow_pickle: File ~/software/mambaforge/install/envs/ofe_demo/lib/python3.11/site-packages/openff/interchange/components/interchange.py:97, in interchange_loader(data) 95 tmp["box"] = Quantity(val["val"], unit.Unit(val["unit"])) 96 elif key == "topology": ---> 97 tmp["topology"] = Topology.from_json(val) 98 elif key == "collections": 99 from openff.interchange.smirnoff import ( 100 SMIRNOFFAngleCollection, 101 SMIRNOFFBondCollection, (...) 107 SMIRNOFFVirtualSiteCollection, 108 ) File ~/software/mambaforge/install/envs/ofe_demo/lib/python3.11/site-packages/openff/toolkit/utils/serialization.py:153, in Serializable.from_json(cls, serialized) 150 import json 152 d = json.loads(serialized) --> 153 return cls.from_dict(d) File ~/software/mambaforge/install/envs/ofe_demo/lib/python3.11/site-packages/openff/toolkit/topology/topology.py:1261, in Topology.from_dict(cls, topology_dict) 1246 """ 1247 Create a new Topology from a dictionary representation 1248 (...) 1258 1259 """ 1260 topology = cls() -> 1261 topology._initialize_from_dict(topology_dict) 1262 return topology File ~/software/mambaforge/install/envs/ofe_demo/lib/python3.11/site-packages/openff/toolkit/topology/topology.py:1285, in Topology._initialize_from_dict(self, topology_dict) 1282 self.box_vectors = Quantity(box_vectors_unitless, box_vectors_unit) 1284 for molecule_dict in topology_dict["molecules"]: -> 1285 new_mol = Molecule.from_dict(molecule_dict) 1286 self._add_molecule_keep_cache(new_mol) 1287 self._invalidate_cached_properties() File ~/software/mambaforge/install/envs/ofe_demo/lib/python3.11/site-packages/openff/toolkit/topology/molecule.py:1234, in FrozenMolecule.from_dict(cls, molecule_dict) 1232 # This implementation is a compromise to let this remain as a classmethod 1233 mol = cls() -> 1234 mol._initialize_from_dict(molecule_dict) 1235 return mol File ~/software/mambaforge/install/envs/ofe_demo/lib/python3.11/site-packages/openff/toolkit/topology/molecule.py:1249, in FrozenMolecule._initialize_from_dict(self, molecule_dict) 1246 # TODO: Provide useful exception messages if there are any failures 1248 self._initialize() -> 1249 self.name = molecule_dict["name"] 1250 for atom_dict in molecule_dict["atoms"]: 1251 self._add_atom(**atom_dict, invalidate_cache=False) KeyError: 'name' ````

Software versions

IAlibay commented 2 months ago

Note: it's not isolated to SystemGenerator

from openff.interchange import Interchange
from openff.toolkit import Molecule, ForceField
from openff.units import unit
import numpy as np
import os
os.environ['INTERCHANGE_EXPERIMENTAL'] = "1"

mol = Molecule.from_smiles("CC")
mol.generate_conformers()
sage = ForceField("openff-2.0.0.offxml")
cubic_box = unit.Quantity(30 * np.eye(3), unit.angstrom)

interchange = Interchange.from_smirnoff(topology=[mol], force_field=sage, box=cubic_box)
interchange.positions = mol.conformers[0]

interchange2 = Interchange.from_openmm(interchange.to_openmm(), interchange.to_openmm_topology())
Interchange.parse_raw(interchange2.json())
mattwthompson commented 2 months ago

This looks similar to https://github.com/openforcefield/openff-toolkit/issues/1783 ... which in that thread I insist was fixed in the version you're using ... hmm

The issue here is with the toolkit / how information is passed between Interchange and the toolkit since Topology didn't always play nicely with _SimpleMolecules that are necessary to use when importing from OpenMM

mattwthompson commented 2 weeks ago

This is unlikely to land before 0.4.0 but I hope to take a closer look at this soon.

I assume the use case is simple on the surface (save state out to JSON, trust that it'll be loaded back and represent the same state) but I anticipate the work will mostly be in gathering up and testing various corner cases.