<jupyter_start><jupyter_text>Theory of operationThis notebook outlines in Python a new approach to signal processing.It is mocked up here in Python but is intended to be re-implementedin a combination of C and Python using the established patternsin the plaster.run folders. Glossary- **Calibration Runs / calibration**    Jobs that are run through our scope using control dyes for the purposes    of establishing various parameters such as regional illumination and PSF.    These runs are run less frequenrly than experimental runs and usually    on an as-needed basis.- **Experiment runs / analysis**    Jobs on the scope that run real experiments. Typically require a paired    calibration run for later analysis.- **Field**    One position of the microscope stage. Real runs have on the order of    10-1000 fields. Fields are all (in principal) independent and thus a    natural point of parallelization.- **Channel**    One wavelength of image. 1-5 channels are acquired on a given run.    These wavelengths are acquired simultanously on different cameras    and will have a fixed spatial relationship with each other.    For now, we are assuming that we can remove any rotational    component of that relationship and have only translational    terms. We might have to revisit this which will have significant    impacts on the performace.- **Cycle**    One chemical cycle. Real experimennts have 10-20 cycles.- **Measure (mea) vs Dimension (dim)**    When a one-dimensional size is required (example: width / height    of a square) this is called a "measure" and abbreviated "mea".    When a muti-dimension size is requried (example: width, height of    a rectangle) this is called a "dim".- **Choreograhy**    The scope acquires images following this pseudo-code. This is called the    choreography of the instrument.        ```    for cycle_i in range(n_cycles):         Tell the motors and pumps to wash in various chemicals         Incubate those chemical for about 1 hour        time.sleep(incubation_time)        for field_i in range(n_fields):            move_stage_to_position(field_i)            for channel_i in range(n_channels):                acquire_image(channel_i)                save_image(field_i, channel_i, cycle_i)                 Note this tuple order (fl, ch, cy) is a standard used throughout the code    ```- **PSF**    A "Point Spread Function" is a small image (~ 10x10 pixels) of the    idealized spot image we expect to see. The following code is highly    optimized around the assumptions that:    1. The PSF has an axis of symmetry as is the case for a 2D Gaussian    2. The PSF varies slowly enough over the field that we partition    the images in to subregions inside of which the PSF can be assumed constant.- **Illumination**    The optics in the instrument do not illuminate the entire field uniformly.    The map of this regional differnce is called the "Regional Illumination Map"    That map is modelled as a 2D exponential but is assumed to be rendered into    a pixel-map for fast lookup.- **Regional properties**    Illumination and PSF are examples of "regional" properites in that their    value varies over the coordinate of the image.- **Channel Alignment**    The Camera for each channel are not expected to be perfectly aligned    translationally and therefore calibration runs are required that    used cross-channel dyes to allow us to map the alignment offsets    relative to channel 0.- **Cycle alignment**    When the stage moves from field to field during the acquisition of a cycle,    the stage can not return to exactly the original location within nanometer    precision thus there is a translational alignment required using the     dyes as their own fiducial markers. It is assumed that there is only translational    variation in this, not rotation.- **Sub-pixel alignment**    Measurment of the radiometry requires aligning the PSF with the peak to     sub-pixel precision. A variety of the complexity of the following algoirthm    is to find and apply these sub-pixel offsets. There are times in the algoirthm    where a 1-2 pixel approimation is made and then refined to sub-pixel precision.- **Spot v Peak v Loc**    Spots and peaks are synonyms for "a bright place that is likely a dye of interest"    A "loc" is the location of a spot/peak. Often "locs" are tracked as the vector    of peak locations and end up synonymous with "peaks" and "spots"- **Multi channel peak finding**    There is no requirement that a given peak is visible in all channels.    Multi-channel-peak-finding is the process of finding candidate peaks that    are visible in *any* channel and this is complicated by the fact that    channel intensities are not expected to be balanced. Therefore the     Multi-channel-peak-finder must perform various balancing and call    the peaks across channel using a binary mask "candiadte peal pixel present in any channel?"    This algorithm is expected to find peak onlyo within a few pixels. Sub-pixel    refinement is defered into the frequency domain.- **Radiometry**    The combined measurements acquired from signal prococess (signal and noise).    Radiometriy comes from the "measurement of the photons from the molecule" and    is essentialyl synponymous with "brightrness" or "intensity"- **Signal**    The measurment of interest: the intensity of each "peak" at each channel and cycle.    The result fills a "radiomtry matrix" or "radmat" with shape (n_peaks, n_channels, n_cycles).    Sometimes we collapse the last two dimensions for visualization so that we can    render the radmat as a (n_peaks, n_channels * n_cycles) heatmap image.- **Noise**    A measurement of the variation of a spot versus the expected PSF. This is    used to find and filter "near-misses" ie when two spots have come close enough    to each other to affect each other's radiomtry. - **Near-miss versus collision**    A near miss is defined as two peaks that are close enough to affect each other's    measurement but far enough away to detect the interference via the noise metric.    A colliosion is a collision that is so close that it can not be distibgioushed    from a single peak of combined intensity. We consider this acceptable confusion    and model it further down-stream.- **FFT coordinate shifting**    TODO    - **FFT Phase Offsets**    For sub-pixel alignments, an non-obious Frequency-domain property is    exploited that a multiplying each element of the freq. domain image    by a constant phasor will shift the image when inverse FFT.- **Priors / hyper_**    TODOTODONaming conventions:_if_ Operational Flow```During calibration:    Assume:     - All channels are rotationally aligned within tolerance       (What that tolerance is is not clear yet.)     - The PSF has some axis of symmetry (eg Elliptical Gaussian)     - Do **NOT** assume that channels are translationally aligned.    Obtain:     - RegPSF per channel        These regions are the same as the sub-division we will make of the field        We will compute an FFT for these PSFs per region along with scale factors.     - RegIllum per channel        Which includes balance, bias, threshold determined by sweeping        over these parameters to maximize recall and these operations are        then amortized over all fields.     - ChannelAlignment        A sub-pixel translation constant for each channel relative to channel 0During analysis:    Cycle alignment:        if cycle == 0:            for several smallish "alignment regions" across the channel 0 image:                FFT those regions and store these FFT values for cycle 1+ alignment        else:            for the same alignment regions stored during cycle 0:                FFT the region and multiply that FFT from cycle 0:                 The location of the maximum of those multiplies results in a                 1-pixel alignment from cycle n to cycle 0.                Gradient descend to revise to sub-pixel precision            if all alignment regions have approx same answer:                take the mean of those alignments and proceed            else:                Fallback to a "low density" aligner where the entire image                is (slowly) used for alignment. This is only for rare debugging                occasions where the density is extremely low.    For each regional subdivision:        if cycle == 0:            multichannel_peak_find() using RegIllum parameters from calib                This returns 1-2 pixel location of each peak that appears                in any channel.        for each peak (loaded from cycle 0 if needed):            for each channel:                FFT a (peak_mea, peak_mea) sub-image from each channel                (this is expected to be done in-place in C, ie without                an extra copy from the larger buffers to a contiguous                 power of 2 FFT buffer).                This FFT is now within a few pixel of the correct alignment                if cycle == 0:                    Gradient-descend to find sub-pixel centering for each channel.                    Some channels may have no data and this GD will                    either fail or produce an invalid answer.                    TODO: PROTOTYPE an algorithm to detect this.                                    Inot the sub-pixel centering from cycle 0, add:                    the sub-pix channel alignment                    the sub-pix cycle alignment (zero if cycle == 0)                eval signal:                    ((f_phase_shifted.real).sum() - f_phase_shifted[0,0]) / f_psf_squared_sum_without_dc                                    eval noise (or similar for collision detection)                    TODO: What is this function?                store sub-pix x, y, sig, noise```<jupyter_code># Helpers

def coord_shift_im(im):
    """
    Re-arrange the pixels of an image such that it's center becomes the 
    lower-left pixel and it wraps around.
    This simplifies a variety of FFT operations.
    In a C implementation we will do this implicitly during
    the sampling from a full-sized image.
    """
    assert im.shape[0] == im.shape[1] and im.ndim == 2
    mea = im.shape[0]
    mea = int(mea)
    space = range(-mea // 2 + 1, mea // 2 + 1)
    # space is now like [-3, -2, -1, 0, 1, 2, 3, 4]
    x, y = np.meshgrid(space, space)
    result = np.zeros((mea, mea))
    result[y, x] = im
    return result

def gauss2_rho_form_lower_left_wrap(amp, std_x, std_y, rho, mea):
    """
    Like gauss2_rho_form but sampled with wrap-around
    and with no half-pixel offset because this is not for visualization!
    """
    
    assert utils.is_power_of_2(mea)
    
    mea = int(mea)
    space = range(-mea // 2 + 1, mea // 2 + 1)
    x, y = np.meshgrid(space, space)

    cov = np.array([
        [std_x ** 2, rho * std_x * std_y],
        [rho * std_x * std_y, std_y ** 2],
    ])

    det_cov = np.linalg.det(cov)
    inv_cov = np.linalg.inv(cov)

    radius_vectors = (np.array([x.flatten(), y.flatten()]).T).T

    r = np.sum(np.multiply((radius_vectors.T @ inv_cov), radius_vectors.T), axis=1)

    exp_term = amp * np.exp(-0.5 * r)
    norm_term = np.sqrt(((np.pi * 2.0) ** 2) * det_cov)
    
    result = coord_shift_im((exp_term / norm_term).reshape((mea, mea)))
    return result
n_peaks = 500
mu = 10_000
sigma = 0.16
im_mea = 512
bg_mean = 100
bg_std = 50
halo_size = 20
halo_inensity = 2
n_cycles = 2
n_channels = 2
hyper_peak_mea = 16
hyper_peak_sigma = 1.8
hyper_n_divs = 4

x1 = np.arange(0, hyper_peak_mea // 2)
x2 = np.arange(-hyper_peak_mea // 2, 0, 1)
phase_j, phase_i = np.meshgrid(np.hstack((x1, x2)), np.hstack((x1, x2)))

reg_psf = PriorsMLEFixtures.reg_psf_uniform(
    hyper_im_mea=im_mea,
    hyper_peak_mea=hyper_peak_mea,
    hyper_n_divs=hyper_n_divs,
    hyper_peak_sigma=hyper_peak_sigma,
)

psf_im, (cy, cx) = reg_psf.render_at_loc((im_mea/2.0, im_mea/2.0))
z.im(psf_im, f_title="Example PSF")

# Transform PSF
# TODO: Operate on all the PSF in a psf reg priors for now: just one.

psf_im = gauss2_rho_form_lower_left_wrap(1.0, hyper_peak_sigma, hyper_peak_sigma, 0.0, hyper_peak_mea)
f_psf_with_dc = np.fft.fft2(psf_im)
f_psf_dc = f_psf_with_dc[0, 0].real

# Zero-out the DC term to eliminate background (but hoold on to it for later scaling)
f_psf_without_dc = f_psf_with_dc.copy()
f_psf_without_dc[0, 0] = 0
f_psf_squared_sum_without_dc = ( ( f_psf_without_dc * np.conjugate(f_psf_without_dc) ).real ).sum()

with z(_cols=2, _cspan=(-1, 1)):
    z.im(f_psf_with_dc.real, f_title="Real component of f_psf")
    z.im(f_psf_with_dc.imag, f_title="Imaginary component of f_psf should be all zero")

def generate_donut_psf():
    psf = np.zeros((hyper_peak_mea, hyper_peak_mea))
    for i in range(-hyper_peak_mea // 2 + 1, hyper_peak_mea // 2 + 1, 1):
        for j in range(-hyper_peak_mea // 2 + 1, hyper_peak_mea // 2 + 1, 1):
            radius = np.sqrt(i**2 + j**2)
            argument = ((radius - 3.0)/3.0)**2
            psf[i, j] = np.exp(-argument)
    return psf

donut_psf = generate_donut_psf()
f_psf_donut_with_dc = np.fft.fft2(donut_psf)
f_psf_donut_without_dc = f_psf_donut_with_dc.copy()
f_psf_donut_without_dc[0, 0] = 0.0

z.im(donut_psf)
# GENERATE synthetic images

reg_illum = PriorsMLEFixtures.reg_illum_variable(hyper_im_mea=im_mea, falloff=0.5)

with synth.Synth(dim=(im_mea, im_mea), n_cycles=n_cycles, n_channels=n_channels) as s:
    s.channel_aln_offsets( np.array([(0.0, 0.0), (2.3, 5.7)]) )
    # s.zero_aln_offsets()

    peaks = (
        synth.PeaksModelPSF(n_peaks=n_peaks, reg_psf=reg_psf)
        .bleach(0.3)
        .gain_constant(mu)
        .locs_randomize()
        #.locs_randomize_no_subpixel()
        .channel_scale_factor((1.0, 0.6,))
    )
    synth.CameraModel(bg_mean=bg_mean, bg_std=bg_std)
    synth.HaloModel(halo_size, halo_inensity)
    synth.RegIllumModel(reg_illum)
    chcy_ims = s.render_chcy()

    # Offset and force everything positive and store as uint 16 like the camera does
    bias = np.min(chcy_ims)
    chcy_ims = np.ascontiguousarray((chcy_ims - bias).astype(np.uint16))
    
    print("Channel alignment")
    print(s.ch_aln)
    
    print("Cycle alignment")
    print(s.aln_offsets)
    
with z(_cols=n_channels, _cspan=(0, 1000), _size=400):
    for cy_i in range(n_cycles):
        for ch_i in range(n_channels):
            z.im(chcy_ims[ch_i, cy_i], f_title=f"{ch_i=} {cy_i=}")

assert np.all(chcy_ims >= 0)

# np.save("_test_psf_im.npy", psf_im)
def find_peaks_and_allocate_buffers_on_cycle_0(ch_ims):
    # PEAK FIND
    bal_ones = np.ascontiguousarray(np.ones_like(ch_ims))
    locs = peak_find_multichannel(
        ch_ims=ch_ims,
        ch_bal_ims=bal_ones,
        ch_thresh=np.array([450, 450]),  # TODO: Get correct channel val
        ch_bias=np.array([-bias, -bias]),  # TODO: Get correct channel val
        field_i=0,
        cycle_i=0
    )
    
    n_locs = locs.shape[0]

    sig = np.full((n_locs, n_channels, n_cycles), np.nan)
    noi = np.full((n_locs, n_channels, n_cycles), np.nan)
    clr = np.full((n_locs, n_channels, n_cycles), np.nan)
    loc_best_fit_center = np.full((n_locs, 2), np.nan)
    
    return locs, sig, noi, clr, loc_best_fit_center


def get_shift_phasors(offset):
    y = offset[0]
    x = offset[1]
    shift_y = -2.0 * np.pi * y / hyper_peak_mea
    shift_x = -2.0 * np.pi * x / hyper_peak_mea
    return np.exp( complex(0.0, 1.0) * (phase_i * shift_y + phase_j * shift_x) )


def resample_with_phase_shift(data, offset):
    return data * get_shift_phasors(offset)


found_disagree = False
found_agree = False


def objective_sci(f_conv, offset):
    f_phase_shifted = resample_with_phase_shift(f_conv, offset)
    return (f_phase_shifted.imag**2).sum()


def objective_jhd(f_conv, offset):
    y = offset[0]
    x = offset[1]
    shift_x = -2.0 * np.pi * x / hyper_peak_mea
    shift_y = -2.0 * np.pi * y / hyper_peak_mea

    temp_r = f_conv.real
    temp_i = f_conv.imag
    phasor_r = np.cos(phase_i * shift_y  +  phase_j * shift_x)
    phasor_i = np.sin(phase_i * shift_y  +  phase_j * shift_x)

    phasor_r_shift_x = -phase_j * phasor_i
    phasor_r_shift_x_shift_x = - phase_j * phase_j * phasor_r
    phasor_r_shift_x_shift_y = - phase_j * phase_i * phasor_r

    phasor_i_shift_x = phase_j * phasor_r
    phasor_i_shift_x_shift_x = -phase_j * phase_j * phasor_i
    phasor_i_shift_x_shift_y = -phase_j * phase_i * phasor_i

    phasor_r_shift_y = -phase_i * phasor_i
    phasor_r_shift_y_shift_y = -phase_i * phase_i * phasor_r
    phasor_r_shift_y_shift_x = -phase_i * phase_j * phasor_r

    phasor_i_shift_y = phase_i * phasor_r
    phasor_i_shift_y_shift_y = -phase_i * phase_i * phasor_i
    phasor_i_shift_y_shift_x = -phase_i * phase_j * phasor_i

    temp_imag = (temp_r * phasor_i + temp_i * phasor_r)**2
    temp_imag_shift_y = 2.0*(temp_r * phasor_i + temp_i * phasor_r) * (temp_r * phasor_i_shift_y + temp_i * phasor_r_shift_y)
    temp_imag_shift_x = 2.0*(temp_r * phasor_i + temp_i * phasor_r) * (temp_r * phasor_i_shift_x + temp_i * phasor_r_shift_x)
    temp_imag_shift_y_shift_y = (
        2.0 * (
              (temp_r * phasor_i_shift_y + temp_i * phasor_r_shift_y) * (temp_r * phasor_i_shift_y + temp_i * phasor_r_shift_y)
            + (temp_r * phasor_i + temp_i * phasor_r) * (temp_r * phasor_i_shift_y_shift_y + temp_i * phasor_r_shift_y_shift_y)
        )
    )
    temp_imag_shift_y_shift_x = (
        2.0 * (
              (temp_r * phasor_i_shift_x + temp_i * phasor_r_shift_x) * (temp_r * phasor_i_shift_y + temp_i * phasor_r_shift_y)
            + (temp_r * phasor_i + temp_i * phasor_r) * (temp_r * phasor_i_shift_y_shift_x + temp_i * phasor_r_shift_y_shift_x)
        )
    )
    temp_imag_shift_x_shift_x = (
        2.0 * (
              (temp_r * phasor_i_shift_x + temp_i * phasor_r_shift_x) * (temp_r * phasor_i_shift_x + temp_i * phasor_r_shift_x)
            + (temp_r * phasor_i + temp_i * phasor_r) * (temp_r * phasor_i_shift_x_shift_x + temp_i * phasor_r_shift_x_shift_x)
        )
    )

    grad = np.array([temp_imag_shift_y.sum(), temp_imag_shift_x.sum()])
    hessian = np.array([
        [temp_imag_shift_y_shift_y.sum(), temp_imag_shift_y_shift_x.sum()],
        [temp_imag_shift_y_shift_x.sum(), temp_imag_shift_x_shift_x.sum()]
    ])
    
    # Given this is a 2x2 we should be able to use Cramer's rule. https://en.wikipedia.org/wiki/Cramer%27s_rule
    # instead of the general linear algebra solver.
#     new_shift = np.linalg.solve(hessian, grad)
#     return temp_imag.sum(), new_shift   


    new_shift_descent = np.linalg.solve(hessian, grad) * hyper_peak_mea / (2.0*np.pi)
    mag_new_shift_descent = np.sqrt(np.inner(new_shift_descent, new_shift_descent))
    mag_grad = np.sqrt(np.inner(grad, grad))

    # CLean up...
    if (new_shift_descent * grad)[0] < -0.00001:
        # Really less than minus floating epsilon
        new_shift_descent[0] = -new_shift_descent[0]
    if (new_shift_descent * grad)[1] < -0.00001:
        # Really less than minus floating epsilon
        new_shift_descent[1] = -new_shift_descent[1]

    if(mag_new_shift_descent > 2.0):
        new_shift = 2 * new_shift_descent / mag_new_shift_descent
        # new_shift now a unit vector
    else:
        new_shift = new_shift_descent

    fcn_value = temp_imag.sum()
    return fcn_value, new_shift    


def analyze_peak(im, loc_i, loc, cy_i, ch_i, bal_im, best_fit_center):
    global found_disagree, found_agree
    peak_im = imops.crop(im, off=YX(loc), dim=HW(hyper_peak_mea, hyper_peak_mea), center=True)
    
    sig, noi, clr = np.nan, np.nan, np.nan
    if peak_im.shape == (hyper_peak_mea, hyper_peak_mea):
        f_peak_im = np.fft.fft2(coord_shift_im(peak_im))
        f_conv = f_psf_without_dc * f_peak_im

        def minimize_sci():
            res = minimize(
                lambda x: objective_sci(f_conv, x),
                (0.0, 0.0),
                method="L-BFGS-B",
                bounds=np.array([(-4.0, 4.0), (-4.0, 4.0)]),
                options=dict(eps=0.01),
                tol=1e-5,
            )
            best_fit_center = res.x
            return best_fit_center

        def minimize_jhd():
            best_fit_center = np.nan
            loc = np.array([0.0, 0.0])
            history = [loc.copy()]

            f_donut_conv = f_psf_donut_without_dc * f_peak_im
            which_f_conv = f_donut_conv
            
            shift_to_switch_to_correct_psf = 0.75
            
            for i in range(10):
                obj, new_shift = objective_jhd(which_f_conv, loc)
                
                new_shift_mag = np.sqrt(np.sum(new_shift**2))
                if new_shift_mag < shift_to_switch_to_correct_psf:
                    which_f_conv = f_conv
                    
                if new_shift_mag < 1e-6:
                    best_fit_center = loc
                    break
                
                loc += new_shift
                history += [loc.copy()]
            else:
                # TODO: Decide what to do on failuer. NAN?
                # best_fit_center = loc
                # Keep as nan
                pass
            
            return best_fit_center, np.vstack(history)

        if cy_i == 0:
            best_fit_center_sci = minimize_sci()
            best_fit_center_jhd, history_jhd = minimize_jhd()

            def save_state(name):
                np.save(f"_{name}_peak.npy", peak_im)
                np.save(f"_{name}_fit_sci.npy", best_fit_center_sci)
                np.save(f"_{name}_fit_jhd.npy", best_fit_center_jhd)
                np.save(f"_{name}_f_conv.npy", f_conv)
                np.save(f"_{name}_f_conv_donut.npy", f_psf_donut_without_dc * f_peak_im)
                np.save(f"_{name}_history_jhd.npy", history_jhd)
            
            # COMPARE the two results and dump first that disagrees
            dif_mag = np.sqrt(np.sum((best_fit_center_jhd - best_fit_center_sci)**2))
            if dif_mag > 1.0 and not found_disagree:
                found_disagree = True
                save_state("disagree")

            if dif_mag < 0.1 and not found_agree:
                found_agree = True
                save_state("agree")

            # CHANGE which minimize result we use..
            # best_fit_center = best_fit_center_sci  # Works.
            best_fit_center = best_fit_center_jhd  # Works better.
    
        if np.all(~np.isnan(best_fit_center)):
            f_phase_conv_shifted = resample_with_phase_shift(f_conv, best_fit_center)
            peak_amplitude = np.real(
                ((f_phase_conv_shifted.real).sum() - f_phase_conv_shifted[0, 0].real) / f_psf_squared_sum_without_dc
            )
            
            bal = bal_im[int(loc[0]), int(loc[1])]
            sig = (peak_amplitude * f_psf_dc) / bal

            f_conv_phase_shifted = resample_with_phase_shift(f_conv, best_fit_center)
            peak_amplitude = np.real( ((f_conv_phase_shifted.real).sum() - f_conv_phase_shifted[0, 0]) / f_psf_squared_sum_without_dc )
            bal = bal_im[int(loc[0]), int(loc[1])]
            sig = (peak_amplitude * f_psf_dc) / bal

            f_data_phase_shifted = resample_with_phase_shift(f_peak_im, best_fit_center)
            f_data_phase_shifted[0, 0] = 0.0
            f_baseline = f_data_phase_shifted - peak_amplitude * f_psf_without_dc
            noise_symm_var = (f_baseline.real ** 2).sum() / (hyper_peak_mea**4)
            noise_asym_var = (f_baseline.imag ** 2).sum() / (hyper_peak_mea**4)
            noi = np.sqrt((noise_symm_var**2 + noise_asym_var**2) / (f_psf_without_dc.real**2).sum())
            clr = np.log(noise_symm_var / noise_asym_var)

    return sig, noi, clr, best_fit_center

locs, sig, noi, clr, loc_best_fit_center = None, None, None, None, None

def analyze_cycle(cy_i):
    global locs, sig, noi, clr, loc_best_fit_center
    
    hd("h2", f"Cycle {cy_i}")
    
    if cy_i == 0:
        ch_ims = np.ascontiguousarray(chcy_ims[:, cy_i])
        locs, sig, noi, clr, loc_best_fit_center = find_peaks_and_allocate_buffers_on_cycle_0(ch_ims)

    # In real code we will run a cycle alignment here
    # For this demo we're just pulling the true offset
    # from the synthetic data.
    cy_aln = s.aln_offsets[cy_i]  # TODO: Rename aln_offsets in base synth
    locs += cy_aln
    
    # Similarly calibration would give us ch_aln and in this
    # demo I'm pulling it from the synthetic.
    ch_aln = s.ch_aln

    ch_bal_ims = reg_illum.render()
    # ch_bal_ims = np.ones_like(ch_bal_ims)  # Prove ch bal is working
    ch_bal_ims = np.repeat(ch_bal_ims[None, :, :], n_channels, axis=0)
    # TODO: Have a truely ch_bal instead of duplicating ch 0 to other channels

    # HACK!
    # locs = locs[0:10]
    
    for loc_i, loc in enumerate(locs):
        for ch_i in range(n_channels):
            im = np.ascontiguousarray(chcy_ims[ch_i, cy_i])
            (
                sig[loc_i, ch_i, cy_i],
                noi[loc_i, ch_i, cy_i],
                clr[loc_i, ch_i, cy_i],
                best_fit_center
            ) = analyze_peak(
                im, loc_i, loc + ch_aln[ch_i], cy_i, im, ch_bal_ims[ch_i], loc_best_fit_center[loc_i]
            )
            
            if cy_i == 0:
                # Only updating the best_fit_center on cycle 0 because it is re-used on subsequent
                loc_best_fit_center[loc_i] = best_fit_center

    with z(_cols=4, _size=220):
        for ch_i in range(n_channels):
            
            mask = (~np.isnan(sig[:, ch_i, cy_i])) & (-0.5 < clr[:, ch_i, cy_i]) & (clr[:, ch_i, cy_i] < 0.5)
            print(f"Keeping {100 * mask.sum() / sig.shape[0]:2.1f}%")
            
            im = np.ascontiguousarray(chcy_ims[ch_i, cy_i])
            z.im(circle_locs(im, locs[mask] + ch_aln[ch_i], inner_radius=5, outer_radius=7), _nan_color="red", f_title=f"{ch_i=} {cy_i=}")

            z.hist(sig[mask, ch_i, cy_i], f_title=f"Sig {ch_i=} {cy_i=}")
            z.hist(noi[mask, ch_i, cy_i], f_title=f"Noi {ch_i=} {cy_i=}")
            z.hist(clr[mask, ch_i, cy_i], f_title=f"Clr {ch_i=} {cy_i=}", _bins=(-2, 2, 50))
#             z.hist2(sig[mask, ch_i, cy_i], noi[mask, ch_i, cy_i], __bins_x=(-1000, 15_000, 100), __bins_y=(-200, 1500, 100), f_title="Sig v Noi")
#             z.hist2(sig[mask, ch_i, cy_i], clr[mask, ch_i, cy_i], __bins_x=(-1000, 15_000, 100), __bins_y=(-4, 4, 100), f_title="Sig v clr")
#             z.scat(x=noi[mask, ch_i, cy_i], y=clr[mask, ch_i, cy_i], alpha=0.1)
#             z.scat(x=noi[mask, ch_i, cy_i], y=clr[mask, ch_i, cy_i], alpha=0.1)

for cy_i in range(n_cycles):
    analyze_cycle(cy_i)
# def contour_objective(f_conv, obj_fn):
#     max_pixel_shift = 2
#     n_steps = 31
#     contour = np.zeros((n_steps, n_steps))
#     for yi, y in enumerate(np.linspace(-max_pixel_shift, max_pixel_shift, n_steps)):
#         for xi, x in enumerate(np.linspace(-max_pixel_shift, max_pixel_shift, n_steps)):
#             obj = obj_fn(f_conv, np.array([y, x]))
#             contour[yi, xi] = obj
#     return contour
# # Inspect agree, disagree
# names = ("agree", "disagree")
# for name in names:
#     with z(_cols=4, _size=250):
#         with z(_merge=True, f_title=f"{name}"):
#             peak_im = np.load(f"_{name}_peak.npy")
#             z.im(peak_im)

#             history_jhd = np.load(f"_{name}_history_jhd.npy")

#             cen = hyper_peak_mea / 2
#             off_sci = np.load(f"_{name}_fit_sci.npy")
#             off_jhd = np.load(f"_{name}_fit_jhd.npy")
#             z.scat(x=[cen - off_sci[1]], y=[cen - off_sci[0]], color="red")
#             z.scat(x=[cen - off_jhd[1]], y=[cen - off_jhd[0]], color="orange")

#             for step in history_jhd:
#                 z.scat(x=[cen - step[1]], y=[cen - step[0]], color="cyan")

#             z.scat(x=[cen - off_jhd[1]], y=[cen - off_jhd[0]], color="orange")
            
#         f_conv_donut = np.load(f"_{name}_f_conv_donut.npy")
#         c = contour_objective(f_conv_donut, objective_sci)
#         z.im(c, f_title=f"Donut contour objective sci on {name}")

#         f_conv = np.load(f"_{name}_f_conv.npy")
#         c = contour_objective(f_conv, objective_sci)
#         z.im(c, f_title=f"Contour objective sci on {name}")

#         c = contour_objective(f_conv, lambda *args: objective_jhd(*args)[0])
#         with z(_merge=True, f_title=f"Contour objective jhd on {name} with jhd history"):
#             z.im(c)
            
#             def convert(p):
#                 return (p + 2) * (31/4)
                
#             for step in history_jhd:
#                 z.scat(x=[convert(step[1])], y=[convert(step[0])], color="cyan")
#             # spy(history_jhd)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.09517996651785714, 'fraction_numerical': 0.013985770089285714, 'mean_word_length': 2.6470363775120833, 'pattern_counts': {'":': 0, '<': 12, '<?xml version=': 0, '>': 8, 'https://': 1, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '1', 'dir': 'jupyter-structured-clean-dedup', 'id': '394738', 'license': 'MIT', 'n_tokens_mistral': 10233, 'n_tokens_neox': 9704, 'n_words': 2744, 'path': 'plaster/explore/cooley.ipynb', 'repo_name': 'erisyon/plaster', 'size': '28672'}