Merge pull request #136 from Exabyte-io/feature/SOF-7386

feature/SOF 7386

pyproject.toml

@@ -16,7 +16,8 @@ classifiers = [
 ]
 dependencies = [
     # add requirements here
-    "numpy",
+    # new verison of numpy==2.0.0 is not handled by pymatgen yet
+    "numpy<=1.26.4",
     "mat3ra-utils",
     "mat3ra-esse",
     "mat3ra-code",

src/py/mat3ra/made/basis/basis.py

@@ -0,0 +1,92 @@
+import json
+from typing import Dict, List, Optional
+
+from mat3ra.code.constants import AtomicCoordinateUnits
+from mat3ra.utils.mixins import RoundNumericValuesMixin
+from pydantic import BaseModel
+
+from ..cell.cell import Cell
+from ..utils import ArrayWithIds
+
+
+class Basis(RoundNumericValuesMixin, BaseModel):
+    elements: ArrayWithIds = ArrayWithIds(values=["Si"])
+    coordinates: ArrayWithIds = ArrayWithIds(values=[0, 0, 0])
+    units: str = AtomicCoordinateUnits.crystal
+    cell: Optional[Cell] = None
+    labels: Optional[ArrayWithIds] = ArrayWithIds(values=[])
+    constraints: Optional[ArrayWithIds] = ArrayWithIds(values=[])
+
+    @classmethod
+    def from_dict(
+        cls,
+        elements: List[Dict],
+        coordinates: List[Dict],
+        units: str,
+        labels: Optional[List[Dict]] = None,
+        cell: Optional[Dict] = None,
+        constraints: Optional[List[Dict]] = None,
+    ) -> "Basis":
+        return Basis(
+            elements=ArrayWithIds.from_list_of_dicts(elements),
+            coordinates=ArrayWithIds.from_list_of_dicts(coordinates),
+            units=units,
+            cell=Cell.from_nested_array(cell) if cell else None,
+            labels=ArrayWithIds.from_list_of_dicts(labels) if labels else ArrayWithIds(values=[]),
+            constraints=ArrayWithIds.from_list_of_dicts(constraints) if constraints else ArrayWithIds(values=[]),
+        )
+
+    def to_json(self, skip_rounding=False):
+        json_value = {
+            "elements": self.elements.to_json(),
+            "coordinates": self.coordinates.to_json(skip_rounding=skip_rounding),
+            "units": self.units,
+            "cell": self.cell.to_json(skip_rounding=skip_rounding) if self.cell else None,
+            "labels": self.labels.to_json(),
+        }
+        return json.loads(json.dumps(json_value))
+
+    def clone(self):
+        return Basis(
+            elements=self.toJSON()["elements"],
+            coordinates=self.toJSON()["coordinates"],
+            units=self.units,
+            cell=self.cell,
+            isEmpty=False,
+            labels=self.labels,
+        )
+
+    @property
+    def is_in_crystal_units(self):
+        return self.units == AtomicCoordinateUnits.crystal
+
+    @property
+    def is_in_cartesian_units(self):
+        return self.units == AtomicCoordinateUnits.cartesian
+
+    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.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.units = AtomicCoordinateUnits.crystal
+
+    def add_atom(self, element="Si", coordinate=[0.5, 0.5, 0.5]):
+        self.elements.add_item(element)
+        self.coordinates.add_item(coordinate)
+
+    def remove_atom_by_id(self, id=None):
+        self.elements.remove_item(id)
+        self.coordinates.remove_item(id)
+        self.labels.remove_item(id)
+
+    def filter_atoms_by_ids(self, ids):
+        self.elements.filter_by_ids(ids)
+        self.coordinates.filter_by_ids(ids)
+        if self.labels is not None:
+            self.labels.filter_by_ids(ids)

src/py/mat3ra/made/cell/cell.py

@@ -0,0 +1,58 @@
+from typing import List
+
+import numpy as np
+from mat3ra.esse.models.core.primitive.array_of_3_numbers import ArrayOf3NumberElementsSchema
+from mat3ra.utils.mixins import RoundNumericValuesMixin
+from pydantic import BaseModel
+
+
+class Cell(RoundNumericValuesMixin, BaseModel):
+    # TODO: figure out how to use
+    vector1: ArrayOf3NumberElementsSchema = [1, 0, 0]
+    vector2: ArrayOf3NumberElementsSchema = [0, 1, 0]
+    vector3: ArrayOf3NumberElementsSchema = [0, 0, 1]
+    __round_precision__ = 1e-6
+
+    @classmethod
+    def from_nested_array(cls, nested_array):
+        if not nested_array:
+            nested_array = [cls.vector1, cls.vector2, cls.vector3]
+        return cls(vector1=nested_array[0], vector2=nested_array[1], vector3=nested_array[2])
+
+    def __init__(self, vector1=[1, 0, 0], vector2=[0, 1, 0], vector3=[0, 0, 1]):
+        super().__init__(**{"vector1": vector1, "vector2": vector2, "vector3": vector3})
+
+    @property
+    def vectors_as_array(self, skip_rounding=False) -> List[ArrayOf3NumberElementsSchema]:
+        if skip_rounding:
+            return [self.vector1, self.vector2, self.vector3]
+        return self.round_array_or_number([self.vector1, self.vector2, self.vector3])
+
+    def to_json(self, skip_rounding=False):
+        _ = self.round_array_or_number
+        if skip_rounding:
+            return {
+                "vector1": _(self.vector1) if skip_rounding else self.vector1,
+                "vector2": _(self.vector2) if skip_rounding else self.vector2,
+                "vector3": _(self.vector3) if skip_rounding else self.vector3,
+            }
+
+    def clone(self):
+        return self.from_nested_array(self.vectors_as_array)
+
+    def clone_and_scale_by_matrix(self, matrix):
+        new_cell = self.clone()
+        new_cell.scale_by_matrix(matrix)
+        return new_cell
+
+    def convert_point_to_cartesian(self, point):
+        np_vector = np.array(self.vectors_as_array)
+        return np.dot(point, np_vector)
+
+    def convert_point_to_fractional(self, point):
+        np_vector = np.array(self.vectors_as_array)
+        return np.dot(point, np.linalg.inv(np_vector))
+
+    def scale_by_matrix(self, matrix):
+        np_vector = np.array(self.vectors_as_array)
+        self.vector1, self.vector2, self.vector3 = np.dot(matrix, np_vector).tolist()

src/py/mat3ra/made/lattice/lattice.py

@@ -0,0 +1,83 @@
+import math
+from typing import Any, Dict, List
+
+import numpy as np
+from mat3ra.utils.mixins import RoundNumericValuesMixin
+from pydantic import BaseModel
+
+from ..cell.cell import Cell
+
+HASH_TOLERANCE = 3
+
+
+class Lattice(RoundNumericValuesMixin, BaseModel):
+    a: float = 1.0
+    b: float = a
+    c: float = a
+    alpha: float = 90.0
+    beta: float = 90.0
+    gamma: float = 90.0
+    units: Dict[str, str] = {
+        "length": "angstrom",
+        "angle": "degree",
+    }
+    type: str = "TRI"
+
+    @property
+    def vectors(self) -> List[List[float]]:
+        a = self.a
+        b = self.b
+        c = self.c
+        # Convert degrees to radians for trigonometric functions
+        alpha_rad = math.radians(self.alpha)
+        beta_rad = math.radians(self.beta)
+        gamma_rad = math.radians(self.gamma)
+
+        # Calculate cosines and sines of the angles
+        cos_alpha = math.cos(alpha_rad)
+        cos_beta = math.cos(beta_rad)
+        cos_gamma = math.cos(gamma_rad)
+        sin_alpha = math.sin(alpha_rad)
+        sin_beta = math.sin(beta_rad)
+
+        # Compute gamma star (used in matrix calculation)
+        gamma_star = math.acos((cos_alpha * cos_beta - cos_gamma) / (sin_alpha * sin_beta))
+        cos_gamma_star = math.cos(gamma_star)
+        sin_gamma_star = math.sin(gamma_star)
+
+        # Return the lattice matrix using the derived trigonometric values
+        return [
+            [a * sin_beta, 0.0, a * cos_beta],
+            [-b * sin_alpha * cos_gamma_star, b * sin_alpha * sin_gamma_star, b * cos_alpha],
+            [0.0, 0.0, c],
+        ]
+
+    def to_json(self, skip_rounding: bool = False) -> Dict[str, Any]:
+        __round__ = RoundNumericValuesMixin.round_array_or_number
+        round_func = __round__ if not skip_rounding else lambda x: x
+        return {
+            "a": round_func(self.a),
+            "b": round_func(self.b),
+            "c": round_func(self.c),
+            "alpha": round_func(self.alpha),
+            "beta": round_func(self.beta),
+            "gamma": round_func(self.gamma),
+            "units": self.units,
+            "type": self.type,
+            "vectors": self.vectors,
+        }
+
+    def clone(self, extra_context: Dict[str, Any]) -> "Lattice":
+        return Lattice(**{**self.to_json(), **extra_context})
+
+    @property
+    def vector_arrays(self) -> List[List[float]]:
+        return self.vectors
+
+    @property
+    def cell(self) -> Cell:
+        return Cell.from_nested_array(self.vector_arrays)
+
+    def volume(self) -> float:
+        np_vector = np.array(self.vector_arrays)
+        return abs(np.linalg.det(np_vector))

src/py/mat3ra/made/material.py

@@ -3,7 +3,9 @@
 from mat3ra.code.constants import AtomicCoordinateUnits, Units
 from mat3ra.code.entity import HasDescriptionHasMetadataNamedDefaultableInMemoryEntity
 from mat3ra.esse.models.material import MaterialSchema
-from mat3ra.made.utils import map_array_with_id_value_to_array
+
+from .basis.basis import Basis
+from .lattice.lattice import Lattice
 
 defaultMaterialConfig = {
     "name": "Silicon FCC",
@@ -61,4 +63,20 @@ def to_json(self, exclude: List[str] = []) -> MaterialSchemaJSON:
 
     @property
     def coordinates_array(self) -> List[List[float]]:
-        return map_array_with_id_value_to_array(self.basis["coordinates"])
+        return self.basis.coordinates.values
+
+    @property
+    def basis(self) -> Basis:
+        return Basis.from_dict(**self.get_prop("basis"))
+
+    @basis.setter
+    def basis(self, basis: Basis) -> None:
+        self.set_prop("basis", basis.to_json())
+
+    @property
+    def lattice(self) -> Lattice:
+        return Lattice(**self.get_prop("lattice"))
+
+    @lattice.setter
+    def lattice(self, lattice: Lattice) -> None:
+        self.set_prop("lattice", lattice.to_json())

src/py/mat3ra/made/tools/build/defect/builders.py

@@ -13,7 +13,6 @@
 from ...convert import PymatgenStructure, to_pymatgen
 from ..mixins import ConvertGeneratedItemsPymatgenStructureMixin
 from .configuration import PointDefectConfiguration
-from mat3ra.made.utils import get_array_with_id_value_element_value_by_index
 
 
 class PointDefectBuilderParameters(BaseModel):
@@ -32,8 +31,8 @@ class PointDefectBuilder(ConvertGeneratedItemsPymatgenStructureMixin, BaseBuilde
     _generator: Callable
 
     def _get_species(self, configuration: BaseBuilder._ConfigurationType):
-        crystal_elements = configuration.crystal.basis["elements"]
-        placeholder_specie = get_array_with_id_value_element_value_by_index(crystal_elements, 0)
+        crystal_elements = configuration.crystal.basis.elements.values
+        placeholder_specie = crystal_elements[0]
         return configuration.chemical_element or placeholder_specie
 
     def _generate(self, configuration: BaseBuilder._ConfigurationType) -> List[_GeneratedItemType]:

src/py/mat3ra/made/tools/material.py

@@ -0,0 +1,40 @@
+import numpy as np
+from mat3ra.code.constants import AtomicCoordinateUnits
+
+
+def scale_one_lattice_vector(material, key="a", factor=1.0):
+    """
+    Scales one lattice vector for the given material in place.
+
+    Args:
+        material: The material to scale.
+        key: The key of the lattice vector to scale.
+        factor: The factor to scale the lattice vector by
+    """
+    material.to_cartesian()
+
+    lattice = getattr(material, "lattice")
+    setattr(lattice, key, getattr(lattice, key) * factor)
+
+    setattr(material, "lattice", lattice)
+
+    material.to_crystal()
+
+
+def get_basis_config_translated_to_center(material):
+    """
+    Updates the basis of a material by translating the coordinates
+    so that the center of the material and lattice are aligned.
+    Args:
+        material: The material to update.
+    """
+    original_units = material.basis.units
+    material.to_cartesian()
+    updated_basis = material.basis
+    center_of_coordinates = updated_basis.center_of_coordinates_point()
+    center_of_lattice = 0.5 * np.sum(material.lattice.vector_arrays, axis=0)
+    translation_vector = center_of_lattice - center_of_coordinates
+    updated_basis.translate_by_vector(translation_vector)
+    material.set_basis(updated_basis.to_json())
+    if original_units != AtomicCoordinateUnits.cartesian:
+        material.to_crystal()

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

@@ -1,7 +1,6 @@
 from typing import Union
 
 from mat3ra.made.material import Material
-from mat3ra.made.utils import filter_array_with_id_value_by_ids, filter_array_with_id_value_by_values
 from pymatgen.analysis.structure_analyzer import SpacegroupAnalyzer
 from pymatgen.core.structure import Structure
 
@@ -21,11 +20,11 @@ def filter_by_label(material: Material, label: Union[int, str]) -> Material:
         Material: The filtered material object.
     """
     new_material = material.clone()
-    labels = material.basis["labels"]
-    filtered_labels = filter_array_with_id_value_by_values(labels, label)
-    filtered_label_ids = [item["id"] for item in filtered_labels]
-    for key in ["coordinates", "elements", "labels"]:
-        new_material.basis[key] = filter_array_with_id_value_by_ids(new_material.basis[key], filtered_label_ids)
+    labels_array = new_material.basis.labels.to_array_of_values_with_ids()
+    filtered_label_ids = [_label.id for _label in labels_array if _label.value == label]
+    new_basis = new_material.basis
+    new_basis.filter_atoms_by_ids(filtered_label_ids)
+    new_material.basis = new_basis
     return new_material