Add estimate transmission to udc strategy pipeline (#367)

pjnaughton · pblowey · web-flow · commit dc931246b9ce · 2026-05-26T14:01:18.000+01:00
Estimate transmission wrapper was added in #366. This PR updates the strategy service to take the suggested maximum transmission from the estimate transmission wrapper and apply it to the UDC strategy. Some other fixes/features added to the UDC strategy pipeline are also applied to the strategy service in this PR: * Include limits from recipe_config.yaml, which are limits that beamline staff can set for certain recipes * Only apply resolution based scale to dose (wavelength scale was erroneously being applied to dose). * Populate the database differently for beamlines that work with dose: * For beamlines that work with dose, populate dosetotal and transmission columns but not exposure time column in ScreeningStrategySubWedge table. Transmission column is populated with recommended maximum transmission (from estimate transmission). * For beamlines that don't work with dose, populate transmission with scaled value along with scaled exposure time but do not populate dosetotal column in ScreeningStrategySubWedge table. This PR also adds a check in the trigger service to ensure that the strategy pipeline is only triggered once for each dcid, running on whichever upstream pipeline finishes first. The trigger function is also updated to only allow the pipeline to run on I04 for initial rollout of the pipeline. Also removes the now obsolete xia2-overload wrapper. --------- Co-authored-by: Phil Blowey <philip.blowey@diamond.ac.uk>
diff --git a/src/dlstbx/services/strategy.py b/src/dlstbx/services/strategy.py
@@ -9,6 +9,8 @@
 
 from dlstbx.util import ChainMapWithReplacement
 
+type Limits = tuple[float, float]
+
 
 def scale_parameter(
     value: float, scale_factor: float, limits: tuple[float, float] | None = None
@@ -72,12 +74,46 @@ class AgamemnonParameters(BaseModel):
     wavelength: float = Field(gt=0)
 
 
+class AgamemnonLimits(BaseModel):
+    transmission: Limits
+    exposure_time: Limits
+    dose: Limits
+
+
 def parse_agamemnon_recipe(recipe_path: Path) -> list[AgamemnonParameters]:
     with open(recipe_path, "r") as f:
         recipe = yaml.safe_load(f)
     return [AgamemnonParameters(**step) for step in recipe]
 
 
+def parse_agamemnon_config(recipe_config_path: Path) -> dict[str, AgamemnonLimits]:
+    inf = float("inf")
+    default_limits = AgamemnonLimits(
+        transmission=(-inf, inf), exposure_time=(-inf, inf), dose=(-inf, inf)
+    )
+
+    def parse_limits(limits):
+        limits = limits[0]
+        params = (limits.get("min", -inf), limits.get("max", inf))
+        return params
+
+    with open(recipe_config_path, "r") as f:
+        recipe_conf = yaml.safe_load(f)
+
+    agamemnon_limits = {}
+    for strategy, parameters in recipe_conf.items():
+        if strategy == "globals":
+            continue
+        if not isinstance(parameters, dict):
+            continue
+
+        agamemnon_limits[strategy] = default_limits.model_copy(
+            update={param: parse_limits(limits) for param, limits in parameters.items()}
+        )
+
+    return agamemnon_limits
+
+
 def parse_config_file(config_file: Path) -> dict:
     config = {}
 
@@ -152,9 +188,11 @@ def generate_strategy(
         """Generate a strategy from the results of an upstream pipeline"""
         self.log.info("Received strategy request, generating strategy")
 
+        recipe_params = rw.recipe_step["parameters"]
         parameters = ChainMapWithReplacement(
             message.get("parameters", {}) if isinstance(message, dict) else {},
-            rw.recipe_step["parameters"].get("ispyb_parameters", {}),
+            recipe_params.get("ispyb_parameters", {}),
+            recipe_params,
             substitutions=rw.environment,
         )
         self.log.info(f"Received parameters for strategy generation:\n{parameters}")
@@ -177,6 +215,7 @@ def generate_strategy(
             if isinstance(parameters["resolution"], list)
             else float(parameters["resolution"])
         )
+        dc_transmission = float(parameters.get("transmission", 100)) / 100
         resolution_offset = 0.5
         min_resolution = 0.9
         resolution = max((resolution_estimate) - resolution_offset, min_resolution)
@@ -193,31 +232,24 @@ def generate_strategy(
             )
         beamline_config = parse_config_file(beamline_config_file)
 
-        scaled_transmission = parameters.get("scaled_transmission", 1.0) 
-
-        transmission_limits = (
-            get_beamline_param(beamline_config, ("gda.mx.udc.minTransmission",), 0.0),
-            min(get_beamline_param(beamline_config, ("gda.mx.udc.maxTransmission",), 1.0), 
-                scaled_transmission)
+        recommended_max_transmission = parameters.get("scaled_transmission", 1.0)
+        base_recipe_home = Path(f"/dls_sw/{beamline}/etc/agamemnon-recipes")
+        agamemnon_recipe_config = base_recipe_home / "recipe_config.yaml"
+        agamemnon_limits: dict[str, AgamemnonLimits] = parse_agamemnon_config(
+            agamemnon_recipe_config
         )
-        exposure_time_limits = (
-            get_beamline_param(
-                beamline_config,
-                ("gda.mx.udc.minExposureTime", "gda.exptTableModel.minExposureTime"),
-                0.002,
-            ),
-            get_beamline_param(
-                beamline_config,
-                ("gda.mx.udc.maxExposureTime", "gda.exptTableModel.maxExposureTime"),
-                float("inf"),
-            ),
+        recipes = (
+            ("OSC.yaml", "OSC", "Native"),
+            ("Ligand binding.yaml", "Ligand binding", "Ligand"),
+            ("SAD.yaml", "SAD", "Phasing"),
         )
-
-        recipes = {"OSC.yaml": "Native", "Ligand binding.yaml": "Ligand"}
         ispyb_command_list = []
 
-        for recipe, recipe_alias in recipes.items():
-            recipe_path = Path(f"/dls_sw/{beamline}/etc/agamemnon-recipes/{recipe}")
+        beamline_wants_dose_displayed = beamline in [
+            "i04",
+        ]
+        for recipe_file, recipe_name, recipe_alias in recipes:
+            recipe_path = base_recipe_home / recipe_file
             if not recipe_path.is_file():
                 self.log.error(
                     f"Recipe file {recipe_path} not found, terminating strategy generation"
@@ -250,41 +282,19 @@ def generate_strategy(
             }
             ispyb_command_list.append(d)
 
+            agamemnon_recipe_limits = agamemnon_limits[recipe_name]
             for n_step, recipe_step in enumerate(recipe_steps, start=1):
                 scale = 1.0
                 default_wavelength = recipe_step.wavelength
-                scale *= get_wavelength_scale(wavelength, default_wavelength)
                 scale *= get_resolution_scale(resolution)
-
-                dose, _ = scale_parameter(recipe_step.dose, scale)
-
+                dose, _ = scale_parameter(
+                    recipe_step.dose, scale, limits=agamemnon_recipe_limits.dose
+                )
                 rotation_axis = recipe_step.scan_axis
                 rotation_start = recipe_step.__getattribute__(f"{rotation_axis}_start")
                 rotation_increment = recipe_step.__getattribute__(
                     f"{rotation_axis}_increment"
                 )
-                transmission = recipe_step.transmission
-                exposure_time = recipe_step.exposure_time
-
-                # Runs twice to ensure that limits are applied correctly to both parameters, as they are interdependent - is this necessary?
-                for _ in range(2):
-                    if scale > 1.0:
-                        transmission, scale = scale_parameter(
-                            transmission, scale, limits=transmission_limits
-                        )
-                        exposure_time, scale = scale_parameter(
-                            exposure_time, scale, limits=exposure_time_limits
-                        )
-                    else:
-                        exposure_time, scale = scale_parameter(
-                            exposure_time, scale, limits=exposure_time_limits
-                        )
-                        transmission, scale = scale_parameter(
-                            transmission, scale, limits=transmission_limits
-                        )
-                    self.log.debug(
-                        f"Exposure time scaled to {exposure_time:.3f} s, transmission scaled to {transmission:.3f}, scale factor now {scale:.3f}"
-                    )
 
                 # Step 3: Store screeningStrategyWedge results, linked to the screeningStrategyId
                 #         Keep the screeningStrategyWedgeId
@@ -303,31 +313,100 @@ def generate_strategy(
                 }
                 ispyb_command_list.append(d)
 
-                # Convert transmission to percentage for ISPyB
-                transmission_pct = transmission * 100
-
                 axis_end = (
                     rotation_start + rotation_increment * recipe_step.number_of_images
                 )
 
-                # Step 4: Store second screeningStrategySubWedge results, linked to the screeningStrategyWedgeId
-                #         Keep the screeningStrategyWedgeId
-                d = {
+                # Partial command for the ScreeningSubWedge table is updated depending on what fields the beamline wants to be shown in Synchweb
+                screening_sub_wedge_command = {
                     "subwedgenumber": 1,
                     "rotationaxis": recipe_step.scan_axis,
                     "axisstart": rotation_start,
                     "axisend": axis_end,
-                    "exposuretime": exposure_time,
-                    "transmission": transmission_pct,
                     "oscillationrange": rotation_increment,
                     "noimages": recipe_step.number_of_images,
                     "resolution": resolution,
                     "ispyb_command": "insert_screening_strategy_sub_wedge",
                     "screening_strategy_wedge_id": f"$ispyb_screening_strategy_wedge_id_{n_step}",
                     "store_result": f"ispyb_screening_strategy_sub_wedge_id_{n_step}",
                 }
-                ispyb_command_list.append(d)
+                transmission_limits = (
+                    max(
+                        get_beamline_param(
+                            beamline_config, ("gda.mx.udc.minTransmission",), 0.0
+                        ),
+                        agamemnon_recipe_limits.transmission[0],
+                    ),
+                    min(
+                        get_beamline_param(
+                            beamline_config, ("gda.mx.udc.maxTransmission",), 1.0
+                        ),
+                        recommended_max_transmission,
+                        agamemnon_recipe_limits.transmission[1],
+                    ),
+                )
+                if beamline_wants_dose_displayed:
+                    screening_sub_wedge_command.update(
+                        dosetotal=dose, transmission=(transmission_limits[1] * 100)
+                    )
+                    ispyb_command_list.append(screening_sub_wedge_command)
+                    continue
+
+                scale *= get_wavelength_scale(wavelength, default_wavelength)
+                transmission = recipe_step.transmission
+                exposure_time = recipe_step.exposure_time
+
+                exposure_time_limits = (
+                    max(
+                        get_beamline_param(
+                            beamline_config,
+                            (
+                                "gda.mx.udc.minExposureTime",
+                                "gda.exptTableModel.minExposureTime",
+                            ),
+                            0.002,
+                        ),
+                        agamemnon_recipe_limits.exposure_time[0],
+                    ),
+                    min(
+                        get_beamline_param(
+                            beamline_config,
+                            (
+                                "gda.mx.udc.maxExposureTime",
+                                "gda.exptTableModel.maxExposureTime",
+                            ),
+                            float("inf"),
+                        ),
+                        agamemnon_recipe_limits.exposure_time[1],
+                    ),
+                )
+                # Runs twice to ensure that limits are applied correctly to both parameters, as they are interdependent - is this necessary?
+                for _ in range(2):
+                    if scale > 1.0:
+                        transmission, scale = scale_parameter(
+                            transmission, scale, limits=transmission_limits
+                        )
+                        exposure_time, scale = scale_parameter(
+                            exposure_time, scale, limits=exposure_time_limits
+                        )
+                    else:
+                        exposure_time, scale = scale_parameter(
+                            exposure_time, scale, limits=exposure_time_limits
+                        )
+                        transmission, scale = scale_parameter(
+                            transmission, scale, limits=transmission_limits
+                        )
+                    self.log.debug(
+                        f"Exposure time scaled to {exposure_time:.3f} s, transmission scaled to {transmission:.3f}, scale factor now {scale:.3f}"
+                    )
+
+                # Convert transmission to percentage for ISPyB
+                relative_transmission_pct = transmission / dc_transmission * 100
 
+                screening_sub_wedge_command.update(
+                    exposuretime=exposure_time, transmission=relative_transmission_pct
+                )
+                ispyb_command_list.append(screening_sub_wedge_command)
         # Send results onwards
         rw.set_default_channel("ispyb")
         rw.send_to("ispyb", {"ispyb_command_list": ispyb_command_list}, transaction=txn)
diff --git a/src/dlstbx/services/trigger.py b/src/dlstbx/services/trigger.py
@@ -2905,12 +2905,29 @@ def trigger_strategy(
             )
             return {"success": True}
 
-        if parameters.beamline not in ["i03", "i04", "i04-1"]:
+        if parameters.beamline not in ["i04"]:
             self.log.info(
                 f"Skipping strategy trigger: beamline {parameters.beamline} not supported"
             )
             return {"success": True}
 
+        udc_strategy_previously_triggered = (
+            session.query(AutoProcProgram.processingPrograms)
+            .join(
+                ProcessingJob,
+                AutoProcProgram.processingJobId == ProcessingJob.processingJobId,
+            )
+            .filter(ProcessingJob.dataCollectionId == parameters.dcid)
+            .filter(AutoProcProgram.processingPrograms == "UDC strategy")
+            .all()
+        )
+
+        if udc_strategy_previously_triggered:
+            self.log.info(
+                f"Skipping strategy trigger: UDC Strategy has already been triggered for dcid={parameters.dcid}."
+            )
+            return {"success": True}
+
         # Get resolution estimate from ispyb records for upstream pipeline - returns None if not found.
         resolution = (
             session.query(AutoProcScalingStatistics.resolutionLimitHigh)
diff --git a/src/dlstbx/wrapper/xia2_overload.py b/src/dlstbx/wrapper/xia2_overload.py
