Merge branch 'w2list' into 'master'

W2list

See merge request !8

km3buu/output.py

@@ -20,7 +20,7 @@ from os.path import isfile, join, abspath
 from tempfile import TemporaryDirectory
 import awkward as ak
 import uproot
-from scipy.interpolate import UnivariateSpline
+from scipy.interpolate import UnivariateSpline, interp1d
 from scipy.spatial.transform import Rotation
 import scipy.constants as constants
 import mendeleev
@@ -102,7 +102,8 @@ EMPTY_KM3NET_HEADER_DICT = {
     "coord_origin": "0 0 0",
     "norma": "0 0",
     "tgen": "0",
-    "simul": ""
+    "simul": "",
+    "primary": "0"
 }
 
 PARTICLE_MC_STATUS = {
@@ -123,6 +124,29 @@ PARTICLE_MC_STATUS = {
     "NucleonClusterTarget": 16
 }
 
+W2LIST_LOOKUP = {
+    "PS": 0,
+    "EG": 1,
+    "XSEC_MEAN": 2,
+    "COLUMN_DEPTH": 3,
+    "P_EARTH": 4,
+    "WATER_INT_LEN": 5,
+    "P_SCALE": 6,
+    "BX": 7,
+    "BY": 8,
+    "ICHAN": 9,
+    "CC": 10,
+    "DISTAMAX": 11,
+    "WATERXSEC": 12,
+    "XSEC": 13,
+    "TARGETA": 15,
+    "TARGETZ": 16,
+    "VERINCAN": 17,
+    "LEPINCAN": 18,
+}
+
+W2LIST_LENGTH = len(W2LIST_LOOKUP)
+
 
 def read_nu_abs_xsection(filepath):
     """
@@ -208,6 +232,34 @@ class GiBUUOutput:
         xsec = np.divide(total_events * weights, n_files)
         return xsec
 
+    @property
+    def mean_xsec(self):
+        root_tupledata = self.arrays
+        energies = np.array(root_tupledata.lepIn_E)
+        weights = self._event_xsec(root_tupledata)
+        Emin = np.min(energies)
+        Emax = np.max(energies)
+        xsec, energy_bins = np.histogram(energies,
+                                         weights=weights,
+                                         bins=np.logspace(
+                                             np.log10(Emin), np.log10(Emax),
+                                             15))
+        deltaE = np.mean(self.flux_data["energy"][1:] -
+                         self.flux_data["energy"][:-1])
+        bin_events = np.array([
+            self.flux_interpolation.integral(energy_bins[i],
+                                             energy_bins[i + 1]) / deltaE
+            for i in range(len(energy_bins) - 1)
+        ])
+        x = (energy_bins[1:] + energy_bins[:-1]) / 2
+        y = xsec / bin_events / x
+        xsec_interp = interp1d(x,
+                               y,
+                               kind="linear",
+                               fill_value=(y[0], y[-1]),
+                               bounds_error=False)
+        return lambda e: xsec_interp(e) * e
+
     def w2weights(self, volume, target_density, solid_angle):
         """
         Calculate w2weights
@@ -451,11 +503,14 @@ def write_detector_file(gibuu_output,
     timestamp = datetime.now()
     header_dct["simul"] = "KM3BUU {} {}".format(
         version, timestamp.strftime("%Y%m%d %H%M%S"))
+    header_dct["primary"] = "{:d}".format(nu_type)
 
     for k, v in header_dct.items():
         head.set_line(k, v)
     head.Write("Head")
 
+    mean_xsec_func = gibuu_output.mean_xsec
+
     for mc_event_id, event in enumerate(event_data):
         evt.clear()
         evt.id = mc_event_id
@@ -464,6 +519,19 @@ def write_detector_file(gibuu_output,
         evt.w.push_back(geometry.volume)  #w1 (can volume)
         evt.w.push_back(w2[mc_event_id])  #w2
         evt.w.push_back(-1.0)  #w3 (= w2*flux)
+        # Event Information (w2list)
+        evt.w2list.resize(W2LIST_LENGTH)
+        evt.w2list[W2LIST_LOOKUP["XSEC_MEAN"]] = mean_xsec_func(event.lepIn_E)
+        evt.w2list[W2LIST_LOOKUP["XSEC"]] = event.xsec
+        evt.w2list[W2LIST_LOOKUP["TARGETA"]] = gibuu_output.A
+        evt.w2list[W2LIST_LOOKUP["TARGETZ"]] = gibuu_output.Z
+        evt.w2list[W2LIST_LOOKUP["BX"]] = bjorkenx[mc_event_id]
+        evt.w2list[W2LIST_LOOKUP["BY"]] = bjorkeny[mc_event_id]
+        evt.w2list[W2LIST_LOOKUP["CC"]] = ichan
+        evt.w2list[W2LIST_LOOKUP["ICHAN"]] = event.evType
+        evt.w2list[W2LIST_LOOKUP["VERINCAN"]] = 1
+        evt.w2list[W2LIST_LOOKUP["LEPINCAN"]] = 1
+
         # Vertex Position
         vtx_pos = np.array(geometry.random_pos())
         # Direction
@@ -507,12 +575,6 @@ def write_detector_file(gibuu_output,
             lep_out_trk.status = PARTICLE_MC_STATUS["StableFinalState"]
             evt.mc_trks.push_back(lep_out_trk)
 
-        # bjorken_y = 1.0 - float(event.lepOut_E / event.lepIn_E)
-        nu_in_trk.setusr('bx', bjorkenx[mc_event_id])
-        nu_in_trk.setusr('by', bjorkeny[mc_event_id])
-        nu_in_trk.setusr('ichan', ichan)
-        nu_in_trk.setusr("cc", is_cc)
-
         evt.mc_trks.push_back(nu_in_trk)
 
         if tau_secondaries is not None:

km3buu/tests/test_output.py

@@ -120,11 +120,13 @@ class TestAANET(unittest.TestCase):
         np.testing.assert_allclose(evt.mc_tracks.pos_y, -122.54421157)
         np.testing.assert_allclose(evt.mc_tracks.pos_z, 208.57726764)
         usr = evt.mc_tracks.usr[0]
+        # XSEC
+        np.testing.assert_almost_equal(evt.w2list[13], 40.62418521597373)
         # Bx
-        np.testing.assert_almost_equal(usr[0], 0.35479262672400624)
+        np.testing.assert_almost_equal(evt.w2list[7], 0.35479262672400624)
         # By
-        np.testing.assert_almost_equal(usr[1], 0.8167222969153614)
+        np.testing.assert_almost_equal(evt.w2list[8], 0.8167222969153614)
         # iChannel
-        np.testing.assert_equal(usr[2], 2)
+        np.testing.assert_equal(evt.w2list[9], 34)
         # CC/NC
-        np.testing.assert_equal(usr[3], 1)
+        np.testing.assert_equal(evt.w2list[10], 2)