Merge pull request #234 from Yong2Sheng/source_injector

Source injector for `cosipy`

cosipy/__init__.py

@@ -13,3 +13,5 @@
 from .spacecraftfile import SpacecraftFile
 
 from .ts_map import FastTSMap
+
+from .source_injector import SourceInjector

cosipy/source_injector/__init__.py

@@ -0,0 +1 @@
+from .source_injector import SourceInjector

cosipy/source_injector/source_injector.py

@@ -0,0 +1,161 @@
+from histpy import Histogram
+from pathlib import Path
+from cosipy.response import FullDetectorResponse
+import h5py as h5
+from histpy import Histogram, Axis, Axes
+from cosipy.response import PointSourceResponse
+import sys
+from mhealpy import HealpixMap
+
+class SourceInjector():
+
+    def __init__(self, response_path, response_frame = "local"):
+
+        """
+        `SourceInjector` convolve response, source model(s) and orientation to produce a mocked simulated data. The data can be saved for data anlysis with cosipy.
+
+        Parameters
+        ----------
+        response : str or pathlib.Path
+            The path to the response file
+        response_frame : str, optional
+            The frame of the Compton data space (CDS) of the response. It only accepts `local` or "galactic". (the default is `local`, which means the CDS is in the local detector frame.
+        """
+
+        self.response_path =  response_path
+
+        if response_frame == "local" or response_frame == "galactic":
+
+            self.response_frame = response_frame
+
+        else:
+            raise ValueError("The response frame can only be `local` or `galactic`!")
+
+
+    @staticmethod
+    def get_psr_in_galactic(coordinate, response_path, spectrum):
+
+        """
+        Get the point source response (psr) in galactic. Please be aware that you must use a galactic response!
+        To do: to make the weight parameter not hardcoded
+        
+        Parameters
+        ----------
+        coordinate : astropy.coordinates.SkyCoord
+            The coordinate.
+        response_path : str or path.lib.Path
+            The path to the response.
+        spectrum : astromodels.functions
+            The spectrum of the source to be placed at the hypothesis coordinate.
+        
+        Returns
+        -------
+        psr : histpy.Histogram
+            The point source response of the spectrum at the hypothesis coordinate.
+        """
+
+        # Open the response
+        # Notes from Israel: Inside it contains a single histogram with all the regular axes for a Compton Data Space (CDS) analysis, in galactic coordinates. Since there is no class yet to handle it, this is how to read in the HDF5 manually.
+
+        with h5.File(response_path) as f:
+
+            axes_group = f['hist/axes']
+            axes = []
+            for axis in axes_group.values():
+                # Get class. Backwards compatible with version
+                # with only Axis
+                axis_cls = Axis
+                if '__class__' in axis.attrs:
+                    class_module, class_name = axis.attrs['__class__']
+                    axis_cls = getattr(sys.modules[class_module], class_name)
+                axes += [axis_cls._open(axis)]
+        axes = Axes(axes)
+
+        # get the pixel number of the hypothesis coordinate
+        map_temp = HealpixMap(base = axes[0])
+        coordinate_pix_number = map_temp.ang2pix(coordinate)
+
+        # get the expectation for the hypothesis coordinate (a point source)
+        with h5.File(response_path) as f:
+            pix = coordinate_pix_number
+            psr = PointSourceResponse(axes[1:], f['hist/contents'][pix+1], unit = f['hist'].attrs['unit'])
+
+        return psr
+
+
+    def inject_point_source(self, spectrum, coordinate, orientation = None, source_name = "point_source",
+                            make_spectrum_plot = False, data_save_path = None, project_axes = None):
+
+        """
+        Get the expected counts for a point source.
+
+        Parameters
+        ----------
+        spectrum : astromodels.functions
+            The spectrum model defined from `astromodels`.
+        coordinate : astropy.coordinates.SkyCoord
+            The coordinate of the point source.
+        orientation : cosipy.spacecraftfile.SpacecraftFile, option
+            The orientation of the telescope during the mock simulation. This is needed when using a detector response. (the default is `None`, which means a galactic response is used.
+        source_name : str, optional
+            The name of the source (the default is `point_source`).
+        make_spectrum_plot : bool, optional
+            Set `True` to make the plot of the injected spectrum.
+        data_save_path : str or pathlib.Path, optional
+            The path to save the injected data to a `.h5` file. This should include the file name. (the default is `None`, which means the injected data won't be saved.
+        project_axes : list, optional
+            The axes to project before saving the data file (the default is `None`, which means the data won't be projected).
+            
+        Returns
+        -------
+        histpy.Histogram
+            The `Histogram object of the injected spectrum.`
+        """
+
+
+        # get the point source response in local frame
+        if self.response_frame == "local":
+
+            if orientation == None:
+                raise TypeError("The when the data are binned in local frame, orientation must be provided to compute the expected counts.")
+
+            # get the dwell time map
+            coord_in_sc_frame = orientation.get_target_in_sc_frame(target_name = source_name, 
+                                                                   target_coord = coordinate, 
+                                                                   quiet = True)
+
+            # get the dwell time map in the detector frame
+            dwell_time_map = orientation.get_dwell_map(response = self.response_path)
+
+            with FullDetectorResponse.open(self.response_path) as response:
+                psr = response.get_point_source_response(dwell_time_map)
+
+        # get the point source response in galactic frame
+        elif self.response_frame == "galactic":
+
+            psr = SourceInjector.get_psr_in_galactic(coordinate = coordinate, response_path = self.response_path, spectrum = spectrum)
+
+        injected = psr.get_expectation(spectrum)
+        # setting the Em and Ei scale to linear to match the simulated data
+        # The linear scale of Em and Ei is the default for COSI data
+        injected.axes["Em"].axis_scale = "linear"
+        injected.axes["Ei"].axis_scale = "linear"
+
+        if project_axes is not None:
+            injected = injected.project(project_axes)
+
+        if make_spectrum_plot is True:
+            ax, plot = injected.project("Em").draw(label = "Injected point source", color = "green")
+            ax.legend()
+            ax.set_xscale("log")
+            ax.set_yscale("log")
+            ax.set_ylabel("Counts")
+
+        if data_save_path is not None:
+            injected.write(data_save_path)
+
+        return injected
+
+
+
+

docs/api/index.rst

@@ -18,6 +18,7 @@ If you are instead interested in an overview on how to use cosipy, see out `tuto
    threeml
    ts_map
    image_deconvolution
-   util          
+   util
+   source_injector          
 
 

docs/api/source_injector.rst

@@ -0,0 +1,12 @@
+Source injector
+=================
+
+Convolve a source model, response and orientation to get mock data without simulations.
+
+.. warning::
+    The `source injector` is still under development. Now it only supports point sources. More types of source will be supported (e.g. extenced sources).
+
+.. automodule:: cosipy.source_injector
+   :imported-members:
+   :members:
+   :undoc-members:

docs/tutorials/index.rst

@@ -59,8 +59,9 @@ List of tutorials and contents, as a link to the corresponding Python notebook i
   - Analyze data in the Compton data space with galactic coordinates.
   - Link to a notebook using Scatt binning which shows its advantages/disadvantages.
 
-9. TODO: Source injector
-  - Nice to have: allow theorist to test the sensitivity of their models
+9. Source injector `(ipynb) <https://github.com/cositools/cosipy/tree/main/docs/tutorials/source_injector/Point_source_injector.ipynb>`_
+  - Convolve the response, point source model and orientation to obtain the mock data.
+  - More types of source (e,g. extended source and polarization) will be suppored.
 
 .. warning::
    Under construction. Some of the explanations described above might be missing. However, the notebooks are fully functional. If you have a question not yet covered by the tutorials, please discuss `issue <https://github.com/cositools/cosipy/discussions>`_ so we can prioritize it.
@@ -76,5 +77,6 @@ List of tutorials and contents, as a link to the corresponding Python notebook i
    Fitting the spectrum of the Crab <spectral_fits/continuum_fit/crab/SpectralFit_Crab.ipynb>
    Extended source model fitting <spectral_fits/extended_source_fit/diffuse_511_spectral_fit.ipynb>
    Image deconvolution <image_deconvolution/511keV/ScAttBinning/511keV-DC2-ScAtt-ImageDeconvolution.ipynb>
+   Source injector <source_injector/Point_source_injector.ipynb>
 
 

docs/tutorials/source_injector/GRB_source_injector.ipynb

@@ -30,7 +30,13 @@
    "cell_type": "code",
    "execution_count": 1,
    "id": "1863fe19-1d2b-4d9d-aacc-6e1eba99b882",
-   "metadata": {},
+   "metadata": {
+    "collapsed": true,
+    "jupyter": {
+     "outputs_hidden": true
+    },
+    "scrolled": true
+   },
    "outputs": [
     {
      "name": "stderr",
@@ -441,9 +447,7 @@
    "cell_type": "code",
    "execution_count": 4,
    "id": "d1e043e0",
-   "metadata": {
-    "scrolled": false
-   },
+   "metadata": {},
    "outputs": [
     {
      "data": {
@@ -27324,9 +27328,9 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "python-env",
+   "display_name": "cosipy_spacecraftfile_new",
    "language": "python",
-   "name": "python-env"
+   "name": "cosipy_spacecraftfile_new"
   },
   "language_info": {
    "codemirror_mode": {
@@ -27338,7 +27342,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.10.9"
+   "version": "3.10.13"
   }
  },
  "nbformat": 4,