Manage TensorReporter with open file

.github/workflows/gh-ci.yaml

@@ -25,7 +25,7 @@ jobs:
         matrix:
           os: [macOS-latest, ubuntu-latest]
           environment:
-            - py312amber
+            - py312
             - py312openeye
 
     steps:

pixi.toml

@@ -6,12 +6,8 @@ platforms = [ "linux-64", "osx-arm64" ]
 [tasks]
 run_mypy = "python -m mypy -p 'tyff'"
 run_tests = {
-    cmd = "python -m pytest -n logical tyff/_tests",
-    env = {OPENMM_DEFAULT_PLATFORM="CPU"},
-}
-dev_tests = {
-    cmd = "python -m pytest -x --lf tyff/_tests",
-    env = {OPENMM_DEFAULT_PLATFORM="CPU"},
+    cmd = "python -m pytest --durations=10 --cov=tyff --cov-config=pyproject.toml --cov-append --cov-report=xml --color=yes tyff/_tests",
+    env = {OPENMM_DEFAULT_PLATFORM="CPU"}
 }
 
 [pypi-dependencies]
@@ -21,10 +17,7 @@ numpy = "*"
 setuptools_scm = "*"
 distributed = "*"
 openmm = "==8.4.*"
-symengine = "*"
 rdkit = "==2023.*"
-matplotlib = "*"
-panel = "*"
 sqlmodel = "*"
 dask = "*"
 dask-jobqueue = "*"
@@ -41,29 +34,21 @@ openff-interchange = "0.5.*"
 openff-nagl-base = "*"
 openff-nagl-models = "*"
 openmmtools = "*"
-pydantic-units = "*"
 pytorch = "*"
-datasets = "*"
-bokeh = "3.9.*"
-# probably optional-ish below here
-jupyter = "*"
-nbconvert = "*"
-tensorboard = "*"
-tensorboardx = "*"
-uncertainties = "*"
-openff-packmol = "0.1.*"
-openff-evaluator-base = "*"
+smee = "*"
+descent = "*"
 # shims
 pyarrow = "20.*"
 rdkit = "2023.*"
 
 [environments]
-py312amber = [ "py312", "ambertools", "test", "typing" ]
-py312openeye = [ "py312", "openeye", "packmol", "test", "typing" ]
+py312 = [ "py312", "test", "typing" ]
+py312openeye = [ "py312", "openeye", "test", "typing" ]
 dev = [ "py312", "openeye", "test", "typing", "dev" ]
-hpc3 = ["py312", "openeye", "cu130", "packmol" ]
 
 [feature.dev.dependencies]
+openff-packmol = "0.1.*"
+openff-evaluator-base = "*"
 pymol-open-source = "*"
 
 [feature.dev.pypi-dependencies]
@@ -73,28 +58,19 @@ ipython = "*"
 ipdb = "*"
 ruff = "*"
 
-[feature.ambertools.dependencies]
-ambertools = "*"
-
 [feature.openeye]
 channels = [ "openeye", "conda-forge" ]
 dependencies = { openeye-toolkits = "*" }
 
-[feature.packmol.dependencies]
-packmol = ">=20"
-
-[feature.cu130.dependencies]
-cuda-version = "13.0"
-
 [feature.py312.dependencies]
 python = "3.12.*"
 
 [feature.test.pypi-dependencies]
 pytest = "*"
 pytest-cov = "*"
 pytest-xdist = "*"
-pytest-mock = "*"
 pytest-randomly = "*"
+pytest-mock = "*"
 
 [feature.typing.pypi-dependencies]
 mypy = "*"

tyff/compute/_equilibrate.py

@@ -5,6 +5,7 @@
 
 from parsl import File
 
+import tyff.mm
 from tyff.compute._files import (
     EquilibrationFiles,
     MinimizationFiles,
@@ -129,11 +130,25 @@ def _run_equilibration(
 
     simulation.reporters.append(dcd_reporter)
 
+    # type hints imply I can pass these in as openmm.unit.Quantity and let it deal with conversions
+    tensor_reporter = tyff.mm.TensorReporter(
+        output_file=files["msgpack_trajectory"].filepath,
+        report_interval=1000,
+        beta=1.0 / openmm.unit.kilocalories_per_mole,
+        pressure=pressure * openmm.unit.kilopascal,
+    )
+
+    simulation.reporters.append(tensor_reporter)
+
     simulation.context.setVelocitiesToTemperature(
         compute_config["temperature"],  # kelvin, but as float
         compute_config["replicate_index"] + 1,
     )
 
+    platform_name = simulation.context.getPlatform().getName()
+
+    logger.info(f"Using OpenMM platform: {platform_name}")
+
     logger.info("Running 10,000 steps of MD")
 
     detected_platform = simulation.context.getPlatform().getName()

tyff/compute/_produce.py

@@ -7,6 +7,7 @@
 import openmm.app
 from parsl import File
 
+import tyff.mm
 from tyff.compute._files import (
     EquilibrationFiles,
     ProductionFiles,
@@ -129,14 +130,25 @@ def _run_production(
     if pressure is None:
         raise PressureNotDefinedError("Trying to set up NPT simulation but no pressure defined.")
 
-    logger.info("Running 100,000 steps of MD")
+    tensor_reporter = tyff.mm.TensorReporter(
+        output_file=files["msgpack_trajectory"].filepath,
+        report_interval=1000,
+        beta=1.0 / openmm.unit.kilocalories_per_mole,
+        pressure=pressure * openmm.unit.kilopascal,
+    )
+
+    simulation.reporters.append(tensor_reporter)
 
-    platform = simulation.context.getPlatform()
-    platform_name = platform.getName()
+    platform_name = simulation.context.getPlatform().getName()
 
-    # Log or print the result
     logger.info(f"Using OpenMM platform: {platform_name}")
 
+    logger.info("Running 100,000 steps of MD")
+
+    for index in range(10):
+        logger.info(f"Running from step {index * 10_000} to step {(index + 1) * 10_000}")
+        simulation.step(10_000)
+
     for index in range(10):
         logger.info(f"Running from step {index * 10_000} to step {(index + 1) * 10_000}")
         simulation.step(10_000)