Fix sampling

resources/processes/DarkNewsTables/DarkNewsDecay.py

@@ -13,6 +13,158 @@
 from siren import dataclasses
 from siren.dataclasses import Particle
 
+# DarkNews methods
+import DarkNews
+from DarkNews.processes import FermionDileptonDecay, FermionSinglePhotonDecay
+from DarkNews import processes as proc
+from DarkNews import Cfourvec as Cfv
+from DarkNews import phase_space
+
+def get_decay_momenta_from_vegas_samples(vsamples, MC_case, decay_case, PN_LAB):
+    """
+    Construct the four momenta of all final state particles in the decay process from the
+    vegas weights.
+
+    Args:
+            vsamples (np.ndarray): integration samples obtained from vegas
+                            as hypercube coordinates. Always in the interval [0,1].
+
+            MC_case (DarkNews.process.dec_case): the decay class of DarkNews
+
+            PN_LAB (np.ndarray): four-momentum of the upscattered N in the lab frame: [E, pX, pY, pZ]
+
+    Returns:
+            dict: each key corresponds to a set of four momenta for a given particle involved,
+                    so the values are 2D np.ndarrays with each row a different event and each column a different
+                    four momentum component. Contains also the weights.
+    """
+
+    four_momenta = {}
+
+    # N boost parameters
+    boost_scattered_N = {
+        "EP_LAB": PN_LAB.T[0],
+        "costP_LAB": Cfv.get_cosTheta(PN_LAB),
+        "phiP_LAB": np.arctan2(PN_LAB.T[2], PN_LAB.T[1]),
+    }
+
+    #######################
+    # DECAY PROCESSES
+
+    if type(decay_case) == proc.FermionDileptonDecay:
+
+        mh = decay_case.m_parent
+        mf = decay_case.m_daughter
+        mm = decay_case.mm
+        mp = decay_case.mm
+
+        if decay_case.vector_on_shell or decay_case.scalar_on_shell:
+
+            if decay_case.vector_on_shell and decay_case.scalar_off_shell:
+                m_mediator = decay_case.mzprime
+            elif decay_case.vector_off_shell and decay_case.scalar_on_shell:
+                m_mediator = decay_case.mhprime
+            else:
+                raise NotImplementedError("Both mediators on-shell is not yet implemented.")
+
+            ########################
+            ### HNL decay
+            N_decay_samples = {"unit_cost": np.array(vsamples[0])}
+            # Ni (k1) --> Nj (k2)  Z' (k3)
+            masses_decay = {
+                "m1": mh,  # Ni
+                "m2": mf,  # Nj
+                "m3": m_mediator,  # Z'
+            }
+            # Phnl, Phnl_daughter, Pz'
+            P1LAB_decay, P2LAB_decay, P3LAB_decay = phase_space.two_body_decay(N_decay_samples, boost=boost_scattered_N, **masses_decay, rng=MC_case.rng)
+
+            # Z' boost parameters
+            boost_Z = {
+                "EP_LAB": P3LAB_decay.T[0],
+                "costP_LAB": Cfv.get_cosTheta(P3LAB_decay),
+                "phiP_LAB": np.arctan2(P3LAB_decay.T[2], P3LAB_decay.T[1]),
+            }
+
+            ########################
+            ### Z' decay
+            Z_decay_samples = {}  # all uniform
+            # Z'(k1) --> ell- (k2)  ell+ (k3)
+            masses_decay = {
+                "m1": m_mediator,  # Ni
+                "m2": mp,  # \ell+
+                "m3": mm,  # \ell-
+            }
+            # PZ', pe-, pe+
+            P1LAB_decayZ, P2LAB_decayZ, P3LAB_decayZ = phase_space.two_body_decay(Z_decay_samples, boost=boost_Z, **masses_decay, rng=MC_case.rng)
+
+            four_momenta["P_decay_N_parent"] = P1LAB_decay
+            four_momenta["P_decay_N_daughter"] = P2LAB_decay
+            four_momenta["P_decay_ell_minus"] = P2LAB_decayZ
+            four_momenta["P_decay_ell_plus"] = P3LAB_decayZ
+
+        elif decay_case.vector_off_shell and decay_case.scalar_off_shell:
+
+            ########################
+            # HNL decay
+            N_decay_samples = {
+                "unit_t": vsamples[0],
+                "unit_u": vsamples[1],
+                "unit_c3": vsamples[2],
+                "unit_phi34": vsamples[3],
+            }
+
+            # Ni (k1) --> ell-(k2)  ell+(k3)  Nj(k4)
+            masses_decay = {
+                "m1": mh,  # Ni
+                "m2": mm,  # ell-
+                "m3": mp,  # ell+
+                "m4": mf,
+            }  # Nj
+            # Phnl, pe-, pe+, pnu
+            (
+                P1LAB_decay,
+                P2LAB_decay,
+                P3LAB_decay,
+                P4LAB_decay,
+            ) = phase_space.three_body_decay(N_decay_samples, boost=boost_scattered_N, **masses_decay, rng=MC_case.rng)
+
+            four_momenta["P_decay_N_parent"] = P1LAB_decay
+            four_momenta["P_decay_ell_minus"] = P2LAB_decay
+            four_momenta["P_decay_ell_plus"] = P3LAB_decay
+            four_momenta["P_decay_N_daughter"] = P4LAB_decay
+
+    elif type(decay_case) == proc.FermionSinglePhotonDecay:
+
+        mh = decay_case.m_parent
+        mf = decay_case.m_daughter
+
+        ########################
+        ### HNL decay
+        N_decay_samples = {"unit_cost": np.array(vsamples[0])}
+        # Ni (k1) --> Nj (k2)  gamma (k3)
+        masses_decay = {
+            "m1": mh,  # Ni
+            "m2": mf,  # Nj
+            "m3": 0.0,  # gamma
+        }
+        # Phnl, Phnl', Pgamma
+        P1LAB_decay, P2LAB_decay, P3LAB_decay = phase_space.two_body_decay(N_decay_samples, boost=boost_scattered_N, **masses_decay, rng=MC_case.rng)
+
+        four_momenta["P_decay_N_parent"] = P1LAB_decay
+        four_momenta["P_decay_N_daughter"] = P2LAB_decay
+        four_momenta["P_decay_photon"] = P3LAB_decay
+
+    return four_momenta
+
+
+class _FakeMCInterface:
+    def __init__(self, random):
+        self.random = random
+        self.rng_func = np.frompyfunc(lambda x: self.random.Uniform(0, 1), 1, 1)
+        self.rng = lambda x: self.rng_func(np.empty(x)).astype(float)
+
+
 # A class representing a single decay_case DarkNews class
 # Only handles methods concerning the decay part
 class PyDarkNewsDecay(DarkNewsDecay):
@@ -267,11 +419,12 @@ def SampleRecordFromDarkNews(self, record, random):
         # Expand dims required to call DarkNews function on signle sample
         four_momenta = get_decay_momenta_from_vegas_samples(
             np.expand_dims(PS, 0),
+            _FakeMCInterface(random),
             self.dec_case,
             np.expand_dims(np.array(record.primary_momentum), 0),
         )
 
-        secondaries = record.GetSecondaryParticleRecords()
+        secondaries = record.get_secondary_particle_records()
 
         if isinstance(self.dec_case, FermionSinglePhotonDecay):
             gamma_idx = 0