Upgraded metadata lookup for backwards compatibility, added option to input metadata manually

src/PyHyperScattering/SST1RSoXSDB.py

@@ -20,8 +20,10 @@
     from httpx import HTTPStatusError
     import tiled
     import dask
-    try: from bluesky_tiled_plugins.queries import RawMongo, Key, FullText, Contains, Regex ## Intended to handle database navigation for 2025 onwards
-    except ImportError: from databroker.queries import RawMongo, Key, FullText, Contains, Regex
+    try: 
+        from bluesky_tiled_plugins.queries import RawMongo, Key, FullText, Contains, Regex # Bluesky changed the location of these queries in 2025, this is new location
+    except ImportError: 
+        from databroker.queries import RawMongo, Key, FullText, Contains, Regex # old location, in case dependencies aren't updated
 except Exception:
     print(
         "Imports failed.  Are you running on a machine with proper libraries for tiled, etc.?"
@@ -42,18 +44,35 @@ class SST1RSoXSDB:
     pix_size_1 = 0.06
     pix_size_2 = 0.06
 
+    # List of metadata key names used historically at SST1 RSoXS
     md_lookup = {
-        "sam_x": "RSoXS Sample Outboard-Inboard",
-        "sam_y": "RSoXS Sample Up-Down",
-        "sam_z": "RSoXS Sample Downstream-Upstream",
-        "sam_th": "RSoXS Sample Rotation",
-        "polarization": "en_polarization_setpoint",
-        "energy": "en_energy_setpoint",
-        "exposure": "RSoXS Shutter Opening Time (ms)",  # md['detector']+'_cam_acquire_time'
-    }
-    md_secondary_lookup = {
-        "energy": "en_monoen_setpoint",
+        "sam_x": ["solid_sample_x", # Started using ~March 2025
+                  "manipulator_x", # Started using ~January 2025
+                 "RSoXS Sample Outboard-Inboard",
+                 ],
+        "sam_y": ["solid_sample_y", # Started using ~March 2025
+                  "manipulator_y", # Started using ~January 2025
+                  "RSoXS Sample Up-Down",
+                 ],
+        "sam_z": ["solid_sample_z", # Started using ~March 2025
+                  "manipulator_z", # Started using ~January 2025
+                  "RSoXS Sample Downstream-Upstream",
+                 ],
+        "sam_th": ["solid_sample_r", # Started using ~March 2025
+                   "manipulator_r", # Started using ~January 2025
+                   "RSoXS Sample Rotation",
+                  ],
+        "polarization": ["en_polarization_setpoint",
+                        ],
+        "energy": ["en_energy_setpoint",
+                   "en_monoen_setpoint",
+                  ],
+        "exposure": ["RSoXS Shutter Opening Time (ms)", 
+                    ],
+        "time": ["time",
+                     ],
     }
+
 
     def __init__(
         self,
@@ -557,7 +576,7 @@ def loadRun(
         coords={},
         return_dataset=False,
         useMonitorShutterThinning=True,
-    ):
+     ):
         """
         Loads a run entry from a catalog result into a raw xarray.
 
@@ -587,6 +606,7 @@ def loadRun(
             )
 
         md = self.loadMd(run)
+
 
         monitors = self.loadMonitors(run)
 
@@ -649,10 +669,10 @@ def loadRun(
 
                 # next, construct the reverse lookup table - best mapping we can make of key to pyhyper word
                 # we start with the lookup table used by loadMd()
-                reverse_lut = {v: k for k, v in self.md_lookup.items()}
-                reverse_lut_secondary = {v: k for k, v in self.md_secondary_lookup.items()}
-                reverse_lut.update(reverse_lut_secondary)
-
+                reverse_lut = {md_key_beamline: md_key_PHS 
+                               for md_key_PHS, md_key_beamline_list in self.md_lookup.items() 
+                               for md_key_beamline in md_key_beamline_list}
+# this creates a reverse mapping where the keys are the possible names in 'beamline language' of a parameter, and the values are the name in 'PyHyper language'.  This exploits the fact that dicts are ordered to provide rank-sorted mapping of parameters.
                 # here, we broaden the table to make a value that default sources from '_setpoint' actually match on either
                 # the bare value or the readback value.
                 reverse_lut_adds = {}
@@ -670,32 +690,6 @@ def loadRun(
                         pyhyper_axes_to_use.append(x)
                 dims = pyhyper_axes_to_use
 
-        """
-        elif dims == None:
-            # use the dim tols to define the dimensions
-            # dims = []
-            # dim_tols = {'en_polarization': 0.5, 'sam_x': 0.05, 'sam_y':0.05, 'en_energy':0.05, 'exposure': 1., 'sam_th': 0.05} # set the amount dims are allowed to change; could make this user-chosen in the future
-            dims = ['en_energy','time'] # I think this always needs to be an axis due to the way that the integrator is set up
-            dim_tols = {'en_polarization': 0.5, 'sam_x': 0.05, 'sam_y':0.05, 'exposure': 1., 'sam_th': 0.05} # set the amount dims are allowed to change; could make this user-chosen in the future
-            if 'spiral' in md['start']['plan_name']:
-                dims = ['energy','time']
-                dim_tols = {'polarization': 0.5, 'sam_x': 0.05, 'sam_y':0.05, 'exposure': 1., 'sam_th': 0.05} # set the amount dims are allowed to change; could make this user-chosen in the future
-            for k in dim_tols.keys():
-                dim = md[k]
-                dim_std = np.std(dim)
-                if dim_std > dim_tols[k]:
-                    dims.append(k)
-        else: # if the user has already specified the dims; user may frequently specify 'energy' or 'polarization', so just changing that so it's readable to access correct metadata (without en_, they are just setpoints)
-            dims = dims
-            for i in range(0,len(dims)):
-                if dims[i] == 'polarization':
-                    dims[i] = 'en_polarization'
-                if dims[i] == 'energy':
-                    dims[i] = 'en_energy'
-            if len(dims) == 0:
-                raise NotImplementedError('You have not entered any dimensions; please enter at least one, or use None rather than an empty list')
-        """
-
         data = run["primary"]["data"][md["detector"] + "_image"]
         if isinstance(data,tiled.client.array.ArrayClient):
             data = run["primary"]["data"].read()[md["detector"] + "_image"]
@@ -747,7 +741,6 @@ def subtract_dark(img, pedestal=100, darks=None):
         # handle the edge case of a partly-finished scan
         if len(index) != len(data["time"]):
             index = index[: len(data["time"])]
-        actual_exposure = md["exposure"] * len(data.dim_0)
         mindex_coords = xr.Coordinates.from_pandas_multiindex(index, 'system')
         retxr = (
             data.sum("dim_0")
@@ -783,9 +776,13 @@ def subtract_dark(img, pedestal=100, darks=None):
             )
         retxr.attrs.update(md)
 
-        retxr.attrs["exposure"] = (
-            len(data.dim_0) * retxr.attrs["exposure"]
-        )  # patch for multi exposures
+        exposure = retxr.attrs.get("exposure")
+
+        if exposure is not None:
+            retxr.attrs["exposure"] = (len(data.dim_0) * exposure)
+        else:
+            retxr.attrs["exposure"] = None  # or 0, or skip setting it  # patch for multi exposures
+
         # now do corrections:
         frozen_attrs = retxr.attrs
         if self.corr_mode == "i0":
@@ -1011,7 +1008,7 @@ def loadMd(self, run):
                 md["beamcenter_y"] = run.start["RSoXS_SAXS_BCY"]
                 md["sdd"] = run.start["RSoXS_SAXS_SDD"]
 
-        elif start["RSoXS_Config"] == "WAXS":
+        elif "WAXS" in start["RSoXS_Config"]:
             md["rsoxs_config"] = "waxs"
             if (meas_time > datetime.datetime(2020, 11, 16)) and (
                 meas_time < datetime.datetime(2021, 1, 15)
@@ -1056,57 +1053,46 @@ def loadMd(self, run):
             primary = run["primary"]["data"]
         except (KeyError, HTTPStatusError):
             raise Exception(
-                "No primary stream --> probably you caught run before image was written.  Try"
-                " again."
+                "No primary stream --> probably you caught run before image was written.  Try again."
             )
 
         md_lookup = copy.deepcopy(self.md_lookup)
-        md_secondary_lookup = copy.deepcopy(self.md_secondary_lookup)
-
+        # Add additional metadata not included in lookup dictionary
+        md_key_names_beamline = []
+
+        for key in md_lookup.keys(): # Make a single list with all historical keys, in the order to check for them.
+            md_key_names_beamline = md_key_names_beamline + md_lookup[key] 
+
         for key in primary.keys():
-            if key not in md_lookup.values():
+            if key not in md_key_names_beamline:
                 if "_image" not in key:
-                    md_lookup[key] = key
-
-        for phs, rsoxs in md_lookup.items():
-            try:
-                md[phs] = primary[rsoxs].read()
-                # print(f'Loading from primary: {phs}, value {primary[rsoxs].values}')
-            except (KeyError, HTTPStatusError):
-                try:
-                    blval = baseline[rsoxs].__array__()
-                    md[phs] = blval.mean().round(4)
-                    if blval.var() > 1e-4*abs(blval.mean()):
-                        warnings.warn(
-                            (
-                                f"{phs} changed during scan: {blval}."
-                            ),
-                            stacklevel=2,
-                        )
+                    md_lookup[key] = [key]
+
+        # Find metadata from Tiled primary and store in PyHyperScattering metadata dictionary
+        for key_name_PHS, key_names_beamline in md_lookup.items(): # first iterate over PyHyper 'words' and ordered lists to check
+            for key_name_beamline in key_names_beamline: # next, go through that list in order
+                if (key_name_PHS in md
+                    and md[key_name_PHS] is not None): 
+                    break # If the pyhyper key is already filled in, no need to try other beamline keys, so stop processing this PyHyper key.
+                try: 
+                    md[key_name_PHS] = primary[key_name_beamline].read() # first try finding metadata in primary stream
                 except (KeyError, HTTPStatusError):
                     try:
-                        md[phs] = primary[md_secondary_lookup[phs]].read()
-                    except (KeyError, HTTPStatusError):
-                        try:
-                            blval = baseline[md_secondary_lookup[phs]].__array__()
-                            md[phs] = blval.mean().round(4)
-                            if blval.var() > 1e-4*abs(blval.mean()):
-                                warnings.warn(
-                                    (
-                                        f"{phs} changed during scan: {blval}."
-                                    ),
-                                    stacklevel=2,
-                                )
-                        except (KeyError, HTTPStatusError):
-                            warnings.warn(
-                                (
-                                    f"Could not find {rsoxs} in either baseline or primary while"
-                                    f" looking for {phs}.   Setting to None."
-                                ),
-                                stacklevel=2,
-                            )
-                            md[phs] = None
-        md["epoch"] = md["meas_time"].timestamp()
+                        baseline_value = baseline[key_name_beamline] # Next, try finding metadata in baseline
+
+                        if isinstance(baseline_value, (tiled.client.array.ArrayClient, tiled.client.array.DaskArrayClient)): 
+                            baseline_value = baseline_value.read() # For tiled_client.array data types, need to use .read() to get the values
+
+                        md[key_name_PHS] = baseline_value.mean().round(4) # Rounded for stacking purposes, to avoid slightly different values when not meaningful
+
+                        if baseline_value.var() > 1e-4*abs(baseline_value.mean()):                             
+                            warnings.warn(f"{key_name_PHS} changed during scan: {baseline_value}.",stacklevel=2)
+                    except (KeyError, HTTPStatusError): 
+                        md[key_name_PHS] = None
+            if md[key_name_PHS] is None:
+                warnings.warn(f"Could not find any of {key_names_beamline} in either baseline or primary. Setting {key_name_PHS} to None.",stacklevel=2)
+
+        md["epoch"] = md["meas_time"].timestamp() # Epoch = the time the entire run started, used for multi-scan stacking
 
         # looking at exposure tests in the stream and issuing warnings
         if "Wide Angle CCD Detector_under_exposed" in md:
@@ -1147,7 +1133,6 @@ def loadMd(self, run):
             warnings.warn(
                 "'Wide Angle CCD Detector_saturated' not found in stream."
             )
-        md["epoch"] = md["meas_time"].timestamp()
 
         try:
             md["wavelength"] = 1.239842e-6 / md["energy"]

tests/test_SST1DBLoader.py

@@ -84,6 +84,40 @@ def test_SST1DB_load_SingleEnergy2Polarizations_scan_explicit_dims(sstdb):
     assert 'polarization' in run.indexes
 
 
+
+@must_have_tiled
+def test_SST1DB_load_energy_scan_20241209(sstdb):
+    run = sstdb.loadRun(91175).unstack('system')
+    assert type(run) == xr.DataArray
+    assert 'energy' in run.indexes
+    assert 'polarization' in run.indexes
+
+@must_have_tiled
+def test_SST1DB_load_energy_scan_20250213(sstdb):
+    run = sstdb.loadRun(92202).unstack('system')
+    assert type(run) == xr.DataArray
+    assert 'energy' in run.indexes
+
+@must_have_tiled
+def test_SST1DB_load_spiral_scan_20250221(sstdb):
+    run = sstdb.loadRun(92770).unstack('system')
+    assert type(run) == xr.DataArray
+    assert 'sam_x' in run.indexes
+    assert 'sam_y' in run.indexes
+
+@must_have_tiled
+def test_SST1DB_load_count_scan_20250222(sstdb):
+    run = sstdb.loadRun(92849,dims=['time','polarization']).unstack('system')
+    assert type(run) == xr.DataArray
+    assert 'time' in run.indexes
+    assert 'polarization' in run.indexes
+
+@must_have_tiled
+def test_SST1DB_load_energy_scan_20250223(sstdb):
+    run = sstdb.loadRun(93065).unstack('system')
+    assert type(run) == xr.DataArray
+    assert 'energy' in run.indexes
+
 @must_have_tiled
 def test_SST1DB_exposurewarnings(sstdb):
     with pytest.warns(UserWarning, match="Wide Angle CCD Detector is reported as underexposed"):