VACF calculation does not seem to respect batching

[jwjeffr]

I've been trying to use the properties.correlations.VelocityAutoCorrelation calculator in batched MD simulations, but it seems that only one VACF trajectory is reported. I think it's intuitive that the vacf property would, in this case, be a tensor of shape (*, n_systems).

Is this intended behavior? If not, I'm happy to write up a small example. If so, I think this would be a very nice enhancement. I have a custom implementation for this I'm also happy to show off, but it's a little hacky.

[CompRhys]

If it's summing multiple systems into a single result then that's definitely not intended behaviour. Please share any relevant examples and then if you have a fix we can see whether or not there's a clean way to do it to fix in prod. I don't think this code will have been looked at much/at-all since initial release.

[jwjeffr]

See the example below. Batched MD for different temperatures, where I'd expect that the vacf would have a temperature index

r"""
Script is computing the VACF for bulk tungsten at different temperatures using a MACE-MP potential

Environment
-----------
Python: 3.13.13
GPU/CUDA: NVIDIA H100 80GB HBM3, CUDA version 13.0
Packages:
    torch-sim-atomstic==0.6.1
    ase==3.29.0
    mace-torch==0.3.16
"""

from ase.build import bulk
from mace.calculators.foundations_models import mace_mp
from torch_sim.models.mace import MaceModel
from torch_sim.properties.correlations import VelocityAutoCorrelation
import torch_sim as ts


DEVICE: str = "cuda"
SAMPLE_STRIDE: int = 10
NUM_STEPS: int = 1_000
TEMPERATURES: list[int] = [300, 600, 900, 1200]


def main():

    vacf_calc = VelocityAutoCorrelation(window_size=100, device=DEVICE, use_running_average=True)
    reporter = ts.TrajectoryReporter(
        filenames=[f".vacf_window_counter{temp:0d}.h5" for temp in TEMPERATURES],
        prop_calculators={SAMPLE_STRIDE: {"vacf": vacf_calc}}
    )

    mace = mace_mp(model="large", return_raw_model=True)
    tungsten = bulk("W", cubic=True).repeat((4, 4, 4))
    _ = ts.integrate(
        system=[tungsten.copy()] * len(TEMPERATURES),
        model=MaceModel(model=mace, device=DEVICE),
        n_steps=NUM_STEPS,
        timestep=1.0e-3,
        temperature=TEMPERATURES,
        integrator=ts.Integrator.nvt_langevin,
        trajectory_reporter=reporter,
    )

    # intuitively, vacf should be 2D, with either rows/columns corresponding to temperatures
    vacf = vacf_calc.vacf.detach().cpu().numpy()
    shape = vacf.shape

    if shape == (len(TEMPERATURES), NUM_STEPS // SAMPLE_STRIDE):
        print("Good! VACF is 2 dimensional")
    elif shape == (NUM_STEPS // SAMPLE_STRIDE,):
        print("Bad. VACF is not batched")
    else:
        raise ValueError("shape is something else")


if __name__ == "__main__":

    main()

The resulting shape does indeed match (NUM_STEPS // SAMPLE_STRIDE,), i.e. what I'd expect from a single trajectory. My workaround is very hacky and is quite long haha, so will put in a separate comment.

[jwjeffr]

My workaround thus far has been to create a reporter-like class that has a list of calculators as an attribute, and set up an extra attribute before running the MD:

import torch
from torch_sim.properties.correlations import VelocityAutoCorrelation

class PerSystemVACF:
    """Maintain one independent VACF calculator per simulated system."""

    def __init__(
        self,
        n_systems: int,
        *,
        window_size: int,
        device: torch.device,
        use_running_average: bool = True,
        normalize: bool = True,
    ) -> None:
        
        self.calculators = [
            VelocityAutoCorrelation(
                window_size=window_size,
                device=device,
                use_running_average=use_running_average,
                normalize=normalize,
            )
            for _ in range(n_systems)
        ]

    def __call__(self, state, model=None) -> torch.Tensor:
        # `vacf_id` is preserved when the batched state is split.
        system_id = int(state.vacf_id.item())

        # This updates only the calculator belonging to this system.
        # The returned tensor is TorchSim's window count.
        return self.calculators[system_id](state, model)

    @property
    def vacfs(self) -> torch.Tensor:
        """Shape: (n_systems, window_size)."""
        return torch.stack(
            [calculator.vacf for calculator in self.calculators],
            dim=0,
        )

# set up systems
...

# attach extra attribute
systems.system_extras["vacf_id"] = torch.arange(
    systems.n_systems,
    device=systems.device,
    dtype=torch.long
)

# create per system vacf calculator
vacf_manager = PerSystemVACF(
    n_systems=systems.n_systems,
    window_size=100,
    device=systems.device,
    use_running_average=True,
    normalize=True
)

# use in prop_calculators dict as normal
...