# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3 (ipykernel)
#     language: python
#     name: python3
# ---

# +
import mbuild as mb
import mdtraj as md
import nglview
from foyer.forcefield import Forcefield as FoyerForcefield
from openff.toolkit.topology import Molecule, Topology
from openff.toolkit.typing.engines.smirnoff.forcefield import ForceField
from openff.units import unit

from openff.interchange.components.interchange import Interchange
from openff.interchange.components.mbuild import offmol_to_compound
from openff.interchange.components.mdtraj import _combine_topologies, _OFFBioTop
from openff.interchange.drivers import get_openmm_energies
# -

nanoparticle = mb.load("silica_nanoparticle.json")

oplsaa_silica = FoyerForcefield("oplsaa_switchable.xml")

topology = _OFFBioTop(mdtop=md.Topology.from_openmm(nanoparticle.to_parmed().topology))

nanoparticle_interchange = Interchange.from_foyer(
    topology=topology,
    force_field=oplsaa_silica,
)

# +
n_solvent = 800

solvent = Molecule.from_smiles("C1CCOC1")
solvent.generate_conformers(n_conformers=1)
# -

sage = ForceField("openff_unconstrained-2.0.0.offxml")

solvent_interchange = Interchange.from_smirnoff(
    sage,
    _OFFBioTop(
        mdtop=md.Topology.from_openmm(
            Topology.from_molecules(n_solvent * [solvent]).to_openmm()
        )
    ),
)

combined = solvent_interchange + nanoparticle_interchange

thf_compound = offmol_to_compound(solvent)

solvated_compound = mb.packing.solvate(
    solute=nanoparticle, solvent=thf_compound, n_solvent=n_solvent, box=mb.Box(3 * [8])
)

combined.topology = _combine_topologies(
    solvent_interchange.topology,
    nanoparticle_interchange.topology,
)
combined.positions = solvated_compound.xyz * unit.nanometer
combined.box = 3 * [8] * unit.nanometer

combined.to_pdb("out.pdb")
view = nglview.show_mdtraj(md.load("out.pdb"))
view.clear_representations()
view.add_representation(
    "spacefill", selection=[*range(combined.topology.mdtop.n_atoms)]
)
view

get_openmm_energies(combined, combine_nonbonded_forces=True)

StarcoderdataJupyterScriptsDedupFiltered