from __future__ import absolute_import
from unittest import TestCase
import numpy as np
import chimenea.pbcor as pbcor

class TestRadiusMap(TestCase):
    def test_radius_from_origin(self):
        shape = (3,3)
        centre = (0,0)
        map = pbcor._pixel_radius_map(shape,centre)
        self.assertEqual(map[0][0],0)
        self.assertEqual(map[0][1],1)
        self.assertEqual(map[1][1],np.sqrt(2))
        self.assertEqual(map[2][2],np.sqrt(8))
        # print
        # print map
    def test_radius_from_centre_odd_dims(self):
        shape = (5,5)
        centre = pbcor._central_position(shape)
        map = pbcor._pixel_radius_map(shape,centre)
        self.assertEqual(map[(2,2)], 0.)

    def test_radius_from_centre_even_dims(self):
        shape = (6,6)
        centre = pbcor._central_position(shape)
        map = pbcor._pixel_radius_map(shape,centre)
        self.assertEqual(map[(2,2)], np.sqrt(0.5))

class TestCorrectionMap(TestCase):
    def test_gaussian(self):
        # shape = (13,13)
        # centre = (6,6)
        shape = (3,3)
        centre = (1,1)
        curve = pbcor.GaussianCurve(sigma=1.)
        map = pbcor._correction_map(curve.correction,
                                    shape,centre)

        self.assertEqual(map[centre], 1.0)
        self.assertEqual(map[0][centre[1]], curve.correction(1))

class TestMask(TestCase):
    def test_gaussian(self):
        shape = (13,13)
        centre = (6,6)

        mask = pbcor.make_mask(shape,centre,radius_pix=1)
        print "MASK:"
        print mask

@@ -1101,35 +1101,55 @@
 e = 
-pbcor.GaussianCurve(sigma=1
+lambda radius_pix: np.exp( -(radius_pix)**2 / 2
 .)%0A 
@@ -1188,27 +1188,16 @@
 ap(curve
-.correction
 ,%0A      
@@ -1337,19 +1337,8 @@
 urve
-.correction
 (1))
@@ -1486,16 +1486,23 @@
 ,centre,
+cutoff_
 radius_p
@@ -1510,24 +1510,26 @@
 x=1)%0A       
+ #
  print %22MASK
@@ -1538,16 +1538,18 @@
 %0A       
+ #
  print m

Fix pbcor unittests