@@ -0,0 +1,2879 @@
+# Licensed under a 3-clause BSD style license - see LICENSE.rst%0Afrom __future__ import (absolute_import, division, print_function,%0A                        unicode_literals)%0Aimport numpy as np%0Afrom numpy.testing import assert_allclose%0Afrom astropy.tests.helper import pytest%0Afrom astropy.modeling import models%0Afrom ..core import detect_sources%0Afrom ..deblend import deblend_sources%0A%0Atry:%0A    import scipy%0A    HAS_SCIPY = True%0Aexcept ImportError:%0A    HAS_SCIPY = False%0A%0Atry:%0A    import skimage%0A    HAS_SKIMAGE = True%0Aexcept ImportError:%0A    HAS_SKIMAGE = False%0A%0A%0A@pytest.mark.skipif('not HAS_SCIPY')%0A@pytest.mark.skipif('not HAS_SKIMAGE')%0Aclass TestDeblendSources(object):%0A    def setup_class(self):%0A        g1 = models.Gaussian2D(100, 50, 50, 5, 5)%0A        g2 = models.Gaussian2D(100, 35, 50, 5, 5)%0A        g3 = models.Gaussian2D(30, 70, 50, 5, 5)%0A        y, x = np.mgrid%5B0:100, 0:100%5D%0A        self.data = g1(x, y) + g2(x, y)%0A        self.data3 = self.data + g3(x, y)%0A        self.threshold = 10%0A        self.npixels = 5%0A        self.segm = detect_sources(self.data, self.threshold, self.npixels)%0A        self.segm3 = detect_sources(self.data3, self.threshold, self.npixels)%0A%0A    @pytest.mark.parametrize('mode', %5B'exponential', 'linear'%5D)%0A    def test_deblend_sources(self, mode):%0A        result = deblend_sources(self.data, self.segm, self.npixels,%0A                                 mode=mode)%0A        assert result.nlabels == 2%0A        mask1 = (result.data == 1)%0A        mask2 = (result.data == 2)%0A        assert_allclose(len(result.data%5Bmask1%5D), len(result.data%5Bmask2%5D))%0A        assert_allclose(np.sum(self.data%5Bmask1%5D), np.sum(self.data%5Bmask2%5D))%0A        assert_allclose(np.nonzero(self.segm), np.nonzero(result))%0A%0A    @pytest.mark.parametrize('mode', %5B'exponential', 'linear'%5D)%0A    def test_deblend_three_sources(self, mode):%0A        result = deblend_sources(self.data3, self.segm3, self.npixels,%0A                                 mode=mode)%0A        assert result.nlabels == 3%0A        assert_allclose(np.nonzero(self.segm3), np.nonzero(result))%0A%0A    def test_deblend_sources_segm_array(self):%0A        result = deblend_sources(self.data, self.segm.data, self.npixels)%0A        assert result.nlabels == 2%0A%0A    def test_segment_img_badshape(self):%0A        segm_wrong = np.zeros((2, 2))%0A        with pytest.raises(ValueError):%0A            deblend_sources(self.data, segm_wrong, self.npixels)%0A%0A    def test_invalid_nlevels(self):%0A        with pytest.raises(ValueError):%0A            deblend_sources(self.data, self.segm, self.npixels, nlevels=0)%0A%0A    def test_invalid_contrast(self):%0A        with pytest.raises(ValueError):%0A            deblend_sources(self.data, self.segm, self.npixels, contrast=-1)%0A%0A    def test_invalid_mode(self):%0A        with pytest.raises(ValueError):%0A            deblend_sources(self.data, self.segm, self.npixels,%0A                            mode='invalid')%0A

Add tests for deblend_sources