Refactoring cxg utility classes in preparation for CXG conversion tooling (#1739)

server/common/annotations/annotations.py

@@ -2,9 +2,9 @@ from abc import ABCMeta, abstractmethod
 
 import fastobo
 import fsspec
-from server.common.errors import OntologyLoadFailure
 
-from server.common.utils import series_to_schema
+from server.common.errors import OntologyLoadFailure
+from server.common.utils.type_conversion_utils import get_schema_type_hint_of_array
 
 
 class Annotations(metaclass=ABCMeta):
@@ -44,7 +44,7 @@ class Annotations(metaclass=ABCMeta):
         if labels is not None and not labels.empty:
             for col in labels.columns:
                 col_schema = dict(name=col, writable=True)
-                col_schema.update(series_to_schema(labels[col]))
+                col_schema.update(get_schema_type_hint_of_array(labels[col]))
                 schema.append(col_schema)
 
         return schema

server/common/annotations/local_file_csv.py

@@ -14,7 +14,6 @@ from server.common.errors import AnnotationsError
 
 
 class AnnotationsLocalFile(Annotations):
-
     CXG_ANNO_COLLECTION = "cxg_anno_collection"
 
     def __init__(self, output_dir, output_file):

server/common/app_config.py

@@ -1,23 +1,23 @@
-from server import display_version as cellxgene_display_version
-from flatten_dict import flatten, unflatten
-import os
-from os.path import splitext, basename, isdir
-import sys
-from urllib.parse import urlparse, quote_plus
-import yaml
 import copy
-
-from server.common.default_config import get_default_config
-from server.common.errors import ConfigurationError, DatasetAccessError, OntologyLoadFailure
-from server.data_common.matrix_loader import MatrixDataLoader, MatrixDataCacheManager, MatrixDataType
-from server.common.utils import find_available_port, is_port_available
+import os
+import sys
 import warnings
+from os.path import splitext, basename, isdir
+from urllib.parse import urlparse, quote_plus
+
+import yaml
+from flatten_dict import flatten, unflatten
+
+import server.compute.diffexp_cxg as diffexp_tiledb
+from server import display_version as cellxgene_display_version
+from server.auth.auth import AuthTypeFactory
 from server.common.annotations.hosted_tiledb import AnnotationsHostedTileDB
 from server.common.annotations.local_file_csv import AnnotationsLocalFile
-from server.common.utils import custom_format_warning
-import server.compute.diffexp_cxg as diffexp_tiledb
 from server.common.data_locator import discover_s3_region_name
-from server.auth.auth import AuthTypeFactory
+from server.common.default_config import get_default_config
+from server.common.errors import ConfigurationError, DatasetAccessError, OntologyLoadFailure
+from server.common.utils.utils import custom_format_warning, find_available_port, is_port_available
+from server.data_common.matrix_loader import MatrixDataLoader, MatrixDataCacheManager, MatrixDataType
 from server.db.db_utils import DbUtils
 
 DEFAULT_SERVER_PORT = 5005
@@ -150,7 +150,6 @@ class AppConfig(object):
         parameters is done"""
 
         if messagefn is None:
-
             def noop(message):
                 pass
 
@@ -284,7 +283,7 @@ class AppConfig(object):
             if auth.requires_client_login():
                 config["authentication"].update({
                     "login": auth.get_login_url(data_adaptor),
-                    "logout" : auth.get_logout_url(data_adaptor),
+                    "logout": auth.get_logout_url(data_adaptor),
                 })
 
         return c
@@ -748,7 +747,8 @@ class DatasetConfig(BaseConfig):
             self.user_annotations__ontology__enable = dc["user_annotations"]["ontology"]["enable"]
             self.user_annotations__ontology__obo_location = dc["user_annotations"]["ontology"]["obo_location"]
             self.user_annotations__hosted_tiledb_array__db_uri = dc["user_annotations"]["hosted_tiledb_array"]["db_uri"]
-            self.user_annotations__hosted_tiledb_array__hosted_file_directory = dc["user_annotations"]["hosted_tiledb_array"]["hosted_file_directory"]  # noqa E501
+            self.user_annotations__hosted_tiledb_array__hosted_file_directory = \
+            dc["user_annotations"]["hosted_tiledb_array"]["hosted_file_directory"]  # noqa E501
 
             self.embeddings__names = dc["embeddings"]["names"]
             self.embeddings__enable_reembedding = dc["embeddings"]["enable_reembedding"]
@@ -896,7 +896,7 @@ class DatasetConfig(BaseConfig):
         server_config = self.app_config.server_config
         if server_config.single_dataset__datapath:
             with server_config.matrix_data_cache_manager.data_adaptor(
-                self.tag, server_config.single_dataset__datapath, self.app_config
+                    self.tag, server_config.single_dataset__datapath, self.app_config
             ) as data_adaptor:
                 if self.diffexp__enable and data_adaptor.parameters.get("diffexp_may_be_slow", False):
                     context["messagefn"](

server/common/utils/matrix_utils.py

@@ -0,0 +1,112 @@
+import logging
+
+import numpy as np
+from scipy.stats import mode
+
+
+def is_matrix_sparse(matrix: np.ndarray, sparse_threshold):
+    """
+    Returns whether `matrix` is sparse or not (i.e. dense). This is determined by figuring out whether the matrix has
+    a sparsity percentage below the sparse_threshold, returning the number of non-zeros encountered and number of
+    elements evaluated.  This function may return before evaluating the whole matrix if it can be determined that matrix
+    is not sparse enough.
+    """
+
+    if sparse_threshold == 100.0:
+        return True
+    if sparse_threshold == 0.0:
+        return False
+
+    total_number_of_rows = matrix.shape[0]
+    total_number_of_columns = matrix.shape[1]
+    total_number_of_matrix_elements = total_number_of_rows * total_number_of_columns
+
+    # For efficiency, we count the number of non-zero elements in chunks of the matrix at a time until we hit the
+    # maximum number of non zero values allowed before the matrix is deemed "dense." This allows the function the
+    # quit early for large dense matrices.
+    row_stride = min(int(np.power(10, np.around(np.log10(1e9 / total_number_of_columns)))), 10_000)
+
+    maximum_number_of_non_zero_elements_in_matrix = int(
+        total_number_of_rows * total_number_of_columns * sparse_threshold / 100
+    )
+    number_of_non_zero_elements = 0
+
+    for start_row_index in range(0, total_number_of_rows, row_stride):
+        end_row_index = min(start_row_index + row_stride, total_number_of_rows)
+
+        matrix_subset = matrix[start_row_index:end_row_index, :]
+        if not isinstance(matrix_subset, np.ndarray):
+            matrix_subset = matrix_subset.toarray()
+
+        number_of_non_zero_elements += np.count_nonzero(matrix_subset)
+        if number_of_non_zero_elements > maximum_number_of_non_zero_elements_in_matrix:
+            if end_row_index != total_number_of_rows:
+                percentage_of_non_zero_elements = 100 * number_of_non_zero_elements / (
+                    end_row_index * total_number_of_columns)
+                logging.info(
+                    f"Matrix is not sparse. Percentage of non-zero elements (estimate): "
+                    f"{percentage_of_non_zero_elements:6.2f}")
+            else:
+                percentage_of_non_zero_elements = 100 * number_of_non_zero_elements / total_number_of_matrix_elements
+                logging.info(
+                    f"Matrix is not sparse. Percentage of non-zero elements (exact): "
+                    f"{percentage_of_non_zero_elements:6.2f}")
+            return False
+
+    is_sparse = (100.0 * number_of_non_zero_elements / total_number_of_matrix_elements) < sparse_threshold
+    return is_sparse
+
+
+def get_column_shift_encode_for_matrix(matrix, sparse_threshold):
+    """
+    Returns a column shift if there is a column shift that allows the given matrix to be considered as sparse. Column
+    shift encoding works by taking the most common value in each column, then subtracting that value from each  element
+    of the column.  If each column mostly contains its most common value, then the resulting matrix can be very sparse.
+
+    This function determines if column shift encoding can be used to transform the matrix into a sparse matrix with a
+    sparsity below the sparse_threshold. If so, returns the array that stores this encoding. This function also returns
+    the number of non-zeros encountered and number of elements evaluated.  This function may return before evaluating
+    the whole matrix if it can be determined that the matrix cannot benefit from column shift encoding.
+    """
+
+    total_number_of_rows = matrix.shape[0]
+    total_number_of_columns = matrix.shape[1]
+    total_number_of_matrix_elements = total_number_of_rows * total_number_of_columns
+
+    stride = max(1, 128_000_000 // total_number_of_rows)
+    column_shift = np.zeros(total_number_of_columns)
+
+    maximum_number_of_non_zero_elements_in_matrix = int(
+        total_number_of_rows * total_number_of_columns * sparse_threshold / 100
+    )
+    number_of_non_zero_elements = 0
+
+    for start_column_index in range(0, total_number_of_columns, stride):
+        end_column_index = min(start_column_index + stride, total_number_of_columns)
+
+        matrix_subset = matrix[:, start_column_index:end_column_index]
+        if not isinstance(matrix_subset, np.ndarray):
+            matrix_subset = matrix_subset.toarray()
+
+        matrix_subset_mode = mode(matrix_subset)
+
+        column_shift[start_column_index:end_column_index] = matrix_subset_mode.mode
+        number_of_non_zero_elements += total_number_of_rows * (end_column_index - start_column_index) - np.sum(
+            matrix_subset_mode.count
+        )
+
+        if number_of_non_zero_elements > maximum_number_of_non_zero_elements_in_matrix:
+            if end_column_index != total_number_of_columns:
+                logging.info(
+                    "Matrix is not sparse even with column shift. Percentage of non-zero elements (estimate): %6.2f"
+                    % (100 * number_of_non_zero_elements / end_column_index * total_number_of_rows)
+                )
+            else:
+                logging.info(
+                    "Matrix is not sparse even with column shift. Percentage of non-zero elements (exact): %6.2f"
+                    % (100 * number_of_non_zero_elements / total_number_of_matrix_elements)
+                )
+            return None
+
+    is_sparse = (100.0 * number_of_non_zero_elements / total_number_of_matrix_elements) < sparse_threshold
+    return column_shift if is_sparse else None

server/common/utils/sanitization_utils.py

@@ -0,0 +1,40 @@
+import re
+
+
+def sanitize_values_in_list(list_of_keys: list):
+    """
+    Returns a dictionary mapping of the old keys in the list of `list_of_keys` to its new, clean name that is both
+    safe and unique.
+    """
+
+    if not all([isinstance(key, str) for key in list_of_keys]):
+        raise Exception("List of keys to sanitize must contain all strings.")
+
+    # Mask out [~/.] and anything outside the ASCII range.
+    mask = re.compile(r"[^ -\-0-\[\]-\}]")
+    clean_keys_list = [mask.sub("_", key) for key in list_of_keys]
+
+    # Dedupe the clean keys list
+    deduped_clean_keys_list = []
+    for index, clean_key in enumerate(clean_keys_list):
+        total_occurrences_of_clean_key = clean_keys_list.count(clean_key)
+        total_occurrences_up_until_current_index = clean_keys_list[:index].count(clean_key)
+        deduped_clean_keys_list.append(
+            clean_key + "_" + str(total_occurrences_up_until_current_index + 1)
+            if total_occurrences_of_clean_key > 1
+            else clean_key
+        )
+
+    return dict(zip(list_of_keys, deduped_clean_keys_list))
+
+
+def sanitize_keys_in_dictionary(dict_to_sanitize: dict):
+    """
+    Clean and dedupe the keys in the given dictionary.
+    """
+
+    clean_keys = sanitize_values_in_list(dict_to_sanitize.keys())
+    for original_key, sanitized_key in clean_keys.items():
+        if original_key != sanitized_key:
+            dict_to_sanitize[sanitized_key] = dict_to_sanitize[original_key]
+            del dict_to_sanitize[original_key]

server/common/utils/type_conversion_utils.py

@@ -0,0 +1,93 @@
+import logging
+
+import numpy as np
+import pandas as pd
+
+
+def get_dtype_of_array(array: pd.Series):
+    return get_dtype_and_schema_of_array(array)[0]
+
+
+def get_schema_type_hint_of_array(array: pd.Series):
+    return get_dtype_and_schema_of_array(array)[1]
+
+
+def get_dtype_and_schema_of_array(array: pd.Series):
+    return (get_dtype_from_dtype(array.dtype, array_values=array),
+            get_schema_type_hint_from_dtype(array.dtype, array_values=array))
+
+
+def get_dtype_from_dtype(dtype, array_values=None):
+    """
+    Given a data type, finds the equivalent data type that the array should be encoded as. Notably, this is relevant
+    for 64 bit values which will get downcast to 32 bit.
+    """
+
+    dtype_name = dtype.name
+    dtype_kind = dtype.kind
+
+    if dtype == np.float32 or dtype == np.int32:
+        return dtype
+    if dtype_name == "bool":
+        return np.uint8
+    if dtype_name == "object" and dtype_kind == "O":
+        return np.unicode
+    if dtype_name == "category":
+        return get_dtype_from_dtype(dtype.categories.dtype, dtype.categories)
+
+    if can_cast_to_float32(dtype):
+        return np.float32
+    if can_cast_to_int32(dtype, array_values):
+        return np.int32
+
+    raise TypeError(f"Annotations of type {dtype} are unsupported.")
+
+
+def get_schema_type_hint_from_dtype(dtype, array_values=None):
+    """
+    Returns a dictionary that contains type hints about the data type given, especially if the data type is 64 bit
+    and will be downcast to 32 bit.
+    """
+
+    dtype_name = dtype.name
+    dtype_kind = dtype.kind
+
+    if dtype == np.float32 or dtype == np.int32:
+        return {"type": dtype_name}
+    if dtype_name == "bool":
+        return {"type": "boolean"}
+    if dtype_name == "object" and dtype_kind == "O":
+        return {"type": "string"}
+    if dtype_name == "category":
+        return {"type": "categorical", "categories": dtype.categories.tolist()}
+
+    if can_cast_to_float32(dtype):
+        return {"type": "float32"}
+    if can_cast_to_int32(dtype, array_values):
+        return {"type": "int32"}
+
+    raise TypeError(f"Annotations of type {dtype} are unsupported.")
+
+
+def can_cast_to_float32(dtype):
+    if dtype.kind == "f":
+        if not np.can_cast(dtype, np.float32):
+            logging.warning(f"Type {dtype.name} will be converted to 32 bit float and may lose precision.")
+        return True
+    return False
+
+
+def can_cast_to_int32(dtype, array_values=None):
+    """
+    A type can be cast to 32 bit, overriding the numpy `cast_cast` function if the values in the array that are of
+    the higher precision type has values that are entirely within the range of the downcast type.
+    """
+
+    if dtype.kind in ["i", "u"]:
+        if np.can_cast(dtype, np.int32):
+            return True
+        ii32 = np.iinfo(np.int32)
+        if not array_values.empty and (
+                array_values.min() >= ii32.min and array_values.max() <= ii32.max) or array_values.empty:
+            return True
+    return False

server/common/utils/utils.py

@@ -5,12 +5,11 @@ import logging
 import os
 import pkgutil
 import socket
-import warnings
-
-from flask import json
 from urllib.parse import urlsplit, urljoin
+
 import numpy as np
-import pandas as pd
+from flask import json
+
 from server.common.errors import ConfigurationError
 
 
@@ -94,61 +93,6 @@ def jsonify_numpy(data):
     return json.dumps(data, cls=Float32JSONEncoder, allow_nan=False)
 
 
-def dtype_to_schema(dtype):
-    schema = {}
-    if dtype == np.float32:
-        schema["type"] = "float32"
-    elif dtype == np.int32:
-        schema["type"] = "int32"
-    elif dtype == np.bool_:
-        schema["type"] = "boolean"
-    elif dtype == np.str:
-        schema["type"] = "string"
-    elif dtype == "category":
-        schema["type"] = "categorical"
-        schema["categories"] = dtype.categories.tolist()
-    else:
-        raise TypeError(f"Annotations of type {dtype} are unsupported.")
-    return schema
-
-
-def can_cast_to_float32(array):
-    if array.dtype.kind == "f":
-        if not np.can_cast(array.dtype, np.float32):
-            warnings.warn(f"Annotation {array.name} will be converted to 32 bit float and may lose precision.")
-        return True
-    return False
-
-
-def can_cast_to_int32(array):
-    if array.dtype.kind in ["i", "u"]:
-        if np.can_cast(array.dtype, np.int32):
-            return True
-        ii32 = np.iinfo(np.int32)
-        if array.min() >= ii32.min and array.max() <= ii32.max:
-            return True
-    return False
-
-
-def series_to_schema(array):
-    assert type(array) == pd.Series
-    try:
-        return dtype_to_schema(array.dtype)
-    except TypeError:
-        dtype = array.dtype
-        data_kind = dtype.kind
-        schema = {}
-        if can_cast_to_float32(array):
-            schema["type"] = "float32"
-        elif can_cast_to_int32(array):
-            schema["type"] = "int32"
-        elif data_kind == "O" and dtype == "object":
-            schema["type"] = "string"
-        else:
-            raise TypeError(f"Annotations of type {dtype} are unsupported.")
-        return schema
-
-
 def import_plugins(plugin_module):
     """
     Load optional plugin modules from server.common.plugins