sparse column shift encoding. (#1502)

Many of our matrices are log normalized, which tends to eliminate
the number of non zero values (if there were any).  This prevents
the matrix from being stored as a sparse matrix.  The solution here
is to use a simple transformation to make it sparse again.  The most
common value from each column is subtracted from that column.  These
values that were subtracted are saved in an array called X_col_shift.

The cellxgene code needs to understand how to undo the transformation when
operating over the X matrix.

- added script to create a synthetic dataset for testing
- added a script to convert an existing CXG dataset to a sparse CXG dataset

server/test/test_diffexp.py

@@ -4,9 +4,10 @@ from server.test import PROJECT_ROOT, app_config
 import server.compute.diffexp_cxg as diffexp_cxg
 import server.compute.diffexp_generic as diffexp_generic
 from server.converters.cxgtool import write_cxg
+from server.test.create_test_matrix import create_test_h5ad
+from server.data_common.fbs.matrix import encode_matrix_fbs, decode_matrix_fbs
 import numpy as np
 import tempfile
-import anndata
 import os
 
 
@@ -14,8 +15,8 @@ class DiffExpTest(unittest.TestCase):
     """Tests the diffexp returns the expected results for one test case, using different
     adaptor types and different algorithms."""
 
-    def load_dataset(self, path):
-        config = app_config(path)
+    def load_dataset(self, path, extra={}):
+        config = app_config(path, extra=extra)
         loader = MatrixDataLoader(path)
         adaptor = loader.open(config)
         return adaptor
@@ -28,6 +29,14 @@ class DiffExpTest(unittest.TestCase):
         mask[sel] = True
         return mask
 
+    def compare_diffexp_results(self, results, expects):
+        self.assertEqual(len(results), len(expects))
+        for result, expect in zip(results, expects):
+            self.assertEqual(result[0], expect[0])
+            self.assertTrue(np.isclose(result[1], expect[1], 1e-6, 1e-4))
+            self.assertTrue(np.isclose(result[2], expect[2], 1e-6, 1e-4))
+            self.assertTrue(np.isclose(result[3], expect[3], 1e-6, 1e-4))
+
     def check_1_10_2_10(self, results):
         """Checks the results for a specific set of rows selections"""
         expects = [
@@ -42,12 +51,7 @@ class DiffExpTest(unittest.TestCase):
             [1575, 1.0317602, 0.007830310753043345, 1.0],
             [576, 0.97873515, 0.008272092578813124, 1.0],
         ]
-        self.assertEqual(len(results), len(expects))
-        for result, expect in zip(results, expects):
-            self.assertEqual(result[0], expect[0])
-            self.assertTrue(np.isclose(result[1], expect[1], 1e-6, 1e-6))
-            self.assertTrue(np.isclose(result[2], expect[2], 1e-6, 1e-6))
-            self.assertTrue(np.isclose(result[3], expect[3], 1e-6, 1e-6))
+        self.compare_diffexp_results(results, expects)
 
     def get_X_col(self, adaptor, cols):
         varmask = np.zeros(adaptor.get_shape()[1], dtype=bool)
@@ -86,33 +90,62 @@ class DiffExpTest(unittest.TestCase):
         self.check_1_10_2_10(results)
 
     def test_cxg_sparse(self):
+        self.sparse_diffexp(False)
+
+    def test_cxg_sparse_col_shift(self):
+        self.sparse_diffexp(True)
+
+    def sparse_diffexp(self, apply_col_shift):
         with tempfile.TemporaryDirectory() as dirname:
+            # create a sparse matrix
+            h5adfile = os.path.join(dirname, "sparse.h5ad")
+            create_test_h5ad(h5adfile, 2000, 2000, 10, apply_col_shift)
+            adaptor_anndata = self.load_dataset(h5adfile, extra=dict(embeddings__names=[]))
+            adata = adaptor_anndata.data
+
             sparsename = os.path.join(dirname, "sparse.cxg")
-            densename = os.path.join(dirname, "dense.cxg")
-            source_h5ad = anndata.read_h5ad(f"{PROJECT_ROOT}/example-dataset/pbmc3k.h5ad")
-            # create a cxg sparse array
-            write_cxg(adata=source_h5ad, container=sparsename, title="pbmc3k", sparse_threshold=100)
-            write_cxg(adata=source_h5ad, container=densename, title="pbmc3k", sparse_threshold=0)
+            write_cxg(adata=adata, container=sparsename, title="sparse", sparse_threshold=11)
             adaptor_sparse = self.load_dataset(sparsename)
+            assert adaptor_sparse.open_array("X").schema.sparse
+            assert adaptor_sparse.has_array("X_col_shift") == apply_col_shift
+
+            densename = os.path.join(dirname, "dense.cxg")
+            write_cxg(adata=adata, container=densename, title="dense", sparse_threshold=0)
             adaptor_dense = self.load_dataset(densename)
+            assert not adaptor_dense.open_array("X").schema.sparse
+            assert not adaptor_dense.has_array("X_col_shift")
 
-            col_results = []
-            for adaptor in (adaptor_sparse, adaptor_dense):
-                maskA = self.get_mask(adaptor, 1, 10)
-                maskB = self.get_mask(adaptor, 2, 10)
-                diffexp_results = diffexp_cxg.diffexp_ttest(adaptor, maskA, maskB, 10)
-                self.check_1_10_2_10(diffexp_results)
-                topcols = [x[0] for x in diffexp_results]
-                cols = self.get_X_col(adaptor, topcols)
-                assert cols.shape[0] == adaptor.get_shape()[0]
-                assert cols.shape[1] == len(diffexp_results)
-                col_results.append(cols)
+            maskA = self.get_mask(adaptor_anndata, 1, 10)
+            maskB = self.get_mask(adaptor_anndata, 2, 10)
 
-                x = adaptor.get_X_array()
-                print(x)
+            diffexp_results_anndata = diffexp_generic.diffexp_ttest(adaptor_anndata, maskA, maskB, 10)
+            diffexp_results_sparse = diffexp_cxg.diffexp_ttest(adaptor_sparse, maskA, maskB, 10)
+            diffexp_results_dense = diffexp_cxg.diffexp_ttest(adaptor_dense, maskA, maskB, 10)
 
-        for row in range(col_results[0].shape[0]):
-            for col in range(col_results[0].shape[1]):
-                sval = col_results[0][row][col]
-                dval = col_results[1][row][col]
-                self.assertTrue(np.isclose(sval, dval, 1e-6, 1e-6))
+            self.compare_diffexp_results(diffexp_results_anndata, diffexp_results_sparse)
+            self.compare_diffexp_results(diffexp_results_anndata, diffexp_results_dense)
+
+            topcols = np.array([x[0] for x in diffexp_results_anndata])
+            cols_anndata = self.get_X_col(adaptor_anndata, topcols)
+            cols_sparse = self.get_X_col(adaptor_sparse, topcols)
+            cols_dense = self.get_X_col(adaptor_dense, topcols)
+            assert cols_anndata.shape[0] == adaptor_sparse.get_shape()[0]
+            assert cols_anndata.shape[1] == len(diffexp_results_anndata)
+
+            def convert(mat, cols):
+                return decode_matrix_fbs(encode_matrix_fbs(mat, col_idx=cols)).to_numpy()
+
+            cols_anndata = convert(cols_anndata, topcols)
+            cols_sparse = convert(cols_sparse, topcols)
+            cols_dense = convert(cols_dense, topcols)
+
+            x = adaptor_sparse.get_X_array()
+            assert x.shape == adaptor_sparse.get_shape()
+
+            for row in range(cols_anndata.shape[0]):
+                for col in range(cols_anndata.shape[1]):
+                    vanndata = cols_anndata[row][col]
+                    vsparse = cols_sparse[row][col]
+                    vdense = cols_dense[row][col]
+                    self.assertTrue(np.isclose(vanndata, vsparse, 1e-6, 1e-6))
+                    self.assertTrue(np.isclose(vanndata, vdense, 1e-6, 1e-6))