jetclass.yaml

# @package _global_

# to execute this experiment run:
# python train.py experiment=jetclass

defaults:
  - override /data: jetclass.yaml
  - override /model: flow_matching.yaml
  - override /callbacks: jetclass.yaml
  - override /trainer: gpu.yaml

# all parameters below will be merged with parameters from default configurations set above
# this allows you to overwrite only specified parameters

# add here checkpoint to continue training
# ckpt_path: XXX/checkpoints/last-EMA.ckpt

tags: ["flow_matching", "JetClass", "uncond"]
run_note: ""

seed: 12345

trainer:
  min_epochs: 1
  max_epochs: 5000
  gradient_clip_val: 0.5

model:
  num_particles: 128
  global_cond_dim: 0 # needs to be calculated when using conditioning (= number of jet types + additional conditioning variables)
  local_cond_dim: 0
  features: 8 # = 3 + number of `additional_part_features` (see below)
  scheduler:
    warmup: 500
    max_iters: 5000

data:
  # preprocessing
  # list of the following: QCD, Hbb, Hcc, Hgg, H4q, Hqql, Zqq, Wqq, Tbqq, Tbl
  used_jet_types: [Tbqq] # null means all jet types
  number_of_used_jets: 3_000_000
  normalize: True
  normalize_sigma: 5
  # conditioning
  conditioning_pt: False
  conditioning_eta: False
  conditioning_mass: False
  conditioning_num_particles: False
  conditioning_jet_type: False
  additional_part_features:
    - part_d0val
    - part_d0err
    - part_dzval
    - part_dzerr
    - part_charge
    # - part_isChargedHadron
    # - part_isNeutralHadron
    # - part_isPhoton
    # - part_isElectron
    # - part_isMuon

callbacks:
  ema:
    decay: 0.999
    apply_ema_every_n_steps: 1
    start_step: 0
    save_ema_weights_in_callback_state: True
    evaluate_ema_weights_instead: True
  jetclass_eval:
    every_n_epochs: 50 # evaluate every n epochs
    additional_eval_epochs: [1, 30, 75] # evaluate at these epochs as well
    num_jet_samples: 50_000 # jet samples to generate
  jetclass_eval_test:
    num_jet_samples: 200_000 # jet samples to generate

  #early_stopping:
  #  monitor: "val/loss"
  #  patience: 2000
  #  mode: "min"

task_name: "jetclass"

logger:
  wandb:
    tags: ${tags}
    group: "flow_matching_jetclass"
    name: ${task_name}
  comet:
    experiment_name: null
    project_name: "flow-matching"