Add the gain replacement functions in jungfraulib

src/cal_tools/jungfraulib.py

@@ -209,3 +209,80 @@ def sort_runs_by_gain(
                          "Please verify the selected 3 dark runs to process.")
 
     return runs
+
+
+def replace_pixels_wrong_gain_badpixel(
+    gain_arr,
+    new_gain_value,
+    badpixel_map,
+    memory_cells,
+):
+    """_summary_
+
+    Args:
+        gain_arr (ndarray): The raw gain array to be updated.
+        new_gain_value (int): The new gain value to replace the gain value
+            for the pixels with WRONG_GAIN_VALUE badpixel.
+        badpixel_map (ndarray): The badpixel map for a cell.
+        memory_cells (int): Number of memory cells for the processed detector.
+
+    Returns:
+        ndarray: The updated gain_arr
+    """
+    import numpy as np
+    from cal_tools.enums import BadPixels
+
+    # replace gain with WRONG_GAIN_VALUE pixels in the badpixel mask
+    wrong_gain_mask = BadPixels.WRONG_GAIN_VALUE.value
+
+    # Create a boolean mask to identify the elements that meet the condition
+    condition_mask = (badpixel_map[..., new_gain_value] & wrong_gain_mask) == wrong_gain_mask
+
+    if memory_cells == 1:
+        condition_mask = condition_mask[np.newaxis, :, :]  # gain_arr.shape = [1, 512, 1024]
+
+    gain_arr[condition_mask] = new_gain_value
+    return gain_arr
+
+
+def replace_pixels_gain_values(
+    gain_arr,
+    new_gain_value,
+    wrong_gain_pixels,
+    karabo_da,
+    badpixel_map,
+    memory_cells
+):
+    """Replace gain values with the chosen new gain value
+
+    Args:
+        gain_arr (ndarray): The raw gain array to be updated.
+        new_gain_value (int): _description_
+        wrong_gain_pixels (list): [module_index, row1, row12, col1, col2].
+            e.g. [4, 0, 255, 896, 1024] for JNGFR04 slicing module [0:255, 896:1024].
+        karabo_da (str): _description_
+        badpixel_map (ndarray): The badpixel map for one cell.
+        memory_cells (int): Number of memory cells.
+
+    Returns:
+        (ndarray): The update gain_arr.
+
+    """
+    if wrong_gain_pixels[0] > -1:
+        if wrong_gain_pixels[0] == int(karabo_da[-2:]):
+            x1 = wrong_gain_pixels[1]
+            x2 = wrong_gain_pixels[2]
+            y1 = wrong_gain_pixels[3]
+            y2 = wrong_gain_pixels[4]
+            gain_arr[:, x1:x2, y1:y2] = new_gain_value
+            return gain_arr
+        else:  # No gain replacement
+            return gain_arr
+
+    else: # replace gain with WRONG_GAIN_VALUE pixels in the badpixel mask
+        return replace_pixels_wrong_gain_badpixel(
+            gain_arr,
+            new_gain_value,
+            badpixel_map,
+            memory_cells,
+        )
\ No newline at end of file