Feature/model attr access

CHANGELOG.md

@@ -4,6 +4,9 @@ Project Changelog
 Release 1.6.0 (TBD)
 -------------------
 
+API changes:
+* Add emission model attribute access to line and lineshape . (#294)
+
 New:
 * Add Function6D framework. (#478)
 * Add e_field attribute to Plasma object for electric field vector. (#465)

cherab/core/atomic/line.pyx

@@ -33,6 +33,10 @@ cdef class Line:
       specify the n-levels with integers (e.g. (3,2)). For all other ions the full spectroscopic
       configuration string should be specified for both states. It is up to the atomic data
       provider package to define the exact notation.
+
+    :ivar Element element: See parameter 'element'.
+    :ivar int charge: See parameter 'charge'.
+    :ivar tuple transition: See parameter 'transition'.
 
     .. code-block:: pycon
 

cherab/core/atomic/tests/test_line.py

@@ -0,0 +1,27 @@
+import unittest
+
+from cherab.core.atomic import Line, deuterium
+
+
+class TestLine(unittest.TestCase):
+
+    def test_initialisation(self):
+        line = Line(deuterium, 0, (3, 2))
+        self.assertEqual(line.element, deuterium)
+        self.assertEqual(line.charge, 0)
+        self.assertEqual(line.transition, (3, 2))
+
+        # test invalid charge
+        with self.assertRaises(ValueError):
+            Line(deuterium, 2, (3, 2))
+        with self.assertRaises(ValueError):
+            Line(deuterium, -1, (3, 2))
+
+    def test_properties(self):
+        element = deuterium
+        charge = 0
+        transition = (3, 2)
+        line = Line(element, charge, transition)
+        self.assertEqual(line.element, element)
+        self.assertEqual(line.charge, charge)
+        self.assertEqual(line.transition, transition)

cherab/core/model/plasma/impact_excitation.pyx

@@ -45,8 +45,10 @@ cdef class ExcitationLine(PlasmaModel):
     :param object lineshape_args: A list of line shape model arguments. Default is None.
     :param object lineshape_kwargs: A dictionary of line shape model keyword arguments. Default is None.
 
-    :ivar Plasma plasma: The plasma to which this emission model is attached.
-    :ivar AtomicData atomic_data: The atomic data provider for this model.
+    :ivar Plasma plasma: See parameter 'plasma'.
+    :ivar AtomicData atomic_data: See parameter 'atomic_data'.
+    :ivar Line line: The emission line object.
+    :ivar LineShapeModel lineshape: The line shape model.
     """
 
     def __init__(self, Line line, Plasma plasma=None, AtomicData atomic_data=None, object lineshape=None,
@@ -75,6 +77,14 @@ cdef class ExcitationLine(PlasmaModel):
     def __repr__(self):
         return '<ExcitationLine: element={}, charge={}, transition={}>'.format(self._line.element.name, self._line.charge, self._line.transition)
 
+    @property
+    def line(self) -> Line:
+        return self._line
+
+    @property
+    def lineshape(self) -> LineShapeModel:
+        return self._lineshape
+
     cpdef Spectrum emission(self, Point3D point, Vector3D direction, Spectrum spectrum):
 
         cdef double ne, ni, te, radiance

cherab/core/model/plasma/recombination.pyx

@@ -45,8 +45,10 @@ cdef class RecombinationLine(PlasmaModel):
     :param object lineshape_args: A list of line shape model arguments. Default is None.
     :param object lineshape_kwargs: A dictionary of line shape model keyword arguments. Default is None.
 
-    :ivar Plasma plasma: The plasma to which this emission model is attached.
-    :ivar AtomicData atomic_data: The atomic data provider for this model.
+    :ivar Plasma plasma: See parameter 'plasma'.
+    :ivar AtomicData atomic_data: See parameter 'atomic_data'.
+    :ivar Line line: The emission line object.
+    :ivar LineShapeModel lineshape: The line shape model.
     """
 
     def __init__(self, Line line, Plasma plasma=None, AtomicData atomic_data=None, object lineshape=None,
@@ -75,6 +77,14 @@ cdef class RecombinationLine(PlasmaModel):
     def __repr__(self):
         return '<RecombinationLine: element={}, charge={}, transition={}>'.format(self._line.element.name, self._line.charge, self._line.transition)
 
+    @property
+    def line(self) -> Line:
+        return self._line
+
+    @property
+    def lineshape(self) -> LineShapeModel:
+        return self._lineshape
+
     cpdef Spectrum emission(self, Point3D point, Vector3D direction, Spectrum spectrum):
 
         cdef double ne, ni, te, radiance

cherab/core/model/plasma/tests/test_models.py

@@ -0,0 +1,105 @@
+import unittest
+from unittest.mock import patch
+
+import numpy as np
+
+from raysect.optical import Point3D, Vector3D, Spectrum
+
+from cherab.core.model.plasma import (
+    ExcitationLine,
+    RecombinationLine,
+    TotalRadiatedPower,
+)
+from cherab.core.atomic import Line, hydrogen
+from cherab.core.model import GaussianLine
+from cherab.tools.plasmas.slab import build_slab_plasma
+from cherab.openadas import OpenADAS
+
+
+class TestPlasmaModels(unittest.TestCase):
+    # make a slab plasma
+
+    plasma = build_slab_plasma(peak_density=5e19)
+    plasma.atomic_data = OpenADAS(permit_extrapolation=True)
+    balmer_alpha = Line(hydrogen, 0, (3, 2))
+
+    def setUp(self):
+        # setup mock to avoid reading the data from the repository
+        self.patcher_excitation = patch(
+            "cherab.openadas.openadas.repository.get_pec_excitation_rate",
+            return_value={
+                "ne": np.linspace(1e18, 1e20, 10),
+                "te": np.linspace(1, 1e3, 12),
+                "rate": np.ones((10, 12)),
+            },
+        )
+        self.mock_get_excitation = self.patcher_excitation.start()
+
+        self.patcher_recombination = patch(
+            "cherab.openadas.openadas.repository.get_pec_recombination_rate",
+            return_value={
+                "ne": np.linspace(1e18, 1e20, 10),
+                "te": np.linspace(1, 1e3, 12),
+                "rate": np.ones((10, 12)),
+            },
+        )
+        self.mock_get_recombination = self.patcher_recombination.start()
+
+        self.patcher_wl = patch(
+            "cherab.openadas.openadas.repository.get_wavelength", return_value=656.28
+        )
+        self.mock_get_wavelength = self.patcher_wl.start()
+
+    def tearDown(self):
+        # stop the mocks after a test is run
+        self.patcher_excitation.stop()
+        self.patcher_recombination.stop()
+        self.patcher_wl.stop()
+
+    def test_excitation(self):
+        exc = ExcitationLine(self.balmer_alpha)
+        self.plasma.models = [exc]
+
+        # sample emission to trigger the caching mechanism
+        exc.emission(Point3D(0, 0, 0), Vector3D(0, 0, 1), Spectrum(300, 1000, 1000))
+
+        # check exc has the correct line
+        self.assertEqual(exc.line, self.balmer_alpha)
+
+        # check exc has the correct lineshape
+        self.assertIsInstance(exc.lineshape, GaussianLine)
+
+        # check the mock was called
+        self.mock_get_excitation.assert_called_once()
+        self.assertEqual(self.mock_get_wavelength.call_count, 2)
+
+    def test_recombination(self):
+        rec = RecombinationLine(self.balmer_alpha)
+        self.plasma.models = [rec]
+
+        # sample emission to trigger the caching mechanism
+        rec.emission(Point3D(0, 0, 0), Vector3D(0, 0, 1), Spectrum(300, 1000, 1000))
+
+        # check rec has the correct line
+        self.assertEqual(rec.line, self.balmer_alpha)
+
+        # check rec has the correct lineshape
+        self.assertIsInstance(rec.lineshape, GaussianLine)
+
+        # check the mock was called
+        self.mock_get_recombination.assert_called_once()
+        self.assertEqual(self.mock_get_wavelength.call_count, 2)
+
+    def test_total_radiated_power(self):
+        trp = TotalRadiatedPower(hydrogen, 0)
+        self.plasma.models = [trp]
+
+        # check initialisation
+        with self.assertRaises(ValueError):
+            TotalRadiatedPower(hydrogen, 2)
+        with self.assertRaises(ValueError):
+            TotalRadiatedPower(hydrogen, -1)
+
+        # check trp has the correct element and charge
+        self.assertEqual(trp.element, hydrogen)
+        self.assertEqual(trp.charge, 0)

cherab/core/model/plasma/thermal_cx.pyx

@@ -47,8 +47,10 @@ cdef class ThermalCXLine(PlasmaModel):
     :param object lineshape_args: A list of line shape model arguments. Default is None.
     :param object lineshape_kwargs: A dictionary of line shape model keyword arguments. Default is None.
 
-    :ivar Plasma plasma: The plasma to which this emission model is attached.
-    :ivar AtomicData atomic_data: The atomic data provider for this model.
+    :ivar Plasma plasma: See parameter 'plasma'.
+    :ivar AtomicData atomic_data: See parameter 'atomic_data'.
+    :ivar Line line: The emission line object.
+    :ivar LineShapeModel lineshape: The line shape model.
     """
 
     def __init__(self, Line line, Plasma plasma=None, AtomicData atomic_data=None, object lineshape=None,
@@ -77,6 +79,14 @@ cdef class ThermalCXLine(PlasmaModel):
     def __repr__(self):
         return '<ThermalCXLine: element={}, charge={}, transition={}>'.format(self._line.element.name, self._line.charge, self._line.transition)
 
+    @property
+    def line(self) -> Line:
+        return self._line
+
+    @property
+    def lineshape(self) -> LineShapeModel:
+        return self._lineshape
+
     cpdef Spectrum emission(self, Point3D point, Vector3D direction, Spectrum spectrum):
 
         cdef:

cherab/core/model/plasma/total_radiated_power.pyx

@@ -53,6 +53,9 @@ cdef class TotalRadiatedPower(PlasmaModel):
     :param int charge: The charge state of the element/isotope.
     :param Plasma plasma: The plasma to which this emission model is attached. Default is None.
     :param AtomicData atomic_data: The atomic data provider for this model. Default is None.
+
+    :ivar Element element: See parameter 'element'.
+    :ivar int charge: See parameter 'charge'.
     """
 
     def __init__(self, Element element, int charge, Plasma plasma=None, AtomicData atomic_data=None):
@@ -68,6 +71,14 @@ cdef class TotalRadiatedPower(PlasmaModel):
 
         # ensure that cache is initialised
         self._change()
+
+    @property
+    def element(self) -> Element:
+        return self._element
+
+    @property
+    def charge(self) -> int:
+        return self._charge
 
     cpdef Spectrum emission(self, Point3D point, Vector3D direction, Spectrum spectrum):