visualastro 0.0.2__py3-none-any.whl

visualastro/io.py

@@ -0,0 +1,251 @@
+from astropy.io import fits
+import matplotlib as mpl
+import matplotlib.pyplot as plt
+import numpy as np
+from tqdm import tqdm
+from .numerical_utils import check_is_array
+from .visual_classes import FitsFile
+
+
+# Fits File I/O Operations
+# ––––––––––––––––––––––––
+def load_fits(filepath, header=True, error=True,
+              print_info=True, transpose=False, dtype=None):
+    '''
+    Load a FITS file and return its data and optional header.
+    Parameters
+    ––––––––––
+    filepath : str
+        Path to the FITS file to load.
+    header : bool, default=True
+        If True, return the FITS header along with the data
+        as a FitsFile object.
+        If False, only the data is returned.
+    error : bool, default=True
+        If True, return the 'ERR' extention of the fits file.
+    print_info : bool, default=True
+        If True, print HDU information using 'hdul.info()'.
+    transpose : bool, default=False
+        If True, transpose the data array before returning.
+    dtype : np.dtype, default=None
+        Data type to convert the FITS data to. If None,
+        determines the dtype from the data.
+    Returns
+    –––––––
+    FitsFile
+        If header is True, returns an object containing:
+        - data: 'np.ndarray' of the FITS data
+        - header: 'astropy.io.fits.Header' if 'header=True', else 'None'
+    data : np.ndarray
+        If header is False, returns just the data component.
+    '''
+    # print fits file info
+    with fits.open(filepath) as hdul:
+        if print_info:
+            hdul.info()
+        # extract data and optionally the header from the file
+        # if header is not requested, return None
+        result = fits.getdata(filepath, header=header)
+        data, fits_header = result if isinstance(result, tuple) else (result, None)
+
+        dt = get_dtype(data, dtype)
+        data = data.astype(dt, copy=False)
+        if transpose:
+            data = data.T
+
+        errors = get_errors(hdul, dt, transpose)
+
+    if header or error:
+        return FitsFile(data, fits_header, errors)
+    else:
+        return data
+
+
+def get_dtype(data, dtype=None, default_dtype=np.float64):
+    '''
+    Returns the dtype from the provided data. Promotes
+    integers to floats if needed.
+    Parameters
+    ––––––––––
+    data : array-like
+        Input array whose dtype will be checked.
+    dtype : data-type, optional, default=None
+        If provided, this dtype is returned directly.
+        If None, returns `data.dtype` if floating or
+        `np.float64` if integer or unsigned.
+    default_dtype : data-type, optional, default=np.float64
+        Float type to use if `data` is integer or unsigned.
+    Returns
+    –––––––
+    dtype : np.dtype
+        NumPy dtype object: user dtype if given, otherwise the array's
+        float dtype or `default_dtype` if array is integer/unsigned.
+    '''
+    # return user dtype if passed in
+    if dtype is not None:
+        return np.dtype(dtype)
+
+    data = check_is_array(data)
+    # by default use data dtype if floating
+    # if unsigned or int use default_dtype
+    if np.issubdtype(data.dtype, np.floating):
+        return np.dtype(data.dtype)
+    else:
+        return np.dtype(default_dtype)
+
+
+def get_errors(hdul, dtype=None, transpose=False):
+    '''
+    Return the error array from an HDUList, falling back to square root
+    of variance if needed.
+    Parameters
+    ––––––––––
+    hdul : astropy.io.fits.HDUList
+        The HDUList object containing FITS extensions to search for errors or variance.
+    dtype : data-type, optional, default=np.float64
+        The desired NumPy dtype of the returned error array.
+    Returns
+    –––––––
+    errors : np.ndarray or None
+        The error array if found, or None if no suitable extension is present.
+    '''
+    errors = None
+    for hdu in hdul[1:]:
+        extname = hdu.header.get('EXTNAME', '').upper()
+        if extname in {'ERR', 'ERROR', 'UNCERT'}:
+            dt = get_dtype(hdu.data, dtype)
+            errors = hdu.data.astype(dt, copy=False)
+            break
+    # fallback to variance if no explicit errors
+    if errors is None:
+        for hdu in hdul[1:]:
+            extname = hdu.header.get('EXTNAME', '').upper()
+            if extname in {'VAR', 'VARIANCE', 'VAR_POISSON', 'VAR_RNOISE'}:
+                dt = get_dtype(hdu.data, dtype)
+                errors = np.sqrt(hdu.data.astype(dtype, copy=False))
+                break
+    if transpose and errors is not None:
+        errors = errors.T
+
+    return errors
+
+
+def write_cube_2_fits(cube, filename, overwrite=False):
+    '''
+    Write a 3D data cube to a series of FITS files.
+    Parameters
+    ––––––––––
+    cube : ndarray (N_frames, N, M)
+        Data cube containing N_frames images of shape (N, M).
+    filename : str
+        Base filename (without extension). Each
+        output file will be saved as "{filename}_i.fits".
+    overwrite : bool, optional, default=False
+        If True, existing files with the same name
+        will be overwritten.
+    Notes
+    –––––
+    Prints a message indicating the number of
+    frames and the base filename.
+    '''
+    N_frames, N, M = cube.shape
+    print(f"Writing {N_frames} fits files to {filename}_i.fits")
+    for i in tqdm(range(N_frames)):
+        output_name = filename + f"_{i}.fits"
+        fits.writeto(output_name, cube[i], overwrite=overwrite)
+
+
+# Figure I/O Operations
+# –––––––––––––––––––––
+def get_kwargs(kwargs, *names, default=None):
+    '''
+    Return the first matching kwarg value from a list of possible names.
+    Parameters
+    ––––––––––
+    kwargs : dict
+            Dictionary of keyword arguments, typically taken from ``**kwargs``.
+    *names : str
+        One or more possible keyword names to search for. The first name found
+        in ``kwargs`` with a non-None value is returned.
+    default : any, optional, default=None
+        Value to return if none of the provided names are found in ``kwargs``.
+        Default is None.
+    Returns
+    –––––––
+    value : any
+        The value of the first matching keyword argument, or `default` if
+        none are found.
+    '''
+    for name in names:
+        if name in kwargs and kwargs[name] is not None:
+            return kwargs[name]
+    return default
+
+
+def save_figure_2_disk(dpi=600, pdf_compression=6, transparent=False, bbox_inches='tight', **kwargs):
+    '''
+    Saves current figure to disk as a
+    eps, pdf, png, or svg, and prompts
+    user for a filename and format.
+    Parameters
+    ––––––––––
+    dpi : float or int, optional, default=600
+        Resolution in dots per inch.
+    pdf_compression : int, optional, default=False
+        'Pdf.compression' value for matplotlib.rcParams.
+        Accepts integers from 0-9, with 0 meaning no
+        compression.
+    transparent : bool, optional, default=False
+        If True, the Axes patches will all be transparent;
+        the Figure patch will also be transparent unless
+        facecolor and/or edgecolor are specified via kwargs.
+    bbox_inches : str or Bbox, default='tight'
+        Bounding box in inches: only the given portion of the
+        figure is saved. If 'tight', try to figure out the
+        tight bbox of the figure.
+
+    **kwargs : dict, optional
+        Additional parameters.
+
+        Supported keyword arguments include:
+
+        - `facecolorcolor` : str, default='auto'
+            The facecolor of the figure. If 'auto',
+            use the current figure facecolor.
+        - `edgecolorcolor` : str, default='auto'
+            The edgecolor of the figure. If 'auto',
+            use the current figure edgecolor.
+    '''
+    # –––– KWARGS ––––
+    facecolor = get_kwargs(kwargs, 'facecolor', 'fc', default='auto')
+    edgecolor = get_kwargs(kwargs, 'edgecolor', 'ec', default='auto')
+    allowed_formats = {'eps', 'pdf', 'png', 'svg'}
+    # prompt user for filename, and extract extension
+    filename = input("Input filename for image (ex: myimage.pdf): ").strip()
+    basename, *extension = filename.rsplit(".", 1)
+    # if extension exists, and is allowed, extract extension from list
+    if extension and extension[0].lower() in allowed_formats:
+        extension = extension[0]
+    # else prompt user to input a valid extension
+    else:
+        extension = ""
+        while extension not in allowed_formats:
+            extension = (
+                input(f"Please choose a format from ({', '.join(allowed_formats)}): ")
+                .strip()
+                .lower()
+            )
+    # construct complete filename
+    filename = f"{basename}.{extension}"
+
+    with plt.rc_context(rc={'pdf.compression': int(pdf_compression)} if extension == 'pdf' else {}):
+        # save figure
+        plt.savefig(
+            fname=filename,
+            format=extension,
+            transparent=transparent,
+            bbox_inches=bbox_inches,
+            facecolor=facecolor,
+            edgecolor=edgecolor,
+            dpi=dpi
+        )

visualastro/numerical_utils.py

@@ -0,0 +1,285 @@
+import warnings
+from astropy.io.fits import Header
+from astropy import units as u
+from astropy.units import Quantity, spectral, Unit, UnitConversionError
+import numpy as np
+from scipy.interpolate import interp1d, CubicSpline
+from spectral_cube import SpectralCube
+from .visual_classes import DataCube, ExtractedSpectrum, FitsFile
+
+
+# Type Checking Arrays and Objects
+# ––––––––––––––––––––––––––––––––
+def check_is_array(data, keep_units=False):
+    '''
+    Ensure array input is np.ndarray.
+    Parameters
+    ––––––––––
+    data : np.ndarray, DataCube, FitsFile, or Quantity
+        Array or DataCube object.
+    keep_inits : bool, optional, default=False
+        If True, keep astropy units attached if present.
+    Returns
+    –––––––
+    data : np.ndarray
+        Array or 'data' component of DataCube.
+    '''
+    if isinstance(data, DataCube):
+        data = data.value
+    elif isinstance(data, FitsFile):
+        data = data.data
+    if isinstance(data, Quantity):
+        if keep_units:
+            return data
+        else:
+            data = data.value
+
+    return np.asarray(data)
+
+
+def check_units_consistency(datas):
+    '''
+    Check that all input objects have the same units and warn if they differ.
+    Additionally ensure that the input is iterable by wrapping in a list.
+    Parameters
+    ----------
+    datas : object or list/tuple of objects
+        Objects to check. Can be Quantity, SpectralCube, DataCube, etc.
+    Returns
+    -------
+    datas : list
+        The input objects as a list.
+    '''
+    datas = datas if isinstance(datas, (list, tuple)) else [datas]
+
+    first_unit = get_units(datas[0])
+    for i, obj in enumerate(datas[1:], start=1):
+        unit = get_units(obj)
+        if unit != first_unit:
+            warnings.warn(
+                f"\nInput at index {i} has unit `{unit}`, which differs from unit `{first_unit}`."
+                f"at index 0."
+            )
+
+    return datas
+
+
+def get_data(obj):
+    '''
+    Extract the underlying data attribute from a DataCube or FitsFile object.
+    Parameters
+    ––––––––––
+    obj : DataCube or FitsFile or np.ndarray
+        The object from which to extract the data. If a raw array is provided,
+        it is returned unchanged.
+    Returns
+    –––––––
+    np.ndarray, or data extension
+        The data attribute contained in the object, or the input array itself
+        if it is not a DataCube or FitsFile.
+    '''
+    if isinstance(obj, DataCube):
+        obj = obj.data
+    elif isinstance(obj, FitsFile):
+        obj = obj.data
+
+    return obj
+
+
+def get_units(obj):
+    '''
+    Extract the unit from an object, if it exists.
+    Parameters
+    ––––––––––
+    obj : Quantity, SpectralCube, FITS-like object, or any
+        The input object from which to extract a unit. This can be:
+        - an astropy.units.Quantity
+        - a SpectralCube
+        - a DataCube or FitsFile
+        - a FITS-like object with a header containing a 'BUNIT' keyword
+        - any other object (returns None if no unit is found)
+    Returns
+    –––––––
+    astropy.units.Unit or None
+        The unit associated with the input object, if it exists.
+        Returns None if the object has no unit or if the unit cannot be parsed.
+    '''
+    # check if object is DataCube or FitsFile
+    data = get_data(obj)
+    # check if unit extension exists
+    if isinstance(data, (DataCube, FitsFile, Quantity, SpectralCube)):
+        return data.unit
+    if isinstance(obj, ExtractedSpectrum):
+        try:
+            return obj.spectrum1d.unit
+        except:
+            try:
+                return obj.flux.unit
+            except:
+                return None
+
+    # try to extract unit from header
+    # use either header extension or obj if obj is a header
+    header = getattr(obj, 'header', obj if isinstance(obj, Header) else None)
+    if isinstance(header, Header) and 'BUNIT' in header:
+        try:
+            return Unit(header['BUNIT'])
+        except Exception:
+            return None
+
+    return None
+
+
+def return_array_values(array):
+    '''
+    Extract the numerical values from an 'astropy.units.Quantity'
+    or return the array as-is.
+    Parameters
+    ––––––––––
+    array : astropy.units.Quantity or array-like
+        The input array. If it is a Quantity, the numerical values are extracted.
+        Otherwise, the input is returned unchanged.
+    Returns
+    –––––––
+    np.ndarray or array-like
+        The numerical values of the array, without units if input was a Quantity,
+        or the original array if it was not a Quantity.
+    '''
+    array = array.value if isinstance(array, Quantity) else array
+
+    return array
+
+
+# Science Operation Functions
+# –––––––––––––––––––––––––––
+def convert_units(quantity, unit):
+    '''
+    Convert an Astropy Quantity to a specified unit, with a fallback if conversion fails.
+    Parameters
+    ––––––––––
+    quantity : astropy.units.Quantity
+        The input quantity to convert.
+    unit : str, astropy.units.Unit, or None
+        The unit to convert to. If None, no conversion is performed.
+    Returns
+    –––––––
+    astropy.units.Quantity
+        The quantity converted to the requested unit if possible; otherwise,
+        the original quantity with its existing unit.
+    Notes
+    –––––
+    - Uses 'spectral()' equivalencies to allow conversions between
+        wavelength, frequency, and velocity units.
+    - If conversion fails, prints a warning and returns the original quantity.
+    '''
+    if unit is None:
+        return quantity
+    try:
+        # convert string unit to Unit if necessary
+        target_unit = Unit(unit) if isinstance(unit, str) else unit
+        return quantity.to(target_unit, equivalencies=spectral())
+    except UnitConversionError:
+        print(
+            f'Could not convert to unit: {unit}.'
+            f'Defaulting to unit: {quantity.unit}.'
+            )
+        return quantity
+
+
+def shift_by_radial_vel(spectral_axis, radial_vel):
+    '''
+    Shift spectral axis to rest frame using a radial velocity.
+    Parameters
+    ––––––––––
+    spectral_axis : astropy.units.Quantity
+        The spectral axis to shift. Can be in frequency or wavelength units.
+    radial_vel : float, astropy.units.Quantity or None
+        Radial velocity in km/s (astropy units are optional). Positive values
+        correspond to a redshift (moving away). If None, no shift is applied.
+    Returns
+    –––––––
+    shifted_axis : astropy.units.Quantity
+        The spectral axis shifted to the rest frame according to the given
+        radial velocity. If the input is in frequency units, the classical
+        Doppler formula for frequency is applied; otherwise, the classical
+        formula for wavelength is applied.
+    '''
+    # speed of light in km/s in vacuum
+    c = 299792.458 # [km/s]
+    if radial_vel is not None:
+        if isinstance(radial_vel, Quantity):
+            radial_vel = radial_vel.to(u.km/u.s).value # type: ignore
+        # if spectral axis in units of frequency
+        if spectral_axis.unit.is_equivalent(u.Unit('Hz')):
+            spectral_axis /= (1 - radial_vel / c)
+        # if spectral axis in units of wavelength
+        else:
+            spectral_axis /= (1 + radial_vel / c)
+
+    return spectral_axis
+
+
+# Numerical Operation Functions
+# –––––––––––––––––––––––––––––
+def interpolate_arrays(xp, yp, x_range, N_samples, method='linear'):
+    '''
+    Interpolate a 1D array over a specified range.
+    Parameters
+    ––––––––––
+    xp : array-like
+        The x-coordinates of the data points.
+    yp : array-like
+        The y-coordinates of the data points.
+    x_range : tuple of float
+        The (min, max) range over which to interpolate.
+    N_samples : int
+        Number of points in the interpolated output.
+    method : str, default='linear'
+        Interpolation method. Options:
+        - 'linear' : linear interpolation
+        - 'cubic' : cubic interpolation using 'interp1d'
+        - 'cubic_spline' : cubic spline interpolation using 'CubicSpline'
+    Returns
+    –––––––
+    x_interp : np.ndarray
+        The evenly spaced x-coordinates over the specified range.
+    y_interp : np.ndarray
+        The interpolated y-values corresponding to 'x_interp'.
+    '''
+    # generate new interpolation samples
+    x_interp = np.linspace(x_range[0], x_range[1], N_samples)
+    # get interpolation method
+    if method == 'cubic_spline':
+        f_interp = CubicSpline(xp, yp)
+    else:
+        # fallback to linear if method is unknown
+        kind = method if method in ['linear', 'cubic'] else 'linear'
+        f_interp = interp1d(xp, yp, kind=kind)
+    # interpolate over new samples
+    y_interp = f_interp(x_interp)
+
+    return x_interp, y_interp
+
+
+def mask_within_range(x, xlim=None):
+    '''
+    Return a boolean mask for values of x within the given limits.
+    Parameters
+    ––––––––––
+    x : array-like
+        Data array (e.g., wavelength or flux values)
+    xlim : tuple or list, optional
+        (xmin, xmax) range. If None, uses the min/max of x.
+    Returns
+    –––––––
+    mask : ndarray of bool
+        True where x is within the limits.
+    '''
+    x = return_array_values(x)
+    xlim = return_array_values(xlim)
+
+    xmin = xlim[0] if xlim is not None else np.nanmin(x)
+    xmax = xlim[1] if xlim is not None else np.nanmax(x)
+    mask = (x > xmin) & (x < xmax)
+
+    return mask

visualastro/photometry.py

@@ -0,0 +1,40 @@
+import numpy as np
+from scipy.ndimage import center_of_mass
+from tqdm import tqdm
+from .numerical_utils import check_is_array
+from .visual_plots import va
+
+def compute_flux(cube, target_pixel_loc, star_radius, sky_radii=None, window_half_width=100, plot=False):
+    cube = check_is_array(cube)
+    # initial target location guess
+    x_pixel, y_pixel = target_pixel_loc
+    # array to store target and sky flux
+    star_flux = np.zeros((cube.shape[0]))
+    sky_flux = np.zeros_like(star_flux)
+    # loop through each image in cube
+    for i in tqdm(range(len(cube))):
+        # extract subimage centered around target
+        xmin, xmax = x_pixel - window_half_width, x_pixel + window_half_width
+        ymin, ymax = y_pixel - window_half_width, y_pixel + window_half_width
+        sub_image = cube[i][xmin:xmax, ymin:ymax]
+        # compute approximate center pixels of target
+        cenx, ceny = center_of_mass(sub_image)
+        # compute distance between target center each pixel in subimage
+        x, y = np.indices(sub_image.shape)
+        distance_from_center = np.sqrt( (x - cenx)**2 + (y - ceny)**2)
+        # mask out pixels outside star aperture
+        star_aperture = distance_from_center < star_radius
+        star_flux[i] = np.nansum(sub_image[star_aperture])
+        if sky_radii is not None:
+            sky_inner_r, sky_outer_r = sky_radii
+            sky_aperture = (distance_from_center < sky_outer_r) & (distance_from_center > sky_inner_r)
+            sky_flux[i] = np.nanmedian(sub_image[sky_aperture])
+            star_flux[i] -= sky_flux[i]
+        if plot:
+            circles = [[cenx, ceny, star_radius]]
+            if sky_radii is not None:
+                for r in sky_radii:
+                    circles.append([cenx, ceny, r])
+            va.imshow(sub_image, circles=circles)
+
+    return star_flux, sky_flux