|
| 1 | +import numpy as np |
| 2 | +from gwcs import WCS as GWCS |
| 3 | +from gwcs import coordinate_frames as cf |
| 4 | +from ndcube import NDCube |
| 5 | + |
| 6 | +import astropy.units as u |
| 7 | +from astropy.coordinates import SpectralCoord |
| 8 | +from astropy.modeling.tabular import Tabular1D |
| 9 | +from astropy.utils import lazyproperty |
| 10 | + |
| 11 | +__all__ = ['gwcs_from_array', 'SpectralAxis', 'Spectrum', 'CountSpectrum'] |
| 12 | + |
| 13 | + |
| 14 | +def gwcs_from_array(array): |
| 15 | + """ |
| 16 | + Create a new WCS from provided tabular data. This defaults to being |
| 17 | + a GWCS object. |
| 18 | + """ |
| 19 | + orig_array = u.Quantity(array) |
| 20 | + |
| 21 | + coord_frame = cf.CoordinateFrame(naxes=1, |
| 22 | + axes_type=('SPECTRAL',), |
| 23 | + axes_order=(0,)) |
| 24 | + spec_frame = cf.SpectralFrame(unit=array.unit, axes_order=(0,)) |
| 25 | + |
| 26 | + # In order for the world_to_pixel transformation to automatically convert |
| 27 | + # input units, the equivalencies in the lookup table have to be extended |
| 28 | + # with spectral unit information. |
| 29 | + SpectralTabular1D = type("SpectralTabular1D", (Tabular1D,), |
| 30 | + {'input_units_equivalencies': {'x0': u.spectral()}}) |
| 31 | + |
| 32 | + forward_transform = SpectralTabular1D(np.arange(len(array)), |
| 33 | + lookup_table=array) |
| 34 | + # If our spectral axis is in descending order, we have to flip the lookup |
| 35 | + # table to be ascending in order for world_to_pixel to work. |
| 36 | + if len(array) == 0 or array[-1] > array[0]: |
| 37 | + forward_transform.inverse = SpectralTabular1D( |
| 38 | + array, lookup_table=np.arange(len(array))) |
| 39 | + else: |
| 40 | + forward_transform.inverse = SpectralTabular1D( |
| 41 | + array[::-1], lookup_table=np.arange(len(array))[::-1]) |
| 42 | + |
| 43 | + class SpectralGWCS(GWCS): |
| 44 | + def pixel_to_world(self, *args, **kwargs): |
| 45 | + if orig_array.unit == '': |
| 46 | + return u.Quantity(super().pixel_to_world_values(*args, **kwargs)) |
| 47 | + return super().pixel_to_world(*args, **kwargs).to( |
| 48 | + orig_array.unit, equivalencies=u.spectral()) |
| 49 | + |
| 50 | + tabular_gwcs = SpectralGWCS(forward_transform=forward_transform, |
| 51 | + input_frame=coord_frame, |
| 52 | + output_frame=spec_frame) |
| 53 | + |
| 54 | + # Store the intended unit from the origin input array |
| 55 | + # tabular_gwcs._input_unit = orig_array.unit |
| 56 | + |
| 57 | + return tabular_gwcs |
| 58 | + |
| 59 | + |
| 60 | +class SpectralAxis(SpectralCoord): |
| 61 | + """ |
| 62 | + Coordinate object representing spectral values corresponding to a specific |
| 63 | + spectrum. Overloads SpectralCoord with additional information (currently |
| 64 | + only bin edges). |
| 65 | +
|
| 66 | + Parameters |
| 67 | + ---------- |
| 68 | + bin_specification: str, optional |
| 69 | + Must be "edges" or "centers". Determines whether specified axis values |
| 70 | + are interpreted as bin edges or bin centers. Defaults to "centers". |
| 71 | + """ |
| 72 | + |
| 73 | + _equivalent_unit = SpectralCoord._equivalent_unit + (u.pixel,) |
| 74 | + |
| 75 | + def __new__(cls, value, *args, bin_specification="centers", **kwargs): |
| 76 | + |
| 77 | + # Enforce pixel axes are ascending |
| 78 | + if ((type(value) is u.quantity.Quantity) and |
| 79 | + (value.size > 1) and |
| 80 | + (value.unit is u.pix) and |
| 81 | + (value[-1] <= value[0])): |
| 82 | + raise ValueError("u.pix spectral axes should always be ascending") |
| 83 | + |
| 84 | + # Convert to bin centers if bin edges were given, since SpectralCoord |
| 85 | + # only accepts centers |
| 86 | + if bin_specification == "edges": |
| 87 | + bin_edges = value |
| 88 | + value = SpectralAxis._centers_from_edges(value) |
| 89 | + |
| 90 | + obj = super().__new__(cls, value, *args, **kwargs) |
| 91 | + |
| 92 | + if bin_specification == "edges": |
| 93 | + obj._bin_edges = bin_edges |
| 94 | + |
| 95 | + return obj |
| 96 | + |
| 97 | + @staticmethod |
| 98 | + def _edges_from_centers(centers, unit): |
| 99 | + """ |
| 100 | + Calculates interior bin edges based on the average of each pair of |
| 101 | + centers, with the two outer edges based on extrapolated centers added |
| 102 | + to the beginning and end of the spectral axis. |
| 103 | + """ |
| 104 | + a = np.insert(centers, 0, 2*centers[0] - centers[1]) |
| 105 | + b = np.append(centers, 2*centers[-1] - centers[-2]) |
| 106 | + edges = (a + b) / 2 |
| 107 | + return edges*unit |
| 108 | + |
| 109 | + @staticmethod |
| 110 | + def _centers_from_edges(edges): |
| 111 | + """ |
| 112 | + Calculates the bin centers as the average of each pair of edges |
| 113 | + """ |
| 114 | + return (edges[1:] + edges[:-1]) / 2 |
| 115 | + |
| 116 | + @lazyproperty |
| 117 | + def bin_edges(self): |
| 118 | + """ |
| 119 | + Calculates bin edges if the spectral axis was created with centers |
| 120 | + specified. |
| 121 | + """ |
| 122 | + if hasattr(self, '_bin_edges'): |
| 123 | + return self._bin_edges |
| 124 | + else: |
| 125 | + return self._edges_from_centers(self.value, self.unit) |
| 126 | + |
| 127 | + |
| 128 | +class Spectrum(NDCube): |
| 129 | + r""" |
| 130 | + Spectrum container for data with one spectral axis. |
| 131 | +
|
| 132 | + Note that "1D" in this case refers to the fact that there is only one |
| 133 | + spectral axis. `Spectrum` can contain "vector 1D spectra" by having the |
| 134 | + ``flux`` have a shape with dimension greater than 1. |
| 135 | +
|
| 136 | + Notes |
| 137 | + ----- |
| 138 | + A stripped down version of `Spectrum1D` from `specutils`. |
| 139 | +
|
| 140 | + Parameters |
| 141 | + ---------- |
| 142 | + data : `~astropy.units.Quantity` |
| 143 | + The data for this spectrum. This can be a simple `~astropy.units.Quantity`, |
| 144 | + or an existing `~Spectrum1D` or `~ndcube.NDCube` object. |
| 145 | + uncertainty : `~astropy.nddata.NDUncertainty` |
| 146 | + Contains uncertainty information along with propagation rules for |
| 147 | + spectrum arithmetic. Can take a unit, but if none is given, will use |
| 148 | + the unit defined in the flux. |
| 149 | + spectral_axis : `~astropy.units.Quantity` or `~specutils.SpectralAxis` |
| 150 | + Dispersion information with the same shape as the dimension specified by spectral_dimension |
| 151 | + of shape plus one if specifying bin edges. |
| 152 | + spectral_dimension : `int` default 0 |
| 153 | + The dimension of the data which represents the spectral information default to first dimension index 0. |
| 154 | + mask : `~numpy.ndarray`-like |
| 155 | + Array where values in the flux to be masked are those that |
| 156 | + ``astype(bool)`` converts to True. (For example, integer arrays are not |
| 157 | + masked where they are 0, and masked for any other value.) |
| 158 | + meta : dict |
| 159 | + Arbitrary container for any user-specific information to be carried |
| 160 | + around with the spectrum container object. |
| 161 | + """ |
| 162 | + |
| 163 | + def __init__(self, data, *, uncertainty=None, spectral_axis=None, |
| 164 | + spectral_dimension=0, mask=None, meta=None, **kwargs): |
| 165 | + |
| 166 | + # If the flux (data) argument is already a Spectrum (as it would |
| 167 | + # be for internal arithmetic operations), avoid setup entirely. |
| 168 | + if isinstance(data, Spectrum): |
| 169 | + super().__init__(data) |
| 170 | + return |
| 171 | + |
| 172 | + # Ensure that the flux argument is an astropy quantity |
| 173 | + if data is not None: |
| 174 | + if not isinstance(data, u.Quantity): |
| 175 | + raise ValueError("Flux must be a `Quantity` object.") |
| 176 | + elif data.isscalar: |
| 177 | + data = u.Quantity([data]) |
| 178 | + |
| 179 | + # Ensure that the unit information codified in the quantity object is |
| 180 | + # the One True Unit. |
| 181 | + kwargs.setdefault('unit', data.unit if isinstance(data, u.Quantity) |
| 182 | + else kwargs.get('unit')) |
| 183 | + |
| 184 | + # If flux and spectral axis are both specified, check that their lengths |
| 185 | + # match or are off by one (implying the spectral axis stores bin edges) |
| 186 | + if data is not None and spectral_axis is not None: |
| 187 | + if spectral_axis.shape[0] == data.shape[spectral_dimension]: |
| 188 | + bin_specification = "centers" |
| 189 | + elif spectral_axis.shape[0] == data.shape[spectral_dimension]+1: |
| 190 | + bin_specification = "edges" |
| 191 | + else: |
| 192 | + raise ValueError( |
| 193 | + "Spectral axis length ({}) must be the same size or one " |
| 194 | + "greater (if specifying bin edges) than that of the spextral" |
| 195 | + "axis ({})".format(spectral_axis.shape[0], |
| 196 | + data.shape[spectral_dimension])) |
| 197 | + |
| 198 | + # Attempt to parse the spectral axis. If none is given, try instead to |
| 199 | + # parse a given wcs. This is put into a GWCS object to |
| 200 | + # then be used behind-the-scenes for all operations. |
| 201 | + if spectral_axis is not None: |
| 202 | + # Ensure that the spectral axis is an astropy Quantity |
| 203 | + if not isinstance(spectral_axis, u.Quantity): |
| 204 | + raise ValueError("Spectral axis must be a `Quantity` or " |
| 205 | + "`SpectralAxis` object.") |
| 206 | + |
| 207 | + # If a spectral axis is provided as an astropy Quantity, convert it |
| 208 | + # to a SpectralAxis object. |
| 209 | + if not isinstance(spectral_axis, SpectralAxis): |
| 210 | + if spectral_axis.shape[0] == data.shape[spectral_dimension] + 1: |
| 211 | + bin_specification = "edges" |
| 212 | + else: |
| 213 | + bin_specification = "centers" |
| 214 | + self._spectral_axis = SpectralAxis( |
| 215 | + spectral_axis, |
| 216 | + bin_specification=bin_specification) |
| 217 | + |
| 218 | + wcs = gwcs_from_array(self._spectral_axis) |
| 219 | + |
| 220 | + super().__init__( |
| 221 | + data=data.value if isinstance(data, u.Quantity) else data, |
| 222 | + wcs=wcs, mask=mask, meta=meta, uncertainty=uncertainty, **kwargs) |
| 223 | + |
| 224 | + |
| 225 | +class CountSpectrum(Spectrum): |
| 226 | + r""" |
| 227 | + Spectrum container for count data with one spectral axis. |
| 228 | +
|
| 229 | + Specifically, the data must be supplied as counts or convertable to counts by multiplying by the provided |
| 230 | + normalisation. |
| 231 | +
|
| 232 | + Parameters |
| 233 | + ---------- |
| 234 | + data : `~astropy.units.Quantity` |
| 235 | + The data for this spectrum. This can be a simple `~astropy.units.Quantity`, |
| 236 | + or an existing `~Spectrum1D` or `~ndcube.NDCube` object. |
| 237 | + norm : `~astropy.units.Quantity` |
| 238 | + The normalisation if the data unit is not counts then the product the unit of data and norm must be counts. |
| 239 | + uncertainty : `~astropy.nddata.NDUncertainty` |
| 240 | + Contains uncertainty information along with propagation rules for spectrum arithmetic. Can take a unit, but if |
| 241 | + none is given, will use the unit defined in the flux. |
| 242 | + spectral_axis : `~astropy.units.Quantity` or `~specutils.SpectralAxis` |
| 243 | + Dispersion information with the same shape as the dimension specified by spectral_dimension |
| 244 | + of shape plus one if specifying bin edges. |
| 245 | + spectral_dimension : `int` default 0 |
| 246 | + The dimension of the data which represents the spectral information default to first dimension index 0. |
| 247 | + mask : `~numpy.ndarray`-like |
| 248 | + Array where values in the flux to be masked are those that |
| 249 | + ``astype(bool)`` converts to True. (For example, integer arrays are not |
| 250 | + masked where they are 0, and masked for any other value.) |
| 251 | + meta : dict |
| 252 | + Arbitrary container for any user-specific information to be carried |
| 253 | + around with the spectrum container object. |
| 254 | + """ |
| 255 | + |
| 256 | + def __init__(self, data, norm=None, **kwargs): |
| 257 | + if data.unit != u.ct: |
| 258 | + data_norm_unit = (data.unit * norm.unit).decompose().bases[0] |
| 259 | + if data_norm_unit != u.ct: |
| 260 | + raise ValueError('Data must be supplied in counts or the product of the norm and data has units of ' |
| 261 | + 'counts') |
| 262 | + data = data * norm |
| 263 | + self.norm = norm |
| 264 | + super().__init__(data, **kwargs) |
0 commit comments