Feature/SOF-7782 Maxwell displacement

pyproject.toml

@@ -21,7 +21,6 @@ dependencies = [
     "mat3ra-utils",
     "mat3ra-esse",
     "mat3ra-code"
-
 ]
 
 [project.optional-dependencies]
@@ -31,6 +30,7 @@ tools = [
     "pymatgen==2024.4.13",
     "ase",
     "pymatgen-analysis-defects==2024.4.23",
+    "mat3ra-periodic-table>=2025.12.26",
 ]
 dev = [
     "pre-commit",

src/py/mat3ra/made/tools/build_components/operations/core/modifications/perturb/__init__.py

@@ -1,7 +1,15 @@
-from .functions import FunctionHolder, PerturbationFunctionHolder, SineWavePerturbationFunctionHolder
+from .functions import (
+    AtomicMassDependentFunctionHolder,
+    FunctionHolder,
+    MaxwellBoltzmannDisplacementHolder,
+    PerturbationFunctionHolder,
+    SineWavePerturbationFunctionHolder,
+)
 
 __all__ = [
+    "AtomicMassDependentFunctionHolder",
     "FunctionHolder",
+    "MaxwellBoltzmannDisplacementHolder",
     "PerturbationFunctionHolder",
     "SineWavePerturbationFunctionHolder",
 ]

src/py/mat3ra/made/tools/build_components/operations/core/modifications/perturb/functions/__init__.py

@@ -1,9 +1,13 @@
+from .atomic_mass_dependent_function_holder import AtomicMassDependentFunctionHolder
 from .function_holder import FunctionHolder
+from .maxwell_boltzmann import MaxwellBoltzmannDisplacementHolder
 from .perturbation_function_holder import PerturbationFunctionHolder
 from .sine_wave_perturbation_function_holder import SineWavePerturbationFunctionHolder
 
 __all__ = [
+    "AtomicMassDependentFunctionHolder",
     "FunctionHolder",
+    "MaxwellBoltzmannDisplacementHolder",
     "PerturbationFunctionHolder",
     "SineWavePerturbationFunctionHolder",
 ]

src/py/mat3ra/made/tools/build_components/operations/core/modifications/perturb/functions/atomic_mass_dependent_function_holder.py

@@ -0,0 +1,32 @@
+from typing import Any, List, Optional, Union
+
+import sympy as sp
+from mat3ra.periodic_table.helpers import get_atomic_mass_from_element
+
+from .function_holder import FunctionHolder
+
+
+class AtomicMassDependentFunctionHolder(FunctionHolder):
+    variables: List[str] = ["x", "y", "z", "m"]
+
+    def __init__(
+        self,
+        function: Union[sp.Expr, str] = sp.Symbol("f"),
+        variables: Optional[List[str]] = None,
+        **data: Any,
+    ):
+        if variables is None:
+            expr = self._to_expr(function)
+            vs = sorted(expr.free_symbols, key=lambda s: s.name)
+            variables = [str(v) for v in vs] or ["x", "y", "z", "m"]
+
+        super().__init__(function=function, variables=variables, **data)
+
+    @staticmethod
+    def get_atomic_mass(coordinate: List[float], material) -> float:
+        if material is None:
+            raise ValueError("Material is required to extract atomic mass")
+
+        atom_id = material.basis.coordinates.get_element_id_by_value(coordinate)
+        element = material.basis.elements.get_element_value_by_index(atom_id)
+        return get_atomic_mass_from_element(element)

src/py/mat3ra/made/tools/build_components/operations/core/modifications/perturb/functions/function_holder.py

@@ -52,7 +52,7 @@ def _to_expr(expr_or_str: Union[sp.Expr, str]) -> sp.Expr:
     def function_str(self) -> str:
         return str(self.function)
 
-    def apply_function(self, coordinate: List[float]) -> float:
+    def apply_function(self, coordinate: List[float], **kwargs: Any) -> float:
         values = [coordinate[AXIS_TO_INDEX_MAP[var]] for var in self.variables]
         return self.function_numeric(*values)
 

src/py/mat3ra/made/tools/build_components/operations/core/modifications/perturb/functions/maxwell_boltzmann.py

@@ -0,0 +1,39 @@
+from typing import Any, List, Optional
+
+import numpy as np
+from pydantic import Field, model_validator
+
+from .atomic_mass_dependent_function_holder import AtomicMassDependentFunctionHolder
+
+DEFAULT_DISORDER_PARAMETER = 1.0
+
+
+class MaxwellBoltzmannDisplacementHolder(AtomicMassDependentFunctionHolder):
+    disorder_parameter: float = Field(
+        default=DEFAULT_DISORDER_PARAMETER,
+        exclude=True,
+        description="Disorder parameter. Can be viewed as effective temperature in eV.",
+    )
+    random_seed: Optional[int] = Field(default=None, exclude=True)
+    random_state: Any = Field(default=None, exclude=True)
+    is_mass_used: bool = Field(default=True, exclude=True)
+
+    @model_validator(mode="after")
+    def setup_random_state(self):
+        if self.random_state is None:
+            self.random_state = np.random.RandomState(self.random_seed) if self.random_seed is not None else np.random
+        return self
+
+    def apply_function(self, coordinate, material=None) -> List[float]:
+        if material is None:
+            raise ValueError("MaxwellBoltzmannDisplacementHolder requires 'material' kwargs")
+
+        if self.is_mass_used:
+            mass = self.get_atomic_mass(coordinate, material)
+            variance = self.disorder_parameter / mass
+        else:
+            variance = self.disorder_parameter
+
+        std_dev = np.sqrt(variance)
+        displacement = self.random_state.normal(0.0, std_dev, size=3)
+        return displacement.tolist()

src/py/mat3ra/made/tools/build_components/operations/core/modifications/perturb/helpers.py

@@ -1,13 +1,14 @@
-from typing import Union
+from typing import Optional, Union
 
 import sympy as sp
-from mat3ra.made.material import Material
 
-from ..... import MaterialWithBuildMetadata
+from mat3ra.made.material import Material
 from .builders.base import PerturbationBuilder
 from .builders.isometric import IsometricPerturbationBuilder
 from .configuration import PerturbationConfiguration
 from .functions import PerturbationFunctionHolder
+from .functions.maxwell_boltzmann import MaxwellBoltzmannDisplacementHolder
+from ..... import MaterialWithBuildMetadata
 
 
 def create_perturbation(
@@ -39,3 +40,44 @@ def create_perturbation(
     else:
         builder = PerturbationBuilder()
     return builder.get_material(configuration)
+
+
+def create_maxwell_displacement(
+    material: Union[Material, MaterialWithBuildMetadata],
+    disorder_parameter: float,
+    random_seed: Optional[int] = None,
+    is_mass_used: bool = True,
+) -> Material:
+    """
+    Apply Maxwell-Boltzmann random displacements to a material.
+
+    Generates random 3D displacement vectors where each component follows a normal
+    distribution with variance proportional to disorder_parameter/m (if is_mass_used=True)
+    or disorder_parameter (if is_mass_used=False), where m is atomic mass.
+
+    Args:
+        material: The material to be perturbed.
+        disorder_parameter: Disorder parameter controlling displacement magnitude,
+                            can be viewed as effective temperature in eV.
+        random_seed: Optional random seed for reproducibility for the same material and parameters.
+        is_mass_used: If True, displacement variance is disorder_parameter/m (mass-dependent).
+                     If False, displacement variance is disorder_parameter (mass-independent).
+
+    Returns:
+        Material with applied Maxwell-Boltzmann displacements.
+    """
+    displacement_holder = MaxwellBoltzmannDisplacementHolder(
+        disorder_parameter=disorder_parameter,
+        random_seed=random_seed,
+        is_mass_used=is_mass_used,
+    )
+
+    configuration = PerturbationConfiguration(
+        material=material,
+        perturbation_function_holder=displacement_holder,
+        use_cartesian_coordinates=True,
+    )
+
+    builder = PerturbationBuilder()
+
+    return builder.get_material(configuration)

src/py/mat3ra/made/tools/operations/core/unary.py

@@ -84,12 +84,11 @@ def perturb(
     perturbed_coordinates: List[List[float]] = []
 
     for coordinate in original_coordinates:
-        # If func_holder returns a scalar, assume z-axis; otherwise vector
-        displacement = perturbation_function.apply_function(coordinate)
+        displacement = perturbation_function.apply_function(coordinate, material=new_material)
+
         if isinstance(displacement, (list, tuple, np.ndarray)):
             delta = np.array(displacement)
         else:
-            # scalar: apply to z-axis
             delta = np.array([0.0, 0.0, displacement])
 
         new_coordinate = np.array(coordinate) + delta

tests/py/unit/test_tools_build_maxwell_disorder.py

@@ -0,0 +1,123 @@
+import numpy as np
+import pytest
+from mat3ra.made.material import Material
+from mat3ra.made.tools.build_components.operations.core.modifications.perturb.functions.maxwell_boltzmann import (
+    MaxwellBoltzmannDisplacementHolder,
+)
+from mat3ra.made.tools.build_components.operations.core.modifications.perturb.helpers import create_maxwell_displacement
+from mat3ra.made.tools.helpers import create_supercell
+from mat3ra.periodic_table.helpers import get_atomic_mass_from_element
+
+from .fixtures.bulk import BULK_Si_PRIMITIVE
+from .fixtures.slab import SI_CONVENTIONAL_SLAB_001
+
+DISORDER_PARAMETER = 1.0  # Temperature-like
+RANDOM_SEED = 42
+NUM_SAMPLES_FOR_MSD = 1000
+
+
+@pytest.mark.parametrize("random_seed", [None, 42, 123])
+def test_maxwell_displacement_deterministic(random_seed):
+    material = Material.create(BULK_Si_PRIMITIVE)
+    displacement_func1 = MaxwellBoltzmannDisplacementHolder(
+        disorder_parameter=DISORDER_PARAMETER, random_seed=random_seed
+    )
+    displacement_func2 = MaxwellBoltzmannDisplacementHolder(
+        disorder_parameter=DISORDER_PARAMETER, random_seed=random_seed
+    )
+
+    coord = [0.0, 0.0, 0.0]
+
+    if random_seed is not None:
+        disp1 = displacement_func1.apply_function(coord, material=material)
+        disp2 = displacement_func2.apply_function(coord, material=material)
+        assert np.allclose(disp1, disp2)
+
+        # Different seed should give different results
+        displacement_func3 = MaxwellBoltzmannDisplacementHolder(
+            disorder_parameter=DISORDER_PARAMETER, random_seed=random_seed + 1
+        )
+        disp3 = displacement_func3.apply_function(coord, material=material)
+        assert not np.allclose(disp1, disp3) or np.allclose(disp1, [0, 0, 0], atol=1e-10)
+    else:
+        # No seed: different instances should give different results (non-deterministic)
+        disp1 = displacement_func1.apply_function(coord, material=material)
+        disp2 = displacement_func2.apply_function(coord, material=material)
+        assert not np.allclose(disp1, disp2) or np.allclose(disp1, [0, 0, 0], atol=1e-10)
+
+
+def test_maxwell_displacement_perturb_integration():
+    material = Material.create(BULK_Si_PRIMITIVE)
+    original_coords = [coord[:] for coord in material.basis.coordinates.values]
+
+    perturbed_material = create_maxwell_displacement(
+        material, disorder_parameter=DISORDER_PARAMETER, random_seed=RANDOM_SEED
+    )
+
+    assert len(perturbed_material.basis.coordinates.values) == len(original_coords)
+    for i, (orig, pert) in enumerate(zip(original_coords, perturbed_material.basis.coordinates.values)):
+        delta = np.array(pert) - np.array(orig)
+        assert np.linalg.norm(delta) > 0 or np.allclose(delta, 0, atol=1e-10)
+
+
+def test_maxwell_displacement_msd_expectation():
+    material = Material.create(BULK_Si_PRIMITIVE)
+    si_mass = get_atomic_mass_from_element("Si")
+    disorder_parameter = DISORDER_PARAMETER
+    expected_variance = disorder_parameter / si_mass
+    expected_msd = 3 * expected_variance
+
+    displacements = []
+    coord = [0.0, 0.0, 0.0]
+    for _ in range(NUM_SAMPLES_FOR_MSD):
+        displacement_func = MaxwellBoltzmannDisplacementHolder(disorder_parameter=disorder_parameter, random_seed=None)
+        disp = displacement_func.apply_function(coord, material=material)
+        displacements.append(disp)
+
+    displacements_array = np.array(displacements)
+    msd = np.mean(np.sum(displacements_array**2, axis=1))
+
+    assert abs(msd - expected_msd) / expected_msd < 0.3
+
+
+@pytest.mark.parametrize(
+    "slab_config, temperature_k, random_seed",
+    [
+        (SI_CONVENTIONAL_SLAB_001, 1300.0, 42),
+        (SI_CONVENTIONAL_SLAB_001, 1300.0, 42),
+    ],
+)
+def test_maxwell_boltzmann_on_slab(slab_config, temperature_k, random_seed):
+    material = Material.create(slab_config)
+    material = create_supercell(material, scaling_factor=[4, 4, 1])
+    original_coords = [coord[:] for coord in material.basis.coordinates.values]
+    original_lattice = material.lattice.vector_arrays.copy()
+
+    perturbed_material = create_maxwell_displacement(
+        material, disorder_parameter=temperature_k, random_seed=random_seed
+    )
+
+    assert len(perturbed_material.basis.coordinates.values) == len(original_coords)
+    assert len(perturbed_material.basis.elements.values) == len(material.basis.elements.values)
+
+    coordinate_changes = []
+    for i, (orig, pert) in enumerate(zip(original_coords, perturbed_material.basis.coordinates.values)):
+        delta = np.array(pert) - np.array(orig)
+        displacement_magnitude = np.linalg.norm(delta)
+        coordinate_changes.append(displacement_magnitude)
+
+    max_displacement = max(coordinate_changes)
+    mean_displacement = np.mean(coordinate_changes)
+
+    assert max_displacement > 0
+    assert mean_displacement > 0
+
+    si_mass = get_atomic_mass_from_element("Si")
+    expected_std = np.sqrt(temperature_k / si_mass)
+
+    assert mean_displacement < 5 * expected_std
+
+    assert np.allclose(perturbed_material.lattice.vector_arrays, original_lattice, atol=1e-10)
+
+    for i, element in enumerate(material.basis.elements.values):
+        assert perturbed_material.basis.elements.values[i] == element