reverse shock correction

Author: YihanWangAstro

README.md

@@ -224,7 +224,7 @@ times = np.logspace(2, 8, 100)
 bands = np.array([1e9, 1e14, 1e17])  
 
 # 3. Calculate the afterglow emission at each time and frequency
-results = model.specific_flux_matrix(times, bands)
+results = model.specific_flux(times, bands)
 
 # 4. Visualize the multi-wavelength light curves
 plt.figure(figsize=(4.8, 3.6),dpi=200)
@@ -269,7 +269,7 @@ frequencies = np.logspace(5, 22, 100)
 epochs = np.array([1e2, 1e3, 1e4, 1e5 ,1e6, 1e7, 1e8])
 
 # 3. Calculate spectra at each epoch
-results = model.specific_flux_matrix(epochs, frequencies)
+results = model.specific_flux(epochs, frequencies)
 
 # 4. Plot broadband spectra at each epoch
 plt.figure(figsize=(4.8, 3.6),dpi=200)

docs/source/examples.rst

@@ -39,7 +39,7 @@ Setting up a simple afterglow model
     bands = np.array([1e9, 1e14, 1e17])  
 
     # Calculate the afterglow emission at each time and frequency
-    results = model.specific_flux_matrix(times, bands)
+    results = model.specific_flux(times, bands)
 
     # Visualize the multi-wavelength light curves
     plt.figure(figsize=(4.8, 3.6),dpi=200)
@@ -48,7 +48,7 @@ Setting up a simple afterglow model
     for i, nu in enumerate(bands):
         exp = int(np.floor(np.log10(nu)))
         base = nu / 10**exp
-    plt.loglog(times, results['syn'][i,:], label=fr'${base:.1f} \times 10^{{{exp}}}$ Hz')
+        plt.loglog(times, results['syn'][i,:], label=fr'${base:.1f} \times 10^{{{exp}}}$ Hz')
 
     plt.xlabel('Time (s)')
     plt.ylabel('Flux Density (erg/cm²/s/Hz)')
@@ -61,7 +61,7 @@ Setting up a simple afterglow model
     epochs = np.array([1e2, 1e3, 1e4, 1e5 ,1e6, 1e7, 1e8])
 
     # Calculate spectra at each epoch
-    results = model.specific_flux_matrix(epochs, frequencies)
+    results = model.specific_flux(epochs, frequencies)
 
     # Plot broadband spectra at each epoch
     plt.figure(figsize=(4.8, 3.6),dpi=200)
@@ -235,84 +235,85 @@ Those profiles are optional and will be set to zero function if not provided.
 Radiation Processes
 -------------------
 
-Synchrotron Self-Compton
+Inverse Compton Cooling
 ^^^^^^^^^^^^^^^^^^^^^^^^    
 
 .. code-block:: python
 
-    from VegasAfterglow import SynchrotronSelfCompton
+    from VegasAfterglow import Radiation
 
-    # Create a model with synchrotron self-Compton
-    ssc = SynchrotronSelfCompton(
-        epsilon_e=0.1,
-        epsilon_B=1e-3,  # Lower magnetization favors IC
-        p=2.2,
-        include_KN=True  # Include Klein-Nishina effects
-    )
-    
-    # Update the model
-    model.set_radiation(ssc)
-    
-    # Calculate over a broader frequency range to capture IC component
-    frequencies_broad = np.logspace(9, 24, 50)  # Radio to gamma-rays
-    
-    # Calculate spectrum at a specific time
-    t_spec = 1e4  # 10,000 seconds
-    spectrum = model.calculate_spectrum(t_spec, frequencies_broad)
-    
-    # Plot the spectrum with components
-    plt.figure(figsize=(10, 6))
-    plt.loglog(frequencies_broad, spectrum, 'b-', label='Total')
-    plt.loglog(frequencies_broad, model.get_synchrotron_spectrum(), 'r--', label='Synchrotron')
-    plt.loglog(frequencies_broad, model.get_ic_spectrum(), 'g--', label='Inverse Compton')
-    
-    plt.xlabel('Frequency (Hz)')
-    plt.ylabel('Flux Density (erg/cm²/s/Hz)')
-    plt.legend()
-    plt.title(f'GRB Afterglow Spectrum at t = {t_spec} s')
-    plt.grid(True, which='both', linestyle='--', alpha=0.5)
-    plt.show()
+    # Create a radiation model with inverse Compton cooling (with Klein-Nishina correction) on synchrotron radiation
+    rad = Radiation(eps_e=1e-1, eps_B=1e-3, p=2.3, IC_cooling=True, KN=False)
 
-Advanced Features
------------------
+    #..other settings
+    model = Model(forward_rad=rad, ...)
+
+Self-Synchrotron Compton Radiation
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+.. code-block:: python
+
+    from VegasAfterglow import Radiation
+
+    # Create a radiation model with self-Compton radiation
+    rad = Radiation(eps_e=1e-1, eps_B=1e-3, p=2.3, SSC=True, KN=True, IC_cooling=True)
+
+    #..other settings
+    model = Model(forward_rad=rad, ...)
+
+    times = np.logspace(2, 8, 200)  
+    bands = np.array([1e9, 1e14, 1e17])  
+
+    results = model.specific_flux(times, bands)
+
+    plt.figure(figsize=(4.8, 3.6),dpi=200)
+
+    # Plot each frequency band 
+    for i, nu in enumerate(bands):
+        exp = int(np.floor(np.log10(nu)))
+        base = nu / 10**exp
+        plt.loglog(times, results['syn'][i,:], label=fr'${base:.1f} \times 10^{{{exp}}}$ Hz')#synchrotron
+        plt.loglog(times, results['IC'][i,:], label=fr'${base:.1f} \times 10^{{{exp}}}$ Hz')#SSC
+
+.. note::
+    (IC_cooling = False, KN = False, SSC = True): The IC radiation is calculated based on synchrotron spectrum without IC cooling.
+    (IC_cooling = True, KN = False, SSC = True): The IC radiation is calculated based on synchrotron spectrum with IC cooling without Klein-Nishina correction.
+    (IC_cooling = True, KN = True, SSC = True): The IC radiation is calculated based on synchrotron spectrum with IC cooling and Klein-Nishina correction.
 
 Reverse Shock
 ^^^^^^^^^^^^^
 
 .. code-block:: python
+    from VegasAfterglow import Radiation
 
-    # Create a model with reverse shock component
-    model_with_rs = Model(
-        jet=jet, 
-        medium=medium, 
-        radiation=radiation,
-        include_reverse_shock=True
-    )
-    
-    # Set reverse shock parameters
-    model_with_rs.set_reverse_shock_parameters(
-        RB=0.1,  # Magnetic field ratio between reverse and forward shock
-        Re=1.0   # Electron energy ratio between reverse and forward shock
-    )
-    
-    # Calculate light curves including reverse shock
-    results_with_rs = model_with_rs.calculate_light_curves(times, frequencies)
+    # Create a radiation model with self-Compton radiation
+    fwd_rad = Radiation(eps_e=1e-1, eps_B=1e-3, p=2.3, SSC=True, KN=True, IC_cooling=True)
+    rvs_rad = Radiation(eps_e=1e-2, eps_B=1e-4, p=2.4, SSC=False, KN=False, IC_cooling=False)
+
+    #..other settings
+    model = Model(forward_rad=fwd_rad, reverse_rad=rvs_rad, ...)
+
+    times = np.logspace(2, 8, 200)  
+    bands = np.array([1e9, 1e14, 1e17])  
+
+    results = model.specific_flux(times, bands)
     
-    # Plot forward vs reverse shock components
-    plt.figure(figsize=(10, 6))
-    for i, nu in enumerate(frequencies):
-        plt.loglog(times, results_with_rs[:, i], label=f'Total {nu:.1e} Hz')
-        plt.loglog(times, model_with_rs.get_forward_shock_light_curve(i), '--', 
-                  label=f'FS {nu:.1e} Hz')
-        plt.loglog(times, model_with_rs.get_reverse_shock_light_curve(i), ':', 
-                  label=f'RS {nu:.1e} Hz')
+    plt.figure(figsize=(4.8, 3.6),dpi=200)
+
+    # Plot each frequency band 
+    for i, nu in enumerate(bands):
+        exp = int(np.floor(np.log10(nu)))
+        base = nu / 10**exp
+        plt.loglog(times, results['syn'][i,:], label=fr'${base:.1f} \times 10^{{{exp}}}$ Hz')
+        plt.loglog(times, results['IC'][i,:], label=fr'${base:.1f} \times 10^{{{exp}}}$ Hz')
+        plt.loglog(times, results['syn_rvs'][i,:], label=fr'${base:.1f} \times 10^{{{exp}}}$ Hz')#reverse shock synchrotron
+
+.. note::
+    You may increase the resolution of the grid to improve the accuracy of the reverse shock synchrotron radiation.
 
-    plt.xlabel('Time (s)')
-    plt.ylabel('Flux Density (erg/cm²/s/Hz)')
-    plt.legend()
-    plt.title('GRB Afterglow with Reverse Shock')
-    plt.grid(True, which='both', linestyle='--', alpha=0.5)
-    plt.show()
+Advanced Features
+-----------------
+
 
 MCMC Parameter Fitting
 ^^^^^^^^^^^^^^^^^^^^^^

docs/source/quickstart.rst

@@ -63,7 +63,7 @@ Now, let's compute and plot multi-wavelength light curves to see how the aftergl
     bands = np.array([1e9, 1e14, 1e17])  
 
     # 3. Calculate the afterglow emission at each time and frequency
-    results = model.specific_flux_matrix(times, bands)
+    results = model.specific_flux(times, bands)
 
     # 4. Visualize the multi-wavelength light curves
     plt.figure(figsize=(4.8, 3.6), dpi=200)
@@ -109,17 +109,16 @@ We can also examine how the broadband spectrum evolves at different times after
     epochs = np.array([1e2, 1e3, 1e4, 1e5, 1e6, 1e7, 1e8])
 
     # 3. Calculate spectra at each epoch
-    results = model.spectra(frequencies, epochs)
+    results = model.specific_flux(epochs, frequencies)
 
     # 4. Plot broadband spectra at each epoch
-    plt.figure(figsize=(4.8, 3.6), dpi=200)
-    colors = plt.cm.viridis(np.linspace(0, 1, len(epochs)))
+    plt.figure(figsize=(4.8, 3.6),dpi=200)
+    colors = plt.cm.viridis(np.linspace(0,1,len(epochs)))
 
     for i, t in enumerate(epochs):
         exp = int(np.floor(np.log10(t)))
         base = t / 10**exp
-        plt.loglog(frequencies, results['syn'][i,:], color=colors[i], 
-                  label=fr'${base:.1f} \times 10^{{{exp}}}$ s')
+        plt.loglog(frequencies, results['syn'][:,i], color=colors[i], label=fr'${base:.1f} \times 10^{{{exp}}}$ s')
 
     # 5. Add vertical lines marking the bands from the light curve plot
     for i, band in enumerate(bands):
@@ -212,7 +211,7 @@ The ``Setups`` class defines the global properties and environment for your mode
     cfg.z = 1.58               # Redshift
 
     # Physical model configuration
-    cfg.medium = "wind"        # Ambient medium: "wind", "ISM" (Interstellar Medium) or "user" (user-defined)
+    cfg.medium = "wind"        # Ambient medium: "wind", "ism" (Interstellar Medium) or "user" (user-defined)
     cfg.jet = "powerlaw"       # Jet structure: "powerlaw", "gaussian", "tophat" or "user" (user-defined)
 
 

include/observer.h

@@ -100,7 +100,7 @@ class Observer {
      * <!-- ************************************************************************************** -->
      */
     template <typename... PhotonGrid>
-    Array point_specific_flux(Array const& t_obs, Array const& nu_obs, PhotonGrid const&... photons);
+    Array specific_flux_series(Array const& t_obs, Array const& nu_obs, PhotonGrid const&... photons);
 
     /**
      * <!-- ************************************************************************************** -->
@@ -327,7 +327,7 @@ MeshGrid Observer::specific_flux(Array const& t_obs, Array const& nu_obs, Photon
                 iterate_to(lg2_t(i, j, 0), lg2_t_obs, t_idx);
 
                 for (size_t k = 0; k < t_grid - 1 && t_idx < t_obs_len; k++) {
-                    Real const lg2_t_hi = lg2_t(i, j, k + 1);
+                    Real lg2_t_hi = lg2_t(i, j, k + 1);
                     if (lg2_t_hi < lg2_t_obs(t_idx)) {
                         continue;
                     } else {
@@ -352,7 +352,7 @@ MeshGrid Observer::specific_flux(Array const& t_obs, Array const& nu_obs, Photon
 }
 
 template <typename... PhotonGrid>
-Array Observer::point_specific_flux(Array const& t_obs, Array const& nu_obs, PhotonGrid const&... photons) {
+Array Observer::specific_flux_series(Array const& t_obs, Array const& nu_obs, PhotonGrid const&... photons) {
     size_t t_obs_len = t_obs.size();
     size_t nu_len = nu_obs.size();
 
@@ -365,21 +365,12 @@ Array Observer::point_specific_flux(Array const& t_obs, Array const& nu_obs, Pho
 
     Array F_nu({t_obs_len}, 0);
 
-    InterpState state;
-
-    // Loop over effective phi and theta grid points.
     for (size_t i = 0; i < eff_phi_grid; i++) {
         for (size_t j = 0; j < theta_grid; j++) {
-            // Skip observation times that are below the grid's start time
-            size_t t_idx = 0;
-            iterate_to(lg2_t(i, j, 0), lg2_t_obs, t_idx);
-
-            for (size_t k = 0; k < t_grid - 1 && t_idx < t_obs_len; k++) {
-                Real const lg2_t_hi = lg2_t(i, j, k + 1);
-                if (lg2_t_hi < lg2_t_obs(t_idx)) {
-                    continue;
-                } else {
-                    for (; t_idx < t_obs_len && lg2_t_obs(t_idx) <= lg2_t_hi; t_idx++) {
+            for (size_t t_idx = 0; t_idx < t_obs_len; t_idx++) {
+                for (size_t k = 0; k < t_grid - 1; k++) {
+                    if (lg2_t_obs(t_idx) <= lg2_t(i, j, k + 1)) {
+                        InterpState state;
                         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));
                         }

pybind/pybind.cpp

@@ -60,8 +60,12 @@ PYBIND11_MODULE(VegasAfterglowC, m) {
              py::arg("jet"), py::arg("medium"), py::arg("observer"), py::arg("forward_rad"),
              py::arg("reverse_rad") = py::none(), py::arg("resolutions") = std::make_tuple(0.5, 3., 5.),
              py::arg("rtol") = 1e-5, py::arg("axisymmetric") = true)
-        .def("specific_flux", &PyModel::specific_flux, py::arg("t"), py::arg("nu"))
-        .def("specific_flux_matrix", &PyModel::specific_flux_matrix, py::arg("t"), py::arg("nu"));
+        .def("specific_flux", &PyModel::specific_flux, py::arg("t"), py::arg("nu"),
+             py::call_guard<py::gil_scoped_release>())
+        .def("specific_flux_series", &PyModel::specific_flux_series, py::arg("t"), py::arg("nu"),
+             py::call_guard<py::gil_scoped_release>())
+        .def("specific_flux_sorted_series", &PyModel::specific_flux_sorted_series, py::arg("t"), py::arg("nu"),
+             py::call_guard<py::gil_scoped_release>());
 
     //========================================================================================================
     //                                 MCMC bindings

pybind/pymodel.cpp

@@ -88,7 +88,7 @@ Medium PyWind(Real A_star) {
 
 void PyModel::single_shock_emission(Shock const& shock, Coord const& coord, Array const& t_obs, Array const& nu_obs,
                                     Observer& obs, PyRadiation rad, FluxDict& flux_dict, std::string suffix,
-                                    bool return_trace) {
+                                    bool serilized) {
     obs.observe(coord, shock, obs_setup.lumi_dist, obs_setup.z);
 
     auto syn_e = generate_syn_electrons(shock);
@@ -106,20 +106,20 @@ void PyModel::single_shock_emission(Shock const& shock, Coord const& coord, Arra
     if (rad.SSC) {
         auto IC_ph = generate_IC_photons(syn_e, syn_ph, rad.KN);
 
-        if (return_trace) {
-            flux_dict["IC" + suffix] = obs.point_specific_flux(t_obs, nu_obs, IC_ph) / unit::flux_den_cgs;
+        if (serilized) {
+            flux_dict["IC" + suffix] = obs.specific_flux_series(t_obs, nu_obs, IC_ph) / unit::flux_den_cgs;
         } else {
             flux_dict["IC" + suffix] = obs.specific_flux(t_obs, nu_obs, IC_ph) / unit::flux_den_cgs;
         }
     }
-    if (return_trace) {
-        flux_dict["syn" + suffix] = obs.point_specific_flux(t_obs, nu_obs, syn_ph) / unit::flux_den_cgs;
+    if (serilized) {
+        flux_dict["syn" + suffix] = obs.specific_flux_series(t_obs, nu_obs, syn_ph) / unit::flux_den_cgs;
     } else {
         flux_dict["syn" + suffix] = obs.specific_flux(t_obs, nu_obs, syn_ph) / unit::flux_den_cgs;
     }
 }
 
-auto PyModel::compute_specific_flux(Array const& t_obs, Array const& nu_obs, bool return_trace) -> FluxDict {
+auto PyModel::compute_specific_flux(Array const& t_obs, Array const& nu_obs, bool serilized) -> FluxDict {
     Coord coord = auto_grid(jet, t_obs, this->theta_w, obs_setup.theta_obs, obs_setup.z, phi_resol, theta_resol,
                             t_resol, axisymmetric);
 
@@ -130,36 +130,50 @@ auto PyModel::compute_specific_flux(Array const& t_obs, Array const& nu_obs, boo
     if (!rvs_rad_opt) {
         auto fwd_shock = generate_fwd_shock(coord, medium, jet, fwd_rad.rad, rtol);
 
-        single_shock_emission(fwd_shock, coord, t_obs, nu_obs, observer, fwd_rad, flux_dict, "", return_trace);
+        single_shock_emission(fwd_shock, coord, t_obs, nu_obs, observer, fwd_rad, flux_dict, "", serilized);
 
         return flux_dict;
     } else {
         auto rvs_rad = *rvs_rad_opt;
         auto [fwd_shock, rvs_shock] = generate_shock_pair(coord, medium, jet, fwd_rad.rad, rvs_rad.rad, rtol);
 
-        single_shock_emission(fwd_shock, coord, t_obs, nu_obs, observer, fwd_rad, flux_dict, "", return_trace);
+        single_shock_emission(fwd_shock, coord, t_obs, nu_obs, observer, fwd_rad, flux_dict, "", serilized);
 
-        single_shock_emission(rvs_shock, coord, t_obs, nu_obs, observer, rvs_rad, flux_dict, "_rvs", return_trace);
+        single_shock_emission(rvs_shock, coord, t_obs, nu_obs, observer, rvs_rad, flux_dict, "_rvs", serilized);
 
         return flux_dict;
     }
 }
 
-auto PyModel::specific_flux(PyArray const& t, PyArray const& nu) -> FluxDict {
+auto PyModel::specific_flux_sorted_series(PyArray const& t, PyArray const& nu) -> FluxDict {
     Array t_obs = t * unit::sec;
     Array nu_obs = nu * unit::Hz;
-    bool return_trace = true;
+    bool serilized = true;
 
     if (t_obs.size() != nu_obs.size()) {
         throw std::invalid_argument(
             "time and frequency arrays must have the same size\n"
-            "If you intend to get matrix output, use `specific_flux_matrix` instead");
+            "If you intend to get matrix-like output, use the generic `specific_flux` instead");
     }
 
-    return compute_specific_flux(t_obs, nu_obs, return_trace);
+    return compute_specific_flux(t_obs, nu_obs, serilized);
 }
 
-auto PyModel::specific_flux_matrix(PyArray const& t, PyArray const& nu) -> FluxDict {
+auto PyModel::specific_flux_series(PyArray const& t, PyArray const& nu) -> FluxDict {
+    Array t_obs = t * unit::sec;
+    Array nu_obs = nu * unit::Hz;
+    bool serilized = true;
+
+    if (t_obs.size() != nu_obs.size()) {
+        throw std::invalid_argument(
+            "time and frequency arrays must have the same size\n"
+            "If you intend to get matrix-like output, use the generic `specific_flux` instead");
+    }
+
+    return compute_specific_flux(t_obs, nu_obs, serilized);
+}
+
+auto PyModel::specific_flux(PyArray const& t, PyArray const& nu) -> FluxDict {
     Array t_obs = t * unit::sec;
     Array nu_obs = nu * unit::Hz;
     bool return_trace = false;