Optionally include hydrogens in RMSDs

yammbs/_store.py

@@ -704,6 +704,7 @@ def get_rmsd(
         force_field: str,
         molecule_ids: list[int] | None = None,
         skip_check: bool = False,
+        include_hydrogens: bool = False,
     ) -> RMSDCollection:
         if not molecule_ids:
             molecule_ids = self.get_molecule_ids()
@@ -737,7 +738,12 @@ def get_rmsd(
                 rmsds.append(
                     RMSD(
                         qcarchive_id=id,
-                        rmsd=get_rmsd(molecule, qm, mm),
+                        rmsd=get_rmsd(
+                            molecule,
+                            qm,
+                            mm,
+                            include_hydrogens=include_hydrogens,
+                        ),
                     ),
                 )
 
@@ -857,14 +863,18 @@ def get_outputs(
 
     def get_metrics(
         self,
+        include_hydrogens: bool = False,
     ) -> MetricCollection:
         metrics = MetricCollection()
         import pandas
 
         # TODO: Optimize this for speed
         for force_field in self.get_force_fields():
             ddes = self.get_dde(force_field=force_field).to_dataframe()
-            rmsds = self.get_rmsd(force_field=force_field).to_dataframe()
+            rmsds = self.get_rmsd(
+                force_field=force_field,
+                include_hydrogens=include_hydrogens,
+            ).to_dataframe()
             tfds = self.get_tfd(force_field=force_field).to_dataframe()
             icrmsds = self.get_internal_coordinate_rmsd(
                 force_field=force_field,

yammbs/_tests/unit_tests/test_analysis.py

@@ -14,12 +14,14 @@
 
 
 class TestAnalysis:
-    def test_rmsd(self, allicin, conformers):
+    @pytest.mark.parametrize("include_hydrogens", [True, False])
+    def test_rmsd(self, allicin, conformers, include_hydrogens):
         # Passing the same conformers should return 0.0
         last_last = get_rmsd(
             molecule=allicin,
             reference=conformers[-1],
             target=conformers[-1],
+            include_hydrogens=include_hydrogens,
         )
 
         assert last_last == 0.0
@@ -28,6 +30,7 @@ def test_rmsd(self, allicin, conformers):
             molecule=allicin,
             reference=conformers[0],
             target=conformers[-1],
+            include_hydrogens=include_hydrogens,
         )
 
         assert isinstance(first_last, float)
@@ -36,6 +39,7 @@ def test_rmsd(self, allicin, conformers):
             molecule=allicin,
             reference=conformers[0],
             target=conformers[1],
+            include_hydrogens=include_hydrogens,
         )
 
         assert first_second != first_last
@@ -44,10 +48,28 @@ def test_rmsd(self, allicin, conformers):
             molecule=allicin,
             reference=conformers[-1],
             target=conformers[0],
+            include_hydrogens=include_hydrogens,
         )
 
         assert last_first == pytest.approx(first_last)
 
+    def test_rmsd_all_atom_rmsd_differs(self, allicin, conformers):
+        heavy_atom_rmsd = get_rmsd(
+            molecule=allicin,
+            reference=conformers[0],
+            target=conformers[1],
+            include_hydrogens=False,
+        )
+
+        all_atom_rmsd = get_rmsd(
+            molecule=allicin,
+            reference=conformers[0],
+            target=conformers[1],
+            include_hydrogens=True,
+        )
+
+        assert heavy_atom_rmsd != all_atom_rmsd
+
     def test_tfd(self, allicin, conformers):
         # Passing the same conformers should return 0.0
         last_last = get_tfd(

yammbs/_tests/unit_tests/torsion/test_store.py

@@ -89,6 +89,8 @@ def test_minimize_basic(single_torsion_dataset, tmp_path):
         database_name=tmp_path / "test.sqlite",
     )
 
+    TORSION_ID = 119466834
+
     store.optimize_mm(
         force_field="openff-2.2.0",
         n_processes=os.cpu_count(),
@@ -110,12 +112,10 @@ def test_minimize_basic(single_torsion_dataset, tmp_path):
     assert len(metrics["metrics"]["openff-2.2.0"]) == 1
 
     expected_metrics = {
-        "rmsd": 0.07475493617511018,
         "rmse": 0.8193199571663233,
         "mean_error": -0.35170719027937586,
         "js_distance": (0.3168201337322116, 500.0),
     }
-    TORSION_ID = 119466834
 
     assert len(expected_metrics) == len(metrics["metrics"]["openff-2.2.0"][TORSION_ID])
 

yammbs/analysis.py

@@ -124,6 +124,7 @@ def get_rmsd(
     molecule: Molecule,
     reference: Array,
     target: Array,
+    include_hydrogens: bool = False,
 ) -> float:
     """Compute the RMSD between two sets of coordinates."""
     from openeye import oechem
@@ -144,9 +145,9 @@ def get_rmsd(
     return oechem.OERMSD(
         molecule1.to_openeye(),
         molecule2.to_openeye(),
-        True,
-        True,
-        True,
+        True,  # automorph
+        not include_hydrogens,  # heavyOnly
+        True,  # overlay
     )
 
 

yammbs/torsion/_store.py

@@ -32,6 +32,7 @@
     RMSDCollection,
     RMSECollection,
     _normalize,
+    get_rmsd,
 )
 from yammbs.torsion.inputs import QCArchiveTorsionDataset
 from yammbs.torsion.models import (
@@ -385,13 +386,16 @@ def get_rmsd(
         force_field: str,
         torsion_ids: list[int] | None = None,
         skip_check: bool = False,
+        include_hydrogens: bool = False,
         restraint_k: float = 0.0,
     ) -> RMSDCollection:
         """Get the RMSD summed over the torsion profile."""
+        from openff.toolkit import Molecule
+
         if not torsion_ids:
             torsion_ids = self.get_torsion_ids()
 
-        if not skip_check:
+        if skip_check is None:
             # TODO: Copy this into each get_* method?
             LOGGER.info("Calling optimize_mm from inside of get_log_sse.")
             self.optimize_mm(force_field=force_field, restraint_k=restraint_k)
@@ -407,7 +411,21 @@ def get_rmsd(
                 allow_undefined_stereo=True,
             )
 
-            rmsds.append(RMSD.from_data(torsion_id, molecule, qm_points, mm_points))
+            rmsds.append(
+                RMSD(
+                    id=torsion_id,
+                    rmsd=sum(
+                        get_rmsd(
+                            molecule,
+                            qm_points[key],
+                            mm_points[key],
+                            include_hydrogens=include_hydrogens,
+                        )
+                        for key in qm_points
+                    )
+                    / len(qm_points),
+                ),
+            )
 
         return rmsds
 
@@ -541,6 +559,7 @@ def get_outputs(self) -> MinimizedTorsionDataset:
 
     def get_metrics(
         self,
+        include_hydrogens: bool = False,
         force_fields: Iterable[str] | None = None,
         js_temperature: float = 500.0,
         restraint_k: float = 0.0,
@@ -550,6 +569,8 @@ def get_metrics(
 
         Parameters
         ----------
+        include_hydrogens
+            Whether RMSDs should include hydrogens (all-atom) or not (heavy atom)
         force_fields : Iterable[str] | None
             Iterable of force fields to compute metrics for. If None, compute for all available.
         js_temperature : float
@@ -583,21 +604,19 @@ def get_metrics(
         # TODO: Optimize this for speed
         for force_field in force_fields:
             rmses = self.get_rmse(force_field=force_field, skip_check=True).to_dataframe()
-            rmsds = self.get_rmsd(force_field=force_field, skip_check=True).to_dataframe()
             mean_errors = self.get_mean_error(force_field=force_field, skip_check=True).to_dataframe()
             js_distances = self.get_js_distance(
                 force_field=force_field,
                 skip_check=True,
                 temperature=js_temperature,
             ).to_dataframe()
 
-            dataframe = rmses.join(rmsds).join(mean_errors).join(js_distances)
+            dataframe = rmses.join(mean_errors).join(js_distances)
 
             dataframe = dataframe.replace({pandas.NA: numpy.nan})
 
             metrics.metrics[force_field] = {
                 id: Metric(  # type: ignore[misc]
-                    rmsd=row["rmsd"],
                     rmse=row["rmse"],
                     mean_error=row["mean_error"],
                     js_distance=(row["js_distance"], row["js_temperature"]),
@@ -829,7 +848,6 @@ def get_summary_df(
         rows = []
         metrics = self.get_metrics().metrics
         metrics_to_plot = {
-            "RMSD / A": lambda x: x.rmsd,
             "RMSE / kcal mol-1": lambda x: x.rmse,
             "Mean Error / kcal mol-1": lambda x: x.mean_error,
             "JS Distance": lambda x: x.js_distance[0],

yammbs/torsion/outputs.py

@@ -48,7 +48,6 @@ class MinimizedTorsionDataset(ImmutableModel):
 class Metric(ImmutableModel):
     """Summary metrics for a given force field and QM reference."""
 
-    rmsd: float
     rmse: float
     mean_error: float
     js_distance: tuple[float, float]