from scvi.data import synthetic_iid
from scvi.model import SCVI
from scvi.core.utils.differential import DifferentialComputation
from functools import partial
import numpy as np


def test_differential_computation(save_path):

    n_latent = 5
    adata = synthetic_iid()
    model = SCVI(adata, n_latent=n_latent)
    model.train(1)

    model_fn = partial(model.get_normalized_expression, return_numpy=True)
    dc = DifferentialComputation(model_fn, adata)

    cell_idx1 = np.asarray(adata.obs.labels == "label_1")
    cell_idx2 = ~cell_idx1

    dc.get_bayes_factors(cell_idx1, cell_idx2, mode="vanilla", use_permutation=True)
    dc.get_bayes_factors(cell_idx1, cell_idx2, mode="change", use_permutation=False)
    dc.get_bayes_factors(cell_idx1, cell_idx2, mode="change", cred_interval_lvls=[0.75])

    delta = 0.5

    def change_fn_test(x, y):
        return x - y

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

    dc.get_bayes_factors(
        cell_idx1,
        cell_idx2,
        mode="change",
        m1_domain_fn=m1_domain_fn_test,
        change_fn=change_fn_test,
    )

@@ -171,16 +171,30 @@
 y as np%0A
+import pytest%0A
 %0A%0Adef te
@@ -1127,8 +1127,271 @@
 ,%0A    )%0A
+%0A    # should fail if just one batch%0A    with pytest.raises(ValueError):%0A        model.differential_expression(adata%5B:20%5D, groupby=%22batch%22)%0A%0A    # test view%0A    model.differential_expression(%0A        adata%5Badata.obs%5B%22labels%22%5D == %22label_1%22%5D, groupby=%22batch%22%0A    )%0A

add some tests