Source code for mangadap.util.trace

# Licensed under a 3-clause BSD style license - see LICENSE.rst
# -*- coding: utf-8 -*-
"""
Provides a class to handle an SDSS trace set; cf. pydl TraceSet.

----

.. include license and copyright
.. include:: ../include/copy.rst

----

.. include common links, assuming primary doc root is up one directory
.. include:: ../include/links.rst
"""
import numpy

[docs] class TraceSet: """ Interact with a trace set written by IDLUTILS. Args: func (:obj:`str`): xmin (scalar-like): xmax (scalar-like): coeff (array-like): """ def __init__(self, func, xmin, xmax, coeff): if func not in self.allowed_functions(): raise ValueError('{0} is not a valid function for use with TraceSet.'.format(func)) self.func = func self.val, self.vander = self._get_val_and_vander() self.coeff = numpy.atleast_2d(coeff) if self.coeff.ndim != 2: raise ValueError('Coefficients array can have no more than two dimensions.') # The number of traces and the order of each trace self.size, self.order = self.coeff.shape self.order -= 1 self.xmin = xmin self.xmax = xmax self.Dx = xmax-xmin self.xmid = xmin + self.Dx/2
[docs] @classmethod def from_fits_table(cls, data): return cls(data['FUNC'][0], data['XMIN'][0], data['XMAX'][0], data['COEFF'][0])
[docs] @staticmethod def allowed_functions(): return ['poly', 'legendre', 'chebyshev']
[docs] def _get_val_and_vander(self): if self.func == 'poly': return numpy.polynomial.polynomial.polyval, numpy.polynomial.polynomial.polyvander if self.func == 'legendre': return numpy.polynomial.legendre.legval, numpy.polynomial.legendre.legvander if self.func == 'chebyshev': return numpy.polynomial.chebyshev.chebval, numpy.polynomial.chebyshev.chebvander raise ValueError('{0} is not a valid function for use with TraceSet.'.format(self.func))
[docs] def _normalized_coordinates(self, x): return 2 * (x-self.xmid)/self.Dx
[docs] def sample(self, x, index=None): """ Sample the trace at the provided coordinates. Args: x (array-like): One-dimensional array with coordinates to sample. index (:obj:`int`, optional): Sample a single trace with the provided index. Default is to return an array sampling all traces. Returns: numpy.ndarray: A 1 or 2D array with the samples of one or more traces. """ if index is not None: return self.val(self._normalized_coordinates(x), self.coeff[index,:]) return numpy.dot(coeff, self.vander(self._normalized_coordinates(x), self.order))