Non-omm ase

python-libraries/nanover-ase/src/nanover/ase/converter.py

@@ -50,6 +50,15 @@
 }
 
 
+def ase_atoms_to_frame_data(
+    ase_atoms: Atoms,
+    *,
+    topology: bool,
+    **kwargs,
+) -> FrameData:
+    return ase_to_frame_data(ase_atoms, topology=topology, **kwargs)
+
+
 def ase_to_frame_data(
     ase_atoms: Atoms,
     positions=True,

python-libraries/nanover-ase/src/nanover/ase/frame_adaptor.py

@@ -4,8 +4,9 @@
 
 from typing import Callable
 from ase import Atoms  # type: ignore
+
+from nanover.ase.converter import ase_atoms_to_frame_data
 from nanover.trajectory import FramePublisher
-from nanover.ase import ase_to_frame_data
 
 
 def send_ase_frame(
@@ -39,8 +40,10 @@ def send_ase_frame(
 
     def send():
         nonlocal frame_index
-        frame = ase_to_frame_data(
+        include_topology = frame_index == 0
+        frame = ase_atoms_to_frame_data(
             ase_atoms,
+            topology=include_topology,
             include_velocities=include_velocities,
             include_forces=include_forces,
         )

python-libraries/nanover-ase/src/nanover/ase/openmm/frame_adaptor.py

@@ -0,0 +1,48 @@
+from ase import Atoms
+
+from nanover.ase import ase_to_frame_data
+from nanover.openmm.converter import add_openmm_topology_to_frame_data
+from nanover.trajectory import FramePublisher, FrameData
+
+
+def openmm_ase_frame_adaptor(
+    ase_atoms: Atoms,
+    frame_publisher: FramePublisher,
+    **kwargs,
+):
+    """
+    Generates and sends frames for a simulation using an :class: OpenMMCalculator.
+    """
+
+    frame_index = 0
+
+    def send():
+        nonlocal frame_index
+        include_topology = frame_index == 0
+        frame_data = openmm_ase_atoms_to_frame_data(
+            ase_atoms, topology=include_topology, **kwargs
+        )
+        frame_publisher.send_frame(frame_index, frame_data)
+        frame_index += 1
+
+    return send
+
+
+def openmm_ase_atoms_to_frame_data(
+    ase_atoms: Atoms,
+    *,
+    topology: bool,
+    **kwargs,
+) -> FrameData:
+    frame_data = ase_to_frame_data(
+        ase_atoms,
+        topology=False,
+        **kwargs,
+    )
+
+    if topology:
+        imd_calculator = ase_atoms.calc
+        topology = imd_calculator.calculator.topology
+        add_openmm_topology_to_frame_data(frame_data, topology)
+
+    return frame_data

python-libraries/nanover-ase/src/nanover/ase/openmm/runner.py

@@ -6,21 +6,19 @@
 from pathlib import Path
 from typing import Optional, List
 
-from ase import units, Atoms  # type: ignore
+from ase import units  # type: ignore
 from ase.md import MDLogger, Langevin
 from ase.md.velocitydistribution import MaxwellBoltzmannDistribution
 from attr import dataclass
 from nanover.app import NanoverImdApplication, NanoverRunner
 from nanover.app.app_server import DEFAULT_NANOVER_PORT
+from nanover.ase.openmm.frame_adaptor import openmm_ase_frame_adaptor
 from nanover.core import NanoverServer, DEFAULT_SERVE_ADDRESS
 from nanover.ase import TrajectoryLogger
 from nanover.essd import DiscoveryServer
-from nanover.openmm import openmm_to_frame_data, serializer
-from nanover.trajectory.frame_publisher import FramePublisher
-from openmm.app import Simulation, Topology
+from nanover.openmm import serializer
+from openmm.app import Simulation
 
-from nanover.ase import ase_to_frame_data
-from nanover.ase.converter import add_ase_positions_to_frame_data
 from nanover.ase.imd import NanoverASEDynamics
 from nanover.ase.openmm.calculator import OpenMMCalculator
 from nanover.ase.wall_constraint import VelocityWallConstraint
@@ -35,41 +33,6 @@
 )
 
 
-def openmm_ase_frame_adaptor(
-    ase_atoms: Atoms,
-    frame_publisher: FramePublisher,
-    include_velocities=False,
-    include_forces=False,
-):
-    """
-    Generates and sends frames for a simulation using an :class: OpenMMCalculator.
-    """
-
-    frame_index = 0
-    topology: Optional[Topology] = None
-
-    def send():
-        nonlocal frame_index, topology
-        # generate topology frame using OpenMM converter.
-        if frame_index == 0:
-            imd_calculator = ase_atoms.calc
-            topology = imd_calculator.calculator.topology
-            frame = openmm_to_frame_data(
-                state=None,
-                topology=topology,
-                include_velocities=include_velocities,
-                include_forces=include_forces,
-            )
-            add_ase_positions_to_frame_data(frame, ase_atoms.get_positions())
-        # from then on, just send positions and state.
-        else:
-            frame = ase_to_frame_data(ase_atoms, topology=False)
-        frame_publisher.send_frame(frame_index, frame)
-        frame_index += 1
-
-    return send
-
-
 @dataclass
 class ImdParams:
     """

python-libraries/nanover-omni/src/nanover/omni/ase.py

@@ -0,0 +1,210 @@
+from dataclasses import dataclass
+from typing import Optional, Any, Protocol
+
+import numpy as np
+from ase import Atoms
+from ase.calculators.calculator import Calculator
+from ase.md import MDLogger
+from ase.md.md import MolecularDynamics
+
+from nanover.app import NanoverImdApplication
+from nanover.ase.converter import EV_TO_KJMOL, ase_atoms_to_frame_data
+from nanover.ase.imd_calculator import ImdCalculator
+from nanover.ase.wall_constraint import VelocityWallConstraint
+from nanover.trajectory import FrameData
+from nanover.utilities.event import Event
+
+
+@dataclass
+class InitialState:
+    positions: Any
+    velocities: Any
+    cell: Any
+
+
+class ASEAtomsToFrameData(Protocol):
+    def __call__(self, ase_atoms: Atoms, *, topology: bool, **kwargs) -> FrameData: ...
+
+
+class ASESimulation:
+    """
+    A wrapper for ASE simulations so they can be run inside the OmniRunner.
+    """
+
+    @classmethod
+    def from_ase_dynamics(
+        cls,
+        dynamics: MolecularDynamics,
+        *,
+        name: Optional[str] = None,
+        ase_atoms_to_frame_data: ASEAtomsToFrameData = ase_atoms_to_frame_data
+    ):
+        """
+        Construct this from an existing ASE dynamics.
+        :param dynamics: An existing ASE Dynamics
+        :param name: An optional name for the simulation instead of default
+        :param ase_atoms_to_frame_data: An optional callback to extra frames from the system
+        """
+        sim = cls(name)
+        sim.dynamics = dynamics
+        sim.ase_atoms_to_frame_data = ase_atoms_to_frame_data
+        sim.initial_calc = dynamics.atoms.calc
+
+        return sim
+
+    @property
+    def atoms(self):
+        if self.dynamics is None:
+            return None
+        else:
+            return self.dynamics.atoms
+
+    def __init__(self, name: Optional[str] = None):
+        self.name = name or "Unnamed ASE OpenMM Simulation"
+
+        self.app_server: Optional[NanoverImdApplication] = None
+
+        self.on_reset_energy_exceeded = Event()
+
+        self.verbose = False
+        self.use_walls = False
+        self.reset_energy: Optional[float] = None
+        self.frame_interval = 5
+        self.include_velocities = False
+        self.include_forces = False
+
+        self.dynamics: Optional[MolecularDynamics] = None
+        self.checkpoint: Optional[InitialState] = None
+        self.initial_calc: Optional[Calculator] = None
+
+        self.frame_index = 0
+        self.ase_atoms_to_frame_data = ase_atoms_to_frame_data
+
+    def load(self):
+        """
+        Load and set up the simulation if it isn't done already.
+        """
+        assert self.dynamics is not None
+
+        if self.use_walls:
+            self.atoms.constraints.append(VelocityWallConstraint())
+
+        self.checkpoint = InitialState(
+            positions=self.atoms.get_positions(),
+            velocities=self.atoms.get_velocities(),
+            cell=self.atoms.get_cell(),
+        )
+
+    def reset(self, app_server: NanoverImdApplication):
+        """
+        Reset the simulation to its initial conditions, reset IMD interactions, and reset frame stream to begin with
+        topology and continue.
+        :param app_server: The app server hosting the frame publisher and imd state
+        """
+        assert (
+            self.dynamics is not None
+            and self.atoms is not None
+            and self.checkpoint is not None
+            and self.initial_calc is not None
+        )
+
+        self.app_server = app_server
+
+        # reset atoms to initial state
+        self.atoms.set_positions(self.checkpoint.positions)
+        self.atoms.set_velocities(self.checkpoint.velocities)
+        self.atoms.set_cell(self.checkpoint.cell)
+
+        # setup calculator
+        self.atoms.calc = ImdCalculator(
+            self.app_server.imd,
+            self.initial_calc,
+            dynamics=self.dynamics,
+        )
+
+        # send the initial topology frame
+        frame_data = self.make_topology_frame()
+        self.app_server.frame_publisher.send_frame(0, frame_data)
+        self.frame_index = 1
+
+        # TODO: deal with this when its clear if dynamics should be reconstructed or not..
+        if self.verbose:
+            self.dynamics.attach(
+                MDLogger(
+                    self.dynamics,
+                    self.atoms,
+                    "-",
+                    header=True,
+                    stress=False,
+                    peratom=False,
+                ),
+                interval=100,
+            )
+
+    def advance_by_one_step(self):
+        """
+        Advance the simulation to the next point a frame should be reported, and send that frame.
+        """
+        self.advance_to_next_report()
+
+    def advance_by_seconds(self, dt: float):
+        """
+        Advance playback time by some seconds, and advance the simulation to the next frame output.
+        :param dt: Time to advance playback by in seconds (ignored)
+        """
+        self.advance_to_next_report()
+
+    def advance_to_next_report(self):
+        """
+        Step the simulation to the next point a frame should be reported, and send that frame.
+        """
+        assert self.dynamics is not None and self.app_server is not None
+
+        # determine step count for next frame
+        steps_to_next_frame = (
+            self.frame_interval
+            - self.dynamics.get_number_of_steps() % self.frame_interval
+        )
+
+        # advance the simulation
+        self.dynamics.run(steps_to_next_frame)
+
+        # generate the next frame
+        frame_data = self.make_regular_frame()
+
+        # send the next frame
+        self.app_server.frame_publisher.send_frame(self.frame_index, frame_data)
+        self.frame_index += 1
+
+        # check if excessive energy requires sim reset
+        if self.reset_energy is not None and self.app_server is not None:
+            energy = self.atoms.get_total_energy() * EV_TO_KJMOL
+            if not np.isfinite(energy) or energy > self.reset_energy:
+                self.on_reset_energy_exceeded.invoke()
+                self.reset(self.app_server)
+
+    def make_topology_frame(self):
+        """
+        Make a NanoVer FrameData corresponding to the current particle positions and topology of the simulation.
+        """
+        assert self.atoms is not None
+
+        return self.ase_atoms_to_frame_data(
+            self.atoms,
+            topology=True,
+            include_velocities=self.include_velocities,
+            include_forces=self.include_forces,
+        )
+
+    def make_regular_frame(self):
+        """
+        Make a NanoVer FrameData corresponding to the current state of the simulation.
+        """
+        assert self.atoms is not None
+
+        return self.ase_atoms_to_frame_data(
+            self.atoms,
+            topology=False,
+            include_velocities=self.include_velocities,
+            include_forces=self.include_forces,
+        )