Source code for bioniumx.preprocessing.normalizer

"""
Continuum normalization for transmission and emission spectra.
"""
import numpy as np
from bioniumx.core import BioniumXObject
from bioniumx.spectra import TransmissionSpectrum, EmissionSpectrum


[docs] def continuum_normalize(spectrum: BioniumXObject, method: str = "polynomial", degree: int = 2): """ Normalize a spectrum by dividing by its continuum. The continuum represents the baseline flux/depth devoid of absorption features. This function fits a model (e.g., polynomial) to the pseudo-continuum points and divides the entire spectrum by this model. Parameters ---------- spectrum : BioniumXObject The input spectrum. method : str, optional Method for fitting the continuum. Currently supports "polynomial" or "mean". Default is "polynomial". degree : int, optional Degree of the polynomial if method is "polynomial". Default is 2. Returns ------- normalized : BioniumXObject A new spectrum object normalized to its continuum. Raises ------ ValueError If an unsupported method is specified. Examples -------- >>> spec_norm = continuum_normalize(spec, method="polynomial", degree=2) """ if hasattr(spectrum, 'transit_depth'): y = spectrum.transit_depth elif hasattr(spectrum, 'flux'): y = spectrum.flux else: raise TypeError("Object must have either 'transit_depth' or 'flux' attribute.") x = spectrum.wavelength if method == "mean": continuum = np.ones_like(y) * np.mean(y) elif method == "polynomial": # Fit polynomial to the entire spectrum (in practice, one should mask lines) coeffs = np.polyfit(x, y, degree) continuum = np.polyval(coeffs, x) else: raise ValueError(f"Unknown method: {method}") # Avoid division by zero continuum = np.where(continuum == 0, 1.0, continuum) y_norm = y / continuum # Error propagation err_norm = spectrum.err / continuum if isinstance(spectrum, TransmissionSpectrum): return TransmissionSpectrum( wavelength=x, transit_depth=y_norm, err=err_norm, **spectrum.meta ) elif isinstance(spectrum, EmissionSpectrum): return EmissionSpectrum( wavelength=x, flux=y_norm, err=err_norm, **spectrum.meta ) else: raise TypeError("Spectrum must be TransmissionSpectrum or EmissionSpectrum.")