From cf6244c7fe7d00444b9eb86431d70c064e9287bc Mon Sep 17 00:00:00 2001
From: VsevolodX <79542055+VsevolodX@users.noreply.github.com>
Date: Mon, 2 Sep 2024 14:20:20 -0700
Subject: [PATCH 01/17] feat: add ASE calculator with docstring and required
 forces

---
 src/py/mat3ra/made/tools/calculate.py | 64 ++++++++++++++++++++++++++-
 1 file changed, 63 insertions(+), 1 deletion(-)

diff --git a/src/py/mat3ra/made/tools/calculate.py b/src/py/mat3ra/made/tools/calculate.py
index a0e01146..f974ebc4 100644
--- a/src/py/mat3ra/made/tools/calculate.py
+++ b/src/py/mat3ra/made/tools/calculate.py
@@ -7,10 +7,13 @@
 from .analyze import get_surface_area, get_surface_atom_indices
 from .build.interface.utils import get_slab
 from .build.passivation.enums import SurfaceTypes
-from .convert import decorator_convert_material_args_kwargs_to_atoms
+from .convert import decorator_convert_material_args_kwargs_to_atoms, from_ase
 from .third_party import ASEAtoms, ASECalculator, ASECalculatorEMT
 from .utils import calculate_norm_of_distances_between_coordinates, decorator_handle_periodic_boundary_conditions
 
+# TODO: import from thrid_party
+from ase.calculators.calculator import all_changes
+
 
 @decorator_convert_material_args_kwargs_to_atoms
 def calculate_total_energy(atoms: ASEAtoms, calculator: ASECalculator):
@@ -167,3 +170,62 @@ def calculate_norm_of_distances(material: Material, shadowing_radius: float = 2.
     substrate_coordinates_values = np.array(substrate_atoms_surface_coordinates.values)
 
     return calculate_norm_of_distances_between_coordinates(film_coordinates_values, substrate_coordinates_values)
+
+
+class SurfaceDistanceCalculator(ASECalculator):
+    """
+    ASE calculator that computes the norm of distances between interfacial gap facing atoms
+    of the film and the substrate.
+
+    Example usage:
+    ```python
+    from ase.optimize import BFGS
+    atoms = to_ase(material)
+    calc = SurfaceDistanceCalculator(shadowing_radius=2.5)
+
+    atoms.calc = calc
+    opt = BFGS(atoms)
+    opt.run(fmax=0.05)
+    ```
+    Args:
+        shadowing_radius (float): Radius for atoms shadowing underlying from being treated as a surface, in Angstroms.
+        force_constant (float): The force constant for the finite difference approximation of the
+    Note:
+        Built following: https://wiki.fysik.dtu.dk/ase/development/calculators.html
+
+        The calculate method is responsible for computing the energy and forces (if requested).
+        Forces are estimated using a finite difference method, which is a simple approximation
+        and might not be the most accurate or efficient for all cases.
+    """
+
+    implemented_properties = ["energy", "forces"]
+
+    def __init__(self, shadowing_radius: float = 2.5, force_constant: float = 1.0, **kwargs):
+        super().__init__(**kwargs)
+        self.shadowing_radius = shadowing_radius
+        self.force_constant = force_constant
+
+    @decorator_convert_material_args_kwargs_to_atoms
+    def calculate(self, atoms: Optional[ASEAtoms] = None, properties=["energy"], system_changes=all_changes):
+        if atoms is None:
+            atoms = self.atoms.copy()
+
+        ASECalculator.calculate(self, atoms, properties, system_changes)
+        material = Material(from_ase(atoms))
+        norm = calculate_norm_of_distances(material, self.shadowing_radius)
+        self.results = {"energy": norm}
+
+        if "forces" in properties:
+            forces = np.zeros((len(atoms), 3))
+            dx = 0.01
+            for i in range(len(atoms)):
+                for j in range(3):
+                    atoms_plus = atoms.copy()
+                    atoms_plus.positions[i, j] += dx
+                    material_plus = Material(from_ase(atoms_plus))
+                    norm_plus = calculate_norm_of_distances(material_plus, self.shadowing_radius)
+
+                    # Finite difference approximation of the force
+                    forces[i, j] = -self.force_constant * (norm_plus - norm) / dx
+
+            self.results["forces"] = forces

From e695e33c539b52a9ade7c7068d82c106a5434ef5 Mon Sep 17 00:00:00 2001
From: VsevolodX <79542055+VsevolodX@users.noreply.github.com>
Date: Mon, 2 Sep 2024 18:19:14 -0700
Subject: [PATCH 02/17] update: add constraints

---
 src/py/mat3ra/made/tools/calculate.py | 40 ++++++++++++++++++++++++---
 1 file changed, 36 insertions(+), 4 deletions(-)

diff --git a/src/py/mat3ra/made/tools/calculate.py b/src/py/mat3ra/made/tools/calculate.py
index f974ebc4..d07827d2 100644
--- a/src/py/mat3ra/made/tools/calculate.py
+++ b/src/py/mat3ra/made/tools/calculate.py
@@ -1,4 +1,4 @@
-from typing import Optional
+from typing import Optional, List, Union
 
 import numpy as np
 from mat3ra.made.tools.convert.utils import InterfacePartsEnum
@@ -13,6 +13,7 @@
 
 # TODO: import from thrid_party
 from ase.calculators.calculator import all_changes
+from ase.constraints import FixAtoms, FixedPlane
 
 
 @decorator_convert_material_args_kwargs_to_atoms
@@ -200,24 +201,47 @@ class SurfaceDistanceCalculator(ASECalculator):
 
     implemented_properties = ["energy", "forces"]
 
-    def __init__(self, shadowing_radius: float = 2.5, force_constant: float = 1.0, **kwargs):
+    def __init__(
+        self,
+        shadowing_radius: float = 2.5,
+        force_constant: float = 1.0,
+        fix_substrate: bool = True,
+        fix_z: bool = True,
+        symprec: float = 0.01,
+        **kwargs,
+    ):
         super().__init__(**kwargs)
         self.shadowing_radius = shadowing_radius
         self.force_constant = force_constant
+        self.fix_substrate = fix_substrate
+        self.fix_z = fix_z
+        self.symprec = symprec
 
     @decorator_convert_material_args_kwargs_to_atoms
     def calculate(self, atoms: Optional[ASEAtoms] = None, properties=["energy"], system_changes=all_changes):
         if atoms is None:
             atoms = self.atoms.copy()
 
+        # Apply constraints
+        constraints: List[Union[FixAtoms, FixedPlane]] = []
+        if self.fix_substrate:
+            substrate_indices = [i for i, tag in enumerate(atoms.get_tags()) if tag == 0]
+            constraints.append(FixAtoms(indices=substrate_indices))
+        if self.fix_z:
+            all_indices = list(range(len(atoms)))
+            constraints.append(FixedPlane(indices=all_indices, direction=[0, 0, 1]))
+
+        atoms.set_constraint(constraints)
+
         ASECalculator.calculate(self, atoms, properties, system_changes)
         material = Material(from_ase(atoms))
         norm = calculate_norm_of_distances(material, self.shadowing_radius)
+
         self.results = {"energy": norm}
 
         if "forces" in properties:
             forces = np.zeros((len(atoms), 3))
-            dx = 0.01
+            dx = 0.01  # Small displacement for finite difference
             for i in range(len(atoms)):
                 for j in range(3):
                     atoms_plus = atoms.copy()
@@ -225,7 +249,15 @@ def calculate(self, atoms: Optional[ASEAtoms] = None, properties=["energy"], sys
                     material_plus = Material(from_ase(atoms_plus))
                     norm_plus = calculate_norm_of_distances(material_plus, self.shadowing_radius)
 
-                    # Finite difference approximation of the force
                     forces[i, j] = -self.force_constant * (norm_plus - norm) / dx
 
+            for constraint in constraints:
+                constraint.adjust_forces(atoms, forces)
+
             self.results["forces"] = forces
+
+    def get_potential_energy(self, atoms=None, force_consistent=False):
+        return self.get_property("energy", atoms)
+
+    def get_forces(self, atoms=None):
+        return self.get_property("forces", atoms)

From f1f4eaa3639fa2770737b51aed3352ac73fa8ba8 Mon Sep 17 00:00:00 2001
From: VsevolodX <79542055+VsevolodX@users.noreply.github.com>
Date: Mon, 2 Sep 2024 19:19:29 -0700
Subject: [PATCH 03/17] update: optimize code

---
 src/py/mat3ra/made/tools/calculate.py | 57 ++++++++++++++-------------
 1 file changed, 30 insertions(+), 27 deletions(-)

diff --git a/src/py/mat3ra/made/tools/calculate.py b/src/py/mat3ra/made/tools/calculate.py
index d07827d2..c7ce3adb 100644
--- a/src/py/mat3ra/made/tools/calculate.py
+++ b/src/py/mat3ra/made/tools/calculate.py
@@ -1,4 +1,4 @@
-from typing import Optional, List, Union
+from typing import List, Optional, Union
 
 import numpy as np
 from mat3ra.made.tools.convert.utils import InterfacePartsEnum
@@ -8,13 +8,9 @@
 from .build.interface.utils import get_slab
 from .build.passivation.enums import SurfaceTypes
 from .convert import decorator_convert_material_args_kwargs_to_atoms, from_ase
-from .third_party import ASEAtoms, ASECalculator, ASECalculatorEMT
+from .third_party import ASEAtoms, ASECalculator, ASECalculatorEMT, ASEFixAtoms, ASEFixedPlane, ase_all_changes
 from .utils import calculate_norm_of_distances_between_coordinates, decorator_handle_periodic_boundary_conditions
 
-# TODO: import from thrid_party
-from ase.calculators.calculator import all_changes
-from ase.constraints import FixAtoms, FixedPlane
-
 
 @decorator_convert_material_args_kwargs_to_atoms
 def calculate_total_energy(atoms: ASEAtoms, calculator: ASECalculator):
@@ -217,21 +213,39 @@ def __init__(
         self.fix_z = fix_z
         self.symprec = symprec
 
-    @decorator_convert_material_args_kwargs_to_atoms
-    def calculate(self, atoms: Optional[ASEAtoms] = None, properties=["energy"], system_changes=all_changes):
-        if atoms is None:
-            atoms = self.atoms.copy()
-
-        # Apply constraints
-        constraints: List[Union[FixAtoms, FixedPlane]] = []
+    def _add_constraints(self, atoms: ASEAtoms) -> ASEAtoms:
+        constraints: List[Union[ASEFixAtoms, ASEFixedPlane]] = []
         if self.fix_substrate:
             substrate_indices = [i for i, tag in enumerate(atoms.get_tags()) if tag == 0]
-            constraints.append(FixAtoms(indices=substrate_indices))
+            constraints.append(ASEFixAtoms(indices=substrate_indices))
         if self.fix_z:
             all_indices = list(range(len(atoms)))
-            constraints.append(FixedPlane(indices=all_indices, direction=[0, 0, 1]))
+            constraints.append(ASEFixedPlane(indices=all_indices, direction=[0, 0, 1]))
 
         atoms.set_constraint(constraints)
+        return atoms
+
+    def _calculate_forces(self, atoms: ASEAtoms, norm: float) -> np.ndarray:
+        forces = np.zeros((len(atoms), 3))
+        dx = 0.01
+        for i in range(len(atoms)):
+            for j in range(3):
+                atoms_plus = atoms.copy()
+                atoms_plus.positions[i, j] += dx
+                material_plus = Material(from_ase(atoms_plus))
+                norm_plus = calculate_norm_of_distances(material_plus, self.shadowing_radius)
+
+                forces[i, j] = -self.force_constant * (norm_plus - norm) / dx
+
+        return forces
+
+    @decorator_convert_material_args_kwargs_to_atoms
+    def calculate(self, atoms: Optional[ASEAtoms] = None, properties=["energy"], system_changes=ase_all_changes):
+        if atoms is None:
+            atoms = self.atoms.copy()
+
+        atoms = self._add_constraints(atoms)
+        constraints = atoms.constraints
 
         ASECalculator.calculate(self, atoms, properties, system_changes)
         material = Material(from_ase(atoms))
@@ -240,20 +254,9 @@ def calculate(self, atoms: Optional[ASEAtoms] = None, properties=["energy"], sys
         self.results = {"energy": norm}
 
         if "forces" in properties:
-            forces = np.zeros((len(atoms), 3))
-            dx = 0.01  # Small displacement for finite difference
-            for i in range(len(atoms)):
-                for j in range(3):
-                    atoms_plus = atoms.copy()
-                    atoms_plus.positions[i, j] += dx
-                    material_plus = Material(from_ase(atoms_plus))
-                    norm_plus = calculate_norm_of_distances(material_plus, self.shadowing_radius)
-
-                    forces[i, j] = -self.force_constant * (norm_plus - norm) / dx
-
+            forces = self._calculate_forces(atoms, norm)
             for constraint in constraints:
                 constraint.adjust_forces(atoms, forces)
-
             self.results["forces"] = forces
 
     def get_potential_energy(self, atoms=None, force_consistent=False):

From db916423572ab68190d90514deddfd94ee9075af Mon Sep 17 00:00:00 2001
From: VsevolodX <79542055+VsevolodX@users.noreply.github.com>
Date: Mon, 2 Sep 2024 19:19:51 -0700
Subject: [PATCH 04/17] chore: import from 3rd party

---
 src/py/mat3ra/made/tools/third_party.py | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/src/py/mat3ra/made/tools/third_party.py b/src/py/mat3ra/made/tools/third_party.py
index 9013c7de..9858d917 100644
--- a/src/py/mat3ra/made/tools/third_party.py
+++ b/src/py/mat3ra/made/tools/third_party.py
@@ -2,7 +2,10 @@
 from ase.build import add_vacuum as ase_add_vacuum
 from ase.build.supercells import make_supercell as ase_make_supercell
 from ase.calculators.calculator import Calculator as ASECalculator
+from ase.calculators.calculator import all_changes as ase_all_changes
 from ase.calculators.emt import EMT as ASECalculatorEMT
+from ase.constraints import FixAtoms as ASEFixAtoms
+from ase.constraints import FixedPlane as ASEFixedPlane
 from pymatgen.analysis.defects.core import Interstitial as PymatgenInterstitial
 from pymatgen.analysis.defects.core import Substitution as PymatgenSubstitution
 from pymatgen.analysis.defects.core import Vacancy as PymatgenVacancy
@@ -25,6 +28,9 @@
     "ASEAtoms",
     "ASECalculator",
     "ASECalculatorEMT",
+    "ASEFixAtoms",
+    "ASEFixedPlane",
+    "ase_all_changes",
     "PymatgenLattice",
     "PymatgenStructure",
     "PymatgenIStructure",

From ce42e81dd4d472c9db6511036a7459eee8053121 Mon Sep 17 00:00:00 2001
From: VsevolodX <79542055+VsevolodX@users.noreply.github.com>
Date: Fri, 13 Sep 2024 15:26:45 -0700
Subject: [PATCH 05/17] chore: docstring

---
 src/py/mat3ra/made/tools/calculate.py | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/src/py/mat3ra/made/tools/calculate.py b/src/py/mat3ra/made/tools/calculate.py
index 8bb8dbfd..3f90afb9 100644
--- a/src/py/mat3ra/made/tools/calculate.py
+++ b/src/py/mat3ra/made/tools/calculate.py
@@ -180,6 +180,13 @@ class SurfaceDistanceCalculator(ASECalculator):
     ASE calculator that computes the norm of distances between interfacial gap facing atoms
     of the film and the substrate.
 
+    Args:
+        shadowing_radius (float): The radius for atom to shadow underlying from being considered surface, in Angstroms.
+        force_constant (float): The force constant for the finite difference approximation of the forces.
+        fix_substrate (bool): Whether to fix the substrate atoms.
+        fix_z (bool): Whether to fix atoms movement in the z direction.
+        symprec (float): The symmetry precision for the ASE calculator.
+
     Example usage:
     ```python
     from ase.optimize import BFGS

From c4daf7411a4201ca9cde65d1a317738a7a9cd3f2 Mon Sep 17 00:00:00 2001
From: VsevolodX <79542055+VsevolodX@users.noreply.github.com>
Date: Fri, 13 Sep 2024 17:33:19 -0700
Subject: [PATCH 06/17] chore: rename norm -> energy

---
 src/py/mat3ra/made/tools/calculate.py | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/src/py/mat3ra/made/tools/calculate.py b/src/py/mat3ra/made/tools/calculate.py
index 3f90afb9..46e055d1 100644
--- a/src/py/mat3ra/made/tools/calculate.py
+++ b/src/py/mat3ra/made/tools/calculate.py
@@ -238,7 +238,7 @@ def _add_constraints(self, atoms: ASEAtoms) -> ASEAtoms:
         atoms.set_constraint(constraints)
         return atoms
 
-    def _calculate_forces(self, atoms: ASEAtoms, norm: float) -> np.ndarray:
+    def _calculate_forces(self, atoms: ASEAtoms, energy: float) -> np.ndarray:
         forces = np.zeros((len(atoms), 3))
         dx = 0.01
         for i in range(len(atoms)):
@@ -246,9 +246,9 @@ def _calculate_forces(self, atoms: ASEAtoms, norm: float) -> np.ndarray:
                 atoms_plus = atoms.copy()
                 atoms_plus.positions[i, j] += dx
                 material_plus = Material(from_ase(atoms_plus))
-                norm_plus = calculate_film_substrate_interaction_metric(material_plus, self.shadowing_radius)
+                energy_plus = calculate_film_substrate_interaction_metric(material_plus, self.shadowing_radius)
 
-                forces[i, j] = -self.force_constant * (norm_plus - norm) / dx
+                forces[i, j] = -self.force_constant * (energy_plus - energy) / dx
 
         return forces
 
@@ -262,12 +262,12 @@ def calculate(self, atoms: Optional[ASEAtoms] = None, properties=["energy"], sys
 
         ASECalculator.calculate(self, atoms, properties, system_changes)
         material = Material(from_ase(atoms))
-        norm = calculate_film_substrate_interaction_metric(material, self.shadowing_radius)
+        energy = calculate_film_substrate_interaction_metric(material, self.shadowing_radius)
 
-        self.results = {"energy": norm}
+        self.results = {"energy": energy}
 
         if "forces" in properties:
-            forces = self._calculate_forces(atoms, norm)
+            forces = self._calculate_forces(atoms, energy)
             for constraint in constraints:
                 constraint.adjust_forces(atoms, forces)
             self.results["forces"] = forces

From 0daf4ceb5848f0f018b69375c4f602637ad04f7b Mon Sep 17 00:00:00 2001
From: VsevolodX <79542055+VsevolodX@users.noreply.github.com>
Date: Mon, 16 Sep 2024 15:20:41 -0700
Subject: [PATCH 07/17] update: use m calculator

---
 .../mat3ra/made/tools/calculate/__init__.py   | 109 +-----------------
 .../made/tools/calculate/calculators.py       | 109 +++++++++++++++++-
 2 files changed, 109 insertions(+), 109 deletions(-)

diff --git a/src/py/mat3ra/made/tools/calculate/__init__.py b/src/py/mat3ra/made/tools/calculate/__init__.py
index e04e3952..88b728f5 100644
--- a/src/py/mat3ra/made/tools/calculate/__init__.py
+++ b/src/py/mat3ra/made/tools/calculate/__init__.py
@@ -1,13 +1,10 @@
 from typing import List, Optional, Union
 
-import numpy as np
-
 from ...material import Material
 from ..analyze import get_surface_area
 from ..build.interface.utils import get_slab
 from ..convert import decorator_convert_material_args_kwargs_to_atoms, from_ase
-from ..enums import SurfaceTypes
-from ..third_party import ASEAtoms, ASECalculator, ASECalculatorEMT, ASEFixAtoms, ASEFixedPlane, ase_all_changes
+from ..third_party import ASEAtoms, ASECalculator, ASECalculatorEMT
 from .interaction_functions import sum_of_inverse_distances_squared
 
 
@@ -129,107 +126,3 @@ def calculate_interfacial_energy(
     surface_energy_film = calculate_surface_energy(film_slab, film_bulk, calculator)
     adhesion_energy = calculate_adhesion_energy(interface, substrate_slab, film_slab, calculator)
     return surface_energy_film + surface_energy_substrate - adhesion_energy
-
-
-class SurfaceDistanceCalculator(ASECalculator):
-    """
-    ASE calculator that computes the norm of distances between interfacial gap facing atoms
-    of the film and the substrate.
-
-    Args:
-        shadowing_radius (float): The radius for atom to shadow underlying from being considered surface, in Angstroms.
-        force_constant (float): The force constant for the finite difference approximation of the forces.
-        fix_substrate (bool): Whether to fix the substrate atoms.
-        fix_z (bool): Whether to fix atoms movement in the z direction.
-        symprec (float): The symmetry precision for the ASE calculator.
-
-    Example usage:
-    ```python
-    from ase.optimize import BFGS
-    atoms = to_ase(material)
-    calc = SurfaceDistanceCalculator(shadowing_radius=2.5)
-
-    atoms.calc = calc
-    opt = BFGS(atoms)
-    opt.run(fmax=0.05)
-    ```
-    Args:
-        shadowing_radius (float): Radius for atoms shadowing underlying from being treated as a surface, in Angstroms.
-        force_constant (float): The force constant for the finite difference approximation of the
-    Note:
-        Built following: https://wiki.fysik.dtu.dk/ase/development/calculators.html
-
-        The calculate method is responsible for computing the energy and forces (if requested).
-        Forces are estimated using a finite difference method, which is a simple approximation
-        and might not be the most accurate or efficient for all cases.
-    """
-
-    implemented_properties = ["energy", "forces"]
-
-    def __init__(
-        self,
-        shadowing_radius: float = 2.5,
-        force_constant: float = 1.0,
-        fix_substrate: bool = True,
-        fix_z: bool = True,
-        symprec: float = 0.01,
-        **kwargs,
-    ):
-        super().__init__(**kwargs)
-        self.shadowing_radius = shadowing_radius
-        self.force_constant = force_constant
-        self.fix_substrate = fix_substrate
-        self.fix_z = fix_z
-        self.symprec = symprec
-
-    def _add_constraints(self, atoms: ASEAtoms) -> ASEAtoms:
-        constraints: List[Union[ASEFixAtoms, ASEFixedPlane]] = []
-        if self.fix_substrate:
-            substrate_indices = [i for i, tag in enumerate(atoms.get_tags()) if tag == 0]
-            constraints.append(ASEFixAtoms(indices=substrate_indices))
-        if self.fix_z:
-            all_indices = list(range(len(atoms)))
-            constraints.append(ASEFixedPlane(indices=all_indices, direction=[0, 0, 1]))
-
-        atoms.set_constraint(constraints)
-        return atoms
-
-    def _calculate_forces(self, atoms: ASEAtoms, energy: float) -> np.ndarray:
-        forces = np.zeros((len(atoms), 3))
-        dx = 0.01
-        for i in range(len(atoms)):
-            for j in range(3):
-                atoms_plus = atoms.copy()
-                atoms_plus.positions[i, j] += dx
-                material_plus = Material(from_ase(atoms_plus))
-                energy_plus = calculate_film_substrate_interaction_metric(material_plus, self.shadowing_radius)
-
-                forces[i, j] = -self.force_constant * (energy_plus - energy) / dx
-
-        return forces
-
-    @decorator_convert_material_args_kwargs_to_atoms
-    def calculate(self, atoms: Optional[ASEAtoms] = None, properties=["energy"], system_changes=ase_all_changes):
-        if atoms is None:
-            atoms = self.atoms.copy()
-
-        atoms = self._add_constraints(atoms)
-        constraints = atoms.constraints
-
-        ASECalculator.calculate(self, atoms, properties, system_changes)
-        material = Material(from_ase(atoms))
-        energy = calculate_film_substrate_interaction_metric(material, self.shadowing_radius)
-
-        self.results = {"energy": energy}
-
-        if "forces" in properties:
-            forces = self._calculate_forces(atoms, energy)
-            for constraint in constraints:
-                constraint.adjust_forces(atoms, forces)
-            self.results["forces"] = forces
-
-    def get_potential_energy(self, atoms=None, force_consistent=False):
-        return self.get_property("energy", atoms)
-
-    def get_forces(self, atoms=None):
-        return self.get_property("forces", atoms)
diff --git a/src/py/mat3ra/made/tools/calculate/calculators.py b/src/py/mat3ra/made/tools/calculate/calculators.py
index 963096e3..df796eff 100644
--- a/src/py/mat3ra/made/tools/calculate/calculators.py
+++ b/src/py/mat3ra/made/tools/calculate/calculators.py
@@ -1,13 +1,15 @@
-from typing import Callable
+from typing import Callable, List, Optional, Union
 
 import numpy as np
 from mat3ra.made.material import Material
 from pydantic import BaseModel
 
 from ..analyze import get_surface_atom_indices
+from ..convert import from_ase
 from ..convert.utils import InterfacePartsEnum
 from ..enums import SurfaceTypes
 from ..modify import get_interface_part
+from ..third_party import ASEAtoms, ASECalculator, ASEFixAtoms, ASEFixedPlane, ase_all_changes
 from .interaction_functions import sum_of_inverse_distances_squared
 
 
@@ -120,3 +122,108 @@ def get_energy(
         substrate_coordinates_values = np.array(substrate_surface_atom_coordinates.values)
 
         return interaction_function(film_coordinates_values, substrate_coordinates_values)
+
+
+class SurfaceDistanceCalculator(ASECalculator):
+    """
+    ASE calculator that computes the norm of distances between interfacial gap facing atoms
+    of the film and the substrate.
+
+    Args:
+        shadowing_radius (float): The radius for atom to shadow underlying from being considered surface, in Angstroms.
+        force_constant (float): The force constant for the finite difference approximation of the forces.
+        fix_substrate (bool): Whether to fix the substrate atoms.
+        fix_z (bool): Whether to fix atoms movement in the z direction.
+        symprec (float): The symmetry precision for the ASE calculator.
+
+    Example usage:
+    ```python
+    from ase.optimize import BFGS
+    atoms = to_ase(material)
+    calc = SurfaceDistanceCalculator(shadowing_radius=2.5)
+
+    atoms.calc = calc
+    opt = BFGS(atoms)
+    opt.run(fmax=0.05)
+    ```
+    Args:
+        shadowing_radius (float): Radius for atoms shadowing underlying from being treated as a surface, in Angstroms.
+        force_constant (float): The force constant for the finite difference approximation of the
+    Note:
+        Built following: https://wiki.fysik.dtu.dk/ase/development/calculators.html
+
+        The calculate method is responsible for computing the energy and forces (if requested).
+        Forces are estimated using a finite difference method, which is a simple approximation
+        and might not be the most accurate or efficient for all cases.
+    """
+
+    implemented_properties = ["energy", "forces"]
+
+    def __init__(
+        self,
+        shadowing_radius: float = 2.5,
+        force_constant: float = 1.0,
+        fix_substrate: bool = True,
+        fix_z: bool = True,
+        symprec: float = 0.01,
+        calculator_parameters: MaterialCalculatorParameters = MaterialCalculatorParameters(),
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.shadowing_radius = shadowing_radius
+        self.force_constant = force_constant
+        self.fix_substrate = fix_substrate
+        self.fix_z = fix_z
+        self.symprec = symprec
+        self.material_calculator = MaterialCalculator(calculator_parameters=calculator_parameters)
+
+    def _add_constraints(self, atoms: ASEAtoms) -> ASEAtoms:
+        constraints: List[Union[ASEFixAtoms, ASEFixedPlane]] = []
+        if self.fix_substrate:
+            substrate_indices = [i for i, tag in enumerate(atoms.get_tags()) if tag == 0]
+            constraints.append(ASEFixAtoms(indices=substrate_indices))
+        if self.fix_z:
+            all_indices = list(range(len(atoms)))
+            constraints.append(ASEFixedPlane(indices=all_indices, direction=[0, 0, 1]))
+
+        atoms.set_constraint(constraints)
+        return atoms
+
+    def _calculate_forces(self, atoms: ASEAtoms, energy: float) -> np.ndarray:
+        forces = np.zeros((len(atoms), 3))
+        dx = 0.01
+        for i in range(len(atoms)):
+            for j in range(3):
+                atoms_plus = atoms.copy()
+                atoms_plus.positions[i, j] += dx
+                material_plus = Material(from_ase(atoms_plus))
+                energy_plus = self.material_calculator.get_energy(material_plus)
+
+                forces[i, j] = -self.force_constant * (energy_plus - energy) / dx
+
+        return forces
+
+    def calculate(self, atoms: Optional[ASEAtoms] = None, properties=["energy"], system_changes=ase_all_changes):
+        if atoms is None:
+            atoms = self.atoms.copy()
+
+        atoms = self._add_constraints(atoms)
+        constraints = atoms.constraints
+
+        ASECalculator.calculate(self, atoms, properties, system_changes)
+        material = Material(from_ase(atoms))
+        energy = self.material_calculator.get_energy(material)
+
+        self.results = {"energy": energy}
+
+        if "forces" in properties:
+            forces = self._calculate_forces(atoms, energy)
+            for constraint in constraints:
+                constraint.adjust_forces(atoms, forces)
+            self.results["forces"] = forces
+
+    def get_potential_energy(self, atoms=None, force_consistent=False):
+        return self.get_property("energy", atoms)
+
+    def get_forces(self, atoms=None):
+        return self.get_property("forces", atoms)

From 81dd780fc97669e509187e15da7fe342e918c356 Mon Sep 17 00:00:00 2001
From: VsevolodX <79542055+VsevolodX@users.noreply.github.com>
Date: Mon, 16 Sep 2024 16:08:24 -0700
Subject: [PATCH 08/17] update: adjustemtns in calculator

---
 .../made/tools/calculate/calculators.py       | 25 +++++++++++--------
 1 file changed, 14 insertions(+), 11 deletions(-)

diff --git a/src/py/mat3ra/made/tools/calculate/calculators.py b/src/py/mat3ra/made/tools/calculate/calculators.py
index df796eff..0b017445 100644
--- a/src/py/mat3ra/made/tools/calculate/calculators.py
+++ b/src/py/mat3ra/made/tools/calculate/calculators.py
@@ -132,9 +132,12 @@ class SurfaceDistanceCalculator(ASECalculator):
     Args:
         shadowing_radius (float): The radius for atom to shadow underlying from being considered surface, in Angstroms.
         force_constant (float): The force constant for the finite difference approximation of the forces.
-        fix_substrate (bool): Whether to fix the substrate atoms.
-        fix_z (bool): Whether to fix atoms movement in the z direction.
-        symprec (float): The symmetry precision for the ASE calculator.
+        is_substrate_fixed (bool): Whether to fix the substrate atoms.
+        is_z_fixed (bool): Whether to fix atoms movement in the z direction.
+        symprec (float): The symmetry precision for the ASE calculator. This parameter determines the tolerance
+                         for symmetry operations, affecting the identification of equivalent atoms and the overall
+                         symmetry of the system. For more details, refer to the ASE documentation:
+                         https://wiki.fysik.dtu.dk/ase/ase/constraints.html#ase.constraints.FixSymmetry
 
     Example usage:
     ```python
@@ -150,7 +153,7 @@ class SurfaceDistanceCalculator(ASECalculator):
         shadowing_radius (float): Radius for atoms shadowing underlying from being treated as a surface, in Angstroms.
         force_constant (float): The force constant for the finite difference approximation of the
     Note:
-        Built following: https://wiki.fysik.dtu.dk/ase/development/calculators.html
+        Built following https://wiki.fysik.dtu.dk/ase/development/calculators.html
 
         The calculate method is responsible for computing the energy and forces (if requested).
         Forces are estimated using a finite difference method, which is a simple approximation
@@ -163,19 +166,19 @@ def __init__(
         self,
         shadowing_radius: float = 2.5,
         force_constant: float = 1.0,
-        fix_substrate: bool = True,
-        fix_z: bool = True,
+        is_substrate_fixed: bool = True,
+        is_z_fixed: bool = True,
         symprec: float = 0.01,
-        calculator_parameters: MaterialCalculatorParameters = MaterialCalculatorParameters(),
+        material_calculator: Union[MaterialCalculator, InterfaceMaterialCalculator] = InterfaceMaterialCalculator(),
         **kwargs,
     ):
         super().__init__(**kwargs)
         self.shadowing_radius = shadowing_radius
         self.force_constant = force_constant
-        self.fix_substrate = fix_substrate
-        self.fix_z = fix_z
+        self.fix_substrate = is_substrate_fixed
+        self.fix_z = is_z_fixed
         self.symprec = symprec
-        self.material_calculator = MaterialCalculator(calculator_parameters=calculator_parameters)
+        self.material_calculator = material_calculator
 
     def _add_constraints(self, atoms: ASEAtoms) -> ASEAtoms:
         constraints: List[Union[ASEFixAtoms, ASEFixedPlane]] = []
@@ -210,7 +213,7 @@ def calculate(self, atoms: Optional[ASEAtoms] = None, properties=["energy"], sys
         atoms = self._add_constraints(atoms)
         constraints = atoms.constraints
 
-        ASECalculator.calculate(self, atoms, properties, system_changes)
+        super().calculate(self, atoms, properties, system_changes)
         material = Material(from_ase(atoms))
         energy = self.material_calculator.get_energy(material)
 

From 4761da89ce2a5937a77e3ce94ea6688781132fc5 Mon Sep 17 00:00:00 2001
From: VsevolodX <79542055+VsevolodX@users.noreply.github.com>
Date: Mon, 16 Sep 2024 18:46:44 -0700
Subject: [PATCH 09/17] chore: small adjustmnets

---
 .../mat3ra/made/tools/calculate/calculators.py  | 17 ++++++++---------
 src/py/mat3ra/made/tools/third_party.py         |  2 ++
 2 files changed, 10 insertions(+), 9 deletions(-)

diff --git a/src/py/mat3ra/made/tools/calculate/calculators.py b/src/py/mat3ra/made/tools/calculate/calculators.py
index 0b017445..94814597 100644
--- a/src/py/mat3ra/made/tools/calculate/calculators.py
+++ b/src/py/mat3ra/made/tools/calculate/calculators.py
@@ -9,7 +9,7 @@
 from ..convert.utils import InterfacePartsEnum
 from ..enums import SurfaceTypes
 from ..modify import get_interface_part
-from ..third_party import ASEAtoms, ASECalculator, ASEFixAtoms, ASEFixedPlane, ase_all_changes
+from ..third_party import ASEAtoms, ASECalculator, ASEFixAtoms, ASEFixBondLengths, ASEFixedPlane, ase_all_changes
 from .interaction_functions import sum_of_inverse_distances_squared
 
 
@@ -124,16 +124,15 @@ def get_energy(
         return interaction_function(film_coordinates_values, substrate_coordinates_values)
 
 
-class SurfaceDistanceCalculator(ASECalculator):
+class FilmSubstrateDistanceCalculator(ASECalculator):
     """
-    ASE calculator that computes the norm of distances between interfacial gap facing atoms
-    of the film and the substrate.
+    ASE calculator that calculates the interaction energy between a film and substrate in an interface material.
 
     Args:
         shadowing_radius (float): The radius for atom to shadow underlying from being considered surface, in Angstroms.
         force_constant (float): The force constant for the finite difference approximation of the forces.
         is_substrate_fixed (bool): Whether to fix the substrate atoms.
-        is_z_fixed (bool): Whether to fix atoms movement in the z direction.
+        is_z_axis_fixed (bool): Whether to fix atoms movement in the z direction.
         symprec (float): The symmetry precision for the ASE calculator. This parameter determines the tolerance
                          for symmetry operations, affecting the identification of equivalent atoms and the overall
                          symmetry of the system. For more details, refer to the ASE documentation:
@@ -167,7 +166,7 @@ def __init__(
         shadowing_radius: float = 2.5,
         force_constant: float = 1.0,
         is_substrate_fixed: bool = True,
-        is_z_fixed: bool = True,
+        is_z_axis_fixed: bool = True,
         symprec: float = 0.01,
         material_calculator: Union[MaterialCalculator, InterfaceMaterialCalculator] = InterfaceMaterialCalculator(),
         **kwargs,
@@ -176,12 +175,12 @@ def __init__(
         self.shadowing_radius = shadowing_radius
         self.force_constant = force_constant
         self.fix_substrate = is_substrate_fixed
-        self.fix_z = is_z_fixed
+        self.fix_z = is_z_axis_fixed
         self.symprec = symprec
         self.material_calculator = material_calculator
 
     def _add_constraints(self, atoms: ASEAtoms) -> ASEAtoms:
-        constraints: List[Union[ASEFixAtoms, ASEFixedPlane]] = []
+        constraints: List[Union[ASEFixAtoms, ASEFixedPlane, ASEFixBondLengths]] = []
         if self.fix_substrate:
             substrate_indices = [i for i, tag in enumerate(atoms.get_tags()) if tag == 0]
             constraints.append(ASEFixAtoms(indices=substrate_indices))
@@ -213,7 +212,7 @@ def calculate(self, atoms: Optional[ASEAtoms] = None, properties=["energy"], sys
         atoms = self._add_constraints(atoms)
         constraints = atoms.constraints
 
-        super().calculate(self, atoms, properties, system_changes)
+        super().calculate(atoms, properties, system_changes)
         material = Material(from_ase(atoms))
         energy = self.material_calculator.get_energy(material)
 
diff --git a/src/py/mat3ra/made/tools/third_party.py b/src/py/mat3ra/made/tools/third_party.py
index 9858d917..37606b46 100644
--- a/src/py/mat3ra/made/tools/third_party.py
+++ b/src/py/mat3ra/made/tools/third_party.py
@@ -5,6 +5,7 @@
 from ase.calculators.calculator import all_changes as ase_all_changes
 from ase.calculators.emt import EMT as ASECalculatorEMT
 from ase.constraints import FixAtoms as ASEFixAtoms
+from ase.constraints import FixBondLengths as ASEFixBondLengths
 from ase.constraints import FixedPlane as ASEFixedPlane
 from pymatgen.analysis.defects.core import Interstitial as PymatgenInterstitial
 from pymatgen.analysis.defects.core import Substitution as PymatgenSubstitution
@@ -30,6 +31,7 @@
     "ASECalculatorEMT",
     "ASEFixAtoms",
     "ASEFixedPlane",
+    "ASEFixBondLengths",
     "ase_all_changes",
     "PymatgenLattice",
     "PymatgenStructure",

From 3617623b04d722ad34c60f8dd64b043f039efb39 Mon Sep 17 00:00:00 2001
From: VsevolodX <79542055+VsevolodX@users.noreply.github.com>
Date: Mon, 16 Sep 2024 19:40:53 -0700
Subject: [PATCH 10/17] feat: add rigid xy constraints with help from o1

---
 .../made/tools/calculate/calculators.py       | 72 ++++++++++++++++++-
 1 file changed, 70 insertions(+), 2 deletions(-)

diff --git a/src/py/mat3ra/made/tools/calculate/calculators.py b/src/py/mat3ra/made/tools/calculate/calculators.py
index 94814597..9a2ec36c 100644
--- a/src/py/mat3ra/made/tools/calculate/calculators.py
+++ b/src/py/mat3ra/made/tools/calculate/calculators.py
@@ -9,7 +9,7 @@
 from ..convert.utils import InterfacePartsEnum
 from ..enums import SurfaceTypes
 from ..modify import get_interface_part
-from ..third_party import ASEAtoms, ASECalculator, ASEFixAtoms, ASEFixBondLengths, ASEFixedPlane, ase_all_changes
+from ..third_party import ASEAtoms, ASECalculator, ASEFixAtoms, ASEFixedPlane, ase_all_changes
 from .interaction_functions import sum_of_inverse_distances_squared
 
 
@@ -124,6 +124,70 @@ def get_energy(
         return interaction_function(film_coordinates_values, substrate_coordinates_values)
 
 
+class FixFilmRigidXY:
+    """
+    Custom constraint to allow only rigid translation in x and y for film atoms.
+    """
+
+    def __init__(self, film_indices):
+        """
+        Initialize the constraint with the indices of film atoms.
+
+        Args:
+            film_indices (list): List of atom indices that belong to the film.
+        """
+        self.film_indices = film_indices
+        self.initial_positions = None
+        self.reference_center = None
+
+    def adjust_positions(self, atoms, new_positions):
+        """
+        Adjust the positions of film atoms to maintain rigidity in x and y.
+
+        Args:
+            atoms (ase.Atoms): The ASE Atoms object.
+            new_positions (numpy.ndarray): The new positions to adjust in place.
+        """
+        if self.initial_positions is None:
+            self.initial_positions = atoms.positions[self.film_indices].copy()
+            self.reference_center = self.initial_positions.mean(axis=0)
+
+        current_center = new_positions[self.film_indices].mean(axis=0)
+        desired_translation = current_center - self.reference_center
+        desired_translation[2] = 0.0
+
+        # Update positions to maintain rigid body translation in x and y
+        for i, idx in enumerate(self.film_indices):
+            new_positions[idx, 0] = self.initial_positions[i, 0] + desired_translation[0]
+            new_positions[idx, 1] = self.initial_positions[i, 1] + desired_translation[1]
+
+    def adjust_forces(self, forces):
+        """
+        Adjust the forces on film atoms to ensure rigidity.
+
+        Args:
+            forces (numpy.ndarray): The forces array to adjust in place.
+        """
+        if self.initial_positions is None:
+            return  # No adjustment needed if initial positions are not set
+
+        # Project forces onto x and y for collective translation
+        net_force_x = np.sum(forces[self.film_indices, 0])
+        net_force_y = np.sum(forces[self.film_indices, 1])
+
+        # Distribute the net force equally to all film atoms to maintain rigidity
+        num_film_atoms = len(self.film_indices)
+        if num_film_atoms == 0:
+            return  # Avoid division by zero
+
+        distributed_force_x = net_force_x / num_film_atoms
+        distributed_force_y = net_force_y / num_film_atoms
+
+        for idx in self.film_indices:
+            forces[idx, 0] = distributed_force_x
+            forces[idx, 1] = distributed_force_y
+
+
 class FilmSubstrateDistanceCalculator(ASECalculator):
     """
     ASE calculator that calculates the interaction energy between a film and substrate in an interface material.
@@ -180,7 +244,7 @@ def __init__(
         self.material_calculator = material_calculator
 
     def _add_constraints(self, atoms: ASEAtoms) -> ASEAtoms:
-        constraints: List[Union[ASEFixAtoms, ASEFixedPlane, ASEFixBondLengths]] = []
+        constraints: List[Union[ASEFixAtoms, ASEFixedPlane, FixFilmRigidXY]] = []
         if self.fix_substrate:
             substrate_indices = [i for i, tag in enumerate(atoms.get_tags()) if tag == 0]
             constraints.append(ASEFixAtoms(indices=substrate_indices))
@@ -188,6 +252,10 @@ def _add_constraints(self, atoms: ASEAtoms) -> ASEAtoms:
             all_indices = list(range(len(atoms)))
             constraints.append(ASEFixedPlane(indices=all_indices, direction=[0, 0, 1]))
 
+        film_indices = [i for i, tag in enumerate(atoms.get_tags()) if tag == 1]
+        if film_indices:
+            constraints.append(FixFilmRigidXY(film_indices))
+
         atoms.set_constraint(constraints)
         return atoms
 

From 56b5bf9c1ce1ececd04e49db00f6e460f8b09b2d Mon Sep 17 00:00:00 2001
From: VsevolodX <79542055+VsevolodX@users.noreply.github.com>
Date: Mon, 16 Sep 2024 19:42:02 -0700
Subject: [PATCH 11/17] chore: remove unused imports

---
 src/py/mat3ra/made/tools/calculate/calculators.py | 2 ++
 src/py/mat3ra/made/tools/third_party.py           | 2 --
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/py/mat3ra/made/tools/calculate/calculators.py b/src/py/mat3ra/made/tools/calculate/calculators.py
index 9a2ec36c..30772d0e 100644
--- a/src/py/mat3ra/made/tools/calculate/calculators.py
+++ b/src/py/mat3ra/made/tools/calculate/calculators.py
@@ -127,6 +127,8 @@ def get_energy(
 class FixFilmRigidXY:
     """
     Custom constraint to allow only rigid translation in x and y for film atoms.
+
+    Created following https://wiki.fysik.dtu.dk/ase/ase/constraints.html#making-your-own-constraint-class
     """
 
     def __init__(self, film_indices):
diff --git a/src/py/mat3ra/made/tools/third_party.py b/src/py/mat3ra/made/tools/third_party.py
index 37606b46..9858d917 100644
--- a/src/py/mat3ra/made/tools/third_party.py
+++ b/src/py/mat3ra/made/tools/third_party.py
@@ -5,7 +5,6 @@
 from ase.calculators.calculator import all_changes as ase_all_changes
 from ase.calculators.emt import EMT as ASECalculatorEMT
 from ase.constraints import FixAtoms as ASEFixAtoms
-from ase.constraints import FixBondLengths as ASEFixBondLengths
 from ase.constraints import FixedPlane as ASEFixedPlane
 from pymatgen.analysis.defects.core import Interstitial as PymatgenInterstitial
 from pymatgen.analysis.defects.core import Substitution as PymatgenSubstitution
@@ -31,7 +30,6 @@
     "ASECalculatorEMT",
     "ASEFixAtoms",
     "ASEFixedPlane",
-    "ASEFixBondLengths",
     "ase_all_changes",
     "PymatgenLattice",
     "PymatgenStructure",

From 4079d8a5cc70e2e5ffaa06a6819a70fb8ae62673 Mon Sep 17 00:00:00 2001
From: VsevolodX <79542055+VsevolodX@users.noreply.github.com>
Date: Mon, 16 Sep 2024 19:52:50 -0700
Subject: [PATCH 12/17] chore: add types :

---
 .../made/tools/calculate/calculators.py       | 25 ++++++-------------
 1 file changed, 8 insertions(+), 17 deletions(-)

diff --git a/src/py/mat3ra/made/tools/calculate/calculators.py b/src/py/mat3ra/made/tools/calculate/calculators.py
index 30772d0e..3f6771c2 100644
--- a/src/py/mat3ra/made/tools/calculate/calculators.py
+++ b/src/py/mat3ra/made/tools/calculate/calculators.py
@@ -124,25 +124,18 @@ def get_energy(
         return interaction_function(film_coordinates_values, substrate_coordinates_values)
 
 
-class FixFilmRigidXY:
+class FixFilmRigidXY(BaseModel):
     """
     Custom constraint to allow only rigid translation in x and y for film atoms.
 
     Created following https://wiki.fysik.dtu.dk/ase/ase/constraints.html#making-your-own-constraint-class
     """
 
-    def __init__(self, film_indices):
-        """
-        Initialize the constraint with the indices of film atoms.
-
-        Args:
-            film_indices (list): List of atom indices that belong to the film.
-        """
-        self.film_indices = film_indices
-        self.initial_positions = None
-        self.reference_center = None
+    film_indices: List[int] = []
+    initial_positions: Optional[np.ndarray] = None
+    reference_center: Optional[np.ndarray] = None
 
-    def adjust_positions(self, atoms, new_positions):
+    def adjust_positions(self, atoms: ASEAtoms, new_positions: np.ndarray) -> None:
         """
         Adjust the positions of film atoms to maintain rigidity in x and y.
 
@@ -163,7 +156,7 @@ def adjust_positions(self, atoms, new_positions):
             new_positions[idx, 0] = self.initial_positions[i, 0] + desired_translation[0]
             new_positions[idx, 1] = self.initial_positions[i, 1] + desired_translation[1]
 
-    def adjust_forces(self, forces):
+    def adjust_forces(self, forces: np.ndarray) -> None:
         """
         Adjust the forces on film atoms to ensure rigidity.
 
@@ -230,7 +223,6 @@ class FilmSubstrateDistanceCalculator(ASECalculator):
     def __init__(
         self,
         shadowing_radius: float = 2.5,
-        force_constant: float = 1.0,
         is_substrate_fixed: bool = True,
         is_z_axis_fixed: bool = True,
         symprec: float = 0.01,
@@ -239,7 +231,6 @@ def __init__(
     ):
         super().__init__(**kwargs)
         self.shadowing_radius = shadowing_radius
-        self.force_constant = force_constant
         self.fix_substrate = is_substrate_fixed
         self.fix_z = is_z_axis_fixed
         self.symprec = symprec
@@ -256,7 +247,7 @@ def _add_constraints(self, atoms: ASEAtoms) -> ASEAtoms:
 
         film_indices = [i for i, tag in enumerate(atoms.get_tags()) if tag == 1]
         if film_indices:
-            constraints.append(FixFilmRigidXY(film_indices))
+            constraints.append(FixFilmRigidXY(film_indices=film_indices))
 
         atoms.set_constraint(constraints)
         return atoms
@@ -271,7 +262,7 @@ def _calculate_forces(self, atoms: ASEAtoms, energy: float) -> np.ndarray:
                 material_plus = Material(from_ase(atoms_plus))
                 energy_plus = self.material_calculator.get_energy(material_plus)
 
-                forces[i, j] = -self.force_constant * (energy_plus - energy) / dx
+                forces[i, j] = -(energy_plus - energy) / dx
 
         return forces
 

From 736be41eabf0f150ad292c690dca18e94fa2fcc2 Mon Sep 17 00:00:00 2001
From: VsevolodX <79542055+VsevolodX@users.noreply.github.com>
Date: Mon, 16 Sep 2024 19:57:27 -0700
Subject: [PATCH 13/17] chore: arbitrary types allowed

---
 src/py/mat3ra/made/tools/calculate/calculators.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/src/py/mat3ra/made/tools/calculate/calculators.py b/src/py/mat3ra/made/tools/calculate/calculators.py
index 3f6771c2..e39fcee3 100644
--- a/src/py/mat3ra/made/tools/calculate/calculators.py
+++ b/src/py/mat3ra/made/tools/calculate/calculators.py
@@ -131,6 +131,9 @@ class FixFilmRigidXY(BaseModel):
     Created following https://wiki.fysik.dtu.dk/ase/ase/constraints.html#making-your-own-constraint-class
     """
 
+    class Config:
+        arbitrary_types_allowed = True
+
     film_indices: List[int] = []
     initial_positions: Optional[np.ndarray] = None
     reference_center: Optional[np.ndarray] = None

From 09b2acd6dcd0b38cbeb33699653e25d0eeaf264e Mon Sep 17 00:00:00 2001
From: VsevolodX <79542055+VsevolodX@users.noreply.github.com>
Date: Tue, 17 Sep 2024 11:10:50 -0700
Subject: [PATCH 14/17] update: inherit from constraint class

---
 .../made/tools/calculate/calculators.py       | 37 +++++++++++++++----
 1 file changed, 30 insertions(+), 7 deletions(-)

diff --git a/src/py/mat3ra/made/tools/calculate/calculators.py b/src/py/mat3ra/made/tools/calculate/calculators.py
index e39fcee3..07a9796f 100644
--- a/src/py/mat3ra/made/tools/calculate/calculators.py
+++ b/src/py/mat3ra/made/tools/calculate/calculators.py
@@ -124,16 +124,39 @@ def get_energy(
         return interaction_function(film_coordinates_values, substrate_coordinates_values)
 
 
-class FixFilmRigidXY(BaseModel):
+class InterfaceConstraint(BaseModel):
     """
-    Custom constraint to allow only rigid translation in x and y for film atoms.
-
-    Created following https://wiki.fysik.dtu.dk/ase/ase/constraints.html#making-your-own-constraint-class
+    Based on  https://wiki.fysik.dtu.dk/ase/ase/constraints.html#making-your-own-constraint-class
     """
 
     class Config:
         arbitrary_types_allowed = True
 
+    def adjust_positions(self, atoms: ASEAtoms, new_positions: np.ndarray) -> None:
+        """
+        Adjust the positions of the atoms in the system.
+
+        Args:
+            atoms (ASEAtoms): The ASE Atoms object.
+            new_positions (numpy.ndarray): The new positions to adjust in place.
+        """
+        raise NotImplementedError
+
+    def adjust_forces(self, forces: np.ndarray) -> None:
+        """
+        Adjust the forces on the atoms in the system.
+
+        Args:
+            forces (numpy.ndarray): The forces array to adjust in place.
+        """
+        raise NotImplementedError
+
+
+class RigidFilmXYConstraint(InterfaceConstraint):
+    """
+    Custom constraint to allow only rigid translation in x and y for film atoms.
+    """
+
     film_indices: List[int] = []
     initial_positions: Optional[np.ndarray] = None
     reference_center: Optional[np.ndarray] = None
@@ -186,7 +209,7 @@ def adjust_forces(self, forces: np.ndarray) -> None:
             forces[idx, 1] = distributed_force_y
 
 
-class FilmSubstrateDistanceCalculator(ASECalculator):
+class FilmSubstrateDistanceASECalculator(ASECalculator):
     """
     ASE calculator that calculates the interaction energy between a film and substrate in an interface material.
 
@@ -240,7 +263,7 @@ def __init__(
         self.material_calculator = material_calculator
 
     def _add_constraints(self, atoms: ASEAtoms) -> ASEAtoms:
-        constraints: List[Union[ASEFixAtoms, ASEFixedPlane, FixFilmRigidXY]] = []
+        constraints: List[Union[ASEFixAtoms, ASEFixedPlane, RigidFilmXYConstraint]] = []
         if self.fix_substrate:
             substrate_indices = [i for i, tag in enumerate(atoms.get_tags()) if tag == 0]
             constraints.append(ASEFixAtoms(indices=substrate_indices))
@@ -250,7 +273,7 @@ def _add_constraints(self, atoms: ASEAtoms) -> ASEAtoms:
 
         film_indices = [i for i, tag in enumerate(atoms.get_tags()) if tag == 1]
         if film_indices:
-            constraints.append(FixFilmRigidXY(film_indices=film_indices))
+            constraints.append(RigidFilmXYConstraint(film_indices=film_indices))
 
         atoms.set_constraint(constraints)
         return atoms

From aaf49462133d06b2be67bf512b1f8a23e80dab07 Mon Sep 17 00:00:00 2001
From: VsevolodX <79542055+VsevolodX@users.noreply.github.com>
Date: Tue, 17 Sep 2024 11:14:15 -0700
Subject: [PATCH 15/17] chore: move constraints to a separate file

---
 .../made/tools/calculate/calculators.py       | 86 +-----------------
 .../made/tools/calculate/constraints.py       | 91 +++++++++++++++++++
 2 files changed, 92 insertions(+), 85 deletions(-)
 create mode 100644 src/py/mat3ra/made/tools/calculate/constraints.py

diff --git a/src/py/mat3ra/made/tools/calculate/calculators.py b/src/py/mat3ra/made/tools/calculate/calculators.py
index 07a9796f..6c963b7d 100644
--- a/src/py/mat3ra/made/tools/calculate/calculators.py
+++ b/src/py/mat3ra/made/tools/calculate/calculators.py
@@ -10,6 +10,7 @@
 from ..enums import SurfaceTypes
 from ..modify import get_interface_part
 from ..third_party import ASEAtoms, ASECalculator, ASEFixAtoms, ASEFixedPlane, ase_all_changes
+from .constraints import RigidFilmXYConstraint
 from .interaction_functions import sum_of_inverse_distances_squared
 
 
@@ -124,91 +125,6 @@ def get_energy(
         return interaction_function(film_coordinates_values, substrate_coordinates_values)
 
 
-class InterfaceConstraint(BaseModel):
-    """
-    Based on  https://wiki.fysik.dtu.dk/ase/ase/constraints.html#making-your-own-constraint-class
-    """
-
-    class Config:
-        arbitrary_types_allowed = True
-
-    def adjust_positions(self, atoms: ASEAtoms, new_positions: np.ndarray) -> None:
-        """
-        Adjust the positions of the atoms in the system.
-
-        Args:
-            atoms (ASEAtoms): The ASE Atoms object.
-            new_positions (numpy.ndarray): The new positions to adjust in place.
-        """
-        raise NotImplementedError
-
-    def adjust_forces(self, forces: np.ndarray) -> None:
-        """
-        Adjust the forces on the atoms in the system.
-
-        Args:
-            forces (numpy.ndarray): The forces array to adjust in place.
-        """
-        raise NotImplementedError
-
-
-class RigidFilmXYConstraint(InterfaceConstraint):
-    """
-    Custom constraint to allow only rigid translation in x and y for film atoms.
-    """
-
-    film_indices: List[int] = []
-    initial_positions: Optional[np.ndarray] = None
-    reference_center: Optional[np.ndarray] = None
-
-    def adjust_positions(self, atoms: ASEAtoms, new_positions: np.ndarray) -> None:
-        """
-        Adjust the positions of film atoms to maintain rigidity in x and y.
-
-        Args:
-            atoms (ase.Atoms): The ASE Atoms object.
-            new_positions (numpy.ndarray): The new positions to adjust in place.
-        """
-        if self.initial_positions is None:
-            self.initial_positions = atoms.positions[self.film_indices].copy()
-            self.reference_center = self.initial_positions.mean(axis=0)
-
-        current_center = new_positions[self.film_indices].mean(axis=0)
-        desired_translation = current_center - self.reference_center
-        desired_translation[2] = 0.0
-
-        # Update positions to maintain rigid body translation in x and y
-        for i, idx in enumerate(self.film_indices):
-            new_positions[idx, 0] = self.initial_positions[i, 0] + desired_translation[0]
-            new_positions[idx, 1] = self.initial_positions[i, 1] + desired_translation[1]
-
-    def adjust_forces(self, forces: np.ndarray) -> None:
-        """
-        Adjust the forces on film atoms to ensure rigidity.
-
-        Args:
-            forces (numpy.ndarray): The forces array to adjust in place.
-        """
-        if self.initial_positions is None:
-            return  # No adjustment needed if initial positions are not set
-
-        # Project forces onto x and y for collective translation
-        net_force_x = np.sum(forces[self.film_indices, 0])
-        net_force_y = np.sum(forces[self.film_indices, 1])
-
-        # Distribute the net force equally to all film atoms to maintain rigidity
-        num_film_atoms = len(self.film_indices)
-        if num_film_atoms == 0:
-            return  # Avoid division by zero
-
-        distributed_force_x = net_force_x / num_film_atoms
-        distributed_force_y = net_force_y / num_film_atoms
-
-        for idx in self.film_indices:
-            forces[idx, 0] = distributed_force_x
-            forces[idx, 1] = distributed_force_y
-
-
 class FilmSubstrateDistanceASECalculator(ASECalculator):
     """
     ASE calculator that calculates the interaction energy between a film and substrate in an interface material.
diff --git a/src/py/mat3ra/made/tools/calculate/constraints.py b/src/py/mat3ra/made/tools/calculate/constraints.py
new file mode 100644
index 00000000..4581528f
--- /dev/null
+++ b/src/py/mat3ra/made/tools/calculate/constraints.py
@@ -0,0 +1,91 @@
+from typing import List, Optional
+
+import numpy as np
+from pydantic import BaseModel
+
+from ..third_party import ASEAtoms
+
+
+class InterfaceConstraint(BaseModel):
+    """
+    Based on  https://wiki.fysik.dtu.dk/ase/ase/constraints.html#making-your-own-constraint-class
+    """
+
+    class Config:
+        arbitrary_types_allowed = True
+
+    def adjust_positions(self, atoms: ASEAtoms, new_positions: np.ndarray) -> None:
+        """
+        Adjust the positions of the atoms in the system.
+
+        Args:
+            atoms (ASEAtoms): The ASE Atoms object.
+            new_positions (numpy.ndarray): The new positions to adjust in place.
+        """
+        raise NotImplementedError
+
+    def adjust_forces(self, forces: np.ndarray) -> None:
+        """
+        Adjust the forces on the atoms in the system.
+
+        Args:
+            forces (numpy.ndarray): The forces array to adjust in place.
+        """
+        raise NotImplementedError
+
+
+class RigidFilmXYConstraint(InterfaceConstraint):
+    """
+    Custom constraint to allow only rigid translation in x and y for film atoms.
+    """
+
+    film_indices: List[int] = []
+    initial_positions: Optional[np.ndarray] = None
+    reference_center: Optional[np.ndarray] = None
+
+    def adjust_positions(self, atoms: ASEAtoms, new_positions: np.ndarray) -> None:
+        """
+        Adjust the positions of film atoms to maintain rigidity in x and y.
+
+        Args:
+            atoms (ase.Atoms): The ASE Atoms object.
+            new_positions (numpy.ndarray): The new positions to adjust in place.
+        """
+        if self.initial_positions is None:
+            self.initial_positions = atoms.positions[self.film_indices].copy()
+            self.reference_center = self.initial_positions.mean(axis=0)
+
+        current_center = new_positions[self.film_indices].mean(axis=0)
+        desired_translation = current_center - self.reference_center
+        desired_translation[2] = 0.0
+
+        # Update positions to maintain rigid body translation in x and y
+        for i, idx in enumerate(self.film_indices):
+            new_positions[idx, 0] = self.initial_positions[i, 0] + desired_translation[0]
+            new_positions[idx, 1] = self.initial_positions[i, 1] + desired_translation[1]
+
+    def adjust_forces(self, forces: np.ndarray) -> None:
+        """
+        Adjust the forces on film atoms to ensure rigidity.
+
+        Args:
+            forces (numpy.ndarray): The forces array to adjust in place.
+        """
+        if self.initial_positions is None:
+            return  # No adjustment needed if initial positions are not set
+
+        # Project forces onto x and y for collective translation
+        net_force_x = np.sum(forces[self.film_indices, 0])
+        net_force_y = np.sum(forces[self.film_indices, 1])
+
+        # Distribute the net force equally to all film atoms to maintain rigidity
+        num_film_atoms = len(self.film_indices)
+        if num_film_atoms == 0:
+            return  # Avoid division by zero
+
+        distributed_force_x = net_force_x / num_film_atoms
+        distributed_force_y = net_force_y / num_film_atoms
+
+        for idx in self.film_indices:
+            forces[idx, 0] = distributed_force_x
+            forces[idx, 1] = distributed_force_y

From da49cf45b336a740ed1569d90ca63b1160b4d565 Mon Sep 17 00:00:00 2001
From: VsevolodX <79542055+VsevolodX@users.noreply.github.com>
Date: Wed, 18 Sep 2024 18:44:25 -0700
Subject: [PATCH 16/17] update: isolate ase-related calculator

---
 .../made/tools/calculate/ase/__init__.py      | 118 ++++++++++++++++++
 .../tools/calculate/{ => ase}/constraints.py  |   4 +-
 .../made/tools/calculate/calculators.py       | 114 +----------------
 3 files changed, 121 insertions(+), 115 deletions(-)
 create mode 100644 src/py/mat3ra/made/tools/calculate/ase/__init__.py
 rename src/py/mat3ra/made/tools/calculate/{ => ase}/constraints.py (97%)

diff --git a/src/py/mat3ra/made/tools/calculate/ase/__init__.py b/src/py/mat3ra/made/tools/calculate/ase/__init__.py
new file mode 100644
index 00000000..647e9b49
--- /dev/null
+++ b/src/py/mat3ra/made/tools/calculate/ase/__init__.py
@@ -0,0 +1,118 @@
+from typing import List, Optional, Union
+
+import numpy as np
+from mat3ra.made.material import Material
+
+from ...convert import from_ase
+from ...third_party import ASEAtoms, ASECalculator, ASEFixAtoms, ASEFixedPlane, ase_all_changes
+from ..calculators import InterfaceMaterialCalculator, MaterialCalculator
+from .constraints import RigidFilmXYInterfaceConstraint
+
+
+class FilmSubstrateDistanceASECalculator(ASECalculator):
+    """
+    ASE calculator that calculates the interaction energy between a film and substrate in an interface material.
+
+    Args:
+        shadowing_radius (float): The radius for atom to shadow underlying from being considered surface, in Angstroms.
+        force_constant (float): The force constant for the finite difference approximation of the forces.
+        is_substrate_fixed (bool): Whether to fix the substrate atoms.
+        is_z_axis_fixed (bool): Whether to fix atoms movement in the z direction.
+        symprec (float): The symmetry precision for the ASE calculator. This parameter determines the tolerance
+                         for symmetry operations, affecting the identification of equivalent atoms and the overall
+                         symmetry of the system. For more details, refer to the ASE documentation:
+                         https://wiki.fysik.dtu.dk/ase/ase/constraints.html#ase.constraints.FixSymmetry
+
+    Example usage:
+    ```python
+    from ase.optimize import BFGS
+    atoms = to_ase(material)
+    calc = SurfaceDistanceCalculator(shadowing_radius=2.5)
+
+    atoms.calc = calc
+    opt = BFGS(atoms)
+    opt.run(fmax=0.05)
+    ```
+    Args:
+        shadowing_radius (float): Radius for atoms shadowing underlying from being treated as a surface, in Angstroms.
+        force_constant (float): The force constant for the finite difference approximation of the
+    Note:
+        Built following https://wiki.fysik.dtu.dk/ase/development/calculators.html
+
+        The calculate method is responsible for computing the energy and forces (if requested).
+        Forces are estimated using a finite difference method, which is a simple approximation
+        and might not be the most accurate or efficient for all cases.
+    """
+
+    implemented_properties = ["energy", "forces"]
+
+    def __init__(
+        self,
+        shadowing_radius: float = 2.5,
+        is_substrate_fixed: bool = True,
+        is_z_axis_fixed: bool = True,
+        symprec: float = 0.01,
+        material_calculator: Union[MaterialCalculator, InterfaceMaterialCalculator] = InterfaceMaterialCalculator(),
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.shadowing_radius = shadowing_radius
+        self.fix_substrate = is_substrate_fixed
+        self.fix_z = is_z_axis_fixed
+        self.symprec = symprec
+        self.material_calculator = material_calculator
+
+    def _add_constraints(self, atoms: ASEAtoms) -> ASEAtoms:
+        constraints: List[Union[ASEFixAtoms, ASEFixedPlane, RigidFilmXYInterfaceConstraint]] = []
+        if self.fix_substrate:
+            substrate_indices = [i for i, tag in enumerate(atoms.get_tags()) if tag == 0]
+            constraints.append(ASEFixAtoms(indices=substrate_indices))
+        if self.fix_z:
+            all_indices = list(range(len(atoms)))
+            constraints.append(ASEFixedPlane(indices=all_indices, direction=[0, 0, 1]))
+
+        film_indices = [i for i, tag in enumerate(atoms.get_tags()) if tag == 1]
+        if film_indices:
+            constraints.append(RigidFilmXYInterfaceConstraint(film_indices=film_indices))
+
+        atoms.set_constraint(constraints)
+        return atoms
+
+    def _calculate_forces(self, atoms: ASEAtoms, energy: float) -> np.ndarray:
+        forces = np.zeros((len(atoms), 3))
+        dx = 0.01
+        for i in range(len(atoms)):
+            for j in range(3):
+                atoms_plus = atoms.copy()
+                atoms_plus.positions[i, j] += dx
+                material_plus = Material(from_ase(atoms_plus))
+                energy_plus = self.material_calculator.get_energy(material_plus)
+
+                forces[i, j] = -(energy_plus - energy) / dx
+
+        return forces
+
+    def calculate(self, atoms: Optional[ASEAtoms] = None, properties=["energy"], system_changes=ase_all_changes):
+        if atoms is None:
+            atoms = self.atoms.copy()
+
+        atoms = self._add_constraints(atoms)
+        constraints = atoms.constraints
+
+        super().calculate(atoms, properties, system_changes)
+        material = Material(from_ase(atoms))
+        energy = self.material_calculator.get_energy(material)
+
+        self.results = {"energy": energy}
+
+        if "forces" in properties:
+            forces = self._calculate_forces(atoms, energy)
+            for constraint in constraints:
+                constraint.adjust_forces(atoms, forces)
+            self.results["forces"] = forces
+
+    def get_potential_energy(self, atoms=None, force_consistent=False):
+        return self.get_property("energy", atoms)
+
+    def get_forces(self, atoms=None):
+        return self.get_property("forces", atoms)
diff --git a/src/py/mat3ra/made/tools/calculate/constraints.py b/src/py/mat3ra/made/tools/calculate/ase/constraints.py
similarity index 97%
rename from src/py/mat3ra/made/tools/calculate/constraints.py
rename to src/py/mat3ra/made/tools/calculate/ase/constraints.py
index 4581528f..bb39cf5f 100644
--- a/src/py/mat3ra/made/tools/calculate/constraints.py
+++ b/src/py/mat3ra/made/tools/calculate/ase/constraints.py
@@ -3,7 +3,7 @@
 import numpy as np
 from pydantic import BaseModel
 
-from ..third_party import ASEAtoms
+from ...third_party import ASEAtoms
 
 
 class InterfaceConstraint(BaseModel):
@@ -34,7 +34,7 @@ def adjust_forces(self, forces: np.ndarray) -> None:
         raise NotImplementedError
 
 
-class RigidFilmXYConstraint(InterfaceConstraint):
+class RigidFilmXYInterfaceConstraint(InterfaceConstraint):
     """
     Custom constraint to allow only rigid translation in x and y for film atoms.
     """
diff --git a/src/py/mat3ra/made/tools/calculate/calculators.py b/src/py/mat3ra/made/tools/calculate/calculators.py
index 6c963b7d..963096e3 100644
--- a/src/py/mat3ra/made/tools/calculate/calculators.py
+++ b/src/py/mat3ra/made/tools/calculate/calculators.py
@@ -1,16 +1,13 @@
-from typing import Callable, List, Optional, Union
+from typing import Callable
 
 import numpy as np
 from mat3ra.made.material import Material
 from pydantic import BaseModel
 
 from ..analyze import get_surface_atom_indices
-from ..convert import from_ase
 from ..convert.utils import InterfacePartsEnum
 from ..enums import SurfaceTypes
 from ..modify import get_interface_part
-from ..third_party import ASEAtoms, ASECalculator, ASEFixAtoms, ASEFixedPlane, ase_all_changes
-from .constraints import RigidFilmXYConstraint
 from .interaction_functions import sum_of_inverse_distances_squared
 
 
@@ -123,112 +120,3 @@ def get_energy(
         substrate_coordinates_values = np.array(substrate_surface_atom_coordinates.values)
 
         return interaction_function(film_coordinates_values, substrate_coordinates_values)
-
-
-class FilmSubstrateDistanceASECalculator(ASECalculator):
-    """
-    ASE calculator that calculates the interaction energy between a film and substrate in an interface material.
-
-    Args:
-        shadowing_radius (float): The radius for atom to shadow underlying from being considered surface, in Angstroms.
-        force_constant (float): The force constant for the finite difference approximation of the forces.
-        is_substrate_fixed (bool): Whether to fix the substrate atoms.
-        is_z_axis_fixed (bool): Whether to fix atoms movement in the z direction.
-        symprec (float): The symmetry precision for the ASE calculator. This parameter determines the tolerance
-                         for symmetry operations, affecting the identification of equivalent atoms and the overall
-                         symmetry of the system. For more details, refer to the ASE documentation:
-                         https://wiki.fysik.dtu.dk/ase/ase/constraints.html#ase.constraints.FixSymmetry
-
-    Example usage:
-    ```python
-    from ase.optimize import BFGS
-    atoms = to_ase(material)
-    calc = SurfaceDistanceCalculator(shadowing_radius=2.5)
-
-    atoms.calc = calc
-    opt = BFGS(atoms)
-    opt.run(fmax=0.05)
-    ```
-    Args:
-        shadowing_radius (float): Radius for atoms shadowing underlying from being treated as a surface, in Angstroms.
-        force_constant (float): The force constant for the finite difference approximation of the
-    Note:
-        Built following https://wiki.fysik.dtu.dk/ase/development/calculators.html
-
-        The calculate method is responsible for computing the energy and forces (if requested).
-        Forces are estimated using a finite difference method, which is a simple approximation
-        and might not be the most accurate or efficient for all cases.
-    """
-
-    implemented_properties = ["energy", "forces"]
-
-    def __init__(
-        self,
-        shadowing_radius: float = 2.5,
-        is_substrate_fixed: bool = True,
-        is_z_axis_fixed: bool = True,
-        symprec: float = 0.01,
-        material_calculator: Union[MaterialCalculator, InterfaceMaterialCalculator] = InterfaceMaterialCalculator(),
-        **kwargs,
-    ):
-        super().__init__(**kwargs)
-        self.shadowing_radius = shadowing_radius
-        self.fix_substrate = is_substrate_fixed
-        self.fix_z = is_z_axis_fixed
-        self.symprec = symprec
-        self.material_calculator = material_calculator
-
-    def _add_constraints(self, atoms: ASEAtoms) -> ASEAtoms:
-        constraints: List[Union[ASEFixAtoms, ASEFixedPlane, RigidFilmXYConstraint]] = []
-        if self.fix_substrate:
-            substrate_indices = [i for i, tag in enumerate(atoms.get_tags()) if tag == 0]
-            constraints.append(ASEFixAtoms(indices=substrate_indices))
-        if self.fix_z:
-            all_indices = list(range(len(atoms)))
-            constraints.append(ASEFixedPlane(indices=all_indices, direction=[0, 0, 1]))
-
-        film_indices = [i for i, tag in enumerate(atoms.get_tags()) if tag == 1]
-        if film_indices:
-            constraints.append(RigidFilmXYConstraint(film_indices=film_indices))
-
-        atoms.set_constraint(constraints)
-        return atoms
-
-    def _calculate_forces(self, atoms: ASEAtoms, energy: float) -> np.ndarray:
-        forces = np.zeros((len(atoms), 3))
-        dx = 0.01
-        for i in range(len(atoms)):
-            for j in range(3):
-                atoms_plus = atoms.copy()
-                atoms_plus.positions[i, j] += dx
-                material_plus = Material(from_ase(atoms_plus))
-                energy_plus = self.material_calculator.get_energy(material_plus)
-
-                forces[i, j] = -(energy_plus - energy) / dx
-
-        return forces
-
-    def calculate(self, atoms: Optional[ASEAtoms] = None, properties=["energy"], system_changes=ase_all_changes):
-        if atoms is None:
-            atoms = self.atoms.copy()
-
-        atoms = self._add_constraints(atoms)
-        constraints = atoms.constraints
-
-        super().calculate(atoms, properties, system_changes)
-        material = Material(from_ase(atoms))
-        energy = self.material_calculator.get_energy(material)
-
-        self.results = {"energy": energy}
-
-        if "forces" in properties:
-            forces = self._calculate_forces(atoms, energy)
-            for constraint in constraints:
-                constraint.adjust_forces(atoms, forces)
-            self.results["forces"] = forces
-
-    def get_potential_energy(self, atoms=None, force_consistent=False):
-        return self.get_property("energy", atoms)
-
-    def get_forces(self, atoms=None):
-        return self.get_property("forces", atoms)

From 5aa64542e0ed668b721787b07915642557b34078 Mon Sep 17 00:00:00 2001
From: VsevolodX <79542055+VsevolodX@users.noreply.github.com>
Date: Wed, 18 Sep 2024 19:09:30 -0700
Subject: [PATCH 17/17] update: isolate filtering

---
 .../made/tools/calculate/ase/__init__.py      | 19 +++++++++++++++++--
 1 file changed, 17 insertions(+), 2 deletions(-)

diff --git a/src/py/mat3ra/made/tools/calculate/ase/__init__.py b/src/py/mat3ra/made/tools/calculate/ase/__init__.py
index 647e9b49..4ba38758 100644
--- a/src/py/mat3ra/made/tools/calculate/ase/__init__.py
+++ b/src/py/mat3ra/made/tools/calculate/ase/__init__.py
@@ -4,11 +4,26 @@
 from mat3ra.made.material import Material
 
 from ...convert import from_ase
+from ...convert.utils import INTERFACE_LABELS_MAP, InterfacePartsEnum
 from ...third_party import ASEAtoms, ASECalculator, ASEFixAtoms, ASEFixedPlane, ase_all_changes
 from ..calculators import InterfaceMaterialCalculator, MaterialCalculator
 from .constraints import RigidFilmXYInterfaceConstraint
 
 
+def get_interface_part_indices(atoms: ASEAtoms, part: InterfacePartsEnum) -> List[int]:
+    """
+    Get the indices of the atoms in the interface part.
+
+    Args:
+        atoms (ASEAtoms): The ASE Atoms object.
+        part (str): The interface part to get the indices of.
+
+    Returns:
+        List[int]: The indices of the atoms in the interface part.
+    """
+    return [i for i, tag in enumerate(atoms.get_tags()) if tag == INTERFACE_LABELS_MAP[part]]
+
+
 class FilmSubstrateDistanceASECalculator(ASECalculator):
     """
     ASE calculator that calculates the interaction energy between a film and substrate in an interface material.
@@ -65,13 +80,13 @@ def __init__(
     def _add_constraints(self, atoms: ASEAtoms) -> ASEAtoms:
         constraints: List[Union[ASEFixAtoms, ASEFixedPlane, RigidFilmXYInterfaceConstraint]] = []
         if self.fix_substrate:
-            substrate_indices = [i for i, tag in enumerate(atoms.get_tags()) if tag == 0]
+            substrate_indices = get_interface_part_indices(atoms, InterfacePartsEnum.SUBSTRATE)
             constraints.append(ASEFixAtoms(indices=substrate_indices))
         if self.fix_z:
             all_indices = list(range(len(atoms)))
             constraints.append(ASEFixedPlane(indices=all_indices, direction=[0, 0, 1]))
 
-        film_indices = [i for i, tag in enumerate(atoms.get_tags()) if tag == 1]
+        film_indices = get_interface_part_indices(atoms, InterfacePartsEnum.FILM)
         if film_indices:
             constraints.append(RigidFilmXYInterfaceConstraint(film_indices=film_indices))