Skip to content

Schema for Active Learning Configuration

The schema to configure the input file for active learning process.

Schema:

### Schema for YAML configuration

train:                              ### ANCHOR: Trainning ML model
    type: dict
    required: True
    schema:
        num_models:                 # Number of models to train. Default is 1
            type: integer
        init_data_paths:            # List of paths to initial data.
            type: list
            required: True

        trainset_ratio:             # Ratio of training set. Default is 0.9
            type: float
        validset_ratio:             # Ratio of validation set. Default is 0.1
            type: float
        num_cores_buildgraph:       # number of cores for building graph data
            type: integer

        continue_train:             # continue checkpoints from previous iteration. Default is True
            type: boolean

        max_num_updates:            # Maximum number of updates to guess num_epochs. Default is None
            type: integer

        distributed:
            type: dict
            schema:
                distributed_backend: # choices: 'mpi' or 'nccl'  'gloo'
                    type: string
                cluster_type:       # choices: 'slurm' or 'sge'
                    type: string
                gpu_per_node:       # only need in SGE batch type. Default is 1
                    type: integer

        mlp_engine:                 # ML engine. Default is 'sevenn'. Choices: 'sevenn'
            type: string
        sevenn_args:                ### See: https://github.com/MDIL-SNU/SevenNet/blob/main/example_inputs/training/input_full.yaml
            type: dict
            schema:
                model:
                    type: dict
                train:
                    type: dict
                data:
                    type: dict


md:                                 ### ANCHOR: Run MD exploration
    type: dict
    required: True
    schema:
        committee_std:              # Committee standard deviation for each property
            type: dict
            required: True
            schema:
                e_std_lo:           # Lower bound of committee standard deviation for energy. E.g., 0.05
                    type: float
                e_std_hi:           # Upper bound of committee standard deviation for energy. E.g., 0.15
                    type: float
                f_std_lo:           # Lower bound of committee standard deviation for force. E.g., 0.05
                    type: float
                f_std_hi:           # Upper bound of committee standard deviation for force. E.g., 0.15
                    type: float
                s_std_lo:           # Lower bound of committee standard deviation for stress. E.g., 0.05
                    type: float
                s_std_hi:           # Upper bound of committee standard deviation for stress. E.g., 0.2
                    type: float

        init_struct_paths:          # List of paths of initial structures. Accept wildcard, e.g., "path/to/dir/*.extxyz"
            type: list
            required: True

        model_idx:                  # indices of models in the committee used for MD run. Default is None; means using average of all models.
            type: list  # list[int]
        # clean_traj: true

        common_md_args:             # define common parameters for MD. Accept all keyword fors [ASE MD schema](https://thang.eu.org/alff_doc/schema/config_ase/)
                                    # These common parameters will be overwritten by the parameters in `sampling_spaces`.
            type: dict
            allow_unknown: True
            schema:
                dt:                 # timestep in fs. Default is 1.0 fs
                    type: float
                thermostat:         # define common thermostat. default is 'Nose_Hoover_chain'
                    type: dict
                barostat:           # define common barostat. default is 'Iso_Nose_Hoover_chain'
                    type: dict

        sampling_spaces:            ### list of dicts of sampling spaces for MD run. For each dict:
                                    # If set both 'stresses' and 'temperatures', it will use NPT ensemble.
                                    # If set only 'temperatures', it will use NVT ensemble.
                                    # If none of them, it will use NVE ensemble.
            type: list
            required: True
            schema:
                type: dict
                allow_unknown: True
                schema:
                    init_struct_idx: # list of index of `init_struct_paths`, e.g., [0, 1, 2, 3], can be in short form [1, 2, "6-10", "3-5:2"]
                        type: list
                    temps:          # list of temperatures in K, e.g., [1, 300, 1500]
                        type: list
                    stresses:       # list of # external stress in GPa. Default is None if not set. Accept float, 6-vector, or 3x3 matrix.
                        type: [float, list]
                    # Other arguments as keywords for [ASE MD schema](https://thang.eu.org/alff_doc/schema/config_ase/), except these keys: 'temperature', 'stress'.


dft:                            ### ANCHOR: DFT calculators
  type: dict
  required: True
  schema:
    calc:                       # accept all keywords for ['ase.calc'](https://thang.eu.org/alff_doc/schema/config_ase/)
        type: dict
        required: True
        schema:
            gpaw:               # GPAW calculator parameters
                type: dict
                required: True
            dftd3:              # DFT-D3 calculator for Van der Waals correction
                type: dict

Example configuration:

##### ANCHOR: Trainning
train:
    num_models: 4
    init_data_paths:
        - ../1_gendata/*/*/02_gendata/data_label.extxyz
        - ../1_gendata/1_iteration_data

    trainset_ratio: 0.9
    validset_ratio: 0.1
    num_cores_buildgraph: 10       # number of cores for building graph data

    continue_train: True            # continue checkpoints from previous iteration. Default is True
    # max_num_updates: 100000       # Maximum number of updates to guess num_epochs. Default is None

    distributed:
        distributed_backend: 'nccl'     # choices: 'mpi' or 'nccl'  'gloo'
        cluster_type: 'slurm'           # choices: 'slurm' or 'sge'
        # gpu_per_node: 1               # only need in sge

    mlp_engine: sevenn
    sevenn_args:  # Updated: Dec 17, 2024. See: https://github.com/MDIL-SNU/SevenNet/blob/main/example_inputs/training/input_full.yaml
        model:
            chemical_species: ['Mo', 'W', 'S', 'Se', 'Te']  # Elements model should know. [ 'Univ' | 'Auto' | manual_user_input ]
            cutoff: 5.0                                   # Cutoff radius in Angstroms. If two atoms are within the cutoff, they are connected.
            channel: 32                                   # The multiplicity(channel) of node features.
            lmax: 2                                       # Maximum order of irreducible representations (rotation order).
            num_convolution_layer: 3                      # The number of message passing layers.

            #irreps_manual:                               # Manually set irreps of the model in each layer
                #- "128x0e"
                #- "128x0e+64x1e+32x2e"
                #- "128x0e+64x1e+32x2e"
                #- "128x0e+64x1e+32x2e"
                #- "128x0e+64x1e+32x2e"
                #- "128x0e"

            weight_nn_hidden_neurons: [64, 64]            # Hidden neurons in convolution weight neural network
            radial_basis:                                 # Function and its parameters to encode radial distance
                radial_basis_name: 'bessel'               # Only 'bessel' is currently supported
                bessel_basis_num: 8
            cutoff_function:                              # Envelop function, multiplied to radial_basis functions to init edge featrues
                cutoff_function_name: 'poly_cut'          # {'poly_cut' and 'poly_cut_p_value'} or {'XPLOR' and 'cutoff_on'}
                poly_cut_p_value: 6

            act_gate: {'e': 'silu', 'o': 'tanh'}          # Equivalent to 'nonlinearity_gates' in nequip
            act_scalar: {'e': 'silu', 'o': 'tanh'}        # Equivalent to 'nonlinearity_scalars' in nequip

            is_parity: False                              # Pairy True (E(3) group) or False (to SE(3) group)

            self_connection_type: linear                # Default is 'nequip'. 'linear' is used for SevenNet-0.
            interaction_type: nequip

            conv_denominator: "avg_num_neigh"             # Valid options are "avg_num_neigh*", "sqrt_avg_num_neigh", or float
            train_denominator: False                      # Enable training for denominator in convolution layer
            train_shift_scale: False                      # Enable training for shift & scale in output layer

        train:
            random_seed: 1
            train_shuffle: True
            is_train_stress: True                         # Includes stress in the loss function
            epoch: 200                                    # Ends training after this number of epochs
            per_epoch:  20                                # Generate checkpoints every this epoch

            # loss: 'Huber'                                # Default is 'mse' (mean squared error)
            # loss_param:
            #     delta: 0.01

            # Each optimizer and scheduler have different available parameters.
            # You can refer to sevenn/train/optim.py for supporting optimizer & schedulers
            optimizer: 'adam'                             # Options available are 'sgd', 'adagrad', 'adam', 'adamw', 'radam'
            optim_param:
                lr: 1.0e-4

            scheduler: linearlr
            scheduler_param:
                start_factor: 1.0
                total_iters: 200
                end_factor: 1.0e-9

            # scheduler: 'reducelronplateau'            # One of 'steplr', 'multisteplr', 'exponentiallr', 'cosineannealinglr', 'reducelronplateau', 'linearlr'
            # scheduler_param:
            #     factor: 0.75
            #     patience: 2
            #     threshold: 5.0e-5   # only changes large than this value will be considered as a change
            #     min_lr: 1.0e-12      # minimum learning rate

            force_loss_weight: 1.0     # 57.9     # 0.1                       # Coefficient for force loss
            stress_loss_weight: 0.001      # 1.0e-6                   # Coefficient for stress loss (to kbar unit), kbar = 0.1 GPa

            # ['target y', 'metric']
            # Target y: TotalEnergy, Energy, Force, Stress, Stress_GPa, TotalLoss
            # Metric  : RMSE, MAE, or Loss
            error_record:
                - ['Energy', 'RMSE']
                - ['Force', 'RMSE']
                - ['Stress', 'RMSE']
                # - ['Stress_GPa', 'RMSE']
                - ['Energy', 'Loss']
                - ['Force', 'Loss']
                - ['Stress', 'Loss']
                - ['TotalLoss', 'None']
            best_metric: TotalLoss

            ### THANG: do not use this, just set `init_checkpoints` above
            # Continue training model from given checkpoint, or pre-trained model checkpoint for fine-tuning
            #continue:
                #checkpoint: 'checkpoint_best.pth'       # Checkpoint of pre-trained model or a model want to continue training.
                #reset_optimizer: False                  # Set True for fine-tuning
                #reset_scheduler: False                  # Set True for fine-tuning

        data:
            batch_size: 200    # 250                     # Per GPU batch size.

            shift: 'per_atom_energy_mean'                # One of 'per_atom_energy_mean*', 'elemwise_reference_energies', float
            scale: 'force_rms'                           # One of 'force_rms*', 'per_atom_energy_std', 'elemwise_force_rms', float

            data_format: 'ase'                   # Default is 'ase'. Choices are 'ase', 'structure_list', '.sevenn_data'
            data_format_args:                    # Paramaters, will be passed to ase.io.read
                energy_key: 'ref_energy'                 # Key for energy in extxyz file
                force_key: 'ref_forces'                  # Key for force in extxyz file
                stress_key: 'ref_stress'                 # Key for stress in extxyz file


##### ANCHOR: Run MD
md:
    committee_std:
        e_std_lo: 0.05
        e_std_hi: 0.15
        f_std_lo: 0.05
        f_std_hi: 0.15
        s_std_lo: 0.05
        s_std_hi: 0.15
    common_md_args:
        dt: 1                # unit: fs
        thermostat:
            name: 'Nose_Hoover_chain'
        barostat:
            name: 'Iso_Nose_Hoover_chain'  # 'Parrinello_Rahman'
    # model_idx: [1]           # index of models used for MD

    init_struct_paths:
        - _init_struct/bulk_MoX2/MoS2_mx2_2H_01x01x02    # bulk MoX2, [0-5]
        - _init_struct/bulk_MoX2/MoS2_mx2_1T_01x01x02
        - _init_struct/bulk_MoX2/MoSe2_mx2_2H_01x01x02
        - _init_struct/bulk_MoX2/MoSe2_mx2_1T_01x01x02
        - _init_struct/bulk_MoX2/MoTe2_mx2_2H_01x01x02
        - _init_struct/bulk_MoX2/MoTe2_mx2_1T_01x01x02

        - _init_struct/layer_MoX2/MoS2_mx2_2H_01x01x02   # layer MoX2, [6-11]
        - _init_struct/layer_MoX2/MoS2_mx2_1T_01x01x02
        - _init_struct/layer_MoX2/MoSe2_mx2_2H_01x01x02
        - _init_struct/layer_MoX2/MoSe2_mx2_1T_01x01x02
        - _init_struct/layer_MoX2/MoTe2_mx2_2H_01x01x02
        - _init_struct/layer_MoX2/MoTe2_mx2_1T_01x01x02

    sampling_spaces:        ### If set both stress and temperature, it will use NPT ensemble.
                            ### If set only temperature, it will use NVT ensemble.
                            ### If none of them, it will use NVE ensemble.
        ### MoX2
        - {}
        - init_struct_idx: [0,1,2,3,4,5]        ### NVT sampling BULK
          equil_steps: 200
          traj_freq: 4
          num_frames: 10
          temps: [1, 300, 600, 900, 1200, 1500]          # temperatures in K

        - init_struct_idx: [6,7,8,9,10,11]      ### NVT sampling LAYER
          equil_steps: 200
          traj_freq: 4
          num_frames: 10
          temps: [1, 300, 600, 900, 1200, 1500]

        - init_struct_idx: [0,1,2,3,4,5]        ### NPT sampling BULK
          equil_steps: 200
          traj_freq: 4
          num_frames: 10
          temps: [1, 300 ]               # temperatures in K
          stresses: [0, 10, 15]          # stresses in GPa

        - init_struct_idx: [6,7,8,9,10,11]      ### NPT sampling SURFACE
          equil_steps: 200
          traj_freq: 4
          num_frames: 10
          temps: [1, 300 ]               # temperatures in K
          stresses: [0, 10, 15]          # stresses in GPa
          barostat:
            mask: [0,1,1]               # disable x-direction barostat
            name: 'Parrinello_Rahman'


##### ANCHOR: DFT calculation
dft:
    calc:
        gpaw:           # accept GPAW parameters
            mode:
                name: 'pw'          # use PlaneWave method energy cutoff in eV
                ecut: 500
            xc: "PBE"               # exchange-correlation functional
            kpts:
                density: 6
                gamma: False        # if not set `kpts`, then only Gamma-point is used
            parallel:
                sl_auto: True       # enable ScaLAPACK parallelization
                use_elpa: True      # enable Elpa eigensolver
                augment_grids: True # use all cores for XC/Poisson solver
                gpu: True           # enable GPU acceleration

        dftd3:           # DFT-D3 method for Van der Waals correction
            damping: "d3zero"   # use DFT-D3 damping. Default is "d3zero" (zero-damping). Choices: "d3bj","d3zero","d3bjm","d3zerom","d3op".