Improve plotting latent light curve and spectrum

.gitignore

@@ -111,3 +111,9 @@ venv.bak/
 /docs/notebooks/*.rsp
 /docs/notebooks/*.fits
 /docs/notebooks/*.ipynb
+
+# hdf5
+*.h5
+
+# Dask worker space
+dask-worker-space/

cosmogrb/grb/grb.py

@@ -134,6 +134,26 @@ def display_energy_dependent_light_curve(
         ].display_energy_dependent_light_curve(
             time=time, energy=energy, ax=ax, cmap=cmap, **kwargs
         )
+
+    def display_time_dependent_spectrum(
+        self, time, energy, ax=None, cmap="viridis", **kwargs
+    ):
+        """FIXME! briefly describe function
+
+        :param time:
+        :param energy:
+        :param ax:
+        :param cmap:
+        :returns:
+        :rtype:
+
+        """
+
+        list(self._lightcurves.values())[
+            0
+        ].display_time_dependent_spectrum(
+            time=time, energy=energy, ax=ax, cmap=cmap, **kwargs
+        )
 
     def display_energy_integrated_light_curve(self, time, ax=None, **kwargs):
         """FIXME! briefly describe function

cosmogrb/io/grb_save.py

@@ -18,6 +18,7 @@ def __init__(
         dec,
         duration,
         z,
+        time_adjustment,
         lightcurves,
         responses,
         source_params,
@@ -72,7 +73,7 @@ def __init__(
                 lightcurve=lightcurves[key], response=responses[key]
             )
 
-        #
+
 
         super(GRBSave, self).__init__(data)
 
@@ -91,6 +92,7 @@ def __init__(
         self._dec = dec
         self._z = z
         self._duration = duration
+        self._time_adjustment = time_adjustment
 
     def __setitem__(self, key, value):
 
@@ -135,6 +137,10 @@ def duration(self):
     def T0(self):
         return self._T0
 
+    @property
+    def time_adjustment(self):
+        return self._time_adjustment
+
     @property
     def extra_info(self):
         return self._extra_info
@@ -247,6 +253,7 @@ def from_file(cls, file_name):
             ra=ra,
             dec=dec,
             z=z,
+            time_adjustment = time_adjustment,
             duration=duration,
             lightcurves=lightcurves,
             responses=responses,
@@ -256,7 +263,7 @@ def from_file(cls, file_name):
 
     def __repr__(self):
 
-        return self._output().to_string()
+        return self._output().to_string(float_format=lambda x: "{:.3e}".format(x))
 
     def info(self):
 

cosmogrb/lightcurve/light_curve_storage.py

@@ -473,7 +473,7 @@ def _bin_spectrum(self, tmin, tmax, times, pha):
 
         """
 
-        idx = np.ones_like(times, dtype=bool)
+        idx = np.zeros_like(times, dtype=bool)
 
         # filter times
 

cosmogrb/lightcurve/lightcurve.py

@@ -220,6 +220,24 @@ def display_energy_dependent_light_curve(
             time=time, energy=energy, ax=ax, cmap=cmap, **kwargs
         )
 
+    def display_time_dependent_spectrum(
+        self, time, energy, ax=None, cmap="viridis", **kwargs
+    ):
+        """FIXME! briefly describe function
+
+        :param time:
+        :param energy:
+        :param ax:
+        :param cmap:
+        :returns:
+        :rtype:
+
+        """
+
+        self._source.display_time_dependent_spectrum(
+            time=time, energy=energy, ax=ax, cmap=cmap, **kwargs
+        )
+
     def display_energy_integrated_light_curve(self, time, ax=None, **kwargs):
         """FIXME! briefly describe function
 

cosmogrb/response/response.py

@@ -260,7 +260,7 @@ def to_fits(
             instrument_name,
         )
 
-        fits_file.writeto(filename, clobber=overwrite)
+        fits_file.writeto(filename, overwrite=overwrite)
 
 
 __all__ = ["Response"]

cosmogrb/sampler/source.py

@@ -77,6 +77,24 @@ def display_energy_dependent_light_curve(
             time, energy, ax, cmap, **kwargs
         )
 
+    def display_time_dependent_spectrum(
+        self, time, energy, ax=None, cmap="viridis", **kwargs
+    ):
+        """FIXME! briefly describe function
+
+        :param time:
+        :param energy:
+        :param ax:
+        :param cmap:
+        :returns:
+        :rtype:
+
+        """
+
+        self._source_function.display_time_dependent_spectrum(
+            time, energy, ax, cmap, **kwargs
+        )
+
     def set_response(self, response):
         """
         called if there is no response upon creation

cosmogrb/sampler/source_function.py

@@ -1,6 +1,10 @@
 import abc
 
+import numpy as np
+
+import matplotlib as mpl
 import matplotlib.pyplot as plt
+from mpl_toolkits.axes_grid1 import make_axes_locatable
 
 from cosmogrb.utils.array_to_cmap import array_to_cmap
 
@@ -97,9 +101,11 @@ def sample_energy(self, times):
 
         pass
 
-    def display_energy_integrated_light_curve(self, time, ax=None, **kwargs):
+    def display_energy_integrated_light_curve(self, time,ax=None, **kwargs):
         """
         plot the latent light curve integrated over energy
+        in the observer frame
+        set F_peak is given in burst frame -> peak is in observer frame at F/(1+z^2)
 
         :param time:
         :param ax:
@@ -116,20 +122,30 @@ def display_energy_integrated_light_curve(self, time, ax=None, **kwargs):
 
             fig = ax.get_figure()
 
-        y = [self.energy_integrated_evolution(t) for t in time]
+        n_energies = 75
+
+        energy_grid = np.power(
+            10.0, np.linspace(np.log10(self._emin*(1+self._z)), np.log10(self._emax*(1+self._z)), n_energies)
+        )
+
+        y = np.zeros(len(time))
+        for i,t in enumerate(time):
+            # flux at index (time, energy)
+            energy_slice = self.evolution(energy_grid, t)
+            y[i] = np.trapz(energy_slice[0, :]*energy_grid, energy_grid)
 
         ax.plot(time, y, **kwargs)
 
-        ax.set_xlabel("time")
-        ax.set_ylabel("flux")
+        ax.set_xlabel("Time [s]")
+        ax.set_ylabel(r"Flux [keV/s/cm$^2$]")
 
         return fig
 
     def display_energy_dependent_light_curve(
-        self, time, energy, ax=None, cmap="viridis", **kwargs
+        self, time, energy, ax=None, cmap="viridis", uselog=True, **kwargs
     ):
         """
-        plot the latent light curve integrated over energy
+        plot the latent light curve for different energies
 
         :param time:
         :param ax:
@@ -146,20 +162,85 @@ def display_energy_dependent_light_curve(
 
             fig = ax.get_figure()
 
-        # index (time, flux)
+        # flux at index (time, energy)
         grid = self.evolution(energy, time)
 
-        _, colors = array_to_cmap(energy, cmap=cmap, use_log=True)
+        _, colors = array_to_cmap(energy, cmap=cmap, use_log=uselog)
+
+        cmap = mpl.colormaps[cmap]
+
+        if uselog:
+            norm = mpl.colors.LogNorm(vmin=min(energy), vmax=max(energy))
+        else:
+            norm = mpl.colors.Normalize(vmin=min(energy), vmax=max(energy))
+
+        #Add an axis for colorbar on right
+        divider = make_axes_locatable(ax)
+        cax = divider.append_axes("right", size="5%", pad=0.05)
+
+        fig.colorbar(mpl.cm.ScalarMappable(norm=norm, cmap=cmap),
+                    cax=cax, label='Energy [keV]')
 
         for i, lc in enumerate(grid.T):
 
             ax.plot(time, lc, color=colors[i], **kwargs)
 
-        ax.set_xlabel("time")
+        ax.set_xlabel("Time [s]")
         ax.set_ylabel("flux")
 
         return fig
 
+    def display_time_dependent_spectrum(
+        self, time, energy, ax=None, cmap="viridis", uselog=False, **kwargs
+    ):
+        """
+        plot the latent nu *F_nu spectrum at different times
+
+        :param time:
+        :param ax:
+        :returns:
+        :rtype:
+
+        """
+
+        if ax is None:
+
+            fig, ax = plt.subplots()
+
+        else:
+
+            fig = ax.get_figure()
+
+        # index (time, flux)
+        grid = self.evolution(energy, time)
+
+        _, colors = array_to_cmap(time, cmap=cmap, use_log=uselog)
+
+        cmap = mpl.colormaps[cmap]
+
+        if uselog:
+            norm = mpl.colors.LogNorm(vmin=min(time), vmax=max(time))
+        else:
+            mpl.colors.Normalize(vmin=min(time), vmax=max(time))
+
+        #Add an axis for colorbar on right
+        divider = make_axes_locatable(ax)
+        cax = divider.append_axes("right", size="5%", pad=0.05)
+
+        fig.colorbar(mpl.cm.ScalarMappable(norm=norm, cmap=cmap),
+                    cax=cax, label='Time [s]')
+
+        for i, lc in enumerate(grid):
+
+            ax.plot(energy, energy**2*lc, color=colors[i], **kwargs)
+
+        ax.set_xlabel("Energy [keV]")
+        ax.set_ylabel(r"$\nu F_{\nu}$ [keV$^2$/s/cm$^2$/keV]")
+        ax.set_xscale('log')
+        ax.set_yscale('log')
+
+        return fig
+
     @property
     def index(self):
         return self._index