dataset:
  train:
    - class: CoDesignDataset
      mmap_dir: ./datasets/pepbdb/processed
      specify_index: ./datasets/pepbdb/processed/train_index.txt
      backbone_only: false
      cluster: ./datasets/pepbdb/train.cluster
      use_covariance_matrix: true
  valid:
    class: CoDesignDataset
    mmap_dir: ./datasets/pepbdb/processed
    specify_index: ./datasets/pepbdb/processed/valid_index.txt
    backbone_only: false
    use_covariance_matrix: true

dataloader:
  shuffle: true
  num_workers: 4
  wrapper:
    class: DynamicBatchWrapper
    complexity: n**2
    ubound_per_batch: 60000  # batch size ~32

trainer:
  class: LDMTrainer
  criterion: Loss
  config:
    max_epoch: 500 # the best checkpoint should be obtained at around epoch 380
    save_topk: 10
    val_freq: 10
    save_dir: ./ckpts/LDM_codesign
    patience: 10
    metric_min_better: true

    optimizer:
      class: AdamW
      lr: 1.0e-4

    scheduler:
      class: ReduceLROnPlateau
      factor: 0.6
      patience: 3
      mode: min
      frequency: val_epoch
      min_lr: 5.0e-6

model:
  class: LDMPepDesign
  autoencoder_ckpt: ""
  autoencoder_no_randomness: true
  hidden_size: 128
  num_steps: 100
  n_layers: 3
  n_rbf: 32
  cutoff: 3.0 # the coordinates are in standard space
  dist_rbf: 32
  dist_rbf_cutoff: 7.0
  diffusion_opt:
    trans_seq_type: Diffusion
    trans_pos_type: Diffusion
  max_gen_position: 60