code cleanup

new mcmc move strategy
model.flux unit fix.
RS: hardcut Gamma

Author: YihanWangAstro

VegasAfterglow/runner.py

@@ -159,7 +159,7 @@ def validate_parameters(self, param_defs: Sequence[ParamDef]) -> None:
 
         # Check reverse shock parameters
         if self.config.rvs_shock:
-            rvs_required = {"p_r", "eps_e_r", "eps_B_r"}
+            rvs_required = {"p_r", "eps_e_r", "eps_B_r", "tau"}
             missing_rvs = rvs_required - param_names
             if missing_rvs:
                 missing_params.extend(

VegasAfterglow/sampler.py

@@ -7,7 +7,6 @@
 
 import emcee
 import numpy as np
-from emcee.moves import DEMove, DESnookerMove
 
 from .types import FitResult, ModelParams, ObsData, Setups, VegasMC
 
@@ -90,22 +89,23 @@ def run(
         """
         Run coarse MCMC + optional stretch-move refinement at higher resolution.
         """
-        # 1) configure coarse grid
+
         cfg = self._make_cfg(base_cfg, *resolution)
 
-        # 2) prepare log-prob
         log_prob = _log_prob(
             data, cfg, self.to_params, self.pl, self.pu, self.model_cls
         )
 
-        # 3) initialize walker positions
         spread = 0.05 * (self.pu - self.pl)
         pos = self.init + spread * np.random.randn(self.nwalkers, self.ndim)
         pos = np.clip(pos, self.pl + 1e-8, self.pu - 1e-8)
 
-        # 4) default moves
         if moves is None:
-            moves = [(DEMove(), 0.8), (DESnookerMove(), 0.2)]
+            moves = [
+                (emcee.moves.StretchMove(a=2.0), 0.6),
+                (emcee.moves.DEMove(gamma=None), 0.3),
+                (emcee.moves.DESnookerMove(gamma=None), 0.1),
+            ]
 
         logger.info(
             "Running coarse MCMC at resolution %s for %d steps",
@@ -118,18 +118,14 @@ def run(
             )
             sampler.run_mcmc(pos, total_steps, progress=True)
 
-        # 5) extract & filter
         burn = int(burn_frac * total_steps)
         chain = sampler.get_chain(discard=burn, thin=thin)
         logp = sampler.get_log_prob(discard=burn, thin=thin)
         chain, logp, _ = self._filter_bad_walkers(chain, logp)
 
-        # 6) flatten & find top k fits
         flat_chain = chain.reshape(-1, self.ndim)
         flat_logp = logp.reshape(-1)
 
-        # Find top k unique parameter combinations
-        # Sort by log prob (descending)
         sorted_idx = np.argsort(flat_logp)[::-1]
 
         # Round parameters to avoid floating point precision issues
@@ -151,7 +147,6 @@ def run(
             top_k_log_probs[-1],
         )
 
-        # 8) return FitResult
         return FitResult(
             samples=chain,
             log_probs=logp,

docs/source/examples.rst

@@ -97,7 +97,7 @@ Suppose you want to calculate the flux at specific time-frequency pairs (t_i, nu
     # Define observing frequencies (must be the same length as times)
     bands = np.logspace(9,17, 200)
 
-    results = model.flux_density(times, bands) #times array could be random order
+    results = model.flux_density(times, bands) #times array must be in ascending order
 
     # the returned results is a FluxDict object with arrays of the same shape as the input times and bands.
 
@@ -222,7 +222,7 @@ User-Defined Medium
     # Define a custom density profile function
     def density(phi, theta, r):# r in cm, phi and theta in radians
         return mp # n_ism =  1 cm^-3
-        #return whatever density profile (cm^-3) you want as a function of phi, theta, and r
+        #return whatever density profile (g*cm^-3) you want as a function of phi, theta, and r
 
     # Create a user-defined medium
     medium = Medium(rho=density)

docs/source/mcmc_fitting.rst

@@ -93,7 +93,7 @@ For large datasets or densely sampled observations, using all available data poi
     flux_err_dense = 0.1 * flux_dense
 
     # Subsample using logarithmic screening
-    # This selects ~50-100 representative points across 5 decades in time
+    # This selects ~5*5=25 representative points across 5 decades in time
     indices = ObsData.logscale_screen(t_dense, num_order=5)
 
     # Add only the selected subset

include/jet.h

@@ -342,7 +342,7 @@ namespace math {
             if (theta <= theta_c) {
                 return 0.;
             } else {
-                return height / (1 + fast_pow(theta / theta_c, k));
+                return height * fast_pow(theta / theta_c, -k);
             }
         };
     }
@@ -352,7 +352,7 @@ namespace math {
             if (theta <= theta_c) {
                 return 1.;
             } else {
-                return height / (1 + fast_pow(theta / theta_c, k)) + 1;
+                return height * fast_pow(theta / theta_c, -k) + 1;
             }
         };
     }
@@ -362,7 +362,7 @@ namespace math {
             if (theta <= theta_c) {
                 return height_c;
             } else {
-                return height_w / (1 + fast_pow(theta / theta_c, k));
+                return height_w * fast_pow(theta / theta_c, -k);
             }
         };
     }
@@ -372,7 +372,7 @@ namespace math {
             if (theta <= theta_c) {
                 return height_c + 1;
             } else {
-                return height_w / (1 + fast_pow(theta / theta_c, k)) + 1;
+                return height_w * fast_pow(theta / theta_c, -k) + 1;
             }
         };
     }

include/observer.h

@@ -70,8 +70,8 @@ class Observer {
      * @return Array of specific flux values at each observation time
      * <!-- ************************************************************************************** -->
      */
-    template <typename... PhotonGrid>
-    Array specific_flux(Array const& t_obs, Real nu_obs, PhotonGrid&... photons);
+    template <typename PhotonGrid>
+    Array specific_flux(Array const& t_obs, Real nu_obs, PhotonGrid& photons);
 
     /**
      * <!-- ************************************************************************************** -->
@@ -86,8 +86,8 @@ class Observer {
      *          relativistic beaming and cosmological effects.
      * <!-- ************************************************************************************** -->
      */
-    template <typename... PhotonGrid>
-    MeshGrid specific_flux(Array const& t_obs, Array const& nu_obs, PhotonGrid&... photons);
+    template <typename PhotonGrid>
+    MeshGrid specific_flux(Array const& t_obs, Array const& nu_obs, PhotonGrid& photons);
 
     /**
      * <!-- ************************************************************************************** -->
@@ -99,8 +99,8 @@ class Observer {
      * @return Array of specific flux values at each observation time for a single observed frequency
      * <!-- ************************************************************************************** -->
      */
-    template <typename... PhotonGrid>
-    Array specific_flux_series(Array const& t_obs, Array const& nu_obs, PhotonGrid&... photons);
+    template <typename PhotonGrid>
+    Array specific_flux_series(Array const& t_obs, Array const& nu_obs, PhotonGrid& photons);
 
     /**
      * <!-- ************************************************************************************** -->
@@ -115,8 +115,8 @@ class Observer {
      * @return Array of integrated flux values at each observation time
      * <!-- ************************************************************************************** -->
      */
-    template <typename... PhotonGrid>
-    Array flux(Array const& t_obs, Array const& band_freq, PhotonGrid&... photons);
+    template <typename PhotonGrid>
+    Array flux(Array const& t_obs, Array const& band_freq, PhotonGrid& photons);
 
     /**
      * <!-- ************************************************************************************** -->
@@ -130,8 +130,8 @@ class Observer {
      * @return 2D grid of spectra (frequency × time)
      * <!-- ************************************************************************************** -->
      */
-    template <typename... PhotonGrid>
-    MeshGrid spectra(Array const& freqs, Array const& t_obs, PhotonGrid&... photons);
+    template <typename PhotonGrid>
+    MeshGrid spectra(Array const& freqs, Array const& t_obs, PhotonGrid& photons);
 
     /**
      * <!-- ************************************************************************************** -->
@@ -145,8 +145,8 @@ class Observer {
      * @return Array containing the spectrum at the given time
      * <!-- ************************************************************************************** -->
      */
-    template <typename... PhotonGrid>
-    Array spectrum(Array const& freqs, Real t_obs, PhotonGrid&... photons);
+    template <typename PhotonGrid>
+    Array spectrum(Array const& freqs, Real t_obs, PhotonGrid& photons);
 
     /**
      * <!-- ************************************************************************************** -->
@@ -253,9 +253,8 @@ class Observer {
      * @return True if both lower and upper boundaries are valid for interpolation, false otherwise
      * <!-- ************************************************************************************** -->
      */
-    template <typename... PhotonGrid>
-    bool set_boundaries(InterpState& state, size_t i, size_t j, size_t k, Real log2_nu,
-                        PhotonGrid&... photons) noexcept;
+    template <typename PhotonGrid>
+    bool set_boundaries(InterpState& state, size_t i, size_t j, size_t k, Real log2_nu, PhotonGrid& photons) noexcept;
 };
 
 //========================================================================================================
@@ -277,9 +276,9 @@ inline void iterate_to(Real value, Array const& arr, size_t& it) noexcept {
     }
 }
 
-template <typename... PhotonGrid>
+template <typename PhotonGrid>
 bool Observer::set_boundaries(InterpState& state, size_t i, size_t j, size_t k, Real lg2_nu_obs,
-                              PhotonGrid&... photons) noexcept {
+                              PhotonGrid& photons) noexcept {
     if (state.last_hi == k + 1 && state.last_lg2_nu == lg2_nu_obs) {
         if (!std::isfinite(state.slope)) {
             return false;
@@ -298,12 +297,12 @@ bool Observer::set_boundaries(InterpState& state, size_t i, size_t j, size_t k,
         state.lg2_I_nu_lo = state.lg2_I_nu_hi;
     } else {
         Real lg2_nu_lo = lg2_one_plus_z + lg2_nu_obs - lg2_doppler(i, j, k);
-        state.lg2_I_nu_lo = 3 * lg2_doppler(i, j, k) + (photons(eff_i, j, k).compute_log2_I_nu(lg2_nu_lo) + ...) +
-                            lg2_emission_area(i, j, k);
+        state.lg2_I_nu_lo =
+            3 * lg2_doppler(i, j, k) + photons(eff_i, j, k).compute_log2_I_nu(lg2_nu_lo) + lg2_emission_area(i, j, k);
     }
 
     Real lg2_nu_hi = lg2_one_plus_z + lg2_nu_obs - lg2_doppler(i, j, k + 1);
-    state.lg2_I_nu_hi = 3 * lg2_doppler(i, j, k + 1) + (photons(eff_i, j, k + 1).compute_log2_I_nu(lg2_nu_hi) + ...) +
+    state.lg2_I_nu_hi = 3 * lg2_doppler(i, j, k + 1) + photons(eff_i, j, k + 1).compute_log2_I_nu(lg2_nu_hi) +
                         lg2_emission_area(i, j, k + 1);
 
     state.slope = (state.lg2_I_nu_hi - state.lg2_I_nu_lo) / lg2_t_ratio;
@@ -317,8 +316,8 @@ bool Observer::set_boundaries(InterpState& state, size_t i, size_t j, size_t k,
     return true;
 }
 
-template <typename... PhotonGrid>
-MeshGrid Observer::specific_flux(Array const& t_obs, Array const& nu_obs, PhotonGrid&... photons) {
+template <typename PhotonGrid>
+MeshGrid Observer::specific_flux(Array const& t_obs, Array const& nu_obs, PhotonGrid& photons) {
     size_t t_obs_len = t_obs.size();
     size_t nu_len = nu_obs.size();
 
@@ -341,7 +340,7 @@ MeshGrid Observer::specific_flux(Array const& t_obs, Array const& nu_obs, Photon
                     if (lg2_t_hi < lg2_t_obs(t_idx)) {
                         continue;
                     } else {
-                        if (set_boundaries(state, i, j, k, lg2_nu[l], photons...)) {
+                        if (set_boundaries(state, i, j, k, lg2_nu[l], photons)) {
                             for (; t_idx < t_obs_len && lg2_t_obs(t_idx) <= lg2_t_hi; t_idx++) {
                                 F_nu(l, t_idx) += interpolate(state, i, j, k, lg2_t_obs(t_idx));
                             }
@@ -361,8 +360,8 @@ MeshGrid Observer::specific_flux(Array const& t_obs, Array const& nu_obs, Photon
     return F_nu;
 }
 
-template <typename... PhotonGrid>
-Array Observer::specific_flux_series(Array const& t_obs, Array const& nu_obs, PhotonGrid&... photons) {
+template <typename PhotonGrid>
+Array Observer::specific_flux_series(Array const& t_obs, Array const& nu_obs, PhotonGrid& photons) {
     size_t t_obs_len = t_obs.size();
     size_t nu_len = nu_obs.size();
 
@@ -388,7 +387,7 @@ Array Observer::specific_flux_series(Array const& t_obs, Array const& nu_obs, Ph
                     k++;
                     continue;
                 } else {
-                    if (set_boundaries(state, i, j, k, lg2_nu[t_idx], photons...)) {
+                    if (set_boundaries(state, i, j, k, lg2_nu(t_idx), photons)) {
                         F_nu(t_idx) += interpolate(state, i, j, k, lg2_t_obs(t_idx));
                     }
                     t_idx++;
@@ -404,25 +403,25 @@ Array Observer::specific_flux_series(Array const& t_obs, Array const& nu_obs, Ph
     return F_nu;
 }
 
-template <typename... PhotonGrid>
-Array Observer::specific_flux(Array const& t_obs, Real nu_obs, PhotonGrid&... photons) {
-    return xt::view(specific_flux(t_obs, Array({nu_obs}), photons...), 0);
+template <typename PhotonGrid>
+Array Observer::specific_flux(Array const& t_obs, Real nu_obs, PhotonGrid& photons) {
+    return xt::view(specific_flux(t_obs, Array({nu_obs}), photons), 0);
 }
 
-template <typename... PhotonGrid>
-Array Observer::spectrum(Array const& freqs, Real t_obs, PhotonGrid&... photons) {
-    return xt::view(spectra(freqs, Array({t_obs}), photons...), 0);
+template <typename PhotonGrid>
+Array Observer::spectrum(Array const& freqs, Real t_obs, PhotonGrid& photons) {
+    return xt::view(spectra(freqs, Array({t_obs}), photons), 0);
 }
 
-template <typename... PhotonGrid>
-MeshGrid Observer::spectra(Array const& freqs, Array const& t_obs, PhotonGrid&... photons) {
-    return xt::transpose(specific_flux(t_obs, freqs, photons...));
+template <typename PhotonGrid>
+MeshGrid Observer::spectra(Array const& freqs, Array const& t_obs, PhotonGrid& photons) {
+    return xt::transpose(specific_flux(t_obs, freqs, photons));
 }
 
-template <typename... PhotonGrid>
-Array Observer::flux(Array const& t_obs, Array const& band_freq, PhotonGrid&... photons) {
+template <typename PhotonGrid>
+Array Observer::flux(Array const& t_obs, Array const& band_freq, PhotonGrid& photons) {
     Array nu_obs = boundary_to_center_log(band_freq);
-    MeshGrid F_nu = specific_flux(t_obs, nu_obs, photons...);
+    MeshGrid F_nu = specific_flux(t_obs, nu_obs, photons);
     Array flux({t_obs.size()}, 0);
     for (size_t i = 0; i < nu_obs.size(); ++i) {
         Real dnu = band_freq(i + 1) - band_freq(i);

pybind/mcmc.cpp

@@ -83,8 +83,8 @@ double FluxData::estimate_chi2() const {
     double chi_square = 0;
     for (size_t i = 0; i < t.size(); ++i) {
         double error = Fv_err(i);
-        if (error == 0)
-            continue;
+        //if (error == 0)
+        //    continue;
         double diff = Fv_obs(i) - Fv_model(i);
         chi_square += weights(i) * (diff * diff) / (error * error);
     }
@@ -95,14 +95,15 @@ double MultiBandData::estimate_chi2() const {
     double chi_square = 0;
     for (size_t i = 0; i < times.size(); ++i) {
         double error = errors(i);
-        if (error == 0)
-            continue;
+        //if (error == 0)
+        //    continue;
         double diff = fluxes(i) - model_fluxes(i);
         chi_square += weights(i) * (diff * diff) / (error * error);
     }
     for (auto& d : flux_data) {
         chi_square += d.estimate_chi2();
     }
+
     return chi_square;
 }