From 8d3a4d4270339460a473744cafca0e3aa6946975 Mon Sep 17 00:00:00 2001
From: Jacan Chaplais <jacan.chaplais@gmail.com>
Date: Thu, 31 Aug 2023 12:50:25 +0100
Subject: [PATCH] added tests for MomentumArray #143

---
 .github/workflows/tests.yml | 31 +++++++++++++++++
 tests/test_data.py          | 67 +++++++++++++++++++++++++++++++++----
 2 files changed, 92 insertions(+), 6 deletions(-)
 create mode 100644 .github/workflows/tests.yml

diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml
new file mode 100644
index 0000000..44736e0
--- /dev/null
+++ b/.github/workflows/tests.yml
@@ -0,0 +1,31 @@
+name: tests
+
+on:
+  - push
+
+jobs:
+  test:
+    runs-on: ${{ matrix.os }}
+    strategy:
+      matrix:
+        os: [ubuntu-latest, macos-latest, windows-latest]
+        pyver: ["3.8", "3.9", "3.10", "3.11"]
+    steps:
+      - uses: actions/checkout@v3
+
+      - name: Install Conda environment from environment.yml
+        uses: mamba-org/provision-with-micromamba@main
+        with:
+          cache-env: true
+          extra-specs: |
+            python=${{ matrix.pyver }}
+            pytest=6.2.5
+            hypothesis=6.62.0
+
+      - name: Install graphicle
+        shell: bash -l {0}
+        run: pip install .
+
+      - name: Run tests
+        shell: bash -l {0}
+        run: pytest
diff --git a/tests/test_data.py b/tests/test_data.py
index b291f01..ca00a48 100644
--- a/tests/test_data.py
+++ b/tests/test_data.py
@@ -1,13 +1,68 @@
+import cmath
+import dataclasses as dc
+import math
+
 import numpy as np
-import typicle
 
-import graphicle
+import graphicle as gcl
+
+ZERO_TOL = 1.0e-10  # absolute tolerance for detecting zero-values
+
+
+@dc.dataclass
+class MomentumExample:
+    """Dataclass for representing the four-momentum of a particle, based
+    on its cylindrical beamline coordinates.
+
+    Parameters
+    ----------
+    pt, phi, pz, energy : float
+        Components of the four-momentum in cylindrical coordinates.
+        Note that ``phi`` is given in units of pi.
+    """
+
+    pt: float = 3.0
+    phi: float = 0.75
+    pz: float = 4.0
+    energy: float = 5.0
 
+    @property
+    def _polar(self) -> complex:
+        return self.pt * cmath.exp(complex(0, self.phi * math.pi))
 
-types_ = typicle.Types()
+    @property
+    def px(self) -> float:
+        """x-component of four momentum."""
+        return self._polar.real
+
+    @property
+    def py(self) -> float:
+        """y-component of four momentum."""
+        return self._polar.imag
+
+    def to_momentum_array(self) -> gcl.MomentumArray:
+        """Formats the components in a ``MomentumArray``."""
+        return gcl.MomentumArray([(self.px, self.py, self.pz, self.energy)])
 
 
 def test_pdgs():
-    pdg_vals = np.arange(1, 7, dtype=types_.int)
-    pdgs = graphicle.PdgArray(pdg_vals)
-    assert list(pdgs.name) == ["d", "u", "s", "c", "b", "t"]
+    pdg_vals = np.arange(1, 7, dtype=np.int32)
+    pdgs = gcl.PdgArray(pdg_vals)
+    assert pdgs.name.tolist() == ["d", "u", "s", "c", "b", "t"]
+
+
+def test_pmu_coords() -> None:
+    momentum = MomentumExample()
+    pmu = momentum.to_momentum_array()
+    assert math.isclose(pmu.pt.item(), momentum.pt)
+    assert math.isclose(pmu.phi.item(), momentum.phi * math.pi)
+    assert math.isclose(pmu.mass.item(), 0.0, abs_tol=ZERO_TOL)
+    assert math.isclose(pmu.theta.item(), math.atan(momentum.pt / momentum.pz))
+
+
+def test_pmu_transform() -> None:
+    momentum = MomentumExample()
+    pmu = momentum.to_momentum_array()
+    zero_transverse = pmu.shift_phi(math.pi) + pmu
+    assert math.isclose(0.0, zero_transverse.pt.item(), abs_tol=ZERO_TOL)
+    assert math.isclose(zero_transverse.mass.item(), 6.0)