From 7b74c61a0ea61a6436e9a91f0eeb63870eaab665 Mon Sep 17 00:00:00 2001
From: Johannes Kasimir <johanneskasimir@gmail.com>
Date: Tue, 24 Sep 2024 13:46:45 +0200
Subject: [PATCH] refactor test and some docs

---
 tests/absorption/correction_test.py | 60 ++++++++++++++---------------
 1 file changed, 29 insertions(+), 31 deletions(-)

diff --git a/tests/absorption/correction_test.py b/tests/absorption/correction_test.py
index 3cff9310..8e183520 100644
--- a/tests/absorption/correction_test.py
+++ b/tests/absorption/correction_test.py
@@ -7,6 +7,25 @@
 from scippneutron.absorption import Cylinder, Material, compute_transmission_map
 
 
+def transmission_fraction_case1(effective_attenuation_factor):
+    '''Transmission fraction at a point on the z-axis
+    infinitely far away when the sample is a cylinder
+    symmetrical around y, with 1mm radius and 1mm height.
+    '''
+    return (
+        2
+        / np.pi
+        * quad(
+            lambda x: (
+                np.sqrt(1 - x**2)
+                * np.exp(-2 * effective_attenuation_factor * np.sqrt(1 - x**2))
+            ),
+            -1,
+            1,
+        )[0]
+    )
+
+
 @pytest.mark.parametrize('scattering_cross_section', [0.1, 0.5, 1.0])
 def test_compute_transmission_map(scattering_cross_section):
     material = Material(
@@ -21,20 +40,6 @@ def test_compute_transmission_map(scattering_cross_section):
         sc.scalar(1.0, unit='mm'),
     )
 
-    def transmission_fraction(effective_attenuation_factor):
-        return (
-            2
-            / np.pi
-            * quad(
-                lambda x: (
-                    np.sqrt(1 - x**2)
-                    * np.exp(-2 * effective_attenuation_factor * np.sqrt(1 - x**2))
-                ),
-                -1,
-                1,
-            )[0]
-        )
-
     tm = compute_transmission_map(
         cylinder,
         material,
@@ -45,7 +50,7 @@ def transmission_fraction(effective_attenuation_factor):
     )
     assert_allclose(
         tm['wavelength', 0]['x', 0].data,
-        sc.scalar(transmission_fraction(scattering_cross_section)),
+        sc.scalar(transmission_fraction_case1(scattering_cross_section)),
         rtol=sc.scalar(1e-3),
     )
 
@@ -53,6 +58,7 @@ def transmission_fraction(effective_attenuation_factor):
 @pytest.mark.parametrize('scattering_cross_section', [0.1, 0.5, 1.0])
 def test_compute_transmission_map_wavelength_dependent(scattering_cross_section):
     material = Material(
+        # Modify the absorption cross section
         sc.scalar(scattering_cross_section / 5, unit='mm**2'),
         sc.scalar(0, unit='mm**2'),
         sc.scalar(1.0, unit='1/mm**3'),
@@ -64,32 +70,24 @@ def test_compute_transmission_map_wavelength_dependent(scattering_cross_section)
         sc.scalar(1.0, unit='mm'),
     )
 
-    def transmission_fraction(effective_attenuation_factor):
-        return (
-            2
-            / np.pi
-            * quad(
-                lambda x: (
-                    np.sqrt(1 - x**2)
-                    * np.exp(-2 * effective_attenuation_factor * np.sqrt(1 - x**2))
-                ),
-                -1,
-                1,
-            )[0]
-        )
-
     tm = compute_transmission_map(
         cylinder,
         material,
         beam_direction=sc.vector([0, 0, 1]),
         wavelength=sc.linspace(
-            'wavelength', 5 * 1.7982, 5 * 1.7982, 1, unit='angstrom'
+            # Set wavelength so that it cancels the
+            # modified absorption cross section
+            'wavelength',
+            5 * 1.7982,
+            5 * 1.7982,
+            1,
+            unit='angstrom',
         ),
         detector_position=sc.vectors(dims='x', values=[[0, 0, 1]], unit='m'),
         quadrature_kind='expensive',
     )
     assert_allclose(
         tm['wavelength', 0]['x', 0].data,
-        sc.scalar(transmission_fraction(scattering_cross_section)),
+        sc.scalar(transmission_fraction_case1(scattering_cross_section)),
         rtol=sc.scalar(1e-3),
     )