class MyModel(Fittable1DModel):
a = Parameter(default=1)
b = Parameter(default=2)
@staticmethod
def evaluate(x, a, b):
return a * x + b
class GaussianAbsorption(Fittable1DModel):
"""
The Gaussian absorption profile.
This is defined in terms of the emission Gaussian1D model.
"""
amplitude = Parameter(default=1., min=0.)
mean = Parameter(default=1.)
stddev = Parameter(default=1.)
@staticmethod
def evaluate(x, amplitude, mean, stddev):
"""
GaussianAbsorption model function.
"""
return models.Gaussian1D.evaluate(x, -amplitude, mean, stddev)
@staticmethod
def fit_deriv(x, amplitude, mean, stddev):
"""
GaussianAbsorption model function derivatives.
"""
import operator
return list(
map(operator.neg,
models.Gaussian1D.fit_deriv(x, -amplitude, mean, stddev)))
class SetterModel(FittableModel):
inputs = ('x', 'y')
outputs = ('z',)
xc = Parameter(default=1, setter=setter1)
yc = Parameter(default=1, setter=setter2)
def evaluate(self, x, y, xc, yc):
return ((x - xc)**2 + (y - yc)**2)
def do_something(self, v):
pass
def __init__(self, xc, yc, p):
self.p = p # p is a value intended to be used by the setter
super().__init__()
self.xc = xc
self.yc = yc
def evaluate(self, x, y, xc, yc):
return ((x - xc)**2 + (y - yc)**2)
def do_something(self, v):
pass
class _ConstraintsTestA(Model):
stddev = Parameter(default=0, min=0, max=0.3)
mean = Parameter(default=0, fixed=True)
@staticmethod
def evaluate(stddev, mean):
return stddev, mean
def test_std(self):
param = Parameter(name='test', default=[1, 2, 3, 4])
assert param.std is None
assert param._std is None
param.std = 5
assert param.std == 5 == param._std
class MyModel(Fittable1DModel):
a = Parameter(default=1)
b = Parameter(default=0, min=0, fixed=True)
@staticmethod
def evaluate(x, a, b):
return x * a + b
class AngleFromGratingEquation(Model):
"""
Grating Equation Model. Computes the diffracted/refracted angle.
Parameters
----------
groove_density : int
Grating ruling density.
order : int
Spectral order.
"""
separable = False
inputs = ("lam", "alpha_in", "beta_in", "z")
outputs = ("alpha_out", "beta_out", "zout")
groove_density = Parameter()
order = Parameter(default=-1)
def evaluate(self, lam, alpha_in, beta_in, z, groove_density, order):
if alpha_in.shape != beta_in.shape != z.shape:
raise ValueError("Expected input arrays to have the same shape")
orig_shape = alpha_in.shape or (1, )
xout = -alpha_in - groove_density * order * lam
yout = -beta_in
zout = np.sqrt(1 - xout**2 - yout**2)
xout.shape = yout.shape = zout.shape = orig_shape
return xout, yout, zout
class TModel_1_1(Fittable1DModel):
p1 = Parameter()
p2 = Parameter()
@staticmethod
def evaluate(x, p1, p2):
return x + p1 + p2
class CustomInputNamesModel(Fittable1DModel):
n_inputs = 1
n_outputs = 1
a = Parameter(default=1.0)
b = Parameter(default=1.0)
def __init__(self, a=a, b=b):
super().__init__(a=a, b=b)
self.inputs = ('inn',)
self.outputs = ('out',)
@staticmethod
def evaluate(inn, a, b):
return a * inn + b
@property
def input_units(self):
if self.a.unit is None and self.b.unit is None:
return None
else:
return {'inn': self.b.unit / self.a.unit}
def _parameter_units_for_data_units(self, inputs_unit, outputs_unit):
return {'a': outputs_unit['out'] / inputs_unit['inn'],
'b': outputs_unit['out']
}
class TestModel(Model):
p1 = Parameter()
p2 = Parameter()
p3 = Parameter()
def evaluate(self, *args):
return
def test_quantity(self):
param = Parameter(name='test', default=[1, 2, 3])
assert param.unit is None
assert param.quantity is None
param = Parameter(name='test', default=[1, 2, 3], unit=u.m)
assert param.unit == u.m
assert (param.quantity == np.array([1, 2, 3]) * u.m).all()
class ModelDefault(Model):
slope = Parameter()
intercept = Parameter()
_separable = False
@staticmethod
def evaluate(x, slope, intercept):
return slope * x + intercept
def __setattr__(self, attr, value):
if attr[0] != '_' and self._param_names and attr in self._param_names:
param = Parameter(attr, default=0.0, model=self)
# This is a little hackish, but we can actually reuse the
# Parameter.__set__ method here
param.__set__(self, value)
else:
super(SinglePSF, self).__setattr__(attr, value)
def __setattr__(self, attr, value):
if attr[0] != '_' and self._param_names and attr in self._param_names:
param = Parameter(attr, default=0.0, model=self)
# This is a little hackish, but we can actually reuse the
# Parameter.__set__ method here
param.__set__(self, value)
else:
super(SinglePSF, self).__setattr__(attr, value)
def test_size(self):
param = Parameter(name='test', default=[1, 2, 3, 4])
assert param.size == 4
param = Parameter(name='test', default=[1])
assert param.size == 1
param = Parameter(name='test', default=1)
assert param.size == 1
class classmodel(FittableModel):
f = Parameter(default=1)
x = Parameter(default=0)
y = Parameter(default=2)
def __init__(self, f=f.default, x=x.default, y=y.default):
super().__init__(f, x, y)
def evaluate(self):
pass
class Spline1D(Fittable1DModel):
degree = Parameter(default=3, fixed=True)
smooth = Parameter(default=1, fixed=True) # default=None crashes the app
def evaluate(self, x, degree, smooth):
_f = UnivariateSpline(self.wave, self.flux,
k=degree,
s=smooth)
return _f(x)
def test_raw_value(self):
param = Parameter(name='test', default=[1, 2, 3, 4])
# Normal case
assert (param._raw_value == param.value).all()
# Bad setter
param._setter = True
param._internal_value = 4
assert param._raw_value == 4
def test_parameter_properties():
"""Test if getting / setting of Parameter properties works."""
p = Parameter('alpha', default=1)
assert p.name == 'alpha'
# Parameter names are immutable
with pytest.raises(AttributeError):
p.name = 'beta'
assert p.fixed is False
p.fixed = True
assert p.fixed is True
assert p.tied is False
p.tied = lambda _: 0
p.tied = False
assert p.tied is False
assert p.min is None
p.min = 42
assert p.min == 42
p.min = None
assert p.min is None
assert p.max is None
p.max = 41
assert p.max == 41
class TModel_1_2(FittableModel):
inputs = ('x',)
outputs = ('y', 'z')
p1 = Parameter()
p2 = Parameter()
p3 = Parameter()
@staticmethod
def evaluate(x, p1, p2, p3):
return (x + p1 + p2, x + p1 + p2 + p3)
class subclassmodel(classmodel):
f = Parameter(default=3, fixed=True)
x = Parameter(default=10)
y = Parameter(default=12)
h = Parameter(default=5)
def __init__(self, f=f.default, x=x.default, y=y.default, h=h.default):
super().__init__(f, x, y)
def evaluate(self):
pass
class TModel_1_2(FittableModel):
n_inputs = 1
n_outputs = 2
p1 = Parameter()
p2 = Parameter()
p3 = Parameter()
@staticmethod
def evaluate(x, p1, p2, p3):
return (x + p1 + p2, x + p1 + p2 + p3)
def test_fixed(self):
param = Parameter(name='test', default=[1, 2, 3, 4])
assert param.fixed == False == param._fixed
# Set error
with pytest.raises(ValueError) as err:
param.fixed = 3
assert str(err.value) == \
"Value must be boolean"
# Set
param.fixed = True
assert param.fixed == True == param._fixed
class TParModel(Model):
"""
A toy model to test parameters machinery
"""
coeff = Parameter()
e = Parameter()
def __init__(self, coeff, e, **kwargs):
super().__init__(coeff=coeff, e=e, **kwargs)
@staticmethod
def evaluate(coeff, e):
pass
class Sersic2DAsym(Sersic2D):
r"""
Two dimensional Sersic profile with asymmetry.
Parameters are same as Sersic2D, plus:
----------
asym_strength : float, optional
Strength of asymmetry.
asym_angle : float, optional
Position angle of maximum asymmetry.
Notes
-----
Asymmetry introduced by multiplying Sersic2D profile by
(1 - asym_strength * cosine(azimuthal angle - asym_angle)
"""
asym_strength = Parameter(default=0)
asym_angle = Parameter(default=0)
@classmethod
def evaluate(cls, x, y, amplitude, r_eff, n, x_0, y_0, ellip, theta,
asym_strength, asym_angle):
"""Two dimensional Sersic profile function with asymmetry."""
if cls._gammaincinv is None:
try:
from scipy.special import gammaincinv
cls._gammaincinv = gammaincinv
except ValueError:
raise ImportError('Sersic2D model requires scipy > 0.11.')
bn = cls._gammaincinv(2. * n, 0.5)
a, b = r_eff, (1 - ellip) * r_eff
cos_theta, sin_theta = np.cos(theta), np.sin(theta)
x_maj = (x - x_0) * cos_theta + (y - y_0) * sin_theta
x_min = -(x - x_0) * sin_theta + (y - y_0) * cos_theta
z = np.sqrt((x_maj / a)**2 + (x_min / b)**2)
eps = 1e-32
angle = np.arctan(x_maj / (x_min + eps))
angle += np.pi * (x_min < 0) - np.pi / 2
angle[np.isnan(angle)] = 0
asym = (1 - asym_strength * np.cos(theta - asym_angle - angle))
return amplitude * asym * np.exp(-bn * (z**(1 / n) - 1))
def _parameter_units_for_data_units(self, inputs_unit, outputs_unit):
par_unit = super()._parameter_units_for_data_units(
inputs_unit, outputs_unit)
return par_unit + {'asym_angle': u.rad}
class BlackBody(Fittable1DModel):
"""
Produce a blackbody flux spectrum.
Note that the wave and flux arrays used to Quantity
Notes
-----
See `~astropy.modeling.Fittable1DModel`
for further details on modeling and all
possible parameters that can be passed in.
Description of the blackbody function itself is described in
`~astropy.analytic_functions.blackbody`
"""
temp = Parameter(default=5000, min=10.)
norm = Parameter(default=1.)
def evaluate(self, x, temp, norm):
"""
Evaluate the blackbody for a given temperature over a wavelength range.
Parameters
----------
x: numpy.ndarray
The wavelengths to evaulate over.
temp: float
The temperature to evualate at.
norm: float
The normalization factor.
Returns
-------
blackbody_flux: numpy.ndarray
The blackbody flux.
"""
# x is passed as a bare numpy array; must be
# converted back to Quantity before calling
# astropy's black body functions.
_x_u = x * self.wave.unit
# convert result of the Planck function to
# flux density in the same units as the data.
_flux = (blackbody_lambda(_x_u, temp) * u.sr).to(self.flux.unit)
# normalize and return just the values,
# to conform to the Model API.
return (norm * _flux).value
class Spectrum1DLinearWCS(BaseSpectrum1DWCS):
"""
A simple linear wcs
"""
dispersion0 = Parameter('dispersion0')
dispersion_delta = Parameter('dispersion_delta')
@deprecated('0.dev??', message='please use Spectrum1DPolynomialWCS')
def __init__(self, dispersion0, dispersion_delta, pixel_index, unit):
super(Spectrum1DLinearWCS, self).__init__()
#### Not clear what to do about units of dispersion0 and dispersion_delta.
# dispersion0 should have units like angstrom, whereas dispersion_delta should have units like angstrom/pix
# for now I assume pixels don't have units and both dispersion0 and dispersion_delta should have the same unit
dispersion0 = u.Quantity(dispersion0, unit)
dispersion_delta = u.Quantity(dispersion_delta, unit)
check_valid_unit(dispersion0.unit)
check_valid_unit(dispersion_delta.unit)
##### Quick fix - needs to be fixed in modelling ###
if unit is None:
unit = dispersion0.unit
self.unit = unit
self.dispersion0 = dispersion0.value
self.dispersion_delta = dispersion_delta.value
self.pixel_index = pixel_index
def __call__(self, pixel_indices):
if misc.isiterable(pixel_indices) and not isinstance(
pixel_indices, basestring):
pixel_indices = np.array(pixel_indices)
return (self.dispersion0 + self.dispersion_delta *
(pixel_indices - self.pixel_index)) * self.unit
def invert(self, dispersion_values):
if not hasattr(dispersion_values, 'unit'):
raise u.UnitsException(
'Must give a dispersion value with a valid unit (i.e. quantity 5 * u.Angstrom)'
)
if misc.isiterable(dispersion_values) and not isinstance(
dispersion_values, basestring):
dispersion_values = np.array(dispersion_values)
return float((dispersion_values - self.dispersion0) /
self.dispersion_delta) + self.pixel_index
class RefractionIndexFromPrism(Model):
"""
Compute the refraction index of a prism (NIRSpec).
Parameters
----------
prism_angle : float
Prism angle in deg.
"""
standard_broadcasting = False
inputs = (
"alpha_in",
"beta_in",
"alpha_out",
)
outputs = ("n")
prism_angle = Parameter(setter=np.deg2rad, getter=np.rad2deg)
def __init__(self, prism_angle, name=None):
super(RefractionIndexFromPrism, self).__init__(prism_angle=prism_angle,
name=name)
def evaluate(self, alpha_in, beta_in, alpha_out, prism_angle):
sangle = math.sin(prism_angle)
cangle = math.cos(prism_angle)
nsq = ((alpha_out + alpha_in * (1 - 2 * sangle**2)) / (2 * sangle * cangle)) **2 + \
alpha_in ** 2 + beta_in ** 2
return np.sqrt(nsq)
class TParModel(Model):
"""
A toy model to test parameters machinery
"""
# standard_broadasting = False
inputs = ('x', )
outputs = ('x', )
coeff = Parameter()
e = Parameter()
def __init__(self, coeff, e, **kwargs):
super().__init__(coeff=coeff, e=e, **kwargs)
@staticmethod
def evaluate(x, coeff, e):
return x * coeff + e
class Snell(Model):
"""Computes the Prism Snell refraction through a surface.
Parameters
----------
n: float
refraction index as calculated for a given wavelenth
"""
inputs = ("xin", "yin", "zin")
outputs = ("xout", "yout", "zout")
n = Parameter(default=1.0)
def __init__(self, n=n, name=None):
super(Snell, self).__init__(n=n, name=name)
def evaluate(self, x, y, z, n):
"""Compute Snell's refraction law from the front surface."""
xout = x / n
yout = y / n
zout = np.sqrt(1.0 - xout**2 - yout**2)
return xout, yout, zout
def inverse(self, x, y, z, n):
"""Compute Snell's refraction law from the back surface."""
xout = x * n
yout = y * n
zout = np.sqrt(1.0 - xout**2 - yout**2)
return xout, yout, zout
def test_param_repr_oneline(self):
# Single value no units
param = Parameter(name='test', default=1)
assert param_repr_oneline(param) == '1.'
# Vector value no units
param = Parameter(name='test', default=[1, 2, 3, 4])
assert param_repr_oneline(param) == '[1., 2., 3., 4.]'
# Single value units
param = Parameter(name='test', default=1*u.m)
assert param_repr_oneline(param) == '1. m'
# Vector value units
param = Parameter(name='test', default=[1, 2, 3, 4] * u.m)
assert param_repr_oneline(param) == '[1., 2., 3., 4.] m'