Add ChEMBL3D explicit-H epochs 200/300/400 and implicit-H epochs 100/200/300 with inference bundles

README.md

@@ -0,0 +1,97 @@
+---
+library_name: pytorch
+tags:
+  - chemistry
+  - molecular-generation
+  - flow-matching
+  - chembl3d
+---
+
+# Morph ChEMBL3D checkpoints
+
+Six checkpoints for unconditional 3D molecular generation, with matching
+vocabularies, training distributions, portable inference configurations and
+the frozen Morph runtime from each training run.
+
+| Variant | Completed training epochs | Directory |
+| --- | --- | --- |
+| Explicit hydrogens | 200, 300, 400 | [chembl3d/explicit_h](chembl3d/explicit_h) |
+| Implicit hydrogens (heavy-atom graphs) | 100, 200, 300 | [chembl3d/implicit_h](chembl3d/implicit_h) |
+
+The `epoch_200.ckpt` name means **200 completed epochs**. The original Lightning
+checkpoint stores `epoch=199`. Files are byte-for-byte copies of the selected
+training checkpoints, including current weights, EMA weights, and optimizer
+state. Sampling uses EMA weights by default.
+
+Each variant contains:
+
+- `epoch_NNN.ckpt`: the three selected checkpoints.
+- `atom_tokens.txt`, `edge_tokens.txt`, `charge_tokens.txt`: ordered vocabularies.
+- `distributions.pt`: the original training statistics, including atom counts,
+  type frequencies, coordinate scale, and geometry histograms.
+- `distributions.json`: a readable export of the same statistics.
+- `config.yaml`: portable sampling config with dataset loading disabled.
+- `training_config.yaml`: the original Hydra training configuration, for provenance.
+- `runtime/`: frozen source/configs and dependency specifications for this run.
+- `metadata.json`: epochs, optimizer steps, original filenames, and SHA-256 hashes.
+- `validation.json`: strict checkpoint loading, current/EMA forward checks and
+  a one-molecule end-to-end sampling check. This is a functionality check, not
+  a validity benchmark.
+
+## Download and generate
+
+Download the variant you need (authenticate with `hf auth login` if required):
+
+```bash
+hf download maltefranke/morph --local-dir morph-models \
+  --include 'chembl3d/implicit_h/*' --include sample.py --include README.md --include manifest.json
+cd morph-models
+```
+
+Install that variant's runtime in an environment with Python >=3.11 and a
+compatible PyTorch 2.9 installation. A CUDA-capable PyTorch build is needed to
+use `--device cuda`. The tested package versions are recorded in
+`runtime/requirements-tested.txt`; the runtime's `pyproject.toml` specifies the
+dependency constraints.
+
+```bash
+python -m pip install ./chembl3d/implicit_h/runtime
+python sample.py --variant implicit_h --epoch 300 \
+  --num-mols 64 --batch-size 16 --device cuda --output implicit_h.sdf
+```
+
+For explicit hydrogens, download `chembl3d/explicit_h/*` instead, install its
+runtime, and run:
+
+```bash
+python sample.py --variant explicit_h --epoch 400 \
+  --num-mols 64 --batch-size 16 --device cuda --output explicit_h.sdf
+```
+
+The script imports the matching bundled source automatically. Each invocation
+uses one variant in a fresh process. No ChEMBL3D dataset or cluster paths are
+needed. `--device cpu` is supported, but use a compute worker for substantial
+sampling. Compilation is disabled by default; attention uses PyTorch SDPA.
+
+Generated coordinates are in Angstrom. The SDF contains valid molecules; the
+adjacent JSON records every attempted molecule and its validity result, including
+failures. Generation uses 100 integration steps and seed 42 by default.
+
+## Hydrogen handling and training data
+
+Both variants use the same split membership: 1,474,008 training, 184,295 validation,
+and 184,242 test molecules. The dataset was filtered to **at most 100 total atoms,
+including hydrogens**, before removing hydrogens for the implicit variant.
+
+The explicit variant predicts H atoms as graph vertices. The implicit variant
+predicts heavy atoms, bonds and formal charges; RDKit infers hydrogens during
+conversion and validity checking. Its bundled runtime preserves that behavior.
+Do not mix vocabularies or distributions between variants. Geometry statistics
+and coordinate normalization were recomputed for the heavy-atom representation.
+
+The 100-atom cap describes the training filter, not a hard limit imposed on
+generated molecules. Sample counts follow the saved training prior settings.
+
+`manifest.json` lists file sizes and SHA-256 hashes. The original training configs
+contain historical cluster paths and should be treated as provenance; use
+`sample.py` with the portable `config.yaml` for inference.

chembl3d/explicit_h/atom_tokens.txt

@@ -0,0 +1,13 @@
+B
+Br
+C
+Cl
+F
+H
+I
+N
+O
+P
+S
+Se
+Si

chembl3d/explicit_h/charge_tokens.txt

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

chembl3d/explicit_h/config.yaml

@@ -0,0 +1,329 @@
+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: -1
+    accumulate_grad_batches: 1
+    num_sanity_val_steps: 0
+    gradient_clip_val: null
+    gradient_clip_algorithm: norm
+    check_val_every_n_epoch: 5
+    num_val_samples: 4096
+    limit_val_batches: null
+    limit_predict_batches: 4
+    max_time: null
+  do_train: false
+  do_validate: false
+  do_predict: true
+  checkpoint_path: null
+  final_checkpoint_path: final.ckpt
+callbacks: {}
+logging:
+  enabled: false
+  project: morph-${data.dataset_name}
+  name: chembl3d-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: chembl3d
+  root_path: .
+  preprocess_workers: 0
+  n_atoms_strategy: uniform
+  uniform_n_atoms_padding: 5
+  optimal_transport: equivariant
+  ins_noise_scale: 0.65
+  move_noise_scale: 0.5
+  allow_charged: true
+  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: null
+  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}
+    atom_type_noise_prob: 0.0
+    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
+  source_dir: null
+  max_total_atoms: 100
+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: 128.0
+    max_period: 1024.0
+    max_count: 96.0

chembl3d/explicit_h/distributions.pt

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

chembl3d/explicit_h/edge_tokens.txt

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

chembl3d/explicit_h/epoch_200.ckpt

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

chembl3d/explicit_h/epoch_300.ckpt

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

chembl3d/explicit_h/epoch_400.ckpt

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

chembl3d/explicit_h/metadata.json

@@ -0,0 +1,71 @@
+{
+  "dataset": "ChEMBL3D",
+  "variant": "explicit_h",
+  "includes_explicit_hydrogens": true,
+  "source_max_total_atoms_including_h": 100,
+  "training_run": "chembl3d-max100-2x4090-15313268",
+  "source_snapshot": "chembl3d_recovery_20260927_v2",
+  "checkpoint_format": "Original Lightning checkpoint, current + EMA + optimizer state",
+  "epoch_numbering": "Completed epochs (one-based); original checkpoint epoch is zero-based",
+  "split_molecules": {
+    "train": 1474008,
+    "val": 184295,
+    "test": 184242
+  },
+  "checkpoints": [
+    {
+      "filename": "epoch_200.ckpt",
+      "original_filename": "epoch=199-step=460600.ckpt",
+      "completed_epochs": 200,
+      "checkpoint_epoch_index": 199,
+      "bytes": 509357065,
+      "sha256": "087f67faf38847d34780d061b0d192e61f8b4365d50cab6c32e0cc7dc5368743",
+      "global_step": 460600
+    },
+    {
+      "filename": "epoch_300.ckpt",
+      "original_filename": "epoch=299-step=690900.ckpt",
+      "completed_epochs": 300,
+      "checkpoint_epoch_index": 299,
+      "bytes": 509357129,
+      "sha256": "ad7b4535cf8d40ead915c59e525fbe3ec555e6f03885b78ed54c45497793f32b",
+      "global_step": 690900
+    },
+    {
+      "filename": "epoch_400.ckpt",
+      "original_filename": "epoch=399-step=921200.ckpt",
+      "completed_epochs": 400,
+      "checkpoint_epoch_index": 399,
+      "bytes": 509357129,
+      "sha256": "efb5c2f85972df19eca85d2d69cb56f6f52453d91ee26293b3e68b299e2bac8f",
+      "global_step": 921200
+    }
+  ],
+  "coordinate_std_angstrom": 2.707319498062134,
+  "atom_tokens": [
+    "B",
+    "Br",
+    "C",
+    "Cl",
+    "F",
+    "H",
+    "I",
+    "N",
+    "O",
+    "P",
+    "S",
+    "Se",
+    "Si"
+  ],
+  "edge_tokens": [
+    "<NO_BOND>",
+    "1",
+    "2",
+    "3"
+  ],
+  "charge_tokens": [
+    "-1",
+    "0",
+    "1"
+  ]
+}

chembl3d/explicit_h/runtime/README.md

@@ -0,0 +1,3 @@
+# Morph runtime for ChEMBL3D explicit_h
+
+Frozen training source. See the model repository README and sample.py.

chembl3d/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

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

@@ -0,0 +1,36 @@
+# LoQI Release 2 -> Morph LMDB, using the shared drug-like training setup.
+defaults:
+  - /data/geom/geom@_here_
+  - _self_
+
+dataset_name: chembl3d
+# Optional directory with the downloaded train_h.pt / val_h.pt / test_h.pt.
+# When unset, missing files are downloaded to <root>/raw/.
+source_dir: ${oc.env:CHEMBL3D_SOURCE_DIR,null}
+vocab: null
+
+preprocessing:
+  num_workers: ${data.preprocess_workers}
+  root: ${data.root_path}/chembl3d
+
+datamodule:
+  distributions: null
+  datasets:
+    train:
+      _target_: morph.dataset.unconditional.chembl3d.FlowMatchingChEMBL3DDataset
+      source_dir: ${data.source_dir}
+      distributions: null
+    val:
+      - _target_: morph.dataset.unconditional.chembl3d.FlowMatchingChEMBL3DDataset
+        root: ${data.preprocessing.root}
+        source_dir: ${data.source_dir}
+        preprocess_workers: ${data.preprocess_workers}
+        vocab: ${data.vocab}
+        distributions: null
+    test:
+      - _target_: morph.dataset.unconditional.chembl3d.FlowMatchingChEMBL3DDataset
+        root: ${data.preprocessing.root}
+        source_dir: ${data.source_dir}
+        preprocess_workers: ${data.preprocess_workers}
+        vocab: ${data.vocab}
+        distributions: null

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

@@ -0,0 +1,8 @@
+# Dataset cache built by morph.dataset.unconditional.chembl3d_filter.
+defaults:
+  - /data/chembl3d/chembl3d@_here_
+  - _self_
+
+max_total_atoms: 100 # Inclusive, with explicit hydrogens retained.
+preprocessing:
+  root: ${data.root_path}/chembl3d_max100

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

@@ -0,0 +1,113 @@
+## GEOM-Drugs data and preprocessing configuration.
+
+dataset_name: geom
+root_path: ${oc.env:PROJECT_ROOT}/data
+preprocess_workers: 10
+
+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
+
+# GEOM-Drugs contains charged species; allow them when scoring validity.
+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}/geom
+  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: 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}
+  # 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: None
+  representation: ${representation}
+
+  datasets:
+    train:
+      _target_: morph.dataset.unconditional.geom.FlowMatchingGEOMDataset
+      root: ${data.preprocessing.root}
+      preprocess_workers: ${data.preprocess_workers}
+      vocab: ${data.vocab}
+      distributions: None
+
+    val:
+      - _target_: morph.dataset.unconditional.geom.FlowMatchingGEOMDataset
+        root: ${data.preprocessing.root}
+        preprocess_workers: ${data.preprocess_workers}
+        vocab: ${data.vocab}
+        distributions: None
+
+    test:
+      - _target_: morph.dataset.unconditional.geom.FlowMatchingGEOMDataset
+        root: ${data.preprocessing.root}
+        preprocess_workers: ${data.preprocess_workers}
+        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: 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}

chembl3d/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}

chembl3d/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}

chembl3d/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}

chembl3d/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}

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

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

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

@@ -0,0 +1,41 @@
+# @package _global_
+defaults:
+  - core
+  - /data: chembl3d/max100
+  - /embeddings: geom
+  - _self_
+
+# Preserve the base recipe's effective global batch: 2 GPUs * 80 * 4 = 640.
+data:
+  datamodule:
+    train_prefetch_factor: 1
+    batch_size:
+      train: 80
+      val: 32
+      test: 32
+    num_workers:
+      train: 20
+      val: 2
+      test: 2
+
+trainer:
+  trainer:
+    accelerator: gpu
+    devices: 2
+    accumulate_grad_batches: 4
+    max_epochs: -1
+    check_val_every_n_epoch: 5
+    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:
+  runtime_monitor:
+    _target_: morph.utils.training_monitor.TrainingMonitor
+    every_n_batches: 100
+  every_n_epochs_checkpoint:
+    every_n_epochs: 5
+    save_top_k: -1
+    save_last: link

chembl3d/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

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

@@ -0,0 +1,6 @@
+# @package _global_
+# GEOM preset: shared core + GEOM data and embeddings.
+defaults:
+  - core
+  - /data: geom/geom
+  - /embeddings: geom

chembl3d/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_

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

@@ -0,0 +1,6 @@
+# @package _global_
+# QM9 preset: shared core + QM9 data and embeddings.
+defaults:
+  - core
+  - /data: qm9/qm9
+  - /embeddings: qm9

chembl3d/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_

chembl3d/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

chembl3d/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

chembl3d/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

chembl3d/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}}

chembl3d/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}}

chembl3d/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}

chembl3d/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

chembl3d/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

chembl3d/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

chembl3d/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

chembl3d/explicit_h/runtime/configs/loss_weights/uniform.yaml

@@ -0,0 +1,15 @@
+# @package model.module.loss_weights
+# Uniform 1.0 weight per loss component. The key set matches
+# LightningModuleRates.LOSS_GROUPS exactly (the 11 actual loss terms
+# emitted by shared_step).
+do_sub_a: 1.0
+sub_a_class: 1.0
+do_sub_e: 1.0
+sub_e_class: 1.0
+do_del: 1.0
+ins_rate: 1.0
+ins_state: 1.0
+ins_e: 1.0
+ins_e_ii: 1.0
+x: 1.0
+c: 1.0

chembl3d/explicit_h/runtime/configs/model/backbones/dit.yaml

@@ -0,0 +1,123 @@
+## Molecular DiT: the repository's first-class generative backbone.
+##
+## This intentionally keeps charge out of the state and adds neither an
+## equivariant lane nor an evolving dense pair state.  It retains the atom-count
+## condition: time plus current graph size disambiguates otherwise identical
+## local states and is especially useful while insertion/deletion is active.
+defaults:
+  - /optimizer: muon_aux_adam
+  - /loss_weights: tuned_weights
+  - _self_
+
+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
+
+module:
+  _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}
+    # Sum the existing incident edge embeddings into each atom's input, excluding
+    # NO_BOND, through a zero-initialized projection. Uses the existing heads.
+    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}
+
+head_node_dim: ${model.backbone.out_node_nf}
+head_edge_dim: ${model.backbone.out_edge_nf}
+
+# Shared adapters avoid seven separate full-width MLPs.  At the >150M preset,
+# node/insertion adapters widen while pair readouts remain compact.
+head_node_adapter_dim: 256
+head_edge_adapter_dim: 256
+head_insertion_adapter_dim: 256
+head_insertion_edge_dim: 256
+head_independent_adapters: true
+
+# Categorical biases are replaced at module construction with exact dataset
+# marginals.  There is no dataset-level transition-frequency statistic for the
+# binary/rate heads, so their defaults deliberately preserve the old neutral
+# predictions instead of inventing a corpus-independent edit prior.
+head_initialization:
+  output_std: 1.0e-3
+  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
+
+# Optional low-rank adaptation of the attention and SwiGLU projections. The
+# adapter residual is zero at initialization, so enabling this while loading a
+# pretrained checkpoint preserves its outputs exactly.
+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

chembl3d/explicit_h/runtime/configs/model/finetune/lora_finetune.yaml

@@ -0,0 +1,13 @@
+# Parameter-efficient fine-tuning of the unconditional DiT.
+defaults:
+  - /model/backbones/dit@_here_
+  - _self_
+
+lora:
+  enabled: true
+
+module:
+  _target_: morph.model.finetune.base.LightningModuleRatesFinetune
+  pretrained_ckpt: ${oc.env:MORPH_PRETRAINED_CKPT}
+  pretrained_prefer_ema: true
+  freeze_pretrained_backbone: true

chembl3d/explicit_h/runtime/configs/model/finetune/scaffold_finetune.yaml

@@ -0,0 +1,134 @@
+# DiT backbone augmented with Flamingo-style gated cross-attention for
+# scaffold-decoration fine-tuning.
+#
+# Conditioning (all derived from mol_t.keep_mask):
+#   - node_cond : per-atom scaffold membership            (always on)
+#   - edge_cond : per-edge scaffold class                 (always on)
+#   - cross-attn: scaffold subgraph encoded by a shared LightweightContextEncoder
+#                 (OPTIONAL — toggle via backbone.use_cross_attn):
+#                   * false → option 1: node+edge only
+#                   * true  → option 2: node+edge+cross-attn
+#
+# Inherits everything from configs/model/backbones/dit.yaml and overrides:
+#   - module._target_ → LightningModuleRatesScaffoldFinetune (loads a pretrained
+#       unconditional DiT checkpoint and adds scaffold loss/inference masking)
+#   - model._target_  → ScaffoldConditionalDiTBackboneWithHeads (auto-builds the
+#       node/edge cond and, when use_cross_attn, the scaffold cross-attn tokens)
+#   - backbone._target_ → ConditionalDiTBackbone (DiT + gated cross-attn blocks)
+#
+# Provides a ``pretrained_ckpt`` field for the lightning module to load.
+defaults:
+  - /model/backbones/dit@_here_
+  - _self_
+
+module:
+  _target_: morph.model.finetune.scaffold.LightningModuleRatesScaffoldFinetune
+  # Path to the unconditional DiT checkpoint to load weights from.
+  # ConditionalDiT keys are auto-retargeted (backbone.* → backbone.dit.*).
+  pretrained_ckpt: ${oc.env:MORPH_PRETRAINED_CKPT}
+  pretrained_prefer_ema: true
+  # If true, dispatches to ``self.model.freeze_pretrained()`` — for the
+  # cross-attn scaffold model this freezes the entire clean lane
+  # (embedding, ``backbone.dit``, all heads). Only the gated cross-attn
+  # blocks and edge gate plus the scaffold encoder train.
+  freeze_pretrained_backbone: true
+  optimizer_config:
+    optimizer:
+      lr: 1e-3
+    aux_optimizer:
+      lr: 1e-4
+  # Whether to exclude scaffold atoms from sub / x losses. These atoms are
+  # frozen at inference, so training them adds no useful signal.
+  exclude_scaffold_from_sub_loss: true
+  exclude_scaffold_from_x_loss: true
+  # Pin kept coordinates during integration and restore the input coordinate
+  # frame on output. To train flowing scaffold geometry, set this and the x-loss
+  # exclusion above to false.
+  freeze_keep_positions: true
+
+  # Override the model target to the scaffold-aware cross-attn wrapper.
+  model:
+    _target_: morph.model.finetune.scaffold.ScaffoldConditionalDiTBackboneWithHeads
+    _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}
+
+    backbone_model_args: ${model.backbone}
+
+    heads_args: ${heads.heads}
+    insertion_head_args: ${insertion_head.model}
+    ins_edge_head_args: ${heads.ins_edge_head}
+
+    # Cross-attention context encoder (option 2). Required because the backbone
+    # below has use_cross_attn: true. For option 1 (node+edge only), set
+    # backbone.use_cross_attn=false AND null this (the two must agree).
+    xattn_encoder_args: ${model.xattn_encoder}
+    # Geometric correspondence prior (query atom ↔ co-located scaffold token).
+    use_distance_bias: true
+
+# Cross-attention context encoder (option 2): encodes the scaffold subgraph of
+# mol_t (extracted via filter_nodes) into context tokens. ctx_dim must match
+# model.backbone.ctx_dim.
+# Node/edge conditioning is NOT configured here — the wrapper builds those directly
+# from backbone.d_model / backbone.in_edge_nf. Null this (and set
+# backbone.use_cross_attn=false) for node+edge-only scaffold conditioning.
+xattn_encoder:
+  _target_: morph.model.finetune.base.LightweightContextEncoder
+  num_atom_types: ${len:${data.vocab.atom_tokens}}
+  num_charges: ${len:${data.vocab.charge_tokens}}
+  atom_embed_dim: 64
+  charge_embed_dim: 16
+  pos_embed_dim: 32
+  ctx_dim: ${model.backbone.ctx_dim}
+  num_attn_layers: 4
+  nhead: 8
+  mlp_ratio: 4.0
+  dropout: 0.0
+  # Edge-based pair bias on the encoder's self-attention (RBF distance +
+  # edge-type of the scaffold's own bonds) so tokens mix with geometry/topology
+  # awareness, like the DiT backbone. num_edge_types/edge_embed_dim only used here.
+  use_pair_bias: true
+  num_edge_types: ${len:${data.vocab.edge_tokens}}
+  edge_embed_dim: 32
+  num_rbf: 32
+  rbf_dmax: 5.0
+  pair_bias_hidden_dim: 32
+
+# Override the backbone to use ConditionalDiTBackbone.
+backbone:
+  _target_: morph.model.backbones.conditional_dit.ConditionalDiTBackbone
+  _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}
+  ctx_dim: 128
+  use_cross_attn: true   # cross-attention (cond_tokens) port
+  # Per-node conditioning: scaffold membership (which atoms to keep).
+  use_node_cond: true
+  # Per-edge conditioning: the scaffold encoder emits edge_cond at in_edge_nf.
+  use_edge_cond: true
+  rbf_embedding_args:
+    _target_: morph.model.layers.embedding.RBFEmbedding
+    num_rbf: 128
+    rbf_dmax: 3.5
+    out_dim: 64
+    dropout: 0.0
+    trainable: false

chembl3d/explicit_h/runtime/configs/optimizer/muon_aux_adam.yaml

@@ -0,0 +1,39 @@
+# @package model.module.optimizer_config
+# Muon + auxiliary AdamW preset for the molecular DiT. Per-run overrides
+# go on ``model.module.optimizer_config``:
+#   model.module.optimizer_config.optimizer.lr=0.005 ...
+implementation: torch
+
+optimizer:
+  # no _target_ -> triggers Muon path in configure_optimizers
+  # The RTX 4090 Euler template selects native-arithmetic CUDA graph replay.
+  execution: ${oc.env:MORPH_MUON_EXECUTION,native}
+  lr: 0.01
+  momentum: 0.95
+  weight_decay: 0.0
+  nesterov: true
+  ns_steps: 5
+  eps: 1e-7
+  adjust_lr_fn: original
+
+aux_optimizer:
+  lr: 3e-4
+  betas: [0.9, 0.95]
+  eps: 1e-10
+  weight_decay: 0.0
+
+# Lightning's two-optimizer path uses manual optimization, so clipping is
+# configured here and applied once to the full model gradient.
+gradient_clip_val: 1.0
+gradient_clip_algorithm: norm
+
+scheduler:
+  _target_: morph.utils.lr_schedulers.MultiGroupCosineWarmupLR
+  # configure_optimizers injects each native optimizer here.
+  warmup_steps: 2000        # Typically 2-5% of total_steps
+  total_steps: 200000        # The total number of training iterations
+  min_lr_fraction: 0.05      # Decays to 5% of each optimizer's peak LR
+  last_epoch: -1
+
+interval: step
+monitor: null

chembl3d/explicit_h/runtime/configs/schedule/beta.yaml

@@ -0,0 +1,3 @@
+_target_: morph.flow_matching.schedules.BetaSchedule
+k_alpha: 1.0
+k_beta: 1.5

chembl3d/explicit_h/runtime/configs/schedule/cubic.yaml

@@ -0,0 +1 @@
+_target_: morph.flow_matching.schedules.CubicSchedule

chembl3d/explicit_h/runtime/configs/schedule/fast_power.yaml

@@ -0,0 +1,3 @@
+# Default kappa schedule: FastPowerSchedule with beta=2.5
+_target_: morph.flow_matching.schedules.FastPowerSchedule
+beta: 2.0

chembl3d/explicit_h/runtime/configs/schedule/linear.yaml

@@ -0,0 +1 @@
+_target_: morph.flow_matching.schedules.LinearSchedule

chembl3d/explicit_h/runtime/configs/schedule/smoothstep.yaml

@@ -0,0 +1,2 @@
+_target_: morph.flow_matching.schedules.SmoothstepSchedule
+shift: 0.8

chembl3d/explicit_h/runtime/configs/trainer/trainer.yaml

@@ -0,0 +1,37 @@
+## Put your trainer args here
+trainer:
+  deterministic: warn
+  devices: auto
+  accelerator: auto
+  #precision: 32
+  strategy: ddp_find_unused_parameters_false
+  precision: "bf16-mixed"
+  max_epochs: 3000
+  # Effective global batch = devices * batch_size.train * accumulation. The
+  # native optimizer mixin handles this because it steps Muon + AdamW.
+  accumulate_grad_batches: 1
+  # A validation batch performs full generative integration, so Lightning's
+  # pre-training sanity pass is expensive and does not exercise optimization.
+  num_sanity_val_steps: 0
+  # Muon + AdamW use Lightning manual optimization; clipping is configured in
+  # configs/optimizer/muon_aux_adam.yaml and applied to both branches together.
+  gradient_clip_val: null
+  gradient_clip_algorithm: norm
+
+  # Validation performs full generative integration. The established QM9
+  # training recipes use this cadence so sampling does not dominate fitting.
+  check_val_every_n_epoch: 50
+ # profiler: simple
+
+  # Fixed GLOBAL validation population. run.py converts this to a per-rank
+  # limit, giving the same 4096 molecules on 1, 2 or 4 GPUs.
+  num_val_samples: 4096
+  # Explicit profiling escape hatch; a non-null value takes precedence over
+  # num_val_samples and is interpreted by Lightning per rank.
+  limit_val_batches: null
+  limit_predict_batches: 4
+
+do_train: true
+do_validate: true
+do_predict: true
+checkpoint_path: null

chembl3d/explicit_h/runtime/pyproject.toml

@@ -0,0 +1,136 @@
+[project]
+name = "morph"
+version = "0.1.0"
+description = "Flow-matching molecular generation and scaffold decoration with a molecular DiT"
+readme = "README.md"
+keywords = ["molecular-generation", "flow-matching", "scaffold-decoration"]
+requires-python = ">=3.11"
+dependencies = [
+    "click>=8.3.1",
+    "hydra-core>=1.3.2",
+    "lightning>=2.5.5",
+    "omegaconf>=2.3.0",
+    "numpy>=2.0",
+    "pandas>=2.3.3",
+    "rdkit>=2025.9.1",
+    "scipy>=1.15.0",
+    # Pre-release allowed: the NVIDIA 25.10 container ships a CUDA-enabled
+    # 2.9.0a0 build. Preserve that container build instead of replacing it with
+    # a mismatched PyPI Torch/torchvision pair.
+    "torch>=2.9.0a0, <2.9.1",
+    "torch-geometric>=2.7.0",
+    "torchmetrics>=1.8.0",
+    "wandb>=0.22.2",
+    "posebusters>=0.6.0",
+    "matplotlib>=3.10.8",
+    "lmdb>=2.2.0",
+    "py3dmol>=2.5.4",
+    "ipykernel>=7.2.0",
+    "ase>=3.28.0",
+]
+
+[dependency-groups]
+dev = ["pytest>=9.1.1", "ruff>=0.16.5"]
+
+[project.optional-dependencies]
+flash-attn = ["flash-attn>=2.8.3, <2.8.4"]
+
+[project.scripts]
+morph = "morph.cli:cli"
+
+[tool.uv.extra-build-dependencies]
+flash-attn = [{ requirement = "torch", match-runtime = true }]
+
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[tool.hatch.build.targets.wheel]
+packages = ["src/morph"]
+
+[tool.hatch.build.targets.wheel.force-include]
+"configs" = "morph/configs"
+
+[tool.ruff]
+# Allow lines to be as long as 88.
+line-length = 88
+# Keep attributed third-party implementations close to their upstream sources.
+extend-exclude = ["src/morph/_vendor"]
+
+[tool.ruff.format]
+docstring-code-format = true
+
+# Correctness, readability, and public API documentation.
+[tool.ruff.lint]
+# consider selecting "ISC", for implicit string concat
+# consider selecting "T201", for print in production code
+select = [
+    "B",    # flake8-bugbear
+    "C4",   # flake8-comprehensions
+    "D",    # pydocstyle
+    "E",    # pycodestyle error
+    "E501", # line too long
+    "EXE",  # flake8-executable
+    "F",    # pyflakes
+    "FLY",  # flynt
+    "I",    # isort
+    "ICN",  # flake8-import-conventions
+    "PD",   # pandas-vet
+    "PERF", # perflint
+    "PIE",  # flake8-pie
+    "PL",   # pylint
+    "PT",   # flake8-pytest-style
+    "PYI",  # flakes8-pyi
+    "Q",    # flake8-quotes
+    "RSE",  # flake8-raise
+    "RUF",  # Ruff-specific rules
+    "SIM",  # flake8-simplify
+    "SLOT", # flakes8-slot
+    "TCH",  # flake8-type-checking
+    "TID",  # flake8-tidy-imports
+    "UP",   # pyupgrade
+    "W",    # pycodestyle warning
+    "W505", # line too long
+    "YTT",  # flake8-2020
+]
+# also consider ignoring "C901", for high mccabe function complexity
+# also consider ignoring "ERA001", for found commented out code
+ignore = [
+    "EXE002",   # shebang is not a python interpreter
+    "RUF012",  # missing classvar type
+    "RUF013",  # Optional prohibited by PEP484
+    "B006",    # Do not use mutable data structures for argument defaults
+    "ANN003",  # missing type annotations for kwargs
+    "B019",    # Use of functools.lru_cache on methods can lead to memory leaks
+    "C408",    # unnecessary-collection-call
+    "COM812",  # trailing comma missing
+    "D205",    # 1 blank line required between summary line and description
+    "D415",    # first line of docstring must end with punctuation
+    "DTZ005",  # use of datetime.now() without timezone
+    "E731",    # do not assign a lambda expression, use a def
+    "EM",      # flake8-errmsg
+    "FBT001",  # Boolean positional argument in function
+    "FBT002",  # Boolean keyword argument in function
+    "NPY002",  # TODO replace legacy np.random.seed
+    "PLR",     # pylint refactor
+    "PLW2901", # Outer for loop variable overwritten by inner assignment target
+    "PT006",   # pytest-parametrize-names-wrong-type
+    "PT011",   # pytest-raises-too-broad
+    "PT013",   # pytest-incorrect-pytest-import
+    "PT019",   # pytest-fixture-param-without-value
+    "PTH",     # prefer Path to os.path
+    "S301",    # pickle can be unsafe
+    "S310",    # url open functions can be unsafe
+    "TRY003",  # long exception messages not defined in the exception class itself
+    "UP015",   # unnecessary "r" in open call
+    "D107",    # docstring for __init__ 
+    "D212",    # remove whitespace after opening quotes
+    "D202",    # no blank lines after docstring
+    "D301", 
+    "D100",
+    "D103",
+]
+pydocstyle.convention = "google"
+
+[env]
+PROJECT_ROOT = "." 

chembl3d/explicit_h/runtime/requirements-tested.txt

@@ -0,0 +1,19 @@
+torch==2.9.0
+numpy==2.4.2
+rdkit==2025.9.6
+lightning==2.6.1
+pytorch-lightning==2.6.1
+torch-geometric==2.7.0
+torchmetrics==1.8.2
+hydra-core==1.3.2
+omegaconf==2.3.0
+scipy==1.17.1
+pandas==3.0.1
+wandb==0.25.0
+posebusters==0.6.5
+matplotlib==3.10.9
+lmdb==2.2.0
+click==8.3.1
+py3dmol==2.5.4
+ipykernel==7.2.0
+ase==3.28.0