cubespa 0.0.9__tar.gz

cubespa-0.0.9/PKG-INFO

@@ -0,0 +1,20 @@
+Metadata-Version: 2.1
+Name: cubespa
+Version: 0.0.9
+Summary: Cube SPectral Analysis
+Home-page: https://github.com/HSouch/cubespa
+Author: Harrison Souchereau
+Author-email: harrison.souchereau@yale.edu
+License: BSD-3-Clause
+Keywords: galaxies datacubes spectra
+Classifier: Development Status :: 4 - Beta
+Classifier: Intended Audience :: Developers
+Classifier: Topic :: Software Development :: Build Tools
+Classifier: Programming Language :: Python :: 3.6
+Classifier: Programming Language :: Python :: 3.7
+Classifier: Programming Language :: Python :: 3.9
+Description-Content-Type: text/markdown
+
+# CUBE SPectral Analysis
+
+Fully-featured, aesthetically pleasing 2D and 3D analysis of astronomical datacubes. 

cubespa-0.0.9/README.md

@@ -0,0 +1,3 @@
+# CUBE SPectral Analysis
+
+Fully-featured, aesthetically pleasing 2D and 3D analysis of astronomical datacubes. 

cubespa-0.0.9/cubespa/__init__.py

@@ -0,0 +1,15 @@
+
+from .analysis import *         # Plots for flux analysis
+from .modeling import *
+
+from .cubespa import *          # Module with all CubeSPA object functions
+from .data import *             # Data handling modules
+from .moment_maps import *      # Module for generating moment maps using masking
+from .overlays import *         # Module for reprojection and creating overlays
+from .pvds import *             # Module for position velocity diagrams (PVDs)
+from .spectra import *          # Module for generating and analyzing spectra
+from .utils import *            # All miscellaneous utility modules
+
+
+__version__ = "0.0.9"
+

cubespa-0.0.9/cubespa/analysis/__init__.py

@@ -0,0 +1,2 @@
+
+from .flux_analysis import *

cubespa-0.0.9/cubespa/analysis/flux_analysis.py

@@ -0,0 +1,145 @@
+import numpy as np
+from .. import utils
+
+from matplotlib import pyplot as plt
+
+from photutils.aperture import EllipticalAnnulus
+
+
+def quadrants(q1_angle, offset = 90, clockwise=False):
+    mult = -1 if clockwise else 1
+    return np.array([(q1_angle + i * 90 * mult + offset) for i in range(4)])
+
+
+def halfline_points(cent, length, angle):
+
+    p1 = (cent[0] , cent[0] + length * np.cos(angle))
+    p2 = (cent[1], cent[1] + length * np.sin(angle))
+    
+    return (p1, p2)
+
+
+def line_points(cent, length, angle):
+    p1 = (cent[0] - length * np.cos(angle), cent[0] + length * np.cos(angle))
+    p2 = (cent[1] - length * np.sin(angle), cent[1] + length * np.sin(angle))
+
+    
+    return (p1, p2)
+
+
+def theta_list(theta_0, clockwise=False, n_samples=90, plot=False, **kwargs):
+    outname = utils.check_kwarg("outname", None, kwargs)
+
+    theta_max = theta_0 + 2*np.pi
+
+    if clockwise:
+        _ = theta_0
+        theta_0, theta_max = theta_max, _
+    
+    print(theta_0, theta_max)
+
+    angles = np.linspace(theta_0, theta_max, n_samples) % (2 * np.pi)
+
+    if plot:
+        thetas_norm = np.linspace(0, 2*np.pi, n_samples)
+        xs, ys = np.cos(angles), np.sin(angles)
+        plt.scatter(xs, ys, c=np.rad2deg(thetas_norm))
+        plt.colorbar()
+        plt.tight_layout()
+        if outname is not None:
+            plt.savefig(outname, dpi=100)
+            plt.close()
+        else:
+            plt.show()
+    
+    return angles
+
+
+def theta_range(shape, x0, y0, theta, dtheta):
+    ys, xs = np.mgrid[:shape[0], :shape[1]].astype(float)
+    xs -= x0
+    ys -= y0
+    
+    thetas = (np.arctan2(ys, xs) ) % (2 * np.pi)
+
+    theta_mask = np.ones(thetas.shape)
+
+    theta_min = (theta - dtheta / 2) % (2 * np.pi)
+    theta_max = (theta + dtheta / 2) % (2 * np.pi)
+    
+    # Adjust for cases where theta_max < theta_min due to wrapping around 2pi
+    if theta_max < theta_min:
+        # In this case, we need to handle the wrap around.
+        theta_mask = np.logical_or(thetas >= theta_min, thetas <= theta_max)
+    else:
+        # Regular case where the range does not wrap around
+        theta_mask = np.logical_and(thetas >= theta_min, thetas <= theta_max)
+
+    return theta_mask.astype(int)
+
+
+def process_azimuth(data, radius, pa, eps, thetas, stellar_x=0, stellar_y=0, annulus_width=50, trim_nans=False):
+
+    
+    ell_annulus = EllipticalAnnulus((stellar_x, stellar_y), radius, radius + annulus_width, 
+                                    (radius + annulus_width) * (1-eps), theta = pa)
+    mask = ell_annulus.to_mask(method='center').to_image(shape=data.shape)
+    
+    data_init_mask = np.copy(data)
+    data_init_mask[mask == 0] = np.nan
+
+    fluxes, flux_stds = [], []
+
+    for i, theta in enumerate(thetas[:]):
+    
+        theta_mask = theta_range(data.shape, stellar_x, stellar_y, theta, 5 * np.pi / 180)
+    
+        this_data = np.copy(data_init_mask)
+        this_data[theta_mask == 0] = np.nan
+        try:
+            med, std = np.nanmedian(this_data), np.nanstd(this_data)
+        except RuntimeWarning:
+            med, std = np.nan, np.nan
+            
+        fluxes.append(med)
+        flux_stds.append(std)
+
+    if trim_nans:
+        fluxes = np.array(fluxes)
+        fluxes[np.isnan(fluxes)] = 0
+
+        flux_stds = np.array(flux_stds)
+        flux_stds[np.isnan(flux_stds)] = 0
+    
+    return (np.array(fluxes), np.array(flux_stds), ell_annulus)
+
+
+def segment_sum(cubespa_obj, data, radius, annulus_width, theta_0, plot=False, clockwise=False):
+    stellar_x, stellar_y = cubespa_obj.center
+    
+    eps, pa = cubespa_obj.eps, cubespa_obj.position_angle
+
+    # print(stellar_x, stellar_y, eps, pa)
+    ell_annulus = EllipticalAnnulus((stellar_x, stellar_y), 
+                                    radius, radius + annulus_width, 
+                                    (radius + annulus_width) * (1-eps), 
+                                    theta = pa)
+    mask = ell_annulus.to_mask(method='center').to_image(shape=data.shape)
+    
+    data_init_mask = np.copy(data)
+    data_init_mask[mask == 0] = np.nan
+
+    theta_mask = theta_range(data.shape, stellar_x, stellar_y, np.deg2rad(theta_0) + np.pi/4, np.pi / 2)
+
+    data_init_mask[theta_mask == 0] = np.nan
+    pixel_sum = np.nansum(data_init_mask)
+
+    if plot:
+        plt.imshow(data, origin="lower", cmap="Greys")
+        plt.imshow(data_init_mask, cmap="viridis", origin="lower", alpha=0.7)
+        plt.imshow(theta_mask, alpha=0.2, cmap="Greys", origin="lower")
+        ell_annulus.plot(color="black")
+        plt.title(f'{radius:.1f} {theta_0}\n{pixel_sum:.2f}', fontsize=10)
+        plt.show()
+
+    return pixel_sum

cubespa-0.0.9/cubespa/cubespa.py

@@ -0,0 +1,184 @@
+from . import data, utils, spectra
+
+from . import plotting
+
+import os, logging
+
+import numpy as np
+
+
+class CubeSPA:
+    """ Base input class for a CubeSPA object
+    """
+
+    def __init__(self, cube, data_index=0, 
+                 psf = None,
+                 mom_maps=None, additional_maps = [],
+                 center = None, position_angle = 0, eps=0,
+                 vsys = 0,
+                 limits = None, 
+                 plot_dir = None,
+                 **kwargs) -> None:
+        """ CubeSPA object
+
+        Args:
+            cube (ndarray or str): Input cube or filename that points to cube.
+            data_index (int, optional): If loading cube, index of cube data. Defaults to 0.
+            mom_maps (str, optional): maskmoment-formatted moment map filename prefix. Defaults to None.
+            additional_maps (list, optional): Additional data.DataSet objects. Defaults to [].
+            center (tuple, optional): Central region, typically defined by stellar light isophotal
+                center. Defaults to None.
+            position_angle (float, optional): Position angle of disk. Defaults to None.
+            eps (float, optional): Ellipticity (1 - b/a) of disk. Defaults to None.
+            limits (array or str, optional): Bounding box containing relevant data. Defaults to None.
+                If "auto", will try to automatically generate from moment maps.
+            plot_dir (str, optional): Directory to place plots. Defaults to None.
+        """
+        
+        self.cube = data.handle_data(cube, handler=data.load_data, data_index=data_index)
+        
+        if self.cube is not None:
+            self.cube_noise_level, self.cube_rms = utils.estimate_rms(self.cube.data, 
+                                                                    utils.check_kwarg("cmin", 5, kwargs), 
+                                                                    utils.check_kwarg("cmax", 5, kwargs))
+
+        self.psf = data.handle_data(psf, handler=data.load_data, data_index=data_index) if psf is not None else psf
+
+        # TODO implement better handling of cx and cy. Right now assuming a 3D cube with Stokes. Can be better!
+        self.center = center
+        if self.center is not None:
+            center_x, center_y, center_type = center
+            if center_type == "radec":
+                cx, cy = self.cube.wcs.wcs_world2pix(center_x, center_y, 0, 0, 0)[:2]
+                self.center = (float(cx), float(cy))
+            else:
+                self.center = (center_x, center_y)
+
+        # Position angle and ellipticty (1 - b/a) of the stellar disk, for flux analysis
+        self.position_angle = np.deg2rad(position_angle)
+        self.eps = eps
+
+        self.velocities = self.velocities_from_wcs(vsys=vsys)
+        self.vsys = vsys
+        
+        # TODO Add this to a separate function with error handling
+        try:
+            self.beam = [self.cube.header["BMIN"], self.cube.header["BMAJ"], self.cube.header["BPA"]]
+            self.beam_pix = self.get_beam_pix()
+
+            self.beam_area = self.get_beam_area()
+            self.beam_area_arcsec = self.get_beam_area(in_pixels=False)
+        except KeyError: 
+            logging.warn(f"Unable to get beam for {cube}")
+            self.beam = self.beam_pix = self.beam_area_arcsec = None
+
+        try:
+            self.mom_maps = data.handle_data(mom_maps, handler=data.load_moment_maps, data_index=data_index)
+            self.additional_maps = additional_maps
+        except Exception as e:
+            logging.warn(f'Failed to load moment maps for {cube}\n{e}')
+
+        if limits == "auto":
+            if "padding" in kwargs.keys():
+                padding = kwargs["padding"]
+            else:
+                padding = 10
+
+            self.limits = utils.bounds_from_moment_map(self.mom_maps.mom0.data, padding = padding)
+        
+        if plot_dir is not None:
+            self.plot_dir = plot_dir
+            self.load_dir(plot_dir)
+    
+
+    def load_dir(path):
+        if not os.path.isdir(path):
+            os.makedirs(path, exist_ok=True)
+        else:
+            pass
+
+    
+    def velocities_from_wcs(self, vsys=0):
+        vmin = self.cube.wcs.wcs.crval[2]
+        vdelt = self.cube.wcs.wcs.cdelt[2]
+
+        return (np.array([vmin + i * vdelt for i in range(len(self.cube.data))]) / 1000) - vsys
+    
+
+    def get_beam_area(self, in_pixels=True):
+        try:
+            header, w = self.cube.header, self.cube.wcs
+            
+            delt = np.mean(np.abs(w.wcs.cdelt[:2]))
+
+            if in_pixels:
+                bmin, bmaj = header["BMIN"] / delt , header["BMAJ"] / delt
+            else:
+                bmin, bmaj = header["BMIN"] * 3600 , header["BMAJ"] * 3600
+
+            area = np.pi * bmaj * bmin / (4 * np.log(2))
+        except Exception as e:
+            print("Failed to get beam area:", e)
+            area = 1
+            
+        return area
+
+    def get_beam_pix(self):
+        try:
+            header, w = self.cube.header, self.cube.wcs
+            
+            delt = np.mean(np.abs(w.wcs.cdelt[:2]))
+            bmaj, bmin, bpa = self.beam
+            bmaj /= delt
+            bmin /= delt
+            return bmaj, bmin, bpa
+        except:
+            return 1, 1
+
+    def get_beam_coords(self, x0=0, y0=0):
+        try:
+            header, w = self.cube.header, self.cube.wcs
+            delt = np.mean(np.abs(w.wcs.cdelt[:2]))
+
+            bmin, bmaj, bpa = header["BMIN"] / delt , header["BMAJ"] / delt, np.deg2rad(header["BPA"] + 90)
+            return utils.ellipse_coords(x0, y0, bmaj / 2, bmin / 2, bpa)
+        
+        except Exception as e:
+            print("Failed to get beam coordinates", e)
+            
+    # UTILITY FUNCTIONS
+    def plot_moment_maps(self, use_limits=True, **kwargs):
+        outname = utils.check_kwarg("outname", None, kwargs)
+
+        plotting.moment_map_plot(self, use_limits=use_limits, outname=outname, kwargs=kwargs)
+
+    
+    def create_spectra(self, position, size, return_products=False, plot=False):
+        aper = spectra.create_aperture(self, position, size)
+        spectrum = spectra.get_spectra(self.cube.data, aper)
+        if plot:
+            plotting.spectra_plot(self, aper, spectrum)
+        if return_products:
+            return aper, spectrum
+
+
+class CubeComparison:
+
+    def __init__(self, cube1: CubeSPA, cube2: CubeSPA) -> None:
+        self.cube1 = cube1
+        self.cube2 = cube2
+
+
+    def create_spectra(self, position, size, return_products=False, plot=False):
+        aper = spectra.create_aperture(self, position, size)
+        spectrum = spectra.get_spectra(self.cube1.data, aper)
+
+        if plot:
+            plotting.spectra_plot(self, aper, spectrum)
+        if return_products:
+            return aper, spectrum
+
+    
+    def readout(self):
+        print(f'Cube 1: {self.cube1.cube.data.shape}    Cube 2: {self.cube2.cube.data.shape}')
+

cubespa-0.0.9/cubespa/data.py

@@ -0,0 +1,173 @@
+import numpy as np
+import logging
+import copy
+
+from . import utils
+
+from astropy.wcs import WCS
+from astropy.convolution import Tophat2DKernel
+from astropy.io import fits
+
+from scipy.signal import convolve2d
+
+from . import plotting
+
+
+class DataSet:
+    def __init__(self, data=None, wcs=None, header=None, label=None, dtype=None):
+        self.data = data
+        self.wcs = wcs
+        self.header = header
+        self.label = label
+
+        self.dtype = dtype
+
+
+class RGBImage(DataSet):
+    def __init__(self, data=None, wcs=None, header=None, label=None):
+        super().__init__(data, wcs, header, label)
+    
+
+
+class MomentMaps:
+    def __init__(self, mom0=None, mom1=None, mom2=None, data_index=0):
+        if type(mom0) == str:
+            self.mom0 = load_data(mom0, data_index=data_index)
+        if type(mom1) == str:
+            self.mom1 = load_data(mom1, data_index=data_index)
+        if type(mom2) == str:
+            self.mom2 = load_data(mom2, data_index=data_index)
+
+
+def load_data(filename, data_index=0, rgb_index = None, label=None):
+    try:
+        with fits.open(filename) as HDUList:
+            hdu = HDUList[data_index]
+            
+            if rgb_index is None:
+                data = hdu.data
+                if len(data.shape) == 4:        # Drop Stokes parameter (could remove this later maybe)
+                    data = data[0]
+                wcs = WCS(hdu.header)
+            else:
+                data = hdu.data[rgb_index]
+                wcs = WCS(hdu.header).dropaxis(2)
+
+            out = DataSet(data=data, wcs=wcs, header=hdu.header, label=label)
+            return out
+    except Exception as e:
+        print(f"Error loading data: {e}")
+        return None
+
+
+def load_moment_maps(topdir, data_index=0):
+    try:
+        mom_0_fn = topdir + ".mom0.fits.gz"
+        mom_1_fn = topdir + ".mom1.fits.gz"
+        mom_2_fn = topdir + ".mom2.fits.gz"
+
+        moment_maps = MomentMaps(mom0=mom_0_fn, mom1=mom_1_fn, mom2=mom_2_fn, data_index=data_index)
+        return moment_maps
+    except Exception as e:
+        print(f"Error loading Moment Maps: {e}")
+        return None
+
+
+
+def handle_data(data, handler, data_index=0):
+    """ Handle incoming data for the CubeSPA object. 
+    If not a string, CubeSPA will return the data enclosed in a dataset without additional info.
+    Args:
+        data (str or cubespa.DataSet): Incoming data. If str, CubeSPA will automatically load it in.
+        handler (method): Method to handle data, either set to load_data or load_moment_maps
+        data_index (int, optional): Index to find data in. Defaults to 0.
+    Returns:
+        utils.DataSet : The output DataSet object. 
+    """
+    if type(data) == str:
+        return handler(data, data_index=data_index)
+    elif type(data) == np.ndarray:
+        logging.warn("Naked data loaded. No WCS or header included!")
+        return utils.DataSet(data=data)
+
+
+def gen_cutout(data, limits):
+    """ Generates a cutout from an image based on a set of limits.
+
+    Args:
+        data (ndarray): Input image data array
+        limits (arr): Limits, assuming [xmin, xmax, ymin, ymax] format
+
+    Returns:
+        _type_: _description_
+    """
+    
+    xmin, xmax, ymin, ymax = limits
+
+    return data[ymin:ymax, xmin:xmax]
+
+
+def gen_cutout(cubespa_obj, cent, size, show_bbox=False):
+
+    # TODO: change WCS coordinates to match new cutout location
+
+    cutout = copy.deepcopy(cubespa_obj)
+
+    if type(size) is int:
+        size = (size, size)
+    
+    xmin, xmax = cent[1] - size[0], cent[1] + size[0]
+    ymin, ymax = cent[0] - size[1], cent[0] + size[1]
+
+    if show_bbox:
+        plotting.plot_bbox(cubespa_obj, [ymin, ymax, xmin, xmax])
+
+    new_cube = cubespa_obj.cube.data[:, xmin:xmax, ymin:ymax]
+    if cubespa_obj.mom_maps is not None:
+        new_mom0 = cubespa_obj.mom_maps.mom0.data[xmin:xmax, ymin:ymax]
+        new_mom1 = cubespa_obj.mom_maps.mom1.data[xmin:xmax, ymin:ymax]
+        new_mom2 = cubespa_obj.mom_maps.mom2.data[xmin:xmax, ymin:ymax]
+
+    for i, additional in enumerate(cubespa_obj.additional_maps):
+        cutout.additional_maps[i].data = additional.data[xmin:xmax, ymin:ymax]
+
+
+    cutout.cube.data = new_cube
+
+    cutout.limits = [ymin, ymax, xmin, xmax]
+
+    cutout.mom_maps.mom0.data = new_mom0
+    cutout.mom_maps.mom1.data = new_mom1
+    cutout.mom_maps.mom2.data = new_mom2
+
+    return cutout
+
+
+def gen_bg_mask(a, thresh_min=None, dilate=None, dilate_thresh = 0.7):
+    """Generate mask from a map, assuming emission is either nan or below some threshold value.
+    Mask can be dilated using the dilate parameter, and be further adjusted with dilate_thresh
+
+    Args:
+        a (ndarray): Input data array
+        thresh_min (float, optional): Minimum threshold to keep unmasked. Defaults to None.
+        dilate (int, optional): Side of tophat kernel to dilate mask. Defaults to None.
+        dilate_thresh (float, optional): Threshold (between 0 and 1) to include "edge areas" in the dilation mask. 
+        A higher value corresponds to a stronger dilation. Defaults to 0.7.
+
+    Returns:
+        ndarray: Boolean mask with same shape as input array.
+    """
+
+    mask = np.zeros(a.shape)
+    mask[np.isnan(a)] = 1
+
+    if thresh_min is not None:
+        mask[a < thresh_min] = 1
+
+    if dilate is not None:
+        kernel = Tophat2DKernel(dilate)
+        mask = convolve2d(mask, kernel, mode="same")
+        mask[mask < dilate_thresh] = 0
+        mask[mask >= dilate_thresh] = 1
+
+    return mask.astype(bool)

cubespa-0.0.9/cubespa/modeling/__init__.py

@@ -0,0 +1,3 @@
+
+from .velocity_map import *
+

cubespa-0.0.9/cubespa/modeling/velocity_map.py

@@ -0,0 +1,102 @@
+import numpy as np
+from .. import utils
+
+from astropy.io import fits
+
+
+def brandt_curve(rs, kwargs):
+    Rmax = utils.check_kwarg("rmax", 50, kwargs)
+    Vmax = utils.check_kwarg("vmax", 50, kwargs)
+    n = utils.check_kwarg("n", 1, kwargs)
+
+    R_norm = rs/Rmax
+    
+    return Vmax * R_norm * (1/3 + (2/3)*(R_norm ** n)) ** (-3 / (2 * n))
+
+
+def df_to_of(x, y, alpha, i): 
+    """ Convert the disk frame to the observer's frame.
+
+    Args:
+        x (_type_): _description_
+        y (_type_): _description_
+        alpha (_type_): _description_
+        i (_type_): _description_
+
+    Returns:
+        _type_: _description_
+    """
+    X = x * np.cos(alpha) - y * np.sin(alpha) * np.cos(i)
+    Y = x * np.sin(alpha) + y * np.cos(alpha) * np.cos(i)
+    return X, Y
+
+def of_to_df(X, Y, alpha, i):
+    """ Convert the observer's frame to the disk frame.
+
+    Args:
+        X (_type_): _description_
+        Y (_type_): _description_
+        alpha (_type_): _description_
+        i (_type_): _description_
+
+    Returns:
+        _type_: _description_
+    """
+    x = X * np.cos(alpha) + Y * np.sin(alpha)
+    y = -X * np.sin(alpha) * np.cos(i) + Y * np.cos(alpha) * np.cos(i)
+    
+    return x, y
+
+
+class VelocityModel:
+    def __init__(self, shape, x0=None, y0=None, pa=None, inc=None, vsys=0, rmax=50):
+        
+        if type(shape) == int:
+            shape = (shape, shape)
+
+        self.shape = shape
+
+        self.x0 = x0 if x0 is not None else int(shape[1] / 2)
+        self.y0 = y0 if y0 is not None else int(shape[0] / 2)
+        self.pa = np.deg2rad(pa) if pa is not None else 0
+        self.inc = np.deg2rad(inc) if inc is not None else np.deg2rad(30)
+        self.vsys = vsys if vsys is not None else 0
+        self.rmax = rmax
+
+        self.vr = None
+
+
+
+    def gen_model(self, velocity_model, fill_value=np.nan, **vmod_kwargs):
+        xs, ys = np.arange(0, self.shape[0]).astype(float), np.arange(0, self.shape[1]).astype(float)
+        Y, X = np.meshgrid(ys, xs)
+
+        X -= self.x0
+        Y -= self.y0
+    
+        x, y = of_to_df(X, Y, self.pa, self.inc)
+    
+        r = np.sqrt(x ** 2 + y ** 2)
+
+        vcirc = velocity_model(r, vmod_kwargs)
+        theta = np.arctan2(x, y)
+        
+        vr = self.vsys + vcirc * np.sin(self.inc) * np.cos(theta)
+        vr[r > self.rmax * 2] = fill_value
+        self.vr = vr
+        return self.vr
+    
+    def write_model(self, outname, to_fits=False, overwrite=True):
+        """ Save the velocity model to a FITS file or Pickle file, depending on user preferences
+
+        Args:
+            filename (str): Output filename
+            to_fits (bool, optional): Whether to save to FITS format. Defaults to False.
+            overwrite (bool, optional): If saving to fits format, whether to overwrite output if already saved. Defaults to True.
+        """
+        if to_fits:
+            hdul = fits.HDUList()
+            hdul.append(fits.ImageHDU(data=self.vr))
+            hdul.writeto(outname, overwrite=overwrite)
+        else:
+            utils.save_pickle(self.vr, outname)