Fitting Interface

Fitting Interface#

class jaxqsofit.JAXQSOFit(config)[source]#

Bases: object

Config-first spectral fitting interface for quasar spectra.

Parameters:

config (jaxqsofit.config.FitConfig)

__init__(config)[source]#

Initialize a config-first JAXQSOFit spectral fitter.

Parameters:

config (jaxqsofit.config.FitConfig)

property numpyro_samples#

Posterior samples mirrored from the internal posterior state.

property pred_out#

Posterior predictive outputs mirrored from the internal posterior state.

property pred_bands#

Posterior uncertainty bands mirrored from the internal posterior state.

property fig#

Main fitted-spectrum figure mirrored from the internal posterior state.

property trace_fig#

Trace figure mirrored from the internal posterior state.

property corner_fig#

Corner figure mirrored from the internal posterior state.

classmethod from_arrays(*, lam, flux, err=None, z=0.0, ra=None, dec=None, filename=None, output_path=None, wdisp=None, psf_mags=None, psf_mag_errs=None, psf_bands=None)[source]#

Build a config-first fitter from raw arrays.

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

Persist posterior samples plus minimal metadata for compact reloads.

Parameters:

_state (_PosteriorState | None)

save(path=None, *, save_name=None, _state=None)[source]#

Persist posterior samples and fit metadata to a compact bundle.

Parameters:

_state (_PosteriorState | None)

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 JAXQSOFit object.

classmethod load(filename=None, output_path=None, save_name=None, plot_fig=True, plot_diagnostics=True, kwargs_plot=None, diagnostics_kwargs=None)#

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

classmethod load_result(*args, **kwargs)[source]#

Load a posterior bundle and wrap it in a FitResult.

Return type:

FitResult

fit(*, verbose=True, kwargs_plot=None)[source]#

Run preprocessing, inference, persistence, and plotting.

The public API is configuration-first: construct JAXQSOFit with a jaxqsofit.config.FitConfig, then call fit(). Model choices, preprocessing, inference settings, output behavior, PSF recalibration data, and priors all live on the config object.

Parameters:

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

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

Returns:

Result object exposing samples, medians, persistence, and plotting helpers while the fitter keeps mirrored posterior state.

Return type:

FitResult

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=1e-2, 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, plot_init=False)[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.

  • plot_init (bool, optional) – If True, plot and store the stage-1 continuum/host warm-start model before starting the full model stage.

run_fsps_optax_nuts_fit(optax_steps=2000, optax_learning_rate=1e-2, 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, plot_init=False)[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.

  • plot_init (bool, optional) – If True, plot and store the stage-1 Optax warm-start model before starting the full Optax stage and NUTS.

reconstruct_posterior_spectrum(wave_out=None, wave_min=2500.0, wave_max=None, n_draws=None, return_components=True, _state=None)[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.

  • _state (_PosteriorState | None)

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=1549.06, vmin=3000.0, vmax=25000.0, min_width=2000.0, depth_threshold=0.1, include_line_emission=True, return_details=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.

Parameters:

  • line_center (float)

  • vmin (float)

  • vmax (float)

  • min_width (float)

  • depth_threshold (float)

  • include_line_emission (bool)

  • return_details (bool)

line_profile_from_components(line_key)[source]#

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

Parameters:

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

Return type:

ndarray

line_profile_from_draw(draw_index, line_key)[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').

Return type:

ndarray

line_props(profile, wave=None)[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.

Return type:

tuple[float, float]

line_props_from_profile(wave, profile)[source]#

Compatibility wrapper for line_props().

Parameters:

  • wave (ndarray) – Wavelength array.

  • profile (ndarray) – Line profile values.

Return type:

tuple[float, float]

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.

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.

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.

plot_spectrum(**kwargs)[source]#

Plot the fitted spectrum, model components, and residuals.

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.

class jaxqsofit.FitResult(fitter, samples, median, method, summary=None, path=None, figure=None, _state=None)[source]#

Bases: object

High-level result object returned by jaxqsofit fits.

Parameters:

  • fitter (Any)

  • samples (Mapping[str, Any] | None)

  • median (Mapping[str, Any])

  • method (str)

  • summary (Mapping[str, Any] | None)

  • path (Path | None)

  • figure (Any)

  • _state (_PosteriorState | None)

fitter: Any#
samples: Mapping[str, Any] | None#
median: Mapping[str, Any]#
method: str#
summary: Mapping[str, Any] | None = None#
path: Path | None = None#
figure: Any = None#
predict(**kwargs)[source]#

Reconstruct posterior spectral components for this fit.

Return type:

PredictionResult

save(path=None, **kwargs)[source]#

Save the result with the fitter’s native persistence format.

Parameters:

path (str | Path | None)

Return type:

Path

plot_corner(**kwargs)[source]#

Plot posterior samples with the fitter’s corner-plot helper.

plot_trace(**kwargs)[source]#

Plot posterior samples with the fitter’s trace-plot helper.

__init__(fitter, samples, median, method, summary=None, path=None, figure=None, _state=None)#
Parameters:

  • fitter (Any)

  • samples (Mapping[str, Any] | None)

  • median (Mapping[str, Any])

  • method (str)

  • summary (Mapping[str, Any] | None)

  • path (Path | None)

  • figure (Any)

  • _state (_PosteriorState | None)

Return type:

None

class jaxqsofit.PredictionResult(data, fitter)[source]#

Bases: object

Dict-like posterior prediction or reconstruction result.

Parameters:

  • data (Mapping[str, Any])

  • fitter (Any)

data: Mapping[str, Any]#
fitter: Any#
property median: dict[str, Any]#

Median predictive values over the leading posterior axis.

keys()[source]#
items()[source]#
get(key, default=None)[source]#
Parameters:

  • key (str)

  • default (Any)

Return type:

Any

__init__(data, fitter)#
Parameters:

  • data (Mapping[str, Any])

  • fitter (Any)

Return type:

None

jaxqsofit.load_from_samples(output_path=None, save_name=None, plot_fig=True, plot_diagnostics=True, kwargs_plot=None, diagnostics_kwargs=None)#

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