Project new data on existing model

[lucas-diedrich]

Description of feature

Project other datasets onto a trained factor model by freezing the weights of the decoder and re-training the encoder.

Options

1. sc-Arches like architectural surgery
2. Fully retrain encoder network

Option 2

from sccoral.model import SCCORAL
import sccoral
import pandas as pd 
import numpy as np 
import anndata as ad 

from sccoral.train._callbacks import PretrainingFreezeWeights

import torch 

def module_equal_state(m1, m2):
    """ Test if decoder weights remain the same """
    for p1, p2 in zip(m1.parameters(), m2.parameters()):
        if p1.data.ne(p2.data).sum() > 0:
            return False
    return True


adata1 = synthetic_data()
adata2 = synthetic_data(batch_size=20)

SCCORAL.setup_anndata(adata1)
model1 = SCCORAL(adata1, n_latent=5)
model1.train(max_epochs=100, pretraining=False)


model1.setup_anndata(adata2)
model2 = SCCORAL(adata2, n_latent=5)

decoder_state = model1.module.decoder.state_dict()

model2.module.decoder.load_state_dict(decoder_state)

freeze_decoder_weights = PretrainingFreezeWeights(submodule="decoder",
                                                                                               n_pretraining_epochs=1000, 
                                                                                               early_stopping=False, 
                                                                                               lr=1e-3, 
                                                                                               train_batch_norm=False
                                                    )
model2.train(
    max_epochs=100,
    pretraining=False,
    trainer_kwargs={
        'callbacks': [freeze_decoder_weights]
    }
)

module_equal_state(model1.module.decoder.factor_loading, model2.module.decoder.factor_loading)

[lucas-diedrich]

Allow some dimensions of the decoder to adjust while freezing the rest

class LinearDecoder(nn.Module):
    """Linear Decoder :cite:p:~Svensson20"""

    def __init__(
        self,
        n_input: int,
        n_cat: int,
        n_con: int,
        n_output: int,
        n_cat_list: Optional[Iterable[int]] = None,
        use_batch_norm: bool = False,
        use_layer_norm: bool = False,
        bias: bool = False,
        **kwargs,
    ):
        super().__init__()
        self.n_input = n_input
        self.n_cat = n_cat if isinstance(n_cat, int) else None
        self.n_con = n_con if isinstance(n_con, int) else None

        self.factor_loading = FCLayers(
            n_in=n_input,
            n_out=n_output,
            n_cat_list=n_cat_list,
            n_layers=1,
            use_activation=False,
            use_batch_norm=use_batch_norm,  # None
            use_layer_norm=use_layer_norm,
            bias=bias,
            dropout_rate=0,
            **kwargs,
        )

        self.factor_loading_cat = FCLayers(
            n_in=n_cat,
            n_out=n_output,
            n_cat_list=n_cat_list,
            n_layers=1,
            use_activation=False,
            use_batch_norm=use_batch_norm,  # None
            use_layer_norm=use_layer_norm,
            bias=bias,
            dropout_rate=0,
            **kwargs,
        )

        self.factor_loading_con = FCLayers(
            n_in=n_con,
            n_out=n_output,
            n_cat_list=n_cat_list,
            n_layers=1,
            use_activation=False,
            use_batch_norm=use_batch_norm,  # None
            use_layer_norm=use_layer_norm,
            bias=bias,
            dropout_rate=0,
            **kwargs,
        )



        self.px_dropout_decoder = FCLayers(
            n_in=n_input+n_cat+n_con,
            n_out=n_output,
            n_cat_list=n_cat_list,  # None
            n_layers=1,
            use_activation=False,
            use_batch_norm=use_batch_norm,
            use_layer_norm=use_layer_norm,
            bias=bias,
            dropout_rate=0,
            **kwargs,
        )

    def forward(self, dispersion: str, z: torch.Tensor, library: torch.Tensor):
        raw_px_scale_input = self.factor_loading(z[:, :self.n_input])
        raw_px_scale = raw_px_scale_input

        raw_px_scale_cat = None
        if (self.n_cat is not None) and (self.n_cat > 0):
            if (self.n_con is None) or (self.n_con == 0): 
                raw_px_scale_cat = self.factor_loading_cat(z[:, -self.n_cat:])
            else:
                raw_px_scale_cat = self.factor_loading_cat(z[:, -self.n_cat-self.n_con:-self.n_con])
            
            raw_px_scale = torch.add(raw_px_scale, raw_px_scale_cat)

        if (self.n_con is not None) & (self.n_con > 0):
            raw_px_scale_con = self.factor_loading_con(z[:, -self.n_con:])
            raw_px_scale = torch.add(raw_px_scale, raw_px_scale_con)

    
        px_scale = torch.softmax(raw_px_scale, dim=-1)
        px_dropout = self.px_dropout_decoder(z)
        px_rate = torch.exp(library) * px_scale
        px_r = None

        return px_scale, px_r, px_rate, px_dropout