# Buggy code
def get_bayes_factors(
    self,
    idx1: Union[List[bool], np.ndarray],
    idx2: Union[List[bool], np.ndarray],
    mode: Literal["vanilla", "change"] = "vanilla",
    batchid1: Optional[Union[List[int], np.ndarray]] = None,
    batchid2: Optional[Union[List[int], np.ndarray]] = None,
    use_observed_batches: Optional[bool] = False,
    n_samples: int = 5000,
    use_permutation: bool = False,
    m_permutation: int = 10000,
    change_fn: Optional[Union[str, Callable]] = None,
    m1_domain_fn: Optional[Callable] = None,
    delta: Optional[float] = 0.5,
    cred_interval_lvls: Optional[Union[List[float], np.ndarray]] = None,
) -> Dict[str, np.ndarray]:
    r"""
    A unified method for differential expression inference.

    Two modes coexist:

    - the `"vanilla"` mode follows protocol described in [Lopez18]_ and [Xu19]_
    In this case, we perform hypothesis testing based on the hypotheses

    .. math::
        M_1: h_1 > h_2 ~\text{and}~ M_2: h_1 \leq h_2.

    DE can then be based on the study of the Bayes factors

    .. math::
        \log p(M_1 | x_1, x_2) / p(M_2 | x_1, x_2).

    - the `"change"` mode (described in [Boyeau19]_)
    This mode consists of estimating an effect size random variable (e.g., log fold-change) and
    performing Bayesian hypothesis testing on this variable.
    The `change_fn` function computes the effect size variable :math:`r` based on two inputs
    corresponding to the posterior quantities (e.g., normalized expression) in both populations.

    Hypotheses:

    .. math::
        M_1: r \in R_1 ~\text{(effect size r in region inducing differential expression)}

    .. math::
        M_2: r  \notin R_1 ~\text{(no differential expression)}

    To characterize the region :math:`R_1`, which induces DE, the user has two choices.

    1. A common case is when the region :math:`[-\delta, \delta]` does not induce differential
       expression. If the user specifies a threshold delta, we suppose that :math:`R_1 = \mathbb{R} \setminus [-\delta, \delta]`
    2. Specify an specific indicator function:

    .. math::
        f: \mathbb{R} \mapsto \{0, 1\} ~\text{s.t.}~ r \in R_1 ~\text{iff.}~ f(r) = 1.

    Decision-making can then be based on the estimates of

    .. math::
        p(M_1 \mid x_1, x_2).

    Both modes require to sample the posterior distributions.
    To that purpose, we sample the posterior in the following way:

    1. The posterior is sampled `n_samples` times for each subpopulation.
    2. For computational efficiency (posterior sampling is quite expensive), instead of
       comparing the obtained samples element-wise, we can permute posterior samples.
       Remember that computing the Bayes Factor requires sampling :math:`q(z_A \mid x_A)` and :math:`q(z_B \mid x_B)`.

    Currently, the code covers several batch handling configurations:

    1. If ``use_observed_batches=True``, then batch are considered as observations
       and cells' normalized means are conditioned on real batch observations.
    2. If case (cell group 1) and control (cell group 2) are conditioned on the same
       batch ids. This requires ``set(batchid1) == set(batchid2)`` or ``batchid1 == batchid2 === None``.
    3. If case and control are conditioned on different batch ids that do not intersect
       i.e., ``set(batchid1) != set(batchid2)`` and ``len(set(batchid1).intersection(set(batchid2))) == 0``.

    This function does not cover other cases yet and will warn users in such cases.

    Parameters
    ----------
    mode
        one of ["vanilla", "change"]
    idx1
        bool array masking subpopulation cells 1. Should be True where cell is
        from associated population
    idx2
        bool array masking subpopulation cells 2. Should be True where cell is
        from associated population
    batchid1
        List of batch ids for which you want to perform DE Analysis for
        subpopulation 1. By default, all ids are taken into account
    batchid2
        List of batch ids for which you want to perform DE Analysis for
        subpopulation 2. By default, all ids are taken into account
    use_observed_batches
        Whether posterior values are conditioned on observed
        batches
    n_samples
        Number of posterior samples
    use_permutation
        Activates step 2 described above.
        Simply formulated, pairs obtained from posterior sampling
        will be randomly permuted so that the number of pairs used
        to compute Bayes Factors becomes `m_permutation`.
    m_permutation
        Number of times we will "mix" posterior samples in step 2.
        Only makes sense when `use_permutation=True`
    change_fn
        function computing effect size based on both posterior values
    m1_domain_fn
        custom indicator function of effect size regions
        inducing differential expression
    delta
        specific case of region inducing differential expression.
        In this case, we suppose that :math:`R \setminus [-\delta, \delta]` does not induce differential expression
        (LFC case)
    cred_interval_lvls
        List of credible interval levels to compute for the posterior
        LFC distribution

    Returns
    -------
    Differential expression properties

    """
    # if not np.array_equal(self.indices, np.arange(len(self.dataset))):
    #     logger.warning(
    #         "Differential expression requires a Posterior object created with all indices."
    #     )

    eps = 1e-8  # used for numerical stability
    # Normalized means sampling for both populations
    scales_batches_1 = self.scale_sampler(
        selection=idx1,
        batchid=batchid1,
        use_observed_batches=use_observed_batches,
        n_samples=n_samples,
    )
    scales_batches_2 = self.scale_sampler(
        selection=idx2,
        batchid=batchid2,
        use_observed_batches=use_observed_batches,
        n_samples=n_samples,
    )

    px_scale_mean1 = scales_batches_1["scale"].mean(axis=0)
    px_scale_mean2 = scales_batches_2["scale"].mean(axis=0)

    # Sampling pairs
    # The objective of code section below is to ensure than the samples of normalized
    # means we consider are conditioned on the same batch id
    batchid1_vals = np.unique(scales_batches_1["batch"])
    batchid2_vals = np.unique(scales_batches_2["batch"])

    create_pairs_from_same_batches = (
        set(batchid1_vals) == set(batchid2_vals)
    ) and not use_observed_batches
    if create_pairs_from_same_batches:
        # First case: same batch normalization in two groups
        logger.debug("Same batches in both cell groups")
        n_batches = len(set(batchid1_vals))
        n_samples_per_batch = (
            m_permutation // n_batches if m_permutation is not None else None
        )
        scales_1 = []
        scales_2 = []
        for batch_val in set(batchid1_vals):
            # Select scale samples that originate from the same batch id
            scales_1_batch = scales_batches_1["scale"][
                scales_batches_1["batch"] == batch_val
            ]
            scales_2_batch = scales_batches_2["scale"][
                scales_batches_2["batch"] == batch_val
            ]

            # Create more pairs
            scales_1_local, scales_2_local = pairs_sampler(
                scales_1_batch,
                scales_2_batch,
                use_permutation=use_permutation,
                m_permutation=n_samples_per_batch,
            )
            scales_1.append(scales_1_local)
            scales_2.append(scales_2_local)
        scales_1 = np.concatenate(scales_1, axis=0)
        scales_2 = np.concatenate(scales_2, axis=0)
    else:
        logger.debug("Ignoring batch conditionings to compare means")
        if len(set(batchid1_vals).intersection(set(batchid2_vals))) >= 1:
            warnings.warn(
                "Batchids of cells groups 1 and 2 are different but have an non-null "
                "intersection. Specific handling of such situations is not implemented "
                "yet and batch correction is not trustworthy."
            )
        scales_1, scales_2 = pairs_sampler(
            scales_batches_1["scale"],
            scales_batches_2["scale"],
            use_permutation=use_permutation,
            m_permutation=m_permutation,
        )

    # Core of function: hypotheses testing based on the posterior samples we obtained above
    if mode == "vanilla":
        logger.debug("Differential expression using vanilla mode")
        proba_m1 = np.mean(scales_1 > scales_2, 0)
        proba_m2 = 1.0 - proba_m1
        res = dict(
            proba_m1=proba_m1,
            proba_m2=proba_m2,
            bayes_factor=np.log(proba_m1 + eps) - np.log(proba_m2 + eps),
            scale1=px_scale_mean1,
            scale2=px_scale_mean2,
        )

    elif mode == "change":
        logger.debug("Differential expression using change mode")

        # step 1: Construct the change function
        def lfc(x, y):
            return np.log2(x) - np.log2(y)

        if change_fn == "log-fold" or change_fn is None:
            change_fn = lfc
        elif not callable(change_fn):
            raise ValueError("'change_fn' attribute not understood")

        # step2: Construct the DE area function
        if m1_domain_fn is None:
            delta = delta if delta is not None else 0.5

            def m1_domain_fn(samples):
                return np.abs(samples) >= delta

        change_fn_specs = inspect.getfullargspec(change_fn)
        domain_fn_specs = inspect.getfullargspec(m1_domain_fn)
        if (len(change_fn_specs.args) != 2) | (len(domain_fn_specs.args) != 1):
            raise ValueError(
                "change_fn should take exactly two parameters as inputs; m1_domain_fn one parameter."
            )
        try:
            change_distribution = change_fn(scales_1, scales_2)
            is_de = m1_domain_fn(change_distribution)
        except TypeError:
            raise TypeError(
                "change_fn or m1_domain_fn have has wrong properties."
                "Please ensure that these functions have the right signatures and"
                "outputs and that they can process numpy arrays"
            )
        proba_m1 = np.mean(is_de, 0)
        change_distribution_props = describe_continuous_distrib(
            samples=change_distribution,
            credible_intervals_levels=cred_interval_lvls,
        )
        change_distribution_props = {
            "lfc_" + key: val for (key, val) in change_distribution_props.items()
        }

        res = dict(
            proba_de=proba_m1,
            proba_not_de=1.0 - proba_m1,
            bayes_factor=np.log(proba_m1 + eps) - np.log(1.0 - proba_m1 + eps),
            scale1=px_scale_mean1,
            scale2=px_scale_mean2,
            **change_distribution_props,
        )
    else:
        raise NotImplementedError("Mode {mode} not recognized".format(mode=mode))

    return res


# Traceback (ValueError)
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-143-b1df721788f5> in <module>
2 for tissue in np.unique(adata.obs['tissue']):
3     sub_adata = adata[adata.obs['tissue']==tissue]
----> 4     de_celltype[tissue] = vae.differential_expression(sub_adata, groupby = 'Propagated.Annotation', batch_correction=True)

/data/yosef2/users/chenling/miniconda3/envs/scvi-tools/lib/python3.8/site-packages/scvi/core/models/rnamixin.py in differential_expression(self, adata, groupby, group1, group2, idx1, idx2, mode, delta, batch_size, all_stats, batch_correction, batchid1, batchid2, **kwargs)
187             batch_size=batch_size,
188         )
--> 189         result = _de_core(
190             adata,
191             model_fn,

/data/yosef2/users/chenling/miniconda3/envs/scvi-tools/lib/python3.8/site-packages/scvi/core/models/_utils.py in _de_core(adata, model_fn, groupby, group1, group2, idx1, idx2, all_stats, all_stats_fn, col_names, mode, batchid1, batchid2, delta, batch_correction, **kwargs)
59             cell_idx2 = (adata.obs[groupby] == group2).to_numpy().ravel()
60
---> 61         all_info = dc.get_bayes_factors(
62             cell_idx1,
63             cell_idx2,

/data/yosef2/users/chenling/miniconda3/envs/scvi-tools/lib/python3.8/site-packages/scvi/core/utils/differential.py in get_bayes_factors(self, idx1, idx2, mode, batchid1, batchid2, use_observed_batches, n_samples, use_permutation, m_permutation, change_fn, m1_domain_fn, delta, cred_interval_lvls)
168         eps = 1e-8  # used for numerical stability
169         # Normalized means sampling for both populations
--> 170         scales_batches_1 = self.scale_sampler(
171             selection=idx1,
172             batchid=batchid1,

/data/yosef2/users/chenling/miniconda3/envs/scvi-tools/lib/python3.8/site-packages/torch/autograd/grad_mode.py in decorate_no_grad(*args, **kwargs)
47         def decorate_no_grad(*args, **kwargs):
48             with self:
---> 49                 return func(*args, **kwargs)
50         return decorate_no_grad
51

/data/yosef2/users/chenling/miniconda3/envs/scvi-tools/lib/python3.8/site-packages/scvi/core/utils/differential.py in scale_sampler(self, selection, n_samples, n_samples_per_cell, batchid, use_observed_batches, give_mean)
392             idx = np.random.choice(np.arange(self.adata.shape[0])[selection], n_samples)
393             px_scales.append(
--> 394                 self.model_fn(self.adata, indices=idx, transform_batch=batch_idx)
395             )
396             batch_idx = batch_idx if batch_idx is not None else np.nan

/data/yosef2/users/chenling/miniconda3/envs/scvi-tools/lib/python3.8/site-packages/torch/autograd/grad_mode.py in decorate_no_grad(*args, **kwargs)
47         def decorate_no_grad(*args, **kwargs):
48             with self:
---> 49                 return func(*args, **kwargs)
50         return decorate_no_grad
51

/data/yosef2/users/chenling/miniconda3/envs/scvi-tools/lib/python3.8/site-packages/scvi/core/models/rnamixin.py in get_normalized_expression(self, adata, indices, transform_batch, gene_list, library_size, n_samples, batch_size, return_mean, return_numpy)
82         scdl = self._make_scvi_dl(adata=adata, indices=indices, batch_size=batch_size)
83         if transform_batch is not None:
---> 84             transform_batch = _get_batch_code_from_category(adata, transform_batch)
85
86         if gene_list is None:

/data/yosef2/users/chenling/miniconda3/envs/scvi-tools/lib/python3.8/site-packages/scvi/model/_utils.py in _get_batch_code_from_category(adata, category)
131     batch_mappings = categorical_mappings["_scvi_batch"]["mapping"]
132     if category not in batch_mappings:
--> 133         raise ValueError('"{}" not a valid batch category.'.format(category))
134     return np.where(batch_mappings == category)[0][0]

ValueError: "0" not a valid batch category.```

#### Versions:
<!-- Output of scvi.__version__ -->
VERSION

<!-- Relevant screenshots -->

# Fixed code
def get_bayes_factors(
    self,
    idx1: Union[List[bool], np.ndarray],
    idx2: Union[List[bool], np.ndarray],
    mode: Literal["vanilla", "change"] = "vanilla",
    batchid1: Optional[Sequence[Union[Number, str]]] = None,
    batchid2: Optional[Sequence[Union[Number, str]]] = None,
    use_observed_batches: Optional[bool] = False,
    n_samples: int = 5000,
    use_permutation: bool = False,
    m_permutation: int = 10000,
    change_fn: Optional[Union[str, Callable]] = None,
    m1_domain_fn: Optional[Callable] = None,
    delta: Optional[float] = 0.5,
    cred_interval_lvls: Optional[Union[List[float], np.ndarray]] = None,
) -> Dict[str, np.ndarray]:
    r"""
    A unified method for differential expression inference.

    Two modes coexist:

    - the `"vanilla"` mode follows protocol described in [Lopez18]_ and [Xu19]_
    In this case, we perform hypothesis testing based on the hypotheses

    .. math::
        M_1: h_1 > h_2 ~\text{and}~ M_2: h_1 \leq h_2.

    DE can then be based on the study of the Bayes factors

    .. math::
        \log p(M_1 | x_1, x_2) / p(M_2 | x_1, x_2).

    - the `"change"` mode (described in [Boyeau19]_)
    This mode consists of estimating an effect size random variable (e.g., log fold-change) and
    performing Bayesian hypothesis testing on this variable.
    The `change_fn` function computes the effect size variable :math:`r` based on two inputs
    corresponding to the posterior quantities (e.g., normalized expression) in both populations.

    Hypotheses:

    .. math::
        M_1: r \in R_1 ~\text{(effect size r in region inducing differential expression)}

    .. math::
        M_2: r  \notin R_1 ~\text{(no differential expression)}

    To characterize the region :math:`R_1`, which induces DE, the user has two choices.

    1. A common case is when the region :math:`[-\delta, \delta]` does not induce differential
       expression. If the user specifies a threshold delta, we suppose that :math:`R_1 = \mathbb{R} \setminus [-\delta, \delta]`
    2. Specify an specific indicator function:

    .. math::
        f: \mathbb{R} \mapsto \{0, 1\} ~\text{s.t.}~ r \in R_1 ~\text{iff.}~ f(r) = 1.

    Decision-making can then be based on the estimates of

    .. math::
        p(M_1 \mid x_1, x_2).

    Both modes require to sample the posterior distributions.
    To that purpose, we sample the posterior in the following way:

    1. The posterior is sampled `n_samples` times for each subpopulation.
    2. For computational efficiency (posterior sampling is quite expensive), instead of
       comparing the obtained samples element-wise, we can permute posterior samples.
       Remember that computing the Bayes Factor requires sampling :math:`q(z_A \mid x_A)` and :math:`q(z_B \mid x_B)`.

    Currently, the code covers several batch handling configurations:

    1. If ``use_observed_batches=True``, then batch are considered as observations
       and cells' normalized means are conditioned on real batch observations.
    2. If case (cell group 1) and control (cell group 2) are conditioned on the same
       batch ids. This requires ``set(batchid1) == set(batchid2)`` or ``batchid1 == batchid2 === None``.
    3. If case and control are conditioned on different batch ids that do not intersect
       i.e., ``set(batchid1) != set(batchid2)`` and ``len(set(batchid1).intersection(set(batchid2))) == 0``.

    This function does not cover other cases yet and will warn users in such cases.

    Parameters
    ----------
    mode
        one of ["vanilla", "change"]
    idx1
        bool array masking subpopulation cells 1. Should be True where cell is
        from associated population
    idx2
        bool array masking subpopulation cells 2. Should be True where cell is
        from associated population
    batchid1
        List of batch ids for which you want to perform DE Analysis for
        subpopulation 1. By default, all ids are taken into account
    batchid2
        List of batch ids for which you want to perform DE Analysis for
        subpopulation 2. By default, all ids are taken into account
    use_observed_batches
        Whether posterior values are conditioned on observed
        batches
    n_samples
        Number of posterior samples
    use_permutation
        Activates step 2 described above.
        Simply formulated, pairs obtained from posterior sampling
        will be randomly permuted so that the number of pairs used
        to compute Bayes Factors becomes `m_permutation`.
    m_permutation
        Number of times we will "mix" posterior samples in step 2.
        Only makes sense when `use_permutation=True`
    change_fn
        function computing effect size based on both posterior values
    m1_domain_fn
        custom indicator function of effect size regions
        inducing differential expression
    delta
        specific case of region inducing differential expression.
        In this case, we suppose that :math:`R \setminus [-\delta, \delta]` does not induce differential expression
        (LFC case)
    cred_interval_lvls
        List of credible interval levels to compute for the posterior
        LFC distribution

    Returns
    -------
    Differential expression properties

    """
    # if not np.array_equal(self.indices, np.arange(len(self.dataset))):
    #     logger.warning(
    #         "Differential expression requires a Posterior object created with all indices."
    #     )

    eps = 1e-8  # used for numerical stability
    # Normalized means sampling for both populations
    scales_batches_1 = self.scale_sampler(
        selection=idx1,
        batchid=batchid1,
        use_observed_batches=use_observed_batches,
        n_samples=n_samples,
    )
    scales_batches_2 = self.scale_sampler(
        selection=idx2,
        batchid=batchid2,
        use_observed_batches=use_observed_batches,
        n_samples=n_samples,
    )

    px_scale_mean1 = scales_batches_1["scale"].mean(axis=0)
    px_scale_mean2 = scales_batches_2["scale"].mean(axis=0)

    # Sampling pairs
    # The objective of code section below is to ensure than the samples of normalized
    # means we consider are conditioned on the same batch id
    batchid1_vals = np.unique(scales_batches_1["batch"])
    batchid2_vals = np.unique(scales_batches_2["batch"])

    create_pairs_from_same_batches = (
        set(batchid1_vals) == set(batchid2_vals)
    ) and not use_observed_batches
    if create_pairs_from_same_batches:
        # First case: same batch normalization in two groups
        logger.debug("Same batches in both cell groups")
        n_batches = len(set(batchid1_vals))
        n_samples_per_batch = (
            m_permutation // n_batches if m_permutation is not None else None
        )
        scales_1 = []
        scales_2 = []
        for batch_val in set(batchid1_vals):
            # Select scale samples that originate from the same batch id
            scales_1_batch = scales_batches_1["scale"][
                scales_batches_1["batch"] == batch_val
            ]
            scales_2_batch = scales_batches_2["scale"][
                scales_batches_2["batch"] == batch_val
            ]

            # Create more pairs
            scales_1_local, scales_2_local = pairs_sampler(
                scales_1_batch,
                scales_2_batch,
                use_permutation=use_permutation,
                m_permutation=n_samples_per_batch,
            )
            scales_1.append(scales_1_local)
            scales_2.append(scales_2_local)
        scales_1 = np.concatenate(scales_1, axis=0)
        scales_2 = np.concatenate(scales_2, axis=0)
    else:
        logger.debug("Ignoring batch conditionings to compare means")
        if len(set(batchid1_vals).intersection(set(batchid2_vals))) >= 1:
            warnings.warn(
                "Batchids of cells groups 1 and 2 are different but have an non-null "
                "intersection. Specific handling of such situations is not implemented "
                "yet and batch correction is not trustworthy."
            )
        scales_1, scales_2 = pairs_sampler(
            scales_batches_1["scale"],
            scales_batches_2["scale"],
            use_permutation=use_permutation,
            m_permutation=m_permutation,
        )

    # Core of function: hypotheses testing based on the posterior samples we obtained above
    if mode == "vanilla":
        logger.debug("Differential expression using vanilla mode")
        proba_m1 = np.mean(scales_1 > scales_2, 0)
        proba_m2 = 1.0 - proba_m1
        res = dict(
            proba_m1=proba_m1,
            proba_m2=proba_m2,
            bayes_factor=np.log(proba_m1 + eps) - np.log(proba_m2 + eps),
            scale1=px_scale_mean1,
            scale2=px_scale_mean2,
        )

    elif mode == "change":
        logger.debug("Differential expression using change mode")

        # step 1: Construct the change function
        def lfc(x, y):
            return np.log2(x) - np.log2(y)

        if change_fn == "log-fold" or change_fn is None:
            change_fn = lfc
        elif not callable(change_fn):
            raise ValueError("'change_fn' attribute not understood")

        # step2: Construct the DE area function
        if m1_domain_fn is None:
            delta = delta if delta is not None else 0.5

            def m1_domain_fn(samples):
                return np.abs(samples) >= delta

        change_fn_specs = inspect.getfullargspec(change_fn)
        domain_fn_specs = inspect.getfullargspec(m1_domain_fn)
        if (len(change_fn_specs.args) != 2) | (len(domain_fn_specs.args) != 1):
            raise ValueError(
                "change_fn should take exactly two parameters as inputs; m1_domain_fn one parameter."
            )
        try:
            change_distribution = change_fn(scales_1, scales_2)
            is_de = m1_domain_fn(change_distribution)
        except TypeError:
            raise TypeError(
                "change_fn or m1_domain_fn have has wrong properties."
                "Please ensure that these functions have the right signatures and"
                "outputs and that they can process numpy arrays"
            )
        proba_m1 = np.mean(is_de, 0)
        change_distribution_props = describe_continuous_distrib(
            samples=change_distribution,
            credible_intervals_levels=cred_interval_lvls,
        )
        change_distribution_props = {
            "lfc_" + key: val for (key, val) in change_distribution_props.items()
        }

        res = dict(
            proba_de=proba_m1,
            proba_not_de=1.0 - proba_m1,
            bayes_factor=np.log(proba_m1 + eps) - np.log(1.0 - proba_m1 + eps),
            scale1=px_scale_mean1,
            scale2=px_scale_mean2,
            **change_distribution_props,
        )
    else:
        raise NotImplementedError("Mode {mode} not recognized".format(mode=mode))

    return res