Add code for scoring lattice strain

mofa/model.py

@@ -81,6 +81,13 @@ class MOFRecord:
     structure: str | None = field(default=None, repr=False)
     """A representative 3D structure of the MOF in POSCAR format"""
 
+    # Detailed outputs from simulations
+    md_trajectory: dict[str, list[str]] = field(default_factory=dict, repr=False)
+    """Structures of the molecule produced during molecular dynamics simulations.
+
+    Key is the name of the level of accuracy for the MD computational (e.g., "uff"),
+    values are the structure in POSCAR format"""
+
     # Properties
     gas_storage: dict[tuple[str, float], float] = field(default_factory=dict, repr=False)
     """Storage capacity of the MOF for different gases and pressures"""

mofa/scoring/geometry.py

@@ -1,8 +1,13 @@
 """Metrics for screening a MOF or linker based on its geometry"""
+from io import StringIO
+
 import numpy as np
 import ase
+from ase.io.vasp import read_vasp
 
-from mofa.scoring.base import MOFScorer
+from mofa.model import MOFRecord
+from dataclasses import dataclass
+from mofa.scoring.base import MOFScorer, Scorer
 
 
 class MinimumDistance(MOFScorer):
@@ -12,3 +17,28 @@ def __call__(self, linker: ase.Atoms) -> float:
         dists: np.ndarray = linker.get_all_distances(mic=True)
         inds = np.triu_indices(dists.shape[0], k=1)
         return np.min(dists[inds])
+
+
+@dataclass
+class LatticeParameterChange(Scorer):
+    """Score the stability of a MOF based on the lattice parameter change"""
+
+    md_level: str = 'uff'
+    """Level of accuracy used for the molecular dynamics simulation"""
+
+    def score_mof(self, record: MOFRecord) -> float:
+        # Get the trajectory from the record
+        if self.md_level not in record.md_trajectory:
+            raise ValueError(f'No data available for MD simulations at level: "{self.md_level}"')
+        traj = record.md_trajectory[self.md_level]
+
+        # Get the initial and final structures
+        init_strc = read_vasp(StringIO(traj[0]))
+        final_strc: ase.Atoms = read_vasp(StringIO(traj[1]))
+
+        # Compute the maximum principal strain
+        #  Following: https://www.cryst.ehu.es/cgi-bin/cryst/programs/nph-strain
+        strain = np.matmul(init_strc.cell.array, np.linalg.inv(final_strc.cell.array)) - np.eye(3)
+        strain = 0.5 * (strain + strain.T)
+        strains = np.linalg.eigvals(strain)
+        return np.abs(strains).max()

tests/scoring/test_geometry.py

@@ -1,6 +1,10 @@
+from ase.io.vasp import write_vasp
+from pytest import raises
 from ase import Atoms
+import numpy as np
+from six import StringIO
 
-from mofa.scoring.geometry import MinimumDistance
+from mofa.scoring.geometry import MinimumDistance, LatticeParameterChange
 
 
 def test_distance(example_record):
@@ -10,3 +14,25 @@ def test_distance(example_record):
 
     assert MinimumDistance()(atoms) == 1.
     assert MinimumDistance().score_mof(example_record) > 0
+
+
+def test_strain(example_record):
+    scorer = LatticeParameterChange()
+
+    # Ensure it throws an error
+    with raises(ValueError):
+        scorer.score_mof(example_record)
+
+    # Make a fake MD trajectory with no change
+    example_record.md_trajectory['uff'] = [example_record.structure] * 2
+    assert np.isclose(scorer.score_mof(example_record), 0)
+
+    # Make sure it compute stresses correctly if the volume shears
+    final_atoms = example_record.atoms
+    final_atoms.set_cell(final_atoms.cell.lengths().tolist() + [80, 90, 90])
+    sio = StringIO()
+    write_vasp(sio, final_atoms)
+    example_record.md_trajectory['uff'][1] = sio.getvalue()
+
+    max_strain = scorer.score_mof(example_record)
+    assert np.isclose(max_strain, 0.09647)  # Checked against https://www.cryst.ehu.es/cryst/strain.html