cortical-tools 0.0.1__py3-none-any.whl

cortical_tools/__init__.py

@@ -0,0 +1 @@
+__version__ = "0.0.1"

cortical_tools/common.py

@@ -0,0 +1,473 @@
+import datetime
+from typing import TYPE_CHECKING, Literal, Optional
+
+if TYPE_CHECKING:
+    from meshparty.meshwork import Meshwork
+
+import numpy as np
+import numpy.typing as npt
+import pandas as pd
+import pcg_skel
+import standard_transform
+from caveclient import CAVEclient
+from caveclient.frameworkclient import CAVEclientFull
+from nglui import statebuilder as sb
+
+from .files import TableExportClient
+from .mesh import MeshClient
+
+
+def null_function_factory(arguments_to_set=[]):
+    def null_function(*args, **kwargs):
+        """
+        A placeholder function that does nothing.
+        """
+        raise NotImplementedError(
+            "This function is a placeholder. Arguments must be set in the main class: {}".format(
+                arguments_to_set
+            )
+        )
+
+
+def cell_id_to_root_id_factory(
+    default_datastack_name,
+    default_server_address,
+    lookup_view_name,
+):
+    def cell_id_to_root_id(
+        cell_ids: list[int],
+        client: Optional[CAVEclient] = None,
+        timestamp: Optional[datetime.datetime] = None,
+        materialization_version: Optional[int] = None,
+        filter_empty: bool = True,
+    ) -> npt.NDArray:
+        """
+        Convert cell IDs to root IDs using the CAVEclient.
+
+        Parameters
+        ----------
+        cell_ids : list[int]
+            List of cell IDs to convert.
+        client : CAVEclient, optional
+            CAVEclient instance, by default None.
+        timestamp : datetime.datetime, optional
+            Timestamp for the query, by default current time.
+        materialization_version : int, optional
+            Materialization version, by default None.
+
+        Returns
+        -------
+        pd.Series
+            Series containing the root IDs with cell IDs as index.
+            Cell ids that do not map to a root id will have NaN values.
+        """
+        if client is None:
+            client = CAVEclient(
+                datastack_name=default_datastack_name,
+                server_address=default_server_address,
+            )
+        view_name = lookup_view_name
+        nuc_df = (
+            client.materialize.views[view_name](
+                id=cell_ids,
+            )
+            .query(
+                split_positions=True, materialization_version=materialization_version
+            )
+            .set_index("id")
+        )
+
+        if timestamp is not None:
+            nuc_df["pt_root_id"] = client.chunkedgraph.get_roots(
+                nuc_df["pt_supervoxel_id"], timestamp=timestamp
+            )
+
+        cell_id_df = pd.DataFrame(
+            index=cell_ids,
+        )
+
+        cell_id_df = cell_id_df.merge(
+            nuc_df[["pt_root_id"]],
+            left_index=True,
+            right_index=True,
+            how="inner",
+        ).sort_index()
+        add_back_index = np.setdiff1d(cell_ids, cell_id_df.index.values)
+        if len(add_back_index) > 0 and not filter_empty:
+            cell_id_df = pd.concat(
+                [
+                    cell_id_df,
+                    pd.DataFrame(index=add_back_index, data={"pt_root_id": -1}),
+                ]
+            )
+            cell_id_df = cell_id_df.loc[cell_ids]
+        return cell_id_df["pt_root_id"].rename("root_id")
+
+    return cell_id_to_root_id
+
+
+def _selective_lookup(
+    query_idx: pd.Index,
+    client: CAVEclient,
+    timestamp: datetime.datetime,
+    main_table: str,
+    alt_tables: list,
+):
+    lookup_df_main = client.materialize.tables[main_table](
+        pt_root_id=query_idx.values
+    ).live_query(timestamp=timestamp, split_positions=True)
+    lookup_df_main = lookup_df_main[["id", "pt_root_id"]].set_index("pt_root_id")
+
+    if len(lookup_df_main) < len(query_idx):
+        lookup_df_alts = []
+        for alt_table in alt_tables:
+            lookup_df_alt = client.materialize.tables[alt_table](
+                pt_ref_root_id=query_idx.values,
+            ).live_query(timestamp=timestamp, split_positions=True)
+            if len(lookup_df_alt) > 0:
+                lookup_df_alts.append(lookup_df_alt[["id", "pt_root_id"]])
+        if len(lookup_df_alts) > 0:
+            lookup_df_alt_concat = pd.concat(lookup_df_alts)[
+                ["id", "pt_root_id"]
+            ].set_index("pt_root_id")
+        else:
+            lookup_df_alt_concat = pd.DataFrame()
+
+        lookup_df = pd.concat([lookup_df_main, lookup_df_alt_concat])
+    else:
+        lookup_df = lookup_df_main
+    return lookup_df
+
+
+def root_id_to_cell_id_factory(
+    default_datastack_name: str,
+    default_server_address: str,
+    main_table: str,
+    alt_tables: list[str],
+):
+    def root_id_to_cell_id(
+        root_ids: list[int],
+        client: Optional[CAVEclient] = None,
+        filter_empty: bool = False,
+    ):
+        """
+        Lookup the cell id for a list of root ids in the microns dataset
+        """
+        if client is None:
+            client = CAVEclient(
+                datastack_name=default_datastack_name,
+                server_address=default_server_address,
+            )
+
+        root_ids = np.unique(root_ids)
+        all_cell_df = pd.DataFrame(
+            index=root_ids,
+            data={"cell_id": -1, "done": False},
+        )
+        all_cell_df["cell_id"] = all_cell_df["cell_id"].astype(int)
+        earliest_timestamp = client.chunkedgraph.get_root_timestamps(
+            root_ids, latest=False
+        )
+        latest_timestamp = client.chunkedgraph.get_root_timestamps(
+            root_ids, latest=True
+        )
+        all_cell_df["ts0"] = earliest_timestamp
+        all_cell_df["ts1"] = latest_timestamp
+
+        while not np.all(all_cell_df["done"].values):
+            ts = all_cell_df.query("done == False").ts1.iloc[0]
+            qry_idx = all_cell_df[
+                (all_cell_df.ts0 < ts) & (all_cell_df.ts1 >= ts)
+            ].index
+
+            lookup_df = _selective_lookup(
+                qry_idx,
+                client,
+                ts,
+                main_table=main_table,
+                alt_tables=alt_tables,
+            )
+
+            # Update the pt root ids of found cells, but the done status of all queried cells
+            all_cell_df.loc[lookup_df.index, "cell_id"] = lookup_df["id"].astype(int)
+            all_cell_df.loc[qry_idx, "done"] = True
+
+        if filter_empty:
+            all_cell_df = all_cell_df.query("cell_id != -1")
+        return all_cell_df["cell_id"].astype(int).loc[root_ids]
+
+    return root_id_to_cell_id
+
+
+class DatasetClient:
+    def __init__(
+        self,
+        datastack_name: Optional[str] = None,
+        server_address: Optional[str] = None,
+        caveclient: Optional[CAVEclient] = None,
+        *,
+        materialization_version: Optional[int] = None,
+        cell_id_lookup_view: Optional[str] = None,
+        root_id_lookup_main_table: Optional[str] = None,
+        root_id_lookup_alt_tables: Optional[list[str]] = None,
+        dataset_transform: Optional[standard_transform.datasets.Dataset] = None,
+        static_table_cloudpath: Optional[str] = None,
+    ):
+        if caveclient is None:
+            caveclient = CAVEclient(
+                datastack_name=datastack_name,
+                server_address=server_address,
+                version=materialization_version,
+            )
+        else:
+            datastack_name = caveclient.datastack_name
+            server_address = caveclient.server_address
+
+        self._client = caveclient
+        self._client.version = (
+            int(caveclient.materialize.version)
+            if materialization_version is None
+            else int(materialization_version)
+        )
+
+        self._datastack_name = datastack_name
+        self._server_address = server_address
+
+        if cell_id_lookup_view is None:
+            self.cell_id_to_root_id = null_function_factory(
+                arguments_to_set=["cell_id_lookup_view"]
+            )
+        else:
+            self.cell_id_to_root_id = cell_id_to_root_id_factory(
+                default_datastack_name=datastack_name,
+                default_server_address=server_address,
+                lookup_view_name=cell_id_lookup_view,
+            )
+
+        if root_id_lookup_main_table is None:
+            self.root_id_to_cell_id = null_function_factory(
+                arguments_to_set=["root_id_lookup_main_table"]
+            )
+        else:
+            self.root_id_to_cell_id = root_id_to_cell_id_factory(
+                default_datastack_name=datastack_name,
+                default_server_address=server_address,
+                main_table=root_id_lookup_main_table,
+                alt_tables=root_id_lookup_alt_tables,
+            )
+
+        if dataset_transform is None:
+            self._dataset_transform = null_function_factory(
+                arguments_to_set=["dataset_transform"]
+            )
+        else:
+            self._dataset_transform = dataset_transform
+
+        self._mesh_client = MeshClient(caveclient=caveclient)
+
+        self.tables = self.cave.materialize.tables
+        self.views = self.cave.materialize.views
+
+        if static_table_cloudpath is None:
+            self.exports = null_function_factory(
+                arguments_to_set=["static_table_cloudpath"]
+            )
+        else:
+            self.exports = TableExportClient(static_table_cloudpath)
+
+    def set_export_cloudpath(self, cloudpath: str):
+        """
+        Set the cloud path for static table exports.
+        """
+        self.exports = TableExportClient(cloudpath)
+
+    @property
+    def cave(self) -> CAVEclientFull:
+        """
+        Get the CAVEclient instance for this CortexClient.
+        """
+        return self._client
+
+    @property
+    def datastack_name(self) -> str:
+        """
+        Get the name of the datastack associated with this CortexClient.
+        """
+        return self._datastack_name
+
+    @property
+    def server_address(self) -> str:
+        """
+        Get the server address associated with this CortexClient.
+        """
+        return self._server_address
+
+    @property
+    def dataset_transform(self) -> standard_transform.datasets.Dataset:
+        """
+        Get the dataset transform associated with this CortexClient.
+        """
+        return self._dataset_transform
+
+    @property
+    def mesh(self) -> MeshClient:
+        """
+        Get the MeshClient instance for this CortexClient.
+        """
+        return self._mesh_client
+
+    @property
+    def space(self) -> standard_transform.datasets.Dataset:
+        """
+        Get the dataset transform for this CortexClient.
+        """
+        return self._dataset_transform
+
+    @property
+    def version(self) -> int:
+        """
+        Get the materialization version of the CAVEclient.
+        """
+        return self.cave.materialize.version
+
+    @version.setter
+    def version(self, value: int):
+        """
+        Set the materialization version of the CAVEclient.
+        """
+        self.cave.materialize.version = value
+
+    def get_l2_ids(self, root_id: int) -> list[int]:
+        return self.cave.chunkedgraph.get_roots(root_id, stop_layer=2)
+
+    def get_skeleton(
+        self,
+        root_id: int,
+        synapses: bool = True,
+        restore_graph: bool = False,
+        restore_properties: bool = True,
+        synapse_reference_tables: Optional[dict] = None,
+        skeleton_version: Optional[int] = None,
+        transform: Optional[Literal["rigid", "streamline"]] = None,
+    ) -> "Meshwork":
+        if skeleton_version is None:
+            skeleton_version = 4
+        nrn = pcg_skel.get_meshwork_from_client(
+            client=self.cave,
+            root_id=root_id,
+            synapses=synapses,
+            restore_graph=restore_graph,
+            restore_properties=restore_properties,
+            synapse_reference_tables=synapse_reference_tables,
+            skeleton_version=skeleton_version,
+        )
+        if transform == "rigid":
+            self.space.transform_nm.apply_meshwork_vertices(nrn, inplace=True)
+            if synapses:
+                space_cols = [
+                    x for x in nrn.anno.pre_syn.df.columns if "pt_position" in x
+                ]
+                anno_dict = {"pre_syn": space_cols, "post_syn": space_cols}
+                self.space.transform_nm.apply_meshwork_annotations(
+                    nrn, anno_dict, inplace=True
+                )
+        elif transform == "streamline":
+            self.space.streamline_nm.apply_meshwork_vertices(nrn, inplace=True)
+            if synapses:
+                space_cols = [
+                    x for x in nrn.anno.pre_syn.df.columns if "pt_position" in x
+                ]
+                anno_dict = {"pre_syn": space_cols, "post_syn": space_cols}
+                self.space.streamline_nm.apply_meshwork_annotations(
+                    nrn, anno_dict, inplace=True
+                )
+        return nrn
+
+    @staticmethod
+    def now() -> datetime.datetime:
+        """
+        Get the current time in UTC timezone.
+        """
+        return datetime.datetime.now(datetime.timezone.utc)
+
+    def version_timestamp(self, version: Optional[int] = None) -> datetime.datetime:
+        """
+        Get the timestamp for a specific materialization version.
+
+        Parameters
+        ----------
+        version : int, optional
+            The materialization version to get the timestamp for, by default None (uses current version).
+
+        Returns
+        -------
+        datetime.datetime
+            The timestamp of the specified materialization version.
+        """
+        if version is None:
+            version = self.cave.materialize.version
+        return self.cave.materialize.get_version_timestamp(version)
+
+    def latest_valid_timestamp(
+        self,
+        root_ids: list[int],
+    ) -> npt.NDArray:
+        """
+        Get the latest valid timestamps for a list of root IDs.
+        If the root ID is out of date, it will return the last timestamp at which it was valid and could be used in queries.
+        If the root ID is up to date, it will return the current timestamp at the request time, which is still ensured to be valid.
+
+        Parameters
+        ----------
+        root_ids : list[int]
+            The list of root IDs to get the latest valid timestamps for.
+
+        Returns
+        -------
+        npt.NDArray
+            The latest valid timestamps for the specified root IDs.
+        """
+        return self.cave.chunkedgraph.get_root_timestamps(root_ids, latest=True)
+
+    def neuroglancer_url(
+        self,
+        target_url: Optional[str] = None,
+        clipboard=False,
+        shorten=False,
+    ) -> str:
+        """
+        Get the Neuroglancer URL for the current datastack and version.
+
+        Parameters
+        ----------
+        target_url : str, optional
+            The base URL for Neuroglancer, by default None (uses default server address).
+
+        Returns
+        -------
+        str
+            The Neuroglancer URL.
+        """
+        vs = sb.ViewerState(client=self.cave).add_layers_from_client()
+        if clipboard:
+            return vs.to_clipboard(
+                target_url=target_url,
+                shorten=shorten,
+            )
+        else:
+            return vs.to_url(
+                target_url=target_url,
+                shorten=shorten,
+            )
+
+    def __repr__(self) -> str:
+        return f"DatasetClient(datastack_name={self.datastack_name}, version={self.cave.materialize.version})"
+
+    def _repr_mimebundle_(self, include=None, exclude=None):
+        """Necessary for IPython to detect _repr_html_ for subclasses."""
+        return {"text/html": self.__repr_html__()}, {}
+
+    def __repr_html__(self) -> str:
+        neuroglancer_url = self.neuroglancer_url()
+        html_str = f"<html><body><a href='{neuroglancer_url}'>{self.__repr__()}</a></body></html>"
+        return html_str

cortical_tools/datasets/microns_prod.py

@@ -0,0 +1,31 @@
+from standard_transform.datasets import minnie_ds
+
+from ..common import DatasetClient
+
+
+class MicronsProdClient(DatasetClient):
+    datastack_name = "minnie65_phase3_v1"
+    server_address = "https://global.daf-apis.com"
+    _cell_id_lookup_view = "nucleus_detection_lookup_v1"
+    _root_id_lookup_main_table = "nucleus_detection_v0"
+    _root_id_lookup_alt_tables = ["nucleus_alternative_points"]
+
+    def __init__(self):
+        super().__init__(
+            datastack_name=self.datastack_name,
+            server_address=self.server_address,
+            cell_id_lookup_view=self._cell_id_lookup_view,
+            root_id_lookup_main_table=self._root_id_lookup_main_table,
+            root_id_lookup_alt_tables=self._root_id_lookup_alt_tables,
+            dataset_transform=minnie_ds,
+        )
+
+    def __repr__(self):
+        return f"MicronsPublicClient(datastack_name={self.datastack_name}, version={self.cave.materialize.version})"
+
+    def _repr_mimebundle_(self, include=None, exclude=None):
+        """Necessary for IPython to detect _repr_html_ for subclasses."""
+        return {"text/html": self.__repr_html__()}, {}
+
+
+client = MicronsProdClient()

cortical_tools/datasets/microns_public.py

@@ -0,0 +1,31 @@
+from standard_transform.datasets import minnie_ds
+
+from ..common import DatasetClient
+
+
+class MinniePublicClient(DatasetClient):
+    datastack_name = "minnie65_public"
+    server_address = "https://global.daf-apis.com"
+    _cell_id_lookup_view = "nucleus_detection_lookup_v1"
+    _root_id_lookup_main_table = "nucleus_detection_v0"
+    _root_id_lookup_alt_tables = ["nucleus_alternative_points"]
+
+    def __init__(self):
+        super().__init__(
+            datastack_name=self.datastack_name,
+            server_address=self.server_address,
+            cell_id_lookup_view=self._cell_id_lookup_view,
+            root_id_lookup_main_table=self._root_id_lookup_main_table,
+            root_id_lookup_alt_tables=self._root_id_lookup_alt_tables,
+            dataset_transform=minnie_ds,
+        )
+
+    def __repr__(self):
+        return f"MinniePublicClient(datastack_name={self.datastack_name}, version={self.cave.materialize.version})"
+
+    def _repr_mimebundle_(self, include=None, exclude=None):
+        """Necessary for IPython to detect _repr_html_ for subclasses."""
+        return {"text/html": self.__repr_html__()}, {}
+
+
+client = MinniePublicClient()

cortical_tools/datasets/v1dd.py

@@ -0,0 +1,31 @@
+from standard_transform.datasets import v1dd_ds
+
+from ..common import DatasetClient
+
+
+class V1ddClient(DatasetClient):
+    datastack_name = "v1dd"
+    server_address = "https://global.em.brain.allentech.org"
+    _cell_id_lookup_view = "nucleus_alternative_lookup"
+    _root_id_lookup_main_table = "nucleus_detection_v0"
+    _root_id_lookup_alt_tables = ["nucleus_alternative_points"]
+
+    def __init__(self):
+        super().__init__(
+            datastack_name=self.datastack_name,
+            server_address=self.server_address,
+            cell_id_lookup_view=self._cell_id_lookup_view,
+            root_id_lookup_main_table=self._root_id_lookup_main_table,
+            root_id_lookup_alt_tables=self._root_id_lookup_alt_tables,
+            dataset_transform=v1dd_ds,
+        )
+
+    def __repr__(self):
+        return f"V1ddClient(datastack_name={self.datastack_name}, version={self.cave.materialize.version})"
+
+    def _repr_mimebundle_(self, include=None, exclude=None):
+        """Necessary for IPython to detect _repr_html_ for subclasses."""
+        return {"text/html": self.__repr_html__()}, {}
+
+
+client = V1ddClient()

cortical_tools/datasets/v1dd_public.py

@@ -0,0 +1,31 @@
+from standard_transform.datasets import v1dd_ds
+
+from ..common import DatasetClient
+
+
+class V1ddPublicClient(DatasetClient):
+    datastack_name = "v1dd_public"
+    server_address = "https://global.em.brain.allentech.org"
+    _cell_id_lookup_view = "nucleus_alternative_lookup"
+    _root_id_lookup_main_table = "nucleus_detection_v0"
+    _root_id_lookup_alt_tables = ["nucleus_alternative_points"]
+
+    def __init__(self):
+        super().__init__(
+            datastack_name=self.datastack_name,
+            server_address=self.server_address,
+            cell_id_lookup_view=self._cell_id_lookup_view,
+            root_id_lookup_main_table=self._root_id_lookup_main_table,
+            root_id_lookup_alt_tables=self._root_id_lookup_alt_tables,
+            dataset_transform=v1dd_ds,
+        )
+
+    def __repr__(self):
+        return f"V1ddPublicClient(datastack_name={self.datastack_name}, version={self.cave.materialize.version})"
+
+    def _repr_mimebundle_(self, include=None, exclude=None):
+        """Necessary for IPython to detect _repr_html_ for subclasses."""
+        return {"text/html": self.__repr_html__()}, {}
+
+
+client = V1ddPublicClient()