jaxqsofit.QSOFit

jaxqsofit.QSOFit#

class jaxqsofit.QSOFit(lam, flux, err=None, z=0.0, ra=-999, dec=-999, filename=None, output_path=None, wdisp=None, psf_mags=None, psf_mag_errs=None, psf_bands=None)[source]#

Bases: object

save_posterior_bundle(save_name=None, save_path=None)[source]#

Persist posterior samples plus minimal metadata for compact reloads.

classmethod load_from_samples(filename=None, output_path=None, save_name=None, plot_fig=True, plot_diagnostics=True, kwargs_plot=None, diagnostics_kwargs=None)[source]#

Load a compressed HDF5 posterior bundle and return a QSOFit object.

fit(name=None, deredden=True, wave_range=None, wave_mask=None, save_fits_name=None, fit_lines=True, save_result=True, plot_fig=True, save_fig=True, show_plot=False, decompose_host=True, fit_pl=True, fit_fe=True, fit_bc=False, fit_bal=False, fit_poly=True, fit_reddening=True, fit_poly_order=2, mask_lya_forest=True, fit_method='optax+nuts', verbose=True, fsps_age_grid=(0.01, 0.03, 0.1, 0.3, 1.0, 3.0, 10.0), fsps_logzsol_grid=(-1.0, -0.5, 0.0, 0.2), prior_config=None, dsps_ssp_fn='tempdata.h5', nuts_warmup=50, nuts_samples=50, nuts_chains=1, nuts_target_accept=0.9, optax_steps=600, optax_lr=0.01, psf_mags=None, psf_mag_errs=None, psf_bands=None, use_psf_phot=False, custom_components=None, custom_line_components=None, kwargs_plot=None)[source]#

Run end-to-end preprocessing, fitting, and optional plotting/saving.

Parameters:

  • name (str or None, optional) – Optional override basename used for figure/output naming.

  • deredden (bool, optional) – If True, apply Galactic dereddening with dustmaps SFD.

  • wave_range (tuple[float, float] or None, optional) – Rest-frame wavelength range to keep.

  • wave_mask (array-like or None, optional) – Rest-frame wavelength intervals to mask.

  • save_fits_name (str or None, optional) – Basename for saved result table.

  • fit_lines (bool, optional) – Enable/disable emission-line model.

  • save_result (bool, optional) – If True, write summary results to CSV.

  • plot_fig (bool, optional) – If True, render decomposition figure.

  • save_fig (bool, optional) – If True, save rendered figures.

  • show_plot (bool, optional) – If True, call plt.show() for the decomposition figure. Default is False to avoid interactive pop-ups in notebook/pipeline runs.

  • decompose_host (bool, optional) – Enable/disable host SPS decomposition.

  • fit_pl (bool, optional) – Enable/disable AGN power-law component.

  • fit_fe (bool, optional) – Enable/disable FeII components.

  • fit_bc (bool, optional) – Enable/disable Balmer continuum.

  • fit_bal (bool, optional) – If True, append built-in BAL trough custom components for N V, Si IV, C IV, C III], FeLoBAL troughs near 2000/2200 A, and Mg II.

  • fit_poly (bool, optional) – Enable/disable multiplicative polynomial tilt.

  • fit_poly_order (int, optional) – Polynomial order for the multiplicative continuum tilt. Uses coefficients poly_c1 through poly_cN.

  • fit_reddening (bool, optional) – If True, apply a built-in smooth SMC-like reddening curve to the AGN power-law continuum.

  • mask_lya_forest (bool, optional) – If True, mask pixels with rest-frame wavelength below Ly-alpha (1215.67 Angstrom) before fitting.

  • fit_method ({'nuts', 'optax', 'optax+nuts'}, optional) – Fitting backend. 'optax' runs staged MAP optimization with Optax/Adam and stores a point estimate, but does not produce posterior samples. 'nuts' runs NumPyro NUTS directly and stores posterior samples. 'optax+nuts' first runs the staged Optax warm start and then initializes NUTS from that solution; this is the default and is usually the most robust posterior-fitting mode.

  • verbose (bool, optional) – Verbose optimizer output where applicable.

  • fsps_age_grid (sequence of float, optional) – SSP age grid in Gyr.

  • fsps_logzsol_grid (sequence of float, optional) – SSP metallicity grid in log(Z/Zsun).

  • prior_config (dict or None, optional) – Prior/config dictionary. If None, defaults are auto-built.

  • dsps_ssp_fn (str, optional) – Path to DSPS SSP HDF5 template file.

  • nuts_warmup (int, optional) – NUTS warmup and posterior sample counts.

  • nuts_samples (int, optional) – NUTS warmup and posterior sample counts.

  • nuts_chains (int, optional) – Number of MCMC chains.

  • nuts_target_accept (float, optional) – Target accept probability for NUTS.

  • optax_steps (int, optional) – Number of SVI/Optax warm-start steps.

  • optax_lr (float, optional) – Learning rate for Optax Adam.

  • psf_mags (array-like or None, optional) – Optional PSF-aperture photometry magnitudes and 1-sigma errors.

  • psf_mag_errs (array-like or None, optional) – Optional PSF-aperture photometry magnitudes and 1-sigma errors.

  • psf_bands (sequence of str or None, optional) – Band labels for psf_mags. Defaults to SDSS ugriz order.

  • use_psf_phot (bool, optional) – If True, add a PSF-photometry likelihood term and infer PSF/fiber scaling plus host leakage.

  • custom_components (sequence[CustomComponentSpec] or None, optional) – Optional additive continuum components. Use make_custom_component for general user-defined components or make_template_component for a template convenience wrapper.

  • custom_line_components (sequence[CustomLineComponentSpec] or None, optional) – Optional additive emission-line components. Each component provides its own evaluator and is tagged as broad or narrow.

  • kwargs_plot (dict or None, optional) – Extra keyword arguments passed to plot_fig().

run_fsps_numpyro_fit(num_warmup=500, num_samples=1000, num_chains=1, target_accept_prob=0.9, age_grid_gyr=(0.1, 0.3, 1.0, 3.0, 10.0), logzsol_grid=(-1.0, -0.5, 0.0, 0.2), prior_config=None, dsps_ssp_fn='tempdata.h5', use_lines=True, decompose_host=True, fit_pl=True, fit_fe=True, fit_bc=True, fit_poly=False, fit_reddening=False, fit_poly_order=2, psf_mags=None, psf_mag_errs=None, psf_filter_curves=None, use_psf_phot=False, custom_components=None, custom_line_components=None, init_values=None)[source]#

Fit the full model using NUTS MCMC and store posterior summaries.

Parameters:

  • num_warmup (int, optional) – MCMC warmup and posterior sample counts.

  • num_samples (int, optional) – MCMC warmup and posterior sample counts.

  • num_chains (int, optional) – Number of MCMC chains.

  • target_accept_prob (float, optional) – Target acceptance probability for NUTS.

  • age_grid_gyr (sequence of float, optional) – SSP age grid in Gyr.

  • logzsol_grid (sequence of float, optional) – SSP metallicity grid in log(Z/Zsun).

  • prior_config (dict or None, optional) – Prior/config dictionary for model blocks.

  • dsps_ssp_fn (str, optional) – DSPS SSP template HDF5 path.

  • use_lines (bool, optional) – Component toggles for model blocks.

  • decompose_host (bool, optional) – Component toggles for model blocks.

  • fit_pl (bool, optional) – Component toggles for model blocks.

  • fit_fe (bool, optional) – Component toggles for model blocks.

  • fit_bc (bool, optional) – Component toggles for model blocks.

  • fit_poly (bool, optional) – Component toggles for model blocks.

  • fit_reddening (bool, optional) – Component toggles for model blocks.

  • fit_poly_order (int, optional) – Polynomial order for the multiplicative continuum tilt.

  • init_values (dict or None, optional) – Optional initial values for init_to_value.

run_fsps_optax_fit(num_steps=2000, learning_rate=0.01, age_grid_gyr=(0.1, 0.3, 1.0, 3.0, 10.0), logzsol_grid=(-1.0, -0.5, 0.0, 0.2), prior_config=None, dsps_ssp_fn='tempdata.h5', use_lines=True, decompose_host=True, fit_pl=True, fit_fe=True, fit_bc=True, fit_poly=False, fit_reddening=False, fit_poly_order=2, psf_mags=None, psf_mag_errs=None, psf_filter_curves=None, use_psf_phot=False, custom_components=None, custom_line_components=None)[source]#

Fit a MAP approximation using staged SVI with an Optax optimizer.

Parameters:

  • num_steps (int, optional) – Total SVI steps across all stages.

  • learning_rate (float, optional) – Adam learning rate.

  • age_grid_gyr (sequence of float, optional) – SSP age grid in Gyr.

  • logzsol_grid (sequence of float, optional) – SSP metallicity grid in log(Z/Zsun).

  • prior_config (dict or None, optional) – Prior/config dictionary for model blocks.

  • dsps_ssp_fn (str, optional) – DSPS SSP template HDF5 path.

  • use_lines (bool, optional) – Component toggles for model blocks.

  • decompose_host (bool, optional) – Component toggles for model blocks.

  • fit_pl (bool, optional) – Component toggles for model blocks.

  • fit_fe (bool, optional) – Component toggles for model blocks.

  • fit_bc (bool, optional) – Component toggles for model blocks.

  • fit_poly (bool, optional) – Component toggles for model blocks.

  • fit_reddening (bool, optional) – Component toggles for model blocks.

  • fit_poly_order (int, optional) – Polynomial order for the multiplicative continuum tilt.

run_fsps_optax_nuts_fit(optax_steps=2000, optax_learning_rate=0.01, num_warmup=500, num_samples=1000, num_chains=1, target_accept_prob=0.9, age_grid_gyr=(0.1, 0.3, 1.0, 3.0, 10.0), logzsol_grid=(-1.0, -0.5, 0.0, 0.2), prior_config=None, dsps_ssp_fn='tempdata.h5', use_lines=True, decompose_host=True, fit_pl=True, fit_fe=True, fit_bc=True, fit_poly=False, fit_reddening=False, fit_poly_order=2, psf_mags=None, psf_mag_errs=None, psf_filter_curves=None, use_psf_phot=False, custom_components=None, custom_line_components=None)[source]#

Warm-start with Optax MAP, then run NUTS as final inference.

Parameters:

  • optax_steps (int, optional) – Number of SVI/Optax warm-start steps.

  • optax_learning_rate (float, optional) – Learning rate for SVI warm-start.

  • num_warmup (int, optional) – NUTS warmup and posterior sample counts.

  • num_samples (int, optional) – NUTS warmup and posterior sample counts.

  • num_chains (int, optional) – Number of MCMC chains.

  • target_accept_prob (float, optional) – Target acceptance probability for NUTS.

  • age_grid_gyr (sequence of float, optional) – SSP age grid in Gyr.

  • logzsol_grid (sequence of float, optional) – SSP metallicity grid in log(Z/Zsun).

  • prior_config (dict or None, optional) – Prior/config dictionary for model blocks.

  • dsps_ssp_fn (str, optional) – DSPS SSP template HDF5 path.

  • use_lines (bool, optional) – Component toggles for model blocks.

  • decompose_host (bool, optional) – Component toggles for model blocks.

  • fit_pl (bool, optional) – Component toggles for model blocks.

  • fit_fe (bool, optional) – Component toggles for model blocks.

  • fit_bc (bool, optional) – Component toggles for model blocks.

  • fit_poly (bool, optional) – Component toggles for model blocks.

  • fit_reddening (bool, optional) – Component toggles for model blocks.

  • fit_poly_order (int, optional) – Polynomial order for the multiplicative continuum tilt.

reconstruct_posterior_spectrum(wave_out=None, wave_min=2500.0, wave_max=None, n_draws=None, return_components=True)[source]#

Rebuild posterior component draws on a requested rest-frame grid.

Parameters:

  • wave_out (array-like or None, optional) – Explicit rest-frame wavelength grid. If None, build a grid from min(wave_min, self.wave.min()) to wave_max or self.wave.max() using the median native wavelength spacing.

  • wave_min (float or None, optional) – Bounds for the auto-generated grid when wave_out is None.

  • wave_max (float or None, optional) – Bounds for the auto-generated grid when wave_out is None.

  • n_draws (int or None, optional) – If provided, use at most the first n_draws posterior samples.

  • return_components (bool, optional) – If True, include per-component draws and medians in the return value. This includes any fitted custom components.

component_fraction_at_wave(component='host', wave0=2500.0, reference='continuum', reconstruct=False, n_draws=None)[source]#

Return component/reference flux fraction at a requested wavelength.

Parameters:

  • component (str, optional) – Component names. Supported reconstructed names are host, PL, Fe_uv, Fe_op, Balmer_cont, and continuum. Any fitted custom component names are also accepted.

  • reference (str, optional) – Component names. Supported reconstructed names are host, PL, Fe_uv, Fe_op, Balmer_cont, and continuum. Any fitted custom component names are also accepted.

  • wave0 (float, optional) – Rest-frame wavelength in Angstrom.

  • reconstruct (bool, optional) – If True, rebuild posterior components on a grid that reaches wave0. Returns (median, err) from the posterior draws.

  • n_draws (int or None, optional) – Maximum number of posterior draws to use in the reconstruction.

balnicity_index(component_names=None, line_center: float = 1549.06, vmin: float = 3000.0, vmax: float = 25000.0, min_width: float = 2000.0, depth_threshold: float = 0.1, include_line_emission: bool = True, return_details: bool = False)[source]#

Return a simple BALnicity-style index from the summed fitted BAL model.

The BAL model is defined as the sum of selected negative custom components, typically names beginning with bal_. The reference model is the BAL-free AGN continuum, optionally plus the fitted emission-line model. The returned BI uses the standard-style integrand 1 - f_norm / 0.9 over contiguous troughs at least min_width wide and deeper than depth_threshold.

line_profile_from_components(line_key: str) → ndarray[source]#

Build a line-only profile from posterior-median Gaussian components.

Parameters:

line_key (str) – Line-name prefix (for example 'Hb_br').

line_profile_from_draw(draw_index: int, line_key: str) → ndarray[source]#

Build a line-only profile for one posterior draw index.

Parameters:

  • draw_index (int) – Posterior draw index.

  • line_key (str) – Line-name prefix (for example 'Hb_br').

line_props(profile: ndarray, wave: ndarray | None = None) → tuple[float, float][source]#

Return (fwhm_kms, integrated_area) from a line profile.

Parameters:

  • profile (ndarray) – Line profile values.

  • wave (ndarray or None, optional) – Wavelength array. If None, self.wave is used.

line_props_from_profile(wave: ndarray, profile: ndarray) → tuple[float, float][source]#

Compatibility wrapper for line_props().

Parameters:

  • wave (ndarray) – Wavelength array.

  • profile (ndarray) – Line profile values.

save_result(conti_result, conti_result_type, conti_result_name, line_result, line_result_type, line_result_name, save_fits_name)[source]#

Write continuum+line summary table to a pandas CSV file.

Parameters:

  • conti_result (ndarray) – Continuum and line result values.

  • line_result (ndarray) – Continuum and line result values.

  • conti_result_type (ndarray) – Legacy dtype tags (stored but not enforced).

  • line_result_type (ndarray) – Legacy dtype tags (stored but not enforced).

  • conti_result_name (ndarray) – Column names for continuum and line outputs.

  • line_result_name (ndarray) – Column names for continuum and line outputs.

  • save_fits_name (str) – Output basename for CSV.

plot_trace(param_names=None, max_vector_elems=2, save_fig_path=None, save_fig_name=None, show_plot=False)[source]#

Plot posterior trace series for selected parameters.

Parameters:

  • param_names (list[str] | str | None, optional) – Parameter selector. Use 'all' to include all posterior keys.

  • max_vector_elems (int or None, optional) – Maximum number of vector elements to expand per key.

  • save_fig_path (str or None, optional) – Output directory when saving figures. If None, uses self.output_path (or '.' when unset).

  • save_fig_name (str or None, optional) – Output filename override.

  • show_plot (bool, optional) – If True, display the figure interactively with plt.show(). Defaults to False so diagnostics are safe to call in headless terminals.

plot_corner(param_names=None, max_vector_elems=2, bins=30, max_points=5000, save_fig_path=None, save_fig_name=None, show_plot=False)[source]#

Plot posterior projections with corner.corner.

Parameters:

  • param_names (list[str] | str | None, optional) – Parameter selector. Use 'all' to include all posterior keys.

  • max_vector_elems (int or None, optional) – Maximum number of vector elements to expand per key.

  • bins (int, optional) – Histogram bin count.

  • max_points (int, optional) – Maximum posterior draws to plot (subsampled if needed).

  • save_fig_path (str or None, optional) – Output directory when saving figures. If None, uses self.output_path (or '.' when unset).

  • save_fig_name (str or None, optional) – Output filename override.

  • show_plot (bool, optional) – If True, display the figure interactively with plt.show(). Defaults to False so diagnostics are safe to call in headless terminals.

plot_mcmc_diagnostics(do_trace=True, do_corner=True, param_names=None, max_vector_elems=2, corner_bins=30, corner_max_points=2000, save_fig_path=None, show_plot=False)[source]#

Plot trace and/or corner diagnostics in a single convenience call.

Parameters:

  • do_trace (bool, optional) – If True, render trace plot.

  • do_corner (bool, optional) – If True, render corner plot.

  • param_names (list[str] | str | None) – Parameter selector shared by both trace and corner plots. Use 'all' to include all posterior parameters.

  • max_vector_elems (int or None, optional) – Maximum number of vector elements to expand per key.

  • corner_bins (int, optional) – Histogram bin count for corner plot.

  • corner_max_points (int, optional) – Maximum posterior draws to use in corner plot.

  • save_fig_path (str or None, optional) – Output directory when saving figures. If None, uses self.output_path (or '.' when unset).

  • show_plot (bool, optional) – If True, display each enabled diagnostics figure interactively with plt.show(). Defaults to False so diagnostics are safe to call in headless terminals.

plot_fig(save_fig_path=None, broad_fwhm=1200, plot_legend=True, ylims=None, plot_residual=True, show_title=True, plot_1sigma=True, sigma_alpha=0.12, show_plot=True, plot_psf_space=False, plot_intrinsic_powerlaw=False)[source]#

Plot data, model components, line decomposition, and residuals.

Parameters:

  • save_fig_path (str or None, optional) – Output directory when saving figures. If None, uses self.output_path (or '.' when unset).

  • broad_fwhm (float, optional) – Reserved broad-line threshold (kept for compatibility).

  • plot_legend (bool, optional) – If True, draw a legend.

  • ylims (tuple[float, float] or None, optional) – Optional y-axis limits for the spectrum panel.

  • plot_residual (bool, optional) – If True, draw residual panel below the spectrum.

  • show_title (bool, optional) – Reserved title toggle kept for compatibility.

  • plot_1sigma (bool, optional) – If True, draw 16-84% posterior bands for available components.

  • sigma_alpha (float, optional) – Alpha transparency of 1-sigma bands.

  • show_plot (bool, optional) – If True, call plt.show(). If False, figure is created/saved without display.

  • plot_psf_space (bool, optional) – If True, plot the PSF-space model/components. In this mode the residual panel is suppressed because the observed spectrum remains on the fiber scale.

  • plot_intrinsic_powerlaw (bool, optional) – If True, overlay the intrinsic AGN power law before multiplicative polynomial tilt and additive edge corrections.