Merge pull request #12 from materialsvirtuallab/eos-and-elasticity-ca…

…lculator-improvements

Improvements to EOS and Elasticity calculators

.github/workflows/testing.yml

@@ -38,7 +38,7 @@ jobs:
         run: |
           pip install numpy
           pip install --quiet -r requirements.txt -r requirements-ci.txt
-          pip install -e .
+          pip install -e '.[models]'
       - name: pytest
         run: |
           pytest --cov=matcalc tests --color=yes

.pre-commit-config.yaml

@@ -6,7 +6,7 @@ ci:
 
 repos:
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.1.3
+    rev: v0.1.6
     hooks:
       - id: ruff
         args: [--fix]
@@ -21,7 +21,7 @@ repos:
       - id: trailing-whitespace
 
   - repo: https://github.com/pre-commit/mirrors-mypy
-    rev: v1.6.1
+    rev: v1.7.1
     hooks:
       - id: mypy
 
@@ -30,11 +30,12 @@ repos:
     hooks:
       - id: codespell
         stages: [commit, commit-msg]
-        exclude_types: [html]
+        exclude_types: [html, svg, javascript]
+        exclude: ^.+\.lock$ # ignore lock files
         additional_dependencies: [tomli] # needed to read pyproject.toml below py3.11
 
   - repo: https://github.com/MarcoGorelli/cython-lint
-    rev: v0.15.0
+    rev: v0.16.0
     hooks:
       - id: cython-lint
         args: [--no-pycodestyle]

docs/conf.py

@@ -215,7 +215,7 @@
 
 # -- Options for LaTeX output ------------------------------------------------
 
-latex_elements = {
+latex_elements: dict[str, str] = {
     # The paper size ('letterpaper' or 'a4paper').
     # 'papersize': 'letterpaper',
     # The font size ('10pt', '11pt' or '12pt').
@@ -317,7 +317,7 @@
 # The format is a list of tuples containing the path and title.
 # epub_pre_files = []
 
-# HTML files shat should be inserted after the pages created by sphinx.
+# HTML files that should be inserted after the pages created by sphinx.
 # The format is a list of tuples containing the path and title.
 # epub_post_files = []
 

matcalc/elasticity.py

@@ -1,11 +1,12 @@
-"""Calculator for phonon properties."""
+"""Calculator for elastic properties."""
 
 from __future__ import annotations
 
 from typing import TYPE_CHECKING
 
+import numpy as np
 from pymatgen.analysis.elasticity import DeformedStructureSet, ElasticTensor, Strain
-from pymatgen.io.ase import AseAtomsAdaptor
+from pymatgen.analysis.elasticity.elastic import get_strain_state_dict
 
 from .base import PropCalc
 from .relaxation import RelaxCalc
@@ -21,29 +22,46 @@ class ElasticityCalc(PropCalc):
     def __init__(
         self,
         calculator: Calculator,
-        norm_strains: float = 0.01,
-        shear_strains: float = 0.01,
+        norm_strains: tuple[float, ...] | float = (-0.01, -0.005, 0.005, 0.01),
+        shear_strains: tuple[float, ...] | float = (-0.06, -0.03, 0.03, 0.06),
         fmax: float = 0.1,
         relax_structure: bool = True,
+        use_equilibrium: bool = True,
     ) -> None:
         """
         Args:
             calculator: ASE Calculator to use.
-
-            fmax: maximum force in the relaxed structure (if relax_structure).
-            norm_strains: strain value to apply to each normal mode.
-            shear_strains: strain value to apply to each shear mode.
+            norm_strains: single or multiple strain values to apply to each normal mode.
+                Defaults to (-0.01, -0.005, 0.005, 0.01).
+            shear_strains: single or multiple strain values to apply to each shear mode.
+                Defaults to (-0.06, -0.03, 0.03, 0.06).
+            fmax: maximum force in the relaxed structure (if relax_structure). Defaults to 0.1.
             relax_structure: whether to relax the provided structure with the given calculator.
+                Defaults to True.
+            use_equilibrium: whether to use the equilibrium stress and strain. Ignored and set
+                to True if either norm_strains or shear_strains has length 1 or is a float.
+                Defaults to True.
         """
         self.calculator = calculator
-        self.norm_strains = norm_strains
-        self.shear_strains = shear_strains
+        self.norm_strains = tuple(np.array([1]) * np.asarray(norm_strains))
+        self.shear_strains = tuple(np.array([1]) * np.asarray(shear_strains))
+        if len(self.norm_strains) == 0:
+            raise ValueError("norm_strains must be nonempty")
+        if len(self.shear_strains) == 0:
+            raise ValueError("shear_strains must be nonempty")
+        if 0 in self.norm_strains or 0 in self.shear_strains:
+            raise ValueError("Strains must be nonzero")
         self.relax_structure = relax_structure
         self.fmax = fmax
+        if len(self.norm_strains) > 1 and len(self.shear_strains) > 1:
+            self.use_equilibrium = use_equilibrium
+        else:
+            self.use_equilibrium = True
 
-    def calc(self, structure: Structure) -> dict:
+    def calc(self, structure: Structure) -> dict[str, float | ElasticTensor | Structure]:
         """
-        Calculates elastic properties of Pymatgen structure with units determined by the calculator.
+        Calculates elastic properties of Pymatgen structure with units determined by the calculator,
+        (often the stress_weight).
 
         Args:
             structure: Pymatgen structure.
@@ -53,33 +71,69 @@ def calc(self, structure: Structure) -> dict:
             shear_modulus_vrh: Voigt-Reuss-Hill shear modulus based on elastic tensor (in eV/A^3),
             bulk_modulus_vrh: Voigt-Reuss-Hill bulk modulus based on elastic tensor (in eV/A^3),
             youngs_modulus: Young's modulus based on elastic tensor (in eV/A^3),
+            residuals_sum: Sum of squares of all residuals in the linear fits of the
+            calculation of the elastic tensor,
+            structure: The equilibrium structure used for the computation.
         }
         """
         if self.relax_structure:
-            rcalc = RelaxCalc(self.calculator, fmax=self.fmax)
-            structure = rcalc.calc(structure)["final_structure"]
+            relax_calc = RelaxCalc(self.calculator, fmax=self.fmax)
+            structure = relax_calc.calc(structure)["final_structure"]
 
-        adaptor = AseAtomsAdaptor()
         deformed_structure_set = DeformedStructureSet(
             structure,
-            [self.norm_strains],
-            [self.shear_strains],
+            self.norm_strains,
+            self.shear_strains,
         )
         stresses = []
         for deformed_structure in deformed_structure_set:
-            atoms = adaptor.get_atoms(deformed_structure)
+            atoms = deformed_structure.to_ase_atoms()
             atoms.calc = self.calculator
             stresses.append(atoms.get_stress(voigt=False))
 
         strains = [Strain.from_deformation(deformation) for deformation in deformed_structure_set.deformations]
-        atoms = adaptor.get_atoms(structure)
+        atoms = structure.to_ase_atoms()
         atoms.calc = self.calculator
-        elastic_tensor = ElasticTensor.from_independent_strains(
-            strains, stresses, eq_stress=atoms.get_stress(voigt=False)
+        elastic_tensor, residuals_sum = self._elastic_tensor_from_strains(
+            strains,
+            stresses,
+            eq_stress=atoms.get_stress(voigt=False) if self.use_equilibrium else None,
         )
         return {
             "elastic_tensor": elastic_tensor,
             "shear_modulus_vrh": elastic_tensor.g_vrh,
             "bulk_modulus_vrh": elastic_tensor.k_vrh,
             "youngs_modulus": elastic_tensor.y_mod,
+            "residuals_sum": residuals_sum,
+            "structure": structure,
         }
+
+    def _elastic_tensor_from_strains(
+        self,
+        strains,
+        stresses,
+        eq_stress=None,
+        tol: float = 1e-7,
+    ):
+        """
+        Slightly modified version of Pymatgen function
+        pymatgen.analysis.elasticity.elastic.ElasticTensor.from_independent_strains;
+        this is to give option to discard eq_stress,
+        which (if the structure is relaxed) tends to sometimes be
+        much lower than neighboring points.
+        Also has option to return the sum of the squares of the residuals
+        for all of the linear fits done to compute the entries of the tensor.
+        """
+        strain_states = [tuple(ss) for ss in np.eye(6)]
+        ss_dict = get_strain_state_dict(strains, stresses, eq_stress=eq_stress, add_eq=self.use_equilibrium)
+        c_ij = np.zeros((6, 6))
+        residuals_sum = 0
+        for ii in range(6):
+            strain = ss_dict[strain_states[ii]]["strains"]
+            stress = ss_dict[strain_states[ii]]["stresses"]
+            for jj in range(6):
+                fit = np.polyfit(strain[:, ii], stress[:, jj], 1, full=True)
+                c_ij[ii, jj] = fit[0][0]
+                residuals_sum += fit[1][0] if len(fit[1]) > 0 else 0.0
+        elastic_tensor = ElasticTensor.from_voigt(c_ij)
+        return elastic_tensor.zeroed(tol), residuals_sum

matcalc/eos.py

@@ -6,6 +6,7 @@
 
 import numpy as np
 from pymatgen.analysis.eos import BirchMurnaghan
+from sklearn.metrics import r2_score
 
 from .base import PropCalc
 from .relaxation import RelaxCalc
@@ -56,10 +57,12 @@ def calc(self, structure: Structure) -> dict:
 
         Returns: {
             eos: {
-                volumes: list[float] in Angstrom^3,
-                energies: list[float] in eV,
+                volumes: tuple[float] in Angstrom^3,
+                energies: tuple[float] in eV,
             },
             bulk_modulus_bm: Birch-Murnaghan bulk modulus in GPa.
+            r2_score_bm: R squared of Birch-Murnaghan fit of energies predicted by model to help detect erroneous
+            calculations. This value should be at least around 1 - 1e-4 to 1 - 1e-5.
         }
         """
         if self.relax_structure:
@@ -70,16 +73,35 @@ def calc(self, structure: Structure) -> dict:
         relaxer = RelaxCalc(
             self.calculator, optimizer=self.optimizer, fmax=self.fmax, max_steps=self.max_steps, relax_cell=False
         )
-        for idx in np.linspace(-self.max_abs_strain, self.max_abs_strain, self.n_points):
-            structure_strained = structure.copy()
-            structure_strained.apply_strain([idx, idx, idx])
+
+        temp_structure = structure.copy()
+        for idx in np.linspace(-self.max_abs_strain, self.max_abs_strain, self.n_points)[self.n_points // 2 :]:
+            structure_strained = temp_structure.copy()
+            structure_strained.apply_strain(
+                (((1 + idx) ** 3 * structure.volume) / (structure_strained.volume)) ** (1 / 3) - 1
+            )
             result = relaxer.calc(structure_strained)
             volumes.append(result["final_structure"].volume)
             energies.append(result["energy"])
+            temp_structure = result["final_structure"]
+
+        for idx in np.flip(np.linspace(-self.max_abs_strain, self.max_abs_strain, self.n_points)[: self.n_points // 2]):
+            structure_strained = temp_structure.copy()
+            structure_strained.apply_strain(
+                (((1 + idx) ** 3 * structure.volume) / (structure_strained.volume)) ** (1 / 3) - 1
+            )
+            result = relaxer.calc(structure_strained)
+            volumes.append(result["final_structure"].volume)
+            energies.append(result["energy"])
+            temp_structure = result["final_structure"]
+
         bm = BirchMurnaghan(volumes=volumes, energies=energies)
         bm.fit()
 
+        volumes, energies = map(list, zip(*sorted(zip(volumes, energies), key=lambda i: i[0])))
+
         return {
             "eos": {"volumes": volumes, "energies": energies},
             "bulk_modulus_bm": bm.b0_GPa,
+            "r2_score_bm": r2_score(energies, bm.func(volumes)),
         }

pyproject.toml

@@ -5,7 +5,6 @@ build-backend = "setuptools.build_meta"
 [project]
 name = "matcalc"
 authors = [
-
     { name = "Eliott Liu", email = "[email protected]" },
     { name = "Janosh Riebesell", email = "[email protected]" },
     { name = "Ji Qi", email = "[email protected]" },
@@ -97,11 +96,11 @@ ignore = [
 ]
 pydocstyle.convention = "google"
 isort.required-imports = ["from __future__ import annotations"]
-extend-exclude = ["tests"]
 
 [tool.ruff.per-file-ignores]
 "__init__.py" = ["F401"]
 "tasks.py" = ["D"]
+"tests/*" = ["D"]
 
 [tool.pytest.ini_options]
 addopts = "--durations=30 --quiet -rXs --color=yes -p no:warnings"

requirements-ci.txt

@@ -5,5 +5,5 @@ coveralls
 mypy
 ruff
 black
-matgl==0.9.0
+matgl==0.9.1
 chgnet==0.2.0

requirements.txt

@@ -1,4 +1,5 @@
-pymatgen==2023.7.20
+pymatgen==2023.11.12
 ase==3.22.1
 joblib==1.3.1
 phonopy==2.20.0
+scikit-learn>=1.3.0

tests/conftest.py

@@ -10,26 +10,30 @@
 """
 from __future__ import annotations
 
-import pytest
 import matgl
-
-from pymatgen.util.testing import PymatgenTest
+import pytest
 from matgl.ext.ase import M3GNetCalculator
+from pymatgen.util.testing import PymatgenTest
 
 matgl.clear_cache(confirm=False)
 
 
 @pytest.fixture(scope="session")
 def LiFePO4():
+    """LiFePO4 structure as session-scoped fixture (don't modify in-place,
+    will affect other tests).
+    """
     return PymatgenTest.get_structure("LiFePO4")
 
 
 @pytest.fixture(scope="session")
 def Li2O():
+    """Li2O structure as session-scoped fixture."""
     return PymatgenTest.get_structure("Li2O")
 
 
 @pytest.fixture(scope="session")
 def M3GNetCalc():
+    """M3GNet calculator as session-scoped fixture."""
     potential = matgl.load_model("M3GNet-MP-2021.2.8-PES")
     return M3GNetCalculator(potential=potential, stress_weight=0.01)