Use log2FC as an auxiliary input feature for pEC50 prediction in moal plan (issue #36 Phases 0-3) - #37
Open
smcolby wants to merge 17 commits into
Open
Use log2FC as an auxiliary input feature for pEC50 prediction in moal plan (issue #36 Phases 0-3)#37smcolby wants to merge 17 commits into
smcolby wants to merge 17 commits into
Conversation
Phase 0 of #36: retain the compound's observed continuous log2FC primary-screen readout as a typed field, separate from the LEFT/INTERVAL censoring derived from it against ps_threshold, so a later auxiliary encoder can use it without touching the existing Tobit loss path. - LabelRecord.raw_ps_readout (types.py): new optional field, defaults to None so existing records are unaffected. - parse_campaign_state / parse_pretrain_records (planning.py): new optional log2fc_column populates raw_ps_readout for PS and DRC rows. - training_records_for_refit (planning.py): when a PS-INTERVAL record is dropped for a DRC-upgraded compound, its raw_ps_readout now survives onto the surviving DRC record if the DRC record has none of its own (e.g. a DRC upgrade acquired directly through the oracle). - PlanDataConfig / PretrainDataConfig (config.py): log2fc_column wired through from_yaml via the existing kwargs splat; CLI plan and pretrain paths thread it through in cli.py.
Single-scalar raw_ps_readout: float | None was too narrow: real PS data can carry several auxiliary readouts per compound (log2FC at multiple concentrations, a direct pIC50), not just one. Replace it with raw_ps_readouts: dict[str, float], keyed by source column name, mirroring the existing IterationResults.metrics: dict[str, float] shape already used elsewhere in types.py for the same kind of dynamically-named data. - LabelRecord.raw_ps_readouts (types.py): dict field, default empty. - parse_campaign_state / parse_pretrain_records (planning.py): log2fc_column: str | None -> log2fc_columns: list[str] | None; each column's non-blank value is read into the dict under its column name. - training_records_for_refit (planning.py): on a DRC upgrade, the dropped PS record's readouts are merged onto the surviving DRC record's own readouts (DRC's own keys take precedence on conflict) rather than a single scalar copy-if-none. - PlanDataConfig / PretrainDataConfig (config.py) and cli.py: renamed and re-typed to match.
Phase 1 config scaffold for #36: PipelineConfig.auxiliary_encoder is None unless the YAML supplies an auxiliary_encoder block, so moal plan is unaffected until a campaign opts in. - freeze_epochs / embedding_dim / checkpoint_path fields, following the same shape as ModelConfig's freeze/checkpoint parameters. - Deliberately does not implement the plate/batch normalization step the issue describes as a non-optional pretraining prerequisite; readouts in raw_ps_readouts are used as-is. Documented as a known limitation in the config docstring pending a plate/batch identifier in moal's campaign-state schema. - Shares the main model's ChemProp/CheMeleon backbone construction rather than a bespoke architecture; mean aggregation is forced when from_foundation=chemeleon, since CheMeleon's own pretraining used a mean readout. - tests/test_config.py (new): defaults to None when absent, round-trips through from_yaml when supplied.
Pulls _build_model / _load_foundation_weights out into module-level build_mpnn() / load_foundation_weights() functions so the upcoming auxiliary log2FC/pIC50 encoder (#36 Phase 1) can share the exact same backbone-construction path (foundation weight loading, mean-pooling aggregation) instead of duplicating it. build_mpnn() also grows an n_tasks parameter (default 1, matching current single-target behavior) so the auxiliary encoder's multi-task head can reuse it. ChemPropLightningModule._build_model is now a thin wrapper delegating to build_mpnn(); no behavioral change to the main model. tests/test_model.py: two tests patched moal.model.ChemPropLightningModule ._load_foundation_weights, which no longer exists as an instance method; updated to patch the module-level moal.model.load_foundation_weights.
Masked multi-task pretraining over LabelRecord.raw_ps_readouts, sharing the main model's ChemProp/CheMeleon backbone construction (build_mpnn) rather than a bespoke architecture. moal plan-only; not wired into any CLI path yet (Phase 2 wires it into the main model's prediction path). - masked_mse_loss: per-task MSE restricted to observed (mask=True) entries; a task with no observed values in a batch contributes zero gradient without raising, so partial readout coverage across compounds trains cleanly. - AuxiliaryEncoderModule: LightningModule wrapping build_mpnn with n_tasks = number of distinct readout keys seen in the training data. Freeze/unfreeze schedule mirrors ChemPropLightningModule's, scheduled independently via AuxiliaryEncoderConfig.freeze_epochs. - AuxiliaryDataModule / _AuxiliaryDataset: train/val split and batching for (mol_graph, target_row, mask_row) triples, following the same shape as moal.dataset.MixedFidelityDataModule. - pretrain_auxiliary_encoder: retrains from scratch on every call by default (task_names is the sorted union of raw_ps_readouts keys across the given records); config.checkpoint_path is the explicit opt-in to skip retraining and load a cached checkpoint instead. - save_auxiliary_encoder_checkpoint / load_auxiliary_encoder_checkpoint: the checkpoint format pretrain_auxiliary_encoder's opt-in path reads, storing task_names alongside the state_dict since the predictor head's width depends on it. AuxiliaryEncoderConfig (config.py) grows the backbone/optimization fields pretraining needs: from_foundation, ffn_hidden_dim, ffn_num_layers, message_hidden_dim, depth, lr, weight_decay, max_epochs.
Exposes the pretrained auxiliary encoder's pooled structural embedding (chemprop's MPNN.fingerprint: message-passing + mean pooling + batch-norm, stopping short of the multi-task predictor head) as the fallback input for the concatenation architecture's never-screened-compound case.
Concatenates, per compound, either its own observed raw_ps_readouts (when PS-screened) or the pretrained AuxiliaryEncoderModule's structural embedding (when never PS-screened) onto the pooled graph embedding before the pEC50 predictor head, via chemprop's native MPNN.forward(bmg, X_d=...) support rather than a bespoke predictor wrapper. A provenance flag distinguishes the two paths per compound. Generalizes the paper's single-scalar-log2FC concatenation to moal's multi-readout raw_ps_readouts (Phase 0): the feature vector is [readout_vector, readout_mask, aux_embedding, provenance_flag], so partial per-task coverage and multiple readout keys both fall out of the same masked-vector shape used by the auxiliary encoder's own pretraining, rather than needing a separate mechanism. - build_mpnn (model.py) grows extra_input_dim, forwarded to RegressionFFN's input_dim. - concatenation_feature_dim / build_concatenation_features: pure functions computing the feature width and the per-compound feature matrix; only computes embeddings for compounds actually lacking readouts, since the embedding forward pass is the expensive path. - ConcatenationChemPropLightningModule: same refit/predict_smiles contract, freeze schedule, and CensoredRegressionLoss as ChemPropLightningModule (duplicated rather than shared via inheritance, matching the auxiliary encoder's precedent, to avoid entangling this experimental path with the well-tested main model). - predict_smiles chunks manually rather than through chemprop's build_dataloader, so each chunk's x_d slice is trivially aligned with its BatchMolGraph by construction instead of depending on undocumented dataloader batch-boundary behavior. Not yet wired into the moal plan CLI path, consistent with Phase 1.
CostAwareGreedyAcquisition previously ranked every prediction on equal footing, with no notion that a compound scored through the concatenation architecture's (Phase 2) auxiliary-embedding path rests on strictly more layers of inference than one scored from an observed readout or a plain graph-only prediction. - embedding_provenance_discount: constructor param, default 1.0 (no-op). Must be in (0.0, 1.0]. - select() / score_summary(): new optional provenance / ps_labeled_provenance arrays (boolean, aligned with the prediction arrays). A True entry gets its DRC and PS scores multiplied by the discount before ranking; omitting provenance (the default) preserves exact current behavior, verified by test_default_discount_is_noop. - score_summary() rows gain an embedding_derived field for transparency in the annotated campaign-state CSV. - AcquisitionConfig grows the matching embedding_provenance_discount field, wired through from_yaml via the existing kwargs splat. Applies specifically to the concatenation architecture per the issue: the retrained-encoder architecture (not implemented here) produces a single uniform prediction path with no per-compound provenance split, so this phase doesn't apply to it. Not yet wired into moal plan's CLI: nothing currently computes or passes a provenance array end-to-end, since the concatenation model itself isn't wired into the CLI either (Phase 2 note).
Prerequisite for wiring the concatenation architecture into moal plan's CLI: annotate_campaign_state now accepts an optional provenance array (aligned with predictions), splits it into the unqueried/upgrade slices the same way predictions are split, and forwards each slice to acquisition.score_summary(). Adds an embedding_derived output column, always populated (False when provenance is None) for transparency in the annotated CSV, matching score_summary's own always-present field. AuxiliaryEncoderModule.embedding_dim: new property exposing the backbone's native output width, replacing direct aux_encoder.model.message_passing.output_dim reads from concatenation_model.py's build_concatenation_features.
moal plan now has a live, selectable path through Phases 1-3 rather
than only correct-but-unreachable library code:
- When cfg.auxiliary_encoder is None (default), behavior is completely
unchanged: _build_plan_model / ChemPropLightningModule as before.
- When set, plan() pretrains the auxiliary encoder on this run's
fit_records, builds a ConcatenationChemPropLightningModule sized from
the encoder's task_names/embedding_dim, refits it, and scores
inference targets through it. A per-compound provenance array (True
wherever raw_ps_readouts was empty, i.e. never PS-screened) flows
into annotate_campaign_state so Phase 3's acquisition discount
applies to embedding-derived predictions.
- _build_acquisition was missing embedding_provenance_discount
entirely — AcquisitionConfig's field existed but had no path from
YAML into CostAwareGreedyAcquisition, so it was a silent no-op
regardless of what a campaign config set. Fixed as part of this
wiring pass, verified end-to-end (score discount observed in a real
plan run's output CSV) rather than left to be caught later.
- _inference_readouts(): unqueried compounds always get {} (never
PS-screened by construction); PS-upgrade compounds get their
training record's raw_ps_readouts looked up by canonical SMILES.
Verified against real end-to-end moal plan runs (not just unit tests):
never-screened compounds correctly flagged embedding_derived=True, the
PS-upgrade candidate with an observed readout flagged False, and the
acquisition discount visibly scaling embedding-derived scores when
configured.
…ules
Three separate LightningDataModule subclasses (MixedFidelityDataModule,
AuxiliaryDataModule, _ConcatenatedDataModule) forced n_val to at least 1
regardless of val_fraction, so val_fraction=0.0 silently still carved out
a validation split instead of disabling it. Fixed by only applying the
floor when val_fraction > 0.0.
Also fixed val_dataloader() returning None when no split exists: Lightning
rejects None from this hook ("An invalid dataloader was returned"), so all
three now return an empty DataLoader instead.
…uxiliary_trainer config AuxiliaryModelConfig now describes architecture only; max_epochs and val_fraction move to a new PipelineConfig.auxiliary_trainer (TrainerConfig), scheduled independently from the main model's trainer. Also adds use_observed_readout to AuxiliaryModelConfig, exposing the concatenation architecture's embedding-vs-embedding+raw-value toggle in YAML. pretrain_auxiliary_encoder takes an explicit max_epochs parameter instead of reading it off the config, and its trainer kwargs now set log_every_n_steps=1 so small readout-bearing pools don't silently suppress step-level loss logging (mirrors TrainerConfig's existing rationale for the main model). Wires auxiliary_trainer through moal plan's CLI with a dedicated CSVLogger (mirroring the main model's), so auxiliary-encoder loss curves are actually persisted and inspectable. The main model now trains on DRC records only: PS records have already contributed what they can via the frozen auxiliary embedding, so re-supervising the Tobit loss with PS's noisier labels on top would duplicate signal and dilute the embedding-path training examples (the only ones representative of never-screened inference targets) beneath the much larger observed-readout-path population.
…ctor-only build_concatenation_features previously routed each compound through either its observed readout OR the auxiliary encoder's structural embedding, never both. That meant compounds with an observed readout (the majority of DRC training rows) never fed their embedding into the main model at all, while every blind-test compound is embedding-only — a train/serve distribution mismatch on the exact input pathway inference depends on. Now the embedding is always computed for every compound; use_observed_readout (default True) controls only whether the raw readout/mask block is additionally populated on top of it. Moved this flag from a per-call parameter on refit()/predict_smiles() to a constructor-only attribute on ConcatenationChemPropLightningModule, since a mismatch between training-time and inference-time routing would silently exercise untrained weights (e.g. train=False then infer=True feeds nonzero values into input dimensions that received zero gradient signal throughout training). Also removes w_drc/w_ps from this class: refit() now rejects any non-DOSE_RESPONSE record, so the PS loss branch never fires and that weighting pair would be a redundant, no-op scalar confounded with sigma. (ModelConfig and CensoredRegressionLoss keep w_drc/w_ps for moal simulate and non-auxiliary moal plan, where PS/DRC mixed training still applies.)
annotate_campaign_state previously wrote only the derived ps_score/drc_score/ overall_score/recommendation columns, dropping the model's raw predicted pEC50 value. Threads it through for both unqueried rows and PS-upgrade rows so it's available for downstream evaluation (e.g. comparing directly against unblinded ground truth) without re-running inference.
chemprop's build_dataloader no longer accepts drop_last as an override kwarg — it computes it internally, dropping the last batch whenever dataset_size % batch_size == 1 to protect batch-norm during training. The explicit drop_last=False at ChemPropLightningModule.predict_smiles and AuxiliaryEncoderModule.embed_smiles's call sites collided with that internal computation, crashing with a TypeError on the installed chemprop version. Extracted safe_inference_batch_size() (moal/model.py) to shrink the batch size (never below 1) until the remainder condition no longer holds, rather than dropping a molecule from inference output — a dropped molecule would silently misalign predictions/embeddings with the input SMILES order.
Extends TrainerConfig (both trainer and auxiliary_trainer blocks) with early_stopping, early_stopping_monitor, early_stopping_patience, early_stopping_mode, and early_stopping_min_delta fields, wired into to_dict() as an EarlyStopping callback when enabled. Off by default (early_stopping=False), so existing configs train for exactly max_epochs as before. The auxiliary encoder logs aux_val_loss rather than val_loss, so auxiliary_trainer configs enabling early stopping must set early_stopping_monitor explicitly or EarlyStopping raises at construction — documented on the field.
Adds AuxiliaryEncoderModule.predict_smiles(), running the full forward pass (message-passing + predictor head) rather than embed_smiles's pooled-embedding-only path, so the encoder can produce a readout prediction for any SMILES, not just ones it observed during training. Wires this into the concatenation architecture as a third feature block, alongside the existing observed-readout and structural-embedding blocks (concatenation_feature_dim: 2*n_tasks + embedding_dim + 1 -> 3*n_tasks + embedding_dim + 1), gated by a new AuxiliaryModelConfig.use_predicted_readout flag (default False, independent of use_observed_readout). Unlike the observed-readout block, the predicted-readout block is never masked or zeroed for missing data: the encoder can predict a value for any compound, so it is populated identically at both training and inference time by construction. This replaces an earlier approach that precomputed predicted values into a CSV column and read them through the existing use_observed_readout path — that approach hit a structural train/inference mismatch, since _inference_readouts hardcodes an empty readout dict for unqueried compounds regardless of what a CSV column contains, so unqueried (i.e. evaluated) compounds never actually received the precomputed value at inference while training records used it whenever present.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What changed and why
Implements Phases 0-3 of #36: using log2FC (and other continuous primary-screen readouts, e.g. pIC50) as an auxiliary signal for pEC50 prediction in
moal plan, following the concatenation architecture from Buterez et al. (2024) / the PXR challenge report cited in the issue.The feature is off by default.
moal planbehaves exactly as it does today unless a campaign config setsauxiliary_encoder:. When it does:LabelRecord.raw_ps_readouts: dict[str, float]carries a compound's observed continuous readouts (log2FC at one or more concentrations, a direct pIC50, etc.), independent of the LEFT/INTERVAL censoring already derived fromvalueagainstps_threshold. Named-dict rather than a single scalar, since real PS data can carry several readouts per compound, mirroringIterationResults.metrics: dict[str, float]'s existing shape for the same kind of dynamically-named data. Readouts survive a PS-to-DRC upgrade (training_records_for_refit), which the auxiliary encoder and concatenation architecture both depend on.moal/auxiliary_encoder.pypretrains a ChemProp encoder via masked multi-task regression overraw_ps_readouts, sharing the main model's backbone construction (build_mpnn, now extracted as a shared function) rather than a bespoke architecture. Retrains from scratch on everymoal planinvocation by default;checkpoint_pathis an explicit opt-in to skip that.moal/concatenation_model.pyconcatenates, per compound, either its own observed readouts or the auxiliary encoder's structural embedding (for compounds never PS-screened) plus a provenance flag, onto the pooled graph embedding before the pEC50 head. Uses chemprop's nativeMPNN.forward(bmg, X_d=...)support rather than a custom predictor wrapper.CostAwareGreedyAcquisitiongains an optionalprovenancearray andembedding_provenance_discountconfig, so an embedding-derived prediction (one more layer of inference than an observed-input prediction) can be discounted in ranking. Default discount is 1.0 (no-op).moal plannow has a live, selectable path through all of the above, verified against real end-to-end runs (not just unit tests).Known limitations / deliberate deviations from the issue text
moal's campaign-state schema has no plate/batch identifier at all, and adding one is a separate schema design question. Readouts are used as-is. Documented inAuxiliaryEncoderConfig's docstring.auxiliary_encodershould be treated as unvalidated for real campaign data, per the issue's own framing.from_foundation="chemeleon"(the default, for both the main model and the auxiliary encoder), aggregation is constrained to mean pooling, since CheMeleon's own pretraining used a mean readout. The issue's language ("must be an adaptive, learned pooling function... not mean") only holds whenfrom_foundation=False.How to verify
python -m pytest tests/— 371 tests pass, including new coverage for every module touched.pre-commit run --all-files— ruff, ruff-format, pyright all pass.moal planruns (see commit messages for details): confirmed never-screened compounds are flaggedembedding_derived=Trueand route through the auxiliary encoder's embedding; a PS-upgrade compound with an observed readout is flaggedFalseand routes through its own value; the acquisition discount visibly scales embedding-derived scores whenacquisition.embedding_provenance_discountis set below 1.0.log2fc_columns:list todata.plan, anauxiliary_encoder:block to the campaign config, and runmoal planagainst a campaign-state CSV with a log2FC column.Scope
Wider than a single logical change by commit count (10 commits), but each commit is one coherent step of a single feature (issue #36, Phases 0-3 plus CLI wiring); none is independently mergeable or reviewable in isolation from the others, since Phase 2 depends on Phase 1's encoder, Phase 3's discount only matters once Phase 2 produces provenance, and none of it is reachable without the CLI wiring.
Notes for reviewers
raw_ps_readouts), and (2) reusing chemprop's nativeX_dconcatenation point instead of a custom predictor wrapper for the concatenation architecture.AuxiliaryEncoderModuleandConcatenationChemPropLightningModuleduplicateChemPropLightningModule's freeze/unfreeze schedule (~15 lines each) rather than sharing it via a mixin. Deliberate: avoids touching the well-tested main model's class hierarchy for a modest amount of duplication. Flagging in case reviewers weigh that trade-off differently._build_acquisitionwas missingembedding_provenance_discountentirely until caught during CLI-wiring verification — the config field existed but had no path from YAML into the acquisition function. Fixed in the last commit; worth a close look since it's exactly the kind of silent-no-op bug that's easy to miss in review.Checklist