emmet 0.86.1 compatibility

.github/workflows/deploy.yml

@@ -21,7 +21,7 @@ jobs:
 
       - uses: actions/setup-python@v5
         with:
-          python-version: "3.10"
+          python-version: "3.11"
           cache: pip
           cache-dependency-path: pyproject.toml
 
@@ -51,7 +51,7 @@ jobs:
 
       - uses: actions/setup-python@v5
         with:
-          python-version: "3.10"
+          python-version: "3.11"
 
       - name: Build
         run: |

.github/workflows/docs.yml

@@ -26,7 +26,7 @@ jobs:
 
       - uses: actions/setup-python@v5
         with:
-          python-version: "3.10"
+          python-version: "3.11"
           cache: pip
           cache-dependency-path: pyproject.toml
 

.github/workflows/testing.yml

@@ -17,7 +17,7 @@ jobs:
 
       - uses: actions/setup-python@v5
         with:
-          python-version: "3.10"
+          python-version: "3.11"
           cache: pip
           cache-dependency-path: pyproject.toml
 
@@ -39,7 +39,7 @@ jobs:
         shell: bash -l {0} # enables conda/mamba env activation by reading bash profile
     strategy:
       matrix:
-        python-version: ["3.10", "3.11", "3.12"]
+        python-version: ["3.11", "3.12"]
         split: [1, 2, 3]
 
     steps:
@@ -48,10 +48,11 @@ jobs:
 
       - name: Set up micromamba
         uses: mamba-org/setup-micromamba@main
-
-      - name: Create mamba environment
-        run: |
-          micromamba create -n a2 python=${{ matrix.python-version }} --yes
+        with:
+          environment-name: a2
+          cache-environment: false
+          create-args: >-
+            python=${{ matrix.python-version }}
 
       - name: Install uv
         run: micromamba run -n a2 pip install uv
@@ -91,7 +92,7 @@ jobs:
 
 
       - uses: codecov/codecov-action@v1
-        if: matrix.python-version == '3.10' && github.repository == 'materialsproject/atomate2'
+        if: matrix.python-version == '3.11' && github.repository == 'materialsproject/atomate2'
         with:
           token: ${{ secrets.CODECOV_TOKEN }}
           name: coverage${{ matrix.split }}
@@ -113,18 +114,19 @@ jobs:
         shell: bash -l {0} # enables conda/mamba env activation by reading bash profile
     strategy:
       matrix:
-        python-version: ["3.10","3.11","3.12"]
+        python-version: ["3.11","3.12"]
 
     steps:
       - name: Check out repo
         uses: actions/checkout@v4
 
       - name: Set up micromamba
         uses: mamba-org/setup-micromamba@main
-
-      - name: Create mamba environment
-        run: |
-          micromamba create -n a2 python=${{ matrix.python-version }} --yes
+        with:
+          environment-name: a2
+          cache-environment: false
+          create-args: >-
+            python=${{ matrix.python-version }}
 
       - name: Install uv
         run: micromamba run -n a2 pip install uv
@@ -174,18 +176,19 @@ jobs:
         shell: bash -l {0} # enables conda/mamba env activation by reading bash profile
     strategy:
       matrix:
-        python-version: ["3.10", "3.11", "3.12"]
+        python-version: ["3.11", "3.12"]
 
     steps:
       - name: Check out repo
         uses: actions/checkout@v4
 
       - name: Set up micromamba
         uses: mamba-org/setup-micromamba@main
-
-      - name: Create mamba environment
-        run: |
-          micromamba create -n a2 python=${{ matrix.python-version }} --yes
+        with:
+          environment-name: a2
+          cache-environment: false
+          create-args: >-
+            python=${{ matrix.python-version }}
 
       - name: Install uv
         run: micromamba run -n a2 pip install uv
@@ -220,7 +223,7 @@ jobs:
           pytest -n auto --splits 1 --group 1 --cov=atomate2 --cov-report=xml tests/ase
 
       - uses: codecov/codecov-action@v1
-        if: matrix.python-version == '3.10' && github.repository == 'materialsproject/atomate2'
+        if: matrix.python-version == '3.11' && github.repository == 'materialsproject/atomate2'
         with:
           token: ${{ secrets.CODECOV_TOKEN }}
           file: ./coverage.xml
@@ -241,18 +244,19 @@ jobs:
         shell: bash -l {0} # enables conda/mamba env activation by reading bash profile
     strategy:
       matrix:
-        python-version: ["3.10", "3.11", "3.12"]
+        python-version: ["3.11", "3.12"]
 
     steps:
       - name: Check out repo
         uses: actions/checkout@v4
 
       - name: Set up micromamba
         uses: mamba-org/setup-micromamba@main
-
-      - name: Create mamba environment
-        run: |
-          micromamba create -n a2 python=${{ matrix.python-version }} --yes
+        with:
+          environment-name: a2
+          cache-environment: false
+          create-args: >-
+            python=${{ matrix.python-version }}
 
       - name: Install uv
         run: micromamba run -n a2 pip install uv
@@ -292,7 +296,7 @@ jobs:
 
 
       - uses: codecov/codecov-action@v1
-        if: matrix.python-version == '3.10' && github.repository == 'materialsproject/atomate2'
+        if: matrix.python-version == '3.11' && github.repository == 'materialsproject/atomate2'
         with:
           token: ${{ secrets.CODECOV_TOKEN }}
           name: coverage
@@ -308,7 +312,7 @@ jobs:
 
       - uses: actions/setup-python@v5
         with:
-          python-version: "3.10"
+          python-version: "3.11"
           cache: pip
           cache-dependency-path: pyproject.toml
 

.github/workflows/update-precommit.yml

@@ -15,7 +15,7 @@ jobs:
       - name: Set up Python
         uses: actions/setup-python@v5
         with:
-          python-version: "3.10"
+          python-version: "3.11"
 
       - name: Install pre-commit
         run: pip install pre-commit

.gitignore

@@ -8,6 +8,9 @@ docs_build/*
 docs/tutorials/*
 docs/tutorials/*/*
 
+# tutorial output files
+tutorials/*.json.gz
+
 # C extensions
 *.so
 

docs/user/install.md

@@ -354,19 +354,34 @@ and then return to this tutorial.
 ### Materials Project API key
 
 You can get an API key from the [Materials Project] by logging in and going to your
-[Dashboard](materials project). Add this also to
-your `~/.config/.pmgrc.yaml` so that it looks like the following
+[Dashboard](materials project).
 
-```yaml
-PMG_VASP_PSP_DIR: <<INSTALL_DIR>>/pps
-PMG_MAPI_KEY: <<YOUR_API_KEY>>
+Add this to your `.bashrc`, `.zshrc`, etc. as an environment variable:
+```console
+export MP_API_KEY=your_api_key_here
 ```
 
-You can generate this file and set these values using the `pymatgen` CLI:
+## Emmet-core setup
 
-```bash
-pmg config --add PMG_VASP_PSP_DIR /abs/path/to/psp PMG_MAPI_KEY your_api_key
+The `emmet-core` package is used to define data schemas for parsing outputs of workflows.
+It is also used in building the Materials Project data, therefore its use in `atomate2` is to broadly ensure compatibility with the Materials Project's data structures.
+`emmet-core` allows you to use either `pymatgen` or `emmet-core`-defined models for larger data objects, such as charge densities (`CHGCAR`, `AECCAR*`), or trajectories (relaxation, MD, etc.).
+The `pymatgen` objects have long been the default in workflows, and are structured to be output as JSON files.
+The `emmet-core` objects have been designed with both JSON and Apache parquet as storage formats.
+
+If you will be storing data in the cloud, or would like to use these newer data models which may use less storage, you can either add a line to your `atomate2.yaml` file:
+```yaml
+VASP_USE_EMMET_MODELS: true
 ```
+or use an environment variable in your `.bashrc`:
+```console
+export ATOMATE2_VASP_USE_EMMET_MODEL=true
+```
+Note that there is an equivalent `emmet-core` setting, which can be set by the environment variable `EMMET_USE_EMMET_MODELS`.
+
+For ASE and machine learning forcefield jobs, you can use the `ASE_FORCEFIELD_USE_EMMET_MODELS` flag in `atomate2.yaml` to toggle the same functionality.
+
+The default in `atomate2` is to use `pymatgen` models.
 
 [materials project]: https://materialsproject.org/dashboard
 

pyproject.toml

@@ -28,7 +28,7 @@ dependencies = [
     "PyYAML",
     "click",
     "custodian>=2024.4.18",
-    "emmet-core>=v0.84.10",
+    "emmet-core>=0.86.1",
     "jobflow>=0.1.11",
     "monty>=2024.12.10",
     "numpy",
@@ -57,11 +57,10 @@ forcefields = [
     "chgnet>=0.2.2",
     "mace-torch>=0.3.3",
     "matgl>=1.2.1",
-    "torchdata<=0.7.1",
+    "torchdata<=0.7.1", # TODO: remove when issue fixed
     # quippy-ase support for py3.12 tracked in https://github.com/libAtoms/QUIP/issues/645
     "quippy-ase>=0.9.14; python_version < '3.12'",
     "sevenn>=0.9.3",
-    "torchdata<=0.7.1",                            # TODO: remove when issue fixed
 ]
 approxneb = ["pymatgen-analysis-diffusion>=2024.7.15"]
 ase = ["ase>=3.26.0"]
@@ -107,14 +106,15 @@ strict-forcefields = [
     "calorine==3.1; python_version < '3.12'",
     "chgnet==0.3.8",
     "mace-torch==0.3.14",
-    "matgl==1.3.0",
+    "matgl==2.0.2",
     "quippy-ase==0.9.14; python_version < '3.12'",
     "sevenn==0.10.4",
     "torch==2.2.0",
     "torchdata==0.7.1",                            # TODO: remove when issue fixed
 ]
 strict = [
     "atomate2[strict-forcefields, docs, cclib, phonons, lobster, openmm, mp, defects, ase, ase-ext]",
+    "numpy<2.0",
 ]
 
 [project.scripts]

src/atomate2/abinit/schemas/calculation.py

@@ -11,11 +11,7 @@
 from abipy.flowtk import events
 from abipy.flowtk.utils import File
 from emmet.core.math import Matrix3D, Vector3D
-
-try:
-    from emmet.core.types.enums import ValueEnum
-except ImportError:
-    from emmet.core.utils import ValueEnum
+from emmet.core.types.enums import ValueEnum
 from pydantic import BaseModel, Field
 from pymatgen.core import Structure
 from typing_extensions import Self

src/atomate2/aims/schemas/calculation.py

@@ -12,11 +12,7 @@
 import numpy as np
 from ase.spectrum.band_structure import BandStructure
 from emmet.core.math import Matrix3D, Vector3D
-
-try:
-    from emmet.core.types.enums import ValueEnum
-except ImportError:
-    from emmet.core.utils import ValueEnum
+from emmet.core.types.enums import ValueEnum
 from pydantic import BaseModel, Field
 from pymatgen.core import Molecule, Structure
 from pymatgen.core.trajectory import Trajectory

src/atomate2/ase/jobs.py

@@ -8,15 +8,9 @@
 from dataclasses import dataclass, field
 from typing import TYPE_CHECKING
 
-from ase.io import Trajectory as AseTrajectory
-
-try:
-    from emmet.core.types.enums import StoreTrajectoryOption
-except ImportError:
-    from emmet.core.vasp.calculation import StoreTrajectoryOption
+from emmet.core.types.enums import StoreTrajectoryOption
 from jobflow import Maker, job
 from pymatgen.core import Molecule, Structure
-from pymatgen.core.trajectory import Trajectory as PmgTrajectory
 from pymatgen.io.ase import AseAtomsAdaptor
 
 from atomate2.ase.schemas import AseResult, AseTaskDoc
@@ -31,7 +25,7 @@
 
     from atomate2.ase.schemas import AseMoleculeTaskDoc, AseStructureTaskDoc
 
-_ASE_DATA_OBJECTS = [PmgTrajectory, AseTrajectory]
+_ASE_DATA_OBJECTS = ["trajectory"]
 
 
 @dataclass

src/atomate2/ase/md.py

@@ -24,16 +24,13 @@
     Stationary,
     ZeroRotation,
 )
-
-try:
-    from emmet.core.types.enums import StoreTrajectoryOption
-except ImportError:
-    from emmet.core.vasp.calculation import StoreTrajectoryOption
+from emmet.core.types.enums import StoreTrajectoryOption
 from jobflow import job
 from pymatgen.core.structure import Molecule, Structure
 from pymatgen.io.ase import AseAtomsAdaptor
 from scipy.interpolate import interp1d
 
+from atomate2 import SETTINGS
 from atomate2.ase.jobs import _ASE_DATA_OBJECTS, AseMaker
 from atomate2.ase.schemas import AseResult, AseTaskDoc
 from atomate2.ase.utils import TrajectoryObserver
@@ -165,6 +162,9 @@ class AseMDMaker(AseMaker, ABC):
         Whether to initialize the atomic velocities with zero angular momentum
     verbose : bool = False
         Whether to print stdout to screen during the MD run.
+    use_emmet_models : bool = False
+        Whether to use emmet-core (True) or pymatgen (False)
+        data models for larger objects, e.g., trajectories.
     """
 
     name: str = "ASE MD"
@@ -185,6 +185,7 @@ class AseMDMaker(AseMaker, ABC):
     zero_linear_momentum: bool = False
     zero_angular_momentum: bool = False
     verbose: bool = False
+    use_emmet_models: bool = SETTINGS.ASE_FORCEFIELD_USE_EMMET_MODELS
 
     def __post_init__(self) -> None:
         """Ensure that ensemble is an enum."""
@@ -433,7 +434,12 @@ def _callback(dyn: MolecularDynamics = md_runner) -> None:
 
         return AseResult(
             final_mol_or_struct=mol_or_struct,
-            trajectory=md_observer.to_emmet_trajectory(filename=None),
+            trajectory=getattr(
+                md_observer,
+                "to_emmet_trajectory"
+                if self.use_emmet_models
+                else "to_pymatgen_trajectory",
+            )(filename=None),
             dir_name=os.getcwd(),
             elapsed_time=t_f - t_i,
         )