From 97f483ae177dee4df8e55ceca13b5ee2e2e2fd70 Mon Sep 17 00:00:00 2001
From: Curtis Chong <curtischong5@gmail.com>
Date: Sat, 20 Jun 2026 13:50:27 -0700
Subject: [PATCH 1/6] wip fix 579

add run cmd

simplify repro

new repro

repro fix

use their model

integrator fixes

revert to sevennet

cleanup

get structures in memory

better integrator definitions

cleanup

better test names

more cleanup

cleanup

more cleanup

cleanup

cleanup

cleanup

delete repro file

fix other params

better name

use constants

cleanup
---
 tests/test_runners.py             | 61 +++++++++++++++++++++++++++++
 torch_sim/integrators/__init__.py | 65 ++++++++++++++++++++++++++++++-
 torch_sim/integrators/md.py       |  8 ++++
 torch_sim/integrators/npt.py      | 46 ++++++++++++----------
 torch_sim/integrators/nvt.py      |  8 ++--
 torch_sim/runners.py              | 20 +++++++---
 6 files changed, 176 insertions(+), 32 deletions(-)

diff --git a/tests/test_runners.py b/tests/test_runners.py
index 8a817c1f6..7e03d7e54 100644
--- a/tests/test_runners.py
+++ b/tests/test_runners.py
@@ -108,6 +108,67 @@ 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_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:
diff --git a/torch_sim/integrators/__init__.py b/torch_sim/integrators/__init__.py
index e6f09f534..3af9ac8fe 100644
--- a/torch_sim/integrators/__init__.py
+++ b/torch_sim/integrators/__init__.py
@@ -74,13 +74,25 @@
 """
 
 # 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 .md import (
+    InversePressureArg,
+    InverseTimeArg,
+    MDState,
+    PressureArg,
+    TimeArg,
+    initialize_momenta,
+    momentum_step,
+    position_step,
+    velocity_verlet,
+)
 from .npt import (
     NPTLangevinAnisotropicState,
     NPTLangevinIsotropicState,
@@ -202,3 +214,52 @@ class Integrator(StrEnum):
         npt_crescale_triclinic_step,
     ),
 }
+
+
+@dataclass(frozen=True, slots=True)
+class UnitKwarg:
+    """Metadata for integrator parameters carrying physical units.
+
+    ``factor`` converts the parameter to internal units (multiply by it).
+    ``channel`` is the function the parameter is used in (either ``"init"``
+    or ``"step"``).
+    """
+
+    factor: float
+    channel: Literal["init", "step"]
+
+
+_UNIT_FACTORS: Final = frozenset(
+    arg.__metadata__[0]
+    for arg in (TimeArg, InverseTimeArg, PressureArg, InversePressureArg)
+)
+
+
+def _collect_unit_kwargs(integrator: Integrator) -> dict[str, UnitKwarg]:
+    """Read the unit-annotated kwargs off an integrator's init/step signatures."""
+    init_fn, step_fn = INTEGRATOR_REGISTRY[integrator]
+    unit_kwargs: dict[str, UnitKwarg] = {}
+    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():
+            for meta in getattr(hint, "__metadata__", ()):
+                if isinstance(meta, float) and meta in _UNIT_FACTORS:
+                    if name in unit_kwargs:
+                        raise TypeError(
+                            f"{integrator}: unit kwarg {name!r} is annotated in both "
+                            f"the init and step functions, making its `integrate` "
+                            f"channel ambiguous"
+                        )
+                    unit_kwargs[name] = UnitKwarg(float(meta), channel)
+    return unit_kwargs
+
+
+#: 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_UNIT_KWARGS: Final[dict[Integrator, dict[str, UnitKwarg]]] = {
+    integrator: unit_kwargs
+    for integrator in Integrator
+    if (unit_kwargs := _collect_unit_kwargs(integrator))
+}
diff --git a/torch_sim/integrators/md.py b/torch_sim/integrators/md.py
index a88685746..1b82d660e 100644
--- a/torch_sim/integrators/md.py
+++ b/torch_sim/integrators/md.py
@@ -4,6 +4,7 @@
 import warnings
 from collections.abc import Callable
 from dataclasses import dataclass
+from typing import Annotated
 
 import torch
 
@@ -17,6 +18,13 @@
 logger = logging.getLogger(__name__)
 
 
+# The metadata is the factor that converts the kwarg to internal units
+TimeArg = Annotated[float | torch.Tensor | None, MetalUnits.time]
+InverseTimeArg = Annotated[float | torch.Tensor | None, 1 / MetalUnits.time]
+PressureArg = Annotated[float | torch.Tensor, MetalUnits.pressure]
+InversePressureArg = Annotated[float | torch.Tensor | None, 1 / MetalUnits.pressure]
+
+
 @dataclass(kw_only=True)
 class MDState(SimState):
     """State information for molecular dynamics simulations.
diff --git a/torch_sim/integrators/npt.py b/torch_sim/integrators/npt.py
index 35cacd32f..f12fc39a5 100644
--- a/torch_sim/integrators/npt.py
+++ b/torch_sim/integrators/npt.py
@@ -11,9 +11,13 @@
 import torch_sim as ts
 from torch_sim._duecredit import dcite
 from torch_sim.integrators.md import (
+    InversePressureArg,
+    InverseTimeArg,
     MDState,
     NoseHooverChain,
     NoseHooverChainFns,
+    PressureArg,
+    TimeArg,
     construct_nose_hoover_chain,
     initialize_momenta,
     momentum_step,
@@ -407,9 +411,9 @@ def npt_langevin_anisotropic_init(
     *,
     kT: float | torch.Tensor,
     dt: float | torch.Tensor,
-    alpha: float | torch.Tensor | None = None,
-    cell_alpha: float | torch.Tensor | None = None,
-    b_tau: float | torch.Tensor | None = None,
+    alpha: InverseTimeArg = None,
+    cell_alpha: InverseTimeArg = None,
+    b_tau: TimeArg = None,
     **_kwargs: Any,
 ) -> NPTLangevinAnisotropicState:
     """Initialize NPT Langevin state with independent per-dimension cell lengths.
@@ -511,7 +515,7 @@ def npt_langevin_anisotropic_step(
     *,
     dt: float | torch.Tensor,
     kT: float | torch.Tensor,
-    external_pressure: float | torch.Tensor,
+    external_pressure: PressureArg,
 ) -> NPTLangevinAnisotropicState:
     r"""Perform one NPT Langevin step with independent per-dimension cell lengths.
 
@@ -886,9 +890,9 @@ def npt_langevin_isotropic_init(
     *,
     kT: float | torch.Tensor,
     dt: float | torch.Tensor,
-    alpha: float | torch.Tensor | None = None,
-    cell_alpha: float | torch.Tensor | None = None,
-    b_tau: float | torch.Tensor | None = None,
+    alpha: InverseTimeArg = None,
+    cell_alpha: InverseTimeArg = None,
+    b_tau: TimeArg = None,
     **_kwargs: Any,
 ) -> NPTLangevinIsotropicState:
     """Initialize an NPT Langevin state using logarithmic strain coordinate.
@@ -986,7 +990,7 @@ def npt_langevin_isotropic_step(
     *,
     dt: float | torch.Tensor,
     kT: float | torch.Tensor,
-    external_pressure: float | torch.Tensor,
+    external_pressure: PressureArg,
 ) -> NPTLangevinIsotropicState:
     r"""Perform one NPT Langevin step using logarithmic strain coordinate.
 
@@ -1643,8 +1647,8 @@ def npt_nose_hoover_isotropic_init(
     chain_length: int = 3,
     chain_steps: int = 2,
     sy_steps: int = 3,
-    t_tau: float | torch.Tensor | None = None,
-    b_tau: float | torch.Tensor | None = None,
+    t_tau: TimeArg = None,
+    b_tau: TimeArg = None,
     **kwargs: Any,
 ) -> NPTNoseHooverIsotropicState:
     """Initialize the NPT Nose-Hoover state.
@@ -1810,7 +1814,7 @@ def npt_nose_hoover_isotropic_step(
     *,
     dt: float | torch.Tensor,
     kT: float | torch.Tensor,
-    external_pressure: float | torch.Tensor,
+    external_pressure: PressureArg,
 ) -> NPTNoseHooverIsotropicState:
     r"""Perform a complete NPT integration step with Nose-Hoover chain thermostats.
 
@@ -2515,8 +2519,8 @@ def npt_crescale_triclinic_step(
     *,
     dt: float | torch.Tensor,
     kT: float | torch.Tensor,
-    external_pressure: float | torch.Tensor,
-    tau: float | torch.Tensor | None = None,
+    external_pressure: PressureArg,
+    tau: TimeArg = None,
 ) -> NPTCRescaleState:
     r"""Perform one NPT integration step with anisotropic stochastic cell rescaling.
 
@@ -2643,8 +2647,8 @@ def npt_crescale_anisotropic_step(
     *,
     dt: float | torch.Tensor,
     kT: float | torch.Tensor,
-    external_pressure: float | torch.Tensor,
-    tau: float | torch.Tensor | None = None,
+    external_pressure: PressureArg,
+    tau: TimeArg = None,
 ) -> NPTCRescaleState:
     """Perform one NPT integration step with cell rescaling barostat.
 
@@ -2719,8 +2723,8 @@ def npt_crescale_triclinic_average_step(
     *,
     dt: float | torch.Tensor,
     kT: float | torch.Tensor,
-    external_pressure: float | torch.Tensor,
-    tau: float | torch.Tensor | None = None,
+    external_pressure: PressureArg,
+    tau: TimeArg = None,
 ) -> NPTCRescaleState:
     """Perform one NPT integration step with cell rescaling barostat.
 
@@ -2795,8 +2799,8 @@ def npt_crescale_isotropic_step(
     *,
     dt: float | torch.Tensor,
     kT: float | torch.Tensor,
-    external_pressure: float | torch.Tensor,
-    tau: float | torch.Tensor | None = None,
+    external_pressure: PressureArg,
+    tau: TimeArg = None,
 ) -> NPTCRescaleState:
     r"""Perform one NPT integration step with isotropic stochastic cell rescaling.
 
@@ -2902,8 +2906,8 @@ def npt_crescale_init(
     *,
     kT: float | torch.Tensor,
     dt: float | torch.Tensor,
-    tau_p: float | torch.Tensor | None = None,
-    isothermal_compressibility: float | torch.Tensor | None = None,
+    tau_p: TimeArg = None,
+    isothermal_compressibility: InversePressureArg = None,
 ) -> NPTCRescaleState:
     """Initialize the NPT cell rescaling state.
 
diff --git a/torch_sim/integrators/nvt.py b/torch_sim/integrators/nvt.py
index 644b10897..9007e63dc 100644
--- a/torch_sim/integrators/nvt.py
+++ b/torch_sim/integrators/nvt.py
@@ -8,9 +8,11 @@
 import torch_sim as ts
 from torch_sim._duecredit import dcite
 from torch_sim.integrators.md import (
+    InverseTimeArg,
     MDState,
     NoseHooverChain,
     NoseHooverChainFns,
+    TimeArg,
     construct_nose_hoover_chain,
     initialize_momenta,
     momentum_step,
@@ -143,7 +145,7 @@ def nvt_langevin_step(
     *,
     dt: float | torch.Tensor,
     kT: float | torch.Tensor,
-    gamma: float | torch.Tensor | None = None,
+    gamma: InverseTimeArg = None,
 ) -> MDState:
     r"""Perform one complete Langevin dynamics integration step using the BAOAB scheme.
 
@@ -291,7 +293,7 @@ def nvt_nose_hoover_init(
     *,
     kT: float | torch.Tensor,
     dt: float | torch.Tensor,
-    tau: float | torch.Tensor | None = None,
+    tau: TimeArg = None,
     chain_length: int = 3,
     chain_steps: int = 3,
     sy_steps: int = 3,
@@ -711,7 +713,7 @@ def nvt_vrescale_step(
     *,
     dt: float | torch.Tensor,
     kT: float | torch.Tensor,
-    tau: float | torch.Tensor | None = None,
+    tau: TimeArg = None,
 ) -> NVTVRescaleState:
     r"""Perform one complete V-Rescale (CSVR) dynamics integration step.
 
diff --git a/torch_sim/runners.py b/torch_sim/runners.py
index cd46cc7ce..2f7bb7ccd 100644
--- a/torch_sim/runners.py
+++ b/torch_sim/runners.py
@@ -18,7 +18,7 @@
 
 import torch_sim as ts
 from torch_sim.autobatching import BinningAutoBatcher, InFlightAutoBatcher
-from torch_sim.integrators import INTEGRATOR_REGISTRY, Integrator
+from torch_sim.integrators import INTEGRATOR_REGISTRY, INTEGRATOR_UNIT_KWARGS, Integrator
 from torch_sim.integrators.md import MDState
 from torch_sim.models.interface import ModelInterface
 from torch_sim.optimizers import OPTIM_REGISTRY, FireState, Optimizer, OptimState
@@ -249,7 +249,7 @@ def _write_initial_state(
             trajectory_reporter.report(state, 0, model=model)
 
 
-def integrate[T: SimState](  # noqa: C901
+def integrate[T: SimState](  # noqa: C901, PLR0915
     system: StateLike,
     model: ModelInterface,
     *,
@@ -287,7 +287,7 @@ def integrate[T: SimState](  # noqa: C901
             it's passed to `tqdm` as kwargs.
         init_kwargs (dict[str, Any], optional): Additional keyword arguments for
             integrator init function.
-        **integrator_kwargs: Additional keyword arguments for integrator init function
+        **integrator_kwargs: Additional keyword arguments for integrator step function
 
     Returns:
         T: Final state after integration
@@ -320,6 +320,16 @@ def integrate[T: SimState](  # noqa: C901
             f"integrator must be key from Integrator or a tuple of "
             f"(init_func, step_func), got {type(integrator)}"
         )
+
+    # Like `timestep` above, unit-carrying kwargs (e.g. `tau`, `gamma`,
+    # `external_pressure`) must be converted to internal units per
+    # INTEGRATOR_UNIT_KWARGS
+    init_kwargs = dict(init_kwargs or {})
+    channels = {"init": init_kwargs, "step": integrator_kwargs}
+    for key, meta in INTEGRATOR_UNIT_KWARGS.get(integrator, {}).items():
+        kwargs = channels[meta.channel]
+        if kwargs.get(key) is not None:
+            kwargs[key] = kwargs[key] * meta.factor
     # batch_iterator will be a list if autobatcher is False
     batch_iterator = _configure_batches_iterator(
         initial_state, model, autobatcher=autobatcher
@@ -351,9 +361,7 @@ def integrate[T: SimState](  # noqa: C901
         batch_kT = (
             kTs[:, system_indices] if (system_indices and len(kTs.shape) == 2) else kTs
         )
-        state = init_func(
-            state=state, model=model, kT=batch_kT[0], dt=dt, **init_kwargs or {}
-        )
+        state = init_func(state=state, model=model, kT=batch_kT[0], dt=dt, **init_kwargs)
 
         # set up trajectory reporters
         if autobatcher and trajectory_reporter is not None and og_filenames is not None:

From 36e2f3612662fb1ac826e9a46a0878d21406b489 Mon Sep 17 00:00:00 2001
From: Curtis Chong <curtischong5@gmail.com>
Date: Fri, 3 Jul 2026 12:12:21 -0700
Subject: [PATCH 2/6] fix lint

---
 torch_sim/units.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/torch_sim/units.py b/torch_sim/units.py
index 19cd6c4fc..146adab22 100644
--- a/torch_sim/units.py
+++ b/torch_sim/units.py
@@ -15,7 +15,7 @@ class BaseConstant(float, Enum):
 
     References:
         http://arxiv.org/pdf/1507.07956.pdf
-        https://wiki.fysik.dtu.dk/ase/_modules/ase/units.html#create_units
+        https://docs.ase-lib.org/_modules/ase/units.html#create_units
     """
 
     def __new__(cls, value: float) -> Self:

From 43cf14c0182e812ee314dcd0c24fedea25cac35f Mon Sep 17 00:00:00 2001
From: Rhys Goodall <rhys.goodall@outlook.com>
Date: Thu, 23 Jul 2026 11:36:16 +0100
Subject: [PATCH 3/6] refactor: replace scalar type hints with specialized
 physical quantity annotations across MD integrators

---
 torch_sim/integrators/md.py  | 26 ++++++-----
 torch_sim/integrators/npt.py | 83 ++++++++++++++++++------------------
 torch_sim/integrators/nve.py |  6 ++-
 torch_sim/integrators/nvt.py | 37 ++++++++--------
 4 files changed, 77 insertions(+), 75 deletions(-)

diff --git a/torch_sim/integrators/md.py b/torch_sim/integrators/md.py
index 1b82d660e..b0c832269 100644
--- a/torch_sim/integrators/md.py
+++ b/torch_sim/integrators/md.py
@@ -19,10 +19,12 @@
 
 
 # The metadata is the factor that converts the kwarg to internal units
-TimeArg = Annotated[float | torch.Tensor | None, MetalUnits.time]
-InverseTimeArg = Annotated[float | torch.Tensor | None, 1 / MetalUnits.time]
+TemperatureArg = Annotated[float | torch.Tensor, MetalUnits.temperature]
+EnergyArg = Annotated[float | torch.Tensor, MetalUnits.energy]
+TimeArg = Annotated[float | torch.Tensor, MetalUnits.time]
+InverseTimeArg = Annotated[float | torch.Tensor, 1 / MetalUnits.time]
 PressureArg = Annotated[float | torch.Tensor, MetalUnits.pressure]
-InversePressureArg = Annotated[float | torch.Tensor | None, 1 / MetalUnits.pressure]
+InversePressureArg = Annotated[float | torch.Tensor, 1 / MetalUnits.pressure]
 
 
 @dataclass(kw_only=True)
@@ -111,7 +113,7 @@ def initialize_momenta(
     positions: torch.Tensor,
     masses: torch.Tensor,
     system_idx: torch.Tensor,
-    kT: float | torch.Tensor,
+    kT: EnergyArg,
     generator: torch.Generator | None = None,
 ) -> torch.Tensor:
     """Initialize particle momenta based on temperature.
@@ -164,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
@@ -185,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
@@ -206,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
@@ -245,9 +245,7 @@ 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:
+def velocity_verlet[T: MDState](state: T, dt: TimeArg, 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)
@@ -346,11 +344,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.
 
diff --git a/torch_sim/integrators/npt.py b/torch_sim/integrators/npt.py
index f12fc39a5..0de3a39be 100644
--- a/torch_sim/integrators/npt.py
+++ b/torch_sim/integrators/npt.py
@@ -11,6 +11,7 @@
 import torch_sim as ts
 from torch_sim._duecredit import dcite
 from torch_sim.integrators.md import (
+    EnergyArg,
     InversePressureArg,
     InverseTimeArg,
     MDState,
@@ -409,11 +410,11 @@ def npt_langevin_anisotropic_init(
     state: SimState,
     model: ModelInterface,
     *,
-    kT: float | torch.Tensor,
-    dt: float | torch.Tensor,
-    alpha: InverseTimeArg = None,
-    cell_alpha: InverseTimeArg = None,
-    b_tau: TimeArg = None,
+    kT: EnergyArg,
+    dt: TimeArg,
+    alpha: InverseTimeArg | None = None,
+    cell_alpha: InverseTimeArg | None = None,
+    b_tau: TimeArg | None = None,
     **_kwargs: Any,
 ) -> NPTLangevinAnisotropicState:
     """Initialize NPT Langevin state with independent per-dimension cell lengths.
@@ -513,8 +514,8 @@ def npt_langevin_anisotropic_step(
     state: NPTLangevinAnisotropicState,
     model: ModelInterface,
     *,
-    dt: float | torch.Tensor,
-    kT: float | torch.Tensor,
+    dt: TimeArg,
+    kT: EnergyArg,
     external_pressure: PressureArg,
 ) -> NPTLangevinAnisotropicState:
     r"""Perform one NPT Langevin step with independent per-dimension cell lengths.
@@ -702,8 +703,8 @@ def volume(self) -> torch.Tensor:
 
 def _compute_isotropic_cell_force(
     state: NPTLangevinIsotropicState,
-    external_pressure: float | torch.Tensor,
-    kT: float | torch.Tensor,
+    external_pressure: PressureArg,
+    kT: EnergyArg,
 ) -> torch.Tensor:
     """Compute force on the strain coordinate ε.
 
@@ -888,11 +889,11 @@ def npt_langevin_isotropic_init(
     state: SimState,
     model: ModelInterface,
     *,
-    kT: float | torch.Tensor,
-    dt: float | torch.Tensor,
-    alpha: InverseTimeArg = None,
-    cell_alpha: InverseTimeArg = None,
-    b_tau: TimeArg = None,
+    kT: EnergyArg,
+    dt: TimeArg,
+    alpha: InverseTimeArg | None = None,
+    cell_alpha: InverseTimeArg | None = None,
+    b_tau: TimeArg | None = None,
     **_kwargs: Any,
 ) -> NPTLangevinIsotropicState:
     """Initialize an NPT Langevin state using logarithmic strain coordinate.
@@ -988,8 +989,8 @@ def npt_langevin_isotropic_step(
     state: NPTLangevinIsotropicState,
     model: ModelInterface,
     *,
-    dt: float | torch.Tensor,
-    kT: float | torch.Tensor,
+    dt: TimeArg,
+    kT: EnergyArg,
     external_pressure: PressureArg,
 ) -> NPTLangevinIsotropicState:
     r"""Perform one NPT Langevin step using logarithmic strain coordinate.
@@ -1642,13 +1643,13 @@ def npt_nose_hoover_isotropic_init(
     state: SimState,
     model: ModelInterface,
     *,
-    kT: float | torch.Tensor,
-    dt: float | torch.Tensor,
+    kT: EnergyArg,
+    dt: TimeArg,
     chain_length: int = 3,
     chain_steps: int = 2,
     sy_steps: int = 3,
-    t_tau: TimeArg = None,
-    b_tau: TimeArg = None,
+    t_tau: TimeArg | None = None,
+    b_tau: TimeArg | None = None,
     **kwargs: Any,
 ) -> NPTNoseHooverIsotropicState:
     """Initialize the NPT Nose-Hoover state.
@@ -1812,8 +1813,8 @@ def npt_nose_hoover_isotropic_step(
     state: NPTNoseHooverIsotropicState,
     model: ModelInterface,
     *,
-    dt: float | torch.Tensor,
-    kT: float | torch.Tensor,
+    dt: TimeArg,
+    kT: EnergyArg,
     external_pressure: PressureArg,
 ) -> NPTNoseHooverIsotropicState:
     r"""Perform a complete NPT integration step with Nose-Hoover chain thermostats.
@@ -2493,10 +2494,10 @@ def _crescale_isotropic_barostat_step(
 
 def _coerce_crescale_step_inputs(
     state: NPTCRescaleState,
-    dt: float | torch.Tensor,
-    kT: float | torch.Tensor,
+    dt: TimeArg,
+    kT: EnergyArg,
     external_pressure: float | torch.Tensor,
-    tau: float | torch.Tensor | None,
+    tau: TimeArg | None,
 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
     """Normalize scalar-or-tensor C-rescale step parameters to state tensors."""
     device, dtype = state.device, state.dtype
@@ -2517,10 +2518,10 @@ def npt_crescale_triclinic_step(
     state: NPTCRescaleState,
     model: ModelInterface,
     *,
-    dt: float | torch.Tensor,
-    kT: float | torch.Tensor,
+    dt: TimeArg,
+    kT: EnergyArg,
     external_pressure: PressureArg,
-    tau: TimeArg = None,
+    tau: TimeArg | None = None,
 ) -> NPTCRescaleState:
     r"""Perform one NPT integration step with anisotropic stochastic cell rescaling.
 
@@ -2645,10 +2646,10 @@ def npt_crescale_anisotropic_step(
     state: NPTCRescaleState,
     model: ModelInterface,
     *,
-    dt: float | torch.Tensor,
-    kT: float | torch.Tensor,
+    dt: TimeArg,
+    kT: EnergyArg,
     external_pressure: PressureArg,
-    tau: TimeArg = None,
+    tau: TimeArg | None = None,
 ) -> NPTCRescaleState:
     """Perform one NPT integration step with cell rescaling barostat.
 
@@ -2721,10 +2722,10 @@ def npt_crescale_triclinic_average_step(
     state: NPTCRescaleState,
     model: ModelInterface,
     *,
-    dt: float | torch.Tensor,
-    kT: float | torch.Tensor,
+    dt: TimeArg,
+    kT: EnergyArg,
     external_pressure: PressureArg,
-    tau: TimeArg = None,
+    tau: TimeArg | None = None,
 ) -> NPTCRescaleState:
     """Perform one NPT integration step with cell rescaling barostat.
 
@@ -2797,10 +2798,10 @@ def npt_crescale_isotropic_step(
     state: NPTCRescaleState,
     model: ModelInterface,
     *,
-    dt: float | torch.Tensor,
-    kT: float | torch.Tensor,
+    dt: TimeArg,
+    kT: EnergyArg,
     external_pressure: PressureArg,
-    tau: TimeArg = None,
+    tau: TimeArg | None = None,
 ) -> NPTCRescaleState:
     r"""Perform one NPT integration step with isotropic stochastic cell rescaling.
 
@@ -2904,10 +2905,10 @@ def npt_crescale_init(
     state: SimState,
     model: ModelInterface,
     *,
-    kT: float | torch.Tensor,
-    dt: float | torch.Tensor,
-    tau_p: TimeArg = None,
-    isothermal_compressibility: InversePressureArg = None,
+    kT: EnergyArg,
+    dt: TimeArg,
+    tau_p: TimeArg | None = None,
+    isothermal_compressibility: InversePressureArg | None = None,
 ) -> NPTCRescaleState:
     """Initialize the NPT cell rescaling state.
 
diff --git a/torch_sim/integrators/nve.py b/torch_sim/integrators/nve.py
index 6aa3053c6..40dde09f7 100644
--- a/torch_sim/integrators/nve.py
+++ b/torch_sim/integrators/nve.py
@@ -5,7 +5,9 @@
 import torch
 
 from torch_sim.integrators.md import (
+    EnergyArg,
     MDState,
+    TimeArg,
     initialize_momenta,
     momentum_step,
     position_step,
@@ -18,7 +20,7 @@ def nve_init(
     state: SimState,
     model: ModelInterface,
     *,
-    kT: float | torch.Tensor,
+    kT: EnergyArg,
     **_kwargs: Any,
 ) -> MDState:
     """Initialize an NVE state from input data.
@@ -68,7 +70,7 @@ def nve_init(
 
 
 def nve_step(
-    state: MDState, model: ModelInterface, *, dt: float | torch.Tensor, **_kwargs: Any
+    state: MDState, model: ModelInterface, *, dt: TimeArg, **_kwargs: Any
 ) -> MDState:
     r"""Perform one complete NVE (microcanonical) integration step.
 
diff --git a/torch_sim/integrators/nvt.py b/torch_sim/integrators/nvt.py
index 9007e63dc..eac1f5c66 100644
--- a/torch_sim/integrators/nvt.py
+++ b/torch_sim/integrators/nvt.py
@@ -8,6 +8,7 @@
 import torch_sim as ts
 from torch_sim._duecredit import dcite
 from torch_sim.integrators.md import (
+    EnergyArg,
     InverseTimeArg,
     MDState,
     NoseHooverChain,
@@ -25,8 +26,8 @@
 
 def _ou_step(
     state: MDState,
-    dt: float | torch.Tensor,
-    kT: float | torch.Tensor,
+    dt: TimeArg,
+    kT: EnergyArg,
     gamma: float | torch.Tensor,
 ) -> MDState:
     """Apply stochastic noise and friction for Langevin dynamics.
@@ -89,7 +90,7 @@ def nvt_langevin_init(
     state: SimState,
     model: ModelInterface,
     *,
-    kT: float | torch.Tensor,
+    kT: EnergyArg,
     **_kwargs: Any,
 ) -> MDState:
     """Initialize an NVT state from input data for Langevin dynamics.
@@ -143,9 +144,9 @@ def nvt_langevin_step(
     state: MDState,
     model: ModelInterface,
     *,
-    dt: float | torch.Tensor,
-    kT: float | torch.Tensor,
-    gamma: InverseTimeArg = None,
+    dt: TimeArg,
+    kT: EnergyArg,
+    gamma: InverseTimeArg | None = None,
 ) -> MDState:
     r"""Perform one complete Langevin dynamics integration step using the BAOAB scheme.
 
@@ -291,9 +292,9 @@ def nvt_nose_hoover_init(
     state: SimState,
     model: ModelInterface,
     *,
-    kT: float | torch.Tensor,
-    dt: float | torch.Tensor,
-    tau: TimeArg = None,
+    kT: EnergyArg,
+    dt: TimeArg,
+    tau: TimeArg | None = None,
     chain_length: int = 3,
     chain_steps: int = 3,
     sy_steps: int = 3,
@@ -379,8 +380,8 @@ def nvt_nose_hoover_step(
     state: NVTNoseHooverState,
     model: ModelInterface,
     *,
-    dt: float | torch.Tensor,
-    kT: float | torch.Tensor,
+    dt: TimeArg,
+    kT: EnergyArg,
 ) -> NVTNoseHooverState:
     r"""Perform one complete Nose-Hoover chain (NHC) integration step.
 
@@ -590,9 +591,9 @@ def get_number_of_degrees_of_freedom(self) -> torch.Tensor:
 
 def _vrescale_update[T: MDState](
     state: T,
-    tau: float | torch.Tensor,
-    kT: float | torch.Tensor,
-    dt: float | torch.Tensor,
+    tau: TimeArg,
+    kT: EnergyArg,
+    dt: TimeArg,
 ) -> T:
     """Update the momentum by a scaling factor as described by Eq.A7 Bussi et al.
 
@@ -655,7 +656,7 @@ def nvt_vrescale_init(
     state: SimState,
     model: ModelInterface,
     *,
-    kT: float | torch.Tensor,
+    kT: EnergyArg,
     **_kwargs: Any,
 ) -> NVTVRescaleState:
     """Initialize an NVT state from input data for velocity rescaling dynamics.
@@ -711,9 +712,9 @@ def nvt_vrescale_step(
     model: ModelInterface,
     state: NVTVRescaleState,
     *,
-    dt: float | torch.Tensor,
-    kT: float | torch.Tensor,
-    tau: TimeArg = None,
+    dt: TimeArg,
+    kT: EnergyArg,
+    tau: TimeArg | None = None,
 ) -> NVTVRescaleState:
     r"""Perform one complete V-Rescale (CSVR) dynamics integration step.
 

From 35445871d3b83c1d7a2aacb07218c674f709307b Mon Sep 17 00:00:00 2001
From: Curtis Chong <curtischong5@gmail.com>
Date: Thu, 23 Jul 2026 11:16:49 -0700
Subject: [PATCH 4/6] make init kwargs copy more explicit

---
 torch_sim/runners.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/torch_sim/runners.py b/torch_sim/runners.py
index e7bcba8b6..23243e643 100644
--- a/torch_sim/runners.py
+++ b/torch_sim/runners.py
@@ -321,10 +321,12 @@ def integrate[T: SimState](  # noqa: C901, PLR0915
             f"(init_func, step_func), got {type(integrator)}"
         )
 
+    # explicitly copy the kwargs since we modify the copied dict below
+    init_kwargs = {} if init_kwargs is None else init_kwargs.copy()
+
     # Like `timestep` above, unit-carrying kwargs (e.g. `tau`, `gamma`,
     # `external_pressure`) must be converted to internal units per
     # INTEGRATOR_UNIT_KWARGS
-    init_kwargs = dict(init_kwargs or {})
     channels = {"init": init_kwargs, "step": integrator_kwargs}
     for key, meta in INTEGRATOR_UNIT_KWARGS.get(integrator, {}).items():
         kwargs = channels[meta.channel]

From 5f65e31f882c7e046d86e897fe9e2f2230161f5b Mon Sep 17 00:00:00 2001
From: Rhys Goodall <rhys.goodall@outlook.com>
Date: Fri, 24 Jul 2026 12:47:07 +0100
Subject: [PATCH 5/6] refactor: generalize unit conversion metadata to support
 multi-channel integrator parameters

---
 tests/test_runners.py             | 14 +++++++++++
 torch_sim/integrators/__init__.py | 40 ++++++++++++-------------------
 torch_sim/integrators/md.py       | 16 +++++++------
 torch_sim/runners.py              | 23 +++++++++---------
 torch_sim/units.py                | 25 +++++++++++++++++++
 5 files changed, 74 insertions(+), 44 deletions(-)

diff --git a/tests/test_runners.py b/tests/test_runners.py
index 7e03d7e54..b70443289 100644
--- a/tests/test_runners.py
+++ b/tests/test_runners.py
@@ -132,6 +132,20 @@ def test_integrate_converts_init_kwarg(
     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,
diff --git a/torch_sim/integrators/__init__.py b/torch_sim/integrators/__init__.py
index 3af9ac8fe..e86b0cd98 100644
--- a/torch_sim/integrators/__init__.py
+++ b/torch_sim/integrators/__init__.py
@@ -81,6 +81,7 @@
 from typing import Any, Final, Literal
 
 import torch_sim as ts
+from torch_sim.units import unit_factor_from_annotation
 
 from .md import (
     InversePressureArg,
@@ -217,48 +218,37 @@ class Integrator(StrEnum):
 
 
 @dataclass(frozen=True, slots=True)
-class UnitKwarg:
+class IntegratorKwargUnits:
     """Metadata for integrator parameters carrying physical units.
 
-    ``factor`` converts the parameter to internal units (multiply by it).
-    ``channel`` is the function the parameter is used in (either ``"init"``
-    or ``"step"``).
+    ``factors`` maps each registered call channel to the factor that converts
+    that channel's parameter to internal units (multiply by it).
     """
 
-    factor: float
-    channel: Literal["init", "step"]
+    factors: tuple[tuple[Literal["init", "step"], float], ...]
 
 
-_UNIT_FACTORS: Final = frozenset(
-    arg.__metadata__[0]
-    for arg in (TimeArg, InverseTimeArg, PressureArg, InversePressureArg)
-)
-
-
-def _collect_unit_kwargs(integrator: Integrator) -> dict[str, UnitKwarg]:
+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, UnitKwarg] = {}
+    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():
-            for meta in getattr(hint, "__metadata__", ()):
-                if isinstance(meta, float) and meta in _UNIT_FACTORS:
-                    if name in unit_kwargs:
-                        raise TypeError(
-                            f"{integrator}: unit kwarg {name!r} is annotated in both "
-                            f"the init and step functions, making its `integrate` "
-                            f"channel ambiguous"
-                        )
-                    unit_kwargs[name] = UnitKwarg(float(meta), channel)
-    return unit_kwargs
+            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_UNIT_KWARGS: Final[dict[Integrator, dict[str, UnitKwarg]]] = {
+INTEGRATOR_KWARG_UNITS: Final[dict[Integrator, dict[str, IntegratorKwargUnits]]] = {
     integrator: unit_kwargs
     for integrator in Integrator
     if (unit_kwargs := _collect_unit_kwargs(integrator))
diff --git a/torch_sim/integrators/md.py b/torch_sim/integrators/md.py
index b0c832269..b12ca27d1 100644
--- a/torch_sim/integrators/md.py
+++ b/torch_sim/integrators/md.py
@@ -12,19 +12,21 @@
 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, MetalUnits.temperature]
-EnergyArg = Annotated[float | torch.Tensor, MetalUnits.energy]
-TimeArg = Annotated[float | torch.Tensor, MetalUnits.time]
-InverseTimeArg = Annotated[float | torch.Tensor, 1 / MetalUnits.time]
-PressureArg = Annotated[float | torch.Tensor, MetalUnits.pressure]
-InversePressureArg = Annotated[float | torch.Tensor, 1 / MetalUnits.pressure]
+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)
diff --git a/torch_sim/runners.py b/torch_sim/runners.py
index 23243e643..7329db38b 100644
--- a/torch_sim/runners.py
+++ b/torch_sim/runners.py
@@ -18,7 +18,7 @@
 
 import torch_sim as ts
 from torch_sim.autobatching import BinningAutoBatcher, InFlightAutoBatcher
-from torch_sim.integrators import INTEGRATOR_REGISTRY, INTEGRATOR_UNIT_KWARGS, Integrator
+from torch_sim.integrators import INTEGRATOR_KWARG_UNITS, INTEGRATOR_REGISTRY, Integrator
 from torch_sim.integrators.md import MDState
 from torch_sim.models.interface import ModelInterface
 from torch_sim.optimizers import OPTIM_REGISTRY, FireState, Optimizer, OptimState
@@ -176,7 +176,7 @@ def _normalize_temperature_tensor(
     """Turn the temperature into a tensor of shape (n_steps,) or (n_steps, n_systems).
 
     Args:
-        temperature (float | int | list | torch.Tensor): Temperature input
+        temperature (float | list | torch.Tensor): Temperature input
         n_steps (int): Number of integration steps
         initial_state (SimState): Initial simulation state for dtype and device
     Returns:
@@ -303,8 +303,6 @@ def integrate[T: SimState](  # noqa: C901, PLR0915
     )
     dtype, device = initial_state.dtype, initial_state.device
     kTs = _normalize_temperature_tensor(temperature, n_steps, initial_state)
-    kTs = kTs * unit_system.temperature
-    dt = torch.as_tensor(timestep * unit_system.time, dtype=dtype, device=device)
 
     # Handle both string names and direct function tuples
     if isinstance(integrator, Integrator):
@@ -321,17 +319,18 @@ def integrate[T: SimState](  # noqa: C901, PLR0915
             f"(init_func, step_func), got {type(integrator)}"
         )
 
-    # explicitly copy the kwargs since we modify the copied dict below
+    # Convert unit-carrying kwargs to internal units per INTEGRATOR_KWARG_UNIT
     init_kwargs = {} if init_kwargs is None else init_kwargs.copy()
 
-    # Like `timestep` above, unit-carrying kwargs (e.g. `tau`, `gamma`,
-    # `external_pressure`) must be converted to internal units per
-    # INTEGRATOR_UNIT_KWARGS
+    kTs = kTs * unit_system.temperature
+    dt = torch.as_tensor(timestep * unit_system.time, dtype=dtype, device=device)
     channels = {"init": init_kwargs, "step": integrator_kwargs}
-    for key, meta in INTEGRATOR_UNIT_KWARGS.get(integrator, {}).items():
-        kwargs = channels[meta.channel]
-        if kwargs.get(key) is not None:
-            kwargs[key] = kwargs[key] * meta.factor
+    for key, meta in INTEGRATOR_KWARG_UNITS.get(integrator, {}).items():
+        for channel, factor in meta.factors:
+            kwargs = channels[channel]
+            if kwargs.get(key) is not None:
+                kwargs[key] = kwargs[key] * factor
+
     # batch_iterator will be a list if autobatcher is False
     batch_iterator = _configure_batches_iterator(
         initial_state, model, autobatcher=autobatcher
diff --git a/torch_sim/units.py b/torch_sim/units.py
index 146adab22..fa3919e9f 100644
--- a/torch_sim/units.py
+++ b/torch_sim/units.py
@@ -5,11 +5,36 @@
 Units are defined similar to https://docs.lammps.org/units.html.
 """
 
+import typing
+from dataclasses import dataclass
 from enum import Enum
 from math import pi, sqrt
 from typing import Self
 
 
+@dataclass(frozen=True, slots=True)
+class UnitAnnotation:
+    """Typed :class:`typing.Annotated` metadata for a unit conversion factor."""
+
+    factor: float
+
+
+def unit_factor_from_annotation(hint: object) -> float | None:
+    """Return a unit factor from an annotation, including optional annotations."""
+    factors = {
+        float(meta.factor)
+        for meta in getattr(hint, "__metadata__", ())
+        if isinstance(meta, UnitAnnotation)
+    }
+    for argument in typing.get_args(hint):
+        if (factor := unit_factor_from_annotation(argument)) is not None:
+            factors.add(factor)
+
+    if len(factors) > 1:
+        raise TypeError(f"Annotation {hint!r} contains multiple unit factors")
+    return next(iter(factors), None)
+
+
 class BaseConstant(float, Enum):
     """CODATA Recommended Values of the Fundamental Physical Constants: 2014.
 

From c767fbf4c04432a724b6a97d75c0e3d8e1ef6ec9 Mon Sep 17 00:00:00 2001
From: Rhys Goodall <rhys.goodall@outlook.com>
Date: Fri, 24 Jul 2026 20:17:21 +0100
Subject: [PATCH 6/6] refactor: remove deprecated velocity_verlet and
 standardize kT tensor initialization in integrators

---
 torch_sim/integrators/__init__.py |  2 +-
 torch_sim/integrators/md.py       | 12 +-----------
 torch_sim/integrators/nve.py      |  4 +++-
 torch_sim/integrators/nvt.py      |  8 ++++++--
 4 files changed, 11 insertions(+), 15 deletions(-)

diff --git a/torch_sim/integrators/__init__.py b/torch_sim/integrators/__init__.py
index e86b0cd98..8683f032e 100644
--- a/torch_sim/integrators/__init__.py
+++ b/torch_sim/integrators/__init__.py
@@ -92,7 +92,7 @@
     initialize_momenta,
     momentum_step,
     position_step,
-    velocity_verlet,
+    velocity_verlet_step,
 )
 from .npt import (
     NPTLangevinAnisotropicState,
diff --git a/torch_sim/integrators/md.py b/torch_sim/integrators/md.py
index b12ca27d1..ee1755085 100644
--- a/torch_sim/integrators/md.py
+++ b/torch_sim/integrators/md.py
@@ -1,7 +1,6 @@
 """Core molecular dynamics state and operations."""
 
 import logging
-import warnings
 from collections.abc import Callable
 from dataclasses import dataclass
 from typing import Annotated
@@ -115,7 +114,7 @@ def initialize_momenta(
     positions: torch.Tensor,
     masses: torch.Tensor,
     system_idx: torch.Tensor,
-    kT: EnergyArg,
+    kT: torch.Tensor,
     generator: torch.Generator | None = None,
 ) -> torch.Tensor:
     """Initialize particle momenta based on temperature.
@@ -137,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]
 
@@ -247,14 +245,6 @@ def velocity_verlet_step[T: MDState](state: T, dt: TimeArg, model: ModelInterfac
     return momentum_step(state, dt_2)
 
 
-def velocity_verlet[T: MDState](state: T, dt: TimeArg, 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.
diff --git a/torch_sim/integrators/nve.py b/torch_sim/integrators/nve.py
index 40dde09f7..1cd453c76 100644
--- a/torch_sim/integrators/nve.py
+++ b/torch_sim/integrators/nve.py
@@ -49,13 +49,15 @@ def nve_init(
     """
     model_output = model(state)
 
+    kT_tensor = torch.as_tensor(kT, device=state.device, dtype=state.dtype)
+
     momenta = getattr(state, "momenta", None)
     if momenta is None:
         momenta = initialize_momenta(
             state.positions,
             state.masses,
             state.system_idx,
-            kT,
+            kT_tensor,
             state.rng,
         )
 
diff --git a/torch_sim/integrators/nvt.py b/torch_sim/integrators/nvt.py
index eac1f5c66..1667cf3ca 100644
--- a/torch_sim/integrators/nvt.py
+++ b/torch_sim/integrators/nvt.py
@@ -120,13 +120,15 @@ def nvt_langevin_init(
     """
     model_output = model(state)
 
+    kT_tensor = torch.as_tensor(kT, device=state.device, dtype=state.dtype)
+
     momenta = getattr(state, "momenta", None)
     if momenta is None:
         momenta = initialize_momenta(
             state.positions,
             state.masses,
             state.system_idx,
-            kT,
+            kT_tensor,
             state.rng,
         )
     md_state = MDState.from_state(
@@ -687,13 +689,15 @@ def nvt_vrescale_init(
     """
     model_output = model(state)
 
+    kT_tensor = torch.as_tensor(kT, device=state.device, dtype=state.dtype)
+
     momenta = getattr(state, "momenta", None)
     if momenta is None:
         momenta = initialize_momenta(
             state.positions,
             state.masses,
             state.system_idx,
-            kT,
+            kT_tensor,
             state.rng,
         )