From 24e157ff4de639f2f7aa5abb835f12783ce499ed Mon Sep 17 00:00:00 2001
From: Jacob Jeffries <88559006+jwjeffr@users.noreply.github.com>
Date: Fri, 14 Aug 2026 15:42:35 +0000
Subject: [PATCH] added optional mass weighting in VACF

---
 tests/test_correlations.py           | 51 +++++++++++++++++++++++++++-
 torch_sim/properties/correlations.py | 29 +++++++++++++---
 2 files changed, 74 insertions(+), 6 deletions(-)

diff --git a/tests/test_correlations.py b/tests/test_correlations.py
index 5b3cbb0b8..f55c6ae08 100644
--- a/tests/test_correlations.py
+++ b/tests/test_correlations.py
@@ -28,9 +28,19 @@ class MockState:
     the components needed for correlation calculations.
     """
 
-    def __init__(self, velocities: torch.Tensor, device: torch.device) -> None:
+    def __init__(
+        self,
+        velocities: torch.Tensor,
+        device: torch.device,
+        masses: torch.Tensor | None = None,
+    ) -> None:
         """Initialize mock state with provided data."""
         self.velocities = velocities
+        self.masses = (
+            torch.ones(velocities.shape[0], device=device, dtype=velocities.dtype)
+            if masses is None
+            else masses
+        )
         self.device = device
         # Required for TrajectoryReporter
         self.n_systems = 1
@@ -438,6 +448,45 @@ def test_velocity_autocorrelation(mock_state_factory: Callable) -> None:
     assert torch.min(vacf) >= -1.0 - 1e-2
 
 
+@pytest.mark.parametrize("normalize", [False, True])
+def test_mass_weighted_velocity_autocorrelation(normalize) -> None:
+    """Test that atomic masses weight the raw ACF before normalization."""
+    window_size = 8
+    masses = torch.tensor([1.0, 4.0], device=DEVICE)
+    t = torch.arange(window_size, device=DEVICE)
+
+    velocity_history = torch.zeros(window_size, 2, 3, device=DEVICE)
+    velocity_history[:, 0] = torch.cos(2 * math.pi * t / 4).unsqueeze(-1)
+    velocity_history[:, 1] = torch.cos(2 * math.pi * t / 8).unsqueeze(-1)
+
+    vacf_calc = VelocityAutoCorrelation(
+        window_size=window_size,
+        device=DEVICE,
+        use_running_average=False,
+        normalize=normalize,
+        mass_weighted=True,
+    )
+
+    for velocities in velocity_history:
+        vacf_calc(MockState(velocities, DEVICE, masses))
+
+    centered = velocity_history - velocity_history.mean(dim=0, keepdim=True)
+    expected_acf = torch.stack(
+        [
+            torch.sum(centered[: window_size - lag] * centered[lag:], dim=0)
+            for lag in range(window_size)
+        ]
+    )
+    expected = torch.sum(
+        expected_acf.mean(dim=2) * masses.unsqueeze(0), dim=1
+    ) / masses.sum()
+    if normalize:
+        expected = expected / expected[0]
+
+    assert vacf_calc.vacf is not None
+    assert torch.allclose(vacf_calc.vacf, expected, atol=1e-5)
+
+
 def test_velocity_autocorrelation_with_trajectory_reporter(
     mock_state_factory: Callable,
 ) -> None:
diff --git a/torch_sim/properties/correlations.py b/torch_sim/properties/correlations.py
index 3614254b2..5138e901f 100644
--- a/torch_sim/properties/correlations.py
+++ b/torch_sim/properties/correlations.py
@@ -393,7 +393,7 @@ class VelocityAutoCorrelation:
     """Calculator for velocity autocorrelation function (VACF).
 
     Computes VACF by averaging over atoms and dimensions, with optional
-    running average across correlation windows.
+    mass weighting and running average across correlation windows.
 
 
     Using ``VelocityAutoCorrelation`` with
@@ -422,6 +422,7 @@ def __init__(
         device: torch.device,
         use_running_average: bool = True,
         normalize: bool = True,
+        mass_weighted: bool = False,
     ) -> None:
         """Initialize VACF calculator.
 
@@ -429,15 +430,21 @@ def __init__(
             window_size: Number of steps in correlation window
             device: Computation device
             use_running_average: Whether to compute running average across windows
-            normalize: Whether to normalize correlation functions to [0,1]
+            normalize: Whether to normalize correlations at zero lag
+            mass_weighted: Whether to weight each atom's raw VACF by its mass
+                before normalization
         """
+        # Mass weighting must be applied before normalization. Normalizing each
+        # atom/component first would cancel the multiplicative mass weights.
         self.corr_calc = CorrelationCalculator(
             window_size=window_size,
             properties={"velocity": lambda s: s.velocities},
             device=device,
-            normalize=normalize,
+            normalize=normalize and not mass_weighted,
         )
         self.use_running_average = use_running_average
+        self.normalize = normalize
+        self.mass_weighted = mass_weighted
         self._window_count = 0
         self._avg = torch.zeros(window_size, device=device)
 
@@ -455,8 +462,20 @@ def __call__(self, state: SimState, _: Any = None) -> torch.Tensor:
 
         if self.corr_calc.buffers["velocity"].count == self.corr_calc.window_size:
             correlations = self.corr_calc.get_auto_correlations()
-            # dims: (natoms, ndims)
-            vacf = torch.mean(correlations["velocity"], dim=(1, 2))
+            velocity_acf = correlations["velocity"]
+            # velocity_acf shape: (window_size, n_atoms, n_dimensions)
+            if self.mass_weighted:
+                masses = state.masses.to(
+                    device=velocity_acf.device, dtype=velocity_acf.dtype
+                )
+                atom_vacf = torch.mean(velocity_acf, dim=2)
+                vacf = torch.sum(atom_vacf * masses.unsqueeze(0), dim=1) / torch.sum(
+                    masses
+                )
+                if self.normalize and vacf[0] > 1e-10:
+                    vacf = vacf / vacf[0]
+            else:
+                vacf = torch.mean(velocity_acf, dim=(1, 2))
 
             self._window_count += 1