SpecProc 1.0.0__py3-none-any.whl

specproc/core/de_blazing.py

@@ -0,0 +1,210 @@
+"""
+De-blazing correction stage.
+
+Corrects for the blaze function of the spectrograph grating
+to flatten the spectral response.
+"""
+# 根据平场闪耀轮廓校正提取光谱响应。
+
+import numpy as np
+import logging
+from pathlib import Path
+from typing import Optional
+from specproc.core.data_structures import SpectraSet, FlatField
+from specproc.plotting.spectra_plotter import plot_spectrum_to_file, plot_spectra_to_pdf
+
+logger = logging.getLogger(__name__)
+
+
+# 用平场闪耀轮廓校正光谱。
+def apply_de_blazing(spectra_set: SpectraSet, flat_field: FlatField) -> SpectraSet:
+    """
+    Apply de-blazing correction to extracted spectra.
+
+    Divides each spectrum by its corresponding blaze profile
+    to remove wavelength-dependent efficiency variations.
+
+    Args:
+        spectra_set: Extracted spectra to correct
+        flat_field: Flat field data containing blaze profiles
+
+    Returns:
+        De-blazed spectra set
+    """
+    if flat_field.blaze_profiles is None or len(flat_field.blaze_profiles) == 0:
+        logger.warning("No blaze profiles available, skipping de-blazing")
+        return spectra_set
+
+    corrected_spectra = SpectraSet()
+
+    for spectrum in spectra_set.spectra.values():
+        aperture_id = spectrum.aperture
+
+        if aperture_id in flat_field.blaze_profiles:
+            blaze_profile = flat_field.blaze_profiles[aperture_id]
+
+            # Ensure blaze profile matches spectrum length
+            if len(blaze_profile) == len(spectrum.flux):
+                blaze_peak = np.nanmax(blaze_profile)
+                if not np.isfinite(blaze_peak) or blaze_peak <= 0:
+                    blaze_peak = 1.0
+                
+                blaze_threshold = 0.02 * blaze_peak
+
+                # --- FIX: Contiguous valid region ---
+                # Only keep the contiguous region around the peak that is above the threshold.
+                # This strictly ignores any artificial "up-swings" or polynomial ringing at the extreme edges.
+                bad_blaze_mask = np.ones_like(blaze_profile, dtype=bool)
+                peak_idx = int(np.nanargmax(blaze_profile))
+                
+                # Walk left from the peak
+                left_bound = peak_idx
+                while left_bound > 0 and blaze_profile[left_bound - 1] >= blaze_threshold:
+                    left_bound -= 1
+                    
+                # Walk right from the peak
+                right_bound = peak_idx
+                while right_bound < len(blaze_profile) - 1 and blaze_profile[right_bound + 1] >= blaze_threshold:
+                    right_bound += 1
+                    
+                # The contiguous region from left_bound to right_bound is good.
+                bad_blaze_mask[left_bound:right_bound + 1] = False
+
+                # Avoid division by zero or very small numbers by setting them to NaN
+                blaze_for_division = np.where(bad_blaze_mask, np.nan, blaze_profile)
+
+                # Apply de-blazing, ignoring division warnings for NaNs
+                with np.errstate(divide='ignore', invalid='ignore'):
+                    corrected_flux = spectrum.flux / blaze_for_division
+                    corrected_error = spectrum.error / blaze_for_division if spectrum.error is not None else None
+
+                # In regions of low blaze, the result is unreliable. Set flux and error to NaN.
+                corrected_flux[bad_blaze_mask] = np.nan
+                if corrected_error is not None:
+                    corrected_error[bad_blaze_mask] = np.nan
+                
+                # Clean up any other non-finite values that might have resulted from the division
+                corrected_flux[~np.isfinite(corrected_flux)] = np.nan
+                if corrected_error is not None:
+                    corrected_error[~np.isfinite(corrected_error)] = np.nan
+
+                # Create corrected spectrum
+                corrected_spectrum = spectrum.copy()
+                corrected_spectrum.flux = corrected_flux
+                corrected_spectrum.error = corrected_error
+                corrected_spectra.add_spectrum(corrected_spectrum)
+            else:
+                logger.warning(f"Blaze profile length mismatch for aperture {aperture_id}, skipping")
+                corrected_spectra.add_spectrum(spectrum)
+        else:
+            logger.warning(f"No blaze profile for aperture {aperture_id}, skipping de-blazing")
+            corrected_spectra.add_spectrum(spectrum)
+
+    logger.info(f"✓ De-blazing complete: {len(corrected_spectra.spectra)} spectra processed")
+    return corrected_spectra
+
+
+# 将去闪耀光谱保存为FITS表格。
+def save_deblazed_spectra(output_path: str, spectra_set: SpectraSet, 
+                          content_desc: str = 'De-blazed calibrated spectra',
+                          log_desc: str = 'de-blazed'):
+    """Save de-blazed spectra to FITS file."""
+    if spectra_set is None:
+        raise RuntimeError("No de-blazed spectra to save")
+
+    from astropy.io import fits
+    from pathlib import Path
+
+    output_path = Path(output_path)
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+
+    # Create structured array for table
+    norders = spectra_set.norders
+    max_pixels = max([s.npixel() for s in spectra_set.spectra.values()])
+
+    # Create binary table
+    cols = []
+    for aperture_id in spectra_set.get_orders():
+        spec = spectra_set.get_spectrum(aperture_id)
+
+        # Ensure arrays are floating-point so NaN padding is always valid.
+        wl_arr = np.asarray(spec.wavelength, dtype=np.float64)
+        fl_arr = np.asarray(spec.flux, dtype=np.float64)
+
+        # Pad arrays to same length
+        wl = np.pad(wl_arr, (0, max_pixels - len(wl_arr)), constant_values=np.nan)
+        fl = np.pad(fl_arr, (0, max_pixels - len(fl_arr)), constant_values=np.nan)
+
+        cols.append(fits.Column(name=f'WAV_{aperture_id:02d}', format='D', array=wl))
+        cols.append(fits.Column(name=f'FLUX_{aperture_id:02d}', format='D', array=fl))
+
+        if spec.error is not None:
+            err_arr = np.asarray(spec.error, dtype=np.float64)
+            err = np.pad(err_arr, (0, max_pixels - len(err_arr)), constant_values=np.nan)
+            cols.append(fits.Column(name=f'ERR_{aperture_id:02d}', format='D', array=err))
+
+    coldefs = fits.ColDefs(cols)
+    table_hdu = fits.BinTableHDU.from_columns(coldefs)
+
+    # Create header
+    header = fits.Header()
+    header['NORDERS'] = norders
+    header['NPX'] = max_pixels
+    header['CONTENTS'] = content_desc
+    header['COMMENT'] = 'Wavelength in Angstroms, Flux in ADU'
+
+    primary_hdu = fits.PrimaryHDU(header=header)
+
+    hdul = fits.HDUList([primary_hdu, table_hdu])
+    hdul.writeto(str(output_path), overwrite=True)
+
+    logger.info(f"Saved {norders} {log_desc} spectra to {output_path}")
+
+
+# 执行去闪耀校正并保存结果。
+def process_de_blazing_stage(spectra_set: SpectraSet,
+                           output_dir_base: str,
+                           flat_field: Optional[FlatField] = None,
+                           save_output: bool = True,
+                           save_deblaze: bool = True,
+                           output_filename: str = 'deblazed_spectra.fits',
+                           plot_prefix: str = 'deblazed',
+                           save_plots: bool = True,
+                           fig_format: str = 'png') -> SpectraSet:
+    """
+    Execute de-blazing correction stage.
+
+    Args:
+        spectra_set: Extracted spectra
+        output_dir_base: Base output directory
+        flat_field: Flat field data with blaze profiles
+        save_output: Whether to save de-blazed spectra to file
+        save_deblaze: Whether to save de-blazed intermediate products
+        output_filename: Filename for the saved FITS file
+        plot_prefix: Prefix for saved plot files
+        save_plots: Whether to save diagnostic plots
+        fig_format: Format for diagnostic plots
+
+    Returns:
+        De-blazed spectra set
+    """
+    if flat_field is None:
+        logger.warning("No flat field provided, skipping de-blazing")
+        return spectra_set
+
+    deblazed_spectra = apply_de_blazing(spectra_set, flat_field)
+
+    # Save de-blazed spectra if enabled
+    if save_output and save_deblaze:
+        deblazed_file = Path(output_dir_base) / 'step6_deblazing' / output_filename
+        deblazed_file.parent.mkdir(parents=True, exist_ok=True)
+        save_deblazed_spectra(str(deblazed_file), deblazed_spectra)
+
+        # Save diagnostic plots if enabled
+        if save_plots:
+            out_dir = deblazed_file.parent
+            pdf_name = f"{plot_prefix}.pdf"
+            pdf_path = out_dir / pdf_name
+            plot_spectra_to_pdf(deblazed_spectra, str(pdf_path), title_prefix="De-blazed Spectrum (pixels)", xlabel="Pixel")
+
+    return deblazed_spectra

specproc/core/extraction.py

@@ -0,0 +1,375 @@
+"""
+Spectrum extraction stage of spectral reduction pipeline.
+
+Handles 1D spectrum extraction using sum or optimal (Gaussian) methods.
+"""
+# 从二维图像中提取求和或最优一维光谱。
+
+import numpy as np
+import logging
+from pathlib import Path
+from typing import Tuple, Optional
+from scipy.optimize import curve_fit
+from specproc.core.data_structures import Spectrum, SpectraSet, ApertureSet, WaveCalib, FlatField
+from specproc.plotting.spectra_plotter import plot_spectrum_to_file, plot_spectra_to_pdf
+
+logger = logging.getLogger(__name__)
+
+
+# 负责从二维图像提取一维光谱。
+class SpectrumExtractor:
+    """Handles 1D spectrum extraction from 2D images."""
+
+    # 设置最优提取阈值与默认方法。
+    def __init__(self, optimal_sigma: float = 3.0, method: str = 'optimal'):
+        """
+        Initialize spectrum extractor.
+
+        Args:
+            optimal_sigma: Gaussian sigma for optimal extraction profile.
+            method: Default extraction method ('optimal' or 'sum').
+        """
+        self.optimal_sigma = optimal_sigma
+        self.method = method
+        self.extracted_spectra = None
+
+    # 对孔径内像素求和提取光谱。
+    def sum_extraction(self, image: np.ndarray, aperture_positions: np.ndarray,
+                      lower_bounds: np.ndarray = None, upper_bounds: np.ndarray = None) -> np.ndarray:
+        """
+        Extract spectrum using simple aperture sum.
+
+        Args:
+            image: 2D image array
+            aperture_positions: Aperture center positions (per column)
+            lower_bounds: Lower boundary positions (per column, absolute y)
+            upper_bounds: Upper boundary positions (per column, absolute y)
+
+        Returns:
+            1D spectrum array
+        """
+        logger.debug("Performing sum extraction...")
+
+        rows, cols = image.shape
+
+        if len(aperture_positions) != cols:
+            raise ValueError("Aperture positions length must match image width")
+
+        spectrum = np.zeros(cols)
+
+        for x in range(cols):
+            center = aperture_positions[x]
+            if lower_bounds is not None and upper_bounds is not None:
+                y_min = max(0, int(np.floor(lower_bounds[x])))
+                y_max = min(rows, int(np.ceil(upper_bounds[x])) + 1)
+            else:
+                y_min = max(0, int(center - 5))
+                y_max = min(rows, int(center + 5) + 1)
+
+            if y_max > y_min:
+                spectrum[x] = np.sum(image[y_min:y_max, x])
+            else:
+                spectrum[x] = 0.0
+
+        logger.debug(f"Sum extraction complete: {len(spectrum)} pixels")
+
+        return spectrum
+
+    # 按空间光强剖面最优提取光谱。
+    def optimal_extraction(self, image: np.ndarray, aperture_positions: np.ndarray,
+                          spatial_profile_2d: Optional[np.ndarray] = None,
+                          lower_bounds: np.ndarray = None,
+                          upper_bounds: np.ndarray = None) -> Tuple[np.ndarray, np.ndarray]:
+        """
+        Extract spectrum using optimal extraction.
+
+        Uses a profile-weighted approach. If a flat-field-derived profile is
+        provided, use it; otherwise fallback to Gaussian profile.
+
+        Args:
+            image: 2D image array
+            aperture_positions: Aperture center positions (per column)
+            profile: Optional spatial profile shape (relative weights)
+            lower_bounds: Lower boundary positions (per column, absolute y)
+            upper_bounds: Upper boundary positions (per column, absolute y)
+
+        Returns:
+            Tuple of (spectrum, error)
+        """
+        logger.debug("Performing optimal extraction...")
+
+        rows, cols = image.shape
+
+        if len(aperture_positions) != cols:
+            raise ValueError("Aperture positions length must match image width")
+
+        spectrum = np.zeros(cols)
+        error = np.zeros(cols)
+
+        for x in range(cols):
+            center = aperture_positions[x]
+            if lower_bounds is not None and upper_bounds is not None:
+                y_min = max(0, int(np.floor(lower_bounds[x])))
+                y_max = min(rows, int(np.ceil(upper_bounds[x])) + 1)
+            else:
+                y_min = max(0, int(center - 5))
+                y_max = min(rows, int(center + 5) + 1)
+
+            if y_max <= y_min:
+                spectrum[x] = 0.0
+                error[x] = 0.0
+                continue
+
+            sub_image = image[y_min:y_max, x]
+            local_y = np.arange(y_min, y_max) - center
+
+            weights = None
+            if spatial_profile_2d is not None:
+                p_col = spatial_profile_2d[y_min:y_max, x].astype(np.float64)
+                sum_p = np.sum(p_col)
+                if sum_p > 0:
+                    weights = p_col / sum_p
+
+            if weights is None:
+                sigma = self.optimal_sigma
+                weights = np.exp(-0.5 * (local_y / sigma) ** 2)
+                weights /= np.sum(weights) if np.sum(weights) > 0 else 1.0
+
+            weighted = sub_image * weights
+            spectrum[x] = np.sum(weighted)
+            error[x] = np.sqrt(np.sum((sub_image * weights) ** 2))
+
+        logger.debug(f"Optimal extraction complete: {len(spectrum)} pixels")
+
+        return spectrum, error
+
+    # 对全部孔径执行光谱提取。
+    def extract_aperture_set(self, image: np.ndarray, apertures: ApertureSet,
+                           wavelength_calib: Optional[WaveCalib] = None,
+                           flat_field: Optional[FlatField] = None,
+                           method_override: Optional[str] = None) -> SpectraSet:
+        """
+        Extract all apertures from 2D image.
+
+        Args:
+            image: 2D science image
+            apertures: ApertureSet with aperture definitions
+            wavelength_calib: Optional wavelength calibration
+
+        Returns:
+            SpectraSet with extracted spectra
+        """
+        logger.info(f"Extracting {apertures.norders} apertures...")
+
+        spectra_set = SpectraSet()
+
+        method = method_override if method_override is not None else self.method
+
+        _, cols = image.shape
+
+        for aperture_id in apertures.get_orders():
+            aperture = apertures.get_aperture(aperture_id)
+            if aperture is None:
+                continue
+
+            x_pix = np.arange(cols)
+            center_pos = aperture.get_position(x_pix)
+            lower_bounds = aperture.get_lower(x_pix)
+            upper_bounds = aperture.get_upper(x_pix)
+
+            spatial_profile_2d = None
+            if flat_field is not None:
+                if flat_field.smoothed_model is not None:
+                    spatial_profile_2d = flat_field.smoothed_model
+
+            if method == 'optimal':
+                flux, error = self.optimal_extraction(image, center_pos,
+                                                      spatial_profile_2d=spatial_profile_2d,
+                                                      lower_bounds=lower_bounds,
+                                                      upper_bounds=upper_bounds)
+            else:
+                flux = self.sum_extraction(image, center_pos,
+                                          lower_bounds=lower_bounds,
+                                          upper_bounds=upper_bounds)
+                error = np.sqrt(np.abs(flux))
+
+            # Apply wavelength calibration if available
+            if wavelength_calib is not None:
+                wavelength = wavelength_calib.apply_to_pixel(x_pix,
+                                                            aperture_id * np.ones_like(x_pix))
+            else:
+                wavelength = x_pix
+
+            # Create Spectrum object
+            spec = Spectrum(
+                aperture=aperture_id,
+                order=aperture.order,
+                wavelength=wavelength,
+                flux=flux,
+                error=error,
+                background=None,
+                mask=None
+            )
+
+            spectra_set.add_spectrum(spec)
+
+        self.extracted_spectra = spectra_set
+        logger.info(f"Extraction complete: {spectra_set.norders} spectra")
+
+        return spectra_set
+
+    # 将提取的光谱集合写入FITS。
+    def save_spectra(self, output_path: str, spectra_set: SpectraSet):
+        """Save extracted spectra to FITS file."""
+        if spectra_set is None:
+            raise RuntimeError("No spectra to save")
+
+        from astropy.io import fits
+        from pathlib import Path
+
+        output_path = Path(output_path)
+        output_path.parent.mkdir(parents=True, exist_ok=True)
+
+        # Create structured array for table
+        norders = spectra_set.norders
+        max_pixels = max([s.npixel() for s in spectra_set.spectra.values()])
+
+
+        # Create binary table
+        cols = []
+        for aperture_id in spectra_set.get_orders():
+            spec = spectra_set.get_spectrum(aperture_id)
+
+            # Pad arrays to same length
+            wl = np.pad(spec.wavelength, (0, max_pixels - len(spec.wavelength)))
+            fl = np.pad(spec.flux, (0, max_pixels - len(spec.flux)))
+
+            cols.append(fits.Column(name=f'WAV_{aperture_id:02d}', format='D', array=wl))
+            cols.append(fits.Column(name=f'FLUX_{aperture_id:02d}', format='D', array=fl))
+
+            if spec.error is not None:
+                err = np.pad(spec.error, (0, max_pixels - len(spec.error)))
+                cols.append(fits.Column(name=f'ERR_{aperture_id:02d}', format='D', array=err))
+
+        # Create HDUs
+        primary_hdu = fits.PrimaryHDU()
+        table_hdu = fits.BinTableHDU.from_columns(cols)
+        table_hdu.header['NORDERS'] = norders
+
+        hdul = fits.HDUList([primary_hdu, table_hdu])
+        hdul.writeto(str(output_path), overwrite=True)
+
+        logger.info(f"Saved {norders} spectra to {output_path}")
+
+
+# 执行光谱提取并保存相关产物。
+def process_extraction_stage(science_image: np.ndarray,
+                            apertures: ApertureSet,
+                            output_dir_base: str,
+                            wavelength_calib: Optional[WaveCalib] = None,
+                            flat_field: Optional[FlatField] = None,
+                            method_override: Optional[str] = None,
+                            output_filename: str = 'extracted_spectra.fits',
+                            plot_prefix: str = 'extracted',
+                            optimal_sigma: float = 3.0,
+                            extraction_method: str = 'optimal',
+                            save_plots: bool = True,
+                            fig_format: str = 'png') -> SpectraSet:
+    """
+    Execute spectrum extraction stage.
+
+    Args:
+        science_image: 2D science image
+        apertures: Detected apertures
+        output_dir_base: Base output directory
+        wavelength_calib: Wavelength calibration
+        flat_field: Flat field data
+        method_override: Override default extraction method
+        output_filename: Name for output FITS file
+        plot_prefix: Prefix for diagnostic plots
+        optimal_sigma: Sigma for Gaussian profile in optimal extraction
+        extraction_method: Default extraction method
+        save_plots: Whether to save diagnostic plots
+        fig_format: Format for diagnostic plots
+
+    Returns:
+        SpectraSet with extracted 1D spectra
+    """
+    extractor = SpectrumExtractor(optimal_sigma=optimal_sigma, method=extraction_method)
+
+    # Extract spectra
+    spectra_set = extractor.extract_aperture_set(science_image, apertures,
+                                                wavelength_calib, flat_field,
+                                                method_override=method_override)
+
+    # Save spectra
+    spectra_file = Path(output_dir_base) / 'step5_extraction' / output_filename
+    spectra_file.parent.mkdir(parents=True, exist_ok=True)
+    extractor.save_spectra(str(spectra_file), spectra_set)
+
+    # Save diagnostic plots if enabled
+    if save_plots:
+        out_dir = spectra_file.parent
+        if wavelength_calib is not None:
+            pdf_path = out_dir / f"{plot_prefix}_all_apertures.pdf"
+            plot_spectra_to_pdf(spectra_set, str(pdf_path), title_prefix="Extracted Spectrum")
+        else:
+            # If no wavelength calibration, plot raw pixels
+            pdf_name = f"{plot_prefix}.pdf"
+            pdf_path = out_dir / pdf_name
+            plot_spectra_to_pdf(spectra_set, str(pdf_path), title_prefix="Extracted Spectrum (pixels)", xlabel="Pixel")
+
+    return spectra_set
+
+# 从FITS文件载入提取光谱。
+def load_extracted_spectra(filepath: str) -> SpectraSet:
+    """Load extracted spectra from FITS file."""
+    from astropy.io import fits
+    from specproc.core.data_structures import Spectrum, SpectraSet
+    
+    spectra_set = SpectraSet()
+    
+    with fits.open(filepath) as hdul:
+        if len(hdul) < 2:
+            raise ValueError("FITS file does not contain a binary table")
+            
+        table_hdu = hdul[1]
+        
+        # Extract unique aperture IDs from column names
+        aperture_ids = set()
+        for col in table_hdu.columns:
+            if col.name.startswith('FLUX_'):
+                try:
+                    aperture_ids.add(int(col.name.split('_')[1]))
+                except ValueError:
+                    pass
+                
+        for ap_id in sorted(list(aperture_ids)):
+            wav_col = f'WAV_{ap_id:02d}'
+            flux_col = f'FLUX_{ap_id:02d}'
+            err_col = f'ERR_{ap_id:02d}'
+            
+            if wav_col in table_hdu.data.columns.names and flux_col in table_hdu.data.columns.names:
+                wav = table_hdu.data[wav_col]
+                flux = table_hdu.data[flux_col]
+                err = table_hdu.data[err_col] if err_col in table_hdu.data.columns.names else None
+                
+                if len(wav.shape) > 1 and wav.shape[0] == 1:
+                    wav = wav[0]
+                if len(flux.shape) > 1 and flux.shape[0] == 1:
+                    flux = flux[0]
+                if err is not None and len(err.shape) > 1 and err.shape[0] == 1:
+                    err = err[0]
+                
+                spec = Spectrum(
+                    aperture=ap_id,
+                    order=ap_id, # order is same as aperture for now
+                    wavelength=wav,
+                    flux=flux,
+                    error=err,
+                    background=None,
+                    mask=None
+                )
+                spectra_set.add_spectrum(spec)
+                
+    return spectra_set