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/auxiliary_encoder.py b/moal/auxiliary_encoder.py new file mode 100644 index 0000000..73ec1c3 --- /dev/null +++ b/moal/auxiliary_encoder.py @@ -0,0 +1,663 @@ +"""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.AuxiliaryModelConfig` +for why). +""" + +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 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 + +from moal.config import AuxiliaryModelConfig +from moal.model import build_mpnn, safe_inference_batch_size +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 = 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( + "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: + """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 + Non-shuffled DataLoader over the validation split, or an empty + ``DataLoader`` if no split was formed. + """ + if self._val_dataset is None: + return DataLoader([], batch_size=self.batch_size) # pyright: ignore[reportArgumentType] + 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 : AuxiliaryModelConfig + Backbone architecture, freeze schedule, and optimization + hyperparameters. + """ + + def __init__(self, task_names: list[str], config: AuxiliaryModelConfig) -> 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() + + @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 + # ------------------------------------------------------------------ + + 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) + + # ------------------------------------------------------------------ + # 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 + ``AuxiliaryModelConfig.embedding_dim``. + """ + 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) + + @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], + config: AuxiliaryModelConfig, + max_epochs: int = 30, + 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 : 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``. + 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": 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) + 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: AuxiliaryModelConfig +) -> AuxiliaryEncoderModule: + """Load an auxiliary encoder checkpoint written by :func:`save_auxiliary_encoder_checkpoint`. + + Parameters + ---------- + path : str or Path + Path to the checkpoint file. + 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. + + 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/cli.py b/moal/cli.py index 165b897..a371f4c 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, @@ -29,6 +30,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 +43,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, @@ -335,6 +342,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_columns=cfg.data.plan.log2fc_columns, is_canonical=cfg.data.plan.is_canonical, expected_ps_threshold=cfg.oracle.ps_threshold, ) @@ -391,32 +399,97 @@ 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_model is not None: + # Concatenation architecture (issue #36 Phase 2): pretrain the + # 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_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 + 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( + drc_records, + aux_encoder=aux_encoder, + 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) + 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 @@ -508,6 +581,7 @@ def _load_pretrain_records( relation_column=pretrain_cfg.relation_column, value_column=pretrain_cfg.value_column, weight_column=pretrain_cfg.weight_column, + log2fc_columns=pretrain_cfg.log2fc_columns, is_canonical=pretrain_cfg.is_canonical, expected_ps_threshold=cfg.oracle.ps_threshold, ) @@ -670,6 +744,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, ) @@ -766,5 +841,83 @@ 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`. + ``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. + + Returns + ------- + ConcatenationChemPropLightningModule + Configured model ready for ``refit()`` and ``predict_smiles()``. + ``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. + """ + 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, + 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, + 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, + 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/moal/concatenation_model.py b/moal/concatenation_model.py new file mode 100644 index 0000000..130e858 --- /dev/null +++ b/moal/concatenation_model.py @@ -0,0 +1,797 @@ +"""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, QueryType + +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 + ``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 3 * n_tasks + embedding_dim + 1 + + +def build_concatenation_features( + canonical_smiles: list[str], + readouts: list[dict[str, float]], + 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. + + 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. + + 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] + 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 ``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, 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/prediction forward passes. 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) + readout_used = np.zeros((n, 1), dtype=np.float32) + + 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 + 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, predicted_vec, embeddings, readout_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. + 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. + 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 + 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, + use_predicted_readout: bool = False, + 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, + 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.use_observed_readout = use_observed_readout + self.use_predicted_readout = use_predicted_readout + 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": [], + } + + # 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=1.0, w_ps=1.0, 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``. + Forwarded to :func:`build_concatenation_features`. + aux_encoder : AuxiliaryEncoderModule + Pretrained auxiliary encoder supplying both the readout-key + order and the structural embedding. + 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, + 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) + + # 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). + + 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, + use_predicted_readout=self.use_predicted_readout, + ) + 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 = 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( + "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: + """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 + Non-shuffled DataLoader over the validation split, or an empty + ``DataLoader`` if no split was formed. + """ + if self._val_dataset is None: + return DataLoader([], batch_size=self.batch_size) # pyright: ignore[reportArgumentType] + 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/config.py b/moal/config.py index 0dfc664..826ef4b 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) @@ -122,6 +123,109 @@ class ModelConfig: from_foundation: str | bool = "chemeleon" +@dataclass(frozen=True) +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 + 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 + ---------- + 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. + 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. + 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. + 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" + 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 + embedding_dim: int = 300 + checkpoint_path: str | None = None + use_observed_readout: bool = True + use_predicted_readout: bool = False + + @dataclass(frozen=True) class AcquisitionConfig: """Acquisition function hyper-parameters. @@ -135,11 +239,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) @@ -183,6 +299,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 @@ -195,6 +333,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. @@ -211,6 +354,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, @@ -224,6 +369,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]: @@ -299,6 +453,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_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. @@ -309,6 +469,7 @@ class PretrainDataConfig: relation_column: str = "relation" value_column: str = "value" weight_column: str | None = None + log2fc_columns: list[str] | None = None is_canonical: bool = False @@ -389,6 +550,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_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. @@ -400,6 +568,7 @@ class PlanDataConfig: relation_column: str = "relation" value_column: str = "value" weight_column: str | None = None + log2fc_columns: list[str] | None = None is_canonical: bool = False @@ -472,6 +641,14 @@ class PipelineConfig: Command-specific dataset and I/O settings. active_learning_loop : ActiveLearningLoopConfig Parameters controlling the active learning iteration loop. + 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. """ @@ -483,6 +660,8 @@ class PipelineConfig: dashboard: DashboardConfig = field(default_factory=DashboardConfig) data: DataConfig = field(default_factory=DataConfig) active_learning_loop: ActiveLearningLoopConfig = field(default_factory=ActiveLearningLoopConfig) + auxiliary_model: AuxiliaryModelConfig | None = None + auxiliary_trainer: TrainerConfig = field(default_factory=TrainerConfig) seed: int = 42 @@ -505,6 +684,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_model_raw = raw.get("auxiliary_model", None) return cls( oracle=OracleConfig(**raw.get("oracle", {})), model=ModelConfig(**raw.get("model", {})), @@ -520,6 +700,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_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/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/moal/model.py b/moal/model.py index dc7c732..378554d 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"}) @@ -70,6 +99,125 @@ 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, + extra_input_dim: int = 0, +) -> 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. + 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 + ------- + 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 + extra_input_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 +363,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 +382,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 @@ -539,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 = [] @@ -563,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], diff --git a/moal/planning.py b/moal/planning.py index 7895c3e..9a34689 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_columns: list[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_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 @@ -105,6 +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_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]] = [] @@ -176,6 +190,24 @@ def parse_campaign_state( f"Row {csv_row}: weight must be finite and positive, got {weight}." ) + 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( smiles=raw_smiles, @@ -187,6 +219,7 @@ def parse_campaign_state( cost=cost_drc, iteration=_PLAN_MODE_ITERATION, weight=weight, + raw_ps_readouts=raw_ps_readouts, ) else: if expected_ps_threshold is not None and not math.isclose( @@ -206,6 +239,7 @@ def parse_campaign_state( cost=cost_ps, iteration=_PLAN_MODE_ITERATION, weight=weight, + raw_ps_readouts=raw_ps_readouts, ) # PS hits are DRC-upgrade inference targets in addition to training records if relation == ">=": @@ -248,6 +282,7 @@ def parse_pretrain_records( relation_column: str = "relation", value_column: str = "value", weight_column: str | None = None, + log2fc_columns: list[str] | None = None, is_canonical: bool = False, expected_ps_threshold: float | None = None, ) -> list[LabelRecord]: @@ -278,6 +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_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 @@ -300,6 +339,7 @@ def parse_pretrain_records( relation_column=relation_column, value_column=value_column, weight_column=weight_column, + log2fc_columns=log2fc_columns, is_canonical=is_canonical, expected_ps_threshold=expected_ps_threshold, ) @@ -318,7 +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. + 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 ---------- @@ -336,15 +380,29 @@ 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_readouts 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_readouts + } + + 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.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 def annotate_campaign_state( @@ -352,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. @@ -359,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 @@ -367,6 +426,11 @@ 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 + - ``predicted_pec50`` — the raw model prediction from ``predictions``, + unmodified by acquisition scoring; NaN for training-only rows Parameters ---------- @@ -378,11 +442,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) @@ -396,22 +466,36 @@ 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 + result["predicted_pec50"] = np.nan 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) - for (row_idx, _), summary in zip(state.unqueried_rows, summaries, strict=False): + summaries = acquisition.score_summary( + unqueried_canonical, unqueried_preds, provenance=unqueried_provenance + ) + 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) @@ -420,16 +504,24 @@ 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"] + 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: upgrade_canonical = [smi for _, smi in state.ps_upgrade_rows] - summaries = acquisition.score_summary(upgrade_canonical, upgrade_preds) - for (row_idx, _), summary in zip(state.ps_upgrade_rows, summaries, strict=False): + summaries = acquisition.score_summary( + upgrade_canonical, upgrade_preds, provenance=upgrade_provenance + ) + 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 diff --git a/moal/types.py b/moal/types.py index 30328f3..5d80560 100644 --- a/moal/types.py +++ b/moal/types.py @@ -80,6 +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_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 ``AuxiliaryModelConfig``) can still use it. + Empty when the compound has never been PS-screened. """ smiles: str @@ -95,6 +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_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_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]), + ) diff --git a/tests/test_auxiliary_encoder.py b/tests/test_auxiliary_encoder.py new file mode 100644 index 0000000..9c6a1b8 --- /dev/null +++ b/tests/test_auxiliary_encoder.py @@ -0,0 +1,232 @@ +"""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 AuxiliaryModelConfig +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) -> AuxiliaryModelConfig: + defaults = { + "from_foundation": False, + "message_hidden_dim": 16, + "ffn_hidden_dim": 16, + "depth": 1, + "freeze_epochs": 0, + } + defaults.update(overrides) + return AuxiliaryModelConfig(**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()) + + 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) + + 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.""" + + 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, max_epochs=1) + + 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, max_epochs=1) + 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, max_epochs=1) + 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 be empty.""" + records = _records_with_readouts()[:1] + dm = AuxiliaryDataModule(records, ["log2fc_1um"], val_fraction=0.1) + dm.setup() + + assert len(dm.val_dataloader()) == 0 diff --git a/tests/test_cli.py b/tests/test_cli.py index b9263e2..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: @@ -367,6 +368,88 @@ 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_model:\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 + + # 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 + 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" diff --git a/tests/test_concatenation_model.py b/tests/test_concatenation_model.py new file mode 100644 index 0000000..68645f5 --- /dev/null +++ b/tests/test_concatenation_model.py @@ -0,0 +1,234 @@ +"""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 AuxiliaryModelConfig +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 = AuxiliaryModelConfig( + 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_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/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.""" + 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, 3 * n_tasks : 3 * n_tasks + _EMBEDDING_DIM] + flag = features[0, -1] + + assert readout_block[0] == 2.5 + assert list(mask_block) == [1.0, 0.0] + assert not np.all(embedding_block == 0.0) + assert flag == 1.0 + + 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 + + readout_mask_block = features[0, : 2 * n_tasks] + embedding_block = features[0, 3 * n_tasks : 3 * 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 == 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, 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): + """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,) + + 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) diff --git a/tests/test_config.py b/tests/test_config.py new file mode 100644 index 0000000..8c470fd --- /dev/null +++ b/tests/test_config.py @@ -0,0 +1,102 @@ +"""Tests for pipeline configuration loading.""" + +from __future__ import annotations + +from pathlib import Path + +import yaml +from lightning.pytorch.callbacks import EarlyStopping + +from moal.config import AuxiliaryModelConfig, PipelineConfig, TrainerConfig + + +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 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_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_model is None + + def test_round_trips_through_from_yaml(self, tmp_path): + """An explicit auxiliary_model block must populate a matching AuxiliaryModelConfig.""" + path = _write_yaml( + tmp_path, + { + "auxiliary_model": { + "freeze_epochs": 3, + "embedding_dim": 128, + "checkpoint_path": "aux_encoder.pt", + } + }, + ) + + cfg = PipelineConfig.from_yaml(path) + + 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) + + +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 diff --git a/tests/test_model.py b/tests/test_model.py index a274c20..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. @@ -471,7 +476,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 +489,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" ) ), ) @@ -538,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 diff --git a/tests/test_planning.py b/tests/test_planning.py index 54b214d..f10401e 100644 --- a/tests/test_planning.py +++ b/tests/test_planning.py @@ -348,6 +348,88 @@ def test_refit_records_drop_upgraded_interval_ps_rows(self, preprocessor): for r in fit_records ) + 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_1um": -1.2, "pic50": ""}, + {"smiles": "CCN", "relation": ">=", "value": 5.0, "log2fc_1um": 3.4, "pic50": 6.8}, + ) + + state = parse_campaign_state( + df, + cost_ps=1.0, + cost_drc=10.0, + upper_bound=11.0, + preprocessor=preprocessor, + log2fc_columns=["log2fc_1um", "pic50"], + expected_ps_threshold=5.0, + ) + + 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_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( + df, + cost_ps=1.0, + cost_drc=10.0, + upper_bound=11.0, + preprocessor=preprocessor, + log2fc_columns=["log2fc"], + expected_ps_threshold=5.0, + ) + + assert state.training_records[0].raw_ps_readouts == {} + + def test_missing_log2fc_column_raises(self, preprocessor): + """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_columns"): + parse_campaign_state( + df, + cost_ps=1.0, + cost_drc=10.0, + upper_bound=11.0, + preprocessor=preprocessor, + log2fc_columns=["log2fc"], + ) + + 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", + 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_readouts={"log2fc_1um": 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, + ) + + 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_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.""" df = pd.DataFrame( @@ -542,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(