diff --git a/.github/workflows/test.yaml b/.github/workflows/test.yaml index d451e328..4f6e4b70 100644 --- a/.github/workflows/test.yaml +++ b/.github/workflows/test.yaml @@ -73,6 +73,10 @@ jobs: cache-dependency-glob: pyproject.toml - name: create hatch environment run: uvx hatch env create ${{ matrix.env.name }} + - name: Show installed packages + if: matrix.env.python == '3.13' + run: | + uvx hatch run hatch-test.py3.13:uv pip freeze - name: run tests using hatch env: MPLBACKEND: agg diff --git a/.gitignore b/.gitignore index 63faed52..99b24a24 100644 --- a/.gitignore +++ b/.gitignore @@ -21,12 +21,12 @@ __pycache__/ /docs/generated/ /docs/_build/ /docs/generated/ -docs/api/ehrdata -docs/api/data -docs/api/tools -docs/api/io -docs/api/plot -!docs/api/api.md +/docs/api/ehrdata.*.rst +/docs/api/data +/docs/api/tools +/docs/api/io +/docs/api/plot +!/docs/api/api.md ehrapy_data test.ipynb diff --git a/.readthedocs.yaml b/.readthedocs.yaml index 1b8283c2..f12e003d 100644 --- a/.readthedocs.yaml +++ b/.readthedocs.yaml @@ -10,5 +10,5 @@ build: - asdf global uv latest build: html: - - uvx hatch run docs:build + - uvx hatch run docs:build -W # do not remove -W, fix warnings. - mv docs/_build $READTHEDOCS_OUTPUT diff --git a/CHANGELOG.md b/CHANGELOG.md index 647ff04f..c17307c6 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -14,6 +14,7 @@ and this project adheres to [Semantic Versioning][]. ### Maintenance - Enhanced {doc}`tutorials/getting_started` ([#184](https://github.com/theislab/ehrdata/pull/184)) @eroell +- Move from zarr<3 to zarr>=3 ([#185](https://github.com/theislab/ehrdata/pull/185)) @eroell ### Fixed @@ -21,7 +22,8 @@ and this project adheres to [Semantic Versioning][]. - `EHRData` drops the `.R` field in favor of using `.layers` for any 3D data arrays ([#184](https://github.com/theislab/ehrdata/pull/184)) @eroell - `EHRData`'s shape property will always return a 3 dimensional shape. If an `EHRData` object has flat arrays only, the third dimension will be 1. ([#184](https://github.com/theislab/ehrdata/pull/184)) @eroell - The following functions now take a `layer` argument: {func}`~ehrdata.io.read_csv`, {func}`~ehrdata.io.from_pandas`, {func}`~ehrdata.io.to_pandas`, {func}`~ehrdata.io.omop.setup_variables`, {func}`~ehrdata.io.omop.setup_interval_variables`, {func}`~ehrdata.dt.ehrdata_blobs`, {func}`~ehrdata.dt.physionet2012`. If it is let to its default, `None`, the `.X` field of `EHRData` is used. Since `.X` is 2D in this release, in cases with 3D data, the `layer` argument needs to be used. ([#184](https://github.com/theislab/ehrdata/pull/184)) @eroell - +- {func}`~ehrdata.io.write_zarr` now writes an `EHRData` specific store encoding, with `anndata` as a substore. This change allows to use `AnnData` with its change to consolidated Zarr metadata, and better isolates `AnnData`'s io. ([#185](https://github.com/theislab/ehrdata/pull/185)) @eroell +- {func}`~ehrdata.io.read_zarr` is adapted to read the new store encoding, and can also deal with `AnnData` stores. ([#185](https://github.com/theislab/ehrdata/pull/185)) @eroell ## [0.0.9] diff --git a/docs/conf.py b/docs/conf.py index b8cd6cc1..615bd1e5 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -81,16 +81,18 @@ } intersphinx_mapping = { - "python": ("https://docs.python.org/3", None), "anndata": ("https://anndata.readthedocs.io/en/stable", None), + "fsspec": ("https://filesystem-spec.readthedocs.io/en/stable", None), "h5py": ("https://docs.h5py.org/en/latest", None), + "lamin": ("https://docs.lamin.ai", None), "numpy": ("https://numpy.org/doc/stable", None), - "torch": ("https://pytorch.org/docs/main/", None), - "scipy": ("https://docs.scipy.org/doc/scipy", None), "pandas": ("https://pandas.pydata.org/docs", None), - "zarr": ("https://zarr.readthedocs.io/en/stable", None), + "python": ("https://docs.python.org/3", None), + "scanpy": ("https://scanpy.readthedocs.io/en/stable", None), + "scipy": ("https://docs.scipy.org/doc/scipy", None), + "torch": ("https://docs.pytorch.org/docs/main", None), "vitessce": ("https://python-docs.vitessce.io", None), - "lamin": ("https://docs.lamin.ai", None), + "zarr": ("https://zarr.readthedocs.io/en/stable", None), } # List of patterns, relative to source directory, that match files and @@ -105,7 +107,7 @@ html_static_path = ["_static"] html_css_files = ["css/custom.css"] html_title = project_name -html_logo = "_static/tutorial_images/logo.png" +html_logo = "_static/tutorial_images/ehrdata_logo.png" html_theme_options = { "repository_url": repository_url, "use_repository_button": True, @@ -127,6 +129,10 @@ nitpick_ignore = [ ("py:class", "pathlib._local.Path"), ("py:class", "types.EllipsisType"), + # TODO: remove once https://github.com/sphinx-doc/sphinx/pull/13508 is released + ("py:class", "ehrdata._types.TypeAliasType"), + # typing.Union fails in tutorials/tutorial_time_series_with_pypots + ("py:data", "typing.Union"), # https://github.com/duckdb/duckdb-web/issues/3806 ("py:class", "duckdb.duckdb.DuckDBPyConnection"), ("py:class", "_duckdb.DuckDBPyConnection"), @@ -136,18 +142,18 @@ ("py:class", "awkward.highlevel.Array"), ("py:class", "h5py._hl.dataset.Dataset"), ("py:class", "zarr.core.Array"), + ("py:class", "zarr.core.buffer.core.Buffer"), ("py:class", "ehrdata._compat.ZappyArray"), ("py:class", "dask.array.core.Array"), ("py:class", "anndata.compat.CupyArray"), ("py:class", "anndata.compat.CupySparseMatrix"), ("py:class", "sparse.numba_backend._coo.core.COO"), ("py:class", "sparse._coo.core.COO"), - ("py:data", "typing.Union"), # typing.Union fails in tutorials/tutorial_time_series_with_pypots ] # Redirect broken parameter annotation classes qualname_overrides = { - "zarr._storage.store.Store": "zarr.storage.MemoryStore", + "zarr.storage._common.StorePath": "zarr.storage.StorePath", "zarr.core.group.Group": "zarr.group.Group", "lnschema_core.models.Artifact": "lamindb.Artifact", } diff --git a/pyproject.toml b/pyproject.toml index bec413cd..2ec40c07 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -30,7 +30,10 @@ dependencies = [ "requests", "rich", "xarray", - "zarr<3", + "zarr>=3", +] +optional-dependencies.dask = [ + "dask", ] optional-dependencies.dev = [ "pre-commit", @@ -38,7 +41,7 @@ optional-dependencies.dev = [ ] optional-dependencies.doc = [ "docutils>=0.8,!=0.18.*,!=0.19.*", - "ehrdata[lamin,torch,vitessce]", + "ehrdata[lamin,torch]", # do not include vitessce[all] in doc: https://github.com/keller-mark/esbuild-py/issues/19 "ipykernel", "ipython", "myst-nb>=1.1", @@ -51,17 +54,18 @@ optional-dependencies.doc = [ "sphinx-tabs", "sphinxcontrib-bibtex>=1", "sphinxext-opengraph", + "vitessce", ] optional-dependencies.ehrapy = [ "ehrapy", ] optional-dependencies.lamin = [ - "lamindb", + "lamindb>=1.1", #"omop", # Not needed and causing an issue in the readthedocs build https://github.com/laminlabs/lamindb/issues/3081 ] optional-dependencies.test = [ "coverage", - "ehrdata[lamin,torch,vitessce]", + "ehrdata[dask,lamin,torch,vitessce]", "pytest", "pytest-cov", ] @@ -71,7 +75,6 @@ optional-dependencies.torch = [ optional-dependencies.vitessce = [ "vitessce[all]>=3.4", # the actual dependency - "zarr>=2.18.7,<3", # vitessce does not support zarr>=3 ] # https://docs.pypi.org/project_metadata/#project-urls urls.Documentation = "https://ehrdata.readthedocs.io/" diff --git a/src/ehrdata/core/constants.py b/src/ehrdata/core/constants.py index 18c6a633..3dc91930 100644 --- a/src/ehrdata/core/constants.py +++ b/src/ehrdata/core/constants.py @@ -11,6 +11,8 @@ DEFAULT_TEM_LAYER_NAME = "tem_data" +EHRDATA_ZARR_ENCODING_VERSION = "0.0.1" + # Missing values # -------------- # These values if encountered as strings are considered to represent missing values in the data diff --git a/src/ehrdata/core/ehrdata.py b/src/ehrdata/core/ehrdata.py index 5be3e80f..c7478d4b 100644 --- a/src/ehrdata/core/ehrdata.py +++ b/src/ehrdata/core/ehrdata.py @@ -171,7 +171,7 @@ def _get_layers_3d_dim(layers: Mapping[str, Any] | None) -> int: class EHRData(AnnData): """Model two and three dimensional electronic health record data. - .. figure:: ../../_static/tutorial_images/logo.png + .. figure:: /_static/tutorial_images/ehrdata_logo.png :width: 260px :align: right :class: dark-light diff --git a/src/ehrdata/integrations/vitessce/_config.py b/src/ehrdata/integrations/vitessce/_config.py index 7826dc64..f90b3c07 100644 --- a/src/ehrdata/integrations/vitessce/_config.py +++ b/src/ehrdata/integrations/vitessce/_config.py @@ -11,13 +11,13 @@ from lamindb import Artifact from vitessce import VitessceConfig - from zarr.storage import Store + from zarr.storage import StoreLike def gen_config( path: Path | None = None, *, - store: Path | Store | None = None, + store: Path | StoreLike | None = None, url: str | None = None, artifact: Artifact | None = None, # arguments not about how the store goes in: diff --git a/src/ehrdata/io/zarr.py b/src/ehrdata/io/zarr.py index 574c0a85..a4a869e8 100644 --- a/src/ehrdata/io/zarr.py +++ b/src/ehrdata/io/zarr.py @@ -1,15 +1,20 @@ from __future__ import annotations +import warnings +from functools import wraps from pathlib import Path -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any, Literal import anndata as ad import zarr +import ehrdata as ed from ehrdata._logger import logger +from ehrdata.core.constants import EHRDATA_ZARR_ENCODING_VERSION from ehrdata.io._array_casting import _cast_arrays_dtype_to_float_or_str_if_nonnumeric_object, _cast_variables_to_float if TYPE_CHECKING: + from collections.abc import Callable, Mapping from os import PathLike from ehrdata import EHRData @@ -23,6 +28,8 @@ def read_zarr( ) -> EHRData: """Read a zarr store into an :class:`~ehrdata.EHRData` object. + Can also read :class:`~anndata.AnnData` Zarr stores. In this case, a default `.tem` field is created in the `ehrdata object`. + Args: filename: The filename, or a Zarr storage class. harmonize_missing_values: Whether to call `ehrdata.harmonize_missing_values` on all detected layers. @@ -36,7 +43,6 @@ def read_zarr( >>> ed.io.write_zarr("mimic_2.zarr", edata) >>> edata_from_zarr = ed.io.read_zarr("mimic_2.zarr") """ - import ehrdata as ed from ehrdata import EHRData if isinstance(filename, Path): @@ -44,7 +50,29 @@ def read_zarr( f = filename if isinstance(filename, zarr.Group) else zarr.open(filename, mode="r") - dictionary_for_init = {k: ad.io.read_elem(f[k]) for k, v in dict(f).items() if not k.startswith("raw.")} + if "encoding-type" not in f.attrs: + err = "The zarr store does not contain an encoding-type attribute." + raise ValueError(err) + + if f.attrs["encoding-type"] == "ehrdata": + if "anndata" in f: + dictionary_for_init = { + k: ad.io.read_elem(f["anndata"][k]) for k, v in dict(f["anndata"]).items() if not k.startswith("raw.") + } + else: + err = "The zarr store does not contain the 'anndata' group." + raise ValueError(err) + if "tem" in f: + dictionary_for_init["tem"] = ad.io.read_elem(f["tem"]) + else: + warnings.warn("The zarr store does not contain the 'tem' group.", stacklevel=2) + + elif f.attrs["encoding-type"] == "anndata": + dictionary_for_init = {k: ad.io.read_elem(f[k]) for k, v in dict(f).items() if not k.startswith("raw.")} + + else: + err = f"Unkown encoding-type '{f.attrs['encoding-type']}'." + raise ValueError(err) edata = EHRData(**dictionary_for_init) @@ -62,23 +90,33 @@ def read_zarr( return edata +def _allow_write_nullable_strings[T, **P](f: Callable[P, T]) -> Callable[P, T]: + @wraps(f) + def wrapped(*args: P.args, **kwargs: P.kwargs): + with ad.settings.override(allow_write_nullable_strings=True): + return f(*args, **kwargs) + + return wrapped + + +@_allow_write_nullable_strings def write_zarr( edata: EHRData, filename: str | Path, *, - chunks: bool | int | tuple[int, ...] | None = None, + chunks: Literal["auto" | "ehrdata_auto"] = "auto", convert_strings_to_categoricals: bool = True, ) -> None: """Write :class:`~ehrdata.EHRData` objects to disk. - To write to a `.zarr` file, `X` and `layers` cannot be written as `object` dtype. + To write to a `.zarr` file, `X`, and `layers` cannot be written as `object` dtype. If any of these fields is of `object` dtype, it this function will attempt to cast it to a numeric dtype; if this fails, the field will be casted to a `str` dtype. Args: edata: Central data object. filename: Name of the output file, can also be prefixed with relative or absolute path to save the file to. - chunks: Chunk shape, passed to :meth:`zarr.Group.create_dataset` for `Zarr` version 2, or to :meth:`zarr.Group.create_array` for `Zarr` version 3. - convert_strings_to_categoricals: Convert columns of `str` dtype in `.obs` and `.var` to `categorical` dtype. + chunks: Specify strategy of how data should be chunked. For simplicity, currently only 2 options are available: `"auto"` will write the data with :func:`~anndata.io.write_elem`'s default settings. `"ehrdata_auto"` will write the data chunked (and sharded) based on a heuristic that loosely speaking writes slightly smaller chunks. + convert_strings_to_categoricals: Convert columns of `str` dtype in `.obs` and `.var` and `.tem` to `categorical` dtype. Examples: >>> import ehrdata as ed @@ -86,13 +124,52 @@ def write_zarr( >>> ed.io.write_zarr("mimic_2.zarr", edata) """ filename = Path(filename) - edata = _cast_arrays_dtype_to_float_or_str_if_nonnumeric_object(edata) - ad.AnnData(edata).write_zarr( - filename, - chunks=chunks, - convert_strings_to_categoricals=convert_strings_to_categoricals, - ) - f = zarr.open(filename, mode="a") - ad.io.write_elem(f, "tem", edata.tem) + store = zarr.open_group(filename, mode="a", use_consolidated=False, zarr_format=3) + + adata = ad.AnnData(edata) + + if convert_strings_to_categoricals: + adata.strings_to_categoricals(adata.obs) + adata.strings_to_categoricals(adata.var) + adata.strings_to_categoricals(edata.tem) + + # write_sharded this is a slightly modified version from https://anndata.readthedocs.io/en/stable/tutorials/zarr-v3.html + # write_sharded is intended as a future blueprint of implementing better chunking defaults for ehrdata based based on real usecases + def write_sharded(group: zarr.Group, adata: ad.AnnData): + def callback( + func: ad.experimental.Write, + g: zarr.Group, + k: str, + elem: ad.typing.RWAble, + dataset_kwargs: Mapping[str, Any], + iospec: ad.experimental.IOSpec, + ): + if iospec.encoding_type in {"array"} and not isinstance(elem, list): + dataset_kwargs = { + "shards": tuple(int(2 ** (16 / len(elem.shape))) for _ in elem.shape), + **dataset_kwargs, + } + dataset_kwargs["chunks"] = tuple(i // 2 for i in dataset_kwargs["shards"]) + elif iospec.encoding_type in {"csr_matrix", "csc_matrix"}: + dataset_kwargs = {"shards": (2**16,), "chunks": (2**8,), **dataset_kwargs} + func(g, k, elem, dataset_kwargs=dataset_kwargs) + + return ad.experimental.write_dispatched(group, "/", adata, callback=callback) + + if chunks == "auto": + ad.io.write_elem(store, "anndata", adata) + elif chunks == "ehrdata_auto": + anndata_group = store.create_group("anndata") + write_sharded(anndata_group, adata) + else: + err = ( + f"chunks={chunks} is not implemented. Currently, only chunks='auto' and chunks='ehrdata_auto' is supported." + ) + raise NotImplementedError(err) + + ad.io.write_elem(store, "tem", edata.tem) + + store.attrs["encoding-version"] = EHRDATA_ZARR_ENCODING_VERSION + store.attrs["encoding-type"] = "ehrdata" diff --git a/tests/conftest.py b/tests/conftest.py index 0ebac26e..0e5c8c51 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -110,6 +110,19 @@ def obs_31(): return pd.DataFrame({"obs_col_1": [1, 2, 3]}, index=["obs1", "obs2", "obs3"]) +@pytest.fixture +def obs_32(): + df = pd.DataFrame( + { + "obs_col_1": [1, 2, 3], + "obs_col_2": ["a", "a", "c"], + }, + index=["obs1", "obs2", "obs3"], + ) + df = df.astype({"obs_col_2": "string"}) + return df + + @pytest.fixture def var_21(): return pd.DataFrame({"var_col_1": [1, 2]}, index=["var1", "var2"]) @@ -120,6 +133,19 @@ def var_31(): return pd.DataFrame({"var_col_1": [1, 2, 3]}, index=["var1", "var2", "var3"]) +@pytest.fixture +def var_32(): + df = pd.DataFrame( + { + "var_col_1": [1, 2, 3], + "var_col_2": ["x", "x", "z"], + }, + index=["var1", "var2", "var3"], + ) + df = df.astype({"var_col_2": "string"}) + return df + + @pytest.fixture def tem_11(): return pd.DataFrame({"tem_col_1": [1]}, index=["t1"]) @@ -135,11 +161,29 @@ def tem_31(): return pd.DataFrame({"tem_col_1": [1, 2, 3]}, index=["t1", "t2", "t3"]) +@pytest.fixture +def tem_32(): + df = pd.DataFrame( + { + "tem_col_1": [1, 2, 3], + "tem_col_2": ["l", "l", "n"], + }, + index=["t1", "t2", "t3"], + ) + df = df.astype({"tem_col_2": "string"}) + return df + + @pytest.fixture def edata_333(X_numpy_33, X_numpy_333, obs_31, var_31, tem_31): return EHRData(X=X_numpy_33, layers={DEFAULT_TEM_LAYER_NAME: X_numpy_333}, obs=obs_31, var=var_31, tem=tem_31) +@pytest.fixture +def edata_333_larger_obs_var_tem(X_numpy_33, X_numpy_333, obs_32, var_32, tem_32): + return EHRData(X=X_numpy_33, layers={DEFAULT_TEM_LAYER_NAME: X_numpy_333}, obs=obs_32, var=var_32, tem=tem_32) + + @pytest.fixture def edata_330(X_numpy_33, obs_31, var_31): return EHRData(X=X_numpy_33, obs=obs_31, var=var_31) @@ -208,15 +252,17 @@ def edata_nonnumeric_missing_330(obs_31, var_31): def edata_basic_with_tem_full(): edata_basic_with_tem_dict = { "X": np.ones((5, 4)), - "obs": pd.DataFrame({"survival": [1, 2, 3, 4, 5]}), - "var": pd.DataFrame({"variables": ["var_1", "var_2", "var_3", "var_4"]}), + "obs": pd.DataFrame({"survival": [1, 2, 3, 4, 5]}, index=["obs1", "obs2", "obs3", "obs4", "obs5"]), + "var": pd.DataFrame( + {"variables": ["var_1", "var_2", "var_3", "var_4"]}, index=["var1", "var2", "var3", "var4"] + ), "obsm": {"obs_level_representation": np.ones((5, 2))}, "varm": {"var_level_representation": np.ones((4, 2))}, "layers": {DEFAULT_TEM_LAYER_NAME: np.ones((5, 4)), "other_layer": np.ones((5, 4))}, "obsp": {"obs_level_connectivities": np.ones((5, 5))}, "varp": {"var_level_connectivities": np.random.randn(4, 4)}, "uns": {"information": ["info1"]}, - "tem": pd.DataFrame({"timestep": ["t1"]}), + "tem": pd.DataFrame({"timestep": ["t1"]}, index=["t1"]), } return EHRData(**edata_basic_with_tem_dict) @@ -271,3 +317,24 @@ def _assert_io_read(edata: EHRData): TEST_DATA_PATH = Path(__file__).parent / "data" + + +def _check_aligned_anndata_parts_equal(edata: EHRData, edata_read: EHRData | ad.AnnData): + pd.testing.assert_frame_equal(edata.obs.iloc[:, :1], edata_read.obs.iloc[:, :1]) + pd.testing.assert_frame_equal(edata.var.iloc[:, :1], edata_read.var.iloc[:, :1]) + + for key in edata.obsm: + assert key in edata_read.obsm + assert np.array_equal(edata.obsm[key], edata_read.obsm[key]) + for key in edata.varm: + assert key in edata_read.varm + assert np.array_equal(edata.varm[key], edata_read.varm[key]) + for key in edata.obsp: + assert key in edata_read.obsp + assert np.array_equal(edata.obsp[key], edata_read.obsp[key]) + for key in edata.varp: + assert key in edata_read.varp + assert np.array_equal(edata.varp[key], edata_read.varp[key]) + for key in edata.uns: + assert key in edata_read.uns + assert np.array_equal(edata.uns[key], edata_read.uns[key]) diff --git a/tests/data/toy_zarr/adata_basic.zarr/.zmetadata b/tests/data/toy_zarr/adata_basic.zarr/.zmetadata new file mode 100644 index 00000000..653be232 --- /dev/null +++ b/tests/data/toy_zarr/adata_basic.zarr/.zmetadata @@ -0,0 +1,382 @@ +{ + "metadata": { + ".zattrs": { + "encoding-type": "anndata", + "encoding-version": "0.1.0" + }, + ".zgroup": { + "zarr_format": 2 + }, + "X/.zarray": { + "chunks": [ + 5, + 4 + ], + "compressor": { + "blocksize": 0, + "clevel": 5, + "cname": "lz4", + "id": "blosc", + "shuffle": 1 + }, + "dtype": "