From 69e5576da29d7e109bcbcb415f590808dc488954 Mon Sep 17 00:00:00 2001
From: joeb42 <70596457+joeb42@users.noreply.github.com>
Date: Sun, 29 Oct 2023 19:35:11 +0000
Subject: [PATCH] add tests

---
 src/drift_chamber/driftchamber.py | 15 ++++++++----
 src/drift_chamber/muon.py         |  8 +++----
 tests/test_muon.py                | 39 +++++++++++++++++++++++++++++++
 3 files changed, 53 insertions(+), 9 deletions(-)
 create mode 100644 tests/test_muon.py

diff --git a/src/drift_chamber/driftchamber.py b/src/drift_chamber/driftchamber.py
index 06784d9..9d60f18 100644
--- a/src/drift_chamber/driftchamber.py
+++ b/src/drift_chamber/driftchamber.py
@@ -31,7 +31,14 @@ class DriftChamber:
     D: numerical diffusivity
     """
 
-    def __init__(self, spacing: float = 1.0, E: float = 1e5, D: float = 0.1, tau: float = 1e-6, mat: bool = True) -> None:
+    def __init__(
+        self,
+        spacing: float = 1.0,
+        E: float = 1e5,
+        D: float = 0.1,
+        tau: float = 1e-6,
+        mat: bool = True,
+    ) -> None:
         """
         Contructor initialises class attributes as well as charge grid and drift diffusion matrix (also lu factorises).
         Takes positive numerical spacing, numerical E (electric field), positive numerical D (diffusivity)
@@ -108,7 +115,7 @@ def initial_ion(self, muon: Muon) -> None:
                 j = math.floor(y)
             i = math.floor(z)
             # Only take into account energy losses so take negative losses as 0 ionisation
-            self.grid[i, j] = max(n_electrons / (self.spacing ** 2), 0)
+            self.grid[i, j] = max(n_electrons / (self.spacing**2), 0)
 
     @property
     def lu(self):
@@ -164,7 +171,7 @@ def __init(self) -> list[np.ndarray]:
         self.im.set_data(np.zeros(self.grid.shape))
         return [self.im]
 
-    def __ani(self, i: int) -> list[np.ndarray]: 
+    def __ani(self, i: int) -> list[np.ndarray]:
         """
         Private method updates animation at each frame. Calls drift_diff method at each frame.
         """
@@ -173,7 +180,7 @@ def __ani(self, i: int) -> list[np.ndarray]:
         self.fig.suptitle(f"t = {round(i * self.tau, 6)} seconds")
         return [self.im]
 
-    def animate(self, vmin: float = 1e-4, vmax: float = 1.) -> None:
+    def animate(self, vmin: float = 1e-4, vmax: float = 1.0) -> None:
         """
         Plots a Matplotlib funcanimation of charge drift-diffusion.
         Takes positive numerical vmin and vmax arg as min, max values for log scale colour normalisation.
diff --git a/src/drift_chamber/muon.py b/src/drift_chamber/muon.py
index 2d86cfd..b0995ba 100644
--- a/src/drift_chamber/muon.py
+++ b/src/drift_chamber/muon.py
@@ -14,6 +14,8 @@ class Muon:
     height: height of plane above detector in cm (vertical direction in relation to diagram above). Must be >= 30 (height of detector)
     """
 
+    POSITIVE_CHARGE_PROBABILITY = 0.53
+
     def __init__(self, width: int = 100, length: int = 80, height: int = 60) -> None:
         """
         Initialises random muon attributes and takes horizontal dimensions and height of plane above detector
@@ -21,17 +23,13 @@ def __init__(self, width: int = 100, length: int = 80, height: int = 60) -> None
         """
         self.energy: float = np.random.lognormal(6.55, 1.8)
         r = np.random.random()
-        self.charge: int = 1 if r > 0.43 else -1
+        self.charge: int = 1 if r > (1 - Muon.POSITIVE_CHARGE_PROBABILITY) else -1
         # Accept/reject method to get random zenith angle, distributed proportional to cos squared
         x, y = np.pi * np.random.random() + 0.5 * np.pi, np.random.random()
         while y >= (np.cos(x)) ** 2:
             x, y = np.pi * np.random.random() + 0.5 * np.pi, np.random.random()
         self.zen: float = x
         self.azi: float = 2 * np.pi * np.random.random()
-        if height < 30:
-            raise Exception(
-                f"{height} arg is too small! Must be > 30 as muons must start above detector."
-            )
         self.height = height
         self.x_coord: float = width * np.random.random() - 0.5 * width
         self.y_coord: float = length * np.random.random() - 0.5 * length + 25
diff --git a/tests/test_muon.py b/tests/test_muon.py
new file mode 100644
index 0000000..43a60c7
--- /dev/null
+++ b/tests/test_muon.py
@@ -0,0 +1,39 @@
+import sys
+from unittest.mock import MagicMock
+sys.modules["drift_chamber.widget"] = MagicMock()
+
+import numpy as np
+import pytest
+from pytest_mock import MockerFixture
+
+from drift_chamber import Muon
+
+
+def test_init(mocker):
+    mock_log_normal = mocker.patch.object(np.random, "lognormal")
+    mock_cos = mocker.patch.object(np, "cos")
+    mock_cos.return_value = 1
+    mock_random = mocker.patch.object(np.random, "random")
+    zen = 1
+    r = 0
+    azi = 1
+    x_coord = 1
+    y_coord = 1
+    charge_random_var = 0
+    mock_random.side_effect = [
+        charge_random_var,
+        (zen - 0.5 * np.pi) / np.pi,
+        r,
+        azi / (2 * np.pi),
+        x_coord,
+        y_coord
+    ]
+    muon = Muon()
+    mock_log_normal.assert_called_once_with(6.55, 1.8)
+    mock_cos.assert_called_once_with(1.0)
+    assert mock_random.call_count == 6
+    assert muon.energy == mock_log_normal.return_value
+    assert muon.zen == zen
+    assert muon.azi == azi
+    assert muon.x_coord == x_coord * 100 - 50 
+    assert muon.y_coord == y_coord * 80 - 15