Source code for MapLines.tools.mcmc

#!/usr/bin/env python
"""
MapLines.tools.mcmc
===================

MCMC utilities for spectral fitting in MapLines.

This module provides helper routines to run Markov Chain Monte Carlo
sampling and to evaluate posterior model realizations from the resulting
chains.

The main functionality includes:

- execution of the ``emcee`` ensemble sampler
- optional multiprocessing support
- burn-in and production runs
- extraction of representative model realizations from posterior samples

These routines are used by the fitting functions in
``MapLines.tools.line_fit`` and rely on the spectral models defined in
``MapLines.tools.models``.

Notes
-----
The sampler is designed for flexible use with the posterior probability
functions defined in ``MapLines.tools.priors`` and supports both serial
and multiprocessing execution modes.
"""
import numpy as np
import emcee
import MapLines.tools.models as mod

[docs] def mcmc(p0,nwalkers,niter,ndim,lnprob,data,verbose=False,multi=True,tim=False,ncpu=10): """ Run an MCMC sampler using the ``emcee`` ensemble sampler. This function performs parameter estimation for spectral models using a Markov Chain Monte Carlo (MCMC) approach. It supports both serial and multiprocessing execution. Parameters ---------- p0 : array_like Initial positions of the walkers in parameter space. nwalkers : int Number of walkers used by the sampler. niter : int Number of iterations in the production run. ndim : int Number of model parameters. lnprob : callable Log-probability function used by the sampler. data : tuple Additional arguments passed to ``lnprob``. verbose : bool, optional Print progress messages. multi : bool, optional Enable multiprocessing. tim : bool, optional Measure execution time. ncpu : int, optional Number of CPUs used when multiprocessing. Returns ------- sampler : emcee.EnsembleSampler MCMC sampler object. pos : ndarray Final walker positions. prob : ndarray Log-probability values. state : object Internal sampler state. Notes ----- This function performs a burn-in phase before the production run. """ if tim: import time if multi: from multiprocessing import Pool from multiprocessing import cpu_count ncput=cpu_count() if ncpu > ncput: ncpu=ncput if ncpu == 0: ncpu=None with Pool(ncpu) as pool: #with Pool() as pool: sampler = emcee.EnsembleSampler(nwalkers, ndim, lnprob, args=data,pool=pool) if tim: start = time.time() if verbose: print("Running burn-in...") p0, _, _ = sampler.run_mcmc(p0, 1000) sampler.reset() if verbose: print("Running production...") pos, prob, state = sampler.run_mcmc(p0, niter) if tim: end = time.time() multi_time = end - start print("Multiprocessing took {0:.1f} seconds".format(multi_time)) else: sampler = emcee.EnsembleSampler(nwalkers, ndim, lnprob, args=data) if tim: start = time.time() if verbose: print("Running burn-in...") p0, _, _ = sampler.run_mcmc(p0, 1000) sampler.reset() if verbose: print("Running production...") pos, prob, state = sampler.run_mcmc(p0, niter) if tim: end = time.time() serial_time = end - start print("Serial took {0:.1f} seconds".format(serial_time)) return sampler, pos, prob, state
[docs] def sample_walkers(nsamples,flattened_chain,waves0,fac0,facN0,velfac0,velfacN0,fwhfac0,fwhfacN0, names0,n_lines,vals,x=0,skew=False,lorentz=False,outflow=False,powlaw=False,feii=False, data=None): """ Generate model realizations from the posterior MCMC samples. This function randomly samples walkers from the flattened MCMC chain and computes the corresponding spectral models. Parameters ---------- nsamples : int Number of models to draw from the chain. flattened_chain : ndarray Flattened MCMC chain. waves0 : array_like Rest wavelengths of spectral lines. fac0, facN0 : array_like Flux normalization factors. velfac0, velfacN0 : array_like Velocity scaling factors. fwhfac0, fwhfacN0 : array_like Line-width scaling factors. names0 : list Names of the spectral components. n_lines : int Number of emission lines. vals : dict Additional model parameters. x : array_like, optional Spectral axis. skew, lorentz, outflow, powlaw, feii : bool Flags controlling model components. data : optional Additional data passed to the model. Returns ------- med_model : ndarray Median model spectrum. spread : ndarray Standard deviation of the sampled models. Notes ----- The function uses the spectral model defined in ``MapLines.tools.models``. :contentReference[oaicite:0]{index=0} """ models = [] draw = np.floor(np.random.uniform(0,len(flattened_chain),size=nsamples)).astype(int) thetas = flattened_chain[draw] for i in thetas: modt = mod.line_model(i, waves0, fac0, facN0, velfac0, velfacN0, fwhfac0, fwhfacN0, names0, n_lines, vals, x=x, skew=skew, lorentz=lorentz, outflow=outflow, powlaw=powlaw, feii=feii, data=data) models.append(modt) spread = np.std(models,axis=0) med_model = np.median(models,axis=0) return med_model,spread