diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 57663ebd75..463f628f28 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -2,7 +2,7 @@ default_language_version: python: python3 repos: - repo: https://github.com/charliermarsh/ruff-pre-commit - rev: v0.1.11 + rev: v0.1.13 hooks: - id: ruff args: [--fix] @@ -35,7 +35,6 @@ repos: files: ^src/ additional_dependencies: - tokenize-rt==4.1.0 - - types-pkg_resources==0.1.2 - types-paramiko - repo: https://github.com/codespell-project/codespell rev: v2.2.6 diff --git a/README.md b/README.md index b2dc2e7b48..175263c0a7 100644 --- a/README.md +++ b/README.md @@ -53,7 +53,7 @@ calculations will be performed: 4. A non-self-consistent calculation on a high symmetry k-point path (for the line mode band structure). -```python +```py from atomate2.vasp.flows.core import RelaxBandStructureMaker from jobflow import run_locally from pymatgen.core import Structure diff --git a/docs/dev/vasp_tests.md b/docs/dev/vasp_tests.md index c7272a5746..2b8db62d44 100644 --- a/docs/dev/vasp_tests.md +++ b/docs/dev/vasp_tests.md @@ -59,7 +59,7 @@ The script should also contain some additional code that will allow `atm dev vasp-test-data` to process the reference data. Below we give an example used to generate the elastic constant workflow test data. -```python +```py from atomate2.vasp.flows.elastic import ElasticMaker from atomate2.vasp.powerups import update_user_kpoints_settings from pymatgen.core import Structure @@ -161,7 +161,7 @@ The most important part is the section that mocks VASP and configures which chec to perform on the input files. For the elastic constant workflow, it looks something like this: -```python +```py # mapping from job name to directory containing test files ref_paths = { "elastic relax 1/6": "Si_elastic/elastic_relax_1_6", @@ -214,7 +214,7 @@ the k-point density during the test. Finally, you should add `assert` statements to validate the workflow outputs. As an example, the full elastic workflow test is reproduced below. -```python +```py def test_elastic(mock_vasp, clean_dir): import numpy as np from jobflow import run_locally diff --git a/docs/dev/workflow_tutorial.md b/docs/dev/workflow_tutorial.md index 3f6f8264ae..8c7adb266f 100644 --- a/docs/dev/workflow_tutorial.md +++ b/docs/dev/workflow_tutorial.md @@ -6,7 +6,7 @@ Every `atomate2` workflow is an instance of jobflow's `Flow ` class, which is a In the context of computational materials science, `Flow ` objects are most easily created by a `Maker`, which contains a factory method make() that produces a `Flow `, given certain inputs. Typically, the input to `Maker`.make() includes atomic coordinate information in the form of a `pymatgen` `Structure` or `Molecule` object. So the basic signature looks like this: -```python +```py class ExampleMaker(Maker): def make(self, coordinates: Structure) -> Flow: # take the input coordinates and return a `Flow ` @@ -18,7 +18,7 @@ The `Maker` class usually contains most of the calculation parameters and other One common task encountered in almost any materials science calculation is writing calculation input files to disk so they can be executed by the underlying software (e.g., VASP, Q-Chem, CP2K, etc.). This is preferably done via a `pymatgen` `InputSet` class. `InputSet` is essentially a dict-like container that specifies the files that need to be written, and their contents. Similarly to the way that `Maker` classes generate `Flow`s, `InputSet`s are most easily created by `InputGenerator` classes. `InputGenerator` have a method `get_input_set()` that typically takes atomic coordinates (e.g., a `Structure` or `Molecule` object) and produce an `InputSet`, e.g., -```python +```py class ExampleInputGenerator(InputGenerator): def get_input_set(self, coordinates: Structure) -> InputSet: # take the input coordinates, determine appropriate @@ -30,7 +30,7 @@ class ExampleInputGenerator(InputGenerator): **TODO - the code block below needs refinement. Not exactly sure how write_inputs() fits into a`Job`** -```python +```py class ExampleMaker(Maker): input_set_generator: ExampleInputGenerator = field( default_factory=ExampleInputGenerator diff --git a/docs/user/atomate-1-vs-2.md b/docs/user/atomate-1-vs-2.md index 3aa317dfe9..8a54a1103f 100644 --- a/docs/user/atomate-1-vs-2.md +++ b/docs/user/atomate-1-vs-2.md @@ -2,6 +2,7 @@ This document contains introductory context for people coming from atomate 1. One of atomate2's core ideas is to allow scaling from a single material, to 100 materials, or 100,000 materials. Therefore, both local submission options and a connection to workflow managers such as FireWorks exist. We plan to support more workflow managers in the future to further ease job submission. + ## Relation between managers running the actual jobs and the workflow as written in `atomate2` There is no leakage between job manager and the workflow definition in `atomate2`. For example, Fireworks is not a required dependency of `atomate2` or `jobflow`. Any `atomate2` workflow can be run using the local manager or an extensible set of external providers, Fireworks just one among them. E.g. all tests are run with mocked calls to the executables using the local manager. diff --git a/docs/user/codes/vasp.md b/docs/user/codes/vasp.md index c11664a7d4..bd26dab91c 100644 --- a/docs/user/codes/vasp.md +++ b/docs/user/codes/vasp.md @@ -260,7 +260,7 @@ LOBSTER_CMD: <> The corresponding flow could, for example, be started with the following code: -```Python +```py from jobflow import SETTINGS from jobflow import run_locally from pymatgen.core.structure import Structure @@ -287,7 +287,7 @@ It is, however, computationally very beneficial to define two different types o Specifically, you might want to change the `_fworker` for the LOBSTER runs and define a separate `lobster` worker within FireWorks: -```python +```py from fireworks import LaunchPad from jobflow.managers.fireworks import flow_to_workflow from pymatgen.core.structure import Structure @@ -321,7 +321,7 @@ lpad.add_wf(wf) Outputs from the automatic analysis with LobsterPy can easily be extracted from the database and also plotted: -```python +```py from jobflow import SETTINGS from pymatgen.electronic_structure.cohp import Cohp from pymatgen.electronic_structure.plotter import CohpPlotter @@ -339,7 +339,7 @@ result = store.query_one( ) for number, (key, cohp) in enumerate( - result["output"]["lobsterpy_data"]["cohp_plot_data"].items() + result["output"]["lobsterpy_data"]["cohp_plot_data"]["data"].items() ): plotter = CohpPlotter() cohp = Cohp.from_dict(cohp) @@ -347,7 +347,7 @@ for number, (key, cohp) in enumerate( plotter.save_plot(f"plots_all_bonds{number}.pdf") for number, (key, cohp) in enumerate( - result["output"]["lobsterpy_data_cation_anion"]["cohp_plot_data"].items() + result["output"]["lobsterpy_data_cation_anion"]["cohp_plot_data"]["data"].items() ): plotter = CohpPlotter() cohp = Cohp.from_dict(cohp) @@ -363,7 +363,7 @@ The inputs for a calculation can be modified in several ways. Every VASP job takes a {obj}`.VaspInputGenerator` as an argument (`input_set_generator`). One option is to specify an alternative input set generator: -```python +```py from atomate2.vasp.sets.core import StaticSetGenerator from atomate2.vasp.jobs.core import StaticMaker @@ -382,7 +382,7 @@ The second approach is to edit the job after it has been made. All VASP jobs hav the `input_set_generator` attribute maker will update the input set that gets written: -```python +```py static_job.maker.input_set_generator.user_incar_settings["LOPTICS"] = True ``` @@ -392,7 +392,7 @@ functions called "powerups" that can apply settings updates to all VASP jobs in These powerups also contain filters for the name of the job and the maker used to generate them. -```python +```py from atomate2.vasp.powerups import update_user_incar_settings from atomate2.vasp.flows.elastic import ElasticMaker from atomate2.vasp.flows.core import DoubleRelaxMaker @@ -485,7 +485,7 @@ There are two options when chaining workflows: These two examples are illustrated in the code below, where we chain a relaxation calculation and a static calculation. -```python +```py from jobflow import Flow from atomate2.vasp.jobs.core import RelaxMaker, StaticMaker from pymatgen.core.structure import Structure diff --git a/docs/user/fireworks.md b/docs/user/fireworks.md index defad2ca5a..07e216e6be 100644 --- a/docs/user/fireworks.md +++ b/docs/user/fireworks.md @@ -13,9 +13,10 @@ FireWorks workflow using the {obj}`~jobflow.managers.fireworks.flow_to_workflow` The workflow can then be submitted to the launchpad in the usual way. For example, to submit an MgO band structure workflow using FireWorks: -```python +```py from fireworks import LaunchPad from atomate2.vasp.flows.core import RelaxBandStructureMaker +from atomate2.vasp.powerups import add_metadata_to_flow from jobflow.managers.fireworks import flow_to_workflow from pymatgen.core import Structure @@ -29,6 +30,14 @@ mgo_structure = Structure( # make a band structure flow to optimise the structure and obtain the band structure bandstructure_flow = RelaxBandStructureMaker().make(mgo_structure) +# (Optional) add metadata to the flow task document. +# Could be useful to filter specific results from the database. +# For e.g., adding material project ID for the compound, use following lines +bandstructure_flow = add_metadata_to_flow( + flow=bandstructure_flow, + additional_fields={"mp_id": "mp-190"}, +) + # convert the flow to a fireworks WorkFlow object wf = flow_to_workflow(bandstructure_flow) diff --git a/docs/user/install.md b/docs/user/install.md index d1ec9ecc39..60dd7d5e36 100644 --- a/docs/user/install.md +++ b/docs/user/install.md @@ -5,14 +5,14 @@ ## Introduction This guide will get you up and running in an environment for running high-throughput -workflows with atomate2. atomate2 is built on the pymatgen, custodian, jobflow, and -FireWorks libraries. Briefly: +workflows with atomate2. atomate2 is built on the `pymatgen`, `custodian`, `jobflow`, and +`FireWorks` libraries. Briefly: -- [pymatgen] is used to create input files and analyze the output of materials science codes. -- [custodian] runs your simulation code (e.g., VASP) and performs error checking/handling +- [`pymatgen`] is used to create input files and analyze the output of materials science codes. +- [`custodian`] runs your simulation code (e.g., VASP) and performs error checking/handling and checkpointing. -- [jobflow] is used to design computational workflows. -- [FireWorks] (optional) is used to manage and execute workflows on HPC machines. +- [`jobflow`] is used to design computational workflows. +- [`FireWorks`] (optional) is used to manage and execute workflows on HPC machines. Running and writing your own workflows are covered in later tutorials. For now, these topics will be covered in enough depth to get you set up and to help you know where to @@ -21,10 +21,10 @@ troubleshoot if you're having problems. Note that this installation tutorial is VASP-centric since almost all functionality currently in atomate2 pertains to VASP. -[pymatgen]: http://pymatgen.org -[custodian]: https://materialsproject.github.io/custodian -[fireworks]: https://materialsproject.github.io/fireworks -[jobflow]: https://materialsproject.github.io/jobflow +[`pymatgen`]: http://pymatgen.org +[`custodian`]: https://materialsproject.github.io/custodian +[`fireworks`]: https://materialsproject.github.io/fireworks +[`jobflow`]: https://materialsproject.github.io/jobflow ### Objectives @@ -76,16 +76,16 @@ MongoDB must be running and available to accept connections whenever you're runn workflows. Thus, it is strongly recommended that you have a server to run MongoDB or (simpler) use a hosting service. Your options are: -- Use a commercial service to host your MongoDB instance. These are typically the +1. Use a commercial service to host your MongoDB instance. These are typically the easiest to use and offer high-quality service but require payment for larger databases. [MongoDB Atlas](https://www.mongodb.com/cloud/atlas) offers a free 500 MB - server which is certainly enough to get started for small to medium size projects, and + server which is certainly enough to get started for small to medium-sized projects, and it is easy to upgrade or migrate your database if you exceed the free allocation. -- Contact your supercomputing center to see if they offer MongoDB hosting (e.g., NERSC +1. Contact your supercomputing center to see if they offer MongoDB hosting (e.g., NERSC has this, Google "request NERSC MongoDB database"). -- Self-host a MongoDB server. +1. Self-host a MongoDB server. -If you're just starting, we suggest the first (with a free plan) or second option +If you're just starting, we suggest option 1 (with a free plan) or 2 (if available to you). The third option will require you to open up network settings to accept outside connections properly which can sometimes be tricky. @@ -113,7 +113,7 @@ issues later in this tutorial, some options are: ## Create a directory scaffold for atomate2 -Installing atomate2 includes installation of codes, configuration files, and various +Installing atomate2 includes the installation of codes, configuration files, and various binaries and libraries. Thus, it is useful to create a directory structure that organizes all these items. @@ -140,7 +140,7 @@ Make sure to create a Python 3.8+ environment as recent versions of atomate2 onl support Python 3.8 and higher. ``` -We highly recommended that you organize your installation of the atomate2 and the other +We highly recommend that you organize your installation of the atomate2 and the other Python codes using a conda virtual environment. Some of the main benefits are: - Different Python projects that have conflicting packages can coexist on the same @@ -194,7 +194,7 @@ that are unique only to these databases. Create the following files in `<>/config`. -### jobflow.yaml +### `jobflow.yaml` The `jobflow.yaml` file contains the credentials of the MongoDB server that will store calculation outputs. The `jobflow.yaml` file requires you to enter the basic database @@ -351,14 +351,13 @@ workflows. ### Define the workflow Workflows are written using the `jobflow` software. Essentially, individual stages of -a workflow are simple Python functions. Jobflow provides a way to connect jobs together -in a natural way. For more details on connecting jobs together see: -[](connecting_vasp_jobs). +a workflow are simple Python functions. Jobflow provides a way to connect jobs in a natural way. +For more details on connecting jobs see: [](connecting_vasp_jobs). Go to the directory where you would like your calculations to run (i.e., your scratch or work directory) and create a file called `relax.py` containing: -```python +```py from atomate2.vasp.jobs.core import RelaxMaker from jobflow import run_locally from pymatgen.core import Structure @@ -416,7 +415,7 @@ on the Grid Engine scheduler, this would be using `qsub job.sh`. Once the job is finished, you can connect to the output database and check the job output. -```python +```py from jobflow import SETTINGS store = SETTINGS.JOB_STORE diff --git a/docs/user/running-workflows.md b/docs/user/running-workflows.md index db164cb03f..28322571af 100644 --- a/docs/user/running-workflows.md +++ b/docs/user/running-workflows.md @@ -7,7 +7,7 @@ Once you have a working installation of atomate2, you'll want to jump in and start running workflows. Atomate2 includes many workflows with reasonable settings that can get you started. This tutorial will quickly guide you through customizing and running a -workflow to calculate the bandstructure of MgO. +workflow to calculate the band structure of MgO. ### Objectives @@ -16,7 +16,7 @@ workflow to calculate the bandstructure of MgO. ### Prerequisites -In order for you to complete this tutorial you need +For you to complete this tutorial you need * A working installation of atomate2. @@ -60,7 +60,7 @@ workflow. Create a Python script named `mgo_bandstructure.py` with the following contents: -```python +```py from atomate2.vasp.flows.core import RelaxBandStructureMaker from jobflow import run_locally from pymatgen.core import Structure @@ -116,7 +116,7 @@ results will be in your database. Finally, we'll plot the results that we calculated. Simply run the following Python code, either as a script or on the Python prompt. -```python +```py from jobflow import SETTINGS from pymatgen.electronic_structure.plotter import DosPlotter, BSPlotter from pymatgen.electronic_structure.dos import CompleteDos diff --git a/pyproject.toml b/pyproject.toml index 51e0a99594..ce32149a4e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -42,7 +42,7 @@ amset = ["amset>=0.4.15", "pydash"] cclib = ["cclib"] mp = ["mp-api>=0.37.5"] phonons = ["phonopy>=1.10.8", "seekpath"] -lobster = ["ijson>=3.2.2", "lobsterpy>=0.3.4"] +lobster = ["ijson>=3.2.2", "lobsterpy>=0.3.6"] defects = ["dscribe>=1.2.0", "pymatgen-analysis-defects>=2022.11.30"] forcefields = [ "ase>=3.22.1", @@ -55,7 +55,7 @@ docs = [ "FireWorks==2.0.3", "autodoc_pydantic==2.0.1", "furo==2023.9.10", - "ipython==8.19.0", + "ipython==8.20.0", "jsonschema[format]", "myst_parser==2.0.0", "numpydoc==1.6.0", @@ -75,10 +75,10 @@ strict = [ "click==8.1.7", "custodian==2023.10.9", "dscribe==2.1.0", - "emmet-core==0.75.2", + "emmet-core==0.76.2", "ijson==3.2.3", "jobflow==0.1.16", - "lobsterpy==0.3.4", + "lobsterpy==0.3.6", "mace-torch>=0.3.3", "matgl==0.9.1", "monty==2023.9.25", @@ -87,7 +87,7 @@ strict = [ "phonopy==2.21.0", "pydantic-settings==2.1.0", "pydantic==2.4.2", - "pymatgen-analysis-defects==2023.12.14", + "pymatgen-analysis-defects==2024.1.12", "pymatgen==2023.12.18", "quippy-ase==0.9.14", "seekpath==2.1.0", @@ -157,7 +157,6 @@ ignore = [ "PT006", # pytest-parametrize-names-wrong-type "RUF013", # implicit-optional # TODO remove PT011, pytest.raises() should always check err msg - "ANN001", # TODO remove this ignore "ANN002", "ANN003", "ANN101", # missing self type annotation diff --git a/src/atomate2/amset/schemas.py b/src/atomate2/amset/schemas.py index 6bf2a07c9f..875ae67568 100644 --- a/src/atomate2/amset/schemas.py +++ b/src/atomate2/amset/schemas.py @@ -117,7 +117,7 @@ class AmsetTaskDocument(StructureMetadata): completed_at: str = Field( None, description="Timestamp for when this task was completed" ) - input: dict = Field(None, description="The input settings") # noqa: A003 + input: dict = Field(None, description="The input settings") transport: TransportData = Field(None, description="The transport results") usage_stats: UsageStats = Field(None, description="Timing and memory usage") mesh: MeshData = Field(None, description="Full AMSET mesh data") diff --git a/src/atomate2/cli/dev.py b/src/atomate2/cli/dev.py index 085a1436d5..a51b905d83 100644 --- a/src/atomate2/cli/dev.py +++ b/src/atomate2/cli/dev.py @@ -1,7 +1,15 @@ """Module containing command line scripts for developers.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + import click +if TYPE_CHECKING: + from pathlib import Path + @click.group(context_settings={"help_option_names": ["-h", "--help"]}) def dev() -> None: @@ -10,7 +18,7 @@ def dev() -> None: @dev.command(context_settings={"help_option_names": ["-h", "--help"]}) @click.argument("test_dir") -def vasp_test_data(test_dir) -> None: +def vasp_test_data(test_dir: Path) -> None: """Generate test data for VASP unit tests. This script expects there is an outputs.json file and job folders in the current @@ -156,7 +164,7 @@ def test_my_flow(mock_vasp, clean_dir, si_structure): print(test_function_str) # noqa: T201 -def _potcar_to_potcar_spec(potcar_filename, output_filename) -> None: +def _potcar_to_potcar_spec(potcar_filename: str | Path, output_filename: Path) -> None: """Convert a POTCAR file to a POTCAR.spec file.""" from pymatgen.io.vasp import Potcar diff --git a/src/atomate2/common/flows/defect.py b/src/atomate2/common/flows/defect.py index f6d8bce15a..871d18d9cf 100644 --- a/src/atomate2/common/flows/defect.py +++ b/src/atomate2/common/flows/defect.py @@ -23,6 +23,7 @@ from pathlib import Path import numpy.typing as npt + from emmet.core.tasks import TaskDoc from pymatgen.analysis.defects.core import Defect from pymatgen.core.structure import Structure from pymatgen.entries.computed_entries import ComputedStructureEntry @@ -361,7 +362,7 @@ def sc_entry_and_locpot_from_prv( """ @abstractmethod - def get_planar_locpot(self, task_doc) -> dict: + def get_planar_locpot(self, task_doc: TaskDoc) -> dict: """Get the Planar Locpot from the TaskDoc. This is needed just in case the planar average locpot is stored in different diff --git a/src/atomate2/common/jobs/defect.py b/src/atomate2/common/jobs/defect.py index 2517ac36a3..c8ad3a4dc4 100644 --- a/src/atomate2/common/jobs/defect.py +++ b/src/atomate2/common/jobs/defect.py @@ -269,8 +269,8 @@ def bulk_supercell_calculation( """ if get_planar_locpot is None: - def get_planar_locpot(tdoc) -> NDArray: - return tdoc.calcs_reversed[0].output.locpot + def get_planar_locpot(task_doc: TaskDoc) -> NDArray: + return task_doc.calcs_reversed[0].output.locpot logger.info("Running bulk supercell calculation. Running...") sc_mat = get_sc_fromstruct(uc_structure) if sc_mat is None else sc_mat diff --git a/src/atomate2/common/jobs/phonons.py b/src/atomate2/common/jobs/phonons.py index 968192e201..92609aa812 100644 --- a/src/atomate2/common/jobs/phonons.py +++ b/src/atomate2/common/jobs/phonons.py @@ -248,9 +248,9 @@ def generate_frequencies_eigenvectors( @job(data=["forces", "displaced_structures"]) def run_phonon_displacements( - displacements, + displacements: list[Structure], structure: Structure, - supercell_matrix, + supercell_matrix: Matrix3D, phonon_maker: BaseVaspMaker | ForceFieldStaticMaker = None, prev_dir: str | Path = None, ) -> Flow: @@ -281,16 +281,16 @@ def run_phonon_displacements( "dirs": [], } - for i, displacement in enumerate(displacements): + for idx, displacement in enumerate(displacements): if prev_dir is not None: phonon_job = phonon_maker.make(displacement, prev_dir=prev_dir) else: phonon_job = phonon_maker.make(displacement) - phonon_job.append_name(f" {i + 1}/{len(displacements)}") + phonon_job.append_name(f" {idx + 1}/{len(displacements)}") # we will add some meta data info = { - "displacement_number": i, + "displacement_number": idx, "original_structure": structure, "supercell_matrix": supercell_matrix, "displaced_structure": displacement, @@ -302,7 +302,7 @@ def run_phonon_displacements( ) phonon_jobs.append(phonon_job) - outputs["displacement_number"].append(i) + outputs["displacement_number"].append(idx) outputs["uuids"].append(phonon_job.output.uuid) outputs["dirs"].append(phonon_job.output.dir_name) outputs["forces"].append(phonon_job.output.output.forces) diff --git a/src/atomate2/common/powerups.py b/src/atomate2/common/powerups.py new file mode 100644 index 0000000000..b4f8bddb2d --- /dev/null +++ b/src/atomate2/common/powerups.py @@ -0,0 +1,76 @@ +"""Utilities for modifying workflow.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from jobflow import Flow, Maker + + +def add_metadata_to_flow( + flow: Flow, additional_fields: dict, class_filter: Maker +) -> Flow: + """ + Return the flow with additional field(metadata) to the task doc. + + This allows adding metadata to the task-docs, could be useful + to query results from DB. + + Parameters + ---------- + flow: + additional_fields : dict + A dict with metadata. + class_filter: .Maker + The Maker to which additional metadata needs to be added + + Returns + ------- + Flow + Flow with added metadata to the task-doc. + """ + flow.update_maker_kwargs( + { + "_set": { + f"task_document_kwargs->additional_fields->{field}": value + for field, value in additional_fields.items() + } + }, + dict_mod=True, + class_filter=class_filter, + ) + + return flow + + +def update_custodian_handlers( + flow: Flow, custom_handlers: tuple, class_filter: Maker +) -> Flow: + """ + Return the flow with custom custodian handlers for VASP jobs. + + This allows user to selectively set error correcting handlers for VASP jobs + or completely unset error handlers. + + Parameters + ---------- + flow: + custom_handlers : tuple + A tuple with custodian handlers. + class_filter: .Maker + The Maker to which custom custodian handler needs to be added + + Returns + ------- + Flow + Flow with modified custodian handlers. + """ + code = class_filter.name.split(" ")[1] + flow.update_maker_kwargs( + {"_set": {f"run_{code}_kwargs->handlers": custom_handlers}}, + dict_mod=True, + class_filter=class_filter, + ) + + return flow diff --git a/src/atomate2/common/schemas/cclib.py b/src/atomate2/common/schemas/cclib.py index ad5f0480ca..391d2ec31d 100644 --- a/src/atomate2/common/schemas/cclib.py +++ b/src/atomate2/common/schemas/cclib.py @@ -256,7 +256,7 @@ def from_logfile( @requires(cclib, "cclib_calculate requires cclib to be installed.") def cclib_calculate( - cclib_obj, + cclib_obj: Any, method: str, cube_file: Union[Path, str], proatom_dir: Union[Path, str], @@ -310,28 +310,28 @@ def cclib_calculate( vol = volume.read_from_cube(str(cube_file)) if method == "bader": - m = Bader(cclib_obj, vol) + _method = Bader(cclib_obj, vol) elif method == "bickelhaupt": - m = Bickelhaupt(cclib_obj) + _method = Bickelhaupt(cclib_obj) elif method == "cpsa": - m = CSPA(cclib_obj) + _method = CSPA(cclib_obj) elif method == "ddec6": - m = DDEC6(cclib_obj, vol, str(proatom_dir)) + _method = DDEC6(cclib_obj, vol, str(proatom_dir)) elif method == "density": - m = Density(cclib_obj) + _method = Density(cclib_obj) elif method == "hirshfeld": - m = Hirshfeld(cclib_obj, vol, str(proatom_dir)) + _method = Hirshfeld(cclib_obj, vol, str(proatom_dir)) elif method == "lpa": - m = LPA(cclib_obj) + _method = LPA(cclib_obj) elif method == "mbo": - m = MBO(cclib_obj) + _method = MBO(cclib_obj) elif method == "mpa": - m = MPA(cclib_obj) + _method = MPA(cclib_obj) else: - raise ValueError(f"{method} is not supported.") + raise ValueError(f"{method=} is not supported.") try: - m.calculate() + _method.calculate() except AttributeError: return None @@ -351,20 +351,20 @@ def cclib_calculate( ] calc_attributes = {} for attribute in avail_attributes: - if hasattr(m, attribute): - calc_attributes[attribute] = getattr(m, attribute) + if hasattr(_method, attribute): + calc_attributes[attribute] = getattr(_method, attribute) return calc_attributes def _get_homos_lumos( - moenergies: list[list[float]], homo_indices: list[int] + mo_energies: list[list[float]], homo_indices: list[int] ) -> tuple[list[float], Optional[list[float]], Optional[list[float]]]: """ Calculate the HOMO, LUMO, and HOMO-LUMO gap energies in eV. Parameters ---------- - moenergies + mo_energies List of MO energies. For restricted calculations, List[List[float]] is length one. For unrestricted, it is length two. homo_indices @@ -380,13 +380,13 @@ def _get_homos_lumos( The HOMO-LUMO gaps (eV), calculated as LUMO_alpha-HOMO_alpha and LUMO_beta-HOMO_beta """ - homo_energies = [moenergies[i][h] for i, h in enumerate(homo_indices)] + homo_energies = [mo_energies[i][h] for i, h in enumerate(homo_indices)] # Make sure that the HOMO+1 (i.e. LUMO) is in moenergies (sometimes virtual # orbitals aren't printed in the output) for i, h in enumerate(homo_indices): - if len(moenergies[i]) < h + 2: + if len(mo_energies[i]) < h + 2: return homo_energies, None, None - lumo_energies = [moenergies[i][h + 1] for i, h in enumerate(homo_indices)] + lumo_energies = [mo_energies[i][h + 1] for i, h in enumerate(homo_indices)] homo_lumo_gaps = [ lumo_energies[i] - homo_energies[i] for i in range(len(homo_energies)) ] diff --git a/src/atomate2/common/schemas/defects.py b/src/atomate2/common/schemas/defects.py index 04861cb719..097df3fbde 100644 --- a/src/atomate2/common/schemas/defects.py +++ b/src/atomate2/common/schemas/defects.py @@ -1,6 +1,7 @@ """General schemas for defect workflow outputs.""" import logging +from collections.abc import Sequence from typing import Any, Callable, Optional, Union import numpy as np @@ -239,8 +240,8 @@ def from_task_outputs( def get_ent( struct: Structure, energy: float, - dir_name, - uuid, + dir_name: str, + uuid: str, ) -> ComputedStructureEntry: return ComputedStructureEntry( structure=struct, @@ -283,14 +284,16 @@ def from_entries( """ - def find_entry(entries, uuid) -> tuple[int, ComputedStructureEntry]: + def find_entry( + entries: Sequence[ComputedStructureEntry], uuid: str + ) -> tuple[int, ComputedStructureEntry]: """Find the entry with the given UUID.""" - for itr, entry in enumerate(entries): + for idx, entry in enumerate(entries): if entry.data["uuid"] == uuid: - return itr, entry + return idx, entry raise ValueError(f"Could not find entry with UUID: {uuid}") - def dQ_entries(e1, e2) -> float: # noqa: N802 + def dQ_entries(e1: ComputedStructureEntry, e2: ComputedStructureEntry) -> float: # noqa: N802 """Get the displacement between two entries.""" return get_dQ(e1.structure, e2.structure) @@ -338,23 +341,23 @@ def dQ_entries(e1, e2) -> float: # noqa: N802 relaxed_index2=idx2, ) - def get_taskdocs(self) -> list[list[TaskDoc]]: + def get_taskdocs(self) -> tuple[list[TaskDoc], list[TaskDoc]]: """Get the distorted task documents.""" - def remove_host_name(dir_name) -> str: + def remove_host_name(dir_name: str) -> str: return dir_name.split(":")[-1] - return [ - [ - TaskDoc.from_directory(remove_host_name(dir_name)) - for dir_name in self.static_dirs1 - ], - [ - TaskDoc.from_directory(remove_host_name(dir_name)) - for dir_name in self.static_dirs2 - ], + static1_task_docs = [ + TaskDoc.from_directory(remove_host_name(dir_name)) + for dir_name in self.static_dirs1 + ] + static2_task_docs = [ + TaskDoc.from_directory(remove_host_name(dir_name)) + for dir_name in self.static_dirs2 ] + return static1_task_docs, static2_task_docs + def sort_pos_dist( list_in: list[Any], diff --git a/src/atomate2/common/schemas/phonons.py b/src/atomate2/common/schemas/phonons.py index e0ac0ddacd..cf48880662 100644 --- a/src/atomate2/common/schemas/phonons.py +++ b/src/atomate2/common/schemas/phonons.py @@ -2,6 +2,7 @@ import copy import logging +from pathlib import Path from typing import Optional, Union import numpy as np @@ -98,6 +99,9 @@ class PhononJobDirs(BaseModel): optimization_run_job_dir: Optional[str] = Field( None, description="Directory where optimization run was performed." ) + taskdoc_run_job_dir: Optional[str] = Field( + None, description="Directory where taskdoc was generated." + ) class PhononBSDOSDoc(StructureMetadata, extra="allow"): # type: ignore[call-arg] @@ -460,6 +464,7 @@ def from_forces_born( "static_run_job_dir": kwargs["static_run_job_dir"], "born_run_job_dir": kwargs["born_run_job_dir"], "optimization_run_job_dir": kwargs["optimization_run_job_dir"], + "taskdoc_run_job_dir": str(Path.cwd()), }, uuids={ "displacements_uuids": displacement_data["uuids"], diff --git a/src/atomate2/cp2k/powerups.py b/src/atomate2/cp2k/powerups.py index 35e696d781..205b30c39f 100644 --- a/src/atomate2/cp2k/powerups.py +++ b/src/atomate2/cp2k/powerups.py @@ -8,6 +8,8 @@ from jobflow import Flow, Job, Maker from pymatgen.io.vasp import Kpoints +from atomate2.common.powerups import add_metadata_to_flow as base_add_metadata_to_flow +from atomate2.common.powerups import update_custodian_handlers as base_custodian_handler from atomate2.cp2k.jobs.base import BaseCp2kMaker @@ -45,12 +47,14 @@ def update_user_input_settings( A copy of the input flow/job/maker modified to use the updated input settings. """ - # Convert nested dictionary updates for cp2k inpt settings + # Convert nested dictionary updates for cp2k input settings # into dict_mod update format - def nested_to_dictmod(d, kk="input_set_generator->user_input_settings") -> dict: + def nested_to_dictmod( + dct: dict, kk: str = "input_set_generator->user_input_settings" + ) -> dict: d2 = {} - for k, v in d.items(): - k2 = kk + f"->{k}" + for k, v in dct.items(): + k2 = f"{kk}->{k}" if isinstance(v, dict): d2.update(nested_to_dictmod(v, kk=k2)) else: @@ -138,3 +142,57 @@ def update_user_kpoints_settings( dict_mod=True, ) return updated_flow + + +def add_metadata_to_flow( + flow: Flow, additional_fields: dict, class_filter: Maker = BaseCp2kMaker +) -> Flow: + """ + Return the Cp2k flow with additional field(metadata) to the task doc. + + This allows adding metadata to the task-docs, could be useful + to query results from DB. + + Parameters + ---------- + flow: + additional_fields : dict + A dict with metadata. + class_filter: .BaseCp2kMaker + The Maker to which additional metadata needs to be added + + Returns + ------- + Flow + Flow with added metadata to the task-doc. + """ + return base_add_metadata_to_flow( + flow=flow, class_filter=class_filter, additional_fields=additional_fields + ) + + +def update_cp2k_custodian_handlers( + flow: Flow, custom_handlers: tuple, class_filter: Maker = BaseCp2kMaker +) -> Flow: + """ + Return the flow with custom custodian handlers for Cp2k jobs. + + This allows user to selectively set error correcting handlers for Cp2k jobs + or completely unset error handlers. + + Parameters + ---------- + flow: + custom_handlers : tuple + A tuple with custodian handlers. + class_filter: .BaseCp2kMaker + The Maker to which custom custodian handler needs to be added + + Returns + ------- + Flow + Flow with modified custodian handlers. + """ + return base_custodian_handler( + flow=flow, custom_handlers=custom_handlers, class_filter=class_filter + ) diff --git a/src/atomate2/cp2k/schemas/calc_types/_generate.py b/src/atomate2/cp2k/schemas/calc_types/_generate.py index f4efdeb190..8881a144c3 100644 --- a/src/atomate2/cp2k/schemas/calc_types/_generate.py +++ b/src/atomate2/cp2k/schemas/calc_types/_generate.py @@ -2,6 +2,7 @@ from itertools import product from pathlib import Path +from typing import Any from monty.serialization import loadfn @@ -24,14 +25,13 @@ _RUN_TYPES.append(f"{rt}{vdw}{u}") # noqa: PERF401 -def get_enum_source(enum_name, doc, items) -> str: +def get_enum_source(enum_name: str, doc: str, items: dict[str, Any]) -> str: header = f""" class {enum_name}(ValueEnum): \"\"\" {doc} \"\"\"\n """ - items = [f' {const} = "{val}"' for const, val in items.items()] - return header + "\n".join(items) + return header + "\n".join(f' {key} = "{val}"' for key, val in items.items()) run_type_enum = get_enum_source( diff --git a/src/atomate2/cp2k/schemas/calc_types/utils.py b/src/atomate2/cp2k/schemas/calc_types/utils.py index 57f9f86e59..d3ec9ae1de 100644 --- a/src/atomate2/cp2k/schemas/calc_types/utils.py +++ b/src/atomate2/cp2k/schemas/calc_types/utils.py @@ -1,6 +1,6 @@ """Module to define various calculation types as Enums for CP2K.""" -from collections.abc import Iterable +from collections.abc import Iterable, Sequence from pathlib import Path from monty.serialization import loadfn @@ -22,7 +22,7 @@ def run_type(inputs: dict) -> RunType: """ dft = inputs.get("dft") - def _variant_equal(v1, v2) -> bool: + def _variant_equal(v1: Sequence, v2: Sequence) -> bool: """Determine if two run_types are equal.""" if isinstance(v1, str) and isinstance(v2, str): return v1.strip().upper() == v2.strip().upper() diff --git a/src/atomate2/cp2k/schemas/calculation.py b/src/atomate2/cp2k/schemas/calculation.py index 8e5bae8354..668cb21630 100644 --- a/src/atomate2/cp2k/schemas/calculation.py +++ b/src/atomate2/cp2k/schemas/calculation.py @@ -258,7 +258,7 @@ class Calculation(BaseModel): has_cp2k_completed: Status = Field( None, description="Whether CP2K completed the calculation successfully" ) - input: CalculationInput = Field( # noqa: A003 + input: CalculationInput = Field( None, description="CP2K input settings for the calculation" ) output: CalculationOutput = Field(None, description="The CP2K calculation output") diff --git a/src/atomate2/cp2k/schemas/task.py b/src/atomate2/cp2k/schemas/task.py index a8c60d6306..5998fbc014 100644 --- a/src/atomate2/cp2k/schemas/task.py +++ b/src/atomate2/cp2k/schemas/task.py @@ -244,7 +244,7 @@ class TaskDocument(StructureMetadata, MoleculeMetadata): completed_at: Optional[str] = Field( None, description="Timestamp for when this task was completed" ) - input: Optional[InputSummary] = Field( # noqa: A003 + input: Optional[InputSummary] = Field( None, description="The input to the first calculation" ) output: Optional[OutputSummary] = Field( diff --git a/src/atomate2/cp2k/sets/base.py b/src/atomate2/cp2k/sets/base.py index f6954abc72..870f906a66 100644 --- a/src/atomate2/cp2k/sets/base.py +++ b/src/atomate2/cp2k/sets/base.py @@ -5,13 +5,13 @@ import os from copy import deepcopy from dataclasses import dataclass, field +from importlib.resources import files as get_mod_path from pathlib import Path from typing import Any import numpy as np from monty.io import zopen from monty.serialization import loadfn -from pkg_resources import resource_filename from pymatgen.core.structure import Molecule, Structure from pymatgen.io.core import InputGenerator, InputSet from pymatgen.io.cp2k.inputs import ( @@ -29,8 +29,8 @@ from atomate2 import SETTINGS -_BASE_CP2K_SET = loadfn(resource_filename("atomate2.cp2k.sets", "BaseCp2kSet.yaml")) -_BASE_GAPW_SET = loadfn(resource_filename("atomate2.cp2k.sets", "BaseAllSet.yaml")) +_BASE_CP2K_SET = loadfn(get_mod_path("atomate2.cp2k.sets") / "BaseCp2kSet.yaml") +_BASE_GAPW_SET = loadfn(get_mod_path("atomate2.cp2k.sets") / "BaseAllSet.yaml") class Cp2kInputSet(InputSet): diff --git a/src/atomate2/forcefields/schemas.py b/src/atomate2/forcefields/schemas.py index 00b1d3adc2..28d0e013db 100644 --- a/src/atomate2/forcefields/schemas.py +++ b/src/atomate2/forcefields/schemas.py @@ -80,7 +80,7 @@ class ForceFieldTaskDocument(StructureMetadata): None, description="Final output structure from the task" ) - input: InputDoc = Field( # noqa: A003 + input: InputDoc = Field( None, description="The inputted information used to run this job." ) diff --git a/src/atomate2/vasp/powerups.py b/src/atomate2/vasp/powerups.py index a050654628..c73813bbbb 100644 --- a/src/atomate2/vasp/powerups.py +++ b/src/atomate2/vasp/powerups.py @@ -8,6 +8,8 @@ from jobflow import Flow, Job, Maker from pymatgen.io.vasp import Kpoints +from atomate2.common.powerups import add_metadata_to_flow as base_add_metadata_to_flow +from atomate2.common.powerups import update_custodian_handlers as base_custodian_handler from atomate2.vasp.jobs.base import BaseVaspMaker @@ -277,3 +279,58 @@ def use_auto_ispin( name_filter=name_filter, class_filter=class_filter, ) + + +def add_metadata_to_flow( + flow: Flow, additional_fields: dict, class_filter: Maker = BaseVaspMaker +) -> Flow: + """ + Return the VASP flow with additional field(metadata) to the task doc. + + This allows adding metadata to the task-docs, could be useful + to query results from DB. + + Parameters + ---------- + flow : Flow + The flow to which to add metadata. + additional_fields : dict + A dict with metadata. + class_filter: .BaseVaspMaker + The Maker to which additional metadata needs to be added + + Returns + ------- + Flow + Flow with added metadata to the task-doc. + """ + return base_add_metadata_to_flow( + flow=flow, class_filter=class_filter, additional_fields=additional_fields + ) + + +def update_vasp_custodian_handlers( + flow: Flow, custom_handlers: tuple, class_filter: Maker = BaseVaspMaker +) -> Flow: + """ + Return the flow with custom custodian handlers for VASP jobs. + + This allows user to selectively set error correcting handlers for VASP jobs + or completely unset error handlers. + + Parameters + ---------- + flow: + custom_handlers : tuple + A tuple with custodian handlers. + class_filter: .Maker + The Maker to which custom custodian handler needs to be added + + Returns + ------- + Flow + Flow with modified custodian handlers. + """ + return base_custodian_handler( + flow=flow, custom_handlers=custom_handlers, class_filter=class_filter + ) diff --git a/src/atomate2/vasp/sets/base.py b/src/atomate2/vasp/sets/base.py index e3fa21d298..a06b8f5066 100644 --- a/src/atomate2/vasp/sets/base.py +++ b/src/atomate2/vasp/sets/base.py @@ -7,6 +7,7 @@ import warnings from copy import deepcopy from dataclasses import dataclass, field +from importlib.resources import files as get_mod_path from itertools import groupby from pathlib import Path from typing import TYPE_CHECKING, Any @@ -14,7 +15,6 @@ import numpy as np from monty.io import zopen from monty.serialization import loadfn -from pkg_resources import resource_filename from pymatgen.electronic_structure.core import Magmom from pymatgen.io.core import InputGenerator, InputSet from pymatgen.io.vasp import Incar, Kpoints, Outcar, Poscar, Potcar, Vasprun @@ -34,7 +34,7 @@ from pymatgen.core import Structure -_BASE_VASP_SET = loadfn(resource_filename("atomate2.vasp.sets", "BaseVaspSet.yaml")) +_BASE_VASP_SET = loadfn(get_mod_path("atomate2.vasp.sets") / "BaseVaspSet.yaml") class VaspInputSet(InputSet): @@ -236,8 +236,9 @@ class VaspInputGenerator(InputGenerator): so these keys can be defined in one of two ways, e.g. either {"LDAUU":{"O":{"Fe":5}}} to set LDAUU for Fe to 5 in an oxide, or {"LDAUU":{"Fe":5}} to set LDAUU to 5 regardless of the input structure. - To set magmoms, pass a dict mapping element symbols to magnetic moments, e.g. - {"MAGMOM": {"Co": 1}}. + To set magmoms, pass a dict mapping the strings of species to magnetic + moments, e.g. {"MAGMOM": {"Co": 1}} or {"MAGMOM": {"Fe2+,spin=4": 3.7}} in the + case of a site with Species("Fe2+", spin=4). If None is given, that key is unset. For example, {"ENCUT": None} will remove ENCUT from the incar settings. user_kpoints_settings @@ -655,7 +656,7 @@ def _get_incar( if k == "MAGMOM": incar[k] = _get_magmoms( structure, - magmoms=self.user_incar_settings.get("MAGMOMS", {}), + magmoms=self.user_incar_settings.get("MAGMOM", {}), config_magmoms=config_magmoms, ) elif k in ("LDAUU", "LDAUJ", "LDAUL") and incar_settings.get("LDAU", False): @@ -726,7 +727,9 @@ def _get_incar( # Finally, re-apply `self.user_incar_settings` to make sure any accidentally # overwritten settings are changed back to the intended values. - _apply_incar_updates(incar, self.user_incar_settings) + # skip dictionary parameters to avoid dictionaries appearing in the INCAR + skip = ["LDAUU", "LDAUJ", "LDAUL", "MAGMOM"] + _apply_incar_updates(incar, self.user_incar_settings, skip=skip) return incar diff --git a/src/atomate2/vasp/sets/matpes.py b/src/atomate2/vasp/sets/matpes.py index 67dc3410ae..736f72edd3 100644 --- a/src/atomate2/vasp/sets/matpes.py +++ b/src/atomate2/vasp/sets/matpes.py @@ -7,10 +7,10 @@ from __future__ import annotations from dataclasses import dataclass, field +from importlib.resources import files as get_mod_path from typing import TYPE_CHECKING from monty.serialization import loadfn -from pkg_resources import resource_filename from atomate2.vasp.sets.base import VaspInputGenerator @@ -21,10 +21,10 @@ # POTCAR section comes from PARENT but atomate2 does not support inheritance yet _BASE_MATPES_PBE_STATIC_SET_NO_POTCAR = loadfn( - resource_filename("pymatgen.io.vasp", "MatPESStaticSet.yaml") + get_mod_path("pymatgen.io.vasp") / "MatPESStaticSet.yaml" ) _POTCAR_BASE_FILE = f"{_BASE_MATPES_PBE_STATIC_SET_NO_POTCAR['PARENT']}.yaml" -_POTCAR_SET = loadfn(resource_filename("pymatgen.io.vasp", _POTCAR_BASE_FILE)) +_POTCAR_SET = loadfn(get_mod_path("pymatgen.io.vasp") / _POTCAR_BASE_FILE) _BASE_MATPES_PBE_STATIC_SET = {**_POTCAR_SET, **_BASE_MATPES_PBE_STATIC_SET_NO_POTCAR} diff --git a/src/atomate2/vasp/sets/mp.py b/src/atomate2/vasp/sets/mp.py index 3e4abc8a47..8ff1029c6d 100644 --- a/src/atomate2/vasp/sets/mp.py +++ b/src/atomate2/vasp/sets/mp.py @@ -9,10 +9,10 @@ from __future__ import annotations from dataclasses import dataclass, field +from importlib.resources import files as get_mod_path from typing import TYPE_CHECKING from monty.serialization import loadfn -from pkg_resources import resource_filename from atomate2.vasp.sets.core import RelaxSetGenerator, StaticSetGenerator @@ -21,10 +21,10 @@ from pymatgen.io.vasp import Outcar, Vasprun _BASE_MP_GGA_RELAX_SET = loadfn( - resource_filename("atomate2.vasp.sets", "BaseMPGGASet.yaml") + get_mod_path("atomate2.vasp.sets") / "BaseMPGGASet.yaml" ) _BASE_MP_R2SCAN_RELAX_SET = loadfn( - resource_filename("atomate2.vasp.sets", "BaseMPR2SCANRelaxSet.yaml") + get_mod_path("atomate2.vasp.sets") / "BaseMPR2SCANRelaxSet.yaml" ) @@ -33,7 +33,9 @@ class MPGGARelaxSetGenerator(RelaxSetGenerator): """Class to generate MP-compatible VASP GGA relaxation input sets.""" config_dict: dict = field(default_factory=lambda: _BASE_MP_GGA_RELAX_SET) + auto_ismear: bool = False auto_kspacing: bool = True + inherit_incar: bool | None = False @dataclass @@ -41,7 +43,9 @@ class MPGGAStaticSetGenerator(StaticSetGenerator): """Class to generate MP-compatible VASP GGA static input sets.""" config_dict: dict = field(default_factory=lambda: _BASE_MP_GGA_RELAX_SET) + auto_ismear: bool = False auto_kspacing: bool = True + inherit_incar: bool | None = False def get_incar_updates( self, @@ -87,8 +91,10 @@ class MPMetaGGAStaticSetGenerator(StaticSetGenerator): """Class to generate MP-compatible VASP GGA static input sets.""" config_dict: dict = field(default_factory=lambda: _BASE_MP_R2SCAN_RELAX_SET) + auto_ismear: bool = False auto_kspacing: bool = True bandgap_tol: float = 1e-4 + inherit_incar: bool | None = False def get_incar_updates( self, @@ -145,7 +151,9 @@ class MPMetaGGARelaxSetGenerator(RelaxSetGenerator): config_dict: dict = field(default_factory=lambda: _BASE_MP_R2SCAN_RELAX_SET) bandgap_tol: float = 1e-4 + auto_ismear: bool = False auto_kspacing: bool = True + inherit_incar: bool | None = False def get_incar_updates( self, diff --git a/tests/cp2k/test_powerups.py b/tests/cp2k/test_powerups.py index e3b74bb819..587bff8e8d 100644 --- a/tests/cp2k/test_powerups.py +++ b/tests/cp2k/test_powerups.py @@ -60,3 +60,43 @@ def test_update_user_settings(powerup, attribute, settings): getattr(flow.jobs[1].function.__self__.input_set_generator, attribute) != settings ) + + +@pytest.mark.parametrize( + "powerup,settings", + [ + ("add_metadata_to_flow", {"mp-id": "mp-xxx"}), + ("add_metadata_to_flow", {"mp-id": "mp-170", "composition": "NaCl"}), + ], +) +def test_add_metadata_to_flow(powerup, settings): + from atomate2.cp2k import powerups + from atomate2.cp2k.flows.core import DoubleRelaxMaker + + powerup_func = getattr(powerups, powerup) + + # test flow + drm = DoubleRelaxMaker() + flow = drm.make(1) + flow = powerup_func(flow, settings) + assert ( + flow.jobs[0].function.__self__.task_document_kwargs["additional_fields"] + == settings + ) + + +@pytest.mark.parametrize( + "powerup, settings", + [("update_cp2k_custodian_handlers", ())], +) +def test_update_cp2k_custodian_handlers(powerup, settings): + from atomate2.cp2k import powerups + from atomate2.cp2k.flows.core import DoubleRelaxMaker + + powerup_func = getattr(powerups, powerup) + + # test flow + drm = DoubleRelaxMaker() + flow = drm.make(1) + flow = powerup_func(flow, settings) + assert flow.jobs[0].function.__self__.run_cp2k_kwargs["handlers"] == settings diff --git a/tests/forcefields/flows/test_elastic.py b/tests/forcefields/flows/test_elastic.py index 9285a0d1c8..1f41671fcb 100644 --- a/tests/forcefields/flows/test_elastic.py +++ b/tests/forcefields/flows/test_elastic.py @@ -1,3 +1,4 @@ +import pytest from jobflow import run_locally from pymatgen.symmetry.analyzer import SpacegroupAnalyzer @@ -6,27 +7,28 @@ from atomate2.forcefields.jobs import MACERelaxMaker -def test_elastic_wf_with_mace(clean_dir, si_structure): +def test_elastic_wf_with_mace(clean_dir, si_structure, test_dir): si_prim = SpacegroupAnalyzer(si_structure).get_primitive_standard_structure() + model_path = f"{test_dir}/forcefields/mace/MACE.model" + common_kwds = dict( + model=model_path, + relax_kwargs={"fmax": 0.00001}, + model_kwargs={"default_dtype": "float64"}, + ) - job = ElasticMaker( - bulk_relax_maker=MACERelaxMaker( - relax_cell=True, relax_kwargs={"fmax": 0.00001} - ), - elastic_relax_maker=MACERelaxMaker( - relax_cell=False, relax_kwargs={"fmax": 0.00001} - ), + flow = ElasticMaker( + bulk_relax_maker=MACERelaxMaker(**common_kwds, relax_cell=True), + elastic_relax_maker=MACERelaxMaker(**common_kwds, relax_cell=False), ).make(si_prim) # run the flow or job and ensure that it finished running successfully - responses = run_locally(job, create_folders=True, ensure_success=True) - elastic_output = responses[job.jobs[-1].uuid][1].output + responses = run_locally(flow, create_folders=True, ensure_success=True) + elastic_output = responses[flow[-1].uuid][1].output assert isinstance(elastic_output, ElasticDocument) - # TODO (@janosh) uncomment below asserts once no longer failing with crazy values - # (3101805 instead of 118). started happening in v0.9.0 release of matgl. reached - # out to Shyue Ping and his group to look into this. - # assert_allclose(elastic_output.derived_properties.k_voigt, 74.698317, atol=1e-1) - # assert_allclose(elastic_output.derived_properties.g_voigt, 27.197747, atol=1e-1) - assert elastic_output.derived_properties.k_voigt > 0 - assert elastic_output.derived_properties.g_voigt > 0 + assert elastic_output.derived_properties.k_voigt == pytest.approx( + 9.7005429, abs=0.01 + ) + assert elastic_output.derived_properties.g_voigt == pytest.approx( + 0.002005039, abs=0.01 + ) assert elastic_output.chemsys == "Si" diff --git a/tests/test_data/vasp/Si_mp_gga_relax/GGA_Relax_1/inputs/INCAR b/tests/test_data/vasp/Si_mp_gga_relax/GGA_Relax_1/inputs/INCAR index 8368b4e030..6e1fbb6731 100644 --- a/tests/test_data/vasp/Si_mp_gga_relax/GGA_Relax_1/inputs/INCAR +++ b/tests/test_data/vasp/Si_mp_gga_relax/GGA_Relax_1/inputs/INCAR @@ -3,7 +3,7 @@ EDIFF = 0.0001 ENCUT = 520 IBRION = 2 ISIF = 3 -ISMEAR = 0 +ISMEAR = -5 ISPIN = 2 LASPH = True LORBIT = 11 @@ -14,4 +14,4 @@ MAGMOM = 2*0.6 NELM = 100 NSW = 99 PREC = Accurate -SIGMA = 0.2 +SIGMA = 0.05 diff --git a/tests/test_data/vasp/Si_mp_meta_gga_relax/pbesol_pre_relax/inputs/INCAR b/tests/test_data/vasp/Si_mp_meta_gga_relax/pbesol_pre_relax/inputs/INCAR index f682301652..189d6d470f 100644 --- a/tests/test_data/vasp/Si_mp_meta_gga_relax/pbesol_pre_relax/inputs/INCAR +++ b/tests/test_data/vasp/Si_mp_meta_gga_relax/pbesol_pre_relax/inputs/INCAR @@ -22,4 +22,4 @@ MAGMOM = 2*0.6 NELM = 200 NSW = 99 PREC = Accurate -SIGMA = 0.2 +SIGMA = 0.05 diff --git a/tests/test_data/vasp/Si_mp_meta_gga_relax/r2scan_relax/inputs/INCAR b/tests/test_data/vasp/Si_mp_meta_gga_relax/r2scan_relax/inputs/INCAR index 7a2c6583db..70878e6b66 100644 --- a/tests/test_data/vasp/Si_mp_meta_gga_relax/r2scan_relax/inputs/INCAR +++ b/tests/test_data/vasp/Si_mp_meta_gga_relax/r2scan_relax/inputs/INCAR @@ -5,7 +5,7 @@ ENAUG = 1360.0 ENCUT = 680.0 IBRION = 2 ISIF = 3 -ISMEAR = -5 +ISMEAR = 0 ISPIN = 2 KSPACING = 0.28253269576667883 LAECHG = True diff --git a/tests/vasp/flows/test_matpes.py b/tests/vasp/flows/test_matpes.py index 46940041a5..f6b98d8eb0 100644 --- a/tests/vasp/flows/test_matpes.py +++ b/tests/vasp/flows/test_matpes.py @@ -9,7 +9,7 @@ def test_matpes_static_flow_maker(mock_vasp, clean_dir, vasp_test_dir): - # map from job name to directory containing reference output files + # map from job name to directory containing reference input/output files pre_relax_dir = "matpes_static_flow/pbe_static" ref_paths = { "MatPES GGA static": pre_relax_dir, diff --git a/tests/vasp/flows/test_mp.py b/tests/vasp/flows/test_mp.py index 77e12bd5c8..930c83cad1 100644 --- a/tests/vasp/flows/test_mp.py +++ b/tests/vasp/flows/test_mp.py @@ -45,7 +45,7 @@ def test_mp_meta_gga_relax_custom_values( def test_mp_meta_gga_double_relax_static(mock_vasp, clean_dir, vasp_test_dir): - # map from job name to directory containing reference output files + # map from job name to directory containing reference input/output files pre_relax_dir = "Si_mp_meta_gga_relax/pbesol_pre_relax" ref_paths = { "MP pre-relax 1": pre_relax_dir, @@ -75,7 +75,7 @@ def test_mp_meta_gga_double_relax_static(mock_vasp, clean_dir, vasp_test_dir): def test_mp_gga_double_relax_static(mock_vasp, clean_dir, vasp_test_dir): - # map from job name to directory containing reference output files + # map from job name to directory containing reference input/output files pre_relax_dir = "Si_mp_gga_relax/GGA_Relax_1" ref_paths = { "MP GGA relax 1": pre_relax_dir, @@ -99,7 +99,7 @@ def test_mp_gga_double_relax_static(mock_vasp, clean_dir, vasp_test_dir): def test_mp_gga_double_relax(mock_vasp, clean_dir, vasp_test_dir): - # map from job name to directory containing reference output files + # map from job name to directory containing reference input/output files pre_relax_dir = "Si_mp_gga_relax/GGA_Relax_1" ref_paths = { "MP GGA relax 1": pre_relax_dir, diff --git a/tests/vasp/jobs/test_mp.py b/tests/vasp/jobs/test_mp.py index 6b828f5fa5..cfc7a5349f 100644 --- a/tests/vasp/jobs/test_mp.py +++ b/tests/vasp/jobs/test_mp.py @@ -44,7 +44,7 @@ def test_mp_meta_gga_relax_maker_default_values(): def test_mp_meta_gga_static_maker(mock_vasp, clean_dir, vasp_test_dir): - # map from job name to directory containing reference output files + # map from job name to directory containing reference input/output files ref_paths = { "MP meta-GGA static": "Si_mp_meta_gga_relax/r2scan_final_static", } @@ -71,7 +71,7 @@ def test_mp_meta_gga_static_maker(mock_vasp, clean_dir, vasp_test_dir): def test_mp_meta_gga_relax_maker(mock_vasp, clean_dir, vasp_test_dir): - # map from job name to directory containing reference output files + # map from job name to directory containing reference input/output files ref_paths = { "MP meta-GGA relax": "Si_mp_meta_gga_relax/r2scan_relax", } @@ -100,7 +100,7 @@ def test_mp_meta_gga_relax_maker(mock_vasp, clean_dir, vasp_test_dir): def test_mp_gga_relax_maker(mock_vasp, clean_dir, vasp_test_dir): - # map from job name to directory containing reference output files + # map from job name to directory containing reference input/output files ref_paths = { "MP GGA relax": "Si_mp_gga_relax/GGA_Relax_1", } diff --git a/tests/vasp/sets/test_mp.py b/tests/vasp/sets/test_mp.py index 086fb7b5c6..709e3e3c15 100644 --- a/tests/vasp/sets/test_mp.py +++ b/tests/vasp/sets/test_mp.py @@ -51,7 +51,7 @@ def test_mp_sets(set_generator: VaspInputGenerator) -> None: else "PBE" ) assert mp_set.inherit_incar is False - assert mp_set.auto_ismear is True + assert mp_set.auto_ismear is False assert mp_set.auto_kspacing is True assert mp_set.force_gamma is True assert mp_set.auto_lreal is False diff --git a/tests/vasp/test_powerups.py b/tests/vasp/test_powerups.py index 6b1164d2a3..6d61dfa7e3 100644 --- a/tests/vasp/test_powerups.py +++ b/tests/vasp/test_powerups.py @@ -63,3 +63,37 @@ def test_update_user_settings(powerup, attribute, settings): getattr(flow.jobs[1].function.__self__.input_set_generator, attribute) != settings ) + + +@pytest.mark.parametrize( + "powerup,settings", + [ + ("add_metadata_to_flow", {"mp-id": "mp-xxx"}), + ("add_metadata_to_flow", {"mp-id": "mp-161", "composition": "NaCl"}), + ], +) +def test_add_metadata_to_flow(powerup, settings): + powerup_func = getattr(powerups, powerup) + + # test flow + drm = DoubleRelaxMaker() + flow = drm.make(1) + flow = powerup_func(flow, settings) + assert ( + flow.jobs[0].function.__self__.task_document_kwargs["additional_fields"] + == settings + ) + + +@pytest.mark.parametrize( + "powerup, settings", + [("update_vasp_custodian_handlers", ())], +) +def test_update_vasp_custodian_handlers(powerup, settings): + powerup_func = getattr(powerups, powerup) + + # test flow + drm = DoubleRelaxMaker() + flow = drm.make(1) + flow = powerup_func(flow, settings) + assert flow.jobs[0].function.__self__.run_vasp_kwargs["handlers"] == settings diff --git a/tests/vasp/test_sets.py b/tests/vasp/test_sets.py index a4aefb0116..37afc6ece0 100644 --- a/tests/vasp/test_sets.py +++ b/tests/vasp/test_sets.py @@ -74,6 +74,10 @@ def test_user_incar_settings(): "NSW": 5_000, "PREC": 10, # wrong type, should be string. "SIGMA": 20, + "LDAUU": {"H": 5.0}, + "LDAUJ": {"H": 6.0}, + "LDAUL": {"H": 3.0}, + "LDAUTYPE": 2, } static_set_generator = StaticSetGenerator(user_incar_settings=uis) @@ -82,6 +86,8 @@ def test_user_incar_settings(): for key in uis: if isinstance(incar[key], str): assert incar[key].lower() == uis[key].lower() + elif isinstance(uis[key], dict): + assert incar[key] == [uis[key][str(site.specie)] for site in structure] else: assert incar[key] == uis[key] @@ -92,9 +98,9 @@ def test_user_incar_settings(): ("struct_no_magmoms", {}), ("struct_with_magmoms", {}), ("struct_with_spin", {}), - ("struct_no_magmoms", {"MAGMOM": [3.7, 0.8]}), - ("struct_with_magmoms", {"MAGMOM": [3.7, 0.8]}), - ("struct_with_spin", {"MAGMOM": [3.7, 0.8]}), + ("struct_no_magmoms", {"MAGMOM": {"Fe": 3.7, "O": 0.8}}), + ("struct_with_magmoms", {"MAGMOM": {"Fe": 3.7, "O": 0.8}}), + ("struct_with_spin", {"MAGMOM": {"Fe2+,spin=4": 3.7, "O2-,spin=0.63": 0.8}}), ], ) def test_incar_magmoms_precedence(structure, user_incar_settings, request) -> None: @@ -121,7 +127,9 @@ def test_incar_magmoms_precedence(structure, user_incar_settings, request) -> No has_struct_spin = getattr(structure.species[0], "spin", None) is not None if user_incar_settings: # case 1 - assert incar_magmom == user_incar_settings["MAGMOM"] + assert incar_magmom == [ + user_incar_settings["MAGMOM"][str(site.specie)] for site in structure + ] elif has_struct_magmom: # case 2 assert incar_magmom == structure.site_properties["magmom"] elif has_struct_spin: # case 3