|
| 1 | +import mudata |
| 2 | +import anndata |
| 3 | +import pandas as pd |
| 4 | +import numpy as np |
| 5 | +from scipy.sparse import issparse |
| 6 | +from mudata import MuData |
| 7 | +from pathlib import Path |
| 8 | +from pandas.testing import assert_frame_equal |
| 9 | +from typing import Literal |
| 10 | +from .typing import AnnotationObjectOrPathLike |
| 11 | + |
| 12 | + |
| 13 | +def _read_if_needed(anndata_mudata_path_or_obj): |
| 14 | + if isinstance(anndata_mudata_path_or_obj, (str, Path)): |
| 15 | + return mudata.read(anndata_mudata_path_or_obj) |
| 16 | + if isinstance(anndata_mudata_path_or_obj, (mudata.MuData, anndata.AnnData)): |
| 17 | + return anndata_mudata_path_or_obj.copy() |
| 18 | + raise AssertionError("Expected 'Path', 'str' to MuData/AnnData " |
| 19 | + "file or MuData/AnnData object.") |
| 20 | + |
| 21 | +def _assert_same_annotation_object_class(left, right): |
| 22 | + assert type(left) == type(right), (f"Two objects are not of the same class:" |
| 23 | + f"\n[Left]:{type(left)}\n[right]:{type(right)}") |
| 24 | + |
| 25 | + |
| 26 | +def assert_mudata_modality_keys_equal(left, right): |
| 27 | + left_keys = set(left.mod.keys()) |
| 28 | + right_keys = set(right.mod.keys()) |
| 29 | + if left_keys!= right_keys: |
| 30 | + raise AssertionError("MuData modalities differ:" |
| 31 | + f"\n[left]:{left_keys}\n[right]:{right_keys}") |
| 32 | + |
| 33 | +def assert_shape_equal(left: AnnotationObjectOrPathLike, right: AnnotationObjectOrPathLike): |
| 34 | + left, right = _read_if_needed(left), _read_if_needed(right) |
| 35 | + _assert_same_annotation_object_class(left, right) |
| 36 | + if left.shape != right.shape: |
| 37 | + raise AssertionError(f"{type(left).__name__} shapes differ:" |
| 38 | + f"\n[left]:{left.shape}\n[right]:{right.shape}") |
| 39 | + if isinstance(left, MuData): |
| 40 | + assert_mudata_modality_keys_equal(left, right) |
| 41 | + for mod_name, modality in left.mod.items(): |
| 42 | + assert_shape_equal(modality, right[mod_name]) |
| 43 | + |
| 44 | + |
| 45 | +def assert_obs_names_equal(left: AnnotationObjectOrPathLike, right: AnnotationObjectOrPathLike, |
| 46 | + *args, **kwargs): |
| 47 | + left, right = _read_if_needed(left), _read_if_needed(right) |
| 48 | + _assert_same_annotation_object_class(left, right) |
| 49 | + pd.testing.assert_index_equal(left.obs_names, right.obs_names, *args, **kwargs) |
| 50 | + if isinstance(left, MuData): |
| 51 | + assert_mudata_modality_keys_equal(left, right) |
| 52 | + for mod_name, modality in left.mod.items(): |
| 53 | + assert_obs_names_equal(modality, right[mod_name]) |
| 54 | + |
| 55 | + |
| 56 | +def assert_var_names_equal(left: AnnotationObjectOrPathLike, right: AnnotationObjectOrPathLike, |
| 57 | + *args, **kwargs): |
| 58 | + left, right = _read_if_needed(left), _read_if_needed(right) |
| 59 | + _assert_same_annotation_object_class(left, right) |
| 60 | + pd.testing.assert_index_equal(left.var_names, right.var_names, *args, **kwargs) |
| 61 | + if isinstance(left, MuData): |
| 62 | + assert_mudata_modality_keys_equal(left, right) |
| 63 | + for mod_name, modality in left.mod.items(): |
| 64 | + assert_var_names_equal(modality, right[mod_name]) |
| 65 | + |
| 66 | + |
| 67 | +def assert_annotation_frame_equal(annotation_attr: Literal["obs", "var"], |
| 68 | + left: AnnotationObjectOrPathLike, right: AnnotationObjectOrPathLike, |
| 69 | + sort=False, *args, **kwargs): |
| 70 | + if not annotation_attr in ("obs", "var"): |
| 71 | + raise ValueError("annotation_attr should be 'obs', or 'var'") |
| 72 | + left, right = _read_if_needed(left), _read_if_needed(right) |
| 73 | + _assert_same_annotation_object_class(left, right) |
| 74 | + left_frame, right_frame = getattr(left, annotation_attr), getattr(right, annotation_attr) |
| 75 | + if sort: |
| 76 | + left_frame, right_frame = left_frame.sort_index(inplace=False), right_frame.sort_index(inplace=False) |
| 77 | + assert_frame_equal(left_frame, right_frame, *args, **kwargs) |
| 78 | + if isinstance(left, MuData): |
| 79 | + assert_mudata_modality_keys_equal(left, right) |
| 80 | + for mod_name, modality in left.mod.items(): |
| 81 | + assert_annotation_frame_equal(annotation_attr, modality, |
| 82 | + right[mod_name], sort=sort, *args, **kwargs) |
| 83 | + |
| 84 | +def _assert_layer_equal(left, right): |
| 85 | + if issparse(left): |
| 86 | + if not issparse(right): |
| 87 | + raise AssertionError("Layers differ:\n[left]: sparse\n[right]: not sparse") |
| 88 | + if left.getformat() != right.getformat(): |
| 89 | + raise AssertionError("Layers format differ:" |
| 90 | + f"\n[left]:{left.getformat()}\n[right]: {right.getformat()}") |
| 91 | + assert np.all(left.indices == right.indices), "Layers differ: indices are not the same" |
| 92 | + assert np.all(left.indptr == right.indptr), "Layers differ: index pointers are not the same" |
| 93 | + np.testing.assert_allclose(left.data, right.data, |
| 94 | + err_msg="Layers data differs.", equal_nan=True) |
| 95 | + else: |
| 96 | + if issparse(right): |
| 97 | + raise AssertionError("Layers differ:\n[left]: not sparse\n[right]: sparse") |
| 98 | + np.testing.assert_allclose(left, right, |
| 99 | + err_msg="Layers data differs.", equal_nan=True) |
| 100 | + |
| 101 | + |
| 102 | +def assert_layers_equal(left: AnnotationObjectOrPathLike, |
| 103 | + right: AnnotationObjectOrPathLike): |
| 104 | + left, right = _read_if_needed(left), _read_if_needed(right) |
| 105 | + _assert_same_annotation_object_class(left, right) |
| 106 | + if left.raw is not None: |
| 107 | + _assert_layer_equal(left.raw, right.raw) |
| 108 | + else: |
| 109 | + if right.raw: |
| 110 | + raise AssertionError("Layer .raw differs: " |
| 111 | + f"\n[left]:{left.raw}\n[right]:{right}") |
| 112 | + if left.X is not None: |
| 113 | + _assert_layer_equal(left.X, right.X) |
| 114 | + if left.layers: |
| 115 | + assert right.layers and (left.layers.keys() == right.layers.keys()), \ |
| 116 | + "Avaiable layers differ:" \ |
| 117 | + f"\n[left]:{left.layers}\n[right]{right.layers}" |
| 118 | + for layer_name, layer in left.layers.items(): |
| 119 | + _assert_layer_equal(layer, right.layers[layer_name]) |
| 120 | + if isinstance(left, MuData): |
| 121 | + assert_mudata_modality_keys_equal(left, right) |
| 122 | + for mod_name, modality in left.mod.items(): |
| 123 | + assert_layers_equal(modality, right[mod_name]) |
| 124 | + |
| 125 | + |
| 126 | +def assert_annotation_objects_equal(left: AnnotationObjectOrPathLike, |
| 127 | + right: AnnotationObjectOrPathLike, |
| 128 | + check_data=True): |
| 129 | + left, right = _read_if_needed(left), _read_if_needed(right) |
| 130 | + _assert_same_annotation_object_class(left, right) |
| 131 | + assert_shape_equal(left, right) |
| 132 | + assert_annotation_frame_equal("obs", left, right) |
| 133 | + assert_annotation_frame_equal("var", left, right) |
| 134 | + if check_data: |
| 135 | + assert_layers_equal(left, right) |
0 commit comments