From 3a5d4e6e380bf4d6f0a931b88be6c0ed4d8c1594 Mon Sep 17 00:00:00 2001 From: Nishant Sharma Date: Fri, 23 Oct 2026 23:31:52 +0530 Subject: [PATCH] fix: Stop MatterSimCalculator instances sharing args_dict MatterSimCalculator.__init__ and the deprecated load() used a mutable `{}` default for args_dict and stored it directly on the instance. calculate() writes batch_size and only_inference into self.args_dict, so running one calculator changed the args_dict of every other calculator created with the default, and also mutated any dict passed by the caller. Default to None and keep a private copy of the given mapping. Add regression tests for both cases. Fixes #177 Co-Authored-By: Claude Opus 5.5 --- src/mattersim/forcefield/potential.py | 8 +++-- tests/forcefield/test_calculator_args_dict.py | 33 +++++++++++++++++++ 2 files changed, 38 insertions(+), 3 deletions(-) create mode 100644 tests/forcefield/test_calculator_args_dict.py diff --git a/src/mattersim/forcefield/potential.py b/src/mattersim/forcefield/potential.py index d67f61f..1df6e6f 100644 --- a/src/mattersim/forcefield/potential.py +++ b/src/mattersim/forcefield/potential.py @@ -1169,7 +1169,7 @@ class MatterSimCalculator(Calculator): def __init__( self, potential: Potential = None, - args_dict: dict = {}, + args_dict: Optional[dict] = None, compute_stress: bool = True, stress_weight: float = GPa, device: str = "cuda" if torch.cuda.is_available() else "cpu", @@ -1223,7 +1223,9 @@ def __init__( self.potential = potential self.compute_stress = compute_stress self.stress_weight = stress_weight - self.args_dict = args_dict + # calculate() writes batch settings into args_dict, so keep a private + # copy: a shared default or a caller-owned dict must never be mutated. + self.args_dict = dict(args_dict) if args_dict is not None else {} self.device = device self.batch_converter = batch_converter self._use_direct_graph = direct_graph or compile @@ -1289,7 +1291,7 @@ def load( device: str = "cuda" if torch.cuda.is_available() else "cpu", args: Dict = None, load_training_state: bool = True, - args_dict: dict = {}, + args_dict: Optional[dict] = None, compute_stress: bool = True, stress_weight: float = GPa, **kwargs, diff --git a/tests/forcefield/test_calculator_args_dict.py b/tests/forcefield/test_calculator_args_dict.py new file mode 100644 index 0000000..874c987 --- /dev/null +++ b/tests/forcefield/test_calculator_args_dict.py @@ -0,0 +1,33 @@ +"""Tests that MatterSimCalculator instances do not share args_dict state.""" + +from mattersim.forcefield import MatterSimCalculator + + +class TestCalculatorArgsDict: + """calculate() writes batch settings into args_dict, so every calculator + must own its dictionary (see microsoft/mattersim#177).""" + + def test_default_args_dict_is_not_shared(self, mattersim_potential_cpu, si_diamond): + """Using one calculator must not change another's args_dict.""" + calc1 = MatterSimCalculator(potential=mattersim_potential_cpu, device="cpu") + calc2 = MatterSimCalculator(potential=mattersim_potential_cpu, device="cpu") + assert calc1.args_dict is not calc2.args_dict + + si_diamond.calc = calc2 + si_diamond.get_potential_energy() + + assert calc1.args_dict == {} + assert calc2.args_dict == {"batch_size": 1, "only_inference": 1} + + def test_caller_args_dict_is_not_mutated(self, mattersim_potential_cpu, si_diamond): + """A dictionary passed by the caller keeps its original contents.""" + args = {"batch_size": 4} + calc = MatterSimCalculator( + potential=mattersim_potential_cpu, args_dict=args, device="cpu" + ) + + si_diamond.calc = calc + si_diamond.get_potential_energy() + + assert args == {"batch_size": 4} + assert calc.args_dict == {"batch_size": 1, "only_inference": 1}