test_numtraits.py

import pytest

import numpy as np
from numtraits import NumericalTrait

from traitlets import HasTraits, TraitError

class ScalarProperties(HasTraits):

    a = NumericalTrait(ndim=0)
    b = NumericalTrait(ndim=0, domain='positive')
    c = NumericalTrait(ndim=0, domain='strictly-positive')
    d = NumericalTrait(ndim=0, domain='negative')
    e = NumericalTrait(ndim=0, domain='strictly-negative')
    f = NumericalTrait(ndim=0, domain=(3, 4))

class TestScalar(object):

    def setup_method(self, method):

        # Make sure we have an instance of ScalarProperties each time
        # a testing method is called.
        self.sp = ScalarProperties()

    def test_simple(self):
        self.sp.a = 1.
        assert self.sp.a == 1.

    def test_scalar(self):
        with pytest.raises(TraitError) as exc:
            self.sp.a = [1, 2]
        assert exc.value.args[0] == "a should be a scalar value"

    def test_numerical(self):
        with pytest.raises(TraitError) as exc:
            self.sp.a = 'a'
        assert exc.value.args[0] == "a should be a numerical value"

    def test_positive(self):
        self.sp.b = 3.
        self.sp.b = 0.
        with pytest.raises(TraitError) as exc:
            self.sp.b = -2
        assert exc.value.args[0] == "b should be positive"

    def test_strictly_positive(self):
        self.sp.c = 3.
        with pytest.raises(TraitError) as exc:
            self.sp.c = 0.
        assert exc.value.args[0] == "c should be strictly positive"
        with pytest.raises(TraitError) as exc:
            self.sp.c = -2
        assert exc.value.args[0] == "c should be strictly positive"

    def test_negative(self):
        self.sp.d = -3.
        self.sp.d = 0.
        with pytest.raises(TraitError) as exc:
            self.sp.d = 2
        assert exc.value.args[0] == "d should be negative"

    def test_strictly_negative(self):
        self.sp.e = -2
        with pytest.raises(TraitError) as exc:
            self.sp.e = 0.
        assert exc.value.args[0] == "e should be strictly negative"
        with pytest.raises(TraitError) as exc:
            self.sp.e = 2
        assert exc.value.args[0] == "e should be strictly negative"

    def test_range(self):
        self.sp.f = 3.
        self.sp.f = 3.5
        self.sp.f = 4.
        with pytest.raises(TraitError) as exc:
            self.sp.f = 0.
        assert exc.value.args[0] == "f should be in the range [3:4]"
        with pytest.raises(TraitError) as exc:
            self.sp.f = 7
        assert exc.value.args[0] == "f should be in the range [3:4]"


class ArrayProperties(HasTraits):

    a = NumericalTrait(shape=(3,))
    b = NumericalTrait(domain='positive', ndim=1)
    c = NumericalTrait(domain='strictly-positive', ndim=1)
    d = NumericalTrait(domain='negative', ndim=1)
    e = NumericalTrait(domain='strictly-negative', ndim=1)
    f = NumericalTrait(domain=(3, 4), ndim=1)
    g = NumericalTrait(shape=(3, 4))

class TestArray(object):

    def setup_method(self, method):

        self.ap = ArrayProperties()

    def test_simple(self):
        self.ap.a = (1, 2, 3)
        self.ap.b = (1, 2, 3, 4)
        np.testing.assert_allclose(self.ap.a, (1, 2, 3))
        np.testing.assert_allclose(self.ap.b, (1, 2, 3, 4))
        assert isinstance(self.ap.a, np.ndarray)
        assert isinstance(self.ap.b, np.ndarray)

    def test_shape(self):
        with pytest.raises(TraitError) as exc:
            self.ap.a = (1, 2, 3, 4)
        assert exc.value.args[0] == "a has incorrect length (expected 3 but found 4)"

    def test_ndim(self):
        with pytest.raises(TraitError) as exc:
            self.ap.a = np.ones((3, 3))
        assert exc.value.args[0] == "a should be a 1-d sequence"

    def test_positive(self):
        self.ap.b = (0., 2., 3.)
        with pytest.raises(TraitError) as exc:
            self.ap.b = (0., -1., 3.)
        assert exc.value.args[0] == "All values of b should be positive"

    def test_strictly_positive(self):
        self.ap.c = (1., 2., 3.)
        with pytest.raises(TraitError) as exc:
            self.ap.c = (0., 2., 3.)
        assert exc.value.args[0] == "All values of c should be strictly positive"
        with pytest.raises(TraitError) as exc:
            self.ap.c = (0., -1., 3.)
        assert exc.value.args[0] == "All values of c should be strictly positive"

    def test_negative(self):
        self.ap.d = (-1., -2., 0.)
        with pytest.raises(TraitError) as exc:
            self.ap.d = (-1., -2., 1.)
        assert exc.value.args[0] == "All values of d should be negative"

    def test_strictly_negative(self):
        self.ap.e = (-1., -2., -3.)
        with pytest.raises(TraitError) as exc:
            self.ap.e = (-1., -2., 0.)
        assert exc.value.args[0] == "All values of e should be strictly negative"
        with pytest.raises(TraitError) as exc:
            self.ap.e = (-1., -2., 1.)
        assert exc.value.args[0] == "All values of e should be strictly negative"

    def test_range(self):
        self.ap.f = (3., 3.5, 4.)
        with pytest.raises(TraitError) as exc:
            self.ap.f = (0., 3.5, 4.)
        assert exc.value.args[0] == "All values of f should be in the range [3:4]"
        with pytest.raises(TraitError) as exc:
            self.ap.f = (0., 3.5, 7.)
        assert exc.value.args[0] == "All values of f should be in the range [3:4]"

    def test_shape_2d(self):
        self.ap.g = np.ones((3, 4))
        with pytest.raises(TraitError) as exc:
            self.ap.g = np.ones((3, 6))
        assert exc.value.args[0] == "g has incorrect shape (expected (3, 4) but found (3, 6))"

    def test_ndim_2d(self):
        with pytest.raises(TraitError) as exc:
            self.ap.g = np.ones((3, 6, 3))
        assert exc.value.args[0] == "g should be a 2-d array"

    def test_invalid(self):
        with pytest.raises(TraitError) as exc:
            self.ap.a = [[1.], [1., 2.]]
        assert exc.value.args[0] == "Could not convert value of a to a Numpy array (Exception: setting an array element with a sequence.)"

# Need to decide on behavior if passing a unit-ed quantity to a trait
# with no convertible_to argument.

try:
    from astropy import units as u
    from pint import UnitRegistry
    import quantities as pq
except ImportError:
    pass
else:

    ureg = UnitRegistry()

    class AstropyUnitsProperties(HasTraits):

        a = NumericalTrait(convertible_to=u.m)
        b = NumericalTrait(convertible_to=u.cm / u.s)

    class TestAstropyUnits(object):

        def setup_method(self, method):

            self.aup = AstropyUnitsProperties()

        def test_valid(self):

            self.aup.a = 3 * u.m
            self.aup.a = [1, 2, 3] * u.cm
            self.aup.a = np.ones((2, 2)) * u.pc

            self.aup.b = 3 * u.m / u.yr
            self.aup.b = [1, 2, 3] * u.cm / u.s
            self.aup.b = np.ones((2, 2)) * u.pc / u.Myr

        def test_invalid_type(self):

            with pytest.raises(TraitError) as exc:
                self.aup.a = 5
            assert exc.value.args[0] == 'a should be given as an Astropy Quantity instance'

            with pytest.raises(TraitError) as exc:
                self.aup.b = np.ones((2, 5))
            assert exc.value.args[0] == 'b should be given as an Astropy Quantity instance'

        def test_invalid_framework(self):

            # pint
            with pytest.raises(TraitError) as exc:
                self.aup.a = 5 * ureg.m
            assert exc.value.args[0] == 'a should be given as an Astropy Quantity instance'

            # quantities
            with pytest.raises(TraitError) as exc:
                self.aup.b = 3 * pq.cm
            assert exc.value.args[0] == 'b should be given as an Astropy Quantity instance'

        def test_invalid_units(self):

            with pytest.raises(TraitError) as exc:
                self.aup.a = 5 * u.s
            assert exc.value.args[0] == 'a should be in units convertible to m'

            with pytest.raises(TraitError) as exc:
                self.aup.b = np.ones((2, 5)) * u.s
            assert exc.value.args[0] == 'b should be in units convertible to cm / s'

    class PintUnitsProperties(HasTraits):

        a = NumericalTrait(convertible_to=ureg.m)
        b = NumericalTrait(convertible_to=ureg.cm / ureg.s)

    class TestPintUnits(object):

        def setup_method(self, method):

            self.pup = PintUnitsProperties()

        def test_valid(self):

            from astropy import units as u

            self.pup.a = 3 * ureg.m
            self.pup.a = [1, 2, 3] * ureg.cm
            self.pup.a = np.ones((2, 2)) * ureg.pc

            self.pup.b = 3 * ureg.m / ureg.year
            self.pup.b = [1, 2, 3] * ureg.cm / ureg.s
            self.pup.b = np.ones((2, 2)) * ureg.pc / ureg.megayear

        def test_invalid_type(self):

            with pytest.raises(TraitError) as exc:
                self.pup.a = 5
            assert exc.value.args[0] == 'a should be given as a Pint Quantity instance'

            with pytest.raises(TraitError) as exc:
                self.pup.b = np.ones((2, 5))
            assert exc.value.args[0] == 'b should be given as a Pint Quantity instance'

        def test_invalid_framework(self):

            # astropy.units
            with pytest.raises(TraitError) as exc:
                self.pup.b = 5 * u.m
            assert exc.value.args[0] == 'b should be given as a Pint Quantity instance'

            # quantitites
            with pytest.raises(TraitError) as exc:
                self.pup.b = 3 * pq.cm
            assert exc.value.args[0] == 'b should be given as a Pint Quantity instance'

        def test_invalid_units(self):

            from astropy import units as u

            with pytest.raises(TraitError) as exc:
                self.pup.a = 5 * ureg.s
            assert exc.value.args[0] == 'a should be in units convertible to meter'

            with pytest.raises(TraitError) as exc:
                self.pup.b = np.ones((2, 5)) * ureg.s
            assert exc.value.args[0] == 'b should be in units convertible to centimeter / second'

    class QuantitiesUnitsProperties(HasTraits):

        a = NumericalTrait(convertible_to=pq.m)
        b = NumericalTrait(convertible_to=pq.cm / pq.s)

    class TestQuantitiesUnits(object):

        def setup_method(self, method):

            self.qup = QuantitiesUnitsProperties()

        def test_valid(self):

            from astropy import units as u

            self.qup.a = 3 * pq.m
            self.qup.a = [1, 2, 3] * pq.cm
            self.qup.a = np.ones((2, 2)) * pq.pc

            self.qup.b = 3 * pq.m / pq.year
            self.qup.b = [1, 2, 3] * pq.cm / pq.s
            self.qup.b = np.ones((2, 2)) * pq.pc / pq.s

        def test_invalid_type(self):

            with pytest.raises(TraitError) as exc:
                self.qup.a = 5
            assert exc.value.args[0] == 'a should be given as a quantities Quantity instance'

            with pytest.raises(TraitError) as exc:
                self.qup.b = np.ones((2, 5))
            assert exc.value.args[0] == 'b should be given as a quantities Quantity instance'

        def test_invalid_framework(self):

            # astropy.units
            with pytest.raises(TraitError) as exc:
                self.qup.a = 5 * u.m
            assert exc.value.args[0] == 'a should be given as a quantities Quantity instance'

            # pint
            with pytest.raises(TraitError) as exc:
                self.qup.b = 3 * ureg.cm
            assert exc.value.args[0] == 'b should be given as a quantities Quantity instance'

        def test_invalid_units(self):

            from astropy import units as u

            with pytest.raises(TraitError) as exc:
                self.qup.a = 5 * pq.s
            assert exc.value.args[0] == 'a should be in units convertible to m'

            with pytest.raises(TraitError) as exc:
                self.qup.b = np.ones((2, 5)) * pq.s
            assert exc.value.args[0] == 'b should be in units convertible to cm/s'



# TODO: add test for domain with units

def test_inconsistent_ndim_shape():

    with pytest.raises(TraitError) as exc:
        a = NumericalTrait(ndim=3, shape=(3, 3))
    assert exc.value.args[0] == "shape=(3, 3) and ndim=3 are inconsistent"


def test_invalid_unit_framework():

    with pytest.raises(TraitError) as exc:
        a = NumericalTrait(convertible_to='m')
    assert exc.value.args[0] == "Could not identify unit framework for target unit of type str"