gnn_train.yaml

stage: edge_classifier
model: InteractionGNN
input_dir:  /home/mcamp/acorn/examples/pennylane/module_map/ # Should contain the files used in training and inference
stage_dir:  /home/mcamp/acorn/examples/Example_1/gnn/ # A directory to be created that will store logs, artifacts, and output data
project: CF_Example_1_GNN # Used in logging
gpus: 1
nodes: 1

# Dataset parameters
data_split: [80,10,10] # Number of [training, validation, testing] examples
num_workers: [0,0,0]
# Truth and weighting parameters. Syntax is...
# Single value: track[key] == value
# List of 2 floats: value1 <= track[key] <= value2
# List with a string as the first value: First value defines the boolean test, the second value is a list or value. Boolean tests are: (in, not_in, is, is_not, within, not_within)
# All other tracks follow the default: If a true edge, weight=0. If a fake edge, weight=1.
weighting:
  - weight: 2.
    conditions:
      y: True
      pt: [1000, .inf]
      nhits: [3, .inf]
      primary: True
      pdgId: [not_in, [11, -11]]

# A list of hard cut conditions. Any hits not passing these conditions will be removed from the training and validation dataset.
# By definition, noise will not pass any of these cuts, so any hard cuts will remove noise.
hard_cuts:
  pt: [5000,.inf]

# Model parameters
edge_cut: 0.5
undirected: True
node_features: [r, phi, z]
node_scales:   [1000, 3.14, 1000]
hidden: 128
n_graph_iters: 8
nb_node_layer: 2
nb_edge_layer: 2
emb_channels: 0
layernorm: True
batchnorm: False
aggregation: [sum, mean, max, std]
hidden_activation: SiLU
output_activation: Tanh

# Training parameters
warmup: 5
lr: 0.001
factor: 0.9
patience: 10
max_epochs: 100