Updated to export TES data and grab both element and edge if they are…

… available

end_of_run_workflow.py

@@ -15,8 +15,8 @@ def log_completion():
 def end_of_run_workflow(stop_doc):
     uid = stop_doc["run_start"]
     general_data_validation(uid)
+    process_tes(uid)
     # Here is where exporters could be added
     general_data_export(uid)
 
-    process_tes(uid)
     log_completion()

export_to_xdi.py

@@ -4,6 +4,7 @@
 from datetime import datetime
 from prefect import get_run_logger
 
+
 def get_config(config, keys, default=None):
     item = get_with_fallbacks(config, keys)
     try:
@@ -34,16 +35,18 @@ def get_xdi_run_header(run):
     metadata["Scan.type"] = run.start.get("scantype", "unknown")
     metadata["Scan.motors"] = run.start.get("motors", ["time"])[0]
 
-    metadata["Element.symbol"] = run.start.get("edge", "")
+    metadata["Element.symbol"] = run.start.get("element", "")
     metadata["Element.edge"] = run.start.get("edge", "")
+    # This is just a kludge for re-export of old data where we used edge, not element in run.start
+    if metadata["Element.symbols"] == "" and metadata["Element.edge"] != "":
+        metadata["Element.symbols"] = metadata["Element.edge"]
 
     proposal = run.start.get("proposal", {})
     metadata["Proposal.id"] = proposal.get("proposal_id", "")
     metadata["Proposal.pi"] = proposal.get("pi_name", "")
     metadata["Proposal.cycle"] = run.start.get("cycle", "")
     metadata["Proposal.start"] = run.start.get("start_datetime", "")
 
-
     metadata["Motors.exslit"] = float(
         get_with_fallbacks(baseline, "eslit", "Exit Slit of Mono Vertical Gap", default=[0])[0]
     )
@@ -100,7 +103,7 @@ def exportToXDI(
     None
     """
     logger = get_run_logger()
-    
+
     metadata = get_xdi_run_header(run)
     metadata.update(headerUpdates)
     logger.info("Got XDI Metadata")
@@ -116,7 +119,10 @@ def exportToXDI(
     if verbose:
         print(f"Exporting to {filename}")
 
-    columns, run_data = get_run_data(run, omit=["tes_scan_point_start", "tes_scan_point_end"])
+    columns, run_data, tes_rois = get_run_data(run, omit=["tes_scan_point_start", "tes_scan_point_end"])
+    for c in columns:
+        if c in tes_rois:
+            metadata[f"{c}.roi"] = "{:.2f} {:.2f}".format(*tes_rois[c])
     logger.info("Got XDI Data")
     # Rename energy columns if present
     if "en_energy" in columns:
@@ -139,4 +145,3 @@ def exportToXDI(
         f.write(header_string)
         f.write("\n")
         np.savetxt(f, data, fmt="%8.8e", delimiter=" ")
-

export_tools.py

@@ -1,6 +1,8 @@
 import datetime
 import numpy as np
+from os.path import join
 from tiled.client import from_uri
+from autoprocess.statelessAnalysis import get_tes_data
 
 
 def initialize_tiled_client(beamline_acronym):
@@ -37,7 +39,7 @@ def get_with_fallbacks(thing, *possible_names, default=None):
 
 
 def get_header_and_data(run):
-    cols, run_data = get_run_data(run)
+    cols, run_data, rois = get_run_data(run)
     header = get_run_header(run)
     header["channelinfo"]["cols"] = cols
     data = np.vstack(run_data).T
@@ -120,9 +122,16 @@ def get_run_data(run, omit=[]):
             continue
         usekeys.append(key)
     data = run.primary.data.read(usekeys)
+    # Add a try-except here after testing
+    save_directory = join(get_proposal_path(run), "ucal_processing")
+
+    rois, tes_data = get_tes_data(run)
     for key in usekeys:
         if len(data[key].shape) == 1:
-            datadict[key] = data[key].data
+            if key in tes_data:
+                datadict[key] = tes_data[key]
+            else:
+                datadict[key] = data[key].data
     if "seconds" not in datadict:
         datadict["seconds"] = np.zeros_like(datadict[key]) + exposure
     for k in natural_order:
@@ -132,7 +141,7 @@ def get_run_data(run, omit=[]):
         if k not in columns and k not in omit:
             columns.append(k)
     data = [datadict[k] for k in columns]
-    return columns, data
+    return columns, data, rois
 
 
 def add_comment_to_lines(multiline_string, comment_char="#"):

process_tes.py

@@ -1,29 +1,34 @@
 from prefect import flow, get_run_logger, task
 from export_tools import initialize_tiled_client, get_proposal_path
 from autoprocess.statelessAnalysis import get_data, handle_calibration_run, handle_science_run
-from autoprocess.utils import get_filename
-from os.path import join
+from autoprocess.utils import get_filename, get_savename
+from os.path import join, exists
 
 
 @flow(log_prints=True)
-def process_tes(uid, beamline_acronym="ucal"):
+def process_tes(uid, beamline_acronym="ucal", reprocess=False):
     catalog = initialize_tiled_client(beamline_acronym)
     logger = get_run_logger()
 
     logger.info(f"In TES Exporter for {uid}")
 
     run = catalog[uid]
+    save_directory = join(get_proposal_path(run), "ucal_processing")
 
     # Check if run contains TES data
     if "tes" not in run.start.get("detectors", []):
         logger.info("No TES in run, skipping!")
         return False
+    if not reprocess:
+        savename = get_savename(run, save_directory)
+        if exists(savename):
+            logger.info(f"TES Already processed to {savename}, will not reprocess")
+            return True
 
     logger.info(f"Loading TES Data from {get_filename(run)}")
     # Get data files
     data = get_data(run)
     data.verbose = False
-    save_directory = join(get_proposal_path(run), "ucal_processing")
     logger.info("TES Data loaded")
     # Handle calibration runs first
     if run.start.get("scantype", "") == "calibration":