"""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()

[((17, 6, 17, 33), 'logging.getLogger', 'logging.getLogger', ({(17, 24, 17, 32): '__name__'}, {}), '(__name__)', False, 'import logging\n'), ((133, 4, 133, 19), 'unittest.main', 'unittest.main', ({}, {}), '()', False, 'import unittest\n'), ((57, 15, 57, 51), 'processing_components.image.operations.create_empty_image_like', 'create_empty_image_like', ({(57, 39, 57, 50): 'm31original'}, {}), '(m31original)', False, 'from processing_components.image.operations import create_empty_image_like\n'), ((72, 15, 72, 51), 'processing_components.image.operations.create_empty_image_like', 'create_empty_image_like', ({(72, 39, 72, 50): 'm31original'}, {}), '(m31original)', False, 'from processing_components.image.operations import create_empty_image_like\n'), ((88, 15, 88, 51), 'processing_components.image.operations.create_empty_image_like', 'create_empty_image_like', ({(88, 39, 88, 50): 'm31original'}, {}), '(m31original)', False, 'from processing_components.image.operations import create_empty_image_like\n'), ((104, 15, 104, 51), 'processing_components.image.operations.create_empty_image_like', 'create_empty_image_like', ({(104, 39, 104, 50): 'm31original'}, {}), '(m31original)', False, 'from processing_components.image.operations import create_empty_image_like\n'), ((24, 25, 24, 52), 'numpy.abs', 'numpy.abs', ({(24, 35, 24, 51): 'm31original.data'}, {}), '(m31original.data)', False, 'import numpy\n'), ((28, 25, 28, 68), 'processing_components.image.iterators.image_raster_iter', 'image_raster_iter', (), '', False, 'from processing_components.image.iterators import image_raster_iter, image_channel_iter, image_null_iter\n'), ((44, 25, 44, 52), 'numpy.abs', 'numpy.abs', ({(44, 35, 44, 51): 'm31original.data'}, {}), '(m31original.data)', False, 'import numpy\n'), ((56, 25, 56, 52), 'numpy.abs', 'numpy.abs', ({(56, 35, 56, 51): 'm31original.data'}, {}), '(m31original.data)', False, 'import numpy\n'), ((71, 25, 71, 52), 'numpy.abs', 'numpy.abs', ({(71, 35, 71, 51): 'm31original.data'}, {}), '(m31original.data)', False, 'import numpy\n'), ((87, 25, 87, 52), 'numpy.abs', 'numpy.abs', ({(87, 35, 87, 51): 'm31original.data'}, {}), '(m31original.data)', False, 'import numpy\n'), ((103, 25, 103, 52), 'numpy.abs', 'numpy.abs', ({(103, 35, 103, 51): 'm31original.data'}, {}), '(m31original.data)', False, 'import numpy\n'), ((121, 24, 121, 72), 'processing_components.image.iterators.image_channel_iter', 'image_channel_iter', (), '', False, 'from processing_components.image.iterators import image_raster_iter, image_channel_iter, image_null_iter\n'), ((128, 31, 128, 55), 'processing_components.image.iterators.image_null_iter', 'image_null_iter', ({(128, 47, 128, 54): 'm31cube'}, {}), '(m31cube)', False, 'from processing_components.image.iterators import image_raster_iter, image_channel_iter, image_null_iter\n'), ((23, 59, 23, 87), 'data_models.polarisation.PolarisationFrame', 'PolarisationFrame', ({(23, 77, 23, 86): '"""stokesI"""'}, {}), "('stokesI')", False, 'from data_models.polarisation import PolarisationFrame\n'), ((38, 29, 38, 53), 'numpy.abs', 'numpy.abs', ({(38, 39, 38, 52): 'm31model.data'}, {}), '(m31model.data)', False, 'import numpy\n'), ((43, 59, 43, 87), 'data_models.polarisation.PolarisationFrame', 'PolarisationFrame', ({(43, 77, 43, 86): '"""stokesI"""'}, {}), "('stokesI')", False, 'from data_models.polarisation import PolarisationFrame\n'), ((50, 29, 50, 89), 'processing_components.image.iterators.image_raster_iter', 'image_raster_iter', (), '', False, 'from processing_components.image.iterators import image_raster_iter, image_channel_iter, image_null_iter\n'), ((55, 59, 55, 87), 'data_models.polarisation.PolarisationFrame', 'PolarisationFrame', ({(55, 77, 55, 86): '"""stokesI"""'}, {}), "('stokesI')", False, 'from data_models.polarisation import PolarisationFrame\n'), ((61, 41, 61, 101), 'processing_components.image.iterators.image_raster_iter', 'image_raster_iter', (), '', False, 'from processing_components.image.iterators import image_raster_iter, image_channel_iter, image_null_iter\n'), ((62, 41, 62, 97), 'processing_components.image.iterators.image_raster_iter', 'image_raster_iter', (), '', False, 'from processing_components.image.iterators import image_raster_iter, image_channel_iter, image_null_iter\n'), ((66, 29, 66, 53), 'numpy.abs', 'numpy.abs', ({(66, 39, 66, 52): 'm31model.data'}, {}), '(m31model.data)', False, 'import numpy\n'), ((70, 59, 70, 87), 'data_models.polarisation.PolarisationFrame', 'PolarisationFrame', ({(70, 77, 70, 86): '"""stokesI"""'}, {}), "('stokesI')", False, 'from data_models.polarisation import PolarisationFrame\n'), ((76, 41, 77, 74), 'processing_components.image.iterators.image_raster_iter', 'image_raster_iter', (), '', False, 'from processing_components.image.iterators import image_raster_iter, image_channel_iter, image_null_iter\n'), ((78, 41, 78, 97), 'processing_components.image.iterators.image_raster_iter', 'image_raster_iter', (), '', False, 'from processing_components.image.iterators import image_raster_iter, image_channel_iter, image_null_iter\n'), ((82, 29, 82, 53), 'numpy.abs', 'numpy.abs', ({(82, 39, 82, 52): 'm31model.data'}, {}), '(m31model.data)', False, 'import numpy\n'), ((86, 59, 86, 87), 'data_models.polarisation.PolarisationFrame', 'PolarisationFrame', ({(86, 77, 86, 86): '"""stokesI"""'}, {}), "('stokesI')", False, 'from data_models.polarisation import PolarisationFrame\n'), ((92, 41, 93, 77), 'processing_components.image.iterators.image_raster_iter', 'image_raster_iter', (), '', False, 'from processing_components.image.iterators import image_raster_iter, image_channel_iter, image_null_iter\n'), ((94, 41, 94, 97), 'processing_components.image.iterators.image_raster_iter', 'image_raster_iter', (), '', False, 'from processing_components.image.iterators import image_raster_iter, image_channel_iter, image_null_iter\n'), ((98, 29, 98, 53), 'numpy.abs', 'numpy.abs', ({(98, 39, 98, 52): 'm31model.data'}, {}), '(m31model.data)', False, 'import numpy\n'), ((102, 59, 102, 87), 'data_models.polarisation.PolarisationFrame', 'PolarisationFrame', ({(102, 77, 102, 86): '"""stokesI"""'}, {}), "('stokesI')", False, 'from data_models.polarisation import PolarisationFrame\n'), ((108, 41, 109, 73), 'processing_components.image.iterators.image_raster_iter', 'image_raster_iter', (), '', False, 'from processing_components.image.iterators import image_raster_iter, image_channel_iter, image_null_iter\n'), ((110, 41, 110, 97), 'processing_components.image.iterators.image_raster_iter', 'image_raster_iter', (), '', False, 'from processing_components.image.iterators import image_raster_iter, image_channel_iter, image_null_iter\n'), ((114, 29, 114, 53), 'numpy.abs', 'numpy.abs', ({(114, 39, 114, 52): 'm31model.data'}, {}), '(m31model.data)', False, 'import numpy\n'), ((117, 55, 117, 83), 'data_models.polarisation.PolarisationFrame', 'PolarisationFrame', ({(117, 73, 117, 82): '"""stokesI"""'}, {}), "('stokesI')", False, 'from data_models.polarisation import PolarisationFrame\n'), ((118, 50, 118, 80), 'numpy.linspace', 'numpy.linspace', ({(118, 65, 118, 68): '100000000.0', (118, 69, 118, 74): '110000000.0', (118, 76, 118, 79): '128'}, {}), '(100000000.0, 110000000.0, 128)', False, 'import numpy\n'), ((125, 55, 125, 83), 'data_models.polarisation.PolarisationFrame', 'PolarisationFrame', ({(125, 73, 125, 82): '"""stokesI"""'}, {}), "('stokesI')", False, 'from data_models.polarisation import PolarisationFrame\n'), ((126, 46, 126, 77), 'numpy.linspace', 'numpy.linspace', ({(126, 61, 126, 64): '100000000.0', (126, 66, 126, 71): '110000000.0', (126, 73, 126, 76): '128'}, {}), '(100000000.0, 110000000.0, 128)', False, 'import numpy\n'), ((27, 60, 27, 88), 'data_models.polarisation.PolarisationFrame', 'PolarisationFrame', ({(27, 78, 27, 87): '"""stokesI"""'}, {}), "('stokesI')", False, 'from data_models.polarisation import PolarisationFrame\n'), ((39, 29, 39, 44), 'numpy.abs', 'numpy.abs', ({(39, 39, 39, 43): 'diff'}, {}), '(diff)', False, 'import numpy\n'), ((60, 60, 60, 88), 'data_models.polarisation.PolarisationFrame', 'PolarisationFrame', ({(60, 78, 60, 87): '"""stokesI"""'}, {}), "('stokesI')", False, 'from data_models.polarisation import PolarisationFrame\n'), ((75, 60, 75, 88), 'data_models.polarisation.PolarisationFrame', 'PolarisationFrame', ({(75, 78, 75, 87): '"""stokesI"""'}, {}), "('stokesI')", False, 'from data_models.polarisation import PolarisationFrame\n'), ((91, 60, 91, 88), 'data_models.polarisation.PolarisationFrame', 'PolarisationFrame', ({(91, 78, 91, 87): '"""stokesI"""'}, {}), "('stokesI')", False, 'from data_models.polarisation import PolarisationFrame\n'), ((107, 60, 107, 88), 'data_models.polarisation.PolarisationFrame', 'PolarisationFrame', ({(107, 78, 107, 87): '"""stokesI"""'}, {}), "('stokesI')", False, 'from data_models.polarisation import PolarisationFrame\n'), ((49, 64, 49, 92), 'data_models.polarisation.PolarisationFrame', 'PolarisationFrame', ({(49, 82, 49, 91): '"""stokesI"""'}, {}), "('stokesI')", False, 'from data_models.polarisation import PolarisationFrame\n')]