"""Unit tests for image iteration


"""
import logging
import unittest

import numpy

from data_models.polarisation import PolarisationFrame

from processing_components.image.iterators import image_raster_iter, image_channel_iter, image_null_iter

from processing_components.image.operations import create_empty_image_like
from processing_components.simulation.testing_support import create_test_image

log = logging.getLogger(__name__)


class TestImageIterators(unittest.TestCase):
    def test_raster(self):
        
        m31original = create_test_image(polarisation_frame=PolarisationFrame('stokesI'))
        assert numpy.max(numpy.abs(m31original.data)), "Original is empty"
        
        for nraster in [1, 2, 4, 8, 9]:
            m31model = create_test_image(polarisation_frame=PolarisationFrame('stokesI'))
            for patch in image_raster_iter(m31model, facets=nraster):
                assert patch.data.shape[3] == (m31model.data.shape[3] // nraster), \
                    "Number of pixels in each patch: %d not as expected: %d" % (patch.data.shape[3],
                                                                                (m31model.data.shape[3] // nraster))
                assert patch.data.shape[2] == (m31model.data.shape[2] // nraster), \
                    "Number of pixels in each patch: %d not as expected: %d" % (patch.data.shape[2],
                                                                                (m31model.data.shape[2] // nraster))
                patch.data *= 2.0
            
            diff = m31model.data - 2.0 * m31original.data
            assert numpy.max(numpy.abs(m31model.data)), "Raster is empty for %d" % nraster
            assert numpy.max(numpy.abs(diff)) == 0.0, "Raster set failed for %d" % nraster

    def test_raster_exception(self):
    
        m31original = create_test_image(polarisation_frame=PolarisationFrame('stokesI'))
        assert numpy.max(numpy.abs(m31original.data)), "Original is empty"
    
        for nraster, overlap in [(-1, -1), (-1, 0), (2, 128), (1e6, 127)]:
            
            with self.assertRaises(AssertionError) as context:
                m31model = create_test_image(polarisation_frame=PolarisationFrame('stokesI'))
                for patch in image_raster_iter(m31model, facets=nraster, overlap=overlap):
                    patch.data *= 2.0
            
    def test_raster_overlap(self):
    
        m31original = create_test_image(polarisation_frame=PolarisationFrame('stokesI'))
        assert numpy.max(numpy.abs(m31original.data)), "Original is empty"
        flat = create_empty_image_like(m31original)
    
        for nraster, overlap in [(1, 0),  (1, 16), (4, 8), (4, 16), (8, 8), (16, 4), (9, 5)]:
            m31model = create_test_image(polarisation_frame=PolarisationFrame('stokesI'))
            for patch, flat_patch in zip(image_raster_iter(m31model, facets=nraster, overlap=overlap),
                                         image_raster_iter(flat, facets=nraster, overlap=overlap)):
                patch.data *= 2.0
                flat_patch.data[...] += 1.0
        
            assert numpy.max(numpy.abs(m31model.data)), "Raster is empty for %d" % nraster

    def test_raster_overlap_linear(self):
    
        m31original = create_test_image(polarisation_frame=PolarisationFrame('stokesI'))
        assert numpy.max(numpy.abs(m31original.data)), "Original is empty"
        flat = create_empty_image_like(m31original)
    
        for nraster, overlap in [(1, 0), (1, 16), (4, 8), (4, 16), (8, 8), (16, 4), (9, 5)]:
            m31model = create_test_image(polarisation_frame=PolarisationFrame('stokesI'))
            for patch, flat_patch in zip(image_raster_iter(m31model, facets=nraster, overlap=overlap,
                                                           taper='linear'),
                                         image_raster_iter(flat, facets=nraster, overlap=overlap)):
                patch.data *= 2.0
                flat_patch.data[...] += 1.0
        
            assert numpy.max(numpy.abs(m31model.data)), "Raster is empty for %d" % nraster

    def test_raster_overlap_quadratic(self):
    
        m31original = create_test_image(polarisation_frame=PolarisationFrame('stokesI'))
        assert numpy.max(numpy.abs(m31original.data)), "Original is empty"
        flat = create_empty_image_like(m31original)
    
        for nraster, overlap in [(1, 0), (1, 16), (4, 8), (4, 16), (8, 8), (16, 4), (9, 5)]:
            m31model = create_test_image(polarisation_frame=PolarisationFrame('stokesI'))
            for patch, flat_patch in zip(image_raster_iter(m31model, facets=nraster, overlap=overlap,
                                                           taper='quadratic'),
                                         image_raster_iter(flat, facets=nraster, overlap=overlap)):
                patch.data *= 2.0
                flat_patch.data[...] += 1.0
        
            assert numpy.max(numpy.abs(m31model.data)), "Raster is empty for %d" % nraster

    def test_raster_overlap_tukey(self):
    
        m31original = create_test_image(polarisation_frame=PolarisationFrame('stokesI'))
        assert numpy.max(numpy.abs(m31original.data)), "Original is empty"
        flat = create_empty_image_like(m31original)
    
        for nraster, overlap in [(1, 0), (1, 16), (4, 8), (4, 16), (8, 8), (16, 4), (9, 5)]:
            m31model = create_test_image(polarisation_frame=PolarisationFrame('stokesI'))
            for patch, flat_patch in zip(image_raster_iter(m31model, facets=nraster, overlap=overlap,
                                                           taper='tukey'),
                                         image_raster_iter(flat, facets=nraster, overlap=overlap)):
                patch.data *= 2.0
                flat_patch.data[...] += 1.0
        
            assert numpy.max(numpy.abs(m31model.data)), "Raster is empty for %d" % nraster

    def test_channelise(self):
        m31cube = create_test_image(polarisation_frame=PolarisationFrame('stokesI'),
                                        frequency=numpy.linspace(1e8,1.1e8, 128))
        
        for subimages in [128, 16, 8, 2, 1]:
            for slab in image_channel_iter(m31cube, subimages=subimages):
                assert slab.data.shape[0] == 128 // subimages

    def test_null(self):
        m31cube = create_test_image(polarisation_frame=PolarisationFrame('stokesI'),
                                    frequency=numpy.linspace(1e8, 1.1e8, 128))
    
        for i, im in enumerate(image_null_iter(m31cube)):
            assert i<1, "Null iterator returns more than one value"


if __name__ == '__main__':
    unittest.main()

[((409, 436), 'logging.getLogger', 'logging.getLogger', (['__name__'], {}), '(__name__)\n', (426, 436), False, 'import logging\n'), ((6682, 6697), 'unittest.main', 'unittest.main', ([], {}), '()\n', (6695, 6697), False, 'import unittest\n'), ((2604, 2640), 'processing_components.image.operations.create_empty_image_like', 'create_empty_image_like', (['m31original'], {}), '(m31original)\n', (2627, 2640), False, 'from processing_components.image.operations import create_empty_image_like\n'), ((3438, 3474), 'processing_components.image.operations.create_empty_image_like', 'create_empty_image_like', (['m31original'], {}), '(m31original)\n', (3461, 3474), False, 'from processing_components.image.operations import create_empty_image_like\n'), ((4349, 4385), 'processing_components.image.operations.create_empty_image_like', 'create_empty_image_like', (['m31original'], {}), '(m31original)\n', (4372, 4385), False, 'from processing_components.image.operations import create_empty_image_like\n'), ((5259, 5295), 'processing_components.image.operations.create_empty_image_like', 'create_empty_image_like', (['m31original'], {}), '(m31original)\n', (5282, 5295), False, 'from processing_components.image.operations import create_empty_image_like\n'), ((634, 661), 'numpy.abs', 'numpy.abs', (['m31original.data'], {}), '(m31original.data)\n', (643, 661), False, 'import numpy\n'), ((848, 891), 'processing_components.image.iterators.image_raster_iter', 'image_raster_iter', (['m31model'], {'facets': 'nraster'}), '(m31model, facets=nraster)\n', (865, 891), False, 'from processing_components.image.iterators import image_raster_iter, image_channel_iter, image_null_iter\n'), ((1943, 1970), 'numpy.abs', 'numpy.abs', (['m31original.data'], {}), '(m31original.data)\n', (1952, 1970), False, 'import numpy\n'), ((2539, 2566), 'numpy.abs', 'numpy.abs', (['m31original.data'], {}), '(m31original.data)\n', (2548, 2566), False, 'import numpy\n'), ((3373, 3400), 'numpy.abs', 'numpy.abs', (['m31original.data'], {}), '(m31original.data)\n', (3382, 3400), False, 'import numpy\n'), ((4284, 4311), 'numpy.abs', 'numpy.abs', (['m31original.data'], {}), '(m31original.data)\n', (4293, 4311), False, 'import numpy\n'), ((5194, 5221), 'numpy.abs', 'numpy.abs', (['m31original.data'], {}), '(m31original.data)\n', (5203, 5221), False, 'import numpy\n'), ((6216, 6264), 'processing_components.image.iterators.image_channel_iter', 'image_channel_iter', (['m31cube'], {'subimages': 'subimages'}), '(m31cube, subimages=subimages)\n', (6234, 6264), False, 'from processing_components.image.iterators import image_raster_iter, image_channel_iter, image_null_iter\n'), ((6554, 6578), 'processing_components.image.iterators.image_null_iter', 'image_null_iter', (['m31cube'], {}), '(m31cube)\n', (6569, 6578), False, 'from processing_components.image.iterators import image_raster_iter, image_channel_iter, image_null_iter\n'), ((579, 607), 'data_models.polarisation.PolarisationFrame', 'PolarisationFrame', (['"""stokesI"""'], {}), "('stokesI')\n", (596, 607), False, 'from data_models.polarisation import PolarisationFrame\n'), ((1633, 1657), 'numpy.abs', 'numpy.abs', (['m31model.data'], {}), '(m31model.data)\n', (1642, 1657), False, 'import numpy\n'), ((1888, 1916), 'data_models.polarisation.PolarisationFrame', 'PolarisationFrame', (['"""stokesI"""'], {}), "('stokesI')\n", (1905, 1916), False, 'from data_models.polarisation import PolarisationFrame\n'), ((2272, 2332), 'processing_components.image.iterators.image_raster_iter', 'image_raster_iter', (['m31model'], {'facets': 'nraster', 'overlap': 'overlap'}), '(m31model, facets=nraster, overlap=overlap)\n', (2289, 2332), False, 'from processing_components.image.iterators import image_raster_iter, image_channel_iter, image_null_iter\n'), ((2484, 2512), 'data_models.polarisation.PolarisationFrame', 'PolarisationFrame', (['"""stokesI"""'], {}), "('stokesI')\n", (2501, 2512), False, 'from data_models.polarisation import PolarisationFrame\n'), ((2871, 2931), 'processing_components.image.iterators.image_raster_iter', 'image_raster_iter', (['m31model'], {'facets': 'nraster', 'overlap': 'overlap'}), '(m31model, facets=nraster, overlap=overlap)\n', (2888, 2931), False, 'from processing_components.image.iterators import image_raster_iter, image_channel_iter, image_null_iter\n'), ((2974, 3030), 'processing_components.image.iterators.image_raster_iter', 'image_raster_iter', (['flat'], {'facets': 'nraster', 'overlap': 'overlap'}), '(flat, facets=nraster, overlap=overlap)\n', (2991, 3030), False, 'from processing_components.image.iterators import image_raster_iter, image_channel_iter, image_null_iter\n'), ((3149, 3173), 'numpy.abs', 'numpy.abs', (['m31model.data'], {}), '(m31model.data)\n', (3158, 3173), False, 'import numpy\n'), ((3318, 3346), 'data_models.polarisation.PolarisationFrame', 'PolarisationFrame', (['"""stokesI"""'], {}), "('stokesI')\n", (3335, 3346), False, 'from data_models.polarisation import PolarisationFrame\n'), ((3704, 3780), 'processing_components.image.iterators.image_raster_iter', 'image_raster_iter', (['m31model'], {'facets': 'nraster', 'overlap': 'overlap', 'taper': '"""linear"""'}), "(m31model, facets=nraster, overlap=overlap, taper='linear')\n", (3721, 3780), False, 'from processing_components.image.iterators import image_raster_iter, image_channel_iter, image_null_iter\n'), ((3882, 3938), 'processing_components.image.iterators.image_raster_iter', 'image_raster_iter', (['flat'], {'facets': 'nraster', 'overlap': 'overlap'}), '(flat, facets=nraster, overlap=overlap)\n', (3899, 3938), False, 'from processing_components.image.iterators import image_raster_iter, image_channel_iter, image_null_iter\n'), ((4057, 4081), 'numpy.abs', 'numpy.abs', (['m31model.data'], {}), '(m31model.data)\n', (4066, 4081), False, 'import numpy\n'), ((4229, 4257), 'data_models.polarisation.PolarisationFrame', 'PolarisationFrame', (['"""stokesI"""'], {}), "('stokesI')\n", (4246, 4257), False, 'from data_models.polarisation import PolarisationFrame\n'), ((4615, 4694), 'processing_components.image.iterators.image_raster_iter', 'image_raster_iter', (['m31model'], {'facets': 'nraster', 'overlap': 'overlap', 'taper': '"""quadratic"""'}), "(m31model, facets=nraster, overlap=overlap, taper='quadratic')\n", (4632, 4694), False, 'from processing_components.image.iterators import image_raster_iter, image_channel_iter, image_null_iter\n'), ((4796, 4852), 'processing_components.image.iterators.image_raster_iter', 'image_raster_iter', (['flat'], {'facets': 'nraster', 'overlap': 'overlap'}), '(flat, facets=nraster, overlap=overlap)\n', (4813, 4852), False, 'from processing_components.image.iterators import image_raster_iter, image_channel_iter, image_null_iter\n'), ((4971, 4995), 'numpy.abs', 'numpy.abs', (['m31model.data'], {}), '(m31model.data)\n', (4980, 4995), False, 'import numpy\n'), ((5139, 5167), 'data_models.polarisation.PolarisationFrame', 'PolarisationFrame', (['"""stokesI"""'], {}), "('stokesI')\n", (5156, 5167), False, 'from data_models.polarisation import PolarisationFrame\n'), ((5525, 5600), 'processing_components.image.iterators.image_raster_iter', 'image_raster_iter', (['m31model'], {'facets': 'nraster', 'overlap': 'overlap', 'taper': '"""tukey"""'}), "(m31model, facets=nraster, overlap=overlap, taper='tukey')\n", (5542, 5600), False, 'from processing_components.image.iterators import image_raster_iter, image_channel_iter, image_null_iter\n'), ((5702, 5758), 'processing_components.image.iterators.image_raster_iter', 'image_raster_iter', (['flat'], {'facets': 'nraster', 'overlap': 'overlap'}), '(flat, facets=nraster, overlap=overlap)\n', (5719, 5758), False, 'from processing_components.image.iterators import image_raster_iter, image_channel_iter, image_null_iter\n'), ((5877, 5901), 'numpy.abs', 'numpy.abs', (['m31model.data'], {}), '(m31model.data)\n', (5886, 5901), False, 'import numpy\n'), ((6026, 6054), 'data_models.polarisation.PolarisationFrame', 'PolarisationFrame', (['"""stokesI"""'], {}), "('stokesI')\n", (6043, 6054), False, 'from data_models.polarisation import PolarisationFrame\n'), ((6106, 6151), 'numpy.linspace', 'numpy.linspace', (['(100000000.0)', '(110000000.0)', '(128)'], {}), '(100000000.0, 110000000.0, 128)\n', (6120, 6151), False, 'import numpy\n'), ((6409, 6437), 'data_models.polarisation.PolarisationFrame', 'PolarisationFrame', (['"""stokesI"""'], {}), "('stokesI')\n", (6426, 6437), False, 'from data_models.polarisation import PolarisationFrame\n'), ((6485, 6530), 'numpy.linspace', 'numpy.linspace', (['(100000000.0)', '(110000000.0)', '(128)'], {}), '(100000000.0, 110000000.0, 128)\n', (6499, 6530), False, 'import numpy\n'), ((793, 821), 'data_models.polarisation.PolarisationFrame', 'PolarisationFrame', (['"""stokesI"""'], {}), "('stokesI')\n", (810, 821), False, 'from data_models.polarisation import PolarisationFrame\n'), ((1724, 1739), 'numpy.abs', 'numpy.abs', (['diff'], {}), '(diff)\n', (1733, 1739), False, 'import numpy\n'), ((2800, 2828), 'data_models.polarisation.PolarisationFrame', 'PolarisationFrame', (['"""stokesI"""'], {}), "('stokesI')\n", (2817, 2828), False, 'from data_models.polarisation import PolarisationFrame\n'), ((3633, 3661), 'data_models.polarisation.PolarisationFrame', 'PolarisationFrame', (['"""stokesI"""'], {}), "('stokesI')\n", (3650, 3661), False, 'from data_models.polarisation import PolarisationFrame\n'), ((4544, 4572), 'data_models.polarisation.PolarisationFrame', 'PolarisationFrame', (['"""stokesI"""'], {}), "('stokesI')\n", (4561, 4572), False, 'from data_models.polarisation import PolarisationFrame\n'), ((5454, 5482), 'data_models.polarisation.PolarisationFrame', 'PolarisationFrame', (['"""stokesI"""'], {}), "('stokesI')\n", (5471, 5482), False, 'from data_models.polarisation import PolarisationFrame\n'), ((2213, 2241), 'data_models.polarisation.PolarisationFrame', 'PolarisationFrame', (['"""stokesI"""'], {}), "('stokesI')\n", (2230, 2241), False, 'from data_models.polarisation import PolarisationFrame\n')]