Repository: scverse/scanpy

Issue:
sc.pp.highly_variable_genes with flavor='seurat_v3' returns wrong means/variances when batch_key argument is used
- [x] I have checked that this issue has not already been reported.
- [x] I have confirmed this bug exists on the latest version of scanpy.
- [x] (optional) I have confirmed this bug exists on the master branch of scanpy.

---

When working on PR #1715, I noticed a small bug when `sc.pp.highly_variable()` is run with `flavor='seurat_v3'` and the `batch_key` argument is used on a dataset with multiple batches:

The columns in the returned data frame `means` and `variances` do not give the correct gene means and gene variances across the whole dataset, but instead give the means and variances of the batch with the last index.

I think this because here means and variances are first computed within each batch:
https://github.com/theislab/scanpy/blob/a085333ead6ff8a0f64c25734060ffc048be56c0/scanpy/preprocessing/_highly_variable_genes.py#L72-L77
..and then not re-computed for the whole dataset before inserting them here:
https://github.com/theislab/scanpy/blob/a085333ead6ff8a0f64c25734060ffc048be56c0/scanpy/preprocessing/_highly_variable_genes.py#L136-L137

Should be an easy fix, I can also prepare it if you want :)

### Code to reproduce:

```python
import numpy as np
import scanpy as sc
import anndata 

import sys
sys.path.append("scanpy/preprocessing")
from _utils import _get_mean_var

np.random.seed(42)
adata = anndata.AnnData(np.random.randint(0,5,(100,100)))
adata.obs['batch'] = np.random.randint(0,5,(100))

#show that without batch_key, everything works fine
true_mean,true_var = _get_mean_var(adata.X)
result = sc.pp.highly_variable_genes(adata,n_top_genes=2,flavor='seurat_v3',inplace=False)
print('Results correct without batch_key?')
print(np.all(result['means']==true_mean))
print(np.all(result['variances']==true_var))

#show that it goes wrong with batch_key..
result_with_batchkey = sc.pp.highly_variable_genes(adata,n_top_genes=2,flavor='seurat_v3',inplace=False,batch_key='batch')
print('Results correct with batch_key?')
print(np.all(result_with_batchkey['means']==true_mean))
print(np.all(result_with_batchkey['variances']==true_var))

#..and that the wrong result comes from the last batch
adata_batch4 = adata[adata.obs['batch']==4,:].copy()
means_batch4,vars_batch4 = _get_mean_var(adata_batch4.X)
print('Wrong results come from last batch?')
print(np.all(result_with_batchkey['means']==means_batch4))
print(np.all(result_with_batchkey['variances']==vars_batch4))
```

```pytb
Results correct without batch_key?
True
True
Results correct with batch_key?
False
False
Wrong results come from last batch?
True
True

/gpfs01/berens/user/jlause/libs/scanpy/scanpy/preprocessing/_highly_variable_genes.py:146: FutureWarning: Slicing a positional slice with .loc is not supported, and will raise TypeError in a future version.  Use .loc with labels or .iloc with positions instead.
  df.loc[: int(n_top_genes), 'highly_variable'] = True
/gpfs01/berens/user/jlause/libs/scanpy/scanpy/preprocessing/_highly_variable_genes.py:146: FutureWarning: Slicing a positional slice with .loc is not supported, and will raise TypeError in a future version.  Use .loc with labels or .iloc with positions instead.
  df.loc[: int(n_top_genes), 'highly_variable'] = True


```

#### Versions

<details>

-----
anndata     0.7.5
scanpy      1.7.1.dev2+g8c469411
sinfo       0.3.1
-----
PIL                 8.1.0
anndata             0.7.5
anyio               NA
attr                20.3.0
babel               2.9.0
backcall            0.2.0
certifi             2020.12.05
cffi                1.14.4
chardet             4.0.0
constants           NA
cycler              0.10.0
cython_runtime      NA
dateutil            2.8.1
decorator           4.4.2
get_version         2.1
google              NA
h5py                2.10.0
highs_wrapper       NA
idna                2.10
ipykernel           5.4.2
ipython_genutils    0.2.0
jedi                0.17.2
jinja2              2.11.2
joblib              1.0.0
json5               NA
jsonschema          3.2.0
jupyter_server      1.2.1
jupyterlab_server   2.1.1
kiwisolver          1.3.1
legacy_api_wrap     1.2
llvmlite            0.35.0
markupsafe          1.1.1
matplotlib          3.3.3
mpl_toolkits        NA
natsort             7.1.0
nbclassic           NA
nbformat            5.0.8
numba               0.52.0
numexpr             2.7.2
numpy               1.19.5
packaging           20.8
pandas              1.2.0
parso               0.7.1
pexpect             4.8.0
pickleshare         0.7.5
pkg_resources       NA
prometheus_client   NA
prompt_toolkit      3.0.10
ptyprocess          0.7.0
pvectorc            NA
pygments            2.7.3
pyparsing           2.4.7
pyrsistent          NA
pytz                2020.5
requests            2.25.1
scanpy              1.7.1.dev2+g8c469411
scipy               1.6.0
send2trash          NA
sinfo               0.3.1
sitecustomize       NA
six                 1.15.0
sklearn             0.24.0
sniffio             1.2.0
storemagic          NA
tables              3.6.1
tornado             6.1
traitlets           5.0.5
typing_extensions   NA
urllib3             1.26.2
wcwidth             0.2.5
yaml                5.3.1
zmq                 20.0.0
-----
IPython             7.19.0
jupyter_client      6.1.11
jupyter_core        4.7.0
jupyterlab          3.0.3
notebook            6.1.6
-----
Python 3.8.0 (default, Oct 28 2019, 16:14:01) [GCC 8.3.0]
Linux-3.10.0-957.el7.x86_64-x86_64-with-glibc2.27
40 logical CPU cores, x86_64
-----
Session information updated at 2021-03-09 19:18

</details>



Discussion:
Makes sense, thanks for the detailed report!

Would you like to make a bugfix PR for this?

Resolving patch:
diff --git a/scanpy/preprocessing/_highly_variable_genes.py b/scanpy/preprocessing/_highly_variable_genes.py
--- a/scanpy/preprocessing/_highly_variable_genes.py
+++ b/scanpy/preprocessing/_highly_variable_genes.py
@@ -55,14 +55,17 @@ def _highly_variable_genes_seurat_v3(
         raise ImportError(
             'Please install skmisc package via `pip install --user scikit-misc'
         )
-
+    df = pd.DataFrame(index=adata.var_names)
     X = adata.layers[layer] if layer is not None else adata.X
+
     if check_values and not check_nonnegative_integers(X):
         warnings.warn(
             "`flavor='seurat_v3'` expects raw count data, but non-integers were found.",
             UserWarning,
         )
 
+    df['means'], df['variances'] = _get_mean_var(X)
+
     if batch_key is None:
         batch_info = pd.Categorical(np.zeros(adata.shape[0], dtype=int))
     else:
@@ -70,13 +73,11 @@ def _highly_variable_genes_seurat_v3(
 
     norm_gene_vars = []
     for b in np.unique(batch_info):
+        X_batch = X[batch_info == b]
 
-        ad = adata[batch_info == b]
-        X = ad.layers[layer] if layer is not None else ad.X
-
-        mean, var = _get_mean_var(X)
+        mean, var = _get_mean_var(X_batch)
         not_const = var > 0
-        estimat_var = np.zeros(adata.shape[1], dtype=np.float64)
+        estimat_var = np.zeros(X.shape[1], dtype=np.float64)
 
         y = np.log10(var[not_const])
         x = np.log10(mean[not_const])
@@ -85,15 +86,18 @@ def _highly_variable_genes_seurat_v3(
         estimat_var[not_const] = model.outputs.fitted_values
         reg_std = np.sqrt(10 ** estimat_var)
 
-        batch_counts = X.astype(np.float64).copy()
+        batch_counts = X_batch.astype(np.float64).copy()
         # clip large values as in Seurat
-        N = np.sum(batch_info == b)
+        N = X_batch.shape[0]
         vmax = np.sqrt(N)
         clip_val = reg_std * vmax + mean
         if sp_sparse.issparse(batch_counts):
             batch_counts = sp_sparse.csr_matrix(batch_counts)
             mask = batch_counts.data > clip_val[batch_counts.indices]
             batch_counts.data[mask] = clip_val[batch_counts.indices[mask]]
+
+            squared_batch_counts_sum = np.array(batch_counts.power(2).sum(axis=0))
+            batch_counts_sum = np.array(batch_counts.sum(axis=0))
         else:
             clip_val_broad = np.broadcast_to(clip_val, batch_counts.shape)
             np.putmask(
@@ -102,10 +106,6 @@ def _highly_variable_genes_seurat_v3(
                 clip_val_broad,
             )
 
-        if sp_sparse.issparse(batch_counts):
-            squared_batch_counts_sum = np.array(batch_counts.power(2).sum(axis=0))
-            batch_counts_sum = np.array(batch_counts.sum(axis=0))
-        else:
             squared_batch_counts_sum = np.square(batch_counts).sum(axis=0)
             batch_counts_sum = batch_counts.sum(axis=0)
 
@@ -129,22 +129,21 @@ def _highly_variable_genes_seurat_v3(
     ma_ranked = np.ma.masked_invalid(ranked_norm_gene_vars)
     median_ranked = np.ma.median(ma_ranked, axis=0).filled(np.nan)
 
-    df = pd.DataFrame(index=np.array(adata.var_names))
     df['highly_variable_nbatches'] = num_batches_high_var
     df['highly_variable_rank'] = median_ranked
     df['variances_norm'] = np.mean(norm_gene_vars, axis=0)
-    df['means'] = mean
-    df['variances'] = var
 
-    df.sort_values(
-        ['highly_variable_rank', 'highly_variable_nbatches'],
-        ascending=[True, False],
-        na_position='last',
-        inplace=True,
+    sorted_index = (
+        df[['highly_variable_rank', 'highly_variable_nbatches']]
+        .sort_values(
+            ['highly_variable_rank', 'highly_variable_nbatches'],
+            ascending=[True, False],
+            na_position='last',
+        )
+        .index
     )
     df['highly_variable'] = False
-    df.loc[: int(n_top_genes), 'highly_variable'] = True
-    df = df.loc[adata.var_names]
+    df.loc[sorted_index[: int(n_top_genes)], 'highly_variable'] = True
 
     if inplace or subset:
         adata.uns['hvg'] = {'flavor': 'seurat_v3'}


Test patch:
diff --git a/scanpy/tests/test_highly_variable_genes.py b/scanpy/tests/test_highly_variable_genes.py
--- a/scanpy/tests/test_highly_variable_genes.py
+++ b/scanpy/tests/test_highly_variable_genes.py
@@ -246,3 +246,25 @@ def test_highly_variable_genes_batches():
     ]
 
     assert np.all(np.isin(colnames, hvg1.columns))
+
+
+from scanpy.preprocessing._utils import _get_mean_var
+
+
+def test_seurat_v3_mean_var_output_with_batchkey():
+    pbmc = sc.datasets.pbmc3k()
+    pbmc.var_names_make_unique()
+    n_cells = pbmc.shape[0]
+    batch = np.zeros((n_cells), dtype=int)
+    batch[1500:] = 1
+    pbmc.obs["batch"] = batch
+
+    # true_mean, true_var = _get_mean_var(pbmc.X)
+    true_mean = np.mean(pbmc.X.toarray(), axis=0)
+    true_var = np.var(pbmc.X.toarray(), axis=0, dtype=np.float64, ddof=1)
+
+    result_df = sc.pp.highly_variable_genes(
+        pbmc, batch_key='batch', flavor='seurat_v3', n_top_genes=4000, inplace=False
+    )
+    np.testing.assert_allclose(true_mean, result_df['means'], rtol=2e-05, atol=2e-05)
+    np.testing.assert_allclose(true_var, result_df['variances'], rtol=2e-05, atol=2e-05)