fix: keep pick shifts, and default score and id for new picks (v2.0)

.github/workflows/tests.yml

@@ -39,7 +39,7 @@ jobs:
       - name: Run portable tests
         run: >
           uv run --python 3.14 --no-project --with pytest==9.1.1 --with "numpy>=2.0.2"
-          --with "zarr==${{ matrix.zarr-version }}"
+          --with "zarr==${{ matrix.zarr-version }}" --with "copick>=2.0.0a3,<3"
           python -m pytest -q -m "not chimerax and not integration"
 
       - name: Check migration Python

pyproject.toml

@@ -33,6 +33,7 @@ dev = [
     "black",
     "ipython",
     "pre-commit",
+    "pytest",
     "ruff",
 ]
 testing = [

src/misc/pickops.py

@@ -1,8 +1,125 @@
+import colorsys
+from typing import Iterable, List, Sequence, Tuple
+
 import numpy as np
-from copick.models import CopickPicks
+from copick.models import CopickLocation, CopickPicks, CopickPoint
+
+# Values a particle placed in the GUI starts with. ArtiaX fills every key with 0, and a pick saved with score 0
+# looks rejected to every tool that filters on score.
+NEW_PARTICLE_DEFAULTS = {"score": 1.0, "instance_id": 0}
+
+
+def point_pose(point: CopickPoint) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
+    """Split a copick point into ArtiaX's origin, shift and rotation.
+
+    copick's particle centre is ``location + t``, with ``t`` the transform's translation in the tomogram frame (Å).
+    ArtiaX places a particle at ``origin · translation · rotation``, the same model, so the location is the origin
+    and ``t`` is the shift.
+
+    Returns:
+        origin (3,), shift (3,) and rotation (3, 3).
+    """
+    transform = np.asarray(point.transformation, dtype=float)
+    origin = np.array([point.location.x, point.location.y, point.location.z], dtype=float)
+    return origin, transform[:3, 3].copy(), transform[:3, :3].copy()
+
+
+def point_identity(point: CopickPoint) -> Tuple[int, float]:
+    """A point's instance_id and score, with copick's defaults (0 and 1.0) where they are missing."""
+    return _instance_id(point.instance_id), _score(point.score)
+
+
+def point_from_pose(
+    origin: Sequence[float],
+    shift: Sequence[float],
+    rotation: np.ndarray,
+    instance_id=0,
+    score=1.0,
+) -> CopickPoint:
+    """Build a copick point from ArtiaX's origin, shift and rotation (3x3, or a 3x4 whose last column is ignored)."""
+    transform = np.eye(4)
+    transform[:3, :3] = np.asarray(rotation, dtype=float)[:3, :3]
+    transform[:3, 3] = np.asarray(shift, dtype=float)
+    return CopickPoint(
+        location=CopickLocation(x=float(origin[0]), y=float(origin[1]), z=float(origin[2])),
+        transformation_=transform.tolist(),
+        instance_id=_instance_id(instance_id),
+        score=_score(score),
+    )
+
+
+def with_new_particle_defaults(data):
+    """Make ``data.new_particle()`` start particles at :data:`NEW_PARTICLE_DEFAULTS`; returns ``data``.
+
+    Wraps the method on this one instance rather than subclassing the format's ParticleData: ArtiaX restores sessions
+    by resolving classes through the owning bundle, and this bundle has no ``get_class`` for its own.
+    """
+    if getattr(data.new_particle, "_copick_defaults", False):
+        return data
+    original = data.new_particle
+
+    def new_particle(*args, **kwargs):
+        particle = original(*args, **kwargs)
+        for key, value in NEW_PARTICLE_DEFAULTS.items():
+            particle[key] = value
+        return particle
+
+    new_particle._copick_defaults = True
+    data.new_particle = new_particle
+    return data
 
 
-def append_no_duplicates(inp: CopickPicks, out: CopickPicks) -> CopickPicks:
+def is_filament(obj) -> bool:
+    """Whether a pickable object is declared a filament.
+
+    copick >= 1.28 has ``is_filament``; older copick still carries the spec in ``metadata["copick"]["filament"]``.
+    """
+    flag = getattr(obj, "is_filament", None)
+    if isinstance(flag, bool):
+        return flag
+    metadata = getattr(obj, "metadata", None)
+    namespace = metadata.get("copick") if isinstance(metadata, dict) else None
+    return isinstance(namespace, dict) and namespace.get("filament") is not None
+
+
+def offset_instance_ids(points: Iterable[CopickPoint], existing: Iterable[CopickPoint]) -> List[CopickPoint]:
+    """Copies of ``points`` whose instance ids are moved past the largest id in ``existing``.
+
+    Two pick sets both number their filaments from 1, so merging one into the other unchanged would join unrelated
+    filaments. Unassigned points (0) stay 0.
+    """
+    start = max((_instance_id(p.instance_id) for p in existing), default=0)
+    moved = []
+    for point in points:
+        copy = point.model_copy(deep=True)
+        instance_id = _instance_id(copy.instance_id)
+        copy.instance_id = instance_id + start if instance_id > 0 else 0
+        moved.append(copy)
+    return moved
+
+
+def instance_id_colors(instance_ids: Sequence[int], base_rgba: Sequence[int]) -> np.ndarray:
+    """One RGBA colour per particle: each instance its own hue, unassigned (0) particles the object's colour.
+
+    Hues step by the golden ratio so neighbouring ids stay distinguishable.
+
+    Returns:
+        (N, 4) uint8, in the order of ``instance_ids``.
+    """
+    ids = np.asarray([_instance_id(i) for i in instance_ids], dtype=np.int64)
+    colors = np.tile(np.asarray(base_rgba, dtype=np.uint8)[:4], (len(ids), 1))
+    for value in np.unique(ids[ids > 0]):
+        rgb = colorsys.hsv_to_rgb((value * 0.618033988749895) % 1.0, 0.65, 0.95)
+        colors[ids == value, :3] = np.round(np.array(rgb) * 255).astype(np.uint8)
+    return colors
+
+
+def append_no_duplicates(inp: CopickPicks, out: CopickPicks, offset_ids: bool = False) -> CopickPicks:
+    """Append the points of ``inp`` that ``out`` does not already have at the same location.
+
+    With ``offset_ids``, the appended points' instance ids are moved past those already in ``out``
+    (:func:`offset_instance_ids`), which is what merging two sets of filaments needs.
+    """
     # Special cases
     if out.points is None:
         out.points = inp.points if inp.points is not None else []
@@ -29,8 +146,19 @@ def append_no_duplicates(inp: CopickPicks, out: CopickPicks) -> CopickPicks:
         out_arr[idx, :] = [pt.location.x, pt.location.y, pt.location.z]
 
     # If not existing in out, append it
-    for idx, pt in enumerate(inp_arr):
-        if not np.any(np.all(np.isclose(pt, out_arr), axis=1)):
-            out.points.append(inp.points[idx])
+    new_points = [
+        inp.points[idx] for idx, pt in enumerate(inp_arr) if not np.any(np.all(np.isclose(pt, out_arr), axis=1))
+    ]
+    if offset_ids:
+        new_points = offset_instance_ids(new_points, out.points)
+    out.points.extend(new_points)
 
     return out
+
+
+def _instance_id(value) -> int:
+    return 0 if value is None else int(value)
+
+
+def _score(value) -> float:
+    return 1.0 if value is None or not np.isfinite(float(value)) else float(value)

src/tool.py

@@ -25,7 +25,7 @@
 from chimerax.ome_zarr.open import open_ome_zarr_from_store
 from chimerax.ui import MainToolWindow
 from copick.impl.filesystem import CopickTomogramFSSpec
-from copick.models import CopickLocation, CopickMesh, CopickPicks, CopickPoint, CopickSegmentation
+from copick.models import CopickMesh, CopickPicks, CopickSegmentation
 from copick.util.uri import serialize_copick_uri
 from copick_shared_ui.core.thumbnail_cache import set_global_cache_config, set_global_cache_image_interface
 
@@ -36,7 +36,15 @@
 
 from .misc.colorops import palette_from_root
 from .misc.meshops import ensure_mesh
-from .misc.pickops import append_no_duplicates
+from .misc.pickops import (
+    append_no_duplicates,
+    instance_id_colors,
+    is_filament,
+    point_from_pose,
+    point_identity,
+    point_pose,
+    with_new_particle_defaults,
+)
 from .misc.settings import CoPickSettings
 from .misc.spotlight import SpotlightManager
 from .misc.tomostate import apply_view_state, capture_view_state
@@ -268,17 +276,17 @@ def store(self):
             if pick.from_tool or pick.read_only:
                 continue
 
-            points = []
-            for _id, p in pl.data:
-                rotmat = np.eye(4)
-                rotmat[0:3, :] = p.rotation.matrix
-                point = CopickPoint(
-                    location=CopickLocation(x=p["location_x"], y=p["location_y"], z=p["location_z"]),
-                    transformation_=rotmat.tolist(),
+            # The shift is the transform's translation: copick's centre is location + t, ArtiaX's origin + shift.
+            points = [
+                point_from_pose(
+                    origin=p.origin_coord,
+                    shift=p.translation.translation(),
+                    rotation=p.rotation.matrix,
                     instance_id=p["instance_id"],
                     score=p["score"],
                 )
-                points.append(point)
+                for _id, p in pl.data
+            ]
 
             pick.points = points
             pick.store()
@@ -706,16 +714,20 @@ def show_particles_from_picks(self, picks: CopickPicks):
         pick_obj = root.get_object(name)
 
         data = formats["Copick Picks file"].particle_data(self.session, file_name=None, oripix=1, trapix=1)
+        # Particles placed in the GUI start at score 1 and instance 0, not ArtiaX's zero fill.
+        data = with_new_particle_defaults(data)
 
         # Fill the ParticleData before constructing the ParticleList.
         points = picks.points if picks.points is not None else []
         for p in points:
+            origin, shift, rotation = point_pose(p)
+            instance_id, score = point_identity(p)
             part = data.new_particle()
-            part.origin = translation((p.location.x, p.location.y, p.location.z))
-            part.translation = translation((0, 0, 0))
-            part.rotation = Place(matrix=p.transformation[0:3, :])
-            part["score"] = float(p.score) if p.score is not None else 0.0
-            part["instance_id"] = int(p.instance_id) if p.instance_id is not None else 0
+            part.origin = translation(tuple(origin))
+            part.translation = translation(tuple(shift))
+            part.rotation = Place(matrix=np.hstack([rotation, np.zeros((3, 1))]))
+            part["score"] = score
+            part["instance_id"] = instance_id
 
         partlist = ParticleList(name, self.session, data)
         self.picks_map[picks] = partlist
@@ -741,6 +753,13 @@ def show_particles_from_picks(self, picks: CopickPicks):
         if pick_obj.radius is not None:
             partlist.radius = pick_obj.radius
 
+        # Filaments: one colour per filament, so neighbouring filaments can be told apart.
+        if pick_obj is not None and is_filament(pick_obj) and partlist.size > 0:
+            partlist.particle_colors = instance_id_colors(
+                [p["instance_id"] for _id, p in partlist.data],
+                pick_obj.color,
+            )
+
         if volume is not None:
             reg = volume.region
             reg = (reg[0], reg[1], (1, 1, 1))
@@ -1057,7 +1076,12 @@ def take_particles(self, index: QModelIndex):
         cur_picks = req_run.get_picks(user_id=user_id, object_name=req_name)
         if len(cur_picks) > 0:
             np = cur_picks[0]
-            np = append_no_duplicates(item.entity, np)
+            # The list is redrawn from storage below, so unsaved edits in it are stored first.
+            if np in self.picks_map:
+                self.store()
+            # Both sets number their filaments from 1; offsetting keeps unrelated filaments apart.
+            req_obj = self.root.get_object(req_name)
+            np = append_no_duplicates(item.entity, np, offset_ids=req_obj is not None and is_filament(req_obj))
         else:
             np = req_run.new_picks(user_id=user_id, object_name=req_name, session_id="19")
             np.meta.trust_orientation = item.entity.trust_orientation

tests/conftest.py

@@ -120,7 +120,17 @@ def tool_module(monkeypatch):
     bundle_package.__path__ = [str(SOURCE_ROOT)]
     _module(monkeypatch, "portable_bundle.misc.colorops", palette_from_root=lambda _root: "")
     _module(monkeypatch, "portable_bundle.misc.meshops", ensure_mesh=lambda mesh: mesh)
-    _module(monkeypatch, "portable_bundle.misc.pickops", append_no_duplicates=lambda left, _right: left)
+    _module(
+        monkeypatch,
+        "portable_bundle.misc.pickops",
+        append_no_duplicates=lambda left, _right, **_kwargs: left,
+        instance_id_colors=lambda ids, base: [tuple(base)] * len(list(ids)),
+        is_filament=lambda _obj: False,
+        point_from_pose=lambda **kwargs: kwargs,
+        point_identity=lambda point: (0, 1.0),
+        point_pose=lambda point: (point.location, (0.0, 0.0, 0.0), None),
+        with_new_particle_defaults=lambda data: data,
+    )
     _module(monkeypatch, "portable_bundle.misc.settings", CoPickSettings=Placeholder)
     _module(monkeypatch, "portable_bundle.misc.spotlight", SpotlightManager=SpotlightPlaceholder)
     _module(