Convert ts.integrate param units

tests/test_runners.py

@@ -108,6 +108,81 @@ def test_integrate_double_nvt(
     assert not torch.isnan(final_state.energy).any()
 
 
+def test_integrate_converts_init_kwarg(
+    ar_supercell_sim_state: SimState, lj_model: LennardJonesModel
+) -> None:
+    """integrate scales Nose-Hoover `tau` to internal units like `timestep`.
+
+    Otherwise `tau` is ~98x too small for metal units, giving an over-stiff
+    thermostat that diverges on force spikes
+
+    See https://github.com/TorchSim/torch-sim/issues/579 for more info
+    """
+    tau = 0.1  # ps, same convention as `timestep`
+    final = ts.integrate(
+        system=ar_supercell_sim_state,
+        model=lj_model,
+        integrator=ts.Integrator.nvt_nose_hoover,
+        n_steps=1,
+        temperature=100.0,
+        timestep=0.002,
+        init_kwargs={"tau": tau},
+    )
+    expected = tau * ts.units.MetalUnits.time
+    assert torch.allclose(final.chain.tau, torch.full_like(final.chain.tau, expected))
+
+
+def test_integrator_unit_kwargs_support_shared_init_and_step_parameters() -> None:
+    """Shared and optional unit kwargs retain their conversion metadata."""
+    dt = ts.integrators.INTEGRATOR_KWARG_UNITS[ts.Integrator.nvt_nose_hoover]["dt"]
+    tau = ts.integrators.INTEGRATOR_KWARG_UNITS[ts.Integrator.nvt_nose_hoover]["tau"]
+    gamma = ts.integrators.INTEGRATOR_KWARG_UNITS[ts.Integrator.nvt_langevin]["gamma"]
+
+    assert dict(dt.factors) == {
+        "init": ts.units.MetalUnits.time,
+        "step": ts.units.MetalUnits.time,
+    }
+    assert dict(tau.factors) == {"init": ts.units.MetalUnits.time}
+    assert dict(gamma.factors) == {"step": 1 / ts.units.MetalUnits.time}
+
+
+def test_integrate_converts_step_kwarg(
+    ar_supercell_sim_state: SimState,
+    lj_model: LennardJonesModel,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """integrate divides inverse-time kwargs (e.g. Langevin `gamma`) by the unit factor.
+
+    The opposite direction from relaxation times: `gamma` is a rate, so it scales as
+    1/time, not time. It is a step kwarg, so it travels through `**integrator_kwargs`
+    to the step function (issue #579 / InverseTimeArg).
+    """
+    gamma = 10.0  # 1/ps, same time convention as `timestep`
+    # gamma is consumed inside the step function and never stored on the state, so
+    # unlike the tau test in test_integrate_converts_persistent_init_kwarg, we must
+    # spy on the step call to see the converted value.
+    received: dict[str, object] = {}
+    init, real_step = ts.integrators.INTEGRATOR_REGISTRY[ts.Integrator.nvt_langevin]
+
+    def spy_step(**kwargs: object) -> MDState:
+        received["gamma"] = kwargs["gamma"]
+        return real_step(**kwargs)
+
+    monkeypatch.setitem(
+        ts.integrators.INTEGRATOR_REGISTRY, ts.Integrator.nvt_langevin, (init, spy_step)
+    )
+    ts.integrate(
+        system=ar_supercell_sim_state,
+        model=lj_model,
+        integrator=ts.Integrator.nvt_langevin,
+        n_steps=1,
+        temperature=100.0,
+        timestep=0.002,
+        gamma=gamma,
+    )
+    assert received["gamma"] == gamma / ts.units.MetalUnits.time
+
+
 def test_integrate_double_nvt_multiple_temperatures(
     ar_double_sim_state: SimState, lj_model: LennardJonesModel
 ) -> None:

torch_sim/integrators/__init__.py

@@ -74,13 +74,26 @@
 """
 
 # ruff: noqa: F401
+import typing
 from collections.abc import Callable
+from dataclasses import dataclass
 from enum import StrEnum
-from typing import Any, Final
+from typing import Any, Final, Literal
 
 import torch_sim as ts
-
-from .md import MDState, initialize_momenta, momentum_step, position_step, velocity_verlet
+from torch_sim.units import unit_factor_from_annotation
+
+from .md import (
+    InversePressureArg,
+    InverseTimeArg,
+    MDState,
+    PressureArg,
+    TimeArg,
+    initialize_momenta,
+    momentum_step,
+    position_step,
+    velocity_verlet_step,
+)
 from .npt import (
     NPTLangevinAnisotropicState,
     NPTLangevinIsotropicState,
@@ -202,3 +215,41 @@ class Integrator(StrEnum):
         npt_crescale_triclinic_step,
     ),
 }
+
+
+@dataclass(frozen=True, slots=True)
+class IntegratorKwargUnits:
+    """Metadata for integrator parameters carrying physical units.
+
+    ``factors`` maps each registered call channel to the factor that converts
+    that channel's parameter to internal units (multiply by it).
+    """
+
+    factors: tuple[tuple[Literal["init", "step"], float], ...]
+
+
+def _collect_unit_kwargs(integrator: Integrator) -> dict[str, IntegratorKwargUnits]:
+    """Read the unit-annotated kwargs off an integrator's init/step signatures."""
+    init_fn, step_fn = INTEGRATOR_REGISTRY[integrator]
+    unit_kwargs: dict[str, dict[Literal["init", "step"], float]] = {}
+    for channel, fn in (("init", init_fn), ("step", step_fn)):
+        hints = typing.get_type_hints(fn, include_extras=True)
+        for name, hint in hints.items():
+            if (factor := unit_factor_from_annotation(hint)) is None:
+                continue
+            unit_kwargs.setdefault(name, {})[channel] = factor
+    return {
+        name: IntegratorKwargUnits(tuple(channel_factors.items()))
+        for name, channel_factors in unit_kwargs.items()
+    }
+
+
+#: Unit-carrying kwargs of each integrator, derived at import time from the
+#: ``TimeArg``/``InverseTimeArg``/``PressureArg``/``InversePressureArg`` annotations
+#: on the registered init/step functions. :func:`torch_sim.runners.integrate` reads
+#: this to unit-convert kwargs.
+INTEGRATOR_KWARG_UNITS: Final[dict[Integrator, dict[str, IntegratorKwargUnits]]] = {
+    integrator: unit_kwargs
+    for integrator in Integrator
+    if (unit_kwargs := _collect_unit_kwargs(integrator))
+}

torch_sim/integrators/md.py

@@ -1,22 +1,33 @@
 """Core molecular dynamics state and operations."""
 
 import logging
-import warnings
 from collections.abc import Callable
 from dataclasses import dataclass
+from typing import Annotated
 
 import torch
 
 from torch_sim._duecredit import dcite
 from torch_sim.models.interface import ModelInterface
 from torch_sim.quantities import calc_kT
 from torch_sim.state import SimState
-from torch_sim.units import MetalUnits
+from torch_sim.units import MetalUnits, UnitAnnotation
 
 
 logger = logging.getLogger(__name__)
 
 
+# The metadata is the factor that converts the kwarg to internal units
+TemperatureArg = Annotated[float | torch.Tensor, UnitAnnotation(MetalUnits.temperature)]
+EnergyArg = Annotated[float | torch.Tensor, UnitAnnotation(MetalUnits.energy)]
+TimeArg = Annotated[float | torch.Tensor, UnitAnnotation(MetalUnits.time)]
+InverseTimeArg = Annotated[float | torch.Tensor, UnitAnnotation(1 / MetalUnits.time)]
+PressureArg = Annotated[float | torch.Tensor, UnitAnnotation(MetalUnits.pressure)]
+InversePressureArg = Annotated[
+    float | torch.Tensor, UnitAnnotation(1 / MetalUnits.pressure)
+]
+
+
 @dataclass(kw_only=True)
 class MDState(SimState):
     """State information for molecular dynamics simulations.
@@ -103,7 +114,7 @@ def initialize_momenta(
     positions: torch.Tensor,
     masses: torch.Tensor,
     system_idx: torch.Tensor,
-    kT: float | torch.Tensor,
+    kT: torch.Tensor,
     generator: torch.Generator | None = None,
 ) -> torch.Tensor:
     """Initialize particle momenta based on temperature.
@@ -125,7 +136,6 @@ def initialize_momenta(
     device = positions.device
     dtype = positions.dtype
 
-    kT = torch.as_tensor(kT, device=device, dtype=dtype)
     if kT.ndim > 0:
         kT = kT[system_idx]
 
@@ -156,7 +166,7 @@ def initialize_momenta(
     )
 
 
-def momentum_step[T: MDState](state: T, dt: float | torch.Tensor) -> T:
+def momentum_step[T: MDState](state: T, dt: TimeArg) -> T:
     """Update particle momenta using current forces.
 
     This function performs the momentum update step of velocity Verlet integration
@@ -177,7 +187,7 @@ def momentum_step[T: MDState](state: T, dt: float | torch.Tensor) -> T:
     return state
 
 
-def position_step[T: MDState](state: T, dt: float | torch.Tensor) -> T:
+def position_step[T: MDState](state: T, dt: TimeArg) -> T:
     """Update particle positions using current velocities.
 
     This function performs the position update step of velocity Verlet integration
@@ -198,9 +208,7 @@ def position_step[T: MDState](state: T, dt: float | torch.Tensor) -> T:
     return state
 
 
-def velocity_verlet_step[T: MDState](
-    state: T, dt: float | torch.Tensor, model: ModelInterface
-) -> T:
+def velocity_verlet_step[T: MDState](state: T, dt: TimeArg, model: ModelInterface) -> T:
     """Perform one complete velocity Verlet integration step.
 
     This function implements the velocity Verlet algorithm, which provides
@@ -237,16 +245,6 @@ def velocity_verlet_step[T: MDState](
     return momentum_step(state, dt_2)
 
 
-def velocity_verlet[T: MDState](
-    state: T, dt: float | torch.Tensor, model: ModelInterface
-) -> T:
-    """Deprecated alias for velocity_verlet_step."""
-    msg = "velocity_verlet is deprecated. Use velocity_verlet_step instead."
-    warnings.warn(msg, DeprecationWarning, stacklevel=2)
-    logger.warning(msg)
-    return velocity_verlet_step(state=state, dt=dt, model=model)
-
-
 @dataclass
 class NoseHooverChain:
     """State information for a Nose-Hoover chain thermostat.
@@ -338,11 +336,11 @@ class NoseHooverChainFns:
 @dcite("10.2183/pjab.69.161")
 @dcite("10.1016/0375-9601(90)90092-3")
 def construct_nose_hoover_chain(  # noqa: C901 PLR0915
-    dt: float | torch.Tensor,
+    dt: TimeArg,
     chain_length: int,
     chain_steps: int,
     sy_steps: int,
-    tau: float | torch.Tensor,
+    tau: TimeArg,
 ) -> NoseHooverChainFns:
     """Creates functions to simulate a Nose-Hoover Chain thermostat.