solarc-eclipse 0.5.0__py3-none-any.whl

euvst_response/data_processing.py

@@ -0,0 +1,269 @@
+"""
+Data processing functions for atmosphere cubes and spectral resampling.
+"""
+
+from __future__ import annotations
+from pathlib import Path
+import numpy as np
+import astropy.units as u
+import dill
+from ndcube import NDCube
+from astropy.wcs import WCS
+from specutils import Spectrum
+from specutils.manipulation import FluxConservingResampler
+from joblib import Parallel, delayed
+from tqdm import tqdm
+from .utils import tqdm_joblib, distance_to_angle
+
+
+def load_atmosphere(pkl_file: str, metadata_line: str = None) -> NDCube:
+    """
+    Load synthetic atmosphere cube from pickle file.
+    
+    Creates a summed cube from all line cubes in the synthesis results.
+    All line cubes are interpolated onto the wavelength grid of the metadata_line
+    before summing to handle different wavelength grids for different lines.
+    
+    Parameters
+    ----------
+    pkl_file : str
+        Path to the synthesized spectra pickle file.
+    metadata_line : str, optional
+        Name of the line to use for metadata and wavelength grid reference. 
+        If None, uses the first line.
+        
+    Returns
+    -------
+    NDCube
+        Summed cube of all line intensities with proper WCS and metadata.
+        Uses the wavelength grid from the metadata_line.
+    """
+    with open(pkl_file, "rb") as f:
+        tmp = dill.load(f)
+    
+    # Handle new synthesis format
+    if "line_cubes" not in tmp:
+        raise ValueError("File does not contain synthesis results with line_cubes")
+        
+    line_cubes = tmp["line_cubes"]
+    if not line_cubes:
+        raise ValueError("No line cubes found in synthesis results")
+    
+    # Get the line names
+    line_names = list(line_cubes.keys())
+    
+    # Choose metadata source line
+    if metadata_line is None:
+        metadata_line = line_names[0]
+    elif metadata_line not in line_names:
+        raise ValueError(f"Metadata line '{metadata_line}' not found. Available lines: {line_names}")
+    
+    # Use the metadata line's wavelength grid as the reference
+    ref_cube = line_cubes[metadata_line]
+    ref_wavelengths = ref_cube.axis_world_coords(-1)[0]
+    
+    # Get spatial dimensions from reference cube
+    ny, nx, nw = ref_cube.data.shape
+    
+    # Initialize summed data with the reference wavelength grid
+    summed_data = np.zeros((ny, nx, nw))
+    
+    for line_name, cube in tqdm(line_cubes.items(), desc="Summing line cubes", unit="line", leave=False):
+        # Get wavelength grid for this cube
+        cube_wavelengths = cube.axis_world_coords(-1)[0]
+        
+        # Check spatial dimensions match
+        ny_cube, nx_cube, _ = cube.data.shape
+        if ny_cube != ny or nx_cube != nx:
+            raise ValueError(f"Spatial dimensions mismatch for {line_name}: expected ({ny}, {nx}), got ({ny_cube}, {nx_cube})")
+        
+        # Vectorized interpolation for the entire cube
+        # Reshape data to (n_pixels, n_wavelengths) for batch interpolation
+        cube_data_reshaped = cube.data.reshape(-1, len(cube_wavelengths))
+        
+        # Batch interpolation using numpy.interp
+        interpolated = np.array([
+            np.interp(ref_wavelengths.value, cube_wavelengths.value, spectrum, left=0.0, right=0.0)
+            for spectrum in cube_data_reshaped
+        ])
+        
+        # Add to summed data
+        summed_data += interpolated.reshape(ny, nx, len(ref_wavelengths))
+    
+    # Create new metadata combining info from all lines
+    combined_meta = ref_cube.meta.copy()
+    combined_meta.update({
+        "combined_lines": line_names,
+        "n_lines": len(line_names),
+        "metadata_source": metadata_line,
+        "summed_intensity": True
+    })
+    
+    # Create the summed cube using the reference cube's WCS
+    summed_cube = NDCube(
+        summed_data,
+        wcs=ref_cube.wcs,
+        unit=ref_cube.unit,
+        meta=combined_meta
+    )
+    
+    return summed_cube
+
+
+def resample_ndcube_spectral_axis(ndcube, spectral_axis, output_resolution, ncpu=-1):
+    """
+    Resample the spectral axis of an NDCube using FluxConservingResampler.
+
+    Parameters
+    ----------
+    ndcube : NDCube
+        The input NDCube.
+    spectral_axis : int
+        The index of the spectral axis (e.g., 0, 1, or 2).
+    output_resolution : astropy.units.Quantity
+        The desired output spectral resolution (e.g., 0.01 * u.nm).
+    ncpu : int, optional
+        Number of CPU cores to use for parallel processing. Default is -1 (use all cores).
+
+    Returns
+    -------
+    NDCube
+        A new NDCube with the spectral axis resampled.
+    """
+    # Get the world coordinates of the spectral axis
+    spectral_world = ndcube.axis_world_coords(spectral_axis)[0]
+    spectral_world = spectral_world.to(output_resolution.unit)
+
+    # Define new spectral grid
+    new_spec_grid = np.arange(
+        spectral_world.min().value,
+        spectral_world.max().value + output_resolution.value,
+        output_resolution.value
+    ) * output_resolution.unit
+
+    n_spec = len(new_spec_grid)
+
+    # Move spectral axis to last for easier iteration
+    data = np.moveaxis(ndcube.data, spectral_axis, -1)
+    shape = data.shape
+    flat_data = data.reshape(-1, shape[-1])
+
+    resampler = FluxConservingResampler(extrapolation_treatment="zero_fill")
+    resampled = np.zeros((flat_data.shape[0], n_spec))
+
+    def _resample_pixel(i):
+        spec = Spectrum(flux=flat_data[i] * ndcube.unit, spectral_axis=spectral_world)
+        res = resampler(spec, new_spec_grid)
+        return res.flux.value
+
+    with tqdm_joblib(tqdm(total=flat_data.shape[0], desc="Resampling spectral axis", unit="pixel", leave=False)):
+        results = Parallel(n_jobs=ncpu)(
+            delayed(_resample_pixel)(i) for i in range(flat_data.shape[0])
+        )
+    resampled = np.vstack(results)
+
+    # Reshape back to original spatial shape, but with new spectral length
+    new_shape = list(shape[:-1]) + [n_spec]
+    resampled = resampled.reshape(new_shape)
+
+    # Move spectral axis back to original position
+    resampled = np.moveaxis(resampled, -1, spectral_axis)
+
+    # Update WCS for new spectral axis
+    new_wcs = ndcube.wcs.deepcopy()
+
+    wcs_axis = new_wcs.wcs.naxis - 1 - spectral_axis  # Reverse axis order for WCS
+    center_pixel = (n_spec + 1) / 2  # 1-based index (FITS convention)
+    new_wcs.wcs.crpix[wcs_axis] = center_pixel
+    new_wcs.wcs.crval[wcs_axis] = new_spec_grid[int(center_pixel - 1)].to_value(new_wcs.wcs.cunit[wcs_axis])
+    new_wcs.wcs.cdelt[wcs_axis] = (new_spec_grid[1] - new_spec_grid[0]).to_value(new_wcs.wcs.cunit[wcs_axis])
+
+    return NDCube(resampled, wcs=new_wcs, unit=ndcube.unit, meta=ndcube.meta)
+
+
+def reproject_ndcube_heliocentric_to_helioprojective(new_cube_spec, sim, det):
+    """ Reproject an NDCube from heliocentric to helioprojective coordinates. """
+
+    nx, ny, _ = new_cube_spec.shape
+    wcs_hc = new_cube_spec.wcs
+
+    dx = wcs_hc.wcs.cdelt[2] * wcs_hc.wcs.cunit[2]
+    dy = wcs_hc.wcs.cdelt[1] * wcs_hc.wcs.cunit[1]
+    x_angle = distance_to_angle(dx)
+    y_angle = distance_to_angle(dy)
+
+    wcs_hp = WCS(naxis=3)
+    wcs_hp.wcs.ctype = [wcs_hc.wcs.ctype[0], 'HPLT-TAN', 'HPLN-TAN']
+    wcs_hp.wcs.cunit = [wcs_hc.wcs.cunit[0], 'arcsec', 'arcsec']
+    wcs_hp.wcs.crpix = [wcs_hc.wcs.crpix[0],
+                        (ny + 1) / 2,
+                        (nx + 1) / 2]
+    wcs_hp.wcs.crval = [wcs_hc.wcs.crval[0], 0, 0]
+    wcs_hp.wcs.cdelt = [wcs_hc.wcs.cdelt[0], y_angle.to_value(u.arcsec), x_angle.to_value(u.arcsec)]
+    new_cube_spec_hp = NDCube(new_cube_spec.data, wcs=wcs_hp, unit=new_cube_spec.unit, meta=new_cube_spec.meta)
+
+    nx_in, ny_in, nl_in = new_cube_spec_hp.shape
+    fov_x = nx_in * x_angle
+    fov_y = ny_in * y_angle
+    pitch_x = sim.slit_width
+    pitch_y = det.plate_scale_angle
+    nx_out = int(np.floor((fov_x / pitch_x).decompose().value))
+    ny_out = int(np.floor((fov_y / pitch_y).decompose().value))
+    shape_out = [nx_out, ny_out, nl_in]
+
+    crpix_spec = (shape_out[2] + 1) / 2
+    crpix_y = (shape_out[1] + 1) / 2
+    crpix_x = (shape_out[0] + 1) / 2
+
+    wcs_tgt = WCS(naxis=3)
+    wcs_tgt.wcs.ctype = [wcs_hc.wcs.ctype[0], 'HPLT-TAN', 'HPLN-TAN']
+    wcs_tgt.wcs.cunit = [wcs_hc.wcs.cunit[0], 'arcsec', 'arcsec']
+    wcs_tgt.wcs.crpix = [crpix_spec, crpix_y, crpix_x]
+    wcs_tgt.wcs.crval = [wcs_hc.wcs.crval[0], 0, 0]
+    wcs_tgt.wcs.cdelt = [wcs_hc.wcs.cdelt[0],
+                        (det.plate_scale_angle * u.pix).to_value(u.arcsec),
+                        (sim.slit_width).to_value(u.arcsec)]
+
+    new_cube_spec_hp_spat = new_cube_spec_hp.reproject_to(
+        wcs_tgt,
+        shape_out=shape_out,
+        algorithm='interpolation',
+        parallel=True,
+        order='bilinear',
+    ) * new_cube_spec_hp.unit
+
+    return new_cube_spec_hp_spat
+
+
+def rebin_atmosphere(cube_sim, det, sim, use_dask=False):
+    """
+    Rebin synthetic atmosphere cube to instrument resolution and spatial sampling.
+    
+    Parameters
+    ----------
+    cube_sim : NDCube
+        Input synthetic atmosphere cube
+    det : Detector_SWC or Detector_EIS
+        Detector configuration
+    sim : Simulation
+        Simulation configuration
+    use_dask : bool, optional
+        Whether to use Dask for automatic parallelization (default: True)
+        
+    Returns
+    -------
+    NDCube
+        Rebinned cube at instrument resolution
+    """
+    print("  Spectral rebinning to instrument resolution (nx,ny,*nl*)...")
+
+    cube_spec = resample_ndcube_spectral_axis(cube_sim, spectral_axis=2, output_resolution=det.wvl_res*u.pix, ncpu=sim.ncpu)
+
+    print("  Spatially rebinning to plate scale (nx,*ny*,nl) and slit width (*nx*,ny,nl)...")
+    cube_det = reproject_ndcube_heliocentric_to_helioprojective(
+        cube_spec,
+        sim,
+        det
+    )
+
+    return cube_det

euvst_response/fitting.py

@@ -0,0 +1,144 @@
+"""
+Spectral fitting functions for Gaussian line profile analysis.
+"""
+
+from __future__ import annotations
+import warnings
+import numpy as np
+import astropy.units as u
+import astropy.constants as const
+from ndcube import NDCube
+from scipy.optimize import curve_fit, OptimizeWarning
+from joblib import Parallel, delayed
+from tqdm import tqdm
+from .utils import gaussian, tqdm_joblib
+
+
+def _guess_params(wv: np.ndarray, prof: np.ndarray) -> list:
+    """Guess initial parameters for Gaussian fit."""
+    back = prof.min()
+    prof_c = prof - back
+    prof_c[prof_c < 0] = 0
+    peak = prof_c.max()
+    centre = wv[np.nanargmax(prof_c)]
+    if peak == 0:
+        sigma = (wv.max() - wv.min()) / 10
+    else:
+        # Simple FWHM estimate
+        half_max = 0.5 * peak
+        indices = np.where(prof_c >= half_max)[0]
+        if len(indices) > 1:
+            fwhm = wv[indices[-1]] - wv[indices[0]]
+            sigma = fwhm / (2 * np.sqrt(2 * np.log(2)))
+        else:
+            sigma = (wv.max() - wv.min()) / 10
+    return [peak, centre, sigma, back]
+
+
+def _fit_one(wv: np.ndarray, prof: np.ndarray) -> np.ndarray:
+    """Fit single spectrum with Gaussian."""
+    p0 = _guess_params(wv, prof)
+    with warnings.catch_warnings():
+        warnings.simplefilter("ignore", OptimizeWarning)
+        try:
+            popt, _ = curve_fit(gaussian, wv, prof, p0=p0)
+            return popt
+        except:
+            return np.array(p0)
+
+
+def fit_cube_gauss(signal_cube: NDCube, n_jobs: int = -1) -> tuple[np.ndarray, list[u.Unit]]:
+    """
+    Fit a Gaussian to every (slit x wavelength) spectrum.
+
+    Returns a tuple of (data_array, units_list) where:
+    - data_array: shape (n_scan, n_slit, 4) with the parameters
+      [peak, centre, sigma, background] for each spatial pixel (values only)
+    - units_list: list of 4 astropy Unit objects corresponding to the parameters
+    """
+
+    n_scan, n_slit, _ = signal_cube.shape
+    wv = signal_cube.axis_world_coords(2)[0].cgs  # wavelength axis (Quantity)
+
+    def _fit_block(spec_block: np.ndarray) -> np.ndarray:
+        results = np.empty((spec_block.shape[0], 4))
+        for i in range(spec_block.shape[0]):
+            results[i] = _fit_one(wv.value, spec_block[i])
+        return results
+
+    with tqdm_joblib(tqdm(total=n_scan, desc="Fit chunks", leave=False)):
+        results = Parallel(n_jobs=n_jobs)(
+            delayed(_fit_block)(signal_cube.data[i]) for i in range(n_scan)
+        )
+
+    # Stack to (n_scan, n_slit, 4) - values only
+    data_array = np.stack(results, axis=0)
+    
+    # Units for [peak, centre, sigma, background]
+    units_list = [signal_cube.unit, wv.unit, wv.unit, signal_cube.unit]
+
+    return data_array, units_list
+
+
+def velocity_from_fit(fit_arr: u.Quantity | np.ndarray, wl0: u.Quantity, n_jobs: int = -1) -> u.Quantity:
+    """
+    Convert fitted line centres to LOS velocity.
+    Works with either a Quantity array or an object-dtype array whose
+    elements are Quantities. Uses joblib.Parallel for speed.
+    """
+    centres_raw = fit_arr[..., 1]  # (n_scan, n_slit)
+    # Ensure we have a pure Quantity array
+    if isinstance(centres_raw, u.Quantity):
+        centres = centres_raw.to(wl0.unit)
+    else:  # object array of Quantity scalars
+        get_val = np.vectorize(lambda q: q.to_value(wl0.unit))
+        centres = u.Quantity(get_val(centres_raw), wl0.unit)
+
+    n_scan = centres.shape[0]
+
+    def _one_row(i):
+        return ((centres[i] - wl0) / wl0 * const.c).to(u.cm / u.s).value
+
+    with tqdm_joblib(tqdm(total=n_scan, desc="Velocity calc", leave=False)):
+        v_val = np.array(
+            Parallel(n_jobs=n_jobs)(
+                delayed(_one_row)(i) for i in range(n_scan)
+            )
+        )
+
+    v = v_val * (u.cm / u.s)
+    return v
+
+
+def width_from_fit(fit_arr: u.Quantity | np.ndarray, n_jobs: int = -1) -> u.Quantity:
+    """
+    Extract fitted line widths (sigma) from fit results.
+    """
+    widths_raw = fit_arr[..., 2]  # (n_scan, n_slit) - sigma is 3rd parameter
+    # Ensure we have a pure Quantity array
+    if isinstance(widths_raw, u.Quantity):
+        widths = widths_raw
+    else:  # object array of Quantity scalars
+        get_val = np.vectorize(lambda q: q.value)
+        get_unit = widths_raw.flat[0].unit  # Get unit from first element
+        widths = u.Quantity(get_val(widths_raw), get_unit)
+    
+    return widths
+
+
+def analyse(fits_all: u.Quantity | np.ndarray, v_true: u.Quantity, wl0: u.Quantity) -> dict:
+    """
+    Monte-Carlo velocity statistics given pre-computed ground truth.
+    """
+    v_all = velocity_from_fit(fits_all, wl0)
+    w_all = width_from_fit(fits_all)
+    return {
+        "v_mean": v_all.mean(axis=0),
+        "v_std":  v_all.std(axis=0),
+        "v_err":  v_true - v_all.mean(axis=0),
+        "v_samples": v_all,
+        "v_true":    v_true,
+        "w_mean": w_all.mean(axis=0),
+        "w_std":  w_all.std(axis=0),
+        "w_samples": w_all,
+    }