From 0ed03da2c8e58414b081ce365f5fbe552f85cfab Mon Sep 17 00:00:00 2001 From: Sean Martin Date: Sat, 19 Sep 2026 19:19:59 +0100 Subject: [PATCH] feat: add SWC file name as annotation --- .../catmaid_skeleton_batch_import/README.md | 26 +++- .../docs/fmost_191805_demo_import.md | 6 +- .../examples/fmost_191805_request.json | 3 +- .../catmaid_skeleton_batch_import/archive.py | 1 + .../artifacts.py | 14 +-- .../catmaid_skeleton_batch_import/config.py | 52 +++++++- .../catmaid_skeleton_batch_import/database.py | 8 +- .../catmaid_skeleton_batch_import/domain.py | 80 ++++++++++++- .../catmaid_skeleton_batch_import/loader.py | 89 ++++++++++---- .../materialization.py | 94 +++++++++++++-- .../orchestrator.py | 39 ++++-- .../catmaid_skeleton_batch_import/planner.py | 37 +++++- .../catmaid_skeleton_batch_import/project.py | 14 ++- .../verification.py | 47 +++++++- .../tests/test_archive.py | 22 +++- .../tests/test_artifacts.py | 34 ++++++ .../tests/test_catmaid_integration.py | 105 +++++++++++++++- .../tests/test_config.py | 67 ++++++++--- .../tests/test_e2e_postgis.py | 26 ++++ .../tests/test_loader.py | 113 ++++++++++++++++-- .../tests/test_orchestrator.py | 6 +- .../tests/test_planner.py | 55 ++++++++- 22 files changed, 826 insertions(+), 112 deletions(-) diff --git a/scripts/catmaid_skeleton_batch_import/README.md b/scripts/catmaid_skeleton_batch_import/README.md index 5645bbb98e..b304881290 100644 --- a/scripts/catmaid_skeleton_batch_import/README.md +++ b/scripts/catmaid_skeleton_batch_import/README.md @@ -56,7 +56,8 @@ The request schema is deliberately small: "stack": { "dimension": [17797, 30801, 10826], "resolution_nm": [350, 350, 1000] - } + }, + "annotations": ["{stem}"] } ``` @@ -64,6 +65,29 @@ The request schema is deliberately small: columns `member_path,external_skeleton_id`. The mapping enriches the output manifest but never chooses CATMAID IDs. +### Annotations + +`annotations` is an optional list of templates. Each template is rendered once +per SWC member and the result becomes a CATMAID annotation on that member's +neuron, so Neuroglancer can search for it as a segment property. The only +placeholder is `{stem}`, the filename without its extension, verbatim: +`191805/0000001.swc` renders `{stem}` as `0000001` and `swc:{stem}` as +`swc:0000001`. + +Every template must contain `{stem}`, so each rendered name belongs to exactly +one neuron; a constant template would create one duplicate annotation row per +neuron. Format specifications, conversions, other placeholders, empty +templates, and repeated templates are rejected. A rendered name longer than 255 +characters or equal to CATMAID's stable-join annotation `stable` aborts +planning. When two members render the same name, for example `a/0000001.swc` +and `b/0000001.swc`, planning aborts and the error names both members. Such +archives imported cleanly before `annotations` existed; rename the files or +split the archive to import them with annotations. + +The manifest gains an `annotations` column holding a JSON array of the rendered +names. Manifests written before this column existed cannot be combined with new +ones. + Coordinates and non-sentinel radii use the same declared physical unit. An SWC radius of `-1` remains CATMAID's unknown-radius sentinel. Voxel and anisotropic source coordinates are not supported. diff --git a/scripts/catmaid_skeleton_batch_import/docs/fmost_191805_demo_import.md b/scripts/catmaid_skeleton_batch_import/docs/fmost_191805_demo_import.md index 28b03f3ec0..82bf6373d0 100644 --- a/scripts/catmaid_skeleton_batch_import/docs/fmost_191805_demo_import.md +++ b/scripts/catmaid_skeleton_batch_import/docs/fmost_191805_demo_import.md @@ -232,7 +232,11 @@ $STAGING_BASE/completed/batch_00001/manifest.tsv.gz The manifest maps each source SWC member to CATMAID neuron and skeleton IDs. The direct bulk load intentionally skips API-level side effects such as import -log rows, default `Import` annotations, and provenance rows. +log rows and provenance rows. Neuron annotations come only from the request's +`annotations` templates (see the README). `examples/fmost_191805_request.json` +uses `["{stem}"]`, so the neuron imported from `191805/0000001.swc` carries the +annotation `0000001`, which Neuroglancer exposes as a searchable segment +property. By default, the loader leaves database triggers enabled because the normal `catmaid_user` role cannot set `session_replication_role`. A privileged DB role can opt into trigger disabling with `LOAD_DISABLE_TRIGGERS=true`, which also diff --git a/scripts/catmaid_skeleton_batch_import/examples/fmost_191805_request.json b/scripts/catmaid_skeleton_batch_import/examples/fmost_191805_request.json index 50d2246f55..fa14fc85cc 100644 --- a/scripts/catmaid_skeleton_batch_import/examples/fmost_191805_request.json +++ b/scripts/catmaid_skeleton_batch_import/examples/fmost_191805_request.json @@ -7,5 +7,6 @@ "stack": { "dimension": [17797, 30801, 10826], "resolution_nm": [350, 350, 1000] - } + }, + "annotations": ["{stem}"] } diff --git a/scripts/catmaid_skeleton_batch_import/src/catmaid_skeleton_batch_import/archive.py b/scripts/catmaid_skeleton_batch_import/src/catmaid_skeleton_batch_import/archive.py index 104fd707b2..5574e0a085 100644 --- a/scripts/catmaid_skeleton_batch_import/src/catmaid_skeleton_batch_import/archive.py +++ b/scripts/catmaid_skeleton_batch_import/src/catmaid_skeleton_batch_import/archive.py @@ -263,6 +263,7 @@ def scan_archive( crc32=info.CRC, display_name=PurePosixPath(info.filename).stem[:255], cable_length_nm=cable_length_nm, + annotations=request.annotation_names(info.filename), ) ) except InvalidInputError: diff --git a/scripts/catmaid_skeleton_batch_import/src/catmaid_skeleton_batch_import/artifacts.py b/scripts/catmaid_skeleton_batch_import/src/catmaid_skeleton_batch_import/artifacts.py index bc2fe8fca2..36f91c6bda 100644 --- a/scripts/catmaid_skeleton_batch_import/src/catmaid_skeleton_batch_import/artifacts.py +++ b/scripts/catmaid_skeleton_batch_import/src/catmaid_skeleton_batch_import/artifacts.py @@ -5,10 +5,11 @@ import csv import gzip import io +import json from pathlib import Path from typing import Iterable, Sequence -from .domain import PlannedBatch +from .domain import PlannedBatch, iter_member_concept_ids from .errors import ImmutableStateError from .util import atomic_write_bytes, deterministic_gzip_bytes, sha256_file @@ -21,6 +22,7 @@ "node_count", "name", "batch_index", + "annotations", ) @@ -117,20 +119,18 @@ def manifest_rows( batch: PlannedBatch, concept_ids: Sequence[int], ) -> list[dict[str, object]]: - if len(concept_ids) != batch.skeleton_count * 3: - raise ValueError("Concept ID count does not match batch skeleton count") rows: list[dict[str, object]] = [] - for index, member in enumerate(batch.members): - neuron_id, skeleton_id, _link_id = concept_ids[index * 3 : index * 3 + 3] + for member, concepts in iter_member_concept_ids(batch, concept_ids): rows.append( { "member_path": member.member_path, "external_skeleton_id": member.external_skeleton_id or "", - "neuron_id": neuron_id, - "skeleton_id": skeleton_id, + "neuron_id": concepts.neuron_id, + "skeleton_id": concepts.skeleton_id, "node_count": member.node_count, "name": member.display_name, "batch_index": batch.index, + "annotations": json.dumps(list(member.annotations)), } ) return rows diff --git a/scripts/catmaid_skeleton_batch_import/src/catmaid_skeleton_batch_import/config.py b/scripts/catmaid_skeleton_batch_import/src/catmaid_skeleton_batch_import/config.py index 9df14e9eb8..cd38de9aab 100644 --- a/scripts/catmaid_skeleton_batch_import/src/catmaid_skeleton_batch_import/config.py +++ b/scripts/catmaid_skeleton_batch_import/src/catmaid_skeleton_batch_import/config.py @@ -5,11 +5,12 @@ import json import math import os +import string from dataclasses import fields from pathlib import Path from typing import Any -from .domain import ImportSettings, IngestionRequest +from .domain import ANNOTATION_PLACEHOLDER, ImportSettings, IngestionRequest from .errors import InvalidInputError @@ -43,8 +44,14 @@ } -def _expect_keys(value: dict[str, Any], expected: set[str], label: str) -> None: - unknown = set(value) - expected +def _expect_keys( + value: dict[str, Any], + expected: set[str], + label: str, + *, + optional: frozenset[str] = frozenset(), +) -> None: + unknown = set(value) - expected - optional missing = expected - set(value) if unknown: raise InvalidInputError( @@ -81,6 +88,37 @@ def _positive_number_tuple(value: Any, label: str) -> tuple[float, float, float] return tuple(result) # type: ignore[return-value] +def _check_annotation_template(template: str) -> None: + label = f"request.annotations template {template!r}" + try: + parsed = list(string.Formatter().parse(template)) + except ValueError as exc: + raise InvalidInputError(f"{label} is malformed: {exc}") from exc + uses_stem = False + for _literal, field_name, format_spec, conversion in parsed: + if field_name is None: + continue + if field_name != ANNOTATION_PLACEHOLDER or format_spec or conversion: + raise InvalidInputError(f"{label} may only use {{stem}}") + uses_stem = True + if not uses_stem: + raise InvalidInputError(f"{label} must contain {{stem}}") + + +def _annotation_templates(value: Any) -> tuple[str, ...]: + if not isinstance(value, list) or not all(isinstance(item, str) for item in value): + raise InvalidInputError("request.annotations must be an array of strings") + if len(set(value)) != len(value): + raise InvalidInputError("request.annotations must not repeat a template") + for template in value: + if not template: + raise InvalidInputError( + "request.annotations must not contain an empty template" + ) + _check_annotation_template(template) + return tuple(value) + + def load_request(path: Path) -> IngestionRequest: try: with path.open("r", encoding="utf-8") as source: @@ -90,7 +128,12 @@ def load_request(path: Path) -> IngestionRequest: if not isinstance(raw, dict): raise InvalidInputError("request must be a JSON object") - _expect_keys(raw, {"schema_version", "source", "stack"}, "request") + _expect_keys( + raw, + {"schema_version", "source", "stack"}, + "request", + optional=frozenset({"annotations"}), + ) if raw["schema_version"] != 1: raise InvalidInputError("request.schema_version must be 1") if not isinstance(raw["source"], dict): @@ -142,6 +185,7 @@ def load_request(path: Path) -> IngestionRequest: raw["stack"]["resolution_nm"], "request.stack.resolution_nm" ), external_skeleton_id_map_path=mapping_path, + annotations=_annotation_templates(raw.get("annotations", [])), ) diff --git a/scripts/catmaid_skeleton_batch_import/src/catmaid_skeleton_batch_import/database.py b/scripts/catmaid_skeleton_batch_import/src/catmaid_skeleton_batch_import/database.py index ca0d436f73..aae6db2f0d 100644 --- a/scripts/catmaid_skeleton_batch_import/src/catmaid_skeleton_batch_import/database.py +++ b/scripts/catmaid_skeleton_batch_import/src/catmaid_skeleton_batch_import/database.py @@ -316,20 +316,20 @@ def allocate_project_ids(connection: Any | None = None) -> dict[str, int]: def allocate_batch_ids( - skeleton_count: int, + concept_count: int, node_count: int, connection: Any | None = None, ) -> dict[str, tuple[int, ...]]: """Reserve exact concept and location IDs for one whole-skeleton batch.""" - if isinstance(skeleton_count, bool) or skeleton_count <= 0: - raise ValueError("skeleton_count must be positive") + if isinstance(concept_count, bool) or concept_count <= 0: + raise ValueError("concept_count must be positive") if isinstance(node_count, bool) or node_count <= 0: raise ValueError("node_count must be positive") connection = _default_connection(connection) return { "concept_ids": allocate_sequence_ids( - "concept_id_seq", skeleton_count * 3, connection + "concept_id_seq", concept_count, connection ), "location_ids": allocate_sequence_ids( "location_id_seq", node_count, connection diff --git a/scripts/catmaid_skeleton_batch_import/src/catmaid_skeleton_batch_import/domain.py b/scripts/catmaid_skeleton_batch_import/src/catmaid_skeleton_batch_import/domain.py index f42b69c270..d998929494 100644 --- a/scripts/catmaid_skeleton_batch_import/src/catmaid_skeleton_batch_import/domain.py +++ b/scripts/catmaid_skeleton_batch_import/src/catmaid_skeleton_batch_import/domain.py @@ -3,8 +3,11 @@ from __future__ import annotations from dataclasses import dataclass -from pathlib import Path -from typing import Any +from pathlib import Path, PurePosixPath +from typing import Any, Iterator, Sequence + + +ANNOTATION_PLACEHOLDER = "stem" @dataclass(frozen=True) @@ -15,6 +18,7 @@ class IngestionRequest: dimension: tuple[int, int, int] resolution_nm: tuple[float, float, float] external_skeleton_id_map_path: Path | None = None + annotations: tuple[str, ...] = () @property def scale_nm(self) -> float: @@ -27,6 +31,10 @@ def max_nm(self) -> tuple[float, float, float]: for dimension, resolution in zip(self.dimension, self.resolution_nm) ) # type: ignore[return-value] + def annotation_names(self, member_path: str) -> tuple[str, ...]: + values = {ANNOTATION_PLACEHOLDER: PurePosixPath(member_path).stem} + return tuple(template.format_map(values) for template in self.annotations) + def to_dict(self) -> dict[str, Any]: source: dict[str, Any] = { "archive_path": str(self.archive_path), @@ -36,7 +44,7 @@ def to_dict(self) -> dict[str, Any]: source["external_skeleton_id_map_path"] = str( self.external_skeleton_id_map_path ) - return { + result: dict[str, Any] = { "schema_version": self.schema_version, "source": source, "stack": { @@ -44,6 +52,9 @@ def to_dict(self) -> dict[str, Any]: "resolution_nm": list(self.resolution_nm), }, } + if self.annotations: + result["annotations"] = list(self.annotations) + return result @dataclass(frozen=True) @@ -108,6 +119,12 @@ class PlannedMember: cable_length_nm: float external_skeleton_id: str | None = None batch_index: int | None = None + annotations: tuple[str, ...] = () + + @property + def concept_count(self) -> int: + """Neuron, skeleton, model_of link, plus an instance and link per annotation.""" + return 3 + 2 * len(self.annotations) def to_dict(self) -> dict[str, Any]: result: dict[str, Any] = { @@ -122,6 +139,8 @@ def to_dict(self) -> dict[str, Any]: result["external_skeleton_id"] = self.external_skeleton_id if self.batch_index is not None: result["batch_index"] = self.batch_index + if self.annotations: + result["annotations"] = list(self.annotations) return result @classmethod @@ -143,6 +162,7 @@ def from_dict(cls, value: dict[str, Any]) -> "PlannedMember": if value.get("batch_index") is not None else None ), + annotations=tuple(str(name) for name in value.get("annotations", ())), ) @@ -159,6 +179,14 @@ def node_count(self) -> int: def skeleton_count(self) -> int: return len(self.members) + @property + def annotation_count(self) -> int: + return sum(len(member.annotations) for member in self.members) + + @property + def concept_count(self) -> int: + return sum(member.concept_count for member in self.members) + def to_summary_dict(self) -> dict[str, Any]: return { "index": self.index, @@ -168,3 +196,49 @@ def to_summary_dict(self) -> dict[str, Any]: "node_count": self.node_count, "status": "planned", } + + +@dataclass(frozen=True) +class AnnotationConceptIds: + name: str + annotation_id: int + link_id: int + + +@dataclass(frozen=True) +class MemberConceptIds: + neuron_id: int + skeleton_id: int + model_of_link_id: int + annotations: tuple[AnnotationConceptIds, ...] + + +def iter_member_concept_ids( + batch: PlannedBatch, concept_ids: Sequence[int] +) -> Iterator[tuple[PlannedMember, MemberConceptIds]]: + if len(concept_ids) != batch.concept_count: + raise ValueError( + f"batch {batch.index}: expected {batch.concept_count} concept IDs, " + f"found {len(concept_ids)}" + ) + offset = 0 + for member in batch.members: + neuron_id, skeleton_id, model_of_link_id, *annotation_ids = concept_ids[ + offset : offset + member.concept_count + ] + offset += member.concept_count + yield member, MemberConceptIds( + neuron_id=int(neuron_id), + skeleton_id=int(skeleton_id), + model_of_link_id=int(model_of_link_id), + annotations=tuple( + AnnotationConceptIds( + name=name, + annotation_id=int(annotation_id), + link_id=int(link_id), + ) + for name, annotation_id, link_id in zip( + member.annotations, annotation_ids[0::2], annotation_ids[1::2] + ) + ), + ) diff --git a/scripts/catmaid_skeleton_batch_import/src/catmaid_skeleton_batch_import/loader.py b/scripts/catmaid_skeleton_batch_import/src/catmaid_skeleton_batch_import/loader.py index f229458403..05dcdec4dc 100644 --- a/scripts/catmaid_skeleton_batch_import/src/catmaid_skeleton_batch_import/loader.py +++ b/scripts/catmaid_skeleton_batch_import/src/catmaid_skeleton_batch_import/loader.py @@ -6,10 +6,18 @@ import math import tempfile import time -from dataclasses import dataclass +from dataclasses import asdict, dataclass from typing import Callable, Iterator, Sequence, TextIO -from .domain import ImportSettings, IngestionRequest, PlannedBatch, PlannedMember, SwcNode +from .domain import ( + AnnotationConceptIds, + ImportSettings, + IngestionRequest, + PlannedBatch, + PlannedMember, + SwcNode, + iter_member_concept_ids, +) from .errors import InvalidInputError, VerificationError @@ -22,15 +30,24 @@ class ProjectContext: neuron_class_id: int skeleton_class_id: int model_of_relation_id: int + annotation_class_id: int | None = None + annotated_with_relation_id: int | None = None @classmethod def from_dict(cls, value: dict[str, object]) -> "ProjectContext": - return cls(**{field: int(value[field]) for field in cls.__dataclass_fields__}) + return cls( + **{ + field: int(value[field]) + for field in cls.__dataclass_fields__ + if value.get(field) is not None + } + ) def to_dict(self) -> dict[str, int]: return { field: int(getattr(self, field)) for field in self.__dataclass_fields__ + if getattr(self, field) is not None } @@ -43,6 +60,7 @@ class ExpectedSkeleton: name: str node_count: int cable_length_nm: float + annotations: tuple[AnnotationConceptIds, ...] def to_validation_dict(self) -> dict[str, object]: return { @@ -53,6 +71,7 @@ def to_validation_dict(self) -> dict[str, object]: "name": self.name, "node_count": self.node_count, "cable_length_nm": self.cable_length_nm, + "annotations": [asdict(annotation) for annotation in self.annotations], } @@ -132,7 +151,7 @@ def prepare_copy_rows( location_ids: Sequence[int], load_nodes: Callable[[PlannedMember], Sequence[SwcNode]], ) -> PreparedBatch: - expected_concepts = batch.skeleton_count * 3 + expected_concepts = batch.concept_count if len(concept_ids) != expected_concepts: raise InvalidInputError( f"batch {batch.index}: expected {expected_concepts} concept IDs, " @@ -147,6 +166,14 @@ def prepare_copy_rows( raise InvalidInputError(f"batch {batch.index}: concept IDs are not unique") if len(set(location_ids)) != len(location_ids): raise InvalidInputError(f"batch {batch.index}: location IDs are not unique") + if batch.annotation_count and ( + project.annotation_class_id is None + or project.annotated_with_relation_id is None + ): + raise InvalidInputError( + f"batch {batch.index}: the project context predates annotation " + "support and cannot receive the requested annotations" + ) class_file = _spooled_text_file() relationship_file = _spooled_text_file() @@ -155,11 +182,10 @@ def prepare_copy_rows( relationship_writer = csv.writer(relationship_file, lineterminator="\n") treenode_writer = csv.writer(treenode_file, lineterminator="\n") - concept_offset = 0 location_offset = 0 expected_skeletons: list[ExpectedSkeleton] = [] try: - for member in batch.members: + for member, concepts in iter_member_concept_ids(batch, concept_ids): nodes = tuple(load_nodes(member)) if len(nodes) != member.node_count: raise InvalidInputError( @@ -178,10 +204,6 @@ def prepare_copy_rows( f"{member.member_path}: cable length no longer matches the plan" ) - neuron_id, skeleton_id, link_id = concept_ids[ - concept_offset : concept_offset + 3 - ] - concept_offset += 3 member_location_ids = location_ids[ location_offset : location_offset + len(nodes) ] @@ -193,7 +215,7 @@ def prepare_copy_rows( class_writer.writerow( [ - neuron_id, + concepts.neuron_id, project.user_id, project.project_id, project.neuron_class_id, @@ -202,7 +224,7 @@ def prepare_copy_rows( ) class_writer.writerow( [ - skeleton_id, + concepts.skeleton_id, project.user_id, project.project_id, project.skeleton_class_id, @@ -211,14 +233,34 @@ def prepare_copy_rows( ) relationship_writer.writerow( [ - link_id, + concepts.model_of_link_id, project.user_id, project.project_id, project.model_of_relation_id, - skeleton_id, - neuron_id, + concepts.skeleton_id, + concepts.neuron_id, ] ) + for annotation in concepts.annotations: + class_writer.writerow( + [ + annotation.annotation_id, + project.user_id, + project.project_id, + project.annotation_class_id, + annotation.name, + ] + ) + relationship_writer.writerow( + [ + annotation.link_id, + project.user_id, + project.project_id, + project.annotated_with_relation_id, + concepts.neuron_id, + annotation.annotation_id, + ] + ) for node in nodes: parent_id = ( @@ -235,7 +277,7 @@ def prepare_copy_rows( repr(node.z * request.scale_nm), project.user_id, project.user_id, - skeleton_id, + concepts.skeleton_id, repr(_scaled_radius(node.radius, request.scale_nm)), parent_id, ] @@ -244,12 +286,13 @@ def prepare_copy_rows( expected_skeletons.append( ExpectedSkeleton( member_path=member.member_path, - neuron_id=int(neuron_id), - skeleton_id=int(skeleton_id), - link_id=int(link_id), + neuron_id=concepts.neuron_id, + skeleton_id=concepts.skeleton_id, + link_id=concepts.model_of_link_id, name=member.display_name, node_count=len(nodes), cable_length_nm=computed_cable, + annotations=concepts.annotations, ) ) @@ -378,8 +421,10 @@ def execute_batch_transaction( counts = { "skeletons": prepared.batch.skeleton_count, "nodes": prepared.batch.node_count, - "class_instance_rows": prepared.batch.skeleton_count * 2, - "relationship_rows": prepared.batch.skeleton_count, + "class_instance_rows": prepared.batch.skeleton_count * 2 + + prepared.batch.annotation_count, + "relationship_rows": prepared.batch.skeleton_count + + prepared.batch.annotation_count, "treenode_rows": prepared.batch.node_count, } @@ -425,6 +470,8 @@ def execute_batch_transaction( neuron_class_id=project.neuron_class_id, skeleton_class_id=project.skeleton_class_id, model_of_relation_id=project.model_of_relation_id, + annotation_class_id=project.annotation_class_id, + annotated_with_relation_id=project.annotated_with_relation_id, expected_skeletons=[ expected.to_validation_dict() for expected in prepared.expected_skeletons diff --git a/scripts/catmaid_skeleton_batch_import/src/catmaid_skeleton_batch_import/materialization.py b/scripts/catmaid_skeleton_batch_import/src/catmaid_skeleton_batch_import/materialization.py index 145ca48dbe..aca5413c97 100644 --- a/scripts/catmaid_skeleton_batch_import/src/catmaid_skeleton_batch_import/materialization.py +++ b/scripts/catmaid_skeleton_batch_import/src/catmaid_skeleton_batch_import/materialization.py @@ -20,6 +20,7 @@ "name", "node_count", "cable_length_nm", + "annotations", ) @@ -68,6 +69,32 @@ def _normalize_expected_skeletons( f"expected skeleton {index}.cable_length_nm must be finite and non-negative" ) values["cable_length_nm"] = cable_length + annotations: list[dict[str, Any]] = [] + for annotation in skeleton["annotations"]: + if not isinstance(annotation["name"], str) or not annotation["name"]: + raise ValueError( + f"expected skeleton {index} has an empty annotation name" + ) + for key in ("annotation_id", "link_id"): + identifier = annotation[key] + if ( + isinstance(identifier, bool) + or not isinstance(identifier, int) + or identifier <= 0 + ): + raise ValueError( + f"expected skeleton {index} annotation {key} " + "must be positive" + ) + annotations.append( + { + "name": annotation["name"], + "annotation_id": annotation["annotation_id"], + "link_id": annotation["link_id"], + } + ) + concept_ids.extend((annotation["annotation_id"], annotation["link_id"])) + values["annotations"] = tuple(annotations) normalized.append(values) concept_ids.extend( (values["neuron_id"], values["skeleton_id"], values["link_id"]) @@ -127,8 +154,9 @@ def _validate_class_instances( user_id: int, neuron_class_id: int, skeleton_class_id: int, + annotation_class_id: int | None, ) -> None: - expected: dict[int, tuple[int, str]] = {} + expected: dict[int, tuple[int | None, str]] = {} for skeleton in skeletons: expected[int(skeleton["neuron_id"])] = ( neuron_class_id, @@ -138,6 +166,11 @@ def _validate_class_instances( skeleton_class_id, str(skeleton["name"]), ) + for annotation in skeleton["annotations"]: + expected[int(annotation["annotation_id"])] = ( + annotation_class_id, + str(annotation["name"]), + ) by_id = _rows_by_id(rows, tuple(expected), "class_instance") for row_id, row in by_id.items(): _, actual_user, actual_project, actual_class, actual_name = row @@ -167,21 +200,34 @@ def _validate_relationships( project_id: int, user_id: int, model_of_relation_id: int, + annotated_with_relation_id: int | None, ) -> None: - expected = {int(item["link_id"]): item for item in skeletons} + expected: dict[int, tuple[int | None, int, int]] = {} + for item in skeletons: + expected[int(item["link_id"])] = ( + model_of_relation_id, + int(item["skeleton_id"]), + int(item["neuron_id"]), + ) + for annotation in item["annotations"]: + expected[int(annotation["link_id"])] = ( + annotated_with_relation_id, + int(item["neuron_id"]), + int(annotation["annotation_id"]), + ) by_id = _rows_by_id(rows, tuple(expected), "class_instance_class_instance") for row_id, row in by_id.items(): _, actual_user, actual_project, relation_id, class_a, class_b = row - item = expected[row_id] + expected_relation_id, expected_a, expected_b = expected[row_id] if ( int(actual_user) != user_id or int(actual_project) != project_id - or int(relation_id) != model_of_relation_id - or int(class_a) != int(item["skeleton_id"]) - or int(class_b) != int(item["neuron_id"]) + or int(relation_id) != expected_relation_id + or int(class_a) != expected_a + or int(class_b) != expected_b ): _fail( - "model_of relationship ownership or endpoints do not match", + "relationship ownership or endpoints do not match", id=row_id, ) @@ -316,6 +362,8 @@ def validate_batch( model_of_relation_id: int, skeletons: Sequence[Mapping[str, Any]], location_ids: Sequence[int], + annotation_class_id: int | None, + annotated_with_relation_id: int | None, connection: Any | None = None, cable_abs_tolerance_nm: float = DEFAULT_CABLE_ABS_TOLERANCE_NM, cable_rel_tolerance: float = DEFAULT_CABLE_REL_TOLERANCE, @@ -328,14 +376,32 @@ def validate_batch( ) if min(cable_abs_tolerance_nm, cable_rel_tolerance) < 0: raise ValueError("cable tolerances cannot be negative") + annotation_count = sum(len(item["annotations"]) for item in normalized) + if annotation_count and ( + annotation_class_id is None or annotated_with_relation_id is None + ): + raise ValueError( + "annotation identifiers are required when annotations are expected" + ) connection = _default_connection(connection) class_ids = tuple( identifier for item in normalized - for identifier in (item["neuron_id"], item["skeleton_id"]) + for identifier in ( + item["neuron_id"], + item["skeleton_id"], + *(annotation["annotation_id"] for annotation in item["annotations"]), + ) + ) + link_ids = tuple( + identifier + for item in normalized + for identifier in ( + item["link_id"], + *(annotation["link_id"] for annotation in item["annotations"]), + ) ) - link_ids = tuple(item["link_id"] for item in normalized) skeleton_ids = tuple(item["skeleton_id"] for item in normalized) with connection.cursor() as cursor: @@ -403,6 +469,7 @@ def validate_batch( user_id=user_id, neuron_class_id=neuron_class_id, skeleton_class_id=skeleton_class_id, + annotation_class_id=annotation_class_id, ) _validate_relationships( relationship_rows, @@ -410,6 +477,7 @@ def validate_batch( project_id=project_id, user_id=user_id, model_of_relation_id=model_of_relation_id, + annotated_with_relation_id=annotated_with_relation_id, ) parents = _validate_treenodes( treenode_rows, @@ -435,10 +503,12 @@ def validate_batch( skeleton_count = len(normalized) node_count = len(normalized_locations) return { - "class_instance_count": skeleton_count * 2, + "class_instance_count": skeleton_count * 2 + annotation_count, "neuron_count": skeleton_count, "skeleton_count": skeleton_count, + "annotation_count": annotation_count, "relationship_count": skeleton_count, + "annotation_relationship_count": annotation_count, "treenode_count": node_count, "edge_count": node_count, "summary_count": skeleton_count, @@ -455,6 +525,8 @@ def materialize_and_verify_batch( skeletons: Sequence[Mapping[str, Any]] | None = None, expected_skeletons: Sequence[Mapping[str, Any]] | None = None, location_ids: Sequence[int], + annotation_class_id: int | None, + annotated_with_relation_id: int | None, connection: Any | None = None, cursor: Any | None = None, rebuild_edges: Any | None = None, @@ -529,6 +601,8 @@ def materialize_and_verify_batch( model_of_relation_id=model_of_relation_id, skeletons=normalized, location_ids=location_ids, + annotation_class_id=annotation_class_id, + annotated_with_relation_id=annotated_with_relation_id, connection=connection, cable_abs_tolerance_nm=cable_abs_tolerance_nm, cable_rel_tolerance=cable_rel_tolerance, diff --git a/scripts/catmaid_skeleton_batch_import/src/catmaid_skeleton_batch_import/orchestrator.py b/scripts/catmaid_skeleton_batch_import/src/catmaid_skeleton_batch_import/orchestrator.py index b1f3863a27..0e0a800f10 100644 --- a/scripts/catmaid_skeleton_batch_import/src/catmaid_skeleton_batch_import/orchestrator.py +++ b/scripts/catmaid_skeleton_batch_import/src/catmaid_skeleton_batch_import/orchestrator.py @@ -7,6 +7,7 @@ import sys import time import zipfile +from dataclasses import asdict from pathlib import Path from typing import Any, Callable, Mapping, Sequence @@ -25,7 +26,14 @@ load_settings, settings_from_state, ) -from .domain import ImportSettings, IngestionRequest, PlannedBatch, PlannedMember, SwcNode +from .domain import ( + ImportSettings, + IngestionRequest, + PlannedBatch, + PlannedMember, + SwcNode, + iter_member_concept_ids, +) from .errors import ( ImmutableStateError, InvalidInputError, @@ -494,7 +502,7 @@ def _prepare_batch_artifacts( concept_ids = read_id_artifact( store.state_dir, artifacts["concept_ids"], - expected_count=batch.skeleton_count * 3, + expected_count=batch.concept_count, ) else: from .database import allocate_sequence_ids @@ -502,7 +510,7 @@ def _prepare_batch_artifacts( try: concept_ids = list( allocate_sequence_ids( - "concept_id_seq", batch.skeleton_count * 3, connection + "concept_id_seq", batch.concept_count, connection ) ) except Exception as exc: @@ -751,8 +759,8 @@ def _run_batch( batch_skeleton_count=batch.skeleton_count, node_count=batch.node_count, row_counts={ - "class_instance_rows": batch.skeleton_count * 2, - "relationship_rows": batch.skeleton_count, + "class_instance_rows": batch.skeleton_count * 2 + batch.annotation_count, + "relationship_rows": batch.skeleton_count + batch.annotation_count, "treenode_rows": batch.node_count, }, ) @@ -936,25 +944,25 @@ def _expected_batch_data( concept_ids = read_id_artifact( store.state_dir, batch_state["artifacts"]["concept_ids"], - expected_count=batch.skeleton_count * 3, + expected_count=batch.concept_count, ) location_ids = read_id_artifact( store.state_dir, batch_state["artifacts"]["location_ids"], expected_count=batch.node_count, ) - for member_index, member in enumerate(batch.members): - neuron_id, skeleton_id, link_id = concept_ids[ - member_index * 3 : member_index * 3 + 3 - ] + for member, concepts in iter_member_concept_ids(batch, concept_ids): expected_skeletons.append( { - "neuron_id": neuron_id, - "skeleton_id": skeleton_id, - "link_id": link_id, + "neuron_id": concepts.neuron_id, + "skeleton_id": concepts.skeleton_id, + "link_id": concepts.model_of_link_id, "name": member.display_name, "node_count": member.node_count, "cable_length_nm": member.cable_length_nm, + "annotations": [ + asdict(annotation) for annotation in concepts.annotations + ], } ) return expected_skeletons, location_ids @@ -991,6 +999,8 @@ def _verify_database_by_batch( model_of_relation_id=project.model_of_relation_id, skeletons=expected_skeletons, location_ids=location_ids, + annotation_class_id=project.annotation_class_id, + annotated_with_relation_id=project.annotated_with_relation_id, connection=connection, ) if ( @@ -1026,6 +1036,9 @@ def _verify_database_by_batch( model_of_relation_id=project.model_of_relation_id, expected_skeleton_count=plan.member_count, expected_node_count=plan.total_nodes, + annotation_class_id=project.annotation_class_id, + annotated_with_relation_id=project.annotated_with_relation_id, + expected_annotation_count=plan.total_annotations, connection=connection, ) diff --git a/scripts/catmaid_skeleton_batch_import/src/catmaid_skeleton_batch_import/planner.py b/scripts/catmaid_skeleton_batch_import/src/catmaid_skeleton_batch_import/planner.py index da2398e2db..2ad0f728d8 100644 --- a/scripts/catmaid_skeleton_batch_import/src/catmaid_skeleton_batch_import/planner.py +++ b/scripts/catmaid_skeleton_batch_import/src/catmaid_skeleton_batch_import/planner.py @@ -22,6 +22,10 @@ PLAN_SCHEMA_VERSION = 1 +# CATMAID's default stable-join annotation. A neuron carrying it always wins a +# join, and two such neurons cannot be joined at all. +STABLE_JOIN_ANNOTATION = "stable" + @dataclass(frozen=True) class ImportPlan: @@ -44,6 +48,10 @@ def batch_count(self) -> int: def total_nodes(self) -> int: return sum(member.node_count for member in self.members) + @property + def total_annotations(self) -> int: + return sum(len(member.annotations) for member in self.members) + def header_dict(self) -> dict[str, Any]: return { "record_type": "header", @@ -85,6 +93,17 @@ def _validate_member(member: PlannedMember) -> None: raise InvalidInputError( f"Planned member {member.member_path!r} has an empty external ID" ) + for annotation in member.annotations: + if len(annotation) > 255: + raise InvalidInputError( + f"Planned member {member.member_path!r} renders annotation " + f"{annotation[:32]!r}... longer than 255 characters" + ) + if annotation == STABLE_JOIN_ANNOTATION: + raise InvalidInputError( + f"Planned member {member.member_path!r} renders annotation " + f"{annotation!r}, which CATMAID reserves for stable joins" + ) def build_plan( @@ -101,6 +120,7 @@ def build_plan( seen_paths: set[str] = set() seen_external_ids: set[str] = set() + annotation_members: dict[str, str] = {} for member in ordered: _validate_member(member) if member.member_path in seen_paths: @@ -115,6 +135,13 @@ def build_plan( f"{member.external_skeleton_id!r}" ) seen_external_ids.add(member.external_skeleton_id) + for annotation in member.annotations: + if annotation in annotation_members: + raise InvalidInputError( + f"Annotation {annotation!r} is rendered by both " + f"{annotation_members[annotation]!r} and {member.member_path!r}" + ) + annotation_members[annotation] = member.member_path if member.node_count > max_nodes_per_batch: raise InvalidInputError( f"SWC member {member.member_path!r} has {member.node_count} nodes; " @@ -226,7 +253,7 @@ def _member_from_record( "cable_length_nm", "batch_index", } - allowed = required | {"external_skeleton_id"} + allowed = required | {"external_skeleton_id", "annotations"} unknown = set(record) - allowed missing = required - set(record) if unknown or missing or record.get("record_type") != "member": @@ -249,6 +276,13 @@ def _member_from_record( raise InvalidInputError( f"{path}:{line_number}: external_skeleton_id must be a string" ) + annotations = record.get("annotations", []) + if not isinstance(annotations, list) or not all( + isinstance(annotation, str) for annotation in annotations + ): + raise InvalidInputError( + f"{path}:{line_number}: annotations must be an array of strings" + ) member = PlannedMember( member_path=_record_string(record, "member_path", path, line_number), node_count=_record_int(record, "node_count", path, line_number), @@ -260,6 +294,7 @@ def _member_from_record( cable_length_nm=float(cable_length), external_skeleton_id=external_id, batch_index=_record_int(record, "batch_index", path, line_number), + annotations=tuple(annotations), ) _validate_member(member) if member.batch_index <= 0: diff --git a/scripts/catmaid_skeleton_batch_import/src/catmaid_skeleton_batch_import/project.py b/scripts/catmaid_skeleton_batch_import/src/catmaid_skeleton_batch_import/project.py index 974003898d..4587126636 100644 --- a/scripts/catmaid_skeleton_batch_import/src/catmaid_skeleton_batch_import/project.py +++ b/scripts/catmaid_skeleton_batch_import/src/catmaid_skeleton_batch_import/project.py @@ -167,15 +167,15 @@ def create_hidden_project( for permission in SYSTEM_PROJECT_PERMISSIONS: components["assign_perm"](permission, system_user, project) - class_map = components["get_class_to_id_map"]( - project.id, ("neuron", "skeleton") - ) + class_names = ("neuron", "skeleton", "annotation") + relation_names = ("model_of", "annotated_with") + class_map = components["get_class_to_id_map"](project.id, class_names) relation_map = components["get_relation_to_id_map"]( - project.id, ("model_of",) + project.id, relation_names ) missing = { - "classes": sorted({"neuron", "skeleton"} - set(class_map)), - "relations": sorted({"model_of"} - set(relation_map)), + "classes": sorted(set(class_names) - set(class_map)), + "relations": sorted(set(relation_names) - set(relation_map)), } if missing["classes"] or missing["relations"]: raise ProjectSetupError( @@ -192,5 +192,7 @@ def create_hidden_project( "neuron_class_id": int(class_map["neuron"]), "skeleton_class_id": int(class_map["skeleton"]), "model_of_relation_id": int(relation_map["model_of"]), + "annotation_class_id": int(class_map["annotation"]), + "annotated_with_relation_id": int(relation_map["annotated_with"]), "metadata": metadata, } diff --git a/scripts/catmaid_skeleton_batch_import/src/catmaid_skeleton_batch_import/verification.py b/scripts/catmaid_skeleton_batch_import/src/catmaid_skeleton_batch_import/verification.py index 58062e0c44..6e4e7c505f 100644 --- a/scripts/catmaid_skeleton_batch_import/src/catmaid_skeleton_batch_import/verification.py +++ b/scripts/catmaid_skeleton_batch_import/src/catmaid_skeleton_batch_import/verification.py @@ -215,6 +215,9 @@ def verify_project_aggregate_counts( model_of_relation_id: int, expected_skeleton_count: int, expected_node_count: int, + annotation_class_id: int | None, + annotated_with_relation_id: int | None, + expected_annotation_count: int, connection: Any | None = None, ) -> dict[str, int]: """Check whole-project row counts without loading project IDs into memory. @@ -235,10 +238,19 @@ def verify_project_aggregate_counts( for value, label in ( (expected_skeleton_count, "expected_skeleton_count"), (expected_node_count, "expected_node_count"), + (expected_annotation_count, "expected_annotation_count"), ): if isinstance(value, bool) or not isinstance(value, int) or value < 0: raise ValueError(f"{label} must be a non-negative integer") + if expected_annotation_count and ( + annotation_class_id is None or annotated_with_relation_id is None + ): + raise ValueError( + "annotation identifiers are required when annotations are expected" + ) + # Project contexts written before annotation support carry no annotation + # identifiers. Comparing against NULL then counts no rows. connection = _default_connection(connection) with connection.cursor() as cursor: cursor.execute( @@ -248,6 +260,10 @@ def verify_project_aggregate_counts( WHERE project_id = %s AND class_id = %s), (SELECT count(*) FROM class_instance WHERE project_id = %s AND class_id = %s), + (SELECT count(*) FROM class_instance + WHERE project_id = %s AND class_id = %s), + (SELECT count(*) FROM class_instance_class_instance + WHERE project_id = %s AND relation_id = %s), (SELECT count(*) FROM class_instance_class_instance WHERE project_id = %s AND relation_id = %s), (SELECT count(*) FROM treenode WHERE project_id = %s), @@ -260,28 +276,36 @@ def verify_project_aggregate_counts( project_id, skeleton_class_id, project_id, + annotation_class_id, + project_id, model_of_relation_id, project_id, + annotated_with_relation_id, + project_id, project_id, project_id, ), ) row = cursor.fetchone() - if row is None or len(row) != 6: + if row is None or len(row) != 8: _fail("could not read project-wide database counts") actual = { "neuron_count": int(row[0]), "skeleton_count": int(row[1]), - "relationship_count": int(row[2]), - "treenode_count": int(row[3]), - "edge_count": int(row[4]), - "summary_count": int(row[5]), + "annotation_count": int(row[2]), + "relationship_count": int(row[3]), + "annotation_relationship_count": int(row[4]), + "treenode_count": int(row[5]), + "edge_count": int(row[6]), + "summary_count": int(row[7]), } expected = { "neuron_count": expected_skeleton_count, "skeleton_count": expected_skeleton_count, + "annotation_count": expected_annotation_count, "relationship_count": expected_skeleton_count, + "annotation_relationship_count": expected_annotation_count, "treenode_count": expected_node_count, "edge_count": expected_node_count, "summary_count": expected_skeleton_count, @@ -293,7 +317,9 @@ def verify_project_aggregate_counts( actual=actual, ) return { - "class_instance_count": actual["neuron_count"] + actual["skeleton_count"], + "class_instance_count": actual["neuron_count"] + + actual["skeleton_count"] + + actual["annotation_count"], **actual, } @@ -307,6 +333,8 @@ def verify_database( model_of_relation_id: int, skeletons: Sequence[Mapping[str, Any]], location_ids: Sequence[int], + annotation_class_id: int | None, + annotated_with_relation_id: int | None, connection: Any | None = None, cable_abs_tolerance_nm: float = DEFAULT_CABLE_ABS_TOLERANCE_NM, cable_rel_tolerance: float = DEFAULT_CABLE_REL_TOLERANCE, @@ -327,6 +355,8 @@ def verify_database( model_of_relation_id=model_of_relation_id, skeletons=skeletons, location_ids=location_ids, + annotation_class_id=annotation_class_id, + annotated_with_relation_id=annotated_with_relation_id, connection=connection, cable_abs_tolerance_nm=cable_abs_tolerance_nm, cable_rel_tolerance=cable_rel_tolerance, @@ -338,6 +368,11 @@ def verify_database( model_of_relation_id=model_of_relation_id, expected_skeleton_count=len(skeletons), expected_node_count=len(location_ids), + annotation_class_id=annotation_class_id, + annotated_with_relation_id=annotated_with_relation_id, + expected_annotation_count=sum( + len(skeleton["annotations"]) for skeleton in skeletons + ), connection=connection, ) if batch_counts != aggregate_counts: diff --git a/scripts/catmaid_skeleton_batch_import/tests/test_archive.py b/scripts/catmaid_skeleton_batch_import/tests/test_archive.py index fa0e50211c..48bdefef3b 100644 --- a/scripts/catmaid_skeleton_batch_import/tests/test_archive.py +++ b/scripts/catmaid_skeleton_batch_import/tests/test_archive.py @@ -18,7 +18,11 @@ VALID_SWC = b"1 1 0 0 0 -1 -1\n2 3 3 4 0 2 1\n" -def request(archive: Path, mapping: Path | None = None) -> IngestionRequest: +def request( + archive: Path, + mapping: Path | None = None, + annotations: tuple[str, ...] = (), +) -> IngestionRequest: return IngestionRequest( schema_version=1, archive_path=archive, @@ -26,6 +30,7 @@ def request(archive: Path, mapping: Path | None = None) -> IngestionRequest: dimension=(1, 1, 1), resolution_nm=(1.0, 1.0, 1.0), external_skeleton_id_map_path=mapping, + annotations=annotations, ) @@ -83,6 +88,21 @@ def test_scan_archive_hashes_parses_and_maps_exact_member_paths(tmp_path: Path) assert scanned.members[0].display_name == "a" +def test_stem_annotation_is_the_member_filename_stem_verbatim(tmp_path: Path) -> None: + archive = tmp_path / "191805.zip" + write_zip( + archive, + [("191805/0000001.swc", VALID_SWC), ("191805/Cell_07.SWC", VALID_SWC)], + ) + + scanned = scan_archive(request(archive, annotations=("{stem}",)), settings()) + + assert [member.annotations for member in scanned.members] == [ + ("0000001",), + ("Cell_07",), + ] + + def test_rejects_non_file_and_invalid_zip(tmp_path: Path) -> None: with pytest.raises(InvalidInputError, match="regular file"): scan_archive(request(tmp_path), settings()) diff --git a/scripts/catmaid_skeleton_batch_import/tests/test_artifacts.py b/scripts/catmaid_skeleton_batch_import/tests/test_artifacts.py index c34e69b256..87b4fb03ba 100644 --- a/scripts/catmaid_skeleton_batch_import/tests/test_artifacts.py +++ b/scripts/catmaid_skeleton_batch_import/tests/test_artifacts.py @@ -4,9 +4,13 @@ from catmaid_skeleton_batch_import.artifacts import ( artifact_descriptor, + manifest_rows, read_id_artifact, + read_manifest, write_id_artifact, + write_manifest, ) +from catmaid_skeleton_batch_import.domain import PlannedBatch, PlannedMember def test_reserved_id_artifact_is_deterministic_and_verified(tmp_path: Path) -> None: @@ -26,3 +30,33 @@ def test_reserved_id_artifact_is_deterministic_and_verified(tmp_path: Path) -> N maximum=max(ids), ) assert read_id_artifact(tmp_path, descriptor, expected_count=3) == ids + + +def test_manifest_lists_each_member_annotation_as_a_json_array(tmp_path: Path) -> None: + member = PlannedMember( + member_path="191805/0000001.swc", + node_count=2, + uncompressed_bytes=20, + crc32=7, + display_name="0000001", + cable_length_nm=2.0, + batch_index=1, + annotations=("0000001",), + ) + batch = PlannedBatch(index=1, members=(member,)) + path = tmp_path / "manifest.csv.gz" + + write_manifest(path, manifest_rows(batch, [100, 101, 102, 103, 104])) + + assert read_manifest(path) == [ + { + "member_path": "191805/0000001.swc", + "external_skeleton_id": "", + "neuron_id": "100", + "skeleton_id": "101", + "node_count": "2", + "name": "0000001", + "batch_index": "1", + "annotations": '["0000001"]', + } + ] diff --git a/scripts/catmaid_skeleton_batch_import/tests/test_catmaid_integration.py b/scripts/catmaid_skeleton_batch_import/tests/test_catmaid_integration.py index 8f7f00ad52..6e4673c3e3 100644 --- a/scripts/catmaid_skeleton_batch_import/tests/test_catmaid_integration.py +++ b/scripts/catmaid_skeleton_batch_import/tests/test_catmaid_integration.py @@ -139,8 +139,15 @@ def create(self, **kwargs): calls = [] components = { "get_system_user": lambda: system_user, - "get_class_to_id_map": lambda *_args: {"neuron": 31, "skeleton": 32}, - "get_relation_to_id_map": lambda *_args: {"model_of": 41}, + "get_class_to_id_map": lambda *_args: { + "neuron": 31, + "skeleton": 32, + "annotation": 33, + }, + "get_relation_to_id_map": lambda *_args: { + "model_of": 41, + "annotated_with": 42, + }, "validate_project_setup": lambda *args, **kwargs: calls.append( ("validate", args, kwargs) ), @@ -172,6 +179,8 @@ def create(self, **kwargs): assert created["project_stack"]["translation"] == (0.0, 0.0, 0.0) assert result["system_user_id"] == 7 assert result["neuron_class_id"] == 31 + assert result["annotation_class_id"] == 33 + assert result["annotated_with_relation_id"] == 42 assert sum(call[0] == "permission" for call in calls) == len( project.SYSTEM_PROJECT_PERMISSIONS ) @@ -208,6 +217,7 @@ def test_materialize_and_verify_batch_uses_catmaid_functions_and_exact_checks(): "name": "a", "node_count": 2, "cable_length_nm": 5.0, + "annotations": [], } ] @@ -217,6 +227,8 @@ def test_materialize_and_verify_batch_uses_catmaid_functions_and_exact_checks(): neuron_class_id=31, skeleton_class_id=32, model_of_relation_id=41, + annotation_class_id=33, + annotated_with_relation_id=42, skeletons=expected, location_ids=[20, 21], connection=connection, @@ -249,6 +261,7 @@ def test_materialization_requires_origin_and_rejects_singletons(): "name": "a", "node_count": 1, "cable_length_nm": 0, + "annotations": [], } ] with pytest.raises(ValueError, match="one-node"): @@ -258,6 +271,8 @@ def test_materialization_requires_origin_and_rejects_singletons(): neuron_class_id=31, skeleton_class_id=32, model_of_relation_id=41, + annotation_class_id=33, + annotated_with_relation_id=42, skeletons=expected, location_ids=[20], connection=ScriptedConnection(in_atomic_block=True), @@ -265,6 +280,58 @@ def test_materialization_requires_origin_and_rejects_singletons(): ) +def _annotated_batch_responses(): + return [ + [ + (10, 7, 1, 31, "a"), + (11, 7, 1, 32, "a"), + (13, 7, 1, 33, "0000001"), + ], + [(12, 7, 1, 41, 11, 10), (14, 7, 1, 42, 10, 13)], + [ + (20, 1, 7, 7, 11, None), + (21, 1, 7, 7, 11, 20), + ], + [(20, None, 1), (21, 20, 1)], + [(11, 1, 7, 2, 5.0)], + ] + + +def _annotated_expected_skeleton(): + return { + "neuron_id": 10, + "skeleton_id": 11, + "link_id": 12, + "name": "a", + "node_count": 2, + "cable_length_nm": 5.0, + "annotations": [{"name": "0000001", "annotation_id": 13, "link_id": 14}], + } + + +def test_batch_validation_counts_annotation_rows_and_their_neuron_links(): + connection = ScriptedConnection(_annotated_batch_responses()) + + result = materialization.validate_batch( + project_id=1, + user_id=7, + neuron_class_id=31, + skeleton_class_id=32, + model_of_relation_id=41, + annotation_class_id=33, + annotated_with_relation_id=42, + skeletons=[_annotated_expected_skeleton()], + location_ids=[20, 21], + connection=connection, + ) + + assert result["class_instance_count"] == 3 + assert result["annotation_count"] == 1 + assert result["annotation_relationship_count"] == 1 + assert connection.executions[0][1] == ([10, 11, 13],) + assert connection.executions[1][1] == ([12, 14],) + + def test_batch_validation_detects_edge_parent_mismatch(): responses = _valid_batch_responses()[2:] responses[3] = [(20, None, 1), (21, None, 1)] @@ -277,6 +344,7 @@ def test_batch_validation_detects_edge_parent_mismatch(): "name": "a", "node_count": 2, "cable_length_nm": 5.0, + "annotations": [], } ] with pytest.raises(VerificationError, match="edge parent"): @@ -286,6 +354,8 @@ def test_batch_validation_detects_edge_parent_mismatch(): neuron_class_id=31, skeleton_class_id=32, model_of_relation_id=41, + annotation_class_id=33, + annotated_with_relation_id=42, skeletons=expected, location_ids=[20, 21], connection=connection, @@ -451,7 +521,7 @@ def test_cache_verification_checks_exact_grid_and_no_dirty_cells(): def test_project_aggregate_counts_use_one_constant_size_result(): - connection = ScriptedConnection([(3, 3, 3, 10, 10, 3)]) + connection = ScriptedConnection([(3, 3, 3, 3, 3, 10, 10, 3)]) result = verification.verify_project_aggregate_counts( project_id=1, @@ -460,14 +530,19 @@ def test_project_aggregate_counts_use_one_constant_size_result(): model_of_relation_id=41, expected_skeleton_count=3, expected_node_count=10, + annotation_class_id=33, + annotated_with_relation_id=42, + expected_annotation_count=3, connection=connection, ) assert result == { - "class_instance_count": 6, + "class_instance_count": 9, "neuron_count": 3, "skeleton_count": 3, + "annotation_count": 3, "relationship_count": 3, + "annotation_relationship_count": 3, "treenode_count": 10, "edge_count": 10, "summary_count": 3, @@ -476,8 +551,25 @@ def test_project_aggregate_counts_use_one_constant_size_result(): assert "SELECT count(*) FROM treenode" in connection.executions[0][0] +def test_project_aggregate_counts_reject_a_missing_annotation_link(): + connection = ScriptedConnection([(3, 3, 3, 3, 2, 10, 10, 3)]) + with pytest.raises(VerificationError, match="project-wide database counts"): + verification.verify_project_aggregate_counts( + project_id=1, + neuron_class_id=31, + skeleton_class_id=32, + model_of_relation_id=41, + expected_skeleton_count=3, + expected_node_count=10, + annotation_class_id=33, + annotated_with_relation_id=42, + expected_annotation_count=3, + connection=connection, + ) + + def test_project_aggregate_counts_reject_extra_project_rows(): - connection = ScriptedConnection([(3, 3, 4, 10, 10, 3)]) + connection = ScriptedConnection([(3, 3, 0, 4, 0, 10, 10, 3)]) with pytest.raises(VerificationError, match="project-wide database counts"): verification.verify_project_aggregate_counts( project_id=1, @@ -486,5 +578,8 @@ def test_project_aggregate_counts_reject_extra_project_rows(): model_of_relation_id=41, expected_skeleton_count=3, expected_node_count=10, + annotation_class_id=33, + annotated_with_relation_id=42, + expected_annotation_count=0, connection=connection, ) diff --git a/scripts/catmaid_skeleton_batch_import/tests/test_config.py b/scripts/catmaid_skeleton_batch_import/tests/test_config.py index 864eca953d..4d6e5a7ee9 100644 --- a/scripts/catmaid_skeleton_batch_import/tests/test_config.py +++ b/scripts/catmaid_skeleton_batch_import/tests/test_config.py @@ -1,6 +1,7 @@ from __future__ import annotations import json +import re from pathlib import Path import pytest @@ -14,25 +15,25 @@ from catmaid_skeleton_batch_import.errors import InvalidInputError -def write_request(path: Path, **source_overrides: object) -> None: +def write_request( + path: Path, annotations: object = None, **source_overrides: object +) -> None: source = { "archive_path": str(path.parent / "input.zip"), "coordinate_unit": "um", } source.update(source_overrides) - path.write_text( - json.dumps( - { - "schema_version": 1, - "source": source, - "stack": { - "dimension": [10, 20, 30], - "resolution_nm": [4, 5.5, 6], - }, - } - ), - encoding="utf-8", - ) + request = { + "schema_version": 1, + "source": source, + "stack": { + "dimension": [10, 20, 30], + "resolution_nm": [4, 5.5, 6], + }, + } + if annotations is not None: + request["annotations"] = annotations + path.write_text(json.dumps(request), encoding="utf-8") def test_request_is_strict_and_normalized(tmp_path: Path) -> None: @@ -56,6 +57,44 @@ def test_request_rejects_unknown_input_policy(tmp_path: Path) -> None: load_request(request_path) +def test_request_accepts_stem_annotation_template(tmp_path: Path) -> None: + request_path = tmp_path / "request.json" + templates = ["{stem}", "swc:{stem}"] + write_request(request_path, annotations=templates) + + request = load_request(request_path) + + assert request.annotations == tuple(templates) + assert request.to_dict()["annotations"] == templates + assert request.annotation_names("191805/0000001.swc") == ( + "0000001", + "swc:0000001", + ) + + +@pytest.mark.parametrize( + ("annotations", "message"), + [ + (["imported"], "must contain {stem}"), + (["{stem}-{neuron_id}"], "may only use"), + (["{stem:>8}"], "may only use"), + (["{stem!r}"], "may only use"), + (["{stem"], "is malformed"), + (["{stem}", "{stem}"], "must not repeat"), + ([""], "empty template"), + ("{stem}", "array of strings"), + ], +) +def test_request_rejects_unusable_annotation_templates( + tmp_path: Path, annotations: object, message: str +) -> None: + request_path = tmp_path / "request.json" + write_request(request_path, annotations=annotations) + + with pytest.raises(InvalidInputError, match=re.escape(message)): + load_request(request_path) + + def test_settings_use_module_defaults(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.delenv("DJANGO_SETTINGS_MODULE", raising=False) for name in ( diff --git a/scripts/catmaid_skeleton_batch_import/tests/test_e2e_postgis.py b/scripts/catmaid_skeleton_batch_import/tests/test_e2e_postgis.py index 8171cf10a1..709347ecb7 100644 --- a/scripts/catmaid_skeleton_batch_import/tests/test_e2e_postgis.py +++ b/scripts/catmaid_skeleton_batch_import/tests/test_e2e_postgis.py @@ -48,6 +48,7 @@ def test_known_rollback_preserves_prior_batch_and_resumes( "dimension": [100, 100, 100], "resolution_nm": [1000, 1000, 1000], }, + "annotations": ["{stem}"], } ), encoding="utf-8", @@ -87,3 +88,28 @@ def fail_second_batch(**kwargs): "committed", "committed", ] + + from django.db import connection + + with connection.cursor() as cursor: + cursor.execute( + """ + SELECT neuron.name, annotation.name + FROM class_instance neuron + JOIN class neuron_class + ON neuron_class.id = neuron.class_id + AND neuron_class.class_name = 'neuron' + LEFT JOIN class_instance_class_instance link + ON link.class_instance_a = neuron.id + AND link.relation_id = ( + SELECT id FROM relation + WHERE project_id = neuron.project_id + AND relation_name = 'annotated_with' + ) + LEFT JOIN class_instance annotation ON annotation.id = link.class_instance_b + WHERE neuron.project_id = %s + ORDER BY neuron.name + """, + [result["project_id"]], + ) + assert cursor.fetchall() == [("alpha", "alpha"), ("beta", "beta")] diff --git a/scripts/catmaid_skeleton_batch_import/tests/test_loader.py b/scripts/catmaid_skeleton_batch_import/tests/test_loader.py index efe476d5a0..b87ee74e48 100644 --- a/scripts/catmaid_skeleton_batch_import/tests/test_loader.py +++ b/scripts/catmaid_skeleton_batch_import/tests/test_loader.py @@ -3,12 +3,15 @@ import csv from pathlib import Path +import pytest + from catmaid_skeleton_batch_import.domain import ( IngestionRequest, PlannedBatch, PlannedMember, SwcNode, ) +from catmaid_skeleton_batch_import.errors import InvalidInputError from catmaid_skeleton_batch_import.loader import ( ProjectContext, _treenode_row_matches, @@ -16,14 +19,23 @@ ) -def test_copy_rows_scale_coordinates_and_preserve_unknown_radius(tmp_path: Path) -> None: - request = IngestionRequest( +NODES = ( + SwcNode(1, 3, 1.0, 2.0, 3.0, -1.0, -1), + SwcNode(2, 3, 2.0, 2.0, 3.0, 2.0, 1), +) + + +def request(tmp_path: Path) -> IngestionRequest: + return IngestionRequest( schema_version=1, archive_path=tmp_path / "input.zip", coordinate_unit="um", dimension=(10, 10, 10), resolution_nm=(1000, 1000, 1000), ) + + +def batch(annotations: tuple[str, ...] = ()) -> PlannedBatch: member = PlannedMember( member_path="folder/cell.swc", node_count=2, @@ -32,9 +44,13 @@ def test_copy_rows_scale_coordinates_and_preserve_unknown_radius(tmp_path: Path) display_name="cell", cable_length_nm=1000.0, batch_index=1, + annotations=annotations, ) - batch = PlannedBatch(index=1, members=(member,)) - project = ProjectContext( + return PlannedBatch(index=1, members=(member,)) + + +def project(**annotation_ids: int) -> ProjectContext: + return ProjectContext( project_id=5, stack_id=6, project_stack_id=7, @@ -42,19 +58,18 @@ def test_copy_rows_scale_coordinates_and_preserve_unknown_radius(tmp_path: Path) neuron_class_id=8, skeleton_class_id=9, model_of_relation_id=10, - ) - nodes = ( - SwcNode(1, 3, 1.0, 2.0, 3.0, -1.0, -1), - SwcNode(2, 3, 2.0, 2.0, 3.0, 2.0, 1), + **annotation_ids, ) + +def test_copy_rows_scale_coordinates_and_preserve_unknown_radius(tmp_path: Path) -> None: with prepare_copy_rows( - batch, - request, - project, + batch(), + request(tmp_path), + project(), concept_ids=[100, 101, 102], location_ids=[200, 201], - load_nodes=lambda _member: nodes, + load_nodes=lambda _member: NODES, ) as prepared: prepared.treenode_file.seek(0) rows = list(csv.reader(prepared.treenode_file)) @@ -74,6 +89,80 @@ def test_copy_rows_scale_coordinates_and_preserve_unknown_radius(tmp_path: Path) assert rows[1][8:] == ["2000.0", "200"] +def test_stem_annotation_is_written_verbatim_for_each_member(tmp_path: Path) -> None: + with prepare_copy_rows( + batch(annotations=("0000001",)), + request(tmp_path), + project(annotation_class_id=11, annotated_with_relation_id=12), + concept_ids=[100, 101, 102, 103, 104], + location_ids=[200, 201], + load_nodes=lambda _member: NODES, + ) as prepared: + class_rows = list(csv.reader(prepared.class_file)) + relationship_rows = list(csv.reader(prepared.relationship_file)) + expected = prepared.expected_skeletons[0] + + assert class_rows == [ + ["100", "1", "5", "8", "cell"], + ["101", "1", "5", "9", "cell"], + ["103", "1", "5", "11", "0000001"], + ] + assert relationship_rows == [ + ["102", "1", "5", "10", "101", "100"], + ["104", "1", "5", "12", "100", "103"], + ] + assert expected.to_validation_dict()["annotations"] == [ + {"name": "0000001", "annotation_id": 103, "link_id": 104} + ] + + +def test_copy_rows_reject_a_concept_count_that_omits_annotation_ids( + tmp_path: Path, +) -> None: + with pytest.raises(InvalidInputError, match="expected 5 concept IDs, found 3"): + prepare_copy_rows( + batch(annotations=("0000001",)), + request(tmp_path), + project(annotation_class_id=11, annotated_with_relation_id=12), + concept_ids=[100, 101, 102], + location_ids=[200, 201], + load_nodes=lambda _member: NODES, + ) + + +def test_project_context_written_before_annotation_support_still_loads() -> None: + context = ProjectContext.from_dict( + { + "schema_version": 1, + "project_id": 5, + "stack_id": 6, + "project_stack_id": 7, + "user_id": 1, + "neuron_class_id": 8, + "skeleton_class_id": 9, + "model_of_relation_id": 10, + } + ) + + assert context.annotation_class_id is None + assert context.annotated_with_relation_id is None + assert "annotation_class_id" not in context.to_dict() + + +def test_annotations_are_refused_for_a_project_context_without_annotation_ids( + tmp_path: Path, +) -> None: + with pytest.raises(InvalidInputError, match="predates annotation support"): + prepare_copy_rows( + batch(annotations=("0000001",)), + request(tmp_path), + project(), + concept_ids=[100, 101, 102, 103, 104], + location_ids=[200, 201], + load_nodes=lambda _member: NODES, + ) + + def test_treenode_row_comparison_allows_only_driver_float_rounding() -> None: expected = (161, 4, 3033851.0, 4584272.0, 8379523.999999999, 2, 2, 140, 1000.0, 160) actual = (161, 4, 3033851.0, 4584272.0, 8379524.0, 2, 2, 140, 1000.0, 160) diff --git a/scripts/catmaid_skeleton_batch_import/tests/test_orchestrator.py b/scripts/catmaid_skeleton_batch_import/tests/test_orchestrator.py index 39b9a9a6ea..e9c1b91cd6 100644 --- a/scripts/catmaid_skeleton_batch_import/tests/test_orchestrator.py +++ b/scripts/catmaid_skeleton_batch_import/tests/test_orchestrator.py @@ -17,6 +17,7 @@ def _member(path: str, node_count: int) -> PlannedMember: crc32=node_count, display_name=Path(path).stem, cable_length_nm=float(node_count), + annotations=(Path(path).stem,), ) @@ -35,6 +36,8 @@ def test_final_database_verification_loads_then_validates_one_batch_at_a_time( neuron_class_id=31, skeleton_class_id=32, model_of_relation_id=41, + annotation_class_id=33, + annotated_with_relation_id=42, ) events = [] @@ -88,6 +91,7 @@ def aggregate_counts(**kwargs): "aggregate", kwargs["expected_skeleton_count"], kwargs["expected_node_count"], + kwargs["expected_annotation_count"], ) ) return aggregate @@ -110,5 +114,5 @@ def aggregate_counts(**kwargs): ("validate", 1, 3), ("load", 3, 2), ("validate", 1, 2), - ("aggregate", 3, 7), + ("aggregate", 3, 7, 3), ] diff --git a/scripts/catmaid_skeleton_batch_import/tests/test_planner.py b/scripts/catmaid_skeleton_batch_import/tests/test_planner.py index 8a9c10ad00..0ddfd688bd 100644 --- a/scripts/catmaid_skeleton_batch_import/tests/test_planner.py +++ b/scripts/catmaid_skeleton_batch_import/tests/test_planner.py @@ -16,7 +16,12 @@ ) -def member(path: str, nodes: int, external_id: str | None = None) -> PlannedMember: +def member( + path: str, + nodes: int, + external_id: str | None = None, + annotations: tuple[str, ...] = (), +) -> PlannedMember: return PlannedMember( member_path=path, node_count=nodes, @@ -25,9 +30,15 @@ def member(path: str, nodes: int, external_id: str | None = None) -> PlannedMemb display_name=Path(path).stem, cable_length_nm=float(nodes), external_skeleton_id=external_id, + annotations=annotations, ) +def write_gzip_text(path: Path, payload: str) -> None: + with gzip.open(path, "wt", encoding="utf-8") as target: + target.write(payload) + + def test_sorts_exact_paths_and_greedily_builds_hard_batches() -> None: plan = build_plan( [member("c.swc", 80), member("a.swc", 100), member("b.swc", 120)], @@ -53,6 +64,48 @@ def test_rejects_oversized_and_duplicate_members() -> None: ) +def test_same_annotation_in_two_folders_aborts_planning_and_names_both() -> None: + with pytest.raises( + InvalidInputError, + match=r"'0000001' is rendered by both 'a/0000001.swc' and 'b/0000001.swc'", + ): + build_plan( + [ + member("b/0000001.swc", 2, annotations=("0000001",)), + member("a/0000001.swc", 2, annotations=("0000001",)), + ], + 10, + ) + + +def test_stable_join_annotation_is_rejected() -> None: + with pytest.raises(InvalidInputError, match="reserves for stable joins"): + build_plan([member("stable.swc", 2, annotations=("stable",))], 10) + + +def test_annotation_longer_than_255_characters_is_rejected() -> None: + with pytest.raises(InvalidInputError, match="longer than 255 characters"): + build_plan([member("long.swc", 2, annotations=("x" * 300,))], 10) + + +def test_plans_written_before_annotation_support_still_load(tmp_path: Path) -> None: + path = tmp_path / "old-plan.jsonl.gz" + write_gzip_text( + path, + '{"batch_count":1,"max_nodes_per_batch":4,"member_count":1,' + '"record_type":"header","schema_version":1,"total_nodes":2}\n' + '{"batch_index":1,"cable_length_nm":2.0,"crc32":2,"display_name":"a",' + '"member_path":"a.swc","node_count":2,"record_type":"member",' + '"uncompressed_bytes":20}\n', + ) + + plan = read_plan(path) + + assert plan.members[0].annotations == () + assert plan.total_annotations == 0 + assert plan.batches[0].concept_count == 3 + + def test_gzip_jsonl_is_deterministic_and_round_trips(tmp_path: Path) -> None: plan = build_plan( [member("b.swc", 3, "external-b"), member("a.swc", 2, "external-a")],