Merge pull request #362 from GallegoSav/issue_361

Adding the Compton Sequence variable during tra file reading

Author: ckarwin

cosipy/data_io/ReadTraTest.py

@@ -54,7 +54,8 @@ def read_tra_old(self, output_name, make_plots=True):
                         'Psi local':psi_loc,\
                         'Distance':dist,\
                         'Chi galactic':chi_gal,\
-                        'Psi galactic':psi_gal}
+                        'Psi galactic':psi_gal,\
+                        'Compton Seq':CO_seq}
         
         Note
         ----
@@ -103,7 +104,10 @@ def read_tra_old(self, output_name, make_plots=True):
         chi_gal = []
         # Measured gal angle psi (lat direction)
         psi_gal = [] 
+        # Compton seq
+        CO_seq = []
 
+        
         # Browse through tra file, select events, and sort into corresponding list:
         # Note: The Reader class from MEGAlib knows where an event starts and ends and
         # returns the Event object which includes all information of an event.
@@ -143,12 +147,17 @@ def read_tra_old(self, output_name, make_plots=True):
             chi_gal.append((Event.GetGalacticPointingRotationMatrix()*Event.Dg()).Phi())
             # Gal longitude angle corresponding to chi:
             psi_gal.append((Event.GetGalacticPointingRotationMatrix()*Event.Dg()).Theta())
-                
+            # Compton sequence (nb of interaction) 
+            CO_seq.append(Event.SequenceLength())
+
+            
         # Initialize arrays:
         erg = np.array(erg)
         tt = np.array(tt)
         et = np.array(et)
 
+        CO_seq = np.array(CO_seq)
+        
         latX = np.array(latX)
         lonX = np.array(lonX)
         # Change longitudes to from 0..360 deg to -180..180 deg
@@ -201,7 +210,8 @@ def read_tra_old(self, output_name, make_plots=True):
                         'Psi local':self.psi_loc_old,
                         'Distance':dist,
                         'Chi galactic':self.chi_gal_old,
-                        'Psi galactic':self.psi_gal_old} 
+                        'Psi galactic':self.psi_gal_old,
+                       'Compton Seq':CO_seq} 
         self.cosi_dataset = cosi_dataset
 
         # Write unbinned data to file (either fits or hdf5):

cosipy/data_io/UnBinnedData.py

@@ -78,7 +78,8 @@ def read_tra(self, input_name=None, output_name=None, run_test=False, use_ori=Fa
                         'Psi local':psi_loc,\
                         'Distance':dist,\
                         'Chi galactic':chi_gal,\
-                        'Psi galactic':psi_gal}\
+                        'Psi galactic':psi_gal,\
+                        'CO seq':CO_seq}\
             Arrays contain unbinned photon data. 
         
         Notes
@@ -126,6 +127,10 @@ def read_tra(self, input_name=None, output_name=None, run_test=False, use_ori=Fa
         dg_y = []
         dg_z = []
 
+        # Compton seq
+        CO_seq = []
+
+        
         # Define electron rest energy, which is used in calculation
         # of Compton scattering angle.
         c_E0 = 510.9989500015 # keV
@@ -210,6 +215,7 @@ def read_tra(self, input_name=None, output_name=None, run_test=False, use_ori=Fa
                         erg.pop()
                         phi.pop()
                         tt.pop()
+						CO_seq.pop()
                         # Not all sims include ori info,
                         # so also need to check before pop:
                         if len(lonX) != 0:
@@ -280,6 +286,11 @@ def read_tra(self, input_name=None, output_name=None, run_test=False, use_ori=Fa
                 lonZ.append(this_lonZ)
                 latZ.append(this_latZ)
 
+            
+            #number of interaction
+            if this_line[0] == "SQ":
+                CO_seq.append(int(this_line[1]))
+            
             # Interaction position information: 
             if (this_line[0] == "CH"):
 
@@ -323,7 +334,8 @@ def read_tra(self, input_name=None, output_name=None, run_test=False, use_ori=Fa
         dg_x = np.array(dg_x)
         dg_y = np.array(dg_y)
         dg_z = np.array(dg_z)
- 
+        CO_seq = np.array(CO_seq)
+        
         # Check if the input data has pointing information, 
         # if not, set dummy values:
         if (use_ori == False) & (len(lonZ)==0):
@@ -410,7 +422,8 @@ def read_tra(self, input_name=None, output_name=None, run_test=False, use_ori=Fa
                         'Psi local':psi_loc,
                         'Distance':dist,
                         'Chi galactic':chi_gal,
-                        'Psi galactic':psi_gal} 
+                        'Psi galactic':psi_gal,
+                       'Compton Seq':CO_seq} 
         self.cosi_dataset = cosi_dataset
 
         # Option to write unbinned data to file (either fits or hdf5):
@@ -550,7 +563,7 @@ def write_unbinned_output(self, output_name):
 
         # Data units:
         units=['keV','s','rad','rad',
-                'rad','rad','rad','rad','cm','deg','deg']
+                'rad','rad','rad','rad','cm','deg','deg','']
 
         # For fits output: 
         if self.unbinned_output == 'fits':
@@ -723,6 +736,47 @@ def select_data_energy(self, emin, emax, output_name=None, unbinned_data=None):
 
         return
 
+    def select_data_COseq(self, seqmin, seqmax, output_name=None, unbinned_data=None):
+
+        """Applies CO sequence cuts [seqmin,seqmax) to unbinnned data dictionary. 
+        
+        Parameters
+        ----------
+        seqmin :  int
+            Minimum number of interaction.
+        seqmax :  int
+            Maximum number of interaction.
+        unbinned_data : str, optional
+            Name of unbinned dictionary file.
+        output_name : str, optional
+            Prefix of output file (default is None, in which case no 
+            file is saved).
+        """
+        
+        logger.info("Making data selections on Compton Sequence...")
+
+        # Option to read in unbinned data file:
+        if unbinned_data:
+            self.cosi_dataset = self.get_dict(unbinned_data)
+
+        # Get energy cut index:
+        COseq_array = self.cosi_dataset["Compton Seq"]
+        COseq_cut_index = (COseq_array >= seqmin) & (COseq_array < seqmax)
+    
+        # Apply cuts to dictionary:
+        for key in self.cosi_dataset:
+
+            self.cosi_dataset[key] = self.cosi_dataset[key][COseq_cut_index]
+
+        # Write unbinned data to file (either fits or hdf5):
+        if output_name != None:
+            logger.info("Saving file...")
+            self.write_unbinned_output(output_name)
+
+        return
+
+
+        
     def combine_unbinned_data(self, input_files, output_name=None):
 
         """Combines input unbinned data files.

cosipy/test_data/unbinned_data_MEGAlib_calc.hdf5

@@ -94,6 +94,11 @@ def test_unbinned_data_all(tmp_path):
     assert np.amax(analysis.cosi_dataset['Energies']) < 300
     assert np.amin(analysis.cosi_dataset['Energies']) >= 200
 
+    # Test Compton sequence selection method
+    analysis.select_data_COseq(3,4,unbinned_data=tmp_path/"test_h5.hdf5")
+    assert np.amax(analysis.cosi_dataset['Compton Seq']) < 4
+    assert np.amin(analysis.cosi_dataset['Compton Seq']) >= 3
+    
     # Test reading tra with no pointing info:
     analysis.data_file = os.path.join(test_data.path,\
             "GalacticScan.inc1.id1.crab10sec.extracted.testsample.nopointinginfo.tra.gz")

tests/dataIO/test_unbinned_data_all.py

@@ -51,7 +51,8 @@ def test_unbinned_data_with_MEGAlib(tmp_path):
     psi_loc_old = dict_old['Psi local'][:ntestsamples]
     chi_gal_old = dict_old['Chi galactic'][:ntestsamples]
     psi_gal_old = dict_old['Psi galactic'][:ntestsamples]
-
+    CO_seq_old = dict_old['Compton Seq'][:ntestsamples]
+    
     # For comparing chi_loc, psi_loc=0 values are arbitrary,
     # so we exclude them from the comparison.
     psi_zero_index = psi_loc_old == 0
@@ -78,12 +79,13 @@ def test_unbinned_data_with_MEGAlib(tmp_path):
     psi_loc_dict = {"old":psi_loc_old,"new":analysis.psi_loc_test,"name":"psi_loc","units":"rad"}
     chi_gal_dict = {"old":chi_gal_old[~chi_gal_bad_index],"new":analysis.chi_gal_test[~chi_gal_bad_index],"name":"chi_gal","units":"rad"}
     psi_gal_dict = {"old":psi_gal_old,"new":analysis.psi_gal_test,"name":"psi_gal","units":"rad"}
+    CO_seq_dict = {"old":CO_seq_old,"new":analysis.cosi_dataset["Compton Seq"],"name":"Compton Seq","units":""}
 
     # Make comparison:
     print("Comparing to MEGAlib calculation:")
     test_list = [energies_dict,time_dict,phi_dict,\
             dist_dict,lonX_dict,latX_dict,lonZ_dict,latZ_dict,lonY_dict,latY_dict,\
-            chi_loc_dict,psi_loc_dict,chi_gal_dict,psi_gal_dict]
+            chi_loc_dict,psi_loc_dict,chi_gal_dict,psi_gal_dict,CO_seq_dict]
     for each in test_list:
         diff = compare(each["old"],each["new"],each["name"],make_plots=False)
         thresh = 1e-10