fix missing import in examples/LiFePO4-NEB.ipynb plus ruff format

materialsvirtuallab · Jul 28, 2024 · 7ff91a6 · 7ff91a6
1 parent 1a83391
commit 7ff91a6
Show file tree

Hide file tree

Showing 9 changed files with 44 additions and 56 deletions.
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
@@ -6,7 +6,7 @@ ci:
 
 repos:
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.4.8
+    rev: v0.5.5
     hooks:
       - id: ruff
         args: [--fix]
@@ -21,7 +21,7 @@ repos:
       - id: trailing-whitespace
 
   - repo: https://github.com/pre-commit/mirrors-mypy
-    rev: v1.10.0
+    rev: v1.11.0
     hooks:
       - id: mypy
 
@@ -48,6 +48,6 @@ repos:
         args: [--drop-empty-cells, --keep-output]
 
   - repo: https://github.com/RobertCraigie/pyright-python
-    rev: v1.1.366
+    rev: v1.1.373
     hooks:
       - id: pyright
diff --git a/examples/Calculating MLIP properties.ipynb b/examples/Calculating MLIP properties.ipynb
@@ -18,17 +18,20 @@
     }
    ],
    "source": [
+    "from __future__ import annotations\n",
+    "\n",
     "import warnings\n",
-    "from matcalc.relaxation import RelaxCalc\n",
-    "from matcalc.phonon import PhononCalc\n",
-    "from matcalc.eos import EOSCalc\n",
-    "from matcalc.elasticity import ElasticityCalc\n",
-    "from matcalc.utils import get_universal_calculator\n",
-    "from tqdm import tqdm\n",
     "from time import perf_counter\n",
-    "import matplotlib.pyplot as plt\n",
     "\n",
+    "import matplotlib.pyplot as plt\n",
     "from mp_api.client import MPRester\n",
+    "from tqdm import tqdm\n",
+    "\n",
+    "from matcalc.elasticity import ElasticityCalc\n",
+    "from matcalc.eos import EOSCalc\n",
+    "from matcalc.phonon import PhononCalc\n",
+    "from matcalc.relaxation import RelaxCalc\n",
+    "from matcalc.utils import get_universal_calculator\n",
     "\n",
     "warnings.filterwarnings(\"ignore\", category=UserWarning, module=\"matgl\")\n",
     "warnings.filterwarnings(\"ignore\", category=DeprecationWarning, module=\"spglib\")"
@@ -584,8 +587,7 @@
     }
    ],
    "source": [
-    "\n",
-    "fig,axes = plt.subplots(2,2)\n",
+    "fig, axes = plt.subplots(2, 2)\n",
     "axes = axes.flatten()\n",
     "for i, (model_name, model) in enumerate(models):\n",
     "    ax = axes[i]\n",
@@ -595,8 +597,7 @@
     "    ax.set_title(model_name)\n",
     "\n",
     "plt.tight_layout()\n",
-    "plt.show()\n",
-    "\n"
+    "plt.show()\n"
    ]
   },
   {
@@ -617,9 +618,8 @@
     }
    ],
    "source": [
-    "\n",
     "for model_name, _ in models[:2]:\n",
-    "    ax = df_preds[f\"time_total_{model_name}\"].hist(label=model_name,alpha=0.6)\n",
+    "    ax = df_preds[f\"time_total_{model_name}\"].hist(label=model_name, alpha=0.6)\n",
     "\n",
     "ax.set_xlabel(\"Total Time (s)\")\n",
     "ax.set_ylabel(\"Count\")\n",
@@ -644,7 +644,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.10.14"
+   "version": "3.11.7"
   }
  },
  "nbformat": 4,

diff --git a/examples/LiFePO4-NEB.ipynb b/examples/LiFePO4-NEB.ipynb
@@ -29,6 +29,7 @@
     "from __future__ import annotations\n",
     "\n",
     "from pymatgen.ext.matproj import MPRester\n",
+    "from ase.neb import NEB, NEBTools\n",
     "\n",
     "from matcalc.util import get_universal_calculator\n",
     "\n",
@@ -127,7 +128,7 @@
     "results = {}\n",
     "for model_name, model in models.items():\n",
     "    relaxer = RelaxCalc(model, optimizer=\"BFGS\", relax_cell=True, fmax=0.02)\n",
-    "    supercell_LFPO_relaxed = relaxer.calc(s_LFPO)['final_structure']\n",
+    "    supercell_LFPO_relaxed = relaxer.calc(s_LFPO)[\"final_structure\"]\n",
     "    results[model_name] = {\"supercell_LFPO\": supercell_LFPO_relaxed}"
    ]
   },
@@ -163,18 +164,18 @@
     "    # NEB path along b and c directions have the same starting image.\n",
     "    s_LFPO_end_b = supercell_LFPO_relaxed.copy()\n",
     "    s_LFPO_end_b.remove_sites([11])\n",
-    "    s_LFPO_end_b_relaxed = relaxer.calc(s_LFPO_end_b)['final_structure']\n",
+    "    s_LFPO_end_b_relaxed = relaxer.calc(s_LFPO_end_b)[\"final_structure\"]\n",
     "    s_LFPO_end_c = supercell_LFPO_relaxed.copy()\n",
     "    s_LFPO_end_c.remove_sites([4])\n",
-    "    s_LFPO_end_c_relaxed = relaxer.calc(s_LFPO_end_c)['final_structure']\n",
+    "    s_LFPO_end_c_relaxed = relaxer.calc(s_LFPO_end_c)[\"final_structure\"]\n",
     "    s_LFPO_start_bc = supercell_LFPO_relaxed.copy()\n",
     "    s_LFPO_start_bc.remove_sites([5])\n",
-    "    s_LFPO_start_bc_relaxed = relaxer.calc(s_LFPO_start_bc)['final_structure']\n",
+    "    s_LFPO_start_bc_relaxed = relaxer.calc(s_LFPO_start_bc)[\"final_structure\"]\n",
     "    results[model_name].update(\n",
     "        {\n",
-    "            \"supercell_LFPO_end_b\":s_LFPO_end_b_relaxed,\n",
-    "            \"supercell_LFPO_end_c\":s_LFPO_end_c_relaxed,\n",
-    "            \"supercell_LFPO_start_bc\":s_LFPO_start_bc_relaxed,\n",
+    "            \"supercell_LFPO_end_b\": s_LFPO_end_b_relaxed,\n",
+    "            \"supercell_LFPO_end_c\": s_LFPO_end_c_relaxed,\n",
+    "            \"supercell_LFPO_start_bc\": s_LFPO_start_bc_relaxed,\n",
     "        }\n",
     "    )"
    ]
@@ -426,7 +427,7 @@
    ],
    "source": [
     "%%time\n",
-    "for neb_path in 'bc':\n",
+    "for neb_path in \"bc\":\n",
     "    for model_name, model in models.items():\n",
     "        NEBcalc = NEBCalc.from_end_images(\n",
     "            start_struct=results[model_name][\"supercell_LFPO_start_bc\"],\n",
@@ -570,7 +571,7 @@
     "%%time\n",
     "import matplotlib.pyplot as plt\n",
     "\n",
-    "for neb_path in 'bc':\n",
+    "for neb_path in \"bc\":\n",
     "    for model_name, model in models.items():\n",
     "        NEB_tool = NEBTools(results[model_name][f\"NEB_{neb_path}\"].images)\n",
     "        print(f\"Path along {neb_path}, {model_name}: \")\n",
@@ -599,13 +600,12 @@
     "from pymatgen.io.ase import AseAtomsAdaptor\n",
     "\n",
     "\n",
-    "def generate_path_cif_from_images(images: list, filename: str):\n",
+    "def generate_path_cif_from_images(images: list, filename: str) -> None:\n",
     "    \"\"\"Generate a cif file from a list of image atoms.\"\"\"\n",
     "    image_structs = list(map(AseAtomsAdaptor().get_structure, images))\n",
     "    sites = set()\n",
     "    lattice = image_structs[0].lattice\n",
-    "    for site in chain(*(struct for struct in image_structs)):\n",
-    "        sites.add(PeriodicSite(site.species, site.frac_coords, lattice))\n",
+    "    sites.update(PeriodicSite(site.species, site.frac_coords, lattice) for site in chain(*(struct for struct in image_structs)))\n",
     "    neb_path = Structure.from_sites(sorted(sites))\n",
     "    neb_path.to(filename, \"cif\")"
    ]
@@ -627,7 +627,7 @@
    ],
    "source": [
     "%%time\n",
-    "for neb_path in 'bc':\n",
+    "for neb_path in \"bc\":\n",
     "    for model_name, model in models.items():\n",
     "        NEB_tool = NEBTools(results[model_name][f\"NEB_{neb_path}\"].images)\n",
     "        generate_path_cif_from_images(NEB_tool.images, f\"NEB_data/traj_{neb_path}_{model_name}/path_final.cif\")"

diff --git a/matcalc/base.py b/matcalc/base.py
@@ -13,13 +13,12 @@
     from pymatgen.core import Structure
 
 
-class PropCalc(metaclass=abc.ABCMeta):
+class PropCalc(abc.ABC):
     """API for a property calculator."""
 
     @abc.abstractmethod
     def calc(self, structure: Structure) -> dict:
-        """
-        All PropCalc subclasses should implement a calc method that takes in a pymatgen structure
+        """All PropCalc subclasses should implement a calc method that takes in a pymatgen structure
         and returns a dict. The method can return more than one property.
 
         Args:
@@ -32,9 +31,8 @@ def calc(self, structure: Structure) -> dict:
     def calc_many(
         self, structures: Sequence[Structure], n_jobs: None | int = None, **kwargs: Any
     ) -> Generator[dict, None, None]:
-        """
-        Performs calc on many structures. The return type is a generator given that the calc method can potentially be
-        expensive. It is trivial to convert the generator to a list/tuple.
+        """Performs calc on many structures. The return type is a generator given that the calc method can
+        potentially be expensive. It is trivial to convert the generator to a list/tuple.
 
         Args:
             structures: List or generator of Structures.

diff --git a/matcalc/elasticity.py b/matcalc/elasticity.py
@@ -67,8 +67,7 @@ def __init__(
         self.relax_calc_kwargs = relax_calc_kwargs
 
     def calc(self, structure: Structure) -> dict[str, Any]:
-        """
-        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:
@@ -123,8 +122,7 @@ def _elastic_tensor_from_strains(
         eq_stress: ArrayLike = None,
         tol: float = 1e-7,
     ) -> tuple[ElasticTensor, float]:
-        """
-        Slightly modified version of Pymatgen function
+        """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

diff --git a/matcalc/neb.py b/matcalc/neb.py
@@ -70,8 +70,7 @@ def from_end_images(
         autosort_tol: float = 0.5,
         **kwargs: Any,
     ) -> NEBCalc:
-        """
-        Initialize a NEBCalc from end images.
+        """Initialize a NEBCalc from end images.
 
         Args:
             start_struct(Structure): The starting image as a pymatgen Structure.
@@ -98,8 +97,7 @@ def from_end_images(
     def calc(  # type: ignore[override]
         self, fmax: float = 0.1, max_steps: int = 1000
     ) -> tuple[float, float]:
-        """
-        Perform NEB calculation.
+        """Perform NEB calculation.
 
         Args:
             fmax (float): Convergence criteria for NEB calculations defined by Max forces.

diff --git a/matcalc/phonon.py b/matcalc/phonon.py
@@ -80,8 +80,7 @@ def __post_init__(self) -> None:
             setattr(self, key, str({True: default_path, False: ""}.get(val, val)))  # type: ignore[arg-type]
 
     def calc(self, structure: Structure) -> dict:
-        """
-        Calculates thermal properties of Pymatgen structure with phonopy.
+        """Calculates thermal properties of Pymatgen structure with phonopy.
 
         Args:
             structure: Pymatgen structure.
@@ -136,8 +135,7 @@ def calc(self, structure: Structure) -> dict:
 
 
 def _calc_forces(calculator: Calculator, supercell: PhonopyAtoms) -> ArrayLike:
-    """
-    Helper to compute forces on a structure.
+    """Helper to compute forces on a structure.
 
     Args:
         calculator: ASE Calculator

diff --git a/matcalc/relaxation.py b/matcalc/relaxation.py
@@ -29,8 +29,7 @@ class TrajectoryObserver:
     """
 
     def __init__(self, atoms: Atoms) -> None:
-        """
-        Init the Trajectory Observer from a Atoms.
+        """Init the Trajectory Observer from a Atoms.
 
         Args:
             atoms (Atoms): Structure to observe.
@@ -83,8 +82,7 @@ def __init__(
         relax_cell: bool = True,
         cell_filter: Filter = FrechetCellFilter,
     ) -> None:
-        """
-        Args:
+        """Args:
             calculator: ASE Calculator to use.
             optimizer (str | ase Optimizer): The optimization algorithm. Defaults to "FIRE".
             max_steps (int): Max number of steps for relaxation. Defaults to 500.
@@ -109,8 +107,7 @@ def __init__(
         self.cell_filter = cell_filter
 
     def calc(self, structure: Structure) -> dict:
-        """
-        Perform relaxation to obtain properties.
+        """Perform relaxation to obtain properties.
 
         Args:
             structure: Pymatgen structure.

diff --git a/matcalc/utils.py b/matcalc/utils.py
@@ -25,8 +25,7 @@
 
 @functools.lru_cache
 def get_universal_calculator(name: str | Calculator, **kwargs: Any) -> Calculator:
-    """
-    Helper method to get some well-known **universal** calculators.
+    """Helper method to get some well-known **universal** calculators.
     Imports should be inside if statements to ensure that all models are optional dependencies.
     All calculators must be universal, i.e. encompass a wide swath of the periodic table.
     Though matcalc can be used with any MLIP, even custom ones, this function is not meant as