From 9bca2963c9c7cec512b706379f5fbb57bdc0fd05 Mon Sep 17 00:00:00 2001 From: Sean Colby Date: Thu, 23 Jul 2026 15:31:21 -0800 Subject: [PATCH 01/18] Add raw_ps_readout to LabelRecord for log2FC auxiliary signal 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. --- moal/cli.py | 2 + moal/config.py | 13 +++++++ moal/planning.py | 66 ++++++++++++++++++++++++++++++--- moal/types.py | 11 ++++++ tests/test_planning.py | 83 ++++++++++++++++++++++++++++++++++++++++++ 5 files changed, 169 insertions(+), 6 deletions(-) diff --git a/moal/cli.py b/moal/cli.py index 165b897..156778b 100644 --- a/moal/cli.py +++ b/moal/cli.py @@ -335,6 +335,7 @@ def plan(config: Path, output_dir: Path | None, verbose: bool) -> None: relation_column=cfg.data.plan.relation_column, value_column=cfg.data.plan.value_column, weight_column=cfg.data.plan.weight_column, + log2fc_column=cfg.data.plan.log2fc_column, is_canonical=cfg.data.plan.is_canonical, expected_ps_threshold=cfg.oracle.ps_threshold, ) @@ -508,6 +509,7 @@ def _load_pretrain_records( relation_column=pretrain_cfg.relation_column, value_column=pretrain_cfg.value_column, weight_column=pretrain_cfg.weight_column, + log2fc_column=pretrain_cfg.log2fc_column, is_canonical=pretrain_cfg.is_canonical, expected_ps_threshold=cfg.oracle.ps_threshold, ) diff --git a/moal/config.py b/moal/config.py index 0dfc664..4a35525 100644 --- a/moal/config.py +++ b/moal/config.py @@ -299,6 +299,11 @@ class PretrainDataConfig: Optional column name for per-sample loss weights. When set, each labeled row's weight is read from this column (NaN / missing cells default to 1.0). When None (default), all records receive weight=1.0. + log2fc_column : str or None + Optional column name for the observed continuous log2FC primary-screen + readout. When set, populates ``LabelRecord.raw_ps_readout`` for PS + rows. When None (default), ``raw_ps_readout`` is ``None`` for every + record. is_canonical : bool When False (default), SMILES are canonicalized via RDKit during parsing. @@ -309,6 +314,7 @@ class PretrainDataConfig: relation_column: str = "relation" value_column: str = "value" weight_column: str | None = None + log2fc_column: str | None = None is_canonical: bool = False @@ -389,6 +395,12 @@ class PlanDataConfig: Optional column name for per-sample loss weights. When set, each labeled row's weight is read from this column (NaN / missing cells default to 1.0). When None (default), all records receive weight=1.0. + log2fc_column : str or None + Optional column name for the observed continuous log2FC primary-screen + readout. When set, populates ``LabelRecord.raw_ps_readout`` for PS + rows (and DRC rows for upgraded compounds, when the CSV carries the + readout on that row). When None (default), ``raw_ps_readout`` is + ``None`` for every record. is_canonical : bool When False (default), SMILES are canonicalized via RDKit during parsing. @@ -400,6 +412,7 @@ class PlanDataConfig: relation_column: str = "relation" value_column: str = "value" weight_column: str | None = None + log2fc_column: str | None = None is_canonical: bool = False diff --git a/moal/planning.py b/moal/planning.py index 7895c3e..dd3a6bd 100644 --- a/moal/planning.py +++ b/moal/planning.py @@ -54,6 +54,7 @@ def parse_campaign_state( relation_column: str = "relation", value_column: str = "value", weight_column: str | None = None, + log2fc_column: str | None = None, is_canonical: bool = False, expected_ps_threshold: float | None = None, ) -> CampaignState: @@ -87,6 +88,13 @@ def parse_campaign_state( each labeled row's weight is read from this column (NaN / empty cells default to 1.0). Must be a finite positive float when present. When None (default), all records receive ``weight=1.0``. + log2fc_column : str or None + Optional column name for the compound's observed continuous log2FC + readout from a primary-screen assay. When provided, populates + ``LabelRecord.raw_ps_readout`` for PS rows (NaN / empty cells leave + it ``None``). Ignored for DRC-only (``==``) rows, since a compound + with no PS row has no observed log2FC. When None (default), + ``raw_ps_readout`` is ``None`` for every record. is_canonical : bool When True, skip RDKit canonicalization. expected_ps_threshold : float or None @@ -105,6 +113,10 @@ def parse_campaign_state( raise ValueError( f"weight_column {weight_column!r} not found in state CSV, got {sorted(df.columns)}" ) + if log2fc_column is not None and log2fc_column not in df.columns: + raise ValueError( + f"log2fc_column {log2fc_column!r} not found in state CSV, got {sorted(df.columns)}" + ) training_records: list[LabelRecord] = [] unqueried_rows: list[tuple[int, str]] = [] @@ -176,6 +188,20 @@ def parse_campaign_state( f"Row {csv_row}: weight must be finite and positive, got {weight}." ) + raw_ps_readout: float | None = None + if log2fc_column is not None: + log2fc_raw = row.get(log2fc_column, None) + if not (pd.isna(log2fc_raw) or str(log2fc_raw).strip() == ""): + try: + raw_ps_readout = float(log2fc_raw) + except (TypeError, ValueError) as exc: + raise ValueError( + f"Row {csv_row}: log2FC must be a finite numeric readout," + f" got {log2fc_raw!r}." + ) from exc + if not math.isfinite(raw_ps_readout): + raise ValueError(f"Row {csv_row}: log2FC must be finite, got {log2fc_raw!r}.") + if relation == "==": record = LabelRecord( smiles=raw_smiles, @@ -187,6 +213,7 @@ def parse_campaign_state( cost=cost_drc, iteration=_PLAN_MODE_ITERATION, weight=weight, + raw_ps_readout=raw_ps_readout, ) else: if expected_ps_threshold is not None and not math.isclose( @@ -206,6 +233,7 @@ def parse_campaign_state( cost=cost_ps, iteration=_PLAN_MODE_ITERATION, weight=weight, + raw_ps_readout=raw_ps_readout, ) # PS hits are DRC-upgrade inference targets in addition to training records if relation == ">=": @@ -248,6 +276,7 @@ def parse_pretrain_records( relation_column: str = "relation", value_column: str = "value", weight_column: str | None = None, + log2fc_column: str | None = None, is_canonical: bool = False, expected_ps_threshold: float | None = None, ) -> list[LabelRecord]: @@ -278,6 +307,10 @@ def parse_pretrain_records( Optional column name for per-sample loss weights. Forwarded to :func:`parse_campaign_state`. When None (default), all records receive ``weight=1.0``. + log2fc_column : str or None + Optional column name for the observed log2FC readout. Forwarded to + :func:`parse_campaign_state`. When None (default), ``raw_ps_readout`` + is ``None`` for every record. is_canonical : bool When True, skip RDKit canonicalization. expected_ps_threshold : float or None @@ -300,6 +333,7 @@ def parse_pretrain_records( relation_column=relation_column, value_column=value_column, weight_column=weight_column, + log2fc_column=log2fc_column, is_canonical=is_canonical, expected_ps_threshold=expected_ps_threshold, ) @@ -318,7 +352,11 @@ def training_records_for_refit(records: list[LabelRecord]) -> list[LabelRecord]: When a compound has both a PS INTERVAL record (``>=`` hit) and a DRC EXACT record, the PS record is excluded to prevent double-weighting during model training. PS LEFT records (``<`` misses) are always - retained regardless of DRC coverage. + retained regardless of DRC coverage. If the excluded PS record carries + an observed ``raw_ps_readout`` that the surviving DRC record lacks (e.g. + a DRC upgrade acquired directly through the oracle, with no log2FC of + its own), the readout is copied onto the surviving DRC record so it is + not lost for the auxiliary log2FC encoder. Parameters ---------- @@ -336,15 +374,31 @@ def training_records_for_refit(records: list[LabelRecord]) -> list[LabelRecord]: upgraded_smiles = { rec.canonical_smiles for rec in records if rec.fidelity == QueryType.DOSE_RESPONSE } - return [ - rec + upgrade_readouts = { + rec.canonical_smiles: rec.raw_ps_readout for rec in records - if not ( + if rec.fidelity == QueryType.PRIMARY_SCREEN + and rec.censoring_type == CensoringType.INTERVAL + and rec.canonical_smiles in upgraded_smiles + and rec.raw_ps_readout is not None + } + + result = [] + for rec in records: + if ( rec.fidelity == QueryType.PRIMARY_SCREEN and rec.censoring_type == CensoringType.INTERVAL and rec.canonical_smiles in upgraded_smiles - ) - ] + ): + continue + if ( + rec.fidelity == QueryType.DOSE_RESPONSE + and rec.raw_ps_readout is None + and rec.canonical_smiles in upgrade_readouts + ): + rec = replace(rec, raw_ps_readout=upgrade_readouts[rec.canonical_smiles]) + result.append(rec) + return result def annotate_campaign_state( diff --git a/moal/types.py b/moal/types.py index 30328f3..96650ee 100644 --- a/moal/types.py +++ b/moal/types.py @@ -80,6 +80,13 @@ class LabelRecord: Normalized to mean=1.0 within each fidelity class by :func:`~moal.planning.normalize_record_weights` before training. Defaults to 1.0 (uniform weighting). + raw_ps_readout : float or None + The compound's observed continuous log2 fold-change from a + single-concentration primary screen, independent of the LEFT/INTERVAL + censoring derived from it against ``oracle.ps_threshold``. Retained on + the surviving DRC record after a PS-to-DRC upgrade so the auxiliary + log2FC encoder (see ``AuxiliaryEncoderConfig``) can still use it. + ``None`` when the compound has never been PS-screened. """ smiles: str @@ -95,6 +102,10 @@ class LabelRecord: :func:`~moal.planning.normalize_record_weights` before training so the global ``w_drc`` / ``w_ps`` scale relationship is preserved. Default 1.0 is a no-op and preserves backward compatibility.""" + raw_ps_readout: float | None = None + """Observed log2FC from the primary screen, kept separate from the + censored ``value``/``censoring_type`` pair so it survives a PS-to-DRC + upgrade for use as an auxiliary-encoder training input.""" @dataclass diff --git a/tests/test_planning.py b/tests/test_planning.py index 54b214d..f7d96a6 100644 --- a/tests/test_planning.py +++ b/tests/test_planning.py @@ -348,6 +348,89 @@ def test_refit_records_drop_upgraded_interval_ps_rows(self, preprocessor): for r in fit_records ) + def test_log2fc_column_populates_raw_ps_readout_on_ps_rows(self, preprocessor): + """log2fc_column values must land on LabelRecord.raw_ps_readout for PS rows.""" + df = _state_df( + {"smiles": "CCO", "relation": "<", "value": 5.0, "log2fc": -1.2}, + {"smiles": "CCN", "relation": ">=", "value": 5.0, "log2fc": 3.4}, + ) + + state = parse_campaign_state( + df, + cost_ps=1.0, + cost_drc=10.0, + upper_bound=11.0, + preprocessor=preprocessor, + log2fc_column="log2fc", + expected_ps_threshold=5.0, + ) + + readouts = {r.canonical_smiles: r.raw_ps_readout for r in state.training_records} + assert readouts[preprocessor.canonicalize("CCO")] == -1.2 + assert readouts[preprocessor.canonicalize("CCN")] == 3.4 + + def test_log2fc_column_blank_cell_leaves_raw_ps_readout_none(self, preprocessor): + """An empty log2fc cell must leave raw_ps_readout as None rather than raising.""" + df = _state_df({"smiles": "CCO", "relation": "<", "value": 5.0, "log2fc": ""}) + + state = parse_campaign_state( + df, + cost_ps=1.0, + cost_drc=10.0, + upper_bound=11.0, + preprocessor=preprocessor, + log2fc_column="log2fc", + expected_ps_threshold=5.0, + ) + + assert state.training_records[0].raw_ps_readout is None + + def test_missing_log2fc_column_raises(self, preprocessor): + """Requesting a log2fc_column absent from the CSV must raise ValueError.""" + df = _state_df({"smiles": "CCO", "relation": "<", "value": 5.0}) + + with pytest.raises(ValueError, match="log2fc_column"): + parse_campaign_state( + df, + cost_ps=1.0, + cost_drc=10.0, + upper_bound=11.0, + preprocessor=preprocessor, + log2fc_column="log2fc", + ) + + def test_refit_records_carry_raw_ps_readout_onto_surviving_drc_record(self, preprocessor): + """When a DRC-upgrade record has no log2FC of its own, the dropped PS record's readout must be copied onto it.""" + upgraded_smiles = preprocessor.canonicalize("CCO") + ps_record = LabelRecord( + smiles="CCO", + canonical_smiles=upgraded_smiles, + value=5.0, + upper_bound=11.0, + censoring_type=CensoringType.INTERVAL, + fidelity=QueryType.PRIMARY_SCREEN, + cost=1.0, + iteration=0, + raw_ps_readout=3.4, + ) + drc_record = LabelRecord( + smiles="CCO", + canonical_smiles=upgraded_smiles, + value=7.2, + upper_bound=7.2, + censoring_type=CensoringType.EXACT, + fidelity=QueryType.DOSE_RESPONSE, + cost=10.0, + iteration=1, + raw_ps_readout=None, + ) + + fit_records = training_records_for_refit([ps_record, drc_record]) + + assert len(fit_records) == 1 + assert fit_records[0].fidelity == QueryType.DOSE_RESPONSE + assert fit_records[0].raw_ps_readout == 3.4 + def test_custom_column_names_are_supported(self, preprocessor): """Non-default smiles, relation, and value column names must be mapped correctly throughout parsing.""" df = pd.DataFrame( From 018ab3c2e4e1d6975b2a640cbbdaeafc7deff2b7 Mon Sep 17 00:00:00 2001 From: Sean Colby Date: Thu, 23 Jul 2026 15:46:26 -0800 Subject: [PATCH 02/18] Generalize raw_ps_readout to a named multi-readout dict 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. --- moal/cli.py | 4 +- moal/config.py | 28 ++++++------ moal/planning.py | 96 ++++++++++++++++++++++-------------------- moal/types.py | 24 ++++++----- tests/test_planning.py | 39 +++++++++-------- 5 files changed, 99 insertions(+), 92 deletions(-) diff --git a/moal/cli.py b/moal/cli.py index 156778b..f58f26d 100644 --- a/moal/cli.py +++ b/moal/cli.py @@ -335,7 +335,7 @@ def plan(config: Path, output_dir: Path | None, verbose: bool) -> None: relation_column=cfg.data.plan.relation_column, value_column=cfg.data.plan.value_column, weight_column=cfg.data.plan.weight_column, - log2fc_column=cfg.data.plan.log2fc_column, + log2fc_columns=cfg.data.plan.log2fc_columns, is_canonical=cfg.data.plan.is_canonical, expected_ps_threshold=cfg.oracle.ps_threshold, ) @@ -509,7 +509,7 @@ def _load_pretrain_records( relation_column=pretrain_cfg.relation_column, value_column=pretrain_cfg.value_column, weight_column=pretrain_cfg.weight_column, - log2fc_column=pretrain_cfg.log2fc_column, + log2fc_columns=pretrain_cfg.log2fc_columns, is_canonical=pretrain_cfg.is_canonical, expected_ps_threshold=cfg.oracle.ps_threshold, ) diff --git a/moal/config.py b/moal/config.py index 4a35525..7c7b70c 100644 --- a/moal/config.py +++ b/moal/config.py @@ -299,11 +299,12 @@ class PretrainDataConfig: Optional column name for per-sample loss weights. When set, each labeled row's weight is read from this column (NaN / missing cells default to 1.0). When None (default), all records receive weight=1.0. - log2fc_column : str or None - Optional column name for the observed continuous log2FC primary-screen - readout. When set, populates ``LabelRecord.raw_ps_readout`` for PS - rows. When None (default), ``raw_ps_readout`` is ``None`` for every - record. + log2fc_columns : list[str] or None + Optional column names for observed continuous auxiliary readouts + (e.g. log2FC at one or more primary-screen concentrations, a direct + pIC50). When set, populates ``LabelRecord.raw_ps_readouts`` keyed by + column name. When None (default), ``raw_ps_readouts`` is empty for + every record. is_canonical : bool When False (default), SMILES are canonicalized via RDKit during parsing. @@ -314,7 +315,7 @@ class PretrainDataConfig: relation_column: str = "relation" value_column: str = "value" weight_column: str | None = None - log2fc_column: str | None = None + log2fc_columns: list[str] | None = None is_canonical: bool = False @@ -395,12 +396,13 @@ class PlanDataConfig: Optional column name for per-sample loss weights. When set, each labeled row's weight is read from this column (NaN / missing cells default to 1.0). When None (default), all records receive weight=1.0. - log2fc_column : str or None - Optional column name for the observed continuous log2FC primary-screen - readout. When set, populates ``LabelRecord.raw_ps_readout`` for PS - rows (and DRC rows for upgraded compounds, when the CSV carries the - readout on that row). When None (default), ``raw_ps_readout`` is - ``None`` for every record. + log2fc_columns : list[str] or None + Optional column names for observed continuous auxiliary readouts + (e.g. log2FC at one or more primary-screen concentrations, a direct + pIC50). When set, populates ``LabelRecord.raw_ps_readouts`` keyed by + column name for PS rows (and DRC rows for upgraded compounds, when + the CSV carries the readout on that row). When None (default), + ``raw_ps_readouts`` is empty for every record. is_canonical : bool When False (default), SMILES are canonicalized via RDKit during parsing. @@ -412,7 +414,7 @@ class PlanDataConfig: relation_column: str = "relation" value_column: str = "value" weight_column: str | None = None - log2fc_column: str | None = None + log2fc_columns: list[str] | None = None is_canonical: bool = False diff --git a/moal/planning.py b/moal/planning.py index dd3a6bd..e7645d8 100644 --- a/moal/planning.py +++ b/moal/planning.py @@ -54,7 +54,7 @@ def parse_campaign_state( relation_column: str = "relation", value_column: str = "value", weight_column: str | None = None, - log2fc_column: str | None = None, + log2fc_columns: list[str] | None = None, is_canonical: bool = False, expected_ps_threshold: float | None = None, ) -> CampaignState: @@ -88,13 +88,13 @@ def parse_campaign_state( each labeled row's weight is read from this column (NaN / empty cells default to 1.0). Must be a finite positive float when present. When None (default), all records receive ``weight=1.0``. - log2fc_column : str or None - Optional column name for the compound's observed continuous log2FC - readout from a primary-screen assay. When provided, populates - ``LabelRecord.raw_ps_readout`` for PS rows (NaN / empty cells leave - it ``None``). Ignored for DRC-only (``==``) rows, since a compound - with no PS row has no observed log2FC. When None (default), - ``raw_ps_readout`` is ``None`` for every record. + log2fc_columns : list[str] or None + Optional column names for the compound's observed continuous + auxiliary readouts (e.g. log2FC at one or more primary-screen + concentrations, a direct pIC50). When provided, each column's value + is read into ``LabelRecord.raw_ps_readouts`` keyed by column name + (NaN / empty cells are omitted rather than stored). When None + (default), ``raw_ps_readouts`` is empty for every record. is_canonical : bool When True, skip RDKit canonicalization. expected_ps_threshold : float or None @@ -113,10 +113,12 @@ def parse_campaign_state( raise ValueError( f"weight_column {weight_column!r} not found in state CSV, got {sorted(df.columns)}" ) - if log2fc_column is not None and log2fc_column not in df.columns: - raise ValueError( - f"log2fc_column {log2fc_column!r} not found in state CSV, got {sorted(df.columns)}" - ) + if log2fc_columns is not None: + missing = [col for col in log2fc_columns if col not in df.columns] + if missing: + raise ValueError( + f"log2fc_columns {missing!r} not found in state CSV, got {sorted(df.columns)}" + ) training_records: list[LabelRecord] = [] unqueried_rows: list[tuple[int, str]] = [] @@ -188,19 +190,23 @@ def parse_campaign_state( f"Row {csv_row}: weight must be finite and positive, got {weight}." ) - raw_ps_readout: float | None = None - if log2fc_column is not None: - log2fc_raw = row.get(log2fc_column, None) - if not (pd.isna(log2fc_raw) or str(log2fc_raw).strip() == ""): - try: - raw_ps_readout = float(log2fc_raw) - except (TypeError, ValueError) as exc: - raise ValueError( - f"Row {csv_row}: log2FC must be a finite numeric readout," - f" got {log2fc_raw!r}." - ) from exc - if not math.isfinite(raw_ps_readout): - raise ValueError(f"Row {csv_row}: log2FC must be finite, got {log2fc_raw!r}.") + raw_ps_readouts: dict[str, float] = {} + for col in log2fc_columns or (): + readout_raw = row.get(col, None) + if pd.isna(readout_raw) or str(readout_raw).strip() == "": + continue + try: + readout = float(readout_raw) + except (TypeError, ValueError) as exc: + raise ValueError( + f"Row {csv_row}: column {col!r} must be a finite numeric readout," + f" got {readout_raw!r}." + ) from exc + if not math.isfinite(readout): + raise ValueError( + f"Row {csv_row}: column {col!r} must be finite, got {readout_raw!r}." + ) + raw_ps_readouts[col] = readout if relation == "==": record = LabelRecord( @@ -213,7 +219,7 @@ def parse_campaign_state( cost=cost_drc, iteration=_PLAN_MODE_ITERATION, weight=weight, - raw_ps_readout=raw_ps_readout, + raw_ps_readouts=raw_ps_readouts, ) else: if expected_ps_threshold is not None and not math.isclose( @@ -233,7 +239,7 @@ def parse_campaign_state( cost=cost_ps, iteration=_PLAN_MODE_ITERATION, weight=weight, - raw_ps_readout=raw_ps_readout, + raw_ps_readouts=raw_ps_readouts, ) # PS hits are DRC-upgrade inference targets in addition to training records if relation == ">=": @@ -276,7 +282,7 @@ def parse_pretrain_records( relation_column: str = "relation", value_column: str = "value", weight_column: str | None = None, - log2fc_column: str | None = None, + log2fc_columns: list[str] | None = None, is_canonical: bool = False, expected_ps_threshold: float | None = None, ) -> list[LabelRecord]: @@ -307,10 +313,10 @@ def parse_pretrain_records( Optional column name for per-sample loss weights. Forwarded to :func:`parse_campaign_state`. When None (default), all records receive ``weight=1.0``. - log2fc_column : str or None - Optional column name for the observed log2FC readout. Forwarded to - :func:`parse_campaign_state`. When None (default), ``raw_ps_readout`` - is ``None`` for every record. + log2fc_columns : list[str] or None + Optional column names for observed auxiliary readouts. Forwarded to + :func:`parse_campaign_state`. When None (default), + ``raw_ps_readouts`` is empty for every record. is_canonical : bool When True, skip RDKit canonicalization. expected_ps_threshold : float or None @@ -333,7 +339,7 @@ def parse_pretrain_records( relation_column=relation_column, value_column=value_column, weight_column=weight_column, - log2fc_column=log2fc_column, + log2fc_columns=log2fc_columns, is_canonical=is_canonical, expected_ps_threshold=expected_ps_threshold, ) @@ -352,11 +358,11 @@ def training_records_for_refit(records: list[LabelRecord]) -> list[LabelRecord]: When a compound has both a PS INTERVAL record (``>=`` hit) and a DRC EXACT record, the PS record is excluded to prevent double-weighting during model training. PS LEFT records (``<`` misses) are always - retained regardless of DRC coverage. If the excluded PS record carries - an observed ``raw_ps_readout`` that the surviving DRC record lacks (e.g. - a DRC upgrade acquired directly through the oracle, with no log2FC of - its own), the readout is copied onto the surviving DRC record so it is - not lost for the auxiliary log2FC encoder. + retained regardless of DRC coverage. Any ``raw_ps_readouts`` entries on + the excluded PS record are merged onto the surviving DRC record (e.g. a + DRC upgrade acquired directly through the oracle, with no readouts of its + own) so they are not lost for the auxiliary encoder; a key already present + on the DRC record's own readouts takes precedence. Parameters ---------- @@ -375,12 +381,12 @@ def training_records_for_refit(records: list[LabelRecord]) -> list[LabelRecord]: rec.canonical_smiles for rec in records if rec.fidelity == QueryType.DOSE_RESPONSE } upgrade_readouts = { - rec.canonical_smiles: rec.raw_ps_readout + rec.canonical_smiles: rec.raw_ps_readouts for rec in records if rec.fidelity == QueryType.PRIMARY_SCREEN and rec.censoring_type == CensoringType.INTERVAL and rec.canonical_smiles in upgraded_smiles - and rec.raw_ps_readout is not None + and rec.raw_ps_readouts } result = [] @@ -391,12 +397,10 @@ def training_records_for_refit(records: list[LabelRecord]) -> list[LabelRecord]: and rec.canonical_smiles in upgraded_smiles ): continue - if ( - rec.fidelity == QueryType.DOSE_RESPONSE - and rec.raw_ps_readout is None - and rec.canonical_smiles in upgrade_readouts - ): - rec = replace(rec, raw_ps_readout=upgrade_readouts[rec.canonical_smiles]) + if rec.fidelity == QueryType.DOSE_RESPONSE and rec.canonical_smiles in upgrade_readouts: + merged = {**upgrade_readouts[rec.canonical_smiles], **rec.raw_ps_readouts} + if merged != rec.raw_ps_readouts: + rec = replace(rec, raw_ps_readouts=merged) result.append(rec) return result diff --git a/moal/types.py b/moal/types.py index 96650ee..9d40f09 100644 --- a/moal/types.py +++ b/moal/types.py @@ -80,13 +80,14 @@ class LabelRecord: Normalized to mean=1.0 within each fidelity class by :func:`~moal.planning.normalize_record_weights` before training. Defaults to 1.0 (uniform weighting). - raw_ps_readout : float or None - The compound's observed continuous log2 fold-change from a - single-concentration primary screen, independent of the LEFT/INTERVAL - censoring derived from it against ``oracle.ps_threshold``. Retained on - the surviving DRC record after a PS-to-DRC upgrade so the auxiliary - log2FC encoder (see ``AuxiliaryEncoderConfig``) can still use it. - ``None`` when the compound has never been PS-screened. + raw_ps_readouts : dict[str, float] + The compound's observed continuous auxiliary readouts (e.g. log2 fold- + change at one or more primary-screen concentrations, a direct pIC50), + keyed by source column name, independent of the LEFT/INTERVAL + censoring derived from ``value`` against ``oracle.ps_threshold``. + Retained on the surviving DRC record after a PS-to-DRC upgrade so the + auxiliary encoder (see ``AuxiliaryEncoderConfig``) can still use it. + Empty when the compound has never been PS-screened. """ smiles: str @@ -102,10 +103,11 @@ class LabelRecord: :func:`~moal.planning.normalize_record_weights` before training so the global ``w_drc`` / ``w_ps`` scale relationship is preserved. Default 1.0 is a no-op and preserves backward compatibility.""" - raw_ps_readout: float | None = None - """Observed log2FC from the primary screen, kept separate from the - censored ``value``/``censoring_type`` pair so it survives a PS-to-DRC - upgrade for use as an auxiliary-encoder training input.""" + raw_ps_readouts: dict[str, float] = field(default_factory=dict) + """Observed auxiliary readouts (log2FC per concentration, direct pIC50, + etc.), keyed by source column name and kept separate from the censored + ``value``/``censoring_type`` pair so they survive a PS-to-DRC upgrade for + use as auxiliary-encoder training inputs.""" @dataclass diff --git a/tests/test_planning.py b/tests/test_planning.py index f7d96a6..7f08984 100644 --- a/tests/test_planning.py +++ b/tests/test_planning.py @@ -348,11 +348,11 @@ def test_refit_records_drop_upgraded_interval_ps_rows(self, preprocessor): for r in fit_records ) - def test_log2fc_column_populates_raw_ps_readout_on_ps_rows(self, preprocessor): - """log2fc_column values must land on LabelRecord.raw_ps_readout for PS rows.""" + def test_log2fc_columns_populate_raw_ps_readouts_on_ps_rows(self, preprocessor): + """log2fc_columns values must land on LabelRecord.raw_ps_readouts keyed by column name.""" df = _state_df( - {"smiles": "CCO", "relation": "<", "value": 5.0, "log2fc": -1.2}, - {"smiles": "CCN", "relation": ">=", "value": 5.0, "log2fc": 3.4}, + {"smiles": "CCO", "relation": "<", "value": 5.0, "log2fc_1um": -1.2, "pic50": ""}, + {"smiles": "CCN", "relation": ">=", "value": 5.0, "log2fc_1um": 3.4, "pic50": 6.8}, ) state = parse_campaign_state( @@ -361,16 +361,16 @@ def test_log2fc_column_populates_raw_ps_readout_on_ps_rows(self, preprocessor): cost_drc=10.0, upper_bound=11.0, preprocessor=preprocessor, - log2fc_column="log2fc", + log2fc_columns=["log2fc_1um", "pic50"], expected_ps_threshold=5.0, ) - readouts = {r.canonical_smiles: r.raw_ps_readout for r in state.training_records} - assert readouts[preprocessor.canonicalize("CCO")] == -1.2 - assert readouts[preprocessor.canonicalize("CCN")] == 3.4 + readouts = {r.canonical_smiles: r.raw_ps_readouts for r in state.training_records} + assert readouts[preprocessor.canonicalize("CCO")] == {"log2fc_1um": -1.2} + assert readouts[preprocessor.canonicalize("CCN")] == {"log2fc_1um": 3.4, "pic50": 6.8} - def test_log2fc_column_blank_cell_leaves_raw_ps_readout_none(self, preprocessor): - """An empty log2fc cell must leave raw_ps_readout as None rather than raising.""" + def test_log2fc_columns_blank_cell_omits_key(self, preprocessor): + """An empty log2fc cell must be omitted from raw_ps_readouts rather than raising.""" df = _state_df({"smiles": "CCO", "relation": "<", "value": 5.0, "log2fc": ""}) state = parse_campaign_state( @@ -379,28 +379,28 @@ def test_log2fc_column_blank_cell_leaves_raw_ps_readout_none(self, preprocessor) cost_drc=10.0, upper_bound=11.0, preprocessor=preprocessor, - log2fc_column="log2fc", + log2fc_columns=["log2fc"], expected_ps_threshold=5.0, ) - assert state.training_records[0].raw_ps_readout is None + assert state.training_records[0].raw_ps_readouts == {} def test_missing_log2fc_column_raises(self, preprocessor): - """Requesting a log2fc_column absent from the CSV must raise ValueError.""" + """Requesting a log2fc_columns entry absent from the CSV must raise ValueError.""" df = _state_df({"smiles": "CCO", "relation": "<", "value": 5.0}) - with pytest.raises(ValueError, match="log2fc_column"): + with pytest.raises(ValueError, match="log2fc_columns"): parse_campaign_state( df, cost_ps=1.0, cost_drc=10.0, upper_bound=11.0, preprocessor=preprocessor, - log2fc_column="log2fc", + log2fc_columns=["log2fc"], ) - def test_refit_records_carry_raw_ps_readout_onto_surviving_drc_record(self, preprocessor): - """When a DRC-upgrade record has no log2FC of its own, the dropped PS record's readout must be copied onto it.""" + def test_refit_records_merge_raw_ps_readouts_onto_surviving_drc_record(self, preprocessor): + """When a DRC-upgrade record lacks readouts of its own, the dropped PS record's readouts must be merged onto it.""" upgraded_smiles = preprocessor.canonicalize("CCO") ps_record = LabelRecord( smiles="CCO", @@ -411,7 +411,7 @@ def test_refit_records_carry_raw_ps_readout_onto_surviving_drc_record(self, prep fidelity=QueryType.PRIMARY_SCREEN, cost=1.0, iteration=0, - raw_ps_readout=3.4, + raw_ps_readouts={"log2fc_1um": 3.4}, ) drc_record = LabelRecord( smiles="CCO", @@ -422,14 +422,13 @@ def test_refit_records_carry_raw_ps_readout_onto_surviving_drc_record(self, prep fidelity=QueryType.DOSE_RESPONSE, cost=10.0, iteration=1, - raw_ps_readout=None, ) fit_records = training_records_for_refit([ps_record, drc_record]) assert len(fit_records) == 1 assert fit_records[0].fidelity == QueryType.DOSE_RESPONSE - assert fit_records[0].raw_ps_readout == 3.4 + assert fit_records[0].raw_ps_readouts == {"log2fc_1um": 3.4} def test_custom_column_names_are_supported(self, preprocessor): """Non-default smiles, relation, and value column names must be mapped correctly throughout parsing.""" From 0e792583f09d3e92496d4266cfd72286bfb7c9c6 Mon Sep 17 00:00:00 2001 From: Sean Colby Date: Thu, 23 Jul 2026 16:34:59 -0800 Subject: [PATCH 03/18] Add AuxiliaryEncoderConfig, off by default 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. --- moal/config.py | 63 ++++++++++++++++++++++++++++++++++++++++++++ tests/test_config.py | 47 +++++++++++++++++++++++++++++++++ 2 files changed, 110 insertions(+) create mode 100644 tests/test_config.py diff --git a/moal/config.py b/moal/config.py index 7c7b70c..43fba49 100644 --- a/moal/config.py +++ b/moal/config.py @@ -122,6 +122,58 @@ class ModelConfig: from_foundation: str | bool = "chemeleon" +@dataclass(frozen=True) +class AuxiliaryEncoderConfig: + """Auxiliary encoder for the primary-screen readouts in ``LabelRecord.raw_ps_readouts``. + + ``moal plan``-only (see the ``moal simulate`` exclusion in the module + docstring reference, issue #36). Off by default: ``moal plan`` behaves + exactly as it does today unless this config is explicitly set. + + Shares the main model's ChemProp/CheMeleon backbone construction + (``ModelConfig.from_foundation``, mean-pooling readout) rather than a + bespoke architecture, so its embeddings live in the same representation + space. When ``from_foundation="chemeleon"``, the readout is constrained + to mean aggregation to match CheMeleon's own pretraining; the paper's + recommended attentive readout is only reachable with + ``from_foundation=False``. + + Trains a masked multi-task regression head, one output per distinct key + observed across ``raw_ps_readouts`` (e.g. one head per log2FC + concentration, plus a head for a direct pIC50 column when present). + Compounds missing a given key contribute no gradient to that head. + + Readouts are used as-is, with no per-plate/per-batch normalization step. + The design this config implements (issue #36) specified that step as a + named, non-optional prerequisite for pretraining; it is not implemented + here and is a known, documented limitation until `moal`'s campaign-state + schema gains a plate/batch identifier. + + Attributes + ---------- + freeze_epochs : int + Number of warm-up epochs to train only the multi-task FFN head, + analogous to ``ModelConfig.freeze_epochs`` but scheduled + independently for the auxiliary encoder. + embedding_dim : int + Dimensionality of the pooled molecular embedding exposed to the main + model's concatenation architecture (Phase 2). Ignored by the + retrained-encoder architecture, which has no separate embedding + output at inference. + checkpoint_path : str or None + Explicit opt-in path to a cached auxiliary-encoder checkpoint. When + set, pretraining is skipped and this checkpoint is loaded instead. + When None (default), the auxiliary encoder is retrained from scratch + on every ``moal plan`` invocation using the current campaign-state + CSV's ``raw_ps_readouts``, so newly accumulated readouts improve the + next run automatically. + """ + + freeze_epochs: int = 5 + embedding_dim: int = 300 + checkpoint_path: str | None = None + + @dataclass(frozen=True) class AcquisitionConfig: """Acquisition function hyper-parameters. @@ -487,6 +539,10 @@ class PipelineConfig: Command-specific dataset and I/O settings. active_learning_loop : ActiveLearningLoopConfig Parameters controlling the active learning iteration loop. + auxiliary_encoder : AuxiliaryEncoderConfig or None + Optional auxiliary log2FC/pIC50 encoder for ``moal plan`` (issue #36). + ``None`` (default) disables the feature entirely; ``moal plan`` + behaves exactly as it does without this config. seed : int Global random seed for the campaign. """ @@ -498,6 +554,7 @@ class PipelineConfig: dashboard: DashboardConfig = field(default_factory=DashboardConfig) data: DataConfig = field(default_factory=DataConfig) active_learning_loop: ActiveLearningLoopConfig = field(default_factory=ActiveLearningLoopConfig) + auxiliary_encoder: AuxiliaryEncoderConfig | None = None seed: int = 42 @@ -520,6 +577,7 @@ def from_yaml(cls, path: str | Path) -> PipelineConfig: data_raw = raw.get("data", {}) simulate_raw = data_raw.get("simulate", {}) pretrain_raw = simulate_raw.pop("pretrain", {}) if isinstance(simulate_raw, dict) else {} + auxiliary_encoder_raw = raw.get("auxiliary_encoder", None) return cls( oracle=OracleConfig(**raw.get("oracle", {})), model=ModelConfig(**raw.get("model", {})), @@ -535,6 +593,11 @@ def from_yaml(cls, path: str | Path) -> PipelineConfig: plan=PlanDataConfig(**data_raw.get("plan", {})), ), active_learning_loop=ActiveLearningLoopConfig(**raw.get("active_learning_loop", {})), + auxiliary_encoder=( + AuxiliaryEncoderConfig(**auxiliary_encoder_raw) + if auxiliary_encoder_raw is not None + else None + ), seed=raw.get("seed", 42), ) diff --git a/tests/test_config.py b/tests/test_config.py new file mode 100644 index 0000000..57e3f83 --- /dev/null +++ b/tests/test_config.py @@ -0,0 +1,47 @@ +"""Tests for pipeline configuration loading.""" + +from __future__ import annotations + +from pathlib import Path + +import yaml + +from moal.config import AuxiliaryEncoderConfig, PipelineConfig + + +def _write_yaml(tmp_path: Path, raw: dict) -> Path: + path = tmp_path / "config.yaml" + with path.open("w") as f: + yaml.safe_dump(raw, f) + return path + + +class TestAuxiliaryEncoderConfig: + """Tests for AuxiliaryEncoderConfig's default-off behavior and round-trip through from_yaml.""" + + def test_defaults_to_none_when_absent(self, tmp_path): + """auxiliary_encoder must be None when the YAML has no auxiliary_encoder key.""" + path = _write_yaml(tmp_path, {"seed": 1}) + + cfg = PipelineConfig.from_yaml(path) + + assert cfg.auxiliary_encoder is None + + def test_round_trips_through_from_yaml(self, tmp_path): + """An explicit auxiliary_encoder block must populate a matching AuxiliaryEncoderConfig.""" + path = _write_yaml( + tmp_path, + { + "auxiliary_encoder": { + "freeze_epochs": 3, + "embedding_dim": 128, + "checkpoint_path": "aux_encoder.pt", + } + }, + ) + + cfg = PipelineConfig.from_yaml(path) + + assert cfg.auxiliary_encoder == AuxiliaryEncoderConfig( + freeze_epochs=3, embedding_dim=128, checkpoint_path="aux_encoder.pt" + ) From d194beda703509abad9723555c3bd14e532cd180 Mon Sep 17 00:00:00 2001 From: Sean Colby Date: Thu, 23 Jul 2026 16:39:14 -0800 Subject: [PATCH 04/18] Extract shared MPNN construction from ChemPropLightningModule 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. --- moal/model.py | 171 +++++++++++++++++++++++++++++++------------- tests/test_model.py | 9 ++- 2 files changed, 125 insertions(+), 55 deletions(-) diff --git a/moal/model.py b/moal/model.py index dc7c732..223c555 100644 --- a/moal/model.py +++ b/moal/model.py @@ -70,6 +70,116 @@ def _validate_from_foundation(value: str | bool) -> None: ) +def load_foundation_weights(from_foundation: str | bool) -> dict: + """Load pretrained message-passing weights from a named model or local path. + + Shared by :class:`ChemPropLightningModule` and + :class:`~moal.auxiliary_encoder.AuxiliaryEncoderModule` so both draw + from the identical checkpoint-loading path. + + Parameters + ---------- + from_foundation : str or bool + ``"chemeleon"`` downloads (or reuses the cached copy of) the + CheMeleon checkpoint from Zenodo. Any other string is treated as a + local filesystem path. Must not be ``False``; validate with + :func:`_validate_from_foundation` first. + + Returns + ------- + dict + Checkpoint dictionary with ``hyper_parameters`` and ``state_dict`` + keys. + """ + if from_foundation == "chemeleon": + download_chemeleon() + ckpt_path = Path().home() / ".chemprop" / "chemeleon_mp.pt" + else: + ckpt_path = Path(str(from_foundation)) + logger.info("Loading foundation weights from local path: %s", ckpt_path) + return cast(dict[str, Any], torch.load(ckpt_path, weights_only=True)) + + +def build_mpnn( + from_foundation: str | bool, + ffn_hidden_dim: int, + ffn_num_layers: int, + message_hidden_dim: int, + depth: int, + n_tasks: int = 1, +) -> nn.Module: + """Construct a ChemProp MPNN, dispatching on ``from_foundation``. + + Shared by :class:`ChemPropLightningModule` and + :class:`~moal.auxiliary_encoder.AuxiliaryEncoderModule` so both models' + embeddings live in the same representation space rather than each + hand-rolling its own encoder construction. + + Parameters + ---------- + from_foundation : str or bool + ``False`` builds the message-passing encoder with random weights at + ``message_hidden_dim`` / ``depth``. Any other value loads foundation + weights via :func:`load_foundation_weights`, which also supplies the + encoder's architecture (``message_hidden_dim`` and ``depth`` are + ignored in that case). + ffn_hidden_dim : int + Hidden dimension of the FFN predictor head. + ffn_num_layers : int + Number of layers in the FFN predictor head. + message_hidden_dim : int + Message-passing hidden width (``d_h``) for the random-init encoder. + Ignored when a foundation checkpoint supplies the architecture. + depth : int + Number of message-passing steps for the random-init encoder. Ignored + when a foundation checkpoint supplies the architecture. + n_tasks : int, optional + Number of regression targets predicted per compound. Default is 1 + (the main model's single pEC50 target). The auxiliary encoder passes + one task per distinct auxiliary readout key it was trained on. + + Returns + ------- + nn.Module + Fully assembled ``chemprop.models.MPNN``. Aggregation is always + ``MeanAggregation``: CheMeleon's own pretraining used a mean readout, + so any foundation-weights branch is constrained to match it; the + random-init branch keeps the same readout for consistency between + the two initialisation paths rather than introducing an + undocumented behavioural difference. + + Notes + ----- + ``message_hidden_dim`` and ``depth`` apply only on the + ``from_foundation=False`` branch; for a foundation checkpoint the + encoder architecture is read from the checkpoint's stored + ``hyper_parameters`` so the pretrained weights load with ``strict=True``. + """ + if from_foundation is False: + logger.info( + "Building ChemProp encoder with random weights " + "(from_foundation=False, d_h=%d, depth=%d).", + message_hidden_dim, + depth, + ) + mp: nn.Module = BondMessagePassing( # pyright: ignore[reportAbstractUsage] + d_h=message_hidden_dim, depth=depth + ) + else: + foundation_weights = load_foundation_weights(from_foundation) + mp = BondMessagePassing(**foundation_weights["hyper_parameters"]) # pyright: ignore[reportAbstractUsage] + mp.load_state_dict(foundation_weights["state_dict"]) + + agg = MeanAggregation() + ffn = RegressionFFN( # pyright: ignore[reportAbstractUsage] + n_tasks=n_tasks, + input_dim=cast(BondMessagePassing, mp).output_dim, + hidden_dim=ffn_hidden_dim, + n_layers=ffn_num_layers, + ) + return cast(nn.Module, MPNN(message_passing=mp, agg=agg, predictor=ffn)) + + def download_chemeleon() -> None: """Download the CheMeleon checkpoint if not already cached locally. @@ -215,6 +325,9 @@ def _build_model( ) -> nn.Module: """Construct the MPNN, dispatching on ``self._from_foundation``. + Thin wrapper around the shared :func:`build_mpnn`; see that function + for the full construction contract. + Parameters ---------- ffn_hidden_dim : int @@ -231,58 +344,16 @@ def _build_model( Returns ------- nn.Module - Fully assembled ``chemprop.models.MPNN``. - - Notes - ----- - ``message_hidden_dim`` and ``depth`` apply only on the - ``from_foundation=False`` branch; for a foundation checkpoint the - encoder architecture is read from the checkpoint's stored - ``hyper_parameters`` so the pretrained weights load with ``strict=True``. + Fully assembled ``chemprop.models.MPNN`` with a single-task + (``n_tasks=1``) predictor head. """ - if self._from_foundation is False: - logger.info( - "Building ChemProp encoder with random weights " - "(from_foundation=False, d_h=%d, depth=%d).", - message_hidden_dim, - depth, - ) - mp: nn.Module = BondMessagePassing( # pyright: ignore[reportAbstractUsage] - d_h=message_hidden_dim, depth=depth - ) - else: - foundation_weights = self._load_foundation_weights() - mp = BondMessagePassing(**foundation_weights["hyper_parameters"]) # pyright: ignore[reportAbstractUsage] - mp.load_state_dict(foundation_weights["state_dict"]) - - agg = MeanAggregation() - ffn = RegressionFFN( # pyright: ignore[reportAbstractUsage] - input_dim=cast(BondMessagePassing, mp).output_dim, - hidden_dim=ffn_hidden_dim, - n_layers=ffn_num_layers, + return build_mpnn( + from_foundation=self._from_foundation, + ffn_hidden_dim=ffn_hidden_dim, + ffn_num_layers=ffn_num_layers, + message_hidden_dim=message_hidden_dim, + depth=depth, ) - return cast(nn.Module, MPNN(message_passing=mp, agg=agg, predictor=ffn)) - - def _load_foundation_weights(self) -> dict: - """Load pretrained message-passing weights from a named model or local path. - - When ``self._from_foundation == "chemeleon"`` the checkpoint is - downloaded from Zenodo if not already cached. For any other string - value it is treated as a local filesystem path. - - Returns - ------- - dict - Checkpoint dictionary with ``hyper_parameters`` and - ``state_dict`` keys. - """ - if self._from_foundation == "chemeleon": - download_chemeleon() - ckpt_path = Path().home() / ".chemprop" / "chemeleon_mp.pt" - else: - ckpt_path = Path(str(self._from_foundation)) - logger.info("Loading foundation weights from local path: %s", ckpt_path) - return cast(dict[str, Any], torch.load(ckpt_path, weights_only=True)) # ------------------------------------------------------------------ # Freeze / unfreeze schedule diff --git a/tests/test_model.py b/tests/test_model.py index a274c20..ea0010f 100644 --- a/tests/test_model.py +++ b/tests/test_model.py @@ -471,7 +471,7 @@ def test_false_encoder_passes_arch_to_bond_message_passing(self, monkeypatch): We temporarily restore the real _build_model so the False-branch dispatch runs, then verify BondMessagePassing is called once with the configured - d_h and depth and _load_foundation_weights is never invoked. + d_h and depth and load_foundation_weights is never invoked. """ calls = [] @@ -484,11 +484,10 @@ def tracking_bmp(*args, **kwargs): monkeypatch.setattr("moal.model.BondMessagePassing", tracking_bmp) monkeypatch.setattr(ChemPropLightningModule, "_build_model", _REAL_BUILD_MODEL) monkeypatch.setattr( - ChemPropLightningModule, - "_load_foundation_weights", - lambda self: (_ for _ in ()).throw( + "moal.model.load_foundation_weights", + lambda from_foundation: (_ for _ in ()).throw( AssertionError( - "_load_foundation_weights must not be called when from_foundation=False" + "load_foundation_weights must not be called when from_foundation=False" ) ), ) From 120082e705442e49c67c951d9d5a242624b68517 Mon Sep 17 00:00:00 2001 From: Sean Colby Date: Thu, 23 Jul 2026 16:43:45 -0800 Subject: [PATCH 05/18] Add auxiliary log2FC/pIC50 encoder (#36 Phase 1) 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. --- moal/auxiliary_encoder.py | 543 ++++++++++++++++++++++++++++++++ moal/config.py | 31 ++ tests/test_auxiliary_encoder.py | 215 +++++++++++++ 3 files changed, 789 insertions(+) create mode 100644 moal/auxiliary_encoder.py create mode 100644 tests/test_auxiliary_encoder.py diff --git a/moal/auxiliary_encoder.py b/moal/auxiliary_encoder.py new file mode 100644 index 0000000..bd3e679 --- /dev/null +++ b/moal/auxiliary_encoder.py @@ -0,0 +1,543 @@ +"""Auxiliary encoder pretrained on primary-screen readouts (log2FC, pIC50, etc.). + +Phase 1 of issue #36 (``moal plan``-only; excluded from ``moal simulate`` to +avoid an acquisition-endogeneity problem in the live active-learning loop). +Trains a small ChemProp encoder via masked multi-task regression over +``LabelRecord.raw_ps_readouts``, sharing the main model's backbone +construction (:func:`moal.model.build_mpnn`) rather than a bespoke +architecture, so its embeddings live in the same representation space as the +main pEC50 model. Readouts are used as-is: no per-plate/per-batch +normalization is applied (see :class:`~moal.config.AuxiliaryEncoderConfig` +for why). +""" + +from __future__ import annotations + +import logging +from pathlib import Path +from typing import Any, cast + +import lightning as L +import torch +import torch.nn as nn +from chemprop.data import BatchMolGraph, MoleculeDatapoint, MoleculeDataset +from torch import Tensor +from torch.optim import Adam +from torch.utils.data import DataLoader, Dataset, random_split + +from moal.config import AuxiliaryEncoderConfig +from moal.model import build_mpnn +from moal.types import LabelRecord + +logger = logging.getLogger(__name__) + + +def masked_mse_loss(preds: Tensor, targets: Tensor, mask: Tensor) -> Tensor: + """Mean squared error computed over only the masked (observed) task entries. + + Parameters + ---------- + preds : Tensor + Shape ``(batch, n_tasks)`` model predictions. + targets : Tensor + Shape ``(batch, n_tasks)`` targets. Values at positions where + ``mask`` is False are ignored and may hold arbitrary placeholder + values. + mask : Tensor + Boolean tensor of shape ``(batch, n_tasks)``; True where a compound + had an observed readout for that task. + + Returns + ------- + Tensor + Scalar masked MSE, differentiable with respect to ``preds``. When + ``mask`` has no True entries (e.g. a batch with no observed readouts + for any task), returns a zero-valued tensor still connected to + ``preds`` so the training step remains well-defined rather than + raising a division-by-zero. + """ + mask_f = mask.to(preds.dtype) + denom = mask_f.sum() + if denom == 0: + return preds.sum() * 0.0 + return ((preds - targets) ** 2 * mask_f).sum() / denom + + +class _AuxiliaryDataset(Dataset): + """Dataset pairing a molecular graph with a masked multi-task target vector. + + Parameters + ---------- + records : list[LabelRecord] + Records with a non-empty ``raw_ps_readouts``. Records lacking any + readout carry no training signal and should be filtered out by the + caller before construction. + task_names : list[str] + Fixed, ordered list of readout keys; determines target/mask column + order and the auxiliary encoder's output dimensionality. + """ + + def __init__(self, records: list[LabelRecord], task_names: list[str]) -> None: + self.records = records + self.task_names = task_names + self._mol_graphs = MoleculeDataset( + [MoleculeDatapoint.from_smi(r.canonical_smiles) for r in records] # pyright: ignore[reportArgumentType] + ) + self._targets = torch.zeros(len(records), len(task_names), dtype=torch.float32) + self._mask = torch.zeros(len(records), len(task_names), dtype=torch.bool) + for i, rec in enumerate(records): + for j, name in enumerate(task_names): + if name in rec.raw_ps_readouts: + self._targets[i, j] = rec.raw_ps_readouts[name] + self._mask[i, j] = True + + def __len__(self) -> int: + """Return the number of records in the dataset. + + Returns + ------- + int + Total number of readout-bearing records. + """ + return len(self.records) + + def __getitem__(self, idx: int) -> tuple[Any, Tensor, Tensor]: + """Return the (datapoint, target row, mask row) triple at ``idx``. + + Parameters + ---------- + idx : int + Zero-based index into the dataset. + + Returns + ------- + tuple[Any, Tensor, Tensor] + A ``(MoleculeDatapoint, targets, mask)`` triple; ``targets`` and + ``mask`` each have shape ``(n_tasks,)``. + """ + return self._mol_graphs[idx], self._targets[idx], self._mask[idx] + + @staticmethod + def collate_fn(batch: list[tuple[Any, Tensor, Tensor]]) -> tuple[Any, Tensor, Tensor]: + """Collate a list of (datapoint, target row, mask row) into a batch. + + Parameters + ---------- + batch : list[tuple[Any, Tensor, Tensor]] + Items as returned by :meth:`__getitem__`. + + Returns + ------- + tuple[BatchMolGraph, Tensor, Tensor] + Batched molecular graph, stacked targets ``(batch, n_tasks)``, + and stacked mask ``(batch, n_tasks)``. + """ + datapoints, targets, masks = zip(*batch, strict=False) + bmg = BatchMolGraph([dp.mg for dp in datapoints]) + return bmg, torch.stack(list(targets)), torch.stack(list(masks)) + + +class AuxiliaryDataModule(L.LightningDataModule): + """LightningDataModule for masked multi-task auxiliary-encoder pretraining. + + Parameters + ---------- + records : list[LabelRecord] + Records with a non-empty ``raw_ps_readouts``. + task_names : list[str] + Fixed, ordered list of readout keys. + batch_size : int, optional + Number of samples per mini-batch. Default is 64. + val_fraction : float, optional + Fraction of records held out for validation. Default is 0.1. + num_workers : int, optional + DataLoader worker count (0 = main process). Default is 0. + seed : int, optional + Random seed for the train/val split. Default is 42. + """ + + def __init__( + self, + records: list[LabelRecord], + task_names: list[str], + batch_size: int = 64, + val_fraction: float = 0.1, + num_workers: int = 0, + seed: int = 42, + ) -> None: + super().__init__() + self.records = records + self.task_names = task_names + self.batch_size = batch_size + self.val_fraction = val_fraction + self.num_workers = num_workers + self.seed = seed + + self._train_dataset: Dataset | None = None + self._val_dataset: Dataset | None = None + + def setup(self, stage: str | None = None) -> None: + """Create the train and validation dataset splits. + + Parameters + ---------- + stage : str or None, optional + Lightning stage identifier; unused, accepted for interface + compatibility. + """ + n_val = max(1, int(len(self.records) * self.val_fraction)) + n_train = len(self.records) - n_val + if n_train <= 0: + logger.warning( + "Too few readout-bearing records (%d) for a val split; using all for training.", + len(self.records), + ) + n_train, n_val = len(self.records), 0 + + full = _AuxiliaryDataset(self.records, self.task_names) + if n_val > 0: + self._train_dataset, self._val_dataset = random_split( + full, + [n_train, n_val], + generator=torch.Generator().manual_seed(self.seed), + ) + else: + self._train_dataset = full + self._val_dataset = None + + def transfer_batch_to_device( + self, batch: tuple[Any, Tensor, Tensor], device: torch.device, dataloader_idx: int + ) -> tuple[Any, Tensor, Tensor]: + """Move the batched mol graph and target/mask tensors to ``device``. + + Parameters + ---------- + batch : tuple[Any, Tensor, Tensor] + A ``(BatchMolGraph, targets, mask)`` triple. + device : torch.device + Target device. + dataloader_idx : int + Index of the dataloader (required by the Lightning interface). + + Returns + ------- + tuple[Any, Tensor, Tensor] + The same triple moved to ``device``. + """ + mol_graph, targets, mask = batch + mol_graph = super().transfer_batch_to_device(mol_graph, device, dataloader_idx) + return mol_graph, targets.to(device), mask.to(device) + + def train_dataloader(self) -> DataLoader: + """Return the training DataLoader. + + Returns + ------- + DataLoader + Shuffled DataLoader over the training split. + """ + if self._train_dataset is None: + raise RuntimeError("setup() must be called before train_dataloader()") + return DataLoader( + self._train_dataset, + batch_size=self.batch_size, + shuffle=True, + collate_fn=_AuxiliaryDataset.collate_fn, + num_workers=self.num_workers, + persistent_workers=self.num_workers > 0, + drop_last=False, + ) + + def val_dataloader(self) -> DataLoader | None: + """Return the validation DataLoader, or ``None`` when no val split exists. + + Returns + ------- + DataLoader or None + Non-shuffled DataLoader over the validation split, or ``None``. + """ + if self._val_dataset is None: + return None + return DataLoader( + self._val_dataset, + batch_size=self.batch_size, + shuffle=False, + collate_fn=_AuxiliaryDataset.collate_fn, + num_workers=self.num_workers, + persistent_workers=self.num_workers > 0, + ) + + +class AuxiliaryEncoderModule(L.LightningModule): + """ChemProp MPNN trained via masked multi-task regression on auxiliary readouts. + + Parameters + ---------- + task_names : list[str] + Fixed, ordered list of readout keys the model was (or will be) + trained against; determines the predictor head's output width. + config : AuxiliaryEncoderConfig + Backbone architecture, freeze schedule, and optimization + hyperparameters. + """ + + def __init__(self, task_names: list[str], config: AuxiliaryEncoderConfig) -> None: + super().__init__() + if not task_names: + raise ValueError("task_names must be non-empty.") + self.task_names = list(task_names) + self._config = config + self._encoder_frozen = True + self.model = build_mpnn( + from_foundation=config.from_foundation, + ffn_hidden_dim=config.ffn_hidden_dim, + ffn_num_layers=config.ffn_num_layers, + message_hidden_dim=config.message_hidden_dim, + depth=config.depth, + n_tasks=len(self.task_names), + ) + self._freeze_encoder() + + # ------------------------------------------------------------------ + # Freeze / unfreeze schedule + # ------------------------------------------------------------------ + + def _encoder_params(self) -> list[nn.Parameter]: + """Return the trainable parameters of the message-passing encoder. + + Returns + ------- + list[nn.Parameter] + Parameters belonging to ``self.model.message_passing``. + """ + return list(cast(nn.Module, self.model.message_passing).parameters()) + + def _head_params(self) -> list[nn.Parameter]: + """Return the trainable parameters of the aggregation layer and FFN head. + + Returns + ------- + list[nn.Parameter] + Parameters belonging to ``self.model.agg`` and + ``self.model.predictor``, concatenated in that order. + """ + return list(cast(nn.Module, self.model.agg).parameters()) + list( + cast(nn.Module, self.model.predictor).parameters() + ) + + def _freeze_encoder(self) -> None: + """Freeze all message-passing encoder parameters.""" + for p in self._encoder_params(): + p.requires_grad_(False) + self._encoder_frozen = True + + def _unfreeze_encoder(self) -> None: + """Unfreeze the message-passing encoder after the warm-up phase.""" + for p in self._encoder_params(): + p.requires_grad_(True) + self._encoder_frozen = False + + def on_train_epoch_start(self) -> None: + """Lightning hook: unfreeze the encoder once warm-up is complete.""" + if self._encoder_frozen and self.current_epoch >= self._config.freeze_epochs: + self._unfreeze_encoder() + self.trainer.strategy.setup_optimizers(self.trainer) + + # ------------------------------------------------------------------ + # Lightning interface + # ------------------------------------------------------------------ + + def forward(self, batch_mol_graph: Any) -> Tensor: + """Run a forward pass and return multi-task predictions. + + Parameters + ---------- + batch_mol_graph : Any + A batched molecular graph (``chemprop.data.BatchMolGraph``). + + Returns + ------- + Tensor + Shape ``(batch, n_tasks)`` predictions. + """ + return cast(Tensor, self.model(batch_mol_graph)) + + def training_step(self, batch: tuple[Any, Tensor, Tensor], batch_idx: int) -> Tensor: + """Compute and log the masked multi-task training loss for one batch. + + Parameters + ---------- + batch : tuple[Any, Tensor, Tensor] + A ``(mol_graph, targets, mask)`` triple. + batch_idx : int + Index of the batch within the current epoch (unused). + + Returns + ------- + Tensor + Scalar training loss used for the backward pass. + """ + mol_graph, targets, mask = batch + preds = self(mol_graph) + loss = masked_mse_loss(preds, targets, mask) + self.log("aux_train_loss", loss, prog_bar=True, batch_size=targets.shape[0]) + return loss + + def validation_step(self, batch: tuple[Any, Tensor, Tensor], batch_idx: int) -> None: + """Compute and log the masked multi-task validation loss for one batch. + + Parameters + ---------- + batch : tuple[Any, Tensor, Tensor] + A ``(mol_graph, targets, mask)`` triple. + batch_idx : int + Index of the batch within the current validation epoch (unused). + """ + mol_graph, targets, mask = batch + preds = self(mol_graph) + loss = masked_mse_loss(preds, targets, mask) + self.log("aux_val_loss", loss, prog_bar=True, batch_size=targets.shape[0]) + + def configure_optimizers(self) -> Adam: + """Build and return the Adam optimizer for the current freeze state. + + Returns + ------- + Adam + When the encoder is frozen, a single-group Adam optimizer for the + multi-task head at ``config.lr``. After the encoder is unfrozen, + a second param group for the encoder is added at the same + ``config.lr`` (no discriminative rate split, unlike the main + model's ``mpnn_lr`` / ``ffn_lr``). + """ + param_groups = [ + { + "params": self._head_params(), + "lr": self._config.lr, + "weight_decay": self._config.weight_decay, + } + ] + if not self._encoder_frozen: + param_groups.append( + { + "params": self._encoder_params(), + "lr": self._config.lr, + "weight_decay": self._config.weight_decay, + } + ) + return Adam(param_groups) + + +def pretrain_auxiliary_encoder( + records: list[LabelRecord], + config: AuxiliaryEncoderConfig, + trainer_kwargs: dict[str, Any] | None = None, + datamodule_kwargs: dict[str, Any] | None = None, +) -> AuxiliaryEncoderModule: + """Pretrain (or load) the auxiliary encoder from a campaign's labeled records. + + When ``config.checkpoint_path`` is set, pretraining is skipped entirely + and the checkpoint is loaded instead — the explicit opt-in override for + cases where retraining every ``moal plan`` invocation is too expensive. + Otherwise the encoder is trained from scratch on every call, using + whichever ``raw_ps_readouts`` exist in ``records`` at that moment. + + Parameters + ---------- + records : list[LabelRecord] + All labeled records from the campaign state. Records with an empty + ``raw_ps_readouts`` are filtered out before training; they carry no + auxiliary training signal. + config : AuxiliaryEncoderConfig + Backbone, freeze schedule, and optimization hyperparameters. + trainer_kwargs : dict[str, Any], optional + Additional keyword arguments forwarded to ``lightning.Trainer``. + Ignored when loading from ``config.checkpoint_path``. + datamodule_kwargs : dict[str, Any], optional + Passed to :class:`AuxiliaryDataModule` (e.g. ``val_fraction``, + ``seed``). Ignored when loading from ``config.checkpoint_path``. + + Returns + ------- + AuxiliaryEncoderModule + The trained (or loaded) auxiliary encoder. + + Raises + ------ + ValueError + If no record in ``records`` carries any auxiliary readout and + ``config.checkpoint_path`` is unset. + """ + if config.checkpoint_path is not None: + return load_auxiliary_encoder_checkpoint(config.checkpoint_path, config) + + readout_records = [rec for rec in records if rec.raw_ps_readouts] + if not readout_records: + raise ValueError( + "No records carry raw_ps_readouts; cannot pretrain the auxiliary encoder. " + "Set config.checkpoint_path to load a cached checkpoint instead." + ) + + task_names = sorted({key for rec in readout_records for key in rec.raw_ps_readouts}) + logger.info( + "Pretraining auxiliary encoder on %d readout-bearing record(s), tasks=%s", + len(readout_records), + task_names, + ) + + module = AuxiliaryEncoderModule(task_names=task_names, config=config) + dm = AuxiliaryDataModule(readout_records, task_names, **(datamodule_kwargs or {})) + dm.setup() + + kwargs: dict[str, Any] = { + "max_epochs": config.max_epochs, + "enable_progress_bar": False, + "enable_model_summary": False, + } + if trainer_kwargs: + kwargs.update(trainer_kwargs) + kwargs.setdefault("logger", False) + kwargs.setdefault("enable_checkpointing", False) + trainer = L.Trainer(**kwargs) + trainer.fit(module, datamodule=dm) + return module + + +def save_auxiliary_encoder_checkpoint(module: AuxiliaryEncoderModule, path: str | Path) -> None: + """Save an auxiliary encoder to a checkpoint usable by ``config.checkpoint_path``. + + Parameters + ---------- + module : AuxiliaryEncoderModule + A trained auxiliary encoder. + path : str or Path + Destination file path. + """ + torch.save({"task_names": module.task_names, "state_dict": module.state_dict()}, path) + + +def load_auxiliary_encoder_checkpoint( + path: str | Path, config: AuxiliaryEncoderConfig +) -> AuxiliaryEncoderModule: + """Load an auxiliary encoder checkpoint written by :func:`save_auxiliary_encoder_checkpoint`. + + Parameters + ---------- + path : str or Path + Path to the checkpoint file. + config : AuxiliaryEncoderConfig + Backbone architecture the checkpoint was trained with; must match + the checkpoint's own architecture (``from_foundation``, + ``ffn_hidden_dim``, etc.) or ``load_state_dict`` will raise. + + Returns + ------- + AuxiliaryEncoderModule + The restored auxiliary encoder, with the encoder still frozen + (caller-visible state, not resumed training state). + """ + logger.info("Loading cached auxiliary encoder checkpoint from %s (retraining skipped).", path) + ckpt = cast(dict[str, Any], torch.load(path, weights_only=True)) + module = AuxiliaryEncoderModule(task_names=ckpt["task_names"], config=config) + module.load_state_dict(ckpt["state_dict"]) + return module diff --git a/moal/config.py b/moal/config.py index 43fba49..317d739 100644 --- a/moal/config.py +++ b/moal/config.py @@ -151,10 +151,33 @@ class AuxiliaryEncoderConfig: Attributes ---------- + from_foundation : str or bool + Encoder initialisation, forwarded to :func:`moal.model.build_mpnn`. + Same semantics as ``ModelConfig.from_foundation``. Default + ``"chemeleon"`` shares the main model's foundation checkpoint. + ffn_hidden_dim : int + Hidden dimension of the multi-task FFN predictor head. + ffn_num_layers : int + Number of layers in the multi-task FFN predictor head. + message_hidden_dim : int + Message-passing hidden width (``d_h``) for the random-init encoder. + Used only when ``from_foundation=False``. + depth : int + Number of message-passing steps for the random-init encoder. Used + only when ``from_foundation=False``. freeze_epochs : int Number of warm-up epochs to train only the multi-task FFN head, analogous to ``ModelConfig.freeze_epochs`` but scheduled independently for the auxiliary encoder. + lr : float + Learning rate for the multi-task FFN head, and for the message-passing + encoder after unfreezing (no separate discriminative rate, unlike the + main model's ``mpnn_lr`` / ``ffn_lr`` split). + weight_decay : float + L2 weight decay applied to all trainable parameters. + max_epochs : int + Number of pretraining epochs, scheduled independently from the main + model's ``TrainerConfig.max_epochs``. embedding_dim : int Dimensionality of the pooled molecular embedding exposed to the main model's concatenation architecture (Phase 2). Ignored by the @@ -169,7 +192,15 @@ class AuxiliaryEncoderConfig: next run automatically. """ + from_foundation: str | bool = "chemeleon" + ffn_hidden_dim: int = 300 + ffn_num_layers: int = 2 + message_hidden_dim: int = 300 + depth: int = 3 freeze_epochs: int = 5 + lr: float = 1e-4 + weight_decay: float = 0.0 + max_epochs: int = 30 embedding_dim: int = 300 checkpoint_path: str | None = None diff --git a/tests/test_auxiliary_encoder.py b/tests/test_auxiliary_encoder.py new file mode 100644 index 0000000..b966993 --- /dev/null +++ b/tests/test_auxiliary_encoder.py @@ -0,0 +1,215 @@ +"""Tests for the auxiliary log2FC/pIC50 encoder (issue #36 Phase 1). + +All tests use ``from_foundation=False`` (random-init ChemProp encoder) so +no network download or cached CheMeleon checkpoint is required. +""" + +from __future__ import annotations + +import pytest +import torch +from chemprop.data import BatchMolGraph, MoleculeDatapoint, MoleculeDataset + +from moal.auxiliary_encoder import ( + AuxiliaryDataModule, + AuxiliaryEncoderModule, + load_auxiliary_encoder_checkpoint, + masked_mse_loss, + pretrain_auxiliary_encoder, + save_auxiliary_encoder_checkpoint, +) +from moal.config import AuxiliaryEncoderConfig +from moal.types import CensoringType, LabelRecord, QueryType + +_SMILES = ["CCO", "CCN", "CCC", "c1ccccc1", "CCCl", "CCBr", "CCOCC", "CCCC"] + + +def _records_with_readouts() -> list[LabelRecord]: + """Build a small set of LabelRecords with mixed, partially-overlapping readouts.""" + records = [] + for i, smi in enumerate(_SMILES): + readouts = {"log2fc_1um": float(i) - 3.0} + if i % 2 == 0: + readouts["pic50"] = 6.0 + i * 0.1 + records.append( + LabelRecord( + smiles=smi, + canonical_smiles=smi, + value=5.0, + upper_bound=11.0, + censoring_type=CensoringType.LEFT, + fidelity=QueryType.PRIMARY_SCREEN, + cost=1.0, + iteration=0, + raw_ps_readouts=readouts, + ) + ) + return records + + +def _batch(smiles_list: list[str]) -> BatchMolGraph: + """Build a BatchMolGraph for a list of SMILES, mirroring moal.dataset's featurization path.""" + dataset = MoleculeDataset([MoleculeDatapoint.from_smi(s) for s in smiles_list]) + return BatchMolGraph([dataset[i].mg for i in range(len(dataset))]) + + +def _fast_config(**overrides) -> AuxiliaryEncoderConfig: + defaults = { + "from_foundation": False, + "message_hidden_dim": 16, + "ffn_hidden_dim": 16, + "depth": 1, + "freeze_epochs": 0, + "max_epochs": 1, + } + defaults.update(overrides) + return AuxiliaryEncoderConfig(**defaults) + + +class TestMaskedMSELoss: + """Tests for masked_mse_loss: gradient flow and zero-mask handling.""" + + def test_masked_entries_contribute_zero_gradient(self): + """A task masked out for every sample in the batch must receive zero gradient on that task's predictions.""" + preds = torch.tensor([[1.0, 5.0], [2.0, 5.0]], requires_grad=True) + targets = torch.tensor([[0.0, 999.0], [0.0, 999.0]]) + mask = torch.tensor([[True, False], [True, False]]) + + loss = masked_mse_loss(preds, targets, mask) + loss.backward() + + assert preds.grad is not None + assert torch.all(preds.grad[:, 1] == 0.0) + assert torch.any(preds.grad[:, 0] != 0.0) + + def test_fully_masked_batch_returns_zero_without_raising(self): + """A batch with no observed targets at all must return a differentiable zero loss, not raise or NaN.""" + preds = torch.zeros(3, 2, requires_grad=True) + targets = torch.zeros(3, 2) + mask = torch.zeros(3, 2, dtype=torch.bool) + + loss = masked_mse_loss(preds, targets, mask) + + assert loss.item() == 0.0 + loss.backward() + assert preds.grad is not None + + +class TestAuxiliaryEncoderModule: + """Tests for AuxiliaryEncoderModule construction and freeze/unfreeze schedule.""" + + def test_freeze_epochs_zero_starts_unfrozen_after_epoch_start(self): + """With freeze_epochs=0, the encoder must unfreeze at the very first epoch boundary.""" + config = _fast_config(freeze_epochs=0) + module = AuxiliaryEncoderModule(task_names=["log2fc_1um"], config=config) + assert module._encoder_frozen is True + + def test_output_width_matches_task_count(self): + """A forward pass's prediction width must equal len(task_names).""" + config = _fast_config() + module = AuxiliaryEncoderModule(task_names=["log2fc_1um", "pic50"], config=config) + + preds = module(_batch(["CCO", "CCN"])) + + assert preds.shape == (2, 2) + + def test_empty_task_names_raises(self): + """Constructing with an empty task_names list must raise ValueError.""" + with pytest.raises(ValueError, match="task_names"): + AuxiliaryEncoderModule(task_names=[], config=_fast_config()) + + +class TestPretrainAuxiliaryEncoder: + """Tests for pretrain_auxiliary_encoder: training end-to-end and checkpoint opt-in.""" + + def test_trains_and_returns_module_with_expected_tasks(self): + """Pretraining on mixed-readout records must produce a module whose task_names is the sorted union of observed keys.""" + records = _records_with_readouts() + config = _fast_config() + + module = pretrain_auxiliary_encoder(records, config) + + assert module.task_names == ["log2fc_1um", "pic50"] + + def test_records_without_any_readout_raises(self): + """Pretraining with no readout-bearing records and no checkpoint_path must raise ValueError.""" + bare_record = LabelRecord( + smiles="CCO", + canonical_smiles="CCO", + value=5.0, + upper_bound=11.0, + censoring_type=CensoringType.EXACT, + fidelity=QueryType.DOSE_RESPONSE, + cost=10.0, + iteration=0, + ) + with pytest.raises(ValueError, match="raw_ps_readouts"): + pretrain_auxiliary_encoder([bare_record], _fast_config()) + + def test_checkpoint_path_skips_retraining(self, tmp_path, monkeypatch): + """When checkpoint_path is set, pretrain_auxiliary_encoder must load the checkpoint rather than training.""" + records = _records_with_readouts() + config = _fast_config() + trained = pretrain_auxiliary_encoder(records, config) + ckpt_path = tmp_path / "aux_encoder.pt" + save_auxiliary_encoder_checkpoint(trained, ckpt_path) + + def _fail_if_called(*args, **kwargs): + raise AssertionError("Trainer.fit must not be called when checkpoint_path is set") + + monkeypatch.setattr("lightning.Trainer.fit", _fail_if_called) + + loaded_config = _fast_config(checkpoint_path=str(ckpt_path)) + loaded = pretrain_auxiliary_encoder(records, loaded_config) + + assert loaded.task_names == trained.task_names + + +class TestAuxiliaryEncoderCheckpoint: + """Tests for save/load round-tripping.""" + + def test_round_trips_weights_and_task_names(self, tmp_path): + """A saved-then-loaded checkpoint must reproduce identical predictions and task_names.""" + records = _records_with_readouts() + config = _fast_config() + trained = pretrain_auxiliary_encoder(records, config) + trained.eval() + + path = tmp_path / "aux_encoder.pt" + save_auxiliary_encoder_checkpoint(trained, path) + loaded = load_auxiliary_encoder_checkpoint(path, config) + loaded.eval() + + bmg = _batch(["CCO", "CCN"]) + with torch.no_grad(): + preds_trained = trained(bmg) + preds_loaded = loaded(bmg) + + assert loaded.task_names == trained.task_names + assert torch.allclose(preds_trained, preds_loaded) + + +class TestAuxiliaryDataModule: + """Tests for AuxiliaryDataModule train/val splitting and dataloader batch shape.""" + + def test_train_batch_shapes_match_task_count(self): + """A training batch's targets/mask must have shape (batch, n_tasks).""" + records = _records_with_readouts() + task_names = ["log2fc_1um", "pic50"] + dm = AuxiliaryDataModule(records, task_names, batch_size=4, val_fraction=0.25, seed=1) + dm.setup() + + batch = next(iter(dm.train_dataloader())) + _, targets, mask = batch + + assert targets.shape[1] == 2 + assert mask.shape[1] == 2 + assert mask.dtype == torch.bool + + def test_too_few_records_uses_all_for_training(self): + """When the record pool is too small for a val split, val_dataloader must return None.""" + records = _records_with_readouts()[:1] + dm = AuxiliaryDataModule(records, ["log2fc_1um"], val_fraction=0.1) + dm.setup() + + assert dm.val_dataloader() is None From 4c4f25762f32bcadb2ea00698a401ce3a7f660e5 Mon Sep 17 00:00:00 2001 From: Sean Colby Date: Thu, 23 Jul 2026 16:53:45 -0800 Subject: [PATCH 06/18] Add AuxiliaryEncoderModule.embed_smiles for Phase 2 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. --- moal/auxiliary_encoder.py | 49 +++++++++++++++++++++++++++++++++ tests/test_auxiliary_encoder.py | 9 ++++++ 2 files changed, 58 insertions(+) diff --git a/moal/auxiliary_encoder.py b/moal/auxiliary_encoder.py index bd3e679..840eb6b 100644 --- a/moal/auxiliary_encoder.py +++ b/moal/auxiliary_encoder.py @@ -18,9 +18,12 @@ from typing import Any, cast import lightning as L +import numpy as np import torch import torch.nn as nn from chemprop.data import BatchMolGraph, MoleculeDatapoint, MoleculeDataset +from chemprop.data.dataloader import build_dataloader +from chemprop.models import MPNN from torch import Tensor from torch.optim import Adam from torch.utils.data import DataLoader, Dataset, random_split @@ -427,6 +430,52 @@ def configure_optimizers(self) -> Adam: ) return Adam(param_groups) + # ------------------------------------------------------------------ + # Inference helpers + # ------------------------------------------------------------------ + + @torch.no_grad() + def embed_smiles(self, smiles_list: list[str], batch_size: int = 256) -> np.ndarray: + """Return pooled structural embeddings (pre-predictor) for a list of SMILES. + + Used by the concatenation architecture (Phase 2) to supply a + structural fallback for compounds with no observed auxiliary + readout. Uses ``chemprop.models.MPNN.fingerprint``, which applies + message-passing, mean pooling, and batch-norm but stops short of the + multi-task predictor head. + + Parameters + ---------- + smiles_list : list[str] + **Must be RDKit-canonical, salt-stripped SMILES**, matching + :meth:`moal.model.ChemPropLightningModule.predict_smiles`'s + contract. + batch_size : int, optional + Number of molecules processed per forward pass. Default is 256. + + Returns + ------- + np.ndarray + Array of shape ``(N, embedding_dim)``, aligned with + ``smiles_list``. ``embedding_dim`` is the backbone's native + output width (CheMeleon's fixed width, or ``message_hidden_dim`` + for a random-init encoder), not + ``AuxiliaryEncoderConfig.embedding_dim``. + """ + dataset = MoleculeDataset([MoleculeDatapoint.from_smi(s) for s in smiles_list]) # pyright: ignore[reportArgumentType] + dataloader = build_dataloader( + dataset, batch_size=batch_size, shuffle=False, drop_last=False + ) + + all_embeddings = [] + with torch.inference_mode(): + for batch in dataloader: + batch.bmg.to(self.device) + embedding = cast(MPNN, self.model).fingerprint(batch.bmg) + all_embeddings.append(embedding.cpu().numpy()) + + return np.concatenate(all_embeddings, axis=0).astype(np.float32) + def pretrain_auxiliary_encoder( records: list[LabelRecord], diff --git a/tests/test_auxiliary_encoder.py b/tests/test_auxiliary_encoder.py index b966993..585ce5f 100644 --- a/tests/test_auxiliary_encoder.py +++ b/tests/test_auxiliary_encoder.py @@ -118,6 +118,15 @@ def test_empty_task_names_raises(self): with pytest.raises(ValueError, match="task_names"): AuxiliaryEncoderModule(task_names=[], config=_fast_config()) + def test_embed_smiles_returns_backbone_width_aligned_with_input(self): + """embed_smiles must return one embedding row per input SMILES, at the backbone's native width.""" + config = _fast_config(message_hidden_dim=24) + module = AuxiliaryEncoderModule(task_names=["log2fc_1um"], config=config) + + embeddings = module.embed_smiles(["CCO", "CCN", "CCC"]) + + assert embeddings.shape == (3, 24) + class TestPretrainAuxiliaryEncoder: """Tests for pretrain_auxiliary_encoder: training end-to-end and checkpoint opt-in.""" From 37cbae9a1f17ec14291e834c14fb8810e40e9596 Mon Sep 17 00:00:00 2001 From: Sean Colby Date: Thu, 23 Jul 2026 16:59:16 -0800 Subject: [PATCH 07/18] Add concatenation architecture for auxiliary readouts (#36 Phase 2) 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. --- moal/concatenation_model.py | 710 ++++++++++++++++++++++++++++++ moal/model.py | 11 +- tests/test_concatenation_model.py | 169 +++++++ 3 files changed, 889 insertions(+), 1 deletion(-) create mode 100644 moal/concatenation_model.py create mode 100644 tests/test_concatenation_model.py diff --git a/moal/concatenation_model.py b/moal/concatenation_model.py new file mode 100644 index 0000000..9240c15 --- /dev/null +++ b/moal/concatenation_model.py @@ -0,0 +1,710 @@ +"""Concatenation architecture for the auxiliary log2FC/pIC50 signal (#36 Phase 2). + +Concatenates, per compound, either its own observed auxiliary readouts (when +PS-screened) or the pretrained :class:`~moal.auxiliary_encoder.AuxiliaryEncoderModule`'s +structural embedding (when never PS-screened), plus a provenance flag +distinguishing the two, onto the pooled graph embedding before the pEC50 +predictor head. Graph-only prediction is the unconditional fallback: a +compound with neither an observed readout nor (obviously) a missing +embedding never occurs, since the embedding path always has a fallback +value. + +Reuses chemprop's native ``MPNN.forward(bmg, X_d=...)`` concatenation point +(see :func:`moal.model.build_mpnn`'s ``extra_input_dim``) rather than a +bespoke predictor wrapper, and the same ``CensoredRegressionLoss``, +freeze/unfreeze schedule, and refit contract as +:class:`moal.model.ChemPropLightningModule`, so this is a second, +coexisting model path selectable per run rather than a replacement. + +``moal plan``-only, matching :mod:`moal.auxiliary_encoder`. +""" + +from __future__ import annotations + +import logging +from pathlib import Path +from typing import Any, cast + +import lightning as L +import numpy as np +import torch +import torch.nn as nn +from chemprop.data import BatchMolGraph, MoleculeDatapoint, MoleculeDataset +from torch import Tensor +from torch.optim import Adam +from torch.utils.data import DataLoader, Dataset, random_split + +from moal.auxiliary_encoder import AuxiliaryEncoderModule +from moal.loss import CensoredRegressionLoss +from moal.model import _validate_from_foundation, build_mpnn +from moal.planning import normalize_record_weights +from moal.types import LabelRecord + +logger = logging.getLogger(__name__) + + +def concatenation_feature_dim(n_tasks: int, embedding_dim: int) -> int: + """Return the width of the concatenation feature vector. + + Parameters + ---------- + n_tasks : int + Number of distinct auxiliary readout keys (``AuxiliaryEncoderModule.task_names``). + embedding_dim : int + Width of the auxiliary encoder's structural embedding (its + backbone's native output width; see + :meth:`~moal.auxiliary_encoder.AuxiliaryEncoderModule.embed_smiles`). + + Returns + ------- + int + ``2 * n_tasks + embedding_dim + 1``: observed-readout vector, + readout mask, structural embedding, and a single provenance flag. + """ + return 2 * n_tasks + embedding_dim + 1 + + +def build_concatenation_features( + canonical_smiles: list[str], + readouts: list[dict[str, float]], + aux_encoder: AuxiliaryEncoderModule, + batch_size: int = 256, +) -> np.ndarray: + """Build the per-compound concatenation feature matrix. + + For each compound: if ``readouts[i]`` is non-empty, the observed-readout + block is populated (per-task values where present, zero elsewhere) and + the mask block marks which tasks were actually observed; the embedding + block stays zero and the provenance flag is 0. If ``readouts[i]`` is + empty, the observed-readout and mask blocks stay zero, the embedding + block holds the auxiliary encoder's structural embedding for that + compound, and the provenance flag is 1. + + Parameters + ---------- + canonical_smiles : list[str] + RDKit-canonical SMILES, one per compound. + readouts : list[dict[str, float]] + Per-compound ``LabelRecord.raw_ps_readouts``-shaped dict, aligned + with ``canonical_smiles``. An empty dict means "never PS-screened". + aux_encoder : AuxiliaryEncoderModule + Pretrained auxiliary encoder; supplies both ``task_names`` (readout + key order) and the structural embedding fallback. + batch_size : int, optional + Batch size for the embedding forward pass over compounds lacking + readouts. Default is 256. + + Returns + ------- + np.ndarray + Array of shape ``(N, concatenation_feature_dim(...))``, aligned with + ``canonical_smiles``. + + Raises + ------ + ValueError + If ``len(canonical_smiles) != len(readouts)``. + """ + if len(canonical_smiles) != len(readouts): + raise ValueError( + f"canonical_smiles ({len(canonical_smiles)}) and readouts ({len(readouts)}) " + "must be the same length." + ) + + task_names = aux_encoder.task_names + n_tasks = len(task_names) + n = len(canonical_smiles) + + readout_vec = np.zeros((n, n_tasks), dtype=np.float32) + readout_mask = np.zeros((n, n_tasks), dtype=np.float32) + embedding_used = np.zeros((n, 1), dtype=np.float32) + + embed_indices: list[int] = [] + embed_smiles: list[str] = [] + for i, readout in enumerate(readouts): + if readout: + for j, name in enumerate(task_names): + if name in readout: + readout_vec[i, j] = readout[name] + readout_mask[i, j] = 1.0 + else: + embedding_used[i, 0] = 1.0 + embed_indices.append(i) + embed_smiles.append(canonical_smiles[i]) + + embedding_dim = cast(int, cast(Any, aux_encoder.model).message_passing.output_dim) + embeddings = np.zeros((n, embedding_dim), dtype=np.float32) + if embed_smiles: + computed = aux_encoder.embed_smiles(embed_smiles, batch_size=batch_size) + for idx, row in zip(embed_indices, computed, strict=True): + embeddings[idx] = row + + return np.concatenate([readout_vec, readout_mask, embeddings, embedding_used], axis=1) + + +class _ConcatenatedDataset(Dataset): + """Dataset pairing a molecular graph and LabelRecord with a precomputed feature row. + + Parameters + ---------- + records : list[LabelRecord] + Labeled observations. + features : np.ndarray + Precomputed concatenation features, shape ``(len(records), feature_dim)``, + aligned with ``records`` (typically from :func:`build_concatenation_features`). + """ + + def __init__(self, records: list[LabelRecord], features: np.ndarray) -> None: + if len(records) != len(features): + raise ValueError( + f"records ({len(records)}) and features ({len(features)}) must be the same length." + ) + self.records = records + self._features = torch.as_tensor(features, dtype=torch.float32) + self._mol_graphs = MoleculeDataset( + [MoleculeDatapoint.from_smi(r.canonical_smiles) for r in records] # pyright: ignore[reportArgumentType] + ) + + def __len__(self) -> int: + """Return the number of records in the dataset. + + Returns + ------- + int + Total number of labeled observations. + """ + return len(self.records) + + def __getitem__(self, idx: int) -> tuple[Any, Tensor, LabelRecord]: + """Return the (datapoint, feature row, LabelRecord) triple at ``idx``. + + Parameters + ---------- + idx : int + Zero-based index into the dataset. + + Returns + ------- + tuple[Any, Tensor, LabelRecord] + A ``(MoleculeDatapoint, feature row, LabelRecord)`` triple. + """ + return self._mol_graphs[idx], self._features[idx], self.records[idx] + + @staticmethod + def collate_fn( + batch: list[tuple[Any, Tensor, LabelRecord]], + ) -> tuple[Any, Tensor, list[LabelRecord]]: + """Collate a list of (datapoint, feature row, LabelRecord) into a batch. + + Parameters + ---------- + batch : list[tuple[Any, Tensor, LabelRecord]] + Items as returned by :meth:`__getitem__`. + + Returns + ------- + tuple[BatchMolGraph, Tensor, list[LabelRecord]] + Batched molecular graph, stacked feature matrix + ``(batch, feature_dim)``, and corresponding label records. + """ + datapoints, features, records = zip(*batch, strict=False) + bmg = BatchMolGraph([dp.mg for dp in datapoints]) + return bmg, torch.stack(list(features)), list(records) + + +class ConcatenationChemPropLightningModule(L.LightningModule): + """ChemProp MPNN with a concatenated auxiliary-signal input before the pEC50 head. + + Parameters mirror :class:`moal.model.ChemPropLightningModule` exactly, + plus ``concat_feature_dim``; see that class for the shared parameters' + documentation. + + Parameters + ---------- + concat_feature_dim : int + Width of the concatenation feature vector (see + :func:`concatenation_feature_dim`); determines the predictor head's + input width alongside the backbone's own pooled-embedding width. + ffn_hidden_dim, ffn_num_layers, message_hidden_dim, depth, freeze_epochs, + mpnn_lr, ffn_lr, mpnn_weight_decay, ffn_weight_decay, sigma, w_drc, w_ps, + learnable_sigma, from_foundation + See :class:`moal.model.ChemPropLightningModule`. + """ + + def __init__( + self, + concat_feature_dim: int, + ffn_hidden_dim: int = 300, + ffn_num_layers: int = 2, + message_hidden_dim: int = 300, + depth: int = 3, + freeze_epochs: int = 10, + mpnn_lr: float = 1e-5, + ffn_lr: float = 1e-3, + mpnn_weight_decay: float = 0.0, + ffn_weight_decay: float = 0.0, + sigma: float = 0.5, + w_drc: float = 1.0, + w_ps: float = 0.3, + learnable_sigma: bool = False, + from_foundation: str | bool = "chemeleon", + ) -> None: + super().__init__() + _validate_from_foundation(from_foundation) + self._from_foundation = from_foundation + self.concat_feature_dim = concat_feature_dim + self.save_hyperparameters() + + self.freeze_epochs = freeze_epochs + self.mpnn_lr = mpnn_lr + self.ffn_lr = ffn_lr + self.mpnn_weight_decay = mpnn_weight_decay + self.ffn_weight_decay = ffn_weight_decay + self._encoder_frozen = True + + self._epoch_losses: dict[str, list[Tensor]] = { + "train_drc": [], + "train_ps": [], + "val_drc": [], + "val_ps": [], + } + + self.loss_fn = CensoredRegressionLoss( + sigma=sigma, w_drc=w_drc, w_ps=w_ps, learnable_sigma=learnable_sigma + ) + + self.model = build_mpnn( + from_foundation=from_foundation, + ffn_hidden_dim=ffn_hidden_dim, + ffn_num_layers=ffn_num_layers, + message_hidden_dim=message_hidden_dim, + depth=depth, + n_tasks=1, + extra_input_dim=concat_feature_dim, + ) + self._freeze_encoder() + + # ------------------------------------------------------------------ + # Freeze / unfreeze schedule + # ------------------------------------------------------------------ + + def _encoder_params(self) -> list[nn.Parameter]: + """Return the trainable parameters of the message-passing encoder. + + Returns + ------- + list[nn.Parameter] + Parameters belonging to ``self.model.message_passing``. + """ + return list(cast(nn.Module, self.model.message_passing).parameters()) + + def _head_params(self) -> list[nn.Parameter]: + """Return the trainable parameters of the aggregation layer and FFN head. + + Returns + ------- + list[nn.Parameter] + Parameters belonging to ``self.model.agg`` and + ``self.model.predictor``, concatenated in that order. + """ + return list(cast(nn.Module, self.model.agg).parameters()) + list( + cast(nn.Module, self.model.predictor).parameters() + ) + + def _freeze_encoder(self) -> None: + """Freeze all message-passing encoder parameters.""" + for p in self._encoder_params(): + p.requires_grad_(False) + self._encoder_frozen = True + + def _unfreeze_encoder(self) -> None: + """Unfreeze the message-passing encoder after the warm-up phase.""" + for p in self._encoder_params(): + p.requires_grad_(True) + self._encoder_frozen = False + + def on_train_epoch_start(self) -> None: + """Lightning hook: unfreeze the encoder once warm-up is complete.""" + if self._encoder_frozen and self.current_epoch >= self.freeze_epochs: + self._unfreeze_encoder() + self.trainer.strategy.setup_optimizers(self.trainer) + + # ------------------------------------------------------------------ + # Lightning interface + # ------------------------------------------------------------------ + + def forward(self, batch_mol_graph: Any, x_d: Tensor) -> Tensor: + """Run a forward pass and return scalar pEC50 predictions. + + Parameters + ---------- + batch_mol_graph : Any + A batched molecular graph (``chemprop.data.BatchMolGraph``). + x_d : Tensor + Concatenation features, shape ``(batch, concat_feature_dim)``. + + Returns + ------- + Tensor + 1-D tensor of shape ``(N,)`` with predicted pEC50 values. + """ + return cast(Tensor, self.model(batch_mol_graph, X_d=x_d).squeeze(-1)) + + def training_step(self, batch: tuple[Any, Tensor, list[LabelRecord]], batch_idx: int) -> Tensor: + """Compute and log the training loss for one batch. + + Parameters + ---------- + batch : tuple[Any, Tensor, list[LabelRecord]] + A ``(mol_graph, x_d, records)`` triple. + batch_idx : int + Index of the batch within the current epoch (unused). + + Returns + ------- + Tensor + Scalar total training loss used for the backward pass. + """ + mol_graph, x_d, records = batch + predictions = self(mol_graph, x_d) + breakdown = self.loss_fn.forward_with_breakdown(predictions, records) + self.log("train_loss", breakdown.total, prog_bar=True, batch_size=len(records)) + if not breakdown.drc_loss.isnan(): + self._epoch_losses["train_drc"].append(breakdown.drc_loss.detach()) + if not breakdown.ps_loss.isnan(): + self._epoch_losses["train_ps"].append(breakdown.ps_loss.detach()) + return breakdown.total + + def validation_step(self, batch: tuple[Any, Tensor, list[LabelRecord]], batch_idx: int) -> None: + """Compute and log the validation loss for one batch. + + Parameters + ---------- + batch : tuple[Any, Tensor, list[LabelRecord]] + A ``(mol_graph, x_d, records)`` triple. + batch_idx : int + Index of the batch within the current validation epoch (unused). + """ + mol_graph, x_d, records = batch + predictions = self(mol_graph, x_d) + breakdown = self.loss_fn.forward_with_breakdown(predictions, records) + self.log("val_loss", breakdown.total, prog_bar=True, batch_size=len(records)) + if not breakdown.drc_loss.isnan(): + self._epoch_losses["val_drc"].append(breakdown.drc_loss.detach()) + if not breakdown.ps_loss.isnan(): + self._epoch_losses["val_ps"].append(breakdown.ps_loss.detach()) + + def on_train_epoch_end(self) -> None: + """Emit epoch-mean DRC and PS training losses with a fixed key set.""" + self._log_epoch_fidelity_means("train") + + def on_validation_epoch_end(self) -> None: + """Emit epoch-mean DRC and PS validation losses with a fixed key set.""" + self._log_epoch_fidelity_means("val") + + def _log_epoch_fidelity_means(self, stage: str) -> None: + """Log epoch-mean fidelity losses for ``stage`` and reset accumulators. + + Parameters + ---------- + stage : str + Either ``"train"`` or ``"val"``. + """ + for fidelity in ("drc", "ps"): + values = self._epoch_losses[f"{stage}_{fidelity}"] + mean = torch.stack(values).mean() if values else torch.tensor(float("nan")) + self.log(f"{stage}_{fidelity}_loss", mean) + self._epoch_losses[f"{stage}_{fidelity}"] = [] + + def configure_optimizers(self) -> Adam: + """Build and return the Adam optimizer for the current freeze state. + + Returns + ------- + Adam + Same param-group structure as + :meth:`moal.model.ChemPropLightningModule.configure_optimizers`. + """ + param_groups = [ + { + "params": self._head_params(), + "lr": self.ffn_lr, + "weight_decay": self.ffn_weight_decay, + } + ] + if not self._encoder_frozen: + param_groups.append( + { + "params": self._encoder_params(), + "lr": self.mpnn_lr, + "weight_decay": self.mpnn_weight_decay, + } + ) + return Adam(param_groups) + + # ------------------------------------------------------------------ + # Inference helpers + # ------------------------------------------------------------------ + + @torch.no_grad() + def predict_smiles( + self, + smiles_list: list[str], + readouts: list[dict[str, float]], + aux_encoder: AuxiliaryEncoderModule, + batch_size: int = 256, + ) -> np.ndarray: + """Run batch inference over a list of canonical SMILES with concatenated features. + + Parameters + ---------- + smiles_list : list[str] + **Must be RDKit-canonical, salt-stripped SMILES**; see + :meth:`moal.model.ChemPropLightningModule.predict_smiles`. + readouts : list[dict[str, float]] + Per-compound observed readouts, aligned with ``smiles_list``; an + empty dict routes that compound through the auxiliary encoder's + structural embedding. Forwarded to + :func:`build_concatenation_features`. + aux_encoder : AuxiliaryEncoderModule + Pretrained auxiliary encoder supplying both the readout-key + order and the structural-embedding fallback. + batch_size : int, optional + Number of molecules processed per forward pass. Default is 256. + + Returns + ------- + np.ndarray + Array of shape ``(N,)`` with pEC50 point estimates, aligned with + ``smiles_list``. + """ + features = build_concatenation_features( + smiles_list, readouts, aux_encoder, batch_size=batch_size + ) + x_d = torch.as_tensor(features, dtype=torch.float32) + + # Chunk manually (rather than via chemprop's build_dataloader) so each + # chunk's x_d slice is trivially aligned with its BatchMolGraph by + # construction, instead of depending on undocumented batch-boundary + # behavior inside the dataloader. + all_preds = [] + with torch.inference_mode(): + for start in range(0, len(smiles_list), batch_size): + chunk_smiles = smiles_list[start : start + batch_size] + chunk_dataset = MoleculeDataset( + [MoleculeDatapoint.from_smi(s) for s in chunk_smiles] # pyright: ignore[reportArgumentType] + ) + bmg = BatchMolGraph([chunk_dataset[i].mg for i in range(len(chunk_dataset))]) + bmg.to(self.device) + chunk_x_d = x_d[start : start + len(chunk_smiles)].to(self.device) + preds = self(bmg, chunk_x_d).cpu().numpy().tolist() + all_preds.extend(preds) + + return np.array(all_preds, dtype=np.float32) + + def refit( + self, + records: list[LabelRecord], + aux_encoder: AuxiliaryEncoderModule, + max_epochs: int = 30, + enable_progress_bar: bool = False, + enable_model_summary: bool = False, + trainer_kwargs: dict[str, Any] | None = None, + datamodule_kwargs: dict[str, Any] | None = None, + output_dir: str | Path | None = None, + ) -> ConcatenationChemPropLightningModule: + """Refit the model on a (growing) labeled pool, using concatenated features. + + Parameters + ---------- + records : list[LabelRecord] + All labeled records accumulated so far. + aux_encoder : AuxiliaryEncoderModule + Pretrained auxiliary encoder used to build each record's + concatenation features via :func:`build_concatenation_features`. + max_epochs : int, optional + Number of training epochs. Default is 30. + enable_progress_bar : bool, optional + Whether to show the Lightning progress bar. Default is False. + enable_model_summary : bool, optional + Whether to print the model summary at the start of training. + Default is False. + trainer_kwargs : dict[str, Any], optional + Additional keyword arguments forwarded directly to + ``lightning.Trainer``. + datamodule_kwargs : dict[str, Any], optional + Passed to the underlying data module (e.g. ``val_fraction``, + ``seed``). + output_dir : str or Path, optional + Directory used as Lightning's ``default_root_dir``. + + Returns + ------- + ConcatenationChemPropLightningModule + self (for chaining). + """ + records = normalize_record_weights(records) + features = build_concatenation_features( + [rec.canonical_smiles for rec in records], + [rec.raw_ps_readouts for rec in records], + aux_encoder, + ) + dm = _ConcatenatedDataModule(records, features, **(datamodule_kwargs or {})) + dm.setup() + + kwargs: dict[str, Any] = { + "max_epochs": max_epochs, + "enable_progress_bar": enable_progress_bar, + "enable_model_summary": enable_model_summary, + } + if trainer_kwargs: + kwargs.update(trainer_kwargs) + if output_dir is not None and "default_root_dir" not in kwargs: + kwargs["default_root_dir"] = str(output_dir) + kwargs.setdefault("logger", False) + kwargs.setdefault("enable_checkpointing", False) + trainer = L.Trainer(**kwargs) + trainer.fit(self, datamodule=dm) + return self + + +class _ConcatenatedDataModule(L.LightningDataModule): + """LightningDataModule for concatenation-architecture pretraining. + + Same train/val split and device-transfer shape as + :class:`~moal.dataset.MixedFidelityDataModule`, extended to also bundle + each record's precomputed concatenation feature row. Not a subclass of + that class: nearly every method's batch shape differs (an added feature + tensor), so subclassing would mean overriding almost everything anyway. + + Parameters + ---------- + records : list[LabelRecord] + All labeled observations (train + val pool). + features : np.ndarray + Precomputed concatenation features, aligned with ``records``. + batch_size : int, optional + Number of samples per mini-batch. Default is 64. + val_fraction : float, optional + Fraction of records held out for validation. Default is 0.1. + num_workers : int, optional + DataLoader worker count (0 = main process). Default is 0. + seed : int, optional + Random seed for the train/val split. Default is 42. + """ + + def __init__( + self, + records: list[LabelRecord], + features: np.ndarray, + batch_size: int = 64, + val_fraction: float = 0.1, + num_workers: int = 0, + seed: int = 42, + ) -> None: + super().__init__() + self.records = records + self._features = features + self.batch_size = batch_size + self.val_fraction = val_fraction + self.num_workers = num_workers + self.seed = seed + + self._train_dataset: Dataset | None = None + self._val_dataset: Dataset | None = None + + def setup(self, stage: str | None = None) -> None: + """Create the train and validation dataset splits over (record, feature) pairs. + + Parameters + ---------- + stage : str or None, optional + Lightning stage identifier; unused, accepted for interface + compatibility. + """ + n_val = max(1, int(len(self.records) * self.val_fraction)) + n_train = len(self.records) - n_val + if n_train <= 0: + logger.warning( + "Too few records (%d) for a val split; using all for training.", + len(self.records), + ) + n_train, n_val = len(self.records), 0 + + full = _ConcatenatedDataset(self.records, self._features) + if n_val > 0: + self._train_dataset, self._val_dataset = random_split( + full, + [n_train, n_val], + generator=torch.Generator().manual_seed(self.seed), + ) + else: + self._train_dataset = full + self._val_dataset = None + + def transfer_batch_to_device( + self, + batch: tuple[Any, Tensor, list[LabelRecord]], + device: torch.device, + dataloader_idx: int, + ) -> tuple[Any, Tensor, list[LabelRecord]]: + """Move the batched mol graph and feature tensor to ``device``. + + Parameters + ---------- + batch : tuple[Any, Tensor, list[LabelRecord]] + A ``(BatchMolGraph, x_d, records)`` triple. + device : torch.device + Target device. + dataloader_idx : int + Index of the dataloader (required by the Lightning interface). + + Returns + ------- + tuple[Any, Tensor, list[LabelRecord]] + The same triple with the graph and feature tensor moved to + ``device``; the LabelRecord list is returned unchanged. + """ + mol_graph, x_d, records = batch + mol_graph = super().transfer_batch_to_device(mol_graph, device, dataloader_idx) + return mol_graph, x_d.to(device), records + + def train_dataloader(self) -> DataLoader: + """Return the training DataLoader. + + Returns + ------- + DataLoader + Shuffled DataLoader over the training split using + :meth:`_ConcatenatedDataset.collate_fn`. + """ + if self._train_dataset is None: + raise RuntimeError("setup() must be called before train_dataloader()") + return DataLoader( + self._train_dataset, + batch_size=self.batch_size, + shuffle=True, + collate_fn=_ConcatenatedDataset.collate_fn, + num_workers=self.num_workers, + persistent_workers=self.num_workers > 0, + drop_last=False, + ) + + def val_dataloader(self) -> DataLoader | None: + """Return the validation DataLoader, or ``None`` when no val split exists. + + Returns + ------- + DataLoader or None + Non-shuffled DataLoader over the validation split, or ``None``. + """ + if self._val_dataset is None: + return None + return DataLoader( + self._val_dataset, + batch_size=self.batch_size, + shuffle=False, + collate_fn=_ConcatenatedDataset.collate_fn, + num_workers=self.num_workers, + persistent_workers=self.num_workers > 0, + ) diff --git a/moal/model.py b/moal/model.py index 223c555..03b561a 100644 --- a/moal/model.py +++ b/moal/model.py @@ -107,6 +107,7 @@ def build_mpnn( message_hidden_dim: int, depth: int, n_tasks: int = 1, + extra_input_dim: int = 0, ) -> nn.Module: """Construct a ChemProp MPNN, dispatching on ``from_foundation``. @@ -137,6 +138,14 @@ def build_mpnn( Number of regression targets predicted per compound. Default is 1 (the main model's single pEC50 target). The auxiliary encoder passes one task per distinct auxiliary readout key it was trained on. + extra_input_dim : int, optional + Width of an additional per-compound feature vector concatenated onto + the pooled graph embedding before the predictor head, via chemprop's + native ``MPNN.forward(bmg, X_d=...)`` support. Default is 0 (no + concatenation; the main model and auxiliary encoder both use this + default). The concatenation architecture (Phase 2) passes the + combined width of its observed-readout, readout-mask, auxiliary + embedding, and provenance-flag blocks. Returns ------- @@ -173,7 +182,7 @@ def build_mpnn( agg = MeanAggregation() ffn = RegressionFFN( # pyright: ignore[reportAbstractUsage] n_tasks=n_tasks, - input_dim=cast(BondMessagePassing, mp).output_dim, + input_dim=cast(BondMessagePassing, mp).output_dim + extra_input_dim, hidden_dim=ffn_hidden_dim, n_layers=ffn_num_layers, ) diff --git a/tests/test_concatenation_model.py b/tests/test_concatenation_model.py new file mode 100644 index 0000000..f113081 --- /dev/null +++ b/tests/test_concatenation_model.py @@ -0,0 +1,169 @@ +"""Tests for the concatenation architecture (issue #36 Phase 2). + +All tests use ``from_foundation=False`` for both the auxiliary encoder and +the concatenation model, so no network download or cached CheMeleon +checkpoint is required. +""" + +from __future__ import annotations + +import numpy as np +import pytest +import torch +from chemprop.data import BatchMolGraph, MoleculeDatapoint, MoleculeDataset + +from moal.auxiliary_encoder import AuxiliaryEncoderModule +from moal.concatenation_model import ( + ConcatenationChemPropLightningModule, + build_concatenation_features, + concatenation_feature_dim, +) +from moal.config import AuxiliaryEncoderConfig +from moal.types import CensoringType, LabelRecord, QueryType + +_EMBEDDING_DIM = 16 + + +@pytest.fixture +def aux_encoder() -> AuxiliaryEncoderModule: + """Small random-init auxiliary encoder with two tasks.""" + config = AuxiliaryEncoderConfig( + from_foundation=False, + message_hidden_dim=_EMBEDDING_DIM, + ffn_hidden_dim=16, + depth=1, + ) + return AuxiliaryEncoderModule(task_names=["log2fc_1um", "pic50"], config=config) + + +def _fast_model(concat_feature_dim: int, **overrides) -> ConcatenationChemPropLightningModule: + defaults = { + "from_foundation": False, + "message_hidden_dim": 16, + "ffn_hidden_dim": 16, + "depth": 1, + "freeze_epochs": 0, + } + defaults.update(overrides) + return ConcatenationChemPropLightningModule(concat_feature_dim=concat_feature_dim, **defaults) + + +def _records() -> list[LabelRecord]: + smiles = ["CCO", "CCN", "CCC", "c1ccccc1"] + readouts = [{"log2fc_1um": 2.1}, {}, {"pic50": 6.4}, {}] + records = [] + for smi, readout in zip(smiles, readouts, strict=True): + records.append( + LabelRecord( + smiles=smi, + canonical_smiles=smi, + value=6.0, + upper_bound=6.0, + censoring_type=CensoringType.EXACT, + fidelity=QueryType.DOSE_RESPONSE, + cost=10.0, + iteration=0, + raw_ps_readouts=readout, + ) + ) + return records + + +class TestConcatenationFeatureDim: + """Tests for concatenation_feature_dim's arithmetic.""" + + def test_matches_2n_plus_embedding_plus_1(self): + """The formula must be 2 * n_tasks + embedding_dim + 1.""" + assert concatenation_feature_dim(n_tasks=3, embedding_dim=10) == 2 * 3 + 10 + 1 + + +class TestBuildConcatenationFeatures: + """Tests for build_concatenation_features: observed vs embedding routing and shape.""" + + def test_observed_readout_routes_to_readout_block_not_embedding(self, aux_encoder): + """A compound with a non-empty readout dict must populate the readout/mask block and leave the embedding block zero, with flag=0.""" + features = build_concatenation_features(["CCO"], [{"log2fc_1um": 2.5}], aux_encoder) + n_tasks = 2 + + readout_block = features[0, :n_tasks] + mask_block = features[0, n_tasks : 2 * n_tasks] + embedding_block = features[0, 2 * n_tasks : 2 * n_tasks + _EMBEDDING_DIM] + flag = features[0, -1] + + assert readout_block[0] == 2.5 + assert list(mask_block) == [1.0, 0.0] + assert np.all(embedding_block == 0.0) + assert flag == 0.0 + + def test_missing_readout_routes_to_embedding_block(self, aux_encoder): + """A compound with an empty readout dict must leave the readout/mask block zero and populate the embedding block, with flag=1.""" + features = build_concatenation_features(["CCO"], [{}], aux_encoder) + n_tasks = 2 + + readout_mask_block = features[0, : 2 * n_tasks] + embedding_block = features[0, 2 * n_tasks : 2 * n_tasks + _EMBEDDING_DIM] + flag = features[0, -1] + + assert np.all(readout_mask_block == 0.0) + assert not np.all(embedding_block == 0.0) + assert flag == 1.0 + + def test_output_shape_matches_concatenation_feature_dim(self, aux_encoder): + """Output width must equal concatenation_feature_dim(n_tasks, embedding_dim).""" + features = build_concatenation_features( + ["CCO", "CCN", "CCC"], [{"log2fc_1um": 1.0}, {}, {"pic50": 5.0}], aux_encoder + ) + + assert features.shape == (3, concatenation_feature_dim(2, _EMBEDDING_DIM)) + + def test_mismatched_lengths_raises(self, aux_encoder): + """canonical_smiles and readouts of different lengths must raise ValueError.""" + with pytest.raises(ValueError, match="same length"): + build_concatenation_features(["CCO", "CCN"], [{}], aux_encoder) + + +class TestConcatenationChemPropLightningModule: + """Tests for training and prediction through the concatenation architecture.""" + + def test_forward_output_shape(self, aux_encoder): + """A forward pass must return one scalar prediction per input molecule.""" + feat_dim = concatenation_feature_dim(2, _EMBEDDING_DIM) + model = _fast_model(feat_dim) + dataset = MoleculeDataset([MoleculeDatapoint.from_smi(s) for s in ["CCO", "CCN"]]) + bmg = BatchMolGraph([dataset[i].mg for i in range(len(dataset))]) + x_d = torch.zeros(2, feat_dim) + + preds = model(bmg, x_d) + + assert preds.shape == (2,) + + def test_refit_and_predict_smiles_round_trip(self, aux_encoder): + """refit() must train without error and predict_smiles() must return one prediction per input SMILES.""" + feat_dim = concatenation_feature_dim(2, _EMBEDDING_DIM) + model = _fast_model(feat_dim) + records = _records() + + model.refit( + records, + aux_encoder=aux_encoder, + max_epochs=1, + datamodule_kwargs={"val_fraction": 0.25, "seed": 1}, + ) + + smiles = [r.canonical_smiles for r in records] + readouts = [r.raw_ps_readouts for r in records] + preds = model.predict_smiles(smiles, readouts, aux_encoder) + + assert preds.shape == (len(records),) + assert np.all(np.isfinite(preds)) + + def test_predict_smiles_chunks_correctly_across_batch_boundary(self, aux_encoder): + """predict_smiles must produce one prediction per SMILES even when batch_size splits the input into multiple chunks.""" + feat_dim = concatenation_feature_dim(2, _EMBEDDING_DIM) + model = _fast_model(feat_dim) + smiles = ["CCO", "CCN", "CCC", "CCCC", "CCCCC"] + readouts = [{"log2fc_1um": float(i)} for i in range(len(smiles))] + + preds = model.predict_smiles(smiles, readouts, aux_encoder, batch_size=2) + + assert preds.shape == (5,) From d281fbf6c7a53201edf07bf932a1e343af3f6362 Mon Sep 17 00:00:00 2001 From: Sean Colby Date: Thu, 23 Jul 2026 17:06:46 -0800 Subject: [PATCH 08/18] Add embedding-provenance score discount to acquisition (#36 Phase 3) 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). --- moal/acquisition.py | 113 +++++++++++++++++++++++++++++++++++--- moal/config.py | 12 ++++ tests/test_acquisition.py | 91 ++++++++++++++++++++++++++++++ 3 files changed, 207 insertions(+), 9 deletions(-) diff --git a/moal/acquisition.py b/moal/acquisition.py index 3c5301e..373ca75 100644 --- a/moal/acquisition.py +++ b/moal/acquisition.py @@ -104,11 +104,22 @@ class CostAwareGreedyAcquisition: tau : float, optional Sigmoid temperature controlling exploitation sharpness. Smaller τ means more sharply exploit the highest-scoring compounds. Default is 0.5. + embedding_provenance_discount : float, optional + Multiplicative discount applied to a candidate's score when its + prediction is flagged as embedding-derived (see ``provenance`` + arguments on :meth:`select` and :meth:`score_summary`) — the + concatenation architecture's (issue #36 Phase 2) fallback path for + compounds never PS-screened, which rests on strictly more layers of + inference than a prediction from an observed input. Must be in + ``(0.0, 1.0]``. Default is 1.0 (no discount); callers that never pass + a ``provenance`` array see no behavior change regardless of this + value. Raises ------ ValueError - If ``cost_ps`` or ``cost_drc`` is not strictly positive. + If ``cost_ps`` or ``cost_drc`` is not strictly positive, or if + ``embedding_provenance_discount`` is not in ``(0.0, 1.0]``. """ def __init__( @@ -118,14 +129,21 @@ def __init__( ps_threshold: float = 5.0, target_threshold: float = 7.0, tau: float = 0.5, + embedding_provenance_discount: float = 1.0, ) -> None: if cost_ps <= 0 or cost_drc <= 0: raise ValueError("Costs must be positive.") + if not (0.0 < embedding_provenance_discount <= 1.0): + raise ValueError( + "embedding_provenance_discount must be in (0.0, 1.0], " + f"got {embedding_provenance_discount}." + ) self.cost_ps = cost_ps self.cost_drc = cost_drc self.ps_threshold = ps_threshold self.target_threshold = target_threshold self.tau = tau + self.embedding_provenance_discount = embedding_provenance_discount # ------------------------------------------------------------------ # Scoring @@ -173,6 +191,40 @@ def _score_ps(self, predictions: np.ndarray) -> np.ndarray: h = _binary_entropy(p_cross) return h / self.cost_ps + def _apply_provenance_discount( + self, scores: np.ndarray, provenance: np.ndarray | None + ) -> np.ndarray: + """Apply ``embedding_provenance_discount`` to embedding-derived candidates. + + Parameters + ---------- + scores : np.ndarray + Raw acquisition scores, shape ``(N,)``. + provenance : np.ndarray or None + Boolean (or 0/1 float) array, shape ``(N,)``, True/1 where the + prediction is embedding-derived (concatenation architecture, + never-PS-screened fallback). ``None`` means no provenance + information was supplied — every candidate is treated as + observed-input, matching current behavior with no discount. + + Returns + ------- + np.ndarray + ``scores`` unchanged where ``provenance`` is False/0 or where + ``provenance is None``; multiplied by + ``embedding_provenance_discount`` where ``provenance`` is + True/1. + """ + if provenance is None: + return scores + provenance = np.asarray(provenance) + if provenance.shape != scores.shape: + raise ValueError( + f"provenance shape {provenance.shape} must match scores shape {scores.shape}." + ) + discount = np.where(provenance.astype(bool), self.embedding_provenance_discount, 1.0) + return scores * discount + # ------------------------------------------------------------------ # Selection # ------------------------------------------------------------------ @@ -186,6 +238,8 @@ def select( wells_per_drc: int, ps_labeled_smiles: list[str] | None = None, ps_labeled_predictions: np.ndarray | None = None, + provenance: np.ndarray | None = None, + ps_labeled_provenance: np.ndarray | None = None, ) -> list[tuple[str, QueryType]]: """Greedily select queries that fit within a plate well budget. @@ -233,6 +287,16 @@ def select( Model pEC50 estimates, shape ``(M,)``, aligned with ``ps_labeled_smiles``. Required when ``ps_labeled_smiles`` is non-empty. + provenance : np.ndarray, optional + Boolean (or 0/1 float) array, shape ``(N,)``, aligned with + ``unlabeled_smiles``. True/1 marks a prediction as + embedding-derived (concatenation architecture, issue #36 Phase 2); + its DRC and PS scores are multiplied by + ``embedding_provenance_discount``. ``None`` (default) applies no + discount, matching current behavior. + ps_labeled_provenance : np.ndarray, optional + Same semantics as ``provenance``, aligned with + ``ps_labeled_smiles`` instead. Returns ------- @@ -271,8 +335,8 @@ def select( candidates: list[tuple[float, str, QueryType]] = [] if unlabeled_smiles: - scores_drc = self._score_drc(predictions) - scores_ps = self._score_ps(predictions) + scores_drc = self._apply_provenance_discount(self._score_drc(predictions), provenance) + scores_ps = self._apply_provenance_discount(self._score_ps(predictions), provenance) for i, smi in enumerate(unlabeled_smiles): candidates.append((float(scores_drc[i]), smi, QueryType.DOSE_RESPONSE)) candidates.append((float(scores_ps[i]), smi, QueryType.PRIMARY_SCREEN)) @@ -285,7 +349,9 @@ def select( f"ps_labeled_smiles length ({len(ps_labeled_smiles)}) must match " f"ps_labeled_predictions length ({len(psl_preds)})." ) - scores_drc_upgrade = self._score_drc(psl_preds) + scores_drc_upgrade = self._apply_provenance_discount( + self._score_drc(psl_preds), ps_labeled_provenance + ) for j, smi in enumerate(ps_labeled_smiles): candidates.append((float(scores_drc_upgrade[j]), smi, QueryType.DOSE_RESPONSE)) @@ -325,7 +391,12 @@ def select( # Diagnostics # ------------------------------------------------------------------ - def score_summary(self, unlabeled_smiles: list[str], predictions: np.ndarray) -> list[dict]: + def score_summary( + self, + unlabeled_smiles: list[str], + predictions: np.ndarray, + provenance: np.ndarray | None = None, + ) -> list[dict]: """Return per-compound score breakdown for inspection and logging. Parameters @@ -335,26 +406,50 @@ def score_summary(self, unlabeled_smiles: list[str], predictions: np.ndarray) -> predictions : np.ndarray Model pEC50 point estimates, shape ``(N,)``, aligned with ``unlabeled_smiles``. + provenance : np.ndarray, optional + Boolean (or 0/1 float) array, shape ``(N,)``, aligned with + ``unlabeled_smiles``. True/1 marks a prediction as + embedding-derived (concatenation architecture, issue #36 Phase 2); + ``score_drc``/``score_ps`` are multiplied by + ``embedding_provenance_discount`` for that row, matching + :meth:`select`'s ranking. ``None`` (default) applies no discount. Returns ------- list[dict] One dict per compound with keys ``smiles``, ``y_hat``, - ``p_active``, ``p_cross_threshold``, ``score_drc``, ``score_ps``. + ``p_active``, ``p_cross_threshold``, ``score_drc``, ``score_ps``, + ``embedding_derived`` (bool, always present; False when + ``provenance`` is None). """ predictions = np.asarray(predictions, dtype=np.float32) + provenance_arr = ( + np.zeros(len(predictions), dtype=bool) + if provenance is None + else np.asarray(provenance, dtype=bool) + ) + if provenance_arr.shape != predictions.shape: + raise ValueError( + f"provenance shape {provenance_arr.shape} must match " + f"predictions shape {predictions.shape}." + ) rows = [] - for smi, y_hat in zip(unlabeled_smiles, predictions, strict=False): + for smi, y_hat, is_embedding in zip( + unlabeled_smiles, predictions, provenance_arr, strict=False + ): p_active = float(_sigmoid(np.array([y_hat - self.target_threshold]), self.tau)[0]) p_cross = float(_sigmoid(np.array([y_hat - self.ps_threshold]), self.tau)[0]) + discount = self.embedding_provenance_discount if is_embedding else 1.0 rows.append( { "smiles": smi, "y_hat": float(y_hat), "p_active": p_active, "p_cross_threshold": p_cross, - "score_drc": p_active / self.cost_drc, - "score_ps": float(_binary_entropy(np.array([p_cross]))[0]) / self.cost_ps, + "score_drc": (p_active / self.cost_drc) * discount, + "score_ps": (float(_binary_entropy(np.array([p_cross]))[0]) / self.cost_ps) + * discount, + "embedding_derived": bool(is_embedding), } ) return rows diff --git a/moal/config.py b/moal/config.py index 317d739..851883a 100644 --- a/moal/config.py +++ b/moal/config.py @@ -218,11 +218,23 @@ class AcquisitionConfig: Optimization target threshold used by the DRC exploitation score. tau : float Sigmoid temperature. Lower = more exploitative. + embedding_provenance_discount : float + Multiplicative discount applied to a candidate's acquisition score + when its prediction rests on the concatenation architecture's + (issue #36 Phase 2) auxiliary-embedding path rather than an observed + readout — i.e. a compound never PS-screened, scored through one more + layer of inference than an observed-input prediction. Must be in + ``(0.0, 1.0]``; 1.0 (default) is a no-op, so acquisition behavior is + unchanged unless a caller explicitly passes per-candidate provenance + to :meth:`~moal.acquisition.CostAwareGreedyAcquisition.select` or + :meth:`~moal.acquisition.CostAwareGreedyAcquisition.score_summary` + *and* sets this below 1.0. """ ps_threshold: float = 5.0 target_threshold: float = 7.0 tau: float = 0.5 + embedding_provenance_discount: float = 1.0 @dataclass(frozen=True) diff --git a/tests/test_acquisition.py b/tests/test_acquisition.py index 2c424b4..8595fde 100644 --- a/tests/test_acquisition.py +++ b/tests/test_acquisition.py @@ -301,3 +301,94 @@ def test_no_smiles_length_mismatch_assertion(self, acq): ps_labeled_smiles=["A", "B"], ps_labeled_predictions=np.array([1.0]), ) + + +class TestProvenanceDiscount: + """Tests for embedding_provenance_discount: constructor validation and select()/score_summary() effects.""" + + def test_default_discount_is_noop(self, acq): + """The default discount of 1.0 must leave select() output unchanged whether or not provenance is passed.""" + smiles = ["A", "B", "C"] + preds = np.array([9.0, 8.0, 7.5], dtype=np.float32) + provenance = np.array([True, False, True]) + + without_provenance = acq.select( + smiles, preds, plate_size=2, wells_per_ps=1, wells_per_drc=1 + ) + with_provenance = acq.select( + smiles, preds, plate_size=2, wells_per_ps=1, wells_per_drc=1, provenance=provenance + ) + + assert without_provenance == with_provenance + + def test_discount_below_one_can_flip_ranking(self): + """A strong discount on an embedding-derived candidate must let an otherwise-lower-scoring observed candidate outrank it.""" + acq = CostAwareGreedyAcquisition( + cost_ps=1.0, + cost_drc=1.0, + ps_threshold=5.0, + target_threshold=7.0, + tau=0.5, + embedding_provenance_discount=0.01, + ) + smiles = ["A", "B"] + preds = np.array([9.0, 7.1], dtype=np.float32) # A scores higher on raw prediction alone + provenance = np.array([True, False]) # A is embedding-derived, B is observed + + selected = acq.select( + smiles, + preds, + plate_size=1, + wells_per_ps=10, + wells_per_drc=1, + provenance=provenance, + ) + + assert selected[0][0] == "B" + + def test_score_summary_reports_discounted_scores_and_embedding_flag(self): + """score_summary must discount score_drc/score_ps for embedding-derived rows and report embedding_derived.""" + acq = CostAwareGreedyAcquisition( + cost_ps=1.0, + cost_drc=1.0, + ps_threshold=5.0, + target_threshold=7.0, + tau=0.5, + embedding_provenance_discount=0.5, + ) + smiles = ["A", "B"] + preds = np.array([8.0, 8.0], dtype=np.float32) + provenance = np.array([True, False]) + + rows = acq.score_summary(smiles, preds, provenance=provenance) + + assert rows[0]["embedding_derived"] is True + assert rows[1]["embedding_derived"] is False + assert rows[0]["score_drc"] == pytest.approx(rows[1]["score_drc"] * 0.5) + assert rows[0]["score_ps"] == pytest.approx(rows[1]["score_ps"] * 0.5) + + def test_score_summary_without_provenance_marks_all_rows_not_embedding_derived(self, acq): + """Omitting provenance must set embedding_derived=False for every row and apply no discount.""" + rows = acq.score_summary(["A"], np.array([8.0], dtype=np.float32)) + + assert rows[0]["embedding_derived"] is False + + @pytest.mark.parametrize("discount", [0.0, 1.5, -0.1]) + def test_out_of_range_discount_raises(self, discount): + """embedding_provenance_discount outside (0.0, 1.0] must raise ValueError at construction.""" + with pytest.raises(ValueError, match="embedding_provenance_discount"): + CostAwareGreedyAcquisition( + cost_ps=1.0, cost_drc=1.0, embedding_provenance_discount=discount + ) + + def test_provenance_shape_mismatch_raises(self, acq): + """A provenance array of the wrong length must raise ValueError rather than silently misaligning.""" + with pytest.raises(ValueError, match="provenance"): + acq.select( + ["A", "B"], + np.array([8.0, 7.0], dtype=np.float32), + plate_size=2, + wells_per_ps=1, + wells_per_drc=1, + provenance=np.array([True]), + ) From 005ad80a0f894e1df76a21eefb1b389c1ff37b43 Mon Sep 17 00:00:00 2001 From: Sean Colby Date: Thu, 23 Jul 2026 20:57:42 -0800 Subject: [PATCH 09/18] Thread provenance through annotate_campaign_state 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/auxiliary_encoder.py | 15 +++++++++++++++ moal/concatenation_model.py | 3 +-- moal/planning.py | 33 +++++++++++++++++++++++++++++---- tests/test_planning.py | 2 +- 4 files changed, 46 insertions(+), 7 deletions(-) diff --git a/moal/auxiliary_encoder.py b/moal/auxiliary_encoder.py index 840eb6b..621c870 100644 --- a/moal/auxiliary_encoder.py +++ b/moal/auxiliary_encoder.py @@ -301,6 +301,21 @@ def __init__(self, task_names: list[str], config: AuxiliaryEncoderConfig) -> Non ) self._freeze_encoder() + @property + def embedding_dim(self) -> int: + """Width of the backbone's pooled structural embedding. + + Returns + ------- + int + The message-passing encoder's native output width (CheMeleon's + fixed width, or ``config.message_hidden_dim`` for a random-init + encoder) — the row width :meth:`embed_smiles` returns, and the + embedding block width the concatenation architecture (Phase 2) + needs from :func:`moal.concatenation_model.concatenation_feature_dim`. + """ + return cast(int, cast(Any, self.model).message_passing.output_dim) + # ------------------------------------------------------------------ # Freeze / unfreeze schedule # ------------------------------------------------------------------ diff --git a/moal/concatenation_model.py b/moal/concatenation_model.py index 9240c15..0efa062 100644 --- a/moal/concatenation_model.py +++ b/moal/concatenation_model.py @@ -132,8 +132,7 @@ def build_concatenation_features( embed_indices.append(i) embed_smiles.append(canonical_smiles[i]) - embedding_dim = cast(int, cast(Any, aux_encoder.model).message_passing.output_dim) - embeddings = np.zeros((n, embedding_dim), dtype=np.float32) + embeddings = np.zeros((n, aux_encoder.embedding_dim), dtype=np.float32) if embed_smiles: computed = aux_encoder.embed_smiles(embed_smiles, batch_size=batch_size) for idx, row in zip(embed_indices, computed, strict=True): diff --git a/moal/planning.py b/moal/planning.py index e7645d8..f0692e9 100644 --- a/moal/planning.py +++ b/moal/planning.py @@ -410,6 +410,7 @@ def annotate_campaign_state( state: CampaignState, predictions: np.ndarray, acquisition: CostAwareGreedyAcquisition, + provenance: np.ndarray | None = None, ) -> pd.DataFrame: """Annotate the campaign state DataFrame with acquisition scores. @@ -417,7 +418,7 @@ def annotate_campaign_state( ``state.unqueried_rows + state.ps_upgrade_rows`` in that order — the same ordering used when calling ``model.predict_smiles``. - Four columns are appended to a copy of ``df``: + Five columns are appended to a copy of ``df``: - ``ps_score`` — PS exploration score; NaN for non-unqueried rows - ``drc_score`` — DRC exploitation score; NaN for training-only rows @@ -425,6 +426,9 @@ def annotate_campaign_state( ``drc_score`` for PS upgrades, NaN for training-only rows - ``recommendation`` — ``"ps"`` or ``"drc"`` for inference targets; NaN for training-only rows + - ``embedding_derived`` — True where ``provenance`` flagged the prediction + as embedding-derived (see ``provenance`` below); NaN for training-only + rows; always False when ``provenance`` is None Parameters ---------- @@ -436,11 +440,17 @@ def annotate_campaign_state( Model pEC50 predictions aligned with unqueried + ps_upgrade rows. acquisition : CostAwareGreedyAcquisition Acquisition function used to compute per-compound scores. + provenance : np.ndarray, optional + Boolean (or 0/1 float) array aligned with ``predictions``, forwarded + to ``acquisition.score_summary`` so a discount (issue #36 Phase 3) + applies to embedding-derived predictions from the concatenation + architecture. ``None`` (default) applies no discount, matching + current behavior. Returns ------- pd.DataFrame - Annotated copy with four new columns appended. + Annotated copy with five new columns appended. """ predictions = np.asarray(predictions, dtype=np.float32) n_inference = len(state.unqueried_rows) + len(state.ps_upgrade_rows) @@ -454,21 +464,32 @@ def annotate_campaign_state( "predictions must contain only finite values; NaN or inf values " "produce undefined acquisition scores." ) + provenance_arr = None if provenance is None else np.asarray(provenance) + if provenance_arr is not None and len(provenance_arr) != n_inference: + raise ValueError( + f"provenance length ({len(provenance_arr)}) must match the number of " + f"inference targets ({n_inference})." + ) result = df.copy() result["ps_score"] = np.nan result["drc_score"] = np.nan result["overall_score"] = np.nan result["recommendation"] = None # Object dtype so string values can be assigned + result["embedding_derived"] = None # Object dtype so bool values can be assigned n_unqueried = len(state.unqueried_rows) unqueried_preds = predictions[:n_unqueried] upgrade_preds = predictions[n_unqueried:] + unqueried_provenance = None if provenance_arr is None else provenance_arr[:n_unqueried] + upgrade_provenance = None if provenance_arr is None else provenance_arr[n_unqueried:] # Score unqueried compounds — both PS and DRC are valid next actions if state.unqueried_rows: unqueried_canonical = [smi for _, smi in state.unqueried_rows] - summaries = acquisition.score_summary(unqueried_canonical, unqueried_preds) + summaries = acquisition.score_summary( + unqueried_canonical, unqueried_preds, provenance=unqueried_provenance + ) for (row_idx, _), summary in zip(state.unqueried_rows, summaries, strict=False): drc = float(summary["score_drc"]) ps = float(summary["score_ps"]) @@ -478,16 +499,20 @@ def annotate_campaign_state( result.at[row_idx, "drc_score"] = drc result.at[row_idx, "overall_score"] = overall result.at[row_idx, "recommendation"] = rec + result.at[row_idx, "embedding_derived"] = summary["embedding_derived"] # Score PS hits — only DRC upgrade is a valid next action; ps_score stays NaN if state.ps_upgrade_rows: upgrade_canonical = [smi for _, smi in state.ps_upgrade_rows] - summaries = acquisition.score_summary(upgrade_canonical, upgrade_preds) + summaries = acquisition.score_summary( + upgrade_canonical, upgrade_preds, provenance=upgrade_provenance + ) for (row_idx, _), summary in zip(state.ps_upgrade_rows, summaries, strict=False): drc = float(summary["score_drc"]) result.at[row_idx, "drc_score"] = drc result.at[row_idx, "overall_score"] = drc result.at[row_idx, "recommendation"] = "drc" + result.at[row_idx, "embedding_derived"] = summary["embedding_derived"] return result diff --git a/tests/test_planning.py b/tests/test_planning.py index 7f08984..f10401e 100644 --- a/tests/test_planning.py +++ b/tests/test_planning.py @@ -624,7 +624,7 @@ def test_uses_acquisition_score_summary_for_scoring(self, preprocessor): """Scores must come from acquisition.score_summary() to ensure the acquisition strategy drives recommendations.""" acquisition = Mock(spec_set=["score_summary"]) acquisition.score_summary.return_value = [ - {"smiles": "CCO", "score_drc": 0.3, "score_ps": 0.7}, + {"smiles": "CCO", "score_drc": 0.3, "score_ps": 0.7, "embedding_derived": False}, ] df = _state_df({"smiles": "CCO", "relation": "", "value": ""}) state = parse_campaign_state( From 54fd774d043fe3b79b70ea85f31ca9447fae0001 Mon Sep 17 00:00:00 2001 From: Sean Colby Date: Thu, 23 Jul 2026 21:02:55 -0800 Subject: [PATCH 10/18] Wire auxiliary encoder and concatenation architecture into moal plan CLI MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- moal/cli.py | 146 ++++++++++++++++++++++++++++++++++++++++------ tests/test_cli.py | 74 +++++++++++++++++++++++ 2 files changed, 202 insertions(+), 18 deletions(-) diff --git a/moal/cli.py b/moal/cli.py index f58f26d..46fdeaf 100644 --- a/moal/cli.py +++ b/moal/cli.py @@ -29,6 +29,11 @@ ) from moal.acquisition import CostAwareGreedyAcquisition +from moal.auxiliary_encoder import AuxiliaryEncoderModule, pretrain_auxiliary_encoder +from moal.concatenation_model import ( + ConcatenationChemPropLightningModule, + concatenation_feature_dim, +) from moal.config import PipelineConfig from moal.dashboard import LiveDashboard from moal.evaluation import ModelMetric, PipelineEvaluator, scaffold_split @@ -37,6 +42,7 @@ from moal.model import ChemPropLightningModule, NoisyOracleModel from moal.oracle import CostAwareOracle from moal.planning import ( + CampaignState, annotate_campaign_state, parse_campaign_state, parse_pretrain_records, @@ -392,32 +398,66 @@ def plan(config: Path, output_dir: Path | None, verbose: bool) -> None: # Setting all seeds L.seed_everything(cfg.seed, workers=True, verbose=False) - # Build model - model = _build_plan_model(cfg) - - # Train model - model.refit( - records=fit_records, - trainer_kwargs=cfg.trainer.to_dict(), - datamodule_kwargs=cfg.trainer.to_datamodule_kwargs(), - reset_weights=cfg.model.reset_weights_on_refit, - output_dir=out_dir, - ) - progress.advance(task) - - progress.update(task, description=scoring_description) - # Collect SMILES for inference: unqueried compounds and PS hits eligible for upgrade inference_smiles = [smi for _, smi in state.unqueried_rows] + [ smi for _, smi in state.ps_upgrade_rows ] - # Make predictions - predictions = model.predict_smiles(inference_smiles) + if cfg.auxiliary_encoder is not None: + # Concatenation architecture (issue #36 Phase 2): pretrain the + # auxiliary encoder on this run's raw_ps_readouts, then train + # the main model with its structural embedding concatenated in + # for compounds that were never PS-screened. + progress.update( + task, description="[yellow]Pretraining auxiliary encoder[/yellow]" + ) + try: + aux_encoder = pretrain_auxiliary_encoder(fit_records, cfg.auxiliary_encoder) + except ValueError as exc: + raise click.ClickException(str(exc)) from exc + + progress.update(task, description=retraining_description) + model = _build_concatenation_model(cfg, aux_encoder) + model.refit( + fit_records, + aux_encoder=aux_encoder, + trainer_kwargs=cfg.trainer.to_dict(), + datamodule_kwargs=cfg.trainer.to_datamodule_kwargs(), + output_dir=out_dir, + ) + progress.advance(task) + + progress.update(task, description=scoring_description) + inference_readouts = _inference_readouts(state, fit_records) + predictions = model.predict_smiles( + inference_smiles, inference_readouts, aux_encoder + ) + # embedding_derived == True wherever the compound had no observed + # readout, mirroring build_concatenation_features's own routing rule + provenance = np.array( + [not readout for readout in inference_readouts], dtype=bool + ) + else: + # Build model + model = _build_plan_model(cfg) + + # Train model + model.refit( + records=fit_records, + trainer_kwargs=cfg.trainer.to_dict(), + datamodule_kwargs=cfg.trainer.to_datamodule_kwargs(), + reset_weights=cfg.model.reset_weights_on_refit, + output_dir=out_dir, + ) + progress.advance(task) + + progress.update(task, description=scoring_description) + predictions = model.predict_smiles(inference_smiles) + provenance = None try: annotated_df = annotate_campaign_state( - state_df, state, predictions, acquisition + state_df, state, predictions, acquisition, provenance=provenance ) except ValueError as exc: raise click.ClickException(str(exc)) from exc @@ -672,6 +712,7 @@ def _build_acquisition(cfg: PipelineConfig) -> CostAwareGreedyAcquisition: ps_threshold=cfg.acquisition.ps_threshold, target_threshold=cfg.acquisition.target_threshold, tau=cfg.acquisition.tau, + embedding_provenance_discount=cfg.acquisition.embedding_provenance_discount, ) @@ -768,5 +809,74 @@ def _build_plan_model(cfg: PipelineConfig) -> ChemPropLightningModule: ) +def _build_concatenation_model( + cfg: PipelineConfig, aux_encoder: AuxiliaryEncoderModule +) -> ConcatenationChemPropLightningModule: + """Instantiate a ``ConcatenationChemPropLightningModule`` for offline planning. + + Parameters + ---------- + cfg : PipelineConfig + Active campaign configuration. Reuses ``cfg.model``'s backbone and + optimization hyperparameters, same as :func:`_build_plan_model`. + aux_encoder : AuxiliaryEncoderModule + Pretrained auxiliary encoder; supplies ``task_names`` and + ``embedding_dim`` to size the concatenation feature width. + + Returns + ------- + ConcatenationChemPropLightningModule + Configured model ready for ``refit()`` and ``predict_smiles()``. + """ + feature_dim = concatenation_feature_dim(len(aux_encoder.task_names), aux_encoder.embedding_dim) + return ConcatenationChemPropLightningModule( + concat_feature_dim=feature_dim, + ffn_hidden_dim=cfg.model.ffn_hidden_dim, + ffn_num_layers=cfg.model.ffn_num_layers, + message_hidden_dim=cfg.model.message_hidden_dim, + depth=cfg.model.depth, + freeze_epochs=cfg.model.freeze_epochs, + mpnn_lr=cfg.model.mpnn_lr, + ffn_lr=cfg.model.ffn_lr, + mpnn_weight_decay=cfg.model.mpnn_weight_decay, + ffn_weight_decay=cfg.model.ffn_weight_decay, + sigma=cfg.model.sigma, + w_drc=cfg.model.w_drc, + w_ps=cfg.model.w_ps, + learnable_sigma=cfg.model.learnable_sigma, + from_foundation=cfg.model.from_foundation, + ) + + +def _inference_readouts( + state: CampaignState, fit_records: list[LabelRecord] +) -> list[dict[str, float]]: + """Build the per-inference-target readout dict list for the concatenation architecture. + + Unqueried compounds have never been PS-screened by definition, so they + always route through the auxiliary encoder's embedding fallback (empty + dict). PS-upgrade candidates already carry their own observed readouts on + the corresponding training record. + + Parameters + ---------- + state : CampaignState + Parsed campaign state. + fit_records : list[LabelRecord] + Training records used to fit the model, keyed by canonical SMILES to + recover each PS-upgrade candidate's ``raw_ps_readouts``. + + Returns + ------- + list[dict[str, float]] + Aligned with ``state.unqueried_rows + state.ps_upgrade_rows``, same + ordering ``model.predict_smiles`` expects. + """ + readouts_by_smiles = {rec.canonical_smiles: rec.raw_ps_readouts for rec in fit_records} + unqueried_readouts: list[dict[str, float]] = [{} for _ in state.unqueried_rows] + upgrade_readouts = [readouts_by_smiles.get(smi, {}) for _, smi in state.ps_upgrade_rows] + return unqueried_readouts + upgrade_readouts + + if __name__ == "__main__": main() diff --git a/tests/test_cli.py b/tests/test_cli.py index b9263e2..63f5a93 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -367,6 +367,80 @@ def test_plan_writes_annotated_state_csv(self, tmp_path, monkeypatch, progress_r # inference_smiles = [unqueried...] + [ps_upgrade...] model.predict_smiles.assert_called_once_with(["CCN", "CCCC", "CCO"]) + def test_plan_uses_concatenation_architecture_when_auxiliary_encoder_configured( + self, tmp_path, monkeypatch + ): + """When auxiliary_encoder is set, plan must pretrain it, build the concatenation model, and annotate embedding_derived.""" + state_csv = tmp_path / "state.csv" + state_csv.write_text( + "smiles,relation,value,log2fc_1um\nCCO,>=,5.0,3.1\nc1ccccc1,==,8.1,\nCCN,,,\nCCCC,,,\n" + ) + output_csv = tmp_path / "state_out.csv" + cfg = tmp_path / "config.yaml" + cfg.write_text( + "oracle:\n" + " cost_ps: 1.0\n" + " cost_drc: 10.0\n" + " ps_threshold: 5.0\n" + "acquisition:\n" + " ps_threshold: 5.0\n" + " target_threshold: 7.0\n" + " tau: 0.5\n" + + _plan_config( + input_csv=str(state_csv), + output_csv=str(output_csv), + extra=" log2fc_columns: [log2fc_1um]\n", + ) + + "model:\n" + " fast: false\n" + "trainer:\n" + " max_epochs: 1\n" + "dashboard:\n" + " enabled: false\n" + "auxiliary_encoder:\n" + " freeze_epochs: 0\n" + ) + + fake_aux_encoder = Mock(spec_set=["task_names", "embedding_dim"]) + fake_aux_encoder.task_names = ["log2fc_1um"] + fake_aux_encoder.embedding_dim = 8 + + pretrain_mock = Mock(return_value=fake_aux_encoder) + monkeypatch.setattr("moal.cli.pretrain_auxiliary_encoder", pretrain_mock) + + concat_model = Mock(spec_set=["refit", "predict_smiles"]) + # unqueried: CCN, CCCC (2); ps upgrade: CCO (1) -> 3 total predictions + concat_model.predict_smiles.return_value = np.array([5.0, 8.0, 6.5], dtype=np.float32) + monkeypatch.setattr( + "moal.cli._build_concatenation_model", lambda cfg, aux_encoder: concat_model + ) + + runner = CliRunner() + result = runner.invoke( + main, + ["plan", "--config", str(cfg), "--output-dir", str(tmp_path / "out")], + ) + + assert result.exit_code == 0, _result_text(result) + pretrain_mock.assert_called_once() + concat_model.refit.assert_called_once() + assert concat_model.refit.call_args.kwargs["aux_encoder"] is fake_aux_encoder + + # predict_smiles must receive per-compound readouts: empty for unqueried, + # the observed reading for the PS-upgrade candidate + call_args = concat_model.predict_smiles.call_args + smiles_arg, readouts_arg, aux_arg = call_args[0] + assert smiles_arg == ["CCN", "CCCC", "CCO"] + assert readouts_arg == [{}, {}, {"log2fc_1um": 3.1}] + assert aux_arg is fake_aux_encoder + + written = pd.read_csv(output_csv) + assert "embedding_derived" in written.columns + unqueried_rows = written[written["smiles"].isin(["CCN", "CCCC"])] + assert unqueried_rows["embedding_derived"].astype(bool).all() + upgrade_row = written[written["smiles"] == "CCO"] + assert not upgrade_row["embedding_derived"].astype(bool).any() + def test_plan_suppresses_noisy_third_party_warnings(self, tmp_path, monkeypatch): """suppress_noisy_loggers must be called exactly once so third-party warnings do not pollute plan output.""" state_csv = tmp_path / "state.csv" From c865e79a57136d20a43b190502103e5a0b48dfbe Mon Sep 17 00:00:00 2001 From: Sean Colby Date: Mon, 27 Jul 2026 09:10:28 -0800 Subject: [PATCH 11/18] Fix val_fraction=0.0 disabling validation splits in Lightning DataModules 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. --- moal/auxiliary_encoder.py | 19 +++++++++++++------ moal/concatenation_model.py | 19 +++++++++++++------ moal/dataset.py | 22 ++++++++++++++-------- tests/test_auxiliary_encoder.py | 4 ++-- 4 files changed, 42 insertions(+), 22 deletions(-) diff --git a/moal/auxiliary_encoder.py b/moal/auxiliary_encoder.py index 621c870..eb2a779 100644 --- a/moal/auxiliary_encoder.py +++ b/moal/auxiliary_encoder.py @@ -188,7 +188,9 @@ def setup(self, stage: str | None = None) -> None: Lightning stage identifier; unused, accepted for interface compatibility. """ - n_val = max(1, int(len(self.records) * self.val_fraction)) + n_val = int(len(self.records) * self.val_fraction) + if self.val_fraction > 0.0: + n_val = max(1, n_val) n_train = len(self.records) - n_val if n_train <= 0: logger.warning( @@ -251,16 +253,21 @@ def train_dataloader(self) -> DataLoader: drop_last=False, ) - def val_dataloader(self) -> DataLoader | None: - """Return the validation DataLoader, or ``None`` when no val split exists. + def val_dataloader(self) -> DataLoader: + """Return the validation DataLoader, empty when no val split exists. + + Lightning requires a real iterable from this hook (returning ``None`` + raises); an empty ``DataLoader`` yields zero validation batches, + which is the correct behavior for ``val_fraction=0.0``. Returns ------- - DataLoader or None - Non-shuffled DataLoader over the validation split, or ``None``. + DataLoader + Non-shuffled DataLoader over the validation split, or an empty + ``DataLoader`` if no split was formed. """ if self._val_dataset is None: - return None + return DataLoader([], batch_size=self.batch_size) # pyright: ignore[reportArgumentType] return DataLoader( self._val_dataset, batch_size=self.batch_size, diff --git a/moal/concatenation_model.py b/moal/concatenation_model.py index 0efa062..25b5f43 100644 --- a/moal/concatenation_model.py +++ b/moal/concatenation_model.py @@ -621,7 +621,9 @@ def setup(self, stage: str | None = None) -> None: Lightning stage identifier; unused, accepted for interface compatibility. """ - n_val = max(1, int(len(self.records) * self.val_fraction)) + n_val = int(len(self.records) * self.val_fraction) + if self.val_fraction > 0.0: + n_val = max(1, n_val) n_train = len(self.records) - n_val if n_train <= 0: logger.warning( @@ -689,16 +691,21 @@ def train_dataloader(self) -> DataLoader: drop_last=False, ) - def val_dataloader(self) -> DataLoader | None: - """Return the validation DataLoader, or ``None`` when no val split exists. + def val_dataloader(self) -> DataLoader: + """Return the validation DataLoader, empty when no val split exists. + + Lightning requires a real iterable from this hook (returning ``None`` + raises); an empty ``DataLoader`` yields zero validation batches, + which is the correct behavior for ``val_fraction=0.0``. Returns ------- - DataLoader or None - Non-shuffled DataLoader over the validation split, or ``None``. + DataLoader + Non-shuffled DataLoader over the validation split, or an empty + ``DataLoader`` if no split was formed. """ if self._val_dataset is None: - return None + return DataLoader([], batch_size=self.batch_size) # pyright: ignore[reportArgumentType] return DataLoader( self._val_dataset, batch_size=self.batch_size, diff --git a/moal/dataset.py b/moal/dataset.py index 6e9af81..e7aada8 100644 --- a/moal/dataset.py +++ b/moal/dataset.py @@ -139,7 +139,9 @@ def setup(self, stage: str | None = None) -> None: ``"test"``, ``"predict"``). Not used; accepted for interface compatibility. """ - n_val = max(1, int(len(self.records) * self.val_fraction)) + n_val = int(len(self.records) * self.val_fraction) + if self.val_fraction > 0.0: + n_val = max(1, n_val) n_train = len(self.records) - n_val if n_train <= 0: logger.warning( @@ -215,18 +217,22 @@ def train_dataloader(self) -> DataLoader: drop_last=False, ) - def val_dataloader(self) -> DataLoader | None: - """Return the validation DataLoader, or ``None`` when no val split exists. + def val_dataloader(self) -> DataLoader: + """Return the validation DataLoader, empty when no val split exists. + + Lightning requires a real iterable from this hook (returning ``None`` + raises); an empty ``DataLoader`` yields zero validation batches, + which is the correct behavior for ``val_fraction=0.0`` or a record + pool too small to form a split during :meth:`setup`. Returns ------- - DataLoader or None - Non-shuffled DataLoader over the validation split, or ``None`` - if the record pool was too small to form a validation set during - :meth:`setup`. + DataLoader + Non-shuffled DataLoader over the validation split, or an empty + ``DataLoader`` if no split was formed. """ if self._val_dataset is None: - return None + return DataLoader([], batch_size=self.batch_size) # pyright: ignore[reportArgumentType] return DataLoader( self._val_dataset, batch_size=self.batch_size, diff --git a/tests/test_auxiliary_encoder.py b/tests/test_auxiliary_encoder.py index 585ce5f..6dc559e 100644 --- a/tests/test_auxiliary_encoder.py +++ b/tests/test_auxiliary_encoder.py @@ -216,9 +216,9 @@ def test_train_batch_shapes_match_task_count(self): assert mask.dtype == torch.bool def test_too_few_records_uses_all_for_training(self): - """When the record pool is too small for a val split, val_dataloader must return None.""" + """When the record pool is too small for a val split, val_dataloader must be empty.""" records = _records_with_readouts()[:1] dm = AuxiliaryDataModule(records, ["log2fc_1um"], val_fraction=0.1) dm.setup() - assert dm.val_dataloader() is None + assert len(dm.val_dataloader()) == 0 From 48f3886fe244d8512cd7c36da459165af4426ddd Mon Sep 17 00:00:00 2001 From: Sean Colby Date: Mon, 27 Jul 2026 09:12:21 -0800 Subject: [PATCH 12/18] Rename AuxiliaryEncoderConfig to AuxiliaryModelConfig; add parallel auxiliary_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. --- moal/auxiliary_encoder.py | 29 ++++++++++++------- moal/cli.py | 50 +++++++++++++++++++++++++++------ moal/config.py | 45 +++++++++++++++++++---------- moal/types.py | 2 +- tests/test_auxiliary_encoder.py | 13 ++++----- tests/test_config.py | 43 ++++++++++++++++++++++------ 6 files changed, 132 insertions(+), 50 deletions(-) diff --git a/moal/auxiliary_encoder.py b/moal/auxiliary_encoder.py index eb2a779..d8a77fc 100644 --- a/moal/auxiliary_encoder.py +++ b/moal/auxiliary_encoder.py @@ -7,7 +7,7 @@ construction (:func:`moal.model.build_mpnn`) rather than a bespoke architecture, so its embeddings live in the same representation space as the main pEC50 model. Readouts are used as-is: no per-plate/per-batch -normalization is applied (see :class:`~moal.config.AuxiliaryEncoderConfig` +normalization is applied (see :class:`~moal.config.AuxiliaryModelConfig` for why). """ @@ -28,7 +28,7 @@ from torch.optim import Adam from torch.utils.data import DataLoader, Dataset, random_split -from moal.config import AuxiliaryEncoderConfig +from moal.config import AuxiliaryModelConfig from moal.model import build_mpnn from moal.types import LabelRecord @@ -286,12 +286,12 @@ class AuxiliaryEncoderModule(L.LightningModule): task_names : list[str] Fixed, ordered list of readout keys the model was (or will be) trained against; determines the predictor head's output width. - config : AuxiliaryEncoderConfig + config : AuxiliaryModelConfig Backbone architecture, freeze schedule, and optimization hyperparameters. """ - def __init__(self, task_names: list[str], config: AuxiliaryEncoderConfig) -> None: + def __init__(self, task_names: list[str], config: AuxiliaryModelConfig) -> None: super().__init__() if not task_names: raise ValueError("task_names must be non-empty.") @@ -482,7 +482,7 @@ def embed_smiles(self, smiles_list: list[str], batch_size: int = 256) -> np.ndar ``smiles_list``. ``embedding_dim`` is the backbone's native output width (CheMeleon's fixed width, or ``message_hidden_dim`` for a random-init encoder), not - ``AuxiliaryEncoderConfig.embedding_dim``. + ``AuxiliaryModelConfig.embedding_dim``. """ dataset = MoleculeDataset([MoleculeDatapoint.from_smi(s) for s in smiles_list]) # pyright: ignore[reportArgumentType] dataloader = build_dataloader( @@ -501,7 +501,8 @@ def embed_smiles(self, smiles_list: list[str], batch_size: int = 256) -> np.ndar def pretrain_auxiliary_encoder( records: list[LabelRecord], - config: AuxiliaryEncoderConfig, + config: AuxiliaryModelConfig, + max_epochs: int = 30, trainer_kwargs: dict[str, Any] | None = None, datamodule_kwargs: dict[str, Any] | None = None, ) -> AuxiliaryEncoderModule: @@ -519,8 +520,11 @@ def pretrain_auxiliary_encoder( All labeled records from the campaign state. Records with an empty ``raw_ps_readouts`` are filtered out before training; they carry no auxiliary training signal. - config : AuxiliaryEncoderConfig + config : AuxiliaryModelConfig Backbone, freeze schedule, and optimization hyperparameters. + max_epochs : int, optional + Number of pretraining epochs; overridden by ``trainer_kwargs["max_epochs"]`` + when present (e.g. from ``PipelineConfig.auxiliary_trainer``). Default is 30. trainer_kwargs : dict[str, Any], optional Additional keyword arguments forwarded to ``lightning.Trainer``. Ignored when loading from ``config.checkpoint_path``. @@ -561,9 +565,14 @@ def pretrain_auxiliary_encoder( dm.setup() kwargs: dict[str, Any] = { - "max_epochs": config.max_epochs, + "max_epochs": max_epochs, "enable_progress_bar": False, "enable_model_summary": False, + # Small readout-bearing pools can easily have fewer than Lightning's + # default log_every_n_steps=50 batches per epoch, which would silently + # suppress all step-level logs (same rationale as TrainerConfig's + # log_every_n_steps default for the main model). + "log_every_n_steps": 1, } if trainer_kwargs: kwargs.update(trainer_kwargs) @@ -588,7 +597,7 @@ def save_auxiliary_encoder_checkpoint(module: AuxiliaryEncoderModule, path: str def load_auxiliary_encoder_checkpoint( - path: str | Path, config: AuxiliaryEncoderConfig + path: str | Path, config: AuxiliaryModelConfig ) -> AuxiliaryEncoderModule: """Load an auxiliary encoder checkpoint written by :func:`save_auxiliary_encoder_checkpoint`. @@ -596,7 +605,7 @@ def load_auxiliary_encoder_checkpoint( ---------- path : str or Path Path to the checkpoint file. - config : AuxiliaryEncoderConfig + config : AuxiliaryModelConfig Backbone architecture the checkpoint was trained with; must match the checkpoint's own architecture (``from_foundation``, ``ffn_hidden_dim``, etc.) or ``load_state_dict`` will raise. diff --git a/moal/cli.py b/moal/cli.py index 46fdeaf..a93fce4 100644 --- a/moal/cli.py +++ b/moal/cli.py @@ -19,6 +19,7 @@ import lightning as L import numpy as np import pandas as pd +from lightning.pytorch.loggers import CSVLogger from rich.console import Console from rich.progress import ( BarColumn, @@ -403,28 +404,59 @@ def plan(config: Path, output_dir: Path | None, verbose: bool) -> None: smi for _, smi in state.ps_upgrade_rows ] - if cfg.auxiliary_encoder is not None: + if cfg.auxiliary_model is not None: # Concatenation architecture (issue #36 Phase 2): pretrain the - # auxiliary encoder on this run's raw_ps_readouts, then train - # the main model with its structural embedding concatenated in - # for compounds that were never PS-screened. + # auxiliary encoder on this run's raw_ps_readouts (full PS + DRC + # pool), then train the main model on DRC records only. PS records + # have already contributed what they can via the frozen auxiliary + # embedding/readout; re-supervising the main Tobit loss with PS's + # noisier LEFT/INTERVAL labels on top of that would both duplicate + # signal already distilled into the auxiliary encoder and dilute + # the embedding-path training examples (the only ones representative + # of never-screened inference targets) beneath the much larger + # observed-readout-path population. progress.update( task, description="[yellow]Pretraining auxiliary encoder[/yellow]" ) + aux_logger = CSVLogger(save_dir=str(out_dir), name="aux_encoder_logs") try: - aux_encoder = pretrain_auxiliary_encoder(fit_records, cfg.auxiliary_encoder) + aux_encoder = pretrain_auxiliary_encoder( + fit_records, + cfg.auxiliary_model, + trainer_kwargs={ + **cfg.auxiliary_trainer.to_dict(), + "logger": aux_logger, + }, + datamodule_kwargs=cfg.auxiliary_trainer.to_datamodule_kwargs(), + ) except ValueError as exc: raise click.ClickException(str(exc)) from exc - - progress.update(task, description=retraining_description) + logger.info("Auxiliary encoder loss curves written to %s", aux_logger.log_dir) + + drc_records = [ + rec for rec in fit_records if rec.fidelity == QueryType.DOSE_RESPONSE + ] + if not drc_records: + raise click.ClickException( + "No DRC records available to train the main model; the " + "concatenation architecture requires at least one DRC-labeled " + "compound after auxiliary encoder pretraining." + ) + concat_description = ( + f"[yellow]Training model[/yellow] — {len(drc_records)} DRC records " + "(PS records used only for auxiliary encoder pretraining)" + ) + progress.update(task, description=concat_description) model = _build_concatenation_model(cfg, aux_encoder) + main_logger = CSVLogger(save_dir=str(out_dir), name="main_model_logs") model.refit( - fit_records, + drc_records, aux_encoder=aux_encoder, - trainer_kwargs=cfg.trainer.to_dict(), + trainer_kwargs={**cfg.trainer.to_dict(), "logger": main_logger}, datamodule_kwargs=cfg.trainer.to_datamodule_kwargs(), output_dir=out_dir, ) + logger.info("Main model loss curves written to %s", main_logger.log_dir) progress.advance(task) progress.update(task, description=scoring_description) diff --git a/moal/config.py b/moal/config.py index 851883a..e26aa1f 100644 --- a/moal/config.py +++ b/moal/config.py @@ -123,8 +123,8 @@ class ModelConfig: @dataclass(frozen=True) -class AuxiliaryEncoderConfig: - """Auxiliary encoder for the primary-screen readouts in ``LabelRecord.raw_ps_readouts``. +class AuxiliaryModelConfig: + """Auxiliary encoder architecture for the ``LabelRecord.raw_ps_readouts`` signal. ``moal plan``-only (see the ``moal simulate`` exclusion in the module docstring reference, issue #36). Off by default: ``moal plan`` behaves @@ -175,9 +175,6 @@ class AuxiliaryEncoderConfig: main model's ``mpnn_lr`` / ``ffn_lr`` split). weight_decay : float L2 weight decay applied to all trainable parameters. - max_epochs : int - Number of pretraining epochs, scheduled independently from the main - model's ``TrainerConfig.max_epochs``. embedding_dim : int Dimensionality of the pooled molecular embedding exposed to the main model's concatenation architecture (Phase 2). Ignored by the @@ -190,6 +187,18 @@ class AuxiliaryEncoderConfig: on every ``moal plan`` invocation using the current campaign-state CSV's ``raw_ps_readouts``, so newly accumulated readouts improve the next run automatically. + use_observed_readout : bool + Controls the concatenation architecture's main-model input, not the + auxiliary encoder's own pretraining (which always uses whatever + ``raw_ps_readouts`` exist, regardless of this flag). The auxiliary + encoder's structural embedding is always concatenated into the main + model for every compound. When True (default), a compound with an + observed readout *additionally* gets its raw value concatenated + alongside the embedding. When False, every compound is scored from + its embedding alone and the readout/mask blocks stay zero for all + compounds; a constant-zero input column is a mathematical no-op for + a plain linear layer (zero gradient, zero forward contribution), so + this does not degrade model capacity. """ from_foundation: str | bool = "chemeleon" @@ -200,9 +209,9 @@ class AuxiliaryEncoderConfig: freeze_epochs: int = 5 lr: float = 1e-4 weight_decay: float = 0.0 - max_epochs: int = 30 embedding_dim: int = 300 checkpoint_path: str | None = None + use_observed_readout: bool = True @dataclass(frozen=True) @@ -582,10 +591,14 @@ class PipelineConfig: Command-specific dataset and I/O settings. active_learning_loop : ActiveLearningLoopConfig Parameters controlling the active learning iteration loop. - auxiliary_encoder : AuxiliaryEncoderConfig or None - Optional auxiliary log2FC/pIC50 encoder for ``moal plan`` (issue #36). - ``None`` (default) disables the feature entirely; ``moal plan`` - behaves exactly as it does without this config. + auxiliary_model : AuxiliaryModelConfig or None + Optional auxiliary log2FC/pIC50 encoder architecture for ``moal plan`` + (issue #36). ``None`` (default) disables the feature entirely; + ``moal plan`` behaves exactly as it does without this config. + auxiliary_trainer : TrainerConfig + Keyword arguments forwarded to ``lightning.Trainer`` during auxiliary + encoder pretraining, scheduled independently from the main model's + ``trainer``. Unused when ``auxiliary_model`` is None. seed : int Global random seed for the campaign. """ @@ -597,7 +610,8 @@ class PipelineConfig: dashboard: DashboardConfig = field(default_factory=DashboardConfig) data: DataConfig = field(default_factory=DataConfig) active_learning_loop: ActiveLearningLoopConfig = field(default_factory=ActiveLearningLoopConfig) - auxiliary_encoder: AuxiliaryEncoderConfig | None = None + auxiliary_model: AuxiliaryModelConfig | None = None + auxiliary_trainer: TrainerConfig = field(default_factory=TrainerConfig) seed: int = 42 @@ -620,7 +634,7 @@ def from_yaml(cls, path: str | Path) -> PipelineConfig: data_raw = raw.get("data", {}) simulate_raw = data_raw.get("simulate", {}) pretrain_raw = simulate_raw.pop("pretrain", {}) if isinstance(simulate_raw, dict) else {} - auxiliary_encoder_raw = raw.get("auxiliary_encoder", None) + auxiliary_model_raw = raw.get("auxiliary_model", None) return cls( oracle=OracleConfig(**raw.get("oracle", {})), model=ModelConfig(**raw.get("model", {})), @@ -636,11 +650,12 @@ def from_yaml(cls, path: str | Path) -> PipelineConfig: plan=PlanDataConfig(**data_raw.get("plan", {})), ), active_learning_loop=ActiveLearningLoopConfig(**raw.get("active_learning_loop", {})), - auxiliary_encoder=( - AuxiliaryEncoderConfig(**auxiliary_encoder_raw) - if auxiliary_encoder_raw is not None + auxiliary_model=( + AuxiliaryModelConfig(**auxiliary_model_raw) + if auxiliary_model_raw is not None else None ), + auxiliary_trainer=TrainerConfig(**raw.get("auxiliary_trainer", {})), seed=raw.get("seed", 42), ) diff --git a/moal/types.py b/moal/types.py index 9d40f09..5d80560 100644 --- a/moal/types.py +++ b/moal/types.py @@ -86,7 +86,7 @@ class LabelRecord: keyed by source column name, independent of the LEFT/INTERVAL censoring derived from ``value`` against ``oracle.ps_threshold``. Retained on the surviving DRC record after a PS-to-DRC upgrade so the - auxiliary encoder (see ``AuxiliaryEncoderConfig``) can still use it. + auxiliary encoder (see ``AuxiliaryModelConfig``) can still use it. Empty when the compound has never been PS-screened. """ diff --git a/tests/test_auxiliary_encoder.py b/tests/test_auxiliary_encoder.py index 6dc559e..88319c2 100644 --- a/tests/test_auxiliary_encoder.py +++ b/tests/test_auxiliary_encoder.py @@ -18,7 +18,7 @@ pretrain_auxiliary_encoder, save_auxiliary_encoder_checkpoint, ) -from moal.config import AuxiliaryEncoderConfig +from moal.config import AuxiliaryModelConfig from moal.types import CensoringType, LabelRecord, QueryType _SMILES = ["CCO", "CCN", "CCC", "c1ccccc1", "CCCl", "CCBr", "CCOCC", "CCCC"] @@ -53,17 +53,16 @@ def _batch(smiles_list: list[str]) -> BatchMolGraph: return BatchMolGraph([dataset[i].mg for i in range(len(dataset))]) -def _fast_config(**overrides) -> AuxiliaryEncoderConfig: +def _fast_config(**overrides) -> AuxiliaryModelConfig: defaults = { "from_foundation": False, "message_hidden_dim": 16, "ffn_hidden_dim": 16, "depth": 1, "freeze_epochs": 0, - "max_epochs": 1, } defaults.update(overrides) - return AuxiliaryEncoderConfig(**defaults) + return AuxiliaryModelConfig(**defaults) class TestMaskedMSELoss: @@ -136,7 +135,7 @@ def test_trains_and_returns_module_with_expected_tasks(self): records = _records_with_readouts() config = _fast_config() - module = pretrain_auxiliary_encoder(records, config) + module = pretrain_auxiliary_encoder(records, config, max_epochs=1) assert module.task_names == ["log2fc_1um", "pic50"] @@ -159,7 +158,7 @@ def test_checkpoint_path_skips_retraining(self, tmp_path, monkeypatch): """When checkpoint_path is set, pretrain_auxiliary_encoder must load the checkpoint rather than training.""" records = _records_with_readouts() config = _fast_config() - trained = pretrain_auxiliary_encoder(records, config) + trained = pretrain_auxiliary_encoder(records, config, max_epochs=1) ckpt_path = tmp_path / "aux_encoder.pt" save_auxiliary_encoder_checkpoint(trained, ckpt_path) @@ -181,7 +180,7 @@ def test_round_trips_weights_and_task_names(self, tmp_path): """A saved-then-loaded checkpoint must reproduce identical predictions and task_names.""" records = _records_with_readouts() config = _fast_config() - trained = pretrain_auxiliary_encoder(records, config) + trained = pretrain_auxiliary_encoder(records, config, max_epochs=1) trained.eval() path = tmp_path / "aux_encoder.pt" diff --git a/tests/test_config.py b/tests/test_config.py index 57e3f83..0e90af7 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -6,7 +6,7 @@ import yaml -from moal.config import AuxiliaryEncoderConfig, PipelineConfig +from moal.config import AuxiliaryModelConfig, PipelineConfig, TrainerConfig def _write_yaml(tmp_path: Path, raw: dict) -> Path: @@ -16,23 +16,23 @@ def _write_yaml(tmp_path: Path, raw: dict) -> Path: return path -class TestAuxiliaryEncoderConfig: - """Tests for AuxiliaryEncoderConfig's default-off behavior and round-trip through from_yaml.""" +class TestAuxiliaryModelConfig: + """Tests for AuxiliaryModelConfig's default-off behavior and round-trip through from_yaml.""" def test_defaults_to_none_when_absent(self, tmp_path): - """auxiliary_encoder must be None when the YAML has no auxiliary_encoder key.""" + """auxiliary_model must be None when the YAML has no auxiliary_model key.""" path = _write_yaml(tmp_path, {"seed": 1}) cfg = PipelineConfig.from_yaml(path) - assert cfg.auxiliary_encoder is None + assert cfg.auxiliary_model is None def test_round_trips_through_from_yaml(self, tmp_path): - """An explicit auxiliary_encoder block must populate a matching AuxiliaryEncoderConfig.""" + """An explicit auxiliary_model block must populate a matching AuxiliaryModelConfig.""" path = _write_yaml( tmp_path, { - "auxiliary_encoder": { + "auxiliary_model": { "freeze_epochs": 3, "embedding_dim": 128, "checkpoint_path": "aux_encoder.pt", @@ -42,6 +42,33 @@ def test_round_trips_through_from_yaml(self, tmp_path): cfg = PipelineConfig.from_yaml(path) - assert cfg.auxiliary_encoder == AuxiliaryEncoderConfig( + assert cfg.auxiliary_model == AuxiliaryModelConfig( freeze_epochs=3, embedding_dim=128, checkpoint_path="aux_encoder.pt" ) + + +class TestAuxiliaryTrainerConfig: + """Tests for auxiliary_trainer's default and round-trip through from_yaml.""" + + def test_defaults_to_trainer_config_defaults_when_absent(self, tmp_path): + """auxiliary_trainer must default to TrainerConfig() when the YAML has no auxiliary_trainer key.""" + path = _write_yaml(tmp_path, {"seed": 1}) + + cfg = PipelineConfig.from_yaml(path) + + assert cfg.auxiliary_trainer == TrainerConfig() + + def test_round_trips_through_from_yaml(self, tmp_path): + """An explicit auxiliary_trainer block must populate a matching TrainerConfig, independent of trainer.""" + path = _write_yaml( + tmp_path, + { + "trainer": {"max_epochs": 30, "val_fraction": 0.0}, + "auxiliary_trainer": {"max_epochs": 15, "val_fraction": 0.2}, + }, + ) + + cfg = PipelineConfig.from_yaml(path) + + assert cfg.auxiliary_trainer == TrainerConfig(max_epochs=15, val_fraction=0.2) + assert cfg.trainer == TrainerConfig(max_epochs=30, val_fraction=0.0) From 26efd171ae7ae411098afb9f33f97034586c82fd Mon Sep 17 00:00:00 2001 From: Sean Colby Date: Mon, 27 Jul 2026 09:13:00 -0800 Subject: [PATCH 13/18] Always compute auxiliary embedding; make use_observed_readout constructor-only MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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.) --- moal/cli.py | 12 ++- moal/concatenation_model.py | 120 +++++++++++++++++++++--------- tests/test_cli.py | 11 ++- tests/test_concatenation_model.py | 53 ++++++++++--- 4 files changed, 147 insertions(+), 49 deletions(-) diff --git a/moal/cli.py b/moal/cli.py index a93fce4..d967230 100644 --- a/moal/cli.py +++ b/moal/cli.py @@ -851,6 +851,9 @@ def _build_concatenation_model( cfg : PipelineConfig Active campaign configuration. Reuses ``cfg.model``'s backbone and optimization hyperparameters, same as :func:`_build_plan_model`. + ``cfg.model.w_drc``/``w_ps`` are not forwarded: the concatenation + architecture always trains on DRC records only (see the ``plan`` + command), so that fidelity weighting has nothing to differentiate. aux_encoder : AuxiliaryEncoderModule Pretrained auxiliary encoder; supplies ``task_names`` and ``embedding_dim`` to size the concatenation feature width. @@ -859,10 +862,17 @@ def _build_concatenation_model( ------- ConcatenationChemPropLightningModule Configured model ready for ``refit()`` and ``predict_smiles()``. + ``use_observed_readout`` is fixed here at construction (from + ``cfg.auxiliary_model.use_observed_readout``) rather than passed + separately to each call, so training and inference routing cannot + drift apart. """ + if cfg.auxiliary_model is None: + raise ValueError("_build_concatenation_model requires cfg.auxiliary_model to be set") feature_dim = concatenation_feature_dim(len(aux_encoder.task_names), aux_encoder.embedding_dim) return ConcatenationChemPropLightningModule( concat_feature_dim=feature_dim, + use_observed_readout=cfg.auxiliary_model.use_observed_readout, ffn_hidden_dim=cfg.model.ffn_hidden_dim, ffn_num_layers=cfg.model.ffn_num_layers, message_hidden_dim=cfg.model.message_hidden_dim, @@ -873,8 +883,6 @@ def _build_concatenation_model( mpnn_weight_decay=cfg.model.mpnn_weight_decay, ffn_weight_decay=cfg.model.ffn_weight_decay, sigma=cfg.model.sigma, - w_drc=cfg.model.w_drc, - w_ps=cfg.model.w_ps, learnable_sigma=cfg.model.learnable_sigma, from_foundation=cfg.model.from_foundation, ) diff --git a/moal/concatenation_model.py b/moal/concatenation_model.py index 25b5f43..c311225 100644 --- a/moal/concatenation_model.py +++ b/moal/concatenation_model.py @@ -38,7 +38,7 @@ from moal.loss import CensoredRegressionLoss from moal.model import _validate_from_foundation, build_mpnn from moal.planning import normalize_record_weights -from moal.types import LabelRecord +from moal.types import LabelRecord, QueryType logger = logging.getLogger(__name__) @@ -68,17 +68,26 @@ def build_concatenation_features( canonical_smiles: list[str], readouts: list[dict[str, float]], aux_encoder: AuxiliaryEncoderModule, + *, + use_observed_readout: bool = True, batch_size: int = 256, ) -> np.ndarray: """Build the per-compound concatenation feature matrix. - For each compound: if ``readouts[i]`` is non-empty, the observed-readout - block is populated (per-task values where present, zero elsewhere) and - the mask block marks which tasks were actually observed; the embedding - block stays zero and the provenance flag is 0. If ``readouts[i]`` is - empty, the observed-readout and mask blocks stay zero, the embedding - block holds the auxiliary encoder's structural embedding for that - compound, and the provenance flag is 1. + The auxiliary encoder's structural embedding is always computed and + included, for every compound, regardless of whether it has an observed + readout — this keeps the main model's embedding-consuming input pathway + exercised by every training row, so training and inference always share + the same input distribution shape (a compound with an observed readout + at training time looks input-wise like any other compound, differing + only in whether its readout block is also populated). + + When ``use_observed_readout`` is True and ``readouts[i]`` is non-empty, + the observed-readout block is additionally populated (per-task values + where present, zero elsewhere), the mask block marks which tasks were + observed, and the readout-used flag is 1. Otherwise the readout and mask + blocks stay zero and the flag is 0 — the compound is scored from its + structural embedding alone. Parameters ---------- @@ -89,10 +98,15 @@ def build_concatenation_features( with ``canonical_smiles``. An empty dict means "never PS-screened". aux_encoder : AuxiliaryEncoderModule Pretrained auxiliary encoder; supplies both ``task_names`` (readout - key order) and the structural embedding fallback. + key order) and the structural embedding. + use_observed_readout : bool, optional + When True (default), compounds with a non-empty readout also get + their raw value concatenated alongside the embedding. When False, + the readout and mask blocks are always zero and every compound is + scored from its embedding alone, matching + ``AuxiliaryModelConfig.use_observed_readout``. batch_size : int, optional - Batch size for the embedding forward pass over compounds lacking - readouts. Default is 256. + Batch size for the embedding forward pass. Default is 256. Returns ------- @@ -117,28 +131,21 @@ def build_concatenation_features( readout_vec = np.zeros((n, n_tasks), dtype=np.float32) readout_mask = np.zeros((n, n_tasks), dtype=np.float32) - embedding_used = np.zeros((n, 1), dtype=np.float32) + readout_used = np.zeros((n, 1), dtype=np.float32) - embed_indices: list[int] = [] - embed_smiles: list[str] = [] - for i, readout in enumerate(readouts): - if readout: + if use_observed_readout: + for i, readout in enumerate(readouts): + if not readout: + continue for j, name in enumerate(task_names): if name in readout: readout_vec[i, j] = readout[name] readout_mask[i, j] = 1.0 - else: - embedding_used[i, 0] = 1.0 - embed_indices.append(i) - embed_smiles.append(canonical_smiles[i]) + readout_used[i, 0] = 1.0 - embeddings = np.zeros((n, aux_encoder.embedding_dim), dtype=np.float32) - if embed_smiles: - computed = aux_encoder.embed_smiles(embed_smiles, batch_size=batch_size) - for idx, row in zip(embed_indices, computed, strict=True): - embeddings[idx] = row + embeddings = aux_encoder.embed_smiles(canonical_smiles, batch_size=batch_size) - return np.concatenate([readout_vec, readout_mask, embeddings, embedding_used], axis=1) + return np.concatenate([readout_vec, readout_mask, embeddings, readout_used], axis=1) class _ConcatenatedDataset(Dataset): @@ -224,15 +231,32 @@ class ConcatenationChemPropLightningModule(L.LightningModule): Width of the concatenation feature vector (see :func:`concatenation_feature_dim`); determines the predictor head's input width alongside the backbone's own pooled-embedding width. + use_observed_readout : bool, optional + Fixed at construction and used by both :meth:`refit` and + :meth:`predict_smiles` — this determines the input distribution the + model's weights are actually fit against (e.g. when False, the + readout/mask input dimensions are always exactly zero throughout + training, so their weights never receive gradient signal; feeding + them nonzero values at inference would exercise untrained weights). + Deliberately not a per-call parameter on either method, so + training-time and inference-time routing cannot drift apart. Default + is True. ffn_hidden_dim, ffn_num_layers, message_hidden_dim, depth, freeze_epochs, - mpnn_lr, ffn_lr, mpnn_weight_decay, ffn_weight_decay, sigma, w_drc, w_ps, + mpnn_lr, ffn_lr, mpnn_weight_decay, ffn_weight_decay, sigma, learnable_sigma, from_foundation - See :class:`moal.model.ChemPropLightningModule`. + See :class:`moal.model.ChemPropLightningModule`. Note ``w_drc``/``w_ps`` + are deliberately absent here: :meth:`refit` requires every record to + be DOSE_RESPONSE (see below), so the DRC-vs-PS fidelity weighting + that parameter pair controls in + :class:`moal.model.ChemPropLightningModule` has nothing to + differentiate in this class and would be a redundant, no-op scalar + confounded with ``sigma``. """ def __init__( self, concat_feature_dim: int, + use_observed_readout: bool = True, ffn_hidden_dim: int = 300, ffn_num_layers: int = 2, message_hidden_dim: int = 300, @@ -243,8 +267,6 @@ def __init__( mpnn_weight_decay: float = 0.0, ffn_weight_decay: float = 0.0, sigma: float = 0.5, - w_drc: float = 1.0, - w_ps: float = 0.3, learnable_sigma: bool = False, from_foundation: str | bool = "chemeleon", ) -> None: @@ -252,6 +274,7 @@ def __init__( _validate_from_foundation(from_foundation) self._from_foundation = from_foundation self.concat_feature_dim = concat_feature_dim + self.use_observed_readout = use_observed_readout self.save_hyperparameters() self.freeze_epochs = freeze_epochs @@ -268,8 +291,12 @@ def __init__( "val_ps": [], } + # w_drc/w_ps are fixed equal (not exposed as parameters): refit() + # requires every record to be DOSE_RESPONSE, so the PS branch never + # fires and the two weights would otherwise be a redundant, no-op + # scalar confounded with sigma. self.loss_fn = CensoredRegressionLoss( - sigma=sigma, w_drc=w_drc, w_ps=w_ps, learnable_sigma=learnable_sigma + sigma=sigma, w_drc=1.0, w_ps=1.0, learnable_sigma=learnable_sigma ) self.model = build_mpnn( @@ -461,13 +488,11 @@ def predict_smiles( **Must be RDKit-canonical, salt-stripped SMILES**; see :meth:`moal.model.ChemPropLightningModule.predict_smiles`. readouts : list[dict[str, float]] - Per-compound observed readouts, aligned with ``smiles_list``; an - empty dict routes that compound through the auxiliary encoder's - structural embedding. Forwarded to - :func:`build_concatenation_features`. + Per-compound observed readouts, aligned with ``smiles_list``. + Forwarded to :func:`build_concatenation_features`. aux_encoder : AuxiliaryEncoderModule Pretrained auxiliary encoder supplying both the readout-key - order and the structural-embedding fallback. + order and the structural embedding. batch_size : int, optional Number of molecules processed per forward pass. Default is 256. @@ -478,7 +503,11 @@ def predict_smiles( ``smiles_list``. """ features = build_concatenation_features( - smiles_list, readouts, aux_encoder, batch_size=batch_size + smiles_list, + readouts, + aux_encoder, + use_observed_readout=self.use_observed_readout, + batch_size=batch_size, ) x_d = torch.as_tensor(features, dtype=torch.float32) @@ -541,12 +570,29 @@ def refit( ------- ConcatenationChemPropLightningModule self (for chaining). + + Raises + ------ + ValueError + If any record's fidelity is not ``QueryType.DOSE_RESPONSE``. The + loss weighting is fixed equal for DRC/PS (see class docstring), + so PS records would be silently mis-weighted rather than + differentiated; callers should train the concatenation + architecture on DRC records only. """ + non_drc = [rec for rec in records if rec.fidelity != QueryType.DOSE_RESPONSE] + if non_drc: + raise ValueError( + f"ConcatenationChemPropLightningModule.refit() received {len(non_drc)} " + "non-DOSE_RESPONSE record(s); this class trains on DRC records only " + "(see class docstring for why PS/DRC loss weighting is unsupported here)." + ) records = normalize_record_weights(records) features = build_concatenation_features( [rec.canonical_smiles for rec in records], [rec.raw_ps_readouts for rec in records], aux_encoder, + use_observed_readout=self.use_observed_readout, ) dm = _ConcatenatedDataModule(records, features, **(datamodule_kwargs or {})) dm.setup() diff --git a/tests/test_cli.py b/tests/test_cli.py index 63f5a93..7c2daef 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -14,6 +14,7 @@ import moal.cli as cli from moal.cli import main from moal.config import PipelineConfig +from moal.types import QueryType def _result_text(result) -> str: @@ -397,7 +398,7 @@ def test_plan_uses_concatenation_architecture_when_auxiliary_encoder_configured( " max_epochs: 1\n" "dashboard:\n" " enabled: false\n" - "auxiliary_encoder:\n" + "auxiliary_model:\n" " freeze_epochs: 0\n" ) @@ -426,6 +427,14 @@ def test_plan_uses_concatenation_architecture_when_auxiliary_encoder_configured( concat_model.refit.assert_called_once() assert concat_model.refit.call_args.kwargs["aux_encoder"] is fake_aux_encoder + # The main model trains on DRC records only; the PS record (CCO, >=) must be + # excluded even though pretrain_auxiliary_encoder saw it via fit_records + refit_records = concat_model.refit.call_args.args[0] + assert len(refit_records) == 1 + assert refit_records[0].fidelity == QueryType.DOSE_RESPONSE + pretrain_records = pretrain_mock.call_args.args[0] + assert any(rec.fidelity == QueryType.PRIMARY_SCREEN for rec in pretrain_records) + # predict_smiles must receive per-compound readouts: empty for unqueried, # the observed reading for the PS-upgrade candidate call_args = concat_model.predict_smiles.call_args diff --git a/tests/test_concatenation_model.py b/tests/test_concatenation_model.py index f113081..a212aea 100644 --- a/tests/test_concatenation_model.py +++ b/tests/test_concatenation_model.py @@ -18,7 +18,7 @@ build_concatenation_features, concatenation_feature_dim, ) -from moal.config import AuxiliaryEncoderConfig +from moal.config import AuxiliaryModelConfig from moal.types import CensoringType, LabelRecord, QueryType _EMBEDDING_DIM = 16 @@ -27,7 +27,7 @@ @pytest.fixture def aux_encoder() -> AuxiliaryEncoderModule: """Small random-init auxiliary encoder with two tasks.""" - config = AuxiliaryEncoderConfig( + config = AuxiliaryModelConfig( from_foundation=False, message_hidden_dim=_EMBEDDING_DIM, ffn_hidden_dim=16, @@ -80,8 +80,8 @@ def test_matches_2n_plus_embedding_plus_1(self): class TestBuildConcatenationFeatures: """Tests for build_concatenation_features: observed vs embedding routing and shape.""" - def test_observed_readout_routes_to_readout_block_not_embedding(self, aux_encoder): - """A compound with a non-empty readout dict must populate the readout/mask block and leave the embedding block zero, with flag=0.""" + def test_observed_readout_also_gets_embedding(self, aux_encoder): + """A compound with a readout must populate readout/mask AND the embedding block, with flag=1.""" features = build_concatenation_features(["CCO"], [{"log2fc_1um": 2.5}], aux_encoder) n_tasks = 2 @@ -92,11 +92,11 @@ def test_observed_readout_routes_to_readout_block_not_embedding(self, aux_encode assert readout_block[0] == 2.5 assert list(mask_block) == [1.0, 0.0] - assert np.all(embedding_block == 0.0) - assert flag == 0.0 + assert not np.all(embedding_block == 0.0) + assert flag == 1.0 - def test_missing_readout_routes_to_embedding_block(self, aux_encoder): - """A compound with an empty readout dict must leave the readout/mask block zero and populate the embedding block, with flag=1.""" + def test_missing_readout_uses_embedding_only(self, aux_encoder): + """A compound with an empty readout dict must leave the readout/mask block zero, populate the embedding block, and flag=0.""" features = build_concatenation_features(["CCO"], [{}], aux_encoder) n_tasks = 2 @@ -106,7 +106,24 @@ def test_missing_readout_routes_to_embedding_block(self, aux_encoder): assert np.all(readout_mask_block == 0.0) assert not np.all(embedding_block == 0.0) - assert flag == 1.0 + assert flag == 0.0 + + def test_use_observed_readout_false_zeroes_readout_block_but_keeps_embedding(self, aux_encoder): + """With use_observed_readout=False, every compound is embedding-only regardless of its own readout data.""" + with_readout = build_concatenation_features( + ["CCO"], [{"log2fc_1um": 2.5}], aux_encoder, use_observed_readout=False + ) + without_readout = build_concatenation_features( + ["CCO"], [{}], aux_encoder, use_observed_readout=False + ) + n_tasks = 2 + + assert np.all(with_readout[0, : 2 * n_tasks] == 0.0) + assert with_readout[0, -1] == 0.0 + np.testing.assert_allclose( + with_readout[0, 2 * n_tasks : 2 * n_tasks + _EMBEDDING_DIM], + without_readout[0, 2 * n_tasks : 2 * n_tasks + _EMBEDDING_DIM], + ) def test_output_shape_matches_concatenation_feature_dim(self, aux_encoder): """Output width must equal concatenation_feature_dim(n_tasks, embedding_dim).""" @@ -167,3 +184,21 @@ def test_predict_smiles_chunks_correctly_across_batch_boundary(self, aux_encoder preds = model.predict_smiles(smiles, readouts, aux_encoder, batch_size=2) assert preds.shape == (5,) + + def test_refit_rejects_non_drc_records(self, aux_encoder): + """refit() must reject any record whose fidelity is not DOSE_RESPONSE.""" + feat_dim = concatenation_feature_dim(2, _EMBEDDING_DIM) + model = _fast_model(feat_dim) + ps_record = LabelRecord( + smiles="CCO", + canonical_smiles="CCO", + value=5.0, + upper_bound=11.0, + censoring_type=CensoringType.INTERVAL, + fidelity=QueryType.PRIMARY_SCREEN, + cost=1.0, + iteration=0, + ) + + with pytest.raises(ValueError, match="DOSE_RESPONSE"): + model.refit([ps_record], aux_encoder=aux_encoder, max_epochs=1) From 460b43dfc61a479329cec8820fbf405b1b52e2c9 Mon Sep 17 00:00:00 2001 From: Sean Colby Date: Mon, 27 Jul 2026 09:13:24 -0800 Subject: [PATCH 14/18] Write predicted_pec50 to moal plan's annotated output 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. --- moal/planning.py | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/moal/planning.py b/moal/planning.py index f0692e9..9a34689 100644 --- a/moal/planning.py +++ b/moal/planning.py @@ -429,6 +429,8 @@ def annotate_campaign_state( - ``embedding_derived`` — True where ``provenance`` flagged the prediction as embedding-derived (see ``provenance`` below); NaN for training-only rows; always False when ``provenance`` is None + - ``predicted_pec50`` — the raw model prediction from ``predictions``, + unmodified by acquisition scoring; NaN for training-only rows Parameters ---------- @@ -477,6 +479,7 @@ def annotate_campaign_state( result["overall_score"] = np.nan result["recommendation"] = None # Object dtype so string values can be assigned result["embedding_derived"] = None # Object dtype so bool values can be assigned + result["predicted_pec50"] = np.nan n_unqueried = len(state.unqueried_rows) unqueried_preds = predictions[:n_unqueried] @@ -490,7 +493,9 @@ def annotate_campaign_state( summaries = acquisition.score_summary( unqueried_canonical, unqueried_preds, provenance=unqueried_provenance ) - for (row_idx, _), summary in zip(state.unqueried_rows, summaries, strict=False): + for (row_idx, _), summary, pred in zip( + state.unqueried_rows, summaries, unqueried_preds, strict=False + ): drc = float(summary["score_drc"]) ps = float(summary["score_ps"]) overall = max(drc, ps) @@ -500,6 +505,7 @@ def annotate_campaign_state( result.at[row_idx, "overall_score"] = overall result.at[row_idx, "recommendation"] = rec result.at[row_idx, "embedding_derived"] = summary["embedding_derived"] + result.at[row_idx, "predicted_pec50"] = float(pred) # Score PS hits — only DRC upgrade is a valid next action; ps_score stays NaN if state.ps_upgrade_rows: @@ -507,12 +513,15 @@ def annotate_campaign_state( summaries = acquisition.score_summary( upgrade_canonical, upgrade_preds, provenance=upgrade_provenance ) - for (row_idx, _), summary in zip(state.ps_upgrade_rows, summaries, strict=False): + for (row_idx, _), summary, pred in zip( + state.ps_upgrade_rows, summaries, upgrade_preds, strict=False + ): drc = float(summary["score_drc"]) result.at[row_idx, "drc_score"] = drc result.at[row_idx, "overall_score"] = drc result.at[row_idx, "recommendation"] = "drc" result.at[row_idx, "embedding_derived"] = summary["embedding_derived"] + result.at[row_idx, "predicted_pec50"] = float(pred) return result From 1ceb45990f17c9bff01246f772ff496ee02988f8 Mon Sep 17 00:00:00 2001 From: Sean Colby Date: Mon, 27 Jul 2026 18:47:56 -0800 Subject: [PATCH 15/18] Fix build_dataloader drop_last incompatibility in inference paths MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- moal/auxiliary_encoder.py | 7 +++---- moal/model.py | 39 ++++++++++++++++++++++++++++++++------- tests/test_model.py | 27 ++++++++++++++++++++++++++- 3 files changed, 61 insertions(+), 12 deletions(-) diff --git a/moal/auxiliary_encoder.py b/moal/auxiliary_encoder.py index d8a77fc..69e077c 100644 --- a/moal/auxiliary_encoder.py +++ b/moal/auxiliary_encoder.py @@ -29,7 +29,7 @@ from torch.utils.data import DataLoader, Dataset, random_split from moal.config import AuxiliaryModelConfig -from moal.model import build_mpnn +from moal.model import build_mpnn, safe_inference_batch_size from moal.types import LabelRecord logger = logging.getLogger(__name__) @@ -485,9 +485,8 @@ def embed_smiles(self, smiles_list: list[str], batch_size: int = 256) -> np.ndar ``AuxiliaryModelConfig.embedding_dim``. """ dataset = MoleculeDataset([MoleculeDatapoint.from_smi(s) for s in smiles_list]) # pyright: ignore[reportArgumentType] - dataloader = build_dataloader( - dataset, batch_size=batch_size, shuffle=False, drop_last=False - ) + batch_size = safe_inference_batch_size(len(dataset), batch_size) + dataloader = build_dataloader(dataset, batch_size=batch_size, shuffle=False) all_embeddings = [] with torch.inference_mode(): diff --git a/moal/model.py b/moal/model.py index 03b561a..c0a9a36 100644 --- a/moal/model.py +++ b/moal/model.py @@ -34,6 +34,35 @@ logger = logging.getLogger(__name__) + +def safe_inference_batch_size(dataset_size: int, batch_size: int) -> int: + """Return a batch size that avoids chemprop's silent single-molecule drop. + + ``chemprop.data.dataloader.build_dataloader`` drops the last batch + whenever ``dataset_size % batch_size == 1``, to protect batch-norm during + training. At inference that would silently omit a molecule and misalign + predictions/embeddings with the input SMILES order, so this shrinks the + batch size (never below 1) until the remainder condition no longer holds. + + Parameters + ---------- + dataset_size : int + Number of molecules to batch. + batch_size : int + Requested batch size. + + Returns + ------- + int + A batch size no larger than ``dataset_size`` for which + ``dataset_size % batch_size != 1`` (or 1, if no larger value works). + """ + effective = min(batch_size, dataset_size) + while effective > 1 and dataset_size % effective == 1: + effective -= 1 + return effective + + _KNOWN_FOUNDATION_MODELS: frozenset[str] = frozenset({"chemeleon"}) @@ -619,13 +648,9 @@ def predict_smiles(self, smiles_list: list[str], batch_size: int = 256) -> np.nd dataset = MoleculeDataset([MoleculeDatapoint.from_smi(s) for s in smiles_list]) # pyright: ignore[reportArgumentType] # Let the dataloader handle batching and graph collation automatically. - # drop_last=False is explicit: chemprop defaults to dropping the last - # batch when len(dataset) % batch_size == 1 to protect batch-norm - # during training, but at inference that would silently omit a molecule - # and misalign predictions with the input SMILES list. - dataloader = build_dataloader( - dataset, batch_size=batch_size, shuffle=False, drop_last=False - ) + # chemprop's build_dataloader no longer accepts drop_last as an override. + batch_size = safe_inference_batch_size(len(dataset), batch_size) + dataloader = build_dataloader(dataset, batch_size=batch_size, shuffle=False) all_preds = [] diff --git a/tests/test_model.py b/tests/test_model.py index ea0010f..78897f6 100644 --- a/tests/test_model.py +++ b/tests/test_model.py @@ -19,7 +19,12 @@ from chemprop.nn import BondMessagePassing, MeanAggregation, RegressionFFN from moal.loss import CensoredRegressionLoss -from moal.model import _KNOWN_FOUNDATION_MODELS, ChemPropLightningModule, NoisyOracleModel +from moal.model import ( + _KNOWN_FOUNDATION_MODELS, + ChemPropLightningModule, + NoisyOracleModel, + safe_inference_batch_size, +) from moal.types import CensoringType, LabelRecord, QueryType # Capture the real _build_model before any test fixture can patch it. @@ -537,3 +542,23 @@ def test_custom_path_loads_weights(self, tmp_path, monkeypatch): m = ChemPropLightningModule(from_foundation=str(weights_path)) assert m.hparams["from_foundation"] == str(weights_path) assert isinstance(m.model, nn.Module) + + +class TestSafeInferenceBatchSize: + """Tests for safe_inference_batch_size avoiding chemprop's remainder-1 batch drop.""" + + def test_single_molecule_does_not_collapse_to_zero(self): + """A single-molecule dataset must not shrink the batch size to 0.""" + assert safe_inference_batch_size(dataset_size=1, batch_size=256) == 1 + + def test_no_remainder_leaves_batch_size_unchanged(self): + """A batch size that already avoids remainder 1 must be returned as-is.""" + assert safe_inference_batch_size(dataset_size=512, batch_size=256) == 256 + + def test_remainder_one_shrinks_batch_size(self): + """A batch size producing exactly one leftover molecule must shrink by one.""" + assert safe_inference_batch_size(dataset_size=257, batch_size=256) == 255 + + def test_batch_size_larger_than_dataset_is_capped(self): + """A batch size larger than the dataset must be capped to the dataset size.""" + assert safe_inference_batch_size(dataset_size=5, batch_size=256) == 5 From 27fe72a44c81f52c872e9c6e2ac85ae98e2a12f0 Mon Sep 17 00:00:00 2001 From: Sean Colby Date: Mon, 27 Jul 2026 18:48:07 -0800 Subject: [PATCH 16/18] Add early stopping support to TrainerConfig MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- moal/config.py | 39 +++++++++++++++++++++++++++++++++++++++ tests/test_config.py | 28 ++++++++++++++++++++++++++++ 2 files changed, 67 insertions(+) diff --git a/moal/config.py b/moal/config.py index e26aa1f..690cbae 100644 --- a/moal/config.py +++ b/moal/config.py @@ -12,6 +12,7 @@ from typing import Any import yaml +from lightning.pytorch.callbacks import EarlyStopping @dataclass(frozen=True) @@ -287,6 +288,28 @@ class TrainerConfig: Clipping algorithm passed to ``lightning.Trainer`` when ``gradient_clip_val`` is set: ``"norm"`` (default) or ``"value"``. Ignored when ``gradient_clip_val`` is None. + early_stopping : bool + Whether to attach a ``lightning.pytorch.callbacks.EarlyStopping`` + callback. Default is False (train for exactly ``max_epochs``, current + behavior unchanged). Requires a validation split (``val_fraction`` > + 0) so the monitored metric is actually logged each epoch. + early_stopping_monitor : str + Metric name to monitor. Default ``"val_loss"`` matches + ``ChemPropLightningModule``'s logged key; the auxiliary encoder logs + ``"aux_val_loss"`` instead, so ``auxiliary_trainer`` configs must set + this explicitly or ``EarlyStopping`` will raise. Ignored when + ``early_stopping`` is False. + early_stopping_patience : int + Number of epochs with no improvement (beyond ``early_stopping_min_delta``) + before stopping. Ignored when ``early_stopping`` is False. + early_stopping_mode : str + ``"min"`` (default) or ``"max"``, matching whether lower or higher + values of ``early_stopping_monitor`` are better. Ignored when + ``early_stopping`` is False. + early_stopping_min_delta : float + Minimum change in the monitored metric to qualify as an improvement. + Default is 0.0 (any improvement resets patience). Ignored when + ``early_stopping`` is False. """ max_epochs: int = 30 @@ -299,6 +322,11 @@ class TrainerConfig: log_every_n_steps: int = 1 gradient_clip_val: float | None = None gradient_clip_algorithm: str = "norm" + early_stopping: bool = False + early_stopping_monitor: str = "val_loss" + early_stopping_patience: int = 10 + early_stopping_mode: str = "min" + early_stopping_min_delta: float = 0.0 def to_dict(self) -> dict[str, Any]: """Return only the kwargs that ``lightning.Trainer`` accepts. @@ -315,6 +343,8 @@ def to_dict(self) -> dict[str, Any]: ``log_every_n_steps``. ``gradient_clip_val`` (and ``gradient_clip_algorithm``) are added only when clipping is enabled, so the default (None) leaves Lightning's clipping off. + ``callbacks`` (an ``EarlyStopping`` instance) is added only when + ``early_stopping`` is True. """ kwargs: dict[str, Any] = { "max_epochs": self.max_epochs, @@ -328,6 +358,15 @@ def to_dict(self) -> dict[str, Any]: if self.gradient_clip_val is not None: kwargs["gradient_clip_val"] = self.gradient_clip_val kwargs["gradient_clip_algorithm"] = self.gradient_clip_algorithm + if self.early_stopping: + kwargs["callbacks"] = [ + EarlyStopping( + monitor=self.early_stopping_monitor, + patience=self.early_stopping_patience, + mode=self.early_stopping_mode, + min_delta=self.early_stopping_min_delta, + ) + ] return kwargs def to_datamodule_kwargs(self) -> dict[str, Any]: diff --git a/tests/test_config.py b/tests/test_config.py index 0e90af7..8c470fd 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -5,6 +5,7 @@ from pathlib import Path import yaml +from lightning.pytorch.callbacks import EarlyStopping from moal.config import AuxiliaryModelConfig, PipelineConfig, TrainerConfig @@ -72,3 +73,30 @@ def test_round_trips_through_from_yaml(self, tmp_path): assert cfg.auxiliary_trainer == TrainerConfig(max_epochs=15, val_fraction=0.2) assert cfg.trainer == TrainerConfig(max_epochs=30, val_fraction=0.0) + + +class TestTrainerConfigEarlyStopping: + """Tests for TrainerConfig.to_dict()'s early-stopping callback wiring.""" + + def test_omits_callbacks_by_default(self): + """to_dict() must not add a callbacks key when early_stopping is False.""" + kwargs = TrainerConfig().to_dict() + + assert "callbacks" not in kwargs + + def test_adds_early_stopping_callback_when_enabled(self): + """to_dict() must add an EarlyStopping callback configured from the early_stopping_* fields.""" + kwargs = TrainerConfig( + early_stopping=True, + early_stopping_monitor="aux_val_loss", + early_stopping_patience=3, + early_stopping_mode="max", + early_stopping_min_delta=0.01, + ).to_dict() + + [callback] = kwargs["callbacks"] + assert isinstance(callback, EarlyStopping) + assert callback.monitor == "aux_val_loss" + assert callback.patience == 3 + assert callback.mode == "max" + assert callback.min_delta == 0.01 From 55ff08c5a196e2612531d8489c5fbf489b2cb23d Mon Sep 17 00:00:00 2001 From: Sean Colby Date: Mon, 27 Jul 2026 18:48:47 -0800 Subject: [PATCH 17/18] Add predicted-readout concatenation feature to the auxiliary encoder MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- moal/auxiliary_encoder.py | 41 ++++++++++++++++++++++++ moal/cli.py | 5 +-- moal/concatenation_model.py | 53 +++++++++++++++++++++++++------ moal/config.py | 11 +++++++ tests/test_auxiliary_encoder.py | 9 ++++++ tests/test_concatenation_model.py | 46 ++++++++++++++++++++++----- 6 files changed, 146 insertions(+), 19 deletions(-) diff --git a/moal/auxiliary_encoder.py b/moal/auxiliary_encoder.py index 69e077c..73ec1c3 100644 --- a/moal/auxiliary_encoder.py +++ b/moal/auxiliary_encoder.py @@ -497,6 +497,47 @@ def embed_smiles(self, smiles_list: list[str], batch_size: int = 256) -> np.ndar return np.concatenate(all_embeddings, axis=0).astype(np.float32) + @torch.no_grad() + def predict_smiles(self, smiles_list: list[str], batch_size: int = 256) -> np.ndarray: + """Return multi-task readout predictions for a list of SMILES. + + Unlike :meth:`embed_smiles`, this runs the full forward pass + including the predictor head, giving a predicted value per + ``task_names`` entry for every compound, whether or not it was + actually screened at that concentration. Used to backfill a + "predicted readout" feature column that covers the full compound + pool, as opposed to :func:`moal.concatenation_model.build_concatenation_features`'s + observed-readout fallback, which only has a value for screened + compounds. + + Parameters + ---------- + smiles_list : list[str] + **Must be RDKit-canonical, salt-stripped SMILES**, matching + :meth:`moal.model.ChemPropLightningModule.predict_smiles`'s + contract. + batch_size : int, optional + Number of molecules processed per forward pass. Default is 256. + + Returns + ------- + np.ndarray + Array of shape ``(N, len(task_names))``, aligned with + ``smiles_list``. + """ + dataset = MoleculeDataset([MoleculeDatapoint.from_smi(s) for s in smiles_list]) # pyright: ignore[reportArgumentType] + batch_size = safe_inference_batch_size(len(dataset), batch_size) + dataloader = build_dataloader(dataset, batch_size=batch_size, shuffle=False) + + all_preds = [] + with torch.inference_mode(): + for batch in dataloader: + batch.bmg.to(self.device) + preds = self(batch.bmg) + all_preds.append(preds.cpu().numpy()) + + return np.concatenate(all_preds, axis=0).astype(np.float32) + def pretrain_auxiliary_encoder( records: list[LabelRecord], diff --git a/moal/cli.py b/moal/cli.py index d967230..a371f4c 100644 --- a/moal/cli.py +++ b/moal/cli.py @@ -862,8 +862,8 @@ def _build_concatenation_model( ------- ConcatenationChemPropLightningModule Configured model ready for ``refit()`` and ``predict_smiles()``. - ``use_observed_readout`` is fixed here at construction (from - ``cfg.auxiliary_model.use_observed_readout``) rather than passed + ``use_observed_readout``/``use_predicted_readout`` are fixed here at + construction (from ``cfg.auxiliary_model``) rather than passed separately to each call, so training and inference routing cannot drift apart. """ @@ -873,6 +873,7 @@ def _build_concatenation_model( return ConcatenationChemPropLightningModule( concat_feature_dim=feature_dim, use_observed_readout=cfg.auxiliary_model.use_observed_readout, + use_predicted_readout=cfg.auxiliary_model.use_predicted_readout, ffn_hidden_dim=cfg.model.ffn_hidden_dim, ffn_num_layers=cfg.model.ffn_num_layers, message_hidden_dim=cfg.model.message_hidden_dim, diff --git a/moal/concatenation_model.py b/moal/concatenation_model.py index c311225..130e858 100644 --- a/moal/concatenation_model.py +++ b/moal/concatenation_model.py @@ -58,10 +58,14 @@ def concatenation_feature_dim(n_tasks: int, embedding_dim: int) -> int: Returns ------- int - ``2 * n_tasks + embedding_dim + 1``: observed-readout vector, - readout mask, structural embedding, and a single provenance flag. + ``3 * n_tasks + embedding_dim + 1``: observed-readout vector, + observed-readout mask, predicted-readout vector, structural + embedding, and a single provenance flag. Fixed regardless of + ``use_observed_readout``/``use_predicted_readout``; unused blocks + are zeroed by :func:`build_concatenation_features` rather than + omitted, so a model's input width never depends on those flags. """ - return 2 * n_tasks + embedding_dim + 1 + return 3 * n_tasks + embedding_dim + 1 def build_concatenation_features( @@ -70,6 +74,7 @@ def build_concatenation_features( aux_encoder: AuxiliaryEncoderModule, *, use_observed_readout: bool = True, + use_predicted_readout: bool = False, batch_size: int = 256, ) -> np.ndarray: """Build the per-compound concatenation feature matrix. @@ -89,6 +94,15 @@ def build_concatenation_features( blocks stay zero and the flag is 0 — the compound is scored from its structural embedding alone. + When ``use_predicted_readout`` is True, the auxiliary encoder's own + predictor head is additionally run over every compound (via + :meth:`~moal.auxiliary_encoder.AuxiliaryEncoderModule.predict_smiles`) + and concatenated as a third block. Unlike the observed-readout block, + this one is never masked: the encoder can predict a value for any + SMILES, so it is populated identically for every compound at both + training and inference time, with no fallback branch for the two to + drift apart. + Parameters ---------- canonical_smiles : list[str] @@ -97,16 +111,20 @@ def build_concatenation_features( Per-compound ``LabelRecord.raw_ps_readouts``-shaped dict, aligned with ``canonical_smiles``. An empty dict means "never PS-screened". aux_encoder : AuxiliaryEncoderModule - Pretrained auxiliary encoder; supplies both ``task_names`` (readout - key order) and the structural embedding. + Pretrained auxiliary encoder; supplies ``task_names`` (readout key + order), the structural embedding, and (when + ``use_predicted_readout``) the predicted-readout block. use_observed_readout : bool, optional When True (default), compounds with a non-empty readout also get their raw value concatenated alongside the embedding. When False, - the readout and mask blocks are always zero and every compound is - scored from its embedding alone, matching + the readout and mask blocks are always zero, matching ``AuxiliaryModelConfig.use_observed_readout``. + use_predicted_readout : bool, optional + When True, the auxiliary encoder's predicted readout is concatenated + for every compound. When False (default), that block is always + zero, matching ``AuxiliaryModelConfig.use_predicted_readout``. batch_size : int, optional - Batch size for the embedding forward pass. Default is 256. + Batch size for the embedding/prediction forward passes. Default is 256. Returns ------- @@ -143,9 +161,16 @@ def build_concatenation_features( readout_mask[i, j] = 1.0 readout_used[i, 0] = 1.0 + if use_predicted_readout: + predicted_vec = aux_encoder.predict_smiles(canonical_smiles, batch_size=batch_size) + else: + predicted_vec = np.zeros((n, n_tasks), dtype=np.float32) + embeddings = aux_encoder.embed_smiles(canonical_smiles, batch_size=batch_size) - return np.concatenate([readout_vec, readout_mask, embeddings, readout_used], axis=1) + return np.concatenate( + [readout_vec, readout_mask, predicted_vec, embeddings, readout_used], axis=1 + ) class _ConcatenatedDataset(Dataset): @@ -241,6 +266,12 @@ class ConcatenationChemPropLightningModule(L.LightningModule): Deliberately not a per-call parameter on either method, so training-time and inference-time routing cannot drift apart. Default is True. + use_predicted_readout : bool, optional + Also fixed at construction, same rationale as ``use_observed_readout``. + Unlike that flag, this block is never masked when enabled: the + auxiliary encoder predicts a value for every compound, so it is + populated identically at training and inference time by + construction, not just by convention. Default is False. ffn_hidden_dim, ffn_num_layers, message_hidden_dim, depth, freeze_epochs, mpnn_lr, ffn_lr, mpnn_weight_decay, ffn_weight_decay, sigma, learnable_sigma, from_foundation @@ -257,6 +288,7 @@ def __init__( self, concat_feature_dim: int, use_observed_readout: bool = True, + use_predicted_readout: bool = False, ffn_hidden_dim: int = 300, ffn_num_layers: int = 2, message_hidden_dim: int = 300, @@ -275,6 +307,7 @@ def __init__( self._from_foundation = from_foundation self.concat_feature_dim = concat_feature_dim self.use_observed_readout = use_observed_readout + self.use_predicted_readout = use_predicted_readout self.save_hyperparameters() self.freeze_epochs = freeze_epochs @@ -507,6 +540,7 @@ def predict_smiles( readouts, aux_encoder, use_observed_readout=self.use_observed_readout, + use_predicted_readout=self.use_predicted_readout, batch_size=batch_size, ) x_d = torch.as_tensor(features, dtype=torch.float32) @@ -593,6 +627,7 @@ def refit( [rec.raw_ps_readouts for rec in records], aux_encoder, use_observed_readout=self.use_observed_readout, + use_predicted_readout=self.use_predicted_readout, ) dm = _ConcatenatedDataModule(records, features, **(datamodule_kwargs or {})) dm.setup() diff --git a/moal/config.py b/moal/config.py index 690cbae..826ef4b 100644 --- a/moal/config.py +++ b/moal/config.py @@ -200,6 +200,16 @@ class AuxiliaryModelConfig: compounds; a constant-zero input column is a mathematical no-op for a plain linear layer (zero gradient, zero forward contribution), so this does not degrade model capacity. + use_predicted_readout : bool + Also concatenates the auxiliary encoder's own predicted readout + (its multi-task predictor head's output, not just its pooled + embedding) into the main model's input, for every compound. Unlike + ``use_observed_readout``, this block is never masked or zeroed by + missing data: the auxiliary encoder can predict a value for any + SMILES, so training and inference always populate this block + identically, with no fallback branch to drift apart. Independent of + ``use_observed_readout``; both may be enabled together. Default is + False. """ from_foundation: str | bool = "chemeleon" @@ -213,6 +223,7 @@ class AuxiliaryModelConfig: embedding_dim: int = 300 checkpoint_path: str | None = None use_observed_readout: bool = True + use_predicted_readout: bool = False @dataclass(frozen=True) diff --git a/tests/test_auxiliary_encoder.py b/tests/test_auxiliary_encoder.py index 88319c2..9c6a1b8 100644 --- a/tests/test_auxiliary_encoder.py +++ b/tests/test_auxiliary_encoder.py @@ -126,6 +126,15 @@ def test_embed_smiles_returns_backbone_width_aligned_with_input(self): assert embeddings.shape == (3, 24) + def test_predict_smiles_returns_task_count_width_aligned_with_input(self): + """predict_smiles must return one prediction row per input SMILES, at task_names width.""" + config = _fast_config() + module = AuxiliaryEncoderModule(task_names=["log2fc_1um", "pic50"], config=config) + + predictions = module.predict_smiles(["CCO", "CCN", "CCC"]) + + assert predictions.shape == (3, 2) + class TestPretrainAuxiliaryEncoder: """Tests for pretrain_auxiliary_encoder: training end-to-end and checkpoint opt-in.""" diff --git a/tests/test_concatenation_model.py b/tests/test_concatenation_model.py index a212aea..68645f5 100644 --- a/tests/test_concatenation_model.py +++ b/tests/test_concatenation_model.py @@ -72,13 +72,13 @@ def _records() -> list[LabelRecord]: class TestConcatenationFeatureDim: """Tests for concatenation_feature_dim's arithmetic.""" - def test_matches_2n_plus_embedding_plus_1(self): - """The formula must be 2 * n_tasks + embedding_dim + 1.""" - assert concatenation_feature_dim(n_tasks=3, embedding_dim=10) == 2 * 3 + 10 + 1 + def test_matches_3n_plus_embedding_plus_1(self): + """The formula must be 3 * n_tasks + embedding_dim + 1.""" + assert concatenation_feature_dim(n_tasks=3, embedding_dim=10) == 3 * 3 + 10 + 1 class TestBuildConcatenationFeatures: - """Tests for build_concatenation_features: observed vs embedding routing and shape.""" + """Tests for build_concatenation_features: observed/predicted vs embedding routing and shape.""" def test_observed_readout_also_gets_embedding(self, aux_encoder): """A compound with a readout must populate readout/mask AND the embedding block, with flag=1.""" @@ -87,7 +87,7 @@ def test_observed_readout_also_gets_embedding(self, aux_encoder): readout_block = features[0, :n_tasks] mask_block = features[0, n_tasks : 2 * n_tasks] - embedding_block = features[0, 2 * n_tasks : 2 * n_tasks + _EMBEDDING_DIM] + embedding_block = features[0, 3 * n_tasks : 3 * n_tasks + _EMBEDDING_DIM] flag = features[0, -1] assert readout_block[0] == 2.5 @@ -101,7 +101,7 @@ def test_missing_readout_uses_embedding_only(self, aux_encoder): n_tasks = 2 readout_mask_block = features[0, : 2 * n_tasks] - embedding_block = features[0, 2 * n_tasks : 2 * n_tasks + _EMBEDDING_DIM] + embedding_block = features[0, 3 * n_tasks : 3 * n_tasks + _EMBEDDING_DIM] flag = features[0, -1] assert np.all(readout_mask_block == 0.0) @@ -121,8 +121,38 @@ def test_use_observed_readout_false_zeroes_readout_block_but_keeps_embedding(sel assert np.all(with_readout[0, : 2 * n_tasks] == 0.0) assert with_readout[0, -1] == 0.0 np.testing.assert_allclose( - with_readout[0, 2 * n_tasks : 2 * n_tasks + _EMBEDDING_DIM], - without_readout[0, 2 * n_tasks : 2 * n_tasks + _EMBEDDING_DIM], + with_readout[0, 3 * n_tasks : 3 * n_tasks + _EMBEDDING_DIM], + without_readout[0, 3 * n_tasks : 3 * n_tasks + _EMBEDDING_DIM], + ) + + def test_predicted_readout_disabled_by_default(self, aux_encoder): + """With use_predicted_readout unset (default False), the predicted-readout block must stay zero.""" + features = build_concatenation_features(["CCO"], [{"log2fc_1um": 2.5}], aux_encoder) + n_tasks = 2 + + predicted_block = features[0, 2 * n_tasks : 3 * n_tasks] + + assert np.all(predicted_block == 0.0) + + def test_predicted_readout_matches_encoder_prediction_regardless_of_observed_readout( + self, aux_encoder + ): + """use_predicted_readout=True must populate the predicted block from the encoder's own + prediction, identically whether or not the compound has an observed readout. + """ + n_tasks = 2 + expected = aux_encoder.predict_smiles(["CCO"])[0] + + with_observed = build_concatenation_features( + ["CCO"], [{"log2fc_1um": 2.5}], aux_encoder, use_predicted_readout=True + ) + without_observed = build_concatenation_features( + ["CCO"], [{}], aux_encoder, use_predicted_readout=True + ) + + np.testing.assert_allclose(with_observed[0, 2 * n_tasks : 3 * n_tasks], expected, atol=1e-6) + np.testing.assert_allclose( + without_observed[0, 2 * n_tasks : 3 * n_tasks], expected, atol=1e-6 ) def test_output_shape_matches_concatenation_feature_dim(self, aux_encoder): From 0362aa3c52cfb498a6431001bf1668952eea838b Mon Sep 17 00:00:00 2001 From: Sean Colby Date: Fri, 4 Sep 2026 10:23:59 -0800 Subject: [PATCH 18/18] Add SMILES embedding --- moal/model.py | 39 +++++++++++++++++++++++++++++++++++++++ 1 file changed, 39 insertions(+) diff --git a/moal/model.py b/moal/model.py index c0a9a36..378554d 100644 --- a/moal/model.py +++ b/moal/model.py @@ -668,6 +668,45 @@ def predict_smiles(self, smiles_list: list[str], batch_size: int = 256) -> np.nd return np.array(all_preds, dtype=np.float32) + @torch.no_grad() + def embed_smiles(self, smiles_list: list[str], batch_size: int = 256) -> np.ndarray: + """Return pooled structural embeddings (pre-predictor) for a list of SMILES. + + Mirrors :meth:`moal.auxiliary_encoder.AuxiliaryEncoderModule.embed_smiles`. + Uses ``chemprop.models.MPNN.fingerprint``, which applies message-passing, + mean pooling, and batch-norm but stops short of the pEC50 predictor head, + so the returned vectors reflect whatever fine-tuning ``refit`` has done + to the encoder so far. + + Parameters + ---------- + smiles_list : list[str] + **Must be RDKit-canonical, salt-stripped SMILES**, matching + :meth:`predict_smiles`'s contract. + batch_size : int, optional + Number of molecules processed per forward pass. Default is 256. + + Returns + ------- + np.ndarray + Array of shape ``(N, embedding_dim)``, aligned with + ``smiles_list``. ``embedding_dim`` is CheMeleon's fixed native + width (2048) for a foundation checkpoint, or ``message_hidden_dim`` + for a random-init encoder. + """ + dataset = MoleculeDataset([MoleculeDatapoint.from_smi(s) for s in smiles_list]) # pyright: ignore[reportArgumentType] + batch_size = safe_inference_batch_size(len(dataset), batch_size) + dataloader = build_dataloader(dataset, batch_size=batch_size, shuffle=False) + + all_embeddings = [] + with torch.inference_mode(): + for batch in dataloader: + batch.bmg.to(self.device) + embedding = cast(MPNN, self.model).fingerprint(batch.bmg) + all_embeddings.append(embedding.cpu().numpy()) + + return np.concatenate(all_embeddings, axis=0).astype(np.float32) + def refit( self, records: list[LabelRecord],