<gh_stars>0
import numpy as np
from plaster.run.priors import PriorsMLEFixtures
from plaster.run.sigproc_v2 import synth
from plaster.run.sigproc_v2.psf import approximate_psf
from plaster.run.sigproc_v2.align_ims import align_ims, resample_aligned_ims
from plaster.run.sigproc_v2.peak_find import pixel_peak_find_one_im
from plaster.tools.image import imops
from plaster.tools.image.coord import XY
from plaster.tools.zlog.zlog import spy
from plaster.tools.zlog.profile import prof
from plaster.tools.zlog import profile
from zest import zest


def zest_align_cy_ims():
    def _synth(n_peaks=500):
        with synth.Synth(n_channels=1, n_cycles=3, dim=(512, 512)) as s:
            reg_psf = PriorsMLEFixtures.reg_psf_uniform()
            (
                synth.PeaksModelPSF(reg_psf, n_peaks=n_peaks)
                .amps_constant(5000)
                .locs_randomize()
            )
            synth.CameraModel(100, 30)
            synth.HaloModel()
            return s.aln_offsets, s.render_chcy()

    def it_handles_clean():
        true_aln_offsets, chcy_ims = _synth()
        pred_aln_offsets = align_ims(chcy_ims[0])
        assert np.all(np.abs(true_aln_offsets[None, :, :] - pred_aln_offsets) < 0.1)

    @zest.retry(3)
    def it_handles_smudged():
        true_aln_offsets, chcy_ims = _synth()
        smudge_im = np.full((400, 400), 1000)
        imops.accum_inplace(chcy_ims[0, 0], smudge_im, loc=XY(100, 100))
        pred_cy_offsets = align_ims(chcy_ims[0])
        assert np.all(np.abs(true_aln_offsets - pred_cy_offsets) < 0.16)

    def it_handles_unalignable():
        cy_ims = np.zeros((3, 512, 512))
        pred_cy_offsets = align_ims(cy_ims)
        assert np.all(pred_cy_offsets[0, :] == 0.0)
        assert np.all(pred_cy_offsets[1:, :] >= 512.0)

    zest()


def zest_resample_aligned_ims():
    def _synth():
        ch_scale = (1.0, 1.0)
        ch_aln = np.array([0.0, 0.0, -5.0, -5.0]).reshape((-1, 2))
        with synth.Synth(n_channels=2, n_cycles=3, dim=(512, 512)) as s:
            s.channel_aln_offsets(ch_aln)
            (
                synth.PeaksModelGaussianCircular(n_peaks=500)
                .uniform_width_and_heights()
                .amps_constant(5000)
                .locs_randomize()
                .channel_scale_factor(ch_scale)
            )
            return s.aln_offsets, s.render_chcy()

    def it_corrects_channels():
        true_aln_offsets, chcy_ims = _synth()

        ch_offsets = align_ims(chcy_ims[:, 0], None)

        cy_offsets_0 = align_ims(chcy_ims[0, :], None)
        cy_offsets_1 = align_ims(chcy_ims[1, :], None)
        chcy_offsets = np.stack((cy_offsets_0, cy_offsets_1))

        chcy_offsets = chcy_offsets + ch_offsets[:, None, :]

        aligned_chcy_ims = resample_aligned_ims(chcy_ims, chcy_offsets)

        # If they mis-align then there will be a lot of
        # new peaks that weren't there. We generated 500 so the
        # assert is that there should be less than that on a combined image

        sum_im = np.sum(aligned_chcy_ims[:, 0], axis=0)

        # Add a little noise so that the pixel_peak_find_one_im can
        # measure the snr without a divide by zero.
        sum_im = sum_im + np.random.normal(0, 1.0, sum_im.shape)

        locs = pixel_peak_find_one_im(sum_im, approximate_psf())

        assert locs.shape[0] <= 500

    zest()


@zest.skip("TODO: Reinstate this. For now it is too slow and failing")
@zest.group("integration")
def zest_align_speed_tests():
    """
    Used to optimize the aligner.
    The test case is a large many cycle field like a photobleaching movie
    """

    def _synth(im_mea, n_cycles):
        """
        No noise or background to simulate post-filtering images
        """
        dim = (im_mea, im_mea)
        with synth.Synth(n_channels=1, n_cycles=n_cycles, dim=dim) as s:
            (
                synth.PeaksModelGaussianCircular(n_peaks=int(0.005 * im_mea ** 2))
                .uniform_width_and_heights()
                .amps_constant(5000)
                .locs_randomize()
            )
            return s.aln_offsets, s.render_chcy()

    @zest.retry(n_tries=2)
    def profile_aligns():
        true_aln_offsets, chcy_ims = _synth(1024, 10)

        profile.prof_start_record_capture()

        with prof("alignment"):
            with prof("find_offsets"):
                pred_cy_offsets = align_ims(chcy_ims[0])
                pred_chcy_offsets = pred_cy_offsets.reshape(1, *pred_cy_offsets.shape)

            assert np.all(np.abs(pred_chcy_offsets - true_aln_offsets) < 0.06)

            with prof("resample"):
                resample_aligned_ims(chcy_ims, pred_chcy_offsets)

        profile_records = profile.prof_stop_record_capture()
        assert profile_records[-1].name == "alignment"
        assert profile_records[-1].elapsed < 10.0

    zest()

import numpy as np
from plaster.run.priors import PriorsMLEFixtures
from plaster.run.sigproc_v2 import synth
from plaster.run.sigproc_v2.psf import approximate_psf
from plaster.run.sigproc_v2.align_ims import align_ims, resample_aligned_ims
from plaster.run.sigproc_v2.peak_find import pixel_peak_find_one_im
from plaster.tools.image import imops
from plaster.tools.image.coord import XY
from plaster.tools.zlog.zlog import spy
from plaster.tools.zlog.profile import prof
from plaster.tools.zlog import profile
from zest import zest


def zest_align_cy_ims():
    def _synth(n_peaks=500):
        with synth.Synth(n_channels=1, n_cycles=3, dim=(512, 512)) as s:
            reg_psf = PriorsMLEFixtures.reg_psf_uniform()
            (
                synth.PeaksModelPSF(reg_psf, n_peaks=n_peaks)
                .amps_constant(5000)
                .locs_randomize()
            )
            synth.CameraModel(100, 30)
            synth.HaloModel()
            return s.aln_offsets, s.render_chcy()

    def it_handles_clean():
        true_aln_offsets, chcy_ims = _synth()
        pred_aln_offsets = align_ims(chcy_ims[0])
        assert np.all(np.abs(true_aln_offsets[None, :, :] - pred_aln_offsets) < 0.1)

    @zest.retry(3)
    def it_handles_smudged():
        true_aln_offsets, chcy_ims = _synth()
        smudge_im = np.full((400, 400), 1000)
        imops.accum_inplace(chcy_ims[0, 0], smudge_im, loc=XY(100, 100))
        pred_cy_offsets = align_ims(chcy_ims[0])
        assert np.all(np.abs(true_aln_offsets - pred_cy_offsets) < 0.16)

    def it_handles_unalignable():
        cy_ims = np.zeros((3, 512, 512))
        pred_cy_offsets = align_ims(cy_ims)
        assert np.all(pred_cy_offsets[0, :] == 0.0)
        assert np.all(pred_cy_offsets[1:, :] >= 512.0)

    zest()


def zest_resample_aligned_ims():
    def _synth():
        ch_scale = (1.0, 1.0)
        ch_aln = np.array([0.0, 0.0, -5.0, -5.0]).reshape((-1, 2))
        with synth.Synth(n_channels=2, n_cycles=3, dim=(512, 512)) as s:
            s.channel_aln_offsets(ch_aln)
            (
                synth.PeaksModelGaussianCircular(n_peaks=500)
                .uniform_width_and_heights()
                .amps_constant(5000)
                .locs_randomize()
                .channel_scale_factor(ch_scale)
            )
            return s.aln_offsets, s.render_chcy()

    def it_corrects_channels():
        true_aln_offsets, chcy_ims = _synth()

        ch_offsets = align_ims(chcy_ims[:, 0], None)

        cy_offsets_0 = align_ims(chcy_ims[0, :], None)
        cy_offsets_1 = align_ims(chcy_ims[1, :], None)
        chcy_offsets = np.stack((cy_offsets_0, cy_offsets_1))

        chcy_offsets = chcy_offsets + ch_offsets[:, None, :]

        aligned_chcy_ims = resample_aligned_ims(chcy_ims, chcy_offsets)

        # If they mis-align then there will be a lot of
        # new peaks that weren't there. We generated 500 so the
        # assert is that there should be less than that on a combined image

        sum_im = np.sum(aligned_chcy_ims[:, 0], axis=0)

        # Add a little noise so that the pixel_peak_find_one_im can
        # measure the snr without a divide by zero.
        sum_im = sum_im + np.random.normal(0, 1.0, sum_im.shape)

        locs = pixel_peak_find_one_im(sum_im, approximate_psf())

        assert locs.shape[0] <= 500

    zest()


@zest.skip("TODO: Reinstate this. For now it is too slow and failing")
@zest.group("integration")
def zest_align_speed_tests():
    """
    Used to optimize the aligner.
    The test case is a large many cycle field like a photobleaching movie
    """

    def _synth(im_mea, n_cycles):
        """
        No noise or background to simulate post-filtering images
        """
        dim = (im_mea, im_mea)
        with synth.Synth(n_channels=1, n_cycles=n_cycles, dim=dim) as s:
            (
                synth.PeaksModelGaussianCircular(n_peaks=int(0.005 * im_mea ** 2))
                .uniform_width_and_heights()
                .amps_constant(5000)
                .locs_randomize()
            )
            return s.aln_offsets, s.render_chcy()

    @zest.retry(n_tries=2)
    def profile_aligns():
        true_aln_offsets, chcy_ims = _synth(1024, 10)

        profile.prof_start_record_capture()

        with prof("alignment"):
            with prof("find_offsets"):
                pred_cy_offsets = align_ims(chcy_ims[0])
                pred_chcy_offsets = pred_cy_offsets.reshape(1, *pred_cy_offsets.shape)

            assert np.all(np.abs(pred_chcy_offsets - true_aln_offsets) < 0.06)

            with prof("resample"):
                resample_aligned_ims(chcy_ims, pred_chcy_offsets)

        profile_records = profile.prof_stop_record_capture()
        assert profile_records[-1].name == "alignment"
        assert profile_records[-1].elapsed < 10.0

    zest()
# If they mis-align then there will be a lot of # new peaks that weren't there. We generated 500 so the # assert is that there should be less than that on a combined image # Add a little noise so that the pixel_peak_find_one_im can # measure the snr without a divide by zero. Used to optimize the aligner.     The test case is a large many cycle field like a photobleaching movie No noise or background to simulate post-filtering images