Train API¶

def __init__(
    self,
    source_params: dict[str, dict],
    central_params: dict,
    classif_params: dict,
    surv_params: dict,
    train_params: dict,
    device: torch.device | None = None,
) -> None:
    """Model initialization.

    Args:
        source_params:
            Per-source configuration.  Keys are source names; each value is a dict
            with the following keys:

            * `input_dim` (int): number of input features.
            * `hidden_dim` (list of int): hidden layer sizes.
            * `latent_dim` (int): per-source latent dimension.
            * `norm` (bool): whether to use batch normalisation.
            * `dropout` (float): dropout rate in `[0, 1]`.

        central_params:
            Central VAE configuration.  Required keys: `hidden_dim` (list of
            int), `latent_dim` (int), `norm` (bool), `dropout` (float),
            `beta` (float — MMD regularisation weight).

        classif_params:
            Classifier configuration.  Required keys: `n_class` (int, >= 2),
            `lambda` (float — loss weight), `hidden_layers` (list of int),
            `dropout` (float).

        surv_params:
            Survival-predictor configuration.  Required keys: `lambda` (float),
            `dims` (list of int), `activation` (str), `l2_reg` (float),
            `norm` (bool), `dropout` (float).

        train_params:
            Training hyperparameters.  Required keys: `switch` (int — epoch at
            which to enter phase 2) and `lr` (float — learning rate).

        device:
            Torch compute device.
    """
    super().__init__()
    self._validate_params(source_params, central_params, classif_params, train_params)

    self.device = _parse_device(device)
    self.source_names: list[str] = list(source_params.keys())
    self.n_source = len(self.source_names)
    self.beta = central_params["beta"]
    self.num_classes = classif_params["n_class"]
    self.lambda_classif = classif_params["lambda"]
    self.lambda_survival = surv_params["lambda"]
    self.switch_epoch = train_params["switch"]
    self.lr = train_params["lr"]
    self.phase = 1
    self._is_fitted = False

    # Store config dicts so save() can serialise them.
    self._source_params = source_params
    self._central_params = central_params
    self._classif_params = classif_params
    self._surv_params = surv_params
    self._train_params = train_params

    # ------------------------------------------------------------------ #
    # Per-source autoencoders (nn.ModuleList so PyTorch tracks them)
    # ------------------------------------------------------------------ #
    self.autoencoders = nn.ModuleList([
        AutoEncoder(
            encoder=Encoder(
                input_dim=source_params[s]["input_dim"],
                hidden_dim=source_params[s]["hidden_dim"],
                latent_dim=source_params[s]["latent_dim"],
                norm_layer=source_params[s]["norm"],
                dropout=source_params[s]["dropout"],
            ),
            decoder=Decoder(
                latent_dim=source_params[s]["latent_dim"],
                hidden_dim=source_params[s]["hidden_dim"],
                output_dim=source_params[s]["input_dim"],
                norm_layer=source_params[s]["norm"],
                dropout=source_params[s]["dropout"],
            ),
            device=self.device,
        )
        for s in self.source_names
    ])

    # ------------------------------------------------------------------ #
    # Central VAE
    # ------------------------------------------------------------------ #
    self.rep_dim = sum(source_params[s]["latent_dim"] for s in self.source_names)
    self.central_layer = VAE(
        encoder=ProbabilisticEncoder(
            input_dim=self.rep_dim,
            hidden_dim=central_params["hidden_dim"],
            latent_dim=central_params["latent_dim"],
            norm_layer=central_params["norm"],
            dropout=central_params["dropout"],
        ),
        decoder=ProbabilisticDecoder(
            latent_dim=central_params["latent_dim"],
            hidden_dim=central_params["hidden_dim"],
            output_dim=self.rep_dim,
            norm_layer=central_params["norm"],
            dropout=central_params["dropout"],
        ),
        device=self.device,
    )

    # ------------------------------------------------------------------ #
    # Task heads
    # ------------------------------------------------------------------ #
    self.classifier = MultiClassifier(
        n_class=self.num_classes,
        latent_dim=central_params["latent_dim"],
        dropout=classif_params["dropout"],
        class_dim=classif_params["hidden_layers"],
    )
    self.survival_predictor = SurvivalNet({
        "drop": surv_params["dropout"],
        "norm": surv_params["norm"],
        "dims": [central_params["latent_dim"]] + surv_params["dims"] + [1],
        "activation": surv_params["activation"],
    })

    self._relocate()
    self.optimizer = self._build_optimizer()

    # Filled during fit()
    self.history: list[tuple] = []
    self.label_encoder: LabelEncoder | None = None
    self.one_hot_encoder: OneHotEncoder | None = None
    self.baseline = None

def fit(
    self,
    mdata: MuData,
    omics_val: MuData | None = None,
    batch_size: int = 32,
    n_epochs: int = 30,
    verbose: bool = True,
) -> None:
    """Train the customics model.

    Args:
        mdata: Multi-omics object whose `obs` holds the clinical annotations.
        omics_val: Validation omics data; same format as `omics_train`.
        batch_size: Mini-batch size.
        n_epochs: Number of training epochs.
        verbose: Log epoch-level loss when True.

    Raises:
        DataValidationError: If required columns are missing or samples don't overlap.
    """

    label, event, surv_time = self._validate_fit_inputs(mdata)

    self.label_encoder = LabelEncoder().fit(mdata.obs[label].values)
    train_labels = pd.Series(self.label_encoder.transform(mdata.obs[label].values), index=mdata.obs_names)
    # Fit OHE on integer-encoded labels so it can transform integer y_true at eval time
    self.one_hot_encoder = OneHotEncoder(sparse_output=False).fit(train_labels.values.reshape(-1, 1))

    loader_kw: dict = {"num_workers": 2, "pin_memory": True} if self.device.type == "cuda" else {}

    shared_samples_train = get_shared_samples(mdata)
    self.baseline = self._compute_baseline(mdata.obs, shared_samples_train, event, surv_time)
    train_loader = DataLoader(
        MultiOmicsDataset(mdata, shared_samples_train, train_labels),
        batch_size=batch_size,
        shuffle=True,
        **loader_kw,
    )

    val_loader: DataLoader | None = None
    if omics_val is not None:
        shared_samples_val = get_shared_samples(omics_val)
        val_labels = pd.Series(self.label_encoder.transform(omics_val.obs[label].values), index=omics_val.obs_names)
        val_loader = DataLoader(
            MultiOmicsDataset(omics_val, shared_samples_val, val_labels),
            batch_size=batch_size,
            shuffle=False,
            **loader_kw,
        )

    self.history = []
    for epoch in range(n_epochs):
        self._switch_phase(epoch)
        train_loss = self._run_epoch(train_loader, training=True)
        if val_loader is not None:
            val_loss = self._run_epoch(val_loader, training=False)
            self.history.append((train_loss, val_loss))
            if verbose:
                logger.info(
                    "Epoch %d/%d | train=%.4f | val=%.4f",
                    epoch + 1,
                    n_epochs,
                    train_loss,
                    val_loss,
                )
        else:
            self.history.append((train_loss,))
            if verbose:
                logger.info("Epoch %d/%d | train=%.4f", epoch + 1, n_epochs, train_loss)

    self._is_fitted = True

def plot_loss(self, show: bool = True, figsize: tuple[float, float] = (6, 3)) -> None:
    """Plot training (and validation) loss history.

    Args:
        show: Display the figure interactively.
        figsize: Figure size.
    """
    from customics.visualization import plot_loss as _plot

    _plot(self.history, self.switch_epoch, show=show, figsize=figsize)

def get_number_parameters(self) -> int:
    """Return the total number of trainable parameters.

    Returns:
        Parameter count.
    """
    return sum(p.numel() for p in self.parameters() if p.requires_grad)

def save(self, path: str | Path) -> None:
    """Save the model architecture config and trained weights to *path*.

    The checkpoint contains the five parameter dicts needed to reconstruct
    the model plus the `state_dict` and fitted encoders so that inference
    works immediately after [load](api/train/#customics.CustOMICS.load).

    Args:
        path: Destination file (conventionally `*.pt` or `*.pth`).
    """
    checkpoint = {
        "source_params": self._source_params,
        "central_params": self._central_params,
        "classif_params": self._classif_params,
        "surv_params": self._surv_params,
        "train_params": self._train_params,
        "state_dict": self.state_dict(),
        "label_encoder": self.label_encoder,
        "one_hot_encoder": self.one_hot_encoder,
        "history": self.history,
        "_is_fitted": self._is_fitted,
    }

    checkpoint_path = Path(path)
    checkpoint_path.parent.mkdir(parents=True, exist_ok=True)

    torch.save(checkpoint, path)
    logger.info(f"Model saved to {path}")

@classmethod
def load(cls, path: str | Path, device: torch.device | None = None) -> CustOMICS:
    """Load a model previously saved with [save](api/train/#customics.CustOMICS.save).

    Args:
        path: Path to the checkpoint file written by [save](api/train/#customics.CustOMICS.save).
        device: Target device.  Defaults to CPU when not specified.

    Returns:
        A fully initialised, ready-to-use model instance.
    """
    device = _parse_device(device)

    checkpoint = torch.load(path, map_location=device, weights_only=False)
    model = cls(
        source_params=checkpoint["source_params"],
        central_params=checkpoint["central_params"],
        classif_params=checkpoint["classif_params"],
        surv_params=checkpoint["surv_params"],
        train_params=checkpoint["train_params"],
        device=device,
    )
    model.load_state_dict(checkpoint["state_dict"])
    model.label_encoder = checkpoint.get("label_encoder")
    model.one_hot_encoder = checkpoint.get("one_hot_encoder")
    model.history = checkpoint.get("history", [])
    model._is_fitted = checkpoint.get("_is_fitted", True)
    model.to(device)
    logger.info(f"Model loaded from {path}")
    return model