Add GEOM-Drugs epoch 100 and QM9 epochs 1000 and 1500 with inference assets

geom/explicit_h/README.md

@@ -1,12 +1,16 @@
-# GEOM-Drugs with explicit hydrogens — epoch 50
+# GEOM-Drugs with explicit hydrogens
 
-Continued from the best checkpoint after 48 completed epochs. The original Lightning checkpoint includes current weights, EMA weights and optimizer state. The inference loader constructs only the selected network (EMA by default).
+Published completed epochs: 50, 100.
+
+Continued from the best checkpoint after 48 completed epochs. The training subset contains at most 72 atoms including hydrogen.
+
+The original Lightning checkpoints include current weights, EMA weights and optimizer state. The inference loader constructs only the selected network (EMA by default).
 
 ```python
 from morph.inference import load_pretrained
 
-model = load_pretrained("geom-explicit-h", epoch=50, device="cuda")
+model = load_pretrained("geom-explicit-h", epoch=100, device="cuda")
 trajectories, validity, vocab, coordinate_std = model.generate(n_mols=16)
 ```
 
-The training subset contains at most 72 atoms including hydrogen. Coordinates returned by generation are in Angstrom. Config, vocabularies and training distributions are included; no dataset is needed for inference.
+Coordinates returned by generation are in Angstrom. Config, vocabularies and training distributions are included; no dataset is needed for inference. Training settings and source are in `training_bundle.tar.gz`.

geom/explicit_h/epoch_100.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0ab5ad7247b5518858c8cb6aa7c4102352e8fe09cbb6744c33a4c784c85b8587
+size 509181449

geom/explicit_h/manifest.json

@@ -1,7 +1,7 @@
 {
   "README.md": {
-    "bytes": 701,
-    "sha256": "9947149d5b311da2f381b57f135811eb1eade2dba84cdf452ae69f5e11ba1016"
+    "bytes": 790,
+    "sha256": "abba7fd92562707cf7a27e782e2abe357f012a261a57964a035351bf673c17eb"
   },
   "atom_tokens.txt": {
     "bytes": 22,
@@ -32,8 +32,8 @@
     "sha256": "016c4a5bbb4f5ef36d1d62dff772ecf8d384e8ab3e2c94ebeb4c4357183f690a"
   },
   "metadata.json": {
-    "bytes": 1094,
-    "sha256": "b23b8944369765d12b2cbae0f5d48614225e44cb59deff8bb0e49783318ddb31"
+    "bytes": 1404,
+    "sha256": "b4deb28242f9bcefe514ab60288652105e2975bca68c53ab1a740976a53fa099"
   },
   "runtime/configs/callbacks/callbacks.yaml": {
     "bytes": 699,
@@ -596,11 +596,15 @@
     "sha256": "4fd11b415e9b4a599344431218db7f9960ba443dc29787410a4566ce5dacb434"
   },
   "validation.json": {
-    "bytes": 294,
-    "sha256": "5acbbb78be426e03abd58c3dc062b30b000e91e815c329dfa7f1422c9866b740"
+    "bytes": 780,
+    "sha256": "b296c55b86daeaf6e68f77b807fe3965ee6627ab0540da584fa9b0fc16f0f238"
   },
   "runtime/README.md": {
     "bytes": 1041,
     "sha256": "7e32e04cce313c7faad9a3170ab009591aa327a900ee67f980889b7d7e4c82ba"
+  },
+  "epoch_100.ckpt": {
+    "bytes": 509181449,
+    "sha256": "0ab5ad7247b5518858c8cb6aa7c4102352e8fe09cbb6744c33a4c784c85b8587"
   }
 }

geom/explicit_h/metadata.json

@@ -17,6 +17,15 @@
       "global_step": 441050,
       "bytes": 509181385,
       "sha256": "016c4a5bbb4f5ef36d1d62dff772ecf8d384e8ab3e2c94ebeb4c4357183f690a"
+    },
+    {
+      "filename": "epoch_100.ckpt",
+      "original_filename": "epoch=99-step=882100.ckpt",
+      "completed_epochs": 100,
+      "checkpoint_epoch_index": 99,
+      "global_step": 882100,
+      "bytes": 509181449,
+      "sha256": "0ab5ad7247b5518858c8cb6aa7c4102352e8fe09cbb6744c33a4c784c85b8587"
     }
   ],
   "atom_tokens": [

geom/explicit_h/validation.json

@@ -1,17 +1,39 @@
 {
   "dataset_access": false,
-  "completed_epochs": 50,
-  "global_step": 441050,
-  "all_checkpoint_weights_finite": true,
-  "strict_load": {
-    "current": true,
-    "ema": true
-  },
-  "sampling_smoke": {
-    "attempts": 1,
-    "integration_steps": 100,
-    "validity": [
-      true
-    ]
-  }
+  "checkpoints": [
+    {
+      "dataset_access": false,
+      "completed_epochs": 50,
+      "global_step": 441050,
+      "all_checkpoint_weights_finite": true,
+      "strict_load": {
+        "current": true,
+        "ema": true
+      },
+      "sampling_smoke": {
+        "attempts": 1,
+        "integration_steps": 100,
+        "validity": [
+          true
+        ]
+      }
+    },
+    {
+      "dataset_access": false,
+      "completed_epochs": 100,
+      "global_step": 882100,
+      "all_checkpoint_weights_finite": true,
+      "strict_load": {
+        "current": true,
+        "ema": true
+      },
+      "sampling_smoke": {
+        "attempts": 1,
+        "integration_steps": 100,
+        "validity": [
+          true
+        ]
+      }
+    }
+  ]
 }

qm9/explicit_h/README.md

@@ -0,0 +1,16 @@
+# QM9 with explicit hydrogens
+
+Published completed epochs: 1000, 1500.
+
+Trained from scratch on the original QM9 split with explicit hydrogen atoms. Validation runs every epoch and checkpoints are retained every 25 epochs.
+
+The original Lightning checkpoints include current weights, EMA weights and optimizer state. The inference loader constructs only the selected network (EMA by default).
+
+```python
+from morph.inference import load_pretrained
+
+model = load_pretrained("qm9-explicit-h", epoch=1500, device="cuda")
+trajectories, validity, vocab, coordinate_std = model.generate(n_mols=16)
+```
+
+Coordinates returned by generation are in Angstrom. Config, vocabularies and training distributions are included; no dataset is needed for inference. Training settings and source are in `training_bundle.tar.gz`.

qm9/explicit_h/atom_tokens.txt

@@ -0,0 +1,5 @@
+C
+F
+H
+N
+O

qm9/explicit_h/charge_tokens.txt

@@ -0,0 +1,3 @@
+-1
+0
+1

qm9/explicit_h/config.yaml

@@ -0,0 +1,327 @@
+representation: molecule
+seed: 42
+heads:
+  heads:
+    _target_: morph.model.layers.heads.MultiHeadModule
+    output_init_std: ${model.head_initialization.output_std}
+    independent_adapters: ${oc.select:model.head_independent_adapters,false}
+    heads_configs:
+      graph_aggregation: mean
+      node_heads:
+        atom_type_head:
+          input_dim: ${model.head_node_dim}
+          output_dim: ${len:${data.vocab.atom_tokens}}
+          adapter_dim: ${model.head_node_adapter_dim}
+          group: state
+        charge_head:
+          input_dim: ${model.head_node_dim}
+          output_dim: ${len:${data.vocab.charge_tokens}}
+          adapter_dim: ${model.head_node_adapter_dim}
+          group: state
+        ins_rate_head:
+          input_dim: ${model.head_node_dim}
+          output_dim: 1
+          adapter_dim: ${model.head_node_adapter_dim}
+          group: edit
+          initial_rate: ${model.head_initialization.ins_rate}
+        do_del_head:
+          input_dim: ${model.head_node_dim}
+          output_dim: 1
+          adapter_dim: ${model.head_node_adapter_dim}
+          group: edit
+          initial_probability: ${model.head_initialization.do_del_probability}
+        do_sub_a_head:
+          input_dim: ${model.head_node_dim}
+          output_dim: 1
+          adapter_dim: ${model.head_node_adapter_dim}
+          group: edit
+          initial_probability: ${model.head_initialization.do_sub_a_probability}
+      edge_heads:
+        edge_type_head:
+          input_dim: ${model.head_edge_dim}
+          output_dim: ${len:${data.vocab.edge_tokens}}
+          adapter_dim: ${model.head_edge_adapter_dim}
+          group: state
+        do_sub_e_head:
+          input_dim: ${model.head_edge_dim}
+          output_dim: 1
+          adapter_dim: ${model.head_edge_adapter_dim}
+          group: state
+          initial_probability: ${model.head_initialization.do_sub_e_probability}
+  ins_edge_head:
+    _target_: morph.model.layers.heads.InsertionEdgeHead
+    _recursive_: false
+    hidden_dim: ${model.head_node_dim}
+    readout_dim: ${model.head_insertion_edge_dim}
+    output_init_std: ${model.head_initialization.output_std}
+    n_edge_types: ${len:${data.vocab.edge_tokens}}
+    n_atom_types: ${len:${data.vocab.atom_tokens}}
+    n_charge_types: ${len:${data.vocab.charge_tokens}}
+    rbf_embedding_args:
+      _target_: morph.model.layers.embedding.RBFEmbedding
+      num_rbf: 128
+      rbf_dmax: 3.5
+      out_dim: 64
+      dropout: 0.0
+      trainable: false
+insertion_head:
+  mode: unimodal
+  dimension: 3
+  num_components: 1
+  model:
+    _target_: morph.model.layers.heads.UnimodalInsertionHead
+    hidden_dim: ${model.head_node_dim}
+    readout_dim: ${model.head_insertion_adapter_dim}
+    output_init_std: ${model.head_initialization.output_std}
+    independent_adapters: ${oc.select:model.head_independent_adapters,false}
+    N_a: ${len:${data.vocab.atom_tokens}}
+    N_c: ${len:${data.vocab.charge_tokens}}
+model:
+  module:
+    optimizer_config:
+      implementation: torch
+      optimizer:
+        execution: ${oc.env:MORPH_MUON_EXECUTION,native}
+        lr: 0.01
+        momentum: 0.95
+        weight_decay: 0.0
+        nesterov: true
+        ns_steps: 5
+        eps: 1.0e-07
+        adjust_lr_fn: original
+      aux_optimizer:
+        lr: 0.0003
+        betas:
+        - 0.9
+        - 0.95
+        eps: 1.0e-10
+        weight_decay: 0.0
+      gradient_clip_val: 1.0
+      gradient_clip_algorithm: norm
+      scheduler:
+        _target_: morph.utils.lr_schedulers.MultiGroupCosineWarmupLR
+        warmup_steps: 2000
+        total_steps: 200000
+        min_lr_fraction: 0.05
+        last_epoch: -1
+      interval: step
+      monitor: null
+    loss_weights:
+      do_sub_a: 0.2
+      sub_a_class: 0.03
+      do_sub_e: 0.25
+      sub_e_class: 0.5
+      do_del: 0.4
+      ins_rate: 0.15
+      ins_state: 0.01
+      ins_e: 0.4
+      x: 2.0
+      c: 0.5
+      ins_e_ii: 0.03
+    _target_: morph.model.lightning.base.LightningModuleRates
+    insertion_head_mode: ${insertion_head.mode}
+    ins_noise_scale: ${data.ins_noise_scale}
+    ema_decay: 0.999
+    ema_warmup_steps: 5000
+    use_ema_for_eval: true
+    time_dist: ${if:${eq:${model.sample_time_strategy}, log}, ${model.beta_sched},
+      ${model.linear_sched}}
+    integrator: ${model.integrator}
+    vocab: ${data.vocab}
+    use_learnable_loss_weights: false
+    use_time_weights: false
+    model:
+      _target_: morph.model.backbones.dit.DiTBackboneWithHeads
+      _recursive_: false
+      atom_type_embedding_args: ${embeddings.atom_type_embedding}
+      edge_type_embedding_args: ${embeddings.edge_type_embedding}
+      time_embedding_args: ${embeddings.time_embedding}
+      node_count_embedding_args: ${embeddings.node_count_embedding}
+      incident_bond_embedding: false
+      backbone_model_args: ${model.backbone}
+      heads_args: ${heads.heads}
+      insertion_head_args: ${insertion_head.model}
+      ins_edge_head_args: ${heads.ins_edge_head}
+  rotate: true
+  n_augmentations: 1
+  integrator: ${integrator.integrator}
+  beta_sched:
+    _target_: torch.distributions.Beta
+    concentration1: 2.0
+    concentration0: 1.0
+  linear_sched:
+    _target_: torch.distributions.Uniform
+    low: 0.0
+    high: 1.0
+  sample_time_strategy: linear
+  head_node_dim: ${model.backbone.out_node_nf}
+  head_edge_dim: ${model.backbone.out_edge_nf}
+  head_node_adapter_dim: 256
+  head_edge_adapter_dim: 256
+  head_insertion_adapter_dim: 256
+  head_insertion_edge_dim: 256
+  head_independent_adapters: true
+  head_initialization:
+    output_std: 0.001
+    coordinate_output_std: 0.0
+    do_del_probability: 0.5
+    do_sub_a_probability: 0.5
+    do_sub_e_probability: 0.5
+    ins_rate: 0.6931471805599453
+  token_weighting:
+    type_loss_token_weights: uniform
+    weight_alpha: 0.4
+  lora:
+    enabled: false
+    rank: 8
+    alpha: 8.0
+    dropout: 0.0
+    targets:
+    - qkv
+    - proj
+    - mlp.gate_value
+    - mlp.out_proj
+    freeze_base: true
+  backbone:
+    _target_: morph.model.backbones.dit.DiTBackbone
+    _recursive_: false
+    d_input: ${embeddings.atom_type_embedding.out_dim}
+    d_cond: ${oc.eval:"${embeddings.time_embedding.out_dim} + ${embeddings.node_count_embedding.out_dim}"}
+    d_model: 512
+    out_node_nf: 256
+    out_edge_nf: 256
+    num_layers: 8
+    nhead: 16
+    mlp_ratio: 4.0
+    ffn_hidden_dim: null
+    ffn_multiple_of: 64
+    qk_norm: true
+    coordinate_output_init_std: ${model.head_initialization.coordinate_output_std}
+    center_coordinate_output: true
+    in_edge_nf: ${embeddings.edge_type_embedding.out_dim}
+    edge_pair_dim: 256
+    pair_bias_hidden_dim: 32
+    num_atom_types: ${len:${data.vocab.atom_tokens}}
+    lora: ${model.lora}
+    rbf_embedding_args:
+      _target_: morph.model.layers.embedding.RBFEmbedding
+      num_rbf: 128
+      rbf_dmax: 3.5
+      out_dim: 64
+      dropout: 0.0
+      trainable: false
+schedule:
+  _target_: morph.flow_matching.schedules.SmoothstepSchedule
+  shift: 0.8
+trainer:
+  trainer:
+    deterministic: warn
+    devices: 1
+    accelerator: auto
+    strategy: auto
+    precision: 32-true
+    max_epochs: 3000
+    accumulate_grad_batches: 2
+    num_sanity_val_steps: 0
+    gradient_clip_val: null
+    gradient_clip_algorithm: norm
+    check_val_every_n_epoch: 1
+    num_val_samples: 4096
+    limit_val_batches: null
+    limit_predict_batches: 4
+    max_time: null
+    enable_progress_bar: false
+  do_train: false
+  do_validate: false
+  do_predict: true
+  checkpoint_path: null
+  final_checkpoint_path: null
+callbacks: {}
+logging:
+  enabled: false
+  project: morph-${data.dataset_name}
+  name: qm9-explicit-h
+  save_dir: .
+  offline: false
+integrator:
+  integrator:
+    _target_: morph.flow_matching.integration.RateIntegrator
+    vocab: ${data.vocab}
+    distributions: None
+    insertion_head_mode: ${insertion_head.mode}
+    insertion_dimension: ${insertion_head.dimension}
+    insertion_num_components: ${insertion_head.num_components}
+    time_strategy: log
+    time_beta_alpha: 2.0
+    time_beta_beta: 1.0
+    num_integration_steps: 100
+    ins_noise_scale: ${data.ins_noise_scale}
+    renoise_insertions: true
+    exact_schedule_integration: true
+    del_schedule: ${data.interpolator.del_schedule}
+    ins_schedule: ${data.interpolator.ins_schedule}
+    sub_schedule: ${data.interpolator.sub_schedule}
+    sub_e_schedule: ${data.interpolator.sub_e_schedule}
+data:
+  dataset_name: qm9
+  root_path: .
+  preprocess_workers: 0
+  n_atoms_strategy: uniform
+  uniform_n_atoms_padding: 5
+  optimal_transport: equivariant
+  ins_noise_scale: 0.35
+  move_noise_scale: 0.3
+  allow_charged: false
+  preprocessing:
+    _target_: morph.dataset.core.preprocessing.Preprocessing
+    _recursive_: false
+    root: .
+    train_dataset: null
+    atom_tokens_path: ${data.preprocessing.root}/atom_tokens.txt
+    edge_tokens_path: ${data.preprocessing.root}/edge_tokens.txt
+    charge_tokens_path: ${data.preprocessing.root}/charge_tokens.txt
+    distributions_path: ${data.preprocessing.root}/distributions.pt
+    num_workers: 0
+  vocab: None
+  interpolator:
+    _target_: morph.flow_matching.interpolation.Interpolator
+    distributions: None
+    optimal_transport: ${data.optimal_transport}
+    ins_noise_scale: ${data.ins_noise_scale}
+    move_noise_scale: ${data.move_noise_scale}
+    pre_align_iters: 2
+    c_move: 1.0
+    c_sub: 0.0
+    c_ins: 6.125
+    c_del: 6.125
+    del_schedule: ${schedule}
+    ins_schedule: ${schedule}
+    sub_schedule:
+      _target_: morph.flow_matching.schedules.SmoothstepSchedule
+      shift: 1.5
+    sub_e_schedule:
+      _target_: morph.flow_matching.schedules.SmoothstepSchedule
+      shift: 1.5
+  datamodule: null
+embeddings:
+  atom_type_embedding:
+    _target_: morph.model.layers.embedding.Embedding
+    num_embeddings: ${len:${data.vocab.atom_tokens}}
+    embedding_dim: 64
+    out_dim: 128
+  edge_type_embedding:
+    _target_: morph.model.layers.embedding.Embedding
+    num_embeddings: ${len:${data.vocab.edge_tokens}}
+    embedding_dim: 64
+    out_dim: 128
+  time_embedding:
+    _target_: morph.model.layers.embedding.TimeEmbedding
+    embedding_dim: 128
+    out_dim: 128
+  node_count_embedding:
+    _target_: morph.model.layers.embedding.ExtrapolatableCountEmbedding
+    embedding_dim: 64
+    out_dim: 128
+    min_period: 64.0
+    max_period: 512.0
+    max_count: 32.0

qm9/explicit_h/distributions.json

@@ -0,0 +1,56 @@
+{
+  "atom_type_distribution": [
+    0.36011767387390137,
+    0.00031341883004643023,
+    0.5080203413963318,
+    0.05379530042409897,
+    0.07775323837995529
+  ],
+  "edge_type_distribution": [
+    0.8788241147994995,
+    0.10941674560308456,
+    0.00502026965841651,
+    0.0020077268127352,
+    0.004731142893433571
+  ],
+  "charge_type_distribution": [
+    0.00020011454762425274,
+    0.9995997548103333,
+    0.00020011454762425274
+  ],
+  "n_atoms_distribution": [
+    0.0,
+    0.0,
+    0.0,
+    1.9999999494757503e-05,
+    3.9999998989515007e-05,
+    4.999999873689376e-05,
+    0.00011999999696854502,
+    0.0002099999983329326,
+    0.0006699999794363976,
+    0.0016299999551847577,
+    0.00395999988541007,
+    0.007689999882131815,
+    0.015320000238716602,
+    0.029020000249147415,
+    0.05211000144481659,
+    0.08551999926567078,
+    0.11935000121593475,
+    0.14640000462532043,
+    0.14754000306129456,
+    0.1411599963903427,
+    0.09175000339746475,
+    0.08540000021457672,
+    0.026650000363588333,
+    0.03260999917984009,
+    0.0037499999161809683,
+    0.007519999984651804,
+    0.00033999999868683517,
+    0.001069999998435378,
+    0.0,
+    9.999999747378752e-05
+  ],
+  "coordinate_std": 1.373397946357727,
+  "pairwise_distance_histogram": null,
+  "radius_of_gyration_histogram": null
+}

qm9/explicit_h/distributions.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:93d265dd9afcf536270680492e4bf6ba776969d0811a32e918b8a8994189ccc6
+size 2585

qm9/explicit_h/edge_tokens.txt

@@ -0,0 +1,5 @@
+<NO_BOND>
+1
+2
+3
+4

qm9/explicit_h/epoch_1000.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:144410a5d6efcb326c3dfa466bbd3ec1fb203d4278923f2ec7d8a486ef9f8f1e
+size 508845203

qm9/explicit_h/epoch_1500.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:90b418b9aa5de7cdb2109df20bc4ded1c46159b3abd10595750a5b583d0f9f87
+size 508845203

qm9/explicit_h/manifest.json

@@ -0,0 +1,610 @@
+{
+  "README.md": {
+    "bytes": 809,
+    "sha256": "fc9f081b53993af560a3dbeebcab1d35e937100453175bbf6692a582734c6520"
+  },
+  "atom_tokens.txt": {
+    "bytes": 10,
+    "sha256": "f5d5fb17b30809b00de20eeb6fe7b0d93a2a60ffd36339e06e875f5ba882de0c"
+  },
+  "charge_tokens.txt": {
+    "bytes": 7,
+    "sha256": "e8f533e310170a31a7e5e5fe125447fbe0253de205966ae295713c98c6ac531c"
+  },
+  "config.yaml": {
+    "bytes": 10254,
+    "sha256": "b2f460c9863459983689c7db35a8bf912f604b381dc3aeb06f77993196a2b566"
+  },
+  "distributions.json": {
+    "bytes": 1303,
+    "sha256": "309fcb1c579957c26b0db523953449ec76e0c58f2a793c6d4e924ca1a49dbe68"
+  },
+  "distributions.pt": {
+    "bytes": 2585,
+    "sha256": "93d265dd9afcf536270680492e4bf6ba776969d0811a32e918b8a8994189ccc6"
+  },
+  "edge_tokens.txt": {
+    "bytes": 18,
+    "sha256": "45ff09c48d164f6add1fa4b61d61564798d49dfebc697bf99a866dca60506d4c"
+  },
+  "epoch_1000.ckpt": {
+    "bytes": 508845203,
+    "sha256": "144410a5d6efcb326c3dfa466bbd3ec1fb203d4278923f2ec7d8a486ef9f8f1e"
+  },
+  "epoch_1500.ckpt": {
+    "bytes": 508845203,
+    "sha256": "90b418b9aa5de7cdb2109df20bc4ded1c46159b3abd10595750a5b583d0f9f87"
+  },
+  "metadata.json": {
+    "bytes": 1395,
+    "sha256": "6e54fdcd0481b12211232de7d75c050995f64a2d89b6da572b624af0300a06dc"
+  },
+  "runtime/README.md": {
+    "bytes": 302,
+    "sha256": "31e2f5e3dc2b9ededfb6c6d8f0de40092288fb2affcf0ba13fa667f997c374c6"
+  },
+  "runtime/configs/callbacks/callbacks.yaml": {
+    "bytes": 699,
+    "sha256": "dce0adebf5956dce687d64a074238f9c12866373bfe21df6658d967d6bfd1f90"
+  },
+  "runtime/configs/data/chembl3d/chembl3d.yaml": {
+    "bytes": 278,
+    "sha256": "a96974d90353a850c4a1d0bac27de5f1fc8724b81cc5bbc122aab6f76f3051c5"
+  },
+  "runtime/configs/data/chembl3d/implicit_h.yaml": {
+    "bytes": 81,
+    "sha256": "08da196d0288cf1fe4dba8dd407c49d9bee003607504beb59f79e5d6b9eaef0c"
+  },
+  "runtime/configs/data/chembl3d/max100.yaml": {
+    "bytes": 246,
+    "sha256": "052dba7b3f3efce990019518823476bd6bf1381bb5f606822c9fc7a485e6d0bf"
+  },
+  "runtime/configs/data/chembl3d/max100_implicit_h.yaml": {
+    "bytes": 79,
+    "sha256": "152325f11c7d1500469d4ee381100ae12247e6a2e0b90e5ef7501b405144bb2d"
+  },
+  "runtime/configs/data/geom/geom.yaml": {
+    "bytes": 428,
+    "sha256": "736c04e86c029bed65ce3ed022f937dd418ddc70d4b9fc064b38d979bf29ee6d"
+  },
+  "runtime/configs/data/geom/implicit_h.yaml": {
+    "bytes": 73,
+    "sha256": "c6b3ff5496af997de0c28e5aa7265903d6d0c8b59aab001ab21d6a64baeb772a"
+  },
+  "runtime/configs/data/geom/scaffold_decoration.yaml": {
+    "bytes": 1107,
+    "sha256": "d317728a98961e98a1f03bcc3b0f8c513169ca23b68ea07ba807b99066e7ca80"
+  },
+  "runtime/configs/data/molecular.yaml": {
+    "bytes": 5005,
+    "sha256": "f79ade5ddbf7d08f2eecbb2aa9c3e6b182b8a8e7e48960aff5e5edda5601bdc9"
+  },
+  "runtime/configs/data/qm9/qm9.yaml": {
+    "bytes": 3648,
+    "sha256": "3beea93868a1e889b049dcc2ccbc3b7489fa4672e964138654bc512ea01ab9f2"
+  },
+  "runtime/configs/data/qm9/scaffold_decoration.yaml": {
+    "bytes": 890,
+    "sha256": "40b36f2bc25bf4052cdd676a26d9e528cb4f7b13deddbb1660f454f97b159c40"
+  },
+  "runtime/configs/default.yaml": {
+    "bytes": 1294,
+    "sha256": "bcf9744acbf53afa4305ffd6e9cf95e6f99571750488eaa8cad85896a51aca66"
+  },
+  "runtime/configs/defaults/chembl3d/chembl3d.yaml": {
+    "bytes": 90,
+    "sha256": "f09553bc670e6c87ac8c1f8bbd0a8df73f1e83b02a3dcf9d6d06b915ac3bce93"
+  },
+  "runtime/configs/defaults/chembl3d/implicit_h.yaml": {
+    "bytes": 102,
+    "sha256": "a0ee1440984eb4c9d3c96a5d1c18df9f59b8e42c499f8a340ccd8deb8c454cd8"
+  },
+  "runtime/configs/defaults/chembl3d/implicit_h_2x4090.yaml": {
+    "bytes": 331,
+    "sha256": "518571405333a40e1602ee4c77fd153b56fd53e77aaa2aaf530796142f58d4b8"
+  },
+  "runtime/configs/defaults/chembl3d/train_2x4090.yaml": {
+    "bytes": 826,
+    "sha256": "0e6b01b264ab1ae70ae7f144ecd0a03d242e24f02feb087e3d5e2b7bc24790da"
+  },
+  "runtime/configs/defaults/core.yaml": {
+    "bytes": 431,
+    "sha256": "0dcfb68e2fece1dc277f647d3b3769af882cf8175e0c3b479bc50029e2538b1c"
+  },
+  "runtime/configs/defaults/geom/geom.yaml": {
+    "bytes": 425,
+    "sha256": "f447a1bdd39e211b5cf43c4cce56e708da6472d58f79e87885621ec4899c1c8e"
+  },
+  "runtime/configs/defaults/geom/implicit_h.yaml": {
+    "bytes": 98,
+    "sha256": "9d3a188e6bd2a2b01609ba0d652184bb8a2f7e27fd89d4a3031568051cae605e"
+  },
+  "runtime/configs/defaults/geom/scaffold_decoration.yaml": {
+    "bytes": 828,
+    "sha256": "bcce546bce2b0b7bfaf7c0663e05423fd03c952bf6e52cbbf31976723661f0d1"
+  },
+  "runtime/configs/defaults/qm9/qm9.yaml": {
+    "bytes": 429,
+    "sha256": "46070ea792300c0817c0acdff256181b21ef17dfe9f72070bb12da27dc28d534"
+  },
+  "runtime/configs/defaults/qm9/scaffold_decoration.yaml": {
+    "bytes": 348,
+    "sha256": "3bb8b6a757ea4921c31df2c0869949baef703df980b964d9636f52b64ac39267"
+  },
+  "runtime/configs/defaults/qm9/train_2x4090.yaml": {
+    "bytes": 554,
+    "sha256": "43a835246afb7126786e8ffa71c8b80ed9db1ef3a16e5240f1eec7882a4ef5ba"
+  },
+  "runtime/configs/embeddings/geom.yaml": {
+    "bytes": 803,
+    "sha256": "06793f4d446d899197e3db28d064592af9049612ee78cb0718a63e3d6fdc3151"
+  },
+  "runtime/configs/embeddings/qm9.yaml": {
+    "bytes": 849,
+    "sha256": "de57af7f1ff9b44dbb952f852c0177789196105b77f396298ff8e0218c0626b7"
+  },
+  "runtime/configs/heads/head.yaml": {
+    "bytes": 2475,
+    "sha256": "238e2346041e0aedccd16a6ed4066e163004c0e45a2cbc6fd1f6d17e16eb44ed"
+  },
+  "runtime/configs/insertion_head/gmm.yaml": {
+    "bytes": 488,
+    "sha256": "1e26914eed9b4971e0f6ff0b61bbdb34103233d346719a55ba6273dd3d785c05"
+  },
+  "runtime/configs/insertion_head/unimodal.yaml": {
+    "bytes": 495,
+    "sha256": "547a6391c3b341e3e369f8f4c7aab9752c67dfbbb533388d25956e9011b305e8"
+  },
+  "runtime/configs/integrator/integrator.yaml": {
+    "bytes": 1580,
+    "sha256": "8f0d9aafb0b2727cd58154e1fd97df693c8b7d4f5a3fb5efcb9188553f54660f"
+  },
+  "runtime/configs/logging/wandb_logger.yaml": {
+    "bytes": 139,
+    "sha256": "a9a2787e048d4eb70f2ff3c49d4355f1a5a382047c8619203b1f9131d829928d"
+  },
+  "runtime/configs/loss_weights/calibrated_geom.yaml": {
+    "bytes": 644,
+    "sha256": "d4763c84aa9e1eb77fb2b456384332a4b3a84b072d560918b2357222a43f4ca1"
+  },
+  "runtime/configs/loss_weights/calibrated_geom_frozen.yaml": {
+    "bytes": 1141,
+    "sha256": "a08acb61d67a852051dbded08f8bb67f21101888fd0a22a96a522781d28c570d"
+  },
+  "runtime/configs/loss_weights/tuned_weights.yaml": {
+    "bytes": 420,
+    "sha256": "2bb5884465666d6fc638f787fd79257e4126c69a5cc55b55615fc45e02e65894"
+  },
+  "runtime/configs/loss_weights/uniform.yaml": {
+    "bytes": 336,
+    "sha256": "3c14d6443048f69d436fe03792703a33c67f368ecbb349de2fd145da4c96eabc"
+  },
+  "runtime/configs/model/backbones/dit.yaml": {
+    "bytes": 4039,
+    "sha256": "65b0f08be8fb307b8b01e7b460fd083fe60d48148fbece378d7e5a23dd4a9fb0"
+  },
+  "runtime/configs/model/finetune/lora_finetune.yaml": {
+    "bytes": 328,
+    "sha256": "6334a5112499c0f3bc1fe5de224e2b6b3265dc524c8d6b9048113868fb864558"
+  },
+  "runtime/configs/model/finetune/scaffold_finetune.yaml": {
+    "bytes": 5659,
+    "sha256": "1aad0d33e9b5ae95d9f50b5201a29bb8677fe42ea2bcf59e0e61d053ef1dfddb"
+  },
+  "runtime/configs/optimizer/muon_aux_adam.yaml": {
+    "bytes": 1212,
+    "sha256": "3b0a6189f3ba81f85f7d26c6f6ad7f64dcb25098a14e411211d7a62983a7180c"
+  },
+  "runtime/configs/schedule/beta.yaml": {
+    "bytes": 78,
+    "sha256": "75c485edfd70016b255282adcf5de25dfb95c9e32bb6587017fb8935ffe64eaa"
+  },
+  "runtime/configs/schedule/cubic.yaml": {
+    "bytes": 54,
+    "sha256": "a998a2d3390ab5daf4df079852ae8daa907687cf059f55c500409c9493949c0c"
+  },
+  "runtime/configs/schedule/fast_power.yaml": {
+    "bytes": 126,
+    "sha256": "22f36106bd43ca253ffea412e00539f2b07efbc49109b97c1232df6a827a2baf"
+  },
+  "runtime/configs/schedule/linear.yaml": {
+    "bytes": 55,
+    "sha256": "f7881af26080ab5569e67b39f6eae2733c033b9623020e6271959f09fac45fa0"
+  },
+  "runtime/configs/schedule/smoothstep.yaml": {
+    "bytes": 70,
+    "sha256": "f6e8651953d93f0d18c997d40c199b3d80ddd352a2041cc51e29adc9a7f22778"
+  },
+  "runtime/configs/trainer/trainer.yaml": {
+    "bytes": 1402,
+    "sha256": "6f78e83ad46529bee7a2878b20727ff403420a515c1629bd73a608a2fb654f5f"
+  },
+  "runtime/pyproject.toml": {
+    "bytes": 4464,
+    "sha256": "d4854128f81144905cef68d37222d0d06fb372e2bd148e49177e0ff082a7b01f"
+  },
+  "runtime/requirements-tested.txt": {
+    "bytes": 177,
+    "sha256": "fac030b53775fd8ed8f89b17050ea74180f99dfab270b1d75d4ca3c0573ede76"
+  },
+  "runtime/src/morph/__init__.py": {
+    "bytes": 79,
+    "sha256": "bb402ea5c6f79e569205cc2bf4168d06430753ce8499ca4ffb8614743626fd59"
+  },
+  "runtime/src/morph/__main__.py": {
+    "bytes": 149,
+    "sha256": "95a7983479559ca2b15c7afa3881b3142ef4636583501184469b972df7646d4f"
+  },
+  "runtime/src/morph/_vendor/__init__.py": {
+    "bytes": 72,
+    "sha256": "8480cd966a03870a7142e1763094bf00e221d2c3c2b1af1d4d458fb8b7f37a20"
+  },
+  "runtime/src/morph/_vendor/geom_drugs_preprocessing.py": {
+    "bytes": 11328,
+    "sha256": "33a9c21d29df269a5f342a0b0e8bced0fc620559f0cb3a7a4cfabd5c369ceccb"
+  },
+  "runtime/src/morph/_vendor/muon.py": {
+    "bytes": 16405,
+    "sha256": "c61af8c2c4a635cc682b84e279851cb54fedc6fb935d8eb11aafbeb1c5e81fc6"
+  },
+  "runtime/src/morph/cli.py": {
+    "bytes": 9723,
+    "sha256": "27eee3fd6d6410747f12b19b670a74f50a51b842a50b1eff4e75c46d0d08a986"
+  },
+  "runtime/src/morph/config.py": {
+    "bytes": 704,
+    "sha256": "bd3562c030611f87cc514308f155ad93740dd59927732f71b2e15009012d78ff"
+  },
+  "runtime/src/morph/dataset/__init__.py": {
+    "bytes": 620,
+    "sha256": "44ae493354fed0fec4b410f46e0c1fb0351b5fbe5f404ff50f77181e60a39cee"
+  },
+  "runtime/src/morph/dataset/conditional/__init__.py": {
+    "bytes": 238,
+    "sha256": "e46c747069acbf3d82407ccfcf96f2e33608c010a429f9abdb064bb71795d134"
+  },
+  "runtime/src/morph/dataset/conditional/scaffold/__init__.py": {
+    "bytes": 167,
+    "sha256": "6ef414b31cceb7303c174d720e0d0666efd6f6678c97b3cf86d33c799d2153dd"
+  },
+  "runtime/src/morph/dataset/conditional/scaffold/cutters.py": {
+    "bytes": 12282,
+    "sha256": "b29e0d7b1758461cda868bcba56d58b21ce76a45a9365bed73f7a2735103616f"
+  },
+  "runtime/src/morph/dataset/conditional/scaffold/datamodule.py": {
+    "bytes": 1402,
+    "sha256": "ecd4a5819d3654f8f76e3e0d059a60bfc9d0ac2ec343514d45517e7094c7c282"
+  },
+  "runtime/src/morph/dataset/conditional/scaffold/lmdb.py": {
+    "bytes": 4243,
+    "sha256": "54ea075ebb4a8ed481f93f57e9a5e25504918b6595bde106f2489e21e7e73b5b"
+  },
+  "runtime/src/morph/dataset/conditional/scaffold/wrappers.py": {
+    "bytes": 7720,
+    "sha256": "a44d61452b676dc1dfd8d6bd44c09b22d463c8129b23cf98ec3903d10a693f2c"
+  },
+  "runtime/src/morph/dataset/core/__init__.py": {
+    "bytes": 376,
+    "sha256": "7bab9c675f0850efc50e14a9898accfde5f8d1c1e2420414ac8c0986822d27e0"
+  },
+  "runtime/src/morph/dataset/core/atomic_data.py": {
+    "bytes": 708,
+    "sha256": "7d59ec16ed5b3cf82b57b49dd4d7799862709353a65f4d473d4571ba91b86fa6"
+  },
+  "runtime/src/morph/dataset/core/batch.py": {
+    "bytes": 23924,
+    "sha256": "7b2b2afd749b621b1ca890d192e4dc9378b61ead9ef34b90f04dc449e6bad369"
+  },
+  "runtime/src/morph/dataset/core/containers.py": {
+    "bytes": 19637,
+    "sha256": "adae656c3fbbe5cfe32f378a419c7e1076f653163ffe7d2c1d1d0abc71a99cd7"
+  },
+  "runtime/src/morph/dataset/core/conversion.py": {
+    "bytes": 2224,
+    "sha256": "0842565c503048a86a517c202dd0cc7c07e9296a841601b178ec38c43f903a6a"
+  },
+  "runtime/src/morph/dataset/core/flow_matching_wrapper.py": {
+    "bytes": 7568,
+    "sha256": "7672ce6dc416bbccae1471b8e4fb66f44adf7eaf8f0f2fa151f1268e9513a17d"
+  },
+  "runtime/src/morph/dataset/core/insertion_data.py": {
+    "bytes": 14393,
+    "sha256": "7e99e4d59b29ea537708b5f7887ed3e61b4b31ebf80d55fe5a0c6823ff8a439c"
+  },
+  "runtime/src/morph/dataset/core/insertion_targets.py": {
+    "bytes": 6614,
+    "sha256": "2e62e89125310572735d916e97c8a544cfb3eb76fb167ecd134b5f8b0424df8c"
+  },
+  "runtime/src/morph/dataset/core/lightning_datamodule.py": {
+    "bytes": 7057,
+    "sha256": "4d525cbcde4a64d7cac41837e0379c40097cc38fbbc2f54ff30ff7e3ab83de4c"
+  },
+  "runtime/src/morph/dataset/core/lmdb.py": {
+    "bytes": 10111,
+    "sha256": "748b7a10f8fb20f74c267de393cc3e1933ea86debde99d8fdaaef94e0a9afd29"
+  },
+  "runtime/src/morph/dataset/core/preprocessing.py": {
+    "bytes": 20145,
+    "sha256": "da9cffd6beb6da63ee749dd8b60d4d0509bcd937778014739b04ffc599dde066"
+  },
+  "runtime/src/morph/dataset/core/properties.py": {
+    "bytes": 1952,
+    "sha256": "4dc63574f5fde284245078adf50fb2d8a9ac16443b8276154c03b7703cc2e937"
+  },
+  "runtime/src/morph/dataset/core/representation.py": {
+    "bytes": 6591,
+    "sha256": "0ecac55e1faf9748826ec8f4c7f1a274263214ea3f68cc69de7047b55dd6a65f"
+  },
+  "runtime/src/morph/dataset/core/tokenized_decode.py": {
+    "bytes": 6559,
+    "sha256": "483afb3d468113da33cc8699998e6439fcde2264ee7ca479724206360584a2b8"
+  },
+  "runtime/src/morph/dataset/core/vocab.py": {
+    "bytes": 2236,
+    "sha256": "f8618488f7bcdc8a58d8b9f3443b764eac60985a33ead1bcc2d17d3bdc213329"
+  },
+  "runtime/src/morph/dataset/unconditional/__init__.py": {
+    "bytes": 403,
+    "sha256": "d7cd3b85745ad1a4d884c0a11a64b3ab9cedcc64e30f604dbe83f5beb14fa97b"
+  },
+  "runtime/src/morph/dataset/unconditional/chembl3d.py": {
+    "bytes": 9854,
+    "sha256": "d5cfd1ee4ddd3085b95b9a678d761d4f0beef99ac5ecfb21d3dfb3cae59653eb"
+  },
+  "runtime/src/morph/dataset/unconditional/geom.py": {
+    "bytes": 4071,
+    "sha256": "12bc792c2a61b3e6c29ae265b9997ba610765e268408e873918b5a65225945be"
+  },
+  "runtime/src/morph/dataset/unconditional/qm9.py": {
+    "bytes": 9520,
+    "sha256": "06b6086cb1234b283cd3a45daad2e06d3bdb5d45810b873f860c82969ccbfe53"
+  },
+  "runtime/src/morph/flow_matching/__init__.py": {
+    "bytes": 78,
+    "sha256": "674f4630aa8421d7e75c913741c6efb1a163fbc38ce909c6cdb011f9ba9ef48e"
+  },
+  "runtime/src/morph/flow_matching/assignment.py": {
+    "bytes": 16050,
+    "sha256": "b6ad66538f5bd7f14e0eee2a0ffd25a4110d1504e057a0a9eb63a354d9f6aadb"
+  },
+  "runtime/src/morph/flow_matching/integration.py": {
+    "bytes": 33356,
+    "sha256": "f64b7b9ed2a79c11932b72e498e57eef924b5c85a05aa6893ddfe11e2d4595e9"
+  },
+  "runtime/src/morph/flow_matching/interpolation.py": {
+    "bytes": 22169,
+    "sha256": "6ef2e03626fd4b3194c68df6e2410ab6e630f4891f97a8328b4eb0b3e18c7dcd"
+  },
+  "runtime/src/morph/flow_matching/prior.py": {
+    "bytes": 2389,
+    "sha256": "0f9b77274238b23be7a2ad88ee8e6dcbb785817fcfa3a35aa863b489d2a6e05d"
+  },
+  "runtime/src/morph/flow_matching/sampling_step.py": {
+    "bytes": 2806,
+    "sha256": "c768682a47da1b7852838e20b39317528717fc0485cbc35030b3ca71d7d845bb"
+  },
+  "runtime/src/morph/flow_matching/schedules.py": {
+    "bytes": 6776,
+    "sha256": "6fd8396a05052ae7831babc75ca40a7e748fb9edffe82bb2559533afa949b2a9"
+  },
+  "runtime/src/morph/inference/__init__.py": {
+    "bytes": 732,
+    "sha256": "63047340b4c7ed784e969b7d54035ab737282ee2acc6eeb26ba6cb3d52b4f573"
+  },
+  "runtime/src/morph/inference/bundles.py": {
+    "bytes": 6395,
+    "sha256": "f6e2d46e6e428452a58208e78c83da01d6c786f10f7062643a093fb6750dafe6"
+  },
+  "runtime/src/morph/inference/catalog.py": {
+    "bytes": 6215,
+    "sha256": "53a27b86d39641fd4964da82f4b1875598a0cdb1533f981074c6b7a14b4174b7"
+  },
+  "runtime/src/morph/inference/generate.py": {
+    "bytes": 11889,
+    "sha256": "c3181fbf6187cf22b476431ed7c0fdea4c74ea8e97acf270f486318e43646d19"
+  },
+  "runtime/src/morph/inference/hub.py": {
+    "bytes": 9739,
+    "sha256": "fa1f43937d66f0ab4b7b9c11ed749aef8d23444b2df3257c36ff72ecd3b821ca"
+  },
+  "runtime/src/morph/inference/module.py": {
+    "bytes": 3820,
+    "sha256": "415e0f9c624c191855874b8dd88d5e7aa18941632bc403a9c0b685bbf74f1d60"
+  },
+  "runtime/src/morph/inference/registry.py": {
+    "bytes": 4314,
+    "sha256": "8ca3e4246e00707ed0d646e6e760acebd5e3d3443e5e8d5f71f3b54653993fa9"
+  },
+  "runtime/src/morph/model/README.md": {
+    "bytes": 10319,
+    "sha256": "9c06b5b90300f7449048640417116aa8869dc1a75f106cbdd9344a89f6e446d3"
+  },
+  "runtime/src/morph/model/__init__.py": {
+    "bytes": 80,
+    "sha256": "1c7c3875387d5ba7f9166a804ea5a7ef7bd469728c1ad5c5003150d44e33be51"
+  },
+  "runtime/src/morph/model/backbones/__init__.py": {
+    "bytes": 64,
+    "sha256": "a1a1419590c11aa6ed429f8b6b293470e30fe767f5571d2517b0883efb071e0c"
+  },
+  "runtime/src/morph/model/backbones/conditional_dit.py": {
+    "bytes": 25451,
+    "sha256": "411fe1d0fd17c24e03e153d57f299aa0436bc5749e2064d5574681878429795d"
+  },
+  "runtime/src/morph/model/backbones/dit.py": {
+    "bytes": 31262,
+    "sha256": "e2798baaf31bc8004a17c2432c56af9c4b12651d259cc0deaf5f18c83c6ddd33"
+  },
+  "runtime/src/morph/model/build.py": {
+    "bytes": 3442,
+    "sha256": "e69f4231815f09c094d787f11edc49e222cbdc970edc6e416b6b94a9b88d4b21"
+  },
+  "runtime/src/morph/model/finetune/__init__.py": {
+    "bytes": 282,
+    "sha256": "8746b1e29e5a3b3b1805389257270b90f5d30995783f2cf41b9fb533c485fc0c"
+  },
+  "runtime/src/morph/model/finetune/base/__init__.py": {
+    "bytes": 803,
+    "sha256": "6e0ded0b3a5c2a76ccb85b05cf63bcc54e145bd1be07d933a737d6151433d856"
+  },
+  "runtime/src/morph/model/finetune/base/conditional_backbone.py": {
+    "bytes": 2853,
+    "sha256": "1353e96443eacf3ce88473a30d02ab4f1c8b6a9c90763ea61c58d353b806b378"
+  },
+  "runtime/src/morph/model/finetune/base/context_encoder.py": {
+    "bytes": 15124,
+    "sha256": "8d5dfc91ec73120161ed0a51dcc7edac442004ad4d5213ba34d9d11cfd1eee15"
+  },
+  "runtime/src/morph/model/finetune/base/finetune_base.py": {
+    "bytes": 9526,
+    "sha256": "2efa00fd8eed1096a9cf8799dee4abfd8b27de1c0600e22a906dddaa089c4ca0"
+  },
+  "runtime/src/morph/model/finetune/scaffold/__init__.py": {
+    "bytes": 366,
+    "sha256": "6345e71c8b03ee5eb0dcd105c0532acf73a1795bf6360b2847be600bcecd40ba"
+  },
+  "runtime/src/morph/model/finetune/scaffold/scaffold.py": {
+    "bytes": 1984,
+    "sha256": "6de7a6083b4303ae422f03f15b5b40b1233d48bfc2dd1e0aa83a4eaa8adf40dd"
+  },
+  "runtime/src/morph/model/finetune/scaffold/scaffold_context.py": {
+    "bytes": 6086,
+    "sha256": "d9446254770e3c3ea2c634a0636ab82a718514a7cb69760d756e9235cfbe2814"
+  },
+  "runtime/src/morph/model/layers/__init__.py": {
+    "bytes": 189,
+    "sha256": "79671fbdea358c330a45b7fd3d842119b32b1d548804b3d441bbfb7dbd703b69"
+  },
+  "runtime/src/morph/model/layers/attention.py": {
+    "bytes": 21038,
+    "sha256": "7efc416ac8fec0236909661ad0693a0aef613ff97a964fb6d8252cb49c60b5c6"
+  },
+  "runtime/src/morph/model/layers/embedding.py": {
+    "bytes": 14356,
+    "sha256": "efe924190da5904e05495f421dcc313134cebe0dee9bf665060e1e5781cd95ef"
+  },
+  "runtime/src/morph/model/layers/gmm.py": {
+    "bytes": 6753,
+    "sha256": "628e90c73bdbb3e70104214f80bbf8a94896ddbe22a9a3fc40cce34b20d2565f"
+  },
+  "runtime/src/morph/model/layers/heads.py": {
+    "bytes": 30913,
+    "sha256": "7113c29a6be81c93ca980dc19d0847d2015ac36702683e4f1adedd549aed46ae"
+  },
+  "runtime/src/morph/model/layers/kernels.py": {
+    "bytes": 12483,
+    "sha256": "be0c52b69c7612946244755e2865de0491440a41163cf828168eb750a8f929cf"
+  },
+  "runtime/src/morph/model/layers/lora.py": {
+    "bytes": 5947,
+    "sha256": "7194750605e7b5df9a02e4c4a39cf29c982f117c40198c99c73211c85ecdfe4b"
+  },
+  "runtime/src/morph/model/lightning/__init__.py": {
+    "bytes": 339,
+    "sha256": "c53943a4eb1e45568cebfb538655cc5eb7ead46284bbcd515800102c89177e52"
+  },
+  "runtime/src/morph/model/lightning/_optimizer.py": {
+    "bytes": 19578,
+    "sha256": "f77519db8d3458c05aec1543cb4cd487995cf398abcc71b62ea1835cb132165e"
+  },
+  "runtime/src/morph/model/lightning/_sampling.py": {
+    "bytes": 15344,
+    "sha256": "2694ba5cd0b315147285ba268ed5ad246d6091d2426fcd7b21f78d0eb0a8fca1"
+  },
+  "runtime/src/morph/model/lightning/_setup.py": {
+    "bytes": 6537,
+    "sha256": "835f510094110c4dbc590a88d9b94d41474ae87b442785a949328feda5b9afce"
+  },
+  "runtime/src/morph/model/lightning/_training.py": {
+    "bytes": 22479,
+    "sha256": "0451c3c4f1ad46cc5380615a3fbf8452d1fb4a0d27d69c885a411a3612c0997c"
+  },
+  "runtime/src/morph/model/lightning/_validation.py": {
+    "bytes": 11509,
+    "sha256": "03890457e504eed96f23dd8633a6af4317a9ed61e80b3e4ec935aaf9e63038d7"
+  },
+  "runtime/src/morph/model/lightning/base.py": {
+    "bytes": 9294,
+    "sha256": "ce9159e7b387ece0a3ce010b164f675b81980d2f617b5b05d00f96695ee33e28"
+  },
+  "runtime/src/morph/model/losses/__init__.py": {
+    "bytes": 190,
+    "sha256": "bb9de5950496ed184a32e11834d8e6f3b7d406773509ffa6c406a86bb5249fae"
+  },
+  "runtime/src/morph/model/losses/functional.py": {
+    "bytes": 11765,
+    "sha256": "42ae52f604a0cb8b7fe08955cfd273945ca068727994f20b02c80cdee3a00e1b"
+  },
+  "runtime/src/morph/utils/__init__.py": {
+    "bytes": 78,
+    "sha256": "eb76ca86f046931eaea693a696317e8b3ce0b8ad8695a92e3744fd1892f2b9b8"
+  },
+  "runtime/src/morph/utils/checkpoints.py": {
+    "bytes": 2301,
+    "sha256": "750302eb039e3947a5d4c84e4e7d04eb46e974cc0031ef1d197a3573e20eae4b"
+  },
+  "runtime/src/morph/utils/edge_alignment.py": {
+    "bytes": 4896,
+    "sha256": "0edcde93cfa3a2d825e783ea5cf64715aa248e59ffb8c6a23d4490d74976a71e"
+  },
+  "runtime/src/morph/utils/graph_math.py": {
+    "bytes": 6755,
+    "sha256": "2bae326c3538d757aac66dcfc4a0b5d023b17bd987d814d3dab9a57f9ce60552"
+  },
+  "runtime/src/morph/utils/loss_accumulation.py": {
+    "bytes": 15115,
+    "sha256": "0f3eca42145e715290576dd86f4bcfb1553714ec4cabeda9246cfaf1e4f4323c"
+  },
+  "runtime/src/morph/utils/lr_schedulers.py": {
+    "bytes": 5554,
+    "sha256": "b954c12df413ca6676263ef1af960e2a4ff52692acfd4223e72a6e02122a3638"
+  },
+  "runtime/src/morph/utils/metrics/__init__.py": {
+    "bytes": 6730,
+    "sha256": "abec143d21e36343fff54db0ca8042111b9e70da426567b13ec36d6cb05696b6"
+  },
+  "runtime/src/morph/utils/metrics/chemistry.py": {
+    "bytes": 34152,
+    "sha256": "61b598b99fad41571f9334b28f94fbeeeb0282d1b80423b1f2bc8f14180884e8"
+  },
+  "runtime/src/morph/utils/metrics/distribution.py": {
+    "bytes": 3348,
+    "sha256": "0a6c25c4e65843308059a3effc44e1072837e747926074c1d7c47ba7ee72cb46"
+  },
+  "runtime/src/morph/utils/metrics/plotting.py": {
+    "bytes": 4270,
+    "sha256": "686d491d7fbc83cf35f772649a869dc1a271966d2a4fa6162d6e2334db735723"
+  },
+  "runtime/src/morph/utils/metrics/posecheck_strain.py": {
+    "bytes": 6695,
+    "sha256": "60a4751cce0e23d87d8602135de0c9850182fe97f20632cb9095445bdc6b0365"
+  },
+  "runtime/src/morph/utils/metrics/tensor.py": {
+    "bytes": 13503,
+    "sha256": "3fcb867cde5d47f81725acc3b862fc56af5ecce23dee3a11ebe6d05c1772c6fe"
+  },
+  "runtime/src/morph/utils/rdkit_utils.py": {
+    "bytes": 9018,
+    "sha256": "0fddc055c6de60893272fe8c50c92de29c9e72b7819827caac753e4cd0c6f1e5"
+  },
+  "runtime/src/morph/utils/run_config.py": {
+    "bytes": 8025,
+    "sha256": "0c289120887a051379523af528ebc9310f356f3b92efb898ddaf773dbf698aa2"
+  },
+  "runtime/src/morph/utils/vocab_utils.py": {
+    "bytes": 6343,
+    "sha256": "7b5506971b5092865b86311173b24d7fb70a881f123c57fd8696c412a25d618b"
+  },
+  "runtime/uv.lock": {
+    "bytes": 611263,
+    "sha256": "6fd6609500ea661b3b5c9f17d0997c5709295ca2cbbe16c97fc8259178fa931e"
+  },
+  "training_config.yaml": {
+    "bytes": 12028,
+    "sha256": "b95ee96def79a9602c4b635e7432d97a75a1427dc141d98c950ff38d6012ef6a"
+  },
+  "validation.json": {
+    "bytes": 781,
+    "sha256": "5b128cbe9fce32cfbdb1ebe63d62a35b50d1b391bac1c0ab4ae25e0d22fa3fdd"
+  }
+}

qm9/explicit_h/metadata.json

@@ -0,0 +1,55 @@
+{
+  "dataset": "qm9",
+  "variant": "explicit_h",
+  "includes_explicit_hydrogens": true,
+  "training_run": "qm9-explicit-25ep-16529301",
+  "resumed_from_completed_epoch": null,
+  "source_snapshot": "qm9_launch_20261008",
+  "checkpoint_format": "Original Lightning checkpoint, current + EMA + optimizer state",
+  "epoch_numbering": "Completed epochs (one-based); original checkpoint epoch is zero-based",
+  "checkpoints": [
+    {
+      "filename": "epoch_1000.ckpt",
+      "original_filename": "epoch=999-step=24000.ckpt",
+      "completed_epochs": 1000,
+      "checkpoint_epoch_index": 999,
+      "global_step": 24000,
+      "bytes": 508845203,
+      "sha256": "144410a5d6efcb326c3dfa466bbd3ec1fb203d4278923f2ec7d8a486ef9f8f1e"
+    },
+    {
+      "filename": "epoch_1500.ckpt",
+      "original_filename": "epoch=1499-step=36000.ckpt",
+      "completed_epochs": 1500,
+      "checkpoint_epoch_index": 1499,
+      "global_step": 36000,
+      "bytes": 508845203,
+      "sha256": "90b418b9aa5de7cdb2109df20bc4ded1c46159b3abd10595750a5b583d0f9f87"
+    }
+  ],
+  "training_split_sizes": {
+    "train": 100000,
+    "validation": 20000,
+    "test": 12947
+  },
+  "atom_tokens": [
+    "C",
+    "F",
+    "H",
+    "N",
+    "O"
+  ],
+  "edge_tokens": [
+    "<NO_BOND>",
+    "1",
+    "2",
+    "3",
+    "4"
+  ],
+  "charge_tokens": [
+    "-1",
+    "0",
+    "1"
+  ],
+  "coordinate_std_angstrom": 1.373397946357727
+}

qm9/explicit_h/runtime/README.md

@@ -0,0 +1,5 @@
+# Morph inference runtime
+
+Archived source used for the QM9 training run.
+
+Install with `python -m pip install .`, then use `morph models list` or `from morph.inference import load_pretrained`. See the model repository at https://huggingface.co/maltefranke/morph for checkpoint loading and generation.

qm9/explicit_h/runtime/configs/callbacks/callbacks.yaml

@@ -0,0 +1,20 @@
+# @package _global_
+
+## Put your callback args here
+callbacks:
+  monitor_metric: ${if:${oc.eval:"'${representation}' in ('geometric_graph', 'molecule')"},val/validity,val/batch/pair_dist_l1}
+  monitor_metric_mode: ${if:${oc.eval:"'${representation}' in ('geometric_graph', 'molecule')"},max,min}
+
+  #   model_checkpoints         -> best-by-metric (top-1, monitored)
+  #   every_n_epochs_checkpoint -> latest (no monitor; overwrites every epoch)
+  model_checkpoints:
+    save_top_k: 1
+    save_last: False
+    verbose: False
+
+  every_n_epochs_checkpoint:
+    every_n_epochs: 1
+    save_top_k: 1
+    save_last: False
+    verbose: False
+    save_on_train_epoch_end: True # Force per-train-epoch firing

qm9/explicit_h/runtime/configs/data/chembl3d/chembl3d.yaml

@@ -0,0 +1,8 @@
+# LoQI ChEMBL3D source format, without dataset-specific size/element filters.
+defaults:
+  - /data/molecular@_here_
+  - _self_
+
+dataset_name: chembl3d
+dataset_target: morph.dataset.unconditional.chembl3d.FlowMatchingChEMBL3DDataset
+source_dir: ${oc.env:CHEMBL3D_SOURCE_DIR,null}

qm9/explicit_h/runtime/configs/data/chembl3d/implicit_h.yaml

@@ -0,0 +1,5 @@
+defaults:
+  - /data/chembl3d/chembl3d@_here_
+  - _self_
+
+hydrogen_mode: implicit

qm9/explicit_h/runtime/configs/data/chembl3d/max100.yaml

@@ -0,0 +1,7 @@
+# Inclusive total-atom limit, applied before hydrogen removal.
+defaults:
+  - /data/chembl3d/chembl3d@_here_
+  - _self_
+
+max_atoms: 100
+cache_name: ${data.dataset_name}_max${data.max_atoms}${if:${eq:${data.hydrogen_mode},implicit},_implicit_h,''}

qm9/explicit_h/runtime/configs/data/chembl3d/max100_implicit_h.yaml

@@ -0,0 +1,5 @@
+defaults:
+  - /data/chembl3d/max100@_here_
+  - _self_
+
+hydrogen_mode: implicit

qm9/explicit_h/runtime/configs/data/geom/geom.yaml

@@ -0,0 +1,11 @@
+# GEOM-Drugs source format and default molecule subset.
+defaults:
+  - /data/molecular@_here_
+  - _self_
+
+dataset_name: geom
+dataset_target: morph.dataset.unconditional.geom.FlowMatchingGEOMDataset
+source_dir: ${oc.env:GEOM_SOURCE_DIR,null} # Optional prefiltered pickle directory.
+# Preserve the established GEOM subset. Limits apply before hydrogen removal.
+max_atoms: 72
+allowed_atom_symbols: [H, C, N, O, S, F, Cl, Br, P, I]

qm9/explicit_h/runtime/configs/data/geom/implicit_h.yaml

@@ -0,0 +1,5 @@
+defaults:
+  - /data/geom/geom@_here_
+  - _self_
+
+hydrogen_mode: implicit

qm9/explicit_h/runtime/configs/data/geom/scaffold_decoration.yaml

@@ -0,0 +1,24 @@
+# GEOM-Drugs data config for the **scaffold decoration** task.
+#
+# Mirrors configs/data/qm9/scaffold_decoration.yaml but on GEOM-Drugs: source
+# is the selected scaffold of each target; target is the full molecule. The model
+# learns to grow it by inserting decoration atoms around the fixed scaffold.
+#
+# Inherits from configs/data/geom/geom.yaml and overrides:
+#   - datamodule._target_ -> ScaffoldLightningDataModule
+#
+# The scaffold cache is built lazily into ``scaffold_groups_dir`` on first run
+# (one file per split and scaffold kind) and reused afterwards.
+defaults:
+  - /data/geom/geom@_here_
+  - _self_
+
+scaffold_groups_dir: ${data.preprocessing.root}/processed/scaffold_groups
+scaffold_kind: rgroup # murcko, rgroup (BRICS), heavy (generate H), linker
+# Evaluate each distinct molecule/scaffold pair; training always draws randomly.
+enumerate_slicings: false
+datamodule:
+  _target_: morph.dataset.conditional.scaffold.datamodule.ScaffoldLightningDataModule
+  scaffold_groups_dir: ${data.scaffold_groups_dir}
+  scaffold_kind: ${data.scaffold_kind}
+  enumerate_slicings: ${data.enumerate_slicings}

qm9/explicit_h/runtime/configs/data/molecular.yaml

@@ -0,0 +1,136 @@
+# Shared drug-like molecular training and preprocessing configuration.
+
+root_path: ${oc.env:PROJECT_ROOT}/data
+preprocess_workers: 10
+preprocess_chunk_size: 128
+statistics_workers: ${data.preprocess_workers}
+hydrogen_mode: explicit
+# The mode is part of the cache path; changing it discovers fresh vocab/statistics.
+cache_name: ${data.dataset_name}${if:${eq:${data.hydrogen_mode},implicit},_implicit_h,''}
+source_dir: null
+max_atoms: null
+allowed_atom_symbols: null
+
+n_atoms_strategy: "uniform" # "approx", "empirical", "learned"
+# Uniform source-count support padding; explicit so inference and training can
+# be matched or ablated without code changes.
+uniform_n_atoms_padding: 5
+optimal_transport: "equivariant"
+# Physical Angstroms (converted to normalized coordinates internally). These
+# reproduce the former normalized 0.25 / 0.2 scales at GEOM's training-set std.
+# Physical Angstroms. Insertion bridge noise vanishes as (1-t); move noise is
+# the persistent training-time jitter floor for matched atoms and insertions
+# that are already present.
+ins_noise_scale: 0.65
+move_noise_scale: 0.5
+
+# Drug-like datasets contain charged species; allow them in validity metrics.
+allow_charged: true
+
+
+# Preprocessing configuration - computes tokens from data and distributions from training dataset only
+preprocessing:
+  _target_: morph.dataset.core.preprocessing.Preprocessing
+  _recursive_: false
+  root: ${data.root_path}/${data.cache_name}
+  num_workers: ${data.statistics_workers}
+  train_dataset: ${data.datamodule.datasets.train}
+  atom_tokens_path: ${data.preprocessing.root}/processed/atom_tokens.txt
+  edge_tokens_path: ${data.preprocessing.root}/processed/edge_tokens.txt
+  charge_tokens_path: ${data.preprocessing.root}/processed/charge_tokens.txt
+  distributions_path: ${data.preprocessing.root}/processed/distributions.pt
+  # Tokens and distributions will be discovered from data and saved to files
+  # On subsequent runs, they will be loaded from the files
+
+vocab: null
+
+interpolator:
+  _target_: morph.flow_matching.interpolation.Interpolator
+  distributions: null
+  optimal_transport: ${data.optimal_transport}
+  ins_noise_scale: ${data.ins_noise_scale}
+  move_noise_scale: ${data.move_noise_scale}
+  # Training-only categorical corruption; pairs are scaled per molecule to
+  # match the expected number of changed atoms. Clean targets stay unchanged.
+  atom_type_noise_prob: 0.0
+  # OT matching (cost always squared-Euclidean; move cost in Angstrom). c_move
+  # scales the squared move distance (A^2); c_ins/c_del/c_sub are fixed A^2 costs,
+  # directly comparable. At 6.125 each, delete+insert costs 3.5^2, so atoms
+  # farther than 3.5 A apart are left unmatched by default.
+  pre_align_iters: 2   # equivariant-OT inner loop (Hungarian+Kabsch) iters
+  c_move: 1.0
+  c_sub: 0.0
+  c_ins: 6.125
+  c_del: 6.125
+  del_schedule: ${schedule}
+  ins_schedule: ${schedule}
+  sub_schedule:
+    _target_: morph.flow_matching.schedules.SmoothstepSchedule
+    shift: 1.5
+  sub_e_schedule:
+    _target_: morph.flow_matching.schedules.SmoothstepSchedule
+    shift: 1.5
+
+datamodule:
+  _target_: morph.dataset.core.lightning_datamodule.LightningDataModule
+  # tokens and distributions will be passed from preprocessing in run.py
+
+  vocab: ${data.vocab}
+  distributions: null
+  representation: ${representation}
+
+  datasets:
+    train:
+      _target_: ${data.dataset_target}
+      root: ${data.preprocessing.root}
+      preprocess_workers: ${data.preprocess_workers}
+      preprocess_chunk_size: ${data.preprocess_chunk_size}
+      hydrogen_mode: ${data.hydrogen_mode}
+      source_dir: ${data.source_dir}
+      max_atoms: ${data.max_atoms}
+      allowed_atom_symbols: ${data.allowed_atom_symbols}
+      vocab: ${data.vocab}
+      distributions: null
+
+    val:
+      - _target_: ${data.dataset_target}
+        root: ${data.preprocessing.root}
+        preprocess_workers: ${data.preprocess_workers}
+        preprocess_chunk_size: ${data.preprocess_chunk_size}
+        hydrogen_mode: ${data.hydrogen_mode}
+        source_dir: ${data.source_dir}
+        max_atoms: ${data.max_atoms}
+        allowed_atom_symbols: ${data.allowed_atom_symbols}
+        vocab: ${data.vocab}
+        distributions: null
+
+    test:
+      - _target_: ${data.dataset_target}
+        root: ${data.preprocessing.root}
+        preprocess_workers: ${data.preprocess_workers}
+        preprocess_chunk_size: ${data.preprocess_chunk_size}
+        hydrogen_mode: ${data.hydrogen_mode}
+        source_dir: ${data.source_dir}
+        max_atoms: ${data.max_atoms}
+        allowed_atom_symbols: ${data.allowed_atom_symbols}
+        vocab: ${data.vocab}
+        distributions: null
+
+  interpolator: ${data.interpolator}
+
+  rotate: ${model.rotate}
+  n_augmentations: ${model.n_augmentations}
+
+  num_workers:
+    train: 10
+    val: 0
+    test: 0
+
+  batch_size:
+    train: 320
+    val: 128
+    test: 64
+
+  n_atoms_strategy: ${data.n_atoms_strategy}
+  uniform_n_atoms_padding: ${data.uniform_n_atoms_padding}
+  time_dist: ${model.module.time_dist}

qm9/explicit_h/runtime/configs/data/qm9/qm9.yaml

@@ -0,0 +1,106 @@
+## QM9 data and preprocessing configuration.
+
+dataset_name: qm9
+root_path: ${oc.env:PROJECT_ROOT}/data
+preprocess_workers: 0
+
+n_atoms_strategy: "uniform" # "approx", "empirical", "learned"
+# Expand both ends of the empirical count support for the uniform flow source.
+# QM9's minimum is already the hard lower bound of 3, so positive padding is
+# necessarily asymmetric; set to 0 for an unpadded 3--29 diagnostic prior.
+uniform_n_atoms_padding: 5
+optimal_transport: "equivariant"
+# Physical Angstroms (converted to normalized coordinates internally). Insertion
+# bridge noise vanishes as (1-t); move noise is the persistent training-time
+# jitter floor for matched atoms and insertions that are already present.
+ins_noise_scale: 0.35
+move_noise_scale: 0.3
+
+# QM9 contains only neutral molecules; charged species are treated as invalid.
+allow_charged: false
+
+
+# Preprocessing configuration - computes tokens from data and distributions from training dataset only
+preprocessing:
+  _target_: morph.dataset.core.preprocessing.Preprocessing
+  _recursive_: false
+  root: ${data.root_path}/qm9
+  train_dataset: ${data.datamodule.datasets.train}
+  # Canonical (representation-agnostic) cache. One preprocessing per dataset;
+  # representation-specific projection happens at runtime in the wrapper.
+  atom_tokens_path: ${data.preprocessing.root}/processed/atom_tokens.txt
+  edge_tokens_path: ${data.preprocessing.root}/processed/edge_tokens.txt
+  charge_tokens_path: ${data.preprocessing.root}/processed/charge_tokens.txt
+  distributions_path: ${data.preprocessing.root}/processed/distributions.pt
+
+vocab: None
+
+interpolator:
+  _target_: morph.flow_matching.interpolation.Interpolator
+  distributions: None
+  optimal_transport: ${data.optimal_transport}
+  ins_noise_scale: ${data.ins_noise_scale}
+  move_noise_scale: ${data.move_noise_scale}
+  # OT matching (cost always squared-Euclidean; move cost in Angstrom). c_move
+  # scales the squared move distance (A^2); c_ins/c_del/c_sub are fixed A^2 costs,
+  # directly comparable. At 6.125 each, delete+insert costs 3.5^2, so atoms
+  # farther than 3.5 A apart are left unmatched by default.
+  pre_align_iters: 2   # equivariant-OT inner loop (Hungarian+Kabsch) iters
+  c_move: 1.0
+  c_sub: 0.0
+  c_ins: 6.125
+  c_del: 6.125
+  del_schedule: ${schedule}
+  ins_schedule: ${schedule}
+  sub_schedule:
+    _target_: morph.flow_matching.schedules.SmoothstepSchedule
+    shift: 1.5
+  sub_e_schedule:
+    _target_: morph.flow_matching.schedules.SmoothstepSchedule
+    shift: 1.5
+
+datamodule:
+  _target_: morph.dataset.core.lightning_datamodule.LightningDataModule
+  # tokens and distributions will be passed from preprocessing in run.py
+
+  vocab: ${data.vocab}
+  distributions: None
+  representation: ${representation}
+
+  datasets:
+    train:
+      _target_: morph.dataset.unconditional.qm9.FlowMatchingQM9Dataset
+      root: ${data.root_path}/qm9
+      vocab: ${data.vocab}
+      distributions: None
+
+    val:
+      - _target_: morph.dataset.unconditional.qm9.FlowMatchingQM9Dataset
+        root: ${data.root_path}/qm9
+        vocab: ${data.vocab}
+        distributions: None
+
+    test:
+      - _target_: morph.dataset.unconditional.qm9.FlowMatchingQM9Dataset
+        root: ${data.root_path}/qm9
+        vocab: ${data.vocab}
+        distributions: None
+
+  interpolator: ${data.interpolator}
+
+  rotate: ${model.rotate}
+  n_augmentations: ${model.n_augmentations}
+
+  num_workers:
+    train: 10
+    val: 0
+    test: 0
+
+  batch_size:
+    train: 1024
+    val: 128
+    test: 32
+
+  n_atoms_strategy: ${data.n_atoms_strategy}
+  uniform_n_atoms_padding: ${data.uniform_n_atoms_padding}
+  time_dist: ${model.module.time_dist}

qm9/explicit_h/runtime/configs/data/qm9/scaffold_decoration.yaml

@@ -0,0 +1,22 @@
+# QM9 data config for the **scaffold decoration** task.
+#
+# Source is the selected scaffold of each target; target is the full molecule.
+# The model learns to grow the molecule by inserting decoration atoms
+# around the fixed scaffold.
+#
+# Inherits from configs/data/qm9/qm9.yaml and overrides:
+#   - datamodule._target_ → ScaffoldLightningDataModule
+#
+defaults:
+  - /data/qm9/qm9@_here_
+  - _self_
+
+scaffold_groups_dir: ${data.root_path}/qm9/processed/scaffold_groups
+scaffold_kind: rgroup # murcko, rgroup (BRICS), heavy (generate H), linker
+# Evaluate each distinct molecule/scaffold pair; training always draws randomly.
+enumerate_slicings: false
+datamodule:
+  _target_: morph.dataset.conditional.scaffold.datamodule.ScaffoldLightningDataModule
+  scaffold_groups_dir: ${data.scaffold_groups_dir}
+  scaffold_kind: ${data.scaffold_kind}
+  enumerate_slicings: ${data.enumerate_slicings}

qm9/explicit_h/runtime/configs/default.yaml

@@ -0,0 +1,25 @@
+defaults:
+  - _self_
+  # Single-knob presets bundle
+  # `data` + `embeddings` (with dataset-tuned ExtrapolatableCount params)
+  # plus the shared `core` group selections. Override individual groups as
+  # usual, e.g. `defaults=geom/geom model=backbones/dit embeddings=qm9`.
+  - defaults: qm9/qm9
+
+# What the model trains on. Datasets declare their CAPABILITIES; run.py validates
+# the requested mode against them before preprocessing.
+#   pointcloud         — x, a only (charges and edges are dummy singletons)
+#   charged_pointcloud — x, a, c   (edges are dummy singletons)
+#   geometric_graph    — x, a, edge_index, e (charges are dummy neutral)
+#   molecule           — full chemistry (current behavior)
+representation: molecule
+seed: 42
+
+# Run directory: <root>/<dataset>/<cfg_method>/<date>/<wandb_id>. Root is
+# MORPH_RUN_ROOT (scratch, set in .env) so every run's outputs/checkpoints/wandb
+# land off-repo; falls back to <cwd>/outputs for ad-hoc runs with no env set.
+# wandb id is MORPH_RUN_ID (same as WANDB_RUN_ID). Set in the shell before
+# torchrun, or set automatically for single-process runs in run.py before Hydra composes.
+hydra:
+  run:
+    dir: ${oc.env:MORPH_RUN_ROOT,${hydra:runtime.cwd}/outputs}/${hydra:runtime.choices.data}/${now:%Y-%m-%d}/${oc.env:MORPH_RUN_ID}

qm9/explicit_h/runtime/configs/defaults/chembl3d/chembl3d.yaml

@@ -0,0 +1,5 @@
+# @package _global_
+defaults:
+  - core
+  - /data: chembl3d/chembl3d
+  - /embeddings: geom

qm9/explicit_h/runtime/configs/defaults/chembl3d/implicit_h.yaml

@@ -0,0 +1,5 @@
+# @package _global_
+defaults:
+  - /defaults/core
+  - /data: chembl3d/implicit_h
+  - /embeddings: geom

qm9/explicit_h/runtime/configs/defaults/chembl3d/implicit_h_2x4090.yaml

@@ -0,0 +1,10 @@
+# @package _global_
+defaults:
+  - /defaults/chembl3d/train_2x4090
+  - override /data: chembl3d/max100_implicit_h
+  - _self_
+
+# Train from scratch: the atom vocabulary/statistics differ from the H model.
+# Preserve global batch 640, eval every epoch, both checkpoint callbacks every 5.
+logging:
+  project: morph-chembl3d-implicit-h

qm9/explicit_h/runtime/configs/defaults/chembl3d/train_2x4090.yaml

@@ -0,0 +1,40 @@
+# @package _global_
+defaults:
+  - /defaults/core
+  - /data: chembl3d/max100
+  - /embeddings: geom
+  - _self_
+
+# Profiled on max100 ChEMBL: 2 GPUs * 320 * 1 = 640 effective global batch.
+data:
+  datamodule:
+    train_prefetch_factor: 1
+    batch_size:
+      train: 320
+      val: 32
+      test: 32
+    num_workers:
+      train: 20
+      val: 2
+      test: 2
+
+trainer:
+  trainer:
+    accelerator: gpu
+    devices: 2
+    accumulate_grad_batches: 1
+    max_epochs: -1
+    check_val_every_n_epoch: 1
+    num_val_samples: 4096
+  # Validation runs during fit; leave wall-clock headroom for final checkpointing.
+  do_validate: false
+  do_predict: false
+  final_checkpoint_path: final.ckpt
+
+callbacks:
+  model_checkpoints:
+    every_n_epochs: 5
+  every_n_epochs_checkpoint:
+    every_n_epochs: 5
+    save_top_k: -1
+    save_last: link

qm9/explicit_h/runtime/configs/defaults/core.yaml

@@ -0,0 +1,13 @@
+# @package _global_
+# Shared default groups that don't change between QM9 / GEOM runs.
+# Dataset-specific presets (defaults/qm9/qm9.yaml, defaults/geom/geom.yaml) inherit
+# from this and only add `data` + `embeddings`.
+defaults:
+  - /heads: head
+  - /insertion_head: unimodal
+  - /model: backbones/dit
+  - /schedule: smoothstep
+  - /trainer: trainer
+  - /callbacks: callbacks
+  - /logging: wandb_logger
+  - /integrator: integrator

qm9/explicit_h/runtime/configs/defaults/geom/geom.yaml

@@ -0,0 +1,20 @@
+# @package _global_
+# GEOM preset: shared core + GEOM data and embeddings.
+defaults:
+  - core
+  - /data: geom/geom
+  - /embeddings: geom
+  - _self_
+
+# Evaluate every epoch; keep every fifth completed epoch for releases/resuming.
+trainer:
+  trainer:
+    check_val_every_n_epoch: 1
+
+callbacks:
+  model_checkpoints:
+    every_n_epochs: 5
+  every_n_epochs_checkpoint:
+    every_n_epochs: 5
+    save_top_k: -1
+    save_last: link

qm9/explicit_h/runtime/configs/defaults/geom/implicit_h.yaml

@@ -0,0 +1,5 @@
+# @package _global_
+defaults:
+  - /defaults/core
+  - /data: geom/implicit_h
+  - /embeddings: geom

qm9/explicit_h/runtime/configs/defaults/geom/scaffold_decoration.yaml

@@ -0,0 +1,23 @@
+# @package _global_
+# GEOM-Drugs **scaffold decoration** fine-tune preset.
+#
+# Mirrors defaults/qm9/scaffold_decoration.yaml on GEOM-Drugs. Bare-scaffold
+# source -> full-molecule target: the model learns to grow a molecule by
+# inserting decoration atoms around the fixed scaffold.
+#
+# Differences from the QM9 scaffold preset:
+#   - data       -> geom/scaffold_decoration (GEOM base datasets)
+#   - embeddings -> geom
+#
+# Reuses the scaffold_finetune model config (loads the pretrained unconditional
+# DiT into the clean lane + adds the node/edge conditioning and the gated
+# cross-attn blocks).
+#
+# Pair with a GEOM unconditional DiT checkpoint via
+# ``model.module.pretrained_ckpt=...``.
+defaults:
+  - core
+  - /data: geom/scaffold_decoration
+  - /embeddings: geom
+  - override /model: finetune/scaffold_finetune
+  - _self_

qm9/explicit_h/runtime/configs/defaults/qm9/qm9.yaml

@@ -0,0 +1,20 @@
+# @package _global_
+# QM9 preset: shared core + QM9 data and embeddings.
+defaults:
+  - /defaults/core
+  - /data: qm9/qm9
+  - /embeddings: qm9
+  - _self_
+
+trainer:
+  trainer:
+    check_val_every_n_epoch: 1
+
+# Keep the full 25-epoch checkpoint history for model releases and resuming.
+callbacks:
+  model_checkpoints:
+    every_n_epochs: 25
+  every_n_epochs_checkpoint:
+    every_n_epochs: 25
+    save_top_k: -1
+    save_last: link

qm9/explicit_h/runtime/configs/defaults/qm9/scaffold_decoration.yaml

@@ -0,0 +1,11 @@
+# @package _global_
+# QM9 **scaffold decoration** fine-tune preset.
+#
+# Bare-scaffold source → full-molecule target. The model learns to grow
+# a molecule by inserting decoration atoms around the fixed scaffold.
+defaults:
+  - core
+  - /data: qm9/scaffold_decoration
+  - /embeddings: qm9
+  - override /model: finetune/scaffold_finetune
+  - _self_

qm9/explicit_h/runtime/configs/defaults/qm9/train_2x4090.yaml

@@ -0,0 +1,29 @@
+# @package _global_
+defaults:
+  - /defaults/qm9/qm9
+  - _self_
+
+# Preserve the established QM9 effective batch: 2 GPUs * 1024 * 2 = 4096.
+data:
+  datamodule:
+    train_prefetch_factor: 1
+    batch_size:
+      train: 1024
+      val: 128
+      test: 32
+    num_workers:
+      train: 20
+      val: 2
+      test: 2
+
+trainer:
+  trainer:
+    accelerator: gpu
+    devices: 2
+    accumulate_grad_batches: 2
+    max_epochs: 3000
+    check_val_every_n_epoch: 1
+    num_val_samples: 4096
+  do_validate: false
+  do_predict: false
+  final_checkpoint_path: final.ckpt

qm9/explicit_h/runtime/configs/embeddings/geom.yaml

@@ -0,0 +1,28 @@
+# Shared embedding stack tuned for GEOM-Drugs.
+# - node_count_embedding uses ExtrapolatableCountEmbedding with periods/max_count
+#   scaled up for the much larger atom counts in GEOM-Drugs.
+
+atom_type_embedding:
+  _target_: morph.model.layers.embedding.Embedding
+  num_embeddings: ${len:${data.vocab.atom_tokens}}
+  embedding_dim: 64
+  out_dim: 128
+
+edge_type_embedding:
+  _target_: morph.model.layers.embedding.Embedding
+  num_embeddings: ${len:${data.vocab.edge_tokens}}
+  embedding_dim: 64
+  out_dim: 128
+
+time_embedding:
+  _target_: morph.model.layers.embedding.TimeEmbedding
+  embedding_dim: 128
+  out_dim: 128
+
+node_count_embedding:
+  _target_: morph.model.layers.embedding.ExtrapolatableCountEmbedding
+  embedding_dim: 64
+  out_dim: 128
+  min_period: 128.0
+  max_period: 1024.0
+  max_count: 96.0

qm9/explicit_h/runtime/configs/embeddings/qm9.yaml

@@ -0,0 +1,29 @@
+# Shared embedding stack tuned for QM9.
+# - node_count_embedding uses ExtrapolatableCountEmbedding so the model can
+#   generalise slightly outside the training atom-count support.
+# - Defaults sized for QM9 (training counts in [3, 29]).
+
+atom_type_embedding:
+  _target_: morph.model.layers.embedding.Embedding
+  num_embeddings: ${len:${data.vocab.atom_tokens}}
+  embedding_dim: 64
+  out_dim: 128
+
+edge_type_embedding:
+  _target_: morph.model.layers.embedding.Embedding
+  num_embeddings: ${len:${data.vocab.edge_tokens}}
+  embedding_dim: 64
+  out_dim: 128
+
+time_embedding:
+  _target_: morph.model.layers.embedding.TimeEmbedding
+  embedding_dim: 128
+  out_dim: 128
+
+node_count_embedding:
+  _target_: morph.model.layers.embedding.ExtrapolatableCountEmbedding
+  embedding_dim: 64
+  out_dim: 128
+  min_period: 64.0
+  max_period: 512.0
+  max_count: 32.0

qm9/explicit_h/runtime/configs/heads/head.yaml

@@ -0,0 +1,75 @@
+# Standard molecular DiT readouts. The insertion-content head is an independent
+# Hydra choice (``insertion_head=unimodal`` by default, or ``=gmm``).
+
+heads:
+  _target_: morph.model.layers.heads.MultiHeadModule
+  output_init_std: ${model.head_initialization.output_std}
+  independent_adapters: ${oc.select:model.head_independent_adapters,false}
+
+  heads_configs:
+    graph_aggregation: mean
+
+    node_heads:
+      atom_type_head:
+        input_dim: ${model.head_node_dim}
+        output_dim: ${len:${data.vocab.atom_tokens}}
+        adapter_dim: ${model.head_node_adapter_dim}
+        group: state
+
+      charge_head:
+        input_dim: ${model.head_node_dim}
+        output_dim: ${len:${data.vocab.charge_tokens}}
+        adapter_dim: ${model.head_node_adapter_dim}
+        group: state
+
+      ins_rate_head:
+        input_dim: ${model.head_node_dim}
+        output_dim: 1
+        adapter_dim: ${model.head_node_adapter_dim}
+        group: edit
+        initial_rate: ${model.head_initialization.ins_rate}
+
+      do_del_head:
+        input_dim: ${model.head_node_dim}
+        output_dim: 1
+        adapter_dim: ${model.head_node_adapter_dim}
+        group: edit
+        initial_probability: ${model.head_initialization.do_del_probability}
+
+      do_sub_a_head:
+        input_dim: ${model.head_node_dim}
+        output_dim: 1
+        adapter_dim: ${model.head_node_adapter_dim}
+        group: edit
+        initial_probability: ${model.head_initialization.do_sub_a_probability}
+
+    edge_heads:
+      edge_type_head:
+        input_dim: ${model.head_edge_dim}
+        output_dim: ${len:${data.vocab.edge_tokens}}
+        adapter_dim: ${model.head_edge_adapter_dim}
+        group: state
+
+      do_sub_e_head:
+        input_dim: ${model.head_edge_dim}
+        output_dim: 1
+        adapter_dim: ${model.head_edge_adapter_dim}
+        group: state
+        initial_probability: ${model.head_initialization.do_sub_e_probability}
+
+ins_edge_head:
+  _target_: morph.model.layers.heads.InsertionEdgeHead
+  _recursive_: false
+  hidden_dim: ${model.head_node_dim}
+  readout_dim: ${model.head_insertion_edge_dim}
+  output_init_std: ${model.head_initialization.output_std}
+  n_edge_types: ${len:${data.vocab.edge_tokens}}
+  n_atom_types: ${len:${data.vocab.atom_tokens}}
+  n_charge_types: ${len:${data.vocab.charge_tokens}}
+  rbf_embedding_args:
+    _target_: morph.model.layers.embedding.RBFEmbedding
+    num_rbf: 128
+    rbf_dmax: 3.5
+    out_dim: 64
+    dropout: 0.0
+    trainable: false

qm9/explicit_h/runtime/configs/insertion_head/gmm.yaml

@@ -0,0 +1,16 @@
+# Optional mixture-density insertion-content head.
+mode: gmm
+dimension: 3
+num_components: 10
+
+model:
+  _target_: morph.model.layers.heads.GMMInsertionHead
+  hidden_dim: ${model.head_node_dim}
+  readout_dim: ${model.head_insertion_adapter_dim}
+  output_init_std: ${model.head_initialization.output_std}
+  num_components: ${insertion_head.num_components}
+  sigma_min: 0.05
+  sigma_max: 2.0
+  sigma_init: 0.5
+  N_a: ${len:${data.vocab.atom_tokens}}
+  N_c: ${len:${data.vocab.charge_tokens}}

qm9/explicit_h/runtime/configs/insertion_head/unimodal.yaml

@@ -0,0 +1,13 @@
+# Default insertion-content head: one nearest pending atom per spawning node.
+mode: unimodal
+dimension: 3
+num_components: 1
+
+model:
+  _target_: morph.model.layers.heads.UnimodalInsertionHead
+  hidden_dim: ${model.head_node_dim}
+  readout_dim: ${model.head_insertion_adapter_dim}
+  output_init_std: ${model.head_initialization.output_std}
+  independent_adapters: ${oc.select:model.head_independent_adapters,false}
+  N_a: ${len:${data.vocab.atom_tokens}}
+  N_c: ${len:${data.vocab.charge_tokens}}

qm9/explicit_h/runtime/configs/integrator/integrator.yaml

@@ -0,0 +1,30 @@
+integrator:
+  _target_: morph.flow_matching.integration.RateIntegrator
+  vocab: ${data.vocab}
+  distributions: None
+  insertion_head_mode: ${insertion_head.mode}
+  insertion_dimension: ${insertion_head.dimension}
+  insertion_num_components: ${insertion_head.num_components}
+  # strategy used for sampling time during **integration**
+  # NOTE can be different from the time strategy used for the model (see below)
+  time_strategy: log # "linear", "log", or "beta"
+  # shape params for time_strategy: beta (Beta inverse-CDF spacing).
+  # defaults match the model's beta_sched (concentration1=2.0, concentration0=1.0)
+  time_beta_alpha: 2.0
+  time_beta_beta: 1.0
+  num_integration_steps: 100
+  ins_noise_scale: ${data.ins_noise_scale}
+  # When true, freshly-inserted nodes are pushed to the t+dt marginal (kappa-mix
+  # on atom/bond types, (1-t)*ins_noise on x) to match training. Charge remains
+  # a clean prediction target rather than an interpolated state. When false,
+  # insertions enter at the model's clean prediction.
+  renoise_insertions: true
+  # Integrate the known categorical schedule hazard over every time interval
+  # instead of using its left-endpoint Euler approximation.
+  # With either method, the final interval always completes pending discrete
+  # transitions using the learned edit probabilities (no extra model call).
+  exact_schedule_integration: true
+  del_schedule: ${data.interpolator.del_schedule}
+  ins_schedule: ${data.interpolator.ins_schedule}
+  sub_schedule: ${data.interpolator.sub_schedule}
+  sub_e_schedule: ${data.interpolator.sub_e_schedule}

qm9/explicit_h/runtime/configs/logging/wandb_logger.yaml

@@ -0,0 +1,6 @@
+## Put your model args here
+enabled: true
+project: morph-${data.dataset_name}
+name: ${oc.env:MORPH_RUN_ID,base}
+save_dir: .
+offline: false

qm9/explicit_h/runtime/configs/loss_weights/calibrated_geom.yaml

@@ -0,0 +1,18 @@
+# @package model.module.loss_weights
+# Balance-only (geometric-mean=1) calibration weights, scaled down by a global
+# factor of 0.2 (1/5). At the original geometric-mean=1 scale the realized
+# full-model grad-norm was ~5, far above gradient_clip_val=1.0, so clipping
+# fired on every step. The 0.2 factor brings the realized norm to ~1.0 so the
+# clip becomes a spike safety-valve again. Re-running calibration now applies
+# this scaling automatically (see trainer.target_total_grad_norm).
+c: 0.164
+do_del: 0.20
+do_sub_a: 0.30
+do_sub_e: 0.24
+ins_e: 0.21
+ins_e_ii: 0.22
+ins_state: 0.14
+ins_rate: 0.20
+sub_a_class: 0.15
+sub_e_class: 0.19
+x: 0.22

qm9/explicit_h/runtime/configs/loss_weights/calibrated_geom_frozen.yaml

@@ -0,0 +1,25 @@
+# @package model.module.loss_weights
+# Balance-only (geometric-mean = 1) loss weights — ratios only. The learning
+# rate owns step size and gradient_clip_val stays a pure spike catch, so these
+# are NOT scaled to a target grad-norm (unlike the older calibrated_geom.yaml,
+# which baked in a ~0.2 global scale to hit the clip target).
+#
+# Source: frozen init-state calibration (trainer.calibrate_mode=loss,
+# calibrate_freeze_model=true), geom / DiT / uncond,
+# run 2026-06-07 (outputs/geom/uncond/cal_then_lr/2026-06-07_12-29-25/calibrate).
+# Because the model is held at init, every probe batch is an unbiased snapshot of
+# the initial per-term losses — so the three binary do_* losses come out equal
+# (all BCE at p=0.5 = ln2), unlike the previous 2500-step *training* probe
+# (calibrated_geom.yaml) which gave them different weights purely from differing
+# training rates. See notebooks/lr_calibration_analysis.ipynb for the comparison.
+c: 0.745146
+do_del: 1.730107
+do_sub_a: 1.730107
+do_sub_e: 1.730106
+ins_e: 0.867232
+ins_e_ii: 0.867232
+ins_state: 0.519379
+ins_rate: 1.128104
+sub_a_class: 0.522109
+sub_e_class: 0.867232
+x: 1.298792

qm9/explicit_h/runtime/configs/loss_weights/tuned_weights.yaml

@@ -0,0 +1,15 @@
+# @package model.module.loss_weights
+# Canonical tuned weights for molecular training. Backbones can select this
+# preset explicitly, and task-specific experiments may override the group. The
+# package directive mounts the values under model.module.loss_weights.
+do_sub_a: 0.2
+sub_a_class: 0.03
+do_sub_e: 0.25
+sub_e_class: 0.5
+do_del: 0.4
+ins_rate: 0.15
+ins_state: 0.01
+ins_e: 0.4
+x: 2.0 # was 1.0
+c: 0.5
+ins_e_ii: 0.03