Support for gamma-gamma in Pythia and Phojet

src/chromo/common.py

@@ -27,7 +27,7 @@
     standard_projectiles,
 )
 from chromo.decay_handler import Pythia8DecayHandler
-from chromo.kinematics import CompositeTarget, EventKinematics
+from chromo.kinematics import CompositeTarget, EventKinematicsBase
 from chromo.util import (
     Nuclei,
     classproperty,
@@ -144,7 +144,7 @@ class EventData:
 
     generator: (str, str)
         Info about the generator, its name and version.
-    kin: EventKinematics
+    kin: EventKinematicsBase
         Info about initial state.
     nevent: int
         Which event in the sequence.
@@ -186,7 +186,7 @@ class EventData:
     """
 
     generator: Tuple[str, str]
-    kin: EventKinematics
+    kin: EventKinematicsBase
     nevent: int
     impact_parameter: float
     n_wounded: Tuple[int, int]
@@ -760,8 +760,8 @@ def cross_section(self, kin=None):
 
         Parameters
         ----------
-        kin : EventKinematics, optional
-            If provided, calculate cross-section for EventKinematics.
+        kin : EventKinematicsBase, optional
+            If provided, calculate cross-section for EventKinematicsBase.
             Otherwise return values for current setup.
         """
         with self._temporary_kinematics(kin):

src/chromo/kinematics.py

@@ -29,7 +29,8 @@
     "TeV",
     "PeV",
     "EeV",
-    "EventKinematics",
+    "EventKinematicsMassless",
+    "EventKinematicsWithRestframe",
     "CenterOfMass",
     "FixedTarget",
     "TotalEnergy",
@@ -150,8 +151,40 @@ def copy(self):
             self._betagamma_cm,
         )
 
+    def _get_beam_data(self, generator_frame):
+        if self._beam_data is not None:
+            return self._beam_data
+
+        event_like = SimpleNamespace(
+            pz=np.array([self.beams[0][2], self.beams[1][2]]),
+            en=np.array([self.beams[0][3], self.beams[1][3]]),
+        )
+        self.apply_boost(event_like, generator_frame, inverse=True)
+
+        self._beam_data = {
+            "pid": np.array([int(self.p1), int(self.p2)]),
+            "status": np.array([4, 4]),
+            "charge": np.array([self.p1.charge, self.p2.charge]),
+            "px": np.zeros((2,), dtype=np.float64),
+            "py": np.zeros((2,), dtype=np.float64),
+            "pz": event_like.pz,
+            "en": event_like.en,
+            # Note that the masses are from `particle` module:
+            # It may introduce some E/p conservation issues if the internal mass
+            # in the model is too different to that from particle/PDG/Pythia
+            "m": np.array([self.m1, self.m2]),
+            "vx": np.zeros((2,), dtype=np.float64),
+            "vy": np.zeros((2,), dtype=np.float64),
+            "vz": np.zeros((2,), dtype=np.float64),
+            "vt": np.zeros((2,), dtype=np.float64),
+            "mothers": np.array([[-1, -1], [-1, -1]], dtype=np.int32),
+            "daughters": np.array([[-1, -1], [-1, -1]], dtype=np.int32),
+        }
+
+        return self._beam_data
+
 
-class EventKinematics(EventKinematicsBase):
+class EventKinematicsWithRestframe(EventKinematicsBase):
     def __init__(
         self,
         particle1,
@@ -272,42 +305,97 @@ def __init__(
 
         self._beam_data = None
 
-    def _get_beam_data(self, generator_frame):
-        if self._beam_data is not None:
-            return self._beam_data
 
-        event_like = SimpleNamespace(
-            pz=np.array([self.beams[0][2], self.beams[1][2]]),
-            en=np.array([self.beams[0][3], self.beams[1][3]]),
-        )
-        self.apply_boost(event_like, generator_frame, inverse=True)
+class EventKinematicsMassless(EventKinematicsBase):
+    """EventKinematics for massless particles."""
 
-        self._beam_data = {
-            "pid": np.array([int(self.p1), int(self.p2)]),
-            "status": np.array([4, 4]),
-            "charge": np.array([self.p1.charge, self.p2.charge]),
-            "px": np.zeros((2,), dtype=np.float64),
-            "py": np.zeros((2,), dtype=np.float64),
-            "pz": event_like.pz,
-            "en": event_like.en,
-            # Note that the masses are from `particle` module:
-            # It may introduce some E/p conservation issues if the internal mass
-            # in the model is too different to that from particle/PDG/Pythia
-            "m": np.array([self.m1, self.m2]),
-            "vx": np.zeros((2,), dtype=np.float64),
-            "vy": np.zeros((2,), dtype=np.float64),
-            "vz": np.zeros((2,), dtype=np.float64),
-            "vt": np.zeros((2,), dtype=np.float64),
-            "mothers": np.array([[-1, -1], [-1, -1]], dtype=np.int32),
-            "daughters": np.array([[-1, -1], [-1, -1]], dtype=np.int32),
-        }
+    def __init__(
+        self,
+        particle1,
+        particle2,
+        *,
+        ecm=None,
+        beam=None,
+        frame=None,
+    ):
+        # Catch input errors
 
-        return self._beam_data
+        if sum(x is not None for x in [ecm, beam]) != 1:
+            raise ValueError(
+                "Please provide only one of ecm/plab/elab/ekin/beam arguments"
+            )
+
+        if particle1 is None or particle2 is None:
+            raise ValueError("particle1 and particle2 must be set")
+
+        part1 = process_particle(particle1)
+        part2 = process_particle(particle2)
+
+        if (
+            mass(part1) != 0
+            or mass(part2) != 0
+            or is_real_nucleus(part1)
+            or is_real_nucleus(part2)
+        ):
+            raise ValueError("Massless particles required.")
+
+        if isinstance(part1, CompositeTarget) or isinstance(part2, CompositeTarget):
+            raise ValueError("CompositeTarget")
+
+        m1 = mass(part1)
+        m2 = mass(part2)
+
+        assert m1 == 0 and m2 == 0
+
+        beams = (np.zeros(4), np.zeros(4))
+
+        # Input specification in center-of-mass frame
+        if ecm is not None:
+            frame = frame or EventFrame.CENTER_OF_MASS
+            s = ecm**2
+            pcm = ecm / 2
+            beams[0][2] = pcm
+            beams[1][2] = -pcm
+            beams[0][3] = pcm
+            beams[1][3] = pcm
+        # Input specification as 4-vectors
+        elif beam is not None:
+            frame = frame or EventFrame.GENERIC
+            p1, p2 = beam
+            beams[0][2] = p1
+            beams[1][2] = p2
+            beams[0][3] = np.abs(p1)
+            beams[1][3] = np.abs(p2)
+            s = np.sum(beams, axis=0)
+            # See note above
+            ecm = energy2momentum(s[3], s[2])
+
+        else:
+            assert False  # this should never happen
 
+        gamma_cm = (beams[0][3] + beams[1][3]) / ecm
+        betagamma_cm = (beams[0][2] + beams[1][2]) / ecm
+        pcm = ecm / 2
 
-class CenterOfMass(EventKinematics):
-    def __init__(self, ecm, particle1, particle2):
-        super().__init__(ecm=ecm, particle1=particle1, particle2=particle2)
+        # Masses should be zero
+        self.m1 = m1
+        self.m2 = m2
+
+        super().__init__(
+            frame,
+            part1,
+            part2,
+            ecm,
+            pcm,
+            np.nan,
+            np.nan,
+            np.nan,
+            beams,
+            gamma_cm,
+            betagamma_cm,
+        )
+
+        self._beam_data = None
 
 
 class TotalEnergy(TaggedFloat):
@@ -322,22 +410,42 @@ class Momentum(TaggedFloat):
     pass
 
 
-class FixedTarget(EventKinematics):
-    def __init__(self, energy, particle1, particle2):
-        if isinstance(energy, (TotalEnergy, int, float)):
-            super().__init__(
-                elab=float(energy), particle1=particle1, particle2=particle2
-            )
-        elif isinstance(energy, KinEnergy):
-            super().__init__(
-                ekin=float(energy), particle1=particle1, particle2=particle2
-            )
-        elif isinstance(energy, Momentum):
-            super().__init__(
-                plab=float(energy), particle1=particle1, particle2=particle2
-            )
-        else:
-            raise ValueError(
-                f"{energy!r} is neither a number nor one of "
-                "TotalEnergy, KinEnergy, Momentum"
-            )
+def CenterOfMass(ecm, particle1, particle2):
+    """EventKinematics generator for center-of-mass kinematics.
+
+    Accepts all particles including photons and nuclei.
+    """
+    if (
+        mass(process_particle(particle1)) == 0
+        and mass(process_particle(particle2)) == 0
+    ):
+        return EventKinematicsMassless(
+            ecm=ecm, particle1=particle1, particle2=particle2
+        )
+    return EventKinematicsWithRestframe(
+        ecm=ecm, particle1=particle1, particle2=particle2
+    )
+
+
+def FixedTarget(energy, particle1, particle2):
+    """EventKinematics generator for fixed-target kinematics.
+
+    Energy can be specified as TotalEnergy, KinEnergy, or Momentum.
+    """
+    if isinstance(energy, (TotalEnergy, int, float)):
+        return EventKinematicsWithRestframe(
+            elab=float(energy), particle1=particle1, particle2=particle2
+        )
+    elif isinstance(energy, KinEnergy):
+        return EventKinematicsWithRestframe(
+            ekin=float(energy), particle1=particle1, particle2=particle2
+        )
+    elif isinstance(energy, Momentum):
+        return EventKinematicsWithRestframe(
+            plab=float(energy), particle1=particle1, particle2=particle2
+        )
+    else:
+        raise ValueError(
+            f"{energy!r} is neither a number nor one of "
+            "TotalEnergy, KinEnergy, Momentum"
+        )

tests/test_generators.py

@@ -1,7 +1,8 @@
 from chromo.constants import GeV
 from chromo.kinematics import (
     CenterOfMass,
-    EventKinematics,
+    EventKinematicsMassless,
+    EventKinematicsWithRestframe,
     FixedTarget,
     TotalEnergy,
     EventFrame,
@@ -162,9 +163,13 @@ def test_generator(projectile, target, frame, Model):
     elif frame == "ft":
         kin = FixedTarget(TotalEnergy(100 * GeV), p1, p2)
     elif frame == "cms2ft":
-        kin = EventKinematics(p1, p2, ecm=100 * GeV, frame=EventFrame.FIXED_TARGET)
+        kin = EventKinematicsWithRestframe(
+            p1, p2, ecm=100 * GeV, frame=EventFrame.FIXED_TARGET
+        )
     elif frame == "ft2cms":
-        kin = EventKinematics(p1, p2, elab=100 * GeV, frame=EventFrame.CENTER_OF_MASS)
+        kin = EventKinematicsWithRestframe(
+            p1, p2, elab=100 * GeV, frame=EventFrame.CENTER_OF_MASS
+        )
     else:
         assert False  # we should never arrive here
 
@@ -211,3 +216,66 @@ def test_generator(projectile, target, frame, Model):
         pytest.xfail(reason="generated new reference, please check it")
 
     assert p_value >= threshold
+
+
+@pytest.mark.skipif(
+    "CI" in os.environ and platform.system() == "Windows",
+    reason="skip to speed up CI on Windows",
+)
+@pytest.mark.parametrize("frame", ("cms", "generic"))
+@pytest.mark.parametrize(
+    "Model", (im.Pythia8, im.Phojet112, im.Phojet191, im.Phojet193)
+)
+def test_generator_gg(frame, Model):
+    p1 = "gamma"
+    p2 = "gamma"
+
+    if frame == "cms":
+        kin = CenterOfMass(100 * GeV, p1, p2)
+    elif frame == "generic":
+        kin = EventKinematicsMassless(
+            p1, p2, beam=(80, -20), frame=EventFrame.CENTER_OF_MASS
+        )
+    else:
+        assert False  # we should never arrive here
+
+    h = run_in_separate_process(run_model, Model, kin)
+    if h is None:
+        assert abs(kin.p1) not in Model.projectiles or abs(kin.p2) not in Model.targets
+        return
+
+    fn = Path(f"{Model.pyname}_{p1}_{p2}_{frame}")
+    path_ref = REFERENCE_PATH / fn.with_suffix(".pkl.gz")
+
+    reference_generated = False
+    if not path_ref.exists():
+        reference_generated = True
+        # New reference is generated. Check plots to see whether reference makes
+        # any sense before committing it.
+        ref_values = run_in_separate_process(run_model, Model, kin, 50, timeout=10000)
+        val_ref = np.reshape(np.mean(ref_values, axis=0), -1)
+        cov_ref = np.cov(np.transpose([np.reshape(x, -1) for x in ref_values]))
+        with gzip.open(path_ref, "wb") as f:
+            pickle.dump((val_ref, cov_ref), f)
+
+    with gzip.open(path_ref) as f:
+        val_ref, cov_ref = pickle.load(f)
+
+    # histogram sometimes contains extreme spikes
+    # not reflected in cov, this murky formula
+    # protects against these false negatives
+    values = np.reshape(h.values(), -1)
+    for i, v in enumerate(values):
+        cov_ref[i, i] = 0.5 * (cov_ref[i, i] + v)
+
+    p_value = compute_p_value(values, val_ref, cov_ref)
+
+    threshold = 1e-6
+
+    if reference_generated or not (p_value >= threshold) or chromo.debug_level > 0:
+        draw_comparison(fn, p_value, h.axes, values, val_ref, cov_ref)
+
+    if reference_generated:
+        pytest.xfail(reason="generated new reference, please check it")
+
+    assert p_value >= threshold

tests/test_pythia8.py

@@ -173,7 +173,6 @@ def test_gp():
     assert len(evt) > 2
 
 
-@pytest.mark.xfail(reason="no rest frame for gamma-gamma")
 def test_gg():
     evt = run_collision(100 * GeV, "gamma", "gamma")
     assert len(evt) > 2