Merge pull request #140 from lucabaldini/uniform_spectrum

Uniform spectrum

Author: augustocattafesta

docs/release_notes.rst

@@ -3,9 +3,11 @@
 Release notes
 =============
 
+* Uniform spectrum added to the `source.py` module.
 * Minor bug fixes to equalization calibration.
 * Pull requests merged and issues closed:
 
+  - https://github.com/lucabaldini/hexsample/pull/140
   - https://github.com/lucabaldini/hexsample/pull/137
 
 

src/hexsample/cli.py

@@ -348,6 +348,10 @@ def add_source_options(parser: argparse.ArgumentParser) -> None:
                            help="element generating the line forest")
         group.add_argument("--initial_level", type=str, default=source.LineForest.initial_level,
                            help="initial level for the line forest")
+        group.add_argument("--emin", type=float, default=source.UniformSpectrum.emin,
+                           help="minimum energy of the uniform spectrum in eV")
+        group.add_argument("--emax", type=float, default=source.UniformSpectrum.emax,
+                           help="maximum energy of the uniform spectrum in eV")
         # ... morphological part...
         group.add_argument(f"--{source.BeamProxy.key()}", type=str,
                            choices=source.BeamProxy.choices(),

src/hexsample/hexagon.py

@@ -125,65 +125,65 @@ def neighbors_even_q(col: int, row: int) -> tuple:
 
 
 _N_ADC_CHANNELS = 7
-_ADC_SEQUENCE_EVEN = (0, 2, 5, 0, 3, 5, 1, 3, 6, 1, 4, 6, 2, 4)
-_ADC_SEQUENCE_ODD = (0, 3, 5, 1, 3, 6, 1, 4, 6, 2, 4, 0, 2, 5)
+_ADC_SEQUENCE_EVEN = np.array([0, 2, 5, 0, 3, 5, 1, 3, 6, 1, 4, 6, 2, 4])
+_ADC_SEQUENCE_ODD = np.array([0, 3, 5, 1, 3, 6, 1, 4, 6, 2, 4, 0, 2, 5])
 _ADC_SEQUENCE_LENGTH = len(_ADC_SEQUENCE_EVEN)
 
 
-def adc_channel_odd_r(col: int, row: int) -> int:
+def adc_channel_odd_r(col: np.ndarray, row: np.ndarray) -> np.ndarray:
     """Transformation from offset coordinates (col, row) into 7-adc channel label,
     that is an int between 0 and 6, for ODD_R grid layout.
 
     Arguments
     ---------
-    col: int
+    col: array_like
         column pixel logical coordinate
-    row: int
+    row: array_like
         row pixel logical coordinate
     """
     start = _ADC_SEQUENCE_ODD[row % _ADC_SEQUENCE_LENGTH]
     index = (col + start) % _N_ADC_CHANNELS
     return index
 
-def adc_channel_even_r(col: int, row: int) -> int:
+def adc_channel_even_r(col: np.ndarray, row: np.ndarray) -> np.ndarray:
     """Transformation from offset coordinates (col, row) into 7-adc channel label,
     that is an int between 0 and 6, for EVEN_R grid layout.
 
     Arguments
     ---------
-    col: int
+    col: array_like
         column pixel logical coordinate
-    row: int
+    row: array_like
         row pixel logical coordinate
     """
     start = _ADC_SEQUENCE_EVEN[row % _ADC_SEQUENCE_LENGTH]
     index = (col + start) % _N_ADC_CHANNELS
     return index
 
-def adc_channel_odd_q(col: int, row: int) -> int:
+def adc_channel_odd_q(col: np.ndarray, row: np.ndarray) -> np.ndarray:
     """Transformation from offset coordinates (col, row) into 7-adc channel label,
     that is an int between 0 and 6, for ODD_Q grid layout.
 
     Arguments
     ---------
-    col: int
+    col: array_like
         column pixel logical coordinate
-    row: int
+    row: array_like
         row pixel logical coordinate
     """
     start = _ADC_SEQUENCE_ODD[col % _ADC_SEQUENCE_LENGTH]
     index = (row + start) % _N_ADC_CHANNELS
     return index
 
-def adc_channel_even_q(col: int, row: int) -> int:
+def adc_channel_even_q(col: np.ndarray, row: np.ndarray) -> np.ndarray:
     """Transformation from offset coordinates (col, row) into 7-adc channel label,
     that is an int between 0 and 6, for EVEN_Q grid layout.
 
     Arguments
     ---------
-    col: int
+    col: array_like
         column pixel logical coordinate
-    row: int
+    row: array_like
         row pixel logical coordinate
     """
     start = _ADC_SEQUENCE_EVEN[col % _ADC_SEQUENCE_LENGTH]

src/hexsample/source.py

@@ -44,6 +44,7 @@
     "SpectrumProxy",
     "BeamProxy",
     "Source",
+    "UniformSpectrum",
 ]
 
 
@@ -140,6 +141,37 @@ def render(self, axes: matplotlib.axes.Axes, **kwargs) -> None:
         setup_gca(xlabel="Energy [eV]", ylabel="Relative intensity", logy=True, grids=True)
 
 
+@dataclass
+class UniformSpectrum(AbstractSpectrum):
+
+    """Class describing a uniform X-ray energy spectrum between a minimum
+    and a maximum energy.
+
+    Arguments
+    ---------
+    emin : float
+        The minimum energy in eV.
+
+    emax : float
+        The maximum energy in eV.
+    """
+
+    emin: float = 1000.
+    emax: float = 10000.
+
+    def rvs(self, size: int = 1) -> np.ndarray:
+        """Overloaded method.
+        """
+        return rng.generator.uniform(self.emin, self.emax, size=size)
+
+    def render(self, axes: matplotlib.axes.Axes, **kwargs) -> None:
+        """Overloaded method.
+        """
+        kwargs.setdefault("color", "black")
+        axes.hlines(1., self.emin, self.emax, **kwargs)
+        setup_gca(xlabel="Energy [eV]", ylabel="Relative intensity", grids=True)
+
+
 @dataclass
 class AbstractBeam(AbstractRandomGenerator):
 
@@ -432,6 +464,7 @@ def rvs(self, size: int = 1) -> Tuple[np.ndarray, np.ndarray]:
 SpectrumProxy = TypeProxy("spectrum") # pylint: disable=invalid-name
 SpectrumProxy.register("line", Line, default=True)
 SpectrumProxy.register("forest", LineForest)
+SpectrumProxy.register("uniform", UniformSpectrum)
 
 
 # Definition of the type proxies for beam types.

tests/test_source.py

@@ -35,6 +35,7 @@
     Source,
     SquareBeam,
     TriangularBeam,
+    UniformSpectrum,
 )
 
 rng.initialize()
@@ -97,6 +98,17 @@ def test_mn_k_forest():
     _test_forest("Mn", chisq_test=False)
 
 
+def test_uniform_spectrum(emin: float = 1000., emax: float = 10000., num_events: int = 10000):
+    """Test the uniform spectrum.
+    """
+    spectrum = UniformSpectrum(emin=emin, emax=emax)
+    rvs = spectrum.rvs(num_events)
+    assert len(rvs) == num_events
+    assert np.all((rvs >= emin) & (rvs <= emax))
+    plt.figure("Uniform spectrum")
+    spectrum.plot()
+
+
 def test_point_beam(x0 : float = 1., y0 : float = -1., num_photons : int = 1000):
     """Unit test for the point beam.
     """