From databroker (#999)

* Initial implementation of from_databroker

* Put from_databroker in docs toctree

* docstring

* Changelog

* dev requirements

* correct uid source

* `__all__`

CHANGELOG.md

@@ -5,6 +5,9 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/).
 
 ## [Unreleased]
 
+### Added
+- `from_databroker` method to import Data objects from databroker catalogs
+
 ### Fixed
 - Avoid passing both `vmin/vmax` and `norm` to `pcolor*` methods
 

WrightTools/data/__init__.py

@@ -12,6 +12,7 @@
 from ._brunold import *
 from ._colors import *
 from ._data import *
+from ._databroker import *
 from ._jasco import *
 from ._kent import *
 from ._aramis import *
@@ -32,6 +33,7 @@
     # From methods in alphabetic order
     "from_BrunoldrRaman",
     "from_COLORS",
+    "from_databroker",
     "from_JASCO",
     "from_KENT",
     "from_Aramis",

WrightTools/data/_databroker.py

@@ -0,0 +1,39 @@
+__all__ = ["from_databroker"]
+
+from ._data import Data
+
+
+def from_databroker(run, dataset="primary"):
+    """Import a dataset from a bluesky run into the WrightTools Data format.
+
+    Parameters
+    ----------
+    run: BlueskyRun
+        The bluesky run as returned by e.g. catalog["<uid>"]
+    dataset: str
+        The string identifier of the stream to import from the bluesky run.
+        By default "primary" is used, but e.g. "baseline" is also common
+    """
+    describe = run.describe()
+    md = describe["metadata"]
+    start = md["start"]
+    ds = run[dataset].read()
+    shape = start.get("shape", (len(ds.time),))
+
+    detectors = start.get("detectors", [])
+
+    data = Data(name=start["uid"])
+    for var in ds:
+        if var == "uid":
+            continue
+        if var.endswith("_busy"):
+            continue
+        if any(var.startswith(d) for d in detectors):
+            data.create_channel(var, values=ds[var].data.reshape(shape))
+        else:
+            # TODO units, once they are in the dataset metadata
+            data.create_variable(var, values=ds[var].data.reshape(shape))
+
+    transform = [x[0] for x, ds_name in start["hints"]["dimensions"] if ds_name == dataset]
+    data.transform(*transform)
+    return data

docs/api/WrightTools.data.rst

@@ -16,6 +16,7 @@ WrightTools\.data package
    join
    from_BrunoldrRaman
    from_COLORS
+   from_databroker
    from_JASCO
    from_KENT
    from_Aramis

setup.py

@@ -46,7 +46,16 @@ def read(fname):
     ],
     extras_require={
         "docs": docs_require,
-        "dev": ["black", "pre-commit", "pydocstyle", "pytest", "pytest-cov"] + docs_require,
+        "dev": [
+            "black",
+            "pre-commit",
+            "pydocstyle",
+            "pytest",
+            "pytest-cov",
+            "databroker>=1.2",
+            "msgpack",
+        ]
+        + docs_require,
     },
     version=version,
     description="Tools for loading, processing, and plotting multidimensional spectroscopy data.",

tests/data/from_databroker.py

@@ -0,0 +1,27 @@
+import pathlib
+
+from databroker.v2 import temp
+import msgpack
+
+import WrightTools as wt
+
+catalog = temp()
+
+__here__ = pathlib.Path(__file__).parent
+
+for child in (__here__ / "bluesky_data").iterdir():
+    with child.open("rb") as f:
+        unpack = msgpack.Unpacker(f)
+        for item in unpack:
+            catalog.v1.insert(*item)
+
+
+def test_2d_data():
+    run = catalog["3dbdd402-434b-4aac-b004-447a2f026d73"]
+    data = wt.data.from_databroker(run)
+    assert data.shape == (10, 11)
+    assert "d1_readback" in data
+    assert "d1_setpoint" in data
+    assert data.channel_names == ("daq_random_walk",)
+    assert data.d1_readback.min() == 0
+    assert data.d1_readback.max() == 1