Batch export for posteriors

[sezginerr]

Hello,

I have been trying to export posteriors batch by batch with a defined batch size, rather than exporting them in a single batch. This is achieved by concatenating local variables exported from batches along the cell dimension in the scvi library. I believe the scvi library assumes that the variables defined with 'obs_plate' are the local variables, while the global variables (any other variable) are sampled from the last batch. However, for the cell2fate model, local and global variables seem to the same, as the guide method in the cell2fate model is an instance of 'poutine.messenger.Messenger'. In this case, posterior sampling does not use the 'return_sites' information and seems to return everything. Here is the function that returns one posterior sampling in the scvi library:

def _get_one_posterior_sample(
    self,
    args,
    kwargs,
    return_sites: Optional[list] = None,
    return_observed: bool = False,
  ):
    if isinstance(self.module.guide, poutine.messenger.Messenger):
        # This already includes trace-replay behavior.
        sample = self.module.guide(*args, **kwargs)
    else:
        guide_trace = poutine.trace(self.module.guide).get_trace(*args, **kwargs)
        model_trace = poutine.trace(
            poutine.replay(self.module.model, guide_trace)
        ).get_trace(*args, **kwargs)
sample = {
            name: site["value"]
            for name, site in model_trace.nodes.items()
            if (
                (site["type"] == "sample")  # sample statement
                and (
                    (return_sites is None) or (name in return_sites)
                )  # selected in return_sites list
                and (
                    (
                        (not site.get("is_observed", True)) or return_observed
                    )  # don't save observed unless requested
                    or (site.get("infer", False).get("_deterministic", False))
                )  # unless it is deterministic
                and not isinstance(
                    site.get("fn", None), poutine.subsample_messenger._Subsample
                )  # don't save plates
            )
        } 

To adapt batch sampling in the cell2fate model, I made a few changes to the function. Firstly, I removed the global sampling, which would sample everything in the last batch again. I'm not sure if it's correct to do so, as the RF/GO analysis slightly changes when I use single batch sampling with or without global variable sampling. Secondly, I defined the variables that have cell number dimensions and concatenated them in the cell number dimension after sampling. For the other variables, I simply applied element-wise averaging between batches. Here is the adapted function that I use:

Details

def _posterior_samples_minibatch(
        self, use_gpu: bool = None, batch_size: Optional[int] = None, **sample_kwargs
    ):
        samples = dict()

        _, device = parse_use_gpu_arg(use_gpu)

        batch_size = batch_size if batch_size is not None else settings.batch_size

        train_dl = AnnDataLoader(
            self.adata_manager, shuffle=False, batch_size=batch_size
        )
        # sample local parameters
        i = 0
        cell_specific=['t_c', 'T_c', 'mu_expression', 'detection_y_c', 'mu', 'data_target']
        for tensor_dict in track(
            train_dl,
            style="tqdm",
            description="Sampling local variables, batch: ",
        ):
            args, kwargs = self.module._get_fn_args_from_batch(tensor_dict)
            args = [a.to(device) for a in args]
            kwargs = {k: v.to(device) for k, v in kwargs.items()}
            self.to_device(device)
            if i == 0:
                return_observed = getattr(sample_kwargs, "return_observed", False)
                obs_plate_sites = self._get_obs_plate_sites(
                    args, kwargs, return_observed=return_observed
                )
                if len(obs_plate_sites) == 0:
                    # if no local variables - don't sample
                    break
                obs_plate_dim = list(obs_plate_sites.values())[0]
                sample_kwargs_obs_plate = sample_kwargs.copy()
                sample_kwargs_obs_plate[
                    "return_sites"
                ] = self._get_obs_plate_return_sites(
                    sample_kwargs["return_sites"], list(obs_plate_sites.keys())
                )

                sample_kwargs_obs_plate["show_progress"] = False

                samples = self._get_posterior_samples(
                    args, kwargs, **sample_kwargs_obs_plate
                )
            else:
                samples_ = self._get_posterior_samples(
                    args, kwargs, **sample_kwargs_obs_plate
                )

                num_cells_in_batch = samples_['t_c'].shape[1]
                for k in samples.keys():

                    if samples_[k].ndim >1:
                        if k in cell_specific:
                            samples[k] =  np.concatenate([samples[k], samples_[k]], axis=1)
                    else:
                        ratio_cells = num_cells_in_batch / batch_size
                        samples[k] = (samples[k] * i + samples_[k] * ratio_cells) / (i+ratio_cells)

                i += 1

            i += 1

         self.module.to(device)

        return samples

<\details>

This approach seems to work well, as the results with batch export and without batch export are quite similar. However, there is again a slight difference between RF/GO analyses. I believe that difference arises because, even though the seed is set to 0, a change in the sampling strategy will alter the random function's output. I also tried to sample the variables that are not cell-specific from only the last batch, but it does not seem to work well, even though it should work theoretically, as the variables that are not cell-specific are sampled from already learned distributions.

I wanted to ask if my approach in line with the cell2fate assumption. Thank you very much for your help.

[vitkl]

Would be great to add this to scvi-tools directly. Sounds like this is solving an important issue.

[AlexanderAivazidis]

Great! Let's add this particular solution to cell2fate only, because it relies on specifying "cell-specific variables", which may be different for each model. You can make a branch that changes the export_posterior method of the cell2fate_DynamicalModel with your modifications and then make a pull request.
After that we can indeed think about how to add a more general solution to scvi-tools directly!

[vitkl]

Why do you need this bit of code?

cell_specific=['t_c', 'T_c', 'mu_expression', 'detection_y_c', 'mu', 'data_target']

The _get_obs_plate_sites code below should already allow detecting which variables are cell-specific:

obs_plate_sites = self._get_obs_plate_sites(
    args, kwargs, return_observed=return_observed
)

element-wise averaging between batches

This is not correct to do for global variables. This breaks the independence of samples and probably yields very narrow distributions (distribution of posterior samples -> distributions of means). Global variables are sampled using a separate step exactly because we need to get them once using one data batch (data batch is irrelevant for these variables, their distribution is the same for all batches).

In this case, posterior sampling does not use the 'return_sites' information and seems to return everything.

This is a known issue that needs to be solved. I initially thought that this what this addresses @sezginerr.

However, for the cell2fate model, local and global variables seem to the same

Is this really the case? Why is this the case? Statements with obs_plate make variables local variables

cell2fate/cell2fate/_cell2fate_DynamicalModel_module.py

Line 427 in 22917f7

with obs_plate:

and list_obs_plate_vars method

cell2fate/cell2fate/_cell2fate_DynamicalModel_module.py

Line 337 in 22917f7

def list_obs_plate_vars(self):

tells scvi-tools which plate is the minibatch plate.

Does the default posterior sampling code fail completely or simply "does not use the 'return_sites' information and seems to return everything"?

[sezginerr]

The _get_obs_plate_sites method currently returns only the variables defined with obs_plate: and excludes deterministic ones, resulting in only t_c and detection_y_c being returned. However, the remaining deterministic variables (T_c, mu_expression, mu, data_target) also have a cell number dimension that should be concatenated in each batch. Instead of adding them manually, I can also check each variable in the first batch and create a dictionary of keys that have a shape equal to the batch size in axis=1 to create cell specific variable dictionary.

Yes, calculating the mean probably causes a narrow distribution. I checked the code again, and I believe there was a bug in my first code. Now it seems that getting globals from the last batch seems to work. However, there are still slight differences between batch and non-batch exports.

While both local and global variables are defined in cell2fate, the problem lies in the posterior sampling process, which disregards the return_sites information coming from _get_obs_plate_sites, causing duplicate returns of all variables, as you mentioned. So it is not that the variables are same but the code does not use that information.

[vitkl]

Are the deterministic variables wrapped into with obs_plate in the model? They should be, otherwise, they will be ignored by plate detection.

[vitkl]

I have some code to solve these issues once and for all. It also handles multiple plates and arbitrary orders of items within plates. My code works for direct posterior median computation - but not for sampling. It needs some minor modifications and testing. Would you be able to help?

[sezginerr]

No, deterministic variables are not wrapped into with obs_plate in cell2fate. See this for example:

cell2fate/cell2fate/_cell2fate_DynamicalModel_module.py

Lines 427 to 429 in 22917f7

	with obs_plate:
	t_c = pyro.sample('t_c', dist.Normal(t_c_loc, t_c_scale).expand([batch_size, 1, 1]))
	T_c = pyro.deterministic('T_c', t_c*T_max)

Yes, I would like to help with the code. I believe it would be great if there is a general solution for this.

[AlexanderAivazidis]

Thanks for your insights on this!
So for a start I will wrap all cell specific deterministic variables in a plate, so we can detect them as local variables automatically.

Next we should definitely implement Vitalii's solution to do direct posterior median computation.

Finally, I think Sezgin's solution is probably still ok for the moment, because the only time we use something other then the posterior mean is for the time variable "T_c", where we are also interested in the variance and that should be sampled accurately given it is a local variable.

[AlexanderAivazidis]

I wrapped cell specific deterministic variables in plates now, but this has not solved it I think, because for some reason "data_target" is not detected as a local variable. So I think it would be great if @sezginerr adds his approximate solution for now and then we come up with something more general next.

[vitkl]

data_target is an observed variable, right? Not much reason to track it.

[vitkl]

Here is my code for computing the quantiles and medians. It should work for this project with no modifications, except for dataloader options.

This code needs to be modified with sampling statements instead of median/quantiles, including correctly handling the sample dimension. I would really appreciate having the posterior sampling solution for the cell2state model project in addition to this problem.

Details

import numpy as np
import torch
from pyro import poutine
from scvi.dataloaders import AnnDataLoader
from scvi.model._utils import parse_use_gpu_arg
from scvi.utils import track

from cell2state.dataloaders._per_gene_per_cell_chr_dataloader import (
    PerGeneChromatinDataSplitter,
)


def expand_zeros_along_dim(tensor, size, dim):
    shape = np.array(tensor.shape)
    shape[dim] = size
    return np.zeros(shape)


def complete_tensor_along_dim(tensor, indices, dim, value, mode="put"):
    shape = value.shape
    shape = np.ones(len(shape))
    shape[dim] = len(indices)
    shape = shape.astype(int)
    indices = indices.reshape(shape)
    if mode == "take":
        return np.take_along_axis(arr=tensor, indices=indices, axis=dim)
    np.put_along_axis(arr=tensor, indices=indices, values=value, axis=dim)
    return tensor


def _complete_full_tensors_using_plates(
    means_global, means, plate_dict, obs_plate_sites, plate_indices, plate_dim
):
    # complete full sized tensors with minibatch values given minibatch indices
    for k in means_global.keys():
        # find which and how many plates contain this tensor
        plates = [
            plate for plate in plate_dict.keys() if k in obs_plate_sites[plate].keys()
        ]
        if len(plates) == 1:
            # if only one plate contains this tensor, complete it using the plate indices
            means_global[k] = complete_tensor_along_dim(
                means_global[k],
                plate_indices[plates[0]],
                plate_dim[plates[0]],
                means[k],
            )
        elif len(plates) == 2:
            # subset data to index for plate 0 and fill index for plate 1
            means_global_k = complete_tensor_along_dim(
                means_global[k],
                plate_indices[plates[0]],
                plate_dim[plates[0]],
                means[k],
                mode="take",
            )
            means_global_k = complete_tensor_along_dim(
                means_global_k,
                plate_indices[plates[1]],
                plate_dim[plates[1]],
                means[k],
            )
            # fill index for plate 0 in the full data
            means_global[k] = complete_tensor_along_dim(
                means_global[k],
                plate_indices[plates[0]],
                plate_dim[plates[0]],
                means_global_k,
            )
            # TODO add a test - observed variables should be identical if this code works correctly
            # This code works correctly but the test needs to be added eventually
            # np.allclose(
            #     samples['data_chromatin'].squeeze(-1).T,
            #     mod_reg.adata_manager.get_from_registry('X')[
            #         :, ~mod_reg.adata_manager.get_from_registry('gene_bool').ravel()
            #     ].toarray()
            # )
        else:
            NotImplementedError(
                f"Posterior sampling/mean/median/quantile not supported for variables with > 2 plates: {k} has {len(plates)}"
            )
    return means_global


class QuantileMixin:
    """
    This mixin class provides methods for:

    - computing median and quantiles of the posterior distribution using both direct and amortised inference

    """

    def _get_obs_plate_sites_v2(
        self,
        args: list,
        kwargs: dict,
        plate_name: str = None,
        return_observed: bool = False,
        return_deterministic: bool = True,
    ):
        """
        Automatically guess which model sites belong to observation/minibatch plate.
        This function requires minibatch plate name specified in `self.module.list_obs_plate_vars["name"]`.
        Parameters
        ----------
        args
            Arguments to the model.
        kwargs
            Keyword arguments to the model.
        return_observed
            Record samples of observed variables.
        Returns
        -------
        Dictionary with keys corresponding to site names and values to plate dimension.
        """
        if plate_name is None:
            plate_name = self.module.list_obs_plate_vars["name"]

        def try_trace(args, kwargs):
            try:
                trace_ = poutine.trace(self.module.guide).get_trace(*args, **kwargs)
                trace_ = poutine.trace(
                    poutine.replay(self.module.model, trace_)
                ).get_trace(*args, **kwargs)
            except ValueError:
                # if sample is unsuccessful try again
                trace_ = try_trace(args, kwargs)
            return trace_

        trace = try_trace(args, kwargs)

        # find plate dimension
        obs_plate = {
            name: {
                fun.name: fun
                for fun in site["cond_indep_stack"]
                if (fun.name in plate_name) or (fun.name == plate_name)
            }
            for name, site in trace.nodes.items()
            if (
                (site["type"] == "sample")  # sample statement
                and (
                    (
                        (not site.get("is_observed", True)) or return_observed
                    )  # don't save observed unless requested
                    or (
                        site.get("infer", False).get("_deterministic", False)
                        and return_deterministic
                    )
                )  # unless it is deterministic
                and not isinstance(
                    site.get("fn", None), poutine.subsample_messenger._Subsample
                )  # don't save plates
            )
            if any(f.name == plate_name for f in site["cond_indep_stack"])
        }

        return obs_plate

    def _get_dataloader(
        self,
        batch_size,
        gene_batch_size,
        data_loader_indices,
        dna_region_batch_mode="windows",
        shuffle_training=False,
        dl_kwargs={},
    ):
        if self.minibatch_genes_:
            dl_kwargs_ = {
                "shuffle_training": shuffle_training,
                "train_size": 1.0,
                "batch_size": batch_size,
                "gene_batch_size": gene_batch_size,
                "cell_plate_inputs": self.module.model.cell_plate_inputs,
                "var_plate_inputs": self.module.model.var_plate_inputs,
                "square_var_plate_inputs": self.module.model.square_var_plate_inputs,
                "per_site_plate_inputs": self.module.model.per_site_plate_inputs,
                "filter_tensors": self.module.model.filter_tensors,
                "gene_bool": self.adata_manager.get_from_registry(
                    registry_key="gene_bool"
                ).flatten(),
                "drop_last": False,
                **dl_kwargs,
            }
            if dna_region_batch_mode == "gene_batch":
                dl_kwargs_["gene_region_coo"] = self._get_gene_region_coo(
                    use_all_genes=False
                )
            elif dna_region_batch_mode == "windows":
                dl_kwargs_["positions"] = self.adata_manager.get_from_registry(
                    registry_key="positions"
                ).flatten()
                dl_kwargs_["chromosomes"] = self.adata_manager.get_from_registry(
                    registry_key="chromosome"
                ).flatten()
            train_dl = PerGeneChromatinDataSplitter(
                self.adata_manager,
                **dl_kwargs_,
            )
            train_dl.setup()
            train_dl = train_dl.train_dataloader()
        else:
            train_dl = AnnDataLoader(
                self.adata_manager,
                shuffle=False,
                batch_size=batch_size,
                indices=data_loader_indices,
                **dl_kwargs,
            )
        return train_dl

    @torch.no_grad()
    def _posterior_quantile_minibatch(
        self,
        q: float = 0.5,
        batch_size: int = 128,
        gene_batch_size: int = 50,
        use_gpu: bool = None,
        use_median: bool = False,
        return_observed: bool = True,
        exclude_vars: list = None,
        data_loader_indices=None,
        show_progress: bool = True,
        dna_region_batch_mode: str = "gene_batch",  # ['gene_batch', 'windows'],
        n_vars_per_dna_region_batch: int = 420,
    ):
        """
        Compute median of the posterior distribution of each parameter, separating local (minibatch) variable
        and global variables, which is necessary when performing amortised inference.

        Note for developers: requires model class method which lists observation/minibatch plate
        variables (self.module.model.list_obs_plate_vars()).

        Parameters
        ----------
        q
            quantile to compute
        batch_size
            number of observations per batch
        use_gpu
            Bool, use gpu?
        use_median
            Bool, when q=0.5 use median rather than quantile method of the guide

        Returns
        -------
        dictionary {variable_name: posterior quantile}

        """

        _, _, device = parse_use_gpu_arg(use_gpu)

        self.module.eval()

        # Warmup run for one step - to eg cache TF effects
        dl = self._get_dataloader(
            batch_size=batch_size,
            gene_batch_size=gene_batch_size,
            data_loader_indices=data_loader_indices,
            dna_region_batch_mode=dna_region_batch_mode,
            dl_kwargs={
                "n_vars_per_dna_region_batch": n_vars_per_dna_region_batch,
            },
        )
        for tensor_dict in dl:
            args, kwargs = self.module._get_fn_args_from_batch(tensor_dict)
            args = [a.to(device) for a in args]
            kwargs = {k: v.to(device) for k, v in kwargs.items()}
            self.to_device(device)
            self.module.train()  # to cache TF effects TODO: remove this hack as soon as possible
            self.module.guide(*args, **kwargs)
            self.module.eval()
            break

        train_dl = self._get_dataloader(
            batch_size=batch_size,
            gene_batch_size=gene_batch_size,
            data_loader_indices=data_loader_indices,
            dna_region_batch_mode=dna_region_batch_mode,
            dl_kwargs={
                "n_vars_per_dna_region_batch": n_vars_per_dna_region_batch,
            },
        )

        i = 0
        for tensor_dict in track(
            train_dl,
            style="tqdm",
            description=f"Computing posterior quantile {q}, data batch: ",
            disable=not show_progress,
        ):
            args, kwargs = self.module._get_fn_args_from_batch(tensor_dict)
            args = [a.to(device) for a in args]
            kwargs = {k: v.to(device) for k, v in kwargs.items()}
            self.to_device(device)

            if i == 0:
                minibatch_plate_names = self.module.list_obs_plate_vars["name"]
                plates = self.module.model.create_plates(*args, **kwargs)
                if not isinstance(plates, list):
                    plates = [plates]
                # find plate indices & dim
                plate_dict = {
                    plate.name: plate
                    for plate in plates
                    if (
                        (plate.name in minibatch_plate_names)
                        or (plate.name == minibatch_plate_names)
                    )
                }
                plate_size = {name: plate.size for name, plate in plate_dict.items()}
                if data_loader_indices is not None:
                    # set total plate size to the number of indices in DL not total number of observations
                    # this option is not really used
                    plate_size = {
                        name: len(train_dl.indices)
                        for name, plate in plate_dict.items()
                        if plate.name == minibatch_plate_names
                    }
                plate_dim = {name: plate.dim for name, plate in plate_dict.items()}
                plate_indices = {
                    name: plate.indices.detach().cpu().numpy()
                    for name, plate in plate_dict.items()
                }
                # find plate sites
                obs_plate_sites = {
                    plate: self._get_obs_plate_sites_v2(
                        args, kwargs, plate_name=plate, return_observed=return_observed
                    )
                    for plate in plate_dict.keys()
                }
                if use_median and q == 0.5:
                    # use median rather than quantile method
                    def try_median(args, kwargs):
                        try:
                            means_ = self.module.guide.median(*args, **kwargs)
                        except ValueError:
                            # if sample is unsuccessful try again
                            means_ = try_median(args, kwargs)
                        return means_

                    means = try_median(args, kwargs)
                else:

                    def try_quantiles(args, kwargs):
                        try:
                            means_ = self.module.guide.quantiles([q], *args, **kwargs)
                        except ValueError:
                            # if sample is unsuccessful try again
                            means_ = try_quantiles(args, kwargs)
                        return means_

                    means = try_quantiles(args, kwargs)
                means = {
                    k: means[k].detach().cpu().numpy()
                    for k in means.keys()
                    if k not in exclude_vars
                }
                means_global = means.copy()
                for plate in plate_dict.keys():
                    # create full sized tensors according to plate size
                    means_global = {
                        k: (
                            expand_zeros_along_dim(
                                means_global[k], plate_size[plate], plate_dim[plate]
                            )
                            if k in obs_plate_sites[plate].keys()
                            else means_global[k]
                        )
                        for k in means_global.keys()
                    }
                # complete full sized tensors with minibatch values given minibatch indices
                means_global = _complete_full_tensors_using_plates(
                    means_global=means_global,
                    means=means,
                    plate_dict=plate_dict,
                    obs_plate_sites=obs_plate_sites,
                    plate_indices=plate_indices,
                    plate_dim=plate_dim,
                )
                if np.all([len(v) == 0 for v in obs_plate_sites.values()]):
                    # if no local variables - don't sample further - return results now
                    break
            else:
                if use_median and q == 0.5:

                    def try_median(args, kwargs):
                        try:
                            means_ = self.module.guide.median(*args, **kwargs)
                        except ValueError:
                            # if sample is unsuccessful try again
                            means_ = try_median(args, kwargs)
                        return means_

                    means = try_median(args, kwargs)
                else:

                    def try_quantiles(args, kwargs):
                        try:
                            means_ = self.module.guide.quantiles([q], *args, **kwargs)
                        except ValueError:
                            # if sample is unsuccessful try again
                            means_ = try_quantiles(args, kwargs)
                        return means_

                    means = try_quantiles(args, kwargs)
                means = {
                    k: means[k].detach().cpu().numpy()
                    for k in means.keys()
                    if k not in exclude_vars
                }
                # find plate indices & dim
                plate_dict = {
                    plate.name: plate
                    for plate in self.module.model.create_plates(*args, **kwargs)
                    if (
                        (plate.name in minibatch_plate_names)
                        or (plate.name == minibatch_plate_names)
                    )
                }
                plate_indices = {
                    name: plate.indices.detach().cpu().numpy()
                    for name, plate in plate_dict.items()
                }
                # TODO - is this correct to call this function again? find plate sites
                obs_plate_sites = {
                    plate: self._get_obs_plate_sites_v2(
                        args, kwargs, plate_name=plate, return_observed=return_observed
                    )
                    for plate in plate_dict.keys()
                }
                # complete full sized tensors with minibatch values given minibatch indices
                means_global = _complete_full_tensors_using_plates(
                    means_global=means_global,
                    means=means,
                    plate_dict=plate_dict,
                    obs_plate_sites=obs_plate_sites,
                    plate_indices=plate_indices,
                    plate_dim=plate_dim,
                )
            i += 1

        self.module.to(device)

        return means_global

<\details>