Merge branch 'main' into colinmegill/#2042-geneset-tests

backend/test/test_czi_hosted/unit/converters/test_h5ad_data_file.py

@@ -8,6 +8,7 @@ from uuid import uuid4
 import anndata
 import numpy as np
 from pandas import Series, DataFrame
+import tiledb
 
 from backend.czi_hosted.common.corpora import CorporaConstants
 from backend.czi_hosted.converters.h5ad_data_file import H5ADDataFile
@@ -129,25 +130,25 @@ class TestH5ADDataFile(unittest.TestCase):
         h5ad_file = H5ADDataFile(self.sample_h5ad_filename, use_corpora_schema=False)
         h5ad_file.to_cxg(self.sample_output_directory, 100)
 
-        self._validate_expected_generated_list_of_tiledb_files()
+        self._validate_cxg_and_h5ad_content_match(self.sample_h5ad_filename, self.sample_output_directory, True)
 
     def test__to_cxg__simple_anndata_with_corpora_and_sparse(self):
         h5ad_file = H5ADDataFile(self.sample_h5ad_filename)
         h5ad_file.to_cxg(self.sample_output_directory, 100)
 
-        self._validate_expected_generated_list_of_tiledb_files()
+        self._validate_cxg_and_h5ad_content_match(self.sample_h5ad_filename, self.sample_output_directory, True)
 
     def test__to_cxg__simple_anndata_no_corpora_and_dense(self):
         h5ad_file = H5ADDataFile(self.sample_h5ad_filename, use_corpora_schema=False)
         h5ad_file.to_cxg(self.sample_output_directory, 0)
 
-        self._validate_expected_generated_list_of_tiledb_files()
+        self._validate_cxg_and_h5ad_content_match(self.sample_h5ad_filename, self.sample_output_directory, False)
 
     def test__to_cxg__simple_anndata_with_corpora_and_dense(self):
         h5ad_file = H5ADDataFile(self.sample_h5ad_filename)
         h5ad_file.to_cxg(self.sample_output_directory, 0)
 
-        self._validate_expected_generated_list_of_tiledb_files()
+        self._validate_cxg_and_h5ad_content_match(self.sample_h5ad_filename, self.sample_output_directory, False)
 
     def test__to_cxg__with_sparse_column_encoding(self):
         anndata = self._create_sample_anndata_dataset()
@@ -157,69 +158,84 @@ class TestH5ADDataFile(unittest.TestCase):
         h5ad_file = H5ADDataFile(sparse_with_column_shift_filename)
         h5ad_file.to_cxg(self.sample_output_directory, 50)
 
-        self._validate_expected_generated_list_of_tiledb_files(has_column_encoding=True)
+        self._validate_cxg_and_h5ad_content_match(
+            sparse_with_column_shift_filename, self.sample_output_directory, False, has_column_encoding=True
+        )
 
         # Clean up
         remove(sparse_with_column_shift_filename)
 
-    def _validate_expected_generated_list_of_tiledb_files(self, has_column_encoding=False):
+    def _validate_cxg_and_h5ad_content_match(self, h5ad_filename, cxg_directory, is_sparse, has_column_encoding=False):
-        (
+        anndata_object = anndata.read_h5ad(h5ad_filename)
-            expected_directories,
-            expected_obs_files,
-            expected_var_files,
-        ) = self._get_expected_generated_list_of_tiledb_files()
 
-        for directory in expected_directories:
+        # Array locations
-            self.assertTrue(path.isdir(directory))
+        metadata_array_location = f"{cxg_directory}/cxg_group_metadata"
+        main_x_array_location = f"{cxg_directory}/X"
+        embedding_array_location = f"{cxg_directory}/emb"
+        specific_embedding_array_location = f"{self.sample_output_directory}/emb/awesome_embedding"
+        obs_array_location = f"{cxg_directory}/obs"
+        var_array_location = f"{cxg_directory}/var"
+        x_col_shift_array_location = f"{cxg_directory}/X_col_shift"
 
-        for obs_file in expected_obs_files:
+        # Assert CXG structure
-            expected_location_of_obs_file = f"{self.sample_output_directory}/obs/*/{obs_file}"
+        self.assertEqual(tiledb.object_type(cxg_directory), "group")
-            self.assertTrue(path.isfile(glob(expected_location_of_obs_file)[0]))
+        self.assertEqual(tiledb.object_type(obs_array_location), "array")
-
+        self.assertEqual(tiledb.object_type(var_array_location), "array")
-        for var_file in expected_var_files:
+        self.assertEqual(tiledb.object_type(main_x_array_location), "array")
-            expected_location_of_var_file = f"{self.sample_output_directory}/var/*/{var_file}"
+        self.assertEqual(tiledb.object_type(embedding_array_location), "group")
-            self.assertTrue(path.isfile(glob(expected_location_of_var_file)[0]))
+        self.assertEqual(tiledb.object_type(specific_embedding_array_location), "array")
 
         if has_column_encoding:
-            self.assertTrue(path.isdir(f"{self.sample_output_directory}/X_col_shift"))
+            self.assertEqual(tiledb.object_type(x_col_shift_array_location), "array")
 
-    def _get_expected_generated_list_of_tiledb_files(self):
+        # Validate metadata
+        metadata_array = tiledb.DenseArray(metadata_array_location, mode="r")
+        self.assertIn("cxg_version", metadata_array.meta)
 
-        # Expected directories
+        # Validate obs index
-        metadata_directory = f"{self.sample_output_directory}/cxg_group_metadata"
+        obs_array = tiledb.DenseArray(obs_array_location, mode="r")
-        main_x_directory = f"{self.sample_output_directory}/X"
+        expected_index_data = anndata_object.obs.index.to_numpy()
-        overall_embedding_directory = f"{self.sample_output_directory}/emb"
+        index_name = json.loads(obs_array.meta["cxg_schema"])["index"]
-        specific_embedding_directory = f"{self.sample_output_directory}/emb/awesome_embedding"
+        actual_index_data = obs_array.query(attrs=[index_name])[:][index_name]
-        obs_directory = f"{self.sample_output_directory}/obs"
+        self.assertTrue(np.array_equal(expected_index_data, actual_index_data))
-        var_directory = f"{self.sample_output_directory}/var"
 
-        # Obs files
+        # Validate obs columns
-        obs_files = []
+        expected_columns = list(anndata_object.obs.columns.values)
-        obs_files.append("name_0.tdb")
+        for column_name in expected_columns:
-        obs_files.append("name_0_var.tdb")
+            expected_data = anndata_object.obs[column_name].to_numpy()
-        obs_files.append("string_category.tdb")
+            actual_data = obs_array.query(attrs=[column_name])[:][column_name]
-        obs_files.append("string_category_var.tdb")
+            self.assertTrue(np.array_equal(expected_data, actual_data))
-        obs_files.append("float_category.tdb")
 
-        # Var files
+        # Validate var index
-        var_files = []
+        var_array = tiledb.DenseArray(var_array_location, mode="r")
-        var_files.append("name_0.tdb")
+        expected_index_data = anndata_object.var.index.to_numpy()
-        var_files.append("name_0_var.tdb")
+        index_name = json.loads(var_array.meta["cxg_schema"])["index"]
-        var_files.append("bool_category.tdb")
+        actual_index_data = var_array.query(attrs=[index_name])[:][index_name]
-        var_files.append("int_category.tdb")
+        self.assertTrue(np.array_equal(expected_index_data, actual_index_data))
 
-        return (
+        # Validate var columns
-            [
+        expected_columns = anndata_object.var.columns.values
-                metadata_directory,
+        for column_name in expected_columns:
-                main_x_directory,
+            expected_data = anndata_object.var[column_name].to_numpy()
-                overall_embedding_directory,
+            actual_data = var_array.query(attrs=[column_name])[:][column_name]
-                specific_embedding_directory,
+            self.assertTrue(np.array_equal(expected_data, actual_data))
-                obs_directory,
+
-                var_directory,
+        # Validate embedding
-            ],
+        expected_embedding_data = anndata_object.obsm.get("X_awesome_embedding")
-            obs_files,
+        embedding_array = tiledb.DenseArray(specific_embedding_array_location, mode="r")
-            var_files,
+        actual_embedding_data = embedding_array[:, 0:2]
-        )
+        self.assertTrue(np.array_equal(expected_embedding_data, actual_embedding_data))
+
+        # Validate X matrix if not column shifted
+        if not has_column_encoding:
+            expected_x_data = anndata_object.X
+            if is_sparse:
+                x_array = tiledb.SparseArray(main_x_array_location, mode="r")
+                actual_x_data = np.reshape(x_array[:, :][""], expected_x_data.shape)
+            else:
+                x_array = tiledb.DenseArray(main_x_array_location, mode="r")
+                actual_x_data = x_array[:, :]
+            self.assertTrue(np.array_equal(expected_x_data, actual_x_data))
 
     def _write_anndata_to_file(self, anndata):
         temporary_filename = f"{PROJECT_ROOT}/backend/test/fixtures/{uuid4()}.h5ad"