fix: Stop MatterSimCalculator instances sharing args_dict

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,

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}