diff --git a/notebooks/requirements_nb.txt b/notebooks/requirements_nb.txt index c5abae8..9e2695e 100644 --- a/notebooks/requirements_nb.txt +++ b/notebooks/requirements_nb.txt @@ -1,2 +1,8 @@ # JupyterLab or equivalent is needed for the notebooks jupyterlab +anndata +zarr<3 +dask +scanpy +aiohttp +jsonpath-ng diff --git a/src/crate_builder_script/anndata_zarr_tools.py b/src/crate_builder_script/anndata_zarr_tools.py new file mode 100644 index 0000000..5b5ee0f --- /dev/null +++ b/src/crate_builder_script/anndata_zarr_tools.py @@ -0,0 +1,282 @@ +"""Provides functions to build the Croissant file.""" +import logging +import os +import re +from pprint import pformat +from typing import Any + +import mlcroissant as mlc +import pandas as pd +import zarr +from anndata.io import read_elem +from api_calls import asset_url +from fsspec.implementations.asyn_wrapper import AsyncFileSystemWrapper +from fsspec.implementations.zip import ZipFileSystem +from utils import FILESET_KEYWORDS, rel_path_to_id, safe_id + +LOGGER = logging.getLogger(__name__) + +ANNDATA_ZARR_RE = re.compile(r"(anndata-zarr/.*anndata.zarr)(/.zgroup)") +ANNDATA_ZARR_ZIP_RE = re.compile(r"(anndata-zarr/.*anndata.zarr.zip)") + +PANDAS_TO_CROISSANT_TYPES = { + "int64": mlc.DataType.INTEGER, + "uint8": mlc.DataType.INTEGER, + "float64": mlc.DataType.FLOAT, + "bool": mlc.DataType.BOOL, + "object": mlc.DataType.TEXT, + "string": mlc.DataType.TEXT, +} + + +# +# Copied from hubmapconsortium/portal-visualization utils.py +# many thanks to Nick Akhmetov +# +class SafeZipFileSystem(ZipFileSystem): + """ + Define ZipFileSystem with better FileNoteFound behavior. + + zarr 3 treats FileNotFoundError as "key absent"; ZipFileSystem raises KeyError, + which breaks zarr's probes for optional metadata files. + """ + + def _open(self, path, mode="rb", *args, **kwargs): + try: + return super()._open(path, mode, *args, **kwargs) + except KeyError as e: + raise FileNotFoundError(path) from e + + +def open_zip_zarr_store(url, token=None): + """Open a remote zipped zarr file via URL.""" + headers = {"Authorization": f"Bearer {token}"} if token else {} + fs = SafeZipFileSystem( + fo=url, + target_protocol="https", # NOT remote_protocol/remote_options: silently ignored + target_options={"client_kwargs": {"headers": headers}}, + ) + # zarr 3's FsspecStore needs an async filesystem; ZipFileSystem is sync. + store = zarr.storage.FsspecStore( + AsyncFileSystemWrapper(fs, asynchronous=True), read_only=True + ) + return store + + +def anndata_zarr_fs_func(fs_info: dict) -> dict: + """Generate a FileSet for anndata_zarr.""" + rslt = {key: fs_info[key] for key in FILESET_KEYWORDS} + LOGGER.info(f"_anndata_zarr_fs_func returning:\n{pformat(rslt)}") + return rslt + + +def anndata_zarr_rs_func(rs_info: dict) -> list[dict]: + """Generate a RecordSet for anndata_zarr.""" + id_base = rs_info["id"] + obs_dtypes = rs_info.get("obs_dtypes", {}) + X_shape = rs_info.get("X_shape", (-1, -1)) + + # Build the row index RecordSet + row_field = mlc.Field( + id=safe_id(id_base, "/cells/cell-id"), + name="cell_id", + data_types=[mlc.DataType.TEXT], + source=mlc.Source( + file_set=id_base, extract=mlc.Extract(json_path="$.obs.index") + ), + ) + field_list = [row_field] + row_index_rs = { + "id": safe_id(id_base, "/cell-id-index"), + "name": "cell-id-index", + "description": "individual cell ids", + "fields": field_list, + } + + # Build the column index RecordSet + col_field = mlc.Field( + id=safe_id(id_base, "/cells/channel-id"), + name="channel_id", + data_types=[mlc.DataType.TEXT], + source=mlc.Source( + file_set=id_base, extract=mlc.Extract(json_path="$.var.index") + ), + ) + field_list = [col_field] + col_index_rs = { + "id": safe_id(id_base, "/channel-index"), + "name": "cell-index", + "description": "individual channel ids", + "fields": field_list, + } + + # Build the RecordSet for X, referencing the indices + field_list = [ + { + "id": safe_id(id_base, "/cells/expression-matrix"), + "name": "expression-matrix", + "data_types": [mlc.DataType.FLOAT], + "is_array": True, + "array_shape": f"{X_shape[0]},{X_shape[1]}", + "source": mlc.Source( + file_set=id_base, extract=mlc.Extract(json_path="$.X") + ), + }, + { + "id": safe_id(id_base, "/cells/expression-rows"), + "name": "expression-rows", + "description": "expression matrix row IDs", + "data_types": [mlc.DataType.TEXT], + "references": mlc.Source(field=row_field.id), + }, + { + "id": safe_id(id_base, "/cells/expression-columns"), + "name": "expresion-columns", + "description": "expression matrix column IDs", + "data_types": [mlc.DataType.TEXT], + "references": mlc.Source(field=col_field.id), + }, + ] + field_list = [mlc.Field(**elt) for elt in field_list] # convert to fields + X_rs = { + "id": safe_id(id_base, "/expression-matrix"), + "name": "expression-matrix", + "description": "expression data matrix", + "fields": field_list, + } + + # Guild the Record of fields representing the obs data + field_list = [] + for col, dtype in obs_dtypes.items(): + if cr_dtype := PANDAS_TO_CROISSANT_TYPES.get(str(dtype)): + field_list.append( + { + "id": safe_id(id_base, f"/cells/{col}"), + "name": col, + "data_types": [cr_dtype], + "source": mlc.Source( + file_set=id_base, + extract=mlc.Extract(json_path=f"$.obs.'{col}'"), + ), + } + ) + else: + raise RuntimeError( + f"zarr obs column {col} type {dtype} cannot be mapped to croissant" + ) + field_list.append( + { + "id": safe_id(id_base, "/cells/obs-rows"), + "name": "obs-rows", + "description": "observation matrix row IDs", + "data_types": [mlc.DataType.TEXT], + "references": mlc.Source(field=row_field.id), + }, + ) + field_list = [mlc.Field(**elt) for elt in field_list] # convert to fields + obs_rs = { + "id": safe_id(id_base, "/cells-obs"), + "name": "cells-obs", + "description": "individual cells obs data", + "fields": field_list, + } + + rslt = [row_index_rs, col_index_rs, X_rs, obs_rs] + LOGGER.info(f"_anndata_zarr_rs_func returning:\n{pformat(rslt)}") + return rslt + + +def test_zarr_element_not_empty(key: str, elem: Any) -> None: + """Raise a warning/error if a zarr element is unexpectedly not empty.""" + if isinstance(elem, pd.DataFrame): + if len(elem): + LOGGER.warning( + f"AnnData Zarr element {key} (a Dataframe)" + " was detected but not translated" + ) + elif isinstance(elem, dict): + if elem: + LOGGER.warning( + f"AnnData Zarr element {key} (a dict)" + " was detected but not translated" + ) + else: + LOGGER.error( + f"The AnnData Zarr element {key} is unexpectedly" f" of type {type(elem)}" + ) + + +def anndata_zarr_detector(ds_uuid: str, rel_path: str) -> dict | None: + """If rel_path is an anndata_zarr root, return appropriate RecordSet info.""" + re_match_fs = ANNDATA_ZARR_RE.match(rel_path) + re_match_zip = ANNDATA_ZARR_ZIP_RE.match(rel_path) + re_match_any = re_match_fs or re_match_zip + if re_match_any: + # Information which directly describes the FileSet and can + # be used to derive the RecordSets + anchor_id = rel_path_to_id(re_match_any.group(1)) + zarr_url = asset_url(ds_uuid, anchor_id) + auth_tok = os.environ["AUTH_TOK"] + if re_match_fs: + zarr_obs_store = zarr.open( + zarr_url, + mode="r", + use_consolidated=False, + path="obs", + storage_options={"headers": {"Authorization": f"Bearer {auth_tok}"}}, + ) + zarr_obs_elem = read_elem(zarr_obs_store) + zarr_var_store = zarr.open( + zarr_url, + mode="r", + use_consolidated=False, + path="var", + storage_options={"headers": {"Authorization": f"Bearer {auth_tok}"}}, + ) + zarr_var_elem = read_elem(zarr_var_store) + include_str = "../**/*" + # We have no examples of the following, but they may be info we + # ultimately want to capture + for elt in ["varm", "varp", "obsm", "obsp", "uns"]: + store = zarr.open( + zarr_url, + mode="r", + use_consolidated=False, + path=elt, + storage_options={ + "headers": {"Authorization": f"Bearer {auth_tok}"} + }, + ) + elem = read_elem(store) + test_zarr_element_not_empty(elt, elem) + elif re_match_zip: + zarr_store = open_zip_zarr_store(zarr_url, token=auth_tok) + zarr_group = zarr.open_group(zarr_store, mode="r", use_consolidated=False) + zarr_obs_elem = read_elem(zarr_group["obs"]) + zarr_var_elem = read_elem(zarr_group["var"]) + include_str = "*/*" + # We have no examples of the following, but they may be info we + # ultimately want to capture + for elt in ["varm", "varp", "obsm", "obsp", "uns"]: + if elt in zarr_group: + elem = read_elem(zarr_group[elt]) + test_zarr_element_not_empty(elt, elem) + else: + raise RuntimeError("Internal Error opening remote zarr") + + obs_dtypes = zarr_obs_elem.dtypes + var_index = zarr_var_elem.index + return { + "id": safe_id(f"anndata-zarr-dir-{anchor_id}"), + "name": f"anndata_zarr_directory-{anchor_id}", + "contained_in": [rel_path_to_id(rel_path)], + "includes": include_str, + "encoding_formats": ["application/x-zarr"], + "fileset_func": anndata_zarr_fs_func, + "recordset_func": anndata_zarr_rs_func, + "obs_dtypes": obs_dtypes, + "var_index": var_index, + "X_shape": (zarr_obs_elem.shape[0], zarr_var_elem.shape[0]), + } + else: + return None diff --git a/src/crate_builder_script/build_crate_from_dataset.py b/src/crate_builder_script/build_crate_from_dataset.py index 05e42d5..f57702d 100644 --- a/src/crate_builder_script/build_crate_from_dataset.py +++ b/src/crate_builder_script/build_crate_from_dataset.py @@ -14,7 +14,7 @@ fetch_entity_info, fetch_uuid_files_info, ) -from croissant_wrapper import CroissantWrapper +from croissant_wrapper import CroissantWrapper, detect_recordset_root from extractors import WrappedEntity from rocrate.model import ( ComputerLanguage, @@ -62,13 +62,13 @@ # - the workflow_instance lacks start and end dates # - "homo sapiens" is hard-coded into the keywords. That will have to be # fixed for SenNet. -# - unpublished examples: # TARGET_ID = "HBM567.VCBK.562" # TARGET_ID = "HBM487.HJZB.546" # primary dataset # - published examples: # TARGET_ID = "HBM866.VMBK.952" # TARGET_ID = "HBM473.QLDT.264" # primary dataset # TARGET_ID = "HBM748.CHWC.963" # derived from snare-seq +# TARGET_ID = "HBM953.LMWQ.235" # CODEX processed with zipped Zarr ############### @@ -522,6 +522,7 @@ def main() -> None: blk_idx = {} for file_blk in uuid_files: blk_idx[file_blk["path"]] = file_blk + LOGGER.info(f"blk_idx is:\n{pformat(blk_idx, depth=1)}") # The dataset DOIs point to the Portal, which is basically a landing page, which is # forbidden as the direct link for a dataset under FAIR. So we can't use the DOI @@ -586,10 +587,12 @@ def main() -> None: if "files" in ds_entity: # This is a derived dataset- include only data products and qa_qc files for fl in ds_entity["files"]: + recordset_args = detect_recordset_root(ds_entity["uuid"], fl["rel_path"]) if ( fl.get("is_data_product", False) or fl.get("is_qa_qc", False) or include_all_files + or recordset_args ): LOGGER.debug(f"Adding {fl['rel_path']}") crate.add_file( @@ -598,6 +601,8 @@ def main() -> None: wrapped_croissant.add_file( ds_entity["uuid"], fl, blk_idx.get(fl["rel_path"]) ) + if recordset_args: + wrapped_croissant.add_recordset_from_root(recordset_args) else: LOGGER.debug(f"{fl['rel_path']} is not a data product or qa_qc") else: diff --git a/src/crate_builder_script/croissant_wrapper.py b/src/crate_builder_script/croissant_wrapper.py index 6479bd0..01fbcc9 100644 --- a/src/crate_builder_script/croissant_wrapper.py +++ b/src/crate_builder_script/croissant_wrapper.py @@ -6,51 +6,23 @@ from pprint import pformat import mlcroissant as mlc +from anndata_zarr_tools import anndata_zarr_detector from api_calls import asset_url, entity_iri, fetch_entity_info +from edam_tables import EDAM_INFO from extractors import WrappedEntity +from utils import ( + get_sha256_from_url, + rel_path_to_id, +) LOGGER = logging.getLogger(__name__) -EDAM_INFO = { - ("EDAM_1.24.format_3727", None): {"desc": "tiff", "mime": "image/tiff"}, - ("EDAM_1.24.format_3464", None): {"desc": "json", "mime": "application/json"}, - ("EDAM_1.24.format_3508", None): {"desc": "pdf", "mime": "application/pdf"}, - ("EDAM_1.24.format_3590", None): {"desc": "hdf5", "mime": "application/x-hdf5"}, - ("EDAM_1.24.format_3987", None): {"desc": "zip", "mime": "application/zip"}, - ("EDAM_1.24.format_3752", None): {"desc": "csv", "mime": "text/csv"}, - ("EDAM_1.24.format_3755", None): { - "desc": "tsv", - "mime": "text/tab-separated-values", - }, - ("EDAM_1.24.format_3790", None): { - "desc": "h5ad (anndata)", - "mime": "application/x-hdf5", - }, - ("EDAM_1.24.format_3915", None): {"desc": "zarr", "mime": "application/vnd.zarr"}, - ("EDAM_1.24.format_4006", None): { - "desc": "zarr (spatialdata)", - "mime": "application/vnd.zarr", - }, - ("EDAM_1.24.data_3671", None): {"desc": "plain text", "mime": "text/plain"}, - ("EDAM_1.24.format_3916", ".mtx"): { - "desc": "adjacency matrix", - "mime": "text/plain", - }, - (None, ".html"): {"desc": "unknown but html encoded", "mime": "text/html"}, - (None, ".tsv"): { - "desc": "unknown but tsv encoded", - "mime": "text/tab-separated-values", - }, - (None, ".csv"): {"desc": "unknown but csv encoded", "mime": "text/csv"}, - ("EDAM_1.24.format_2333", None): { - "desc": "binary format", - "mime": "application/octet-stream", - }, - ("EDAM_1.24.format_3003", None): { - "desc": "BED (Browser Extensible Data)", - "mime": "text/x-bed", - }, -} + +def detect_recordset_root(ds_uuid: str, rel_path: str) -> dict | None: + """If rel_path is a fileset root, return appropriate RecordSet info.""" + if rslt := anndata_zarr_detector(ds_uuid, rel_path): + return rslt + return None def _ts_to_iso(ms): @@ -303,17 +275,19 @@ def __init__(self, name: str, description: str): def add_file(self, ds_uuid: str, file_info: dict, file_blk: dict | None) -> None: """Define and describe a single file in the dataset.""" args = { - "id": file_info["rel_path"], + "id": rel_path_to_id(file_info["rel_path"]), "name": file_info["rel_path"], "description": file_info["description"], "content_url": asset_url(ds_uuid, file_info["rel_path"]), } + path_ext = Path(file_info["rel_path"]).suffix + LOGGER.info(f"EDAM lookup {file_info.get('edam_term')} <{path_ext}>") if edam := file_info.get("edam_term"): - if (edam, None) in EDAM_INFO: - args["encoding_formats"] = [EDAM_INFO[(edam, None)]["mime"]] - elif (edam, path_ext) in EDAM_INFO: + if (edam, path_ext) in EDAM_INFO: args["encoding_formats"] = [EDAM_INFO[(edam, path_ext)]["mime"]] + elif (edam, None) in EDAM_INFO: + args["encoding_formats"] = [EDAM_INFO[(edam, None)]["mime"]] elif (None, path_ext) in EDAM_INFO: args["encoding_formats"] = [EDAM_INFO[(None, path_ext)]["mime"]] else: @@ -323,10 +297,32 @@ def add_file(self, ds_uuid: str, file_info: dict, file_blk: dict | None) -> None args["encoding_formats"] = ["application/octet-stream"] if file_blk: args["sha256"] = file_blk["sha256_checksum"] + else: + args["sha256"] = get_sha256_from_url(args["content_url"]) + LOGGER.info(f"args for {file_info['rel_path']} info:\n{pformat(args)}") + LOGGER.info(f"file_info is:\n{pformat(file_info)}") self.file_objects.append(mlc.FileObject(**args)) - # def add_record_set(self, record_set: mlc.RecordSet): - # self.record_sets.append(record_set) + def add_recordset_from_root(self, rs_info: dict) -> None: + """ + Add one or more recordsets, and a fileset if necessary. + + rs_info is a dict returned by detect_recordset_root(). + """ + if fileset_func := rs_info.get("fileset_func"): + self._add_fileset(fileset_func(rs_info)) + if recordset_func := rs_info.get("recordset_func"): + rs_info_list = recordset_func(rs_info) + for rs_info in rs_info_list: + self._add_recordset(rs_info) + + def _add_fileset(self, fs_info: dict) -> None: + """Add a FileSet using info returned by detect_recordset_root().""" + self.file_objects.append(mlc.FileSet(**fs_info)) + + def _add_recordset(self, rs_info: dict) -> None: + """Add a RecordSet.""" + self.record_sets.append(mlc.RecordSet(**rs_info)) def write(self, croissant_filename: str): """Write the Croissant file.""" diff --git a/src/crate_builder_script/edam_tables.py b/src/crate_builder_script/edam_tables.py new file mode 100644 index 0000000..e3692b1 --- /dev/null +++ b/src/crate_builder_script/edam_tables.py @@ -0,0 +1,46 @@ +"""Provides file type information based on edam code and extension.""" +EDAM_INFO = { + ("EDAM_1.24.format_3727", None): {"desc": "tiff", "mime": "image/tiff"}, + ("EDAM_1.24.format_3464", None): {"desc": "json", "mime": "application/json"}, + ("EDAM_1.24.format_3508", None): {"desc": "pdf", "mime": "application/pdf"}, + ("EDAM_1.24.format_3590", None): {"desc": "hdf5", "mime": "application/x-hdf5"}, + ("EDAM_1.24.format_3987", None): {"desc": "zip", "mime": "application/zip"}, + ("EDAM_1.24.format_3752", None): {"desc": "csv", "mime": "text/csv"}, + ("EDAM_1.24.format_3755", None): { + "desc": "tsv", + "mime": "text/tab-separated-values", + }, + ("EDAM_1.24.format_3790", None): { + "desc": "h5ad (anndata)", + "mime": "application/x-hdf5", + }, + ("EDAM_1.24.format_3915", None): {"desc": "zarr", "mime": "application/vnd.zarr"}, + ("EDAM_1.24.format_4006", None): { + "desc": "zarr (spatialdata)", + "mime": "application/vnd.zarr", + }, + ("EDAM_1.24.data_3671", None): {"desc": "plain text", "mime": "text/plain"}, + ("EDAM_1.24.format_3916", ".mtx"): { + "desc": "adjacency matrix", + "mime": "text/plain", + }, + (None, ".html"): {"desc": "unknown but html encoded", "mime": "text/html"}, + (None, ".tsv"): { + "desc": "unknown but tsv encoded", + "mime": "text/tab-separated-values", + }, + (None, ".csv"): {"desc": "unknown but csv encoded", "mime": "text/csv"}, + ("EDAM_1.24.format_2333", ""): { + # this entry covers the .zgroup entry at the top of a Zarr dir + "desc": "Zarr root file", + "mime": "application/json", + }, + ("EDAM_1.24.format_2333", None): { + "desc": "binary format", + "mime": "application/octet-stream", + }, + ("EDAM_1.24.format_3003", None): { + "desc": "BED (Browser Extensible Data)", + "mime": "text/x-bed", + }, +} diff --git a/src/crate_builder_script/utils.py b/src/crate_builder_script/utils.py new file mode 100644 index 0000000..4b21d5b --- /dev/null +++ b/src/crate_builder_script/utils.py @@ -0,0 +1,38 @@ +"""Provide small utility calls.""" +import hashlib + +import requests + +FILESET_KEYWORDS = ["id", "name", "contained_in", "includes", "encoding_formats"] + +RECORDSET_KEYWORDS = ["id", "name", "description", "fields"] + + +def safe_id(base: str, key: str | None = None) -> str: + """Work around mlcommons restrictions on id strings.""" + if key: + id = f"{base}/{key}" + else: + id = f"{base}" + id = id.replace(" ", "--") + return id + + +def rel_path_to_id(rel_path: str) -> str: + """Provide a standard mapping from paths to ID strings.""" + return safe_id(rel_path) + + +def get_sha256_from_url(url: str, chunk_size: int = 16384) -> str: + """ + Given a file URL, return the sha256 of the file. + + Thanks to Google Gemini. Small files only, please. + """ + hash = hashlib.sha256() + with requests.get(url, stream=True) as resp: + resp.raise_for_status() + for chunk in resp.iter_content(chunk_size=chunk_size): + if chunk: + hash.update(chunk) + return hash.hexdigest() diff --git a/src/requirements.txt b/src/requirements.txt index f018eba..28144e9 100644 --- a/src/requirements.txt +++ b/src/requirements.txt @@ -7,3 +7,5 @@ bagit mlcroissant[dev]>=1.1.0 git+https://github.com/hubmapconsortium/cwltool.git@docker-gpu#egg=cwltool setuptools<82 +jsonpath-ng +zarr>=3