diff --git a/.github/workflows/python-package.yml b/.github/workflows/python-package.yml index 31bb5e449..ec86bc127 100644 --- a/.github/workflows/python-package.yml +++ b/.github/workflows/python-package.yml @@ -114,6 +114,8 @@ jobs: title: Update test durations file body: Auto update test durations file branch: update-test-durations + labels: | + automated-pr delete-branch: true base: main diff --git a/README.md b/README.md index bd93de9c5..fff8025b3 100644 --- a/README.md +++ b/README.md @@ -95,6 +95,12 @@ 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 .. ``` diff --git a/configs/create_configs.py b/configs/create_configs.py new file mode 100644 index 000000000..0385f470e --- /dev/null +++ b/configs/create_configs.py @@ -0,0 +1,11 @@ +"""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" +) # rss_config_all.yaml contains list of all hypers and supported MLIPs diff --git a/configs/mlip_hypers.json b/configs/mlip_hypers.json new file mode 100644 index 000000000..ba02c7b11 --- /dev/null +++ b/configs/mlip_hypers.json @@ -0,0 +1,306 @@ +{ + "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": "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, + "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": null, + "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": 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, + "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", + "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": 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, + "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 new file mode 100755 index 000000000..304f527fe --- /dev/null +++ b/configs/rss_config.yaml @@ -0,0 +1,207 @@ +tag: SiO2 +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: +- 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 +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: false +dimer_box: +- 20.0 +- 20.0 +- 20.0 +dimer_range: +- 1.0 +- 5.0 +dimer_num: 41 +custom_incar: + ISMEAR: 0 + SIGMA: 0.1 + PREC: Accurate + ADDGRID: .TRUE. + EDIFF: 1e-07 + NELM: 250 + LWAVE: .FALSE. + LCHARG: .FALSE. + ALGO: normal + AMIX: + LREAL: .FALSE. + ISYM: 0 + 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: +- initial +- traj_early +- traj +rss_group: +- traj +test_ratio: 0.0 +regularization: true +retain_existing_sigma: false +scheme: linear-hull +reg_minmax: +- - 0.1 + - 1.0 +- - 0.001 + - 0.1 +- - 0.0316 + - 0.316 +- - 0.0632 + - 0.632 +distillation: false +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 +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.001 +max_iteration_number: 10 +num_groups: 16 +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 +keep_symmetry: false +write_traj: true +num_processes_rss: 128 +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: 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 diff --git a/configs/rss_config_all.yaml b/configs/rss_config_all.yaml new file mode 100644 index 000000000..4b95cdfc7 --- /dev/null +++ b/configs/rss_config_all.yaml @@ -0,0 +1,419 @@ +tag: SiO2 +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: +- 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 +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: false +dimer_box: +- 20.0 +- 20.0 +- 20.0 +dimer_range: +- 1.0 +- 5.0 +dimer_num: 41 +custom_incar: + ISMEAR: 0 + SIGMA: 0.1 + PREC: Accurate + ADDGRID: .TRUE. + EDIFF: 1e-07 + NELM: 250 + LWAVE: .FALSE. + LCHARG: .FALSE. + ALGO: normal + AMIX: + LREAL: .FALSE. + ISYM: 0 + 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: +- initial +- traj_early +- traj +rss_group: +- traj +test_ratio: 0.0 +regularization: true +retain_existing_sigma: false +scheme: linear-hull +reg_minmax: +- - 0.1 + - 1.0 +- - 0.001 + - 0.1 +- - 0.0316 + - 0.316 +- - 0.0632 + - 0.632 +distillation: false +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 +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.001 +max_iteration_number: 10 +num_groups: 16 +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 +keep_symmetry: false +write_traj: true +num_processes_rss: 128 +device_for_rss: cpu +# 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: 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 + 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 + 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_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/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/installation/installation.md b/docs/user/installation/installation.md index b787f6312..572d9e87d 100644 --- a/docs/user/installation/installation.md +++ b/docs/user/installation/installation.md @@ -123,6 +123,11 @@ Both packages rely on the [MongoDB](https://www.mongodb.com/) database manager f We recommend using `jobflow-remote` as it is more flexible to use, especially on clusters where users cannot store their own MongoDB. You can find a more comprehensive `jobflow-remote` tutorial [here](../jobflowremote.md). +> ℹ️ These workflow managers are not included in the standard installation by default, to install autoplex with workflow managers please install `autoplex` using `pip install autoplex[strict,workflow-managers]` + +> ℹ️ If using fireworks to manage your jobs, additionally please update `pymongo` package to v4.11, using +`pip install --upgrade pymongo==4.11` + Submission using `FireWorks`: ```python from fireworks import LaunchPad 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/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 91e78ebe0..e134aa7a9 100644 --- a/docs/user/rss/flow/input/input.md +++ b/docs/user/rss/flow/input/input.md @@ -166,33 +166,66 @@ 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 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 bc38c3fd7..996ae05c1 100644 --- a/docs/user/rss/flow/quick_start/start.md +++ b/docs/user/rss/flow/quick_start/start.md @@ -24,19 +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 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 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 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') +rss_config = RssConfig.from_file('path/to/your/config.yaml') + +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) @@ -46,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").make(config_file='path/to/your/name.yaml') +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")) ``` @@ -58,7 +65,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', @@ -77,18 +85,26 @@ 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', - 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. 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 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 .. ``` diff --git a/pyproject.toml b/pyproject.toml index 13b7bc3dc..1ce0de8f2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -45,9 +45,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 = [ diff --git a/src/autoplex/__init__.py b/src/autoplex/__init__.py index 0aec76f9f..58c8b7c24 100644 --- a/src/autoplex/__init__.py +++ b/src/autoplex/__init__.py @@ -8,3 +8,22 @@ """ from autoplex._version import __version__ +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/auto/phonons/flows.py b/src/autoplex/auto/phonons/flows.py index 0000371b0..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 ( @@ -19,6 +18,7 @@ MPStaticSet, ) +from autoplex import MLIP_HYPERS from autoplex.auto.phonons.jobs import ( complete_benchmark, dft_phonopy_gen_data, @@ -33,10 +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, - load_mlip_hyperparameter_defaults, -) __all__ = [ "CompleteDFTvsMLBenchmarkWorkflow", @@ -166,16 +162,12 @@ 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 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). @@ -224,17 +216,16 @@ 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" - use_defaults_fitting: bool = True run_fits_on_different_cluster: bool = False 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, @@ -252,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. @@ -279,9 +272,8 @@ def make( fit_input = {} bm_outputs = [] - default_hyperparameters = load_mlip_hyperparameter_defaults( - mlip_fit_parameter_file_path=self.path_to_hyperparameters - ) + hyperparameters = hyperparameters.model_copy(deep=True) + default_hyperparameters = hyperparameters.model_dump(by_alias=True) soap_default_dict = next( ( @@ -389,12 +381,10 @@ 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, 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, @@ -408,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, ) @@ -479,7 +470,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, @@ -490,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, ) @@ -824,8 +815,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/auto/rss/flows.py b/src/autoplex/auto/rss/flows.py index 9d2684f32..e69aff28e 100644 --- a/src/autoplex/auto/rss/flows.py +++ b/src/autoplex/auto/rss/flows.py @@ -1,13 +1,11 @@ """RSS (random structure searching) flow for exploring and learning potential energy surfaces from scratch.""" -import os -from dataclasses import dataclass -from pathlib import Path +from dataclasses import dataclass, field from jobflow import Flow, Maker, Response, job -from ruamel.yaml import YAML from autoplex.auto.rss.jobs import do_rss_iterations, initial_rss +from autoplex.settings import RssConfig @dataclass @@ -19,24 +17,21 @@ 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. + 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" - path_to_default_config_parameters: Path | str | None = None + rss_config: RssConfig = field(default_factory=lambda: RssConfig()) @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. @@ -65,7 +60,7 @@ def make(self, config_file: str | None = None, **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 @@ -137,13 +132,16 @@ def make(self, config_file: str | None = None, **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 @@ -230,30 +228,23 @@ def make(self, config_file: str | None = None, **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. """ - 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) + default_config = self.rss_config.model_copy(deep=True) + if kwargs: + default_config.update_parameters(kwargs) - 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_params = default_config.model_dump(by_alias=True, exclude_none=True) - config.update(kwargs) - self._process_hookean_paras(config) + # 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) - config_params = config.copy() + self._process_hookean_paras(config_params) if "train_from_scratch" not in config_params: raise ValueError( @@ -350,8 +341,9 @@ def make(self, config_file: str | None = None, **kwargs): return Response(replace=Flow(rss_flow), output=do_rss_job.output) - def _process_hookean_paras(self, config): - if "hookean_paras" in config: + @staticmethod + def _process_hookean_paras(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/auto/rss/jobs.py b/src/autoplex/auto/rss/jobs.py index 60ed679f5..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. """ @@ -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, @@ -222,8 +224,8 @@ def initial_rss( auto_delta=auto_delta, glue_xml=False, ).make( - database_dir=do_data_preprocessing.output, isolated_atom_energies=do_data_collection.output["isolated_atom_energies"], + database_dir=do_data_preprocessing.output, device=device_for_fitting, **fit_kwargs, ) @@ -236,14 +238,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" ], @@ -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. @@ -603,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, diff --git a/src/autoplex/data/rss/jobs.py b/src/autoplex/data/rss/jobs.py index b4a1e853c..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, + 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 - 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, + 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 - 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 d0101f44b..3ac5df575 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[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 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: 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: list[str] + List of path to the MLIP model. iteration_index: str Index for the current iteration. structures: list[Structure] diff --git a/src/autoplex/fitting/common/flows.py b/src/autoplex/fitting/common/flows.py index 50172c395..c98ed3231 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 @@ -84,10 +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 If true, run fits on different clusters. """ @@ -106,7 +101,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. @@ -114,13 +108,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, device: str = "cpu", @@ -133,6 +127,10 @@ 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 List of element names (strings) involved in the training dataset isolated_atom_energies: dict @@ -177,10 +175,10 @@ 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, - use_defaults=self.use_defaults, device=device, species_list=species_list, database_dict=data_prep_job.output["database_dict"], @@ -199,11 +197,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, @@ -211,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 69d7c4e32..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, @@ -14,16 +15,12 @@ nequip_fitting, ) -current_dir = Path(__file__).absolute().parent -GAP_DEFAULTS_FILE_PATH = current_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, @@ -33,10 +30,10 @@ 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, + hyperparameters: MLIP_HYPERS = MLIP_HYPERS, **fit_kwargs, ): """ @@ -48,10 +45,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 @@ -71,8 +66,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 @@ -80,6 +73,11 @@ 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. fit_kwargs: dict Additional keyword arguments for MLIP fitting. """ @@ -127,7 +125,7 @@ def machine_learning_fit( ).exists(): train_test_error = gap_fitting( db_dir=database_dir, - path_to_hyperparameters=path_to_hyperparameters, + hyperparameters=hyperparameters.GAP, species_list=species_list, num_processes_fit=num_processes_fit, auto_delta=auto_delta, @@ -145,7 +143,7 @@ def machine_learning_fit( elif mlip_type == "J-ACE": train_test_error = jace_fitting( db_dir=database_dir, - path_to_hyperparameters=path_to_hyperparameters, + hyperparameters=hyperparameters.J_ACE, isolated_atom_energies=isolated_atom_energies, ref_energy_name=ref_energy_name, ref_force_name=ref_force_name, @@ -158,7 +156,7 @@ def machine_learning_fit( elif mlip_type == "NEQUIP": train_test_error = nequip_fitting( db_dir=database_dir, - path_to_hyperparameters=path_to_hyperparameters, + hyperparameters=hyperparameters.NEQUIP, isolated_atom_energies=isolated_atom_energies, ref_energy_name=ref_energy_name, ref_force_name=ref_force_name, @@ -171,7 +169,7 @@ def machine_learning_fit( elif mlip_type == "M3GNET": train_test_error = m3gnet_fitting( db_dir=database_dir, - path_to_hyperparameters=path_to_hyperparameters, + hyperparameters=hyperparameters.M3GNET, ref_energy_name=ref_energy_name, ref_force_name=ref_force_name, ref_virial_name=ref_virial_name, @@ -183,11 +181,10 @@ def machine_learning_fit( elif mlip_type == "MACE": train_test_error = mace_fitting( db_dir=database_dir, - path_to_hyperparameters=path_to_hyperparameters, + hyperparameters=hyperparameters.MACE, 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/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" - } -} diff --git a/src/autoplex/fitting/common/utils.py b/src/autoplex/fitting/common/utils.py index 3bd7b89fc..5f1bfd798 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 @@ -15,6 +14,7 @@ import ase import lightning as pl +import matgl import matplotlib.pyplot as plt import numpy as np import pandas as pd @@ -33,6 +33,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 @@ -41,6 +42,7 @@ from scipy.spatial import ConvexHull from scipy.special import comb +from autoplex import GAP_HYPERS, JACE_HYPERS, M3GNET_HYPERS, MACE_HYPERS, NEQUIP_HYPERS from autoplex.data.common.utils import ( data_distillation, plot_energy_forces, @@ -48,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" ) @@ -59,7 +58,7 @@ def gap_fitting( db_dir: Path, species_list: list | None = None, - path_to_hyperparameters: Path | str = MLIP_PHONON_DEFAULTS_FILE_PATH, + hyperparameters: GAP_HYPERS = GAP_HYPERS, num_processes_fit: int = 32, auto_delta: bool = True, glue_xml: bool = False, @@ -80,8 +79,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 +109,7 @@ 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 + 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") @@ -124,22 +122,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_parameters( + { + "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_parameters(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 +269,7 @@ def gap_fitting( ) def jace_fitting( db_dir: str | Path, - path_to_hyperparameters: Path | str = MLIP_RSS_DEFAULTS_FILE_PATH, + hyperparameters: JACE_HYPERS = JACE_HYPERS, isolated_atom_energies: dict | None = None, ref_energy_name: str = "REF_energy", ref_force_name: str = "REF_forces", @@ -290,8 +288,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 + Fit hyperparameters. isolated_atom_energies: dict: mandatory dictionary mapping element numbers to isolated energies. ref_energy_name : str, optional @@ -327,8 +325,7 @@ 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 + 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, ".") @@ -361,20 +358,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_parameters(fit_kwargs) + + jace_hypers = hyperparameters.model_dump(by_alias=True) order = jace_hypers["order"] totaldegree = jace_hypers["totaldegree"] @@ -455,7 +442,7 @@ def jace_fitting( def nequip_fitting( db_dir: Path, - path_to_hyperparameters: Path | str = MLIP_RSS_DEFAULTS_FILE_PATH, + hyperparameters: NEQUIP_HYPERS = NEQUIP_HYPERS, isolated_atom_energies: dict | None = None, ref_energy_name: str = "REF_energy", ref_force_name: str = "REF_forces", @@ -474,8 +461,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 @@ -525,8 +512,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"] @@ -547,150 +534,29 @@ 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_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: - 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) - 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 - """ + nequip_hypers = hyperparameters.model_dump(by_alias=True) - 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( @@ -736,11 +602,12 @@ def nequip_fitting( def m3gnet_fitting( db_dir: Path, - path_to_hyperparameters: Path | str = MLIP_RSS_DEFAULTS_FILE_PATH, + hyperparameters: M3GNET_HYPERS = M3GNET_HYPERS, device: str = "cuda", 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: """ @@ -750,8 +617,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 @@ -760,9 +627,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 ----------------- @@ -780,16 +648,16 @@ 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. 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 ------- @@ -806,36 +674,15 @@ 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 - default_hyperparameters = load_mlip_hyperparameter_defaults( - mlip_fit_parameter_file_path=path_to_hyperparameters - ) - - m3gnet_hypers = default_hyperparameters["M3GNET"] + 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"] - 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) @@ -875,7 +722,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, @@ -892,10 +739,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, @@ -919,7 +766,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, @@ -932,10 +779,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, @@ -975,21 +822,77 @@ 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["foundation_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, + 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"), + 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 a foundation model (pretrained model) + logging.info( + f"Finetuning foundation model: {m3gnet_hypers['foundation_model']}" + ) + m3gnet_nnp = matgl.load_model(m3gnet_hypers["foundation_model"]) + model = m3gnet_nnp.model + 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, + 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")) # Inference mode = False is required for calculating forces, stress in test mode and prediction mode if device == "cuda": @@ -997,7 +900,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, @@ -1006,7 +909,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, @@ -1110,12 +1013,11 @@ def m3gnet_fitting( def mace_fitting( db_dir: Path, - path_to_hyperparameters: Path | str = MLIP_RSS_DEFAULTS_FILE_PATH, + hyperparameters: MACE_HYPERS = MACE_HYPERS, device: str = "cuda", 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: """ @@ -1129,8 +1031,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 @@ -1171,26 +1073,17 @@ 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" ) - if use_defaults: - default_hyperparameters = load_mlip_hyperparameter_defaults( - mlip_fit_parameter_file_path=path_to_hyperparameters - ) + hyperparameters.update_parameters(fit_kwargs) - mace_hypers = default_hyperparameters["MACE"] - else: - mace_hypers = {} - - # 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) boolean_hypers = [ "distributed", @@ -1307,24 +1200,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, diff --git a/src/autoplex/settings.py b/src/autoplex/settings.py new file mode 100644 index 000000000..b9ac8598b --- /dev/null +++ b/src/autoplex/settings.py @@ -0,0 +1,1196 @@ +"""Settings for autoplex.""" + +from __future__ import annotations + +import logging +from typing import Any, Literal + +import numpy as np # noqa: TC002 +from monty.json import MontyDecoder, jsanitize +from monty.serialization import loadfn +from pydantic import BaseModel, ConfigDict, Field +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" +) + +__all__ = [ + "AutoplexBaseModel", + "GAPSettings", + "JACESettings", + "M3GNETSettings", + "MACESettings", + "MLIPHypers", + "NEPSettings", + "NEQUIPSettings", + "RssConfig", +] + + +class AutoplexBaseModel(BaseModel): + """Base class for all models in autoplex.""" + + model_config = ConfigDict( + validate_assignment=True, + protected_namespaces=(), + extra="allow", + arbitrary_types_allowed=True, + ) + + def update_parameters(self, updates: dict[str, Any]): + """ + Update the default parameters of the model instance, including nested fields. + + Args: + 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, self.__class__) and isinstance(value, dict): + # Update nested model + field_value.update_parameters( + value + ) # Recursively call update_parameters + else: + # Update field value + setattr(self, key, value) + + 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): + """ + 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) + + 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 + ) + + @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( + 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(AutoplexBaseModel): + """Model describing two body hyperparameters for the GAP fits.""" + + distance_Nb_order: int = Field( + 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 | 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" + ) + 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(AutoplexBaseModel): + """Model describing threebody hyperparameters for the GAP fits.""" + + distance_Nb_order: int = Field( + 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 | 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" + ) + 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(AutoplexBaseModel): + """Model describing soap hyperparameters for the GAP fits.""" + + 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(AutoplexBaseModel): + """Model describing the hyperparameters for the GAP fits for Phonons.""" + + general: GAPGeneralSettings = Field( + default_factory=GAPGeneralSettings, + 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(AutoplexBaseModel): + """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 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="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 " + "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") + 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 hidden neurons in radial function, smaller is faster", + ) + 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") + 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_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", + ) + + +class M3GNETSettings(AutoplexBaseModel): + """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" + ) + 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 " + "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.", + ) + 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" + ) + 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" + ) + 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" + ) + 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" + ) + dim_edge_embedding: int = Field( + default=128, description="Dimension of edge embedding" + ) + dim_state_embedding: int = Field( + 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" + ) + 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" + ) + units: int = Field(default=128, description="Number of neurons in each MLP layer") + + +class MACESettings(AutoplexBaseModel): + """Model describing the hyperparameters for the MACE fits.""" + + 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="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" + ) + 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" + ) + 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=False, description="Whether to restart the latest model" + ) + seed: int = Field(default=123, description="Seed for the random number generator") + 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(AutoplexBaseModel): + """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(AutoplexBaseModel): + """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", + alias="J-ACE", + ) + 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(AutoplexBaseModel): + """ + A model describing the state information. + + 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. + """ + + 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(AutoplexBaseModel): + """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(AutoplexBaseModel): + """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(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=" 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"] + | 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, + 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") + + +class CustomIncar(AutoplexBaseModel): + """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 RssConfig(AutoplexBaseModel): + """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] | 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.", + ) + 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" + ) + 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" + ) + 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." + ) + 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].update( + {arg: config_params[arg]} + ) + old_config_keys.append(arg) + + for key in old_config_keys: + del config_params[key] + + return cls(**config_params) diff --git a/tests/auto/phonons/test_flows.py b/tests/auto/phonons/test_flows.py index 4ec5084b4..a82bd2779 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", @@ -926,6 +927,51 @@ 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, + "foundation_model": "M3GNet-MP-2021.2.8-DIRECT-PES", + "use_foundation_model_element_refs": True, + }] + ) + + # 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 @@ -994,7 +1040,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=False, ).make( structure_list=[structure], mp_ids=["test"], @@ -1064,7 +1109,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=False, split_ratio=0.3, ).make( structure_list=[structure], @@ -1148,7 +1192,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", }] ) diff --git a/tests/auto/phonons/test_jobs.py b/tests/auto/phonons/test_jobs.py index 06f8dd00a..eed556191 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/auto/rss/test_flows.py b/tests/auto/rss/test_flows.py index da0093cbc..fc7c62f1f 100644 --- a/tests/auto/rss/test_flows.py +++ b/tests/auto/rss/test_flows.py @@ -1,8 +1,10 @@ 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.settings import RssConfig +from autoplex.auto.rss.flows import RssMaker os.environ["OMP_NUM_THREADS"] = "1" @@ -307,3 +309,19 @@ 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_file(test_dir): + + config_model = RssConfig.from_file(test_dir / "rss" / "rss_config.yaml") + + # Test if config is updated as expected + rss = RssMaker(rss_config=config_model) + + 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] + 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..22af90025 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 ) @@ -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", @@ -379,7 +380,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 +413,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..15762c270 100644 --- a/tests/fitting/common/test_jobs.py +++ b/tests/fitting/common/test_jobs.py @@ -11,13 +11,13 @@ 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( + _ = run_locally( gapfit, ensure_success=True, create_folders=True, store=memory_jobstore ) @@ -32,14 +32,14 @@ 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, ) - responses = run_locally( + _ = run_locally( jacefit, ensure_success=True, create_folders=True, store=memory_jobstore ) @@ -53,15 +53,15 @@ 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, device="cpu", ) - responses = run_locally( + _ = run_locally( nequipfit, ensure_success=True, create_folders=True, store=memory_jobstore ) @@ -71,31 +71,32 @@ 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, - database_dir=database_dir, ).make( + database_dir=database_dir, isolated_atom_energies={14: -0.84696938}, cutoff=3.0, threebody_cutoff=2.0, 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", test_equal_to_val=True, ) - responses = run_locally( - nequipfit, ensure_success=True, create_folders=True, store=memory_jobstore + _ = run_locally( + 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): @@ -105,8 +106,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}', @@ -122,7 +123,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 ) @@ -137,11 +138,10 @@ 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, - database_dir=database_dir, ).make( + database_dir=database_dir, name="MACE_final", foundation_model="small", multiheads_finetuning=False, @@ -167,7 +167,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 ) @@ -182,11 +182,10 @@ 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, - database_dir=database_dir, ).make( + database_dir=database_dir, name="MACE_final", foundation_model="small", multiheads_finetuning=False, @@ -212,7 +211,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 ) diff --git a/tests/fitting/common/test_utils.py b/tests/fitting/common/test_utils.py index 0ca145761..8ad176a95 100644 --- a/tests/fitting/common/test_utils.py +++ b/tests/fitting/common/test_utils.py @@ -1,12 +1,11 @@ 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, prepare_fit_environment, - MLIP_PHONON_DEFAULTS_FILE_PATH ) def test_stratified_split(test_dir): @@ -21,11 +20,9 @@ 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 = 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, @@ -57,11 +54,8 @@ 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 = 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, @@ -91,9 +85,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_parameters({"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 +114,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 +125,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 = MLIP_HYPERS.GAP.model_copy(deep=True) + gap_hyper_parameter_dict = gap_hyper_parameter.model_dump(by_alias=True) # three_body_disable @@ -193,11 +188,8 @@ 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 = MLIP_HYPERS.GAP.model_copy(deep=True) + gap_hyper_parameter_dict = gap_hyper_parameter.model_dump(by_alias=True) # check with only soap enabled diff --git a/src/autoplex/auto/rss/rss_default_configuration.yaml b/tests/test_data/rss/rss_config.yaml similarity index 93% rename from src/autoplex/auto/rss/rss_default_configuration.yaml rename to tests/test_data/rss/rss_config.yaml index 89cdc713a..3eca4f9a7 100644 --- a/src/autoplex/auto/rss/rss_default_configuration.yaml +++ b/tests/test_data/rss/rss_config.yaml @@ -1,5 +1,5 @@ # General Parameters -tag: +tag: "test" train_from_scratch: true resume_from_previous_state: test_error: @@ -9,12 +9,12 @@ resume_from_previous_state: # Buildcell Parameters generated_struct_numbers: - - 8000 - - 2000 + - 9000 + - 1000 buildcell_options: fragment_file: null fragment_numbers: null -num_processes_buildcell: 128 +num_processes_buildcell: 64 # Sampling Parameters num_of_initial_selected_structs: @@ -40,15 +40,15 @@ random_seed: null # DFT Labelling Parameters include_isolated_atom: true isolatedatom_box: - - 20.0 - - 20.0 - - 20.0 + - 10.0 + - 10.0 + - 10.0 e0_spin: false include_dimer: true dimer_box: - - 20.0 - - 20.0 - - 20.0 + - 10.0 + - 10.0 + - 10.0 dimer_range: - 1.0 - 5.0 @@ -102,7 +102,7 @@ ref_energy_name: 'REF_energy' ref_force_name: 'REF_forces' ref_virial_name: 'REF_virial' auto_delta: true -num_processes_fit: 32 +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. @@ -140,4 +140,4 @@ hookean_paras: keep_symmetry: false write_traj: true num_processes_rss: 128 -device_for_rss: 'cpu' +device_for_rss: 'cuda'