From 84486c6ff950a48802f26d43392fcfb206911985 Mon Sep 17 00:00:00 2001
From: Lisa Dorofeeva <elizaveta.dorofeeva@desy.de>
Date: Wed, 25 Jun 2025 11:15:43 +0200
Subject: [PATCH] Allow to apply additional mask that is read out of file -
 NumPy or hdf5

---
 README.md                          | 14 +++++++++++-
 dap/algos/__init__.py              |  1 +
 dap/algos/addmaskfile.py           | 35 ++++++++++++++++++++++++++++++
 dap/algos/jfdata.py                |  2 ++
 dap/algos/whitefield_correction.py |  5 ++---
 5 files changed, 53 insertions(+), 4 deletions(-)
 create mode 100644 dap/algos/addmaskfile.py

diff --git a/README.md b/README.md
index 372d1bf..60ab466 100644
--- a/README.md
+++ b/README.md
@@ -116,7 +116,7 @@ options:
        * `'wf_method': 'div'|'sub'` - Method of white field correction - either division or subtraction is supported.
 
      Algorithm Output:
-       * `'is_white_field_corrected': bool` - Indicates whether white field correction took place.
+       * `'white_field_correction_applied': 1/0` - Indicates whether the algorithm ran successfully.
        * Image is changed **in-place**.
 
    * **streakfinder Algorithm**
@@ -231,6 +231,18 @@ options:
 
        Use the `'apply_additional_mask': 0/1`  - Input flag to enable this functionality.
 
+   * **Additional Mask from file**
+
+       Alternative to previous additional masking, mask data is read from specified file. NumPy and HDF5 formats are supported.
+
+     Input parameters:
+       * `'apply_additional_mask': 1/0` - Input flag to enable this functionality.
+       * `'mask_file': str` - Path to the hdf5 file with mask data.
+       * `'mask_ds': str` [Optional] - Name of the dataset containing mask in the hdf5 file, default is `"data/data"`.
+
+     Algorithm Output:
+        * `'mask_from_file_applied': 1/0` - Indicates whether the algorithm ran successfully.
+
    * **Filter based on pulse picker information**
 
        If the event propagation capability is accessible for the detector and the pulse picker information is correctly configured for propagation, the filtration based on pulse picker information becomes feasible by using the
diff --git a/dap/algos/__init__.py b/dap/algos/__init__.py
index 27bcf04..b9bf2a5 100644
--- a/dap/algos/__init__.py
+++ b/dap/algos/__init__.py
@@ -1,5 +1,6 @@
 
 from .addmask import calc_apply_additional_mask
+from .addmaskfile import calc_apply_additional_mask_from_file
 from .aggregation import calc_apply_aggregation
 from .jfdata import JFData
 from .mask import calc_mask_pixels
diff --git a/dap/algos/addmaskfile.py b/dap/algos/addmaskfile.py
new file mode 100644
index 0000000..7787195
--- /dev/null
+++ b/dap/algos/addmaskfile.py
@@ -0,0 +1,35 @@
+import h5py
+import numpy as np
+
+
+def calc_apply_additional_mask_from_file(results, pixel_mask_pf):
+    apply_additional_mask = results.get("apply_additional_mask_from_file", False)
+    if not apply_additional_mask:
+        return
+    results["mask_from_file_applied"] = 0
+    mask_file = results.get("mask_file", None)
+    if not mask_file:
+        return
+    mask_dataset = results.get("mask_ds", "data/data")
+
+    # Support for hdf5 and npy
+    if mask_file.endswith(".npy"):
+        try:
+            mask = np.load(mask_file)
+        except Exception as error:
+            print(f"Error loading mask data from NumPy file {mask_file}:\n{error}")
+            return
+    else:
+        try:
+            with h5py.File(mask_file, "r") as mask_file:
+                mask = np.asarray(mask_file[mask_dataset])
+        except Exception as error:
+            print(f"Error loading mask from hdf5 file {mask_file}:\n{error}")
+            return
+
+    try:
+        np.multiply(pixel_mask_pf, mask, out=pixel_mask_pf)
+    except Exception as error:
+        print(f"Error applying additional mask from file {mask_file}:\n{error}")
+    else:
+        results["mask_from_file_applied"] = 1
diff --git a/dap/algos/jfdata.py b/dap/algos/jfdata.py
index bd2fa1e..f204de8 100644
--- a/dap/algos/jfdata.py
+++ b/dap/algos/jfdata.py
@@ -3,6 +3,7 @@ import numpy as np
 import jungfrau_utils as ju
 
 from .addmask import calc_apply_additional_mask
+from .addmaskfile import calc_apply_additional_mask_from_file
 
 
 class JFData:
@@ -58,6 +59,7 @@ class JFData:
 
         pixel_mask_pf = np.ascontiguousarray(pixel_mask_corrected)
         calc_apply_additional_mask(results, pixel_mask_pf) # changes pixel_mask_pf in place
+        calc_apply_additional_mask_from_file(results, pixel_mask_pf)  # changes pixel_mask_pf in place
 
         self.id_pixel_mask_corrected = new_id_pixel_mask_corrected
         self.pixel_mask_pf = pixel_mask_pf
diff --git a/dap/algos/whitefield_correction.py b/dap/algos/whitefield_correction.py
index ad8e07c..37785e3 100644
--- a/dap/algos/whitefield_correction.py
+++ b/dap/algos/whitefield_correction.py
@@ -29,12 +29,11 @@ def calc_apply_whitefield_correction(results, data):
     """
     In-place white field correction of the detector data
     """
-    results["is_white_field_corrected"] = False
     do_whitefield_correction = results.get("do_whitefield_correction", False)
     if not do_whitefield_correction:
-        print(f"No whitefield correction")
         return
 
+    results["white_field_correction_applied"] = 0
     params_required = [
         "wf_data_file",
         "wf_method",
@@ -70,4 +69,4 @@ def calc_apply_whitefield_correction(results, data):
         print(f"ERROR: White field correction failed.\n"
               f"{error=}")
     else:
-        results["is_white_field_corrected"] = True
+        results["white_field_correction_applied"] = 1