Adapt Makers to use pydantic models instead of yamls

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')

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,24 +21,32 @@ 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 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.
     """
 
     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."""
+        if self.config_file and os.path.exists(self.config_file):
+            new_config = loadfn(self.config_file)
+
+            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
 
     @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.
 
@@ -235,25 +245,10 @@ def make(self, config_file: str | None = None, **kwargs):
             - 'current_iter': int, The current iteration index.
             - '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(

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:

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]
+
+
+

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'