nanopyx.core.transform.binning

 1def rebin_2d(arr, bin_factor, mode="sum"):
 2    """
 3    Bins a 2D array by a given factor. The last two dimensions of the array are binned.
 4    :param arr: numpy array with any shape as long as last two dimensions are y, x (example: time, channel, z, y, x)
 5    :param bin_factor: factor used to bin dimensions
 6    :param mode: can be either sum, mean or max, defaults to sum if not specified or not valid mode
 7    :return: binned array
 8    """
 9    if mode not in ["sum", "mean", "max"]:
10        mode = "sum"
11
12    if arr.ndim < 2:
13        raise ValueError("Array must have at least 2 dimensions")
14
15    if bin_factor == 1:
16        return arr
17
18    if bin_factor < 1:
19        raise ValueError("Binning factor must be greater than 1")
20
21    if arr.shape[-1] % bin_factor != 0 or arr.shape[-2] % bin_factor != 0:
22        raise ValueError("Binning factor must be a divisor of the last two dimensions")
23
24    if mode == "sum":
25        return arr.reshape(
26            arr.shape[:-2]
27            + (
28                arr.shape[-2] // bin_factor,
29                bin_factor,
30                arr.shape[-1] // bin_factor,
31                bin_factor,
32            )
33        ).sum(axis=(-1, -3))
34    elif mode == "mean":
35        return arr.reshape(
36            arr.shape[:-2]
37            + (
38                arr.shape[-2] // bin_factor,
39                bin_factor,
40                arr.shape[-1] // bin_factor,
41                bin_factor,
42            )
43        ).mean(axis=(-1, -3))
44    elif mode == "max":
45        return arr.reshape(
46            arr.shape[:-2]
47            + (
48                arr.shape[-2] // bin_factor,
49                bin_factor,
50                arr.shape[-1] // bin_factor,
51                bin_factor,
52            )
53        ).max(axis=(-1, -3))