# Path: sncosmo/registry.py
# def _get_registry(data_class):
# def register_loader(data_class, name, func, args=None,
#                     version=None, meta=None, force=False):
# def register(instance, name=None, data_class=None, force=False):
# def retrieve(data_class, name, version=None):


from io import StringIO
from numpy.testing import assert_allclose, assert_approx_equal
from sncosmo import registry
import numpy as np
import sncosmo

    def setup_class(self):
        self.source = flatsource()

    def test_flux(self):
        for a in [1., 2.]:
            self.source.set(amplitude=a)
            assert_allclose(self.source.flux(1., 2000.), a)
            assert_allclose(self.source.flux(1., [2000.]), np.array([a]))
            assert_allclose(self.source.flux([1.], 2000.), np.array([[a]]))
            assert_allclose(self.source.flux([1.], [2000.]), np.array([[a]]))

    def test_getsource(self):
        # register the source & retrieve it from the registry
        registry.register(self.source, name="testsource",
                          data_class=sncosmo.Source)
        m = sncosmo.get_source("testsource")

    def test_bandflux(self):
        self.source.set(amplitude=1.0)
        f = self.source.bandflux("bessellb", 0.)

        # Correct answer
        b = sncosmo.get_bandpass("bessellb")
        dwave = np.gradient(b.wave)
        ans = np.sum(b.trans * b.wave * dwave) / sncosmo.models.HC_ERG_AA
        assert_approx_equal(ans, f)

    def test_bandflux_shapes(self):
        # Just check that these work.

# Path: sncosmo/registry.py
# def _get_registry(data_class):
# def register_loader(data_class, name, func, args=None,
#                     version=None, meta=None, force=False):
# def register(instance, name=None, data_class=None, force=False):
# def retrieve(data_class, name, version=None):


# Path: sncosmo/tests/test_models.py
from io import StringIO
from numpy.testing import assert_allclose, assert_approx_equal
from sncosmo import registry
import numpy as np
import sncosmo


    def setup_class(self):
        self.source = flatsource()

    def test_flux(self):
        for a in [1., 2.]:
            self.source.set(amplitude=a)
            assert_allclose(self.source.flux(1., 2000.), a)
            assert_allclose(self.source.flux(1., [2000.]), np.array([a]))
            assert_allclose(self.source.flux([1.], 2000.), np.array([[a]]))
            assert_allclose(self.source.flux([1.], [2000.]), np.array([[a]]))

    def test_getsource(self):
        # register the source & retrieve it from the registry
        registry.register(self.source, name="testsource",
                          data_class=sncosmo.Source)
        m = sncosmo.get_source("testsource")

    def test_bandflux(self):
        self.source.set(amplitude=1.0)
        f = self.source.bandflux("bessellb", 0.)

        # Correct answer
        b = sncosmo.get_bandpass("bessellb")
        dwave = np.gradient(b.wave)
        ans = np.sum(b.trans * b.wave * dwave) / sncosmo.models.HC_ERG_AA
        assert_approx_equal(ans, f)

    def test_bandflux_shapes(self):
        # Just check that these work.

        self.source.bandflux("bessellb", 0., zp=25., zpsys="ab")