From acb670d33b8e231cd302994f3253d3f6c1767125 Mon Sep 17 00:00:00 2001 From: naik-aakash Date: Sun, 22 Dec 2024 15:05:54 +0100 Subject: [PATCH 01/74] load yamls using postin > no need to copy to remote server --- src/autoplex/auto/rss/flows.py | 49 +++++++++++++++++++--------------- 1 file changed, 27 insertions(+), 22 deletions(-) diff --git a/src/autoplex/auto/rss/flows.py b/src/autoplex/auto/rss/flows.py index 00f6cb8f6..4e2c83f23 100644 --- a/src/autoplex/auto/rss/flows.py +++ b/src/autoplex/auto/rss/flows.py @@ -22,21 +22,41 @@ class RssMaker(Maker): path_to_default_config_parameters: Path | str | None Path to the default RSS configuration file 'rss_default_configuration.yaml'. If None, the default path will be used. + config_file: Path | str | None + Path to the configuration file that defines the setup parameters for the whole RSS workflow. + If not provided, the default file 'rss_default_configuration.yaml' will be used. """ name: str = "ml-driven rss" path_to_default_config_parameters: Path | str | None = None + config_file: Path | str | None = None + + def __post_init__(self) -> None: + """Ensure that custom configuration parameters are loaded when the maker is initialized.""" + rss_default_config_path = ( + self.path_to_default_config_parameters + or Path(__file__).absolute().parent / "rss_default_configuration.yaml" + ) + + yaml = YAML(typ="safe", pure=True) + + with open(rss_default_config_path) as f: + config = yaml.load(f) + + if self.config_file and os.path.exists(self.config_file): + with open(self.config_file) as f: + new_config = yaml.load(f) + config.update(new_config) + + self.config = config @job - def make(self, config_file: str | None = None, **kwargs): + def make(self, **kwargs): """ Make a rss workflow using the specified configuration file and additional keyword arguments. Parameters ---------- - config_file: str | None - Path to the configuration file that defines the setup parameters for the whole RSS workflow. - If not provided, the default file 'rss_default_configuration.yaml' will be used. kwargs: dict, optional Additional optional keyword arguments to customize the job execution. @@ -244,25 +264,10 @@ def make(self, config_file: str | None = None, **kwargs): kb_temp: float The temperature (in eV) for Boltzmann sampling. """ - rss_default_config_path = ( - self.path_to_default_config_parameters - or Path(__file__).absolute().parent / "rss_default_configuration.yaml" - ) - - yaml = YAML(typ="safe", pure=True) - - with open(rss_default_config_path) as f: - config = yaml.load(f) - - if config_file and os.path.exists(config_file): - with open(config_file) as f: - new_config = yaml.load(f) - config.update(new_config) - - config.update(kwargs) - self._process_hookean_paras(config) + self.config.update(kwargs) + self._process_hookean_paras(self.config) - config_params = config.copy() + config_params = self.config.copy() if "train_from_scratch" not in config_params: raise ValueError( From 354fbd6935fc486323d3f932dab787d70e73cc67 Mon Sep 17 00:00:00 2001 From: naik-aakash Date: Tue, 7 Jan 2025 18:28:37 +0100 Subject: [PATCH 02/74] add CONFIG as const to maker --- src/autoplex/auto/rss/flows.py | 35 +++++++++++++--------------------- 1 file changed, 13 insertions(+), 22 deletions(-) diff --git a/src/autoplex/auto/rss/flows.py b/src/autoplex/auto/rss/flows.py index 18f56a455..8529aa44b 100644 --- a/src/autoplex/auto/rss/flows.py +++ b/src/autoplex/auto/rss/flows.py @@ -5,10 +5,12 @@ from pathlib import Path from jobflow import Flow, Maker, Response, job -from ruamel.yaml import YAML +from monty.serialization import loadfn from autoplex.auto.rss.jobs import do_rss_iterations, initial_rss +MODULE_DIR = Path(os.path.dirname(__file__)) + @dataclass class RssMaker(Maker): @@ -19,36 +21,25 @@ class RssMaker(Maker): ---------- name: str Name of the flow. - path_to_default_config_parameters: Path | str | None - Path to the default RSS configuration file 'rss_default_configuration.yaml'. - If None, the default path will be used. config_file: Path | str | None Path to the configuration file that defines the setup parameters for the whole RSS workflow. If not provided, the default file 'rss_default_configuration.yaml' will be used. """ name: str = "ml-driven rss" - path_to_default_config_parameters: Path | str | None = None config_file: Path | str | None = None + CONFIG = loadfn(MODULE_DIR / "rss_default_configuration.yaml") def __post_init__(self) -> None: """Ensure that custom configuration parameters are loaded when the maker is initialized.""" - rss_default_config_path = ( - self.path_to_default_config_parameters - or Path(__file__).absolute().parent / "rss_default_configuration.yaml" - ) - - yaml = YAML(typ="safe", pure=True) - - with open(rss_default_config_path) as f: - config = yaml.load(f) - if self.config_file and os.path.exists(self.config_file): - with open(self.config_file) as f: - new_config = yaml.load(f) - config.update(new_config) + new_config = loadfn(self.config_file) - self.config = config + for key, value in new_config.items(): + if key in self.CONFIG and isinstance(value, self.CONFIG[key]): + self.CONFIG[key] = value + else: + raise ValueError(f"Invalid key '{key}' in the configuration file.") @job def make(self, **kwargs): @@ -255,10 +246,10 @@ def make(self, **kwargs): - 'current_iter': int, The current iteration index. - 'kb_temp': float, The temperature (in eV) for Boltzmann sampling. """ - self.config.update(kwargs) - self._process_hookean_paras(self.config) + self.CONFIG.update(kwargs) + self._process_hookean_paras(self.CONFIG) - config_params = self.config.copy() + config_params = self.CONFIG.copy() if "train_from_scratch" not in config_params: raise ValueError( From 0a90c9f6acf8d04f1849b0abfb2cbd57d83fe7f2 Mon Sep 17 00:00:00 2001 From: naik-aakash Date: Wed, 8 Jan 2025 17:36:19 +0100 Subject: [PATCH 03/74] add simple tests for RSSMaker __post_init__ --- src/autoplex/auto/rss/flows.py | 7 +- .../auto/rss/rss_default_configuration.yaml | 2 +- tests/auto/rss/test_flows.py | 19 ++- tests/test_data/rss/rss_config.yaml | 143 ++++++++++++++++++ 4 files changed, 165 insertions(+), 6 deletions(-) create mode 100644 tests/test_data/rss/rss_config.yaml diff --git a/src/autoplex/auto/rss/flows.py b/src/autoplex/auto/rss/flows.py index 8529aa44b..0d5d54056 100644 --- a/src/autoplex/auto/rss/flows.py +++ b/src/autoplex/auto/rss/flows.py @@ -22,7 +22,7 @@ class RssMaker(Maker): name: str Name of the flow. config_file: Path | str | None - Path to the configuration file that defines the setup parameters for the whole RSS workflow. + Path to the custom configuration file that defines the setup parameters for the whole RSS workflow. If not provided, the default file 'rss_default_configuration.yaml' will be used. """ @@ -36,10 +36,9 @@ def __post_init__(self) -> None: new_config = loadfn(self.config_file) for key, value in new_config.items(): - if key in self.CONFIG and isinstance(value, self.CONFIG[key]): + # TODO: Need better defaults in default file or we move to pydantic models + if key in self.CONFIG and isinstance(value, type(self.CONFIG[key])): self.CONFIG[key] = value - else: - raise ValueError(f"Invalid key '{key}' in the configuration file.") @job def make(self, **kwargs): diff --git a/src/autoplex/auto/rss/rss_default_configuration.yaml b/src/autoplex/auto/rss/rss_default_configuration.yaml index 89cdc713a..fd46d013d 100644 --- a/src/autoplex/auto/rss/rss_default_configuration.yaml +++ b/src/autoplex/auto/rss/rss_default_configuration.yaml @@ -1,5 +1,5 @@ # General Parameters -tag: +tag: '' train_from_scratch: true resume_from_previous_state: test_error: diff --git a/tests/auto/rss/test_flows.py b/tests/auto/rss/test_flows.py index da0093cbc..86d415992 100644 --- a/tests/auto/rss/test_flows.py +++ b/tests/auto/rss/test_flows.py @@ -1,8 +1,8 @@ import os -import pytest from pathlib import Path from jobflow import run_locally, Flow from tests.conftest import mock_rss, mock_do_rss_iterations, mock_do_rss_iterations_multi_jobs +from autoplex.auto.rss.flows import RssMaker os.environ["OMP_NUM_THREADS"] = "1" @@ -307,3 +307,20 @@ def test_mock_workflow_multi_node(test_dir, mock_vasp, memory_jobstore, clean_di selected_atoms = job2.output.resolve(memory_jobstore) assert len(selected_atoms) == 3 + +def test_rssmaker_custom_config(test_dir): + + # For now only test if __post_init is working and updating defaults + rss = RssMaker(config_file= test_dir / "rss" / "rss_config.yaml") + + # TODO: test needs to be more robust after updating default config files + assert rss.CONFIG["tag"] == "test" + assert rss.CONFIG["generated_struct_numbers"] == [9000, 1000] + assert rss.CONFIG["num_processes_buildcell"] == 64 + assert rss.CONFIG["num_processes_fit"] == 64 + assert rss.CONFIG["device_for_rss"] == "gpu" + assert rss.CONFIG["isolatedatom_box"] == [10, 10, 10] + assert rss.CONFIG["dimer_box"] == [10, 10, 10] + + + diff --git a/tests/test_data/rss/rss_config.yaml b/tests/test_data/rss/rss_config.yaml new file mode 100644 index 000000000..a146ade12 --- /dev/null +++ b/tests/test_data/rss/rss_config.yaml @@ -0,0 +1,143 @@ +# General Parameters +tag: "test" +train_from_scratch: true +resume_from_previous_state: + test_error: + pre_database_dir: + mlip_path: + isolated_atom_energies: + +# Buildcell Parameters +generated_struct_numbers: + - 9000 + - 1000 +buildcell_options: +fragment_file: null +fragment_numbers: null +num_processes_buildcell: 64 + +# Sampling Parameters +num_of_initial_selected_structs: + - 80 + - 20 +num_of_rss_selected_structs: 100 +initial_selection_enabled: true +rss_selection_method: 'bcur2i' +bcur_params: + soap_paras: + l_max: 12 + n_max: 12 + atom_sigma: 0.0875 + cutoff: 10.5 + cutoff_transition_width: 1.0 + zeta: 4.0 + average: true + species: true + frac_of_bcur: 0.8 + bolt_max_num: 3000 +random_seed: null + +# DFT Labelling Parameters +include_isolated_atom: true +isolatedatom_box: + - 10.0 + - 10.0 + - 10.0 +e0_spin: false +include_dimer: true +dimer_box: + - 10.0 + - 10.0 + - 10.0 +dimer_range: + - 1.0 + - 5.0 +dimer_num: 21 +custom_incar: + ISMEAR: 0 + SIGMA: 0.05 + PREC: 'Accurate' + ADDGRID: '.TRUE.' + EDIFF: 1e-7 + NELM: 250 + LWAVE: '.FALSE.' + LCHARG: '.FALSE.' + ALGO: 'Normal' + AMIX: null + LREAL: '.FALSE.' + ISYM: 0 + ENCUT: 520.0 + KSPACING: 0.20 + GGA: null + KPAR: 8 + NCORE: 16 + LSCALAPACK: '.FALSE.' + LPLANE: '.FALSE.' +custom_potcar: +vasp_ref_file: 'vasp_ref.extxyz' + +# Data Preprocessing Parameters +config_types: + - 'initial' + - 'traj_early' + - 'traj' +rss_group: + - 'traj' +test_ratio: 0.1 +regularization: true +scheme: 'linear-hull' +reg_minmax: + - [0.1, 1] + - [0.001, 0.1] + - [0.0316, 0.316] + - [0.0632, 0.632] +distillation: false +force_max: null +force_label: null +pre_database_dir: null + +# MLIP Parameters +mlip_type: 'GAP' +ref_energy_name: 'REF_energy' +ref_force_name: 'REF_forces' +ref_virial_name: 'REF_virial' +auto_delta: true +num_processes_fit: 64 +device_for_fitting: 'cpu' +##The following hyperparameters are only applicable to GAP. +##If you want to use other models, please replace the corresponding hyperparameters. +twob: + cutoff: 5.0 + n_sparse: 15 + theta_uniform: 1.0 +threeb: + cutoff: 3.0 +soap: + l_max: 10 + n_max: 10 + atom_sigma: 0.5 + n_sparse: 2500 + cutoff: 5.0 +general: + three_body: false + +# RSS Exploration Parameters +scalar_pressure_method: 'uniform' +scalar_exp_pressure: 1 +scalar_pressure_exponential_width: 0.2 +scalar_pressure_low: 0 +scalar_pressure_high: 25 +max_steps: 300 +force_tol: 0.01 +stress_tol: 0.01 +stop_criterion: 0.01 +max_iteration_number: 25 +num_groups: 6 +initial_kt: 0.3 +current_iter_index: 1 +hookean_repul: false +hookean_paras: +keep_symmetry: false +write_traj: true +num_processes_rss: 128 +device_for_rss: 'gpu' From fa842fe909142a351745170a2b2c3464407f2ea7 Mon Sep 17 00:00:00 2001 From: naik-aakash Date: Wed, 8 Jan 2025 17:49:14 +0100 Subject: [PATCH 04/74] update doc --- docs/user/rss/flow/quick_start/start.md | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/docs/user/rss/flow/quick_start/start.md b/docs/user/rss/flow/quick_start/start.md index bc38c3fd7..e13b77fae 100644 --- a/docs/user/rss/flow/quick_start/start.md +++ b/docs/user/rss/flow/quick_start/start.md @@ -28,12 +28,11 @@ Parameters can be specified either through a YAML configuration file or as direc > **Recommendation**: This is currently our recommended approach for setting up and managing RSS workflows. -The RSS workflow can be initiated using a custom YAML configuration file. A comprehensive list of parameters, including default settings and modifiable options, is available in `autoplex/auto/rss/rss_default_configuration.yaml`. When creating a new YAML file, any specified keys will override the corresponding default values. To start a new workflow, pass the path to your YAML file as the `config_file` argument in the `make` method. If you are using remote submission via jobflow-remote, please be aware that the configuration file has to be placed on the remote cluster and the file path has to reflect this as well (i.e., it is a path on the remote cluster). +The RSS workflow can be initiated using a custom YAML configuration file. A comprehensive list of parameters, including default settings and modifiable options, is available in `autoplex/auto/rss/rss_default_configuration.yaml`. When creating a new YAML file, any specified keys will override the corresponding default values. To start a new workflow, pass the path to your YAML file as the `config_file` argument in the `make` method. ```python from autoplex.auto.rss.flows import RssMaker from fireworks import LaunchPad -from jobflow import Flow from jobflow.managers.fireworks import flow_to_workflow rss_job = RssMaker(name="your workflow name").make(config_file='path/to/your/name.yaml') From 777fa8ca9a1e2298baf729a0f1b4841b1eaf8d1f Mon Sep 17 00:00:00 2001 From: naik-aakash Date: Thu, 9 Jan 2025 07:45:08 +0100 Subject: [PATCH 05/74] add strict Path checks --- src/autoplex/auto/rss/flows.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/src/autoplex/auto/rss/flows.py b/src/autoplex/auto/rss/flows.py index 0d5d54056..ad5f6457a 100644 --- a/src/autoplex/auto/rss/flows.py +++ b/src/autoplex/auto/rss/flows.py @@ -32,8 +32,9 @@ class RssMaker(Maker): def __post_init__(self) -> None: """Ensure that custom configuration parameters are loaded when the maker is initialized.""" - if self.config_file and os.path.exists(self.config_file): - new_config = loadfn(self.config_file) + if self.config_file and Path(self.config_file).resolve(strict=True): + print(Path(self.config_file).resolve()) + new_config = loadfn(Path(self.config_file).resolve()) for key, value in new_config.items(): # TODO: Need better defaults in default file or we move to pydantic models From 3ddfb143c91a9fe2573127814b630b80de779c0b Mon Sep 17 00:00:00 2001 From: naik-aakash Date: Thu, 9 Jan 2025 07:46:21 +0100 Subject: [PATCH 06/74] make limitations of config file clear and update code snippets --- docs/user/rss/flow/quick_start/start.md | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/docs/user/rss/flow/quick_start/start.md b/docs/user/rss/flow/quick_start/start.md index e13b77fae..86762ee8a 100644 --- a/docs/user/rss/flow/quick_start/start.md +++ b/docs/user/rss/flow/quick_start/start.md @@ -28,14 +28,19 @@ Parameters can be specified either through a YAML configuration file or as direc > **Recommendation**: This is currently our recommended approach for setting up and managing RSS workflows. -The RSS workflow can be initiated using a custom YAML configuration file. A comprehensive list of parameters, including default settings and modifiable options, is available in `autoplex/auto/rss/rss_default_configuration.yaml`. When creating a new YAML file, any specified keys will override the corresponding default values. To start a new workflow, pass the path to your YAML file as the `config_file` argument in the `make` method. +The RSS workflow can be initiated using a custom YAML configuration file. +A comprehensive list of parameters, including default settings and modifiable options, is available in `autoplex/auto/rss/rss_default_configuration.yaml`. +When initializing the flow with a YAML file, any specified keys will override the corresponding default values. +To start a new workflow, pass the path to your YAML file as the `config_file` argument in the `RssMaker`. + +> ℹ️ Note that, if you are using `RssMaker` as part of another job or flow and executing jobs on a different system (e.g. using remote submission via jobflow-remote), then the configuration file has to be placed on the remote cluster and the file path has to reflect this as (i.e., it is valid a path on the remote cluster). ```python from autoplex.auto.rss.flows import RssMaker from fireworks import LaunchPad from jobflow.managers.fireworks import flow_to_workflow -rss_job = RssMaker(name="your workflow name").make(config_file='path/to/your/name.yaml') +rss_job = RssMaker(name="your workflow name", config_file='path/to/your/config.yaml').make() wf = flow_to_workflow(rss_job) lpad = LaunchPad.auto_load() lpad.add_wf(wf) @@ -48,7 +53,7 @@ The above code is based on [`FireWorks`](https://materialsproject.github.io/fire from autoplex.auto.rss.flows import RssMaker from jobflow_remote import submit_flow -rss_job = RssMaker(name="your workflow name").make(config_file='path/to/your/name.yaml') +rss_job = RssMaker(name="your workflow name", config_file='path/to/your/name.yaml').make() resources = {"nodes": N, "partition": "name", "qos": "name", "time": "8:00:00", "mail_user": "your_email", "mail_type": "ALL", "account": "your account"} print(submit_flow(rss_job, worker="your worker", resources=resources, project="your project name")) ``` From 52a8fdbb4156af4248ba68500c8f198079f60bcc Mon Sep 17 00:00:00 2001 From: naik-aakash Date: Thu, 9 Jan 2025 07:47:42 +0100 Subject: [PATCH 07/74] remove print --- src/autoplex/auto/rss/flows.py | 1 - 1 file changed, 1 deletion(-) diff --git a/src/autoplex/auto/rss/flows.py b/src/autoplex/auto/rss/flows.py index ad5f6457a..4b7d3be2c 100644 --- a/src/autoplex/auto/rss/flows.py +++ b/src/autoplex/auto/rss/flows.py @@ -33,7 +33,6 @@ class RssMaker(Maker): def __post_init__(self) -> None: """Ensure that custom configuration parameters are loaded when the maker is initialized.""" if self.config_file and Path(self.config_file).resolve(strict=True): - print(Path(self.config_file).resolve()) new_config = loadfn(Path(self.config_file).resolve()) for key, value in new_config.items(): From 0205e7d8c72a8356c5e2709997818e57af18e5db Mon Sep 17 00:00:00 2001 From: naik-aakash Date: Thu, 9 Jan 2025 16:21:12 +0100 Subject: [PATCH 08/74] add initial pydantic config for RSS --- src/autoplex/settings.py | 383 +++++++++++++++++++++++++++++++++++++++ 1 file changed, 383 insertions(+) create mode 100644 src/autoplex/settings.py diff --git a/src/autoplex/settings.py b/src/autoplex/settings.py new file mode 100644 index 000000000..b3ee19051 --- /dev/null +++ b/src/autoplex/settings.py @@ -0,0 +1,383 @@ +"""Settings for autoplex.""" + +from __future__ import annotations + +from typing import Literal + +from pydantic import BaseModel, Field + + +class ResumeFromPreviousState(BaseModel): + """ + A model describing the state information. + + It is required to resume a previously interrupted or saved RSS workflow. + When 'train_from_scratch' is set to False, this parameter is mandatory + for the workflow to pick up from a saved state. + """ + + test_error: float | None = Field( + default=None, + description="The test error from the last completed training step.", + ) + pre_database_dir: str | None = Field( + default=None, + description="Path to the directory containing the pre-existing database for resuming", + ) + mlip_path: str | None = Field( + default=None, description="Path to the file of a previous MLIP model." + ) + isolated_atom_energies: dict | None = Field( + default=None, + description="A dictionary with isolated atom energy values mapped to atomic numbers", + ) + + +class SoapParas(BaseModel): + """A model describing the SOAP parameters.""" + + l_max: int = Field(default=12, description="Maximum degree of spherical harmonics") + n_max: int = Field( + default=12, description="Maximum number of radial basis functions" + ) + atom_sigma: float = Field(default=0.0875, description="idth of Gaussian smearing") + cutoff: float = Field(default=10.5, description="Radial cutoff distance") + cutoff_transition_width: float = Field( + default=1.0, description="Width of the transition region for the cutoff" + ) + zeta: float = Field(default=4.0, description="Exponent for dot-product SOAP kernel") + average: bool = Field( + default=True, description="Whether to average the SOAP vectors" + ) + species: bool = Field( + default=True, description="Whether to consider species information" + ) + + +class BcurParams(BaseModel): + """A model describing the parameters for the BCUR method.""" + + soap_paras: SoapParas = Field(default_factory=SoapParas) + frac_of_bcur: float = Field( + default=0.8, description="Fraction of Boltzmann CUR selections" + ) + bolt_max_num: int = Field( + default=3000, description="Maximum number of Boltzmann selections" + ) + + +class BuildcellOptions(BaseModel): + """A model describing the parameters for buildcell.""" + + NFORM: str | None = Field(default=None, description="The number of formula units") + SYMMOPS: str | None = Field(default=None, description="The symmetry operations") + SLACK: float | None = Field(default=None, description="The slack factor") + OVERLAP: float | None = Field(default=None, description="The overlap factor") + MINSEP: str | None = Field(default=None, description="The minimum separation") + + +class CustomIncar(BaseModel): + """A model describing the INCAR parameters.""" + + ISMEAR: int = 0 + SIGMA: float = 0.05 + PREC: str = "Accurate" + ADDGRID: str = ".TRUE." + EDIFF: float = 1e-07 + NELM: int = 250 + LWAVE: str = ".FALSE." + LCHARG: str = ".FALSE." + ALGO: str = "Normal" + AMIX: float | None = None + LREAL: str = ".FALSE." + ISYM: int = 0 + ENCUT: float = 520.0 + KSPACING: float = 0.2 + GGA: str | None = None + KPAR: int = 8 + NCORE: int = 16 + LSCALAPACK: str = ".FALSE." + LPLANE: str = ".FALSE." + + +class Twob(BaseModel): + """A model describing the two-body GAP parameters.""" + + cutoff: float = Field(default=5.0, description="Radial cutoff distance") + n_sparse: int = Field(default=15, description="Number of sparse points") + theta_uniform: float = Field( + default=1.0, description="Width of the uniform distribution for theta" + ) + + +class Threeb(BaseModel): + """A model describing the three-body GAP parameters.""" + + cutoff: float = Field(default=3.0, description="Radial cutoff distance") + + +class Soap(BaseModel): + """A model describing the SOAP GAP parameters.""" + + l_max: int = Field(default=10, description="Maximum degree of spherical harmonics") + n_max: int = Field( + default=10, description="Maximum number of radial basis functions" + ) + atom_sigma: float = Field(default=0.5, description="Width of Gaussian smearing") + n_sparse: int = Field(default=2500, description="Number of sparse points") + cutoff: float = Field(default=5.0, description="Radial cutoff distance") + + +class General(BaseModel): + """A model describing the general GAP parameters.""" + + three_body: bool = Field( + default=False, description="Whether to include three-body terms" + ) + + +class RssConfig(BaseModel): + """A model describing the complete RSS configuration.""" + + tag: str | None = Field( + default=None, + description="Tag of systems. It can also be used for setting up elements " + "and stoichiometry. For example, the tag of 'SiO2' will be recognized " + "as a 1:2 ratio of Si to O and passed into the parameters of buildcell. " + "However, note that this will be overwritten if the stoichiometric ratio " + "of elements is defined in the 'buildcell_options'", + ) + train_from_scratch: bool = Field( + default=True, + description="If True, it starts the workflow from scratch " + "If False, it resumes from a previous state.", + ) + resume_from_previous_state: ResumeFromPreviousState = Field( + default_factory=ResumeFromPreviousState + ) + generated_struct_numbers: list[int, int] = Field( + default_factory=lambda: [8000, 2000], + description="Expected number of generated " + "randomized unit cells by buildcell.", + ) + buildcell_options: list[BuildcellOptions] | None = Field( + default=None, description="Customized parameters for buildcell." + ) + fragment_file: str | None = Field(default=None, description="") + fragment_numbers: list[int] | None = Field( + default=None, + description=" Numbers of each fragment to be included in the random structures. " + "Defaults to 1 for all specified.", + ) + num_processes_buildcell: int = Field( + default=128, description="Number of processes for buildcell." + ) + num_of_initial_selected_structs: list[int, int] = Field( + default_factory=lambda: [80, 20], + description="Number of structures to be sampled directly " + "from the buildcell-generated randomized cells.", + ) + num_of_rss_selected_structs: int = Field( + default=100, + description="Number of structures to be selected from each RSS iteration.", + ) + initial_selection_enabled: bool = Field( + default=True, + description="If true, sample structures from initially generated " + "randomized cells using CUR.", + ) + rss_selection_method: Literal["bcur1s", "bcur2i", None] = Field( + default="bcur2i", + description="Method for selecting samples from the RSS trajectories: " + "Boltzmann flat histogram in enthalpy first, then CUR. Options are as follows", + ) + bcur_params: BcurParams = Field( + default_factory=BcurParams, description="Parameters for the BCUR method." + ) + random_seed: int | None = Field( + default=None, description="A seed to ensure reproducibility of CUR selection." + ) + include_isolated_atom: bool = Field( + default=True, + description="Perform single-point calculations for isolated atoms.", + ) + isolatedatom_box: list[float, float, float] = Field( + default_factory=lambda: [20.0, 20.0, 20.0], + description="List of the lattice constants for an " + "isolated atom configuration.", + ) + e0_spin: bool = Field( + default=False, + description="Include spin polarization in isolated atom and dimer calculations", + ) + include_dimer: bool = Field( + default=True, + description="Perform single-point calculations for dimers only once", + ) + dimer_box: list[float, float, float] = Field( + default_factory=lambda: [20.0, 20.0, 20.0], + description="The lattice constants of a dimer box.", + ) + dimer_range: list[float, float] = Field( + default_factory=lambda: [1.0, 5.0], + description="The range of the dimer distance.", + ) + dimer_num: int = Field( + default=21, + description="Number of different distances to consider for dimer calculations.", + ) + custom_incar: CustomIncar = Field( + default_factory=CustomIncar, + description="Custom VASP input parameters. " + "If provided, will update the default parameters", + ) + custom_potcar: str | None = Field( + default=None, + description="POTCAR settings to update. Keys are element symbols, " + "values are the desired POTCAR labels.", + ) + vasp_ref_file: str = Field( + default="vasp_ref.extxyz", description="Reference file for VASP data" + ) + config_types: list[str] = Field( + default_factory=lambda: ["initial", "traj_early", "traj"], + description="Configuration types for the VASP calculations", + ) + rss_group: list[str] = Field( + default_factory=lambda: ["traj"], + description="Group of configurations for the RSS calculations", + ) + test_ratio: float = Field( + default=0.1, + description="The proportion of the test set after splitting the data", + ) + regularization: bool = Field( + default=True, + description="Whether to apply regularization. This only works for GAP to date.", + ) + scheme: Literal["linear-hull", "volume-stoichiometry", None] = Field( + default="linear-hull", description="Method to use for regularization" + ) + reg_minmax: list[list[float]] = Field( + default_factory=lambda: [ + [0.1, 1], + [0.001, 0.1], + [0.0316, 0.316], + [0.0632, 0.632], + ], + description="List of tuples of (min, max) values for energy, force, " + "virial sigmas for regularization", + ) + distillation: bool = Field( + default=False, description="Whether to apply data distillation" + ) + force_max: float | None = Field( + default=None, description="Maximum force value to exclude structures" + ) + force_label: str | None = Field( + default=None, description="The label of force values to use for distillation" + ) + pre_database_dir: str | None = Field( + default=None, description="Directory where the previous database was saved." + ) + mlip_type: Literal["GAP", "J-ACE", "NEQUIP", "M3GNET", "MACE"] = Field( + default="GAP", description="MLIP to be fitted" + ) + ref_energy_name: str = Field( + default="REF_energy", description="Reference energy name." + ) + ref_force_name: str = Field( + default="REF_forces", description="Reference force name." + ) + ref_virial_name: str = Field( + default="REF_virial", description="Reference virial name." + ) + auto_delta: bool = Field( + default=True, + description="Whether to automatically calculate the delta value for GAP terms.", + ) + num_processes_fit: int = Field( + default=32, description="Number of processes used for fitting" + ) + device_for_fitting: Literal["cpu", "cuda"] = Field( + default="cpu", description="Device to be used for model fitting" + ) + twob: Twob = Field( + default_factory=Twob, + description="Parameters for the two-body descriptor, Applicable on to GAP", + ) + threeb: Threeb = Field( + default_factory=Threeb, + description="Parameters for the three-body descriptor, Applicable on to GAP", + ) + soap: Soap = Field( + default_factory=Soap, + description="Parameters for the SOAP descriptor, Applicable on to GAP", + ) + general: General = Field( + default_factory=General, description="General parameters for the GAP model" + ) + scalar_pressure_method: Literal["exp", "uniform"] = Field( + default="uniform", description="Method for adding external pressures." + ) + scalar_exp_pressure: int = Field( + default=1, description="Scalar exponential pressure" + ) + scalar_pressure_exponential_width: float = Field( + default=0.2, description="Width for scalar pressure exponential" + ) + scalar_pressure_low: int = Field( + default=0, description="Lower limit for scalar pressure" + ) + scalar_pressure_high: int = Field( + default=25, description="Upper limit for scalar pressure" + ) + max_steps: int = Field( + default=300, description="Maximum number of steps for the GAP optimization" + ) + force_tol: float = Field( + default=0.01, description="Force residual tolerance for relaxation" + ) + stress_tol: float = Field( + default=0.01, description="Stress residual tolerance for relaxation." + ) + stop_criterion: float = Field( + default=0.01, description="Convergence criterion for stopping RSS iterations." + ) + max_iteration_number: int = Field( + default=25, description="Maximum number of RSS iterations to perform." + ) + num_groups: int = Field( + default=6, + description="Number of structure groups, used for assigning tasks across multiple nodes." + "For example, if there are 10,000 trajectories to relax and 'num_groups=10'," + "the trajectories will be divided into 10 groups and 10 independent jobs will be created," + "with each job handling 1,000 trajectories.", + ) + initial_kt: float = Field( + default=0.3, description="Initial temperature (in eV) for Boltzmann sampling." + ) + current_iter_index: int = Field( + default=1, description="Current iteration index for the RSS." + ) + hookean_repul: bool = Field( + default=False, description="Whether to apply Hookean repulsion" + ) + hookean_paras: dict | None = Field( + default=None, + description="Parameters for the Hookean repulsion as a " + "dictionary of tuples.", + ) + keep_symmetry: bool = Field( + default=False, description="Whether to preserve symmetry during relaxations." + ) + write_traj: bool = Field( + default=True, + description="Bool indicating whether to write the trajectory files.", + ) + num_processes_rss: int = Field( + default=128, description="Number of processes used for running RSS." + ) + device_for_rss: Literal["cpu", "cuda"] = Field( + default="cpu", description="Device to be used for RSS calculations." + ) From 5bd40bb8ee2d0931172cba3b25b13b089c8e3c29 Mon Sep 17 00:00:00 2001 From: naik-aakash Date: Thu, 9 Jan 2025 16:24:47 +0100 Subject: [PATCH 09/74] update test --- tests/auto/rss/test_flows.py | 28 +++++++++++++++++----------- 1 file changed, 17 insertions(+), 11 deletions(-) diff --git a/tests/auto/rss/test_flows.py b/tests/auto/rss/test_flows.py index 86d415992..3fbb55062 100644 --- a/tests/auto/rss/test_flows.py +++ b/tests/auto/rss/test_flows.py @@ -1,7 +1,10 @@ import os from pathlib import Path -from jobflow import run_locally, Flow +from jobflow import run_locally, Flow, job +from mace.tools.model_script_utils import configure_model + from tests.conftest import mock_rss, mock_do_rss_iterations, mock_do_rss_iterations_multi_jobs +from autoplex.settings import RssConfig from autoplex.auto.rss.flows import RssMaker os.environ["OMP_NUM_THREADS"] = "1" @@ -310,17 +313,20 @@ def test_mock_workflow_multi_node(test_dir, mock_vasp, memory_jobstore, clean_di def test_rssmaker_custom_config(test_dir): - # For now only test if __post_init is working and updating defaults - rss = RssMaker(config_file= test_dir / "rss" / "rss_config.yaml") + from monty.serialization import loadfn + + rss_config = loadfn(test_dir / "rss" / "rss_config.yaml") - # TODO: test needs to be more robust after updating default config files - assert rss.CONFIG["tag"] == "test" - assert rss.CONFIG["generated_struct_numbers"] == [9000, 1000] - assert rss.CONFIG["num_processes_buildcell"] == 64 - assert rss.CONFIG["num_processes_fit"] == 64 - assert rss.CONFIG["device_for_rss"] == "gpu" - assert rss.CONFIG["isolatedatom_box"] == [10, 10, 10] - assert rss.CONFIG["dimer_box"] == [10, 10, 10] + config_model = RssConfig(**rss_config) + # Test if config is updated as expected + rss = RssMaker(config=config_model) + assert rss.config.tag == "test" + assert rss.config.generated_struct_numbers == [9000, 1000] + assert rss.config.num_processes_buildcell == 64 + assert rss.config.num_processes_fit == 64 + assert rss.config.device_for_rss == "cuda" + assert rss.config.isolatedatom_box == [10, 10, 10] + assert rss.config.dimer_box == [10, 10, 10] From 2982dba972392a42c8d1339b51fa9d7de65714ab Mon Sep 17 00:00:00 2001 From: naik-aakash Date: Thu, 9 Jan 2025 16:25:47 +0100 Subject: [PATCH 10/74] update test config file --- tests/test_data/rss/rss_config.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_data/rss/rss_config.yaml b/tests/test_data/rss/rss_config.yaml index a146ade12..3eca4f9a7 100644 --- a/tests/test_data/rss/rss_config.yaml +++ b/tests/test_data/rss/rss_config.yaml @@ -140,4 +140,4 @@ hookean_paras: keep_symmetry: false write_traj: true num_processes_rss: 128 -device_for_rss: 'gpu' +device_for_rss: 'cuda' From 4ee2e2c72dc3fa6781a7cad221a75eedb1f600b1 Mon Sep 17 00:00:00 2001 From: naik-aakash Date: Thu, 9 Jan 2025 16:26:20 +0100 Subject: [PATCH 11/74] restore rss_default config to initial --- src/autoplex/auto/rss/rss_default_configuration.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/autoplex/auto/rss/rss_default_configuration.yaml b/src/autoplex/auto/rss/rss_default_configuration.yaml index fd46d013d..89cdc713a 100644 --- a/src/autoplex/auto/rss/rss_default_configuration.yaml +++ b/src/autoplex/auto/rss/rss_default_configuration.yaml @@ -1,5 +1,5 @@ # General Parameters -tag: '' +tag: train_from_scratch: true resume_from_previous_state: test_error: From 3e77720b8e1efaf93402d20580fcf1276b6bc004 Mon Sep 17 00:00:00 2001 From: naik-aakash Date: Thu, 9 Jan 2025 16:27:26 +0100 Subject: [PATCH 12/74] remove config_file arg and use pydantic model instead --- src/autoplex/auto/rss/flows.py | 34 ++++++++++++---------------------- 1 file changed, 12 insertions(+), 22 deletions(-) diff --git a/src/autoplex/auto/rss/flows.py b/src/autoplex/auto/rss/flows.py index 4b7d3be2c..676d0e2cf 100644 --- a/src/autoplex/auto/rss/flows.py +++ b/src/autoplex/auto/rss/flows.py @@ -1,13 +1,13 @@ """RSS (random structure searching) flow for exploring and learning potential energy surfaces from scratch.""" import os -from dataclasses import dataclass +from dataclasses import dataclass, field from pathlib import Path from jobflow import Flow, Maker, Response, job -from monty.serialization import loadfn from autoplex.auto.rss.jobs import do_rss_iterations, initial_rss +from autoplex.settings import RssConfig MODULE_DIR = Path(os.path.dirname(__file__)) @@ -21,24 +21,13 @@ class RssMaker(Maker): ---------- name: str Name of the flow. - config_file: Path | str | None - Path to the custom configuration file that defines the setup parameters for the whole RSS workflow. - If not provided, the default file 'rss_default_configuration.yaml' will be used. + config: RssConfig | None + Pydantic model that defines the setup parameters for the whole RSS workflow. + If not provided, the defaults from 'autoplex.settings.RssConfig' will be used. """ name: str = "ml-driven rss" - config_file: Path | str | None = None - CONFIG = loadfn(MODULE_DIR / "rss_default_configuration.yaml") - - def __post_init__(self) -> None: - """Ensure that custom configuration parameters are loaded when the maker is initialized.""" - if self.config_file and Path(self.config_file).resolve(strict=True): - new_config = loadfn(Path(self.config_file).resolve()) - - for key, value in new_config.items(): - # TODO: Need better defaults in default file or we move to pydantic models - if key in self.CONFIG and isinstance(value, type(self.CONFIG[key])): - self.CONFIG[key] = value + config: RssConfig | None = field(default_factory=lambda: RssConfig()) @job def make(self, **kwargs): @@ -75,7 +64,7 @@ def make(self, **kwargs): buildcell_options: list[dict] | None Customized parameters for buildcell. Default is None. fragment: Atoms | list[Atoms] | None - Fragment(s) for random structures, e.g., molecules, to be placed indivudally intact. + Fragment(s) for random structures, e.g., molecules, to be placed individually intact. atoms.arrays should have a 'fragment_id' key with unique identifiers for each fragment if in same Atoms. atoms.cell must be defined (e.g., Atoms.cell = np.eye(3)*20). fragment_numbers: list[str] | None @@ -245,10 +234,10 @@ def make(self, **kwargs): - 'current_iter': int, The current iteration index. - 'kb_temp': float, The temperature (in eV) for Boltzmann sampling. """ - self.CONFIG.update(kwargs) - self._process_hookean_paras(self.CONFIG) + updated_config = self.config.model_copy(update=kwargs) + config_params = updated_config.model_dump() - config_params = self.CONFIG.copy() + self._process_hookean_paras(config_params) if "train_from_scratch" not in config_params: raise ValueError( @@ -345,7 +334,8 @@ def make(self, **kwargs): return Response(replace=Flow(rss_flow), output=do_rss_job.output) - def _process_hookean_paras(self, config): + @staticmethod + def _process_hookean_paras(config): if "hookean_paras" in config: config["hookean_paras"] = { tuple(map(int, k.strip("()").split(", "))): tuple(v) From 58e6b2a13e4a5feced9d7dc3d3df3c9918414531 Mon Sep 17 00:00:00 2001 From: naik-aakash Date: Thu, 9 Jan 2025 17:39:03 +0100 Subject: [PATCH 13/74] add MLIP hypers models --- src/autoplex/settings.py | 322 ++++++++++++++++++++++++++++++++++++++- 1 file changed, 321 insertions(+), 1 deletion(-) diff --git a/src/autoplex/settings.py b/src/autoplex/settings.py index b3ee19051..94cedb82b 100644 --- a/src/autoplex/settings.py +++ b/src/autoplex/settings.py @@ -7,11 +7,331 @@ from pydantic import BaseModel, Field +class GeneralSettings(BaseModel): + """Model describing general hyperparameters for the GAP.""" + + at_file: str = Field( + default="train.extxyz", description="Name of the training file" + ) + default_sigma: str = Field( + default="{0.0001 0.05 0.05 0}", description="Default sigma values" + ) + energy_parameter_name: str = Field( + default="REF_energy", description="Name of the energy parameter" + ) + force_parameter_name: str = Field( + default="REF_forces", description="Name of the force parameter" + ) + virial_parameter_name: str = Field( + default="REF_virial", description="Name of the virial parameter" + ) + sparse_jitter: float = Field(default=1.0e-8, description="Sparse jitter") + do_copy_at_file: str = Field(default="F", description="Copy the training file to") + openmp_chunk_size: int = Field(default=10000, description="OpenMP chunk size") + gp_file: str = Field(default="gap_file.xml", description="Name of the GAP file") + e0_offset: float = Field(default=0.0, description="E0 offset") + two_body: bool = Field( + default=False, description="Whether to include two-body terms" + ) + three_body: bool = Field( + default=False, description="Whether to include three-body terms" + ) + soap: bool = Field(default=True, description="Whether to include SOAP terms") + + +class TwobSettings(BaseModel): + """Model describing two body hyperparameters for the GAP.""" + + distance_Nb_order: int = Field( + default=2, description="Distance_Nb order for two-body" + ) + f0: float = Field(default=0.0, description="F0 value for two-body") + add_species: str = Field( + default="T", description="Whether to add species information" + ) + cutoff: float = Field(default=5.0, description="Radial cutoff distance") + n_sparse: int = Field(default=15, description="Number of sparse points") + covariance_type: str = Field( + default="ard_se", description="Covariance type for two-body" + ) + delta: float = Field(default=2.00, description="Delta value for two-body") + theta_uniform: float = Field( + default=0.5, description="Width of the uniform distribution for theta" + ) + sparse_method: str = Field( + default="uniform", description="Sparse method for two-body" + ) + compact_clusters: str = Field( + default="T", description="Whether to compact clusters" + ) + + +class ThreebSettings(BaseModel): + """Model describing threebody hyperparameters for the GAP.""" + + distance_Nb_order: int = Field( + default=3, description="Distance_Nb order for three-body" + ) + f0: float = Field(default=0.0, description="F0 value for three-body") + add_species: str = Field( + default="T", description="Whether to add species information" + ) + cutoff: float = Field(default=3.25, description="Radial cutoff distance") + n_sparse: int = Field(default=100, description="Number of sparse points") + covariance_type: str = Field( + default="ard_se", description="Covariance type for three-body" + ) + delta: float = Field(default=2.00, description="Delta value for three-body") + theta_uniform: float = Field( + default=1.0, description="Width of the uniform distribution for theta" + ) + sparse_method: str = Field( + default="uniform", description="Sparse method for three-body" + ) + compact_clusters: str = Field( + default="T", description="Whether to compact clusters" + ) + + +class SoapSettings(BaseModel): + """Model describing soap hyperparameters for the GAP.""" + + add_species: str = Field( + default="T", description="Whether to add species information" + ) + l_max: int = Field(default=10, description="Maximum degree of spherical harmonics") + n_max: int = Field( + default=12, description="Maximum number of radial basis functions" + ) + atom_sigma: float = Field(default=0.5, description="Width of Gaussian smearing") + zeta: int = Field(default=4, description="Exponent for dot-product SOAP kernel") + cutoff: float = Field(default=5.0, description="Radial cutoff distance") + cutoff_transition_width: float = Field( + default=1.0, description="Width of the transition region for the cutoff" + ) + central_weight: float = Field(default=1.0, description="Weight for central atom") + n_sparse: int = Field(default=6000, description="Number of sparse points") + delta: float = Field(default=1.00, description="Delta value for SOAP") + f0: float = Field(default=0.0, description="F0 value for SOAP") + covariance_type: str = Field( + default="dot_product", description="Covariance type for SOAP" + ) + sparse_method: str = Field( + default="cur_points", description="Sparse method for SOAP" + ) + + +class GAPSettings(BaseModel): + """Model describing the hyperparameters for the GAP fits for Phonons.""" + + general: GeneralSettings = Field( + default_factory=GeneralSettings, + description="General hyperparameters for the GAP fits", + ) + twob: TwobSettings = Field( + default_factory=TwobSettings, + description="Two body hyperparameters for the GAP fits", + ) + threeb: ThreebSettings = Field( + default_factory=ThreebSettings, + description="Three body hyperparameters for the GAP fits", + ) + soap: SoapSettings = Field( + default_factory=SoapSettings, + description="Soap hyperparameters for the GAP fits", + ) + + +class JACESettings(BaseModel): + """Model describing the hyperparameters for the J-ACE fits.""" + + order: int = Field(default=3, description="Order of the J-ACE model") + totaldegree: int = Field(default=6, description="Total degree of the J-ACE model") + cutoff: float = Field(default=2.0, description="Radial cutoff distance") + solver: str = Field(default="BLR", description="Solver for the J-ACE model") + + +class NEQUIPSettings(BaseModel): + """Model describing the hyperparameters for the NEQUIP fits.""" + + r_max: float = Field(default=4.0, description="Radial cutoff distance") + num_layers: int = Field(default=4, description="Number of layers") + l_max: int = Field(default=2, description="Maximum degree of spherical harmonics") + num_features: int = Field(default=32, description="Number of features") + num_basis: int = Field(default=8, description="Number of basis functions") + invariant_layers: int = Field(default=2, description="Number of invariant layers") + invariant_neurons: int = Field( + default=64, description="Number of invariant neurons" + ) + batch_size: int = Field(default=5, description="Batch size") + learning_rate: float = Field(default=0.005, description="Learning rate") + max_epochs: int = Field(default=10000, description="Maximum number of epochs") + default_dtype: str = Field(default="float32", description="Default data type") + + +class M3GNETSettings(BaseModel): + """Model describing the hyperparameters for the M3GNET fits.""" + + exp_name: str = Field(default="training", description="Name of the experiment") + results_dir: str = Field( + default="m3gnet_results", description="Directory to save the results" + ) + cutoff: float = Field(default=5.0, description="Radial cutoff distance") + threebody_cutoff: float = Field( + default=4.0, description="Three-body cutoff distance" + ) + batch_size: int = Field(default=10, description="Batch size") + max_epochs: int = Field(default=1000, description="Maximum number of epochs") + include_stresses: bool = Field( + default=True, description="Whether to include stresses" + ) + hidden_dim: int = Field(default=128, description="Hidden dimension") + num_units: int = Field(default=128, description="Number of units") + max_l: int = Field(default=4, description="Maximum degree of spherical harmonics") + max_n: int = Field( + default=4, description="Maximum number of radial basis functions" + ) + test_equal_to_val: bool = Field( + default=True, description="Whether the test set is equal to the validation set" + ) + + +class MACESettings(BaseModel): + """Model describing the hyperparameters for the MACE fits.""" + + model: str = Field(default="MACE", description="type of the model") + name: str = Field(default="MACE_model", description="Name of the model") + config_type_weights: str = Field( + default="{'Default':1.0}", description="Weights for the configuration types" + ) + hidden_irreps: str = Field(default="128x0e + 128x1o", description="Hidden irreps") + r_max: float = Field(default=5.0, description="Radial cutoff distance") + batch_size: int = Field(default=10, description="Batch size") + max_num_epochs: int = Field(default=1500, description="Maximum number of epochs") + start_swa: int = Field(default=1200, description="Start of the SWA") + ema_decay: float = Field( + default=0.99, description="Exponential moving average decay" + ) + correlation: int = Field(default=3, description="Correlation") + loss: str = Field(default="huber", description="Loss function") + default_dtype: str = Field(default="float32", description="Default data type") + swa: bool = Field(default=True, description="Whether to use SWA") + ema: bool = Field(default=True, description="Whether to use EMA") + amsgrad: bool = Field(default=True, description="Whether to use AMSGrad") + restart_latest: bool = Field( + default=True, description="Whether to restart the latest model" + ) + seed: int = Field(default=123, description="Seed for the random number generator") + device: Literal["cpu", "cuda"] = Field( + default="cpu", description="Device to be used for model fitting" + ) + + +class NEPSettings(BaseModel): + """Model describing the hyperparameters for the NEP fits.""" + + version: int = Field(default=4, description="Version of the NEP model") + type: list[int | str] = Field( + default_factory=lambda: [1, "X"], + description="Mandatory Parameter. Number of atom types and list of " + "chemical species. Number of atom types must be an integer, followed by " + "chemical symbols of species as in periodic table " + "for which model needs to be trained, separated by comma. " + "Default is [1, 'X'] as a placeholder. Example: [2, 'Pb', 'Te']", + ) + type_weight: float = Field( + default=1.0, description="Weights for different chemical species" + ) + model_type: int = Field( + default=0, + description="Type of model that is being trained. " + "Can be 0 (potential), 1 (dipole), " + "2 (polarizability)", + ) + prediction: int = Field( + default=0, description="Mode of NEP run. Set 0 for training and 1 for inference" + ) + cutoff: list[int, int] = Field( + default_factory=lambda: [6, 5], + description="Radial and angular cutoff. First element is for radial cutoff " + "and second element is for angular cutoff", + ) + n_max: list[int, int] = Field( + default_factory=lambda: [4, 4], + description="Number of radial and angular descriptors. First element " + "is for radial and second element is for angular.", + ) + basis_size: list[int, int] = Field( + default_factory=lambda: [8, 8], + description="Number of basis functions that are used to build the radial and angular descriptor. " + "First element is for radial descriptor and second element is for angular descriptor", + ) + l_max: list[int] = Field( + default_factory=lambda: [4, 2, 1], + description="The maximum expansion order for the angular terms. " + "First element is for three-body, second element is for four-body and third element is for five-body", + ) + neuron: int = Field( + default=80, description="Number of neurons in the hidden layer." + ) + lambda_1: float = Field( + default=0.0, description="Weight for the L1 regularization term." + ) + lambda_e: float = Field(default=1.0, description="Weight for the energy loss term.") + lambda_f: float = Field(default=1.0, description="Weight for the force loss term.") + lambda_v: float = Field(default=0.1, description="Weight for the virial loss term.") + force_delta: int = Field( + default=0, + description=" Sets bias the on the loss function to put more emphasis " + "on obtaining accurate predictions for smaller forces.", + ) + batch: int = Field(default=1000, description="Batch size for training.") + population: int = Field( + default=60, description="Size of the population used by the SNES algorithm." + ) + generation: int = Field( + default=100000, description="Number of generations used by the SNES algorithm." + ) + zbl: int = Field( + default=2, + description="Cutoff to use in universal ZBL potential at short distances. " + "Acceptable values are in range 1 to 2.5.", + ) + + +class MLIPHypers(BaseModel): + """Model containing the hyperparameter defaults for supported MLIPs in autoplex.""" + + GAP: GAPSettings = Field( + default_factory=GAPSettings, description="Hyperparameters for the GAP model" + ) + J_ACE: JACESettings = Field( + default_factory=JACESettings, description="Hyperparameters for the J-ACE model" + ) + NEQUIP: NEQUIPSettings = Field( + default_factory=NEQUIPSettings, + description="Hyperparameters for the NEQUIP model", + ) + M3GNET: M3GNETSettings = Field( + default_factory=M3GNETSettings, + description="Hyperparameters for the M3GNET model", + ) + MACE: MACESettings = Field( + default_factory=MACESettings, description="Hyperparameters for the MACE model" + ) + NEP: NEPSettings = Field( + default_factory=NEPSettings, description="Hyperparameters for the NEP model" + ) + + +# RSS Configuration + + class ResumeFromPreviousState(BaseModel): """ A model describing the state information. - It is required to resume a previously interrupted or saved RSS workflow. + Useful to resume a previously interrupted or saved RSS workflow. When 'train_from_scratch' is set to False, this parameter is mandatory for the workflow to pick up from a saved state. """ From b918db8340298d50d5258b6b4d07db6877da68a6 Mon Sep 17 00:00:00 2001 From: naik-aakash Date: Thu, 9 Jan 2025 17:41:46 +0100 Subject: [PATCH 14/74] add MODULE_DIR --- src/autoplex/fitting/common/jobs.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/src/autoplex/fitting/common/jobs.py b/src/autoplex/fitting/common/jobs.py index 016f08092..83457d1ed 100644 --- a/src/autoplex/fitting/common/jobs.py +++ b/src/autoplex/fitting/common/jobs.py @@ -1,5 +1,6 @@ """General fitting jobs using several MLIPs available.""" +import os from pathlib import Path from jobflow import job @@ -13,8 +14,8 @@ nequip_fitting, ) -current_dir = Path(__file__).absolute().parent -GAP_DEFAULTS_FILE_PATH = current_dir / "mlip-phonon-defaults.json" +MODULE_DIR = Path(os.path.dirname(__file__)) +GAP_DEFAULTS_FILE_PATH = MODULE_DIR / "mlip-phonon-defaults.json" @job From b0913e19e22d13ed5b78cf23e426f644ac2b4550 Mon Sep 17 00:00:00 2001 From: naik-aakash Date: Fri, 10 Jan 2025 17:52:47 +0100 Subject: [PATCH 15/74] make RssConfig and MLIPHYPERs at root level --- src/autoplex/__init__.py | 4 ++ src/autoplex/settings.py | 93 ++++++++++++++++++++++++++-------------- 2 files changed, 66 insertions(+), 31 deletions(-) diff --git a/src/autoplex/__init__.py b/src/autoplex/__init__.py index 0aec76f9f..e2e6e9677 100644 --- a/src/autoplex/__init__.py +++ b/src/autoplex/__init__.py @@ -8,3 +8,7 @@ """ from autoplex._version import __version__ +from autoplex.settings import MLIPHypers, RssConfig + +MLIP_HYPERS = MLIPHypers() +RSS_CONFIG = RssConfig() diff --git a/src/autoplex/settings.py b/src/autoplex/settings.py index 94cedb82b..401c726d9 100644 --- a/src/autoplex/settings.py +++ b/src/autoplex/settings.py @@ -2,13 +2,40 @@ from __future__ import annotations -from typing import Literal +from typing import Any, Literal -from pydantic import BaseModel, Field +from pydantic import BaseModel, ConfigDict, Field -class GeneralSettings(BaseModel): - """Model describing general hyperparameters for the GAP.""" +class UpdateBaseModel(BaseModel): + """Base class for all models in autoplex.""" + + model_config = ConfigDict(validate_assignment=True, protected_namespaces=()) + + def update_fields(self, updates: dict[str, Any]): + """ + Update the fields of the model, including nested fields. + + Args: + updates (Dict[str, Any]): A dictionary containing the fields to update. + """ + for key, value in updates.items(): + if hasattr(self, key): + field_value = getattr(self, key) + if isinstance(field_value, UpdateBaseModel) and isinstance(value, dict): + # Update nested model + field_value.update_fields(value) # Recursively call update_fields + else: + # Update field value + setattr(self, key, value) + else: + raise ValueError( + f"Field {key} does not exist in {self.__class__.__name__}." + ) + + +class GeneralSettings(UpdateBaseModel): + """Model describing general hyperparameters for the GAP fits.""" at_file: str = Field( default="train.extxyz", description="Name of the training file" @@ -39,17 +66,19 @@ class GeneralSettings(BaseModel): soap: bool = Field(default=True, description="Whether to include SOAP terms") -class TwobSettings(BaseModel): - """Model describing two body hyperparameters for the GAP.""" +class TwobSettings(UpdateBaseModel): + """Model describing two body hyperparameters for the GAP fits.""" distance_Nb_order: int = Field( - default=2, description="Distance_Nb order for two-body" + default=2, + description="Distance_Nb order for two-body", + alias="distance_Nb order", ) f0: float = Field(default=0.0, description="F0 value for two-body") add_species: str = Field( default="T", description="Whether to add species information" ) - cutoff: float = Field(default=5.0, description="Radial cutoff distance") + cutoff: float | int = Field(default=5.0, description="Radial cutoff distance") n_sparse: int = Field(default=15, description="Number of sparse points") covariance_type: str = Field( default="ard_se", description="Covariance type for two-body" @@ -66,17 +95,19 @@ class TwobSettings(BaseModel): ) -class ThreebSettings(BaseModel): - """Model describing threebody hyperparameters for the GAP.""" +class ThreebSettings(UpdateBaseModel): + """Model describing threebody hyperparameters for the GAP fits.""" distance_Nb_order: int = Field( - default=3, description="Distance_Nb order for three-body" + default=3, + description="Distance_Nb order for three-body", + alias="distance_Nb order", ) f0: float = Field(default=0.0, description="F0 value for three-body") add_species: str = Field( default="T", description="Whether to add species information" ) - cutoff: float = Field(default=3.25, description="Radial cutoff distance") + cutoff: float | int = Field(default=3.25, description="Radial cutoff distance") n_sparse: int = Field(default=100, description="Number of sparse points") covariance_type: str = Field( default="ard_se", description="Covariance type for three-body" @@ -93,8 +124,8 @@ class ThreebSettings(BaseModel): ) -class SoapSettings(BaseModel): - """Model describing soap hyperparameters for the GAP.""" +class SoapSettings(UpdateBaseModel): + """Model describing soap hyperparameters for the GAP fits.""" add_species: str = Field( default="T", description="Whether to add species information" @@ -121,7 +152,7 @@ class SoapSettings(BaseModel): ) -class GAPSettings(BaseModel): +class GAPSettings(UpdateBaseModel): """Model describing the hyperparameters for the GAP fits for Phonons.""" general: GeneralSettings = Field( @@ -142,7 +173,7 @@ class GAPSettings(BaseModel): ) -class JACESettings(BaseModel): +class JACESettings(UpdateBaseModel): """Model describing the hyperparameters for the J-ACE fits.""" order: int = Field(default=3, description="Order of the J-ACE model") @@ -151,7 +182,7 @@ class JACESettings(BaseModel): solver: str = Field(default="BLR", description="Solver for the J-ACE model") -class NEQUIPSettings(BaseModel): +class NEQUIPSettings(UpdateBaseModel): """Model describing the hyperparameters for the NEQUIP fits.""" r_max: float = Field(default=4.0, description="Radial cutoff distance") @@ -169,7 +200,7 @@ class NEQUIPSettings(BaseModel): default_dtype: str = Field(default="float32", description="Default data type") -class M3GNETSettings(BaseModel): +class M3GNETSettings(UpdateBaseModel): """Model describing the hyperparameters for the M3GNET fits.""" exp_name: str = Field(default="training", description="Name of the experiment") @@ -196,7 +227,7 @@ class M3GNETSettings(BaseModel): ) -class MACESettings(BaseModel): +class MACESettings(UpdateBaseModel): """Model describing the hyperparameters for the MACE fits.""" model: str = Field(default="MACE", description="type of the model") @@ -227,7 +258,7 @@ class MACESettings(BaseModel): ) -class NEPSettings(BaseModel): +class NEPSettings(UpdateBaseModel): """Model describing the hyperparameters for the NEP fits.""" version: int = Field(default=4, description="Version of the NEP model") @@ -299,7 +330,7 @@ class NEPSettings(BaseModel): ) -class MLIPHypers(BaseModel): +class MLIPHypers(UpdateBaseModel): """Model containing the hyperparameter defaults for supported MLIPs in autoplex.""" GAP: GAPSettings = Field( @@ -327,7 +358,7 @@ class MLIPHypers(BaseModel): # RSS Configuration -class ResumeFromPreviousState(BaseModel): +class ResumeFromPreviousState(UpdateBaseModel): """ A model describing the state information. @@ -353,7 +384,7 @@ class ResumeFromPreviousState(BaseModel): ) -class SoapParas(BaseModel): +class SoapParas(UpdateBaseModel): """A model describing the SOAP parameters.""" l_max: int = Field(default=12, description="Maximum degree of spherical harmonics") @@ -374,7 +405,7 @@ class SoapParas(BaseModel): ) -class BcurParams(BaseModel): +class BcurParams(UpdateBaseModel): """A model describing the parameters for the BCUR method.""" soap_paras: SoapParas = Field(default_factory=SoapParas) @@ -386,7 +417,7 @@ class BcurParams(BaseModel): ) -class BuildcellOptions(BaseModel): +class BuildcellOptions(UpdateBaseModel): """A model describing the parameters for buildcell.""" NFORM: str | None = Field(default=None, description="The number of formula units") @@ -396,7 +427,7 @@ class BuildcellOptions(BaseModel): MINSEP: str | None = Field(default=None, description="The minimum separation") -class CustomIncar(BaseModel): +class CustomIncar(UpdateBaseModel): """A model describing the INCAR parameters.""" ISMEAR: int = 0 @@ -420,7 +451,7 @@ class CustomIncar(BaseModel): LPLANE: str = ".FALSE." -class Twob(BaseModel): +class Twob(UpdateBaseModel): """A model describing the two-body GAP parameters.""" cutoff: float = Field(default=5.0, description="Radial cutoff distance") @@ -430,13 +461,13 @@ class Twob(BaseModel): ) -class Threeb(BaseModel): +class Threeb(UpdateBaseModel): """A model describing the three-body GAP parameters.""" cutoff: float = Field(default=3.0, description="Radial cutoff distance") -class Soap(BaseModel): +class Soap(UpdateBaseModel): """A model describing the SOAP GAP parameters.""" l_max: int = Field(default=10, description="Maximum degree of spherical harmonics") @@ -448,7 +479,7 @@ class Soap(BaseModel): cutoff: float = Field(default=5.0, description="Radial cutoff distance") -class General(BaseModel): +class General(UpdateBaseModel): """A model describing the general GAP parameters.""" three_body: bool = Field( @@ -456,7 +487,7 @@ class General(BaseModel): ) -class RssConfig(BaseModel): +class RssConfig(UpdateBaseModel): """A model describing the complete RSS configuration.""" tag: str | None = Field( From 30e7cf6a1fc0908a65f0c327eb2a1f4b7494187d Mon Sep 17 00:00:00 2001 From: naik-aakash Date: Fri, 10 Jan 2025 17:56:46 +0100 Subject: [PATCH 16/74] remove load_hypers function and use models instead --- src/autoplex/fitting/common/utils.py | 102 ++++++++------------------- 1 file changed, 30 insertions(+), 72 deletions(-) diff --git a/src/autoplex/fitting/common/utils.py b/src/autoplex/fitting/common/utils.py index 3bd7b89fc..7df78a676 100644 --- a/src/autoplex/fitting/common/utils.py +++ b/src/autoplex/fitting/common/utils.py @@ -1,7 +1,6 @@ """Utility functions for fitting jobs.""" import contextlib -import json import logging import os import re @@ -41,6 +40,7 @@ from scipy.spatial import ConvexHull from scipy.special import comb +from autoplex import MLIP_HYPERS from autoplex.data.common.utils import ( data_distillation, plot_energy_forces, @@ -59,7 +59,7 @@ def gap_fitting( db_dir: Path, species_list: list | None = None, - path_to_hyperparameters: Path | str = MLIP_PHONON_DEFAULTS_FILE_PATH, + hyperparameters: MLIP_HYPERS.GAP = MLIP_HYPERS.GAP, num_processes_fit: int = 32, auto_delta: bool = True, glue_xml: bool = False, @@ -80,8 +80,8 @@ def gap_fitting( Path to database directory. species_list: list List of element names (strings) - path_to_hyperparameters : str or Path. - Path to JSON file containing the GAP hyperparameters. + hyperparameters: MLIP_HYPERS.GAP + fit hyperparameters. num_processes_fit: int Number of processes used for gap_fit auto_delta: bool @@ -110,8 +110,8 @@ def gap_fitting( A dictionary with train_error, test_error, path_to_mlip """ - if path_to_hyperparameters is None: - path_to_hyperparameters = MLIP_PHONON_DEFAULTS_FILE_PATH + if hyperparameters is None: + hyperparameters = MLIP_HYPERS.GAP.model_copy(deep=True) # keep additional pre- and suffixes gap_file_xml = train_name.replace("train", "gap_file").replace(".extxyz", ".xml") quip_train_file = train_name.replace("train", "quip_train") @@ -124,22 +124,22 @@ def gap_fitting( train_data_path = os.path.join(db_dir, train_name) test_data_path = os.path.join(db_dir, test_name) - default_hyperparameters = load_mlip_hyperparameter_defaults( - mlip_fit_parameter_file_path=path_to_hyperparameters - ) - gap_default_hyperparameters = default_hyperparameters["GAP"] + hyperparameters.update_fields( + { + "general": { + "gp_file": gap_file_xml, + "energy_parameter_name": ref_energy_name, + "force_parameter_name": ref_force_name, + "virial_parameter_name": ref_virial_name, + } + } + ) - gap_default_hyperparameters["general"].update({"gp_file": gap_file_xml}) - gap_default_hyperparameters["general"]["energy_parameter_name"] = ref_energy_name - gap_default_hyperparameters["general"]["force_parameter_name"] = ref_force_name - gap_default_hyperparameters["general"]["virial_parameter_name"] = ref_virial_name + if fit_kwargs: + hyperparameters.update_fields(fit_kwargs) - for parameter in gap_default_hyperparameters: - if fit_kwargs: - for arg in fit_kwargs: - if parameter == arg: - gap_default_hyperparameters[parameter].update(fit_kwargs[arg]) + gap_default_hyperparameters = hyperparameters.model_dump(by_alias=True) include_two_body = gap_default_hyperparameters["general"]["two_body"] include_three_body = gap_default_hyperparameters["general"]["three_body"] @@ -271,7 +271,7 @@ def gap_fitting( ) def jace_fitting( db_dir: str | Path, - path_to_hyperparameters: Path | str = MLIP_RSS_DEFAULTS_FILE_PATH, + hyperparameters: MLIP_HYPERS.J_ACE = MLIP_HYPERS.J_ACE, isolated_atom_energies: dict | None = None, ref_energy_name: str = "REF_energy", ref_force_name: str = "REF_forces", @@ -290,8 +290,8 @@ def jace_fitting( ---------- db_dir: str or Path directory containing the training and testing data files. - path_to_hyperparameters : str or Path. - Path to JSON file containing the J-ACE hyperparameters. + hyperparameters: MLIP_HYPERS.J_ACE + J-ACE hyperparameters. isolated_atom_energies: dict: mandatory dictionary mapping element numbers to isolated energies. ref_energy_name : str, optional @@ -327,8 +327,8 @@ def jace_fitting( ------ - ValueError: If the `isolated_atom_energies` dictionary is empty or not provided when required. """ - if path_to_hyperparameters is None: - path_to_hyperparameters = MLIP_RSS_DEFAULTS_FILE_PATH + if hyperparameters is None: + hyperparameters = MLIP_HYPERS.J_ACE.model_copy(deep=True) train_atoms = ase.io.read(os.path.join(db_dir, "train.extxyz"), index=":") source_file_path = os.path.join(db_dir, "test.extxyz") shutil.copy(source_file_path, ".") @@ -361,20 +361,10 @@ def jace_fitting( ] ase.io.write("train_ace.extxyz", train_ace, format="extxyz") - default_hyperparameters = load_mlip_hyperparameter_defaults( - mlip_fit_parameter_file_path=path_to_hyperparameters - ) - jace_hypers = default_hyperparameters["J-ACE"] - if fit_kwargs: - for parameter in jace_hypers: - if parameter in fit_kwargs: - if isinstance(fit_kwargs[parameter], type(jace_hypers[parameter])): - jace_hypers[parameter] = fit_kwargs[parameter] - else: - raise TypeError( - f"The type of {parameter} should be {type(jace_hypers[parameter])}!" - ) + hyperparameters.update_fields(fit_kwargs) + + jace_hypers = hyperparameters.model_dump(by_alias=True) order = jace_hypers["order"] totaldegree = jace_hypers["totaldegree"] @@ -547,11 +537,7 @@ def nequip_fitting( else: raise ValueError("isolated_atom_energies is empty or not defined!") - default_hyperparameters = load_mlip_hyperparameter_defaults( - mlip_fit_parameter_file_path=path_to_hyperparameters - ) - - nequip_hypers = default_hyperparameters["NEQUIP"] + nequip_hypers = MLIP_HYPERS.NEQUIP.model_dump(by_alias=True) if fit_kwargs: for parameter in nequip_hypers: @@ -808,11 +794,8 @@ def m3gnet_fitting( """ if path_to_hyperparameters is None: path_to_hyperparameters = MLIP_RSS_DEFAULTS_FILE_PATH - default_hyperparameters = load_mlip_hyperparameter_defaults( - mlip_fit_parameter_file_path=path_to_hyperparameters - ) - m3gnet_hypers = default_hyperparameters["M3GNET"] + m3gnet_hypers = MLIP_HYPERS.M3GNET.model_dump(by_alias=True) if fit_kwargs: for parameter in m3gnet_hypers: @@ -1179,14 +1162,7 @@ def mace_fitting( atoms=atoms, ref_virial_name=ref_virial_name, out_file_name="train.extxyz" ) - if use_defaults: - default_hyperparameters = load_mlip_hyperparameter_defaults( - mlip_fit_parameter_file_path=path_to_hyperparameters - ) - - mace_hypers = default_hyperparameters["MACE"] - else: - mace_hypers = {} + mace_hypers = MLIP_HYPERS.MACE.model_dump(by_alias=True) if use_defaults else {} # TODO: should we do a type check? not sure # as it will be a lot of work to keep it updated @@ -1307,24 +1283,6 @@ def check_convergence(test_error: float) -> bool: return convergence -def load_mlip_hyperparameter_defaults(mlip_fit_parameter_file_path: str | Path) -> dict: - """ - Load gap fit default parameters from the json file. - - Parameters - ---------- - mlip_fit_parameter_file_path : str or Path. - Path to MLIP default parameter JSON files. - - Returns - ------- - dict - gap fit default parameters. - """ - with open(mlip_fit_parameter_file_path, encoding="utf-8") as f: - return json.load(f) - - def gap_hyperparameter_constructor( gap_parameter_dict: dict, include_two_body: bool = False, From 8cab4fbb0674e09527acc82c0dbbc3fc2dd703ce Mon Sep 17 00:00:00 2001 From: naik-aakash Date: Fri, 10 Jan 2025 17:57:37 +0100 Subject: [PATCH 17/74] adapt machine_learning_fit (GAP, J-ACE) --- src/autoplex/fitting/common/jobs.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/src/autoplex/fitting/common/jobs.py b/src/autoplex/fitting/common/jobs.py index b015c1910..4b6ae1521 100644 --- a/src/autoplex/fitting/common/jobs.py +++ b/src/autoplex/fitting/common/jobs.py @@ -128,7 +128,6 @@ def machine_learning_fit( ).exists(): train_test_error = gap_fitting( db_dir=database_dir, - path_to_hyperparameters=path_to_hyperparameters, species_list=species_list, num_processes_fit=num_processes_fit, auto_delta=auto_delta, @@ -146,7 +145,6 @@ def machine_learning_fit( elif mlip_type == "J-ACE": train_test_error = jace_fitting( db_dir=database_dir, - path_to_hyperparameters=path_to_hyperparameters, isolated_atom_energies=isolated_atom_energies, ref_energy_name=ref_energy_name, ref_force_name=ref_force_name, From b94ba98a01cbf164570facb754e87d724753af7c Mon Sep 17 00:00:00 2001 From: naik-aakash Date: Fri, 10 Jan 2025 17:58:11 +0100 Subject: [PATCH 18/74] update utils tests for gap constructor --- tests/fitting/common/test_utils.py | 36 +++++++++++------------------- 1 file changed, 13 insertions(+), 23 deletions(-) diff --git a/tests/fitting/common/test_utils.py b/tests/fitting/common/test_utils.py index 0ca145761..5be01863f 100644 --- a/tests/fitting/common/test_utils.py +++ b/tests/fitting/common/test_utils.py @@ -1,7 +1,7 @@ import os.path +from autoplex import MLIP_HYPERS from autoplex.fitting.common.utils import ( - load_mlip_hyperparameter_defaults, gap_hyperparameter_constructor, check_convergence, data_distillation, @@ -21,11 +21,8 @@ def test_stratified_split(test_dir): assert len(test) == 3 def test_gap_hyperparameter_constructor(): - hyper_parameter_dict = load_mlip_hyperparameter_defaults( - mlip_fit_parameter_file_path=MLIP_PHONON_DEFAULTS_FILE_PATH - ) - gap_hyper_parameter_dict = hyper_parameter_dict["GAP"] + gap_hyper_parameter_dict = MLIP_HYPERS.GAP.model_dump(by_alias=True) gap_input_list = gap_hyperparameter_constructor( gap_parameter_dict=gap_hyper_parameter_dict, @@ -57,11 +54,7 @@ def test_gap_hyperparameter_constructor(): assert ref_list == gap_input_list - hyper_parameter_dict = load_mlip_hyperparameter_defaults( - mlip_fit_parameter_file_path=MLIP_PHONON_DEFAULTS_FILE_PATH - ) - - gap_hyper_parameter_dict = hyper_parameter_dict["GAP"] + gap_hyper_parameter_dict = MLIP_HYPERS.GAP.model_dump(by_alias=True) gap_input_list = gap_hyperparameter_constructor( gap_parameter_dict=gap_hyper_parameter_dict, @@ -91,9 +84,13 @@ def test_gap_hyperparameter_constructor(): assert ref_list == gap_input_list # test if returned string is changed if passed in dict is updated - gap_hyper_parameter_dict["twob"].update({"cutoff": 8}) - gap_hyper_parameter_dict["threeb"].update({"cutoff": 8, "n_sparse": 100}) - gap_hyper_parameter_dict["soap"].update({"delta": 1.5, "zeta": 2}) + gap_hyper_parameter = MLIP_HYPERS.GAP.model_copy(deep=True) + gap_hyper_parameter.update_fields({"twob": {"cutoff": 8}, + "threeb": {"cutoff": 8.0, + "n_sparse": 100}, + "soap": {"delta": 1.5, + "zeta": 2}}) + gap_hyper_parameter_dict = gap_hyper_parameter.model_dump(by_alias=True) gap_input_list_updated = gap_hyperparameter_constructor( gap_parameter_dict=gap_hyper_parameter_dict, @@ -116,7 +113,7 @@ def test_gap_hyperparameter_constructor(): "gap={distance_Nb order=2 f0=0.0 add_species=T cutoff=8 " "n_sparse=15 covariance_type=ard_se delta=2.0 theta_uniform=0.5 " "sparse_method=uniform compact_clusters=T :distance_Nb order=3 f0=0.0 add_species=T " - "cutoff=8 n_sparse=100 covariance_type=ard_se " + "cutoff=8.0 n_sparse=100 covariance_type=ard_se " "delta=2.0 theta_uniform=1.0 sparse_method=uniform compact_clusters=T :soap " "add_species=T l_max=10 n_max=12 atom_sigma=0.5 zeta=2 cutoff=5.0 " "cutoff_transition_width=1.0 central_weight=1.0 n_sparse=6000 delta=1.5 " @@ -127,11 +124,8 @@ def test_gap_hyperparameter_constructor(): # check disable three_body and two_body - hyper_parameter_dict = load_mlip_hyperparameter_defaults( - mlip_fit_parameter_file_path=MLIP_PHONON_DEFAULTS_FILE_PATH - ) - gap_hyper_parameter_dict = hyper_parameter_dict["GAP"] + gap_hyper_parameter_dict = MLIP_HYPERS.GAP.model_dump(by_alias=True) # three_body_disable @@ -193,11 +187,7 @@ def test_gap_hyperparameter_constructor(): assert ref_list == gap_input_list - hyper_parameter_dict = load_mlip_hyperparameter_defaults( - mlip_fit_parameter_file_path=MLIP_PHONON_DEFAULTS_FILE_PATH - ) - - gap_hyper_parameter_dict = hyper_parameter_dict["GAP"] + gap_hyper_parameter_dict = MLIP_HYPERS.GAP.model_dump(by_alias=True) # check with only soap enabled From 2d79c34c4f54b6ec36894284eb089495b0b3f4d1 Mon Sep 17 00:00:00 2001 From: naik-aakash Date: Fri, 10 Jan 2025 18:31:47 +0100 Subject: [PATCH 19/74] fix imports error --- src/autoplex/auto/phonons/flows.py | 13 ++++++------- 1 file changed, 6 insertions(+), 7 deletions(-) diff --git a/src/autoplex/auto/phonons/flows.py b/src/autoplex/auto/phonons/flows.py index 0000371b0..e9b45f17a 100644 --- a/src/autoplex/auto/phonons/flows.py +++ b/src/autoplex/auto/phonons/flows.py @@ -19,6 +19,7 @@ MPStaticSet, ) +from autoplex import MLIP_HYPERS from autoplex.auto.phonons.jobs import ( complete_benchmark, dft_phonopy_gen_data, @@ -33,10 +34,7 @@ from autoplex.data.phonons.flows import IsoAtomStaticMaker, TightDFTStaticMaker from autoplex.data.phonons.jobs import reduce_supercell_size_job from autoplex.fitting.common.flows import MLIPFitMaker -from autoplex.fitting.common.utils import ( - MLIP_PHONON_DEFAULTS_FILE_PATH, - load_mlip_hyperparameter_defaults, -) +from autoplex.fitting.common.utils import MLIP_PHONON_DEFAULTS_FILE_PATH __all__ = [ "CompleteDFTvsMLBenchmarkWorkflow", @@ -279,9 +277,10 @@ def make( fit_input = {} bm_outputs = [] - default_hyperparameters = load_mlip_hyperparameter_defaults( - mlip_fit_parameter_file_path=self.path_to_hyperparameters - ) + # default_hyperparameters = load_mlip_hyperparameter_defaults( + # mlip_fit_parameter_file_path=self.path_to_hyperparameters + # ) + default_hyperparameters = MLIP_HYPERS.model_dump(by_alias=True) soap_default_dict = next( ( From f87161f190c02d36f07f26ca3ca2371718d65a96 Mon Sep 17 00:00:00 2001 From: naik-aakash Date: Fri, 10 Jan 2025 19:56:59 +0100 Subject: [PATCH 20/74] use deep copy --- src/autoplex/auto/phonons/flows.py | 7 +++---- src/autoplex/fitting/common/utils.py | 6 ++---- tests/fitting/common/test_utils.py | 13 ++++++++----- 3 files changed, 13 insertions(+), 13 deletions(-) diff --git a/src/autoplex/auto/phonons/flows.py b/src/autoplex/auto/phonons/flows.py index e9b45f17a..0084d1bec 100644 --- a/src/autoplex/auto/phonons/flows.py +++ b/src/autoplex/auto/phonons/flows.py @@ -277,10 +277,9 @@ def make( fit_input = {} bm_outputs = [] - # default_hyperparameters = load_mlip_hyperparameter_defaults( - # mlip_fit_parameter_file_path=self.path_to_hyperparameters - # ) - default_hyperparameters = MLIP_HYPERS.model_dump(by_alias=True) + hyper_parameters = MLIP_HYPERS.model_copy(deep=True) + default_hyperparameters = hyper_parameters.model_dump(by_alias=True) + # default_hyperparameters = MLIP_HYPERS.model_copy(deep=True).model_dump(by_alias=True) soap_default_dict = next( ( diff --git a/src/autoplex/fitting/common/utils.py b/src/autoplex/fitting/common/utils.py index 7df78a676..0a9295182 100644 --- a/src/autoplex/fitting/common/utils.py +++ b/src/autoplex/fitting/common/utils.py @@ -110,8 +110,7 @@ def gap_fitting( A dictionary with train_error, test_error, path_to_mlip """ - if hyperparameters is None: - hyperparameters = MLIP_HYPERS.GAP.model_copy(deep=True) + hyperparameters = hyperparameters.model_copy(deep=True) # keep additional pre- and suffixes gap_file_xml = train_name.replace("train", "gap_file").replace(".extxyz", ".xml") quip_train_file = train_name.replace("train", "quip_train") @@ -327,8 +326,7 @@ def jace_fitting( ------ - ValueError: If the `isolated_atom_energies` dictionary is empty or not provided when required. """ - if hyperparameters is None: - hyperparameters = MLIP_HYPERS.J_ACE.model_copy(deep=True) + hyperparameters = hyperparameters.model_copy(deep=True) train_atoms = ase.io.read(os.path.join(db_dir, "train.extxyz"), index=":") source_file_path = os.path.join(db_dir, "test.extxyz") shutil.copy(source_file_path, ".") diff --git a/tests/fitting/common/test_utils.py b/tests/fitting/common/test_utils.py index 5be01863f..43b224c72 100644 --- a/tests/fitting/common/test_utils.py +++ b/tests/fitting/common/test_utils.py @@ -22,7 +22,8 @@ def test_stratified_split(test_dir): def test_gap_hyperparameter_constructor(): - gap_hyper_parameter_dict = MLIP_HYPERS.GAP.model_dump(by_alias=True) + gap_hyper_parameter = MLIP_HYPERS.GAP.model_copy(deep=True) + gap_hyper_parameter_dict = gap_hyper_parameter.model_dump(by_alias=True) gap_input_list = gap_hyperparameter_constructor( gap_parameter_dict=gap_hyper_parameter_dict, @@ -54,7 +55,8 @@ def test_gap_hyperparameter_constructor(): assert ref_list == gap_input_list - gap_hyper_parameter_dict = MLIP_HYPERS.GAP.model_dump(by_alias=True) + gap_hyper_parameter = MLIP_HYPERS.GAP.model_copy(deep=True) + gap_hyper_parameter_dict = gap_hyper_parameter.model_dump(by_alias=True) gap_input_list = gap_hyperparameter_constructor( gap_parameter_dict=gap_hyper_parameter_dict, @@ -124,8 +126,8 @@ def test_gap_hyperparameter_constructor(): # check disable three_body and two_body - - gap_hyper_parameter_dict = MLIP_HYPERS.GAP.model_dump(by_alias=True) + gap_hyper_parameter = MLIP_HYPERS.GAP.model_copy(deep=True) + gap_hyper_parameter_dict = gap_hyper_parameter.model_dump(by_alias=True) # three_body_disable @@ -187,7 +189,8 @@ def test_gap_hyperparameter_constructor(): assert ref_list == gap_input_list - gap_hyper_parameter_dict = MLIP_HYPERS.GAP.model_dump(by_alias=True) + gap_hyper_parameter = MLIP_HYPERS.GAP.model_copy(deep=True) + gap_hyper_parameter_dict = gap_hyper_parameter.model_dump(by_alias=True) # check with only soap enabled From 4faeebd33ba44029df4191ebb5703d5eec1c5ab6 Mon Sep 17 00:00:00 2001 From: naik-aakash Date: Fri, 10 Jan 2025 21:00:30 +0100 Subject: [PATCH 21/74] temp disable ValueError raise behaviour --- src/autoplex/settings.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/src/autoplex/settings.py b/src/autoplex/settings.py index 401c726d9..630acf517 100644 --- a/src/autoplex/settings.py +++ b/src/autoplex/settings.py @@ -28,10 +28,10 @@ def update_fields(self, updates: dict[str, Any]): else: # Update field value setattr(self, key, value) - else: - raise ValueError( - f"Field {key} does not exist in {self.__class__.__name__}." - ) + # else: + # raise ValueError( + # f"Field {key} does not exist in {self.__class__.__name__}." + # ) class GeneralSettings(UpdateBaseModel): From 3b3ee4705dbe55c6e29f0d10d9eaa67be3c8d88a Mon Sep 17 00:00:00 2001 From: naik-aakash Date: Fri, 10 Jan 2025 21:37:10 +0100 Subject: [PATCH 22/74] fix rssmaker issue --- src/autoplex/auto/rss/jobs.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/src/autoplex/auto/rss/jobs.py b/src/autoplex/auto/rss/jobs.py index 60ed679f5..dd0d73c2a 100644 --- a/src/autoplex/auto/rss/jobs.py +++ b/src/autoplex/auto/rss/jobs.py @@ -221,8 +221,8 @@ def initial_rss( apply_data_preprocessing=False, auto_delta=auto_delta, glue_xml=False, - ).make( database_dir=do_data_preprocessing.output, + ).make( isolated_atom_energies=do_data_collection.output["isolated_atom_energies"], device=device_for_fitting, **fit_kwargs, @@ -236,14 +236,12 @@ def initial_rss( do_mlip_fit, ] - (mlip_path,) = do_mlip_fit.output["mlip_path"] - return Response( replace=Flow(job_list), output={ "test_error": do_mlip_fit.output["test_error"], "pre_database_dir": do_data_preprocessing.output, - "mlip_path": mlip_path, + "mlip_path": do_mlip_fit.output["mlip_path"], "isolated_atom_energies": do_data_collection.output[ "isolated_atom_energies" ], From 6013e400da85f85d7505fd4ea6a56df795af7fc2 Mon Sep 17 00:00:00 2001 From: naik-aakash Date: Sun, 12 Jan 2025 20:47:32 +0100 Subject: [PATCH 23/74] add from_file method and rename update_fields>update_parameters --- src/autoplex/settings.py | 25 ++++++++++++++++++++----- 1 file changed, 20 insertions(+), 5 deletions(-) diff --git a/src/autoplex/settings.py b/src/autoplex/settings.py index 630acf517..b010ee213 100644 --- a/src/autoplex/settings.py +++ b/src/autoplex/settings.py @@ -4,6 +4,7 @@ from typing import Any, Literal +from monty.serialization import loadfn from pydantic import BaseModel, ConfigDict, Field @@ -12,27 +13,41 @@ class UpdateBaseModel(BaseModel): model_config = ConfigDict(validate_assignment=True, protected_namespaces=()) - def update_fields(self, updates: dict[str, Any]): + def update_parameters(self, updates: dict[str, Any]): """ - Update the fields of the model, including nested fields. + Update the default parameters of the model instance, including nested fields. Args: - updates (Dict[str, Any]): A dictionary containing the fields to update. + updates (Dict[str, Any]): A dictionary containing the fields as keys to update. """ for key, value in updates.items(): if hasattr(self, key): field_value = getattr(self, key) - if isinstance(field_value, UpdateBaseModel) and isinstance(value, dict): + if isinstance(field_value, self.__class__) and isinstance(value, dict): # Update nested model - field_value.update_fields(value) # Recursively call update_fields + field_value.update_parameters( + value + ) # Recursively call update_fields else: # Update field value setattr(self, key, value) + # else: # raise ValueError( # f"Field {key} does not exist in {self.__class__.__name__}." # ) + @classmethod + def from_file(cls, filename: str): + """ + Load the parameters from a file. + + Args: + filename (str): The name of the file to load the parameters from. + """ + custom_params = loadfn(filename) + return cls(**custom_params) + class GeneralSettings(UpdateBaseModel): """Model describing general hyperparameters for the GAP fits.""" From 1bc2be766fe3306114387ec2a3ae156bb5fdb9d3 Mon Sep 17 00:00:00 2001 From: naik-aakash Date: Sun, 12 Jan 2025 20:49:08 +0100 Subject: [PATCH 24/74] extend use of MLIP hypers model to M3GNET, MACE, NEQUIP --- src/autoplex/auto/phonons/flows.py | 2 - src/autoplex/auto/rss/flows.py | 5 ++- src/autoplex/fitting/common/flows.py | 7 +-- src/autoplex/fitting/common/jobs.py | 15 ++----- src/autoplex/fitting/common/utils.py | 67 +++++++++++----------------- 5 files changed, 37 insertions(+), 59 deletions(-) diff --git a/src/autoplex/auto/phonons/flows.py b/src/autoplex/auto/phonons/flows.py index 0084d1bec..fdeb9f57c 100644 --- a/src/autoplex/auto/phonons/flows.py +++ b/src/autoplex/auto/phonons/flows.py @@ -392,7 +392,6 @@ def make( force_max=self.force_max, pre_xyz_files=pre_xyz_files, pre_database_dir=pre_database_dir, - path_to_hyperparameters=self.path_to_hyperparameters, atomwise_regularization_parameter=self.atomwise_regularization_parameter, force_min=self.force_min, atom_wise_regularization=self.atom_wise_regularization, @@ -477,7 +476,6 @@ def make( force_max=self.force_max, pre_xyz_files=pre_xyz_files, pre_database_dir=pre_database_dir, - path_to_hyperparameters=self.path_to_hyperparameters, atomwise_regularization_parameter=atomwise_reg_parameter, force_min=self.force_min, auto_delta=self.auto_delta, diff --git a/src/autoplex/auto/rss/flows.py b/src/autoplex/auto/rss/flows.py index 676d0e2cf..291c8ee5f 100644 --- a/src/autoplex/auto/rss/flows.py +++ b/src/autoplex/auto/rss/flows.py @@ -234,7 +234,8 @@ def make(self, **kwargs): - 'current_iter': int, The current iteration index. - 'kb_temp': float, The temperature (in eV) for Boltzmann sampling. """ - updated_config = self.config.model_copy(update=kwargs) + default_config = self.config.model_copy(deep=True) + updated_config = default_config.update_parameters(kwargs) config_params = updated_config.model_dump() self._process_hookean_paras(config_params) @@ -336,7 +337,7 @@ def make(self, **kwargs): @staticmethod def _process_hookean_paras(config): - if "hookean_paras" in config: + if "hookean_paras" in config and config["hookean_paras"] is not None: config["hookean_paras"] = { tuple(map(int, k.strip("()").split(", "))): tuple(v) for k, v in config["hookean_paras"].items() diff --git a/src/autoplex/fitting/common/flows.py b/src/autoplex/fitting/common/flows.py index 50172c395..0e317c078 100644 --- a/src/autoplex/fitting/common/flows.py +++ b/src/autoplex/fitting/common/flows.py @@ -9,6 +9,7 @@ import ase.io from jobflow import Flow, Maker, job +from autoplex import MLIP_HYPERS from autoplex.fitting.common.jobs import machine_learning_fit from autoplex.fitting.common.regularization import set_custom_sigma from autoplex.fitting.common.utils import ( @@ -70,8 +71,6 @@ class MLIPFitMaker(Maker): Names of the pre-database train xyz file and test xyz file. pre_database_dir: str or None The pre-database directory. - path_to_hyperparameters : str or Path. - Path to JSON file containing the MLIP hyperparameters. atomwise_regularization_parameter: float Regularization value for the atom-wise force components. atom_wise_regularization: bool @@ -106,7 +105,6 @@ class MLIPFitMaker(Maker): separated: bool = False pre_xyz_files: list[str] | None = None pre_database_dir: str | None = None - path_to_hyperparameters: Path | str | None = None regularization: bool = False # This is only used for GAP. atomwise_regularization_parameter: float = 0.1 # This is only used for GAP. atom_wise_regularization: bool = True # This is only used for GAP. @@ -121,6 +119,7 @@ class MLIPFitMaker(Maker): def make( self, fit_input: dict | None = None, # This is specific to phonon workflow + hyperparameters: MLIP_HYPERS = MLIP_HYPERS, species_list: list | None = None, isolated_atom_energies: dict | None = None, device: str = "cpu", @@ -133,6 +132,8 @@ def make( ---------- fit_input: dict Output from the CompletePhononDFTMLDataGenerationFlow process. + hyperparameters: MLIP_HYPERS + Hyperparameters for the MLIP. species_list: list List of element names (strings) involved in the training dataset isolated_atom_energies: dict diff --git a/src/autoplex/fitting/common/jobs.py b/src/autoplex/fitting/common/jobs.py index 4b6ae1521..893553a8a 100644 --- a/src/autoplex/fitting/common/jobs.py +++ b/src/autoplex/fitting/common/jobs.py @@ -1,6 +1,5 @@ """General fitting jobs using several MLIPs available.""" -import os from pathlib import Path import numpy as np @@ -15,16 +14,12 @@ nequip_fitting, ) -MODULE_DIR = Path(os.path.dirname(__file__)) -GAP_DEFAULTS_FILE_PATH = MODULE_DIR / "mlip-phonon-defaults.json" - @job def machine_learning_fit( database_dir: str | Path, species_list: list, run_fits_on_different_cluster: bool = False, - path_to_hyperparameters: Path | str | None = None, isolated_atom_energies: dict | None = None, num_processes_fit: int = 32, auto_delta: bool = True, @@ -49,10 +44,8 @@ def machine_learning_fit( Path to the directory containing the database. species_list: list List of element names (strings) involved in the training dataset - run_fit_on_different_cluster: bool + run_fits_on_different_cluster: bool Whether to run fitting on different clusters. - path_to_hyperparameters : str or Path. - Path to JSON file containing the MLIP hyperparameters. isolated_atom_energies: dict Dictionary of isolated atoms energies. num_processes_fit: int @@ -81,6 +74,9 @@ def machine_learning_fit( hyperpara_opt: bool Perform hyperparameter optimization using XPOT (XPOT: https://pubs.aip.org/aip/jcp/article/159/2/024803/2901815) + run_fits_on_different_cluster: bool + Indicates if fits are to be run on a different cluster. + If True, the fitting data (train.extxyz, test.extxyz) is stored in the database. fit_kwargs: dict Additional keyword arguments for MLIP fitting. """ @@ -157,7 +153,6 @@ def machine_learning_fit( elif mlip_type == "NEQUIP": train_test_error = nequip_fitting( db_dir=database_dir, - path_to_hyperparameters=path_to_hyperparameters, isolated_atom_energies=isolated_atom_energies, ref_energy_name=ref_energy_name, ref_force_name=ref_force_name, @@ -170,7 +165,6 @@ def machine_learning_fit( elif mlip_type == "M3GNET": train_test_error = m3gnet_fitting( db_dir=database_dir, - path_to_hyperparameters=path_to_hyperparameters, ref_energy_name=ref_energy_name, ref_force_name=ref_force_name, ref_virial_name=ref_virial_name, @@ -182,7 +176,6 @@ def machine_learning_fit( elif mlip_type == "MACE": train_test_error = mace_fitting( db_dir=database_dir, - path_to_hyperparameters=path_to_hyperparameters, ref_energy_name=ref_energy_name, ref_force_name=ref_force_name, ref_virial_name=ref_virial_name, diff --git a/src/autoplex/fitting/common/utils.py b/src/autoplex/fitting/common/utils.py index 0a9295182..e7ed26e0b 100644 --- a/src/autoplex/fitting/common/utils.py +++ b/src/autoplex/fitting/common/utils.py @@ -81,7 +81,7 @@ def gap_fitting( species_list: list List of element names (strings) hyperparameters: MLIP_HYPERS.GAP - fit hyperparameters. + Fit hyperparameters. num_processes_fit: int Number of processes used for gap_fit auto_delta: bool @@ -124,7 +124,7 @@ def gap_fitting( test_data_path = os.path.join(db_dir, test_name) - hyperparameters.update_fields( + hyperparameters.update_parameters( { "general": { "gp_file": gap_file_xml, @@ -136,7 +136,7 @@ def gap_fitting( ) if fit_kwargs: - hyperparameters.update_fields(fit_kwargs) + hyperparameters.update_parameters(fit_kwargs) gap_default_hyperparameters = hyperparameters.model_dump(by_alias=True) @@ -290,7 +290,7 @@ def jace_fitting( db_dir: str or Path directory containing the training and testing data files. hyperparameters: MLIP_HYPERS.J_ACE - J-ACE hyperparameters. + Fit hyperparameters. isolated_atom_energies: dict: mandatory dictionary mapping element numbers to isolated energies. ref_energy_name : str, optional @@ -360,7 +360,7 @@ def jace_fitting( ase.io.write("train_ace.extxyz", train_ace, format="extxyz") if fit_kwargs: - hyperparameters.update_fields(fit_kwargs) + hyperparameters.update_parameters(fit_kwargs) jace_hypers = hyperparameters.model_dump(by_alias=True) @@ -443,7 +443,7 @@ def jace_fitting( def nequip_fitting( db_dir: Path, - path_to_hyperparameters: Path | str = MLIP_RSS_DEFAULTS_FILE_PATH, + hyperparameters: MLIP_HYPERS.NEQUIP = MLIP_HYPERS.NEQUIP, isolated_atom_energies: dict | None = None, ref_energy_name: str = "REF_energy", ref_force_name: str = "REF_forces", @@ -462,8 +462,8 @@ def nequip_fitting( ---------- db_dir: Path directory containing the training and testing data files. - path_to_hyperparameters : str or Path. - Path to JSON file containing the NwquIP hyperparameters. + hyperparameters: MLIP_HYPERS.NEQUIP + Fit hyperparameters. isolated_atom_energies: dict mandatory dictionary mapping element numbers to isolated energies. ref_energy_name : str, optional @@ -513,8 +513,8 @@ def nequip_fitting( """ [TODO] train Nequip on virials """ - if path_to_hyperparameters is None: - path_to_hyperparameters = MLIP_RSS_DEFAULTS_FILE_PATH + hyperparameters = hyperparameters.model_copy(deep=True) + train_data = ase.io.read(os.path.join(db_dir, "train.extxyz"), index=":") train_nequip = [ at for at in train_data if "IsolatedAtom" not in at.info["config_type"] @@ -535,17 +535,10 @@ def nequip_fitting( else: raise ValueError("isolated_atom_energies is empty or not defined!") - nequip_hypers = MLIP_HYPERS.NEQUIP.model_dump(by_alias=True) - if fit_kwargs: - for parameter in nequip_hypers: - if parameter in fit_kwargs: - if isinstance(fit_kwargs[parameter], type(nequip_hypers[parameter])): - nequip_hypers[parameter] = fit_kwargs[parameter] - else: - raise TypeError( - f"The type of {parameter} should be {type(nequip_hypers[parameter])}!" - ) + hyperparameters.update_parameters(fit_kwargs) + + nequip_hypers = hyperparameters.model_dump(by_alias=True) r_max = nequip_hypers["r_max"] num_layers = nequip_hypers["num_layers"] @@ -720,7 +713,7 @@ def nequip_fitting( def m3gnet_fitting( db_dir: Path, - path_to_hyperparameters: Path | str = MLIP_RSS_DEFAULTS_FILE_PATH, + hyperparameters: MLIP_HYPERS.M3GNET = MLIP_HYPERS.M3GNET, device: str = "cuda", ref_energy_name: str = "REF_energy", ref_force_name: str = "REF_forces", @@ -734,8 +727,8 @@ def m3gnet_fitting( ---------- db_dir: Path Directory containing the training and testing data files. - path_to_hyperparameters : str or Path. - Path to JSON file containing the M3GNet hyperparameters. + hyperparameters: MLIP_HYPERS.M3GNET + Fit hyperparameters. device: str Device on which the model will be trained, e.g., 'cuda' or 'cpu'. ref_energy_name : str, optional @@ -790,20 +783,12 @@ def m3gnet_fitting( * Availability: https://matgl.ai/tutorials%2FTraining%20a%20M3GNet%20Potential%20with%20PyTorch%20Lightning.html * License: BSD 3-Clause License """ - if path_to_hyperparameters is None: - path_to_hyperparameters = MLIP_RSS_DEFAULTS_FILE_PATH - - m3gnet_hypers = MLIP_HYPERS.M3GNET.model_dump(by_alias=True) + hyperparameters = hyperparameters.model_copy(deep=True) if fit_kwargs: - for parameter in m3gnet_hypers: - if parameter in fit_kwargs: - if isinstance(fit_kwargs[parameter], type(m3gnet_hypers[parameter])): - m3gnet_hypers[parameter] = fit_kwargs[parameter] - else: - raise TypeError( - f"The type of {parameter} should be {type(m3gnet_hypers[parameter])}!" - ) + hyperparameters.update_parameters(fit_kwargs) + + m3gnet_hypers = hyperparameters.model_dump(by_alias=True) exp_name = m3gnet_hypers["exp_name"] results_dir = m3gnet_hypers["results_dir"] @@ -1091,7 +1076,7 @@ def m3gnet_fitting( def mace_fitting( db_dir: Path, - path_to_hyperparameters: Path | str = MLIP_RSS_DEFAULTS_FILE_PATH, + hyperparameters: MLIP_HYPERS.MACE = MLIP_HYPERS.MACE, device: str = "cuda", ref_energy_name: str = "REF_energy", ref_force_name: str = "REF_forces", @@ -1110,8 +1095,8 @@ def mace_fitting( ---------- db_dir: Path directory containing the training and testing data files. - path_to_hyperparameters : str or Path. - Path to JSON file containing the MACE hyperparameters. + hyperparameters: MLIP_HYPERS.MACE + Fit hyperparameters. device: str specify device to use cuda or cpu. ref_energy_name : str, optional @@ -1152,15 +1137,15 @@ def mace_fitting( A dictionary containing train_error, test_error, and the path to the fitted MLIP. """ - if path_to_hyperparameters is None: - path_to_hyperparameters = MLIP_RSS_DEFAULTS_FILE_PATH + hyperparameters = hyperparameters.model_copy(deep=True) + if ref_virial_name is not None: atoms = read(f"{db_dir}/train.extxyz", index=":") mace_virial_format_conversion( atoms=atoms, ref_virial_name=ref_virial_name, out_file_name="train.extxyz" ) - mace_hypers = MLIP_HYPERS.MACE.model_dump(by_alias=True) if use_defaults else {} + mace_hypers = hyperparameters.model_dump(by_alias=True) if use_defaults else {} # TODO: should we do a type check? not sure # as it will be a lot of work to keep it updated From 5a155641c86b5dbe240e0929606f4c21fee8f2a9 Mon Sep 17 00:00:00 2001 From: naik-aakash Date: Sun, 12 Jan 2025 20:59:36 +0100 Subject: [PATCH 25/74] method rename not applied > fix failing test --- tests/fitting/common/test_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/fitting/common/test_utils.py b/tests/fitting/common/test_utils.py index 43b224c72..4bfc51e42 100644 --- a/tests/fitting/common/test_utils.py +++ b/tests/fitting/common/test_utils.py @@ -87,7 +87,7 @@ def test_gap_hyperparameter_constructor(): # test if returned string is changed if passed in dict is updated gap_hyper_parameter = MLIP_HYPERS.GAP.model_copy(deep=True) - gap_hyper_parameter.update_fields({"twob": {"cutoff": 8}, + gap_hyper_parameter.update_parameters({"twob": {"cutoff": 8}, "threeb": {"cutoff": 8.0, "n_sparse": 100}, "soap": {"delta": 1.5, From a713593f83a14dee341e69069e9bae1816307db4 Mon Sep 17 00:00:00 2001 From: naik-aakash Date: Sun, 12 Jan 2025 21:07:27 +0100 Subject: [PATCH 26/74] method rename not applied > fix failing test --- tests/fitting/common/test_utils.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/fitting/common/test_utils.py b/tests/fitting/common/test_utils.py index 4bfc51e42..8ad176a95 100644 --- a/tests/fitting/common/test_utils.py +++ b/tests/fitting/common/test_utils.py @@ -6,7 +6,6 @@ check_convergence, data_distillation, prepare_fit_environment, - MLIP_PHONON_DEFAULTS_FILE_PATH ) def test_stratified_split(test_dir): From 6e7886d5b80c6ea08df13cc89c93af8317c97095 Mon Sep 17 00:00:00 2001 From: naik-aakash Date: Wed, 15 Jan 2025 17:32:22 +0100 Subject: [PATCH 27/74] add MACE kwargs --- src/autoplex/fitting/common/utils.py | 12 ++- src/autoplex/settings.py | 121 +++++++++++++++++++++++---- 2 files changed, 112 insertions(+), 21 deletions(-) diff --git a/src/autoplex/fitting/common/utils.py b/src/autoplex/fitting/common/utils.py index e7ed26e0b..05ed190ee 100644 --- a/src/autoplex/fitting/common/utils.py +++ b/src/autoplex/fitting/common/utils.py @@ -1105,6 +1105,8 @@ def mace_fitting( Reference force name. ref_virial_name : str, optional Reference virial name. + use_defaults: bool + if True, use the default hyperparameters. fit_kwargs: dict. optional dictionary with parameters for mace fitting with keys same as mlip-rss-defaults.json. @@ -1145,11 +1147,13 @@ def mace_fitting( atoms=atoms, ref_virial_name=ref_virial_name, out_file_name="train.extxyz" ) - mace_hypers = hyperparameters.model_dump(by_alias=True) if use_defaults else {} + hyperparameters.update_parameters(fit_kwargs) - # TODO: should we do a type check? not sure - # as it will be a lot of work to keep it updated - mace_hypers.update(fit_kwargs) + mace_hypers = ( + hyperparameters.model_dump(by_alias=True, exclude_none=True) + if use_defaults + else {} + ) boolean_hypers = [ "distributed", diff --git a/src/autoplex/settings.py b/src/autoplex/settings.py index b010ee213..8ea512d34 100644 --- a/src/autoplex/settings.py +++ b/src/autoplex/settings.py @@ -245,32 +245,119 @@ class M3GNETSettings(UpdateBaseModel): class MACESettings(UpdateBaseModel): """Model describing the hyperparameters for the MACE fits.""" - model: str = Field(default="MACE", description="type of the model") - name: str = Field(default="MACE_model", description="Name of the model") + model: Literal[ + "BOTNet", + "MACE", + "ScaleShiftMACE", + "ScaleShiftBOTNet", + "AtomicDipolesMACE", + "EnergyDipolesMACE", + ] = Field(default="MACE", description="type of the model") + name: str = Field(default="MACE_model", description="Experiment name") + amsgrad: bool = Field(default=True, description="Use amsgrad variant of optimizer") + batch_size: int = Field(default=10, description="Batch size") + compute_avg_num_neighbors: ( + bool | Literal["yes", "true", "t", "y", "1", "no", "false", "f", "n", "0"] + ) = Field(default=True, description="Compute average number of neighbors") + compute_forces: ( + bool | Literal["yes", "true", "t", "y", "1", "no", "false", "f", "n", "0"] + ) = Field(default=True, description="Compute forces") config_type_weights: str = Field( - default="{'Default':1.0}", description="Weights for the configuration types" + default="{'Default':1.0}", + description="String of dictionary containing the weights for each config type", + ) + compute_stress: ( + bool | Literal["yes", "true", "t", "y", "1", "no", "false", "f", "n", "0"] + ) = Field(default=False, description="Compute stress") + compute_statistics: bool = Field(default=False, description="Compute statistics") + correlation: int = Field(default=3, description="Correlation order at each layer") + default_dtype: Literal["float32", "float64"] = Field( + default="float32", description="Default data type" + ) + device: Literal["cpu", "cuda", "mps", "xpu"] = Field( + default="cpu", description="Device to be used for model fitting" ) - hidden_irreps: str = Field(default="128x0e + 128x1o", description="Hidden irreps") - r_max: float = Field(default=5.0, description="Radial cutoff distance") - batch_size: int = Field(default=10, description="Batch size") - max_num_epochs: int = Field(default=1500, description="Maximum number of epochs") - start_swa: int = Field(default=1200, description="Start of the SWA") + distributed: bool = Field( + default=False, description="Train in multi-GPU data parallel mode" + ) + energy_weight: float = Field(default=1.0, description="Weight for the energy loss") + ema: bool = Field(default=True, description="Whether to use EMA") ema_decay: float = Field( default=0.99, description="Exponential moving average decay" ) - correlation: int = Field(default=3, description="Correlation") - loss: str = Field(default="huber", description="Loss function") - default_dtype: str = Field(default="float32", description="Default data type") - swa: bool = Field(default=True, description="Whether to use SWA") - ema: bool = Field(default=True, description="Whether to use EMA") - amsgrad: bool = Field(default=True, description="Whether to use AMSGrad") + E0s: str | None = Field( + default=None, description="Dictionary of isolated atom energies" + ) + forces_weight: float = Field( + default=100.0, description="Weight for the forces loss" + ) + foundation_filter_elements: ( + bool | Literal["yes", "true", "t", "y", "1", "no", "false", "f", "n", "0"] + ) = Field(default=True, description="Filter element during fine-tuning") + foundation_model: str | None = Field( + default=None, description="Path to the foundation model for finetuning" + ) + foundation_model_readout: bool = Field( + default=True, description="Use readout of foundation model for finetuning" + ) + keep_checkpoint: bool = Field(default=False, description="Keep all checkpoints") + keep_isolated_atoms: ( + bool | Literal["yes", "true", "t", "y", "1", "no", "false", "f", "n", "0"] + ) = Field( + default=False, + description="Keep isolated atoms in the dataset, useful for finetuning", + ) + hidden_irreps: str = Field(default="128x0e + 128x1o", description="Hidden irreps") + loss: Literal[ + "ef", + "weighted", + "forces_only", + "virials", + "stress", + "dipole", + "huber", + "universal", + "energy_forces_dipole", + ] = Field(default="huber", description="Loss function") + lr: float = Field(default=0.001, description="Learning rate") + multiheads_finetuning: ( + bool | Literal["yes", "true", "t", "y", "1", "no", "false", "f", "n", "0"] + ) = Field(default=False, description="Multiheads finetuning") + max_num_epochs: int = Field(default=1500, description="Maximum number of epochs") + pair_repulsion: bool = Field( + default=False, description="Use pair repulsion term with ZBL potential" + ) + patience: int = Field( + default=2048, + description="Maximum number of consecutive epochs of increasing loss", + ) + r_max: float = Field(default=5.0, description="Radial cutoff distance") restart_latest: bool = Field( - default=True, description="Whether to restart the latest model" + default=False, description="Whether to restart the latest model" ) seed: int = Field(default=123, description="Seed for the random number generator") - device: Literal["cpu", "cuda"] = Field( - default="cpu", description="Device to be used for model fitting" + save_cpu: bool = Field(default=True, description="Save CPU") + save_all_checkpoints: bool = Field( + default=False, description="Save all checkpoints" + ) + scaling: Literal["std_scaling", "rms_forces_scaling", "no_scaling"] = Field( + default="rms_forces_scaling", description="Scaling" + ) + stress_weight: float = Field(default=1.0, description="Weight for the stress loss") + start_swa: int = Field( + default=1200, description="Start of the SWA", alias="start_stage_two" + ) + swa: bool = Field( + default=True, + description="Use Stage Two loss weight, it will decrease the learning " + "rate and increases the energy weight at the end of the training", + alias="stage_two", + ) + valid_batch_size: int = Field(default=10, description="Validation batch size") + virials_weight: float = Field( + default=1.0, description="Weight for the virials loss" ) + wandb: bool = Field(default=False, description="Use Weights and Biases for logging") class NEPSettings(UpdateBaseModel): From 0b253f72b6403fb50258e6115c82ddb58415b384 Mon Sep 17 00:00:00 2001 From: naik-aakash Date: Wed, 15 Jan 2025 18:36:59 +0100 Subject: [PATCH 28/74] update MACE finetuning tests --- tests/auto/phonons/test_flows.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/auto/phonons/test_flows.py b/tests/auto/phonons/test_flows.py index 4ec5084b4..c69045341 100644 --- a/tests/auto/phonons/test_flows.py +++ b/tests/auto/phonons/test_flows.py @@ -994,7 +994,7 @@ def test_complete_dft_vs_ml_benchmark_workflow_mace_finetuning( volume_custom_scale_factors=[0.975, 1.0, 1.025, 1.05], benchmark_kwargs={"calculator_kwargs": {"device": "cpu"}}, apply_data_preprocessing=True, - use_defaults_fitting=False, + use_defaults_fitting=True, ).make( structure_list=[structure], mp_ids=["test"], @@ -1064,7 +1064,7 @@ def test_complete_dft_vs_ml_benchmark_workflow_mace_finetuning_mp_settings( benchmark_kwargs={"calculator_kwargs": {"device": "cpu"}}, add_dft_rattled_struct=True, apply_data_preprocessing=True, - use_defaults_fitting=False, + use_defaults_fitting=True, split_ratio=0.3, ).make( structure_list=[structure], From 355ff75fcd394ada6639537b994c89354d8529a7 Mon Sep 17 00:00:00 2001 From: naik-aakash Date: Wed, 15 Jan 2025 20:36:17 +0100 Subject: [PATCH 29/74] move database_dir back to make --- src/autoplex/auto/rss/jobs.py | 2 +- src/autoplex/fitting/common/flows.py | 12 ++++++------ 2 files changed, 7 insertions(+), 7 deletions(-) diff --git a/src/autoplex/auto/rss/jobs.py b/src/autoplex/auto/rss/jobs.py index dd0d73c2a..34e630021 100644 --- a/src/autoplex/auto/rss/jobs.py +++ b/src/autoplex/auto/rss/jobs.py @@ -221,9 +221,9 @@ def initial_rss( apply_data_preprocessing=False, auto_delta=auto_delta, glue_xml=False, - database_dir=do_data_preprocessing.output, ).make( isolated_atom_energies=do_data_collection.output["isolated_atom_energies"], + database_dir=do_data_preprocessing.output, device=device_for_fitting, **fit_kwargs, ) diff --git a/src/autoplex/fitting/common/flows.py b/src/autoplex/fitting/common/flows.py index 0e317c078..403a89df6 100644 --- a/src/autoplex/fitting/common/flows.py +++ b/src/autoplex/fitting/common/flows.py @@ -83,8 +83,6 @@ class MLIPFitMaker(Maker): Number of processes for fitting. apply_data_preprocessing: bool Determine whether to preprocess the data. - database_dir: Path | str - Path to the directory containing the database. use_defaults: bool If true, uses default fit parameters run_fits_on_different_cluster: bool @@ -112,13 +110,13 @@ class MLIPFitMaker(Maker): glue_xml: bool = False # This is only used for GAP. num_processes_fit: int | None = None apply_data_preprocessing: bool = True - database_dir: Path | str | None = None use_defaults: bool = True run_fits_on_different_cluster: bool = False def make( self, fit_input: dict | None = None, # This is specific to phonon workflow + database_dir: Path | str | None = None, hyperparameters: MLIP_HYPERS = MLIP_HYPERS, species_list: list | None = None, isolated_atom_energies: dict | None = None, @@ -132,6 +130,8 @@ def make( ---------- fit_input: dict Output from the CompletePhononDFTMLDataGenerationFlow process. + database_dir: Path | str + Path to the directory containing the database. hyperparameters: MLIP_HYPERS Hyperparameters for the MLIP. species_list: list @@ -200,11 +200,11 @@ def make( # this will only run if train.extxyz and test.extxyz files are present in the database_dir # TODO: shouldn't this be the exception rather then the default run?! # TODO: I assume we always want to use data from before? - if isinstance(self.database_dir, str): - self.database_dir = Path(self.database_dir) + if isinstance(database_dir, str): + database_dir = Path(database_dir) mlip_fit_job = machine_learning_fit( - database_dir=self.database_dir, + database_dir=database_dir, isolated_atom_energies=isolated_atom_energies, num_processes_fit=self.num_processes_fit, auto_delta=self.auto_delta, From cb2eaf30d8270d662f09ef12126534b9d0df4b47 Mon Sep 17 00:00:00 2001 From: naik-aakash Date: Wed, 15 Jan 2025 20:40:45 +0100 Subject: [PATCH 30/74] adapt tests --- tests/auto/phonons/test_jobs.py | 2 +- tests/conftest.py | 4 ++-- tests/fitting/common/test_flows.py | 8 ++++---- tests/fitting/common/test_jobs.py | 14 +++++++------- 4 files changed, 14 insertions(+), 14 deletions(-) diff --git a/tests/auto/phonons/test_jobs.py b/tests/auto/phonons/test_jobs.py index 0a7f81559..06ab67606 100644 --- a/tests/auto/phonons/test_jobs.py +++ b/tests/auto/phonons/test_jobs.py @@ -163,10 +163,10 @@ def test_complete_benchmark(clean_dir, test_dir, memory_jobstore): glue_xml=False, apply_data_preprocessing=False, separated=True, - database_dir=database_dir, ).make( twob={"delta": 2.0, "cutoff": 4}, threeb={"n_sparse": 10}, + database_dir=database_dir, **fit_kwargs ) dft_data = loadfn(test_dir / "benchmark" / "phonon_doc_si.json") diff --git a/tests/conftest.py b/tests/conftest.py index 63206345c..c7cd56761 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -180,8 +180,8 @@ def mock_rss(input_dir: str = None, ref_virial_name=ref_virial_name, num_processes_fit=num_processes_fit, apply_data_preprocessing=False, - database_dir=job5.output, - ).make(isolated_atom_energies=job4.output['isolated_atom_energies'], **fit_kwargs) + ).make(isolated_atom_energies=job4.output['isolated_atom_energies'], + database_dir=job5.output, **fit_kwargs) job_list = [job2, job3, job4, job5, job6] return Response( diff --git a/tests/fitting/common/test_flows.py b/tests/fitting/common/test_flows.py index bb8e4372a..30201856c 100644 --- a/tests/fitting/common/test_flows.py +++ b/tests/fitting/common/test_flows.py @@ -137,7 +137,7 @@ def test_mlip_fit_maker(test_dir, clean_dir, memory_jobstore, vasp_test_dir, fit fit_input=fit_input_dict, ) - responses = run_locally( + _ = run_locally( gapfit, ensure_success=True, create_folders=True, store=memory_jobstore ) @@ -167,7 +167,7 @@ def test_mlip_fit_maker_with_kwargs( threeb={"n_sparse": 100}, ) - responses = run_locally( + _ = run_locally( gapfit, ensure_success=True, create_folders=True, store=memory_jobstore ) @@ -379,7 +379,7 @@ def test_mlip_fit_maker_glue_xml( general={"core_param_file": "glue.xml", "core_ip_args": "{IP Glue}"}, ) - responses = run_locally( + _ = run_locally( gapfit, ensure_success=True, create_folders=True, store=memory_jobstore ) @@ -412,7 +412,7 @@ def test_mlip_fit_maker_glue_xml_with_other_name( general={"core_param_file": "glue.xml", "core_ip_args": "{IP Glue}"}, ) - responses = run_locally( + _ = run_locally( gapfit, ensure_success=True, create_folders=True, store=memory_jobstore ) diff --git a/tests/fitting/common/test_jobs.py b/tests/fitting/common/test_jobs.py index 1b1fc2bad..4aeef9ffc 100644 --- a/tests/fitting/common/test_jobs.py +++ b/tests/fitting/common/test_jobs.py @@ -11,10 +11,10 @@ def test_gap_fit_maker(test_dir, memory_jobstore, clean_dir): auto_delta=False, glue_xml=False, apply_data_preprocessing=False, - database_dir=database_dir ).make( twob={"delta": 2.0, "cutoff": 4}, threeb={"n_sparse": 10}, + database_dir=database_dir ) responses = run_locally( @@ -32,8 +32,8 @@ def test_jace_fit_maker(test_dir, memory_jobstore, clean_dir): mlip_type="J-ACE", num_processes_fit=4, apply_data_preprocessing=False, - database_dir=database_dir, ).make( + database_dir=database_dir, isolated_atom_energies={14: -0.84696938}, order=2, totaldegree=4, @@ -53,8 +53,8 @@ def test_nequip_fit_maker(test_dir, memory_jobstore, clean_dir): mlip_type="NEQUIP", num_processes_fit=1, apply_data_preprocessing=False, - database_dir=database_dir, ).make( + database_dir=database_dir, isolated_atom_energies={14: -0.84696938}, r_max=3.14, max_epochs=10, @@ -75,8 +75,8 @@ def test_m3gnet_fit_maker(test_dir, memory_jobstore, clean_dir): mlip_type="M3GNET", num_processes_fit=1, apply_data_preprocessing=False, - database_dir=database_dir, ).make( + database_dir=database_dir, isolated_atom_energies={14: -0.84696938}, cutoff=3.0, threebody_cutoff=2.0, @@ -105,8 +105,8 @@ def test_mace_fit_maker(test_dir, memory_jobstore, clean_dir): mlip_type="MACE", num_processes_fit=1, apply_data_preprocessing=False, - database_dir=database_dir, ).make( + database_dir=database_dir, isolated_atom_energies={14: -0.84696938}, model="MACE", config_type_weights='{"Default":1.0}', @@ -140,8 +140,8 @@ def test_mace_finetuning_maker(test_dir, memory_jobstore, clean_dir): use_defaults=False, num_processes_fit=1, apply_data_preprocessing=False, - database_dir=database_dir, ).make( + database_dir=database_dir, name="MACE_final", foundation_model="small", multiheads_finetuning=False, @@ -185,8 +185,8 @@ def test_mace_finetuning_maker2(test_dir, memory_jobstore, clean_dir): use_defaults=False, num_processes_fit=1, apply_data_preprocessing=False, - database_dir=database_dir, ).make( + database_dir=database_dir, name="MACE_final", foundation_model="small", multiheads_finetuning=False, From 7e8c65e2794b649f1d35c722b758a7071658ac6f Mon Sep 17 00:00:00 2001 From: naik-aakash Date: Wed, 15 Jan 2025 20:59:04 +0100 Subject: [PATCH 31/74] remove use_defaults arg --- src/autoplex/auto/phonons/flows.py | 6 ------ src/autoplex/fitting/common/flows.py | 4 ---- src/autoplex/fitting/common/jobs.py | 4 ---- src/autoplex/fitting/common/utils.py | 9 +-------- tests/auto/phonons/test_flows.py | 2 -- tests/fitting/common/test_jobs.py | 2 -- 6 files changed, 1 insertion(+), 26 deletions(-) diff --git a/src/autoplex/auto/phonons/flows.py b/src/autoplex/auto/phonons/flows.py index fdeb9f57c..2deeade92 100644 --- a/src/autoplex/auto/phonons/flows.py +++ b/src/autoplex/auto/phonons/flows.py @@ -172,8 +172,6 @@ class CompleteDFTvsMLBenchmarkWorkflow(Maker): Use the glue.xml core potential instead of fitting 2b terms. glue_file_path: str Name of the glue.xml file path. - use_defaults_fitting: bool - Use the fit defaults. run_fits_on_different_cluster: bool Allows you to run fits on a different cluster than DFT (will transfer fit database via MongoDB, might be slow). @@ -226,7 +224,6 @@ class CompleteDFTvsMLBenchmarkWorkflow(Maker): summary_filename_prefix: str = "results_" glue_xml: bool = False glue_file_path: str = "glue.xml" - use_defaults_fitting: bool = True run_fits_on_different_cluster: bool = False def make( @@ -387,7 +384,6 @@ def make( mlip_type=ml_model, glue_xml=self.glue_xml, glue_file_path=self.glue_file_path, - use_defaults=self.use_defaults_fitting, split_ratio=self.split_ratio, force_max=self.force_max, pre_xyz_files=pre_xyz_files, @@ -820,8 +816,6 @@ class CompleteDFTvsMLBenchmarkWorkflowMPSettings(CompleteDFTvsMLBenchmarkWorkflo Prefix of the result summary file. glue_file_path: str Name of the glue.xml file path. - use_defaults_fitting: bool - Use the fit defaults. """ phonon_bulk_relax_maker: BaseVaspMaker = field( diff --git a/src/autoplex/fitting/common/flows.py b/src/autoplex/fitting/common/flows.py index 403a89df6..356bbc04a 100644 --- a/src/autoplex/fitting/common/flows.py +++ b/src/autoplex/fitting/common/flows.py @@ -83,8 +83,6 @@ class MLIPFitMaker(Maker): Number of processes for fitting. apply_data_preprocessing: bool Determine whether to preprocess the data. - use_defaults: bool - If true, uses default fit parameters run_fits_on_different_cluster: bool If true, run fits on different clusters. """ @@ -110,7 +108,6 @@ class MLIPFitMaker(Maker): glue_xml: bool = False # This is only used for GAP. num_processes_fit: int | None = None apply_data_preprocessing: bool = True - use_defaults: bool = True run_fits_on_different_cluster: bool = False def make( @@ -181,7 +178,6 @@ def make( ref_energy_name=self.ref_energy_name, ref_force_name=self.ref_force_name, ref_virial_name=self.ref_virial_name, - use_defaults=self.use_defaults, device=device, species_list=species_list, database_dict=data_prep_job.output["database_dict"], diff --git a/src/autoplex/fitting/common/jobs.py b/src/autoplex/fitting/common/jobs.py index 893553a8a..4e5cfa624 100644 --- a/src/autoplex/fitting/common/jobs.py +++ b/src/autoplex/fitting/common/jobs.py @@ -29,7 +29,6 @@ def machine_learning_fit( ref_energy_name: str = "REF_energy", ref_force_name: str = "REF_forces", ref_virial_name: str = "REF_virial", - use_defaults: bool = True, device: str = "cuda", database_dict: dict | None = None, hyperpara_opt: bool = False, @@ -65,8 +64,6 @@ def machine_learning_fit( Reference force name. ref_virial_name: str Reference virial name. - use_defaults: bool - If True, use default fitting parameters device: str Device to be used for model fitting, either "cpu" or "cuda". database_dict: dict @@ -179,7 +176,6 @@ def machine_learning_fit( ref_energy_name=ref_energy_name, ref_force_name=ref_force_name, ref_virial_name=ref_virial_name, - use_defaults=use_defaults, device=device, fit_kwargs=fit_kwargs, ) diff --git a/src/autoplex/fitting/common/utils.py b/src/autoplex/fitting/common/utils.py index 05ed190ee..084cb3d50 100644 --- a/src/autoplex/fitting/common/utils.py +++ b/src/autoplex/fitting/common/utils.py @@ -1081,7 +1081,6 @@ def mace_fitting( ref_energy_name: str = "REF_energy", ref_force_name: str = "REF_forces", ref_virial_name: str = "REF_virial", - use_defaults=True, fit_kwargs: dict | None = None, ) -> dict: """ @@ -1105,8 +1104,6 @@ def mace_fitting( Reference force name. ref_virial_name : str, optional Reference virial name. - use_defaults: bool - if True, use the default hyperparameters. fit_kwargs: dict. optional dictionary with parameters for mace fitting with keys same as mlip-rss-defaults.json. @@ -1149,11 +1146,7 @@ def mace_fitting( hyperparameters.update_parameters(fit_kwargs) - mace_hypers = ( - hyperparameters.model_dump(by_alias=True, exclude_none=True) - if use_defaults - else {} - ) + mace_hypers = hyperparameters.model_dump(by_alias=True, exclude_none=True) boolean_hypers = [ "distributed", diff --git a/tests/auto/phonons/test_flows.py b/tests/auto/phonons/test_flows.py index c69045341..c0a3be30e 100644 --- a/tests/auto/phonons/test_flows.py +++ b/tests/auto/phonons/test_flows.py @@ -994,7 +994,6 @@ def test_complete_dft_vs_ml_benchmark_workflow_mace_finetuning( volume_custom_scale_factors=[0.975, 1.0, 1.025, 1.05], benchmark_kwargs={"calculator_kwargs": {"device": "cpu"}}, apply_data_preprocessing=True, - use_defaults_fitting=True, ).make( structure_list=[structure], mp_ids=["test"], @@ -1064,7 +1063,6 @@ def test_complete_dft_vs_ml_benchmark_workflow_mace_finetuning_mp_settings( benchmark_kwargs={"calculator_kwargs": {"device": "cpu"}}, add_dft_rattled_struct=True, apply_data_preprocessing=True, - use_defaults_fitting=True, split_ratio=0.3, ).make( structure_list=[structure], diff --git a/tests/fitting/common/test_jobs.py b/tests/fitting/common/test_jobs.py index 4aeef9ffc..0f7532bcb 100644 --- a/tests/fitting/common/test_jobs.py +++ b/tests/fitting/common/test_jobs.py @@ -137,7 +137,6 @@ def test_mace_finetuning_maker(test_dir, memory_jobstore, clean_dir): ref_energy_name=None, ref_force_name=None, ref_virial_name=None, - use_defaults=False, num_processes_fit=1, apply_data_preprocessing=False, ).make( @@ -182,7 +181,6 @@ def test_mace_finetuning_maker2(test_dir, memory_jobstore, clean_dir): ref_energy_name=None, ref_force_name=None, ref_virial_name=None, - use_defaults=False, num_processes_fit=1, apply_data_preprocessing=False, ).make( From 37752a33ed24103e663d5efb0ff024717e314157 Mon Sep 17 00:00:00 2001 From: naik-aakash Date: Fri, 17 Jan 2025 15:02:04 +0100 Subject: [PATCH 32/74] update NEQUIP hypers --- src/autoplex/__init__.py | 17 ++- src/autoplex/settings.py | 302 ++++++++++++++++++++++++++++++++++----- 2 files changed, 284 insertions(+), 35 deletions(-) diff --git a/src/autoplex/__init__.py b/src/autoplex/__init__.py index e2e6e9677..58c8b7c24 100644 --- a/src/autoplex/__init__.py +++ b/src/autoplex/__init__.py @@ -8,7 +8,22 @@ """ from autoplex._version import __version__ -from autoplex.settings import MLIPHypers, RssConfig +from autoplex.settings import ( + GAPSettings, + JACESettings, + M3GNETSettings, + MACESettings, + MLIPHypers, + NEPSettings, + NEQUIPSettings, + RssConfig, +) MLIP_HYPERS = MLIPHypers() RSS_CONFIG = RssConfig() +GAP_HYPERS = GAPSettings() +JACE_HYPERS = JACESettings() +M3GNET_HYPERS = M3GNETSettings() +MACE_HYPERS = MACESettings() +NEQUIP_HYPERS = NEQUIPSettings() +NEP_HYPERS = NEPSettings() diff --git a/src/autoplex/settings.py b/src/autoplex/settings.py index 8ea512d34..2ea942c4d 100644 --- a/src/autoplex/settings.py +++ b/src/autoplex/settings.py @@ -2,16 +2,23 @@ from __future__ import annotations +import logging from typing import Any, Literal from monty.serialization import loadfn from pydantic import BaseModel, ConfigDict, Field +logging.basicConfig( + level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s" +) -class UpdateBaseModel(BaseModel): + +class AutoplexBaseModel(BaseModel): """Base class for all models in autoplex.""" - model_config = ConfigDict(validate_assignment=True, protected_namespaces=()) + model_config = ConfigDict( + validate_assignment=True, protected_namespaces=(), extra="allow" + ) def update_parameters(self, updates: dict[str, Any]): """ @@ -27,15 +34,17 @@ def update_parameters(self, updates: dict[str, Any]): # Update nested model field_value.update_parameters( value - ) # Recursively call update_fields + ) # Recursively call update_parameters else: # Update field value setattr(self, key, value) - # else: - # raise ValueError( - # f"Field {key} does not exist in {self.__class__.__name__}." - # ) + else: + logging.warning( + f"Field {key} not found in default {self.__class__.__name__} model." + f"New field has been added. Please ensure the added field contains correct datatype." + ) + setattr(self, key, value) @classmethod def from_file(cls, filename: str): @@ -49,7 +58,7 @@ def from_file(cls, filename: str): return cls(**custom_params) -class GeneralSettings(UpdateBaseModel): +class GeneralSettings(AutoplexBaseModel): """Model describing general hyperparameters for the GAP fits.""" at_file: str = Field( @@ -81,7 +90,7 @@ class GeneralSettings(UpdateBaseModel): soap: bool = Field(default=True, description="Whether to include SOAP terms") -class TwobSettings(UpdateBaseModel): +class TwobSettings(AutoplexBaseModel): """Model describing two body hyperparameters for the GAP fits.""" distance_Nb_order: int = Field( @@ -110,7 +119,7 @@ class TwobSettings(UpdateBaseModel): ) -class ThreebSettings(UpdateBaseModel): +class ThreebSettings(AutoplexBaseModel): """Model describing threebody hyperparameters for the GAP fits.""" distance_Nb_order: int = Field( @@ -139,7 +148,7 @@ class ThreebSettings(UpdateBaseModel): ) -class SoapSettings(UpdateBaseModel): +class SoapSettings(AutoplexBaseModel): """Model describing soap hyperparameters for the GAP fits.""" add_species: str = Field( @@ -167,7 +176,7 @@ class SoapSettings(UpdateBaseModel): ) -class GAPSettings(UpdateBaseModel): +class GAPSettings(AutoplexBaseModel): """Model describing the hyperparameters for the GAP fits for Phonons.""" general: GeneralSettings = Field( @@ -188,7 +197,7 @@ class GAPSettings(UpdateBaseModel): ) -class JACESettings(UpdateBaseModel): +class JACESettings(AutoplexBaseModel): """Model describing the hyperparameters for the J-ACE fits.""" order: int = Field(default=3, description="Order of the J-ACE model") @@ -197,25 +206,250 @@ class JACESettings(UpdateBaseModel): solver: str = Field(default="BLR", description="Solver for the J-ACE model") -class NEQUIPSettings(UpdateBaseModel): - """Model describing the hyperparameters for the NEQUIP fits.""" +class Nonlinearity(AutoplexBaseModel): + """Model describing the nonlinearity to be used for the NEQUIP fits.""" + + e: Literal["silu", "ssp", "tanh", "abs"] = Field( + default="silu", description="Even nonlinearity" + ) + o: Literal["silu", "ssp", "tanh", "abs"] = Field( + default="tanh", description="Odd nonlinearity" + ) + + +class LossCoeff(BaseModel): + """Model describing different weights to use in a weighted loss functions.""" + + forces: int | list[int | str] = Field( + default=1, description="Forces loss coefficient" + ) + total_energy: int | list[int | str] = Field( + default=[1, "PerAtomMSELoss"], description="Total energy loss coefficient" + ) + +class NEQUIPSettings(AutoplexBaseModel): + """Model describing the hyperparameters for the NEQUIP fits. + + References + ---------- + * Defaults taken from https://github.com/mir-group/nequip/blob/main/configs/ + """ + + root: str = Field(default="results", description="Root directory") + run_name: str = Field(default="autoplex", description="Name of the run") + seed: int = Field(default=123, description="Model seed") + dataset_seed: int = Field(default=123, description="Dataset seed") + append: bool = Field( + default=False, + description="When true a restarted run will append to the previous log file", + ) + default_dtype: str = Field(default="float32", description="Default data type") + model_dtype: str = Field(default="float32", description="Model data type") + allow_tf32: bool = Field( + default=True, + description="Consider setting to false if you plan to mix " + "training/inference over any devices that are " + "not NVIDIA Ampere or later", + ) r_max: float = Field(default=4.0, description="Radial cutoff distance") num_layers: int = Field(default=4, description="Number of layers") l_max: int = Field(default=2, description="Maximum degree of spherical harmonics") + parity: bool = Field( + default=True, + description="Whether to include features with odd mirror parity; " + "often turning parity off gives equally good results but faster networks", + ) num_features: int = Field(default=32, description="Number of features") - num_basis: int = Field(default=8, description="Number of basis functions") - invariant_layers: int = Field(default=2, description="Number of invariant layers") + nonlinearity_type: Literal["gate", "norm"] = Field( + default="gate", description="Type of nonlinearity, 'gate' is recommended" + ) + nonlinearity_scalars: Nonlinearity = Field( + default_factory=Nonlinearity, description="Nonlinearity scalars" + ) + nonlinearity_gates: Nonlinearity = Field( + default_factory=Nonlinearity, description="Nonlinearity gates" + ) + num_basis: int = Field( + default=8, description="Number of basis functions used in the radial basis" + ) + besselbasis_trainable: bool = Field( + default=True, + description="If true, train the bessel weights", + alias="BesselBasis_trainable", + ) + polynomialcutoff_p: int = Field( + default=5, + description="p-exponent used in polynomial cutoff function, " + "smaller p corresponds to stronger decay with distance", + alias="PolynomialCutoff_p", + ) + + invariant_layers: int = Field( + default=2, description="Number of radial layers, smaller is faster" + ) invariant_neurons: int = Field( - default=64, description="Number of invariant neurons" + default=64, + description="Number of hidden neurons in radial function, smaller is faster", ) - batch_size: int = Field(default=5, description="Batch size") + avg_num_neighbors: None | Literal["auto"] = Field( + default="auto", + description="Number of neighbors to divide by, " + "None => no normalization, " + "auto computes it based on dataset", + ) + use_sc: bool = Field( + default=True, + description="Use self-connection or not, usually gives big improvement", + ) + dataset: Literal["ase"] = Field( + default="ase", + description="Type of data set, can be npz or ase." + "Note that autoplex only supports ase at this point", + ) + validation_dataset: Literal["ase"] = Field( + default="ase", + description="Type of validation data set, can be npz or ase." + "Note that autoplex only supports ase at this point", + ) + dataset_file_name: str = Field( + default="./train_nequip.extxyz", description="Name of the dataset file" + ) + validation_dataset_file_name: str = Field( + default="./test.extxyz", description="Name of the validation dataset file" + ) + ase_args: dict = Field( + default={"format": "extxyz"}, description="Any arguments needed by ase.io.read" + ) + dataset_key_mapping: dict = Field( + default={"forces": "forces", "energy": "total_energy"}, + description="Mapping of keys in the dataset to the expected keys", + ) + validation_dataset_key_mapping: dict = Field( + default={"forces": "forces", "energy": "total_energy"}, + description="Mapping of keys in the validation dataset to the expected keys", + ) + chemical_symbols: list[str] = Field( + default=[], description="List of chemical symbols" + ) + wandb: bool = Field(default=False, description="Use wandb for logging") + verbose: Literal["debug", "info", "warning", "error", "critical"] = Field( + default="info", description="Verbosity level" + ) + log_batch_freq: int = Field( + default=10, + description="Batch frequency, how often to print training errors within the same epoch", + ) + log_epoch_freq: int = Field( + default=1, description="Epoch frequency, how often to print training errors" + ) + save_checkpoint_freq: int = Field( + default=-1, + description="Frequency to save the intermediate checkpoint. " + "No saving of intermediate checkpoints when the value is not positive.", + ) + save_ema_checkpoint_freq: int = Field( + default=-1, + description="Frequency to save the intermediate EMA checkpoint. " + "No saving of intermediate EMA checkpoints when the value is not positive.", + ) + n_train: int = Field(default=1000, description="Number of training samples") + n_val: int = Field(default=1000, description="Number of validation samples") learning_rate: float = Field(default=0.005, description="Learning rate") + batch_size: int = Field(default=5, description="Batch size") + validation_batch_size: int = Field(default=10, description="Validation batch size") max_epochs: int = Field(default=10000, description="Maximum number of epochs") - default_dtype: str = Field(default="float32", description="Default data type") + shuffle: bool = Field(default=True, description="Shuffle the dataset") + metrics_key: str = Field( + default="validation_loss", + description="Metrics used for scheduling and saving best model", + ) + use_ema: bool = Field( + default=True, + description="Use exponential moving average on weights for val/test", + ) + ema_decay: float = Field( + default=0.99, description="Exponential moving average decay" + ) + ema_use_num_updates: bool = Field( + default=True, description="Use number of updates for EMA decay" + ) + report_init_validation: bool = Field( + default=True, + description="Report the validation error for just initialized model", + ) + early_stopping_patiences: dict = Field( + default={"validation_loss": 50}, + description="Stop early if a metric value stopped decreasing for n epochs", + ) + early_stopping_lower_bounds: dict = Field( + default={"LR": 1.0e-5}, + description="Stop early if a metric value is lower than the given value", + ) + loss_coeffs: LossCoeff = Field( + default_factory=LossCoeff, description="Loss coefficients" + ) + metrics_components: list = Field( + default_factory=lambda: [ + ["forces", "mae"], + ["forces", "rmse"], + ["forces", "mae", {"PerSpecies": True, "report_per_component": False}], + ["forces", "rmse", {"PerSpecies": True, "report_per_component": False}], + ["total_energy", "mae"], + ["total_energy", "mae", {"PerAtom": True}], + ], + description="Metrics components", + ) + optimizer_name: str = Field(default="Adam", description="Optimizer name") + optimizer_amsgrad: bool = Field( + default=True, description="Use AMSGrad variant of Adam" + ) + lr_scheduler_name: str = Field( + default="ReduceLROnPlateau", description="Learning rate scheduler name" + ) + lr_scheduler_patience: int = Field( + default=100, description="Patience for learning rate scheduler" + ) + lr_scheduler_factor: float = Field( + default=0.5, description="Factor for learning rate scheduler" + ) + per_species_rescale_shifts_trainable: bool = Field( + default=False, + description="Whether the shifts are trainable. Defaults to False.", + ) + per_species_rescale_scales_trainable: bool = Field( + default=False, + description="Whether the scales are trainable. Defaults to False.", + ) + per_species_rescale_shifts: ( + float + | list[float] + | Literal[ + "dataset_per_atom_total_energy_mean", + "dataset_per_species_total_energy_mean", + ] + ) = Field( + default="dataset_per_atom_total_energy_mean", + description="The value can be a constant float value, an array for each species, or a string. " + "If float values are prpvided , they must be in the same energy units as the training data", + ) + per_species_rescale_scales: ( + float + | list[float] + | Literal[ + "dataset_forces_absmax", + "dataset_per_atom_total_energy_std", + "dataset_per_species_total_energy_std", + "dataset_per_species_forces_rms", + ] + ) = Field( + default="dataset_forces_rms", + description="The value can be a constant float value, an array for each species, or a string. " + "If float values are prpvided , they must be in the same energy units as the training data", + ) -class M3GNETSettings(UpdateBaseModel): +class M3GNETSettings(AutoplexBaseModel): """Model describing the hyperparameters for the M3GNET fits.""" exp_name: str = Field(default="training", description="Name of the experiment") @@ -242,7 +476,7 @@ class M3GNETSettings(UpdateBaseModel): ) -class MACESettings(UpdateBaseModel): +class MACESettings(AutoplexBaseModel): """Model describing the hyperparameters for the MACE fits.""" model: Literal[ @@ -360,7 +594,7 @@ class MACESettings(UpdateBaseModel): wandb: bool = Field(default=False, description="Use Weights and Biases for logging") -class NEPSettings(UpdateBaseModel): +class NEPSettings(AutoplexBaseModel): """Model describing the hyperparameters for the NEP fits.""" version: int = Field(default=4, description="Version of the NEP model") @@ -432,7 +666,7 @@ class NEPSettings(UpdateBaseModel): ) -class MLIPHypers(UpdateBaseModel): +class MLIPHypers(AutoplexBaseModel): """Model containing the hyperparameter defaults for supported MLIPs in autoplex.""" GAP: GAPSettings = Field( @@ -460,7 +694,7 @@ class MLIPHypers(UpdateBaseModel): # RSS Configuration -class ResumeFromPreviousState(UpdateBaseModel): +class ResumeFromPreviousState(AutoplexBaseModel): """ A model describing the state information. @@ -486,7 +720,7 @@ class ResumeFromPreviousState(UpdateBaseModel): ) -class SoapParas(UpdateBaseModel): +class SoapParas(AutoplexBaseModel): """A model describing the SOAP parameters.""" l_max: int = Field(default=12, description="Maximum degree of spherical harmonics") @@ -507,7 +741,7 @@ class SoapParas(UpdateBaseModel): ) -class BcurParams(UpdateBaseModel): +class BcurParams(AutoplexBaseModel): """A model describing the parameters for the BCUR method.""" soap_paras: SoapParas = Field(default_factory=SoapParas) @@ -519,7 +753,7 @@ class BcurParams(UpdateBaseModel): ) -class BuildcellOptions(UpdateBaseModel): +class BuildcellOptions(AutoplexBaseModel): """A model describing the parameters for buildcell.""" NFORM: str | None = Field(default=None, description="The number of formula units") @@ -529,7 +763,7 @@ class BuildcellOptions(UpdateBaseModel): MINSEP: str | None = Field(default=None, description="The minimum separation") -class CustomIncar(UpdateBaseModel): +class CustomIncar(AutoplexBaseModel): """A model describing the INCAR parameters.""" ISMEAR: int = 0 @@ -553,7 +787,7 @@ class CustomIncar(UpdateBaseModel): LPLANE: str = ".FALSE." -class Twob(UpdateBaseModel): +class Twob(AutoplexBaseModel): """A model describing the two-body GAP parameters.""" cutoff: float = Field(default=5.0, description="Radial cutoff distance") @@ -563,13 +797,13 @@ class Twob(UpdateBaseModel): ) -class Threeb(UpdateBaseModel): +class Threeb(AutoplexBaseModel): """A model describing the three-body GAP parameters.""" cutoff: float = Field(default=3.0, description="Radial cutoff distance") -class Soap(UpdateBaseModel): +class Soap(AutoplexBaseModel): """A model describing the SOAP GAP parameters.""" l_max: int = Field(default=10, description="Maximum degree of spherical harmonics") @@ -581,7 +815,7 @@ class Soap(UpdateBaseModel): cutoff: float = Field(default=5.0, description="Radial cutoff distance") -class General(UpdateBaseModel): +class General(AutoplexBaseModel): """A model describing the general GAP parameters.""" three_body: bool = Field( @@ -589,7 +823,7 @@ class General(UpdateBaseModel): ) -class RssConfig(UpdateBaseModel): +class RssConfig(AutoplexBaseModel): """A model describing the complete RSS configuration.""" tag: str | None = Field( From 72758de6c595fe4f6b065c4dbbbda2c4a7b7c513 Mon Sep 17 00:00:00 2001 From: naik-aakash Date: Fri, 17 Jan 2025 15:13:35 +0100 Subject: [PATCH 33/74] update NEQUIP hypers --- src/autoplex/settings.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/autoplex/settings.py b/src/autoplex/settings.py index 2ea942c4d..ae98da32a 100644 --- a/src/autoplex/settings.py +++ b/src/autoplex/settings.py @@ -443,7 +443,7 @@ class NEQUIPSettings(AutoplexBaseModel): "dataset_per_species_forces_rms", ] ) = Field( - default="dataset_forces_rms", + default="dataset_per_species_forces_rms", description="The value can be a constant float value, an array for each species, or a string. " "If float values are prpvided , they must be in the same energy units as the training data", ) From 618eb563c23badd0d1b4b72f632ba7146cac3420 Mon Sep 17 00:00:00 2001 From: naik-aakash Date: Fri, 17 Jan 2025 15:16:34 +0100 Subject: [PATCH 34/74] add hyperparameter arg --- src/autoplex/fitting/common/flows.py | 2 + src/autoplex/fitting/common/jobs.py | 9 ++ src/autoplex/fitting/common/utils.py | 159 +++++---------------------- 3 files changed, 36 insertions(+), 134 deletions(-) diff --git a/src/autoplex/fitting/common/flows.py b/src/autoplex/fitting/common/flows.py index 356bbc04a..c98ed3231 100644 --- a/src/autoplex/fitting/common/flows.py +++ b/src/autoplex/fitting/common/flows.py @@ -175,6 +175,7 @@ def make( glue_file_path=self.glue_file_path, mlip_type=self.mlip_type, hyperpara_opt=self.hyperpara_opt, + hyperparameters=hyperparameters, ref_energy_name=self.ref_energy_name, ref_force_name=self.ref_force_name, ref_virial_name=self.ref_virial_name, @@ -208,6 +209,7 @@ def make( glue_file_path=self.glue_file_path, mlip_type=self.mlip_type, hyperpara_opt=self.hyperpara_opt, + hyperparameters=hyperparameters, ref_energy_name=self.ref_energy_name, ref_force_name=self.ref_force_name, ref_virial_name=self.ref_virial_name, diff --git a/src/autoplex/fitting/common/jobs.py b/src/autoplex/fitting/common/jobs.py index 4e5cfa624..54ccb6f3e 100644 --- a/src/autoplex/fitting/common/jobs.py +++ b/src/autoplex/fitting/common/jobs.py @@ -5,6 +5,7 @@ import numpy as np from jobflow import job +from autoplex import MLIP_HYPERS from autoplex.fitting.common.utils import ( check_convergence, gap_fitting, @@ -32,6 +33,7 @@ def machine_learning_fit( device: str = "cuda", database_dict: dict | None = None, hyperpara_opt: bool = False, + hyperparameters: MLIP_HYPERS = MLIP_HYPERS, **fit_kwargs, ): """ @@ -71,6 +73,8 @@ def machine_learning_fit( hyperpara_opt: bool Perform hyperparameter optimization using XPOT (XPOT: https://pubs.aip.org/aip/jcp/article/159/2/024803/2901815) + hyperparameters: MLIP_HYPERS + Hyperparameters for MLIP fitting. run_fits_on_different_cluster: bool Indicates if fits are to be run on a different cluster. If True, the fitting data (train.extxyz, test.extxyz) is stored in the database. @@ -121,6 +125,7 @@ def machine_learning_fit( ).exists(): train_test_error = gap_fitting( db_dir=database_dir, + hyperparameters=hyperparameters.GAP, species_list=species_list, num_processes_fit=num_processes_fit, auto_delta=auto_delta, @@ -138,6 +143,7 @@ def machine_learning_fit( elif mlip_type == "J-ACE": train_test_error = jace_fitting( db_dir=database_dir, + hyperparameters=hyperparameters.J_ACE, isolated_atom_energies=isolated_atom_energies, ref_energy_name=ref_energy_name, ref_force_name=ref_force_name, @@ -150,6 +156,7 @@ def machine_learning_fit( elif mlip_type == "NEQUIP": train_test_error = nequip_fitting( db_dir=database_dir, + hyperparameters=hyperparameters.NEQUIP, isolated_atom_energies=isolated_atom_energies, ref_energy_name=ref_energy_name, ref_force_name=ref_force_name, @@ -162,6 +169,7 @@ def machine_learning_fit( elif mlip_type == "M3GNET": train_test_error = m3gnet_fitting( db_dir=database_dir, + hyperparameters=hyperparameters.M3GNET, ref_energy_name=ref_energy_name, ref_force_name=ref_force_name, ref_virial_name=ref_virial_name, @@ -173,6 +181,7 @@ def machine_learning_fit( elif mlip_type == "MACE": train_test_error = mace_fitting( db_dir=database_dir, + hyperparameters=hyperparameters.MACE, ref_energy_name=ref_energy_name, ref_force_name=ref_force_name, ref_virial_name=ref_virial_name, diff --git a/src/autoplex/fitting/common/utils.py b/src/autoplex/fitting/common/utils.py index 084cb3d50..d4828db64 100644 --- a/src/autoplex/fitting/common/utils.py +++ b/src/autoplex/fitting/common/utils.py @@ -32,6 +32,7 @@ from matgl.models import M3GNet from matgl.utils.training import PotentialLightningModule from monty.dev import requires +from monty.serialization import dumpfn from nequip.ase import NequIPCalculator from numpy import ndarray from pymatgen.core import Structure @@ -40,7 +41,7 @@ from scipy.spatial import ConvexHull from scipy.special import comb -from autoplex import MLIP_HYPERS +from autoplex import GAP_HYPERS, JACE_HYPERS, M3GNET_HYPERS, MACE_HYPERS, NEQUIP_HYPERS from autoplex.data.common.utils import ( data_distillation, plot_energy_forces, @@ -59,7 +60,7 @@ def gap_fitting( db_dir: Path, species_list: list | None = None, - hyperparameters: MLIP_HYPERS.GAP = MLIP_HYPERS.GAP, + hyperparameters: GAP_HYPERS = GAP_HYPERS, num_processes_fit: int = 32, auto_delta: bool = True, glue_xml: bool = False, @@ -270,7 +271,7 @@ def gap_fitting( ) def jace_fitting( db_dir: str | Path, - hyperparameters: MLIP_HYPERS.J_ACE = MLIP_HYPERS.J_ACE, + hyperparameters: JACE_HYPERS = JACE_HYPERS, isolated_atom_energies: dict | None = None, ref_energy_name: str = "REF_energy", ref_force_name: str = "REF_forces", @@ -443,7 +444,7 @@ def jace_fitting( def nequip_fitting( db_dir: Path, - hyperparameters: MLIP_HYPERS.NEQUIP = MLIP_HYPERS.NEQUIP, + hyperparameters: NEQUIP_HYPERS = NEQUIP_HYPERS, isolated_atom_energies: dict | None = None, ref_energy_name: str = "REF_energy", ref_force_name: str = "REF_forces", @@ -535,139 +536,29 @@ def nequip_fitting( else: raise ValueError("isolated_atom_energies is empty or not defined!") + nequip_config_updates = { + "dataset_key_mapping": { + f"{ref_energy_name}": "total_energy", + f"{ref_force_name}": "forces", + }, + "validation_dataset_key_mapping": { + f"{ref_energy_name}": "total_energy", + f"{ref_force_name}": "forces", + }, + "chemical_symbols": ele_syms, + "dataset_file_name": "./train_nequip.extxyz", + "validation_dataset_file_name": f"{db_dir}/test.extxyz", + "n_train": num_of_train, + "n_val": num_of_val, + } + hyperparameters.update_parameters(nequip_config_updates) + if fit_kwargs: hyperparameters.update_parameters(fit_kwargs) nequip_hypers = hyperparameters.model_dump(by_alias=True) - r_max = nequip_hypers["r_max"] - num_layers = nequip_hypers["num_layers"] - l_max = nequip_hypers["l_max"] - num_features = nequip_hypers["num_features"] - num_basis = nequip_hypers["num_basis"] - invariant_layers = nequip_hypers["invariant_layers"] - invariant_neurons = nequip_hypers["invariant_neurons"] - batch_size = nequip_hypers["batch_size"] - learning_rate = nequip_hypers["learning_rate"] - max_epochs = nequip_hypers["max_epochs"] - default_dtype = nequip_hypers["default_dtype"] - - nequip_text = f"""root: results -run_name: autoplex -seed: 123 -dataset_seed: 456 -append: true -default_dtype: {default_dtype} - -# network -r_max: {r_max} -num_layers: {num_layers} -l_max: {l_max} -parity: true -num_features: {num_features} -nonlinearity_type: gate - -nonlinearity_scalars: - e: silu - o: tanh - -nonlinearity_gates: - e: silu - o: tanh - -num_basis: {num_basis} -BesselBasis_trainable: true -PolynomialCutoff_p: 6 - -invariant_layers: {invariant_layers} -invariant_neurons: {invariant_neurons} -avg_num_neighbors: auto - -use_sc: true -dataset: ase -validation_dataset: ase -dataset_file_name: ./train_nequip.extxyz -validation_dataset_file_name: {db_dir}/test.extxyz - -ase_args: - format: extxyz -dataset_key_mapping: - {ref_energy_name}: total_energy - {ref_force_name}: forces -validation_dataset_key_mapping: - {ref_energy_name}: total_energy - {ref_force_name}: forces - -chemical_symbols: -{isolated_atom_energies_update} -wandb: False - -verbose: info -log_batch_freq: 10 -log_epoch_freq: 1 -save_checkpoint_freq: -1 -save_ema_checkpoint_freq: -1 - -n_train: {num_of_train} -n_val: {num_of_val} -learning_rate: {learning_rate} -batch_size: {batch_size} -validation_batch_size: 10 -max_epochs: {max_epochs} -shuffle: true -metrics_key: validation_loss -use_ema: true -ema_decay: 0.99 -ema_use_num_updates: true -report_init_validation: true - -early_stopping_patiences: - validation_loss: 50 - -early_stopping_lower_bounds: - LR: 1.0e-5 - -loss_coeffs: - forces: 1 - total_energy: - - 1 - - PerAtomMSELoss - -metrics_components: - - - forces - - mae - - - forces - - rmse - - - forces - - mae - - PerSpecies: True - report_per_component: False - - - forces - - rmse - - PerSpecies: True - report_per_component: False - - - total_energy - - mae - - - total_energy - - mae - - PerAtom: True - -optimizer_name: Adam -optimizer_amsgrad: true - -lr_scheduler_name: ReduceLROnPlateau -lr_scheduler_patience: 100 -lr_scheduler_factor: 0.5 - -per_species_rescale_shifts_trainable: false -per_species_rescale_scales_trainable: false - -per_species_rescale_shifts: dataset_per_atom_total_energy_mean -per_species_rescale_scales: dataset_forces_rms - """ - - with open("nequip.yaml", "w") as file: - file.write(nequip_text) + dumpfn(nequip_hypers, "nequip.yaml") run_nequip("nequip-train nequip.yaml", "nequip_train") run_nequip( @@ -713,7 +604,7 @@ def nequip_fitting( def m3gnet_fitting( db_dir: Path, - hyperparameters: MLIP_HYPERS.M3GNET = MLIP_HYPERS.M3GNET, + hyperparameters: M3GNET_HYPERS = M3GNET_HYPERS, device: str = "cuda", ref_energy_name: str = "REF_energy", ref_force_name: str = "REF_forces", @@ -1076,7 +967,7 @@ def m3gnet_fitting( def mace_fitting( db_dir: Path, - hyperparameters: MLIP_HYPERS.MACE = MLIP_HYPERS.MACE, + hyperparameters: MACE_HYPERS = MACE_HYPERS, device: str = "cuda", ref_energy_name: str = "REF_energy", ref_force_name: str = "REF_forces", From 37e83b5a64b9cd9bae2cadaf975941ae10eef8df Mon Sep 17 00:00:00 2001 From: naik-aakash Date: Fri, 17 Jan 2025 15:17:40 +0100 Subject: [PATCH 35/74] minor change in tests --- tests/auto/rss/test_flows.py | 8 ++------ tests/fitting/common/test_jobs.py | 14 +++++++------- 2 files changed, 9 insertions(+), 13 deletions(-) diff --git a/tests/auto/rss/test_flows.py b/tests/auto/rss/test_flows.py index 3fbb55062..15c70d418 100644 --- a/tests/auto/rss/test_flows.py +++ b/tests/auto/rss/test_flows.py @@ -311,13 +311,9 @@ def test_mock_workflow_multi_node(test_dir, mock_vasp, memory_jobstore, clean_di assert len(selected_atoms) == 3 -def test_rssmaker_custom_config(test_dir): +def test_rssmaker_custom_config_file(test_dir): - from monty.serialization import loadfn - - rss_config = loadfn(test_dir / "rss" / "rss_config.yaml") - - config_model = RssConfig(**rss_config) + config_model = RssConfig.from_file(test_dir / "rss" / "rss_config.yaml") # Test if config is updated as expected rss = RssMaker(config=config_model) diff --git a/tests/fitting/common/test_jobs.py b/tests/fitting/common/test_jobs.py index 0f7532bcb..3ce26f7a7 100644 --- a/tests/fitting/common/test_jobs.py +++ b/tests/fitting/common/test_jobs.py @@ -17,7 +17,7 @@ def test_gap_fit_maker(test_dir, memory_jobstore, clean_dir): database_dir=database_dir ) - responses = run_locally( + _ = run_locally( gapfit, ensure_success=True, create_folders=True, store=memory_jobstore ) @@ -39,7 +39,7 @@ def test_jace_fit_maker(test_dir, memory_jobstore, clean_dir): totaldegree=4, ) - responses = run_locally( + _ = run_locally( jacefit, ensure_success=True, create_folders=True, store=memory_jobstore ) @@ -61,7 +61,7 @@ def test_nequip_fit_maker(test_dir, memory_jobstore, clean_dir): device="cpu", ) - responses = run_locally( + _ = run_locally( nequipfit, ensure_success=True, create_folders=True, store=memory_jobstore ) @@ -91,7 +91,7 @@ def test_m3gnet_fit_maker(test_dir, memory_jobstore, clean_dir): test_equal_to_val=True, ) - responses = run_locally( + _ = run_locally( nequipfit, ensure_success=True, create_folders=True, store=memory_jobstore ) @@ -122,7 +122,7 @@ def test_mace_fit_maker(test_dir, memory_jobstore, clean_dir): device="cpu", ) - responses = run_locally( + _ = run_locally( macefit, ensure_success=True, create_folders=True, store=memory_jobstore ) @@ -166,7 +166,7 @@ def test_mace_finetuning_maker(test_dir, memory_jobstore, clean_dir): seed = 3, ) - responses = run_locally( + _ = run_locally( macefit, ensure_success=True, create_folders=True, store=memory_jobstore ) @@ -210,7 +210,7 @@ def test_mace_finetuning_maker2(test_dir, memory_jobstore, clean_dir): seed = 3, ) - responses = run_locally( + _ = run_locally( macefit, ensure_success=True, create_folders=True, store=memory_jobstore ) From 834d553025493dab8f12871d672b5c4a1a917fe5 Mon Sep 17 00:00:00 2001 From: naik-aakash Date: Fri, 17 Jan 2025 16:43:01 +0100 Subject: [PATCH 36/74] extend m3gnet hypers and make finetuning m3gnet possible --- src/autoplex/fitting/common/utils.py | 78 +++++++++++++++------------- src/autoplex/settings.py | 27 +++++++--- tests/auto/phonons/test_flows.py | 5 +- tests/fitting/common/test_jobs.py | 10 ++-- 4 files changed, 72 insertions(+), 48 deletions(-) diff --git a/src/autoplex/fitting/common/utils.py b/src/autoplex/fitting/common/utils.py index d4828db64..2a49f8dc0 100644 --- a/src/autoplex/fitting/common/utils.py +++ b/src/autoplex/fitting/common/utils.py @@ -14,6 +14,7 @@ import ase import lightning as pl +import matgl import matplotlib.pyplot as plt import numpy as np import pandas as pd @@ -609,6 +610,7 @@ def m3gnet_fitting( ref_energy_name: str = "REF_energy", ref_force_name: str = "REF_forces", ref_virial_name: str = "REF_virial", + test_equal_to_val: bool = True, fit_kwargs: dict | None = None, ) -> dict: """ @@ -628,9 +630,10 @@ def m3gnet_fitting( Reference force name. ref_virial_name : str, optional Reference virial name. + test_equal_to_val: bool + If True, the testing dataset will be the same as the validation dataset. fit_kwargs: dict. - optional dictionary with parameters for m3gnet fitting with keys same as - mlip-rss-defaults.json. + optional dictionary with parameters for M3GNET fitting. Keyword Arguments ----------------- @@ -656,8 +659,6 @@ def m3gnet_fitting( Maximum degree of spherical harmonics. max_n: int Maximum radial function degree. - test_equal_to_val: bool - If True, the testing dataset will be the same as the validation dataset. Returns ------- @@ -683,16 +684,6 @@ def m3gnet_fitting( exp_name = m3gnet_hypers["exp_name"] results_dir = m3gnet_hypers["results_dir"] - cutoff = m3gnet_hypers["cutoff"] - threebody_cutoff = m3gnet_hypers["threebody_cutoff"] - batch_size = m3gnet_hypers["batch_size"] - max_epochs = m3gnet_hypers["max_epochs"] - include_stresses = m3gnet_hypers["include_stresses"] - hidden_dim = m3gnet_hypers["hidden_dim"] - num_units = m3gnet_hypers["num_units"] - max_l = m3gnet_hypers["max_l"] - max_n = m3gnet_hypers["max_n"] - test_equal_to_val = m3gnet_hypers["test_equal_to_val"] os.makedirs(os.path.join(results_dir, exp_name), exist_ok=True) @@ -732,7 +723,7 @@ def m3gnet_fitting( ) = convert_xyz_to_structure( train_m3gnet, include_forces=True, - include_stresses=include_stresses, + include_stresses=m3gnet_hypers.get("include_stresses"), ref_energy_name=ref_energy_name, ref_force_name=ref_force_name, ref_virial_name=ref_virial_name, @@ -749,10 +740,10 @@ def m3gnet_fitting( train_element_types ) # this print has to stay as the stdout is written to the file train_converter = Structure2Graph( - element_types=train_element_types, cutoff=cutoff + element_types=train_element_types, cutoff=m3gnet_hypers.get("cutoff") ) train_datasets = MGLDataset( - threebody_cutoff=threebody_cutoff, + threebody_cutoff=m3gnet_hypers.get("threebody_cutoff"), structures=train_structs, converter=train_converter, labels=train_labels, @@ -776,7 +767,7 @@ def m3gnet_fitting( ) = convert_xyz_to_structure( test_data, include_forces=True, - include_stresses=include_stresses, + include_stresses=m3gnet_hypers.get("include_stresses"), ref_energy_name=ref_energy_name, ref_force_name=ref_force_name, ref_virial_name=ref_virial_name, @@ -789,10 +780,10 @@ def m3gnet_fitting( } test_element_types = get_element_list(test_structs) test_converter = Structure2Graph( - element_types=test_element_types, cutoff=cutoff + element_types=test_element_types, cutoff=m3gnet_hypers.get("cutoff") ) test_dataset = MGLDataset( - threebody_cutoff=threebody_cutoff, + threebody_cutoff=m3gnet_hypers.get("threebody_cutoff"), structures=test_structs, converter=test_converter, labels=test_labels, @@ -832,21 +823,38 @@ def m3gnet_fitting( val_data=val_dataset, test_data=test_dataset, collate_fn=my_collate_fn, - batch_size=batch_size, + batch_size=m3gnet_hypers.get("batch_size"), num_workers=1, ) - model = M3GNet( - element_types=train_element_types, - is_intensive=False, - cutoff=cutoff, - threebody_cutoff=threebody_cutoff, - dim_node_embedding=hidden_dim, - dim_edge_embedding=hidden_dim, - units=num_units, - max_l=max_l, - max_n=max_n, - ) - lit_module = PotentialLightningModule(model=model, include_line_graph=True) + # train from scratch + if not m3gnet_hypers["pretrained_model"]: # train from scratch + model = M3GNet( + element_types=train_element_types, + is_intensive=m3gnet_hypers.get("is_intensive"), + cutoff=m3gnet_hypers.get("cutoff"), + threebody_cutoff=m3gnet_hypers.get("threebody_cutoff"), + dim_node_embedding=m3gnet_hypers.get("dim_node_embedding"), + dim_edge_embedding=m3gnet_hypers.get("dim_edge_embedding"), + units=m3gnet_hypers.get("units"), + max_l=m3gnet_hypers.get("max_l"), + max_n=m3gnet_hypers.get("max_n"), + nblocks=m3gnet_hypers.get("nblocks"), + ) + lit_module = PotentialLightningModule(model=model, include_line_graph=True) + else: # finetune pretrained model + logging.info( + f"Finetuning pretrained model: {m3gnet_hypers['pretrained_model']}" + ) + m3gnet_nnp = matgl.load_model(m3gnet_hypers["pretrained_model"]) + model = m3gnet_nnp.model + property_offset = m3gnet_nnp.element_refs.property_offset + lit_module = PotentialLightningModule( + model=model, + element_refs=property_offset, + lr=1e-4, + include_line_graph=True, + ) + logger = CSVLogger(name=exp_name, save_dir=os.path.join(results_dir, "logs")) # Inference mode = False is required for calculating forces, stress in test mode and prediction mode if device == "cuda": @@ -854,7 +862,7 @@ def m3gnet_fitting( gpu_id = os.environ.get("CUDA_VISIBLE_DEVICES", "0") torch.cuda.set_device(torch.device(f"cuda:{gpu_id}")) trainer = pl.Trainer( - max_epochs=max_epochs, + max_epochs=m3gnet_hypers.get("max_epochs"), accelerator="gpu", logger=logger, inference_mode=False, @@ -863,7 +871,7 @@ def m3gnet_fitting( raise ValueError("CUDA is not available.") else: trainer = pl.Trainer( - max_epochs=max_epochs, + max_epochs=m3gnet_hypers.get("max_epochs"), accelerator="cpu", logger=logger, inference_mode=False, diff --git a/src/autoplex/settings.py b/src/autoplex/settings.py index ae98da32a..906e3dcf2 100644 --- a/src/autoplex/settings.py +++ b/src/autoplex/settings.py @@ -456,24 +456,39 @@ class M3GNETSettings(AutoplexBaseModel): results_dir: str = Field( default="m3gnet_results", description="Directory to save the results" ) - cutoff: float = Field(default=5.0, description="Radial cutoff distance") + pretrained_model: ( + Literal["M3GNet-MP-2021.2.8-PES", "M3GNet-MP-2021.2.8-DIRECT-PES"] | None + ) = Field(default=None, description="Pretrained model") + cutoff: float = Field(default=5.0, description="Cutoff radius of the graph") threebody_cutoff: float = Field( - default=4.0, description="Three-body cutoff distance" + default=4.0, description="Cutoff radius for 3 body interactions" ) batch_size: int = Field(default=10, description="Batch size") max_epochs: int = Field(default=1000, description="Maximum number of epochs") include_stresses: bool = Field( default=True, description="Whether to include stresses" ) - hidden_dim: int = Field(default=128, description="Hidden dimension") - num_units: int = Field(default=128, description="Number of units") + dim_node_embedding: int = Field( + default=128, description="Dimension of node embedding" + ) + dim_edge_embedding: int = Field( + default=128, description="Dimension of edge embedding" + ) + dim_state_embedding: int = Field( + default=0, description="Dimension of state embedding" + ) max_l: int = Field(default=4, description="Maximum degree of spherical harmonics") max_n: int = Field( default=4, description="Maximum number of radial basis functions" ) - test_equal_to_val: bool = Field( - default=True, description="Whether the test set is equal to the validation set" + nblocks: int = Field(default=3, description="Number of blocks") + rbf_type: Literal["Gaussian", "SphericalBessel"] = Field( + default="Gaussian", description="Type of radial basis function" + ) + is_intensive: bool = Field( + default=False, description="Whether the prediction is intensive" ) + units: int = Field(default=128, description="Number of neurons in each MLP layer") class MACESettings(AutoplexBaseModel): diff --git a/tests/auto/phonons/test_flows.py b/tests/auto/phonons/test_flows.py index c0a3be30e..9caeca309 100644 --- a/tests/auto/phonons/test_flows.py +++ b/tests/auto/phonons/test_flows.py @@ -901,8 +901,9 @@ def test_complete_dft_vs_ml_benchmark_workflow_m3gnet( "batch_size": 1, "max_epochs": 3, "include_stresses": True, - "hidden_dim": 8, - "num_units": 8, + "dim_node_embedding": 8, + "dim_edge_embedding": 8, + "units": 8, "max_l": 4, "max_n": 4, "device": "cpu", diff --git a/tests/fitting/common/test_jobs.py b/tests/fitting/common/test_jobs.py index 3ce26f7a7..f4a7cd9e9 100644 --- a/tests/fitting/common/test_jobs.py +++ b/tests/fitting/common/test_jobs.py @@ -71,7 +71,7 @@ def test_nequip_fit_maker(test_dir, memory_jobstore, clean_dir): def test_m3gnet_fit_maker(test_dir, memory_jobstore, clean_dir): database_dir = test_dir / "fitting/rss_training_dataset/" - nequipfit = MLIPFitMaker( + m3gnetfit = MLIPFitMaker( mlip_type="M3GNET", num_processes_fit=1, apply_data_preprocessing=False, @@ -83,8 +83,8 @@ def test_m3gnet_fit_maker(test_dir, memory_jobstore, clean_dir): batch_size=1, max_epochs=3, include_stresses=True, - hidden_dim=8, - num_units=8, + dim_node_embedding=8, + units=8, max_l=4, max_n=4, device="cpu", @@ -92,10 +92,10 @@ def test_m3gnet_fit_maker(test_dir, memory_jobstore, clean_dir): ) _ = run_locally( - nequipfit, ensure_success=True, create_folders=True, store=memory_jobstore + m3gnetfit, ensure_success=True, create_folders=True, store=memory_jobstore ) - assert Path(nequipfit.output["mlip_path"][0].resolve(memory_jobstore)).exists() + assert Path(m3gnetfit.output["mlip_path"][0].resolve(memory_jobstore)).exists() def test_mace_fit_maker(test_dir, memory_jobstore, clean_dir): From 13dfbe8c022fa146cc90f28a2d81157d813b420e Mon Sep 17 00:00:00 2001 From: naik-aakash Date: Fri, 17 Jan 2025 21:08:10 +0100 Subject: [PATCH 37/74] cleanup json/yaml load remanants --- src/autoplex/auto/phonons/flows.py | 15 +++++++-------- src/autoplex/auto/rss/flows.py | 12 ++++-------- src/autoplex/fitting/common/utils.py | 3 --- 3 files changed, 11 insertions(+), 19 deletions(-) diff --git a/src/autoplex/auto/phonons/flows.py b/src/autoplex/auto/phonons/flows.py index 2deeade92..1beb83c54 100644 --- a/src/autoplex/auto/phonons/flows.py +++ b/src/autoplex/auto/phonons/flows.py @@ -3,7 +3,6 @@ import logging import warnings from dataclasses import dataclass, field -from pathlib import Path from atomate2.common.schemas.phonons import PhononBSDOSDoc from atomate2.vasp.flows.mp import ( @@ -34,7 +33,6 @@ from autoplex.data.phonons.flows import IsoAtomStaticMaker, TightDFTStaticMaker from autoplex.data.phonons.jobs import reduce_supercell_size_job from autoplex.fitting.common.flows import MLIPFitMaker -from autoplex.fitting.common.utils import MLIP_PHONON_DEFAULTS_FILE_PATH __all__ = [ "CompleteDFTvsMLBenchmarkWorkflow", @@ -164,8 +162,6 @@ class CompleteDFTvsMLBenchmarkWorkflow(Maker): Settings for supercell generation benchmark_kwargs: dict Keyword arguments for the benchmark flows - path_to_hyperparameters : str or Path. - Path to JSON file containing the MLIP hyperparameters. summary_filename_prefix: str Prefix of the result summary file. glue_xml: bool @@ -220,7 +216,6 @@ class CompleteDFTvsMLBenchmarkWorkflow(Maker): default_factory=lambda: {"min_length": 15, "max_length": 20} ) benchmark_kwargs: dict = field(default_factory=dict) - path_to_hyperparameters: Path | str = MLIP_PHONON_DEFAULTS_FILE_PATH summary_filename_prefix: str = "results_" glue_xml: bool = False glue_file_path: str = "glue.xml" @@ -230,6 +225,7 @@ def make( self, structure_list: list[Structure], mp_ids, + hyperparameters: MLIP_HYPERS = MLIP_HYPERS, dft_references: list[PhononBSDOSDoc] | None = None, benchmark_structures: list[Structure] | None = None, benchmark_mp_ids: list[str] | None = None, @@ -247,6 +243,8 @@ def make( List of pymatgen structures. mp_ids: Materials Project IDs. + hyperparameters: MLIP_HYPERS + Hyperparameters for the MLIP models. dft_references: list[PhononBSDOSDoc] | None List of DFT reference files containing the PhononBSDOCDoc object. Reference files have to refer to a finite displacement of 0.01. @@ -274,9 +272,8 @@ def make( fit_input = {} bm_outputs = [] - hyper_parameters = MLIP_HYPERS.model_copy(deep=True) - default_hyperparameters = hyper_parameters.model_dump(by_alias=True) - # default_hyperparameters = MLIP_HYPERS.model_copy(deep=True).model_dump(by_alias=True) + hyperparameters = hyperparameters.model_copy(deep=True) + default_hyperparameters = hyperparameters.model_dump(by_alias=True) soap_default_dict = next( ( @@ -401,6 +398,7 @@ def make( ).make( species_list=isoatoms.output["species"], isolated_atom_energies=isoatoms.output["energies"], + hyperparameters=hyperparameters, fit_input=fit_input, **fit_kwargs, ) @@ -482,6 +480,7 @@ def make( ).make( species_list=isoatoms.output["species"], isolated_atom_energies=isoatoms.output["energies"], + hyperparameters=hyperparameters, fit_input=fit_input, soap=soap_dict, ) diff --git a/src/autoplex/auto/rss/flows.py b/src/autoplex/auto/rss/flows.py index 291c8ee5f..f3260cd4f 100644 --- a/src/autoplex/auto/rss/flows.py +++ b/src/autoplex/auto/rss/flows.py @@ -1,15 +1,11 @@ """RSS (random structure searching) flow for exploring and learning potential energy surfaces from scratch.""" -import os from dataclasses import dataclass, field -from pathlib import Path from jobflow import Flow, Maker, Response, job +from autoplex import RSS_CONFIG from autoplex.auto.rss.jobs import do_rss_iterations, initial_rss -from autoplex.settings import RssConfig - -MODULE_DIR = Path(os.path.dirname(__file__)) @dataclass @@ -21,13 +17,13 @@ class RssMaker(Maker): ---------- name: str Name of the flow. - config: RssConfig | None + config: RSS_CONFIG Pydantic model that defines the setup parameters for the whole RSS workflow. - If not provided, the defaults from 'autoplex.settings.RssConfig' will be used. + If not explicitly set, the defaults from 'autoplex.settings.RssConfig' will be used. """ name: str = "ml-driven rss" - config: RssConfig | None = field(default_factory=lambda: RssConfig()) + config: RSS_CONFIG = field(default_factory=lambda: RSS_CONFIG) @job def make(self, **kwargs): diff --git a/src/autoplex/fitting/common/utils.py b/src/autoplex/fitting/common/utils.py index 2a49f8dc0..9ac31a24a 100644 --- a/src/autoplex/fitting/common/utils.py +++ b/src/autoplex/fitting/common/utils.py @@ -50,9 +50,6 @@ stratified_dataset_split, ) -current_dir = Path(__file__).absolute().parent -MLIP_PHONON_DEFAULTS_FILE_PATH = current_dir / "mlip-phonon-defaults.json" -MLIP_RSS_DEFAULTS_FILE_PATH = current_dir / "mlip-rss-defaults.json" logging.basicConfig( level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s" ) From e071cef60494da936db1150e888f4c0034f5ca2c Mon Sep 17 00:00:00 2001 From: naik-aakash Date: Fri, 17 Jan 2025 21:08:51 +0100 Subject: [PATCH 38/74] remove unused imports --- tests/auto/rss/test_flows.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/tests/auto/rss/test_flows.py b/tests/auto/rss/test_flows.py index 15c70d418..3e3391ddd 100644 --- a/tests/auto/rss/test_flows.py +++ b/tests/auto/rss/test_flows.py @@ -1,7 +1,6 @@ import os from pathlib import Path -from jobflow import run_locally, Flow, job -from mace.tools.model_script_utils import configure_model +from jobflow import run_locally, Flow from tests.conftest import mock_rss, mock_do_rss_iterations, mock_do_rss_iterations_multi_jobs from autoplex.settings import RssConfig From 09b2f6d821d0258f512f3edbd9d7c17f878a6250 Mon Sep 17 00:00:00 2001 From: naik-aakash Date: Fri, 17 Jan 2025 22:25:49 +0100 Subject: [PATCH 39/74] update docs --- docs/reference/index.rst | 1 + docs/user/phonon/flows/fitting/fitting.md | 23 +++++++++++----- docs/user/rss/flow/input/input.md | 5 ++-- docs/user/rss/flow/quick_start/start.md | 32 +++++++++++++++-------- docs/user/rss/index.md | 2 +- 5 files changed, 43 insertions(+), 20 deletions(-) diff --git a/docs/reference/index.rst b/docs/reference/index.rst index a4e2055e8..572e5f35f 100644 --- a/docs/reference/index.rst +++ b/docs/reference/index.rst @@ -16,3 +16,4 @@ This section gives an overview of the API for autoplex. benchmark data fitting + settings diff --git a/docs/user/phonon/flows/fitting/fitting.md b/docs/user/phonon/flows/fitting/fitting.md index dab14786e..ab5efd8bd 100644 --- a/docs/user/phonon/flows/fitting/fitting.md +++ b/docs/user/phonon/flows/fitting/fitting.md @@ -37,8 +37,20 @@ complete_flow = CompleteDFTvsMLBenchmarkWorkflow( The MLIP model specific settings and hyperparameters setup varies from model to model and is demonstrated in the next sections. Also, [`atomate2`-based MLPhononMaker](https://materialsproject.github.io/atomate2/reference/atomate2.forcefields.jobs.html#module-atomate2.forcefields.jobs) settings can be changed via `benchmark_kwargs` as shown in the code snippet. + +> `autoplex` relies on pydantic models for validating the hyperparameter sets of the supported MLIP architectures. +> Note that all the possible hyperparameters are not yet included. It is upto the user to ensure if any other parameters are supplied +> are in correct format and required datatype. To get an overview of the default hyperparameter sets, +> you can use the following code snippet. + +```python +from autoplex import MLIP_HYPERS + +print(MLIP_HYPERS.model_dump(by_alias=True)) +``` + > ℹ️ Note that `autoplex` provides the most comprehensive features for **GAP**, and more features for the other models will -follow in future versions. +follow in future versions. ## GAP @@ -83,12 +95,12 @@ complete_flow = CompleteDFTvsMLBenchmarkWorkflow( }] ) ``` -`autoplex` provides a JSON dict file containing default GAP fit settings in -`autoplex/fitting/common/mlip-phonon-defaults.json`, +`autoplex` provides a Pydantic model containing default GAP fit settings in +`autoplex.settings.GAPSettings`, that can be overwritten using the fit keyword arguments as demonstrated in the code snippet. `autoplex` follows a certain convention for naming files and labelling the data -(see `autoplex/fitting/common/mlip-phonon-defaults.json`). +(see `autoplex.settings.GAPSettings.GeneralSettings`). ```json "general": { "at_file": "train.extxyz", @@ -222,7 +234,6 @@ Currently, this can only be done by cloning the git-repo and installing it from [https://github.com/ACEsuit/mace/](https://github.com/ACEsuit/mace/). We currently install the main branch from there automatically within autoplex. -It is now important that you switch off the default settings for the fitting procedure (use_defaults_fitting=False). Please be careful with performing very low-data finetuning. Currently, we use a stratified split for splitting the data into train and test data, i.e. there will be at least one data point from the dataset including single displaced cells and one rattled structure. @@ -234,7 +245,7 @@ It can also be used without finetuning option. To finetune optimally, please ada complete_workflow_mace = CompleteDFTvsMLBenchmarkWorkflowMPSettings( ml_models=["MACE"], volume_custom_scale_factors=[0.95,1.00,1.05], rattle_type=0, distort_type=0, - apply_data_preprocessing=True, use_defaults_fitting=False, + apply_data_preprocessing=True, ... ).make( structure_list=[structure], diff --git a/docs/user/rss/flow/input/input.md b/docs/user/rss/flow/input/input.md index 91e78ebe0..474dbf6ff 100644 --- a/docs/user/rss/flow/input/input.md +++ b/docs/user/rss/flow/input/input.md @@ -166,9 +166,10 @@ pre_database_dir: null Regularization is currently only applicable to GAP potentials and is adjusted using the `scheme` parameter. Common schemes include `'linear-hull'` and `'volume-stoichiometry'`. For systems with fixed stoichiometry, `'linear-hull'` is recommended. For systems with varying stoichiometries, `'volume-stoichiometry'` is more appropriate. -## MLIP Parameters +## MLIP parameters -The section defines the settings for training machine learning potentials. Currently supported architectures include GAP, ACE(Julia), NequIP, M3GNet, and MACE. You can specify the desired model using the `mlip_type` argument and tune hyperparameters flexibly by adding key-value pairs. Default and adjustable hyperparameters are available in `autoplex/autoplex/fitting/common/mlip-rss-defaults.json`. +The section defines the settings for training machine learning potentials. Currently supported architectures include GAP, ACE(Julia), NequIP, M3GNet, and MACE. +You can specify the desired model using the `mlip_type` argument and tune hyperparameters flexibly by adding key-value pairs. ```yaml # MLIP Parameters diff --git a/docs/user/rss/flow/quick_start/start.md b/docs/user/rss/flow/quick_start/start.md index 86762ee8a..e3ad02260 100644 --- a/docs/user/rss/flow/quick_start/start.md +++ b/docs/user/rss/flow/quick_start/start.md @@ -24,23 +24,24 @@ The `RssMaker` class in `autoplex` is the core interface for creating ML-RSS pot Parameters can be specified either through a YAML configuration file or as direct arguments in the `make` method. -## Running the workflow with a YAML configuration file +## Running the workflow with a RSSConfig object > **Recommendation**: This is currently our recommended approach for setting up and managing RSS workflows. -The RSS workflow can be initiated using a custom YAML configuration file. -A comprehensive list of parameters, including default settings and modifiable options, is available in `autoplex/auto/rss/rss_default_configuration.yaml`. -When initializing the flow with a YAML file, any specified keys will override the corresponding default values. -To start a new workflow, pass the path to your YAML file as the `config_file` argument in the `RssMaker`. - -> ℹ️ Note that, if you are using `RssMaker` as part of another job or flow and executing jobs on a different system (e.g. using remote submission via jobflow-remote), then the configuration file has to be placed on the remote cluster and the file path has to reflect this as (i.e., it is valid a path on the remote cluster). +The RSSConfig object can be instantiated using a custom YAML configuration file, as illustrated in previous section. +A comprehensive list of parameters, including default settings and modifiable options, is available in `autoplex.settings.RssConfig` pydantic model. +To start a new workflow, create an `RssConfig` object using the YAML file and pass it to the `RSSMaker` class. +When initializing the RssConfig object with a YAML file, any specified keys will override the corresponding default values. ```python +from autoplex.settings import RssConfig from autoplex.auto.rss.flows import RssMaker from fireworks import LaunchPad from jobflow.managers.fireworks import flow_to_workflow -rss_job = RssMaker(name="your workflow name", config_file='path/to/your/config.yaml').make() +config = RssConfig.from_file('path/to/your/config.yaml') + +rss_job = RssMaker(name="your workflow name", config=config).make() wf = flow_to_workflow(rss_job) lpad = LaunchPad.auto_load() lpad.add_wf(wf) @@ -62,7 +63,8 @@ For details on setting up `FireWorks`, see [FireWorks setup](../../../mongodb.md ## Running the workflow with direct parameter specification -As an alternative to using a YAML configuration file, the RSS workflow can be initiated by directly specifying parameters in the `make` method. This approach is ideal for cases where only a few parameters need to be customized. You can override the default settings by passing them as keyword arguments, offering a more flexible and lightweight way to set up the workflow. +As an alternative to using a RssConfig object, the RSS workflow can be initiated by directly specifying parameters in the `make` method. This approach is ideal for cases where only a few parameters need to be customized. +You can override the default settings by passing them as keyword arguments, offering a more flexible and lightweight way to set up the workflow. ```python rss_job = RssMaker(name="your workflow name").make(tag='Si', @@ -81,13 +83,21 @@ If you choose to use the direct parameter specification method, at a minimum, yo > **Recommendation**: We strongly recommend enabling `hookean_repul`, as it applies a strong repulsive force when the distance between two atoms falls below a certain threshold. This ensures that the generated structures are physically reasonable. -> **Note**: If both a YAML file and direct parameter specifications are provided, any overlapping parameters will be overridden by the directly specified values. +> **Note**: If both a custom RssConfig object and direct parameter specifications are provided, any overlapping parameters will be overridden by the directly specified values. ## Building RSS models with various ML potentials -Currently, `RssMaker` supports GAP (Gaussian Approximation Potential), ACE (Atomic Cluster Expansion), and three graph-based network models including NequIP, M3GNet, and MACE. You can specify the desired model using the `mlip_type` argument and adjust relevant hyperparameters within the `make` method. Default and adjustable hyperparameters are available in `autoplex/fitting/common/mlip-rss-defaults.json`. +Currently, `RssMaker` supports GAP (Gaussian Approximation Potential), ACE (Atomic Cluster Expansion), and three graph-based network models including NequIP, M3GNet, and MACE. +You can specify the desired model using the `mlip_type` argument and adjust relevant hyperparameters within the `make` method. +Overview of default and adjustable hyperparameters for each model can be accessed using `MLIP_HYPERS` pydantic model of autoplex. ```python +from autoplex import MLIP_HYPERS +from autoplex.auto.rss.flows import RssMaker + +print(MLIP_HYPERS.MACE) # Eg:- access MACE hyperparameters + +# Intialize the workflow with the desired MLIP model rss_job = RssMaker(name="your workflow name").make(tag='SiO2', ... # Other parameters here mlip_type='MACE', diff --git a/docs/user/rss/index.md b/docs/user/rss/index.md index 255b307c3..6b013c5ab 100644 --- a/docs/user/rss/index.md +++ b/docs/user/rss/index.md @@ -7,7 +7,7 @@ This section provides tutorials for the random structure searching (RSS) workflo ```{toctree} :maxdepth: 3 flow/introduction/intro -flow/quick_start/start flow/input/input +flow/quick_start/start flow/example/example ``` \ No newline at end of file From 07e686b976dbc5535d5e26d8dae6deb5880bcacb Mon Sep 17 00:00:00 2001 From: naik-aakash Date: Fri, 17 Jan 2025 22:39:17 +0100 Subject: [PATCH 40/74] add note to AIRSS install instructions --- README.md | 2 ++ docs/user/setup.md | 3 +++ 2 files changed, 5 insertions(+) diff --git a/README.md b/README.md index bd93de9c5..3dc4bb923 100644 --- a/README.md +++ b/README.md @@ -95,6 +95,8 @@ julia -e 'using Pkg; Pkg.Registry.add("General"); Pkg.Registry.add(Pkg.Registry. Additionally, `buildcell` as a part of `AIRSS` needs to be installed if one wants to use the RSS functionality: +> ℹ️ To be able to build the AIRSS utilities one needs gcc and gfortran version 5 and above. Other compiler families (such as ifort) are not supported. + ```bash curl -O https://www.mtg.msm.cam.ac.uk/files/airss-0.9.3.tgz; tar -xf airss-0.9.3.tgz; rm airss-0.9.3.tgz; cd airss; make ; make install ; make neat; cd .. ``` diff --git a/docs/user/setup.md b/docs/user/setup.md index 0796354b4..8867d74eb 100644 --- a/docs/user/setup.md +++ b/docs/user/setup.md @@ -30,6 +30,9 @@ julia -e 'using Pkg; Pkg.Registry.add("General"); Pkg.Registry.add(Pkg.Registry. Additionally, `buildcell` as a part of `AIRSS` needs to be installed if one wants to use the RSS functionality: +> ℹ️ To be able to build the AIRSS utilities one needs gcc and gfortran version 5 and above. Other compiler families (such as ifort) are not supported. + + ```bash curl -O https://www.mtg.msm.cam.ac.uk/files/airss-0.9.3.tgz; tar -xf airss-0.9.3.tgz; rm airss-0.9.3.tgz; cd airss; make ; make install ; make neat; cd .. ``` From 0eabca30bc47ef3d4898da8ccfc7aa5434b2e22a Mon Sep 17 00:00:00 2001 From: naik-aakash Date: Fri, 17 Jan 2025 23:17:04 +0100 Subject: [PATCH 41/74] add example script to root --- configs/create_configs.py | 9 +++ configs/mlip_hypers.json | 105 ++++++++++++++++++++++++++++ configs/rss_config.yaml | 143 ++++++++++++++++++++++++++++++++++++++ 3 files changed, 257 insertions(+) create mode 100644 configs/create_configs.py create mode 100644 configs/mlip_hypers.json create mode 100644 configs/rss_config.yaml diff --git a/configs/create_configs.py b/configs/create_configs.py new file mode 100644 index 000000000..9c8cee876 --- /dev/null +++ b/configs/create_configs.py @@ -0,0 +1,9 @@ +"""Example script to create hyperparameters and rss config objects using json/yaml files.""" + +from autoplex.settings import MLIPHypers, RssConfig + +# create a custom hyperparameters object using json file +custom_hyperparameters = MLIPHypers.from_file("mlip_hypers.json") + +# create a custom rss config object using json file +custom_rss_config = RssConfig.from_file("rss_config.yaml") diff --git a/configs/mlip_hypers.json b/configs/mlip_hypers.json new file mode 100644 index 000000000..576efd2d8 --- /dev/null +++ b/configs/mlip_hypers.json @@ -0,0 +1,105 @@ +{ + "GAP": { + "general": { + "at_file": "train.extxyz", + "default_sigma": "{0.0001 0.05 0.05 0}", + "energy_parameter_name": "REF_energy", + "force_parameter_name": "REF_forces", + "virial_parameter_name": "REF_virial", + "sparse_jitter": 1.0e-8, + "do_copy_at_file": "F", + "openmp_chunk_size": 10000, + "gp_file": "gap_file.xml", + "e0_offset": 0.0, + "two_body": false, + "three_body": false, + "soap": true + }, + "twob": { + "distance_Nb order": 2, + "f0": 0.0, + "add_species": "T", + "cutoff": 5.0, + "n_sparse": 15, + "covariance_type": "ard_se", + "delta": 2.00 , + "theta_uniform": 0.5 , + "sparse_method": "uniform", + "compact_clusters": "T" + }, + "threeb": { + "distance_Nb order": 3, + "f0": 0.0, + "add_species": "T", + "cutoff": 3.25, + "n_sparse": 100, + "covariance_type": "ard_se", + "delta": 2.00 , + "theta_uniform": 1.0 , + "sparse_method": "uniform", + "compact_clusters": "T" + }, + "soap": { + "add_species": "T", + "l_max": 10, + "n_max": 12, + "atom_sigma": 0.5, + "zeta": 4, + "cutoff": 5.0, + "cutoff_transition_width": 1.0, + "central_weight": 1.0, + "n_sparse": 6000, + "delta": 1.00, + "f0": 0.0, + "covariance_type": "dot_product", + "sparse_method": "cur_points" + } + }, + "NEQUIP": { + "r_max": 4.0, + "num_layers": 4, + "l_max": 2, + "num_features": 32, + "num_basis": 8, + "invariant_layers": 2, + "invariant_neurons": 64, + "batch_size": 5, + "learning_rate": 0.005, + "max_epochs": 10000, + "default_dtype": "float32" + }, + "M3GNET": { + "exp_name": "training", + "results_dir": "m3gnet_results", + "cutoff": 5.0, + "threebody_cutoff": 4.0, + "batch_size": 10, + "max_epochs": 1000, + "include_stresses": true, + "hidden_dim": 128, + "num_units": 128, + "max_l": 4, + "max_n": 4, + "test_equal_to_val": true + }, + "MACE": { + "model": "MACE", + "name": "MACE_model", + "config_type_weights": "{'Default':1.0}", + "hidden_irreps": "128x0e + 128x1o", + "r_max": 5.0, + "batch_size": 10, + "max_num_epochs": 1500, + "start_swa": 1200, + "ema_decay": 0.99, + "correlation": 3, + "loss": "huber", + "default_dtype": "float32", + "swa": true, + "ema": true, + "amsgrad": true, + "restart_latest": true, + "seed": 123, + "device": "cpu" + } +} diff --git a/configs/rss_config.yaml b/configs/rss_config.yaml new file mode 100644 index 000000000..e166b0b79 --- /dev/null +++ b/configs/rss_config.yaml @@ -0,0 +1,143 @@ +# General Parameters +tag: "test" +train_from_scratch: true +resume_from_previous_state: + test_error: + pre_database_dir: + mlip_path: + isolated_atom_energies: + +# Buildcell Parameters +generated_struct_numbers: + - 9000 + - 1000 +buildcell_options: +fragment_file: null +fragment_numbers: null +num_processes_buildcell: 64 + +# Sampling Parameters +num_of_initial_selected_structs: + - 80 + - 20 +num_of_rss_selected_structs: 100 +initial_selection_enabled: true +rss_selection_method: 'bcur2i' +bcur_params: + soap_paras: + l_max: 12 + n_max: 12 + atom_sigma: 0.0875 + cutoff: 10.5 + cutoff_transition_width: 1.0 + zeta: 4.0 + average: true + species: true + frac_of_bcur: 0.8 + bolt_max_num: 3000 +random_seed: null + +# DFT Labelling Parameters +include_isolated_atom: true +isolatedatom_box: + - 10.0 + - 10.0 + - 10.0 +e0_spin: false +include_dimer: true +dimer_box: + - 10.0 + - 10.0 + - 10.0 +dimer_range: + - 1.0 + - 5.0 +dimer_num: 21 +custom_incar: + ISMEAR: 0 + SIGMA: 0.05 + PREC: 'Accurate' + ADDGRID: '.TRUE.' + EDIFF: 1e-7 + NELM: 250 + LWAVE: '.FALSE.' + LCHARG: '.FALSE.' + ALGO: 'Normal' + AMIX: null + LREAL: '.FALSE.' + ISYM: 0 + ENCUT: 520.0 + KSPACING: 0.20 + GGA: null + KPAR: 8 + NCORE: 16 + LSCALAPACK: '.FALSE.' + LPLANE: '.FALSE.' +custom_potcar: +vasp_ref_file: 'vasp_ref.extxyz' + +# Data Preprocessing Parameters +config_types: + - 'initial' + - 'traj_early' + - 'traj' +rss_group: + - 'traj' +test_ratio: 0.1 +regularization: true +scheme: 'linear-hull' +reg_minmax: + - [0.1, 1] + - [0.001, 0.1] + - [0.0316, 0.316] + - [0.0632, 0.632] +distillation: false +force_max: null +force_label: null +pre_database_dir: null + +# MLIP Parameters +mlip_type: 'GAP' +ref_energy_name: 'REF_energy' +ref_force_name: 'REF_forces' +ref_virial_name: 'REF_virial' +auto_delta: true +num_processes_fit: 64 +device_for_fitting: 'cpu' +##The following hyperparameters are only applicable to GAP. +##If you want to use other models, please replace the corresponding hyperparameters. +twob: + cutoff: 5.0 + n_sparse: 15 + theta_uniform: 1.0 +threeb: + cutoff: 3.0 +soap: + l_max: 10 + n_max: 10 + atom_sigma: 0.5 + n_sparse: 2500 + cutoff: 5.0 +general: + three_body: false + +# RSS Exploration Parameters +scalar_pressure_method: 'uniform' +scalar_exp_pressure: 1 +scalar_pressure_exponential_width: 0.2 +scalar_pressure_low: 0 +scalar_pressure_high: 25 +max_steps: 300 +force_tol: 0.01 +stress_tol: 0.01 +stop_criterion: 0.01 +max_iteration_number: 25 +num_groups: 6 +initial_kt: 0.3 +current_iter_index: 1 +hookean_repul: false +hookean_paras: +keep_symmetry: false +write_traj: true +num_processes_rss: 128 +device_for_rss: 'cpu' From aa3e10c902b365067b5c0ebe318896885c6c87e1 Mon Sep 17 00:00:00 2001 From: naik-aakash Date: Fri, 17 Jan 2025 23:17:41 +0100 Subject: [PATCH 42/74] delete unused yaml and jsons from src --- .../auto/rss/rss_default_configuration.yaml | 143 ------------------ .../fitting/common/mlip-phonon-defaults.json | 58 ------- .../fitting/common/mlip-rss-defaults.json | 60 -------- 3 files changed, 261 deletions(-) delete mode 100644 src/autoplex/auto/rss/rss_default_configuration.yaml delete mode 100644 src/autoplex/fitting/common/mlip-phonon-defaults.json delete mode 100644 src/autoplex/fitting/common/mlip-rss-defaults.json diff --git a/src/autoplex/auto/rss/rss_default_configuration.yaml b/src/autoplex/auto/rss/rss_default_configuration.yaml deleted file mode 100644 index 89cdc713a..000000000 --- a/src/autoplex/auto/rss/rss_default_configuration.yaml +++ /dev/null @@ -1,143 +0,0 @@ -# General Parameters -tag: -train_from_scratch: true -resume_from_previous_state: - test_error: - pre_database_dir: - mlip_path: - isolated_atom_energies: - -# Buildcell Parameters -generated_struct_numbers: - - 8000 - - 2000 -buildcell_options: -fragment_file: null -fragment_numbers: null -num_processes_buildcell: 128 - -# Sampling Parameters -num_of_initial_selected_structs: - - 80 - - 20 -num_of_rss_selected_structs: 100 -initial_selection_enabled: true -rss_selection_method: 'bcur2i' -bcur_params: - soap_paras: - l_max: 12 - n_max: 12 - atom_sigma: 0.0875 - cutoff: 10.5 - cutoff_transition_width: 1.0 - zeta: 4.0 - average: true - species: true - frac_of_bcur: 0.8 - bolt_max_num: 3000 -random_seed: null - -# DFT Labelling Parameters -include_isolated_atom: true -isolatedatom_box: - - 20.0 - - 20.0 - - 20.0 -e0_spin: false -include_dimer: true -dimer_box: - - 20.0 - - 20.0 - - 20.0 -dimer_range: - - 1.0 - - 5.0 -dimer_num: 21 -custom_incar: - ISMEAR: 0 - SIGMA: 0.05 - PREC: 'Accurate' - ADDGRID: '.TRUE.' - EDIFF: 1e-7 - NELM: 250 - LWAVE: '.FALSE.' - LCHARG: '.FALSE.' - ALGO: 'Normal' - AMIX: null - LREAL: '.FALSE.' - ISYM: 0 - ENCUT: 520.0 - KSPACING: 0.20 - GGA: null - KPAR: 8 - NCORE: 16 - LSCALAPACK: '.FALSE.' - LPLANE: '.FALSE.' -custom_potcar: -vasp_ref_file: 'vasp_ref.extxyz' - -# Data Preprocessing Parameters -config_types: - - 'initial' - - 'traj_early' - - 'traj' -rss_group: - - 'traj' -test_ratio: 0.1 -regularization: true -scheme: 'linear-hull' -reg_minmax: - - [0.1, 1] - - [0.001, 0.1] - - [0.0316, 0.316] - - [0.0632, 0.632] -distillation: false -force_max: null -force_label: null -pre_database_dir: null - -# MLIP Parameters -mlip_type: 'GAP' -ref_energy_name: 'REF_energy' -ref_force_name: 'REF_forces' -ref_virial_name: 'REF_virial' -auto_delta: true -num_processes_fit: 32 -device_for_fitting: 'cpu' -##The following hyperparameters are only applicable to GAP. -##If you want to use other models, please replace the corresponding hyperparameters. -twob: - cutoff: 5.0 - n_sparse: 15 - theta_uniform: 1.0 -threeb: - cutoff: 3.0 -soap: - l_max: 10 - n_max: 10 - atom_sigma: 0.5 - n_sparse: 2500 - cutoff: 5.0 -general: - three_body: false - -# RSS Exploration Parameters -scalar_pressure_method: 'uniform' -scalar_exp_pressure: 1 -scalar_pressure_exponential_width: 0.2 -scalar_pressure_low: 0 -scalar_pressure_high: 25 -max_steps: 300 -force_tol: 0.01 -stress_tol: 0.01 -stop_criterion: 0.01 -max_iteration_number: 25 -num_groups: 6 -initial_kt: 0.3 -current_iter_index: 1 -hookean_repul: false -hookean_paras: -keep_symmetry: false -write_traj: true -num_processes_rss: 128 -device_for_rss: 'cpu' diff --git a/src/autoplex/fitting/common/mlip-phonon-defaults.json b/src/autoplex/fitting/common/mlip-phonon-defaults.json deleted file mode 100644 index 216ef1ba5..000000000 --- a/src/autoplex/fitting/common/mlip-phonon-defaults.json +++ /dev/null @@ -1,58 +0,0 @@ -{ - "GAP": { - "general": { - "at_file": "train.extxyz", - "default_sigma": "{0.0001 0.05 0.05 0}", - "energy_parameter_name": "REF_energy", - "force_parameter_name": "REF_forces", - "virial_parameter_name": "REF_virial", - "sparse_jitter": 1.0e-8, - "do_copy_at_file": "F", - "openmp_chunk_size": 10000, - "gp_file": "gap_file.xml", - "e0_offset": 0.0, - "two_body": false, - "three_body": false, - "soap": true - }, - "twob": { - "distance_Nb order": 2, - "f0": 0.0, - "add_species": "T", - "cutoff": 5.0, - "n_sparse": 15, - "covariance_type": "ard_se", - "delta": 2.00 , - "theta_uniform": 0.5 , - "sparse_method": "uniform", - "compact_clusters": "T" - }, - "threeb": { - "distance_Nb order": 3, - "f0": 0.0, - "add_species": "T", - "cutoff": 3.25, - "n_sparse": 100, - "covariance_type": "ard_se", - "delta": 2.00 , - "theta_uniform": 1.0 , - "sparse_method": "uniform", - "compact_clusters": "T" - }, - "soap": { - "add_species": "T", - "l_max": 10, - "n_max": 12, - "atom_sigma": 0.5, - "zeta": 4, - "cutoff": 5.0, - "cutoff_transition_width": 1.0, - "central_weight": 1.0, - "n_sparse": 6000, - "delta": 1.00, - "f0": 0.0, - "covariance_type": "dot_product", - "sparse_method": "cur_points" - } - } -} diff --git a/src/autoplex/fitting/common/mlip-rss-defaults.json b/src/autoplex/fitting/common/mlip-rss-defaults.json deleted file mode 100644 index a4acb5dd7..000000000 --- a/src/autoplex/fitting/common/mlip-rss-defaults.json +++ /dev/null @@ -1,60 +0,0 @@ -{ - "GAP": { - "two_body": "True", - "three_body": "False", - "soap": "True" - }, - "J-ACE": { - "order": 3, - "totaldegree": 6, - "cutoff": 2.0, - "solver": "BLR" - }, - "NEQUIP": { - "r_max": 4.0, - "num_layers": 4, - "l_max": 2, - "num_features": 32, - "num_basis": 8, - "invariant_layers": 2, - "invariant_neurons": 64, - "batch_size": 5, - "learning_rate": 0.005, - "max_epochs": 10000, - "default_dtype": "float32" - }, - "M3GNET": { - "exp_name": "training", - "results_dir": "m3gnet_results", - "cutoff": 5.0, - "threebody_cutoff": 4.0, - "batch_size": 10, - "max_epochs": 1000, - "include_stresses": true, - "hidden_dim": 128, - "num_units": 128, - "max_l": 4, - "max_n": 4, - "test_equal_to_val": true - }, - "MACE": { - "model": "MACE", - "name": "MACE_model", - "config_type_weights": "{'Default':1.0}", - "hidden_irreps": "128x0e + 128x1o", - "r_max": 5.0, - "batch_size": 10, - "max_num_epochs": 1500, - "start_swa": 1200, - "ema_decay": 0.99, - "correlation": 3, - "loss": "huber", - "default_dtype": "float32", - "swa": true, - "ema": true, - "amsgrad": true, - "restart_latest": true, - "seed": 123, - "device": "cpu" - } -} From ab719204baf78e7790b0639a6f5221745bf0678f Mon Sep 17 00:00:00 2001 From: naik-aakash Date: Fri, 17 Jan 2025 23:18:33 +0100 Subject: [PATCH 43/74] disable package data section --- pyproject.toml | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index f11b027e1..148235756 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -43,9 +43,9 @@ dependencies = [ [tool.setuptools.packages.find] where = ["src"] -[tool.setuptools.package-data] -"autoplex.fitting.common" = ["*.json"] -"autoplex.auto.rss" = ["*.yaml"] +#[tool.setuptools.package-data] +#"autoplex.fitting.common" = ["*.json"] +#"autoplex.auto.rss" = ["*.yaml"] [project.optional-dependencies] docs = [ From 902b982bae3c316543b6c2489086484747171203 Mon Sep 17 00:00:00 2001 From: naik-aakash Date: Sun, 19 Jan 2025 18:38:40 +0100 Subject: [PATCH 44/74] update doc-strings --- src/autoplex/auto/rss/flows.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/src/autoplex/auto/rss/flows.py b/src/autoplex/auto/rss/flows.py index f3260cd4f..163884a65 100644 --- a/src/autoplex/auto/rss/flows.py +++ b/src/autoplex/auto/rss/flows.py @@ -132,13 +132,16 @@ def make(self, **kwargs): Reference file for VASP data. Default is 'vasp_ref.extxyz'. config_types: list[str] Configuration types for the VASP calculations. Default is None. - rss_group: list[str] + rss_group: list[str] | str Group name for RSS to setting up regularization. test_ratio: float The proportion of the test set after splitting the data. The value is allowed to be set to 0; in this case, the testing error would not be meaningful anymore. regularization: bool If True, apply regularization. This only works for GAP to date. Default is False. + retain_existing_sigma: bool + Whether to keep the current sigma values for specific configuration types. + If set to True, existing sigma values for specific configurations will remain unchanged. scheme: str Method to use for regularization. Options are @@ -232,7 +235,7 @@ def make(self, **kwargs): """ default_config = self.config.model_copy(deep=True) updated_config = default_config.update_parameters(kwargs) - config_params = updated_config.model_dump() + config_params = updated_config.model_dump(by_alias=True, exclude_none=True) self._process_hookean_paras(config_params) From a38b46122cb29e5c25fd0517a6c4d2740d061b06 Mon Sep 17 00:00:00 2001 From: naik-aakash Date: Sun, 19 Jan 2025 18:39:24 +0100 Subject: [PATCH 45/74] update rss config models --- src/autoplex/settings.py | 25 +++++++++++++++++++++++-- 1 file changed, 23 insertions(+), 2 deletions(-) diff --git a/src/autoplex/settings.py b/src/autoplex/settings.py index 906e3dcf2..7db1987b9 100644 --- a/src/autoplex/settings.py +++ b/src/autoplex/settings.py @@ -771,9 +771,25 @@ class BcurParams(AutoplexBaseModel): class BuildcellOptions(AutoplexBaseModel): """A model describing the parameters for buildcell.""" + ABFIX: bool = Field(default=False, description="Whether to fix the lattice vectors") NFORM: str | None = Field(default=None, description="The number of formula units") - SYMMOPS: str | None = Field(default=None, description="The symmetry operations") + SYMMOPS: str | None = Field( + default=None, + description=" Build structures having a specified " + "number of symmetry operations. For crystals, " + "the allowed values are (1,2,3,4,6,8,12,16,24,48). " + "For clusters (indicated with #CLUSTER), the allowed " + "values are (1,2,3,5,4,6,7,8,9,10,11,12,24). " + "Ranges are allowed (e.g., #SYMMOPS=1-4).", + ) + SYSTEM: None | Literal["Rhom", "Tric", "Mono", "Cubi", "Hexa", "Orth", "Tetra"] = ( + Field(default=None, description="Enforce a crystal system") + ) SLACK: float | None = Field(default=None, description="The slack factor") + OCTET: bool = Field( + default=False, + description="Check number of valence electrons is a multiple of eight", + ) OVERLAP: float | None = Field(default=None, description="The overlap factor") MINSEP: str | None = Field(default=None, description="The minimum separation") @@ -945,7 +961,7 @@ class RssConfig(AutoplexBaseModel): default_factory=lambda: ["initial", "traj_early", "traj"], description="Configuration types for the VASP calculations", ) - rss_group: list[str] = Field( + rss_group: list[str] | str = Field( default_factory=lambda: ["traj"], description="Group of configurations for the RSS calculations", ) @@ -957,6 +973,11 @@ class RssConfig(AutoplexBaseModel): default=True, description="Whether to apply regularization. This only works for GAP to date.", ) + retain_existing_sigma: bool = Field( + default=False, + description="Whether to retain the existing sigma values for specific configuration types." + "If True, existing sigma values for specific configurations will remain unchanged", + ) scheme: Literal["linear-hull", "volume-stoichiometry", None] = Field( default="linear-hull", description="Method to use for regularization" ) From 6f311e2053d63ae275d5bdde90a7a644decd0dbd Mon Sep 17 00:00:00 2001 From: naik-aakash Date: Sun, 19 Jan 2025 19:20:00 +0100 Subject: [PATCH 46/74] use config yaml from tutorial --- configs/rss_config.yaml | 100 ++++++++++++++++++++++++---------------- 1 file changed, 60 insertions(+), 40 deletions(-) diff --git a/configs/rss_config.yaml b/configs/rss_config.yaml index e166b0b79..57492c4d7 100644 --- a/configs/rss_config.yaml +++ b/configs/rss_config.yaml @@ -1,5 +1,5 @@ # General Parameters -tag: "test" +tag: 'SiO2' train_from_scratch: true resume_from_previous_state: test_error: @@ -9,12 +9,20 @@ resume_from_previous_state: # Buildcell Parameters generated_struct_numbers: - - 9000 - - 1000 + - 8000 + - 2000 buildcell_options: + - SPECIES: "Si%NUM=1,O%NUM=2" + NFORM: '{2,4,6,8}' + SYMMOPS: '1-4' + MINSEP: '1.5 Si-Si=2.7-3.0 Si-O=1.3-1.6 O-O=2.28-2.58' + - SPECIES: "Si%NUM=1,O%NUM=2" + NFORM: '{3,5,7}' + SYMMOPS: '1-4' + MINSEP: '1.5 Si-Si=2.7-3.0 Si-O=1.3-1.6 O-O=2.28-2.58' fragment_file: null fragment_numbers: null -num_processes_buildcell: 64 +num_processes_buildcell: 128 # Sampling Parameters num_of_initial_selected_structs: @@ -40,39 +48,48 @@ random_seed: null # DFT Labelling Parameters include_isolated_atom: true isolatedatom_box: - - 10.0 - - 10.0 - - 10.0 + - 20.0 + - 20.0 + - 20.0 e0_spin: false -include_dimer: true +include_dimer: false dimer_box: - - 10.0 - - 10.0 - - 10.0 + - 20.0 + - 20.0 + - 20.0 dimer_range: - 1.0 - 5.0 -dimer_num: 21 +dimer_num: 41 custom_incar: + KPAR: 8 + NCORE: 16 + LSCALAPACK: ".FALSE." + LPLANE: ".FALSE." ISMEAR: 0 - SIGMA: 0.05 - PREC: 'Accurate' - ADDGRID: '.TRUE.' - EDIFF: 1e-7 + SIGMA: 0.1 + PREC: "Accurate" + ADDGRID: ".TRUE." + EDIFF: 1E-7 NELM: 250 - LWAVE: '.FALSE.' - LCHARG: '.FALSE.' - ALGO: 'Normal' + LWAVE: ".FALSE." + LCHARG: ".FALSE." + ALGO: "normal" AMIX: null - LREAL: '.FALSE.' + LREAL: ".FALSE." ISYM: 0 - ENCUT: 520.0 - KSPACING: 0.20 + ENCUT: 900.0 + KSPACING: 0.23 GGA: null - KPAR: 8 - NCORE: 16 - LSCALAPACK: '.FALSE.' - LPLANE: '.FALSE.' + AMIX_MAG: null + BMIX: null + BMIX_MAG: null + ISTART: null + LMIXTAU: null + NBANDS: null + NELMDL: null + METAGGA: "SCAN" + LASPH: ".TRUE." custom_potcar: vasp_ref_file: 'vasp_ref.extxyz' @@ -83,7 +100,7 @@ config_types: - 'traj' rss_group: - 'traj' -test_ratio: 0.1 +test_ratio: 0.0 regularization: true scheme: 'linear-hull' reg_minmax: @@ -102,41 +119,44 @@ ref_energy_name: 'REF_energy' ref_force_name: 'REF_forces' ref_virial_name: 'REF_virial' auto_delta: true -num_processes_fit: 64 +num_processes_fit: 32 device_for_fitting: 'cpu' -##The following hyperparameters are only applicable to GAP. -##If you want to use other models, please replace the corresponding hyperparameters. twob: cutoff: 5.0 n_sparse: 15 theta_uniform: 1.0 + delta: 1.0 threeb: - cutoff: 3.0 + cutoff: 3.25 soap: - l_max: 10 - n_max: 10 + l_max: 6 + n_max: 12 atom_sigma: 0.5 - n_sparse: 2500 + n_sparse: 3000 cutoff: 5.0 + delta: 0.2 general: three_body: false # RSS Exploration Parameters -scalar_pressure_method: 'uniform' -scalar_exp_pressure: 1 +scalar_pressure_method: 'exp' +scalar_exp_pressure: 100 scalar_pressure_exponential_width: 0.2 scalar_pressure_low: 0 scalar_pressure_high: 25 max_steps: 300 force_tol: 0.01 stress_tol: 0.01 -stop_criterion: 0.01 -max_iteration_number: 25 -num_groups: 6 +stop_criterion: 0.001 +max_iteration_number: 10 +num_groups: 16 initial_kt: 0.3 current_iter_index: 1 -hookean_repul: false +hookean_repul: true hookean_paras: + '(1, 1)': [1000, 0.6] + '(8, 1)': [1000, 0.4] + '(8, 8)': [1000, 1.0] keep_symmetry: false write_traj: true num_processes_rss: 128 From bc8a05e8034e0c75472a61350792c7d847c39ea0 Mon Sep 17 00:00:00 2001 From: naik-aakash Date: Sun, 19 Jan 2025 19:21:11 +0100 Subject: [PATCH 47/74] add __all__ --- src/autoplex/settings.py | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/src/autoplex/settings.py b/src/autoplex/settings.py index 7db1987b9..bbaf4e0cf 100644 --- a/src/autoplex/settings.py +++ b/src/autoplex/settings.py @@ -12,6 +12,18 @@ level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s" ) +__all__ = [ + "AutoplexBaseModel", + "GAPSettings", + "JACESettings", + "M3GNETSettings", + "MACESettings", + "MLIPHypers", + "NEPSettings", + "NEQUIPSettings", + "RssConfig", +] + class AutoplexBaseModel(BaseModel): """Base class for all models in autoplex.""" From d2f2222d572ed390f60ec2d84c592330e3e341f5 Mon Sep 17 00:00:00 2001 From: naik-aakash Date: Mon, 20 Jan 2025 15:21:34 +0100 Subject: [PATCH 48/74] add pending M3GNET parameters --- src/autoplex/settings.py | 66 +++++++++++++++++++++++++++++++++++++--- 1 file changed, 61 insertions(+), 5 deletions(-) diff --git a/src/autoplex/settings.py b/src/autoplex/settings.py index bbaf4e0cf..3387d98e0 100644 --- a/src/autoplex/settings.py +++ b/src/autoplex/settings.py @@ -3,11 +3,15 @@ from __future__ import annotations import logging -from typing import Any, Literal +from typing import TYPE_CHECKING, Any, Literal from monty.serialization import loadfn from pydantic import BaseModel, ConfigDict, Field +if TYPE_CHECKING: + from torch.optim import Optimizer + from torch.optim.lr_scheduler import LRScheduler + logging.basicConfig( level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s" ) @@ -29,7 +33,10 @@ class AutoplexBaseModel(BaseModel): """Base class for all models in autoplex.""" model_config = ConfigDict( - validate_assignment=True, protected_namespaces=(), extra="allow" + validate_assignment=True, + protected_namespaces=(), + extra="allow", + arbitrary_types_allowed=True, ) def update_parameters(self, updates: dict[str, Any]): @@ -468,9 +475,19 @@ class M3GNETSettings(AutoplexBaseModel): results_dir: str = Field( default="m3gnet_results", description="Directory to save the results" ) - pretrained_model: ( - Literal["M3GNet-MP-2021.2.8-PES", "M3GNet-MP-2021.2.8-DIRECT-PES"] | None - ) = Field(default=None, description="Pretrained model") + pretrained_model: str | None = Field( + default=None, + description="Pretrained model. Can be a Path to locally stored model " + "or name of pretrained PES model available in the " + "matgl (`M3GNet-MP-2021.2.8-PES` or " + "`M3GNet-MP-2021.2.8-DIRECT-PES`). When name of " + "model is provided, ensure system has internet " + "access to be able to download the model." + "If None, the model will be trained from scratch.", + ) + allow_missing_labels: bool = Field( + default=False, description="Allow missing labels" + ) cutoff: float = Field(default=5.0, description="Cutoff radius of the graph") threebody_cutoff: float = Field( default=4.0, description="Cutoff radius for 3 body interactions" @@ -480,6 +497,16 @@ class M3GNETSettings(AutoplexBaseModel): include_stresses: bool = Field( default=True, description="Whether to include stresses" ) + data_mean: float = Field(default=0.0, description="Mean of the training data") + data_std: float = Field( + default=1.0, description="Standard deviation of the training data" + ) + decay_steps: int = Field( + default=1000, description="Number of steps for decaying learning rate" + ) + decay_alpha: float = Field( + default=0.96, description="Parameter determines the minimum learning rate" + ) dim_node_embedding: int = Field( default=128, description="Dimension of node embedding" ) @@ -489,14 +516,43 @@ class M3GNETSettings(AutoplexBaseModel): dim_state_embedding: int = Field( default=0, description="Dimension of state embedding" ) + energy_weight: float = Field(default=1.0, description="Weight for energy loss") + force_weight: float = Field(default=1.0, description="Weight for forces loss") + include_line_graph: bool = Field( + default=True, description="Whether to include line graph" + ) + loss: Literal["mse_loss", "huber_loss", "smooth_l1_loss", "l1_loss"] = Field( + default="mse_loss", description="Loss function used for training" + ) + loss_params: dict | None = Field( + default=None, description="Loss function parameters" + ) + lr: float = Field(default=0.001, description="Learning rate for training") + magmom_target: Literal["absolute", "symbreak"] | None = Field( + default="absolute", + description="Whether to predict the absolute " + "site-wise value of magmoms or adapt the loss " + "function to predict the signed value " + "breaking symmetry. If None " + "given the loss function will be adapted.", + ) + magmom_weight: float = Field(default=0.0, description="Weight for magnetic moments") max_l: int = Field(default=4, description="Maximum degree of spherical harmonics") max_n: int = Field( default=4, description="Maximum number of radial basis functions" ) nblocks: int = Field(default=3, description="Number of blocks") + optimizer: Optimizer | None = Field(default=None, description="Optimizer") rbf_type: Literal["Gaussian", "SphericalBessel"] = Field( default="Gaussian", description="Type of radial basis function" ) + scheduler: LRScheduler | None = Field( + default=None, description="Learning rate scheduler" + ) + stress_weight: float = Field(default=0.0, description="Weight for stress loss") + sync_dist: bool = Field( + default=False, description="Sync logging across all GPU workers" + ) is_intensive: bool = Field( default=False, description="Whether the prediction is intensive" ) From ffb5ee87129345615121cbbd0cabfcddc6d7e5cf Mon Sep 17 00:00:00 2001 From: naik-aakash Date: Mon, 20 Jan 2025 15:22:36 +0100 Subject: [PATCH 49/74] fix kwargs in tests of m3gnet --- tests/fitting/common/test_flows.py | 5 +++-- tests/fitting/common/test_jobs.py | 1 + 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/tests/fitting/common/test_flows.py b/tests/fitting/common/test_flows.py index 30201856c..22af90025 100644 --- a/tests/fitting/common/test_flows.py +++ b/tests/fitting/common/test_flows.py @@ -303,8 +303,9 @@ def test_mlip_fit_maker_m3gnet( batch_size=1, max_epochs=3, include_stresses=True, - hidden_dim=8, - num_units=8, + dim_node_embedding=8, + dim_edge_embedding=8, + units=8, max_l=4, max_n=4, device="cpu", diff --git a/tests/fitting/common/test_jobs.py b/tests/fitting/common/test_jobs.py index f4a7cd9e9..15762c270 100644 --- a/tests/fitting/common/test_jobs.py +++ b/tests/fitting/common/test_jobs.py @@ -84,6 +84,7 @@ def test_m3gnet_fit_maker(test_dir, memory_jobstore, clean_dir): max_epochs=3, include_stresses=True, dim_node_embedding=8, + dim_edge_embedding=8, units=8, max_l=4, max_n=4, From 8545d4dc57cbc122d5fa6f42c37e317ad7cf4ea8 Mon Sep 17 00:00:00 2001 From: naik-aakash Date: Mon, 20 Jan 2025 15:23:34 +0100 Subject: [PATCH 50/74] correct kwargs in m3gnet --- src/autoplex/fitting/common/utils.py | 48 ++++++++++++++++++++++++---- 1 file changed, 42 insertions(+), 6 deletions(-) diff --git a/src/autoplex/fitting/common/utils.py b/src/autoplex/fitting/common/utils.py index 9ac31a24a..7a655b1d6 100644 --- a/src/autoplex/fitting/common/utils.py +++ b/src/autoplex/fitting/common/utils.py @@ -648,9 +648,11 @@ def m3gnet_fitting( Maximum number of training epochs. include_stresses: bool If True, includes stress tensors in the model predictions and training process. - hidden_dim: int - Dimensionality of the hidden layers in the model. - num_units: int + dim_node_embedding: int + Dimension of node embedding. + dim_edge_embedding: int + Dimension of edge embeddings. + units: int Number of units in each dense layer of the model. max_l: int Maximum degree of spherical harmonics. @@ -837,7 +839,26 @@ def m3gnet_fitting( max_n=m3gnet_hypers.get("max_n"), nblocks=m3gnet_hypers.get("nblocks"), ) - lit_module = PotentialLightningModule(model=model, include_line_graph=True) + lit_module = PotentialLightningModule( + model=model, + include_line_graph=m3gnet_hypers.get("include_line_graph"), + allow_missing_labels=m3gnet_hypers.get("allow_missing_labels"), + energy_weight=m3gnet_hypers.get("energy_weight"), + force_weight=m3gnet_hypers.get("force_weight"), + lr=m3gnet_hypers.get("lr"), + loss=m3gnet_hypers.get("loss"), + loss_params=m3gnet_hypers.get("loss_params"), + stress_weight=m3gnet_hypers.get("stress_weight"), + magmom_weight=m3gnet_hypers.get("magmom_weight"), + data_mean=m3gnet_hypers.get("data_mean"), + data_std=m3gnet_hypers.get("data_std"), + decay_alpha=m3gnet_hypers.get("decay_alpha"), + decay_steps=m3gnet_hypers.get("decay_steps"), + sync_dist=m3gnet_hypers.get("sync_dist"), + magmom_target=m3gnet_hypers.get("magmom_target"), + optimizer=m3gnet_hypers.get("optimizer"), + scheduler=m3gnet_hypers.get("scheduler"), + ) else: # finetune pretrained model logging.info( f"Finetuning pretrained model: {m3gnet_hypers['pretrained_model']}" @@ -848,8 +869,23 @@ def m3gnet_fitting( lit_module = PotentialLightningModule( model=model, element_refs=property_offset, - lr=1e-4, - include_line_graph=True, + include_line_graph=m3gnet_hypers.get("include_line_graph"), + allow_missing_labels=m3gnet_hypers.get("allow_missing_labels"), + energy_weight=m3gnet_hypers.get("energy_weight"), + force_weight=m3gnet_hypers.get("force_weight"), + lr=m3gnet_hypers.get("lr"), + loss=m3gnet_hypers.get("loss"), + loss_params=m3gnet_hypers.get("loss_params"), + stress_weight=m3gnet_hypers.get("stress_weight"), + magmom_weight=m3gnet_hypers.get("magmom_weight"), + data_mean=m3gnet_hypers.get("data_mean"), + data_std=m3gnet_hypers.get("data_std"), + decay_alpha=m3gnet_hypers.get("decay_alpha"), + decay_steps=m3gnet_hypers.get("decay_steps"), + sync_dist=m3gnet_hypers.get("sync_dist"), + magmom_target=m3gnet_hypers.get("magmom_target"), + optimizer=m3gnet_hypers.get("optimizer"), + scheduler=m3gnet_hypers.get("scheduler"), ) logger = CSVLogger(name=exp_name, save_dir=os.path.join(results_dir, "logs")) From 99b13d72cf334fa6998a181777c66d6b98f3a864 Mon Sep 17 00:00:00 2001 From: naik-aakash Date: Mon, 20 Jan 2025 15:35:10 +0100 Subject: [PATCH 51/74] remove typechecking --- src/autoplex/settings.py | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/src/autoplex/settings.py b/src/autoplex/settings.py index 3387d98e0..7d7716e74 100644 --- a/src/autoplex/settings.py +++ b/src/autoplex/settings.py @@ -3,14 +3,12 @@ from __future__ import annotations import logging -from typing import TYPE_CHECKING, Any, Literal +from typing import Any, Literal from monty.serialization import loadfn from pydantic import BaseModel, ConfigDict, Field - -if TYPE_CHECKING: - from torch.optim import Optimizer - from torch.optim.lr_scheduler import LRScheduler +from torch.optim import Optimizer # noqa: TC002 +from torch.optim.lr_scheduler import LRScheduler # noqa: TC002 logging.basicConfig( level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s" From c38104c0ce619f0b64367bfbc02857eb5522d99c Mon Sep 17 00:00:00 2001 From: naik-aakash Date: Tue, 21 Jan 2025 09:35:40 +0100 Subject: [PATCH 52/74] add missing kwarg --- src/autoplex/fitting/common/utils.py | 1 + src/autoplex/settings.py | 4 ++++ 2 files changed, 5 insertions(+) diff --git a/src/autoplex/fitting/common/utils.py b/src/autoplex/fitting/common/utils.py index 7a655b1d6..83adc3352 100644 --- a/src/autoplex/fitting/common/utils.py +++ b/src/autoplex/fitting/common/utils.py @@ -841,6 +841,7 @@ def m3gnet_fitting( ) lit_module = PotentialLightningModule( model=model, + element_refs=m3gnet_hypers.get("element_refs"), include_line_graph=m3gnet_hypers.get("include_line_graph"), allow_missing_labels=m3gnet_hypers.get("allow_missing_labels"), energy_weight=m3gnet_hypers.get("energy_weight"), diff --git a/src/autoplex/settings.py b/src/autoplex/settings.py index 7d7716e74..6917ce35d 100644 --- a/src/autoplex/settings.py +++ b/src/autoplex/settings.py @@ -5,6 +5,7 @@ import logging from typing import Any, Literal +import numpy as np # noqa: TC002 from monty.serialization import loadfn from pydantic import BaseModel, ConfigDict, Field from torch.optim import Optimizer # noqa: TC002 @@ -515,6 +516,9 @@ class M3GNETSettings(AutoplexBaseModel): default=0, description="Dimension of state embedding" ) energy_weight: float = Field(default=1.0, description="Weight for energy loss") + element_refs: np.ndarray | None = Field( + default=None, description="Element offset for PES" + ) force_weight: float = Field(default=1.0, description="Weight for forces loss") include_line_graph: bool = Field( default=True, description="Whether to include line graph" From 75db90649b1da94ab8334e803b1b269e8cda9c89 Mon Sep 17 00:00:00 2001 From: naik-aakash Date: Tue, 21 Jan 2025 12:23:24 +0100 Subject: [PATCH 53/74] add missing m3gnet finetuning tests --- tests/auto/phonons/test_flows.py | 44 ++++++++++++++++++++++++++++++++ 1 file changed, 44 insertions(+) diff --git a/tests/auto/phonons/test_flows.py b/tests/auto/phonons/test_flows.py index 9caeca309..1ca23c324 100644 --- a/tests/auto/phonons/test_flows.py +++ b/tests/auto/phonons/test_flows.py @@ -927,6 +927,50 @@ def test_complete_dft_vs_ml_benchmark_workflow_m3gnet( 5.2622804443539355, abs=3.0 # bad fit data, fluctuates between 4 and 7 ) +def test_complete_dft_vs_ml_benchmark_workflow_m3gnet_finetuning( + vasp_test_dir, mock_vasp, test_dir, memory_jobstore, ref_paths4_mpid, fake_run_vasp_kwargs4_mpid, clean_dir +): + path_to_struct = vasp_test_dir / "dft_ml_data_generation" / "POSCAR" + structure = Structure.from_file(path_to_struct) + + complete_workflow_m3gnet = CompleteDFTvsMLBenchmarkWorkflow( + ml_models=["M3GNET"], + symprec=1e-2, supercell_settings={"min_length": 8, "min_atoms": 20}, displacements=[0.01], + volume_custom_scale_factors=[0.975, 1.0, 1.025, 1.05], + apply_data_preprocessing=True, + ).make( + structure_list=[structure], + mp_ids=["test"], + benchmark_mp_ids=["mp-22905"], + pre_xyz_files=["vasp_ref.extxyz"], + pre_database_dir=test_dir / "fitting" / "ref_files", + benchmark_structures=[structure], + fit_kwargs_list=[{ + "batch_size": 1, + "max_epochs": 1, + "include_stresses": True, + "device": "cpu", + "test_equal_to_val": True, + "pretrained_model": "M3GNet-MP-2021.2.8-DIRECT-PES", + }] + ) + + # automatically use fake VASP and write POTCAR.spec during the test + mock_vasp(ref_paths4_mpid, fake_run_vasp_kwargs4_mpid) + + # run the flow or job and ensure that it finished running successfully + responses = run_locally( + complete_workflow_m3gnet, + create_folders=True, + ensure_success=True, + store=memory_jobstore, + ) + assert complete_workflow_m3gnet.jobs[5].name == "complete_benchmark_mp-22905" + assert responses[complete_workflow_m3gnet.jobs[-1].output.uuid][1].output["metrics"][0][0][ + "benchmark_phonon_rmse"] == pytest.approx( + 4.6, abs=0.5, + ) + def test_complete_dft_vs_ml_benchmark_workflow_mace( vasp_test_dir, mock_vasp, test_dir, memory_jobstore, ref_paths4_mpid, fake_run_vasp_kwargs4_mpid, clean_dir From 68798cf779c42b47a13ca197c173cf4b5d82fd2e Mon Sep 17 00:00:00 2001 From: naik-aakash Date: Tue, 21 Jan 2025 15:06:15 +0100 Subject: [PATCH 54/74] small fix for kwargs updates and rename field name --- src/autoplex/auto/rss/flows.py | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/src/autoplex/auto/rss/flows.py b/src/autoplex/auto/rss/flows.py index 163884a65..b9d355ae2 100644 --- a/src/autoplex/auto/rss/flows.py +++ b/src/autoplex/auto/rss/flows.py @@ -4,8 +4,8 @@ from jobflow import Flow, Maker, Response, job -from autoplex import RSS_CONFIG from autoplex.auto.rss.jobs import do_rss_iterations, initial_rss +from autoplex.settings import RssConfig @dataclass @@ -17,13 +17,13 @@ class RssMaker(Maker): ---------- name: str Name of the flow. - config: RSS_CONFIG + rss_config: RssConfig Pydantic model that defines the setup parameters for the whole RSS workflow. If not explicitly set, the defaults from 'autoplex.settings.RssConfig' will be used. """ name: str = "ml-driven rss" - config: RSS_CONFIG = field(default_factory=lambda: RSS_CONFIG) + rss_config: RssConfig = field(default_factory=lambda: RssConfig()) @job def make(self, **kwargs): @@ -233,9 +233,11 @@ def make(self, **kwargs): - 'current_iter': int, The current iteration index. - 'kb_temp': float, The temperature (in eV) for Boltzmann sampling. """ - default_config = self.config.model_copy(deep=True) - updated_config = default_config.update_parameters(kwargs) - config_params = updated_config.model_dump(by_alias=True, exclude_none=True) + default_config = self.rss_config.model_copy(deep=True) + if kwargs: + default_config = default_config.update_parameters(kwargs) + + config_params = default_config.model_dump(by_alias=True, exclude_none=True) self._process_hookean_paras(config_params) From 55490a664a6c62d80debcd5c104c6c3e74b949b7 Mon Sep 17 00:00:00 2001 From: naik-aakash Date: Tue, 21 Jan 2025 15:06:44 +0100 Subject: [PATCH 55/74] adapt test for config read --- tests/auto/rss/test_flows.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/tests/auto/rss/test_flows.py b/tests/auto/rss/test_flows.py index 3e3391ddd..fc7c62f1f 100644 --- a/tests/auto/rss/test_flows.py +++ b/tests/auto/rss/test_flows.py @@ -315,13 +315,13 @@ def test_rssmaker_custom_config_file(test_dir): config_model = RssConfig.from_file(test_dir / "rss" / "rss_config.yaml") # Test if config is updated as expected - rss = RssMaker(config=config_model) + rss = RssMaker(rss_config=config_model) - assert rss.config.tag == "test" - assert rss.config.generated_struct_numbers == [9000, 1000] - assert rss.config.num_processes_buildcell == 64 - assert rss.config.num_processes_fit == 64 - assert rss.config.device_for_rss == "cuda" - assert rss.config.isolatedatom_box == [10, 10, 10] - assert rss.config.dimer_box == [10, 10, 10] + assert rss.rss_config.tag == "test" + assert rss.rss_config.generated_struct_numbers == [9000, 1000] + assert rss.rss_config.num_processes_buildcell == 64 + assert rss.rss_config.num_processes_fit == 64 + assert rss.rss_config.device_for_rss == "cuda" + assert rss.rss_config.isolatedatom_box == [10, 10, 10] + assert rss.rss_config.dimer_box == [10, 10, 10] From 5f6e3e3359f80a01522f3dd1da9eb2889d48c8dd Mon Sep 17 00:00:00 2001 From: naik-aakash Date: Tue, 21 Jan 2025 15:16:56 +0100 Subject: [PATCH 56/74] update doc of rss maker --- docs/user/rss/flow/quick_start/start.md | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/docs/user/rss/flow/quick_start/start.md b/docs/user/rss/flow/quick_start/start.md index e3ad02260..714a1a9af 100644 --- a/docs/user/rss/flow/quick_start/start.md +++ b/docs/user/rss/flow/quick_start/start.md @@ -39,9 +39,9 @@ from autoplex.auto.rss.flows import RssMaker from fireworks import LaunchPad from jobflow.managers.fireworks import flow_to_workflow -config = RssConfig.from_file('path/to/your/config.yaml') +rss_config = RssConfig.from_file('path/to/your/config.yaml') -rss_job = RssMaker(name="your workflow name", config=config).make() +rss_job = RssMaker(name="your workflow name", rss_config=rss_config).make() wf = flow_to_workflow(rss_job) lpad = LaunchPad.auto_load() lpad.add_wf(wf) @@ -51,10 +51,12 @@ The above code is based on [`FireWorks`](https://materialsproject.github.io/fire ```python +from autoplex.settings import RssConfig from autoplex.auto.rss.flows import RssMaker from jobflow_remote import submit_flow -rss_job = RssMaker(name="your workflow name", config_file='path/to/your/name.yaml').make() +rss_config = RssConfig.from_file('path/to/your/config.yaml') +rss_job = RssMaker(name="your workflow name", rss_config=rss_config).make() resources = {"nodes": N, "partition": "name", "qos": "name", "time": "8:00:00", "mail_user": "your_email", "mail_type": "ALL", "account": "your account"} print(submit_flow(rss_job, worker="your worker", resources=resources, project="your project name")) ``` From 224630e6d5052fc5332ff1b3364c7254028e5ddd Mon Sep 17 00:00:00 2001 From: naik-aakash Date: Tue, 21 Jan 2025 16:15:58 +0100 Subject: [PATCH 57/74] fix none type return issue --- src/autoplex/auto/rss/flows.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/autoplex/auto/rss/flows.py b/src/autoplex/auto/rss/flows.py index b9d355ae2..6240daa4a 100644 --- a/src/autoplex/auto/rss/flows.py +++ b/src/autoplex/auto/rss/flows.py @@ -235,7 +235,7 @@ def make(self, **kwargs): """ default_config = self.rss_config.model_copy(deep=True) if kwargs: - default_config = default_config.update_parameters(kwargs) + default_config.update_parameters(kwargs) config_params = default_config.model_dump(by_alias=True, exclude_none=True) From 0ab9fa460a57c89360a1a5bad2656866ffac2a3e Mon Sep 17 00:00:00 2001 From: naik-aakash Date: Thu, 30 Jan 2025 11:20:20 +0100 Subject: [PATCH 58/74] update RssConfig model , add backward compatiblity --- src/autoplex/settings.py | 100 ++++++++++++++++++--------------------- 1 file changed, 46 insertions(+), 54 deletions(-) diff --git a/src/autoplex/settings.py b/src/autoplex/settings.py index 6917ce35d..b3da9c17e 100644 --- a/src/autoplex/settings.py +++ b/src/autoplex/settings.py @@ -73,6 +73,7 @@ def from_file(cls, filename: str): filename (str): The name of the file to load the parameters from. """ custom_params = loadfn(filename) + return cls(**custom_params) @@ -262,8 +263,8 @@ class NEQUIPSettings(AutoplexBaseModel): default=False, description="When true a restarted run will append to the previous log file", ) - default_dtype: str = Field(default="float32", description="Default data type") - model_dtype: str = Field(default="float32", description="Model data type") + default_dtype: str = Field(default="float64", description="Default data type") + model_dtype: str = Field(default="float64", description="Model data type") allow_tf32: bool = Field( default=True, description="Consider setting to false if you plan to mix " @@ -758,7 +759,9 @@ class MLIPHypers(AutoplexBaseModel): default_factory=GAPSettings, description="Hyperparameters for the GAP model" ) J_ACE: JACESettings = Field( - default_factory=JACESettings, description="Hyperparameters for the J-ACE model" + default_factory=JACESettings, + description="Hyperparameters for the J-ACE model", + alias="J-ACE", ) NEQUIP: NEQUIPSettings = Field( default_factory=NEQUIPSettings, @@ -888,42 +891,6 @@ class CustomIncar(AutoplexBaseModel): LPLANE: str = ".FALSE." -class Twob(AutoplexBaseModel): - """A model describing the two-body GAP parameters.""" - - cutoff: float = Field(default=5.0, description="Radial cutoff distance") - n_sparse: int = Field(default=15, description="Number of sparse points") - theta_uniform: float = Field( - default=1.0, description="Width of the uniform distribution for theta" - ) - - -class Threeb(AutoplexBaseModel): - """A model describing the three-body GAP parameters.""" - - cutoff: float = Field(default=3.0, description="Radial cutoff distance") - - -class Soap(AutoplexBaseModel): - """A model describing the SOAP GAP parameters.""" - - l_max: int = Field(default=10, description="Maximum degree of spherical harmonics") - n_max: int = Field( - default=10, description="Maximum number of radial basis functions" - ) - atom_sigma: float = Field(default=0.5, description="Width of Gaussian smearing") - n_sparse: int = Field(default=2500, description="Number of sparse points") - cutoff: float = Field(default=5.0, description="Radial cutoff distance") - - -class General(AutoplexBaseModel): - """A model describing the general GAP parameters.""" - - three_body: bool = Field( - default=False, description="Whether to include three-body terms" - ) - - class RssConfig(AutoplexBaseModel): """A model describing the complete RSS configuration.""" @@ -1095,21 +1062,6 @@ class RssConfig(AutoplexBaseModel): device_for_fitting: Literal["cpu", "cuda"] = Field( default="cpu", description="Device to be used for model fitting" ) - twob: Twob = Field( - default_factory=Twob, - description="Parameters for the two-body descriptor, Applicable on to GAP", - ) - threeb: Threeb = Field( - default_factory=Threeb, - description="Parameters for the three-body descriptor, Applicable on to GAP", - ) - soap: Soap = Field( - default_factory=Soap, - description="Parameters for the SOAP descriptor, Applicable on to GAP", - ) - general: General = Field( - default_factory=General, description="General parameters for the GAP model" - ) scalar_pressure_method: Literal["exp", "uniform"] = Field( default="uniform", description="Method for adding external pressures." ) @@ -1174,3 +1126,43 @@ class RssConfig(AutoplexBaseModel): device_for_rss: Literal["cpu", "cuda"] = Field( default="cpu", description="Device to be used for RSS calculations." ) + mlip_hypers: MLIPHypers = Field( + default_factory=MLIPHypers, description="MLIP hyperparameters" + ) + + @classmethod + def from_file(cls, filename: str): + """Create RSS configuration object from a file.""" + config_params = loadfn(filename) + + # check if config file has the required keys when train_from_scratch is False + train_from_scratch = config_params.get("train_from_scratch") + resume_from_previous_state = config_params.get("resume_from_previous_state") + + if not train_from_scratch: + for key, value in resume_from_previous_state.items(): + if value is None: + raise ValueError( + f"Value for {key} in `resume_from_previous_state` cannot be None when " + f"`train_from_scratch` is set to False" + ) + + # check if mlip arg is in the config file + # Needed for backward compatibility with older config files of RSS workflow + mlip_type = config_params["mlip_type"].replace("-", "_") + mlip_hypers = MLIPHypers().__getattribute__(mlip_type) + + if "mlip_hypers" not in config_params: + config_params["mlip_hypers"] = {config_params["mlip_type"]: {}} + + old_config_keys = [] + for arg in config_params: + mlip_type = config_params["mlip_type"].replace("-", "_") + if arg in mlip_hypers.model_fields: + config_params["mlip_hypers"][mlip_type][arg] = config_params[arg] + old_config_keys.append(arg) + + for key in old_config_keys: + del config_params[key] + + return cls(**config_params) From 3bcd392e44b9f53d0bdc850298be430818bc8999 Mon Sep 17 00:00:00 2001 From: naik-aakash Date: Thu, 30 Jan 2025 11:21:10 +0100 Subject: [PATCH 59/74] adapt config processing to new format --- src/autoplex/auto/rss/flows.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/src/autoplex/auto/rss/flows.py b/src/autoplex/auto/rss/flows.py index 6240daa4a..1e2d06569 100644 --- a/src/autoplex/auto/rss/flows.py +++ b/src/autoplex/auto/rss/flows.py @@ -239,6 +239,11 @@ def make(self, **kwargs): config_params = default_config.model_dump(by_alias=True, exclude_none=True) + # Extract MLIP hyperparameters from the config_params + mlip_hypers = config_params["mlip_hypers"][config_params["mlip_type"]] + del config_params["mlip_hypers"] + config_params.update(mlip_hypers) + self._process_hookean_paras(config_params) if "train_from_scratch" not in config_params: From bb83042f2b7dd1ffbd3d4e43cf8d455391a06376 Mon Sep 17 00:00:00 2001 From: naik-aakash Date: Thu, 30 Jan 2025 11:21:50 +0100 Subject: [PATCH 60/74] adapt rss tutorials to new changes --- docs/user/rss/flow/example/example.md | 225 +++++++++++++----------- docs/user/rss/flow/input/input.md | 72 +++++--- docs/user/rss/flow/quick_start/start.md | 4 +- 3 files changed, 181 insertions(+), 120 deletions(-) diff --git a/docs/user/rss/flow/example/example.md b/docs/user/rss/flow/example/example.md index aba1c5dcf..75e1e0361 100644 --- a/docs/user/rss/flow/example/example.md +++ b/docs/user/rss/flow/example/example.md @@ -10,38 +10,34 @@ This section provides guidance on exploring silica models at different functiona ```yaml # General Parameters -tag: 'SiO2' +tag: SiO2 train_from_scratch: true -resume_from_previous_state: - test_error: - pre_database_dir: - mlip_path: - isolated_atom_energies: - +resume_from_previous_state: {} # Buildcell Parameters generated_struct_numbers: - - 8000 - - 2000 +- 8000 +- 2000 buildcell_options: - - SPECIES: "Si%NUM=1,O%NUM=2" - NFORM: '{2,4,6,8}' - SYMMOPS: '1-4' - MINSEP: '1.5 Si-Si=2.7-3.0 Si-O=1.3-1.6 O-O=2.28-2.58' - - SPECIES: "Si%NUM=1,O%NUM=2" - NFORM: '{3,5,7}' - SYMMOPS: '1-4' - MINSEP: '1.5 Si-Si=2.7-3.0 Si-O=1.3-1.6 O-O=2.28-2.58' -fragment_file: null -fragment_numbers: null +- ABFIX: false + NFORM: '{2,4,6,8}' + SYMMOPS: 1-4 + OCTET: false + MINSEP: 1.5 Si-Si=2.7-3.0 Si-O=1.3-1.6 O-O=2.28-2.58 + SPECIES: Si%NUM=1,O%NUM=2 +- ABFIX: false + NFORM: '{3,5,7}' + SYMMOPS: 1-4 + OCTET: false + MINSEP: 1.5 Si-Si=2.7-3.0 Si-O=1.3-1.6 O-O=2.28-2.58 + SPECIES: Si%NUM=1,O%NUM=2 num_processes_buildcell: 128 - # Sampling Parameters num_of_initial_selected_structs: - - 80 - - 20 +- 80 +- 20 num_of_rss_selected_structs: 100 initial_selection_enabled: true -rss_selection_method: 'bcur2i' +rss_selection_method: bcur2i bcur_params: soap_paras: l_max: 12 @@ -54,103 +50,69 @@ bcur_params: species: true frac_of_bcur: 0.8 bolt_max_num: 3000 -random_seed: null # DFT Labelling Parameters include_isolated_atom: true isolatedatom_box: - - 20.0 - - 20.0 - - 20.0 +- 20.0 +- 20.0 +- 20.0 e0_spin: false include_dimer: false dimer_box: - - 20.0 - - 20.0 - - 20.0 +- 20.0 +- 20.0 +- 20.0 dimer_range: - - 1.0 - - 5.0 +- 1.0 +- 5.0 dimer_num: 41 custom_incar: - KPAR: 8 - NCORE: 16 - LSCALAPACK: ".FALSE." - LPLANE: ".FALSE." ISMEAR: 0 SIGMA: 0.1 - PREC: "Accurate" - ADDGRID: ".TRUE." - EDIFF: 1E-7 + PREC: Accurate + ADDGRID: .TRUE. + EDIFF: 1e-07 NELM: 250 - LWAVE: ".FALSE." - LCHARG: ".FALSE." - ALGO: "normal" - AMIX: null - LREAL: ".FALSE." + LWAVE: .FALSE. + LCHARG: .FALSE. + ALGO: normal + LREAL: .FALSE. ISYM: 0 ENCUT: 900.0 KSPACING: 0.23 - GGA: null - AMIX_MAG: null - BMIX: null - BMIX_MAG: null - ISTART: null - LMIXTAU: null - NBANDS: null - NELMDL: null - METAGGA: "SCAN" - LASPH: ".TRUE." -custom_potcar: -vasp_ref_file: 'vasp_ref.extxyz' + KPAR: 8 + NCORE: 16 + LSCALAPACK: .FALSE. + LPLANE: .FALSE. + METAGGA: SCAN + LASPH: .TRUE. +vasp_ref_file: vasp_ref.extxyz # Data Preprocessing Parameters config_types: - - 'initial' - - 'traj_early' - - 'traj' +- initial +- traj_early +- traj rss_group: - - 'traj' +- traj test_ratio: 0.0 regularization: true -scheme: 'linear-hull' +retain_existing_sigma: false +scheme: linear-hull reg_minmax: - - [0.1, 1] - - [0.001, 0.1] - - [0.0316, 0.316] - - [0.0632, 0.632] +- - 0.1 + - 1.0 +- - 0.001 + - 0.1 +- - 0.0316 + - 0.316 +- - 0.0632 + - 0.632 distillation: false -force_max: null -force_label: null -pre_database_dir: null - -# MLIP Parameters -mlip_type: 'GAP' -ref_energy_name: 'REF_energy' -ref_force_name: 'REF_forces' -ref_virial_name: 'REF_virial' -auto_delta: true -num_processes_fit: 32 -device_for_fitting: 'cpu' -twob: - cutoff: 5.0 - n_sparse: 15 - theta_uniform: 1.0 - delta: 1.0 -threeb: - cutoff: 3.25 -soap: - l_max: 6 - n_max: 12 - atom_sigma: 0.5 - n_sparse: 3000 - cutoff: 5.0 - delta: 0.2 -general: - three_body: false # RSS Exploration Parameters -scalar_pressure_method: 'exp' +scalar_pressure_method: exp scalar_exp_pressure: 100 scalar_pressure_exponential_width: 0.2 scalar_pressure_low: 0 @@ -165,13 +127,80 @@ initial_kt: 0.3 current_iter_index: 1 hookean_repul: true hookean_paras: - '(1, 1)': [1000, 0.6] - '(8, 1)': [1000, 0.4] - '(8, 8)': [1000, 1.0] + (1, 1): + - 1000 + - 0.6 + (8, 1): + - 1000 + - 0.4 + (8, 8): + - 1000 + - 1.0 keep_symmetry: false write_traj: true num_processes_rss: 128 -device_for_rss: 'cpu' +device_for_rss: cpu + +# MLIP Parameters +mlip_type: GAP +ref_energy_name: REF_energy +ref_force_name: REF_forces +ref_virial_name: REF_virial +auto_delta: true +num_processes_fit: 32 +device_for_fitting: cpu +mlip_hypers: + GAP: + general: + at_file: train.extxyz + default_sigma: '{0.0001 0.05 0.05 0}' + energy_parameter_name: REF_energy + force_parameter_name: REF_forces + virial_parameter_name: REF_virial + sparse_jitter: 1e-08 + do_copy_at_file: F + openmp_chunk_size: 10000 + gp_file: gap_file.xml + e0_offset: 0.0 + two_body: false + three_body: false + soap: true + twob: + distance_Nb order: 2 + f0: 0.0 + add_species: T + cutoff: 5.0 + n_sparse: 15 + covariance_type: ard_se + delta: 1.0 + theta_uniform: 1.0 + sparse_method: uniform + compact_clusters: T + threeb: + distance_Nb order: 3 + f0: 0.0 + add_species: T + cutoff: 3.25 + n_sparse: 100 + covariance_type: ard_se + delta: 2.0 + theta_uniform: 1.0 + sparse_method: uniform + compact_clusters: T + soap: + add_species: T + l_max: 6 + n_max: 12 + atom_sigma: 0.5 + zeta: 4 + cutoff: 5.0 + cutoff_transition_width: 1.0 + central_weight: 1.0 + n_sparse: 3000 + delta: 0.2 + f0: 0.0 + covariance_type: dot_product + sparse_method: cur_points ``` To switch from SCAN to PBE, simply remove the `METAGGA` and `LASPH` entries from the `custom_incar` settings. All other parameters remain unchanged. diff --git a/docs/user/rss/flow/input/input.md b/docs/user/rss/flow/input/input.md index 474dbf6ff..e134aa7a9 100644 --- a/docs/user/rss/flow/input/input.md +++ b/docs/user/rss/flow/input/input.md @@ -174,26 +174,58 @@ You can specify the desired model using the `mlip_type` argument and tune hyperp ```yaml # MLIP Parameters mlip_type: 'GAP' -ref_energy_name: 'REF_energy' -ref_force_name: 'REF_forces' -ref_virial_name: 'REF_virial' -auto_delta: true -num_processes_fit: 32 -device_for_fitting: 'cpu' -twob: - cutoff: 10.0 - n_sparse: 30 - theta_uniform: 1.0 -threeb: - cutoff: 3.25 -soap: - l_max: 8 - n_max: 8 - atom_sigma: 0.75 - n_sparse: 2000 - cutoff: 5.0 -general: - three_body: true +mlip_hypers: + GAP: + general: + at_file: train.extxyz + default_sigma: '{0.0001 0.05 0.05 0}' + energy_parameter_name: REF_energy + force_parameter_name: REF_forces + virial_parameter_name: REF_virial + sparse_jitter: 1e-08 + do_copy_at_file: F + openmp_chunk_size: 10000 + gp_file: gap_file.xml + e0_offset: 0.0 + two_body: true + three_body: false + soap: true + twob: + distance_Nb_order: 2 + f0: 0.0 + add_species: T + cutoff: 5.0 + n_sparse: 15 + covariance_type: ard_se + delta: 2.0 + theta_uniform: 0.5 + sparse_method: uniform + compact_clusters: T + threeb: + distance_Nb_order: 3 + f0: 0.0 + add_species: T + cutoff: 3.25 + n_sparse: 100 + covariance_type: ard_se + delta: 2.0 + theta_uniform: 1.0 + sparse_method: uniform + compact_clusters: T + soap: + add_species: T + l_max: 10 + n_max: 12 + atom_sigma: 0.5 + zeta: 4 + cutoff: 5.0 + cutoff_transition_width: 1.0 + central_weight: 1.0 + n_sparse: 6000 + delta: 1.0 + f0: 0.0 + covariance_type: dot_product + sparse_method: cur_points ``` ## RSS Exploration Parameters diff --git a/docs/user/rss/flow/quick_start/start.md b/docs/user/rss/flow/quick_start/start.md index 714a1a9af..996ae05c1 100644 --- a/docs/user/rss/flow/quick_start/start.md +++ b/docs/user/rss/flow/quick_start/start.md @@ -103,8 +103,8 @@ print(MLIP_HYPERS.MACE) # Eg:- access MACE hyperparameters rss_job = RssMaker(name="your workflow name").make(tag='SiO2', ... # Other parameters here mlip_type='MACE', - hidden_irreps="128x0e + 128x1o", - r_max=5.0) + {"MACE": "hidden_irreps":"128x0e + 128x1o","r_max":5.0}, + ) ``` > **Note**: We primarily recommend the GAP-RSS model for now, as GAP has demonstrated great stability with small datasets. Other models have not been thoroughly explored yet. However, we encourage users to experiment with and test other individual models or combinations for potentially interesting results. From 09cffbf8f1670c9651b171f33015a7aeefab6a60 Mon Sep 17 00:00:00 2001 From: naik-aakash Date: Thu, 30 Jan 2025 11:22:12 +0100 Subject: [PATCH 61/74] update example config files --- configs/create_configs.py | 4 +- configs/mlip_hypers.json | 245 ++++++++++++++++++++--- configs/rss_config.yaml | 228 ++++++++++----------- configs/rss_config_all.yaml | 383 ++++++++++++++++++++++++++++++++++++ 4 files changed, 727 insertions(+), 133 deletions(-) mode change 100644 => 100755 configs/rss_config.yaml create mode 100755 configs/rss_config_all.yaml diff --git a/configs/create_configs.py b/configs/create_configs.py index 9c8cee876..0385f470e 100644 --- a/configs/create_configs.py +++ b/configs/create_configs.py @@ -6,4 +6,6 @@ custom_hyperparameters = MLIPHypers.from_file("mlip_hypers.json") # create a custom rss config object using json file -custom_rss_config = RssConfig.from_file("rss_config.yaml") +custom_rss_config = RssConfig.from_file( + "rss_config.yaml" +) # rss_config_all.yaml contains list of all hypers and supported MLIPs diff --git a/configs/mlip_hypers.json b/configs/mlip_hypers.json index 576efd2d8..ba02c7b11 100644 --- a/configs/mlip_hypers.json +++ b/configs/mlip_hypers.json @@ -6,7 +6,7 @@ "energy_parameter_name": "REF_energy", "force_parameter_name": "REF_forces", "virial_parameter_name": "REF_virial", - "sparse_jitter": 1.0e-8, + "sparse_jitter": 1e-08, "do_copy_at_file": "F", "openmp_chunk_size": 10000, "gp_file": "gap_file.xml", @@ -22,8 +22,8 @@ "cutoff": 5.0, "n_sparse": 15, "covariance_type": "ard_se", - "delta": 2.00 , - "theta_uniform": 0.5 , + "delta": 2.0, + "theta_uniform": 0.5, "sparse_method": "uniform", "compact_clusters": "T" }, @@ -34,8 +34,8 @@ "cutoff": 3.25, "n_sparse": 100, "covariance_type": "ard_se", - "delta": 2.00 , - "theta_uniform": 1.0 , + "delta": 2.0, + "theta_uniform": 1.0, "sparse_method": "uniform", "compact_clusters": "T" }, @@ -49,57 +49,258 @@ "cutoff_transition_width": 1.0, "central_weight": 1.0, "n_sparse": 6000, - "delta": 1.00, + "delta": 1.0, "f0": 0.0, "covariance_type": "dot_product", "sparse_method": "cur_points" } }, + "J-ACE": { + "order": 3, + "totaldegree": 6, + "cutoff": 2.0, + "solver": "BLR" + }, "NEQUIP": { + "root": "results", + "run_name": "autoplex", + "seed": 123, + "dataset_seed": 123, + "append": false, + "default_dtype": "float64", + "model_dtype": "float64", + "allow_tf32": true, "r_max": 4.0, "num_layers": 4, "l_max": 2, + "parity": true, "num_features": 32, + "nonlinearity_type": "gate", + "nonlinearity_scalars": { + "e": "silu", + "o": "tanh" + }, + "nonlinearity_gates": { + "e": "silu", + "o": "tanh" + }, "num_basis": 8, + "BesselBasis_trainable": true, + "PolynomialCutoff_p": 5, "invariant_layers": 2, "invariant_neurons": 64, - "batch_size": 5, + "avg_num_neighbors": "auto", + "use_sc": true, + "dataset": "ase", + "validation_dataset": "ase", + "dataset_file_name": "./train_nequip.extxyz", + "validation_dataset_file_name": "./test.extxyz", + "ase_args": { + "format": "extxyz" + }, + "dataset_key_mapping": { + "forces": "forces", + "energy": "total_energy" + }, + "validation_dataset_key_mapping": { + "forces": "forces", + "energy": "total_energy" + }, + "chemical_symbols": [], + "wandb": false, + "verbose": "info", + "log_batch_freq": 10, + "log_epoch_freq": 1, + "save_checkpoint_freq": -1, + "save_ema_checkpoint_freq": -1, + "n_train": 1000, + "n_val": 1000, "learning_rate": 0.005, + "batch_size": 5, + "validation_batch_size": 10, "max_epochs": 10000, - "default_dtype": "float32" + "shuffle": true, + "metrics_key": "validation_loss", + "use_ema": true, + "ema_decay": 0.99, + "ema_use_num_updates": true, + "report_init_validation": true, + "early_stopping_patiences": { + "validation_loss": 50 + }, + "early_stopping_lower_bounds": { + "LR": 1e-05 + }, + "loss_coeffs": { + "forces": 1, + "total_energy": [ + 1, + "PerAtomMSELoss" + ] + }, + "metrics_components": [ + [ + "forces", + "mae" + ], + [ + "forces", + "rmse" + ], + [ + "forces", + "mae", + { + "PerSpecies": true, + "report_per_component": false + } + ], + [ + "forces", + "rmse", + { + "PerSpecies": true, + "report_per_component": false + } + ], + [ + "total_energy", + "mae" + ], + [ + "total_energy", + "mae", + { + "PerAtom": true + } + ] + ], + "optimizer_name": "Adam", + "optimizer_amsgrad": true, + "lr_scheduler_name": "ReduceLROnPlateau", + "lr_scheduler_patience": 100, + "lr_scheduler_factor": 0.5, + "per_species_rescale_shifts_trainable": false, + "per_species_rescale_scales_trainable": false, + "per_species_rescale_shifts": "dataset_per_atom_total_energy_mean", + "per_species_rescale_scales": "dataset_per_species_forces_rms" }, "M3GNET": { "exp_name": "training", "results_dir": "m3gnet_results", + "pretrained_model": null, + "allow_missing_labels": false, "cutoff": 5.0, "threebody_cutoff": 4.0, "batch_size": 10, "max_epochs": 1000, "include_stresses": true, - "hidden_dim": 128, - "num_units": 128, + "data_mean": 0.0, + "data_std": 1.0, + "decay_steps": 1000, + "decay_alpha": 0.96, + "dim_node_embedding": 128, + "dim_edge_embedding": 128, + "dim_state_embedding": 0, + "energy_weight": 1.0, + "element_refs": null, + "force_weight": 1.0, + "include_line_graph": true, + "loss": "mse_loss", + "loss_params": null, + "lr": 0.001, + "magmom_target": "absolute", + "magmom_weight": 0.0, "max_l": 4, "max_n": 4, - "test_equal_to_val": true + "nblocks": 3, + "optimizer": null, + "rbf_type": "Gaussian", + "scheduler": null, + "stress_weight": 0.0, + "sync_dist": false, + "is_intensive": false, + "units": 128 }, "MACE": { "model": "MACE", "name": "MACE_model", - "config_type_weights": "{'Default':1.0}", - "hidden_irreps": "128x0e + 128x1o", - "r_max": 5.0, + "amsgrad": true, "batch_size": 10, - "max_num_epochs": 1500, - "start_swa": 1200, - "ema_decay": 0.99, + "compute_avg_num_neighbors": true, + "compute_forces": true, + "config_type_weights": "{'Default':1.0}", + "compute_stress": false, + "compute_statistics": false, "correlation": 3, - "loss": "huber", "default_dtype": "float32", - "swa": true, + "device": "cpu", + "distributed": false, + "energy_weight": 1.0, "ema": true, - "amsgrad": true, - "restart_latest": true, + "ema_decay": 0.99, + "E0s": null, + "forces_weight": 100.0, + "foundation_filter_elements": true, + "foundation_model": null, + "foundation_model_readout": true, + "keep_checkpoint": false, + "keep_isolated_atoms": false, + "hidden_irreps": "128x0e + 128x1o", + "loss": "huber", + "lr": 0.001, + "multiheads_finetuning": false, + "max_num_epochs": 1500, + "pair_repulsion": false, + "patience": 2048, + "r_max": 5.0, + "restart_latest": false, "seed": 123, - "device": "cpu" + "save_cpu": true, + "save_all_checkpoints": false, + "scaling": "rms_forces_scaling", + "stress_weight": 1.0, + "start_stage_two": 1200, + "stage_two": true, + "valid_batch_size": 10, + "virials_weight": 1.0, + "wandb": false + }, + "NEP": { + "version": 4, + "type": [ + 1, + "X" + ], + "type_weight": 1.0, + "model_type": 0, + "prediction": 0, + "cutoff": [ + 6, + 5 + ], + "n_max": [ + 4, + 4 + ], + "basis_size": [ + 8, + 8 + ], + "l_max": [ + 4, + 2, + 1 + ], + "neuron": 80, + "lambda_1": 0.0, + "lambda_e": 1.0, + "lambda_f": 1.0, + "lambda_v": 0.1, + "force_delta": 0, + "batch": 1000, + "population": 60, + "generation": 100000, + "zbl": 2 } } diff --git a/configs/rss_config.yaml b/configs/rss_config.yaml old mode 100644 new mode 100755 index 57492c4d7..a516d40f7 --- a/configs/rss_config.yaml +++ b/configs/rss_config.yaml @@ -1,36 +1,23 @@ -# General Parameters -tag: 'SiO2' +tag: '' train_from_scratch: true resume_from_previous_state: test_error: pre_database_dir: mlip_path: isolated_atom_energies: - -# Buildcell Parameters generated_struct_numbers: - - 8000 - - 2000 +- 8000 +- 2000 buildcell_options: - - SPECIES: "Si%NUM=1,O%NUM=2" - NFORM: '{2,4,6,8}' - SYMMOPS: '1-4' - MINSEP: '1.5 Si-Si=2.7-3.0 Si-O=1.3-1.6 O-O=2.28-2.58' - - SPECIES: "Si%NUM=1,O%NUM=2" - NFORM: '{3,5,7}' - SYMMOPS: '1-4' - MINSEP: '1.5 Si-Si=2.7-3.0 Si-O=1.3-1.6 O-O=2.28-2.58' -fragment_file: null -fragment_numbers: null +fragment_file: +fragment_numbers: num_processes_buildcell: 128 - -# Sampling Parameters num_of_initial_selected_structs: - - 80 - - 20 +- 80 +- 20 num_of_rss_selected_structs: 100 initial_selection_enabled: true -rss_selection_method: 'bcur2i' +rss_selection_method: bcur2i bcur_params: soap_paras: l_max: 12 @@ -43,121 +30,142 @@ bcur_params: species: true frac_of_bcur: 0.8 bolt_max_num: 3000 -random_seed: null - -# DFT Labelling Parameters +random_seed: include_isolated_atom: true isolatedatom_box: - - 20.0 - - 20.0 - - 20.0 +- 20.0 +- 20.0 +- 20.0 e0_spin: false -include_dimer: false +include_dimer: true dimer_box: - - 20.0 - - 20.0 - - 20.0 +- 20.0 +- 20.0 +- 20.0 dimer_range: - - 1.0 - - 5.0 -dimer_num: 41 +- 1.0 +- 5.0 +dimer_num: 21 custom_incar: - KPAR: 8 - NCORE: 16 - LSCALAPACK: ".FALSE." - LPLANE: ".FALSE." ISMEAR: 0 - SIGMA: 0.1 - PREC: "Accurate" - ADDGRID: ".TRUE." - EDIFF: 1E-7 + SIGMA: 0.05 + PREC: Accurate + ADDGRID: .TRUE. + EDIFF: 1e-07 NELM: 250 - LWAVE: ".FALSE." - LCHARG: ".FALSE." - ALGO: "normal" - AMIX: null - LREAL: ".FALSE." + LWAVE: .FALSE. + LCHARG: .FALSE. + ALGO: Normal + AMIX: + LREAL: .FALSE. ISYM: 0 - ENCUT: 900.0 - KSPACING: 0.23 - GGA: null - AMIX_MAG: null - BMIX: null - BMIX_MAG: null - ISTART: null - LMIXTAU: null - NBANDS: null - NELMDL: null - METAGGA: "SCAN" - LASPH: ".TRUE." + ENCUT: 520.0 + KSPACING: 0.2 + GGA: + KPAR: 8 + NCORE: 16 + LSCALAPACK: .FALSE. + LPLANE: .FALSE. custom_potcar: -vasp_ref_file: 'vasp_ref.extxyz' - -# Data Preprocessing Parameters +vasp_ref_file: vasp_ref.extxyz config_types: - - 'initial' - - 'traj_early' - - 'traj' +- initial +- traj_early +- traj rss_group: - - 'traj' -test_ratio: 0.0 +- traj +test_ratio: 0.1 regularization: true -scheme: 'linear-hull' +retain_existing_sigma: false +scheme: linear-hull reg_minmax: - - [0.1, 1] - - [0.001, 0.1] - - [0.0316, 0.316] - - [0.0632, 0.632] +- - 0.1 + - 1 +- - 0.001 + - 0.1 +- - 0.0316 + - 0.316 +- - 0.0632 + - 0.632 distillation: false -force_max: null -force_label: null -pre_database_dir: null - -# MLIP Parameters -mlip_type: 'GAP' -ref_energy_name: 'REF_energy' -ref_force_name: 'REF_forces' -ref_virial_name: 'REF_virial' +force_max: +force_label: +pre_database_dir: +mlip_type: GAP +ref_energy_name: REF_energy +ref_force_name: REF_forces +ref_virial_name: REF_virial auto_delta: true num_processes_fit: 32 -device_for_fitting: 'cpu' -twob: - cutoff: 5.0 - n_sparse: 15 - theta_uniform: 1.0 - delta: 1.0 -threeb: - cutoff: 3.25 -soap: - l_max: 6 - n_max: 12 - atom_sigma: 0.5 - n_sparse: 3000 - cutoff: 5.0 - delta: 0.2 -general: - three_body: false - -# RSS Exploration Parameters -scalar_pressure_method: 'exp' -scalar_exp_pressure: 100 +device_for_fitting: cpu +scalar_pressure_method: uniform +scalar_exp_pressure: 1 scalar_pressure_exponential_width: 0.2 scalar_pressure_low: 0 scalar_pressure_high: 25 max_steps: 300 force_tol: 0.01 stress_tol: 0.01 -stop_criterion: 0.001 -max_iteration_number: 10 -num_groups: 16 +stop_criterion: 0.01 +max_iteration_number: 25 +num_groups: 6 initial_kt: 0.3 current_iter_index: 1 -hookean_repul: true +hookean_repul: false hookean_paras: - '(1, 1)': [1000, 0.6] - '(8, 1)': [1000, 0.4] - '(8, 8)': [1000, 1.0] keep_symmetry: false write_traj: true num_processes_rss: 128 -device_for_rss: 'cpu' +device_for_rss: cpu +mlip_hypers: + GAP: + general: + at_file: train.extxyz + default_sigma: '{0.0001 0.05 0.05 0}' + energy_parameter_name: REF_energy + force_parameter_name: REF_forces + virial_parameter_name: REF_virial + sparse_jitter: 1e-08 + do_copy_at_file: F + openmp_chunk_size: 10000 + gp_file: gap_file.xml + e0_offset: 0.0 + two_body: true + three_body: false + soap: true + twob: + distance_Nb_order: 2 + f0: 0.0 + add_species: T + cutoff: 5.0 + n_sparse: 15 + covariance_type: ard_se + delta: 2.0 + theta_uniform: 0.5 + sparse_method: uniform + compact_clusters: T + threeb: + distance_Nb_order: 3 + f0: 0.0 + add_species: T + cutoff: 3.25 + n_sparse: 100 + covariance_type: ard_se + delta: 2.0 + theta_uniform: 1.0 + sparse_method: uniform + compact_clusters: T + soap: + add_species: T + l_max: 10 + n_max: 12 + atom_sigma: 0.5 + zeta: 4 + cutoff: 5.0 + cutoff_transition_width: 1.0 + central_weight: 1.0 + n_sparse: 6000 + delta: 1.0 + f0: 0.0 + covariance_type: dot_product + sparse_method: cur_points diff --git a/configs/rss_config_all.yaml b/configs/rss_config_all.yaml new file mode 100755 index 000000000..cd5832e83 --- /dev/null +++ b/configs/rss_config_all.yaml @@ -0,0 +1,383 @@ +tag: '' +train_from_scratch: true +resume_from_previous_state: + test_error: + pre_database_dir: + mlip_path: + isolated_atom_energies: +generated_struct_numbers: +- 8000 +- 2000 +buildcell_options: +fragment_file: +fragment_numbers: +num_processes_buildcell: 128 +num_of_initial_selected_structs: +- 80 +- 20 +num_of_rss_selected_structs: 100 +initial_selection_enabled: true +rss_selection_method: bcur2i +bcur_params: + soap_paras: + l_max: 12 + n_max: 12 + atom_sigma: 0.0875 + cutoff: 10.5 + cutoff_transition_width: 1.0 + zeta: 4.0 + average: true + species: true + frac_of_bcur: 0.8 + bolt_max_num: 3000 +random_seed: +include_isolated_atom: true +isolatedatom_box: +- 20.0 +- 20.0 +- 20.0 +e0_spin: false +include_dimer: true +dimer_box: +- 20.0 +- 20.0 +- 20.0 +dimer_range: +- 1.0 +- 5.0 +dimer_num: 21 +custom_incar: + ISMEAR: 0 + SIGMA: 0.05 + PREC: Accurate + ADDGRID: .TRUE. + EDIFF: 1e-07 + NELM: 250 + LWAVE: .FALSE. + LCHARG: .FALSE. + ALGO: Normal + AMIX: + LREAL: .FALSE. + ISYM: 0 + ENCUT: 520.0 + KSPACING: 0.2 + GGA: + KPAR: 8 + NCORE: 16 + LSCALAPACK: .FALSE. + LPLANE: .FALSE. +custom_potcar: +vasp_ref_file: vasp_ref.extxyz +config_types: +- initial +- traj_early +- traj +rss_group: +- traj +test_ratio: 0.1 +regularization: true +retain_existing_sigma: false +scheme: linear-hull +reg_minmax: +- - 0.1 + - 1 +- - 0.001 + - 0.1 +- - 0.0316 + - 0.316 +- - 0.0632 + - 0.632 +distillation: false +force_max: +force_label: +scalar_pressure_method: uniform +scalar_exp_pressure: 1 +scalar_pressure_exponential_width: 0.2 +scalar_pressure_low: 0 +scalar_pressure_high: 25 +max_steps: 300 +force_tol: 0.01 +stress_tol: 0.01 +stop_criterion: 0.01 +max_iteration_number: 25 +num_groups: 6 +initial_kt: 0.3 +current_iter_index: 1 +hookean_repul: false +hookean_paras: +keep_symmetry: false +write_traj: true +num_processes_rss: 128 +device_for_rss: cpu +mlip_type: GAP +ref_energy_name: REF_energy +ref_force_name: REF_forces +ref_virial_name: REF_virial +auto_delta: true +num_processes_fit: 32 +device_for_fitting: cpu +pre_database_dir: +# One needs to define mlip_hypers only for the MLIP one wishes to fit +# Here comprehensive list is provided with defaults only for reference +mlip_hypers: + GAP: + general: + at_file: train.extxyz + default_sigma: '{0.0001 0.05 0.05 0}' + energy_parameter_name: REF_energy + force_parameter_name: REF_forces + virial_parameter_name: REF_virial + sparse_jitter: 1e-08 + do_copy_at_file: F + openmp_chunk_size: 10000 + gp_file: gap_file.xml + e0_offset: 0.0 + two_body: false + three_body: false + soap: true + twob: + distance_Nb_order: 2 + f0: 0.0 + add_species: T + cutoff: 5.0 + n_sparse: 15 + covariance_type: ard_se + delta: 2.0 + theta_uniform: 0.5 + sparse_method: uniform + compact_clusters: T + threeb: + distance_Nb_order: 3 + f0: 0.0 + add_species: T + cutoff: 3.25 + n_sparse: 100 + covariance_type: ard_se + delta: 2.0 + theta_uniform: 1.0 + sparse_method: uniform + compact_clusters: T + soap: + add_species: T + l_max: 10 + n_max: 12 + atom_sigma: 0.5 + zeta: 4 + cutoff: 5.0 + cutoff_transition_width: 1.0 + central_weight: 1.0 + n_sparse: 6000 + delta: 1.0 + f0: 0.0 + covariance_type: dot_product + sparse_method: cur_points + J_ACE: + order: 3 + totaldegree: 6 + cutoff: 2.0 + solver: BLR + NEQUIP: + root: results + run_name: autoplex + seed: 123 + dataset_seed: 123 + append: false + default_dtype: float32 + model_dtype: float32 + allow_tf32: true + r_max: 4.0 + num_layers: 4 + l_max: 2 + parity: true + num_features: 32 + nonlinearity_type: gate + nonlinearity_scalars: + e: silu + o: tanh + nonlinearity_gates: + e: silu + o: tanh + num_basis: 8 + besselbasis_trainable: true + polynomialcutoff_p: 5 + invariant_layers: 2 + invariant_neurons: 64 + avg_num_neighbors: auto + use_sc: true + dataset: ase + validation_dataset: ase + dataset_file_name: ./train_nequip.extxyz + validation_dataset_file_name: ./test.extxyz + ase_args: + format: extxyz + dataset_key_mapping: + forces: forces + energy: total_energy + validation_dataset_key_mapping: + forces: forces + energy: total_energy + chemical_symbols: [] + wandb: false + verbose: info + log_batch_freq: 10 + log_epoch_freq: 1 + save_checkpoint_freq: -1 + save_ema_checkpoint_freq: -1 + n_train: 1000 + n_val: 1000 + learning_rate: 0.005 + batch_size: 5 + validation_batch_size: 10 + max_epochs: 10000 + shuffle: true + metrics_key: validation_loss + use_ema: true + ema_decay: 0.99 + ema_use_num_updates: true + report_init_validation: true + early_stopping_patiences: + validation_loss: 50 + early_stopping_lower_bounds: + LR: 1e-05 + loss_coeffs: + forces: 1 + total_energy: + - 1 + - PerAtomMSELoss + metrics_components: + - - forces + - mae + - - forces + - rmse + - - forces + - mae + - PerSpecies: true + report_per_component: false + - - forces + - rmse + - PerSpecies: true + report_per_component: false + - - total_energy + - mae + - - total_energy + - mae + - PerAtom: true + optimizer_name: Adam + optimizer_amsgrad: true + lr_scheduler_name: ReduceLROnPlateau + lr_scheduler_patience: 100 + lr_scheduler_factor: 0.5 + per_species_rescale_shifts_trainable: false + per_species_rescale_scales_trainable: false + per_species_rescale_shifts: dataset_per_atom_total_energy_mean + per_species_rescale_scales: dataset_per_species_forces_rms + M3GNET: + exp_name: training + results_dir: m3gnet_results + pretrained_model: + allow_missing_labels: false + cutoff: 5.0 + threebody_cutoff: 4.0 + batch_size: 10 + max_epochs: 1000 + include_stresses: true + data_mean: 0.0 + data_std: 1.0 + decay_steps: 1000 + decay_alpha: 0.96 + dim_node_embedding: 128 + dim_edge_embedding: 128 + dim_state_embedding: 0 + energy_weight: 1.0 + element_refs: + force_weight: 1.0 + include_line_graph: true + loss: mse_loss + loss_params: + lr: 0.001 + magmom_target: absolute + magmom_weight: 0.0 + max_l: 4 + max_n: 4 + nblocks: 3 + optimizer: + rbf_type: Gaussian + scheduler: + stress_weight: 0.0 + sync_dist: false + is_intensive: false + units: 128 + MACE: + model: MACE + name: MACE_model + amsgrad: true + batch_size: 10 + compute_avg_num_neighbors: true + compute_forces: true + config_type_weights: "{'Default':1.0}" + compute_stress: false + compute_statistics: false + correlation: 3 + default_dtype: float32 + device: cpu + distributed: false + energy_weight: 1.0 + ema: true + ema_decay: 0.99 + E0s: + forces_weight: 100.0 + foundation_filter_elements: true + foundation_model: + foundation_model_readout: true + keep_checkpoint: false + keep_isolated_atoms: false + hidden_irreps: 128x0e + 128x1o + loss: huber + lr: 0.001 + multiheads_finetuning: false + max_num_epochs: 1500 + pair_repulsion: false + patience: 2048 + r_max: 5.0 + restart_latest: false + seed: 123 + save_cpu: true + save_all_checkpoints: false + scaling: rms_forces_scaling + stress_weight: 1.0 + start_swa: 1200 + swa: true + valid_batch_size: 10 + virials_weight: 1.0 + wandb: false + NEP: + version: 4 + type: + - 1 + - X + type_weight: 1.0 + model_type: 0 + prediction: 0 + cutoff: + - 6 + - 5 + n_max: + - 4 + - 4 + basis_size: + - 8 + - 8 + l_max: + - 4 + - 2 + - 1 + neuron: 80 + lambda_1: 0.0 + lambda_e: 1.0 + lambda_f: 1.0 + lambda_v: 0.1 + force_delta: 0 + batch: 1000 + population: 60 + generation: 100000 + zbl: 2 From a64bb3553c2d573330f441231d20ec6fdf072e8a Mon Sep 17 00:00:00 2001 From: naik-aakash Date: Thu, 30 Jan 2025 11:23:34 +0100 Subject: [PATCH 62/74] minor change in nequip phonon tests --- tests/auto/phonons/test_flows.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/auto/phonons/test_flows.py b/tests/auto/phonons/test_flows.py index 1ca23c324..826eca8b0 100644 --- a/tests/auto/phonons/test_flows.py +++ b/tests/auto/phonons/test_flows.py @@ -1191,7 +1191,6 @@ def test_complete_dft_vs_ml_benchmark_workflow_nequip( "batch_size": 1, "learning_rate": 0.005, "max_epochs": 1, - "default_dtype": "float32", "device": "cpu", }] ) From 7a60c5dc08e8268279cb9a0fda9b589f221b4fba Mon Sep 17 00:00:00 2001 From: naik-aakash Date: Thu, 30 Jan 2025 11:34:20 +0100 Subject: [PATCH 63/74] update buildcell install instructions --- README.md | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/README.md b/README.md index 3dc4bb923..fff8025b3 100644 --- a/README.md +++ b/README.md @@ -96,6 +96,10 @@ julia -e 'using Pkg; Pkg.Registry.add("General"); Pkg.Registry.add(Pkg.Registry. Additionally, `buildcell` as a part of `AIRSS` needs to be installed if one wants to use the RSS functionality: > ℹ️ To be able to build the AIRSS utilities one needs gcc and gfortran version 5 and above. Other compiler families (such as ifort) are not supported. +> These compilers are usually available on HPCs and one can simply load them if needed. On Ubuntu/Debian systems, one can install the necessary compilers with the following command: +````bash +apt install -y build-essential gfortran +```` ```bash curl -O https://www.mtg.msm.cam.ac.uk/files/airss-0.9.3.tgz; tar -xf airss-0.9.3.tgz; rm airss-0.9.3.tgz; cd airss; make ; make install ; make neat; cd .. From 63450bf89bbba72892a8902e5e58408b5950a15e Mon Sep 17 00:00:00 2001 From: naik-aakash Date: Thu, 30 Jan 2025 18:43:21 +0100 Subject: [PATCH 64/74] fix logic of backward compatiblity with rss config yaml --- src/autoplex/settings.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/autoplex/settings.py b/src/autoplex/settings.py index b3da9c17e..7337366e4 100644 --- a/src/autoplex/settings.py +++ b/src/autoplex/settings.py @@ -1159,7 +1159,9 @@ def from_file(cls, filename: str): for arg in config_params: mlip_type = config_params["mlip_type"].replace("-", "_") if arg in mlip_hypers.model_fields: - config_params["mlip_hypers"][mlip_type][arg] = config_params[arg] + config_params["mlip_hypers"][mlip_type].update( + {arg: config_params[arg]} + ) old_config_keys.append(arg) for key in old_config_keys: From a9dcd30df4ce32723b067d237f98decba892f092 Mon Sep 17 00:00:00 2001 From: naik-aakash Date: Thu, 30 Jan 2025 18:46:26 +0100 Subject: [PATCH 65/74] remove redundant variable creation --- src/autoplex/auto/rss/jobs.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/autoplex/auto/rss/jobs.py b/src/autoplex/auto/rss/jobs.py index 34e630021..7681cc2b2 100644 --- a/src/autoplex/auto/rss/jobs.py +++ b/src/autoplex/auto/rss/jobs.py @@ -171,6 +171,8 @@ def initial_rss( if dimer_box is None: dimer_box = [20.0, 20.0, 20.0] + print(buildcell_options) + do_randomized_structure_generation = BuildMultiRandomizedStructure( generated_struct_numbers=generated_struct_numbers, buildcell_options=buildcell_options, @@ -601,13 +603,11 @@ def do_rss_iterations( if include_dimer: include_dimer = False - (mlip_path,) = do_mlip_fit.output["mlip_path"] - do_iteration = do_rss_iterations( input={ "test_error": do_mlip_fit.output["test_error"], "pre_database_dir": do_data_preprocessing.output, - "mlip_path": mlip_path, + "mlip_path": do_mlip_fit.output["mlip_path"], "isolated_atom_energies": input["isolated_atom_energies"], "current_iter": current_iter, "kt": kt, From 0b18cfd7ffb33e146195e00b71ba649313fa1528 Mon Sep 17 00:00:00 2001 From: naik-aakash Date: Fri, 31 Jan 2025 10:13:08 +0100 Subject: [PATCH 66/74] add as_dict and from_dict method to pydantic models, address review comment --- src/autoplex/settings.py | 25 ++++++++++++++++++++++--- 1 file changed, 22 insertions(+), 3 deletions(-) diff --git a/src/autoplex/settings.py b/src/autoplex/settings.py index 7337366e4..a8ba12ba3 100644 --- a/src/autoplex/settings.py +++ b/src/autoplex/settings.py @@ -7,6 +7,8 @@ import numpy as np # noqa: TC002 from monty.serialization import loadfn +from monty.json import jsanitize, MontyDecoder +from pandas.core.window.doc import kwargs_scipy from pydantic import BaseModel, ConfigDict, Field from torch.optim import Optimizer # noqa: TC002 from torch.optim.lr_scheduler import LRScheduler # noqa: TC002 @@ -76,8 +78,25 @@ def from_file(cls, filename: str): return cls(**custom_params) + def as_dict(self): + """Return the model as a MSONable dictionary.""" + return jsanitize(self.model_copy(deep=True), strict=True, allow_bson=True, enum_values=True) -class GeneralSettings(AutoplexBaseModel): + @classmethod + def from_dict(cls, d: dict): + """Create a model from a MSONable dictionary. + + Args: + d (dict): A MSONable dictionary representation of the Model. + """ + decoded = { + k: MontyDecoder().process_decoded(v) + for k, v in d.items() + if not k.startswith("@") + } + return cls(**decoded) + +class GAPGeneralSettings(AutoplexBaseModel): """Model describing general hyperparameters for the GAP fits.""" at_file: str = Field( @@ -198,8 +217,8 @@ class SoapSettings(AutoplexBaseModel): class GAPSettings(AutoplexBaseModel): """Model describing the hyperparameters for the GAP fits for Phonons.""" - general: GeneralSettings = Field( - default_factory=GeneralSettings, + general: GAPGeneralSettings = Field( + default_factory=GAPGeneralSettings, description="General hyperparameters for the GAP fits", ) twob: TwobSettings = Field( From 484d585b246ba1aa92da88c1035bce5266ada867 Mon Sep 17 00:00:00 2001 From: naik-aakash Date: Fri, 31 Jan 2025 10:13:39 +0100 Subject: [PATCH 67/74] add as_dict and from_dict method to pydantic models, address review comment --- src/autoplex/settings.py | 1 - 1 file changed, 1 deletion(-) diff --git a/src/autoplex/settings.py b/src/autoplex/settings.py index a8ba12ba3..d2b6dc28a 100644 --- a/src/autoplex/settings.py +++ b/src/autoplex/settings.py @@ -8,7 +8,6 @@ import numpy as np # noqa: TC002 from monty.serialization import loadfn from monty.json import jsanitize, MontyDecoder -from pandas.core.window.doc import kwargs_scipy from pydantic import BaseModel, ConfigDict, Field from torch.optim import Optimizer # noqa: TC002 from torch.optim.lr_scheduler import LRScheduler # noqa: TC002 From 36d5888161b896374b469ebca55cdddd525e206d Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 31 Jan 2025 09:14:21 +0000 Subject: [PATCH 68/74] pre-commit auto-fixes --- src/autoplex/settings.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/src/autoplex/settings.py b/src/autoplex/settings.py index d2b6dc28a..81d77c483 100644 --- a/src/autoplex/settings.py +++ b/src/autoplex/settings.py @@ -6,8 +6,8 @@ from typing import Any, Literal import numpy as np # noqa: TC002 +from monty.json import MontyDecoder, jsanitize from monty.serialization import loadfn -from monty.json import jsanitize, MontyDecoder from pydantic import BaseModel, ConfigDict, Field from torch.optim import Optimizer # noqa: TC002 from torch.optim.lr_scheduler import LRScheduler # noqa: TC002 @@ -79,7 +79,9 @@ def from_file(cls, filename: str): def as_dict(self): """Return the model as a MSONable dictionary.""" - return jsanitize(self.model_copy(deep=True), strict=True, allow_bson=True, enum_values=True) + return jsanitize( + self.model_copy(deep=True), strict=True, allow_bson=True, enum_values=True + ) @classmethod def from_dict(cls, d: dict): @@ -95,6 +97,7 @@ def from_dict(cls, d: dict): } return cls(**decoded) + class GAPGeneralSettings(AutoplexBaseModel): """Model describing general hyperparameters for the GAP fits.""" From 112b1e09088c97441bc7c7e3a8e8dfda2a8f009e Mon Sep 17 00:00:00 2001 From: naik-aakash Date: Fri, 31 Jan 2025 10:34:52 +0100 Subject: [PATCH 69/74] update example config to show all options --- configs/rss_config.yaml | 82 +++++++++++++++++++-------- configs/rss_config_all.yaml | 110 ++++++++++++++++++++++++------------ 2 files changed, 132 insertions(+), 60 deletions(-) mode change 100755 => 100644 configs/rss_config_all.yaml diff --git a/configs/rss_config.yaml b/configs/rss_config.yaml index a516d40f7..304f527fe 100755 --- a/configs/rss_config.yaml +++ b/configs/rss_config.yaml @@ -1,4 +1,4 @@ -tag: '' +tag: SiO2 train_from_scratch: true resume_from_previous_state: test_error: @@ -9,6 +9,24 @@ generated_struct_numbers: - 8000 - 2000 buildcell_options: +- ABFIX: false + NFORM: '{2,4,6,8}' + SYMMOPS: 1-4 + SYSTEM: + SLACK: + OCTET: false + OVERLAP: + MINSEP: 1.5 Si-Si=2.7-3.0 Si-O=1.3-1.6 O-O=2.28-2.58 + SPECIES: Si%NUM=1,O%NUM=2 +- ABFIX: false + NFORM: '{3,5,7}' + SYMMOPS: 1-4 + SYSTEM: + SLACK: + OCTET: false + OVERLAP: + MINSEP: 1.5 Si-Si=2.7-3.0 Si-O=1.3-1.6 O-O=2.28-2.58 + SPECIES: Si%NUM=1,O%NUM=2 fragment_file: fragment_numbers: num_processes_buildcell: 128 @@ -37,7 +55,7 @@ isolatedatom_box: - 20.0 - 20.0 e0_spin: false -include_dimer: true +include_dimer: false dimer_box: - 20.0 - 20.0 @@ -45,27 +63,36 @@ dimer_box: dimer_range: - 1.0 - 5.0 -dimer_num: 21 +dimer_num: 41 custom_incar: ISMEAR: 0 - SIGMA: 0.05 + SIGMA: 0.1 PREC: Accurate ADDGRID: .TRUE. EDIFF: 1e-07 NELM: 250 LWAVE: .FALSE. LCHARG: .FALSE. - ALGO: Normal + ALGO: normal AMIX: LREAL: .FALSE. ISYM: 0 - ENCUT: 520.0 - KSPACING: 0.2 + ENCUT: 900.0 + KSPACING: 0.23 GGA: KPAR: 8 NCORE: 16 LSCALAPACK: .FALSE. LPLANE: .FALSE. + AMIX_MAG: + BMIX: + BMIX_MAG: + ISTART: + LMIXTAU: + NBANDS: + NELMDL: + METAGGA: SCAN + LASPH: .TRUE. custom_potcar: vasp_ref_file: vasp_ref.extxyz config_types: @@ -74,13 +101,13 @@ config_types: - traj rss_group: - traj -test_ratio: 0.1 +test_ratio: 0.0 regularization: true retain_existing_sigma: false scheme: linear-hull reg_minmax: - - 0.1 - - 1 + - 1.0 - - 0.001 - 0.1 - - 0.0316 @@ -98,21 +125,30 @@ ref_virial_name: REF_virial auto_delta: true num_processes_fit: 32 device_for_fitting: cpu -scalar_pressure_method: uniform -scalar_exp_pressure: 1 +scalar_pressure_method: exp +scalar_exp_pressure: 100 scalar_pressure_exponential_width: 0.2 scalar_pressure_low: 0 scalar_pressure_high: 25 max_steps: 300 force_tol: 0.01 stress_tol: 0.01 -stop_criterion: 0.01 -max_iteration_number: 25 -num_groups: 6 +stop_criterion: 0.001 +max_iteration_number: 10 +num_groups: 16 initial_kt: 0.3 current_iter_index: 1 -hookean_repul: false +hookean_repul: true hookean_paras: + (1, 1): + - 1000 + - 0.6 + (8, 1): + - 1000 + - 0.4 + (8, 8): + - 1000 + - 1.0 keep_symmetry: false write_traj: true num_processes_rss: 128 @@ -130,22 +166,22 @@ mlip_hypers: openmp_chunk_size: 10000 gp_file: gap_file.xml e0_offset: 0.0 - two_body: true + two_body: false three_body: false soap: true twob: - distance_Nb_order: 2 + distance_Nb order: 2 f0: 0.0 add_species: T cutoff: 5.0 n_sparse: 15 covariance_type: ard_se - delta: 2.0 - theta_uniform: 0.5 + delta: 1.0 + theta_uniform: 1.0 sparse_method: uniform compact_clusters: T threeb: - distance_Nb_order: 3 + distance_Nb order: 3 f0: 0.0 add_species: T cutoff: 3.25 @@ -157,15 +193,15 @@ mlip_hypers: compact_clusters: T soap: add_species: T - l_max: 10 + l_max: 6 n_max: 12 atom_sigma: 0.5 zeta: 4 cutoff: 5.0 cutoff_transition_width: 1.0 central_weight: 1.0 - n_sparse: 6000 - delta: 1.0 + n_sparse: 3000 + delta: 0.2 f0: 0.0 covariance_type: dot_product sparse_method: cur_points diff --git a/configs/rss_config_all.yaml b/configs/rss_config_all.yaml old mode 100755 new mode 100644 index cd5832e83..4b95cdfc7 --- a/configs/rss_config_all.yaml +++ b/configs/rss_config_all.yaml @@ -1,4 +1,4 @@ -tag: '' +tag: SiO2 train_from_scratch: true resume_from_previous_state: test_error: @@ -9,6 +9,24 @@ generated_struct_numbers: - 8000 - 2000 buildcell_options: +- ABFIX: false + NFORM: '{2,4,6,8}' + SYMMOPS: 1-4 + SYSTEM: + SLACK: + OCTET: false + OVERLAP: + MINSEP: 1.5 Si-Si=2.7-3.0 Si-O=1.3-1.6 O-O=2.28-2.58 + SPECIES: Si%NUM=1,O%NUM=2 +- ABFIX: false + NFORM: '{3,5,7}' + SYMMOPS: 1-4 + SYSTEM: + SLACK: + OCTET: false + OVERLAP: + MINSEP: 1.5 Si-Si=2.7-3.0 Si-O=1.3-1.6 O-O=2.28-2.58 + SPECIES: Si%NUM=1,O%NUM=2 fragment_file: fragment_numbers: num_processes_buildcell: 128 @@ -37,7 +55,7 @@ isolatedatom_box: - 20.0 - 20.0 e0_spin: false -include_dimer: true +include_dimer: false dimer_box: - 20.0 - 20.0 @@ -45,27 +63,36 @@ dimer_box: dimer_range: - 1.0 - 5.0 -dimer_num: 21 +dimer_num: 41 custom_incar: ISMEAR: 0 - SIGMA: 0.05 + SIGMA: 0.1 PREC: Accurate ADDGRID: .TRUE. EDIFF: 1e-07 NELM: 250 LWAVE: .FALSE. LCHARG: .FALSE. - ALGO: Normal + ALGO: normal AMIX: LREAL: .FALSE. ISYM: 0 - ENCUT: 520.0 - KSPACING: 0.2 + ENCUT: 900.0 + KSPACING: 0.23 GGA: KPAR: 8 NCORE: 16 LSCALAPACK: .FALSE. LPLANE: .FALSE. + AMIX_MAG: + BMIX: + BMIX_MAG: + ISTART: + LMIXTAU: + NBANDS: + NELMDL: + METAGGA: SCAN + LASPH: .TRUE. custom_potcar: vasp_ref_file: vasp_ref.extxyz config_types: @@ -74,13 +101,13 @@ config_types: - traj rss_group: - traj -test_ratio: 0.1 +test_ratio: 0.0 regularization: true retain_existing_sigma: false scheme: linear-hull reg_minmax: - - 0.1 - - 1 + - 1.0 - - 0.001 - 0.1 - - 0.0316 @@ -90,33 +117,42 @@ reg_minmax: distillation: false force_max: force_label: -scalar_pressure_method: uniform -scalar_exp_pressure: 1 +pre_database_dir: +mlip_type: GAP +ref_energy_name: REF_energy +ref_force_name: REF_forces +ref_virial_name: REF_virial +auto_delta: true +num_processes_fit: 32 +device_for_fitting: cpu +scalar_pressure_method: exp +scalar_exp_pressure: 100 scalar_pressure_exponential_width: 0.2 scalar_pressure_low: 0 scalar_pressure_high: 25 max_steps: 300 force_tol: 0.01 stress_tol: 0.01 -stop_criterion: 0.01 -max_iteration_number: 25 -num_groups: 6 +stop_criterion: 0.001 +max_iteration_number: 10 +num_groups: 16 initial_kt: 0.3 current_iter_index: 1 -hookean_repul: false +hookean_repul: true hookean_paras: + (1, 1): + - 1000 + - 0.6 + (8, 1): + - 1000 + - 0.4 + (8, 8): + - 1000 + - 1.0 keep_symmetry: false write_traj: true num_processes_rss: 128 device_for_rss: cpu -mlip_type: GAP -ref_energy_name: REF_energy -ref_force_name: REF_forces -ref_virial_name: REF_virial -auto_delta: true -num_processes_fit: 32 -device_for_fitting: cpu -pre_database_dir: # One needs to define mlip_hypers only for the MLIP one wishes to fit # Here comprehensive list is provided with defaults only for reference mlip_hypers: @@ -136,18 +172,18 @@ mlip_hypers: three_body: false soap: true twob: - distance_Nb_order: 2 + distance_Nb order: 2 f0: 0.0 add_species: T cutoff: 5.0 n_sparse: 15 covariance_type: ard_se - delta: 2.0 - theta_uniform: 0.5 + delta: 1.0 + theta_uniform: 1.0 sparse_method: uniform compact_clusters: T threeb: - distance_Nb_order: 3 + distance_Nb order: 3 f0: 0.0 add_species: T cutoff: 3.25 @@ -159,19 +195,19 @@ mlip_hypers: compact_clusters: T soap: add_species: T - l_max: 10 + l_max: 6 n_max: 12 atom_sigma: 0.5 zeta: 4 cutoff: 5.0 cutoff_transition_width: 1.0 central_weight: 1.0 - n_sparse: 6000 - delta: 1.0 + n_sparse: 3000 + delta: 0.2 f0: 0.0 covariance_type: dot_product sparse_method: cur_points - J_ACE: + J-ACE: order: 3 totaldegree: 6 cutoff: 2.0 @@ -182,8 +218,8 @@ mlip_hypers: seed: 123 dataset_seed: 123 append: false - default_dtype: float32 - model_dtype: float32 + default_dtype: float64 + model_dtype: float64 allow_tf32: true r_max: 4.0 num_layers: 4 @@ -198,8 +234,8 @@ mlip_hypers: e: silu o: tanh num_basis: 8 - besselbasis_trainable: true - polynomialcutoff_p: 5 + BesselBasis_trainable: true + PolynomialCutoff_p: 5 invariant_layers: 2 invariant_neurons: 64 avg_num_neighbors: auto @@ -345,8 +381,8 @@ mlip_hypers: save_all_checkpoints: false scaling: rms_forces_scaling stress_weight: 1.0 - start_swa: 1200 - swa: true + start_stage_two: 1200 + stage_two: true valid_batch_size: 10 virials_weight: 1.0 wandb: false From a3689208dff7bf67472e52d3f67b64edf7f6427f Mon Sep 17 00:00:00 2001 From: naik-aakash Date: Fri, 31 Jan 2025 14:47:44 +0100 Subject: [PATCH 70/74] add set option to buildcell_options>SYSTEM --- src/autoplex/settings.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/src/autoplex/settings.py b/src/autoplex/settings.py index 81d77c483..0a72d1c1e 100644 --- a/src/autoplex/settings.py +++ b/src/autoplex/settings.py @@ -876,9 +876,11 @@ class BuildcellOptions(AutoplexBaseModel): "values are (1,2,3,5,4,6,7,8,9,10,11,12,24). " "Ranges are allowed (e.g., #SYMMOPS=1-4).", ) - SYSTEM: None | Literal["Rhom", "Tric", "Mono", "Cubi", "Hexa", "Orth", "Tetra"] = ( - Field(default=None, description="Enforce a crystal system") - ) + SYSTEM: ( + None + | Literal["Rhom", "Tric", "Mono", "Cubi", "Hexa", "Orth", "Tetra"] + | set[Literal["Rhom", "Tric", "Mono", "Cubi", "Hexa", "Orth", "Tetra"]] + ) = Field(default=None, description="Enforce a crystal system") SLACK: float | None = Field(default=None, description="The slack factor") OCTET: bool = Field( default=False, From 29fd84b47820d26fa1bee873d763b05dd76781e8 Mon Sep 17 00:00:00 2001 From: naik-aakash Date: Fri, 31 Jan 2025 16:07:05 +0100 Subject: [PATCH 71/74] fix mlip_path errors --- src/autoplex/data/rss/jobs.py | 12 ++++++------ src/autoplex/data/rss/utils.py | 14 ++++++++------ 2 files changed, 14 insertions(+), 12 deletions(-) diff --git a/src/autoplex/data/rss/jobs.py b/src/autoplex/data/rss/jobs.py index b4a1e853c..fc6018bbc 100644 --- a/src/autoplex/data/rss/jobs.py +++ b/src/autoplex/data/rss/jobs.py @@ -414,7 +414,7 @@ def _parallel_process( @job def do_rss_single_node( mlip_type: str, - mlip_path: str, + mlip_path: str | list[str], iteration_index: str, structures: list[Structure], output_file_name: str = "RSS_relax_results", @@ -444,8 +444,8 @@ def do_rss_single_node( mlip_type: str Choose one specific MLIP type: 'GAP' | 'J-ACE' | 'NequIP' | 'M3GNet' | 'MACE'. - mlip_path: str - Path to the MLIP model. + mlip_path: str | list[str] + Path to the MLIP model or List of Path to the MLIP model. iteration_index: str Index for the current iteration. structures: list[Structure] @@ -522,7 +522,7 @@ def do_rss_single_node( @job def do_rss_multi_node( mlip_type: str, - mlip_path: str, + mlip_path: str | list[str], iteration_index: str, structure: list[Structure] | list[list[Structure]] | None = None, structure_paths: str | list[str] | None = None, @@ -553,8 +553,8 @@ def do_rss_multi_node( mlip_type: str Choose one specific MLIP type: 'GAP' | 'J-ACE' | 'NequIP' | 'M3GNet' | 'MACE'. - mlip_path: str - Path to the MLIP model. + mlip_path: str | list[str] + Path to the MLIP model or List of Path to the MLIP model. iteration_index: str Index for the current iteration. structure: list[Structure] diff --git a/src/autoplex/data/rss/utils.py b/src/autoplex/data/rss/utils.py index d0101f44b..4cd4029e1 100644 --- a/src/autoplex/data/rss/utils.py +++ b/src/autoplex/data/rss/utils.py @@ -321,7 +321,7 @@ def __repr__(self): def process_rss( atom: Atoms, mlip_type: str, - mlip_path: str, + mlip_path: str | list[str], output_file_name: str = "RSS_relax_results", scalar_pressure_method: str = "exp", scalar_exp_pressure: float = 100, @@ -348,8 +348,8 @@ def process_rss( mlip_type: str Choose one specific MLIP type: 'GAP' | 'J-ACE' | 'NequIP' | 'M3GNet' | 'MACE'. - mlip_path: str - Path to the MLIP model. + mlip_path: str | list + Path to the MLIP model or List of Path to the MLIP model. output_file_name: str Prefix for the trajectory/log file name. The actual output file name may be composed of this prefix, an index, and file types. @@ -395,6 +395,8 @@ def process_rss( for k, v in hookean_paras.items() } + mlip_path = mlip_path[0] if isinstance(mlip_path, list) else mlip_path + if mlip_type == "GAP": gap_label = os.path.join(mlip_path, "gap_file.xml") gap_control = "Potential xml_label=" + extract_gap_label(gap_label) @@ -554,7 +556,7 @@ def build_traj(): def minimize_structures( mlip_type: str, - mlip_path: str, + mlip_path: str | list[str], iteration_index: str, structures: list[Structure], output_file_name: str = "RSS_relax_results", @@ -583,8 +585,8 @@ def minimize_structures( mlip_type: str Choose one specific MLIP type: 'GAP' | 'J-ACE' | 'NequIP' | 'M3GNet' | 'MACE'. - mlip_path: str - Path to the MLIP model. + mlip_path: str | list[str] + Path to the MLIP model or List of path to the MLIP model. iteration_index: str Index for the current iteration. structures: list[Structure] From 71cb7b42eb2fa45e6fcb01660e8fc015bdd04010 Mon Sep 17 00:00:00 2001 From: naik-aakash Date: Fri, 31 Jan 2025 16:33:41 +0100 Subject: [PATCH 72/74] fix doc-strings --- src/autoplex/auto/rss/flows.py | 2 +- src/autoplex/auto/rss/jobs.py | 8 ++++---- src/autoplex/data/rss/jobs.py | 12 ++++++------ src/autoplex/data/rss/utils.py | 8 ++++---- 4 files changed, 15 insertions(+), 15 deletions(-) diff --git a/src/autoplex/auto/rss/flows.py b/src/autoplex/auto/rss/flows.py index 1e2d06569..e69aff28e 100644 --- a/src/autoplex/auto/rss/flows.py +++ b/src/autoplex/auto/rss/flows.py @@ -228,7 +228,7 @@ def make(self, **kwargs): - 'test_error': float, The test error of the fitted MLIP. - 'pre_database_dir': str, The directory of the latest RSS database. - - 'mlip_path': str, The path to the latest fitted MLIP. + - 'mlip_path': List of path to the latest fitted MLIP. - 'isolated_atom_energies': dict, The isolated energy values. - 'current_iter': int, The current iteration index. - 'kb_temp': float, The temperature (in eV) for Boltzmann sampling. diff --git a/src/autoplex/auto/rss/jobs.py b/src/autoplex/auto/rss/jobs.py index 7681cc2b2..6ab0cb81a 100644 --- a/src/autoplex/auto/rss/jobs.py +++ b/src/autoplex/auto/rss/jobs.py @@ -162,7 +162,7 @@ def initial_rss( - 'test_error': float, The test error of the fitted MLIP. - 'pre_database_dir': str, The directory of the preprocessed database. - - 'mlip_path': str, The path to the fitted MLIP. + - 'mlip_path': List of path to the fitted MLIP. - 'isolated_atom_energies': dict, The isolated energy values. - 'current_iter': int, The current iteration index, set to 0. """ @@ -330,8 +330,8 @@ def do_rss_iterations( The test error of the fitted MLIP. pre_database_dir: str The directory of the preprocessed database. - mlip_path: str - The path to the fitted MLIP. + mlip_path: list[str] + List of path to the fitted MLIP. isolated_atom_energies: dict The isolated energy values. current_iter: int @@ -471,7 +471,7 @@ def do_rss_iterations( - 'test_error': float, The test error of the fitted MLIP. - 'pre_database_dir': str, The directory of the preprocessed database. - - 'mlip_path': str, The path to the fitted MLIP. + - 'mlip_path': List of path to the fitted MLIP. - 'isolated_atom_energies': dict, The isolated energy values. - 'current_iter': int, The current iteration index. - 'kt': float, The temperature (in eV) for Boltzmann sampling. diff --git a/src/autoplex/data/rss/jobs.py b/src/autoplex/data/rss/jobs.py index fc6018bbc..b77fa865e 100644 --- a/src/autoplex/data/rss/jobs.py +++ b/src/autoplex/data/rss/jobs.py @@ -414,7 +414,7 @@ def _parallel_process( @job def do_rss_single_node( mlip_type: str, - mlip_path: str | list[str], + mlip_path: list[str], iteration_index: str, structures: list[Structure], output_file_name: str = "RSS_relax_results", @@ -444,8 +444,8 @@ def do_rss_single_node( mlip_type: str Choose one specific MLIP type: 'GAP' | 'J-ACE' | 'NequIP' | 'M3GNet' | 'MACE'. - mlip_path: str | list[str] - Path to the MLIP model or List of Path to the MLIP model. + mlip_path: list[str] + List of Path to the MLIP model. iteration_index: str Index for the current iteration. structures: list[Structure] @@ -522,7 +522,7 @@ def do_rss_single_node( @job def do_rss_multi_node( mlip_type: str, - mlip_path: str | list[str], + mlip_path: list[str], iteration_index: str, structure: list[Structure] | list[list[Structure]] | None = None, structure_paths: str | list[str] | None = None, @@ -553,8 +553,8 @@ def do_rss_multi_node( mlip_type: str Choose one specific MLIP type: 'GAP' | 'J-ACE' | 'NequIP' | 'M3GNet' | 'MACE'. - mlip_path: str | list[str] - Path to the MLIP model or List of Path to the MLIP model. + mlip_path: list[str] + List of Path to the MLIP model. iteration_index: str Index for the current iteration. structure: list[Structure] diff --git a/src/autoplex/data/rss/utils.py b/src/autoplex/data/rss/utils.py index 4cd4029e1..3ac5df575 100644 --- a/src/autoplex/data/rss/utils.py +++ b/src/autoplex/data/rss/utils.py @@ -348,7 +348,7 @@ def process_rss( mlip_type: str Choose one specific MLIP type: 'GAP' | 'J-ACE' | 'NequIP' | 'M3GNet' | 'MACE'. - mlip_path: str | list + mlip_path: str | list[str] Path to the MLIP model or List of Path to the MLIP model. output_file_name: str Prefix for the trajectory/log file name. The actual output file name @@ -556,7 +556,7 @@ def build_traj(): def minimize_structures( mlip_type: str, - mlip_path: str | list[str], + mlip_path: list[str], iteration_index: str, structures: list[Structure], output_file_name: str = "RSS_relax_results", @@ -585,8 +585,8 @@ def minimize_structures( mlip_type: str Choose one specific MLIP type: 'GAP' | 'J-ACE' | 'NequIP' | 'M3GNet' | 'MACE'. - mlip_path: str | list[str] - Path to the MLIP model or List of path to the MLIP model. + mlip_path: list[str] + List of path to the MLIP model. iteration_index: str Index for the current iteration. structures: list[Structure] From 5c3193663e0de50ab1712174bc3f7919cff245b6 Mon Sep 17 00:00:00 2001 From: naik-aakash Date: Fri, 31 Jan 2025 17:50:35 +0100 Subject: [PATCH 73/74] minor change to m3gnet finetuning arg name --- src/autoplex/fitting/common/utils.py | 14 +++++++++----- src/autoplex/settings.py | 5 ++++- 2 files changed, 13 insertions(+), 6 deletions(-) diff --git a/src/autoplex/fitting/common/utils.py b/src/autoplex/fitting/common/utils.py index 83adc3352..5f1bfd798 100644 --- a/src/autoplex/fitting/common/utils.py +++ b/src/autoplex/fitting/common/utils.py @@ -826,7 +826,7 @@ def m3gnet_fitting( num_workers=1, ) # train from scratch - if not m3gnet_hypers["pretrained_model"]: # train from scratch + if not m3gnet_hypers["foundation_model"]: # train from scratch model = M3GNet( element_types=train_element_types, is_intensive=m3gnet_hypers.get("is_intensive"), @@ -860,13 +860,17 @@ def m3gnet_fitting( optimizer=m3gnet_hypers.get("optimizer"), scheduler=m3gnet_hypers.get("scheduler"), ) - else: # finetune pretrained model + else: # finetune a foundation model (pretrained model) logging.info( - f"Finetuning pretrained model: {m3gnet_hypers['pretrained_model']}" + f"Finetuning foundation model: {m3gnet_hypers['foundation_model']}" ) - m3gnet_nnp = matgl.load_model(m3gnet_hypers["pretrained_model"]) + m3gnet_nnp = matgl.load_model(m3gnet_hypers["foundation_model"]) model = m3gnet_nnp.model - property_offset = m3gnet_nnp.element_refs.property_offset + property_offset = ( + m3gnet_nnp.element_refs.property_offset + if m3gnet_hypers["use_foundation_model_element_refs"] + else None + ) lit_module = PotentialLightningModule( model=model, element_refs=property_offset, diff --git a/src/autoplex/settings.py b/src/autoplex/settings.py index 0a72d1c1e..b9ac8598b 100644 --- a/src/autoplex/settings.py +++ b/src/autoplex/settings.py @@ -496,7 +496,7 @@ class M3GNETSettings(AutoplexBaseModel): results_dir: str = Field( default="m3gnet_results", description="Directory to save the results" ) - pretrained_model: str | None = Field( + foundation_model: str | None = Field( default=None, description="Pretrained model. Can be a Path to locally stored model " "or name of pretrained PES model available in the " @@ -506,6 +506,9 @@ class M3GNETSettings(AutoplexBaseModel): "access to be able to download the model." "If None, the model will be trained from scratch.", ) + use_foundation_model_element_refs: bool = Field( + default=False, description="Use element refs from the foundation model" + ) allow_missing_labels: bool = Field( default=False, description="Allow missing labels" ) From cc538b72f84f05505c8a3e581c3724c14fb55910 Mon Sep 17 00:00:00 2001 From: naik-aakash Date: Fri, 31 Jan 2025 17:52:42 +0100 Subject: [PATCH 74/74] adapt m3gnet finetuning test arg --- tests/auto/phonons/test_flows.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/auto/phonons/test_flows.py b/tests/auto/phonons/test_flows.py index 826eca8b0..a82bd2779 100644 --- a/tests/auto/phonons/test_flows.py +++ b/tests/auto/phonons/test_flows.py @@ -951,7 +951,8 @@ def test_complete_dft_vs_ml_benchmark_workflow_m3gnet_finetuning( "include_stresses": True, "device": "cpu", "test_equal_to_val": True, - "pretrained_model": "M3GNet-MP-2021.2.8-DIRECT-PES", + "foundation_model": "M3GNet-MP-2021.2.8-DIRECT-PES", + "use_foundation_model_element_refs": True, }] )