From b38e2aacc503c061ceb89061009e031797ba776c Mon Sep 17 00:00:00 2001
From: PythonFZ <fabian.zills@web.de>
Date: Thu, 4 Jan 2024 17:55:05 +0100
Subject: [PATCH 01/15] mix calc implementation

---
 ipsuite/calculators/__init__.pyi          |   2 +
 ipsuite/calculators/mix.py                | 132 ++++++++++++++++++++++
 ipsuite/models/ensemble.py                |   4 +-
 ipsuite/nodes.py                          |   1 +
 tests/integration/calculators/test_mix.py |  37 ++++++
 5 files changed, 173 insertions(+), 3 deletions(-)
 create mode 100644 ipsuite/calculators/mix.py
 create mode 100644 tests/integration/calculators/test_mix.py

diff --git a/ipsuite/calculators/__init__.pyi b/ipsuite/calculators/__init__.pyi
index 261d81d6..f4016e09 100644
--- a/ipsuite/calculators/__init__.pyi
+++ b/ipsuite/calculators/__init__.pyi
@@ -18,6 +18,7 @@ from .lammps import LammpsSimulator
 from .orca import OrcaSinglePoint
 from .torch_d3 import TorchD3
 from .xtb import xTBSinglePoint
+from .mix import MixCalculator
 
 __all__ = [
     "CP2KSinglePoint",
@@ -40,4 +41,5 @@ __all__ = [
     "LammpsSimulator",
     "TorchD3",
     "FixedLayerConstraint",
+    "MixCalculator",
 ]
diff --git a/ipsuite/calculators/mix.py b/ipsuite/calculators/mix.py
new file mode 100644
index 00000000..14dfc049
--- /dev/null
+++ b/ipsuite/calculators/mix.py
@@ -0,0 +1,132 @@
+from ase.calculators.calculator import Calculator, all_changes, PropertyNotImplementedError
+from ipsuite import base
+import typing
+from ipsuite.models.base import MLModel
+import zntrack
+import numpy as np
+from ipsuite import base, fields
+from ipsuite.utils.ase_sim import freeze_copy_atoms
+import tqdm
+import ase
+import contextlib
+
+
+class _MixCalculator(Calculator):
+    def __init__(self, calculators: typing.List[Calculator], methods: list, **kwargs):
+        Calculator.__init__(self, **kwargs)
+        self.calculators = calculators
+        self.implemented_properties = self.calculators[0].implemented_properties
+        self.methods = methods
+
+    def calculate(
+        self,
+        atoms=None,
+        properties=None,
+        system_changes=all_changes,
+    ):
+        if properties is None:
+            properties = self.implemented_properties
+
+        Calculator.calculate(self, atoms, properties, system_changes)
+
+        mean_results = []
+        sum_results = []
+
+        for i, calc in enumerate(self.calculators):
+            _atoms = atoms.copy()
+            _atoms.calc = calc
+            if self.methods[i] == "mean":
+                mean_results.append(_atoms)
+            elif self.methods[i] == "sum":
+                sum_results.append(_atoms)
+            else:
+                raise NotImplementedError
+
+        for atoms in mean_results:
+            if "energy" in self.results:
+                self.results["energy"] += atoms.get_potential_energy()
+            else:
+                self.results["energy"] = atoms.get_potential_energy()
+            
+            if "forces" in self.results:
+                self.results["forces"] += atoms.get_forces()
+            else:
+                self.results["forces"] = atoms.get_forces()
+            
+            with contextlib.suppress(PropertyNotImplementedError):
+                if "stress" in self.results:
+                    self.results["stress"] += atoms.get_stress()
+                else:
+                    self.results["stress"] = atoms.get_stress()
+        
+        if "energy" in self.results:
+            self.results["energy"] /= len(mean_results)
+        if "forces" in self.results:
+            self.results["forces"] /= len(mean_results)
+        if "stress" in self.results:
+            self.results["stress"] /= len(mean_results)
+
+        for atoms in sum_results:
+            if "energy" in self.results:
+                self.results["energy"] += atoms.get_potential_energy()
+            else:
+                self.results["energy"] = atoms.get_potential_energy()
+        
+            if "forces" in self.results:
+                self.results["forces"] += atoms.get_forces()
+            else:
+                self.results["forces"] = atoms.get_forces()
+            
+            with contextlib.suppress(PropertyNotImplementedError):
+                if "stress" in self.results:
+                    self.results["stress"] += atoms.get_stress()
+                else:
+                    self.results["stress"] = atoms.get_stress()
+                    
+
+class CalculatorNode(typing.Protocol):
+    def get_calculator(self) -> typing.Type[Calculator]:
+        ...
+
+
+class MixCalculator(base.ProcessAtoms):
+    """Combine multiple models or calculators into one.
+
+    Attributes:
+        calculators: list[CalculatorNode]
+            List of calculators to combine.
+        methods: str|list[str]
+            choose from ['mean', 'sum'] either for all calculators
+            as a string or for each calculator individually as a list.
+            All calculators that are assigned with 'mean' will be
+            computed first, then the calculators assigned with 'sum'
+            will be added.
+    """
+
+    calculators: typing.List[CalculatorNode] = zntrack.deps()
+    methods: str | typing.List[str] = zntrack.params("sum")
+
+    def run(self) -> None:
+        calc = self.get_calculator()
+        self.atoms = []
+        for atoms in tqdm.tqdm(self.get_data(), ncols=70):
+            atoms.calc = calc
+            atoms.get_potential_energy()
+            self.atoms.append(freeze_copy_atoms(atoms))
+
+    def get_calculator(self, **kwargs) -> Calculator:
+        """Property to return a model specific ase calculator object.
+
+        Returns
+        -------
+        calc:
+            ase calculator object
+        """
+        if isinstance(self.methods, str):
+            methods = [self.methods] * len(self.calculators)
+        else:
+            methods = self.methods
+        return _MixCalculator(
+            calculators=[x.get_calculator(**kwargs) for x in self.calculators],
+            methods=methods,
+        )
diff --git a/ipsuite/models/ensemble.py b/ipsuite/models/ensemble.py
index 54a21f02..81f08e3a 100644
--- a/ipsuite/models/ensemble.py
+++ b/ipsuite/models/ensemble.py
@@ -54,10 +54,8 @@ def calculate(
 class EnsembleModel(base.IPSNode):
     models: typing.List[MLModel] = zntrack.deps()
 
-    uuid = zntrack.zn.outs()  # to connect this Node to other Nodes it requires an output.
-
     def run(self) -> None:
-        self.uuid = str(uuid4())
+        pass
 
     def get_calculator(self, **kwargs) -> ase.calculators.calculator.Calculator:
         """Property to return a model specific ase calculator object.
diff --git a/ipsuite/nodes.py b/ipsuite/nodes.py
index 79dfcc58..5be27302 100644
--- a/ipsuite/nodes.py
+++ b/ipsuite/nodes.py
@@ -90,6 +90,7 @@ class _Nodes:
     OrcaSinglePoint = "ipsuite.calculators.OrcaSinglePoint"
     ApaxJaxMD = "ipsuite.calculators.ApaxJaxMD"
     LammpsSimulator = "ipsuite.calculators.LammpsSimulator"
+    MixCalculator = "ipsuite.calculators.MixCalculator"
 
     LangevinThermostat = "ipsuite.calculators.LangevinThermostat"
     NPTThermostat = "ipsuite.calculators.NPTThermostat"
diff --git a/tests/integration/calculators/test_mix.py b/tests/integration/calculators/test_mix.py
new file mode 100644
index 00000000..8d530999
--- /dev/null
+++ b/tests/integration/calculators/test_mix.py
@@ -0,0 +1,37 @@
+import ipsuite as ips
+import numpy.testing as npt
+
+def test_mix_calculators(proj_path, traj_file):
+    with ips.Project(automatic_node_names=True) as proj:
+        data = ips.AddData(traj_file)
+        lj1 = ips.calculators.LJSinglePoint(data=data.atoms)
+        lj2 = ips.calculators.LJSinglePoint(data=data.atoms)
+        lj3 = ips.calculators.LJSinglePoint(data=data.atoms)
+
+        mix1 = ips.calculators.MixCalculator(
+            data=data.atoms,
+            calculators=[lj1, lj2],
+            methods="mean",
+        )
+
+        mix2 = ips.calculators.MixCalculator(
+            data=data.atoms,
+            calculators=[lj1, lj2],
+            methods="sum",
+        )
+    
+    proj.run()
+
+    lj1.load()
+    mix1.load()
+
+    for a, b in zip(lj1.atoms, mix1.atoms):
+        assert a.get_potential_energy() == b.get_potential_energy()
+        npt.assert_almost_equal(a.get_forces(), b.get_forces())
+
+    lj2.load()
+    mix2.load()
+
+    for a, b, c in zip(lj1.atoms, lj2.atoms, mix2.atoms):
+        assert a.get_potential_energy() + b.get_potential_energy() == c.get_potential_energy()
+        npt.assert_almost_equal(a.get_forces() + b.get_forces(), c.get_forces())

From dbde7ea601e4e897165d1a5514f776706753bce8 Mon Sep 17 00:00:00 2001
From: PythonFZ <fabian.zills@web.de>
Date: Thu, 4 Jan 2024 17:58:53 +0100
Subject: [PATCH 02/15] additional test

---
 tests/integration/calculators/test_mix.py | 32 +++++++++++++++++++++--
 1 file changed, 30 insertions(+), 2 deletions(-)

diff --git a/tests/integration/calculators/test_mix.py b/tests/integration/calculators/test_mix.py
index 8d530999..780dd940 100644
--- a/tests/integration/calculators/test_mix.py
+++ b/tests/integration/calculators/test_mix.py
@@ -1,6 +1,7 @@
 import ipsuite as ips
 import numpy.testing as npt
 
+
 def test_mix_calculators(proj_path, traj_file):
     with ips.Project(automatic_node_names=True) as proj:
         data = ips.AddData(traj_file)
@@ -19,7 +20,13 @@ def test_mix_calculators(proj_path, traj_file):
             calculators=[lj1, lj2],
             methods="sum",
         )
-    
+
+        mix3 = ips.calculators.MixCalculator(
+            data=data.atoms,
+            calculators=[lj1, lj2, lj3],
+            methods=["mean", "sum", "mean"],
+        )
+
     proj.run()
 
     lj1.load()
@@ -33,5 +40,26 @@ def test_mix_calculators(proj_path, traj_file):
     mix2.load()
 
     for a, b, c in zip(lj1.atoms, lj2.atoms, mix2.atoms):
-        assert a.get_potential_energy() + b.get_potential_energy() == c.get_potential_energy()
+        assert (
+            a.get_potential_energy() + b.get_potential_energy()
+            == c.get_potential_energy()
+        )
         npt.assert_almost_equal(a.get_forces() + b.get_forces(), c.get_forces())
+
+    lj3.load()
+    mix3.load()
+
+    for a, b, c, d in zip(lj1.atoms, lj2.atoms, lj3.atoms, mix3.atoms):
+
+        # (a + c / 2) + b
+        true_energy = (
+            a.get_potential_energy()
+            + b.get_potential_energy()
+        )
+        true_forces = (
+            a.get_forces()
+            + b.get_forces()
+        )
+
+        assert true_energy == d.get_potential_energy()
+        npt.assert_almost_equal(true_forces, d.get_forces())

From 2cd0d378e1977e4bbed8096e84c936bb1a14ec54 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Thu, 4 Jan 2024 17:00:25 +0000
Subject: [PATCH 03/15] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 ipsuite/calculators/__init__.pyi          |  2 +-
 ipsuite/calculators/mix.py                | 34 +++++++++++++----------
 tests/integration/calculators/test_mix.py | 13 +++------
 3 files changed, 24 insertions(+), 25 deletions(-)

diff --git a/ipsuite/calculators/__init__.pyi b/ipsuite/calculators/__init__.pyi
index f4016e09..d9b38b76 100644
--- a/ipsuite/calculators/__init__.pyi
+++ b/ipsuite/calculators/__init__.pyi
@@ -15,10 +15,10 @@ from .ase_md import (
 from .ase_standard import EMTSinglePoint, LJSinglePoint
 from .cp2k import CP2KSinglePoint, CP2KYaml
 from .lammps import LammpsSimulator
+from .mix import MixCalculator
 from .orca import OrcaSinglePoint
 from .torch_d3 import TorchD3
 from .xtb import xTBSinglePoint
-from .mix import MixCalculator
 
 __all__ = [
     "CP2KSinglePoint",
diff --git a/ipsuite/calculators/mix.py b/ipsuite/calculators/mix.py
index 14dfc049..ea2a9dad 100644
--- a/ipsuite/calculators/mix.py
+++ b/ipsuite/calculators/mix.py
@@ -1,14 +1,19 @@
-from ase.calculators.calculator import Calculator, all_changes, PropertyNotImplementedError
-from ipsuite import base
+import contextlib
 import typing
-from ipsuite.models.base import MLModel
-import zntrack
+
+import ase
 import numpy as np
+import tqdm
+import zntrack
+from ase.calculators.calculator import (
+    Calculator,
+    PropertyNotImplementedError,
+    all_changes,
+)
+
 from ipsuite import base, fields
+from ipsuite.models.base import MLModel
 from ipsuite.utils.ase_sim import freeze_copy_atoms
-import tqdm
-import ase
-import contextlib
 
 
 class _MixCalculator(Calculator):
@@ -47,18 +52,18 @@ def calculate(
                 self.results["energy"] += atoms.get_potential_energy()
             else:
                 self.results["energy"] = atoms.get_potential_energy()
-            
+
             if "forces" in self.results:
                 self.results["forces"] += atoms.get_forces()
             else:
                 self.results["forces"] = atoms.get_forces()
-            
+
             with contextlib.suppress(PropertyNotImplementedError):
                 if "stress" in self.results:
                     self.results["stress"] += atoms.get_stress()
                 else:
                     self.results["stress"] = atoms.get_stress()
-        
+
         if "energy" in self.results:
             self.results["energy"] /= len(mean_results)
         if "forces" in self.results:
@@ -71,22 +76,21 @@ def calculate(
                 self.results["energy"] += atoms.get_potential_energy()
             else:
                 self.results["energy"] = atoms.get_potential_energy()
-        
+
             if "forces" in self.results:
                 self.results["forces"] += atoms.get_forces()
             else:
                 self.results["forces"] = atoms.get_forces()
-            
+
             with contextlib.suppress(PropertyNotImplementedError):
                 if "stress" in self.results:
                     self.results["stress"] += atoms.get_stress()
                 else:
                     self.results["stress"] = atoms.get_stress()
-                    
+
 
 class CalculatorNode(typing.Protocol):
-    def get_calculator(self) -> typing.Type[Calculator]:
-        ...
+    def get_calculator(self) -> typing.Type[Calculator]: ...
 
 
 class MixCalculator(base.ProcessAtoms):
diff --git a/tests/integration/calculators/test_mix.py b/tests/integration/calculators/test_mix.py
index 780dd940..07ee9d96 100644
--- a/tests/integration/calculators/test_mix.py
+++ b/tests/integration/calculators/test_mix.py
@@ -1,6 +1,7 @@
-import ipsuite as ips
 import numpy.testing as npt
 
+import ipsuite as ips
+
 
 def test_mix_calculators(proj_path, traj_file):
     with ips.Project(automatic_node_names=True) as proj:
@@ -52,14 +53,8 @@ def test_mix_calculators(proj_path, traj_file):
     for a, b, c, d in zip(lj1.atoms, lj2.atoms, lj3.atoms, mix3.atoms):
 
         # (a + c / 2) + b
-        true_energy = (
-            a.get_potential_energy()
-            + b.get_potential_energy()
-        )
-        true_forces = (
-            a.get_forces()
-            + b.get_forces()
-        )
+        true_energy = a.get_potential_energy() + b.get_potential_energy()
+        true_forces = a.get_forces() + b.get_forces()
 
         assert true_energy == d.get_potential_energy()
         npt.assert_almost_equal(true_forces, d.get_forces())

From f8311dd1e8625c141e4921cfd5cf169c1aade8c2 Mon Sep 17 00:00:00 2001
From: PythonFZ <fabian.zills@web.de>
Date: Thu, 4 Jan 2024 18:04:24 +0100
Subject: [PATCH 04/15] test as external Node

---
 tests/integration/calculators/test_mix.py | 23 +++++++++++++++++++++++
 1 file changed, 23 insertions(+)

diff --git a/tests/integration/calculators/test_mix.py b/tests/integration/calculators/test_mix.py
index 780dd940..a0df3f75 100644
--- a/tests/integration/calculators/test_mix.py
+++ b/tests/integration/calculators/test_mix.py
@@ -63,3 +63,26 @@ def test_mix_calculators(proj_path, traj_file):
 
         assert true_energy == d.get_potential_energy()
         npt.assert_almost_equal(true_forces, d.get_forces())
+
+def test_mix_calculator_external(proj_path, traj_file):
+    lj1 = ips.calculators.LJSinglePoint(data=None)
+    lj2 = ips.calculators.LJSinglePoint(data=None)
+
+    with ips.Project(automatic_node_names=True) as proj:
+        data = ips.AddData(traj_file)
+        lj3 = ips.calculators.LJSinglePoint(data=data.atoms)
+
+        mix1 = ips.calculators.MixCalculator(
+            data=data.atoms,
+            calculators=[lj1, lj2],
+            methods="mean",
+        )
+    
+    proj.run()
+
+    lj3.load()
+    mix1.load()
+
+    for a, b in zip(lj3.atoms, mix1.atoms):
+        assert a.get_potential_energy() == b.get_potential_energy()
+        npt.assert_almost_equal(a.get_forces(), b.get_forces())

From d734ca2133af8aea262518b46f1610a4c42cf9e3 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Thu, 4 Jan 2024 17:05:25 +0000
Subject: [PATCH 05/15] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 tests/integration/calculators/test_mix.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/tests/integration/calculators/test_mix.py b/tests/integration/calculators/test_mix.py
index 231a848f..fba18817 100644
--- a/tests/integration/calculators/test_mix.py
+++ b/tests/integration/calculators/test_mix.py
@@ -59,6 +59,7 @@ def test_mix_calculators(proj_path, traj_file):
         assert true_energy == d.get_potential_energy()
         npt.assert_almost_equal(true_forces, d.get_forces())
 
+
 def test_mix_calculator_external(proj_path, traj_file):
     lj1 = ips.calculators.LJSinglePoint(data=None)
     lj2 = ips.calculators.LJSinglePoint(data=None)
@@ -72,7 +73,7 @@ def test_mix_calculator_external(proj_path, traj_file):
             calculators=[lj1, lj2],
             methods="mean",
         )
-    
+
     proj.run()
 
     lj3.load()

From 5e140cf7d8c72111ba02b1a41fc1c56720db2aed Mon Sep 17 00:00:00 2001
From: PythonFZ <fabian.zills@web.de>
Date: Thu, 4 Jan 2024 18:08:33 +0100
Subject: [PATCH 06/15] linting

---
 ipsuite/calculators/mix.py | 5 +----
 ipsuite/models/ensemble.py | 1 -
 2 files changed, 1 insertion(+), 5 deletions(-)

diff --git a/ipsuite/calculators/mix.py b/ipsuite/calculators/mix.py
index ea2a9dad..3b7be399 100644
--- a/ipsuite/calculators/mix.py
+++ b/ipsuite/calculators/mix.py
@@ -1,8 +1,6 @@
 import contextlib
 import typing
 
-import ase
-import numpy as np
 import tqdm
 import zntrack
 from ase.calculators.calculator import (
@@ -11,8 +9,7 @@
     all_changes,
 )
 
-from ipsuite import base, fields
-from ipsuite.models.base import MLModel
+from ipsuite import base
 from ipsuite.utils.ase_sim import freeze_copy_atoms
 
 
diff --git a/ipsuite/models/ensemble.py b/ipsuite/models/ensemble.py
index 81f08e3a..71028dec 100644
--- a/ipsuite/models/ensemble.py
+++ b/ipsuite/models/ensemble.py
@@ -1,5 +1,4 @@
 import typing
-from uuid import uuid4
 
 import ase
 import numpy as np

From 6c6fc2d8586553cf4bc77731e6c254e4883be86c Mon Sep 17 00:00:00 2001
From: PythonFZ <fabian.zills@web.de>
Date: Sat, 6 Jan 2024 12:42:30 +0100
Subject: [PATCH 07/15] add dvc-s3; use local dir instead of clone

---
 Dockerfile | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/Dockerfile b/Dockerfile
index 46677ae7..78a598d7 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -1,11 +1,14 @@
 FROM pytorch/pytorch:2.1.2-cuda11.8-cudnn8-runtime
 RUN conda install git
-RUN git clone https://github.com/zincware/ipsuite
+
+COPY . /workspace/ipsuite
 
 WORKDIR /workspace/ipsuite
 RUN pip install .[comparison,gap,nequip,apax,allegro,mace]
 RUN pip install --upgrade torch --extra-index-url https://download.pytorch.org/whl/cu116
 RUN pip install --upgrade "jax[cuda11_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
+RUN pip install git+https://github.com/PythonFZ/torch-dftd.git@patch-2
+RUN pip install dvc-s3
 
 COPY entrypoint.sh entrypoint.sh
 RUN chmod +x entrypoint.sh

From 47d37301eb526419dd8c269cd0139d0cb1c56f78 Mon Sep 17 00:00:00 2001
From: PythonFZ <fabian.zills@web.de>
Date: Sat, 6 Jan 2024 12:42:38 +0100
Subject: [PATCH 08/15] simplify method

---
 ipsuite/calculators/mix.py | 53 ++++++++++++--------------------------
 1 file changed, 16 insertions(+), 37 deletions(-)

diff --git a/ipsuite/calculators/mix.py b/ipsuite/calculators/mix.py
index 3b7be399..f3e1b357 100644
--- a/ipsuite/calculators/mix.py
+++ b/ipsuite/calculators/mix.py
@@ -44,46 +44,24 @@ def calculate(
             else:
                 raise NotImplementedError
 
-        for atoms in mean_results:
-            if "energy" in self.results:
-                self.results["energy"] += atoms.get_potential_energy()
-            else:
-                self.results["energy"] = atoms.get_potential_energy()
-
-            if "forces" in self.results:
-                self.results["forces"] += atoms.get_forces()
-            else:
-                self.results["forces"] = atoms.get_forces()
-
-            with contextlib.suppress(PropertyNotImplementedError):
-                if "stress" in self.results:
-                    self.results["stress"] += atoms.get_stress()
-                else:
-                    self.results["stress"] = atoms.get_stress()
-
+        self.results["energy"] = sum(x.get_potential_energy() for x in mean_results) / len(mean_results)
+        self.results["forces"] = sum(x.get_forces() for x in mean_results) / len(mean_results)
+        with contextlib.suppress(PropertyNotImplementedError):
+            self.results["stress"] = sum(x.get_stress() for x in mean_results) / len(mean_results)
+        
         if "energy" in self.results:
-            self.results["energy"] /= len(mean_results)
+            self.results["energy"] += sum(x.get_potential_energy() for x in sum_results)
+        else:
+            self.results["energy"] = sum(x.get_potential_energy() for x in sum_results)
         if "forces" in self.results:
-            self.results["forces"] /= len(mean_results)
-        if "stress" in self.results:
-            self.results["stress"] /= len(mean_results)
-
-        for atoms in sum_results:
-            if "energy" in self.results:
-                self.results["energy"] += atoms.get_potential_energy()
-            else:
-                self.results["energy"] = atoms.get_potential_energy()
-
-            if "forces" in self.results:
-                self.results["forces"] += atoms.get_forces()
+            self.results["forces"] += sum(x.get_forces() for x in sum_results)
+        else:
+            self.results["forces"] = sum(x.get_forces() for x in sum_results)
+        with contextlib.suppress(PropertyNotImplementedError):
+            if "stress" in self.results:
+                self.results["stress"] += sum(x.get_stress() for x in sum_results)
             else:
-                self.results["forces"] = atoms.get_forces()
-
-            with contextlib.suppress(PropertyNotImplementedError):
-                if "stress" in self.results:
-                    self.results["stress"] += atoms.get_stress()
-                else:
-                    self.results["stress"] = atoms.get_stress()
+                self.results["stress"] = sum(x.get_stress() for x in sum_results)
 
 
 class CalculatorNode(typing.Protocol):
@@ -106,6 +84,7 @@ class MixCalculator(base.ProcessAtoms):
 
     calculators: typing.List[CalculatorNode] = zntrack.deps()
     methods: str | typing.List[str] = zntrack.params("sum")
+    # weights: list = zntrack.params(None) ?
 
     def run(self) -> None:
         calc = self.get_calculator()

From 93d4d9577fb10b5030de936bdc6a6f9ab82de27e Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Sat, 6 Jan 2024 11:42:49 +0000
Subject: [PATCH 09/15] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 ipsuite/calculators/mix.py | 14 ++++++++++----
 1 file changed, 10 insertions(+), 4 deletions(-)

diff --git a/ipsuite/calculators/mix.py b/ipsuite/calculators/mix.py
index f3e1b357..58340a72 100644
--- a/ipsuite/calculators/mix.py
+++ b/ipsuite/calculators/mix.py
@@ -44,11 +44,17 @@ def calculate(
             else:
                 raise NotImplementedError
 
-        self.results["energy"] = sum(x.get_potential_energy() for x in mean_results) / len(mean_results)
-        self.results["forces"] = sum(x.get_forces() for x in mean_results) / len(mean_results)
+        self.results["energy"] = sum(
+            x.get_potential_energy() for x in mean_results
+        ) / len(mean_results)
+        self.results["forces"] = sum(x.get_forces() for x in mean_results) / len(
+            mean_results
+        )
         with contextlib.suppress(PropertyNotImplementedError):
-            self.results["stress"] = sum(x.get_stress() for x in mean_results) / len(mean_results)
-        
+            self.results["stress"] = sum(x.get_stress() for x in mean_results) / len(
+                mean_results
+            )
+
         if "energy" in self.results:
             self.results["energy"] += sum(x.get_potential_energy() for x in sum_results)
         else:

From dc01bd6efeb315b333e45bf23f46c2f8bfb39ebd Mon Sep 17 00:00:00 2001
From: PythonFZ <fabian.zills@web.de>
Date: Sat, 6 Jan 2024 13:10:15 +0100
Subject: [PATCH 10/15] bugfix

---
 ipsuite/calculators/mix.py | 45 +++++++++++++++++---------------------
 1 file changed, 20 insertions(+), 25 deletions(-)

diff --git a/ipsuite/calculators/mix.py b/ipsuite/calculators/mix.py
index 58340a72..b657667e 100644
--- a/ipsuite/calculators/mix.py
+++ b/ipsuite/calculators/mix.py
@@ -13,6 +13,18 @@
 from ipsuite.utils.ase_sim import freeze_copy_atoms
 
 
+def _update_is_exists(results, key, atoms_list, func, mean: bool):
+    with contextlib.suppress(PropertyNotImplementedError):
+        value = sum(func(x) for x in atoms_list)
+        if mean and len(atoms_list) > 0:
+            value /= len(atoms_list)
+
+        if key in results:
+            results[key] += value
+        else:
+            results[key] = value
+
+
 class _MixCalculator(Calculator):
     def __init__(self, calculators: typing.List[Calculator], methods: list, **kwargs):
         Calculator.__init__(self, **kwargs)
@@ -43,31 +55,14 @@ def calculate(
                 sum_results.append(_atoms)
             else:
                 raise NotImplementedError
-
-        self.results["energy"] = sum(
-            x.get_potential_energy() for x in mean_results
-        ) / len(mean_results)
-        self.results["forces"] = sum(x.get_forces() for x in mean_results) / len(
-            mean_results
-        )
-        with contextlib.suppress(PropertyNotImplementedError):
-            self.results["stress"] = sum(x.get_stress() for x in mean_results) / len(
-                mean_results
-            )
-
-        if "energy" in self.results:
-            self.results["energy"] += sum(x.get_potential_energy() for x in sum_results)
-        else:
-            self.results["energy"] = sum(x.get_potential_energy() for x in sum_results)
-        if "forces" in self.results:
-            self.results["forces"] += sum(x.get_forces() for x in sum_results)
-        else:
-            self.results["forces"] = sum(x.get_forces() for x in sum_results)
-        with contextlib.suppress(PropertyNotImplementedError):
-            if "stress" in self.results:
-                self.results["stress"] += sum(x.get_stress() for x in sum_results)
-            else:
-                self.results["stress"] = sum(x.get_stress() for x in sum_results)
+        
+        _update_is_exists(self.results, "energy", mean_results, lambda x: x.get_potential_energy(), True)
+        _update_is_exists(self.results, "forces", mean_results, lambda x: x.get_forces(), True)
+        _update_is_exists(self.results, "stress", mean_results, lambda x: x.get_stress(), True)
+
+        _update_is_exists(self.results, "energy", sum_results, lambda x: x.get_potential_energy(), False)
+        _update_is_exists(self.results, "forces", sum_results, lambda x: x.get_forces(), False)
+        _update_is_exists(self.results, "stress", sum_results, lambda x: x.get_stress(), False)
 
 
 class CalculatorNode(typing.Protocol):

From b725d5235514d589ef08447326320ff7987a7b7b Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Sat, 6 Jan 2024 12:10:26 +0000
Subject: [PATCH 11/15] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 ipsuite/calculators/mix.py | 28 ++++++++++++++++++++--------
 1 file changed, 20 insertions(+), 8 deletions(-)

diff --git a/ipsuite/calculators/mix.py b/ipsuite/calculators/mix.py
index b657667e..a59cbb11 100644
--- a/ipsuite/calculators/mix.py
+++ b/ipsuite/calculators/mix.py
@@ -55,14 +55,26 @@ def calculate(
                 sum_results.append(_atoms)
             else:
                 raise NotImplementedError
-        
-        _update_is_exists(self.results, "energy", mean_results, lambda x: x.get_potential_energy(), True)
-        _update_is_exists(self.results, "forces", mean_results, lambda x: x.get_forces(), True)
-        _update_is_exists(self.results, "stress", mean_results, lambda x: x.get_stress(), True)
-
-        _update_is_exists(self.results, "energy", sum_results, lambda x: x.get_potential_energy(), False)
-        _update_is_exists(self.results, "forces", sum_results, lambda x: x.get_forces(), False)
-        _update_is_exists(self.results, "stress", sum_results, lambda x: x.get_stress(), False)
+
+        _update_is_exists(
+            self.results, "energy", mean_results, lambda x: x.get_potential_energy(), True
+        )
+        _update_is_exists(
+            self.results, "forces", mean_results, lambda x: x.get_forces(), True
+        )
+        _update_is_exists(
+            self.results, "stress", mean_results, lambda x: x.get_stress(), True
+        )
+
+        _update_is_exists(
+            self.results, "energy", sum_results, lambda x: x.get_potential_energy(), False
+        )
+        _update_is_exists(
+            self.results, "forces", sum_results, lambda x: x.get_forces(), False
+        )
+        _update_is_exists(
+            self.results, "stress", sum_results, lambda x: x.get_stress(), False
+        )
 
 
 class CalculatorNode(typing.Protocol):

From 4fee2a4584aa504c2737f45e3d45d6dabc332be9 Mon Sep 17 00:00:00 2001
From: Fabian Zills <fzills@icp.uni-stuttgart.de>
Date: Thu, 11 Jan 2024 18:47:41 +0100
Subject: [PATCH 12/15] typo

---
 ipsuite/calculators/mix.py | 14 +++++++-------
 1 file changed, 7 insertions(+), 7 deletions(-)

diff --git a/ipsuite/calculators/mix.py b/ipsuite/calculators/mix.py
index a59cbb11..db1c990b 100644
--- a/ipsuite/calculators/mix.py
+++ b/ipsuite/calculators/mix.py
@@ -13,7 +13,7 @@
 from ipsuite.utils.ase_sim import freeze_copy_atoms
 
 
-def _update_is_exists(results, key, atoms_list, func, mean: bool):
+def _update_if_exists(results, key, atoms_list, func, mean: bool):
     with contextlib.suppress(PropertyNotImplementedError):
         value = sum(func(x) for x in atoms_list)
         if mean and len(atoms_list) > 0:
@@ -56,23 +56,23 @@ def calculate(
             else:
                 raise NotImplementedError
 
-        _update_is_exists(
+        _update_if_exists(
             self.results, "energy", mean_results, lambda x: x.get_potential_energy(), True
         )
-        _update_is_exists(
+        _update_if_exists(
             self.results, "forces", mean_results, lambda x: x.get_forces(), True
         )
-        _update_is_exists(
+        _update_if_exists(
             self.results, "stress", mean_results, lambda x: x.get_stress(), True
         )
 
-        _update_is_exists(
+        _update_if_exists(
             self.results, "energy", sum_results, lambda x: x.get_potential_energy(), False
         )
-        _update_is_exists(
+        _update_if_exists(
             self.results, "forces", sum_results, lambda x: x.get_forces(), False
         )
-        _update_is_exists(
+        _update_if_exists(
             self.results, "stress", sum_results, lambda x: x.get_stress(), False
         )
 

From bed696a9a4597abafc49736de487f7f94f027df7 Mon Sep 17 00:00:00 2001
From: Fabian Zills <fzills@icp.uni-stuttgart.de>
Date: Thu, 11 Jan 2024 18:55:31 +0100
Subject: [PATCH 13/15] replace EAFP with LBYL

---
 ipsuite/calculators/mix.py | 23 ++++++++++++++++-------
 1 file changed, 16 insertions(+), 7 deletions(-)

diff --git a/ipsuite/calculators/mix.py b/ipsuite/calculators/mix.py
index db1c990b..d01acef9 100644
--- a/ipsuite/calculators/mix.py
+++ b/ipsuite/calculators/mix.py
@@ -13,7 +13,9 @@
 from ipsuite.utils.ase_sim import freeze_copy_atoms
 
 
-def _update_if_exists(results, key, atoms_list, func, mean: bool):
+def _update_if_exists(results, key, atoms_list, func, mean: bool, first_run: bool, existing_properties: list):
+    if key in existing_properties and not first_run:
+        return
     with contextlib.suppress(PropertyNotImplementedError):
         value = sum(func(x) for x in atoms_list)
         if mean and len(atoms_list) > 0:
@@ -24,6 +26,9 @@ def _update_if_exists(results, key, atoms_list, func, mean: bool):
         else:
             results[key] = value
 
+        if first_run:
+            existing_properties.append(key)
+
 
 class _MixCalculator(Calculator):
     def __init__(self, calculators: typing.List[Calculator], methods: list, **kwargs):
@@ -31,6 +36,8 @@ def __init__(self, calculators: typing.List[Calculator], methods: list, **kwargs
         self.calculators = calculators
         self.implemented_properties = self.calculators[0].implemented_properties
         self.methods = methods
+        self._existing_properties = {"sum": [], "mean": []}
+        self._first_run = True
 
     def calculate(
         self,
@@ -57,25 +64,27 @@ def calculate(
                 raise NotImplementedError
 
         _update_if_exists(
-            self.results, "energy", mean_results, lambda x: x.get_potential_energy(), True
+            self.results, "energy", mean_results, lambda x: x.get_potential_energy(), True, self._first_run, self._existing_properties["mean"]
         )
         _update_if_exists(
-            self.results, "forces", mean_results, lambda x: x.get_forces(), True
+            self.results, "forces", mean_results, lambda x: x.get_forces(), True, self._first_run, self._existing_properties["mean"]
         )
         _update_if_exists(
-            self.results, "stress", mean_results, lambda x: x.get_stress(), True
+            self.results, "stress", mean_results, lambda x: x.get_stress(), True, self._first_run, self._existing_properties["mean"]
         )
 
         _update_if_exists(
-            self.results, "energy", sum_results, lambda x: x.get_potential_energy(), False
+            self.results, "energy", sum_results, lambda x: x.get_potential_energy(), False, self._first_run, self._existing_properties["sum"]
         )
         _update_if_exists(
-            self.results, "forces", sum_results, lambda x: x.get_forces(), False
+            self.results, "forces", sum_results, lambda x: x.get_forces(), False, self._first_run, self._existing_properties["sum"]
         )
         _update_if_exists(
-            self.results, "stress", sum_results, lambda x: x.get_stress(), False
+            self.results, "stress", sum_results, lambda x: x.get_stress(), False, self._first_run, self._existing_properties["sum"]
         )
 
+        self._first_run = False
+
 
 class CalculatorNode(typing.Protocol):
     def get_calculator(self) -> typing.Type[Calculator]: ...

From 30b046c8289c0c4143b20a36926a760c559683af Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Thu, 11 Jan 2024 17:56:38 +0000
Subject: [PATCH 14/15] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 ipsuite/calculators/mix.py | 52 +++++++++++++++++++++++++++++++++-----
 1 file changed, 45 insertions(+), 7 deletions(-)

diff --git a/ipsuite/calculators/mix.py b/ipsuite/calculators/mix.py
index d01acef9..e31b5b8f 100644
--- a/ipsuite/calculators/mix.py
+++ b/ipsuite/calculators/mix.py
@@ -13,7 +13,9 @@
 from ipsuite.utils.ase_sim import freeze_copy_atoms
 
 
-def _update_if_exists(results, key, atoms_list, func, mean: bool, first_run: bool, existing_properties: list):
+def _update_if_exists(
+    results, key, atoms_list, func, mean: bool, first_run: bool, existing_properties: list
+):
     if key in existing_properties and not first_run:
         return
     with contextlib.suppress(PropertyNotImplementedError):
@@ -64,23 +66,59 @@ def calculate(
                 raise NotImplementedError
 
         _update_if_exists(
-            self.results, "energy", mean_results, lambda x: x.get_potential_energy(), True, self._first_run, self._existing_properties["mean"]
+            self.results,
+            "energy",
+            mean_results,
+            lambda x: x.get_potential_energy(),
+            True,
+            self._first_run,
+            self._existing_properties["mean"],
         )
         _update_if_exists(
-            self.results, "forces", mean_results, lambda x: x.get_forces(), True, self._first_run, self._existing_properties["mean"]
+            self.results,
+            "forces",
+            mean_results,
+            lambda x: x.get_forces(),
+            True,
+            self._first_run,
+            self._existing_properties["mean"],
         )
         _update_if_exists(
-            self.results, "stress", mean_results, lambda x: x.get_stress(), True, self._first_run, self._existing_properties["mean"]
+            self.results,
+            "stress",
+            mean_results,
+            lambda x: x.get_stress(),
+            True,
+            self._first_run,
+            self._existing_properties["mean"],
         )
 
         _update_if_exists(
-            self.results, "energy", sum_results, lambda x: x.get_potential_energy(), False, self._first_run, self._existing_properties["sum"]
+            self.results,
+            "energy",
+            sum_results,
+            lambda x: x.get_potential_energy(),
+            False,
+            self._first_run,
+            self._existing_properties["sum"],
         )
         _update_if_exists(
-            self.results, "forces", sum_results, lambda x: x.get_forces(), False, self._first_run, self._existing_properties["sum"]
+            self.results,
+            "forces",
+            sum_results,
+            lambda x: x.get_forces(),
+            False,
+            self._first_run,
+            self._existing_properties["sum"],
         )
         _update_if_exists(
-            self.results, "stress", sum_results, lambda x: x.get_stress(), False, self._first_run, self._existing_properties["sum"]
+            self.results,
+            "stress",
+            sum_results,
+            lambda x: x.get_stress(),
+            False,
+            self._first_run,
+            self._existing_properties["sum"],
         )
 
         self._first_run = False

From 329b451944d782eaa892b57091b2e9ed60761295 Mon Sep 17 00:00:00 2001
From: Fabian Zills <fzills@icp.uni-stuttgart.de>
Date: Thu, 11 Jan 2024 19:56:40 +0100
Subject: [PATCH 15/15] rollback breaking changes

---
 ipsuite/calculators/mix.py | 61 +++++---------------------------------
 1 file changed, 7 insertions(+), 54 deletions(-)

diff --git a/ipsuite/calculators/mix.py b/ipsuite/calculators/mix.py
index e31b5b8f..db1c990b 100644
--- a/ipsuite/calculators/mix.py
+++ b/ipsuite/calculators/mix.py
@@ -13,11 +13,7 @@
 from ipsuite.utils.ase_sim import freeze_copy_atoms
 
 
-def _update_if_exists(
-    results, key, atoms_list, func, mean: bool, first_run: bool, existing_properties: list
-):
-    if key in existing_properties and not first_run:
-        return
+def _update_if_exists(results, key, atoms_list, func, mean: bool):
     with contextlib.suppress(PropertyNotImplementedError):
         value = sum(func(x) for x in atoms_list)
         if mean and len(atoms_list) > 0:
@@ -28,9 +24,6 @@ def _update_if_exists(
         else:
             results[key] = value
 
-        if first_run:
-            existing_properties.append(key)
-
 
 class _MixCalculator(Calculator):
     def __init__(self, calculators: typing.List[Calculator], methods: list, **kwargs):
@@ -38,8 +31,6 @@ def __init__(self, calculators: typing.List[Calculator], methods: list, **kwargs
         self.calculators = calculators
         self.implemented_properties = self.calculators[0].implemented_properties
         self.methods = methods
-        self._existing_properties = {"sum": [], "mean": []}
-        self._first_run = True
 
     def calculate(
         self,
@@ -66,63 +57,25 @@ def calculate(
                 raise NotImplementedError
 
         _update_if_exists(
-            self.results,
-            "energy",
-            mean_results,
-            lambda x: x.get_potential_energy(),
-            True,
-            self._first_run,
-            self._existing_properties["mean"],
+            self.results, "energy", mean_results, lambda x: x.get_potential_energy(), True
         )
         _update_if_exists(
-            self.results,
-            "forces",
-            mean_results,
-            lambda x: x.get_forces(),
-            True,
-            self._first_run,
-            self._existing_properties["mean"],
+            self.results, "forces", mean_results, lambda x: x.get_forces(), True
         )
         _update_if_exists(
-            self.results,
-            "stress",
-            mean_results,
-            lambda x: x.get_stress(),
-            True,
-            self._first_run,
-            self._existing_properties["mean"],
+            self.results, "stress", mean_results, lambda x: x.get_stress(), True
         )
 
         _update_if_exists(
-            self.results,
-            "energy",
-            sum_results,
-            lambda x: x.get_potential_energy(),
-            False,
-            self._first_run,
-            self._existing_properties["sum"],
+            self.results, "energy", sum_results, lambda x: x.get_potential_energy(), False
         )
         _update_if_exists(
-            self.results,
-            "forces",
-            sum_results,
-            lambda x: x.get_forces(),
-            False,
-            self._first_run,
-            self._existing_properties["sum"],
+            self.results, "forces", sum_results, lambda x: x.get_forces(), False
         )
         _update_if_exists(
-            self.results,
-            "stress",
-            sum_results,
-            lambda x: x.get_stress(),
-            False,
-            self._first_run,
-            self._existing_properties["sum"],
+            self.results, "stress", sum_results, lambda x: x.get_stress(), False
         )
 
-        self._first_run = False
-
 
 class CalculatorNode(typing.Protocol):
     def get_calculator(self) -> typing.Type[Calculator]: ...