Merge branch 'main' into almostequal

atomistics/workflows/elastic/elastic_moduli.py

@@ -2,6 +2,8 @@
 
 import numpy as np
 
+from atomistics.shared.output import OutputElastic
+
 
 def get_bulkmodul_voigt(elastic_matrix: np.ndarray) -> float:
     return (
@@ -214,3 +216,8 @@ def AVR(self) -> float:
     @cache
     def elastic_matrix_eigval(self):
         return get_elastic_matrix_eigval(elastic_matrix=self.elastic_matrix())
+
+    def to_dict(self, output_keys: tuple = OutputElastic.keys()) -> dict:
+        return OutputElastic(**{k: getattr(self, k) for k in OutputElastic.keys()}).get(
+            output_keys=output_keys
+        )

atomistics/workflows/elastic/helper.py

@@ -4,6 +4,8 @@
 import numpy as np
 import scipy.constants
 
+from atomistics.shared.output import OutputElastic
+from atomistics.workflows.elastic.elastic_moduli import ElasticProperties
 from atomistics.workflows.elastic.symmetry import (
     symmetry_analysis,
     get_C_from_A2,
@@ -89,6 +91,30 @@ def generate_structures_helper(
 
 
 def analyse_structures_helper(
+    output_dict: dict,
+    sym_dict: dict,
+    fit_order: int = 2,
+    zero_strain_job_name: str = "s_e_0",
+    output_keys: tuple = OutputElastic.keys(),
+):
+    elastic_matrix, A2, strain_energy, ene0 = _get_elastic_matrix(
+        output_dict=output_dict,
+        Lag_strain_list=sym_dict["Lag_strain_list"],
+        epss=sym_dict["epss"],
+        v0=sym_dict["v0"],
+        LC=sym_dict["LC"],
+        fit_order=fit_order,
+        zero_strain_job_name=zero_strain_job_name,
+    )
+    sym_dict["strain_energy"] = strain_energy
+    sym_dict["e0"] = ene0
+    sym_dict["A2"] = A2
+    return sym_dict, ElasticProperties(elastic_matrix=elastic_matrix).to_dict(
+        output_keys=output_keys
+    )
+
+
+def _get_elastic_matrix(
     output_dict: dict,
     Lag_strain_list: list[float],
     epss: float,

atomistics/workflows/elastic/workflow.py

@@ -3,7 +3,6 @@
 
 from atomistics.shared.output import OutputElastic
 from atomistics.workflows.interface import Workflow
-from atomistics.workflows.elastic.elastic_moduli import ElasticProperties
 from atomistics.workflows.elastic.helper import (
     generate_structures_helper,
     analyse_structures_helper,
@@ -47,7 +46,7 @@ def generate_structures(self) -> dict:
 
     def analyse_structures(
         self, output_dict: dict, output_keys: tuple = OutputElastic.keys()
-    ):
+    ) -> dict:
         """
 
         Args:
@@ -57,19 +56,11 @@ def analyse_structures(
         Returns:
 
         """
-        elastic_matrix, A2, strain_energy, ene0 = analyse_structures_helper(
+        self._data, elastic_dict = analyse_structures_helper(
             output_dict=output_dict,
-            Lag_strain_list=self._data["Lag_strain_list"],
-            epss=self._data["epss"],
-            v0=self._data["v0"],
-            LC=self._data["LC"],
+            sym_dict=self._data,
             fit_order=self.fit_order,
             zero_strain_job_name=self.zero_strain_job_name,
+            output_keys=output_keys,
         )
-        self._data["strain_energy"] = strain_energy
-        self._data["e0"] = ene0
-        self._data["A2"] = A2
-        elastic = ElasticProperties(elastic_matrix=elastic_matrix)
-        return OutputElastic(
-            **{k: getattr(elastic, k) for k in OutputElastic.keys()}
-        ).get(output_keys=output_keys)
+        return elastic_dict

atomistics/workflows/evcurve/helper.py

@@ -0,0 +1,185 @@
+import numpy as np
+from ase.atoms import Atoms
+from typing import Optional, List
+
+from atomistics.shared.output import OutputEnergyVolumeCurve
+from atomistics.workflows.evcurve.fit import EnergyVolumeFit
+
+
+def _strain_axes(
+    structure: Atoms, volume_strain: float, axes: tuple[str, str, str] = ("x", "y", "z")
+) -> Atoms:
+    """
+    Strain box along given axes to achieve given *volumetric* strain.
+
+    Returns a copy.
+    """
+    axes = np.array([a in axes for a in ("x", "y", "z")])
+    num_axes = sum(axes)
+    # formula calculates the strain along each axis to achieve the overall volumetric strain
+    # beware that: (1+e)**x - 1 != e**x
+    strains = axes * ((1 + volume_strain) ** (1.0 / num_axes) - 1)
+    return apply_strain(structure=structure, epsilon=strains, return_box=True)
+
+
+def apply_strain(
+    structure: Atoms, epsilon: float, return_box: bool = False, mode: str = "linear"
+) -> Atoms:
+    """
+    Apply a given strain on the structure. It applies the matrix `F` in the manner:
+
+    ```
+        new_cell = F @ current_cell
+    ```
+
+    Args:
+        epsilon (float/list/ndarray): epsilon matrix. If a single number is set, the same
+            strain is applied in each direction. If a 3-dim vector is set, it will be
+            multiplied by a unit matrix.
+        return_box (bool): whether to return a box. If set to True, only the returned box will
+            have the desired strain and the original box will stay unchanged.
+        mode (str): `linear` or `lagrangian`. If `linear`, `F` is equal to the epsilon - 1.
+            If `lagrangian`, epsilon is given by `(F^T * F - 1) / 2`. It raises an error if
+            the strain is not symmetric (if the shear components are given).
+    """
+    epsilon = np.array([epsilon]).flatten()
+    if len(epsilon) == 3 or len(epsilon) == 1:
+        epsilon = epsilon * np.eye(3)
+    epsilon = epsilon.reshape(3, 3)
+    if epsilon.min() < -1.0:
+        raise ValueError("Strain value too negative")
+    if return_box:
+        structure_copy = structure.copy()
+    else:
+        structure_copy = structure
+    cell = structure_copy.cell.copy()
+    if mode == "linear":
+        F = epsilon + np.eye(3)
+    elif mode == "lagrangian":
+        if not np.allclose(epsilon, epsilon.T):
+            raise ValueError("Strain must be symmetric if `mode = 'lagrangian'`")
+        E, V = np.linalg.eigh(2 * epsilon + np.eye(3))
+        F = np.einsum("ik,k,jk->ij", V, np.sqrt(E), V)
+    else:
+        raise ValueError("mode must be `linear` or `lagrangian`")
+    cell = np.matmul(F, cell)
+    structure_copy.set_cell(cell, scale_atoms=True)
+    if return_box:
+        return structure_copy
+
+
+def get_energy_lst(output_dict: dict, structure_dict: dict) -> list:
+    return [output_dict["energy"][k] for k in structure_dict.keys()]
+
+
+def get_volume_lst(structure_dict: dict) -> list:
+    return [structure.get_volume() for structure in structure_dict.values()]
+
+
+def fit_ev_curve_internal(
+    volume_lst: np.ndarray, energy_lst: np.ndarray, fit_type: str, fit_order: int
+) -> EnergyVolumeFit:
+    fit_module = EnergyVolumeFit(
+        volume_lst=volume_lst,
+        energy_lst=energy_lst,
+    )
+    fit_module.fit(fit_type=fit_type, fit_order=fit_order)
+    return fit_module
+
+
+def fit_ev_curve(
+    volume_lst: np.ndarray, energy_lst: np.ndarray, fit_type: str, fit_order: int
+) -> dict:
+    return fit_ev_curve_internal(
+        volume_lst=volume_lst,
+        energy_lst=energy_lst,
+        fit_type=fit_type,
+        fit_order=fit_order,
+    ).fit_dict
+
+
+def get_strains(
+    vol_range: Optional[float],
+    num_points: Optional[int],
+    strain_lst: Optional[List[float]] = None,
+) -> np.ndarray:
+    if strain_lst is None:
+        return np.linspace(
+            -vol_range,
+            vol_range,
+            int(num_points),
+        )
+    return strain_lst
+
+
+def generate_structures_helper(
+    structure: Atoms,
+    vol_range: Optional[float],
+    num_points: Optional[int],
+    strain_lst: Optional[List[float]] = None,
+    axes: tuple[str, str, str] = ("x", "y", "z"),
+) -> dict:
+    strain_lst = get_strains(
+        vol_range=vol_range, num_points=num_points, strain_lst=strain_lst
+    )
+    return {
+        1
+        + np.round(strain, 7): _strain_axes(
+            structure=structure, axes=axes, volume_strain=strain
+        )
+        for strain in strain_lst
+    }
+
+
+def analyse_structures_helper(
+    output_dict: dict,
+    structure_dict: dict,
+    fit_type: str = "polynomial",
+    fit_order: int = 3,
+    output_keys: tuple = OutputEnergyVolumeCurve.keys(),
+) -> dict:
+    return EnergyVolumeCurveProperties(
+        fit_module=fit_ev_curve_internal(
+            volume_lst=get_volume_lst(structure_dict=structure_dict),
+            energy_lst=get_energy_lst(
+                output_dict=output_dict, structure_dict=structure_dict
+            ),
+            fit_type=fit_type,
+            fit_order=fit_order,
+        )
+    ).to_dict(output_keys=output_keys)
+
+
+class EnergyVolumeCurveProperties:
+    def __init__(self, fit_module):
+        self._fit_module = fit_module
+
+    def volume_eq(self) -> float:
+        return self._fit_module.fit_dict["volume_eq"]
+
+    def energy_eq(self) -> float:
+        return self._fit_module.fit_dict["energy_eq"]
+
+    def bulkmodul_eq(self) -> float:
+        return self._fit_module.fit_dict["bulkmodul_eq"]
+
+    def b_prime_eq(self) -> float:
+        return self._fit_module.fit_dict["b_prime_eq"]
+
+    def volume(self) -> np.ndarray:
+        return self._fit_module.fit_dict["volume"]
+
+    def energy(self) -> np.ndarray:
+        return self._fit_module.fit_dict["energy"]
+
+    def fit_dict(self) -> dict:
+        return {
+            k: self._fit_module.fit_dict[k]
+            for k in ["fit_type", "least_square_error", "poly_fit", "fit_order"]
+            if k in self._fit_module.fit_dict.keys()
+        }
+
+    def to_dict(self, output_keys: tuple = OutputEnergyVolumeCurve.keys()) -> dict:
+        return OutputEnergyVolumeCurve(
+            **{k: getattr(self, k) for k in OutputEnergyVolumeCurve.keys()}
+        ).get(output_keys=output_keys)