SDSS fit: J102839.11+450009.4 with tight Fe FWHM priors

This notebook fetches the SDSS spectrum for J102839.11+450009.4, overrides Fe UV/optical FWHM priors, and fits with decompose_host=False.

Current default: fit_method="optax" runs staged SVI/Optax MAP optimization (continuum warm start, then full model). fit_method="optax+nuts" uses that staged MAP point to initialize NUTS.

import numpy as np
from astroquery.sdss import SDSS
from astropy.coordinates import SkyCoord
from astropy import units as u

import jaxqsofit
from jaxqsofit import QSOFit

coord = SkyCoord.from_name("J102839.11+450009.4")
xid = SDSS.query_region(coord, spectro=True, radius=5 * u.arcsec)
sp = SDSS.get_spectra(matches=xid[:1])[0]

tb = sp[1].data
lam = np.asarray(10 ** tb["loglam"], dtype=float)
flux = np.asarray(tb["flux"], dtype=float)
ivar = np.asarray(tb["ivar"], dtype=float)

err = np.full_like(flux, 1e-6)
m = np.isfinite(ivar) & (ivar > 0)
err[m] = 1.0 / np.sqrt(ivar[m])

z = float(sp[2].data["z"][0])
ra = float(coord.ra.deg)
dec = float(coord.dec.deg)
plateid = int(sp[0].header.get("plateid", 0))
mjd = int(sp[0].header.get("mjd", 0))
fiberid = int(sp[0].header.get("fiberid", 0))
sdss_filename = f"{plateid:04d}-{mjd}-{fiberid:04d}"

lam.size, z, sdss_filename

flux_for_priors = flux * (1.0 + z)
prior_config = jaxqsofit.defaults.build_default_prior_config(flux=flux_for_priors)
# Tight Fe priors make this notebook emphasize the iron-template behavior.
prior_config["log_Fe_uv_FWHM"] = {"loc": np.log(100.0), "scale": 0.1}
prior_config["log_Fe_op_FWHM"] = {"loc": np.log(100.0), "scale": 0.1}

line_table = prior_config["line"]["table"]

# Increase ngauss by +1 for existing broad Mg II and Hbeta entries
#for name in ("Hb_br"): # "MgII_br",
#    row = next((r for r in line_table if r.get("linename") == name), None)
#    if row is not None:
#        row["ngauss"] = int(row.get("ngauss", 1)) + 1

# quick check
[(r["linename"], r["ngauss"]) for r in line_table if r.get("linename") in ("MgII_br", "Hb_br")]

q = QSOFit(
    lam=lam,
    flux=flux,
    err=err,
    z=z,
    ra=ra,
    dec=dec,
    filename=sdss_filename,
    output_path=".",
)

q.fit(
    deredden=True,
    fit_method="optax+nuts",
    fit_lines=True,
    decompose_host=False,
    fit_pl=True,
    fit_fe=True,
    fit_bc=False,
    fit_poly=True,
    prior_config=prior_config,
    dsps_ssp_fn="../tempdata.h5",
    optax_steps=600,
    optax_lr=1e-2,
    nuts_warmup=50,
    nuts_samples=50,
    nuts_chains=1,
    plot_fig=True,
    save_fig=False,
    save_result=False,
)

q.plot_mcmc_diagnostics(
    param_names='Fe_uv_FWHM',
    do_trace=True,
    do_corner=True,
    max_vector_elems=2,
    corner_bins=25,
    corner_max_points=1500,
)