added architecture

src/metatrain/experimental/alchemical_model/default-hypers.yaml

@@ -1,29 +1,30 @@
-name: experimental.alchemical_model
+architecture:
+  name: experimental.alchemical_model
 
-model:
-  soap:
-    num_pseudo_species: 4
-    cutoff: 5.0
-    basis_cutoff_power_spectrum: 400
-    radial_basis_type: "physical"
-    basis_scale: 3.0
-    trainable_basis: true
-    normalize: true
-    contract_center_species: true
-  bpnn:
-    hidden_sizes: [32, 32]
-    output_size: 1
-  zbl: false
+  model:
+    soap:
+      num_pseudo_species: 4
+      cutoff: 5.0
+      basis_cutoff_power_spectrum: 400
+      radial_basis_type: "physical"
+      basis_scale: 3.0
+      trainable_basis: true
+      normalize: true
+      contract_center_species: true
+    bpnn:
+      hidden_sizes: [32, 32]
+      output_size: 1
+    zbl: false
 
-training:
-  batch_size: 8
-  num_epochs: 100
-  learning_rate: 0.001
-  early_stopping_patience: 200
-  scheduler_patience: 100
-  scheduler_factor: 0.8
-  log_interval: 5
-  checkpoint_interval: 25
-  per_structure_targets: []
-  loss_weights: {}
-  log_mae: False
+  training:
+    batch_size: 8
+    num_epochs: 100
+    learning_rate: 0.001
+    early_stopping_patience: 200
+    scheduler_patience: 100
+    scheduler_factor: 0.8
+    log_interval: 5
+    checkpoint_interval: 25
+    per_structure_targets: []
+    loss_weights: {}
+    log_mae: False

src/metatrain/experimental/gap/default-hypers.yaml

@@ -1,27 +1,28 @@
-name: experimental.gap
+architecture:
+  name: experimental.gap
 
-model:
-  soap:
-    cutoff: 5.0
-    max_radial: 8
-    max_angular: 6
-    atomic_gaussian_width: 0.3
-    radial_basis:
-      Gto: {}
-    center_atom_weight: 1.0
-    cutoff_function:
-      ShiftedCosine:
-        width: 1.0
-    radial_scaling:
-      Willatt2018:
-        rate: 1.0
-        scale: 2.0
-        exponent: 7.0
-  krr:
-    degree: 2
-    num_sparse_points: 500
-  zbl: false
+  model:
+    soap:
+      cutoff: 5.0
+      max_radial: 8
+      max_angular: 6
+      atomic_gaussian_width: 0.3
+      radial_basis:
+        Gto: {}
+      center_atom_weight: 1.0
+      cutoff_function:
+        ShiftedCosine:
+          width: 1.0
+      radial_scaling:
+        Willatt2018:
+          rate: 1.0
+          scale: 2.0
+          exponent: 7.0
+    krr:
+      degree: 2
+      num_sparse_points: 500
+    zbl: false
 
-training:
-  regularizer: 0.001
-  regularizer_forces: null
+  training:
+    regularizer: 0.001
+    regularizer_forces: null

src/metatrain/experimental/pet/default-hypers.yaml

@@ -1,59 +1,60 @@
-name: experimental.pet
+architecture:
+  name: experimental.pet
 
-model:
-  CUTOFF_DELTA: 0.2
-  AVERAGE_POOLING: False
-  TRANSFORMERS_CENTRAL_SPECIFIC: False
-  HEADS_CENTRAL_SPECIFIC: False
-  ADD_TOKEN_FIRST: True
-  ADD_TOKEN_SECOND: True
-  N_GNN_LAYERS: 3
-  TRANSFORMER_D_MODEL: 128
-  TRANSFORMER_N_HEAD: 4
-  TRANSFORMER_DIM_FEEDFORWARD: 512
-  HEAD_N_NEURONS: 128
-  N_TRANS_LAYERS: 3
-  ACTIVATION: silu
-  USE_LENGTH: True
-  USE_ONLY_LENGTH: False
-  R_CUT: 5.0
-  R_EMBEDDING_ACTIVATION: False
-  COMPRESS_MODE: mlp
-  BLEND_NEIGHBOR_SPECIES: False
-  AVERAGE_BOND_ENERGIES: False
-  USE_BOND_ENERGIES: True
-  USE_ADDITIONAL_SCALAR_ATTRIBUTES: False
-  SCALAR_ATTRIBUTES_SIZE: null
-  TRANSFORMER_TYPE: PostLN # PostLN or PreLN
-  USE_LONG_RANGE: False
-  K_CUT: null # should be float; only used when USE_LONG_RANGE is True
-  K_CUT_DELTA: null
-  DTYPE: float32 # float32 or float16 or bfloat16
-  N_TARGETS: 1
-  TARGET_INDEX_KEY: target_index
-  RESIDUAL_FACTOR: 0.5
-  USE_ZBL: False
+  model:
+    CUTOFF_DELTA: 0.2
+    AVERAGE_POOLING: False
+    TRANSFORMERS_CENTRAL_SPECIFIC: False
+    HEADS_CENTRAL_SPECIFIC: False
+    ADD_TOKEN_FIRST: True
+    ADD_TOKEN_SECOND: True
+    N_GNN_LAYERS: 3
+    TRANSFORMER_D_MODEL: 128
+    TRANSFORMER_N_HEAD: 4
+    TRANSFORMER_DIM_FEEDFORWARD: 512
+    HEAD_N_NEURONS: 128
+    N_TRANS_LAYERS: 3
+    ACTIVATION: silu
+    USE_LENGTH: True
+    USE_ONLY_LENGTH: False
+    R_CUT: 5.0
+    R_EMBEDDING_ACTIVATION: False
+    COMPRESS_MODE: mlp
+    BLEND_NEIGHBOR_SPECIES: False
+    AVERAGE_BOND_ENERGIES: False
+    USE_BOND_ENERGIES: True
+    USE_ADDITIONAL_SCALAR_ATTRIBUTES: False
+    SCALAR_ATTRIBUTES_SIZE: null
+    TRANSFORMER_TYPE: PostLN # PostLN or PreLN
+    USE_LONG_RANGE: False
+    K_CUT: null # should be float; only used when USE_LONG_RANGE is True
+    K_CUT_DELTA: null
+    DTYPE: float32 # float32 or float16 or bfloat16
+    N_TARGETS: 1
+    TARGET_INDEX_KEY: target_index
+    RESIDUAL_FACTOR: 0.5
+    USE_ZBL: False
 
-training:
-  INITIAL_LR: 1e-4
-  EPOCH_NUM: 1000
-  EPOCHS_WARMUP: 50
-  SCHEDULER_STEP_SIZE_ATOMIC: 500000000 # structural version is called "SCHEDULER_STEP_SIZE"
-  GLOBAL_AUG: True
-  SLIDING_FACTOR: 0.7
-  ATOMIC_BATCH_SIZE: 850 # structural version is called "STRUCTURAL_BATCH_SIZE"
-  BALANCED_DATA_LOADER: False # if True, use DynamicBatchSampler from torch_geometric
-  MAX_TIME: 234000
-  ENERGY_WEIGHT: 0.1 # only used when fitting MLIP
-  MULTI_GPU: False
-  RANDOM_SEED: 0
-  CUDA_DETERMINISTIC: False
-  MODEL_TO_START_WITH: null
-  SUPPORT_MISSING_VALUES: False
-  USE_WEIGHT_DECAY: False
-  WEIGHT_DECAY: 0.0
-  DO_GRADIENT_CLIPPING: False
-  GRADIENT_CLIPPING_MAX_NORM: null # must be overwritten if DO_GRADIENT_CLIPPING is True
-  USE_SHIFT_AGNOSTIC_LOSS: False # only used when fitting general target. Primary use case: EDOS
-  ENERGIES_LOSS: per_structure # per_structure or per_atom
-  CHECKPOINT_INTERVAL: 100
+  training:
+    INITIAL_LR: 1e-4
+    EPOCH_NUM: 1000
+    EPOCHS_WARMUP: 50
+    SCHEDULER_STEP_SIZE_ATOMIC: 500000000 # structural version is called "SCHEDULER_STEP_SIZE"
+    GLOBAL_AUG: True
+    SLIDING_FACTOR: 0.7
+    ATOMIC_BATCH_SIZE: 850 # structural version is called "STRUCTURAL_BATCH_SIZE"
+    BALANCED_DATA_LOADER: False # if True, use DynamicBatchSampler from torch_geometric
+    MAX_TIME: 234000
+    ENERGY_WEIGHT: 0.1 # only used when fitting MLIP
+    MULTI_GPU: False
+    RANDOM_SEED: 0
+    CUDA_DETERMINISTIC: False
+    MODEL_TO_START_WITH: null
+    SUPPORT_MISSING_VALUES: False
+    USE_WEIGHT_DECAY: False
+    WEIGHT_DECAY: 0.0
+    DO_GRADIENT_CLIPPING: False
+    GRADIENT_CLIPPING_MAX_NORM: null # must be overwritten if DO_GRADIENT_CLIPPING is True
+    USE_SHIFT_AGNOSTIC_LOSS: False # only used when fitting general target. Primary use case: EDOS
+    ENERGIES_LOSS: per_structure # per_structure or per_atom
+    CHECKPOINT_INTERVAL: 100

src/metatrain/experimental/soap_bpnn/default-hypers.yaml

@@ -1,38 +1,39 @@
-name: experimental.soap_bpnn
+architecture:
+  name: experimental.soap_bpnn
 
-model:
-  soap:
-    cutoff: 5.0
-    max_radial: 8
-    max_angular: 6
-    atomic_gaussian_width: 0.3
-    center_atom_weight: 1.0
-    cutoff_function:
-      ShiftedCosine:
-        width: 1.0
-    radial_scaling:
-      Willatt2018:
-        rate: 1.0
-        scale: 2.0
-        exponent: 7.0
-  bpnn:
-    layernorm: true
-    num_hidden_layers: 2
-    num_neurons_per_layer: 32
-  zbl: false
+  model:
+    soap:
+      cutoff: 5.0
+      max_radial: 8
+      max_angular: 6
+      atomic_gaussian_width: 0.3
+      center_atom_weight: 1.0
+      cutoff_function:
+        ShiftedCosine:
+          width: 1.0
+      radial_scaling:
+        Willatt2018:
+          rate: 1.0
+          scale: 2.0
+          exponent: 7.0
+    bpnn:
+      layernorm: true
+      num_hidden_layers: 2
+      num_neurons_per_layer: 32
+    zbl: false
 
-training:
-  distributed: False
-  distributed_port: 39591
-  batch_size: 8
-  num_epochs: 100
-  learning_rate: 0.001
-  early_stopping_patience: 200
-  scheduler_patience: 100
-  scheduler_factor: 0.8
-  log_interval: 5
-  checkpoint_interval: 25
-  fixed_composition_weights: {}
-  per_structure_targets: []
-  loss_weights: {}
-  log_mae: False
+  training:
+    distributed: False
+    distributed_port: 39591
+    batch_size: 8
+    num_epochs: 100
+    learning_rate: 0.001
+    early_stopping_patience: 200
+    scheduler_patience: 100
+    scheduler_factor: 0.8
+    log_interval: 5
+    checkpoint_interval: 25
+    fixed_composition_weights: {}
+    per_structure_targets: []
+    loss_weights: {}
+    log_mae: False