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}