[WIP] kit _bluesky.py

CHANGELOG.md

@@ -52,6 +52,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/).
 ## [3.6.0]
 
 ### Added
+- new kit module `bluesky` for working with `bluesky-in-a-box` data / directories
 - `Data.norm_for_each`: easier norm-by-axis syntax
 - `kit.from_list_of_objects`: convenience method for grabbing an channel/variable/axis/etc. using multiple specifiers
 - `Data.get_var`: get variable object by index, name, or variable itself

WrightTools/kit/__init__.py

@@ -15,3 +15,5 @@
 from ._unicode import *
 from ._utilities import *
 from ._glob import *
+
+from . import _bluesky as bluesky

WrightTools/kit/_bluesky.py

@@ -0,0 +1,278 @@
+"""
+Helpers and containers for data structures in the Wright Group's Bluesky deployment
+https://github.com/wright-group/bluesky-in-a-box/
+"""
+
+import re
+import json
+import datetime
+import pathlib
+import logging
+from numpy import prod
+from typing import NamedTuple, Generator, Iterable
+
+from ._array import joint_shape
+from .._open import open as wt5_open
+
+logger = logging.getLogger(__name__)
+
+__folder_parts__ = [
+    r"(?P<date>\d\d\d\d-\d\d-\d\d)",
+    r"(?P<time>" + r"\d{5}" + ")",
+    r"(?P<plan>\w*)",
+    r"(?P<name>[\s\w\d.-=+-]*)",  # not great...just trying to include all valid characters
+    r"(?P<uid>\w{8})",
+]
+__folder_seed__ = " ".join(__folder_parts__)
+__datetime_seed__ = re.compile(" ".join(__folder_parts__[:3]))
+__fmtseed__ = "{date} {time:05d} {plan} {name} {uid}"
+
+
+class BlueskyFolder:
+    """container class for Bluesky acquisitions"""
+
+    def __init__(self, folder_path: str | pathlib.Path):
+        self.path = pathlib.Path(folder_path)
+        # DDK: better to extract the information from the data inside, rather than relying on the name
+        # self.info = parse_folder_contents(folder_path.name)
+        self.info = parse_folder_name(folder_path.name)
+        if self.info is None:
+            return
+
+        self._primary = None
+        self._baseline = None
+        self.logger = logging.getLogger(self.info.uid)
+        self.logger.info(self.info)
+
+    @property
+    def primary(self):
+        """open procedure based on plan"""
+        if self._primary is None:
+            # TODO: open procedure based on plan
+            if self.info.plan == "grid_scan_wp":
+                self._primary = wt5_open(self.path / "primary.wt5")
+            else:
+                raise NotImplementedError(f"plan {self.info.plan}")
+        return self._primary
+
+    @property
+    def baseline(self):
+        if self._baseline is None:
+            self._baseline = wt5_open(self.path / "primary.wt5")
+        return self._baseline
+
+    @property
+    def baseline_tree(self) -> str:
+        return (self.path / "baseline tree.txt").read_text()
+
+    @property
+    def primary_tree(self) -> str:
+        return (self.path / "primary tree.txt").read_text()
+
+    @property
+    def start(self) -> dict:
+        path = self.path / "bluesky_docs" / "start.json"
+        return json.load(path.open())
+
+    @property
+    def stop(self) -> dict:
+        path = self.path / "bluesky_docs" / "stop.json"
+        return json.load(path.open())
+
+    @property
+    def primary_descriptor(self) -> dict:
+        path = self.path / "bluesky_docs" / "primary descriptor.json"
+        return json.load(path.open())
+
+    @property
+    def baseline_descriptor(self) -> dict:
+        path = self.path / "bluesky_docs" / "baseline descriptor.json"
+        return json.load(path.open())
+
+
+def compress_sensor_mappings(data, sensor_name):
+    """
+    Compress sensor mappings that repetitively broadcast against e.g. all points in a grid_scan.
+    Note: do NOT use if the sensor axes are dynamic. This method performs no redundancy checks checks.
+
+    Parameters
+    ----------
+    data : WrightTools.Data
+        the data object
+
+    sensor_name : str
+        the name of the sensor to inspect for compression
+
+    Returns
+    -------
+    data : the input data object is returned
+    """
+    # determine which axes of ndim are entirely sensor mappings and not part of grid
+    is_sensor_var = lambda name: name.startswith(sensor_name)
+    sensor_vars = {vn for vn in filter(is_sensor_var, data.variable_names)}
+
+    sensor_shape = joint_shape(*[data[vi][:] for vi in sensor_vars])
+    non_sensor_shape = joint_shape(*[data[vi][:] for vi in set(data.variable_names) - sensor_vars])
+
+    exclusive = [prod(s) == max(s) for s in zip(sensor_shape, non_sensor_shape)]
+
+    if any(exclusive):
+        # replace the variables with a slice of the original
+        transform = list(data.axis_expressions)
+        slc = [slice(None) if exi else 0 for exi in exclusive]
+        for sn in sensor_vars:
+            new_shape = tuple(data[sn].shape[i] if exclusive[i] else 1 for i in range(data.ndim))
+            logger.info(f"reshaping {sn}: {data[sn].shape} -> {new_shape}")
+            # now rewrite the array
+            data.create_variable(
+                f"_{sn}", values=data[sn][*slc].reshape(new_shape), units=data[sn].units
+            )
+            # perhaps we should just create a new array and then overwrite
+            if sn in transform:
+                data.transform(*[ai for ai in transform if ai != sn])
+                data.remove_variable(sn)
+                data.rename_variables(**{f"_{sn}": sn}, implied=False, verbose=False)
+                data.transform(*transform)
+    return data
+
+
+def apply_points_axes(data):
+    """
+    Switch to reduced dimensional axes when available.
+    Useful for gridscan plans.
+    """
+    transform = [
+        f"{n}_points" if f"{n}_points" in data.variable_names else n for n in data.axis_names
+    ]
+    data.transform(*transform)
+    return data
+
+
+class FolderInfo(NamedTuple):
+    """Object representation of bluesky folder names"""
+
+    date: datetime.date
+    time: datetime.time
+    plan: str
+    name: str
+    uid: str
+
+    @property
+    def strf(self) -> str:
+        return __fmtseed__.format(
+            date=self.date.strftime("%Y-%m-%d"),
+            time=int(
+                datetime.timedelta(
+                    minutes=self.time.minute, seconds=self.time.second, hours=self.time.hour
+                ).total_seconds()
+            ),
+            plan=self.plan,
+            name=self.name,
+            uid=self.uid,
+        )
+
+
+def filter_bluesky(
+    items: Iterable[pathlib.Path], **bluesky_identifiers
+) -> Generator[pathlib.Path, None, None]:
+    """
+    Filter an iterator of folder names for bluesky folder pattern that match specified values.
+
+    Parameters
+    ----------
+
+    items: pathlikes
+        potential paths of bluesky folders
+
+    kwargs
+    ------
+
+    bluesky_identifiers
+        keys corresponding to FolderInfo properties (e.g. date, plan).
+
+    Yields
+    -------
+
+    pathlib.Path:
+        bluesky folders corresponding to full matches with the bluesky_identifiers
+
+    Examples
+    --------
+    ```
+    # match within a directory
+    spooky_folders = [
+        info for info in filter_bluesky(
+            data_folder.iterdir(),
+            date="2025-10-31"
+        )
+    ]
+    ```
+    """
+    for key in bluesky_identifiers.keys():
+        assert key in FolderInfo._fields
+
+    for item in map(pathlib.Path, items):
+        if (info := parse_folder_name(item.name)) is not None:
+            idict = info._asdict()
+            if all(idict[k] == v for k, v in bluesky_identifiers.items()):
+                yield item
+
+
+def bluesky_paths(dir: pathlib.Path, **bluesky_identifiers) -> list[pathlib.Path]:
+    """
+    walk a directory to find bluesky folder names that match the specified identifiers.
+
+    Parameters
+    ----------
+
+    dir: path-like
+        the directory to iterate through
+
+    kwargs
+    ------
+
+    bluesky_identifiers
+        keys corresponding to FolderInfo properties (e.g. date, plan).
+
+    Returns
+    -------
+    matches: list of BlueskyFolder objects
+        BlueskyFolders corresponding to full matches with the bluesky_identifiers
+    """
+
+    return sorted(
+        [dir / info.folder for info in filter_bluesky(dir.iterdir(), **bluesky_identifiers)]
+    )
+
+
+def parse_folder_name(folder: str) -> FolderInfo | None:
+    """
+    Convert a bluesky-formatted folder name into a structured dictonary-like format.
+
+    Parameters
+    ----------
+    folder : string
+        the folder name
+
+    Returns
+    -------
+    FolderInfo | None
+        if the name is parsed, returns a FolderInfo object.
+        otherwise, returns None.
+    """
+    out = None
+    if ((uid_match := re.fullmatch(r"(?P<uid>\w{8})", folder.split()[-1])) is not None) and (
+        (datetime_match := __datetime_seed__.match(folder)) is not None
+    ):
+        matchdict = uid_match.groupdict() | datetime_match.groupdict()
+        matchdict["name"] = " ".join(folder.split()[3:-1])
+        out = _to_object(matchdict)
+    return out
+
+
+def _to_object(mdict: dict) -> FolderInfo:
+    """convert re match dictionary into FolderInfo object"""
+    date = datetime.date.fromisoformat(mdict.pop("date"))
+    ts = int(mdict.pop("time"))  # total seconds since date start
+    time = datetime.time(hour=ts // 3600, minute=(ts % 3600) // 60, second=ts % 60)
+    return FolderInfo(date=date, time=time, **mdict)

tests/kit/bluesky.py

@@ -0,0 +1,46 @@
+import WrightTools as wt
+import logging
+from WrightTools import datasets
+
+logging.basicConfig(level=logging.INFO)
+
+
+def test_compress_sensor_mappings():
+    data = wt.open(datasets.wt5.v1p0p1_shg_andor)
+    wt.kit.bluesky.compress_sensor_mappings(data, "andor_camera")
+    assert data["andor_camera_x_index"].shape == (1, 1, 60)
+    assert data["andor_camera_y_index"].shape == (1, 60, 1)
+
+
+def test_folderinfo():
+    names = [
+        "2025-10-27 52433 count 2 beam PL spot2 post d7f183b5",
+        "2025-10-27 00622 grid_scan_wp spot 3 spectral 6a45457c",
+        "2025-10-27 00010 custom_plan +-5 --testing-- p3uyqhr7",
+    ]
+
+    fis = [wt.kit.bluesky.parse_folder_name(ni) for ni in names]
+
+    for name, fi in zip(names, fis):
+        assert fi is not None
+        assert fi.strf == name
+
+    assert fis[0].name == "2 beam PL spot2 post"
+    assert fis[1].name == "spot 3 spectral"
+    assert fis[0].plan == "count"
+    assert fis[1].plan == "grid_scan_wp"
+
+
+def test_filter():
+    name1 = "2025-10-27 52433 count 2 beam PL spot2 post d7f183b5"
+    name2 = "2025-10-27 54622 grid_scan_wp spot 3 spectral 6a45457c"
+
+    gridscans = [x for x in wt.kit.bluesky.filter_bluesky([name1, name2], plan="grid_scan_wp")]
+    assert len(gridscans) == 1
+    assert str(gridscans[0]) == name2
+
+
+if __name__ == "__main__":
+    test_compress_sensor_mappings()
+    test_folderinfo()
+    test_filter()