Estimate transmission wrapper

src/dlstbx/services/strategy.py

@@ -151,10 +151,12 @@ 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 +179,9 @@ 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,12 +198,16 @@ def generate_strategy(
             )
         beamline_config = parse_config_file(beamline_config_file)
 
-        scaled_transmission = parameters.get("scaled_transmission", 1.0) 
+        recommended_max_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)
+        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
+                ),
+                recommended_max_transmission
+            ),
         )
         exposure_time_limits = (
             get_beamline_param(
@@ -304,7 +313,7 @@ def generate_strategy(
                 ispyb_command_list.append(d)
 
                 # Convert transmission to percentage for ISPyB
-                transmission_pct = transmission * 100
+                relative_transmission_pct = transmission / dc_transmission * 100
 
                 axis_end = (
                     rotation_start + rotation_increment * recipe_step.number_of_images
@@ -318,7 +327,7 @@ def generate_strategy(
                     "axisstart": rotation_start,
                     "axisend": axis_end,
                     "exposuretime": exposure_time,
-                    "transmission": transmission_pct,
+                    "transmission": relative_transmission_pct,
                     "oscillationrange": rotation_increment,
                     "noimages": recipe_step.number_of_images,
                     "resolution": resolution,

src/dlstbx/wrapper/estimate_transmission.py

@@ -0,0 +1,120 @@
+from __future__ import annotations
+
+import json
+import shutil
+import subprocess
+from collections import Counter
+from pathlib import Path
+from itertools import accumulate
+from dials.array_family import flex
+
+from dlstbx.wrapper import Wrapper
+
+class EstimateTransmissionWrapper(Wrapper):
+    _logger_name = "dlstbx.wrap.estimate_transmission"
+
+    def run(self):
+        assert hasattr(self, "recwrap"), "No recipewrapper object found"
+
+        params = self.recwrap.recipe_step["job_parameters"]
+        working_directory = Path(params["working_directory"])
+        results_directory = Path(params["results_directory"])
+
+        beamline = params["beamline"]
+        pixel_percentile = params["pixel_percentile"].get(beamline, 100) / 100
+        target_countrate_pct = params["target_countrate_pct"].get(beamline, 50) / 100
+        transmission = float(params["transmission"])
+        file = params["input_file"]
+
+        commands = [
+            ("dials.import", ["dials.import", file]),
+            ("dials.find_spots", [ "dials.find_spots", 
+                                  "imported.expt", 
+                                  "ice_rings.filter=True"],
+            ),
+        ]
+
+        for command, script in commands:
+            result = subprocess.run(script, cwd=working_directory,  check=True)
+
+            if result.returncode:
+                self.log.info(f"{command} failed with return code {result.returncode}")
+                self.log.info(result.stderr)
+
+                self.log.debug(f"Command output:\n{result.stdout}")
+                self.log.debug(f"From command: {script}")
+                return False
+
+        experiment_file = working_directory / "imported.expt"
+        with experiment_file.open("r") as f:
+            experiment = json.load(f)
+            trusted_range = experiment["detector"][0]["panels"][0]["trusted_range"][1]
+
+        reflection_file = working_directory / "strong.refl"
+        reflections = flex.reflection_table.from_file(reflection_file)
+        counts_hist = self.build_hist_from_reflections(reflections)
+
+        num_counts = list(counts_hist.keys())
+        num_pixels = list(counts_hist.values())
+
+        index_of_pixel_percentile = self.get_percentile_index(num_pixels, pixel_percentile)
+        counts_at_percentile = int(num_counts[index_of_pixel_percentile])
+
+        pixel_countrate_pct = counts_at_percentile / trusted_range 
+        self.log.info(f"The countrate percentage of the {pixel_percentile}% most intense pixel is {pixel_countrate_pct}")
+        scale_factor =  target_countrate_pct / pixel_countrate_pct
+
+        scaled_transmission = min(1, (transmission * scale_factor) / 100)
+        self.log.info(f"Scaled transmission is : {scaled_transmission}")
+
+        self.recwrap.send_to(
+            "strategy",
+            {"parameters": {"scaled_transmission": float(scaled_transmission)}},
+        )
+
+        results_directory.mkdir(parents=True, exist_ok=True)
+        output_files = ["dials.find_spots.log"]
+        for output_file in output_files:
+            source_file = working_directory / output_file
+            destination = results_directory / output_file
+
+            if not source_file.exists():
+                self.log.info(f"{source_file=} does not exsist")
+                return False
+
+            self.log.info(f"Copying {str(source_file)} to {str(destination)}")
+            shutil.copy(source_file, destination)
+
+        self.save_hist_to_json(counts_hist, trusted_range, results_directory)
+
+        self.log.info("Done.")
+        return True
+
+    def build_hist_from_reflections(self, reflections):
+        "Iterate through the shoeboxes to a reflection and generate a pixel histogram"
+
+        shoeboxes = reflections["shoebox"]
+        counter = Counter()
+        for sbox in shoeboxes:
+            counter.update(sbox.data.as_numpy_array().ravel())
+
+        sorted_counter = sorted(counter.items())
+        return {str(int(k)): v for k, v in sorted_counter}
+
+    def get_percentile_index(self, num_pixels, percentile):
+        threshold = sum(num_pixels) * percentile
+
+        for i, cum_sum in enumerate(accumulate(num_pixels)):
+            if cum_sum >= threshold:
+                return i
+
+        return len(num_pixels)
+
+    def save_hist_to_json(self, hist, max_trusted_value, results_dir):
+        results_path = results_dir / "overload.json"
+        self.log.info(f"Saving counts histogram to {str(results_path)}")
+        with open(results_path, 'w') as f:
+            json.dump({ "counts": hist,
+                        "overload_limit": max_trusted_value}, f, indent=2)
+
+        self.log.info("Saved.")