Merge pull request #62 from atomscale-ai/enhancement/add_similarity_traj_provider

Similarity trajectory provider

Author: munrojm

pyproject.toml

@@ -125,7 +125,6 @@ lint.extend-ignore = [
   "B028",  # No explicit stacklevel
   "EM101", # Exception must not use a string literal
   "EM102", # Exception must not use an f-string literal
-  "PD901", # Avoid using the generic variable name `df` for DataFrames
 ]
 lint.typing-modules = ["mypackage._compat.typing"]
 src = ["src"]

src/atomscale/results/__init__.py

@@ -5,6 +5,7 @@
 from .raman import RamanResult
 from .rheed_image import RHEEDImageCollection, RHEEDImageResult, _get_rheed_image_result
 from .rheed_video import RHEEDVideoResult
+from .similarity_trajectory import SimilarityTrajectoryResult
 from .unknown import UnknownResult
 from .xps import XPSResult
 
@@ -18,6 +19,7 @@
     "RHEEDImageResult",
     "RHEEDVideoResult",
     "RamanResult",
+    "SimilarityTrajectoryResult",
     "UnknownResult",
     "XPSResult",
     "_get_rheed_image_result",

src/atomscale/results/similarity_trajectory.py

@@ -0,0 +1,34 @@
+from __future__ import annotations
+
+from collections.abc import Sequence
+from uuid import UUID
+
+from monty.json import MSONable
+from pandas import DataFrame
+
+
+class SimilarityTrajectoryResult(MSONable):
+    def __init__(
+        self,
+        source_id: UUID | str,
+        workflow: str,
+        window_span: float,
+        timeseries_data: DataFrame,
+        source_data_ids: Sequence[UUID | str] | None = None,
+    ):
+        """Similarity trajectory result
+
+        Args:
+            source_id (UUID | str): Source ID for the similarity trajectory query.
+            workflow (str): Workflow name used for the similarity analysis.
+            window_span (float): Window span parameter used for the trajectory.
+            timeseries_data (DataFrame): Pandas DataFrame with similarity trajectory data.
+            source_data_ids (Sequence[UUID | str] | None): Sequence of source data IDs included in the trajectory.
+        """
+        self.source_id = source_id
+        self.workflow = workflow
+        self.window_span = window_span
+        self.timeseries_data = timeseries_data
+        self.source_data_ids: list[UUID | str] = (
+            list(source_data_ids) if source_data_ids else []
+        )

src/atomscale/timeseries/__init__.py

@@ -4,11 +4,13 @@
 from .provider import TimeseriesProvider
 from .registry import get_provider
 from .rheed import RHEEDProvider
+from .similarity import SimilarityTrajectoryProvider
 
 __all__ = [
     "MetrologyProvider",
     "OpticalProvider",
     "RHEEDProvider",
+    "SimilarityTrajectoryProvider",
     "TimeseriesProvider",
     "align_timeseries",
     "get_provider",

src/atomscale/timeseries/align.py

@@ -71,26 +71,26 @@ def _extract_timeseries(result):
     """Return (data_id, domain, df_with_timeindex) or None for non-timeseries."""
     if isinstance(result, RHEEDVideoResult):
         domain = "rheed"
-        df = result.timeseries_data
+        timeseries = result.timeseries_data
     elif isinstance(result, OpticalResult):
         domain = "optical"
-        df = result.timeseries_data
+        timeseries = result.timeseries_data
     elif isinstance(result, MetrologyResult):
         domain = "metrology"
-        df = result.timeseries_data
+        timeseries = result.timeseries_data
     else:
         return None
 
-    if df is None or df.empty:
+    if timeseries is None or timeseries.empty:
         return None
 
     # Build time index: prefer absolute epochs; fall back to upload_datetime + relative offsets.
     upload_dt = getattr(result, "upload_datetime", None)
 
-    time_index = _infer_absolute_time(df)
+    time_index = _infer_absolute_time(timeseries)
     if time_index is None and upload_dt is not None:
         base = pd.to_datetime(upload_dt, utc=True, errors="coerce")
-        rel = _infer_relative_time(df)
+        rel = _infer_relative_time(timeseries)
         if base is not pd.NaT and rel is not None:
             time_index = base + rel
 
@@ -101,7 +101,7 @@ def _extract_timeseries(result):
     if not valid_mask.any():
         return None
 
-    indexed = df.loc[valid_mask].copy(deep=False)
+    indexed = timeseries.loc[valid_mask].copy(deep=False)
     indexed.index = pd.Index(time_index[valid_mask], name="time")
     indexed = indexed.sort_index()
 
@@ -173,11 +173,11 @@ def align_timeseries(
         if not extracted:
             continue
 
-        data_id, domain, df = extracted
-        df = df.copy(deep=False)
-        df.columns = pd.MultiIndex.from_product([[data_id], [domain], df.columns])
-        frames.append(df)
-        indices.append(df.index)
+        data_id, domain, frame = extracted
+        frame = frame.copy(deep=False)
+        frame.columns = pd.MultiIndex.from_product([[data_id], [domain], frame.columns])
+        frames.append(frame)
+        indices.append(frame.index)
 
     if not frames:
         return pd.DataFrame()
@@ -211,28 +211,28 @@ def align_timeseries(
 
     # Merge compatible metrics across items: if multiple columns share (domain, metric)
     # and never conflict where they overlap, collapse into (shared, domain, metric).
-    def _merge_compatible_metrics(df: pd.DataFrame) -> pd.DataFrame:
-        if not isinstance(df.columns, pd.MultiIndex):
-            return df
-        domains = df.columns.get_level_values(1)
-        metrics = df.columns.get_level_values(2)
+    def _merge_compatible_metrics(data: pd.DataFrame) -> pd.DataFrame:
+        if not isinstance(data.columns, pd.MultiIndex):
+            return data
+        domains = data.columns.get_level_values(1)
+        metrics = data.columns.get_level_values(2)
         new_cols: dict = {}
         drop_cols: list = []
 
         for domain in domains.unique():
             for metric in metrics.unique():
                 cols = [
                     c
-                    for c in df.columns
+                    for c in data.columns
                     if c[1] == domain and c[2] == metric and c[0] != "shared"
                 ]
                 if len(cols) <= 1:
                     continue
 
-                merged = df[cols[0]]
+                merged = data[cols[0]]
                 conflict = False
                 for c in cols[1:]:
-                    other = df[c]
+                    other = data[c]
                     overlap_mask = merged.notna() & other.notna()
                     if (merged[overlap_mask] != other[overlap_mask]).any():
                         conflict = True
@@ -247,10 +247,10 @@ def _merge_compatible_metrics(df: pd.DataFrame) -> pd.DataFrame:
                 drop_cols.extend(cols)
 
         if new_cols:
-            df = df.drop(columns=drop_cols)
+            data = data.drop(columns=drop_cols)
             for col, series in new_cols.items():
-                df[col] = series
-            df = df.sort_index(axis=1)
-        return df
+                data[col] = series
+            data = data.sort_index(axis=1)
+        return data
 
     return _merge_compatible_metrics(aligned)

src/atomscale/timeseries/registry.py

@@ -4,11 +4,13 @@
 from .optical import OpticalProvider
 from .provider import TimeseriesProvider
 from .rheed import RHEEDProvider
+from .similarity import SimilarityTrajectoryProvider
 
 _PROVIDER_CLASSES: dict[str, type[TimeseriesProvider]] = {
     RHEEDProvider.TYPE: RHEEDProvider,
     OpticalProvider.TYPE: OpticalProvider,
     MetrologyProvider.TYPE: MetrologyProvider,
+    SimilarityTrajectoryProvider.TYPE: SimilarityTrajectoryProvider,
 }
 
 

src/atomscale/timeseries/similarity.py

@@ -0,0 +1,113 @@
+from __future__ import annotations
+
+from collections.abc import Mapping, Sequence
+from typing import Any
+from uuid import UUID
+
+from pandas import DataFrame, concat
+
+from atomscale.core import BaseClient
+from atomscale.results.similarity_trajectory import SimilarityTrajectoryResult
+from atomscale.timeseries.provider import TimeseriesProvider
+
+
+class SimilarityTrajectoryProvider(TimeseriesProvider[SimilarityTrajectoryResult]):
+    TYPE = "similarity_trajectory"
+
+    RENAME_MAP: Mapping[str, str] = {
+        "reference_id": "Reference ID",
+        "reference_item_name": "Reference Name",
+        "real_time_seconds": "Time",
+        "similarity_values": "Similarity",
+        "unix_times": "UNIX Timestamp",
+        "is_active": "Active",
+        "averaged_count": "Averaged Count",
+    }
+    INDEX_COLS: Sequence[str] = ["Reference ID", "Time"]
+
+    def fetch_raw(self, client: BaseClient, data_id: str, **kwargs: Any) -> Any:
+        """Fetch similarity trajectory data from the API.
+
+        Args:
+            client: The API client.
+            data_id: The source ID for the similarity query.
+            **kwargs: Must include 'workflow' (required). Optional parameters:
+                window_span, reference_ids, softmax_mode, reference_n_values.
+
+        Returns:
+            Raw API response payload.
+
+        Raises:
+            KeyError: If 'workflow' is not provided in kwargs.
+        """
+        workflow = kwargs.pop("workflow")
+        return client._get(
+            sub_url=f"similarity/{workflow}/{data_id}/trajectory/",
+            params=kwargs,
+        )
+
+    def to_dataframe(self, raw: Any) -> DataFrame:
+        if not raw:
+            return DataFrame(None)
+
+        trajectories = raw.get("trajectories", [])
+        if not trajectories:
+            return DataFrame(None)
+
+        frames: list[DataFrame] = []
+        for traj in trajectories:
+            ref_id = traj.get("reference_id")
+            ref_name = traj.get("reference_item_name")
+            similarity_values = traj.get("similarity_values", [])
+            real_time_seconds = traj.get("real_time_seconds", [])
+            unix_times = traj.get("unix_times", [])
+            is_active = traj.get("is_active")
+            averaged_count = traj.get("averaged_count")
+
+            if not similarity_values:
+                continue
+
+            # Build dataframe from columnar data
+            traj_df = DataFrame(
+                {
+                    "reference_id": ref_id,
+                    "reference_item_name": ref_name,
+                    "similarity_values": similarity_values,
+                    "real_time_seconds": real_time_seconds,
+                    "unix_times": unix_times,
+                    "is_active": is_active,
+                    "averaged_count": averaged_count,
+                }
+            )
+            frames.append(traj_df)
+
+        if not frames:
+            return DataFrame(None)
+
+        df_all = concat(frames, axis=0, ignore_index=True)
+        df_all = df_all.rename(columns=self.RENAME_MAP)
+
+        idx_cols = [c for c in self.INDEX_COLS if c in df_all.columns]
+        if idx_cols:
+            df_all = df_all.set_index(idx_cols)
+
+        return df_all
+
+    def build_result(
+        self,
+        client: BaseClient,  # noqa: ARG002
+        data_id: str,
+        data_type: str,  # noqa: ARG002
+        ts_df: DataFrame,
+        *,
+        workflow: str = "",
+        window_span: float = 0.0,
+        source_data_ids: Sequence[UUID | str] | None = None,
+    ) -> SimilarityTrajectoryResult:
+        return SimilarityTrajectoryResult(
+            source_id=data_id,
+            workflow=workflow,
+            window_span=window_span,
+            timeseries_data=ts_df,
+            source_data_ids=source_data_ids,
+        )

tests/conftest.py

@@ -27,3 +27,5 @@ class ResultIDs:
     metrology = ""
     photoluminescence = ""
     raman = ""
+    similarity_workflow = "rheed_stationary"
+    similarity_source_id = "bb3494b1-b5fb-4f3e-ac50-e4024f8aacf5"