SpecProc 1.0.0__py3-none-any.whl

specproc/utils/fits_io.py

@@ -0,0 +1,269 @@
+"""
+FITS file I/O utilities for reading and writing spectral data.
+
+Provides functions for loading raw FITS images, writing intermediate
+results, and managing FITS headers.
+"""
+# 负责光谱数据 FITS 文件的读取与写入。
+
+import numpy as np
+from astropy.io import fits
+from pathlib import Path
+from typing import Tuple, Dict, Optional
+import logging
+
+logger = logging.getLogger(__name__)
+
+
+# 读取 FITS 图像及其头信息。
+def read_fits_image(filepath: str) -> Tuple[np.ndarray, dict]:
+    """
+    Read a FITS image file.
+
+    Args:
+        filepath: Path to FITS file
+
+    Returns:
+        Tuple of (image_data, header_dict)
+    """
+    try:
+        with fits.open(filepath) as hdul:
+            if len(hdul) == 0:
+                raise ValueError(f"FITS file {filepath} has no HDUs")
+
+            # Get primary data (image)
+            image_data = hdul[0].data
+            if image_data is None:
+                # Try first non-empty HDU
+                for hdu in hdul:
+                    if hdu.data is not None:
+                        image_data = hdu.data
+                        break
+
+            # Get header
+            header = dict(hdul[0].header)
+
+            return image_data, header
+    except Exception as e:
+        logger.error(f"Error reading FITS file {filepath}: {e}")
+        raise
+
+
+# 读取 FITS 表格及其头信息。
+def read_fits_table(filepath: str, hdu: int = 1) -> Tuple[np.ndarray, dict]:
+    """
+    Read structured table from FITS file.
+
+    Args:
+        filepath: Path to FITS file
+        hdu: HDU index (default 1 for binary table)
+
+    Returns:
+        Tuple of (table_data, header_dict)
+    """
+    try:
+        with fits.open(filepath) as hdul:
+            if len(hdul) <= hdu:
+                raise ValueError(f"FITS file {filepath} has no HDU at index {hdu}")
+
+            table_hdu = hdul[hdu]
+            data = table_hdu.data
+            header = dict(table_hdu.header)
+
+            return data, header
+    except Exception as e:
+        logger.error(f"Error reading FITS table from {filepath}: {e}")
+        raise
+
+
+# 将图像数据写入 FITS 文件。
+def write_fits_image(filepath: str, data: np.ndarray, header: Optional[Dict] = None,
+                     overwrite: bool = True, dtype: Optional[str] = None):
+    """
+    Write image to FITS file.
+
+    Args:
+        filepath: Output file path
+        data: Image array
+        header: Optional FITS header dictionary
+        overwrite: Overwrite existing file
+        dtype: Optional data type for output (e.g., 'float32', 'float64')
+    """
+    try:
+        # Convert data to specified dtype if provided
+        if dtype is not None:
+            original_dtype = str(data.dtype)
+            data = data.astype(dtype)
+            logger.debug(f"Converted data from {original_dtype} to {dtype}")
+        
+        # Create primary HDU
+        hdu = fits.PrimaryHDU(data=data)
+
+        # If writing to a float format, remove BZERO and BSCALE to prevent incorrect scaling
+        # by FITS readers. These keywords are for scaling integers.
+        if data.dtype.kind == 'f' and header:
+            if 'BZERO' in header: del header['BZERO']
+            if 'BSCALE' in header: del header['BSCALE']
+
+        # Add header keywords
+        if header:
+            for key, value in header.items():
+                try:
+                    # Skip COMMENT and HISTORY - they often contain special formatting
+                    # that causes issues with astropy, and are not critical for processing
+                    if key.upper() in ('COMMENT', 'HISTORY'):
+                        continue
+
+                    # astropy.io.fits considers a single dot as an illegal value.
+                    # Replace it with an empty string, which is valid.
+                    if isinstance(value, str) and value.strip() == '.':
+                        value = ''
+                    
+                    # Clean up non-ASCII characters for FITS header values
+                    if isinstance(value, str):
+                        # Remove or replace non-ASCII characters
+                        cleaned_value = value.encode('ascii', 'ignore').decode('ascii')
+                        # Truncate if too long (FITS headers have max 80 chars per record)
+                        if len(cleaned_value) > 70:
+                            cleaned_value = cleaned_value[:67] + '...'
+                        hdu.header[key] = cleaned_value
+                    else:
+                        hdu.header[key] = value
+                except Exception as e:
+                    logger.warning(f"Could not add header keyword {key}: {e}")
+
+        # Write to file
+        hdul = fits.HDUList([hdu])
+        hdul.writeto(filepath, overwrite=overwrite)
+
+        logger.info(f"Wrote image to {filepath}")
+    except Exception as e:
+        logger.error(f"Error writing FITS image {filepath}: {e}")
+        raise
+
+
+# 将结构化数组写入 FITS 表格。
+def write_fits_table(filepath: str, data: np.ndarray, header: Optional[Dict] = None,
+                     overwrite: bool = True):
+    """
+    Write structured array as binary table to FITS file.
+
+    Args:
+        filepath: Output file path
+        data: Structured numpy array with named fields
+        header: Optional FITS header
+        overwrite: Overwrite existing file
+    """
+    try:
+        # Create binary table HDU from structured array
+        cols = []
+        for name in data.dtype.names:
+            col_data = data[name]
+            if col_data.dtype == np.object_:
+                # Handle variable-length columns
+                format_str = f'P{len(col_data[0])}'
+                col = fits.Column(name=name, format=format_str, array=col_data)
+            else:
+                col = fits.Column(name=name, format=f'{len(col_data[0]) if col_data.ndim > 1 else 1}D',
+                                array=col_data)
+            cols.append(col)
+
+        table_hdu = fits.BinTableHDU.from_columns(cols)
+
+        # Add header keywords
+        if header:
+            for key, value in header.items():
+                try:
+                    table_hdu.header[key] = value
+                except Exception as e:
+                    logger.warning(f"Could not add header keyword {key}: {e}")
+
+        # Create HDU list with primary + table
+        primary_hdu = fits.PrimaryHDU()
+        hdul = fits.HDUList([primary_hdu, table_hdu])
+        hdul.writeto(filepath, overwrite=overwrite)
+
+        logger.info(f"Wrote table to {filepath}")
+    except Exception as e:
+        logger.error(f"Error writing FITS table {filepath}: {e}")
+        raise
+
+
+# 忽略大小写查找头关键字。
+def get_header_value(header: dict, keyword: str, default: any = None) -> any:
+    """
+    Get header value with case-insensitive lookup.
+
+    Args:
+        header: FITS header dictionary
+        keyword: Keyword name
+        default: Default value if not found
+
+    Returns:
+        Header value or default
+    """
+    # Try exact match
+    if keyword in header:
+        return header[keyword]
+
+    # Try case-insensitive
+    for key, value in header.items():
+        if key.upper() == keyword.upper():
+            return value
+
+    return default
+
+
+# 获取 FITS 图像的基本信息。
+def get_image_info(filepath: str) -> Dict:
+    """
+    Get basic information about a FITS image.
+
+    Args:
+        filepath: Path to FITS file
+
+    Returns:
+        Dictionary with image info (shape, datatype, header keywords)
+    """
+    try:
+        img, header = read_fits_image(filepath)
+        return {
+            'filepath': filepath,
+            'shape': img.shape,
+            'dtype': str(img.dtype),
+            'exptime': get_header_value(header, 'EXPTIME', 0.0),
+            'date_obs': get_header_value(header, 'DATE-OBS', ''),
+            'object': get_header_value(header, 'OBJECT', ''),
+        }
+    except Exception as e:
+        logger.error(f"Error getting info from {filepath}: {e}")
+        return {}
+
+
+# 按指定方法合并多个 FITS 图像。
+def combine_fits_images(*filepaths: str, method: str = 'mean') -> np.ndarray:
+    """
+    Combine multiple FITS images.
+
+    Args:
+        filepaths: Paths to FITS files
+        method: 'mean', 'median', or 'sum'
+
+    Returns:
+        Combined image array
+    """
+    images = []
+    for filepath in filepaths:
+        img, _ = read_fits_image(filepath)
+        images.append(img)
+
+    images = np.array(images)
+
+    if method == 'mean':
+        return np.mean(images, axis=0)
+    elif method == 'median':
+        return np.median(images, axis=0)
+    elif method == 'sum':
+        return np.sum(images, axis=0)
+    else:
+        raise ValueError(f"Unknown combine method: {method}")

specproc/utils/image_processing.py

@@ -0,0 +1,426 @@
+"""
+Image processing utilities for spectral data.
+
+Provides functions for image operations like combining, filtering,
+and detecting bad pixels.
+"""
+# 提供图像合并、滤波和背景处理工具。
+
+import numpy as np
+from typing import Optional, Tuple
+from scipy.ndimage import gaussian_filter
+from scipy import signal, ndimage
+from scipy.interpolate import SmoothBivariateSpline
+import logging
+
+logger = logging.getLogger(__name__)
+
+
+# 合并多幅图像并估计不确定度。
+def combine_images(images: np.ndarray, method: str = 'median',
+                   sigma: float = 3.0) -> Tuple[np.ndarray, np.ndarray]:
+    """
+    Combine multiple 2D images.
+
+    Args:
+        images: Array of shape (n_images, height, width)
+        method: 'mean', 'median', or 'sigma_clip'
+        sigma: Sigma clipping threshold (for sigma_clip method)
+
+    Returns:
+        Tuple of (combined_image, uncertainty_map)
+    """
+    if method == 'mean':
+        combined = np.mean(images, axis=0)
+        uncertainty = np.std(images, axis=0) / np.sqrt(images.shape[0])
+        return combined, uncertainty
+
+    elif method == 'median':
+        combined = np.median(images, axis=0)
+        uncertainty = np.std(images, axis=0)
+        return combined, uncertainty
+
+    elif method == 'sigma_clip':
+        combined = np.zeros_like(images[0], dtype=float)
+        uncertainty = np.zeros_like(images[0], dtype=float)
+
+        for i in range(images.shape[1]):
+            for j in range(images.shape[2]):
+                col = images[:, i, j]
+                mean = np.mean(col)
+                std = np.std(col)
+
+                # Mask outliers
+                mask = np.abs(col - mean) < sigma * std
+                if np.sum(mask) > 0:
+                    combined[i, j] = np.mean(col[mask])
+                    uncertainty[i, j] = std
+                else:
+                    combined[i, j] = mean
+                    uncertainty[i, j] = std
+
+        return combined, uncertainty
+
+    else:
+        raise ValueError(f"Unknown combine method: {method}")
+
+
+# 对二维图像应用 Savitzky-Golay 滤波。
+def savitzky_golay_2d(image: np.ndarray, window_size: int = 5,
+                      order: int = 2) -> np.ndarray:
+    """
+    Apply 2D Savitzky-Golay filter.
+
+    Args:
+        image: 2D image array
+        window_size: Filter window size (must be odd)
+        order: Polynomial order
+
+    Returns:
+        Filtered image
+    """
+    if window_size % 2 == 0:
+        window_size += 1
+
+    # Create 1D filter
+    coeffs = signal.savgol_coeffs(window_size, order)
+
+    # Apply horizontal and vertical
+    horizontal = ndimage.convolve1d(image, coeffs, axis=1)
+    filtered = ndimage.convolve1d(horizontal, coeffs, axis=0)
+
+    return filtered
+
+
+# 按阈值检测图像中的异常像素。
+def find_bad_pixels(image: np.ndarray, threshold: float = 3.0,
+                    method: str = 'sigma') -> np.ndarray:
+    """
+    Find bad (outlier) pixels in an image.
+
+    Args:
+        image: 2D image array
+        threshold: Detection threshold (sigma)
+        method: 'sigma' for sigma-clipping, 'percentile' for percentile-based
+
+    Returns:
+        Boolean mask (True = bad pixel)
+    """
+    if method == 'sigma':
+        mean = np.mean(image)
+        std = np.std(image)
+        bad_pixels = np.abs(image - mean) > threshold * std
+        return bad_pixels
+
+    elif method == 'percentile':
+        lower = np.percentile(image, threshold)
+        upper = np.percentile(image, 100 - threshold)
+        bad_pixels = (image < lower) | (image > upper)
+        return bad_pixels
+
+    else:
+        raise ValueError(f"Unknown method: {method}")
+
+
+# 构建多项式或切比雪夫基矩阵。
+def _build_design_matrix(x: np.ndarray, y: np.ndarray, order: int,
+                         method: str) -> np.ndarray:
+    """Build 2D basis matrix for polynomial or Chebyshev fitting."""
+    if method == 'chebyshev':
+        vander = np.polynomial.chebyshev.chebvander2d(x, y, [order, order])
+        return vander.reshape(x.size, -1)
+
+    # Fallback: regular polynomial basis.
+    ncoeff = (order + 1) * (order + 1)
+    A = np.zeros((x.size, ncoeff), dtype=np.float64)
+    idx = 0
+    for i in range(order + 1):
+        for j in range(order + 1):
+            A[:, idx] = (x ** i) * (y ** j)
+            idx += 1
+    return A
+
+
+# 使用中位绝对偏差估计稳健标准差。
+def _robust_sigma(values: np.ndarray) -> float:
+    """Robust sigma estimate via MAD; returns 0 for empty arrays."""
+    if values.size == 0:
+        return 0.0
+    med = np.median(values)
+    mad = np.median(np.abs(values - med))
+    if mad <= 0:
+        std = np.std(values)
+        return float(std) if np.isfinite(std) else 0.0
+    return float(1.4826 * mad)
+
+
+# 为单幅图像拟合平滑二维背景。
+def _estimate_single_background(image: np.ndarray, order: int = 2,
+                                valid_mask: Optional[np.ndarray] = None,
+                                method: str = 'chebyshev',
+                                smooth_sigma: float = 20.0,
+                                sigma_clip: float = 3.0,
+                                maxiters: int = 4,
+                                bspline_smooth: float = 1.0,
+                                clip_mode: str = 'both') -> np.ndarray:
+    """Estimate smooth 2D background for a single (non-split) image.
+
+    Parameters
+    ----------
+    clip_mode : {'both', 'upper', 'lower'}
+        Sigma-clipping mode.  'both' rejects outliers on both sides
+        (symmetric, the traditional behaviour).  'upper' rejects only
+        positive outliers (pixels brighter than the model) which is
+        recommended for flat-field background estimation where unmasked
+        faint orders contaminate the inter-order region.  'lower' rejects
+        only negative outliers.
+    """
+    rows, cols = image.shape
+    yy, xx = np.meshgrid(np.arange(rows), np.arange(cols), indexing='ij')
+
+    fit_mask = np.isfinite(image)
+    if valid_mask is not None:
+        fit_mask &= valid_mask.astype(bool)
+
+    if np.sum(fit_mask) < max(100, (order + 1) ** 2):
+        logger.warning("Too few valid samples for background fit; returning zeros")
+        return np.zeros_like(image, dtype=np.float64)
+
+    work_data = image.astype(np.float64)
+
+    if method in ('smooth', 'gaussian_smooth'):
+        current = fit_mask.copy()
+        s = max(1.0, float(smooth_sigma))
+
+        for _ in range(max(1, int(maxiters))):
+            use = np.zeros_like(work_data, dtype=np.float64)
+            wgt = np.zeros_like(work_data, dtype=np.float64)
+            use[current] = work_data[current]
+            wgt[current] = 1.0
+
+            num = gaussian_filter(use, sigma=s, mode='nearest')
+            den = gaussian_filter(wgt, sigma=s, mode='nearest')
+            model = num / np.clip(den, 1e-8, None)
+
+            if sigma_clip <= 0:
+                return model
+
+            resid = work_data[current] - model[current]
+            sig = _robust_sigma(resid)
+            if sig <= 0:
+                break
+
+            diff = work_data - model
+            if clip_mode == 'upper':
+                keep = diff <= sigma_clip * sig
+            elif clip_mode == 'lower':
+                keep = diff >= -sigma_clip * sig
+            else:
+                keep = np.abs(diff) <= sigma_clip * sig
+            new_current = current & keep
+            if np.array_equal(new_current, current):
+                break
+            if np.sum(new_current) < max(100, (order + 1) ** 2):
+                break
+            current = new_current
+
+        return model
+
+    x_norm = (xx.astype(np.float64) / max(cols - 1, 1)) * 2.0 - 1.0
+    y_norm = (yy.astype(np.float64) / max(rows - 1, 1)) * 2.0 - 1.0
+
+    xp = x_norm[fit_mask]
+    yp = y_norm[fit_mask]
+    zp = work_data[fit_mask]
+    keep = np.ones_like(zp, dtype=bool)
+
+    for _ in range(max(1, int(maxiters))):
+        if np.sum(keep) < max(100, (order + 1) ** 2):
+            break
+
+        try:
+            if method == 'bspline':
+                kxy = max(1, min(3, int(order)))
+                s = max(0.0, float(bspline_smooth)) * float(np.sum(keep))
+                spline = SmoothBivariateSpline(
+                    xp[keep], yp[keep], zp[keep], kx=kxy, ky=kxy, s=s
+                )
+                fit_all = spline.ev(xp, yp)
+                model = spline.ev(x_norm.ravel(), y_norm.ravel()).reshape(rows, cols)
+            else:
+                fit_kind = 'chebyshev' if method not in ('poly', 'polynomial') else 'poly'
+                A = _build_design_matrix(xp[keep], yp[keep], order, fit_kind)
+                coeffs, _, _, _ = np.linalg.lstsq(A, zp[keep], rcond=None)
+                A_all = _build_design_matrix(xp, yp, order, fit_kind)
+                fit_all = A_all @ coeffs
+                A_img = _build_design_matrix(x_norm.ravel(), y_norm.ravel(), order, fit_kind)
+                model = (A_img @ coeffs).reshape(rows, cols)
+        except Exception as e:
+            logger.warning(f"Background {method} fit failed, fallback to gaussian_smooth: {e}")
+            return _estimate_single_background(
+                image,
+                order=order,
+                valid_mask=valid_mask,
+                method='gaussian_smooth',
+                smooth_sigma=smooth_sigma,
+                sigma_clip=sigma_clip,
+                maxiters=maxiters,
+                bspline_smooth=bspline_smooth,
+                clip_mode=clip_mode,
+            )
+
+        if sigma_clip <= 0:
+            return model
+
+        resid = zp - fit_all
+        sig = _robust_sigma(resid[keep])
+        if sig <= 0:
+            break
+
+        center = np.median(resid[keep])
+        if clip_mode == 'upper':
+            new_keep = (resid - center) <= sigma_clip * sig
+        elif clip_mode == 'lower':
+            new_keep = (resid - center) >= -sigma_clip * sig
+        else:
+            new_keep = np.abs(resid - center) <= sigma_clip * sig
+        if np.array_equal(new_keep, keep):
+            break
+        keep = new_keep
+
+    return model
+
+
+# Public alias: estimate a single region without any internal splitting.
+# Use this when the caller handles detector-half splitting externally.
+estimate_background_region = _estimate_single_background
+# 估计二维背景并可按区域分别拟合。
+def estimate_background_2d(image: np.ndarray, order: int = 2,
+                           split_vertically: bool = True,
+                           valid_mask: Optional[np.ndarray] = None,
+                           method: str = 'chebyshev',
+                           smooth_sigma: float = 20.0,
+                           sigma_clip: float = 3.0,
+                           maxiters: int = 4,
+                           bspline_smooth: float = 1.0) -> np.ndarray:
+    """
+    Estimate smooth 2D background using polynomial fitting.
+    
+    By default, splits the image into upper and lower halves for independent
+    processing (e.g., for CCDs with two readout amplifiers).
+
+    Args:
+        image: 2D image array
+        order: Polynomial order
+        split_vertically: Split image into upper/lower halves for independent 
+                         background fitting (default: True)
+        valid_mask: Optional boolean mask of valid sampling pixels (True = use)
+        method: 'chebyshev', 'bspline', 'gaussian_smooth', or 'poly'
+        smooth_sigma: Gaussian sigma for smooth method
+        sigma_clip: Sigma-clipping threshold for robust rejection
+        maxiters: Maximum sigma-clipping iterations
+        bspline_smooth: Smoothing factor scale for bspline method
+
+    Returns:
+        Background model
+    """
+    rows, cols = image.shape
+    
+    if split_vertically and rows > 1:
+        # Split into top and bottom halves for independent processing
+        # Note: numpy image[0,:] = ds9 row 1 (bottom), image[rows-1,:] = ds9 row rows (top)
+        # image[:mid_row,:] = image_bottom = ds9 bottom half (small row numbers)
+        # image[mid_row:,:] = image_top = ds9 top half (large row numbers)
+        mid_row = rows // 2
+        logger.info(f"Estimating background with vertical split at row {mid_row}")
+        
+        image_bottom = image[:mid_row, :]
+        image_top = image[mid_row:, :]
+        mask_bottom = valid_mask[:mid_row, :] if valid_mask is not None else None
+        mask_top = valid_mask[mid_row:, :] if valid_mask is not None else None
+        
+        background_bottom = _estimate_single_background(
+            image_bottom,
+            order,
+            valid_mask=mask_bottom,
+            method=method,
+            smooth_sigma=smooth_sigma,
+            sigma_clip=sigma_clip,
+            maxiters=maxiters,
+            bspline_smooth=bspline_smooth,
+        )
+        background_top = _estimate_single_background(
+            image_top,
+            order,
+            valid_mask=mask_top,
+            method=method,
+            smooth_sigma=smooth_sigma,
+            sigma_clip=sigma_clip,
+            maxiters=maxiters,
+            bspline_smooth=bspline_smooth,
+        )
+        
+        # Combine backgrounds
+        background = np.zeros_like(image)
+        background[:mid_row, :] = background_bottom
+        background[mid_row:, :] = background_top
+        
+        return background
+    else:
+        return _estimate_single_background(
+            image,
+            order,
+            valid_mask=valid_mask,
+            method=method,
+            smooth_sigma=smooth_sigma,
+            sigma_clip=sigma_clip,
+            maxiters=maxiters,
+            bspline_smooth=bspline_smooth,
+        )
+
+
+# 归一化平场图像至单位响应。
+def normalize_flat(flat_image: np.ndarray, mask: Optional[np.ndarray] = None) -> np.ndarray:
+    """
+    Normalize flat field to unity using sigma clipping.
+
+    Args:
+        flat_image: Flat field image
+        mask: Optional bad pixel mask
+
+    Returns:
+        Normalized flat (clipped to valid range)
+    """
+    from astropy.stats import sigma_clipped_stats
+
+    # Apply mask if provided
+    if mask is not None:
+        flat_array = flat_image.copy()
+        flat_array[mask > 0] = np.nan
+        # Use sigma clipping with masked values
+        mean_val, _, _ = sigma_clipped_stats(flat_array, mask=np.isnan(flat_array), sigma=3.0, maxiters=5)
+    else:
+        # Use sigma clipping without mask
+        mean_val, _, _ = sigma_clipped_stats(flat_image, sigma=3.0, maxiters=5)
+
+    normalized = flat_image / mean_val
+
+    # Clip to avoid extreme values
+    normalized = np.clip(normalized, 0.1, 10.0)
+
+    return normalized
+
+
+# 从图像中截取指定矩形区域。
+def extract_subimage(image: np.ndarray, x_start: int, x_end: int,
+                    y_start: int, y_end: int) -> np.ndarray:
+    """Extract rectangular sub-image."""
+    return image[y_start:y_end, x_start:x_end]
+
+
+# 按边缘宽度裁剪图像。
+def trim_image(image: np.ndarray, trim_left: int = 0, trim_right: int = 0,
+              trim_top: int = 0, trim_bottom: int = 0) -> np.ndarray:
+    """Trim edges of image."""
+    rows, cols = image.shape
+    return image[trim_top:rows-trim_bottom, trim_left:cols-trim_right]