def validate(cls, model, bounding_box):
"""
Validate a given bounding box sequence against the given model (which
may be either a subclass of `~astropy.modeling.Model` or an instance
thereof, so long as the ``.inputs`` attribute is defined.
Currently this just checks that the bounding_box is either a 2-tuple
of lower and upper bounds for 1-D models, or an N-tuple of 2-tuples
for N-D models.
This also returns a normalized version of the bounding_box input to
ensure it is always an N-tuple (even for the 1-D case).
"""
nd = model.n_inputs
if nd == 1:
MESSAGE = "Bounding box for {0} model must be a sequence of length "
"2 consisting of a lower and upper bound, or a 1-tuple "
"containing such a sequence as its sole element.".format(
model.name)
try:
valid_shape = np.shape(bounding_box) in ((2, ), (1, 2))
except TypeError:
# np.shape does not work with lists of Quantities
valid_shape = np.shape([b.to_value() for b in bounding_box
]) in ((2, ), (1, 2))
except ValueError:
raise ValueError(MESSAGE)
if not isiterable(bounding_box) or not valid_shape:
raise ValueError(MESSAGE)
if len(bounding_box) == 1:
return cls((tuple(bounding_box[0]), ))
else:
return cls(tuple(bounding_box))
else:
MESSAGE = "Bounding box for {0} model must be a sequence of length "
"{1} (the number of model inputs) consisting of pairs of "
"lower and upper bounds for those inputs on which to "
"evaluate the model.".format(model.name, nd)
try:
valid_shape = all([len(i) == 2 for i in bounding_box])
except TypeError:
valid_shape = False
if len(bounding_box) != nd:
valid_shape = False
if not isiterable(bounding_box) or not valid_shape:
raise ValueError(MESSAGE)
return cls(tuple(bounds) for bounds in bounding_box)
def validate(cls, model, bounding_box):
"""
Validate a given bounding box sequence against the given model (which
may be either a subclass of `~astropy.modeling.Model` or an instance
thereof, so long as the ``.inputs`` attribute is defined.
Currently this just checks that the bounding_box is either a 2-tuple
of lower and upper bounds for 1-D models, or an N-tuple of 2-tuples
for N-D models.
This also returns a normalized version of the bounding_box input to
ensure it is always an N-tuple (even for the 1-D case).
"""
nd = model.n_inputs
if nd == 1:
MESSAGE = "Bounding box for {0} model must be a sequence of length "
"2 consisting of a lower and upper bound, or a 1-tuple "
"containing such a sequence as its sole element.".format(model.name)
try:
valid_shape = np.shape(bounding_box) in ((2,), (1, 2))
except TypeError:
# np.shape does not work with lists of Quantities
valid_shape = np.shape([b.to_value() for b in bounding_box]) in ((2,), (1, 2))
except ValueError:
raise ValueError(MESSAGE)
if not isiterable(bounding_box) or not valid_shape:
raise ValueError(MESSAGE)
if len(bounding_box) == 1:
return cls((tuple(bounding_box[0]),))
else:
return cls(tuple(bounding_box))
else:
MESSAGE = "Bounding box for {0} model must be a sequence of length "
"{1} (the number of model inputs) consisting of pairs of "
"lower and upper bounds for those inputs on which to "
"evaluate the model.".format(model.name, nd)
try:
valid_shape = all([len(i) == 2 for i in bounding_box])
except TypeError:
valid_shape = False
if len(bounding_box) != nd:
valid_shape = False
if not isiterable(bounding_box) or not valid_shape:
raise ValueError(MESSAGE)
return cls(tuple(bounds) for bounds in bounding_box)
def test_quantity_iterability():
"""Regressiont est for issue #878.
Scalar quantities should not be iterable and should raise a type error on
iteration.
"""
q1 = [15.0, 17.0] * u.m
assert isiterable(q1)
q2 = next(iter(q1))
assert q2 == 15.0 * u.m
assert not isiterable(q2)
pytest.raises(TypeError, iter, q2)
def plot_lnprob(self, tipmag, alphargb, alphaother, fracother, magrng=100, doplot=True, delog=False, **plotkwargs):
"""
Plots (optionally) and returns arrays suitable for plotting the pdf. If
`magrng` is a scalar, it gives the number of samples over the data
domain. If an array, it's used as the x axis.
"""
from copy import copy
from astropy.utils import isiterable
from matplotlib import pyplot as plt
fakemod = copy(self)
if isiterable(magrng):
fakemod.magdata = np.sort(magrng)
else:
fakemod.magdata = np.linspace(self.mindata, self.maxdata, magrng)
if fakemod.magunc is not None:
sorti = np.argsort(self.magdata)
fakemod.magunc = np.interp(fakemod.magdata, self.magdata[sorti], self.magunc[sorti])
lnpb = fakemod.lnprob(tipmag, alphargb, alphaother, fracother)
if delog:
lnpb = np.exp(lnpb - np.min(lnpb))
if doplot:
plt.plot(fakemod.magdata, lnpb, **plotkwargs)
return fakemod.magdata, lnpb
def __call__(self, x, y):
input_values = np.asarray(x), np.asarray(y)
output_values = [np.zeros_like(arg) for arg in input_values]
output_values.append(np.zeros_like(input_values[0]))
result = self.regions_mask[x, y]
if not astu.isiterable(result):
if result != 0:
return self._selector[result[0]](*input_values)
else:
return None #input_values
unique_regions = np.unique(result).tolist()
try:
unique_regions.remove(0)
except ValueError:
pass
try:
unique_regions.remove("")
except ValueError:
pass
ra = []
dec = []
lam = []
for i in unique_regions:
indices = (result==i)
transform = self._selector[i]
xyind = [val[indices] for val in input_values]
r, d, l = transform(x[indices], y[indices])
ra.extend(r)
dec.extend(d)
lam.extend(l)
return np.asarray(ra), np.asarray(dec), np.asarray(lam)
def z_to_array(z):
"""
Convert input scalar or array to an array
Parameters
----------
z: float or ndarray
Redshift
Returns
-------
z: ndarray
flg_z: int
0 -- Input was a scalar
1 -- Input was an array
"""
# float or ndarray?
if not isiterable(z):
z = np.array([z])
flg_z = 0
else:
flg_z = 1
# Return
return z, flg_z
def _tofloat(value):
"""Convert a parameter to float or float array"""
if isiterable(value):
try:
value = np.asanyarray(value, dtype=float)
except (TypeError, ValueError):
# catch arrays with strings or user errors like different
# types of parameters in a parameter set
raise InputParameterError(
"Parameter of {0} could not be converted to "
"float".format(type(value)))
elif isinstance(value, Quantity):
# Quantities are fine as is
pass
elif isinstance(value, np.ndarray):
# A scalar/dimensionless array
value = float(value.item())
elif isinstance(value, (numbers.Number, np.number)):
value = float(value)
elif isinstance(value, bool):
raise InputParameterError(
"Expected parameter to be of numerical type, not boolean")
else:
raise InputParameterError(
"Don't know how to convert parameter of {0} to "
"float".format(type(value)))
return value
def halomass_from_stellarmass(log_mstar, z=0):
""" Halo mass from Stellar mass (Moster+2013).
Inverts the function `stellarmass_from_halomass`
numerically.
Args:
log_mstar (float or numpy.ndarray): log_10 stellar mass
in solar mass units.
z (float, optional): galaxy redshift
Returns:
log_Mhalo (float): log_10 halo mass
in solar mass units.
"""
try:
log_mstar * z
except ValueError:
raise TypeError(
"log_mstar and z can't be broadcast together for root finding. Use numpy arrays of same length or scalar values.")
f = lambda x: stellarmass_from_halomass(x, z=z) - log_mstar
guess = 2 + log_mstar
if isiterable(log_mstar):
return fsolve(f, guess)
else:
return fsolve(f, guess)[0]
def __init__(self, num_axes, ref_system=None, ref_pos=None, units=None, axes_names=None, name=None):
""" Initialize a frame"""
if units is not None and astu.isiterable(units):
if len(units) != num_axes:
raise ValueError("Number of units does not match number of axes")
self._units=units
if axes_names is not None and astu.isiterable(axes_names):
if len(axes_names) != num_axes:
raise ValueError("Number of axes names does not match number of axes")
self._axes_names = axes_names
self._reference_system = ref_system
self._reference_position = ref_pos
if name is None:
self._name = ref_system
else:
self._name = name
def set(self, islice, attr, value):
"""
Set the attribute for a slice of the specobjs
Args:
islice (int, ndarray of bool, slice): Indicates SpecObj to affect
attr (str):
value (anything) : Value of the item
Returns:
"""
sub_sobjs = self.specobjs[islice]
if isiterable(value):
if sub_sobjs.size == len(value): # Assume you want each paired up
for kk, sobj in enumerate(sub_sobjs):
setattr(sobj, attr, value[kk])
return
# Assuming scalar assignment
if isinstance(sub_sobjs, SpecObj):
setattr(sub_sobjs, attr, value)
else:
for sobj in sub_sobjs:
setattr(sobj, attr, value)
return
def pixel_to_world(self, *quantity_axis_list):
"""
Convert a pixel coordinate to a data (world) coordinate by using
`~gwcs.wcs.WCS`.
This method expects input and returns output in the same order as the
array dimensions. (Which is the reverse of the underlying WCS object.)
Parameters
----------
quantity_axis_list : iterable
An iterable of `~astropy.units.Quantity` with unit as pixel `pix`.
Returns
-------
coord : `list`
A list of arrays containing the output coordinates.
"""
world = self.wcs.pixel_to_world(*quantity_axis_list[::-1])
# Convert list to tuple as a more standard return type
if isinstance(world, list):
world = tuple(world)
# If our return is an iterable then reverse it to match pixel dims.
if isiterable(world):
return world[::-1]
return world
def __call__(self, x, y):
input_values = np.asarray(x), np.asarray(y)
output_values = [np.zeros_like(arg) for arg in input_values]
output_values.append(np.zeros_like(input_values[0]))
result = self.regions_mask[x, y]
if not astu.isiterable(result):
if result != 0:
return self._selector[result[0]](*input_values)
else:
return None #input_values
unique_regions = np.unique(result).tolist()
try:
unique_regions.remove(0)
except ValueError:
pass
try:
unique_regions.remove("")
except ValueError:
pass
ra = []
dec = []
lam = []
for i in unique_regions:
indices = (result == i)
transform = self._selector[i]
xyind = [val[indices] for val in input_values]
r, d, l = transform(x[indices], y[indices])
ra.extend(r)
dec.extend(d)
lam.extend(l)
return np.asarray(ra), np.asarray(dec), np.asarray(lam)
def __new__(cls, angle, unit=None, dtype=None, copy=True, **kwargs):
if not isinstance(angle, u.Quantity):
if unit is not None:
unit = cls._convert_unit_to_angle_unit(u.Unit(unit))
if isinstance(angle, tuple):
angle = cls._tuple_to_float(angle, unit)
elif isinstance(angle, str):
angle, angle_unit = util.parse_angle(angle, unit)
if angle_unit is None:
angle_unit = unit
if isinstance(angle, tuple):
angle = cls._tuple_to_float(angle, angle_unit)
if angle_unit is not unit:
# Possible conversion to `unit` will be done below.
angle = u.Quantity(angle, angle_unit, copy=False)
elif (isiterable(angle)
and not (isinstance(angle, np.ndarray)
and angle.dtype.kind not in 'SUVO')):
angle = [Angle(x, unit, copy=False) for x in angle]
return super().__new__(cls,
angle,
unit,
dtype=dtype,
copy=copy,
**kwargs)
def test_freq_recovery(self):
# define a bunch of arrays of times to make sure SuperFreq isn't
# sensitive to the times
ts = [np.linspace(0., 150., 12000),
np.linspace(0., 150., 24414),
np.linspace(0., 150., 42104),
np.linspace(150., 300., 12000),
np.linspace(150., 300., 24414),
np.linspace(150., 300., 42104),
np.linspace(0., 150., 12000) + 50*(2*np.pi/self.omega[0])]
for i,t in enumerate(ts):
print(i, t.min(), t.max(), len(t))
f = self.make_f(t)
nfreq = len(self.omega)
if not isiterable(self.p):
ps = [self.p]
else:
ps = self.p
for p in ps:
print(i, p)
# create SuperFreq object for this time array
sf = SuperFreq(t, p=p)
# solve for the frequencies
w,amp,phi = sf.frecoder(f[:sf.n], break_condition=1E-5)
np.testing.assert_allclose(self.omega, w[:nfreq], rtol=1E-7)
np.testing.assert_allclose(self.A, amp[:nfreq], rtol=1E-5)
np.testing.assert_allclose(self.phi, phi[:nfreq], rtol=1E-3)
def _tofloat(value):
"""Convert a parameter to float or float array"""
if isiterable(value):
try:
value = np.asanyarray(value, dtype=float)
except (TypeError, ValueError):
# catch arrays with strings or user errors like different
# types of parameters in a parameter set
raise InputParameterError(
f"Parameter of {type(value)} could not be converted to float")
elif isinstance(value, Quantity):
# Quantities are fine as is
pass
elif isinstance(value, np.ndarray):
# A scalar/dimensionless array
value = float(value.item())
elif isinstance(value, (numbers.Number, np.number)):
value = float(value)
elif isinstance(value, bool):
raise InputParameterError(
"Expected parameter to be of numerical type, not boolean")
else:
raise InputParameterError(
f"Don't know how to convert parameter of {type(value)} to float")
return value
def __init__(self,
naxes,
axes_type,
axes_order,
reference_frame=None,
reference_position=None,
unit=None,
axes_names=None,
name=None,
axis_physical_types=None):
self._naxes = naxes
self._axes_order = tuple(axes_order)
if isinstance(axes_type, str):
self._axes_type = (axes_type, )
else:
self._axes_type = tuple(axes_type)
self._reference_frame = reference_frame
if unit is not None:
if astutil.isiterable(unit):
unit = tuple(unit)
else:
unit = (unit, )
if len(unit) != naxes:
raise ValueError(
"Number of units does not match number of axes.")
else:
self._unit = tuple([u.Unit(au) for au in unit])
else:
self._unit = tuple("" for na in range(naxes))
if axes_names is not None:
if isinstance(axes_names, str):
axes_names = (axes_names, )
else:
axes_names = tuple(axes_names)
if len(axes_names) != naxes:
raise ValueError(
"Number of axes names does not match number of axes.")
else:
axes_names = tuple([""] * naxes)
self._axes_names = axes_names
if name is None:
self._name = self.__class__.__name__
else:
self._name = name
self._reference_position = reference_position
if len(self._axes_type) != naxes:
raise ValueError(
"Length of axes_type does not match number of axes.")
if len(self._axes_order) != naxes:
raise ValueError(
"Length of axes_order does not match number of axes.")
super(CoordinateFrame, self).__init__()
self._axis_physical_types = self._set_axis_physical_types(
axis_physical_types)
def __init__(self, slits, models):
super(Slit2Msa, self).__init__()
if isiterable(slits[0]):
self._slits = [tuple(s) for s in slits]
self.slit_ids = [s[0] for s in self._slits]
else:
self._slits = list(slits)
self.slit_ids = self._slits
self.models = models
def __init__(self, slits, models):
super(Slit2Msa, self).__init__()
if isiterable(slits[0]):
self._slits = [tuple(s) for s in slits]
self.slit_ids = [s[0] for s in self._slits]
else:
self._slits = list(slits)
self.slit_ids = self._slits
self.models = models
def __init__(self, naxes, axes_order=(0, 1), reference_frame=None, reference_position=None,
unit=None, axes_names=None, name=None):
""" Initialize a frame"""
self._axes_order = axes_order
# map data axis into frame axes - 0-based
self._naxes = naxes
if unit is not None:
if astutil.isiterable(unit):
if len(unit) != naxes:
raise ValueError("Number of units does not match number of axes.")
else:
self._unit = [u.Unit(au) for au in unit]
else:
self._unit = [u.Unit(unit)]
else:
self._unit = reference_frame.representation_component_units.values()
if axes_names is not None and astutil.isiterable(axes_names):
if len(axes_names) != naxes:
raise ValueError("Number of axes names does not match number of axes.")
else:
axes_names = reference_frame.representation_component_names.values()
self._axes_names = axes_names
self._reference_frame = reference_frame
if name is None:
try:
self._name = reference_frame.name
except AttributeError:
self._name = repr(reference_frame)
else:
self._name = name
if reference_position is not None:
self._reference_position = reference_position
else:
try:
self._reference_position = reference_frame.reference_position
except AttributeError:
self._reference_position = None
super(CoordinateFrame, self).__init__()
def _getdata(self, keys):
for idx, (key, axis) in enumerate(zip(keys, self.hdu.shape)):
if isinstance(key, slice):
ks = range(*key.indices(axis))
break
elif isiterable(key):
# Handle both integer and boolean arrays.
ks = np.arange(axis, dtype=int)[key]
break
# This should always break at some point if _getdata is called.
data = [self[keys[:idx] + (k,) + keys[idx + 1:]] for k in ks]
if any(isinstance(key, slice) or isiterable(key)
for key in keys[idx + 1:]):
# data contains multidimensional arrays; combine them.
return np.array(data)
else:
# Only singleton dimensions remain; concatenate in a 1D array.
return np.concatenate([np.atleast_1d(array) for array in data])
def _getdata(self, keys):
for idx, (key, axis) in enumerate(zip(keys, self.hdu.shape)):
if isinstance(key, slice):
ks = range(*key.indices(axis))
break
elif isiterable(key):
# Handle both integer and boolean arrays.
ks = np.arange(axis, dtype=int)[key]
break
# This should always break at some point if _getdata is called.
data = [self[keys[:idx] + (k,) + keys[idx + 1:]] for k in ks]
if any(isinstance(key, slice) or isiterable(key)
for key in keys[idx + 1:]):
# data contains multidimensional arrays; combine them.
return np.array(data)
else:
# Only singleton dimensions remain; concatenate in a 1D array.
return np.concatenate([np.atleast_1d(array) for array in data])
def pythonify(node):
for key, item in node.items():
if hasattr(item, 'items'):
pythonify(item)
else:
if hasattr(item,'tolist'):
node[key] = item.tolist()
elif isiterable(item) and not isinstance(item, str):
node[key] = map(float, list(item))
else:
node[key] = float(item)
def __init__(self,
defaultvalue='',
description=None,
cfgtype=None,
module=None,
aliases=None):
from astropy.utils import isiterable
if module is None:
module = find_current_module(2)
if module is None:
msg1 = 'Cannot automatically determine get_config module, '
msg2 = 'because it is not called from inside a valid module'
raise RuntimeError(msg1 + msg2)
else:
module = module.__name__
self.module = module
self.description = description
self.__doc__ = description
# now determine cfgtype if it is not given
if cfgtype is None:
if (isiterable(defaultvalue)
and not isinstance(defaultvalue, str)):
# it is an options list
dvstr = [str(v) for v in defaultvalue]
cfgtype = 'option(' + ', '.join(dvstr) + ')'
defaultvalue = dvstr[0]
elif isinstance(defaultvalue, bool):
cfgtype = 'boolean'
elif isinstance(defaultvalue, int):
cfgtype = 'integer'
elif isinstance(defaultvalue, float):
cfgtype = 'float'
elif isinstance(defaultvalue, str):
cfgtype = 'string'
defaultvalue = str(defaultvalue)
self.cfgtype = cfgtype
self._validate_val(defaultvalue)
self.defaultvalue = defaultvalue
if aliases is None:
self.aliases = []
elif isinstance(aliases, str):
self.aliases = [aliases]
else:
self.aliases = aliases
def growth_function_carroll(redshift, cosmology):
'''Growth function.
This function returns the growth function as a function of redshift for a
given cosmology as approximated by Carroll, Press & Turner (1992),
equation 29 in [1]_.
Parameters
----------
redshift : (nz,) array_like
Array of redshifts at which to evaluate the growth function.
cosmology : astropy.cosmology.Cosmology
Cosmology object providing methods for the evolution history of
omega_matter and omega_lambda with redshift.
Returns
-------
growth : (nz,) array_like
The growth function evaluated at the input redshifts for the given
cosmology.
Examples
--------
This example returns the growth function for a given array of redshifts
and for the Astropy default cosmology:
>>> import numpy as np
>>> from astropy.cosmology import default_cosmology
>>> redshift = np.array([0, 1, 2])
>>> cosmology = default_cosmology.get()
>>> growth_function_carroll(redshift, cosmology)
array([0.781361..., 0.476280..., 0.327549...])
References
----------
.. [1] Carroll, M. and Press, W. and Turner, E., (1992),
doi : 10.1146/annurev.aa.30.090192.002435
'''
if isiterable(redshift):
redshift = np.asarray(redshift)
if np.any(redshift < 0):
raise ValueError('Redshifts must be non-negative')
Om = cosmology.Om(redshift)
Ode = cosmology.Ode(redshift)
Dz = 2.5 * Om / (1 + redshift)
return Dz / (np.power(Om, 4.0 / 7.0) - Ode + (1 + 0.5 * Om) *
(1.0 + Ode / 70.0))
def __init__(self,
naxes,
axes_type,
axes_order,
reference_frame=None,
reference_position=None,
unit=None,
axes_names=None,
name=None,
wcsobj=None):
self._naxes = naxes
self._axes_order = tuple(axes_order)
if isinstance(axes_type, six.string_types):
self._axes_type = (axes_type, )
else:
self._axes_type = tuple(axes_type)
self._reference_frame = reference_frame
if unit is not None:
if astutil.isiterable(unit):
unit = tuple(unit)
else:
unit = (unit, )
if len(unit) != naxes:
raise ValueError(
"Number of units does not match number of axes.")
else:
self._unit = tuple([u.Unit(au) for au in unit])
if axes_names is not None:
if isinstance(axes_names, six.string_types):
axes_names = (axes_names, )
else:
axes_names = tuple(axes_names)
if len(axes_names) != naxes:
raise ValueError(
"Number of axes names does not match number of axes.")
else:
axes_names = tuple([""] * naxes)
self._axes_names = axes_names
if name is None:
self._name = self.__class__.__name__
else:
self._name = name
if reference_position is not None:
self._reference_position = reference_position
else:
self._reference_position = None
super(CoordinateFrame, self).__init__()
def __init__(self,
num_axes,
ref_system=None,
ref_pos=None,
units=None,
axes_names=None,
name=None):
""" Initialize a frame"""
if units is not None and astu.isiterable(units):
if len(units) != num_axes:
raise ValueError(
"Number of units does not match number of axes")
self._units = units
if axes_names is not None and astu.isiterable(axes_names):
if len(axes_names) != num_axes:
raise ValueError(
"Number of axes names does not match number of axes")
self._axes_names = axes_names
self._reference_system = ref_system
self._reference_position = ref_pos
if name is None:
self._name = ref_system
else:
self._name = name
def coordinate_to_quantity(self, *coords):
# list or tuple
if len(coords) == 1 and astutil.isiterable(coords[0]):
coords = list(coords[0])
elif len(coords) == 2:
coords = list(coords)
else:
raise ValueError("Unexpected number of coordinates in "
"input to frame {} : "
"expected 2, got {}".format(self.name, len(coords)))
for i in range(2):
if not hasattr(coords[i], 'unit'):
coords[i] = coords[i] * self.unit[i]
return tuple(coords)
def sanitize_slices(slices, ndim):
"""
Given a slice as input sanitise it to an easier to parse format.format
This function returns a list ``ndim`` long containing slice objects (or ints).
"""
if not isinstance(slices, (tuple, list)): # We just have a single int
slices = (slices, )
if len(slices) > ndim:
raise ValueError(
f"The dimensionality of the specified slice {slices} can not be greater "
f"than the dimensionality ({ndim}) of the wcs.")
if any((isiterable(s) for s in slices)):
raise IndexError(
"This slice is invalid, only integer or range slices are supported."
)
slices = list(slices)
if Ellipsis in slices:
if slices.count(Ellipsis) > 1:
raise IndexError(
"an index can only have a single ellipsis ('...')")
# Replace the Ellipsis with the correct number of slice(None)s
e_ind = slices.index(Ellipsis)
slices.remove(Ellipsis)
n_e = ndim - len(slices)
for i in range(n_e):
ind = e_ind + i
slices.insert(ind, slice(None))
for i in range(ndim):
if i < len(slices):
slc = slices[i]
if isinstance(slc, slice):
if slc.step and slc.step != 1:
raise IndexError(
"Slicing WCS with a step is not supported.")
elif not isinstance(slc, numbers.Integral):
raise IndexError("Only integer or range slices are accepted.")
else:
slices.append(slice(None))
return slices
def _unpack_params(p):
params = p.copy()
for key, item in p.items():
if '_unit' in key:
continue
if isiterable(item) and not isinstance(item, str):
params[key] = np.array(item).astype(float)
else:
params[key] = float(item)
if key + '_unit' in params:
params[key] = params[key] * u.Unit(params[key + '_unit'])
del params[key + '_unit']
return params
def _unpack_params(p):
params = p.copy()
for key, item in p.items():
if '_unit' in key:
continue
if isiterable(item) and not isinstance(item, str):
params[key] = np.array(item).astype(float)
else:
params[key] = float(item)
if key+'_unit' in params:
params[key] = params[key] * u.Unit(params[key+'_unit'])
del params[key+'_unit']
return params
def get_selector(self, *inputs):
"""
Get the selector value corresponding to this argument
Parameters
----------
*inputs :
All the processed model evaluation inputs.
"""
_selector = inputs[self.index]
if isiterable(_selector):
if len(_selector) == 1:
return _selector[0]
else:
return tuple(_selector)
return _selector
def __init__(self, naxes, axes_type, axes_order, reference_frame=None,
reference_position=None, unit=None, axes_names=None,
name=None, axis_physical_types=None):
self._naxes = naxes
self._axes_order = tuple(axes_order)
if isinstance(axes_type, str):
self._axes_type = (axes_type,)
else:
self._axes_type = tuple(axes_type)
self._reference_frame = reference_frame
if unit is not None:
if astutil.isiterable(unit):
unit = tuple(unit)
else:
unit = (unit,)
if len(unit) != naxes:
raise ValueError("Number of units does not match number of axes.")
else:
self._unit = tuple([u.Unit(au) for au in unit])
else:
self._unit = tuple("" for na in range(naxes))
if axes_names is not None:
if isinstance(axes_names, str):
axes_names = (axes_names,)
else:
axes_names = tuple(axes_names)
if len(axes_names) != naxes:
raise ValueError("Number of axes names does not match number of axes.")
else:
axes_names = tuple([""] * naxes)
self._axes_names = axes_names
if name is None:
self._name = self.__class__.__name__
else:
self._name = name
self._reference_position = reference_position
if len(self._axes_type) != naxes:
raise ValueError("Length of axes_type does not match number of axes.")
if len(self._axes_order) != naxes:
raise ValueError("Length of axes_order does not match number of axes.")
super(CoordinateFrame, self).__init__()
self._axis_physical_types = self._set_axis_physical_types(axis_physical_types)
def __init__(self, defaultvalue='', description=None, cfgtype=None,
module=None, aliases=None):
from astropy.utils import isiterable
if module is None:
module = find_current_module(2)
if module is None:
msg1 = 'Cannot automatically determine get_config module, '
msg2 = 'because it is not called from inside a valid module'
raise RuntimeError(msg1 + msg2)
else:
module = module.__name__
self.module = module
self.description = description
self.__doc__ = description
# now determine cfgtype if it is not given
if cfgtype is None:
if (isiterable(defaultvalue) and not
isinstance(defaultvalue, str)):
# it is an options list
dvstr = [str(v) for v in defaultvalue]
cfgtype = 'option(' + ', '.join(dvstr) + ')'
defaultvalue = dvstr[0]
elif isinstance(defaultvalue, bool):
cfgtype = 'boolean'
elif isinstance(defaultvalue, int):
cfgtype = 'integer'
elif isinstance(defaultvalue, float):
cfgtype = 'float'
elif isinstance(defaultvalue, str):
cfgtype = 'string'
defaultvalue = str(defaultvalue)
self.cfgtype = cfgtype
self._validate_val(defaultvalue)
self.defaultvalue = defaultvalue
if aliases is None:
self.aliases = []
elif isinstance(aliases, str):
self.aliases = [aliases]
else:
self.aliases = aliases
def load_fits(self, fileorhdu, viewer_name='Main Viewer'):
"""Load FITS image into the desired Ginga image viewer.
Parameters
----------
fileorhdu
File or HDU list object.
viewer_name : str
Name of Ginga image viewer to display to.
Raises
------
KeyError
Viewer name does not exist.
ValueError
Invalid file or HDU list object, or HDU list does not contain any
image.
"""
if isinstance(fileorhdu, file):
fileorhdu = fits.HDUList.fromfile(fileorhdu)
if isiterable(fileorhdu):
for hdui in fileorhdu:
if hasattr(hdui, 'is_image') and hdui.is_image:
hdu = hdui
break
else:
raise ValueError(
'fileorhdu was iterable but did not contain any '
'image HDUs')
elif hasattr(fileorhdu, 'data') and hasattr(fileorhdu, 'header'):
# quacks like an HDU - give it a shot
hdu = fileorhdu
else:
raise ValueError('fileorhdu was not a fits file or HDU-ish thing')
viewer = self.viewers[viewer_name]
if viewer.fitsimage.get_image() is None:
aim = AstroImage(logger=self.logger)
aim.load_hdu(hdu)
viewer.fitsimage.set_image(aim)
else:
viewer.fitsimage.get_image().load_hdu(hdu)
def load_fits(self, fileorhdu, viewer_name='Main Viewer'):
"""Load FITS image into the desired Ginga image viewer.
Parameters
----------
fileorhdu
File or HDU list object.
viewer_name : str
Name of Ginga image viewer to display to.
Raises
------
KeyError
Viewer name does not exist.
ValueError
Invalid file or HDU list object, or HDU list does not contain any
image.
"""
if isinstance(fileorhdu, file):
fileorhdu = fits.HDUList.fromfile(fileorhdu)
if isiterable(fileorhdu):
for hdui in fileorhdu:
if hasattr(hdui, 'is_image') and hdui.is_image:
hdu = hdui
break
else:
raise ValueError(
'fileorhdu was iterable but did not contain any '
'image HDUs')
elif hasattr(fileorhdu, 'data') and hasattr(fileorhdu, 'header'):
# quacks like an HDU - give it a shot
hdu = fileorhdu
else:
raise ValueError('fileorhdu was not a fits file or HDU-ish thing')
viewer = self.viewers[viewer_name]
if viewer.fitsimage.get_image() is None:
aim = AstroImage(logger=self.logger)
aim.load_hdu(hdu)
viewer.fitsimage.set_image(aim)
else:
viewer.fitsimage.get_image().load_hdu(hdu)
def _init_knots(self, knots, has_bounds, lower, upper):
if np.issubdtype(type(knots), np.integer):
self._t = np.concatenate((lower, np.zeros(knots), upper))
elif isiterable(knots):
self._user_knots = True
if has_bounds:
self._t = np.concatenate((lower, np.array(knots), upper))
else:
if len(knots) < 2 * (self._degree + 1):
raise ValueError(
f"Must have at least {2*(self._degree + 1)} knots.")
self._t = np.array(knots)
else:
raise ValueError(f"Knots: {knots} must be iterable or value")
# check that knots form a viable spline
self.bspline
def plot_data_and_model(self, samplerorparams, perc=50, datakwargs={}, lfkwargs={}):
from astropy.utils import isiterable
from matplotlib import pyplot as plt
if isiterable(samplerorparams):
ps = samplerorparams
else:
sampler = samplerorparams
ps = np.percentile(sampler.flatchain, perc, axis=0)
self.plot_lnprob(*ps, **lfkwargs)
n, edges = np.histogram(self.magdata, bins=datakwargs.pop('bins', 100))
cens = (edges[1:]+edges[:-1])/2
N = np.trapz(x=cens, y=n)
plt.scatter(cens, np.log(n/N), **datakwargs)
plt.ylabel('log(lf/data)')
def _validate_prepare_time(self, t, pos_c):
"""
Make sure that t is a 1D array and compatible with the C position array.
"""
if hasattr(t, 'unit'):
t = t.decompose(self.units).value
if not isiterable(t):
t = np.atleast_1d(t)
t = np.ascontiguousarray(t.ravel())
if len(t) > 1:
if len(t) != pos_c.shape[0]:
raise ValueError("If passing in an array of times, it must have a shape "
"compatible with the input position(s).")
return t
def _validate_prepare_time(self, t, pos_c):
"""
Make sure that t is a 1D array and compatible with the C position array.
"""
if hasattr(t, 'unit'):
t = t.decompose(self.units).value
if not isiterable(t):
t = np.atleast_1d(t)
t = np.ascontiguousarray(t.ravel())
if len(t) > 1:
if len(t) != pos_c.shape[0]:
raise ValueError("If passing in an array of times, it must have a shape "
"compatible with the input position(s).")
return t
def __init__(self, indx, naxis):
if _is_int(indx):
if 0 <= indx < naxis:
self.npts = 1
self.offset = indx
self.contiguous = True
else:
raise IndexError(f'Index {indx} out of range.')
elif isinstance(indx, slice):
start, stop, step = indx.indices(naxis)
self.npts = (stop - start) // step
self.offset = start
self.contiguous = step == 1
elif isiterable(indx):
self.npts = len(indx)
self.offset = 0
self.contiguous = False
else:
raise IndexError(f'Illegal index {indx}')
def get_tau0(wrest, fosc, N, b):
"""Get the value of the optical depth at the line center,
tau0. Taken from Draine 2011 (see Chapter 9). It neglects stimulated
emission which is fine for IGM or ISM except for radio-frequency
transitions.
Parameters
----------
wrest : Quantity
Rest-frame wavelength of the transition
fosc : float
Oscillator strength of the transition
N : Quantity or Quantity array
Column density
b : Quantity or Quantity array of same shape as N
Doppler parameter
Returns
-------
tau0: float or array
Optical depth at the line center. If N and b are
arrays they must be of same shape.
"""
# check format for N and b
if isiterable(N):
if np.shape(N) != np.shape(b):
raise IOError('If N is array, b must be array of same shape.')
# convert to CGS
b_cgs = b.to('cm/s')
wrest_cgs = wrest.to('cm')
N_cgs = N.to('1/cm2')
# tau0
tau0 = np.sqrt(
np.pi
) * e2_me_c_cgs * N_cgs * fosc * wrest_cgs / b_cgs # eq. 9.8 Draine 2011
# check dimensionless
tau0 = tau0.decompose()
if tau0.unit != u.dimensionless_unscaled:
raise IOError(
'Something went wrong with the units, check input units.')
return tau0.value
def __init__(self, indx, naxis):
if _is_int(indx):
if 0 <= indx < naxis:
self.npts = 1
self.offset = indx
self.contiguous = True
else:
raise IndexError('Index {} out of range.'.format(indx))
elif isinstance(indx, slice):
start, stop, step = indx.indices(naxis)
self.npts = (stop - start) // step
self.offset = start
self.contiguous = step == 1
elif isiterable(indx):
self.npts = len(indx)
self.offset = 0
self.contiguous = False
else:
raise IndexError('Illegal index {}'.format(indx))
def __init__(self, naxes, axes_type, axes_order, reference_frame=None,
reference_position=None, unit=None, axes_names=None,
name=None, wcsobj=None):
self._naxes = naxes
self._axes_order = tuple(axes_order)
if isinstance(axes_type, six.string_types):
self._axes_type = (axes_type,)
else:
self._axes_type = tuple(axes_type)
self._reference_frame = reference_frame
if unit is not None:
if astutil.isiterable(unit):
unit = tuple(unit)
else:
unit = (unit,)
if len(unit) != naxes:
raise ValueError("Number of units does not match number of axes.")
else:
self._unit = tuple([u.Unit(au) for au in unit])
if axes_names is not None:
if isinstance(axes_names, six.string_types):
axes_names = (axes_names,)
else:
axes_names = tuple(axes_names)
if len(axes_names) != naxes:
raise ValueError("Number of axes names does not match number of axes.")
else:
axes_names = tuple([""] * naxes)
self._axes_names = axes_names
if name is None:
self._name = self.__class__.__name__
else:
self._name = name
if reference_position is not None:
self._reference_position = reference_position
else:
self._reference_position = None
super(CoordinateFrame, self).__init__()
def get_tau0(wrest, fosc, N, b):
"""Get the value of the optical depth at the line center,
tau0. Taken from Draine 2011 (see Chapter 9). It neglects stimulated
emission which is fine for IGM or ISM except for radio-frequency
transitions.
Parameters
----------
wrest : Quantity
Rest-frame wavelength of the transition
fosc : float
Oscillator strength of the transition
N : Quantity or Quantity array
Column density
b : Quantity or Quantity array of same shape as N
Doppler parameter
Returns
-------
tau0: float or array
Optical depth at the line center. If N and b are
arrays they must be of same shape.
"""
# check format for N and b
if isiterable(N):
if np.shape(N) != np.shape(b):
raise IOError('If N is array, b must be array of same shape.')
# convert to CGS
b_cgs = b.to('cm/s')
wrest_cgs = wrest.to('cm')
N_cgs = N.to('1/cm2')
# tau0
tau0 = np.sqrt(np.pi) * e2_me_c_cgs * N_cgs * fosc * wrest_cgs / b_cgs # eq. 9.8 Draine 2011
# check dimensionless
tau0 = tau0.decompose()
if tau0.unit != u.dimensionless_unscaled:
raise IOError('Something went wrong with the units, check input units.')
return tau0.value
def plot_lnprob(self, tipmag, alphargb, alphaother, fracother, magrng=100, doplot=True, delog=False):
"""
Plots (optionally) and returns arrays suitable for plotting the pdf. If
`magrng` is a scalar, it gives the number of samples over the data
domain. If an array, it's used as the x axis.
"""
from astropy.utils import isiterable
from matplotlib import pyplot as plt
if isiterable(magrng):
fakemod = self.__class__(magrng)
else:
fakemod = self.__class__(np.linspace(self.mindata, self.maxdata, magrng))
lnpb = fakemod.lnprob(tipmag, alphargb, alphaother, fracother)
if delog:
lnpb = np.exp(lnpb - np.min(lnpb))
if doplot:
plt.plot(fakemod.magdata, lnpb)
return fakemod.magdata, lnpb
def create_distortion_transform(self, distortion):
dist_info = copy.deepcopy(distortion)
try:
dist_info.pop('title')
except KeyError:
pass
trans = []
mods = dist_info['models']
for m in mods:
name = m.pop('model_name')
mclass = getattr(models, name)
for p in m:
if astu.isiterable(m[p]) and 'convert2array' in m[p]:
a = np.array(m[p]['value'])
a.shape = m[p]['convert2array']
m[p] = a
trans.append(mclass(**m))
if len(trans) > 1:
return SCompositeModel(trans, inmap, outmap)
else:
return trans[0]
def _validate_shape(interval):
"""Validate the shape of an interval representation"""
MESSAGE = """An interval must be some sort of sequence of length 2"""
try:
shape = np.shape(interval)
except TypeError:
try:
# np.shape does not work with lists of Quantities
if len(interval) == 1:
interval = interval[0]
shape = np.shape([b.to_value() for b in interval])
except (ValueError, TypeError, AttributeError):
raise ValueError(MESSAGE)
valid_shape = shape in ((2,), (1, 2), (2, 0))
if not valid_shape:
valid_shape = (len(shape) > 0) and (shape[0] == 2) and \
all(isinstance(b, np.ndarray) for b in interval)
if not isiterable(interval) or not valid_shape:
raise ValueError(MESSAGE)
def undistort(self, x, y):
input_values = np.asarray(x), np.asarray(y)
output_values = [np.zeros_like(arg) for arg in input_values]
output_values.append(np.zeros_like(input_values[0]))
result = self.regions_mask[x, y]
if not astu.isiterable(result):
if result != 0:
return self._selector[result[0]](*input_values)
else:
return None #input_values
unique_regions = np.unique(result).tolist()
unique_regions.remove(0)
focx = []
focy = []
for i in unique_regions:
indices = (result==i)
transform = self._selector[i]
xyind = [val[indices] for val in input_values]
r, d = transform.undistort(x[indices], y[indices])
focx.extend(r)
focy.extend(d)
return np.asarray(focx), np.asarray(focy)
def undistort(self, x, y):
input_values = np.asarray(x), np.asarray(y)
output_values = [np.zeros_like(arg) for arg in input_values]
output_values.append(np.zeros_like(input_values[0]))
result = self.regions_mask[x, y]
if not astu.isiterable(result):
if result != 0:
return self._selector[result[0]](*input_values)
else:
return None #input_values
unique_regions = np.unique(result).tolist()
unique_regions.remove(0)
focx = []
focy = []
for i in unique_regions:
indices = (result == i)
transform = self._selector[i]
xyind = [val[indices] for val in input_values]
r, d = transform.undistort(x[indices], y[indices])
focx.extend(r)
focy.extend(d)
return np.asarray(focx), np.asarray(focy)
def __call__(self, f=None, helps=None):
"""Add a helper to a numpy function.
Normally used as a decorator.
If ``helps`` is given, it should be the numpy function helped (or an
iterable of numpy functions helped).
If ``helps`` is not given, it is assumed the function helped is the
numpy function with the same name as the decorated function.
"""
if f is not None:
if helps is None:
helps = getattr(np, f.__name__)
if not isiterable(helps):
helps = (helps, )
for h in helps:
self.assignments[h] = f
return f
elif helps is not None:
return functools.partial(self.__call__, helps=helps)
else: # pragma: no cover
raise ValueError("function_helper requires at least one argument.")
def test_rolling_window(self):
ts = [np.linspace(0.+dd, 100.+dd, 10000) for dd in np.linspace(0,20,64)]
for i,t in enumerate(ts):
print(i, t.min(), t.max(), len(t))
f = self.make_f(t)
nfreq = len(self.omega)
if not isiterable(self.p):
ps = [self.p]
else:
ps = self.p
for p in ps:
print(i, p)
# create SuperFreq object for this time array
sf = SuperFreq(t, p=p)
# try recovering the strongest frequency
w,amp,phi = sf.frecoder(f[:sf.n], break_condition=1E-5)
np.testing.assert_allclose(self.omega, w[:nfreq], rtol=1E-7)
np.testing.assert_allclose(self.A, amp[:nfreq], rtol=1E-5)
np.testing.assert_allclose(self.phi, phi[:nfreq], rtol=1E-4)
def concatenate(coords):
"""
Combine multiple coordinate objects into a single
`~astropy.coordinates.SkyCoord`.
"Coordinate objects" here mean frame objects with data,
`~astropy.coordinates.SkyCoord`, or representation objects. Currently,
they must all be in the same frame, but in a future version this may be
relaxed to allow inhomogenous sequences of objects.
Parameters
----------
coords : sequence of coordinate objects
The objects to concatenate
Returns
-------
cskycoord : SkyCoord
A single sky coordinate with its data set to the concatenation of all
the elements in ``coords``
"""
if getattr(coords, 'isscalar', False) or not isiterable(coords):
raise TypeError('The argument to concatenate must be iterable')
scs = [SkyCoord(coord, copy=False) for coord in coords]
# Check that all frames are equivalent
for sc in scs[1:]:
if not sc.is_equivalent_frame(scs[0]):
raise ValueError("All inputs must have equivalent frames: "
"{0} != {1}".format(sc, scs[0]))
# TODO: this can be changed to SkyCoord.from_representation() for a speed
# boost when we switch to using classmethods
return SkyCoord(concatenate_representations([c.data for c in coords]),
frame=scs[0].frame)
def get_ebv(self, coordinates, interpolate=True, order=1):
"""Get E(B-V) value(s) at given coordinate(s).
Parameters
----------
coordinates : astropy Coordinates object or tuple
If tuple, treated as (RA, Dec) in degrees in the ICRS (e.g.,
"J2000") system. RA and Dec can each be float or list or numpy
array.
interpolate : bool
Interpolate between the map values using
`scipy.ndimage.map_coordinates`. Default is ``True``.
order : int
Interpolation order, if interpolate=True. Default is ``1``.
Returns
-------
ebv : float or `~numpy.ndarray`
Specific extinction E(B-V) at the given locations.
"""
# Parse input
if not isinstance(coordinates, SkyCoord):
ra, dec = coordinates
coordinates = SkyCoord(ra=ra, dec=dec, frame='icrs', unit=u.degree)
# Convert to galactic coordinates.
coordinates = coordinates.galactic
l = coordinates.l.radian
b = coordinates.b.radian
# Check if l, b are scalar. If so, convert to 1-d arrays.
return_scalar = False
if not isiterable(l):
return_scalar = True
l, b = np.array([l]), np.array([b])
# Initialize return array
ebv = np.empty_like(l)
# Treat north (b>0) separately from south (b<0).
for sign, mask, ext in [(1, b >= 0, 'ngp'), (-1, b < 0, 'sgp')]:
if not np.any(mask):
continue
# Only load the FITS file for this hemisphere if it is needed
# and has not been previously loaded. Once loaded, it will be
# kept in memory for subsequent calls.
if self.__dict__[ext] is None:
hdulist = fits.open(self.fname.format(ext))
header = hdulist[0].header
self.__dict__[ext] = {'CRPIX1': header['CRPIX1'],
'CRPIX2': header['CRPIX2'],
'LAM_SCAL': header['LAM_SCAL'],
'data': hdulist[0].data}
hdulist.close()
d = self.__dict__[ext]
# Project from galactic longitude/latitude to lambert pixels.
# (See SFD98).
x = d['CRPIX1']-1. + (d['LAM_SCAL'] * np.cos(l[mask]) *
np.sqrt(1. - sign*np.sin(b[mask])))
y = d['CRPIX2']-1. - sign*(d['LAM_SCAL'] * np.sin(l[mask]) *
np.sqrt(1. - sign*np.sin(b[mask])))
# Get map values at these pixel coordinates.
if interpolate:
ebv[mask] = map_coordinates(d['data'], [y, x], order=order)
else:
x = np.round(x).astype(np.int)
y = np.round(y).astype(np.int)
ebv[mask] = d['data'][y, x]
if return_scalar:
return ebv[0]
return ebv
def plot_density_contours(self, grid, filled=True, ax=None, labels=None,
subplots_kw=dict(), **kwargs):
"""
Plot density contours. Computes the density on a grid
(specified by the array `grid`).
.. warning:: Right now the grid input must be arrays and must already be in
the unit system of the potential. Quantity support is coming...
Parameters
----------
grid : tuple
Coordinate grids or slice value for each dimension. Should be a
tuple of 1D arrays or numbers.
filled : bool (optional)
Use :func:`~matplotlib.pyplot.contourf` instead of
:func:`~matplotlib.pyplot.contour`. Default is ``True``.
ax : matplotlib.Axes (optional)
labels : iterable (optional)
List of axis labels.
subplots_kw : dict
kwargs passed to matplotlib's subplots() function if an axes object
is not specified.
kwargs : dict
kwargs passed to either contourf() or plot().
Returns
-------
fig : `~matplotlib.Figure`
"""
import matplotlib.pyplot as plt
from matplotlib import cm
# figure out which elements are iterable, which are numeric
_grids = []
_slices = []
for ii, g in enumerate(grid):
if isiterable(g):
_grids.append((ii, g))
else:
_slices.append((ii, g))
# figure out the dimensionality
ndim = len(_grids)
# if ndim > 2, don't know how to handle this!
if ndim > 2:
raise ValueError("ndim > 2: you can only make contours on a 2D grid. For other "
"dimensions, you have to specify values to slice.")
if ax is None:
# default figsize
fig, ax = plt.subplots(1, 1, **subplots_kw)
else:
fig = ax.figure
if ndim == 1:
# 1D curve
x1 = _grids[0][1]
r = np.zeros((len(_grids) + len(_slices), len(x1)))
r[_grids[0][0]] = x1
for ii, slc in _slices:
r[ii] = slc
Z = self.density(r*self.units['length']).value
ax.plot(x1, Z, **kwargs)
if labels is not None:
ax.set_xlabel(labels[0])
ax.set_ylabel("potential")
else:
# 2D contours
x1, x2 = np.meshgrid(_grids[0][1], _grids[1][1])
shp = x1.shape
x1, x2 = x1.ravel(), x2.ravel()
r = np.zeros((len(_grids) + len(_slices), len(x1)))
r[_grids[0][0]] = x1
r[_grids[1][0]] = x2
for ii, slc in _slices:
r[ii] = slc
Z = self.density(r*self.units['length']).value
# make default colormap not suck
cmap = kwargs.pop('cmap', cm.Blues)
if filled:
cs = ax.contourf(x1.reshape(shp), x2.reshape(shp), Z.reshape(shp),
cmap=cmap, **kwargs)
else:
cs = ax.contour(x1.reshape(shp), x2.reshape(shp), Z.reshape(shp),
cmap=cmap, **kwargs)
# cs.cmap.set_under('w')
# cs.cmap.set_over('k')
if labels is not None:
ax.set_xlabel(labels[0])
ax.set_ylabel(labels[1])
return fig
def animate_source(source, label=None, fps=30, length=20.,
phase_range=(None, None), wave_range=(None, None),
match_peakphase=True, match_peakflux=True,
peakwave=4000., fname=None, still=False):
"""Animate spectral timeseries of model(s) using matplotlib.animation.
*Note:* Requires matplotlib v1.1 or higher.
Parameters
----------
source : `~sncosmo.Source` or str or iterable thereof
The Source to animate or list of sources to animate.
label : str or list of str, optional
If given, label(s) for Sources, to be displayed in a legend on
the animation.
fps : int, optional
Frames per second. Default is 30.
length : float, optional
Movie length in seconds. Default is 15.
phase_range : (float, float), optional
Phase range to plot (in the timeframe of the first source if multiple
sources are given). `None` indicates to use the maximum extent of the
source(s).
wave_range : (float, float), optional
Wavelength range to plot. `None` indicates to use the maximum extent
of the source(s).
match_peakflux : bool, optional
For multiple sources, scale fluxes so that the peak of the spectrum
at the peak matches that of the first source. Default is
True.
match_peakphase : bool, optional
For multiple sources, shift additional sources so that the source's
reference phase matches that of the first source.
peakwave : float, optional
Wavelength used in match_peakflux and match_peakphase. Default is
4000.
fname : str, optional
If not `None`, save animation to file `fname`. Requires ffmpeg
to be installed with the appropriate codecs: If `fname` has
the extension '.mp4' the libx264 codec is used. If the
extension is '.webm' the VP8 codec is used. Otherwise, the
'mpeg4' codec is used. The first frame is also written to a
png.
still : bool, optional
When writing to a file, also save the first frame as a png file.
This is useful for displaying videos on a webpage.
Returns
-------
ani : `~matplotlib.animation.FuncAnimation`
Animation object that can be shown or saved.
"""
from matplotlib import pyplot as plt
from matplotlib import animation
warn_once('animate_source', '1.4', '2.0')
# Convert input to a list (if it isn't already).
if (not isiterable(source)) or isinstance(source, six.string_types):
sources = [source]
else:
sources = source
# Check that all entries are Source or strings.
for m in sources:
if not (isinstance(m, six.string_types) or isinstance(m, Source)):
raise ValueError('str or Source instance expected for '
'source(s)')
sources = [get_source(m) for m in sources]
# Get the source labels
if label is None:
labels = [None] * len(sources)
elif isinstance(label, six.string_types):
labels = [label]
else:
labels = label
if len(labels) != len(sources):
raise ValueError('if given, length of label must match '
'that of source')
# Get a wavelength array for each source.
waves = [np.arange(m.minwave(), m.maxwave(), 10.) for m in sources]
# Phase offsets needed to match peak phases.
peakphases = [m.peakphase(peakwave) for m in sources]
if match_peakphase:
phase_offsets = [p - peakphases[0] for p in peakphases]
else:
phase_offsets = [0.] * len(sources)
# Determine phase range to display.
minphase, maxphase = phase_range
if minphase is None:
minphase = min([sources[i].minphase() - phase_offsets[i] for
i in range(len(sources))])
if maxphase is None:
maxphase = max([sources[i].maxphase() - phase_offsets[i] for
i in range(len(sources))])
# Determine the wavelength range to display.
minwave, maxwave = wave_range
if minwave is None:
minwave = min([m.minwave() for m in sources])
if maxwave is None:
maxwave = max([m.maxwave() for m in sources])
# source time interval between frames
phase_interval = (maxphase - minphase) / (length * fps)
# maximum flux density of entire spectrum at the peak phase
# for each source
max_fluxes = [np.max(m.flux(phase, w))
for m, phase, w in zip(sources, peakphases, waves)]
# scaling factors
if match_peakflux:
peakfluxes = [m.flux(phase, peakwave) # Not the same as max_fluxes!
for m, phase in zip(sources, peakphases)]
scaling_factors = [peakfluxes[0] / f for f in peakfluxes]
global_max_flux = max_fluxes[0]
else:
scaling_factors = [1.] * len(sources)
global_max_flux = max(max_fluxes)
ymin = -0.06 * global_max_flux
ymax = 1.1 * global_max_flux
# Set up the figure, the axis, and the plot element we want to animate
fig = plt.figure()
ax = plt.axes(xlim=(minwave, maxwave), ylim=(ymin, ymax))
plt.axhline(y=0., c='k')
plt.xlabel('Wavelength ($\\AA$)')
plt.ylabel('Flux Density ($F_\lambda$)')
phase_text = ax.text(0.05, 0.95, '', ha='left', va='top',
transform=ax.transAxes)
empty_lists = 2 * len(sources) * [[]]
lines = ax.plot(*empty_lists, lw=1)
if label is not None:
for line, l in zip(lines, labels):
line.set_label(l)
legend = plt.legend(loc='upper right')
def init():
for line in lines:
line.set_data([], [])
phase_text.set_text('')
return tuple(lines) + (phase_text,)
def animate(i):
current_phase = minphase + phase_interval * i
for j in range(len(sources)):
y = sources[j].flux(current_phase + phase_offsets[j], waves[j])
lines[j].set_data(waves[j], y * scaling_factors[j])
phase_text.set_text('phase = {0:.1f}'.format(current_phase))
return tuple(lines) + (phase_text,)
ani = animation.FuncAnimation(fig, animate, init_func=init,
frames=int(fps*length), interval=(1000./fps),
blit=True)
# Save the animation as an mp4 or webm file.
# This requires that ffmpeg is installed.
if fname is not None:
if still:
i = fname.rfind('.')
stillfname = fname[:i] + '.png'
plt.savefig(stillfname)
ext = fname[i+1:]
codec = {'mp4': 'libx264', 'webm': 'libvpx'}.get(ext, 'mpeg4')
ani.save(fname, fps=fps, codec=codec, extra_args=['-vcodec', codec],
writer='ffmpeg_file', bitrate=1800)
plt.close()
else:
return ani
def _parse_args(args):
# first try pulling out separate coordinates
c1 = args.pop('coord1', None)
c2 = args.pop('coord2', None)
cargs = (c1, c2)
# if either are None, check for a coordstr
if c1 is None or c2 is None:
coordstr = args.pop('coordstr', None)
cargs = (coordstr,)
if (c1 is None or c2 is None) and (coordstr is None):
# TODO: throw a error
raise ValueError("Shitty arguments, yo.")
# default units are degrees
c1u = u.Unit(args.pop('coord1unit', u.deg))
c2u = u.Unit(args.pop('coord2unit', u.deg))
to_system = args.pop('to')
from_system = args.pop('from')
# extra arguments to pass to the coordinate frames
fromargs = dict()
toargs = dict()
for k, v in args.items():
if k.startswith('from_'):
fromargs[k.split('from_')[1]] = v
if k.startswith('to_'):
toargs[k.split('to_')[1]] = v
fromargs['frame'] = from_system.lower()
# get the 'to' frame class
ToFrame = frame_transform_graph.lookup_name(to_system.lower())
c = SkyCoord(*cargs, unit=(c1u, c2u), **fromargs).transform_to(frame=ToFrame(**toargs))
comps = []
for compnm in c.representation_component_names:
compval = getattr(c, compnm)
uout = c.representation_component_units.get(compnm, None)
if uout is not None:
compval.to(uout)
comps.append((compnm, compval))
output = {}
for i, (compnm, comp) in enumerate(comps):
# container for information for this coordinate
this_coord = dict()
ip1s = str(i + 1)
if not isiterable(comp.value):
val = [comp.value]
else:
val = comp.value
this_coord['data'] = list(val)
this_coord['name'] = compnm
this_coord['unit'] = str(comp.unit)
output['coord' + ip1s] = this_coord
for fattr_nm in c.get_frame_attr_names():
output['frame_' + fattr_nm] = str(getattr(c, fattr_nm))
return output
