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 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..7010751 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 + +# 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) 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/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( 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()