config.yaml

net: deep_sets # deep_sets or particle_net

epochs: 100
shuffle_buffer: 100

loss: mean_absolute_error
optimizer: adam
lr: 2.e-3

callbacks:
  reduce_lr_on_plateau:
    factor: 0.2
    min_lr: 1.e-8
    min_delta: 1.e-4
    patience: 5
  early_stopping:
    min_delta: 1.e-4
    patience: 7

model:
  deep_sets:
    type: mlp # mlp or resnet
    activation: relu
    initializer: he_normal
    head:
      units: [1024, 512, 256, 128, 64]
      batch_norm: False
      dropout: 0
    ch:
      units: [64, 128, 256]
      batch_norm: True
      dropout: 0
    ne:
      units: [64, 128, 256]
      batch_norm: True
      dropout: 0
    sv:
      units: [32, 64, 128]
      batch_norm: True
      dropout: 0

  particle_net:
    activation: relu
    initializer: he_normal
    batch_norm: True
    shortcut: True
    pooling: average # average or max
    ch:
      K: 16
      channels: [
        [64, 64, 64],
        [128, 128, 128],
        [256, 256, 256]
      ]
    ne:
      K: 16
      channels: [
        [64, 64, 64],
        [128, 128, 128],
        [256, 256, 256]
      ]
    sv:
      K: 8
      channels: [
        [32, 32, 32],
        [64, 64, 64],
        [128, 128, 128]
      ]
    dropout: 0
    units: [512, 256, 128, 64]

data:
  batch_size: 1024
  train_size: 0.6
  validation_size: 0.2
  test_size: 0.2
  
  features:
    jet: 
      numerical: [
        pt, eta, phi, mass, area, rho, num_pv
      ]
      categorical: []
      synthetic: [log_pt, mult, ptD, axis2]
    ch:
      num_points: 64 # ParticleNet point cutoff
      numerical: [
        ch_pt, ch_eta, ch_phi, ch_dxy, ch_dxy_significance, ch_dz,
        ch_num_hits, ch_num_pixel_hits, ch_lost_hits, ch_norm_chi2
      ]
      categorical: [ch_id, ch_pv_ass]
      synthetic: [ch_rel_pt, ch_rel_eta, ch_rel_phi]
    ne:
      num_points: 64
      numerical: [
        ne_pt, ne_eta, ne_phi, ne_hcal_frac
      ]
      categorical: [ne_id]
      synthetic: [ne_rel_pt, ne_rel_eta, ne_rel_phi]
    sv:
      num_points: 16
      numerical: [
        sv_mass, sv_pt, sv_eta, sv_phi, sv_distance, sv_significance, sv_num_tracks
      ]
      categorical: []
      synthetic: [sv_rel_pt, sv_rel_eta, sv_rel_phi]
  transforms:
    categorical:
      ch_id: [-211, -13, -11, 11, 13, 211]
      ch_pv_ass: [0, 1, 4, 5, 6, 7]
      ne_id: [1, 2, 22, 130]