# Example input.yaml for fine-tuning sevennet-0
# '*' signifies default. You can check log.sevenn for defaults.

model:  # model keys should be consistent except for train_* keys
    chemical_species: 'Auto'
    cutoff: 5.0
    channel: 128
    is_parity: False
    lmax: 2
    num_convolution_layer: 5
    irreps_manual:
        - "128x0e"
        - "128x0e+64x1e+32x2e"
        - "128x0e+64x1e+32x2e"
        - "128x0e+64x1e+32x2e"
        - "128x0e+64x1e+32x2e"
        - "128x0e"

    weight_nn_hidden_neurons: [64, 64]
    radial_basis:
        radial_basis_name: 'bessel'
        bessel_basis_num: 8
    cutoff_function:
        cutoff_function_name: 'XPLOR'
        cutoff_on: 4.5
    self_connection_type: 'linear'

    train_shift_scale: False   # customizable (True | False)
    train_denominator: False   # customizable (True | False)

train:  # Customizable
    random_seed: 1
    is_train_stress: True
    epoch: 100

    loss: 'Huber'  # keeping original loss function give better ft result
    loss_param:
        delta: 0.01

    optimizer: 'adam'
    optim_param:
        lr: 0.004
    scheduler: 'exponentiallr'
    scheduler_param:
        gamma: 0.99

    force_loss_weight: 1.0
    stress_loss_weight: 0.01

    per_epoch: 10  # Generate checkpoints every this epoch

    # ['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']
        - ['TotalLoss', 'None']

    continue:
        reset_optimizer: True
        reset_scheduler: True
        reset_epoch: True
        checkpoint: 'SevenNet-0_11July2024'

data:  # Customizable
    batch_size: 4
    data_divide_ratio: 0.1

    # SevenNet automatically matches data format from its filename.
    # For those not `structure_list` or `.pt` files, assumes it is ASE readable
    # In this case, below arguments are directly passed to `ase.io.read`
    data_format_args:
        index: ':'                                # see `https://wiki.fysik.dtu.dk/ase/ase/io/io.html` for more valid arguments

    # validset is needed if you want '_best.pth' during training. If not, both validset and testset is optional.
    load_trainset_path: ['./train_*.extxyz']  # Example of using ase as data_format, support multiple files and expansion(*)
    load_validset_path: ['./valid.extxyz']
    load_testset_path:  ['./sevenn_data/mydata.pt']  # Graph can be preprocessed using `sevenn_graph_build` and accessible like this
