Merge pull request #142 from Exabyte-io/feature/SOF-7380

Feature/SOF-7380 feature: add vacuum adjustment functions

src/py/mat3ra/made/basis.py

@@ -67,13 +67,13 @@ def is_in_cartesian_units(self):
     def to_cartesian(self):
         if self.is_in_cartesian_units:
             return
-        self.coordinates = self.coordinates.map_array_in_place(self.cell.convert_point_to_cartesian)
+        self.coordinates.map_array_in_place(self.cell.convert_point_to_cartesian)
         self.units = AtomicCoordinateUnits.cartesian
 
     def to_crystal(self):
         if self.is_in_crystal_units:
             return
-        self.coordinates = self.coordinates.map_array_in_place(self.cell.convert_point_to_crystal)
+        self.coordinates.map_array_in_place(self.cell.convert_point_to_crystal)
         self.units = AtomicCoordinateUnits.crystal
 
     def add_atom(self, element="Si", coordinate=[0.5, 0.5, 0.5]):

src/py/mat3ra/made/cell.py

@@ -44,7 +44,7 @@ def convert_point_to_cartesian(self, point):
         np_vector = np.array(self.vectors_as_nested_array)
         return np.dot(point, np_vector)
 
-    def convert_point_to_fractional(self, point):
+    def convert_point_to_crystal(self, point):
         np_vector = np.array(self.vectors_as_nested_array)
         return np.dot(point, np.linalg.inv(np_vector))
 

src/py/mat3ra/made/tools/analyze.py

@@ -1,4 +1,4 @@
-from typing import Callable, List, Optional
+from typing import Callable, List, Optional, Literal
 
 import numpy as np
 
@@ -211,7 +211,7 @@ def get_atom_indices_with_condition_on_coordinates(
     Args:
         material (Material): Material object
         condition (Callable[List[float], bool]): Function that checks if coordinates satisfy the condition.
-        use_cartesian (bool): Whether to use Cartesian coordinates for the condition evaluation.
+        use_cartesian_coordinates (bool): Whether to use Cartesian coordinates for the condition evaluation.
 
     Returns:
         List[int]: List of indices of atoms whose coordinates satisfy the condition.
@@ -229,3 +229,30 @@ def get_atom_indices_with_condition_on_coordinates(
             selected_indices.append(coord.id)
 
     return selected_indices
+
+
+def get_atomic_coordinates_extremum(
+    material: Material,
+    extremum: Literal["max", "min"] = "max",
+    axis: Literal["x", "y", "z"] = "z",
+    use_cartesian_coordinates: bool = False,
+) -> float:
+    """
+    Return minimum or maximum of coordinates along the specified axis.
+
+    Args:
+        material (Material): Material object.
+        extremum (str): "min" or "max".
+        axis (str):  "x", "y", or "z".
+        use_cartesian_coordinates (bool): Whether to use Cartesian coordinates.
+    Returns:
+        float: Minimum or maximum of coordinates along the specified axis.
+    """
+    new_material = material.clone()
+    if use_cartesian_coordinates:
+        new_basis = new_material.basis
+        new_basis.to_cartesian()
+        new_material.basis = new_basis
+    coordinates = new_material.basis.coordinates.to_array_of_values_with_ids()
+    values = [coord.value[{"x": 0, "y": 1, "z": 2}[axis]] for coord in coordinates]
+    return getattr(np, extremum)(values)

src/py/mat3ra/made/tools/convert/__init__.py

@@ -184,6 +184,8 @@ def to_ase(material_or_material_data: Union[Material, Dict[str, Any]]) -> ASEAto
         atoms.set_tags(map_array_with_id_value_to_array(atomic_labels))
     if "metadata" in material_config:
         atoms.info.update({"metadata": material_config["metadata"]})
+
+    atoms.info.update({"name": material_config["name"]})
     return atoms
 
 
@@ -205,6 +207,7 @@ def from_ase(ase_atoms: ASEAtoms) -> Dict[str, Any]:
     ase_metadata = ase_atoms.info.get("metadata", {})
     if ase_metadata:
         material["metadata"].update(ase_metadata)
+    material["name"] = ase_atoms.info.get("name", "")
     return material
 
 

src/py/mat3ra/made/tools/modify.py

@@ -1,17 +1,19 @@
-from typing import Callable, List, Optional, Union
+from typing import Callable, List, Literal, Optional, Union
 
-import numpy as np
 from mat3ra.made.material import Material
 
-from .analyze import get_atom_indices_with_condition_on_coordinates, get_atom_indices_within_radius_pbc
-from .convert import decorator_convert_material_args_kwargs_to_structure
-from .third_party import PymatgenSpacegroupAnalyzer, PymatgenStructure
+from .analyze import (
+    get_atom_indices_with_condition_on_coordinates,
+    get_atom_indices_within_radius_pbc,
+    get_atomic_coordinates_extremum,
+)
+from .convert import decorator_convert_material_args_kwargs_to_structure, from_ase, to_ase
+from .third_party import PymatgenStructure, ase_add_vacuum
 from .utils import (
     is_coordinate_in_box,
     is_coordinate_in_cylinder,
     is_coordinate_in_triangular_prism,
     is_coordinate_within_layer,
-    translate_to_bottom_pymatgen_structure,
 )
 
 
@@ -35,22 +37,54 @@ def filter_by_label(material: Material, label: Union[int, str]) -> Material:
     return new_material
 
 
-@decorator_convert_material_args_kwargs_to_structure
-def translate_to_bottom(structure: PymatgenStructure, use_conventional_cell: bool = True):
+def translate_to_z_level(
+    material: Material, z_level: Optional[Literal["top", "bottom", "center"]] = "bottom"
+) -> Material:
     """
-    Translate atoms to the bottom of the cell (vacuum on top) to allow for the correct consecutive interface generation.
-    If use_conventional_cell is passed, conventional cell is used.
+    Translate atoms to the specified z-level.
 
     Args:
-        structure (Structure): The pymatgen Structure object to normalize.
-        use_conventional_cell: Whether to convert to the conventional cell.
+        material (Material): The material object to normalize.
+        z_level (str): The z-level to translate the atoms to (top, bottom, center)
     Returns:
-        Structure: The normalized pymatgen Structure object.
+        Material: The translated material object.
     """
-    if use_conventional_cell:
-        structure = PymatgenSpacegroupAnalyzer(structure).get_conventional_standard_structure()
-    structure = translate_to_bottom_pymatgen_structure(structure)
-    return structure
+    min_z = get_atomic_coordinates_extremum(material, "min")
+    max_z = get_atomic_coordinates_extremum(material)
+    if z_level == "top":
+        material = translate_by_vector(material, vector=[0, 0, 1 - max_z])
+    elif z_level == "bottom":
+        material = translate_by_vector(material, vector=[0, 0, -min_z])
+    elif z_level == "center":
+        material = translate_by_vector(material, vector=[0, 0, (1 - min_z - max_z) / 2])
+    return material
+
+
+def translate_by_vector(
+    material: Material,
+    vector: Optional[List[float]] = None,
+    use_cartesian_coordinates: bool = False,
+) -> Material:
+    """
+    Translate atoms by a vector.
+
+    Args:
+        material (Material): The material object to normalize.
+        vector (List[float]): The vector to translate the atoms by (in crystal coordinates by default).
+        use_cartesian_coordinates (bool): Whether to use cartesian coordinates.
+    Returns:
+        Material: The translated material object.
+    """
+    if not use_cartesian_coordinates:
+        vector = material.basis.cell.convert_point_to_cartesian(vector)
+
+    if vector is None:
+        vector = [0, 0, 0]
+
+    atoms = to_ase(material)
+    # ASE accepts cartesian coordinates for translation
+    atoms.translate(tuple(vector))
+    return Material(from_ase(atoms))
 
 
 @decorator_convert_material_args_kwargs_to_structure
@@ -140,7 +174,7 @@ def filter_by_layers(
     if central_atom_id is not None:
         center_coordinate = material.basis.coordinates.get_element_value_by_index(central_atom_id)
     vectors = material.lattice.vectors
-    direction_vector = np.array(vectors[2])
+    direction_vector = vectors[2]
 
     def condition(coordinate):
         return is_coordinate_within_layer(coordinate, center_coordinate, direction_vector, layer_thickness)
@@ -323,3 +357,61 @@ def condition(coordinate):
     return filter_by_condition_on_coordinates(
         material, condition, use_cartesian_coordinates=use_cartesian_coordinates, invert_selection=invert_selection
     )
+
+
+def add_vacuum(material: Material, vacuum: float = 5.0, on_top=True, to_bottom=False) -> Material:
+    """
+    Add vacuum to the material along the c-axis.
+    On top, on bottom, or both.
+
+    Args:
+        material (Material): The material object to add vacuum to.
+        vacuum (float): The thickness of the vacuum to add in angstroms.
+        on_top (bool): Whether to add vacuum on top.
+        to_bottom (bool): Whether to add vacuum on bottom.
+
+    Returns:
+        Material: The material object with vacuum added.
+    """
+    new_material_atoms = to_ase(material)
+    vacuum_amount = vacuum * 2 if on_top and to_bottom else vacuum
+    ase_add_vacuum(new_material_atoms, vacuum_amount)
+    new_material = Material(from_ase(new_material_atoms))
+    if to_bottom and not on_top:
+        new_material = translate_to_z_level(new_material, z_level="top")
+    elif on_top and to_bottom:
+        new_material = translate_to_z_level(new_material, z_level="center")
+    return new_material
+
+
+def remove_vacuum(material: Material, from_top=True, from_bottom=True, fixed_padding=1.0) -> Material:
+    """
+    Remove vacuum from the material along the c-axis.
+    From top, from bottom, or from both.
+
+    Args:
+        material (Material): The material object to set the vacuum thickness.
+        from_top (bool): Whether to remove vacuum from the top.
+        from_bottom (bool): Whether to remove vacuum from the bottom.
+        fixed_padding (float): The fixed padding of vacuum to add to avoid collisions in pbc (in angstroms).
+
+    Returns:
+        Material: The material object with the vacuum thickness set.
+    """
+    translated_material = translate_to_z_level(material, z_level="bottom")
+    new_basis = translated_material.basis
+    new_basis.to_cartesian()
+    new_lattice = translated_material.lattice
+    new_lattice.c = get_atomic_coordinates_extremum(translated_material, use_cartesian_coordinates=True) + fixed_padding
+    new_basis.cell.vector3 = new_lattice.vectors[2]
+    new_basis.to_crystal()
+    new_material = material.clone()
+
+    new_material.basis = new_basis
+    new_material.lattice = new_lattice
+
+    if from_top and not from_bottom:
+        new_material = translate_to_z_level(new_material, z_level="top")
+    if from_bottom and not from_top:
+        new_material = translate_to_z_level(new_material, z_level="bottom")
+    return new_material

src/py/mat3ra/made/tools/third_party.py

@@ -1,4 +1,5 @@
 from ase import Atoms as ASEAtoms
+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.emt import EMT as ASECalculatorEMT
@@ -36,6 +37,7 @@
     "PymatgenInterstitial",
     "label_pymatgen_slab_termination",
     "ase_make_supercell",
+    "ase_add_vacuum",
     "PymatgenAseAtomsAdaptor",
     "PymatgenPoscar",
 ]

tests/py/unit/fixtures.py

@@ -1,3 +1,6 @@
+import copy
+from typing import Any, Dict
+
 from ase.build import bulk
 from mat3ra.made.material import Material
 from mat3ra.made.tools.build.interface.termination_pair import TerminationPair
@@ -59,7 +62,7 @@
 INTERFACE_NAME = "Cu4(001)-Si8(001), Interface, Strain 0.062pct"
 
 # TODO: Use fixtures package when available
-SI_CONVENTIONAL_CELL = {
+SI_CONVENTIONAL_CELL: Dict[str, Any] = {
     "name": "Si8",
     "basis": {
         "elements": [
@@ -110,7 +113,7 @@
     "isUpdated": True,
 }
 
-SI_SUPERCELL_2X2X1 = {
+SI_SUPERCELL_2X2X1: Dict[str, Any] = {
     "name": "Si8",
     "basis": {
         "elements": [
@@ -162,7 +165,7 @@
 }
 
 
-SI_SLAB_CONFIGURATION = {
+SI_SLAB_CONFIGURATION: Dict[str, Any] = {
     "type": "SlabConfiguration",
     "bulk": SI_CONVENTIONAL_CELL,
     "miller_indices": (0, 0, 1),
@@ -173,9 +176,7 @@
     "use_orthogonal_z": True,
 }
 
-
-SI_SLAB = {
-    "name": "Si8(001), termination Si_P4/mmm_1, Slab",
+SI_SLAB: Dict[str, Any] = {
     "basis": {
         "elements": [
             {"id": 0, "value": "Si"},
@@ -211,6 +212,7 @@
             "units": "angstrom",
         },
     },
+    "name": "Si8(001), termination Si_P4/mmm_1, Slab",
     "isNonPeriodic": False,
     "_id": "",
     "metadata": {
@@ -220,3 +222,14 @@
     },
     "isUpdated": True,
 }
+
+SI_SLAB_VACUUM = copy.deepcopy(SI_SLAB)
+SI_SLAB_VACUUM["basis"]["coordinates"] = [
+    {"id": 0, "value": [0.5, 0.5, 0.386029718]},
+    {"id": 1, "value": [0.5, 0.0, 0.4718141]},
+    {"id": 2, "value": [0.0, 0.0, 0.557598482]},
+    {"id": 3, "value": [-0.0, 0.5, 0.643382864]},
+]
+SI_SLAB_VACUUM["basis"]["cell"] = [[3.867, 0.0, 0.0], [-0.0, 3.867, 0.0], [0.0, 0.0, 15.937527692]]
+SI_SLAB_VACUUM["lattice"]["c"] = 15.937527692
+SI_SLAB_VACUUM["lattice"]["vectors"]["c"] = [0.0, 0.0, 15.937527692]

tests/py/unit/test_tools_modify.py

@@ -2,16 +2,19 @@
 from mat3ra.made.material import Material
 from mat3ra.made.tools.convert import from_ase
 from mat3ra.made.tools.modify import (
+    add_vacuum,
     filter_by_circle_projection,
     filter_by_label,
     filter_by_layers,
     filter_by_rectangle_projection,
     filter_by_sphere,
     filter_by_triangle_projection,
+    remove_vacuum,
+    translate_to_z_level,
 )
 from mat3ra.utils import assertion as assertion_utils
 
-from .fixtures import SI_CONVENTIONAL_CELL
+from .fixtures import SI_CONVENTIONAL_CELL, SI_SLAB, SI_SLAB_VACUUM
 
 COMMON_PART = {
     "units": "crystal",
@@ -136,3 +139,22 @@ def test_filter_by_triangle_projection():
     cavity = filter_by_triangle_projection(material, [0.4, 0.4], [0.4, 0.5], [0.5, 0.5], invert_selection=True)
     assertion_utils.assert_deep_almost_equal(expected_basis_sphere_cluster, section.basis.to_json())
     assertion_utils.assert_deep_almost_equal(expected_basis_sphere_cavity, cavity.basis.to_json())
+
+
+def test_add_vacuum():
+    material = Material(SI_SLAB)
+    material_with_vacuum = add_vacuum(material, 5.0)
+    assertion_utils.assert_deep_almost_equal(SI_SLAB_VACUUM, material_with_vacuum.to_json())
+
+
+def test_remove_vacuum():
+    material_with_vacuum = Material(SI_SLAB_VACUUM)
+    vacuum = 6.836
+    material_with_no_vacuum = remove_vacuum(material_with_vacuum, from_top=True, from_bottom=True, fixed_padding=0)
+    material_with_set_vacuum = add_vacuum(material_with_no_vacuum, vacuum)
+    # to compare correctly, we need to translate the expected material to the bottom
+    # as it down when setting vacuum to 0
+    material = Material(SI_SLAB)
+    material_down = translate_to_z_level(material, z_level="bottom")
+
+    assertion_utils.assert_deep_almost_equal(material_down.to_json(), material_with_set_vacuum.to_json())