redback 1.0.2__py3-none-any.whl → 1.0.31__py3-none-any.whl

redback/transient/afterglow.py

@@ -3,6 +3,7 @@ from __future__ import annotations
 import os
 from os.path import join
 from typing import Union
+import re
 
 import numpy as np
 import pandas as pd
@@ -26,7 +27,8 @@ class Afterglow(Transient):
             flux_density_err: np.ndarray = None, magnitude: np.ndarray = None, magnitude_err: np.ndarray = None,
             redshift: float = np.nan, photon_index: float = np.nan, frequency: np.ndarray = None,
             bands: np.ndarray = None, system: np.ndarray = None, active_bands: Union[np.ndarray, str] = 'all',
-            use_phase_model: bool = False, optical_data: bool = False, **kwargs: None) -> None:
+            plotting_order: Union[np.ndarray, str] = None, use_phase_model: bool = False,
+            optical_data: bool = False, **kwargs: None) -> None:
 
         """
         This is a general constructor for the Afterglow class. Note that you only need to give data corresponding to
@@ -81,6 +83,8 @@ class Afterglow(Transient):
         :type bands: np.ndarray, optional
         :param active_bands: List or array of active bands to be used in the analysis. Use all available bands if 'all' is given.
         :type active_bands: Union[list, np.ndarray]
+        :param plotting_order: Order in which to plot the bands/and how unique bands are stored.
+        :type plotting_order: Union[np.ndarray, str], optional
         :param kwargs:
             Additional classes that can be customised to fulfil the truncation on flux to luminosity conversion:
             FluxToLuminosityConverter: Conversion class to convert fluxes to luminosities.
@@ -99,7 +103,8 @@ class Afterglow(Transient):
                          Lum50=Lum50, Lum50_err=Lum50_err, flux=flux, flux_err=flux_err, flux_density=flux_density,
                          flux_density_err=flux_density_err, use_phase_model=use_phase_model, optical_data=optical_data,
                          magnitude=magnitude, magnitude_err=magnitude_err, frequency=frequency, redshift=redshift,
-                         photon_index=photon_index, system=system, bands=bands, active_bands=active_bands, **kwargs)
+                         photon_index=photon_index, system=system, bands=bands, active_bands=active_bands,
+                         plotting_order=plotting_order, **kwargs)
         self._set_data()
         self._set_photon_index()
         self._set_t90()
@@ -306,9 +311,13 @@ class Afterglow(Transient):
         :return: The cleaned string converted into a float.
         :rtype: float
         """
-        for r in ["PL", "CPL", ",", "C", "~", " ", 'Gemini:emission', '()']:
-            string = string.replace(r, "")
-        return float(string)
+        try:
+            for r in ["PL", "CPL", ",", "C", "~", " ", 'Gemini:emission', '()']:
+                string = string.replace(r, "")
+                new_float = float(string)
+        except ValueError:
+            new_float = float(re.findall("\d+\.\d+", string)[0])
+        return new_float
 
     def analytical_flux_to_luminosity(self) -> None:
         """Converts flux to luminosity using the analytical method."""

redback/transient/kilonova.py

@@ -20,7 +20,8 @@ class Kilonova(OpticalTransient):
             flux_density: np.ndarray = None, flux_density_err: np.ndarray = None, magnitude: np.ndarray = None,
             magnitude_err: np.ndarray = None, redshift: float = np.nan, photon_index: float = np.nan,
             bands: np.ndarray = None, system: np.ndarray = None, active_bands: Union[np.ndarray, str] = 'all',
-            use_phase_model: bool = False, optical_data: bool = True, **kwargs: None) -> None:
+            plotting_order: Union[np.ndarray, str] = None, use_phase_model: bool = False,
+            optical_data: bool = True, **kwargs: None) -> None:
         """This is a general constructor for the Kilonova class. Note that you only need to give data corresponding to
         the data mode you are using. For luminosity data provide times in the rest frame, if using a phase model
         provide time in MJD, else use the default time (observer frame).
@@ -74,6 +75,8 @@ class Kilonova(OpticalTransient):
         :type bands: np.ndarray, optional
         :param active_bands: List or array of active bands to be used in the analysis. Use all available bands if 'all' is given.
         :type active_bands: Union[list, np.ndarray]
+        :param plotting_order: Order in which to plot the bands/and how unique bands are stored.
+        :type plotting_order: Union[np.ndarray, str], optional
         :param kwargs: Additional callables:
                         bands_to_frequency: Conversion function to convert a list of bands to frequencies. Use
                         redback.utils.bands_to_frequency if not given.
@@ -85,7 +88,7 @@ class Kilonova(OpticalTransient):
                          Lum50_err=Lum50_err, flux_density=flux_density, flux_density_err=flux_density_err,
                          magnitude=magnitude, magnitude_err=magnitude_err, data_mode=data_mode, name=name, bands=bands,
                          system=system, active_bands=active_bands, use_phase_model=use_phase_model, redshift=redshift,
-                         photon_index=photon_index, optical_data=optical_data, **kwargs)
+                         photon_index=photon_index, optical_data=optical_data, plotting_order=plotting_order, **kwargs)
         self.directory_structure = redback.get_data.directory.open_access_directory_structure(
             transient=name, transient_type="kilonova")
         self._set_data()

redback/transient/prompt.py

@@ -8,7 +8,7 @@ import pandas as pd
 from typing import Union
 
 from redback.get_data.utils import get_batse_trigger_from_grb
-from redback.get_data.directory import swift_prompt_directory_structure
+from redback.get_data.directory import swift_prompt_directory_structure, batse_prompt_directory_structure
 from redback.transient.transient import Transient
 
 dirname = os.path.dirname(__file__)
@@ -61,7 +61,12 @@ class PromptTimeSeries(Transient):
         self.channels = channels
         self.instrument = instrument
         self._set_data()
-        self.directory_structure = swift_prompt_directory_structure(grb=self.name, bin_size=self.bin_size)
+        if self.instrument == "batse":
+            self.directory_structure = batse_prompt_directory_structure(grb=self.name)
+        elif self.instrument == "swift":
+            self.directory_structure = swift_prompt_directory_structure(grb=self.name, bin_size=self.bin_size)
+        else:
+            raise ValueError("Instrument must be either 'batse' or 'swift'.")
 
     @classmethod
     def from_batse_grb_name(
@@ -95,8 +100,9 @@ class PromptTimeSeries(Transient):
         :rtype tuple:
         """
         name = f"GRB{name.lstrip('GRB')}"
-        directory_structure = swift_prompt_directory_structure(grb=name)
-        _time_series_data = np.genfromtxt(directory_structure.processed_file_path, delimiter=",")[1:]
+        directory_structure = batse_prompt_directory_structure(grb=name)
+        path = directory_structure.directory_path + name.lstrip('GRB') + '_BATSE_lc.csv'
+        _time_series_data = np.genfromtxt(path, delimiter=",")[1:]
 
         bin_left = _time_series_data[:, 0]
         bin_right = _time_series_data[:, 1]
@@ -150,6 +156,8 @@ class PromptTimeSeries(Transient):
         :type kwargs: None
         """
         plt.step(self.time, self.counts / self.bin_size)
+        plt.ylabel('Counts')
+        plt.xlabel('Time')
         plt.show()
         plt.clf()
 
@@ -177,6 +185,8 @@ class PromptTimeSeries(Transient):
         plt.clf()
         plt.step(self.time, self.counts / self.bin_size)
         plt.plot(self.time, model(self.time, **dict(posterior.iloc[-1])))
+        plt.ylabel('Counts')
+        plt.xlabel('Time')
         plt.show()
         plt.clf()
 

redback/transient/supernova.py

@@ -15,7 +15,8 @@ class Supernova(OpticalTransient):
             flux_density: np.ndarray = None, flux_density_err: np.ndarray = None, magnitude: np.ndarray = None,
             magnitude_err: np.ndarray = None, redshift: float = np.nan, photon_index: float = np.nan,
             bands: np.ndarray = None, system: np.ndarray = None, active_bands: Union[np.ndarray, str] = 'all',
-            use_phase_model: bool = False, optical_data:bool = True, **kwargs: None) -> None:
+            plotting_order: Union[np.ndarray, str] = None, use_phase_model: bool = False,
+            optical_data:bool = True, **kwargs: None) -> None:
         """
         This is a general constructor for the Supernova class. Note that you only need to give data corresponding to
         the data mode you are using. For luminosity data provide times in the rest frame, if using a phase model
@@ -69,6 +70,8 @@ class Supernova(OpticalTransient):
         :type bands: np.ndarray, optional
         :param active_bands: List or array of active bands to be used in the analysis. Use all available bands if 'all' is given.
         :type active_bands: Union[list, np.ndarray]
+        :param plotting_order: Order in which to plot the bands/and how unique bands are stored.
+        :type plotting_order: Union[np.ndarray, str], optional
         :param kwargs: Additional callables:
                         bands_to_frequency: Conversion function to convert a list of bands to frequencies. Use
                         redback.utils.bands_to_frequency if not given.
@@ -80,7 +83,8 @@ class Supernova(OpticalTransient):
                          Lum50_err=Lum50_err, flux_density=flux_density, flux_density_err=flux_density_err,
                          magnitude=magnitude, magnitude_err=magnitude_err, data_mode=data_mode, name=name,
                          use_phase_model=use_phase_model, bands=bands, system=system, active_bands=active_bands,
-                         redshift=redshift, photon_index=photon_index, **kwargs)
+                         redshift=redshift, photon_index=photon_index, optical_data=optical_data,
+                         plotting_order=plotting_order, **kwargs)
         self.directory_structure = redback.get_data.directory.open_access_directory_structure(
             transient=name, transient_type="supernova")
         self._set_data()

redback/transient/tde.py

@@ -15,7 +15,8 @@ class TDE(OpticalTransient):
             flux_density: np.ndarray = None, flux_density_err: np.ndarray = None, magnitude: np.ndarray = None,
             magnitude_err: np.ndarray = None, redshift: float = np.nan, photon_index: float = np.nan,
             bands: np.ndarray = None, system: np.ndarray = None, active_bands: Union[np.ndarray, str] = 'all',
-            use_phase_model: bool = False, optical_data:bool = True, **kwargs: None) -> None:
+            plotting_order: Union[np.ndarray, str] = None, use_phase_model: bool = False,
+            optical_data:bool = True, **kwargs: None) -> None:
         """
 
         :param name: Name of the transient.
@@ -66,6 +67,8 @@ class TDE(OpticalTransient):
         :type bands: np.ndarray, optional
         :param active_bands: List or array of active bands to be used in the analysis. Use all available bands if 'all' is given.
         :type active_bands: Union[list, np.ndarray]
+        :param plotting_order: Order in which to plot the bands/and how unique bands are stored.
+        :type plotting_order: Union[np.ndarray, str], optional
         :param kwargs: Additional callables:
                         bands_to_frequency: Conversion function to convert a list of bands to frequencies. Use
                         redback.utils.bands_to_frequency if not given.
@@ -77,7 +80,7 @@ class TDE(OpticalTransient):
                          magnitude=magnitude, magnitude_err=magnitude_err, data_mode=data_mode, name=name,
                          use_phase_model=use_phase_model, optical_data=optical_data, bands=bands,
                          system=system, active_bands=active_bands,redshift=redshift,
-                         photon_index=photon_index, **kwargs)
+                         photon_index=photon_index, plotting_order=plotting_order, **kwargs)
         self.directory_structure = redback.get_data.directory.open_access_directory_structure(
             transient=self.name, transient_type="tidal_disruption_event")
         self._set_data()

redback/transient/transient.py

@@ -38,7 +38,7 @@ class Transient(object):
             ttes: np.ndarray = None, bin_size: float = None, redshift: float = np.nan, data_mode: str = None,
             name: str = '', photon_index: float = np.nan, use_phase_model: bool = False,
             optical_data: bool = False, frequency: np.ndarray = None, system: np.ndarray = None, bands: np.ndarray = None,
-            active_bands: Union[np.ndarray, str] = None, **kwargs: None) -> None:
+            active_bands: Union[np.ndarray, str] = None, plotting_order: Union[np.ndarray, str] = None, **kwargs: None) -> None:
         """This is a general constructor for the Transient class. Note that you only need to give data corresponding to
         the data mode you are using. For luminosity data provide times in the rest frame, if using a phase model
         provide time in MJD, else use the default time (observer frame).
@@ -98,6 +98,8 @@ class Transient(object):
         :param active_bands: List or array of active bands to be used in the analysis.
                              Use all available bands if 'all' is given.
         :type active_bands: Union[list, np.ndarray], optional
+        :param plotting_order: Order in which to plot the bands/and how unique bands are stored.
+        :type plotting_order: Union[np.ndarray, str], optional
         :param kwargs: Additional callables:
                        bands_to_frequency: Conversion function to convert a list of bands to frequencies.
                                            Use redback.utils.bands_to_frequency if not given.
@@ -142,6 +144,7 @@ class Transient(object):
         self.name = name
         self.use_phase_model = use_phase_model
         self.optical_data = optical_data
+        self.plotting_order = plotting_order
 
         self.meta_data = None
         self.photon_index = photon_index
@@ -179,7 +182,7 @@ class Transient(object):
     @classmethod
     def from_lasair_data(
             cls, name: str, data_mode: str = "magnitude", active_bands: Union[np.ndarray, str] = 'all',
-            use_phase_model: bool = False) -> Transient:
+            use_phase_model: bool = False, plotting_order: Union[np.ndarray, str] = None) -> Transient:
         """Constructor method to built object from LASAIR data.
 
         :param name: Name of the transient.
@@ -189,6 +192,8 @@ class Transient(object):
         :param active_bands: Sets active bands based on array given.
                              If argument is 'all', all unique bands in `self.bands` will be used.
         :type active_bands: Union[np.ndarray, str]
+        :param plotting_order: Order in which to plot the bands/and how unique bands are stored.
+        :type plotting_order: Union[np.ndarray, str], optional
         :param use_phase_model: Whether to use a phase model.
         :type use_phase_model: bool, optional
 
@@ -214,12 +219,12 @@ class Transient(object):
         return cls(name=name, data_mode=data_mode, time=time_days, time_err=None, time_mjd=time_mjd,
                    flux_density=flux_density, flux_density_err=flux_density_err, magnitude=magnitude,
                    magnitude_err=magnitude_err, flux=flux, flux_err=flux_err, bands=bands, active_bands=active_bands,
-                   use_phase_model=use_phase_model, optical_data=True)
+                   use_phase_model=use_phase_model, optical_data=True, plotting_order=plotting_order)
 
     @classmethod
     def from_simulated_optical_data(
             cls, name: str, data_mode: str = "magnitude", active_bands: Union[np.ndarray, str] = 'all',
-            use_phase_model: bool = False) -> Transient:
+            plotting_order: Union[np.ndarray, str] = None, use_phase_model: bool = False) -> Transient:
         """Constructor method to built object from SimulatedOpticalTransient.
 
         :param name: Name of the transient.
@@ -229,6 +234,8 @@ class Transient(object):
         :param active_bands: Sets active bands based on array given.
                              If argument is 'all', all unique bands in `self.bands` will be used.
         :type active_bands: Union[np.ndarray, str]
+        :param plotting_order: Order in which to plot the bands/and how unique bands are stored.
+        :type plotting_order: Union[np.ndarray, str], optional
         :param use_phase_model: Whether to use a phase model.
         :type use_phase_model: bool, optional
 
@@ -250,7 +257,7 @@ class Transient(object):
         return cls(name=name, data_mode=data_mode, time=time_days, time_err=None, time_mjd=time_mjd,
                    flux_density=flux_density, flux_density_err=flux_density_err, magnitude=magnitude,
                    magnitude_err=magnitude_err, flux=flux, flux_err=flux_err, bands=bands, active_bands=active_bands,
-                   use_phase_model=use_phase_model, optical_data=True)
+                   use_phase_model=use_phase_model, optical_data=True, plotting_order=plotting_order)
 
     @property
     def _time_attribute_name(self) -> str:
@@ -501,7 +508,10 @@ class Transient(object):
         :return: All bands that we get from the data, eliminating all duplicates.
         :rtype: np.ndarray
         """
-        return np.unique(self.bands)
+        if self.plotting_order is not None:
+            return self.plotting_order
+        else:
+            return np.unique(self.bands)
 
     @property
     def unique_frequencies(self) -> np.ndarray:
@@ -789,7 +799,8 @@ class OpticalTransient(Transient):
             flux_density_err: np.ndarray = None, magnitude: np.ndarray = None, magnitude_err: np.ndarray = None,
             redshift: float = np.nan, photon_index: float = np.nan, frequency: np.ndarray = None,
             bands: np.ndarray = None, system: np.ndarray = None, active_bands: Union[np.ndarray, str] = 'all',
-            use_phase_model: bool = False, optical_data:bool = True, **kwargs: None) -> None:
+            plotting_order: Union[np.ndarray, str] = None, use_phase_model: bool = False,
+            optical_data:bool = True, **kwargs: None) -> None:
         """This is a general constructor for the Transient class. Note that you only need to give data corresponding to
         the data mode you are using. For luminosity data provide times in the rest frame, if using a phase model
         provide time in MJD, else use the default time (observer frame).
@@ -839,6 +850,8 @@ class OpticalTransient(Transient):
         :param active_bands: List or array of active bands to be used in the analysis.
                              Use all available bands if 'all' is given.
         :type active_bands: Union[list, np.ndarray], optional
+        :param plotting_order: Order in which to plot the bands/and how unique bands are stored.
+        :type plotting_order: Union[np.ndarray, str], optional
         :param use_phase_model: Whether we are using a phase model.
         :type use_phase_model: bool, optional
         :param optical_data: Whether we are fitting optical data, useful for plotting.
@@ -855,14 +868,15 @@ class OpticalTransient(Transient):
                          flux=flux, flux_err=flux_err, redshift=redshift, photon_index=photon_index,
                          flux_density=flux_density, flux_density_err=flux_density_err, magnitude=magnitude,
                          magnitude_err=magnitude_err, data_mode=data_mode, name=name,
-                         use_phase_model=use_phase_model, optical_data=optical_data, system=system, bands=bands,
+                         use_phase_model=use_phase_model, optical_data=optical_data,
+                         system=system, bands=bands, plotting_order=plotting_order,
                          active_bands=active_bands, **kwargs)
         self.directory_structure = None
 
     @classmethod
     def from_open_access_catalogue(
             cls, name: str, data_mode: str = "magnitude", active_bands: Union[np.ndarray, str] = 'all',
-            use_phase_model: bool = False) -> OpticalTransient:
+            plotting_order: Union[np.ndarray, str] = None, use_phase_model: bool = False) -> OpticalTransient:
         """Constructor method to built object from Open Access Catalogue
 
         :param name: Name of the transient.
@@ -873,6 +887,8 @@ class OpticalTransient(Transient):
             Sets active bands based on array given.
             If argument is 'all', all unique bands in `self.bands` will be used.
         :type active_bands: Union[np.ndarray, str]
+        :param plotting_order: Order in which to plot the bands/and how unique bands are stored.
+        :type plotting_order: Union[np.ndarray, str], optional
         :param use_phase_model: Whether to use a phase model.
         :type use_phase_model: bool, optional
 
@@ -890,7 +906,8 @@ class OpticalTransient(Transient):
         return cls(name=name, data_mode=data_mode, time=time_days, time_err=None, time_mjd=time_mjd,
                    flux_density=flux_density, flux_density_err=flux_density_err, magnitude=magnitude,
                    magnitude_err=magnitude_err, bands=bands, system=system, active_bands=active_bands,
-                   use_phase_model=use_phase_model, optical_data=True, flux=flux, flux_err=flux_err)
+                   use_phase_model=use_phase_model, optical_data=True, flux=flux, flux_err=flux_err,
+                   plotting_order=plotting_order)
 
     @property
     def event_table(self) -> str:

redback/transient_models/afterglow_models.py

@@ -1,6 +1,6 @@
 from astropy.cosmology import Planck18 as cosmo  # noqa
 from inspect import isfunction
-from redback.utils import logger, citation_wrapper, calc_ABmag_from_flux_density, lambda_to_nu
+from redback.utils import logger, citation_wrapper, calc_ABmag_from_flux_density, lambda_to_nu, bands_to_frequency
 from redback.constants import day_to_s, speed_of_light, solar_mass, proton_mass, electron_mass, sigma_T
 from redback.sed import get_correct_output_format_from_spectra
 import astropy.units as uu
@@ -285,16 +285,16 @@ class RedbackAfterglows():
         te = np.arcsin(cs / (self.cc * (G2 - 1.) ** 0.5))  # equivalent form for angle due to spreading
         # prepare ex and OmG in this function
         if self.is_expansion:
-            ex = te  # expansion
+            ex = te / (G2) # expansion
             fac = 0.5 * latstep
-            OmG = rotstep * (np.cos(thi - fac) - np.cos(ex + thi + fac))  # equivalent form for linear spacing
+            OmG = rotstep * (np.cos(thi - fac) - np.cos(ex/self.res + thi + fac))  # equivalent form for linear spacing
         else:
-            te = np.ones(te.size)  # no expansion
+            ex = np.ones(te.size)  # no expansion
             fac = 0.5 * latstep
             OmG = rotstep * (np.cos(thi - fac) - np.cos(thi + fac))  # equivalent form for linear spacing
         # prepare R
         size = G.size
-        exponent = ((1 - np.cos(te[0])) / (1 - np.cos(te[:size]))) ** 0.5
+        exponent = ((1 - np.cos(latstep + ex[0])) / (1 - np.cos(latstep + ex[:size]))) ** (1/2)
         R = ((3. - k) * Ne[:size] / (self.fourpi * n)) ** (1. / (3. - k))
         R[1:] = np.diff(R) * exponent[1:size] ** (1. / (3. - k))
         R = np.cumsum(R)
@@ -688,150 +688,10 @@ def _get_kn_dynamics(n0, Eej, Mej):
     # convert from lab frame to observer time. Approximate expression acounting for radial+azimuthal time-of-flight effect
     # (eq. 26 from Nava et al. 2013; see also Waxman 1997)
     tobs = R / (Gamma ** 2 * (1.0 + beta) * speed_of_light) + np.insert(
-        integrate.cumtrapz((1.0 - beta_sh) / beta_sh, x=R) / speed_of_light, 0, 0.0)
+        integrate.cumulative_trapezoid((1.0 - beta_sh) / beta_sh, x=R) / speed_of_light, 0, 0.0)
 
     return t, R, beta, Gamma, eden, tobs, beta_sh, Gamma_sh
 
-@citation_wrapper('https://ui.adsabs.harvard.edu/abs/2022MNRAS.516.4949S/abstract')
-def kilonova_afterglow_redback(time, redshift, loge0, mej, logn0, logepse, logepsb, p,
-                             **kwargs):
-    """
-    Calculate the afterglow emission from a kilonova remnant, following the model of Sarin et al. 2022.
-    This model was modified by Nikhil Sarin following code provided by Ben Margalit.
-
-    :param time: time in observer frame (days) in observer frame
-    :param redshift: source redshift
-    :param loge0: log10 of the initial kinetic energy of the ejecta (erg)
-    :param mej: ejecta mass (solar masses)
-    :param logn0: log10 of the circumburst density (cm^-3)
-    :param logepse: log10 of the fraction of shock energy given to electrons
-    :param logepsb: log10 of the fraction of shock energy given to magnetic field
-    :param p: power-law index of the electron energy distribution
-    :param kwargs: Additional keyword arguments
-    :param zeta_e: fraction of electrons participating in diffusive shock acceleration. Default is 1.
-    :param output_format: Whether to output flux density or AB mag
-    :param frequency: frequency in Hz for the flux density calculation
-    :param cosmology: Cosmology to use for luminosity distance calculation. Defaults to Planck18. Must be a astropy.cosmology object.
-    :return: flux density or AB mag. Note this is going to give the monochromatic magnitude at the effective frequency for the band.
-        For a proper calculation of the magntitude use the sed variant models.
-    """
-    Eej = 10 ** loge0
-    Mej = mej * solar_mass
-    cosmology = kwargs.get('cosmology', cosmo)
-    epsilon_e = 10 ** logepse
-    epsilon_B = 10 ** logepsb
-    n0 = 10 ** logn0
-    zeta_e = kwargs.get('zeta_e', 1.0)
-    qe = 4.803e-10
-
-    dl = cosmology.luminosity_distance(redshift).cgs.value
-    # calculate blast-wave dynamics
-    t, R, beta, Gamma, eden, tobs, beta_sh, Gamma_sh = _get_kn_dynamics(n0=n0, Eej=Eej, Mej=Mej)
-
-    # shock-amplified magnetic field, minimum & cooling Lorentz factors
-    B = (8.0 * np.pi * epsilon_B * eden) ** 0.5
-    gamma_m = 1.0 + (epsilon_e / zeta_e) * ((p - 2.0) / (p - 1.0)) * (proton_mass / electron_mass) * (Gamma - 1.0)
-    gamma_c = 6.0 * np.pi * electron_mass * speed_of_light / (sigma_T * Gamma * t * B ** 2)
-
-    # number of emitting electrons, where zeta_DN is an approximate smooth interpolation between the "standard"
-    # and deep-Newtonian regime discussed by Sironi & Giannios (2013)
-    zeta_DN = (gamma_m - 1.0) / gamma_m
-    Ne = zeta_DN * zeta_e * (4.0 * np.pi / 3.0) * R ** 3 * n0
-
-    # LOS approximation
-    mu = 1.0
-    blueshift = Gamma * (1.0 - beta * mu)
-
-    frequency = kwargs['frequency']
-    if isinstance(frequency, float):
-        frequency = np.ones(len(time)) * frequency
-    fnu_func = {}
-    for nu in frequency:
-        Fnu_opt_thin = _pnu_synchrotron(nu * blueshift * (1.0 + redshift), B, gamma_m, gamma_c, Ne, p) * (1.0 + redshift) / (
-                    4.0 * np.pi * dl ** 2 * blueshift)
-
-        # correct for synchrotron self-absorption (approximate treatment, correct up to factors of order unity)
-        Fnu_opt_thick = Gamma * (8 * np.pi ** 2 * (nu * blueshift * (1.0 + redshift)) ** 2 / speed_of_light ** 2) * R ** 2 * (
-                    1.0 / 3.0) * electron_mass * speed_of_light ** 2 * np.maximum(gamma_m, (
-                    2 * np.pi * electron_mass * speed_of_light * nu * blueshift * (1.0 + redshift) / (qe * B)) ** 0.5) * (1.0 + redshift) / (
-                                    4.0 * np.pi * dl ** 2 * blueshift)
-        # new prescription:
-        Fnu = Fnu_opt_thick * (1e0 - np.exp(-Fnu_opt_thin / Fnu_opt_thick))
-        # add brute-force optically-thin case to avoid roundoff error in 1e0-np.exp(-x) term (above) when x->0
-        Fnu[Fnu == 0.0] = Fnu_opt_thin[Fnu == 0.0]
-
-        fnu_func[nu] = interp1d(tobs/day_to_s, Fnu, bounds_error=False, fill_value='extrapolate')
-
-    # interpolate onto actual observed frequency and time values
-    flux_density = []
-    for freq, tt in zip(frequency, time):
-        flux_density.append(fnu_func[freq](tt))
-
-    fmjy = np.array(flux_density) / 1e-26
-    if kwargs['output_format'] == 'flux_density':
-        return fmjy
-    elif kwargs['output_format'] == 'magnitude':
-        return calc_ABmag_from_flux_density(fmjy).value
-
-@citation_wrapper('https://ui.adsabs.harvard.edu/abs/2011Natur.478...82N/abstract')
-def kilonova_afterglow_nakarpiran(time, redshift, loge0, mej, logn0, logepse, logepsb, p, **kwargs):
-    """
-    A kilonova afterglow model based on Nakar & Piran 2011
-
-    :param time: time in days in the observer frame
-    :param redshift: source redshift
-    :param loge0: initial kinetic energy in erg of ejecta
-    :param mej: mass of ejecta in solar masses
-    :param logn0: log10 of the number density of the circumburst medium in cm^-3
-    :param logepse: log10 of the fraction of energy given to electrons
-    :param logepsb: log10 of the fraction of energy given to the magnetic field
-    :param p: electron power law index
-    :param kwargs: Additional keyword arguments
-    :param output_format: Whether to output flux density or AB mag
-    :param frequency: frequency in Hz for the flux density calculation
-    :param cosmology: Cosmology to use for luminosity distance calculation. Defaults to Planck18. Must be a astropy.cosmology object.
-    :return: flux density or AB mag. Note this is going to give the monochromatic magnitude at the effective frequency for the band.
-        For a proper calculation of the magntitude use the sed variant models.
-    :return:
-    """
-    Eej = 10 ** loge0
-    Mej = mej * solar_mass
-    Gamma0 = 1.0 + Eej / (Mej * speed_of_light ** 2)
-    vej = speed_of_light * (1.0 - Gamma0 ** (-2)) ** 0.5
-    cosmology = kwargs.get('cosmology', cosmo)
-    epsilon_e = 10 ** logepse
-    epsilon_B = 10 ** logepsb
-    n0 = 10 ** logn0
-    dl = cosmology.luminosity_distance(redshift).cgs.value
-
-    # in days
-    t_dec = 30 * (Eej / 1e49) ** (1.0 / 3.0) * (n0 / 1e0) ** (-1.0 / 3.0) * (vej / speed_of_light) ** (-5.0 / 3.0)
-
-    fnu_dec_dict = {}
-    fnu_func = {}
-    temp_time = np.linspace(0.1, 100, 200) * t_dec
-    frequency = kwargs['frequency']
-    if isinstance(frequency, float):
-        frequency = np.ones(len(time)) * frequency
-    for freq in frequency:
-        # Eq. 11 in Nakar & Piran 2011 (in Mjy)
-        fnu_dec_dict[freq] = 0.3 * (Eej / 1e49) * n0 ** (0.25 * (p + 1)) * (epsilon_B / 1e-1) ** (0.25 * (p + 1)) * (
-                epsilon_e / 1e-1) ** (p - 1) * (vej / speed_of_light) ** (0.5 * (5 * p - 7)) * (freq / 1.4e9) ** (
-                          -0.5 * (p - 1)) * (dl / 1e27) ** (-2)
-        fnu = fnu_dec_dict[freq] * (temp_time / t_dec) ** 3
-        fnu[temp_time > t_dec] = fnu_dec_dict[freq] * (temp_time[temp_time > t_dec] / t_dec) ** (-0.3 * (5 * p - 7))
-        fnu_func[freq] = interp1d(temp_time, fnu, bounds_error=False, fill_value='extrapolate')
-
-    # interpolate onto actual observed frequency and time values
-    flux_density = []
-    for freq, tt in zip(frequency, time):
-        flux_density.append(fnu_func[freq](tt))
-    fmjy = flux_density * uu.mJy
-    if kwargs['output_format'] == 'flux_density':
-        return fmjy.value
-    elif kwargs['output_format'] == 'magnitude':
-        return calc_ABmag_from_flux_density(fmjy.value).value
-
 @citation_wrapper('redback, https://ui.adsabs.harvard.edu/abs/2018MNRAS.481.2581L/abstract')
 def tophat_redback(time, redshift, thv, loge0, thc, logn0, p, logepse, logepsb, g0, xiN, **kwargs):
     """
@@ -2253,4 +2113,108 @@ def afterglow_models_sed(time, **kwargs):
                                                       spectra=spectra, lambda_array=lambda_observer_frame,
                                                       **kwargs)
 
+@citation_wrapper('https://ui.adsabs.harvard.edu/abs/2024ApJS..273...17W/abstract')
+def jetsimpy_tophat(time, redshift, thv, loge0, thc, nism, A, p, logepse, logepsb, g0, **kwargs):
+    """
+    A tophat jet model from jetsimpy
+    :param time: time in days in observer frame
+    :param redshift: source redshift
+    :param thv: viewing angle in radians
+    :param loge0: log10 on axis isotropic equivalent energy
+    :param thc: half width of jet core/jet opening angle in radians
+    :param nism: number density of ISM in cm^-3 (ntot = A * (r / 1e17)^-2 + nism (cm^-3))
+    :param A: wind density scale (ntot = A * (r / 1e17)^-2 + nism (cm^-3))
+    :param p: electron distribution power law index.
+    :param logepse: log10 fraction of thermal energy in electrons
+    :param logepsb: log10 fraction of thermal energy in magnetic field
+    :param g0: initial lorentz factor
+    :param kwargs: Additional keyword arguments
+    :param output_format: Whether to output flux density or AB mag
+    :param frequency: frequency in Hz for the flux density calculation
+    :param cosmology: Cosmology to use for luminosity distance calculation. Defaults to Planck18. Must be a astropy.cosmology object.
+    :return: flux density or AB mag. Note this is going to give the monochromatic magnitude at the effective frequency for the band.
+    """
+    import jetsimpy #Can not use models unless jetsimpy is downloaded
+    time = time * day_to_s
+    cosmology = kwargs.get('cosmology', cosmo)
+    dl = cosmology.luminosity_distance(redshift).cgs.value
+    P = dict(Eiso = 10 ** loge0, lf = g0, theta_c = thc, n0 = nism, A = A, eps_e = 10 ** logepse, eps_b = 10 ** logepsb, p = p, theta_v = thv, d = dl*3.24078e-25, z = redshift) #make a param dict
+    if kwargs['output_format'] == 'flux_density':
+        frequency = kwargs['frequency']
+        flux_density = jetsimpy.FluxDensity_tophat(time, frequency, P)
+        return flux_density   
+    else:
+        frequency = bands_to_frequency(kwargs['bands'])       
+        flux_density = jetsimpy.FluxDensity_tophat(time, frequency, P)
+        return calc_ABmag_from_flux_density(flux_density).value
+
+@citation_wrapper('https://ui.adsabs.harvard.edu/abs/2024ApJS..273...17W/abstract')
+def jetsimpy_gaussian(time, redshift, thv, loge0, thc, nism, A, p, logepse, logepsb, g0, **kwargs):
+    """
+    A gaussian jet model from jetsimpy
+    :param time: time in days in observer frame
+    :param redshift: source redshift
+    :param thv: viewing angle in radians
+    :param loge0: log10 on axis isotropic equivalent energy
+    :param thc: half width of jet core/jet opening angle in radians
+    :param nism: number density of ISM in cm^-3 (ntot = A * (r / 1e17)^-2 + nism (cm^-3))
+    :param A: wind density scale (ntot = A * (r / 1e17)^-2 + nism (cm^-3))
+    :param p: electron distribution power law index.
+    :param logepse: log10 fraction of thermal energy in electrons
+    :param logepsb: log10 fraction of thermal energy in magnetic field
+    :param g0: initial lorentz factor
+    :param kwargs: Additional keyword arguments
+    :param output_format: Whether to output flux density or AB mag
+    :param frequency: frequency in Hz for the flux density calculation
+    :param cosmology: Cosmology to use for luminosity distance calculation. Defaults to Planck18. Must be a astropy.cosmology object.
+    :return: flux density or AB mag. Note this is going to give the monochromatic magnitude at the effective frequency for the band.
+    """
+    import jetsimpy #Can not use models unless jetsimpy is downloaded
+    time = time * day_to_s
+    cosmology = kwargs.get('cosmology', cosmo)
+    dl = cosmology.luminosity_distance(redshift).cgs.value
+    P = dict(Eiso = 10 ** loge0, lf = g0, theta_c = thc, n0 = nism, A = A, eps_e = 10 ** logepse, eps_b = 10 ** logepsb, p = p, theta_v = thv, d = dl*3.24078e-25, z = redshift) #make a param dict
+    if kwargs['output_format'] == 'flux_density':
+        frequency = kwargs['frequency']
+        flux_density = jetsimpy.FluxDensity_gaussian(time, frequency, P)
+        return flux_density   
+    else:
+        frequency = bands_to_frequency(kwargs['bands'])       
+        flux_density = jetsimpy.FluxDensity_gaussian(time, frequency, P)
+        return calc_ABmag_from_flux_density(flux_density).value
 
+@citation_wrapper('https://ui.adsabs.harvard.edu/abs/2024ApJS..273...17W/abstract')
+def jetsimpy_powerlaw(time, redshift, thv, loge0, thc, nism, A, p, logepse, logepsb, g0, s, **kwargs):
+    """
+    A power-law jet model from jetsimpy
+    :param time: time in days in observer frame
+    :param redshift: source redshift
+    :param thv: viewing angle in radians
+    :param loge0: log10 on axis isotropic equivalent energy
+    :param thc: half width of jet core/jet opening angle in radians
+    :param nism: number density of ISM in cm^-3 (ntot = A * (r / 1e17)^-2 + nism (cm^-3))
+    :param A: wind density scale (ntot = A * (r / 1e17)^-2 + nism (cm^-3))
+    :param p: electron distribution power law index.
+    :param logepse: log10 fraction of thermal energy in electrons
+    :param logepsb: log10 fraction of thermal energy in magnetic field
+    :param g0: initial lorentz factor
+    :param s: power-law jet slope
+    :param kwargs: Additional keyword arguments
+    :param output_format: Whether to output flux density or AB mag
+    :param frequency: frequency in Hz for the flux density calculation
+    :param cosmology: Cosmology to use for luminosity distance calculation. Defaults to Planck18. Must be a astropy.cosmology object.
+    :return: flux density or AB mag. Note this is going to give the monochromatic magnitude at the effective frequency for the band.
+    """
+    import jetsimpy #Can not use models unless jetsimpy is downloaded
+    time = time * day_to_s
+    cosmology = kwargs.get('cosmology', cosmo)
+    dl = cosmology.luminosity_distance(redshift).cgs.value
+    P = dict(Eiso = 10 ** loge0, lf = g0, theta_c = thc, n0 = nism, A = A, eps_e = 10 ** logepse, eps_b = 10 ** logepsb, p = p, theta_v = thv, d = dl*3.24078e-25, z = redshift, s = s) #make a param dict
+    if kwargs['output_format'] == 'flux_density':
+        frequency = kwargs['frequency']
+        flux_density = jetsimpy.FluxDensity_powerlaw(time, frequency, P)
+        return flux_density   
+    else:
+        frequency = bands_to_frequency(kwargs['bands'])       
+        flux_density = jetsimpy.FluxDensity_powerlaw(time, frequency, P)
+        return calc_ABmag_from_flux_density(flux_density).value