diff --git a/src/mosaic/animation/animations.py b/src/mosaic/animation/animations.py index 31d3263b..02c12b93 100644 --- a/src/mosaic/animation/animations.py +++ b/src/mosaic/animation/animations.py @@ -1,4 +1,5 @@ from abc import ABC, abstractmethod +from dataclasses import dataclass from enum import Enum from typing import Any @@ -11,6 +12,10 @@ class BaseAnimation(ABC): """Base class for all animations""" + #: Whether playback speed is meaningful, i.e. duration derives from a + #: reference length divided by rate. Span-like animations disable this. + has_rate = True + def __init__( self, vtk_widget, @@ -70,6 +75,32 @@ def update_parameters(self, **kwargs) -> None: self.stop_frame = max(1, int(self._base_duration / self.rate)) self.parameters.update(**kwargs) + def _reference_duration(self) -> float: + if self._base_duration is None: + self._base_duration = self.duration * self.rate + return self._base_duration + + def set_duration(self, frames: int) -> None: + """Set the animation length in frames. + + For animations with a rate, the rate is adjusted so that the + reference duration plays back over the requested number of frames. + """ + frames = max(1, int(frames)) + if self.has_rate: + self.rate = self._reference_duration() / frames + self.stop_frame = self.start_frame + frames + return None + + def set_rate(self, rate: float) -> None: + """Set the playback rate, deriving the duration from the reference length.""" + if rate <= 0 or not self.has_rate: + return None + frames = max(1, int(self._reference_duration() / rate)) + self.rate = rate + self.stop_frame = self.start_frame + frames + return None + def reset(self) -> None: """Reset any externally-modified state before a new frame is applied. @@ -79,6 +110,15 @@ def reset(self) -> None: """ pass + def commit(self) -> None: + """Apply state accumulated across all tracks after a frame's updates. + + Called by the composer once every track has been updated, so that + animations sharing targets can combine their contributions and write + to VTK once. The base implementation is a no-op. + """ + return None + def update(self, global_frame: int) -> None: """Update animation state for the given global frame""" if not self.enabled: @@ -137,7 +177,15 @@ def _ease(self, t: float) -> float: class TrajectoryAnimation(BaseAnimation): - """Animation for molecular trajectories""" + """Animation for molecular trajectories. + + Before its first track starts, a trajectory shows that track's first + frame; after a track ends, it holds the last one. When several tracks + animate the same trajectory, the most recently started track drives it. + """ + + # Keyed by trajectory: (rank, local frame, track) of the driving track. + _claims: dict[int, tuple[tuple[bool, int], int, "TrajectoryAnimation"]] = {} def _available_trajectories(self): """Return (name, uuid) pairs for all trajectory objects.""" @@ -213,8 +261,46 @@ def get_settings(self) -> list[dict[str, Any]]: }, ] + def reset(self) -> None: + """Withdraw this track's claim on its trajectory; applied on commit.""" + for key, claim in list(self._claims.items()): + if claim[2] is self: + del self._claims[key] + return None + + def update(self, global_frame: int) -> None: + """Bid for the trajectory, also before the track starts.""" + if not self.enabled or getattr(self, "_trajectory", None) is None: + return None + + local_frame = global_frame - self.global_start_frame + self.start_frame + started = local_frame >= self.start_frame + local_frame = max(self.start_frame, min(self.stop_frame, local_frame)) + local_frame -= (local_frame - self.start_frame) % self.stride + + # Started tracks beat pending ones; among started tracks the latest + # start wins, among pending ones the earliest. Ties go to later tracks. + start = self.global_start_frame + rank = (started, start if started else -start) + + key = id(self._trajectory) + claim = self._claims.get(key) + if claim is None or rank >= claim[0]: + self._claims[key] = (rank, local_frame, self) + return None + + def commit(self) -> None: + trajectory = getattr(self, "_trajectory", None) + if trajectory is None: + return None + + claim = self._claims.get(id(trajectory)) + if claim is not None and claim[2] is self: + self._update(claim[1]) + return None + def _update(self, frame: int) -> None: - if not hasattr(self, "_trajectory"): + if getattr(self, "_trajectory", None) is None: return None progress = self._get_progress(frame) @@ -226,16 +312,33 @@ def _update(self, frame: int) -> None: data_frame = int(data_start + progress * (data_stop - data_start)) data_frame = max(data_start, min(data_stop - 1, data_frame)) + if self._trajectory.current_frame == data_frame: + return None + self._trajectory.display_frame(data_frame) - # TODO: selection restore should move to reset() once the interactor - # preserves highlight state across display_frame calls. - uuids = self.cdata.models._get_selected_uuids() - if uuids: - self.cdata.models.set_selection_by_uuid(uuids) + # display_frame drops the highlight color. Restore it via the container, + # since the interactor would render mid-frame, before later tracks + # have updated, and the scene would flicker between frames. + models = self.cdata.models + uuids = models._get_selected_uuids() + if self._trajectory.uuid in uuids: + models.container.highlight(uuids) class VolumeAnimation(BaseAnimation): - """Volume slicing animation""" + """Volume slicing animation. + + Volume tracks hold their last slice after they end, so several can be + active at once. Each frame, the most recently started track drives the + viewer; the others yield instead of rewriting slice and orientation. + """ + + # Keyed by viewer: (global start frame, local frame, track) of the + # track that drives the viewer this frame. + _claims: dict[int, tuple[int, int, "VolumeAnimation"]] = {} + + # Keyed by viewer: slice visibility before any volume track touched it. + _original_visibility: dict[int, bool] = {} def _init_parameters(self) -> None: self._forward_fill = True @@ -244,7 +347,6 @@ def _init_parameters(self) -> None: self.parameters["projection"] = "Off" self.parameters["axis"] = "Z" self.parameters["hide"] = False - self._original_visibility = None try: self.update_parameters( axis=self.volume_viewer.primary.orientation_selector.currentText().lower() @@ -252,14 +354,30 @@ def _init_parameters(self) -> None: except Exception: pass + @property + def _viewer_key(self) -> int: + return id(self.volume_viewer.primary) + def reset(self) -> None: - """Restore volume visibility to the state before this animation touched it. + """Withdraw this track's claim on the viewer; applied on commit.""" + claim = self._claims.get(self._viewer_key) + if claim is not None and claim[2] is self: + del self._claims[self._viewer_key] + return None + + def commit(self) -> None: + """Apply the winning track's state, or restore the viewer if none.""" + key = self._viewer_key + claim = self._claims.get(key) + if claim is None: + original = self._original_visibility.pop(key, None) + if original is not None: + self.volume_viewer.primary.slice.SetVisibility(original) + return None - Uses the VTK actor directly to avoid triggering intermediate renders - that would cause flickering during playback. - """ - if self._original_visibility is not None: - self.volume_viewer.primary.slice.SetVisibility(self._original_visibility) + if claim[2] is self: + self._apply(claim[1]) + return None def _get_axis_size(self, axis=None): _mapping = {"x": 0, "y": 1, "z": 2} @@ -343,9 +461,16 @@ def update_parameters(self, **kwargs): super().update_parameters(**kwargs) def _update(self, frame: int) -> None: + key = self._viewer_key + claim = self._claims.get(key) + if claim is None or self.global_start_frame >= claim[0]: + self._claims[key] = (self.global_start_frame, frame, self) + return None + + def _apply(self, frame: int) -> None: progress = self._get_progress(frame) if progress is None: - return + return None if self.parameters["direction"] == "backward": progress = 1.0 - progress @@ -357,12 +482,11 @@ def _update(self, frame: int) -> None: viewer = self.volume_viewer.primary - # Track original visibility so reset() can restore it - if self._original_visibility is None: - self._original_visibility = viewer.is_visible - - if self.parameters.get("hide", False): - viewer.slice.SetVisibility(False) + # Uses the VTK actor directly; set_visibility would render mid-frame + original = self._original_visibility.setdefault( + self._viewer_key, viewer.is_visible + ) + viewer.slice.SetVisibility(original and not self.parameters.get("hide", False)) # Suppress intermediate renders viewer._rendering_suspended = True @@ -480,20 +604,23 @@ def _update(self, frame: int) -> None: class ActorSelectionDialog(QDialog): - """Dialog for selecting actors using ContainerTreeWidget.""" + """Pick the objects a visibility track affects. + + Lists clusters and models with the same rows, grouping and section + headers as the object browser sidebar. + """ def __init__(self, cdata, current_selection=None, parent=None): + from qtpy.QtCore import Qt from qtpy.QtWidgets import ( - QFrame, - QGroupBox, + QAbstractItemView, QHBoxLayout, QLabel, - QPushButton, - QTreeWidget, + QToolButton, QVBoxLayout, ) - from mosaic.widgets import DialogFooter + from mosaic.widgets import DialogFooter, SearchWidget from mosaic.widgets.container_list import ( ContainerTreeWidget, StyledTreeWidgetItem, @@ -501,108 +628,188 @@ def __init__(self, cdata, current_selection=None, parent=None): super().__init__(parent) self.setWindowTitle("Select Objects") - self.resize(400, 500) self.setModal(True) + self.resize(360, 520) - self._cdata = cdata - self._trees = [] - self._tree_labels = [] - - layout = QVBoxLayout(self) current_selection = set(current_selection or []) + self._trees = {} - quick_group = QGroupBox("Quick Select") - quick_layout = QHBoxLayout(quick_group) - quick_layout.setContentsMargins(8, 8, 8, 8) - quick_layout.setSpacing(6) - - for label, callback in [ - ("All", self._select_all), - ("Clusters", lambda: self._select_by_type("data")), - ("Models", lambda: self._select_by_type("models")), - ]: - btn = QPushButton(label) - btn.clicked.connect(callback) - quick_layout.addWidget(btn) - - layout.addWidget(quick_group) + layout = QVBoxLayout(self) + layout.setContentsMargins(12, 12, 12, 8) + layout.setSpacing(8) + + search = SearchWidget(placeholder="Filter objects") + search.searchTextChanged.connect(self._filter) + layout.addWidget(search) + + quick = QHBoxLayout() + quick.setContentsMargins(0, 0, 0, 0) + quick.setSpacing(2) + for text, data_types in ( + ("All", ("data", "models")), + ("None", ()), + ("Clusters", ("data",)), + ("Models", ("models",)), + ): + button = QToolButton() + button.setText(text) + button.setToolButtonStyle(Qt.ToolButtonStyle.ToolButtonTextOnly) + button.clicked.connect(lambda _, x=data_types: self._select(x)) + quick.addWidget(button) + quick.addStretch(1) + + self._count = QLabel() + self._count.setStyleSheet( + f"color: {Colors.TEXT_MUTED}; font-size: {Typography.LABEL}px;" + ) + quick.addWidget(self._count) + layout.addLayout(quick) - for label, data_type, interactor in [ + for title, data_type, interactor in ( ("Clusters", "data", cdata.data), ("Models", "models", cdata.models), - ]: + ): objects = {obj.uuid: obj for _, obj in cdata.format_datalist(data_type)} if not objects: continue - header = QLabel(label) - header.setStyleSheet(f"font-weight: 500; font-size: {Typography.LABEL}px;") - layout.addWidget(header) - tree = ContainerTreeWidget() tree.tree_widget.setSelectionMode( - QTreeWidget.SelectionMode.ExtendedSelection + QAbstractItemView.SelectionMode.MultiSelection ) - self._trees.append(tree) - self._tree_labels.append(data_type) + tree.tree_widget.setDragDropMode(QAbstractItemView.DragDropMode.NoDragDrop) - state = interactor.data_list.to_state() uuid_to_item = {} for uuid, obj in objects.items(): - item = StyledTreeWidgetItem( - obj._meta.get("name"), - obj.visible, - {"object_id": id(obj), "data_type": data_type, **obj._meta}, - ) + item = StyledTreeWidgetItem(geometry=obj, visible=obj.visible) + item.metadata["object_id"] = id(obj) + if not obj.visible: + item.setToolTip(0, "Hidden. Stays hidden during playback.") uuid_to_item[uuid] = item + # Mirror the sidebar's order and groups + state = interactor.data_list.to_state() + placed = state.get_all_uuids() tree.apply_state(state, uuid_to_item) + for uuid, item in uuid_to_item.items(): + if uuid not in placed: + tree.tree_widget.addTopLevelItem(item) + tree.tree_widget.expandAll() for uuid, obj in objects.items(): - if uuid in uuid_to_item and id(obj) in current_selection: - uuid_to_item[uuid].setSelected(True) - layout.addWidget(tree) - - separator = QFrame() - separator.setFrameShape(QFrame.Shape.HLine) - separator.setStyleSheet(f"color: {Colors.BORDER_DARK};") - layout.addWidget(separator) - - layout.addWidget(DialogFooter(dialog=self, margin=(0, 10, 0, 0))) - - def _select_all(self): - """Select all items in all trees.""" - for tree in self._trees: - tree.tree_widget.selectAll() - - def _select_by_type(self, data_type: str): - """Select only items of a specific type (data or models).""" - for tree, label in zip(self._trees, self._tree_labels): - if label == data_type: - tree.tree_widget.selectAll() - else: - tree.tree_widget.clearSelection() + uuid_to_item[uuid].setSelected(id(obj) in current_selection) + tree.tree_widget.itemSelectionChanged.connect(self._update_count) + + header = QLabel(title) + header.setStyleSheet( + f""" + font-size: {Typography.LABEL}px; font-weight: 600; + color: {Colors.TEXT_MUTED}; + padding: 8px 0px 4px 0px; + background: transparent; + """ + ) + section = QVBoxLayout() + section.setContentsMargins(0, 0, 0, 0) + section.setSpacing(2) + section.addWidget(header) + section.addWidget(tree, 1) + layout.addLayout(section, 1) + self._trees[data_type] = tree + + layout.addWidget( + DialogFooter( + info_text="Hidden objects stay hidden.", + dialog=self, + margin=(0, 8, 0, 0), + ) + ) + self._update_count() + + def _items(self, data_types=None): + """Yield object rows, skipping group headers.""" + from mosaic.widgets.container_list import StyledTreeWidgetItem + + for data_type, tree in self._trees.items(): + if data_types is not None and data_type not in data_types: + continue + for item, _parent, _index in tree.traverse(): + if isinstance(item, StyledTreeWidgetItem): + yield item + + def _select(self, data_types): + """Select every listed object of the given types, deselect the rest.""" + for data_type, tree in self._trees.items(): + tree.tree_widget.blockSignals(True) + for item in self._items((data_type,)): + item.setSelected(data_type in data_types and not item.isHidden()) + tree.tree_widget.blockSignals(False) + self._update_count() + + def _filter(self, text): + from mosaic.widgets.container_list import GroupTreeWidgetItem + + text = text.strip().lower() + for item in self._items(): + item.setHidden(bool(text) and text not in item.text(0).lower()) + + for tree in self._trees.values(): + for item, _parent, _index in tree.traverse(): + if not isinstance(item, GroupTreeWidgetItem): + continue + children = [item.child(i) for i in range(item.childCount())] + item.setHidden(bool(text) and all(c.isHidden() for c in children)) + + def _update_count(self): + items = list(self._items()) + selected = sum(item.isSelected() for item in items) + self._count.setText(f"{selected} of {len(items)} selected") def get_selected_objects(self): - selected = [] - for tree in self._trees: - selected.extend( - item.metadata["object_id"] for item in tree.selected_items() - ) - return selected + return [ + item.metadata["object_id"] for item in self._items() if item.isSelected() + ] + + +@dataclass +class _FadeState: + """Original appearance of an actor and the track currently fading it.""" + + geometry: Any + actor: Any + opacity: float + #: (global start frame, opacity factor, track) of the driving track. + claim: tuple[int, float, Any] | None = None + hidden: bool = False class VisibilityAnimation(BaseAnimation): - """Visibility fade animation""" + """Hide or fade objects for the span of the track. + + Outside the track objects keep their own appearance. Inside, their + opacity is scaled by a factor that moves from ``start_opacity`` to + ``target_opacity`` over the track, or jumps to the target with instant + easing. Objects faded to zero are switched off rather than drawn fully + transparent, which keeps them out of depth peeling. + """ + + has_rate = False + + # Shared across tracks so overlapping visibility tracks resolve to one + # winner per actor, and the captured opacity is never a value another + # track wrote earlier. + _ledger: dict[int, _FadeState] = {} def _init_parameters(self) -> None: - self._forward_fill = True self.parameters.clear() self.parameters.update( - {"start_opacity": 1.0, "target_opacity": 0.0, "easing": "instant"} + { + "start_opacity": 1.0, + "target_opacity": 0.0, + "easing": "instant", + } ) - self._original_opacities = {} - self._cached_actors = None + self._cached_geometries = None def get_settings(self) -> list[dict[str, Any]]: return [ @@ -612,22 +819,25 @@ def get_settings(self) -> list[dict[str, Any]]: "min": 0.0, "max": 1.0, "default": self.parameters.get("start_opacity", 1.0), - "description": "Start opacity (0.0 for invisible, 1.0 for fully visible)", + "description": "Opacity at the start of the track, relative to " + "each object's own opacity (0 hides objects).", }, { "label": "target_opacity", "type": "float", "min": 0.0, "max": 1.0, - "default": self.parameters.get("target_opacity", 1.0), - "description": "Target opacity (0.0 for invisible, 1.0 for fully visible)", + "default": self.parameters.get("target_opacity", 0.0), + "description": "Opacity at the end of the track, relative to " + "each object's own opacity (0 hides objects).", }, { "label": "easing", "type": "select", - "options": ["linear", "ease-in", "ease-out", "ease-in-out", "instant"], + "options": ["instant", "linear", "ease-in", "ease-out", "ease-in-out"], "default": self.parameters.get("easing", "instant"), - "description": "Animation style (instant for immediate change)", + "description": "Instant holds the target for the whole track; " + "other easings fade from start to target across the track.", }, { "label": "Objects", @@ -646,55 +856,98 @@ def _open_object_selection_dialog(self, _checked=None): cdata=self.cdata, current_selection=current_selection ) if dialog.exec(): - selected_objects = dialog.get_selected_objects() - self._cached_actors = None - self.update_parameters(selected_objects=selected_objects) - + self.update_parameters(selected_objects=dialog.get_selected_objects()) except Exception: pass return False - def _get_actors(self): - if self._cached_actors is not None: - return self._cached_actors + def update_parameters(self, **kwargs) -> None: + if "selected_objects" in kwargs: + self._cached_geometries = None + return super().update_parameters(**kwargs) + + def _get_geometries(self): + if self._cached_geometries is not None: + return self._cached_geometries - object_ids = self.parameters.get("selected_objects", []) + object_ids = set(self.parameters.get("selected_objects", [])) try: - all_objects = {} - for _name, obj in self.cdata.format_datalist("data"): - all_objects[id(obj)] = obj - for _name, obj in self.cdata.format_datalist("models"): - all_objects[id(obj)] = obj - - self._cached_actors = [ - all_objects[x].actor for x in object_ids if x in all_objects + self._cached_geometries = [ + obj + for data_type in ("data", "models") + for _name, obj in self.cdata.format_datalist(data_type) + if id(obj) in object_ids ] except Exception: return [] + return self._cached_geometries - return self._cached_actors + def _fade_factor(self, frame: int): + """Opacity multiplier at the given local frame, or None outside the span.""" + duration = self.duration + elapsed = frame - self.start_frame + if not 0 <= elapsed < duration: + return None + + progress = 1.0 + if self.parameters.get("easing", "instant") != "instant" and duration > 1: + progress = self._ease(elapsed / (duration - 1)) + + start = float(self.parameters.get("start_opacity", 1.0)) + target = float(self.parameters.get("target_opacity", 0.0)) + return start + (target - start) * progress def reset(self) -> None: - """Restore actors to the opacity they had before this animation touched them.""" - for actor in self._get_actors(): - original = self._original_opacities.get(id(actor)) - if original is not None: - actor.GetProperty().SetOpacity(original) + """Withdraw this track's claims; actors are restored on commit.""" + for state in self._ledger.values(): + if state.claim is not None and state.claim[2] is self: + state.claim = None + return None def _update(self, frame: int) -> None: - progress = self._get_progress(frame) - if progress is None: - return + if (factor := self._fade_factor(frame)) is None: + return None + + start = self.global_start_frame + for geometry in self._get_geometries(): + if not geometry.visible: + continue + actor = geometry.actor + state = self._ledger.get(id(actor)) + if state is None: + state = _FadeState(geometry, actor, actor.GetProperty().GetOpacity()) + self._ledger[id(actor)] = state + + # Most recently started track wins; ties go to the later track + if state.claim is None or start >= state.claim[0]: + state.claim = (start, factor, self) + return None + + def commit(self) -> None: + """Write the winning fade of each actor, restoring unclaimed ones. + + VTK setters skip unchanged values, so actors whose fade is constant + between frames are not marked modified. + """ + for key in list(self._ledger): + state = self._ledger[key] + factor = 1.0 if state.claim is None else state.claim[1] + + if factor <= 0: + state.actor.GetProperty().SetOpacity(state.opacity) + state.actor.SetVisibility(False) + state.hidden = True + continue - start_opacity = self.parameters["start_opacity"] - target_opacity = self.parameters["target_opacity"] - current_opacity = start_opacity + (target_opacity - start_opacity) * progress + state.actor.GetProperty().SetOpacity(state.opacity * factor) + if state.hidden: + state.actor.SetVisibility(state.geometry.visible) + state.hidden = False - for actor in self._get_actors(): - if id(actor) not in self._original_opacities: - self._original_opacities[id(actor)] = actor.GetProperty().GetOpacity() - actor.GetProperty().SetOpacity(current_opacity) + if state.claim is None: + del self._ledger[key] + return None class WaypointAnimation(BaseAnimation): @@ -713,7 +966,6 @@ def _init_parameters(self) -> None: "easing": "linear", } ) - self._path_actors = [] self._curve = None self._positions = None self._focal_points = None @@ -739,7 +991,6 @@ def update_parameters(self, **kwargs): if "spline_order" in kwargs: self.parameters["spline_order"] = int(kwargs.pop("spline_order")) self._init_spline() - self._update_path_preview() if "easing" in kwargs: self.parameters["easing"] = kwargs.pop("easing") @@ -756,7 +1007,6 @@ def _capture_waypoint(self): } self.parameters["waypoints"].append(waypoint) self._init_spline() - self._update_path_preview() def _remove_last_waypoint(self): """Remove the most recently added waypoint.""" @@ -764,13 +1014,11 @@ def _remove_last_waypoint(self): if waypoints: waypoints.pop() self._init_spline() - self._update_path_preview() def _clear_waypoints(self): """Remove all waypoints.""" self.parameters["waypoints"].clear() self._init_spline() - self.cleanup_preview() def _init_spline(self): """Initialize the spline curve from waypoints.""" @@ -808,110 +1056,6 @@ def _init_spline(self): [np.interp(t, t_wp, view_ups[:, i]) for i in range(3)] ) - def _update_path_preview(self): - """Visualize the camera path in the 3D viewport. - - Draws from the same sampled positions used by the animation so the - preview exactly matches what plays back. - """ - import vtk - - _, renderer = self._get_rendering_context(return_renderer=True) - - self.cleanup_preview() - - waypoints = self.parameters.get("waypoints", []) - - # Show point markers even for a single waypoint - if len(waypoints) < 2 or self._positions is None: - for wp in waypoints: - actor = self._create_waypoint_marker(wp["position"]) - renderer.AddActor(actor) - self._path_actors.append(actor) - self.vtk_widget.GetRenderWindow().Render() - return - - # Build a polyline from the sampled animation positions - vtk_points = vtk.vtkPoints() - for pos in self._positions: - vtk_points.InsertNextPoint(pos) - - lines = vtk.vtkCellArray() - n_pts = len(self._positions) - lines.InsertNextCell(n_pts) - for i in range(n_pts): - lines.InsertCellPoint(i) - - poly_data = vtk.vtkPolyData() - poly_data.SetPoints(vtk_points) - poly_data.SetLines(lines) - - tube_filter = vtk.vtkTubeFilter() - tube_filter.SetInputData(poly_data) - tube_filter.SetRadius(1.0) - tube_filter.SetNumberOfSides(8) - tube_filter.SetVaryRadiusToVaryRadiusOff() - tube_filter.Update() - - mapper = vtk.vtkPolyDataMapper() - mapper.SetInputConnection(tube_filter.GetOutputPort()) - mapper.SetResolveCoincidentTopologyToPolygonOffset() - mapper.SetScalarVisibility(False) - - path_actor = vtk.vtkActor() - path_actor.SetMapper(mapper) - # Pink #ec4899 → RGB - path_actor.GetProperty().SetColor(0.925, 0.286, 0.6) - path_actor.GetProperty().SetOpacity(0.7) - renderer.AddActor(path_actor) - self._path_actors.append(path_actor) - - for wp in waypoints: - actor = self._create_waypoint_marker(wp["position"]) - renderer.AddActor(actor) - self._path_actors.append(actor) - - self.vtk_widget.GetRenderWindow().Render() - - def _create_waypoint_marker(self, position): - """Create a small sphere actor at the given position.""" - import vtk - - point_data = vtk.vtkPoints() - point_data.InsertNextPoint(position) - - vertices = vtk.vtkCellArray() - vertices.InsertNextCell(1) - vertices.InsertCellPoint(0) - - poly_data = vtk.vtkPolyData() - poly_data.SetPoints(point_data) - poly_data.SetVerts(vertices) - - mapper = vtk.vtkPolyDataMapper() - mapper.SetInputData(poly_data) - mapper.SetResolveCoincidentTopologyToPolygonOffset() - - actor = vtk.vtkActor() - actor.SetMapper(mapper) - # Pink #ec4899 - actor.GetProperty().SetColor(0.925, 0.286, 0.6) - actor.GetProperty().SetPointSize(12) - actor.GetProperty().SetRenderPointsAsSpheres(True) - - return actor - - def cleanup_preview(self): - """Remove all preview actors from the renderer.""" - try: - _, renderer = self._get_rendering_context(return_renderer=True) - for actor in self._path_actors: - renderer.RemoveActor(actor) - self._path_actors.clear() - self.vtk_widget.GetRenderWindow().Render() - except Exception: - self._path_actors.clear() - def get_settings(self) -> list[dict[str, Any]]: n = len(self.parameters.get("waypoints", [])) count_text = f"{n} waypoint{'s' if n != 1 else ''} defined" diff --git a/src/mosaic/animation/compose.py b/src/mosaic/animation/compose.py index ec913fdb..a84fd61d 100644 --- a/src/mosaic/animation/compose.py +++ b/src/mosaic/animation/compose.py @@ -24,8 +24,10 @@ QHBoxLayout, QMessageBox, QPushButton, + QScrollArea, QSizePolicy, QSpinBox, + QSplitter, QToolButton, QVBoxLayout, QWidget, @@ -80,6 +82,7 @@ def __init__(self, vtk_widget, volume_viewer=None, cdata=None, parent=None): self.vtk_widget = vtk_widget self.volume_viewer = volume_viewer + self.setWindowTitle("Animation") self._setup_ui() self._setup_shortcuts() @@ -152,12 +155,18 @@ def _setup_ui(self): self.timeline.content.trackSelected.connect(self._on_track_selected) self.timeline.frameMoved.connect(self._frame_throttle) self.timeline.trackMoved.connect(self._on_track_moved) + self.timeline.trackResized.connect(self._on_track_resized) self.timeline.trackRemoved.connect(self.delete_track) - self.timeline.setMinimumHeight(120) + self.timeline.setMinimumHeight(220) self.timeline.setSizePolicy( QSizePolicy.Policy.Expanding, QSizePolicy.Policy.Expanding ) - main_layout.addWidget(self.timeline, 1) + + # Settings scroll below the timeline instead of squeezing it + settings = QWidget() + settings_layout = QVBoxLayout(settings) + settings_layout.setContentsMargins(0, 0, 0, 0) + settings_layout.setSpacing(8) presets_group = QGroupBox("Presets") presets_layout = QVBoxLayout(presets_group) @@ -181,7 +190,7 @@ def _setup_ui(self): btn.clicked.connect(lambda _, p=preset_type: self._load_preset(p)) presets_row.addWidget(btn) presets_layout.addLayout(presets_row) - main_layout.addWidget(presets_group) + settings_layout.addWidget(presets_group) # Add Animation group (2 rows x 3 cols for compactness) anims_group = QGroupBox("Add Animation") @@ -197,11 +206,35 @@ def _setup_ui(self): btn.setFixedHeight(Colors.WIDGET_HEIGHT) btn.clicked.connect(lambda _, t=anim_type: self.add_animation(t)) anims_layout.addWidget(btn, i // cols, i % cols) - main_layout.addWidget(anims_group) + settings_layout.addWidget(anims_group) self.properties_panel = AnimationSettings() self.properties_panel.animationChanged.connect(self._on_animation_changed) - main_layout.addWidget(self.properties_panel) + settings_layout.addWidget(self.properties_panel) + settings_layout.addStretch() + + settings_scroll = QScrollArea() + settings_scroll.setWidget(settings) + settings_scroll.setWidgetResizable(True) + settings_scroll.setHorizontalScrollBarPolicy( + Qt.ScrollBarPolicy.ScrollBarAlwaysOff + ) + settings_scroll.setStyleSheet( + "QScrollArea { border: none; background: transparent; }" + ) + settings.setObjectName("composerSettings") + settings.setStyleSheet("#composerSettings { background: transparent; }") + + splitter = QSplitter(Qt.Orientation.Vertical) + splitter.addWidget(self.timeline) + splitter.addWidget(settings_scroll) + splitter.setChildrenCollapsible(False) + splitter.setStretchFactor(0, 1) + splitter.setStretchFactor(1, 1) + splitter.setSizes([400, 400]) + splitter.setHandleWidth(6) + splitter.setStyleSheet("QSplitter::handle { background: transparent; }") + main_layout.addWidget(splitter, 1) self.setMinimumWidth(280) self.resize(300, 700) @@ -264,8 +297,7 @@ def _load_reveal_flythrough_preset(self): This creates a sequence that: 1. Slices through the volume (forward) with all actors hidden - 2. Reveals all actors at the midpoint - 3. Slices back through the volume (backward) + 2. Slices back through the volume (backward) with actors revealed """ if self.volume_viewer is None: MosaicMessageBox.warning( @@ -286,36 +318,18 @@ def _load_reveal_flythrough_preset(self): for _, obj in self.cdata.format_datalist("models"): all_objects.append(id(obj)) - hide_anim = self.add_animation(AnimationType.VISIBILITY) - if hide_anim: - hide_anim.name = "Hide Actors" - hide_anim.update_parameters( - start_opacity=1.0, - target_opacity=0.0, - easing="instant", - selected_objects=all_objects, - ) - hide_anim.start_frame = 0 - hide_anim.stop_frame = 1 - forward_slice = self.add_animation(AnimationType.SLICE) if forward_slice: forward_slice.name = "Slice Forward" forward_slice.update_parameters(direction="forward") - reveal_anim = self.add_animation(AnimationType.VISIBILITY) - if reveal_anim: - reveal_anim.name = "Reveal Actors" - reveal_anim.update_parameters( - start_opacity=0.0, - target_opacity=1.0, - easing="instant", - selected_objects=all_objects, - ) - reveal_anim.start_frame = 0 - reveal_anim.stop_frame = 1 + hide_anim = self.add_animation(AnimationType.VISIBILITY) + if hide_anim: + hide_anim.name = "Hide Actors" + hide_anim.update_parameters(selected_objects=all_objects) + hide_anim.global_start_frame = forward_slice.global_start_frame + hide_anim.set_duration(forward_slice.duration) - # 4. Volume slice backward backward_slice = self.add_animation(AnimationType.SLICE) if backward_slice: backward_slice.name = "Slice Backward" @@ -360,11 +374,12 @@ def _get_next_start_frame(self) -> int: ) def _cleanup_tracks(self, tracks=None): - """Reset animation state and remove preview actors for the given tracks.""" - for track in tracks or self.tracks: + """Restore state modified by the given tracks.""" + tracks = tracks or self.tracks + for track in tracks: track.animation.reset() - if hasattr(track.animation, "cleanup_preview"): - track.animation.cleanup_preview() + for track in tracks: + track.animation.commit() def _get_track(self, track_id: str): return next((t for t in self.tracks if t.id == track_id), None) @@ -392,9 +407,22 @@ def _on_track_moved(self, track_id: str, new_frame: int): track.animation.global_start_frame = new_frame self.properties_panel.global_start_spin.setValue(new_frame) self.timeline.update() + self._frame_throttle(self.current_frame) + + def _on_track_resized(self, track_id: str, duration: int): + if (track := self._get_track(track_id)) is None: + return None + + track.animation.set_duration(duration) + if self.selected_track == track_id: + self.properties_panel.sync_timing() + self.timeline.update() + self._frame_throttle(self.current_frame) + return None def _on_animation_changed(self, changes: dict[str, Any]): self.timeline.update() + self._frame_throttle(self.current_frame) def _get_total_frames(self) -> int: """Get the total number of frames across all tracks.""" @@ -431,30 +459,23 @@ def set_current_frame(self, frame: int): track.animation.reset() for track in self.tracks: track.animation.update(self.current_frame) + for track in self.tracks: + track.animation.commit() - self.vtk_widget.GetRenderWindow().Render() + self.cdata.viewport.render_vtk() def _go_to_end(self): """Jump to the last frame.""" self.set_current_frame(self._get_total_frames()) - def _set_preview_visibility(self, visible: bool): - """Show or hide preview actors (e.g. waypoint path) on all tracks.""" - for track in self.tracks: - for actor in getattr(track.animation, "_path_actors", []): - actor.SetVisibility(visible) - self.vtk_widget.GetRenderWindow().Render() - def toggle_play(self): """Toggle playback state.""" self.is_playing = not self.is_playing if self.is_playing: - self._set_preview_visibility(False) self.play_btn.setIcon(icon("ph.pause", role="muted")) self.timer.start(1000 // self.playback_fps) else: - self._set_preview_visibility(True) self.play_btn.setIcon(icon("ph.play", role="primary")) self.timer.stop() @@ -636,7 +657,6 @@ def export_animation(self): n_frames = len(frames) try: - self._set_preview_visibility(False) render_window.SetOffScreenRendering(1) render_window.SetMultiSamples(multisamples) render_window.SetAlphaBitPlanes(1 if transparent_bg else 0) @@ -687,7 +707,6 @@ def export_animation(self): finally: if indicator is not None: indicator.hide_progress() - self._set_preview_visibility(True) render_window.SetAlphaBitPlanes(original_alpha) render_window.SetMultiSamples(original_multisamples) render_window.SetOffScreenRendering(0) diff --git a/src/mosaic/animation/settings.py b/src/mosaic/animation/settings.py index 3cfff9f6..da17e754 100644 --- a/src/mosaic/animation/settings.py +++ b/src/mosaic/animation/settings.py @@ -71,7 +71,17 @@ def setup_ui(self): ) frame_layout.addWidget(self.global_start_spin, 0, 1) - frame_layout.addWidget(QLabel("Rate:"), 0, 2) + frame_layout.addWidget(QLabel("Duration:"), 0, 2) + self.duration_spin = QSpinBox() + self.duration_spin.setRange(1, 2 << 29) + self.duration_spin.setSuffix(" f") + self.duration_spin.setFixedHeight(Colors.WIDGET_HEIGHT) + self.duration_spin.setMinimumWidth(0) + self.duration_spin.valueChanged.connect(self._on_duration_changed) + frame_layout.addWidget(self.duration_spin, 0, 3) + + self._rate_label = QLabel("Rate:") + frame_layout.addWidget(self._rate_label, 1, 0) self.rate_spin = QDoubleSpinBox() self.rate_spin.setRange(0.1, 10.0) self.rate_spin.setValue(1.0) @@ -80,9 +90,7 @@ def setup_ui(self): self.rate_spin.setFixedHeight(Colors.WIDGET_HEIGHT) self.rate_spin.setMinimumWidth(0) self.rate_spin.valueChanged.connect(self._on_rate_changed) - frame_layout.addWidget(self.rate_spin, 0, 3) - - self._base_duration = 100 # Reference duration at rate=1.0 + frame_layout.addWidget(self.rate_spin, 1, 1) main_layout.addLayout(frame_layout) @@ -114,13 +122,37 @@ def _on_theme_changed(self): label.setStyleSheet(muted) def _on_rate_changed(self, rate): - """Handle rate change by updating the derived stop_frame.""" - if rate <= 0: + if not self.animation: + return None + self.animation.set_rate(rate) + self.sync_timing() + self.animationChanged.emit({"rate": rate}) + return None + + def _on_duration_changed(self, frames): + if not self.animation: + return None + self.animation.set_duration(frames) + self.sync_timing() + self.animationChanged.emit({"duration": frames}) + return None + + def sync_timing(self): + """Refresh duration and rate widgets from the current animation.""" + if not self.animation: return None - new_duration = max(1, int(self._base_duration / rate)) - if self.animation: - self.animation.rate = rate - self.on_change(new_duration, "stop_frame") + + for widget in (self.duration_spin, self.rate_spin): + widget.blockSignals(True) + self.duration_spin.setValue(self.animation.duration) + self.rate_spin.setValue(self.animation.rate) + for widget in (self.duration_spin, self.rate_spin): + widget.blockSignals(False) + + has_rate = self.animation.has_rate + self._rate_label.setVisible(has_rate) + self.rate_spin.setVisible(has_rate) + return None def set_animation(self, animation): self.animation = animation @@ -130,25 +162,18 @@ def set_animation(self, animation): self.name_edit, self.global_start_spin, self.enabled_check, - self.rate_spin, ]: widget.blockSignals(True) self.name_edit.setText(animation.name) self.global_start_spin.setValue(animation.global_start_frame) self.enabled_check.setChecked(animation.enabled) - - if animation._base_duration is not None: - self._base_duration = animation._base_duration - else: - self._base_duration = animation.stop_frame * animation.rate - self.rate_spin.setValue(animation.rate) + self.sync_timing() for widget in [ self.name_edit, self.global_start_spin, self.enabled_check, - self.rate_spin, ]: widget.blockSignals(False) @@ -209,10 +234,7 @@ def on_change(self, value, key): # Data range changes update stop_frame; sync the UI if key in ("data_start", "data_stop"): - self._base_duration = self.animation._base_duration or self._base_duration - self.rate_spin.blockSignals(True) - self.rate_spin.setValue(self.animation.rate) - self.rate_spin.blockSignals(False) + self.sync_timing() self.animationChanged.emit({key: value}) diff --git a/src/mosaic/animation/timeline.py b/src/mosaic/animation/timeline.py index 42edcc45..38b4c9ab 100644 --- a/src/mosaic/animation/timeline.py +++ b/src/mosaic/animation/timeline.py @@ -10,14 +10,20 @@ class TimelineContent(QWidget): trackSelected = Signal(str) frameMoved = Signal(int) trackMoved = Signal(str, int) + trackResized = Signal(str, int) trackRemoved = Signal(str) + #: Half-width in pixels of the grab zone around a track's right edge. + resize_handle_width = 5 + def __init__(self, parent=None): super().__init__(parent) self.tracks = [] self.current_frame = 0 self.selected_track = None self.dragging_track = None + self.resizing_track = None + self.drag_active = False self.drag_start_x = 0 self.drag_start_frame = 0 @@ -213,92 +219,113 @@ def paintEvent(self, event): painter.setPen(QPen(QColor(Colors.BORDER_HOVER), 1)) painter.drawLine(int(playhead_x), 0, int(playhead_x), self.height()) - def mousePressEvent(self, event: QMouseEvent): - if event.button() != Qt.MouseButton.LeftButton or not self.tracks: - return super().mousePressEvent(event) + def _hit_test(self, x: float, y: float): + """Return ``(track, region)`` under the cursor, or ``(None, None)``. + Region is one of ``"remove"``, ``"resize"`` or ``"body"``. + """ track_y = self.ruler_height + 10 min_click_width = 30 # Minimum clickable width for narrow tracks + bar_height = self.track_height - 10 for track in self.tracks: start_x = self.frame_to_x(track.animation.global_start_frame) end_x = self.frame_to_x( track.animation.global_start_frame + track.animation.duration ) + if not track_y <= y <= track_y + bar_height: + track_y += self.track_height + continue - # Expand clickable area for narrow tracks - track_width = end_x - start_x - if track_width < min_click_width: - padding = (min_click_width - track_width) / 2 - click_start_x = start_x - padding - click_end_x = end_x + padding - else: - click_start_x = start_x - click_end_x = end_x - - # Check if click is on remove button remove_x = int(end_x - self.remove_button_size - 5) - remove_y = track_y + (self.track_height - 10 - self.remove_button_size) // 2 - + remove_y = track_y + (bar_height - self.remove_button_size) // 2 if ( - remove_x <= event.x() <= remove_x + self.remove_button_size - and remove_y <= event.y() <= remove_y + self.remove_button_size + remove_x <= x <= remove_x + self.remove_button_size + and remove_y <= y <= remove_y + self.remove_button_size and start_x < remove_x < self.width() ): - self.trackRemoved.emit(track.id) + return track, "remove" - # Use expanded click area for track selection - if ( - click_start_x <= event.x() <= click_end_x - and track_y <= event.y() <= track_y + self.track_height - 10 - ): - self.selected_track = track.id - self.dragging_track = track.id - self.drag_start_x = event.x() - self.drag_start_frame = track.animation.global_start_frame - self.trackSelected.emit(track.id) - self.update() - return + # On narrow bars only grab outside the bar so it remains movable + handle = self.resize_handle_width + inner = handle if end_x - start_x >= 3 * handle else 0 + if end_x - inner <= x <= end_x + handle: + return track, "resize" + + padding = max(0, (min_click_width - (end_x - start_x)) / 2) + if start_x - padding <= x <= end_x + padding: + return track, "body" track_y += self.track_height + return None, None + + def _ensure_visible(self, frame: int) -> None: + if frame > self.view_end_frame - 10: + self.view_end_frame = frame + 20 + self.update() + return None + + def mousePressEvent(self, event: QMouseEvent): + if event.button() != Qt.MouseButton.LeftButton or not self.tracks: + return super().mousePressEvent(event) + + track, region = self._hit_test(event.x(), event.y()) + if region == "remove": + self.trackRemoved.emit(track.id) + return None + + if track is not None: + self.selected_track = track.id + self.drag_active = False + self.drag_start_x = event.x() + self.drag_start_frame = track.animation.global_start_frame + if region == "resize": + self.resizing_track = track.id + else: + self.dragging_track = track.id + self.trackSelected.emit(track.id) + self.update() + return None frame = self.x_to_frame(event.x()) self.current_frame = max(0, frame) self.frameMoved.emit(self.current_frame) self.update() + return None def mouseMoveEvent(self, event: QMouseEvent): - # Check if we're hovering over a remove button if not (event.buttons() & Qt.MouseButton.LeftButton): - hovered_button = None - track_y = self.ruler_height + 10 - for track in self.tracks: - start_x = self.frame_to_x(track.animation.global_start_frame) - end_x = self.frame_to_x( - track.animation.global_start_frame + track.animation.duration - ) - - remove_x = int(end_x - self.remove_button_size - 5) - remove_y = ( - track_y + (self.track_height - 10 - self.remove_button_size) // 2 - ) - - if ( - remove_x <= event.x() <= remove_x + self.remove_button_size - and remove_y <= event.y() <= remove_y + self.remove_button_size - and start_x < remove_x < self.width() - ): - hovered_button = track.id - break - - track_y += self.track_height + track, region = self._hit_test(event.x(), event.y()) + hovered_button = track.id if region == "remove" else None if hovered_button != self.hovered_remove_button: self.hovered_remove_button = hovered_button self.update() - draggable = self.dragging_track and self.tracks - if not draggable or not (event.buttons() & Qt.MouseButton.LeftButton): + if region == "resize": + self.setCursor(Qt.CursorShape.SizeHorCursor) + else: + self.unsetCursor() + return super().mouseMoveEvent(event) + + # Clicking to select should never nudge a track by a frame or two + if (self.resizing_track or self.dragging_track) and not self.drag_active: + if abs(event.x() - self.drag_start_x) < QApplication.startDragDistance(): + return None + self.drag_active = True + + if self.resizing_track: + track = next((t for t in self.tracks if t.id == self.resizing_track), None) + if track is None: + return None + + start = track.animation.global_start_frame + duration = max(1, self.x_to_frame(event.x()) - start) + self.trackResized.emit(track.id, duration) + self._ensure_visible(start + duration) + return None + + if not (self.dragging_track and self.tracks): return super().mouseMoveEvent(event) frame = self.x_to_frame(event.x()) @@ -311,13 +338,13 @@ def mouseMoveEvent(self, event: QMouseEvent): if track is None: return None - track_end = new_frame + track.animation.duration - if track_end > self.view_end_frame - 10: - self.view_end_frame = track_end + 20 - self.update() + self._ensure_visible(new_frame + track.animation.duration) + return None def mouseReleaseEvent(self, event: QMouseEvent): - self.dragging_track = False + self.dragging_track = None + self.resizing_track = None + self.drag_active = False return self.update() def wheelEvent(self, event): @@ -358,6 +385,7 @@ class TimelineWidget(QScrollArea): trackSelected = Signal(str) frameMoved = Signal(int) trackMoved = Signal(str, int) + trackResized = Signal(str, int) trackRemoved = Signal(str) def __init__(self, parent=None): @@ -368,6 +396,7 @@ def __init__(self, parent=None): self.content.trackSelected.connect(self.trackSelected.emit) self.content.frameMoved.connect(self.frameMoved.emit) self.content.trackMoved.connect(self.trackMoved.emit) + self.content.trackResized.connect(self.trackResized.emit) self.content.trackRemoved.connect(self.trackRemoved.emit) self.setWidget(self.content) diff --git a/src/mosaic/data/tab-caret-left.svg b/src/mosaic/data/tab-caret-left.svg new file mode 100644 index 00000000..2113b407 --- /dev/null +++ b/src/mosaic/data/tab-caret-left.svg @@ -0,0 +1,4 @@ + + + + diff --git a/src/mosaic/data/tab-caret-right.svg b/src/mosaic/data/tab-caret-right.svg new file mode 100644 index 00000000..3bb9462f --- /dev/null +++ b/src/mosaic/data/tab-caret-right.svg @@ -0,0 +1,4 @@ + + + + diff --git a/src/mosaic/dialogs/histogram.py b/src/mosaic/dialogs/histogram.py index a0195549..8054295a 100644 --- a/src/mosaic/dialogs/histogram.py +++ b/src/mosaic/dialogs/histogram.py @@ -238,6 +238,7 @@ class HistogramDialog(QDialog): def __init__(self, cdata, parent=None): super().__init__(parent) self.cdata = cdata + self.setWindowTitle("Histogram") layout = QVBoxLayout(self) group = QGroupBox("Select") diff --git a/src/mosaic/dts/configure_panel.py b/src/mosaic/dts/configure_panel.py index ec192789..c1aec762 100644 --- a/src/mosaic/dts/configure_panel.py +++ b/src/mosaic/dts/configure_panel.py @@ -496,6 +496,9 @@ def _reset_fields(self): widget.setValue(widget.minimum()) widget.blockSignals(False) + set_widget_value(self._param_widgets["scale_factor"], 1.0) + self._param_widgets["offset"].setText("0,0,0") + self._volume_screen_cb.blockSignals(True) self._volume_screen_cb.setChecked(False) self._volume_screen_cb.blockSignals(False) diff --git a/src/mosaic/geometry.py b/src/mosaic/geometry.py index d5586c4e..402a084d 100644 --- a/src/mosaic/geometry.py +++ b/src/mosaic/geometry.py @@ -2220,12 +2220,17 @@ class GeometryTrajectory(Geometry): def __init__(self, trajectory: list[dict], **kwargs): super().__init__(**kwargs) self._trajectory = trajectory + self._current_frame = None def __getstate__(self): state = super().__getstate__() state.update({"trajectory": self._trajectory}) return state + def swap_data(self, *args, **kwargs): + self._current_frame = None + return super().swap_data(*args, **kwargs) + @property def frames(self): """ @@ -2268,8 +2273,21 @@ def display_frame(self, frame_idx: int) -> bool: meta=meta, model=model, ) + self._current_frame = frame_idx return True + @property + def current_frame(self): + """ + Index of the displayed frame. + + Returns + ------- + int or None + Frame last shown via :py:meth:`display_frame`, or None if unknown. + """ + return getattr(self, "_current_frame", None) + # For backwards compatibility class PointCloud(Geometry): diff --git a/src/mosaic/gui.py b/src/mosaic/gui.py index 69e4c5bc..38abbe91 100644 --- a/src/mosaic/gui.py +++ b/src/mosaic/gui.py @@ -38,6 +38,7 @@ QPushButton, QSplitter, QStackedWidget, + QTabWidget, QVBoxLayout, QWidget, ) @@ -80,6 +81,9 @@ class App(QMainWindow): def __init__(self): super().__init__() self.setWindowState(Qt.WindowNoState) + self.setTabPosition( + Qt.DockWidgetArea.AllDockWidgetAreas, QTabWidget.TabPosition.North + ) central_widget = QWidget() self.setCentralWidget(central_widget) diff --git a/src/mosaic/properties.py b/src/mosaic/properties.py index 4a73807f..8734fb45 100644 --- a/src/mosaic/properties.py +++ b/src/mosaic/properties.py @@ -40,15 +40,15 @@ def _aggregate(data, mode: str): mode = mode.lower() if data.ndim == 2: if mode == "mean": - data = data.mean(axis=1) + data = np.nanmean(data, axis=1) elif mode == "min": - data = data.min(axis=1) + data = np.nanmin(data, axis=1) elif mode == "max": - data = data.max(axis=1) + data = np.nanmax(data, axis=1) elif mode == "median": - data = np.median(data, axis=1) + data = np.nanmedian(data, axis=1) elif mode == "std": - data = np.std(data, axis=1) + data = np.nanstd(data, axis=1) return data @@ -213,19 +213,22 @@ def distance( fetch_k = k + 1 if is_self_query else k dist, _ = find_closest_points(query, geometry.points, k=fetch_k) + + # KDTree pads neighbors beyond the query size with inf + dist = np.where(np.isinf(dist), np.nan, dist) if is_self_query: - dist = dist[:, 1:] if dist.ndim == 2 else dist + dist = dist[:, 1:] - if distance is None: - distance = dist - distance = np.minimum(distance, dist) + dist = np.asarray(dist).reshape(geometry.points.shape[0], -1) + distance = dist if distance is None else np.fmin(distance, dist) if distance is None: return None - if distance.ndim == 2: - distance = distance[:, k_start - 1 : k] - return _aggregate(distance, mode=aggregation) + distance = distance[:, k_start - 1 : k] + with warnings.catch_warnings(): + warnings.simplefilter("ignore", RuntimeWarning) + return _aggregate(distance, mode=aggregation) def box_size(geometry, axis: int = None): diff --git a/src/mosaic/stylesheets.py b/src/mosaic/stylesheets.py index 3aa3ecf1..1986bb86 100644 --- a/src/mosaic/stylesheets.py +++ b/src/mosaic/stylesheets.py @@ -658,6 +658,56 @@ def _build_QDockWidget_style(): subcontrol-position: right center; right: 24px; }} + QMainWindow > QTabBar {{ + background: transparent; + border: none; + qproperty-drawBase: 0; + qproperty-elideMode: ElideNone; + qproperty-usesScrollButtons: 1; + }} + QMainWindow > QTabBar::tab {{ + border: none; + background: transparent; + border-radius: 6px; + padding: 6px 14px; + margin: 4px 2px; + font-size: {Typography.BODY}px; + color: {Colors.TEXT_MUTED}; + }} + QMainWindow > QTabBar::tab:selected {{ + background: {"rgba(255, 255, 255, 15)" if Colors.is_dark() else "rgba(0, 0, 0, 10)"}; + color: {Colors.TEXT_PRIMARY}; + }} + QMainWindow > QTabBar::tab:hover:!selected {{ + color: {Colors.TEXT_SECONDARY}; + }} + QMainWindow > QTabBar::tear {{ + width: 0px; + border: none; + image: none; + }} + QMainWindow > QTabBar::scroller {{ + width: 48px; + }} + QMainWindow > QTabBar QToolButton {{ + background: {Colors.SURFACE}; + border: none; + border-radius: 6px; + margin: 4px 0px; + }} + QMainWindow > QTabBar QToolButton:hover {{ + background: {Colors.BG_HOVER}; + }} + QMainWindow > QTabBar QToolButton::left-arrow {{ + image: url('{_get_resource_path("tab-caret-left.svg")}'); + width: 14px; + height: 14px; + }} + QMainWindow > QTabBar QToolButton::right-arrow {{ + image: url('{_get_resource_path("tab-caret-right.svg")}'); + width: 14px; + height: 14px; + }} QDockWidget::close-button:hover, QDockWidget::float-button:hover {{ background: {Colors.BG_HOVER}; border: 1px solid {Colors.BORDER_DARK}; @@ -835,6 +885,24 @@ def _apply_macos_titlebar(ns_win, hide_title=False): pass +def _set_macos_titlebar_transparent(ns_win): + """Re-assert a transparent title bar on a native window.""" + try: + import ctypes + import ctypes.util + + objc = ctypes.cdll.LoadLibrary(ctypes.util.find_library("objc")) + objc.sel_registerName.restype = ctypes.c_void_p + objc.sel_registerName.argtypes = [ctypes.c_char_p] + send = objc.objc_msgSend + send.argtypes = [ctypes.c_void_p, ctypes.c_void_p, ctypes.c_bool] + send.restype = ctypes.c_void_p + send(ns_win, objc.sel_registerName(b"setTitlebarAppearsTransparent:"), True) + except Exception: + pass + return None + + def _update_macos_titlebar_color(ns_win): """Set an NSWindow's background color to match the current theme surface.""" try: @@ -902,20 +970,38 @@ def install_macos_titlebar_filter(): if sys.platform != "darwin": return None - from qtpy.QtCore import QEvent, QObject - from qtpy.QtWidgets import QApplication + from qtpy.QtCore import QEvent, QObject, QTimer + from qtpy.QtWidgets import QApplication, QMainWindow, QTabBar + + tab_bar_events = ( + QEvent.Type.Show, + QEvent.Type.Hide, + QEvent.Type.Move, + QEvent.Type.Resize, + ) + + def _restore_transparency(window): + if (ns_win := _get_nswindow(window)) is not None: + _set_macos_titlebar_transparent(ns_win) + return None class _MacOSTitleBarFilter(QObject): """Event filter that styles the native title bar of every top-level window.""" def eventFilter(self, obj, event): + # QTabBar geometry changes make Qt's cocoa plugin reset + # titlebarAppearsTransparent on the main window (updateMacBorderMetrics). + if isinstance(obj, QTabBar) and event.type() in tab_bar_events: + window = obj.window() + if isinstance(window, QMainWindow): + QTimer.singleShot(0, lambda: _restore_transparency(window)) + return False + if ( event.type() == QEvent.Type.Show and hasattr(obj, "isWindow") and obj.isWindow() ): - from qtpy.QtWidgets import QMainWindow - if isinstance(obj, QMainWindow): return False ns_win = _get_nswindow(obj) diff --git a/src/mosaic/widgets/dock.py b/src/mosaic/widgets/dock.py index b414c529..f236ced4 100644 --- a/src/mosaic/widgets/dock.py +++ b/src/mosaic/widgets/dock.py @@ -180,7 +180,26 @@ def closeEvent(self, event): if main_window is None: return dialog_widget.show() + dock.setWindowTitle(dialog_widget.windowTitle()) main_window.addDockWidget(dock_area, dock) + + sibling = next( + ( + other + for other in main_window.findChildren(QDockWidget) + if other is not dock + and other.isVisible() + and not other.isFloating() + and isinstance(other.titleBarWidget(), _DockTitleBar) + and main_window.dockWidgetArea(other) == dock_area + ), + None, + ) + if sibling is not None: + width = sibling.width() + main_window.tabifyDockWidget(sibling, dock) + main_window.resizeDocks([dock], [width], Qt.Orientation.Horizontal) + setattr(instance, dock_attr_name, dock) dock.show() dock.raise_() diff --git a/tests/animation/__init__.py b/tests/animation/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/animation/test_object_selection.py b/tests/animation/test_object_selection.py new file mode 100644 index 00000000..f9f2a6d2 --- /dev/null +++ b/tests/animation/test_object_selection.py @@ -0,0 +1,87 @@ +from types import SimpleNamespace + +import numpy as np +import pytest + +from mosaic.animation.animations import ActorSelectionDialog +from mosaic.geometry import Geometry +from mosaic.tree_state import TreeStateData + + +class _Data: + def __init__(self, clusters, models, groups=None): + self._objects = {"data": clusters, "models": models} + self.data = self._interactor(clusters, groups or {}) + self.models = self._interactor(models, {}) + + @staticmethod + def _interactor(objects, groups): + state = TreeStateData() + grouped = {uuid for members in groups.values() for uuid in members} + state.root_items = [g.uuid for g in objects if g.uuid not in grouped] + for name, members in groups.items(): + state.groups[name] = members + state.group_names[name] = name + state.root_items.append(name) + return SimpleNamespace(data_list=SimpleNamespace(to_state=lambda: state)) + + def format_datalist(self, data_type): + return [(g._meta["name"], g) for g in self._objects[data_type]] + + +def _geometry(name, visible=True): + g = Geometry(points=np.random.default_rng(0).random((5, 3)), meta={"name": name}) + g.set_visibility(visible) + return g + + +@pytest.fixture +def scene(qapp): + clusters = [_geometry("membrane"), _geometry("ribosome"), _geometry("noise", False)] + models = [_geometry("mesh")] + return clusters, models + + +def test_preselects_current_selection(scene): + clusters, models = scene + dialog = ActorSelectionDialog( + _Data(clusters, models), current_selection=[id(clusters[1])] + ) + assert dialog.get_selected_objects() == [id(clusters[1])] + assert dialog._count.text() == "1 of 4 selected" + + +@pytest.mark.parametrize( + "data_types, expected", + [ + (("data", "models"), 4), + (("data",), 3), + (("models",), 1), + ((), 0), + ], +) +def test_quick_select(scene, data_types, expected): + dialog = ActorSelectionDialog(_Data(*scene)) + dialog._select(data_types) + assert len(dialog.get_selected_objects()) == expected + + +def test_quick_select_respects_filter(scene): + clusters, models = scene + dialog = ActorSelectionDialog(_Data(clusters, models)) + dialog._filter("mem") + dialog._select(("data", "models")) + assert dialog.get_selected_objects() == [id(clusters[0])] + + +def test_grouped_objects_are_listed_and_filtered(scene): + clusters, models = scene + groups = {"organelles": [clusters[0].uuid, clusters[1].uuid]} + dialog = ActorSelectionDialog(_Data(clusters, models, groups)) + + dialog._select(("data",)) + assert set(dialog.get_selected_objects()) == {id(g) for g in clusters} + + dialog._filter("mesh") + tree = dialog._trees["data"].tree_widget + assert all(tree.topLevelItem(i).isHidden() for i in range(tree.topLevelItemCount())) diff --git a/tests/animation/test_trajectory.py b/tests/animation/test_trajectory.py new file mode 100644 index 00000000..3c43fccc --- /dev/null +++ b/tests/animation/test_trajectory.py @@ -0,0 +1,116 @@ +from unittest.mock import Mock + +import pytest + +from mosaic.animation.animations import TrajectoryAnimation + + +class _Trajectory: + uuid = "traj" + frames = 10 + + def __init__(self): + self.current_frame = None + self.shown = [] + + def display_frame(self, idx): + self.shown.append(idx) + self.current_frame = idx + return True + + +def _track(trajectory, start=0, selected=()): + models = Mock() + models._get_selected_uuids.return_value = list(selected) + cdata = Mock(models=models) + cdata.format_datalist.return_value = [] + + anim = TrajectoryAnimation( + vtk_widget=None, cdata=cdata, volume_viewer=None, global_start_frame=start + ) + anim._trajectory = trajectory + anim.update_parameters(data_start=0, data_stop=10) + return anim, models + + +def _render(tracks, frame): + for t in tracks: + t.reset() + for t in tracks: + t.update(frame) + for t in tracks: + t.commit() + + +@pytest.fixture(autouse=True) +def _clear_claims(): + TrajectoryAnimation._claims.clear() + yield + TrajectoryAnimation._claims.clear() + + +def test_before_start_shows_first_frame(): + trajectory = _Trajectory() + track, _ = _track(trajectory, start=5) + + _render([track], 8) + _render([track], 0) + assert trajectory.current_frame == 0 + + +def test_after_end_holds_last_frame(): + trajectory = _Trajectory() + track, _ = _track(trajectory, start=0) + + _render([track], 2) + _render([track], 50) + assert trajectory.current_frame == 9 + + +def test_latest_started_track_drives_trajectory(): + trajectory = _Trajectory() + first, _ = _track(trajectory, start=0) + second, _ = _track(trajectory, start=20) + + _render([second, first], 3) + assert trajectory.current_frame == 3 + _render([second, first], 22) + assert trajectory.current_frame == 2 + + +def test_earliest_track_fills_before_any_start(): + trajectory = _Trajectory() + late, _ = _track(trajectory, start=20) + early, _ = _track(trajectory, start=10) + early.update_parameters(data_start=4, data_stop=10) + + _render([late, early], 0) + assert trajectory.current_frame == 4 + + +def test_unchanged_frame_is_not_redisplayed(): + trajectory = _Trajectory() + track, _ = _track(trajectory, start=0) + + _render([track], 50) + _render([track], 60) + assert trajectory.shown == [9] + + +@pytest.mark.parametrize("selected", [["other"], ["traj", "other"]]) +def test_update_never_renders_through_interactor(selected): + track, models = _track(_Trajectory(), selected=selected) + _render([track], 3) + models.set_selection_by_uuid.assert_not_called() + models.render_vtk.assert_not_called() + + +def test_highlight_restored_only_when_trajectory_selected(): + track, models = _track(_Trajectory(), selected=["other"]) + _render([track], 3) + models.container.highlight.assert_not_called() + + TrajectoryAnimation._claims.clear() + track, models = _track(_Trajectory(), selected=["traj", "other"]) + _render([track], 3) + models.container.highlight.assert_called_once_with(["traj", "other"]) diff --git a/tests/animation/test_visibility.py b/tests/animation/test_visibility.py new file mode 100644 index 00000000..bce32d04 --- /dev/null +++ b/tests/animation/test_visibility.py @@ -0,0 +1,274 @@ +from types import SimpleNamespace + +import numpy as np +import pytest +from qtpy.QtCore import QPoint, Qt +from qtpy.QtTest import QTest + +from mosaic.animation.animations import VisibilityAnimation +from mosaic.animation.timeline import TimelineContent +from mosaic.geometry import Geometry + + +class _Data: + def __init__(self, geometries): + self.geometries = geometries + + def format_datalist(self, data_type): + if data_type != "data": + return [] + return [(f"obj{i}", g) for i, g in enumerate(self.geometries)] + + +def _geometry(opacity=1.0, visible=True): + geometry = Geometry(points=np.random.default_rng(0).random((10, 3))) + geometry.actor.GetProperty().SetOpacity(opacity) + geometry.set_visibility(visible) + return geometry + + +def _track(cdata, geometries, start=0, duration=10, **params): + anim = VisibilityAnimation( + vtk_widget=None, cdata=cdata, volume_viewer=None, global_start_frame=start + ) + anim.update_parameters(selected_objects=[id(g) for g in geometries], **params) + anim.set_duration(duration) + return anim + + +def _render(tracks, frame): + for t in tracks: + t.reset() + for t in tracks: + t.update(frame) + for t in tracks: + t.commit() + + +@pytest.fixture(autouse=True) +def _clear_ledger(): + VisibilityAnimation._ledger.clear() + yield + VisibilityAnimation._ledger.clear() + + +def test_hidden_objects_are_switched_off_not_transparent(): + geometry = _geometry(opacity=0.4) + track = _track(_Data([geometry]), [geometry], start=5, duration=10) + + _render([track], 7) + assert not geometry.actor.GetVisibility() + assert geometry.actor.GetProperty().GetOpacity() == pytest.approx(0.4) + + +def test_track_restores_objects_after_its_end(): + geometry = _geometry(opacity=0.4) + track = _track(_Data([geometry]), [geometry], start=5, duration=10) + + for frame in (4, 15, 40): + _render([track], 7) + _render([track], frame) + assert geometry.actor.GetVisibility() + assert geometry.actor.GetProperty().GetOpacity() == pytest.approx(0.4) + + +def test_fade_scales_each_objects_own_opacity(): + geometry = _geometry(opacity=0.4) + track = _track(_Data([geometry]), [geometry], duration=20, target_opacity=0.5) + + _render([track], 3) + assert geometry.actor.GetVisibility() + assert geometry.actor.GetProperty().GetOpacity() == pytest.approx(0.2) + + +def test_instant_easing_switches_at_track_edges(): + geometry = _geometry() + track = _track(_Data([geometry]), [geometry], duration=20) + + _render([track], 0) + assert not geometry.actor.GetVisibility() + _render([track], 19) + assert not geometry.actor.GetVisibility() + _render([track], 20) + assert geometry.actor.GetVisibility() + + +def test_eased_fade_moves_from_start_to_target_across_track(): + geometry = _geometry(opacity=0.8) + track = _track( + _Data([geometry]), + [geometry], + duration=11, + start_opacity=0.0, + target_opacity=1.0, + easing="linear", + ) + opacity = geometry.actor.GetProperty().GetOpacity + + _render([track], 0) + assert not geometry.actor.GetVisibility() + _render([track], 5) + assert geometry.actor.GetVisibility() + assert opacity() == pytest.approx(0.4) + _render([track], 10) + assert opacity() == pytest.approx(0.8) + + +def test_instant_easing_ignores_start_opacity(): + geometry = _geometry() + track = _track( + _Data([geometry]), + [geometry], + duration=10, + start_opacity=1.0, + target_opacity=0.5, + ) + + _render([track], 0) + assert geometry.actor.GetProperty().GetOpacity() == pytest.approx(0.5) + + +def test_user_hidden_objects_are_left_alone(): + shown, hidden = _geometry(), _geometry(visible=False) + track = _track(_Data([shown, hidden]), [shown, hidden], duration=10) + + _render([track], 2) + _render([track], 20) + assert shown.actor.GetVisibility() + assert not hidden.actor.GetVisibility() + + +def test_later_started_track_overrides_earlier_one(): + geometry = _geometry(opacity=0.8) + cdata = _Data([geometry]) + hide = _track(cdata, [geometry], start=0, duration=30, target_opacity=0.0) + reveal = _track(cdata, [geometry], start=10, duration=5, target_opacity=0.5) + + _render([reveal, hide], 12) + assert geometry.actor.GetVisibility() + assert geometry.actor.GetProperty().GetOpacity() == pytest.approx(0.4) + + _render([reveal, hide], 20) + assert not geometry.actor.GetVisibility() + + _render([reveal, hide], 40) + assert geometry.actor.GetVisibility() + assert geometry.actor.GetProperty().GetOpacity() == pytest.approx(0.8) + + +def test_later_track_only_overrides_its_own_objects(): + a, b = _geometry(), _geometry() + cdata = _Data([a, b]) + hide = _track(cdata, [a, b], start=0, duration=30) + reveal = _track(cdata, [a], start=10, duration=5, target_opacity=1.0) + + _render([hide, reveal], 12) + assert a.actor.GetVisibility() + assert not b.actor.GetVisibility() + + +def test_same_start_goes_to_later_track(): + geometry = _geometry() + cdata = _Data([geometry]) + first = _track(cdata, [geometry], start=0, duration=10, target_opacity=0.0) + second = _track(cdata, [geometry], start=0, duration=10, target_opacity=0.5) + + _render([first, second], 3) + assert geometry.actor.GetProperty().GetOpacity() == pytest.approx(0.5) + + +def test_constant_frames_do_not_modify_actor(): + geometry = _geometry() + track = _track(_Data([geometry]), [geometry], duration=50) + + _render([track], 2) + mtimes = geometry.actor.GetMTime(), geometry.actor.GetProperty().GetMTime() + _render([track], 3) + assert (geometry.actor.GetMTime(), geometry.actor.GetProperty().GetMTime()) == ( + mtimes + ) + + +def test_deselected_objects_are_restored(): + a, b = _geometry(), _geometry() + track = _track(_Data([a, b]), [a, b], duration=10) + + _render([track], 2) + track.update_parameters(selected_objects=[id(a)]) + _render([track], 2) + assert not a.actor.GetVisibility() + assert b.actor.GetVisibility() + + +def test_visibility_duration_ignores_rate(): + track = _track(_Data([]), [], duration=10) + track.set_duration(37) + assert track.duration == 37 + assert track.rate == 1.0 + + +@pytest.fixture +def timeline(qapp): + anim = SimpleNamespace(global_start_frame=10, duration=50, enabled=True, name="t") + content = TimelineContent() + content.resize(600, 200) + content.set_tracks([SimpleNamespace(id="a", animation=anim, color="#888888")]) + content.view_start_frame, content.view_end_frame = 0, 100 + return content + + +def test_dragging_track_edge_emits_resize(timeline): + resized, moved = [], [] + timeline.trackResized.connect(lambda tid, d: resized.append((tid, d))) + timeline.trackMoved.connect(lambda tid, f: moved.append((tid, f))) + + y = timeline.ruler_height + 20 + edge = int(timeline.frame_to_x(60)) + QTest.mousePress(timeline, Qt.MouseButton.LeftButton, pos=QPoint(edge, y)) + QTest.mouseMove(timeline, QPoint(int(timeline.frame_to_x(80)), y)) + QTest.mouseRelease( + timeline, Qt.MouseButton.LeftButton, pos=QPoint(int(timeline.frame_to_x(80)), y) + ) + + assert resized and resized[-1] == ("a", 70) + assert not moved + + +def test_dragging_track_body_still_moves(timeline): + resized, moved = [], [] + timeline.trackResized.connect(lambda tid, d: resized.append((tid, d))) + timeline.trackMoved.connect(lambda tid, f: moved.append((tid, f))) + + y = timeline.ruler_height + 20 + QTest.mousePress( + timeline, Qt.MouseButton.LeftButton, pos=QPoint(int(timeline.frame_to_x(30)), y) + ) + QTest.mouseMove(timeline, QPoint(int(timeline.frame_to_x(40)), y)) + + assert moved and moved[-1] == ("a", 20) + assert not resized + + +def test_settings_scroll_instead_of_squeezing_timeline(qapp): + from mosaic.animation.compose import AnimationComposerDialog + + dialog = AnimationComposerDialog(vtk_widget=None, cdata=_Data([])) + dialog.resize(300, 500) + dialog.show() + qapp.processEvents() + + assert dialog.timeline.height() >= dialog.timeline.minimumHeight() >= 220 + dialog.close() + + +def test_click_without_moving_does_not_shift_track(timeline): + moved = [] + timeline.trackMoved.connect(lambda tid, f: moved.append((tid, f))) + + y = timeline.ruler_height + 20 + x = int(timeline.frame_to_x(30)) + QTest.mousePress(timeline, Qt.MouseButton.LeftButton, pos=QPoint(x, y)) + QTest.mouseMove(timeline, QPoint(x + 1, y)) + QTest.mouseRelease(timeline, Qt.MouseButton.LeftButton, pos=QPoint(x + 1, y)) + + assert not moved diff --git a/tests/animation/test_volume.py b/tests/animation/test_volume.py new file mode 100644 index 00000000..0203f4bb --- /dev/null +++ b/tests/animation/test_volume.py @@ -0,0 +1,96 @@ +from unittest.mock import Mock + +import pytest +import vtk + +from mosaic.animation.animations import VolumeAnimation + + +class _Viewer: + def __init__(self): + self.slice = vtk.vtkActor() + self.is_visible = True + self.orientation_selector = Mock() + self.orientation_selector.currentText.return_value = "Z" + self.project_selector = Mock() + self._rendering_suspended = False + self.slices = [] + + def get_dimensions(self): + return (10, 10, 10) + + def get_orientation(self): + return "Z" + + def get_projection(self): + return "Off" + + def set_slice(self, n): + self.slices.append(n) + + +def _track(viewer, start, **params): + anim = VolumeAnimation( + vtk_widget=None, + cdata=None, + volume_viewer=Mock(primary=viewer), + global_start_frame=start, + ) + anim.update_parameters(**params) + return anim + + +def _render(tracks, frame): + for t in tracks: + t.reset() + for t in tracks: + t.update(frame) + for t in tracks: + t.commit() + + +@pytest.fixture(autouse=True) +def _clear_state(): + VolumeAnimation._claims.clear() + VolumeAnimation._original_visibility.clear() + yield + VolumeAnimation._claims.clear() + VolumeAnimation._original_visibility.clear() + + +def test_finished_track_does_not_fight_later_track(): + viewer = _Viewer() + first = _track(viewer, start=0) + second = _track(viewer, start=30, direction="backward") + + _render([first, second], 35) + assert viewer.slices == [5] + + +def test_latest_start_wins_regardless_of_track_order(): + viewer = _Viewer() + later = _track(viewer, start=30) + earlier = _track(viewer, start=0) + + _render([later, earlier], 32) + assert viewer.slices == [2] + + +def test_hide_restores_visibility_when_scrubbing_back(): + viewer = _Viewer() + track = _track(viewer, start=10, hide="True") + + _render([track], 12) + assert not viewer.slice.GetVisibility() + _render([track], 0) + assert viewer.slice.GetVisibility() + + +def test_hidden_volume_is_not_recorded_as_original(): + viewer = _Viewer() + hider = _track(viewer, start=0, hide="True") + shower = _track(viewer, start=20) + + _render([hider, shower], 5) + _render([hider, shower], 25) + assert viewer.slice.GetVisibility() diff --git a/tests/core/test_properties.py b/tests/core/test_properties.py index 41874502..9b6d24ab 100644 --- a/tests/core/test_properties.py +++ b/tests/core/test_properties.py @@ -122,6 +122,28 @@ def test_only_self(self): assert result is not None assert result.shape[0] == 30 + def test_only_self_single_point_is_nan(self): + g = Geometry(points=np.zeros((1, 3), dtype=np.float32)) + for k in (1, 2): + assert np.isnan(distance(g, only_self=True, k=k)).all() + + def test_only_self_k_exceeds_neighbors(self): + points = np.array([[0, 0, 0], [3, 4, 0]], dtype=np.float32) + g = Geometry(points=points) + for k in (1, 2, 5): + np.testing.assert_allclose(distance(g, only_self=True, k=k), [5, 5]) + + def test_k_exceeds_query_size(self): + g = Geometry(points=np.zeros((1, 3), dtype=np.float32)) + query = np.array([[3, 4, 0], [6, 8, 0]], dtype=np.float32) + np.testing.assert_allclose(distance(g, queries=[query], k=5), [7.5]) + + def test_self_and_other_query_k1(self): + g = Geometry(points=np.array([[0, 0, 0], [3, 4, 0]], dtype=np.float32)) + other = np.array([[-1, 0, 0], [100, 0, 0]], dtype=np.float32) + result = distance(g, queries=[other, g], include_self=True, k=1) + np.testing.assert_allclose(result, [1, 5], rtol=1e-6) + def test_exclude_self(self): g = _geom(30) result = distance(g, queries=[g], include_self=False, k=1)