From fae5540663bc2dd7adaf84e92141ce3a14cdaa30 Mon Sep 17 00:00:00 2001 From: uermel Date: Fri, 2 Oct 2026 15:43:10 -0700 Subject: [PATCH 1/4] fix: keep pick shifts, and default score and id for new picks Shifts. copick's particle centre is location + t, with t the transform's translation. Loading set ArtiaX's shift to zero and kept only the rotation, so a shifted pick was drawn at its location; saving wrote the rotation alone, so storing a user's pick set erased every t in it. Load now sets the ArtiaX shift from t (ArtiaX places a particle at origin + shift, the same model) and save writes it back. A particle dragged in the GUI has its shift reset by ArtiaX and is saved at its new location with t = 0. New picks. ArtiaX zero-fills every key, so a point placed in the GUI was saved with score 0. The particle data now starts new particles at score 1.0 and instance 0, by wrapping that one object's new_particle rather than subclassing the format (ArtiaX restores sessions through the bundle's get_class, which this bundle lacks). A missing score loads as 1.0, not 0. Taking particles into the user's own set stores the open list first (it is redrawn from storage, which discarded unsaved edits), and for filament objects moves the incoming instance ids past the set's own, so two sets' "filament 1" stay apart. Filament objects are coloured by instance id. Filament detection reads copick's is_filament, or the metadata spec on copick < 1.28. The conversion lives in src/misc/pickops.py and is tested without ChimeraX (pytest tests); checked against ArtiaX in ChimeraX 1.10 headless. Co-Authored-By: Claude Opus 5.5 --- pyproject.toml | 1 + src/misc/pickops.py | 138 +++++++++++++++++++++++++++++++++++-- src/tool.py | 56 ++++++++++----- tests/test_pickops.py | 154 ++++++++++++++++++++++++++++++++++++++++++ 4 files changed, 328 insertions(+), 21 deletions(-) create mode 100644 tests/test_pickops.py diff --git a/pyproject.toml b/pyproject.toml index fbb99fc..1c24227 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -33,6 +33,7 @@ dev = [ "black", "ipython", "pre-commit", + "pytest", "ruff", ] testing = [ diff --git a/src/misc/pickops.py b/src/misc/pickops.py index 9e31f38..338cea4 100644 --- a/src/misc/pickops.py +++ b/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 + +NEW_PARTICLE_DEFAULTS = {"score": 1.0, "instance_id": 0} +"""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.""" + + +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) diff --git a/src/tool.py b/src/tool.py index 377da7d..87228bd 100644 --- a/src/tool.py +++ b/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 diff --git a/tests/test_pickops.py b/tests/test_pickops.py new file mode 100644 index 0000000..55ce790 --- /dev/null +++ b/tests/test_pickops.py @@ -0,0 +1,154 @@ +"""Pick <-> particle conversion, without ChimeraX: run with ``pytest tests``.""" + +import importlib.util +from pathlib import Path +from types import SimpleNamespace + +import numpy as np +import pytest +from copick.models import CopickLocation, CopickPoint, PickableObject + + +def _load_pickops(): + # Importing the bundle package pulls in ChimeraX; the helpers only need copick and numpy. + path = Path(__file__).resolve().parents[1] / "src" / "misc" / "pickops.py" + spec = importlib.util.spec_from_file_location("chimerax_copick_pickops", path) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +pickops = _load_pickops() + + +def _rotation_zyz(a, b, c): + """Rotation by a about Z, then b about the new Y, then c about the new Z (degrees).""" + + def axis(i, deg): + t = np.radians(deg) + c_, s_ = np.cos(t), np.sin(t) + if i == "z": + return np.array([[c_, -s_, 0], [s_, c_, 0], [0, 0, 1]]) + return np.array([[c_, 0, s_], [0, 1, 0], [-s_, 0, c_]]) + + return axis("z", a) @ axis("y", b) @ axis("z", c) + + +def _point(location, shift=(0.0, 0.0, 0.0), euler=(0.0, 0.0, 0.0), instance_id=0, score=1.0): + transform = np.eye(4) + transform[:3, :3] = _rotation_zyz(*euler) + transform[:3, 3] = shift + return CopickPoint( + location=CopickLocation(x=location[0], y=location[1], z=location[2]), + transformation_=transform.tolist(), + instance_id=instance_id, + score=score, + ) + + +def _translation(v): + m = np.eye(4) + m[:3, 3] = v + return m + + +def test_shifted_rotated_point_round_trips(): + point = _point((100.0, 200.0, 300.0), shift=(4.0, -6.0, 2.5), euler=(30, 60, -45), instance_id=7, score=0.25) + origin, shift, rotation = pickops.point_pose(point) + # ArtiaX hands the rotation back as a 3x4 Place matrix with no translation. + back = pickops.point_from_pose(origin, shift, np.hstack([rotation, np.zeros((3, 1))]), 7, 0.25) + assert back.location == point.location + assert np.allclose(back.transformation, point.transformation) + assert (back.instance_id, back.score) == (7, 0.25) + + +def test_artiax_transform_places_the_particle_at_location_plus_shift(): + point = _point((100.0, 200.0, 300.0), shift=(4.0, -6.0, 2.5), euler=(10, 20, 30)) + origin, shift, rotation = pickops.point_pose(point) + rot = np.eye(4) + rot[:3, :3] = rotation + # ArtiaX: full_transform = origin * translation * rotation + full = _translation(origin) @ _translation(shift) @ rot + expected = np.asarray(point.transformation).copy() + expected[:3, 3] += [100.0, 200.0, 300.0] + assert np.allclose(full, expected) + + +def test_missing_identity_takes_copick_defaults(): + point = _point((1.0, 2.0, 3.0)) + point.instance_id, point.score = None, None + assert pickops.point_identity(point) == (0, 1.0) + assert pickops.point_identity(_point((1.0, 2.0, 3.0), score=float("nan"))) == (0, 1.0) + built = pickops.point_from_pose((0, 0, 0), (0, 0, 0), np.eye(3), instance_id=None, score=None) + assert (built.instance_id, built.score) == (0, 1.0) + + +class _FakeParticleData: + """ArtiaX zero-fills every key of a new particle.""" + + def __init__(self): + self.particles = [] + + def new_particle(self): + particle = {"score": 0, "instance_id": 0, "shift_x": 0} + self.particles.append(particle) + return particle + + +def test_new_particles_start_at_score_one(): + data = pickops.with_new_particle_defaults(_FakeParticleData()) + assert pickops.with_new_particle_defaults(data) is data # wrapping twice is a no-op + particle = data.new_particle() + assert particle["score"] == 1.0 and particle["instance_id"] == 0 and particle["shift_x"] == 0 + # Loading overwrites the defaults with the point's own values. + loaded = data.new_particle() + loaded["score"], loaded["instance_id"] = 0.5, 3 + assert data.particles[1] == {"score": 0.5, "instance_id": 3, "shift_x": 0} + + +def test_is_filament_reads_the_property_or_the_metadata(): + assert pickops.is_filament(SimpleNamespace(is_filament=True)) + assert not pickops.is_filament(SimpleNamespace(is_filament=False, metadata={"copick": {"filament": {}}})) + # copick < 1.28: no property, the spec is in metadata + assert pickops.is_filament(SimpleNamespace(metadata={"copick": {"filament": {"polar": True}}})) + assert not pickops.is_filament(SimpleNamespace(metadata={"copick": {"filament": None}})) + assert not pickops.is_filament(SimpleNamespace(metadata={"copick": "something else"})) + assert not pickops.is_filament(SimpleNamespace(metadata=None)) + assert pickops.is_filament( + PickableObject(name="microtubule", is_particle=True, label=1, metadata={"copick": {"filament": {}}}), + ) + assert not pickops.is_filament(PickableObject(name="ribosome", is_particle=True, label=2)) + + +def test_merging_filaments_keeps_them_apart(): + existing = SimpleNamespace(points=[_point((0, 0, 0), instance_id=1), _point((0, 10, 0), instance_id=2)]) + incoming = SimpleNamespace( + points=[ + _point((0, 0, 0), instance_id=1), # already present: skipped + _point((50, 0, 0), instance_id=1), + _point((60, 0, 0), instance_id=1), + _point((70, 0, 0), instance_id=0), + _point((80, 0, 0), instance_id=3), + ], + ) + merged = pickops.append_no_duplicates(incoming, existing, offset_ids=True) + assert [p.instance_id for p in merged.points] == [1, 2, 3, 3, 0, 5] + assert [p.instance_id for p in incoming.points] == [1, 1, 1, 0, 3], "the source set is not modified" + + plain = SimpleNamespace(points=[_point((0, 0, 0), instance_id=1)]) + merged = pickops.append_no_duplicates(SimpleNamespace(points=[_point((5, 0, 0), instance_id=1)]), plain) + assert [p.instance_id for p in merged.points] == [1, 1] + + +def test_instance_colours(): + base = [255, 0, 0, 200] + colors = pickops.instance_id_colors([0, 1, 1, 2, 0], base) + assert colors.shape == (5, 4) and colors.dtype == np.uint8 + assert (colors[0] == base).all() and (colors[4] == base).all() + assert (colors[1] == colors[2]).all() and not (colors[1] == colors[3]).all() + assert (colors[:, 3] == 200).all() + + +@pytest.mark.parametrize("ids", [[], [0, 0]]) +def test_instance_colours_without_instances(ids): + assert (pickops.instance_id_colors(ids, [1, 2, 3, 255]) == [1, 2, 3, 255]).all() From bd0714020b2d3e7835a3accb094d91abca186a76 Mon Sep 17 00:00:00 2001 From: uermel Date: Fri, 2 Oct 2026 17:24:36 -0700 Subject: [PATCH 2/4] style: constant note as a comment, so the module has one docstring Co-Authored-By: Claude Opus 5.5 --- src/misc/pickops.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/misc/pickops.py b/src/misc/pickops.py index 338cea4..7010751 100644 --- a/src/misc/pickops.py +++ b/src/misc/pickops.py @@ -4,9 +4,9 @@ import numpy as np 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} -"""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.""" def point_pose(point: CopickPoint) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: From b2fe8e0e5bfa1c3a177972aa6dfaafa733708919 Mon Sep 17 00:00:00 2001 From: uermel Date: Mon, 5 Oct 2026 23:47:38 -0700 Subject: [PATCH 3/4] test: stub the new pick helpers in the v2.0 tool tests The v2.0 tests replace the bundle's modules with stubs; the tool now imports the pick identity, pose and colour helpers from misc.pickops, so the stub provides them too. Co-Authored-By: Claude Opus 5.5 --- tests/conftest.py | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/tests/conftest.py b/tests/conftest.py index 003d67e..7fe2fd4 100644 --- a/tests/conftest.py +++ b/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( From 0dd26b4cb014473c3c23a194c7286a472a523492 Mon Sep 17 00:00:00 2001 From: uermel Date: Tue, 6 Oct 2026 01:37:50 -0700 Subject: [PATCH 4/4] ci: install copick for the portable tests The pick tests (tests/test_pickops.py) exercise misc/pickops.py, which works on copick's pick models, so the portable test environment needs copick. Co-Authored-By: Claude Opus 5.5 --- .github/workflows/tests.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index b9eaa10..9f7ef97 100644 --- a/.github/workflows/tests.yml +++ b/.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