nmcp-precomputed 3.0.4__py3-none-any.whl

nmcp/__init__.py

@@ -0,0 +1,4 @@
+from .precomputed import SegmentInfo, SegmentProperty, SegmentTagProperty, SomaSegmentTagProperty, NmcpPropertyValues
+from .precomputed import (ensure_bucket_folders, create_from_json_files, create_from_dict, create_from_data,
+                          remove_skeleton, list_skeletons, extract_neuron_properties, SkeletonComponents)
+from .data import RemoteDataClient, PrecomputedEntry

nmcp/__main__.py

@@ -0,0 +1,13 @@
+import argparse
+
+from .precomputed_worker import main
+
+parser = argparse.ArgumentParser()
+
+parser.add_argument("-u", "--url", help="URL of the GraphQL service")
+parser.add_argument("-a", "--authkey", help="authorization header for GraphQL service")
+parser.add_argument("-o", "--output", help="the output cloud volume location")
+
+args = parser.parse_args()
+
+main(args.url, args.authkey, args.output)

nmcp/data/__init__.py

@@ -0,0 +1,2 @@
+from .remote_data_client import RemoteDataClient
+from .precomputed_entry import PrecomputedEntry

nmcp/data/precomputed_entry.py

@@ -0,0 +1,10 @@
+from dataclasses import dataclass
+
+
+@dataclass
+class PrecomputedEntry:
+    id: str
+    skeletonId: int
+    version: int | None
+    reconstructionId: str
+    generatedAt: float | None

nmcp/data/remote_data_client.py

@@ -0,0 +1,466 @@
+import logging
+from datetime import datetime
+from enum import IntEnum
+from typing_extensions import List, Dict, Any, Optional
+
+from gql import Client, gql
+from gql.transport.requests import RequestsHTTPTransport
+
+from .precomputed_entry import PrecomputedEntry
+
+logger = logging.getLogger(__name__)
+
+atlas_pending_query = gql(
+    """
+    query QueryPrecomputed {
+          pendingPrecomputed {
+            id
+            skeletonId
+            version
+            generatedAt
+            reconstructionId
+          }
+    }
+    """
+)
+
+atlas_update_mutation = gql(
+    """
+    mutation UpdatePrecomputed($id: String!, $status: Int! $version: Int!, $generatedAt: Date!) {
+        updatePrecomputed(id: $id, status: $status, version: $version, generatedAt: $generatedAt) {
+            id
+            skeletonId
+            version
+            generatedAt
+            reconstructionId
+        }
+    }
+    """
+)
+
+atlas_reconstruction_data_query = gql(
+    """
+    query reconstructionAsJson($id: String!, $options: PortalReconstructionInput) {
+        reconstructionAsJson(id: $id, options: $options) {
+        comment
+            neurons {
+                id
+                idString
+                DOI
+                soma {
+                    x
+                    y
+                    z
+                    allenId
+                }
+                sample {
+                    genotype
+                }
+                axon {
+                    x
+                    y
+                    z
+                    radius
+                    sampleNumber
+                    parentNumber
+                    allenId
+                    structureIdentifier
+                }
+                axonChunkInfo {
+                    totalCount
+                    offset
+                    limit
+                    hasMore
+                }
+                dendrite {
+                    x
+                    y
+                    z
+                    radius
+                    sampleNumber
+                    parentNumber
+                    allenId
+                    structureIdentifier
+                }
+                dendriteChunkInfo {
+                    totalCount
+                    offset
+                    limit
+                    hasMore
+                }
+            }
+        }
+    }
+    """
+)
+
+specimen_pending_query = gql(
+    """
+    query QuerySpecimenSpacePrecomputed {
+          specimenSpacePendingPrecomputed {
+            id
+            skeletonId
+            version
+            generatedAt
+            reconstructionId
+          }
+    }
+    """
+)
+
+specimen_update_mutation = gql(
+    """
+    mutation UpdateSpecimenSpacePrecomputed($id: String!, $status: Int! $version: Int!, $generatedAt: Date!) {
+        updateSpecimenSpacePrecomputed(id: $id, status: $status, version: $version, generatedAt: $generatedAt) {
+            id
+            skeletonId
+            version
+            generatedAt
+            reconstructionId
+        }
+    }
+    """
+)
+
+specimen_reconstruction_data_query = gql(
+    """
+    query SpecimenSpaceReconstructionAsJson($id: String!) {
+        specimenSpaceReconstructionAsJson(id: $id) {
+        comment
+            neurons {
+                id
+                idString
+                DOI
+                soma {
+                    x
+                    y
+                    z
+                }
+                sample {
+                    genotype
+                }
+                axon {
+                    x
+                    y
+                    z
+                    radius
+                    sampleNumber
+                    parentNumber
+                    structureIdentifier
+                }
+                dendrite {
+                    x
+                    y
+                    z
+                    radius
+                    sampleNumber
+                    parentNumber
+                    structureIdentifier
+                }
+            }
+        }
+    }
+    """
+)
+
+PRECOMPUTED_VERSION = 1
+
+
+class PrecomputedStatus(IntEnum):
+    Initialized = 0,
+    Pending = 100,
+    Passed = 200,
+    FailedToLoad = 300,
+    FailedToGenerate = 400
+
+
+class RemoteDataClient:
+    def __init__(self, url: str, auth_key: str):
+        transport = RequestsHTTPTransport(
+            url=url,
+            verify=True,
+            retries=3,
+            headers={"Content-Type": "application/json", "Authorization": auth_key}
+        )
+
+        self._client = Client(transport=transport, fetch_schema_from_transport=False)
+
+
+    def find_specimen_space_pending(self) -> List[PrecomputedEntry]:
+        pending = list()
+
+        result = self._client.execute(specimen_pending_query)
+
+        for precomputed in result["specimenSpacePendingPrecomputed"]:
+            pending.append(PrecomputedEntry(**precomputed))
+
+        return pending
+
+    def mark_specimen_generated(self, entry_id: str) -> None:
+        params = {"id": entry_id, "status": PrecomputedStatus.Passed, "version": PRECOMPUTED_VERSION,
+                  "generatedAt": datetime.now().timestamp() * 1000}
+        self._client.execute(specimen_update_mutation, variable_values=params)
+
+    def mark_specimen_failed_load(self, entry_id: str) -> None:
+        params = {"id": entry_id, "status": PrecomputedStatus.FailedToLoad, "version": PRECOMPUTED_VERSION,
+                  "generatedAt": datetime.now().timestamp() * 1000}
+        self._client.execute(specimen_update_mutation, variable_values=params)
+
+    def mark_specimen_failed_generate(self, entry_id: str) -> None:
+        params = {"id": entry_id, "status": PrecomputedStatus.FailedToGenerate, "version": PRECOMPUTED_VERSION,
+                  "generatedAt": datetime.now().timestamp() * 1000}
+        self._client.execute(specimen_update_mutation, variable_values=params)
+
+    def find_atlas_pending(self) -> List[PrecomputedEntry]:
+        pending = list()
+
+        result = self._client.execute(atlas_pending_query)
+
+        for precomputed in result["pendingPrecomputed"]:
+            pending.append(PrecomputedEntry(**precomputed))
+
+        return pending
+
+    def mark_atlas_generated(self, entry_id: str) -> None:
+        params = {"id": entry_id, "status": PrecomputedStatus.Passed, "version": PRECOMPUTED_VERSION,
+                  "generatedAt": datetime.now().timestamp() * 1000}
+        self._client.execute(atlas_update_mutation, variable_values=params)
+
+    def mark_atlas_failed_load(self, entry_id: str) -> None:
+        params = {"id": entry_id, "status": PrecomputedStatus.FailedToLoad, "version": PRECOMPUTED_VERSION,
+                  "generatedAt": datetime.now().timestamp() * 1000}
+        self._client.execute(atlas_update_mutation, variable_values=params)
+
+    def mark_atlas_failed_generate(self, entry_id: str) -> None:
+        params = {"id": entry_id, "status": PrecomputedStatus.FailedToGenerate, "version": PRECOMPUTED_VERSION,
+                  "generatedAt": datetime.now().timestamp() * 1000}
+        self._client.execute(atlas_update_mutation, variable_values=params)
+
+    def _get_reconstruction_part(self, result: Dict[str, Any], name: Optional[str] = None) -> Optional[Dict[str, Any]]:
+        if not result or "reconstructionAsJson" not in result:
+            return None
+
+        if not "neurons" in result["reconstructionAsJson"]:
+            return None
+
+        if len(result["reconstructionAsJson"]["neurons"]) == 0:
+            return None
+
+        if name is not None:
+            if name in result["reconstructionAsJson"]["neurons"][0]:
+                return result["reconstructionAsJson"]["neurons"][0][name]
+            else:
+                return None
+
+        return result["reconstructionAsJson"]["neurons"][0]
+
+    def get_reconstruction_header(self, reconstruction_id: str):
+        """Get header information for a reconstruction."""
+        try:
+            # Do not include axon or dendrite data
+            header_input = {
+                "axonLimit": 0,
+                "dendriteLimit": 0
+            }
+            params = {"id": reconstruction_id, "options": header_input}
+            result = self._client.execute(atlas_reconstruction_data_query, variable_values=params)
+
+            return self._get_reconstruction_part(result)
+
+        except Exception as ex:
+            logger.error(f"Error getting reconstruction header for {reconstruction_id}: {ex}")
+
+        return None
+
+    def get_axon_chunks(self, reconstruction_id: str, chunk_size: int = 25000, offset: int = 0, limit: int = None):
+        """Get axon data in chunks for a reconstruction.
+
+        Args:
+            reconstruction_id: The ID of the reconstruction
+            chunk_size: Number of points to retrieve per request
+            offset: Starting offset for retrieval
+            limit: Maximum total number of points to retrieve (None for all)
+
+        Returns:
+            Dict with "data" (list of axon points) and "chunk_info" (pagination info)
+        """
+        try:
+            axon_data = []
+            current_offset = offset
+            remaining_limit = limit
+
+            while True:
+                # Calculate limit for this chunk
+                request_limit = chunk_size
+                if remaining_limit is not None:
+                    request_limit = min(chunk_size, remaining_limit)
+                    if request_limit <= 0:
+                        break
+
+                axon_input = {
+                    "axonOffset": current_offset,
+                    "axonLimit": request_limit
+                }
+                params = {"id": reconstruction_id, "options": axon_input}
+                result = self._client.execute(atlas_reconstruction_data_query, variable_values=params)
+
+                chunk_data = self._get_reconstruction_part(result)
+
+                if chunk_data:
+                    chunk_points = chunk_data["axon"] or []
+                    axon_data.extend(chunk_points)
+
+                    # Update tracking variables
+                    if remaining_limit is not None:
+                        remaining_limit -= len(chunk_points)
+                    current_offset += len(chunk_points)
+
+                    # Check if we have more data and should continue
+                    chunk_info = chunk_data["axonChunkInfo"]
+                    if not chunk_info or not chunk_info["hasMore"] or len(chunk_points) == 0:
+                        break
+
+                    # If we got fewer points than requested, we"re done
+                    if len(chunk_points) < request_limit:
+                        break
+                else:
+                    break
+
+            return {
+                "data": axon_data,
+                "chunk_info": {
+                    "total_retrieved": len(axon_data),
+                    "offset": offset,
+                    "requested_limit": limit
+                }
+            }
+
+        except Exception as ex:
+            logger.error(f"Error getting axon chunks for {reconstruction_id}: {ex}")
+
+        return None
+
+    def get_dendrite_chunks(self, reconstruction_id: str, chunk_size: int = 25000, offset: int = 0, limit: int = None):
+        """Get dendrite data in chunks for a reconstruction.
+
+        Args:
+            reconstruction_id: The ID of the reconstruction
+            chunk_size: Number of points to retrieve per request
+            offset: Starting offset for retrieval
+            limit: Maximum total number of points to retrieve (None for all)
+
+        Returns:
+            Dict with "data" (list of dendrite points) and "chunk_info" (pagination info)
+        """
+        try:
+            dendrite_data = []
+            current_offset = offset
+            remaining_limit = limit
+
+            while True:
+                # Calculate limit for this chunk
+                request_limit = chunk_size
+                if remaining_limit is not None:
+                    request_limit = min(chunk_size, remaining_limit)
+                    if request_limit <= 0:
+                        break
+
+                dendrite_input = {
+                    "dendriteOffset": current_offset,
+                    "dendriteLimit": request_limit
+                }
+                params = {"id": reconstruction_id, "options": dendrite_input}
+                result = self._client.execute(atlas_reconstruction_data_query, variable_values=params)
+
+                chunk_data = self._get_reconstruction_part(result)
+
+                if chunk_data:
+                    chunk_points = chunk_data["dendrite"] or []
+                    dendrite_data.extend(chunk_points)
+
+                    # Update tracking variables
+                    if remaining_limit is not None:
+                        remaining_limit -= len(chunk_points)
+                    current_offset += len(chunk_points)
+
+                    # Check if we have more data and should continue
+                    chunk_info = chunk_data["dendriteChunkInfo"]
+                    if not chunk_info or not chunk_info["hasMore"] or len(chunk_points) == 0:
+                        break
+
+                    # If we got fewer points than requested, we"re done
+                    if len(chunk_points) < request_limit:
+                        break
+                else:
+                    break
+
+            return {
+                "data": dendrite_data,
+                "chunk_info": {
+                    "total_retrieved": len(dendrite_data),
+                    "offset": offset,
+                    "requested_limit": limit
+                }
+            }
+
+        except Exception as ex:
+            logger.error(f"Error getting dendrite chunks for {reconstruction_id}: {ex}")
+
+        return None
+
+    def get_atlas_reconstruction_data(self, reconstruction_id: str):
+        """Get complete reconstruction data using the individual chunk methods.
+
+        Maintains backward compatibility with the original interface.
+        """
+        try:
+            header = self.get_reconstruction_header(reconstruction_id)
+            if not header:
+                return None
+
+            neuron = {
+                "id": header["id"],
+                "idString": header["idString"],
+                "DOI": header["DOI"],
+                "soma": header["soma"],
+                "axon": [], "dendrite": [],
+                "allenInformation": header["allenInformation"]
+            }
+
+            axon_result = self.get_axon_chunks(reconstruction_id)
+            if axon_result and axon_result["data"]:
+                neuron["axon"] = axon_result["data"]
+
+            dendrite_result = self.get_dendrite_chunks(reconstruction_id)
+            if dendrite_result and dendrite_result["data"]:
+                neuron["dendrite"] = dendrite_result["data"]
+
+            return neuron
+
+        except Exception as ex:
+            logger.error(f"Error getting reconstruction data for {reconstruction_id}: {ex}")
+
+        return None
+
+    def get_specimen_space_reconstruction_data(self, reconstruction_id: str):
+        try:
+            params = {"id": reconstruction_id}
+
+            result = self._client.execute(specimen_reconstruction_data_query, variable_values=params)
+
+
+            if not result or "specimenSpaceReconstructionAsJson" not in result:
+                return None
+
+            if not "neurons" in result["specimenSpaceReconstructionAsJson"]:
+                return None
+
+            if len(result["specimenSpaceReconstructionAsJson"]["neurons"]) == 0:
+                return None
+
+            return result["specimenSpaceReconstructionAsJson"]["neurons"][0]
+
+        except Exception as ex:
+            logger.error(f"Error getting reconstruction header for {reconstruction_id}: {ex}")
+
+        return None

nmcp/from_json.py

@@ -0,0 +1,34 @@
+import argparse
+import glob
+import logging
+import os.path
+import sys
+
+from .precomputed import create_from_json_files
+
+logging.basicConfig(level=logging.WARNING)
+
+
+def main():
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument("input", help="the input json file")
+    parser.add_argument("output", help="the output cloud volume location")
+
+    args = parser.parse_args()
+
+    if os.path.isdir(args.input):
+        input_files = glob.glob(f"{args.input}/*.json")
+    else:
+        input_files = [args.input]
+
+    create_from_json_files(input_files, args.output)
+
+    return True
+
+
+if __name__ == '__main__':
+    if main():
+        sys.exit(0)
+    else:
+        sys.exit(1)

nmcp/from_service.py

@@ -0,0 +1,32 @@
+import argparse
+import logging
+
+from .data import RemoteDataClient
+from .precomputed import create_from_dict
+
+logging.basicConfig(level=logging.WARNING)
+
+
+def main():
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument("-o", "--output", help="the output cloud volume location")
+    parser.add_argument("-u", "--url", help="URL of the GraphQL service")
+    parser.add_argument("-a", "--authkey", help="authorization header for GraphQL service")
+
+    args = parser.parse_args()
+
+    client = RemoteDataClient(args.url, args.authkey)
+
+    pending = client.find_atlas_pending()
+
+    print(pending)
+
+    for pend in pending:
+        data = client.get_atlas_reconstruction_data(pend.reconstructionId)
+        data["skeleton_id"] = pend.skeletonSegmentId
+        create_from_dict(data, args.output)
+
+
+if __name__ == "__main__":
+    main()

nmcp/list_skeletons.py

@@ -0,0 +1,24 @@
+import argparse
+import logging
+
+from .precomputed import list_skeletons
+
+logging.basicConfig(level=logging.WARNING)
+
+
+def main():
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument("-o", "--output", help="the output cloud volume location")
+
+    args = parser.parse_args()
+
+    ids = list_skeletons(args.output)
+
+    print(f"{len(ids)} skeletons in {args.output}")
+
+    print(ids)
+
+
+if __name__ == "__main__":
+    main()

nmcp/precomputed/__init__.py

@@ -0,0 +1,5 @@
+from .segment_property import SegmentProperty
+from .segment_tag_property import SegmentTagProperty, SomaSegmentTagProperty
+from .segment_info import SegmentInfo, NmcpPropertyValues
+from .nmcp_precomputed import (ensure_bucket_folders, create_from_json_files, create_from_dict, create_from_data,
+                               remove_skeleton, list_skeletons, extract_neuron_properties, SkeletonComponents)