Interfaces refactoring #369

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft

israelmcmc wants to merge 142 commits into cositools:develop from israelmcmc:interfaces

cosipy/__init__.py

-Original file line number
+Diff line change
@@ Expand Up / @@ -2,15 +2,16 @@ @@
     from .response import DetectorResponse
+    from .spacecraftfile import *
     from .data_io import DataIO
     from .data_io import UnBinnedData
     from .data_io import BinnedData
     from .data_io import ReadTraTest
-    from .threeml import COSILike
     from .threeml import Band_Eflux
-    from .spacecraftfile import SpacecraftFile
+    from .spacecraftfile import SpacecraftHistory
     from .ts_map import FastTSMap, MOCTSMap
@@ Expand Down @@

cosipy/background_estimation/ContinuumEstimation.py

-Original file line number
+Diff line change
@@ Expand Up @@
             ----------
             ori_file : str
                 Full path to orienation file.
-            sc_orientation : cosipy.spacecraftfile.SpacecraftFile
+            sc_orientation : cosipy.spacecraftfile.SpacecraftHistory
                 Spacecraft orientation object.
             detector_response : str
                 Full path to detector response file.
@@ Expand Down @@

cosipy/background_estimation/__init__.py

-Original file line number
+Diff line change
@@ -1,2 +1,3 @@
     from .LineBackgroundEstimation import LineBackgroundEstimation
     from .ContinuumEstimation import ContinuumEstimation
+    from .free_norm_threeml_binned_bkg import *

cosipy/background_estimation/free_norm_threeml_binned_bkg.py

-Original file line number
+Diff line change
@@ -0,0 +1,307 @@
+    import itertools
+    from typing import Dict, Tuple, Union, Any, Type, Optional, Iterable
+    import numpy as np
+    from astromodels import Parameter
+    from astropy.coordinates import SkyCoord, CartesianRepresentation, UnitSphericalRepresentation
+    from astropy.time import Time
+    from histpy import Histogram
+    from histpy import Axes
+    from astropy import units as u
+    from scoords import SpacecraftFrame
+    from cosipy import SpacecraftHistory
+    from cosipy.data_io.EmCDSUnbinnedData import TimeTagEmCDSEventInSCFrame
+    from cosipy.interfaces import BinnedBackgroundInterface, BinnedDataInterface, DataInterface, BackgroundDensityInterface, \
+        BackgroundInterface, EventInterface
+    __all__ = ["FreeNormBinnedBackground"]
+    from cosipy.interfaces.data_interface import TimeTagEmCDSEventDataInSCFrameInterface
+    from cosipy.interfaces.event import TimeTagEmCDSEventInSCFrameInterface
+    from cosipy.util.iterables import itertools_batched
+    class FreeNormBackground(BackgroundInterface):
+        """
+        This must translate to/from regular parameters
+        with arbitrary type from/to 3ML parameters
+        Default to "bkg_norm" is there was a single unlabeled component
+        """
+        _default_label = 'bkg_norm'
+        def __init__(self,
+                     distribution:Union[Histogram, Dict[str, Histogram]],
+                     sc_history:SpacecraftHistory,
+                     copy = True):
+            """
+            Parameters
+            ----------
+            distribution
+            sc_history
+            copy: copy hist distribution
+            """
+            if isinstance(distribution, Histogram):
+                # Single component
+                self._distributions = {self._default_label: distribution}
+                self._norms = np.ones(1) # Hz. Each component
+                self._norm = 1 # Hz. Total
+                self._single_component = True
+            else:
+                # Multiple label components.
+                self._distributions = distribution
+                self._norms = np.ones(self.ncomponents) # Hz Each component
+                self._norm = np.sum(self._norms) # Hz. Total
+                self._single_component = False
+            self._labels = tuple(self._distributions.keys())
+            # Normalize
+            # Unit: second
+            self._livetime = sc_history.cumulative_livetime().to_value(u.s)
+            for label,dist in self._distributions.items():
+                dist_norm = np.sum(dist)
+                if copy:
+                    self._distributions[label] = dist/dist_norm
+                else:
+                    dist /= dist_norm
+            # These will be densify anyway since _expectation is dense
+            # And histpy doesn't yet handle this operation efficiently
+            # See Histogram._inplace_operation_handle_sparse()
+            # Do it once and for all
+            for label, bkg in self._distributions.items():
+                if bkg.is_sparse:
+                    self._distributions[label] = bkg.to_dense()
+            if self.ncomponents == 0:
+                raise ValueError("You need to input at least one components")
+            self._axes = None
+            for bkg in self._distributions.values():
+                if self._axes is None:
+                    self._axes = bkg.axes
+                else:
+                    if self._axes != bkg.axes:
+                        raise ValueError("All background components mus have the same axes")
+        @property
+        def norm(self):
+            """
+            Sum of all rates
+            """
+            return u.Quantity(self._norm, u.Hz)
+        @property
+        def norms(self):
+            if self._single_component:
+                return {self._default_label: u.Quantity(self._norms[0], u.Hz)}
+            else:
+                return {l:u.Quantity(n, u.Hz, copy = False) for l,n in zip(self.labels,self._norms)}
+        @property
+        def ncomponents(self):
+            return len(self._distributions)
+        @property
+        def meausured_axes(self):
+            return self._axes
+        @property
+        def labels(self):
+            return self._labels
+        def set_norm(self, norm: Union[u.Quantity, Dict[str, u.Quantity]]):
+            if self._single_component:
+                if isinstance(norm, dict):
+                    self._norms[0] = norm[self._default_label].to_value(u.Hz)
+                else:
+                    self._norms[0] = norm.to_value(u.Hz)
+            else:
+                # Multiple
+                if not isinstance(norm, dict):
+                    raise TypeError("This a multi-component background. Provide labeled norm values in a dictionary")
+                for label,norm_i in norm.items():
+                    if label not in self.labels:
+                        raise ValueError(f"Norm {label} not in {self.labels}")
+                    self._norms[self.labels.index(label)] = norm_i.to_value(u.Hz)
+            self._norm = sum(n for n in self._norms)
+        def set_parameters(self, **parameters:Dict[str, u.Quantity]) -> None:
+            """
+            Same keys as background components
+            """
+            self.set_norm(parameters)
+        @property
+        def parameters(self) -> Dict[str, u.Quantity]:
+            return self.norms
+    class FreeNormBinnedBackground(FreeNormBackground, BinnedBackgroundInterface):
+        def __init__(self, *args, **kwargs):
+            super().__init__(*args, **kwargs)
+            # Cache
+            self._expectation = None
+            self._last_norm_values = None
+        def expectation(self, axes:Axes, copy:bool = True)->Histogram:
+            """
+            Parameters
+            ----------
+            axes
+            copy:
+                If True, it will return an array that the user if free to modify.
+                Otherwise, it will result a reference, possible to the cache, that
+                the user should not modify
+            Returns
+            -------
+            """
+            if axes != self.meausured_axes:
+                raise ValueError("Requested axes do not match the background component axes")
+            # Check if we can use the cache
+            if self._expectation is None:
+                # First call. Initialize
+                self._expectation = Histogram(self.meausured_axes)
+            elif self.norms == self._last_norm_values:
+                # No changes. Use cache
+                if copy:
+                    return self._expectation.copy()
+                else:
+                    return self._expectation
+            else:
+                # First call or norms have change. Recalculate
+                self._expectation.clear()
+            # Compute expectation
+            for label,bkg in self._distributions.items():
+                norm = self._norms[self.labels.index(label)]
+                self._expectation += bkg * norm * self._livetime
+            # Cache. Regular copy is enough since norm values are float en not mutable
+            self._last_norm_values = self.norms.copy()
+            if copy:
+                return self._expectation.copy()
+            else:
+                return self._expectation
+    class FreeNormBackgroundInterpolatedDensityTimeTagEmCDS(FreeNormBackground, BackgroundDensityInterface):
+        @property
+        def event_type(self) -> Type[EventInterface]:
+            return TimeTagEmCDSEventInSCFrameInterface
+        def __init__(self,
+                     data:TimeTagEmCDSEventDataInSCFrameInterface,
+                     distribution:Union[Histogram, Dict[str, Histogram]],
+                     sc_history:SpacecraftHistory,
+                     copy=True,
+                     batch_size = 100000,
+                     *args,
+                     **kwargs):
+            super().__init__(distribution, sc_history,
+                             copy=copy, *args, **kwargs)
+            # We need the density per phase space for the specific measurement units TimeTagEmCDSEventInSCFrameInterface
+            # Energy: keV
+            # Phi: rad
+            # PsiChi: sr (for the phase space. The axis is a HealpixAxis)
+            # Time: seconds (taken into account by the norm (a rate) unit)
+            psichi_frame = None
+            for label,dist in self._distributions.items():
+                dist = self._distributions[label] = dist.project('Em', 'Phi', 'PsiChi')
+                dist.axes['Em'] = dist.axes['Em'].to(u.keV).to(None, copy=False, update=False)
+                dist.axes['Phi'] = dist.axes['Phi'].to(u.rad).to(None, copy=False, update=False)
+                energy_phase_space = dist.axes['Em'].widths
+                phi_phase_space = dist.axes['Phi'].widths
+                psichi_phase_space = dist.axes['PsiChi'].pixarea().to_value(u.sr)
+                if psichi_frame is None:
+                    psichi_frame = dist.axes['PsiChi'].coordsys
+                else:
+                    if psichi_frame != dist.axes['PsiChi'].coordsys:
+                        raise ValueError("All PsiChi axes must be in the same frame")
+                dist /= dist.axes.expand_dims(energy_phase_space, 'Em')
+                dist /= dist.axes.expand_dims(phi_phase_space, 'Phi')
+                dist /= psichi_phase_space
+            # Compute the probabilities once and for all
+            # TODO: account for livetime
+            self._prob = [[] for _ in range(self.ncomponents)]
+            for events_chunk in itertools_batched(data, batch_size):
+                jd1, jd2, energy,phi, psichi_lon, psichi_lat  = np.asarray([[
+                    event.jd1,
+                    event.jd2,
+                    event.energy_keV,
+                                                                   event.scattering_angle_rad,
+                                                                   event.scattered_lon_rad_sc,
+                                                                   event.scattered_lat_rad_sc]
+                                                   for event in events_chunk], dtype=float).transpose()
+                times = Time(jd1, jd2, format = 'jd')
+                # Transform local to inertial
+                sc_psichi_coord = SkyCoord(psichi_lon, psichi_lat, unit=u.rad, frame=SpacecraftFrame())
+                sc_psichi_vec = sc_psichi_coord.cartesian.xyz.value
+                attitudes = sc_history.interp_attitude(times).transform_to(psichi_frame)
+                inertial_psichi_vec = attitudes.rot.apply(sc_psichi_vec.transpose())
+                inertial_psichi_sph = UnitSphericalRepresentation.from_cartesian(CartesianRepresentation(*inertial_psichi_vec.transpose()))
+                inertial_psichi_coord = SkyCoord(inertial_psichi_sph, frame = psichi_frame)
+                for label,dist in self._distributions.items():
+                    prob = dist.interp(energy, phi, inertial_psichi_coord)
+                    self._prob[self.labels.index(label)].extend(prob)
+            self._prob = np.asarray(self._prob)
+        def expected_counts(self) -> float:
+            """
+            Total expected counts
+            """
+            return self._livetime * self._norm
+        def expectation_density(self) -> Iterable[float]:
+            """
+            Return the expected number of counts density from the start-th event
+            to the stop-th event. This equals the event probabiliy times the number of events
+            """
+            # Multiply each probability by the norm, and then sum
+            return np.tensordot(self._prob, self._norms, axes = (0,0))

cosipy/data_io/BinnedData.py

-Original file line number
+Diff line change
@@ Expand Up / @@ -16,6 +16,11 @@ @@
     from cosipy.data_io import UnBinnedData
     import logging
+    import astropy.units as u
+    from astropy.coordinates import SkyCoord
+    from cosipy.interfaces import BinnedDataInterface
     logger = logging.getLogger(__name__)
@@ Expand Down Expand Up @@
                 d = {"Time[UTC]":self.time_bin_centers,"Rate[ct/s]":self.time_hist/self.time_bin_widths}
                 df = pd.DataFrame(data=d)
                 df.to_csv(f"{output_name}.dat",index=False,sep="\t",columns=["Time[UTC]","Rate[ct/s]"])
+        def get_em_cds(self):
+            return EmCDSBinnedData(self.binned_data.project('Em', 'Phi', 'PsiChi'))
+    class EmCDSBinnedData(BinnedDataInterface):
+        """
+        Measured energy (Em), Compton polar scattering angle (Phi), and the scattering direction (PsiChi).
+        Phi and PsiChi are the Compton Data Space (CDS). No time dependence
+        """
+        def __init__(self, data:Histogram):
+            # Checks
+            if set(data.axes.labels) != {'Em', 'Phi', 'PsiChi'}:
+                raise ValueError(f"Wrong axes. 'Em', 'Psi', 'PsiChi' expected.")
+            if not data.axes['Em'].unit.is_equivalent(u.keV):
+                raise ValueError(f"Em axis should have units of energy")
+            if not data.axes['Phi'].unit.is_equivalent(u.deg):
+                raise ValueError(f"Psi axis should have angle units")
+            if not isinstance(data.axes['PsiChi'],HealpixAxis):
+                raise ValueError(f"PsiChi must be of type {HealpixAxis}.")
+            if data.axes['PsiChi'].coordsys is None:
+                raise ValueError(f"PsiChi axes must have a coordinate system.")
+            self._data = data
+        @property
+        def data(self) -> Histogram:
+            return self._data
+        @property
+        def axes(self) -> Axes:
+            return self._data.axes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Interfaces refactoring #369

Uh oh!

Diff view

Diff view

There are no files selected for viewing

Uh oh!

Uh oh!

Interfaces refactoring #369

Are you sure you want to change the base?

Uh oh!

Interfaces refactoring #369

Uh oh!

Uh oh!

Diff view

Diff view

There are no files selected for viewing

Uh oh!

Uh oh!