constants:
  seed: &seed 42
  raise_train_error: true   # Whether the code should raise an error if it crashes during training

datamodule:
  module_type: "DGLFromSmilesDataModule"
  args:
    df_path: goli/data/micro_ZINC/micro_ZINC.csv
    cache_data_path: goli/data/cache/micro_ZINC/full.cache
    label_cols: ['score']
    smiles_col: SMILES

    # Featurization
    featurization_n_jobs: -1
    featurization_progress: True
    featurization:
      atom_property_list_onehot: [atomic-number, valence]
      atom_property_list_float: [mass, electronegativity, in-ring]
      edge_property_list: [bond-type-onehot, stereo, in-ring]
      add_self_loop: False
      explicit_H: False
      use_bonds_weights: False
      pos_encoding_as_features: &pos_enc
        pos_type: laplacian_eigvec
        num_pos: 3
        normalization: "none"
        disconnected_comp: True
      pos_encoding_as_directions: *pos_enc

    # Train, val, test parameters
    split_val: 0.2
    split_test: 0.2
    split_seed: *seed
    splits_path: null
    batch_size_train_val: 128
    batch_size_test: 128

    # Data loading
    num_workers: 0
    pin_memory: False
    persistent_workers: False  # Keep True on Windows if running multiple workers


architecture:
  model_type: fulldglnetwork
  pre_nn:   # Set as null to avoid a pre-nn network
    out_dim: 32
    hidden_dims: 32
    depth: 1
    activation: relu
    last_activation: none
    dropout: &dropout 0.1
    normalization: &normalization "batch_norm"
    last_normalization: *normalization
    residual_type: none

  pre_nn_edges:   # Set as null to avoid a pre-nn network
    out_dim: 16
    hidden_dims: 16
    depth: 2
    activation: relu
    last_activation: none
    dropout: *dropout
    normalization: *normalization
    last_normalization: *normalization
    residual_type: none

  gnn:  # Set as null to avoid a post-nn network
    out_dim: 32
    hidden_dims: 32
    depth: 4
    activation: relu
    last_activation: none
    dropout: *dropout
    normalization: *normalization
    last_normalization: *normalization
    residual_type: simple
    pooling: [sum, max, dir1]
    virtual_node: 'sum'
    layer_type: 'dgn-msgpass'
    layer_kwargs:
      # num_heads: 3
      aggregators: [mean, max, dir1/dx_abs, dir1/smooth]
      scalers: [identity, amplification, attenuation]

  post_nn:
    out_dim: 1
    hidden_dims: 32
    depth: 2
    activation: relu
    last_activation: none
    dropout: *dropout
    normalization: *normalization
    last_normalization: "none"
    residual_type: none

predictor:
  metrics_on_progress_bar: ["mae", "pearsonr", "f1 > 3", "precision > 3"]
  loss_fun: mse
  random_seed: *seed
  optim_kwargs:
    lr: 1.e-2
    weight_decay: 1.e-7
  lr_reduce_on_plateau_kwargs:
    factor: 0.5
    patience: 7
  scheduler_kwargs:
    monitor: &monitor loss/val
    frequency: 1
  target_nan_mask: 0 # null: no mask, 0: 0 mask, ignore: ignore nan values from loss


metrics:
  - name: mae
    metric: mae
    threshold_kwargs: null

  - name: pearsonr
    metric: pearsonr
    threshold_kwargs: null

  - name: f1 > 3
    metric: f1
    num_classes: 2
    average: micro
    threshold_kwargs: &threshold_3
      operator: greater
      threshold: 3
      th_on_preds: True
      th_on_target: True
      target_to_int: True

  - name: f1 > 5
    metric: f1
    num_classes: 2
    average: micro
    threshold_kwargs:
      operator: greater
      threshold: 5
      th_on_preds: True
      th_on_target: True
      target_to_int: True

  - name: precision > 3
    metric: precision
    class_reduction: micro
    threshold_kwargs: *threshold_3

trainer:
  logger:
    save_dir: logs/micro_ZINC
  early_stopping:
    monitor: *monitor
    min_delta: 0
    patience: 10
    mode: &mode min
  model_checkpoint:
    dirpath: models_checkpoints/micro_ZINC/
    filename: "model"
    monitor: *monitor
    mode: *mode
    save_top_k: 1
    every_n_val_epochs: 1
  trainer:
    max_epochs: 25
    min_epochs: 5
    gpus: 1

