Add 3d integer matrices, string and boolean matrices (#17)

* Generate 3d float matrix

* Add more nd matrices

* Fix merge typo

Author: LouiseDck

src/dummy_anndata/__init__.py

@@ -1,7 +1,7 @@
 from importlib.metadata import version
 
 from .generate_vector import generate_vector, vector_generators
-from .generate_matrix import generate_matrix, matrix_generators
+from .generate_matrix import generate_matrix, matrix_generators, extra_uns_matrix_generators
 from .generate_dict import generate_scalar, generate_dict, scalar_generators
 from .generate_dataframe import generate_dataframe
 from .generate_dataset import generate_dataset
@@ -15,6 +15,7 @@
     "generate_dataset",
     "vector_generators",
     "matrix_generators",
+    "extra_uns_matrix_generators",
     "scalar_generators"
 ]
 

src/dummy_anndata/generate_dataset.py

@@ -4,7 +4,7 @@
 
 from .generate_dataframe import generate_dataframe
 from .generate_dict import generate_dict, scalar_generators
-from .generate_matrix import matrix_generators
+from .generate_matrix import matrix_generators, extra_uns_matrix_generators
 from .generate_vector import vector_generators
 
 
@@ -52,10 +52,11 @@ def generate_dataset(
     varp_types : list of str, optional
         Types of matrices to generate for `varp`. Each type must be a key in `matrix_generators`.
     uns_types : list of str, optional
-        Types of data to generate for `uns`. Each type must be a key in `vector_generators`, `matrix_generators`, or `scalar_generators`.
+        Types of data to generate for `uns`. Each type must be a key in `vector_generators`, `matrix_generators`, or `scalar_generators`
+        or `extra_uns_matrix_generators`.
     nested_uns_types : list of str, optional
         Types of data to generate for the nested `uns` dictionary. They will be a new dictionary at the key `nested`.
-        Each type must be a key in `vector_generators`, `matrix_generators`, or `scalar_generators`.
+        Each type must be a key in `vector_generators`, `matrix_generators`, or `scalar_generators` or `extra_uns_matrix_generators`.
 
     Returns
     -------
@@ -128,10 +129,10 @@ def generate_dataset(
     if varp_types is None:  # varp_types are all matrices
         varp_types = list(matrix_generators.keys())
     if uns_types is None:
-        uns_types = list(vector_generators.keys()) + list(matrix_generators.keys()) + list(scalar_generators.keys())
+        uns_types = list(vector_generators.keys()) + list(matrix_generators.keys()) + list(scalar_generators.keys()) + list(extra_uns_matrix_generators.keys())
     if nested_uns_types is None:
         nested_uns_types = (
-            list(vector_generators.keys()) + list(matrix_generators.keys()) + list(scalar_generators.keys())
+            list(vector_generators.keys()) + list(matrix_generators.keys()) + list(scalar_generators.keys()) + list(extra_uns_matrix_generators.keys())
         )
 
     X = None

src/dummy_anndata/generate_dict.py

@@ -2,7 +2,7 @@
 
 from typing import Union
 
-from .generate_matrix import matrix_generators, generated_matrix_types
+from .generate_matrix import matrix_generators, generated_matrix_types, extra_uns_matrix_generators
 from .generate_vector import vector_generators, generated_vector_types
 from .generate_dataframe import generate_dataframe
 
@@ -32,6 +32,8 @@ def generate_type(type, n_rows, n_cols):
         return vector_generators[type](n_rows)
     if type in matrix_generators:
         return matrix_generators[type](n_rows, n_cols)
+    if type in extra_uns_matrix_generators:
+        return extra_uns_matrix_generators[type](n_rows, n_cols)
     return None
 
 all_types = generated_scalar_types | generated_vector_types | generated_matrix_types

src/dummy_anndata/generate_matrix.py

@@ -9,7 +9,6 @@ def float_mtx(n_obs, n_vars, NAs=False):
         mtx[0, 0] = np.nan
     return mtx
 
-
 def int_mtx(n_obs, n_vars):
     mtx = np.arange(n_obs * n_vars).reshape(n_obs, n_vars)
     return mtx
@@ -21,6 +20,20 @@ def float_mtx_nd(nr_values, dimensions, NAs=False):
         mtx[0, 0] = np.nan
     return mtx
 
+def int_mtx_nd(nr_values, dimensions, NAs=False):
+    mtx = np.arange(nr_values).reshape(dimensions)
+    if NAs:  # numpy matrices do no support pd.NA
+        mtx[0, 0] = np.nan
+    return mtx
+
+def float_mtx_sparse_nd(nr_values, dimensions, row_major=True, NAs=False):
+    mtx = float_mtx_nd(nr_values, dimensions, NAs)
+    if row_major:
+        return sp.sparse.csr_matrix(mtx)
+    else:
+        return sp.sparse.csc_matrix(mtx)
+
+
 # Possible matrix generators
 # integer matrices do not support NAs in Python
 matrix_generators = {
@@ -34,6 +47,20 @@ def float_mtx_nd(nr_values, dimensions, NAs=False):
     "integer_csparse": lambda n_obs, n_vars: sp.sparse.csc_matrix(int_mtx(n_obs, n_vars)),
     "integer_rsparse": lambda n_obs, n_vars: sp.sparse.csr_matrix(int_mtx(n_obs, n_vars)),
     "float_matrix_3d": lambda n_obs, n_vars: float_mtx_nd(n_obs * n_vars * 3, (n_obs, n_vars, 3)),
+    "integer_matrix_3d": lambda n_obs, n_vars: int_mtx_nd(n_obs * n_vars * 3, (n_obs, n_vars, 3)),
+}
+
+def string_matrix_nd(nr_values, dimensions):
+    return np.array(['a' for _ in range(nr_values)]).reshape(dimensions)
+
+def bool_matrix_nd(nr_values, dimensions):
+    return np.array([True if i % 2 else False for i in range(nr_values)]).reshape(dimensions)
+
+extra_uns_matrix_generators = {
+    "string_matrix": lambda n_obs, n_vars: np.array(['a' for _ in range(n_obs * n_vars)]).reshape(n_obs, n_vars),
+    "bool_matrix": lambda n_obs, n_vars: np.array([True for _ in range(n_obs * n_vars)]).reshape(n_obs, n_vars),
+    "string_matrix_3d": lambda n_obs, n_vars: string_matrix_nd(n_obs * n_vars * 3, (n_obs, n_vars, 3)),
+    "bool_matrix_3d": lambda n_obs, n_vars: bool_matrix_nd(n_obs * n_vars * 3, (n_obs, n_vars, 3)),
 }
 
 generated_matrix_types = np.ndarray | sp.sparse.csc_matrix | sp.sparse.csr_matrix
@@ -58,6 +85,10 @@ def generate_matrix(n_obs: int, n_vars: int, matrix_type: str) -> generated_matr
         AssertionError: If the matrix_type is unknown.
 
     """
-    assert matrix_type in matrix_generators, f"Unknown matrix type: {matrix_type}"
+    assert matrix_type in matrix_generators.keys() or matrix_type in extra_uns_matrix_generators.keys(), f"Unknown matrix type: {matrix_type}"
+
+    if matrix_type in matrix_generators:
+        return matrix_generators[matrix_type](n_obs, n_vars)
+    else:
+        return extra_uns_matrix_generators[matrix_type](n_obs, n_vars)
 
-    return matrix_generators[matrix_type](n_obs, n_vars)