iteration vs dilation

iter_vs_dilation.py

import time

import numpy as np

from scipy.ndimage import binary_dilation, binary_propagation
from stdatamodels.jwst.datamodels import dqflags


GOOD = dqflags.group["GOOD"]
DNU = dqflags.group["DO_NOT_USE"]
CHLO = dqflags.group["CHARGELOSS"]
CHLO_DNU = CHLO + DNU


def flag_pixels(data, gdq, signal_threshold):
    n_ints, n_grps, n_rows, n_cols = gdq.shape

    chargeloss_pix = np.where((data > signal_threshold) & (gdq != DNU))
    new_gdq = gdq.copy()

    for k in range(len(chargeloss_pix[0])):
        integ, group = chargeloss_pix[0][k], chargeloss_pix[1][k]
        row, col = chargeloss_pix[2][k], chargeloss_pix[3][k]
        new_gdq[integ, group:, row, col] |= CHLO_DNU

        # North
        if row > 0:
            new_gdq[integ, group:, row-1, col] |= CHLO_DNU

        # South
        if row < (n_rows-1):
            new_gdq[integ, group:, row+1, col] |= CHLO_DNU

        # East
        if col < (n_cols-1):
            new_gdq[integ, group:, row, col+1] |= CHLO_DNU

        # West
        if col > 0:
            new_gdq[integ, group:, row, col-1] |= CHLO_DNU

    return new_gdq


_dilation_mask = np.array([[[
    [0, 1, 0],
    [1, 1, 1],
    [0, 1, 0],
]]], dtype='bool')
def binary_flag(data, gdq, signal_threshold):
    mask = (data > signal_threshold) & (gdq != DNU)
    dilated = binary_dilation(mask, _dilation_mask, output=mask)
    propagated = np.add.accumulate(dilated, axis=1, dtype='bool', out=dilated)
    new_gdq_0 = gdq.copy()
    new_gdq_0[propagated] |= CHLO_DNU
    return new_gdq_0


if __name__ == '__main__':
    for N in [1, 10, 100, 1000]:
        data = np.zeros((10, 10, 100, 100), dtype='f4')

        # set N random pixels
        data.flat[:N] = 1
        # shuffle the array
        np.random.shuffle(data)
        threshold = 0.5

        gdq = np.zeros_like(data, dtype='uint32')

        t0 = time.monotonic()
        for _ in range(100):
            bf = binary_flag(data, gdq, threshold)
        t0 = (time.monotonic() - t0) / 100

        t1 = time.monotonic()
        for _ in range(100):
            fp = flag_pixels(data, gdq, threshold)
        t1 = (time.monotonic() - t1) / 100

        print(t0 / 1000, t1 / 1000)

        assert np.all(bf == fp)