pyreduce-astro 0.7a4__cp314-cp314-win_amd64.whl

pyreduce/cwrappers.py

@@ -0,0 +1,404 @@
+"""
+Wrapper for REDUCE C functions
+
+This module provides access to the extraction algorithms in the
+C libraries and sanitizes the input parameters.
+
+"""
+
+import ctypes
+import logging
+
+import numpy as np
+from scipy.ndimage import median_filter
+
+logger = logging.getLogger(__name__)
+
+try:
+    from .clib._slitfunc_2d import ffi
+    from .clib._slitfunc_2d import lib as slitfunc_2dlib
+    from .clib._slitfunc_bd import lib as slitfunclib
+except ImportError:  # pragma: no cover
+    logger.error(
+        "C libraries could not be found. Compiling them by running build_extract.py"
+    )
+    from .clib import build_extract
+
+    build_extract.build()
+    del build_extract
+
+    from .clib._slitfunc_2d import ffi
+    from .clib._slitfunc_2d import lib as slitfunc_2dlib
+    from .clib._slitfunc_bd import lib as slitfunclib
+
+
+c_double = ctypes.c_double
+c_int = ctypes.c_int
+c_mask = ctypes.c_ubyte
+
+
+def slitfunc(img, ycen, lambda_sp=0, lambda_sf=0.1, osample=1):
+    """Decompose image into spectrum and slitfunction
+
+    This is for horizontal straight orders only, for curved orders use slitfunc_curved instead
+
+    Parameters
+    ----------
+    img : array[n, m]
+        image to decompose, should just contain a small part of the overall image
+    ycen : array[n]
+        traces the center of the order along the image, relative to the center of the image?
+    lambda_sp : float, optional
+        smoothing parameter of the spectrum (the default is 0, which no smoothing)
+    lambda_sf : float, optional
+        smoothing parameter of the slitfunction (the default is 0.1, which )
+    osample : int, optional
+        Subpixel ovsersampling factor (the default is 1, which no oversampling)
+
+    Returns
+    -------
+    sp, sl, model, unc
+        spectrum, slitfunction, model, spectrum uncertainties
+    """
+
+    # Convert input to expected datatypes
+    lambda_sf = float(lambda_sf)
+    lambda_sp = float(lambda_sp)
+    osample = int(osample)
+    img = np.asanyarray(img, dtype=c_double)
+    ycen = np.asarray(ycen, dtype=c_double)
+
+    assert img.ndim == 2, "Image must be 2 dimensional"
+    assert ycen.ndim == 1, "Ycen must be 1 dimensional"
+
+    assert img.shape[1] == ycen.size, (
+        f"Image and Ycen shapes are incompatible, got {img.shape} and {ycen.shape}"
+    )
+
+    assert osample > 0, f"Oversample rate must be positive, but got {osample}"
+    assert lambda_sf >= 0, (
+        f"Slitfunction smoothing must be positive, but got {lambda_sf}"
+    )
+    assert lambda_sp >= 0, f"Spectrum smoothing must be positive, but got {lambda_sp}"
+
+    # Get some derived values
+    nrows, ncols = img.shape
+    ny = osample * (nrows + 1) + 1
+    ycen = ycen - ycen.astype(c_int)
+
+    # Prepare all arrays
+    # Inital guess for slit function and spectrum
+    sp = np.ma.sum(img, axis=0)
+    requirements = ["C", "A", "W", "O"]
+    sp = np.require(sp, dtype=c_double, requirements=requirements)
+
+    sl = np.zeros(ny, dtype=c_double)
+
+    mask = ~np.ma.getmaskarray(img)
+    mask = np.require(mask, dtype=c_int, requirements=requirements)
+
+    img = np.ma.getdata(img)
+    img = np.require(img, dtype=c_double, requirements=requirements)
+
+    pix_unc = np.zeros_like(img)
+    pix_unc = np.require(pix_unc, dtype=c_double, requirements=requirements)
+
+    ycen = np.require(ycen, dtype=c_double, requirements=requirements)
+    model = np.zeros((nrows, ncols), dtype=c_double)
+    unc = np.zeros(ncols, dtype=c_double)
+
+    # Call the C function
+    slitfunclib.slit_func_vert(
+        ffi.cast("int", ncols),
+        ffi.cast("int", nrows),
+        ffi.cast("double *", img.ctypes.data),
+        ffi.cast("double *", pix_unc.ctypes.data),
+        ffi.cast("int *", mask.ctypes.data),
+        ffi.cast("double *", ycen.ctypes.data),
+        ffi.cast("int", osample),
+        ffi.cast("double", lambda_sp),
+        ffi.cast("double", lambda_sf),
+        ffi.cast("double *", sp.ctypes.data),
+        ffi.cast("double *", sl.ctypes.data),
+        ffi.cast("double *", model.ctypes.data),
+        ffi.cast("double *", unc.ctypes.data),
+    )
+    mask = ~mask.astype(bool)
+
+    return sp, sl, model, unc, mask
+
+
+def slitfunc_curved(
+    img, ycen, tilt, shear, lambda_sp, lambda_sf, osample, yrange, maxiter=20, gain=1
+):
+    """Decompose an image into a spectrum and a slitfunction, image may be curved
+
+    Parameters
+    ----------
+    img : array[n, m]
+        input image
+    ycen : array[n]
+        traces the center of the order
+    tilt : array[n]
+        tilt (1st order curvature) of the order along the image, set to 0 if order straight
+    shear : array[n]
+        shear (2nd order curvature) of the order along the image, set to 0 if order straight
+    osample : int
+        Subpixel ovsersampling factor (the default is 1, no oversampling)
+    lambda_sp : float
+        smoothing factor spectrum (the default is 0, no smoothing)
+    lambda_sl : float
+        smoothing factor slitfunction (the default is 0.1, small smoothing)
+    yrange : array[2]
+        number of pixels below and above the central line that have been cut out
+    maxiter : int, optional
+        maximumim number of iterations, by default 20
+    gain : float, optional
+        gain of the image, by default 1
+
+    Returns
+    -------
+    sp, sl, model, unc
+        spectrum, slitfunction, model, spectrum uncertainties
+    """
+
+    # Convert datatypes to expected values
+    lambda_sf = float(lambda_sf)
+    lambda_sp = float(lambda_sp)
+    osample = int(osample)
+    maxiter = int(maxiter)
+    img = np.asanyarray(img, dtype=c_double)
+    ycen = np.asarray(ycen, dtype=c_double)
+    yrange = np.asarray(yrange, dtype=int)
+
+    assert img.ndim == 2, "Image must be 2 dimensional"
+    assert ycen.ndim == 1, "Ycen must be 1 dimensional"
+    assert maxiter > 0, "Maximum iterations must be positive"
+
+    if np.isscalar(tilt):
+        tilt = np.full(img.shape[1], tilt, dtype=c_double)
+    else:
+        tilt = np.asarray(tilt, dtype=c_double)
+    if np.isscalar(shear):
+        shear = np.full(img.shape[1], shear, dtype=c_double)
+    else:
+        shear = np.asarray(shear, dtype=c_double)
+
+    assert img.shape[1] == ycen.size, (
+        f"Image and Ycen shapes are incompatible, got {img.shape} and {ycen.shape}"
+    )
+    assert img.shape[1] == tilt.size, (
+        f"Image and Tilt shapes are incompatible, got {img.shape} and {tilt.shape}"
+    )
+    assert img.shape[1] == shear.size, (
+        f"Image and Shear shapes are incompatible, got {img.shape} and {shear.shape}"
+    )
+
+    assert osample > 0, f"Oversample rate must be positive, but got {osample}"
+    assert lambda_sf >= 0, (
+        f"Slitfunction smoothing must be positive, but got {lambda_sf}"
+    )
+    assert lambda_sp >= 0, f"Spectrum smoothing must be positive, but got {lambda_sp}"
+
+    # assert np.ma.all(np.isfinite(img)), "All values in the image must be finite"
+    assert np.all(np.isfinite(ycen)), "All values in ycen must be finite"
+    assert np.all(np.isfinite(tilt)), "All values in tilt must be finite"
+    assert np.all(np.isfinite(shear)), "All values in shear must be finite"
+
+    assert yrange.ndim == 1, "Yrange must be 1 dimensional"
+    assert yrange.size == 2, "Yrange must have 2 elements"
+    assert yrange[0] + yrange[1] + 1 == img.shape[0], (
+        "Yrange must cover the whole image"
+    )
+    assert yrange[0] >= 0, "Yrange must be positive"
+    assert yrange[1] >= 0, "Yrange must be positive"
+
+    # Retrieve some derived values
+    nrows, ncols = img.shape
+    ny = osample * (nrows + 1) + 1
+
+    ycen_offset = ycen.astype(c_int)
+    ycen_int = ycen - ycen_offset
+    y_lower_lim = int(yrange[0])
+
+    mask = np.ma.getmaskarray(img)
+    img = np.ma.getdata(img)
+    mask2 = ~np.isfinite(img)
+    img[mask2] = 0
+    mask |= ~np.isfinite(img)
+
+    # sp should never be all zero (thats a horrible guess) and leads to all nans
+    # This is a simplified run of the algorithm without oversampling or curvature
+    # But strong smoothing
+    # To remove the most egregious outliers, which would ruin the fit
+    sp = np.sum(img, axis=0)
+    median_filter(sp, 5, output=sp)
+    sl = np.median(img, axis=1)
+    sl /= np.sum(sl)
+
+    model = sl[:, None] * sp[None, :]
+    diff = model - img
+    mask[np.abs(diff) > 10 * diff.std()] = True
+
+    sp = np.sum(img, axis=0)
+
+    mask = np.where(mask, c_int(0), c_int(1))
+    # Determine the shot noise
+    # by converting electrons to photonsm via the gain
+    pix_unc = np.nan_to_num(np.abs(img), copy=False)
+    pix_unc *= gain
+    np.sqrt(pix_unc, out=pix_unc)
+    pix_unc[pix_unc < 1] = 1
+
+    psf_curve = np.zeros((ncols, 3), dtype=c_double)
+    psf_curve[:, 1] = tilt
+    psf_curve[:, 2] = shear
+
+    # Initialize arrays and ensure the correct datatype for C
+    requirements = ["C", "A", "W", "O"]
+    sp = np.require(sp, dtype=c_double, requirements=requirements)
+    mask = np.require(mask, dtype=c_mask, requirements=requirements)
+    img = np.require(img, dtype=c_double, requirements=requirements)
+    pix_unc = np.require(pix_unc, dtype=c_double, requirements=requirements)
+    ycen_int = np.require(ycen_int, dtype=c_double, requirements=requirements)
+    ycen_offset = np.require(ycen_offset, dtype=c_int, requirements=requirements)
+
+    # This memory could be reused between swaths
+    sl = np.zeros(ny, dtype=c_double)
+    model = np.zeros((nrows, ncols), dtype=c_double)
+    unc = np.zeros(ncols, dtype=c_double)
+
+    # Info contains the folowing: sucess, cost, status, iteration, delta_x
+    info = np.zeros(5, dtype=c_double)
+
+    col = np.sum(mask, axis=0) == 0
+    if np.any(col):
+        mask[mask.shape[0] // 2, col] = 1
+    # assert not np.any(np.sum(mask, axis=0) == 0), "At least one mask column is all 0."
+
+    # Call the C function
+    slitfunc_2dlib.slit_func_curved(
+        ffi.cast("int", ncols),
+        ffi.cast("int", nrows),
+        ffi.cast("int", ny),
+        ffi.cast("double *", img.ctypes.data),
+        ffi.cast("double *", pix_unc.ctypes.data),
+        ffi.cast("unsigned char *", mask.ctypes.data),
+        ffi.cast("double *", ycen_int.ctypes.data),
+        ffi.cast("int *", ycen_offset.ctypes.data),
+        ffi.cast("int", y_lower_lim),
+        ffi.cast("int", osample),
+        ffi.cast("double", lambda_sp),
+        ffi.cast("double", lambda_sf),
+        ffi.cast("int", maxiter),
+        ffi.cast("double *", psf_curve.ctypes.data),
+        ffi.cast("double *", sp.ctypes.data),
+        ffi.cast("double *", sl.ctypes.data),
+        ffi.cast("double *", model.ctypes.data),
+        ffi.cast("double *", unc.ctypes.data),
+        ffi.cast("double *", info.ctypes.data),
+    )
+
+    if np.any(np.isnan(sp)):
+        logger.error("NaNs in the spectrum")
+
+    # The decomposition failed
+    if info[0] == 0:
+        status = info[2]
+        if status == 0:
+            msg = "I dont't know what happened"
+        elif status == -1:
+            msg = f"Did not finish convergence after maxiter ({maxiter}) iterations"
+        elif status == -2:
+            msg = "Curvature is larger than the swath. Check the curvature!"
+        else:
+            msg = f"Check the C code, for status = {status}"
+        logger.error(msg)
+        # raise RuntimeError(msg)
+
+    mask = mask == 0
+
+    return sp, sl, model, unc, mask, info
+
+
+# x, y, w
+xi_ref = [("x", c_int), ("y", c_int), ("w", c_double)]
+# x, iy, w
+zeta_ref = [("x", c_int), ("iy", c_int), ("w", c_double)]
+
+
+def xi_zeta_tensors(
+    ncols: int,
+    nrows: int,
+    ycen: np.ndarray,
+    yrange,  # (int, int)
+    osample: int,
+    tilt: np.ndarray,
+    shear: np.ndarray,
+):
+    ncols = int(ncols)
+    nrows = int(nrows)
+    osample = int(osample)
+    ny = osample * (nrows + 1) + 1
+
+    ycen_offset = ycen.astype(c_int)
+    ycen_int = ycen - ycen_offset
+    y_lower_lim = int(yrange[0])
+
+    psf_curve = np.zeros((ncols, 3), dtype=c_double)
+    psf_curve[:, 1] = tilt
+    psf_curve[:, 2] = shear
+
+    requirements = ["C", "A", "W", "O"]
+    ycen_int = np.require(ycen_int, dtype=c_double, requirements=requirements)
+    ycen_offset = np.require(ycen_offset, dtype=c_int, requirements=requirements)
+
+    xi = np.empty((ncols, ny, 4), dtype=xi_ref)
+    zeta = np.empty((ncols, nrows, 3 * (osample + 1)), dtype=zeta_ref)
+    m_zeta = np.empty((ncols, nrows), dtype=c_int)
+
+    slitfunc_2dlib.xi_zeta_tensors(
+        ffi.cast("int", ncols),
+        ffi.cast("int", nrows),
+        ffi.cast("int", ny),
+        ffi.cast("double *", ycen_int.ctypes.data),
+        ffi.cast("int *", ycen_offset.ctypes.data),
+        ffi.cast("int", y_lower_lim),
+        ffi.cast("int", osample),
+        ffi.cast("double *", psf_curve.ctypes.data),
+        ffi.cast("xi_ref *", xi.ctypes.data),
+        ffi.cast("zeta_ref *", zeta.ctypes.data),
+        ffi.cast("int *", m_zeta.ctypes.data),
+    )
+
+    return xi, zeta, m_zeta
+
+
+def create_spectral_model(
+    ncols: int,
+    nrows: int,
+    osample: int,
+    xi: "xi_ref",
+    spec: np.ndarray,
+    slitfunc: np.ndarray,
+):
+    ncols = int(ncols)
+    nrows = int(nrows)
+
+    requirements = ["C", "A", "W", "O"]
+    spec = np.require(spec, dtype=c_double, requirements=requirements)
+    slitfunc = np.require(slitfunc, dtype=c_double, requirements=requirements)
+    xi = np.require(xi, dtype=xi_ref, requirements=requirements)
+
+    img = np.empty((nrows + 1, ncols), dtype=c_double)
+
+    slitfunc_2dlib.create_spectral_model(
+        ffi.cast("int", ncols),
+        ffi.cast("int", nrows),
+        ffi.cast("int", osample),
+        ffi.cast("xi_ref *", xi.ctypes.data),
+        ffi.cast("double *", spec.ctypes.data),
+        ffi.cast("double *", slitfunc.ctypes.data),
+        ffi.cast("double *", img.ctypes.data),
+    )
+    return img

pyreduce/datasets.py

@@ -0,0 +1,238 @@
+"""
+Provides example datasets for the examples
+
+This requires the server to be up and running,
+if data needs to be downloaded
+"""
+
+import logging
+import os
+import tarfile
+from os.path import isfile, join
+
+import wget
+
+logger = logging.getLogger(__name__)
+
+
+def get_data_dir():
+    """Get the default data directory.
+
+    Returns $REDUCE_DATA if set, otherwise ~/REDUCE_DATA
+    """
+    return os.environ.get("REDUCE_DATA", os.path.expanduser("~/REDUCE_DATA"))
+
+
+def load_data_from_server(filename, directory):
+    server = r"http://sme.astro.uu.se/pyreduce/"
+    url = server + filename
+    directory = join(directory, filename)
+    wget.download(url, out=directory)
+
+
+def get_dataset(name, local_dir=None):
+    """Load a dataset
+
+    Note
+    ----
+    This method will not override existing files with the same
+    name, even if they have a different content. Therefore
+    if the files were changed for any reason, the user has to
+    manually delete them from the disk before using this method.
+
+    Parameters
+    ----------
+    name : str
+        Name of the dataset
+    local_dir : str, optional
+        directory to save data at (default: $REDUCE_DATA or ~/REDUCE_DATA)
+
+    Returns
+    -------
+    dataset_dir : str
+        directory where the data was saved
+    """
+
+    if local_dir is None:
+        local_dir = get_data_dir()
+
+    # load data if necessary
+    fname = f"{name}.tar.gz"
+    data_dir = join(local_dir, name)
+    filename = join(data_dir, fname)
+
+    os.makedirs(data_dir, exist_ok=True)
+    if not os.path.isfile(filename):
+        logger.info("Downloading dataset %s", name)
+        logger.info("Data is stored at %s", data_dir)
+        load_data_from_server(fname, data_dir)
+    else:
+        logger.info("Using existing dataset %s", name)
+
+    # Extract the downloaded .tar.gz file
+    with tarfile.open(filename) as file:
+        raw_dir = join(data_dir, "raw")
+        names = [f for f in file if not isfile(join(raw_dir, f.name))]
+        if len(names) != 0:
+            logger.info("Extracting data from tarball")
+            file.extractall(path=raw_dir, members=names)
+
+    return data_dir
+
+
+def UVES(local_dir=None):  # pragma: no cover
+    """Load an example dataset
+    instrument: UVES
+    target: HD132205
+
+    Parameters
+    ----------
+    local_dir : str, optional
+        directory to save data at (default: "./")
+
+    Returns
+    -------
+    dataset_dir : str
+        directory where the data was saved
+    """
+
+    return get_dataset("UVES", local_dir)
+
+
+def HARPS(local_dir=None):  # pragma: no cover
+    """Load an example dataset
+    instrument: HARPS
+    target: HD109200
+
+    Parameters
+    ----------
+    local_dir : str, optional
+        directory to save data at (default: "./")
+
+    Returns
+    -------
+    dataset_dir : str
+        directory where the data was saved
+    """
+
+    return get_dataset("HARPS", local_dir)
+
+
+def LICK_APF(local_dir=None):  # pragma: no cover
+    """Load an example dataset
+    instrument: LICK_APF
+    target: KIC05005618
+
+    Parameters
+    ----------
+    local_dir : str, optional
+        directory to save data at (default: "./")
+
+    Returns
+    -------
+    dataset_dir : str
+        directory where the data was saved
+    """
+
+    return get_dataset("APF", local_dir)
+
+
+def MCDONALD(local_dir=None):  # pragma: no cover
+    """Load an example dataset
+    instrument: JWST_MIRI
+    target: ?
+
+    Data simulated with MIRIsim
+
+    Parameters
+    ----------
+    local_dir : str, optional
+        directory to save data at (default: "./")
+
+    Returns
+    -------
+    dataset_dir : str
+        directory where the data was saved
+    """
+
+    return get_dataset("MCDONALD", local_dir)
+
+
+def JWST_MIRI(local_dir=None):  # pragma: no cover
+    """Load an example dataset
+    instrument: JWST_MIRI
+    target: ?
+
+    Data simulated with MIRIsim
+
+    Parameters
+    ----------
+    local_dir : str, optional
+        directory to save data at (default: "./")
+
+    Returns
+    -------
+    dataset_dir : str
+        directory where the data was saved
+    """
+
+    return get_dataset("MIRI", local_dir)
+
+
+def JWST_NIRISS(local_dir=None):  # pragma: no cover
+    """Load an example dataset
+    instrument: JWST_NIRISS
+    target: ?
+
+    Data simulated with awesimsoss
+
+    Parameters
+    ----------
+    local_dir : str, optional
+        directory to save data at (default: "./")
+
+    Returns
+    -------
+    dataset_dir : str
+        directory where the data was saved
+    """
+
+    return get_dataset("NIRISS", local_dir)
+
+
+def KECK_NIRSPEC(local_dir=None):  # pragma: no cover
+    """Load an example dataset
+    instrument: KECK_NIRSPEC
+    target: GJ1214
+
+    Parameters
+    ----------
+    local_dir : str, optional
+        directory to save data at (default: "./")
+
+    Returns
+    -------
+    dataset_dir : str
+        directory where the data was saved
+    """
+
+    return get_dataset("NIRSPEC", local_dir)
+
+
+def XSHOOTER(local_dir=None):  # pragma: no cover
+    """Load an example dataset
+    instrument: XSHOOTER
+    target: Ux-Ori
+
+    Parameters
+    ----------
+    local_dir : str, optional
+        directory to save data at (default: "./")
+
+    Returns
+    -------
+    dataset_dir : str
+        directory where the data was saved
+    """
+
+    return get_dataset("XSHOOTER", local_dir)