ewoksid02 0.1.0__py3-none-any.whl

ewoksid02/utils/average.py

@@ -0,0 +1,158 @@
+import logging
+from typing import List, Optional, Tuple
+
+import numexpr
+import numpy
+
+logger = logging.getLogger(__name__)
+
+
+def get_centrosymmetric_limits(azimuth_range: List[float]) -> Tuple[List[float], ...]:
+    """
+    Calculate the centrosymmetric limits for a given azimuth range.
+
+    Inputs:
+        azimuth_range (List[float]): The azimuth range as a list of two values.
+    Outputs:
+        Tuple[List[float], ...]: A tuple containing one or two lists of centrosymmetric limits.
+    """
+    if azimuth_range[0] == azimuth_range[1]:
+        return ([-180, 180],)
+
+    lim = numpy.array(azimuth_range)
+    if lim.prod() <= 0:
+        lim = max(abs(lim))
+        return ([max(-lim, -180), min(lim, 180)],)
+    else:
+        a = min(max(abs(lim)), 180)
+        b = max(min(abs(lim)), -180)
+        return ([-a, -b], [b, a])
+
+
+def get_array_limit(
+    azimuth_range: List[float], azimuth_array: numpy.ndarray
+) -> Tuple[int, int]:
+    """
+    Get the array indices corresponding to the azimuth range.
+
+    Inputs:
+        azimuth_range (List[float]): The azimuth range as a list of two values.
+        azimuth_array (numpy.ndarray): The array of azimuth values.
+    Outputs:
+        Tuple[int, int]: The indices corresponding to the azimuth range.
+    """
+    arg_min = numpy.argmin(abs(azimuth_array - azimuth_range[0]))
+    arg_max = numpy.argmin(abs(azimuth_array - azimuth_range[1]))
+    return (arg_min, arg_max)
+
+
+def calculate_average(
+    dataset_intensity: Optional[numpy.ndarray],
+    dataset_sum_signal: Optional[numpy.ndarray] = None,
+    dataset_sum_norm: Optional[numpy.ndarray] = None,
+    dataset_sum_variance: Optional[numpy.ndarray] = None,
+    array_ranges: Optional[List[Tuple[int, int]]] = None,
+    Dummy: int = -10,
+    calculate_variance: bool = False,
+    **kwargs,
+) -> Tuple[
+    Optional[numpy.ndarray],
+    Optional[numpy.ndarray],
+    Optional[numpy.ndarray],
+    Optional[numpy.ndarray],
+]:
+    """
+    Calculate the average intensity, signal normalization, variance, and sigma for a dataset.
+
+    Inputs:
+        dataset_intensity (Optional[numpy.ndarray]): The intensity dataset.
+        dataset_sum_signal (Optional[numpy.ndarray], optional): The sum signal dataset. Defaults to None.
+        dataset_sum_norm (Optional[numpy.ndarray], optional): The sum normalization dataset. Defaults to None.
+        dataset_sum_variance (Optional[numpy.ndarray], optional): The sum variance dataset. Defaults to None.
+        array_ranges (Optional[List[Tuple[int, int]]], optional): The ranges for averaging. Defaults to None.
+        Dummy (int, optional): The dummy value to replace invalid data. Defaults to -10.
+    Outputs:
+        Tuple[Optional[numpy.ndarray], Optional[numpy.ndarray], Optional[numpy.ndarray], Optional[numpy.ndarray]]:
+        The average intensity, signal normalization, variance, and sigma datasets.
+    """
+    sum_signal = None
+    sum_norm = None
+    sum_variance = None
+    dataset_average_intensity = numpy.zeros(
+        (dataset_intensity.shape[0], dataset_intensity.shape[-1])
+    )
+    dataset_average_signal_norm = numpy.zeros_like(dataset_average_intensity)
+    dataset_average_variance = numpy.zeros_like(dataset_average_intensity)
+    dataset_average_sigma = numpy.zeros_like(dataset_average_intensity)
+
+    if dataset_sum_signal is not None:
+        sum_signal = numpy.zeros(
+            (dataset_sum_signal.shape[0], dataset_sum_signal.shape[-1])
+        )
+    if dataset_sum_norm is not None:
+        sum_norm = numpy.zeros((dataset_sum_norm.shape[0], dataset_sum_norm.shape[-1]))
+    if dataset_sum_variance is not None:
+        sum_variance = numpy.zeros(
+            (dataset_sum_variance.shape[0], dataset_sum_variance.shape[-1])
+        )
+    array_ranges = array_ranges or [(0, -1)]
+
+    mask_slice = numpy.zeros_like(dataset_intensity)
+    for array_range in array_ranges:
+        if not isinstance(array_range, (list, tuple)):
+            logger.error(f"Array range {array_range} is not a list or tuple.")
+            continue
+
+        if len(array_range) != 2:
+            logger.error(
+                f"Array range {array_range} is not valid. It should contain two elements."
+            )
+            continue
+
+        lim_0 = int(array_range[0])
+        lim_1 = int(array_range[1])
+
+        mask_slice[:, lim_0:lim_1, :] = 1
+
+    dataset_average_intensity = numexpr.evaluate(
+        "where(mask_slice == 0, Dummy, dataset_intensity)"
+    )
+    dataset_average_intensity[dataset_average_intensity == Dummy] = numpy.nan
+    dataset_average_intensity = numpy.nanmean(dataset_average_intensity, axis=1)
+
+    if dataset_sum_signal is not None:
+        sum_signal = numexpr.evaluate(
+            "where(mask_slice == 0, Dummy, dataset_sum_signal)"
+        )
+        sum_signal[sum_signal == Dummy] = numpy.nan
+        sum_signal = numpy.nansum(sum_signal, axis=1)
+    if dataset_sum_norm is not None:
+        sum_norm = numexpr.evaluate("where(mask_slice == 0, Dummy, dataset_sum_norm)")
+        sum_norm[sum_norm == Dummy] = numpy.nan
+        sum_norm = numpy.nansum(sum_norm, axis=1)
+    if dataset_sum_variance is not None:
+        sum_variance = numexpr.evaluate(
+            "where(mask_slice == 0, Dummy, dataset_sum_variance)"
+        )
+        sum_variance[sum_variance == Dummy] = numpy.nan
+        sum_variance = numpy.nansum(sum_variance, axis=1)
+
+    if sum_signal is not None and sum_norm is not None:
+        dataset_average_signal_norm = numexpr.evaluate(
+            "where(sum_norm <= 0.0, sum_norm, sum_signal / sum_norm)"
+        )
+
+    if sum_variance is not None and sum_norm is not None:
+        dataset_average_variance = numexpr.evaluate(
+            "where(sum_norm <= 0.0, sum_norm, sum_variance / (sum_norm * sum_norm))"
+        )
+        dataset_average_sigma = numexpr.evaluate(
+            "where(sum_norm <= 0.0, sum_norm, sqrt(sum_variance) / sum_norm)"
+        )
+
+    return (
+        dataset_average_intensity,
+        dataset_average_signal_norm,
+        dataset_average_variance,
+        dataset_average_sigma,
+    )