pyreduce-astro 0.7a4__cp314-cp314-win_amd64.whl

pyreduce/echelle.py

@@ -0,0 +1,413 @@
+"""
+Contains functions to read and modify echelle structures, just as in reduce
+
+Mostly for compatibility reasons
+"""
+
+import logging
+
+import astropy.io.fits as fits
+import numpy as np
+import scipy.constants
+
+logger = logging.getLogger(__name__)
+
+
+class Echelle:
+    def __init__(self, head=None, filename="", data=None):
+        if data is None:
+            data = {}
+        if head is None:
+            head = {}
+        self.filename = filename
+        self.header = head
+        self._data = data
+
+    @property
+    def nord(self):
+        if "spec" in self._data.keys():
+            return self._data["spec"].shape[0]
+        return None
+
+    @property
+    def ncol(self):
+        if "spec" in self._data.keys():
+            return self._data["spec"].shape[1]
+        return None
+
+    @property
+    def spec(self):
+        if "spec" in self._data.keys():
+            return self._data["spec"]
+        else:
+            return None
+
+    @spec.setter
+    def spec(self, value):
+        self._data["spec"] = value
+
+    @property
+    def sig(self):
+        if "sig" in self._data.keys():
+            return self._data["sig"]
+        else:
+            return None
+
+    @sig.setter
+    def sig(self, value):
+        self._data["sig"] = value
+
+    @property
+    def wave(self):
+        if "wave" in self._data.keys():
+            return self._data["wave"]
+        else:
+            return None
+
+    @wave.setter
+    def wave(self, value):
+        self._data["wave"] = value
+
+    @property
+    def cont(self):
+        if "cont" in self._data.keys():
+            return self._data["cont"]
+        else:
+            return None
+
+    @cont.setter
+    def cont(self, value):
+        self._data["cont"] = value
+
+    @property
+    def columns(self):
+        if "columns" in self._data.keys():
+            return self._data["columns"]
+        else:
+            return None
+
+    @columns.setter
+    def columns(self, value):
+        self._data["columns"] = value
+
+    @property
+    def mask(self):
+        if "mask" in self._data.keys():
+            return self._data["mask"]
+        else:
+            return None
+
+    @mask.setter
+    def mask(self, value):
+        self._data["mask"] = value
+
+    def __getitem__(self, index):
+        return self._data[index]
+
+    def __setitem__(self, index, value):
+        self._data[index] = value
+
+    def __delitem__(self, index):
+        del self._data[index]
+
+    def __contains__(self, index):
+        return index in self._data.keys()
+
+    @staticmethod
+    def read(
+        fname,
+        extension=1,
+        raw=False,
+        continuum_normalization=True,
+        barycentric_correction=True,
+        radial_velociy_correction=True,
+    ):
+        """
+        Read data from an echelle file
+        Expand wavelength and continuum polynomials
+        Apply barycentric/radial velocity correction
+        Apply continuum normalization
+
+        Will load any fields in the binary table, however special attention is given only to specific names:
+        "SPEC"    : Spectrum
+        "SIG"     : Sigma, i.e. (absolute) uncertainty
+        "CONT"    : Continuum
+        "WAVE"    : Wavelength solution
+        "COLUMNS" : Column range
+
+        Parameters
+        ----------
+        fname : str
+            filename to load
+        extension : int, optional
+            fits extension of the data within the file (default: 1)
+        raw : bool, optional
+            if true apply no corrections to the data (default: False)
+        continuum_normalization : bool, optional
+            apply continuum normalization (default: True)
+        barycentric_correction : bool, optional
+            apply barycentric correction (default: True)
+        radial_velociy_correction : bool, optional
+            apply radial velocity correction (default: True)
+
+        Returns
+        -------
+        ech : obj
+            Echelle structure, with data contained in attributes
+        """
+
+        with fits.open(fname, memmap=False) as hdu:
+            header = hdu[0].header
+            data = hdu[extension].data
+
+        _data = {column.lower(): data[column][0] for column in data.dtype.names}
+        ech = Echelle(filename=fname, head=header, data=_data)
+        nord, ncol = ech.nord, ech.ncol
+
+        if not raw:
+            if "spec" in ech:
+                base_order = header.get("obase", 1)
+                ech["orders"] = np.arange(base_order, base_order + nord)
+
+            # Wavelength
+            if "wave" in ech:
+                ech["wave"] = expand_polynomial(ncol, ech["wave"])
+
+                # Correct for radial velocity and barycentric correction
+                # + : away from observer
+                # - : towards observer
+                velocity_correction = 0
+                if barycentric_correction:
+                    velocity_correction -= header.get("barycorr", 0)
+                    header["barycorr"] = 0
+                if radial_velociy_correction:
+                    velocity_correction += header.get("radvel", 0)
+                    header["radvel"] = 0
+
+                speed_of_light = scipy.constants.speed_of_light * 1e-3
+                ech["wave"] *= 1 + velocity_correction / speed_of_light
+
+            # Continuum
+            if "cont" in ech:
+                ech["cont"] = expand_polynomial(ncol, ech["cont"])
+
+            # Create Mask, based on column range
+            if "columns" in ech:
+                ech["mask"] = np.full((nord, ncol), True)
+                for iord in range(nord):
+                    ech["mask"][
+                        iord, ech["columns"][iord, 0] : ech["columns"][iord, 1]
+                    ] = False
+
+                if "spec" in ech:
+                    ech["spec"] = np.ma.masked_array(ech["spec"], mask=ech["mask"])
+                if "sig" in ech:
+                    ech["sig"] = np.ma.masked_array(ech["sig"], mask=ech["mask"])
+                if "cont" in ech:
+                    ech["cont"] = np.ma.masked_array(ech["cont"], mask=ech["mask"])
+                if "wave" in ech:
+                    ech["wave"] = np.ma.masked_array(ech["wave"], mask=ech["mask"])
+
+            # Apply continuum normalization
+            if continuum_normalization and "cont" in ech:
+                if "spec" in ech:
+                    ech["spec"] /= ech["cont"]
+                if "sig" in ech:
+                    ech["sig"] /= ech["cont"]
+
+        return ech
+
+    def save(self, fname):
+        """
+        Save data in an Echelle fits, i.e. a fits file with a Binary Table in Extension 1
+
+        Parameters
+        ----------
+        fname : str
+            filename
+        """
+        save(fname, self.header, **self._data)
+
+
+def calc_2dpolynomial(solution2d):
+    """Expand a 2d polynomial, where the data is given in a REDUCE make_wave format
+    Note that the coefficients are for order/100 and column/1000 respectively, where the order is counted from order base up
+
+    Parameters
+    ----------
+    solution2d : array
+        data in a REDUCE make_wave format:
+        0: version
+        1: number of columns
+        2: number of orders
+        3: order base, i.e. 0th order number (should not be 0)
+        4-6: empty
+        7: number of cross coefficients
+        8: number of column only coefficients
+        9: number of order only coefficients
+        10: coefficient - constant
+        11-x: column coefficients
+        x-y : order coefficients
+        z-  : cross coefficients (xy, xy2, x2y, x2y2, xy3, x3y), with x = orders, y = columns
+
+    Returns
+    -------
+    poly : array[nord, ncol]
+        expanded polynomial
+    """
+
+    # make wave style 2d fit
+    ncol = int(solution2d[1])
+    nord = int(solution2d[2])
+    order_base = int(solution2d[3])
+    deg_cross, deg_column, deg_order = (
+        int(solution2d[7]),
+        int(solution2d[8]),
+        int(solution2d[9]),
+    )
+    coeff_in = solution2d[10:]
+
+    coeff = np.zeros((deg_order + 1, deg_column + 1))
+    coeff[0, 0] = coeff_in[0]
+    coeff[0, 1:] = coeff_in[1 : 1 + deg_column]
+    coeff[1:, 0] = coeff_in[1 + deg_column : 1 + deg_column + deg_order]
+    if deg_cross in [4, 6]:
+        coeff[1, 1] = coeff_in[deg_column + deg_order + 1]
+        coeff[1, 2] = coeff_in[deg_column + deg_order + 2]
+        coeff[2, 1] = coeff_in[deg_column + deg_order + 3]
+        coeff[2, 2] = coeff_in[deg_column + deg_order + 4]
+    if deg_cross == 6:
+        coeff[1, 3] = coeff_in[deg_column + deg_order + 5]
+        coeff[3, 1] = coeff_in[deg_column + deg_order + 6]
+
+    x = np.arange(order_base, order_base + nord, dtype=float)
+    y = np.arange(ncol, dtype=float)
+
+    poly = np.polynomial.polynomial.polygrid2d(x / 100, y / 1000, coeff) / x[:, None]
+
+    return poly
+
+
+def calc_1dpolynomials(ncol, poly):
+    """Expand a set of 1d polynomials (one per order) seperately
+
+    Parameters
+    ----------
+    ncol : int
+        number of columns
+    poly : array[nord, degree]
+        polynomial coefficients
+
+    Returns
+    -------
+    poly : array[nord, ncol]
+        expanded polynomials
+    """
+
+    nord = poly.shape[0]
+    x = np.arange(ncol)
+    result = np.zeros((nord, ncol))
+    for i, coef in enumerate(poly):
+        result[i] = np.polyval(coef, x)
+    return result
+
+
+def expand_polynomial(ncol, poly):
+    """Checks if and how to expand data poly, then expands the data if necessary
+
+    Parameters
+    ----------
+    ncol : int
+        number of columns in the image
+    poly : array[nord, ...]
+        polynomial coefficients to expand, or already expanded data
+
+    Returns
+    -------
+    poly : array[nord, ncol]
+        expanded data
+    """
+
+    if poly.ndim == 1:
+        poly = calc_2dpolynomial(poly)
+    elif poly.shape[1] < 20:
+        poly = calc_1dpolynomials(ncol, poly)
+    return poly
+
+
+def read(fname, **kwargs):
+    return Echelle.read(fname, **kwargs)
+
+
+def save(fname, header, **kwargs):
+    """Save data in an Echelle fits, i.e. a fits file with a Binary Table in Extension 1
+
+    The data is passed in kwargs, with the name of the binary table column as the key
+    Floating point data is saved as float32 (E), Integer data as int16 (I)
+
+    Parameters
+    ----------
+    fname : str
+        filename
+    header : fits.header
+        FITS header
+    **kwargs : array[]
+        data to be saved in the file
+    """
+
+    if not isinstance(header, fits.Header):
+        header = fits.Header(cards=header)
+
+    primary = fits.PrimaryHDU(header=header)
+
+    columns = []
+    for key, value in kwargs.items():
+        if value is None:
+            continue
+
+        if isinstance(value, list):
+            for i, arr in enumerate(value):
+                if np.issubdtype(arr.dtype, np.floating):
+                    arr = arr.astype(np.float32)
+                    format_code = "E"
+                elif np.issubdtype(arr.dtype, np.integer):
+                    arr = arr.astype(np.int16)
+                    format_code = "I"
+                else:
+                    raise ValueError(
+                        f"Unsupported data type for array {i}: {arr.dtype}"
+                    )
+
+                # Create a column with the appropriate format
+                col = fits.Column(
+                    name=f"slitfu{i}", format=f"{len(arr)}{format_code}", array=[arr]
+                )
+                columns.append(col)
+
+        else:
+            arr = value.ravel()[None, :]
+
+            if np.issubdtype(arr.dtype, np.floating):
+                if key == "wave":
+                    arr = arr.astype(np.float64)
+                    dtype = "D"
+                else:
+                    arr = arr.astype(np.float32)
+                    dtype = "E"
+            elif np.issubdtype(arr.dtype, np.integer):
+                arr = arr.astype(np.int16)
+                dtype = "I"
+            elif np.issubdtype(arr.dtype, np.bool_):
+                arr = arr.astype(np.bool_)
+                dtype = "B"
+            else:
+                raise TypeError(f"Could not understand dtype {arr.dtype}")
+
+            form = "%i%s" % (value.size, dtype)
+            dim = str(value.shape[::-1])
+            columns += [fits.Column(name=key.upper(), array=arr, format=form, dim=dim)]
+
+    table = fits.BinTableHDU.from_columns(columns)
+
+    hdulist = fits.HDUList(hdus=[primary, table])
+    hdulist.writeto(fname, overwrite=True, output_verify="silentfix+ignore")

pyreduce/estimate_background_scatter.py

@@ -0,0 +1,130 @@
+"""
+Module that estimates the background scatter
+"""
+
+import logging
+
+import matplotlib.pyplot as plt
+import numpy as np
+
+from . import util
+from .extract import fix_parameters
+from .util import make_index, polyfit2d
+
+logger = logging.getLogger(__name__)
+
+
+def estimate_background_scatter(
+    img,
+    orders,
+    column_range=None,
+    extraction_width=0.1,
+    scatter_degree=4,
+    sigma_cutoff=2,
+    border_width=10,
+    plot=False,
+    plot_title=None,
+):
+    """
+    Estimate the background by fitting a 2d polynomial to interorder data
+
+    Interorder data is all pixels minus the orders +- the extraction width
+
+    Parameters
+    ----------
+    img : array[nrow, ncol]
+        (flat) image data
+    orders : array[nord, degree]
+        order polynomial coefficients
+    column_range : array[nord, 2], optional
+        range of columns to use in each order (default: None == all columns)
+    extraction_width : float, array[nord, 2], optional
+        extraction width for each order, values below 1.5 are considered fractional, others as number of pixels (default: 0.1)
+    scatter_degree : int, optional
+        polynomial degree of the 2d fit for the background scatter (default: 4)
+    plot : bool, optional
+        wether to plot the fitted polynomial and the data or not (default: False)
+
+    Returns
+    -------
+    array[nord+1, ncol]
+        background scatter between orders
+    array[nord+1, ncol]
+        y positions of the interorder lines, the scatter values are taken from
+    """
+
+    nrow, ncol = img.shape
+    nord, _ = orders.shape
+
+    extraction_width, column_range, orders = fix_parameters(
+        extraction_width,
+        column_range,
+        orders,
+        nrow,
+        ncol,
+        nord,
+        ignore_column_range=True,
+    )
+
+    # Method 1: Select all pixels, but those known to be in orders
+    bw = border_width
+    mask = np.full(img.shape, True)
+    if bw is not None and bw != 0:
+        mask[:bw] = mask[-bw:] = mask[:, :bw] = mask[:, -bw:] = False
+    for i in range(nord):
+        left, right = column_range[i]
+        left -= extraction_width[i, 1] * 2
+        right += extraction_width[i, 0] * 2
+        left = max(0, left)
+        right = min(ncol, right)
+
+        x_order = np.arange(left, right)
+        y_order = np.polyval(orders[i], x_order)
+
+        y_above = y_order + extraction_width[i, 1]
+        y_below = y_order - extraction_width[i, 0]
+
+        y_above = np.floor(y_above)
+        y_below = np.ceil(y_below)
+
+        index = make_index(y_below, y_above, left, right, zero=True)
+        np.clip(index[0], 0, nrow - 1, out=index[0])
+
+        mask[index] = False
+
+    mask &= ~np.ma.getmask(img)
+
+    y, x = np.indices(mask.shape)
+    y, x = y[mask].ravel(), x[mask].ravel()
+    z = np.ma.getdata(img[mask]).ravel()
+
+    mask = z <= np.median(z) + sigma_cutoff * z.std()
+    y, x, z = y[mask], x[mask], z[mask]
+
+    coeff = polyfit2d(x, y, z, degree=scatter_degree, plot=plot, plot_title=plot_title)
+    logger.debug("Background scatter coefficients: %s", str(coeff))
+
+    if plot:  # pragma: no cover
+        # Calculate scatter at interorder positionsq
+        yp, xp = np.indices(img.shape)
+        back = np.polynomial.polynomial.polyval2d(xp, yp, coeff)
+
+        plt.subplot(121)
+        plt.title("Input Image + In-between Order traces")
+        plt.xlabel("x [pixel]")
+        plt.ylabel("y [pixel]")
+        vmin, vmax = np.percentile(img - back, (5, 95))
+        plt.imshow(img - back, vmin=vmin, vmax=vmax, aspect="equal", origin="lower")
+        plt.plot(x, y, ",")
+
+        plt.subplot(122)
+        plt.title("2D fit to the scatter between orders")
+        plt.xlabel("x [pixel]")
+        plt.ylabel("y [pixel]")
+        plt.imshow(back, vmin=0, vmax=abs(np.max(back)), aspect="equal", origin="lower")
+
+        if plot_title is not None:
+            plt.suptitle(plot_title)
+        util.show_or_save("scatter")
+
+    return coeff