diff --git a/docs/guides/analysis-results.rst b/docs/guides/analysis-results.rst index 074ad74c..7d325240 100644 --- a/docs/guides/analysis-results.rst +++ b/docs/guides/analysis-results.rst @@ -51,6 +51,73 @@ Common columns: * - ``cluster_id`` - Pattern cluster assignment +Low-Level Features +------------------ + +RHEED videos expose a larger set of low-level, per-region features (e.g. +``area_0``, ``eccentricity_0``, ``fwhm_0_3``) beyond the standard columns above. +Request them with :meth:`~atomscale.client.Client.get_rheed_timeseries`, which +returns a DataFrame indexed by ``["Angle", "Frame Number"]``: + +.. code-block:: python + + df = client.get_rheed_timeseries(data_id, include_low_level_features=True) + print(df.filter(like="area").columns) + +The low-level columns keep their raw backend names (they are not renamed). + +Segmentation Masks +------------------ + +Each *featurized* frame of a processed RHEED video carries a binary segmentation +mask of the diffraction pattern. Attach the masks to the timeseries — aligned on +the ``Frame Number`` axis, alongside any low-level features — with +``include_masks``: + +.. code-block:: python + + from atomscale.results import decode_mask_rle + + df = client.get_rheed_timeseries( + data_id, + include_low_level_features=True, + include_masks=True, + ) + + # Mask columns: mask_rle (COCO RLE string), mask_height, mask_width. Coverage + # is sparse -- frames without a mask are NA -- so drop those rows first. + row = df.dropna(subset=["mask_rle"]).iloc[0] + mask = decode_mask_rle(row["mask_rle"], row["mask_height"], row["mask_width"]) + print(mask.shape) # (H, W) uint8, values 0/1 + +Fetch masks on their own — optionally decoded and keyed by absolute frame +number — with :meth:`~atomscale.client.Client.get_frame_masks`: + +.. code-block:: python + + masks = client.get_frame_masks(data_id, decode=True) # {frame_number: (H, W) array} + +Embedding Vectors +----------------- + +The similarity pipeline persists Chronos embedding vectors for RHEED data — the +*inputs* to similarity matching, as opposed to the derived similarity-vs-time +trajectory. Fetch them with :meth:`~atomscale.client.Client.get_embeddings`: + +.. code-block:: python + + emb = client.get_embeddings(data_id, window_span=60.0, kind="window") + print(emb.vectors.shape) # (n_windows, dimension) + +To find the RHEED data items most similar to a given one, run a +k-nearest-neighbour query over the embedding index with +:meth:`~atomscale.client.Client.query_rheed_embeddings`: + +.. code-block:: python + + neighbours = client.query_rheed_embeddings(data_id, top_k=10) + print(neighbours[["data_id", "similarity"]]) + Extracted Frames ---------------- diff --git a/src/atomscale/client.py b/src/atomscale/client.py index 2266c57f..7eb4526c 100644 --- a/src/atomscale/client.py +++ b/src/atomscale/client.py @@ -15,6 +15,7 @@ from typing import Any, BinaryIO, Literal import pandas as pd +from numpy.typing import NDArray from pandas import DataFrame from requests.exceptions import RequestException @@ -35,6 +36,7 @@ XPSResult, XRDResult, _get_rheed_image_result, + decode_mask_rle, ) from atomscale.results.group import PhysicalSampleResult, ProjectResult from atomscale.timeseries.align import align_timeseries @@ -48,6 +50,12 @@ # don't have to choose one; this value is available for effectively all RHEED data. _DEFAULT_SIMILARITY_METRIC = "specular_intensity" +# The frame-mask endpoint requires an inclusive upper frame bound. When a caller +# asks for the whole video (``to_frame=None``) we send this sentinel — far larger +# than any real frame count — and let the server clamp it to the artifact's actual +# range, fetching every featurized frame without a preliminary frame-count lookup. +_ALL_FRAMES_SENTINEL = 2**31 - 1 + def _retry_client_call( fn: Callable[..., Any], @@ -479,6 +487,45 @@ def get_similarity_trajectory( window_span=window_span or 0.0, ) + def query_rheed_embeddings( + self, + data_id: str, + *, + workflow: str = "rheed_stationary", + window_span: float = 60.0, + kind: Literal["prototype", "window"] = "prototype", + top_k: int = 10, + ) -> DataFrame: + """Find RHEED data items whose embeddings are most similar to this one. + + Runs k-NN over the embedding index using this item's own vectors and + returns the best-matching *other* data items ("find similar growths"). + + Args: + data_id: Data ID whose vectors seed the query. + workflow: Similarity workflow name. Defaults to "rheed_stationary". + window_span: Embedding window span in seconds (must match an embedded span). + kind: "prototype" (coarse, default) or "window" (finer, more queries). + top_k: Max neighbors to return. The backend caps this at 30. + + Returns: + DataFrame with columns ``data_id``, ``similarity`` (1 = identical), and + locus columns (``source_index``, ``neighbor_index``, ``real_time_seconds``, + ``unix_time_ms``), sorted by descending similarity. Empty when this item + has no embeddings for the given (workflow, window_span). + """ + from atomscale.similarity.embedding_provider import RHEEDEmbeddingProvider + + provider = RHEEDEmbeddingProvider() + params: dict[str, Any] = { + "workflow": workflow, + "window_span": window_span, + "kind": kind, + "top_k": top_k, + } + raw = provider.fetch_neighbors_raw(self, data_id, **params) + return provider.neighbors_to_dataframe(raw) + def get_embeddings( self, data_id: str, @@ -589,6 +636,7 @@ def get_rheed_timeseries( *, property_names: list[str] | None = None, include_low_level_features: bool = False, + include_masks: bool = False, last_n: int | None = None, elapsed_seconds: float | None = None, ) -> DataFrame: @@ -608,6 +656,16 @@ def get_rheed_timeseries( include_low_level_features: When ``True``, include the full set of low-level per-point features as additional columns. Defaults to ``False``. + include_masks: When ``True``, fetch the per-frame RHEED segmentation + masks (see :meth:`get_frame_masks`) and attach them to the + DataFrame as ``mask_rle`` / ``mask_height`` / ``mask_width`` + columns, joined on the ``Frame Number`` axis. Only masks for the + frames the returned series spans are fetched, so this respects any + ``last_n`` / ``elapsed_seconds`` window rather than pulling the whole + video's masks. Coverage is sparse (featurized frames only), so rows + whose frame has no mask — and all rows when the video has no mask + artifact — get NA in those columns. Decode a row's ``mask_rle`` with + :func:`atomscale.results.decode_mask_rle`. Defaults to ``False``. last_n: If set, only return the last ``N`` points. elapsed_seconds: If set, only return points within the last ``elapsed_seconds`` of the recording. @@ -615,7 +673,8 @@ def get_rheed_timeseries( Returns: DataFrame: The RHEED timeseries, indexed by ``["Angle", "Frame Number"]`` when available. Low-level feature columns are included when - ``include_low_level_features=True``. + ``include_low_level_features=True``; mask columns when + ``include_masks=True``. """ provider = get_provider("rheed") raw = provider.fetch_raw( @@ -626,7 +685,24 @@ def get_rheed_timeseries( last_n=last_n, elapsed_seconds=elapsed_seconds, ) - return provider.to_dataframe(raw) + ts_df = provider.to_dataframe(raw) + if include_masks: + # Scope the mask fetch to the frame-number range the returned series + # actually spans, so a windowed query (``last_n`` / ``elapsed_seconds``) + # doesn't pull the whole video's masks only to discard most in the join. + # ``None`` bounds mean an empty series or no frame axis to key on, so + # there is nothing to fetch or attach. + bounds = provider.frame_number_bounds(ts_df) + if bounds is not None: + first_frame, last_frame = bounds + mask_rows = self.get_frame_masks( + data_id, + from_frame=first_frame, + to_frame=last_frame, + decode=False, + ) + ts_df = provider.attach_frame_masks(ts_df, mask_rows) # type: ignore[arg-type] + return ts_df def get_frame( self, @@ -661,6 +737,81 @@ def get_frame( self, {"image_uuid": frame.get("image_uuid"), "metadata": metadata} ) + def get_frame_masks( + self, + data_id: str, + *, + from_frame: int = 0, + to_frame: int | None = None, + decode: bool = False, + ) -> list[dict[str, Any]] | dict[int, NDArray]: + """Fetch per-frame RHEED segmentation masks for a processed video. + + Each *featurized* frame of a processed RHEED video carries a binary + segmentation mask of the diffraction pattern, encoded as a COCO + run-length-encoding (RLE) ``counts`` string (the same format as the + single-frame :meth:`get_frame` mask). ``frame_number`` is the absolute + frame index, keyed identically to the processed video frames and the + RHEED timeseries ``Frame Number`` axis, so a decoded ``masks[frame_number]`` + overlays that frame of the video fetched via :meth:`download`. + + Coverage is **sparse**: masks exist only for featurized frames. For + stationary videos that is every frame; for rotating / per-azimuth videos + it is the sampled subset, so the returned frame numbers are not + necessarily contiguous — frames without a mask are simply absent. + + Args: + data_id: Data ID of the RHEED **video** (the same id used for the + video / timeseries). + from_frame: First absolute frame number to fetch, inclusive. Must be + ``>= 0``. Defaults to ``0``. + to_frame: Last absolute frame number to fetch, inclusive. ``None`` + (default) fetches through the end of the video (every featurized + frame from ``from_frame`` onward). + decode: When ``True``, decode each RLE mask into an ``(H, W)`` uint8 + (0/1) NumPy array and return a dict keyed by frame number. When + ``False`` (default), return the raw rows with the RLE string intact. + + Returns: + list[dict] | dict[int, NDArray]: When ``decode=False``, a list of row + dicts each with ``data_id``, ``processed_data_id``, ``frame_number``, + ``mask_rle``, ``mask_height`` and ``mask_width``. When ``decode=True``, + a dict ``{frame_number: np.ndarray}`` of decoded ``(H, W)`` uint8 masks. + Returns an empty list / dict when the video has no per-frame mask + artifact (e.g. a non-RHEED item, or a video processed before per-frame + masks were persisted). + """ + if from_frame < 0: + raise ValueError(f"from_frame must be >= 0, got {from_frame}") + if to_frame is not None and to_frame < 0: + raise ValueError(f"to_frame must be >= 0, got {to_frame}") + if to_frame is not None and to_frame < from_frame: + raise ValueError( + f"to_frame ({to_frame}) must be >= from_frame ({from_frame})" + ) + + resolved_to = _ALL_FRAMES_SENTINEL if to_frame is None else to_frame + + rows: list[dict] | None = self._get( # type: ignore[assignment] + sub_url=f"rheed/images/{data_id}/frame_masks", + params={"from": from_frame, "to": resolved_to}, + ) + + # `_get` returns None for a 404 ("No frame-mask artifact for this video") + # and for an empty body; both mean "no masks available" here. + if not rows: + return {} if decode else [] + + if not decode: + return rows + + return { + row["frame_number"]: decode_mask_rle( + row["mask_rle"], row["mask_height"], row["mask_width"] + ) + for row in rows + } + def iter_poll_similarity_trajectory( self, source_id: str, diff --git a/src/atomscale/results/__init__.py b/src/atomscale/results/__init__.py index 4b16d0aa..15fe83fb 100644 --- a/src/atomscale/results/__init__.py +++ b/src/atomscale/results/__init__.py @@ -6,7 +6,12 @@ from .optical import OpticalResult from .photoluminescence import PhotoluminescenceResult from .raman import RamanResult -from .rheed_image import RHEEDImageCollection, RHEEDImageResult, _get_rheed_image_result +from .rheed_image import ( + RHEEDImageCollection, + RHEEDImageResult, + _get_rheed_image_result, + decode_mask_rle, +) from .rheed_video import RHEEDVideoResult from .similarity_trajectory import SimilarityTrajectoryResult from .unknown import UnknownResult @@ -31,4 +36,5 @@ "XPSResult", "XRDResult", "_get_rheed_image_result", + "decode_mask_rle", ] diff --git a/src/atomscale/results/rheed_image.py b/src/atomscale/results/rheed_image.py index c213e4c0..b0bedacb 100644 --- a/src/atomscale/results/rheed_image.py +++ b/src/atomscale/results/rheed_image.py @@ -22,6 +22,29 @@ tp.quiet() +def decode_mask_rle( + mask_rle: str | bytes, mask_height: int, mask_width: int +) -> NDArray: + """Decode a COCO RLE ``counts`` string into a binary segmentation mask. + + Shared by every RHEED mask endpoint the SDK consumes (single-frame + ``rheed/images/{id}/mask`` and per-frame ``rheed/images/{id}/frame_masks``), + which all return the same pycocotools ``frString`` format. + + Args: + mask_rle (str | bytes): COCO run-length-encoding ``counts`` string + (pycocotools ``frString``, column-major / Fortran order). + mask_height (int): Mask height ``H`` in pixels. + mask_width (int): Mask width ``W`` in pixels. + + Returns: + (NDArray): An ``(H, W)`` ``uint8`` array with values 0 or 1. + """ + return mask_util.decode( + {"counts": mask_rle, "size": (mask_height, mask_width)} # type: ignore # noqa: PGH003 + ) + + class RHEEDImageResult(MSONable): def __init__( self, @@ -751,14 +774,9 @@ def _get_rheed_image_result( mask_array = None if mask_data is not None and mask_rle is not None: - mask_height = mask_data["mask_height"] - mask_width = mask_data["mask_width"] - - mask_dict = { - "counts": mask_rle, - "size": (mask_height, mask_width), - } - mask_array = mask_util.decode(mask_dict) # type: ignore # noqa: PGH003 + mask_array = decode_mask_rle( + mask_rle, mask_data["mask_height"], mask_data["mask_width"] + ) # Get raw and processed image data image_download: dict[str, str] | None = client._get( # type: ignore # noqa: PGH003 diff --git a/src/atomscale/similarity/embedding_provider.py b/src/atomscale/similarity/embedding_provider.py new file mode 100644 index 00000000..dfab314b --- /dev/null +++ b/src/atomscale/similarity/embedding_provider.py @@ -0,0 +1,58 @@ +"""Provider for RHEED embedding-vector k-NN neighbor queries. + +Wraps the read-only embedding-neighbors endpoint the backend exposes under +``/similarity/{workflow}/{data_id}/embeddings/neighbors/`` (k-NN "find similar"). +Mirrors the fetch/build split of +:class:`~atomscale.similarity.provider.SimilarityTrajectoryProvider`, returning a +tidy neighbors DataFrame. Raw embedding *vectors* are fetched separately via +:meth:`atomscale.Client.get_embeddings`. +""" + +from __future__ import annotations + +from typing import Any + +from pandas import DataFrame + +from atomscale.core import BaseClient + +_NEIGHBOR_COLUMNS = [ + "data_id", + "similarity", + "source_index", + "neighbor_index", + "real_time_seconds", + "unix_time_ms", +] + + +class RHEEDEmbeddingProvider: + TYPE = "rheed_embeddings" + + def fetch_neighbors_raw( + self, client: BaseClient, data_id: str, **kwargs: Any + ) -> Any: + """Fetch k-NN neighbors ("find similar") for a data_id. + + Args: + client: The API client. + data_id: The data ID whose vectors seed the query. + **kwargs: Must include ``workflow``. Optional: ``window_span``, + ``kind`` ("prototype"|"window"), ``top_k``. + """ + workflow = kwargs.pop("workflow") + return client._get( + sub_url=f"similarity/{workflow}/{data_id}/embeddings/neighbors/", + params=kwargs, + ) + + @staticmethod + def neighbors_to_dataframe(raw: Any) -> DataFrame: + """Build a tidy neighbors DataFrame from a fetch_neighbors_raw payload.""" + neighbors = (raw or {}).get("neighbors", []) or [] + if not neighbors: + return DataFrame(columns=_NEIGHBOR_COLUMNS) + frame = DataFrame(neighbors) + ordered = [c for c in _NEIGHBOR_COLUMNS if c in frame.columns] + extra = [c for c in frame.columns if c not in _NEIGHBOR_COLUMNS] + return frame[ordered + extra] diff --git a/src/atomscale/timeseries/rheed.py b/src/atomscale/timeseries/rheed.py index 448e7c6a..45857b68 100644 --- a/src/atomscale/timeseries/rheed.py +++ b/src/atomscale/timeseries/rheed.py @@ -49,6 +49,8 @@ class RHEEDProvider(TimeseriesProvider[RHEEDVideoResult]): "composition_metric", ] INDEX_COLS: Sequence[str] = ["Angle", "Frame Number"] + # Columns added to the timeseries DataFrame when per-frame masks are attached. + MASK_COLS: Sequence[str] = ["mask_rle", "mask_height", "mask_width"] def fetch_raw(self, client: BaseClient, data_id: str, **kwargs) -> Any: return client._get(sub_url=f"rheed/timeseries/{data_id}/", params=kwargs) @@ -119,6 +121,69 @@ def to_dataframe(self, raw: Any) -> DataFrame: return df_all + @classmethod + def attach_frame_masks( + cls, df: DataFrame, mask_rows: Sequence[Mapping[str, Any]] + ) -> DataFrame: + """Attach per-frame RLE segmentation masks to a RHEED timeseries DataFrame. + + Joins each mask row onto the timeseries row(s) with the matching absolute + frame number, adding ``mask_rle`` / ``mask_height`` / ``mask_width`` columns + (see :data:`MASK_COLS`). ``mask_rows`` are the raw rows returned by + :meth:`atomscale.Client.get_frame_masks` (each with ``frame_number``, + ``mask_rle``, ``mask_height``, ``mask_width``). + + Coverage is sparse — masks exist only for featurized frames — so timeseries + rows whose frame has no mask get NA in the mask columns. When ``mask_rows`` + is empty (no mask artifact for the video), the columns are still added and + are all-NA, so a caller that asked for masks always gets the columns. + + Returns ``df`` unchanged (no mask columns) when it has no ``Frame Number`` + axis to key on, since masks cannot be aligned without it. + """ + has_frame_axis = "Frame Number" in (df.index.names or []) or ( + "Frame Number" in df.columns + ) + if df.empty or not has_frame_axis: + return df + + cols = ["frame_number", *cls.MASK_COLS] + mask_df = ( + DataFrame(list(mask_rows), columns=cols) + .rename(columns={"frame_number": "Frame Number"}) + .set_index("Frame Number") + ) + + # Drop any pre-existing mask columns so a re-attach doesn't create + # duplicate/suffixed columns via the join. + clashing = [c for c in cls.MASK_COLS if c in df.columns] + base = df.drop(columns=clashing) if clashing else df + + return base.join(mask_df, on="Frame Number") + + @staticmethod + def frame_number_bounds(df: DataFrame) -> tuple[int, int] | None: + """Inclusive ``(min, max)`` absolute frame numbers present in ``df``. + + Reads the ``Frame Number`` axis (index level or column). Returns ``None`` + when the DataFrame is empty or has no frame-number axis, so callers can fall + back to a whole-video fetch. Used to scope a per-frame mask request to just + the frames a (possibly windowed via ``last_n`` / ``elapsed_seconds``) + timeseries actually spans, rather than fetching the whole video's masks. + """ + if df.empty: + return None + if "Frame Number" in (df.index.names or []): + values = df.index.get_level_values("Frame Number") + elif "Frame Number" in df.columns: + values = df["Frame Number"] + else: + return None + values = values.dropna() + if len(values) == 0: + return None + return int(values.min()), int(values.max()) + def snapshot_url(self, data_id: str) -> str: return f"data_entries/video_single_frames/{data_id}" diff --git a/tests/test_rheed_embedding.py b/tests/test_rheed_embedding.py new file mode 100644 index 00000000..21256012 --- /dev/null +++ b/tests/test_rheed_embedding.py @@ -0,0 +1,65 @@ +import os + +import pytest +from pandas import DataFrame + +from atomscale import Client +from atomscale.similarity.embedding_provider import RHEEDEmbeddingProvider + +from .conftest import ResultIDs + +# --------------------------- pure unit tests (no network) --------------------------- + + +def test_type_constant(): + assert RHEEDEmbeddingProvider.TYPE == "rheed_embeddings" + + +def test_neighbors_to_dataframe(): + raw = { + "neighbors": [ + {"data_id": "x", "similarity": 0.9, "source_index": 0, "neighbor_index": 3, + "real_time_seconds": None, "unix_time_ms": None}, + {"data_id": "y", "similarity": 0.7, "source_index": 1, "neighbor_index": 0, + "real_time_seconds": None, "unix_time_ms": None}, + ] + } + df = RHEEDEmbeddingProvider().neighbors_to_dataframe(raw) + assert list(df.columns)[:2] == ["data_id", "similarity"] + assert df.iloc[0]["data_id"] == "x" + assert df.iloc[0]["similarity"] == 0.9 + + +def test_neighbors_to_dataframe_empty(): + df = RHEEDEmbeddingProvider().neighbors_to_dataframe({"neighbors": []}) + assert isinstance(df, DataFrame) + assert list(df.columns) == [ + "data_id", "similarity", "source_index", "neighbor_index", + "real_time_seconds", "unix_time_ms", + ] + assert df.empty + + +# --------------------------- live-API tests (gated on creds) --------------------------- + + +def _skip_without_api(): + if not os.getenv("AS_API_KEY") and not os.getenv("ATOMSCALE_API_KEY"): + pytest.skip("No API key configured for live embedding tests") + if not ResultIDs.similarity_source_id or not ResultIDs.similarity_workflow: + pytest.skip("No similarity source configured") + + +def test_query_rheed_embeddings_live(): + _skip_without_api() + df = Client().query_rheed_embeddings( + ResultIDs.similarity_source_id, + workflow=ResultIDs.similarity_workflow, + window_span=60.0, + kind="prototype", + top_k=5, + ) + assert isinstance(df, DataFrame) + # the source item must never be returned as its own neighbor + if not df.empty: + assert str(ResultIDs.similarity_source_id) not in set(df["data_id"].astype(str)) diff --git a/tests/test_rheed_image.py b/tests/test_rheed_image.py index 28adc184..71cb0976 100644 --- a/tests/test_rheed_image.py +++ b/tests/test_rheed_image.py @@ -1,13 +1,15 @@ from io import BytesIO +import numpy as np import pytest from pandas import DataFrame from PIL import Image as PILImage from PIL.Image import Image +from pycocotools import mask as mask_util from atomscale import Client from atomscale.results import RHEEDImageResult -from atomscale.results.rheed_image import _get_rheed_image_result +from atomscale.results.rheed_image import _get_rheed_image_result, decode_mask_rle from .conftest import ResultIDs @@ -249,3 +251,110 @@ def test_get_frame_no_frames_returns_none(monkeypatch): unit_client = Client(api_key="key_test", endpoint="http://example.com/") monkeypatch.setattr(unit_client, "_get", lambda *a, **k: None) assert unit_client.get_frame("video-1") is None + + +# -------------------------------------------------------------------------- +# Unit tests (no live API) for per-frame mask fetching / decoding. +# -------------------------------------------------------------------------- + + +def _mask_row(frame_number: int, height: int = 6, width: int = 5) -> dict: + """Build a frame-mask row with a real COCO RLE counts string (as JSON str).""" + mask = np.zeros((height, width), dtype=np.uint8) + # A small filled block so the decoded mask is non-trivial and frame-specific. + mask[1 : 1 + (frame_number % height or 1), 0:2] = 1 + counts = mask_util.encode(np.asfortranarray(mask))["counts"].decode("utf-8") + return { + "data_id": "video-1", + "processed_data_id": "proc-1", + "frame_number": frame_number, + "mask_rle": counts, + "mask_height": height, + "mask_width": width, + } + + +def test_decode_mask_rle_roundtrips(): + row = _mask_row(3) + mask = decode_mask_rle(row["mask_rle"], row["mask_height"], row["mask_width"]) + assert mask.shape == (row["mask_height"], row["mask_width"]) + assert set(np.unique(mask)).issubset({0, 1}) + assert mask.sum() > 0 + + +def test_get_frame_masks_returns_raw_rows(monkeypatch): + unit_client = Client(api_key="key_test", endpoint="http://example.com/") + rows = [_mask_row(0), _mask_row(5)] + captured: dict = {} + + def fake_get(sub_url, params=None, **kwargs): + captured["sub_url"] = sub_url + captured["params"] = params + return rows + + monkeypatch.setattr(unit_client, "_get", fake_get) + + out = unit_client.get_frame_masks("video-1") + + assert out is rows # raw rows passed through untouched + assert captured["sub_url"] == "rheed/images/video-1/frame_masks" + # to_frame=None → whole video via the clamp sentinel + assert captured["params"]["from"] == 0 + assert captured["params"]["to"] == 2**31 - 1 + + +def test_get_frame_masks_explicit_range(monkeypatch): + unit_client = Client(api_key="key_test", endpoint="http://example.com/") + captured: dict = {} + + def fake_get(sub_url, params=None, **kwargs): # noqa: ARG001 + captured["params"] = params + return [] + + monkeypatch.setattr(unit_client, "_get", fake_get) + + unit_client.get_frame_masks("video-1", from_frame=10, to_frame=200) + assert captured["params"] == {"from": 10, "to": 200} + + +def test_get_frame_masks_decode(monkeypatch): + unit_client = Client(api_key="key_test", endpoint="http://example.com/") + rows = [_mask_row(0), _mask_row(7)] + monkeypatch.setattr(unit_client, "_get", lambda *a, **k: rows) + + masks = unit_client.get_frame_masks("video-1", decode=True) + + assert isinstance(masks, dict) + assert set(masks) == {0, 7} + for row in rows: + arr = masks[row["frame_number"]] + assert arr.shape == (row["mask_height"], row["mask_width"]) + expected = decode_mask_rle( + row["mask_rle"], row["mask_height"], row["mask_width"] + ) + assert np.array_equal(arr, expected) + + +def test_get_frame_masks_no_artifact_returns_empty(monkeypatch): + """A 404 on the frame-mask endpoint (→ _get returns None) yields empty.""" + unit_client = Client(api_key="key_test", endpoint="http://example.com/") + monkeypatch.setattr(unit_client, "_get", lambda *a, **k: None) + + assert unit_client.get_frame_masks("video-1") == [] + assert unit_client.get_frame_masks("video-1", decode=True) == {} + + +@pytest.mark.parametrize( + ("from_frame", "to_frame"), + [(-1, None), (0, -5), (10, 5)], +) +def test_get_frame_masks_invalid_range_raises(monkeypatch, from_frame, to_frame): + unit_client = Client(api_key="key_test", endpoint="http://example.com/") + # _get must never be reached when validation fails. + monkeypatch.setattr( + unit_client, + "_get", + lambda *a, **k: pytest.fail("_get should not be called on invalid range"), + ) + with pytest.raises(ValueError): # noqa: PT011 + unit_client.get_frame_masks("video-1", from_frame=from_frame, to_frame=to_frame) diff --git a/tests/test_rheed_timeseries.py b/tests/test_rheed_timeseries.py index 379c02e2..f26ba118 100644 --- a/tests/test_rheed_timeseries.py +++ b/tests/test_rheed_timeseries.py @@ -1,7 +1,7 @@ """Unit tests for Client.get_rheed_timeseries and low-level feature flattening.""" import pytest -from pandas import DataFrame +from pandas import DataFrame, isna from atomscale import Client from atomscale.timeseries.rheed import RHEEDProvider @@ -93,7 +93,7 @@ def test_to_dataframe_low_level_features_missing_points(provider): assert "feat_a" in df.columns values = df["feat_a"].tolist() assert values[0] == 1.0 - assert values[1] != values[1] # NaN + assert isna(values[1]) # NA — no feature for that frame def test_to_dataframe_without_low_level_features_unchanged(provider): @@ -182,3 +182,171 @@ def fake_get(sub_url, params=None, **kwargs): assert captured["params"]["include_low_level_features"] is False assert captured["params"]["property_names"] is None + + +# -------------------------------------------------------------------------- +# Per-frame mask attachment. +# -------------------------------------------------------------------------- + + +def _ts_df(provider, frame_numbers): + """Build a small timeseries DataFrame indexed by (Angle, Frame Number).""" + raw = { + "series_by_angle": [ + { + "angle": 0.0, + "series": [ + {"frame_number": fn, "specular_intensity": float(fn)} + for fn in frame_numbers + ], + } + ] + } + return provider.to_dataframe(raw) + + +def _mask_row(frame_number: int) -> dict: + return { + "data_id": "video-1", + "processed_data_id": "proc-1", + "frame_number": frame_number, + "mask_rle": f"rle-{frame_number}", + "mask_height": 6, + "mask_width": 5, + } + + +def test_attach_frame_masks_sparse_coverage(provider): + """Masks join on Frame Number; frames without a mask get NA.""" + df = _ts_df(provider, [0, 1, 2, 3]) + out = provider.attach_frame_masks(df, [_mask_row(0), _mask_row(2)]) + + assert list(out.index.names) == ["Angle", "Frame Number"] + assert set(provider.MASK_COLS).issubset(out.columns) + by_frame = out["mask_rle"].groupby("Frame Number").first() + assert by_frame[0] == "rle-0" + assert by_frame[2] == "rle-2" + assert isna(by_frame[1]) # NA — no mask for that frame + assert isna(by_frame[3]) # NA — no mask for that frame + # Height/width carried through for the populated frames. + assert out.xs(0, level="Frame Number")["mask_width"].iloc[0] == 5 + + +def test_attach_frame_masks_empty_adds_all_na_columns(provider): + """No mask artifact → columns still present, all NA.""" + df = _ts_df(provider, [0, 1]) + out = provider.attach_frame_masks(df, []) + + assert set(provider.MASK_COLS).issubset(out.columns) + assert out["mask_rle"].isna().all() + assert list(out.index.names) == ["Angle", "Frame Number"] + + +def test_attach_frame_masks_no_frame_axis_passthrough(provider): + """A DataFrame with no Frame Number axis is returned unchanged.""" + df = DataFrame({"a": [1, 2]}) + out = provider.attach_frame_masks(df, [_mask_row(0)]) + assert "mask_rle" not in out.columns + assert out.equals(df) + + +def test_attach_frame_masks_reattach_no_duplicate_columns(provider): + """Re-attaching replaces mask columns rather than suffixing them.""" + df = _ts_df(provider, [0, 1]) + once = provider.attach_frame_masks(df, [_mask_row(0)]) + twice = provider.attach_frame_masks(once, [_mask_row(0), _mask_row(1)]) + + assert list(twice.columns).count("mask_rle") == 1 + assert twice["mask_rle"].groupby("Frame Number").first()[1] == "rle-1" + + +def test_frame_number_bounds(provider): + assert provider.frame_number_bounds(_ts_df(provider, [3, 7, 5])) == (3, 7) + # Empty / no-frame-axis DataFrames yield None so callers can fall back. + assert provider.frame_number_bounds(DataFrame(None)) is None + assert provider.frame_number_bounds(DataFrame({"a": [1, 2]})) is None + + +def test_get_rheed_timeseries_include_masks(client, monkeypatch): + """include_masks=True fetches masks and merges them into the DataFrame.""" + captured: list[str] = [] + + def fake_get(sub_url, params=None, **kwargs): + captured.append(sub_url) + if sub_url.endswith("/frame_masks"): + return [_mask_row(1), _mask_row(2)] + return { + "series_by_angle": [ + { + "angle": 0.0, + "series": [ + {"frame_number": 1, "specular_intensity": 100.0}, + {"frame_number": 2, "specular_intensity": 110.0}, + ], + } + ] + } + + monkeypatch.setattr(client, "_get", fake_get) + + df = client.get_rheed_timeseries("video-1", include_masks=True) + + assert "rheed/timeseries/video-1/" in captured + assert "rheed/images/video-1/frame_masks" in captured + assert "mask_rle" in df.columns + assert df["mask_rle"].groupby("Frame Number").first().to_dict() == { + 1: "rle-1", + 2: "rle-2", + } + + +def test_get_rheed_timeseries_include_masks_scopes_to_window(client, monkeypatch): + """Mask fetch is bounded by the (windowed) series' frame range, not the whole video.""" + mask_params: dict = {} + + def fake_get(sub_url, params=None, **kwargs): + if sub_url.endswith("/frame_masks"): + mask_params.update(params or {}) + return [_mask_row(100), _mask_row(102)] + # A last_n-style window: only frames 100..102 come back from the series. + return { + "series_by_angle": [ + { + "angle": 0.0, + "series": [ + {"frame_number": fn, "specular_intensity": float(fn)} + for fn in (100, 101, 102) + ], + } + ] + } + + monkeypatch.setattr(client, "_get", fake_get) + + df = client.get_rheed_timeseries("video-1", include_masks=True, last_n=3) + + # from/to are clamped to the frames the series spans, not 0..(all frames). + assert mask_params == {"from": 100, "to": 102} + by_frame = df["mask_rle"].groupby("Frame Number").first() + assert by_frame[100] == "rle-100" + assert by_frame[102] == "rle-102" + assert isna(by_frame[101]) # NA — no mask for that frame + + +def test_get_rheed_timeseries_without_masks_makes_no_mask_call(client, monkeypatch): + """Default (include_masks=False) must not hit the frame_masks endpoint.""" + seen: list[str] = [] + + def fake_get(sub_url, params=None, **kwargs): + seen.append(sub_url) + return { + "series_by_angle": [ + {"angle": 0.0, "series": [{"frame_number": 1, "spot_count": 3}]} + ] + } + + monkeypatch.setattr(client, "_get", fake_get) + df = client.get_rheed_timeseries("video-1") + + assert not any(s.endswith("/frame_masks") for s in seen) + assert "mask_rle" not in df.columns diff --git a/tests/test_rheed_video.py b/tests/test_rheed_video.py index 5e733324..4d41cd07 100644 --- a/tests/test_rheed_video.py +++ b/tests/test_rheed_video.py @@ -3,6 +3,7 @@ from atomscale import Client from atomscale.results import RHEEDVideoResult +from atomscale.timeseries.rheed import RHEEDProvider from .conftest import ResultIDs @@ -53,3 +54,60 @@ def test_get_dataframe(result: RHEEDVideoResult): assert isinstance(result.timeseries_data, DataFrame) assert not len(set(result.timeseries_data.keys().values) - column_names) assert result.timeseries_data.index.names == ["Angle", "Frame Number"] + + +def test_to_dataframe_flattens_low_level_features(): + """When include_low_level_features is set, the backend nests a + low_level_features dict per point; the provider flattens it into raw-named + columns (no RENAME_MAP entry) and drops the nested column.""" + raw = { + "series_by_angle": [ + { + "angle": "0", + "series": [ + { + "frame_number": 0, + "relative_time_seconds": 0.0, + "unix_timestamp_ms": 0.0, + "specular_intensity": 5.0, + "low_level_features": {"area_0": 1.2, "eccentricity_0": 0.3}, + }, + { + "frame_number": 1, + "relative_time_seconds": 0.1, + "unix_timestamp_ms": 100.0, + "specular_intensity": 6.0, + "low_level_features": {"area_0": 1.5, "eccentricity_0": 0.4}, + }, + ], + } + ] + } + df = RHEEDProvider().to_dataframe(raw) + assert "area_0" in df.columns + assert "eccentricity_0" in df.columns + assert "low_level_features" not in df.columns + assert "Specular Intensity" in df.columns # RENAME_MAP still applied + assert df["area_0"].tolist() == [1.2, 1.5] + + +def test_to_dataframe_without_low_level_features(): + """No low_level_features key (flag off) -> no extra columns.""" + raw = { + "series_by_angle": [ + { + "angle": "0", + "series": [ + { + "frame_number": 0, + "relative_time_seconds": 0.0, + "unix_timestamp_ms": 0.0, + "specular_intensity": 5.0, + } + ], + } + ] + } + df = RHEEDProvider().to_dataframe(raw) + assert "area_0" not in df.columns + assert "low_level_features" not in df.columns