Feature/thermal cx

cherab/openadas/install.py

@@ -19,7 +19,7 @@
 import urllib
 from cherab.openadas import repository
 from cherab.openadas.parse import *
-
+from cherab.core.utility import RecursiveDict
 
 ADAS_DOWNLOAD_CACHE = os.path.expanduser('~/.cherab/openadas/download_cache')
 OPENADAS_FILE_URL = 'http://open.adas.ac.uk/download/'
@@ -34,6 +34,9 @@ def install_files(configuration, download=False, repository_path=None, adas_path
         if adf.lower() == 'adf11acd':
             for args in configuration[adf]:
                 install_adf11acd(*args, download=download, repository_path=repository_path, adas_path=adas_path)
+        if adf.lower() == 'adf11ccd':
+            for args in configuration[adf]:
+                install_adf11ccd(*args, download=download, repository_path=repository_path, adas_path=adas_path)
         if adf.lower() == 'adf11plt':
             for args in configuration[adf]:
                 install_adf11plt(*args, download=download, repository_path=repository_path, adas_path=adas_path)
@@ -104,6 +107,36 @@ def install_adf11acd(element, file_path, download=False, repository_path=None, a
     repository.update_recombination_rates(rate, repository_path)
 
 
+def install_adf11ccd(donor_element, donor_charge, receiver_element, file_path, download=False,
+                     repository_path=None, adas_path=None):
+    """
+    Adds the thermal charge exchange rate defined in an ADF11 file to the repository.
+
+    :param donor_element: Element donating the electron, for the case of ADF11 files it is
+      neutral hydrogen.
+    :param donor_charge: Charge of the donor atom/ion.
+    :param receiver_element: Element receiving the electron.
+    :param file_path: Path relative to ADAS root.
+    :param download: Attempt to download file if not present (Default=True).
+    :param repository_path: Path to the repository in which to install the rates (optional).
+    :param adas_path: Path to ADAS files repository (optional).
+    """
+
+    print('Installing {}...'.format(file_path))
+    path = _locate_adas_file(file_path, download, adas_path)
+    if not path:
+        raise ValueError('Could not locate the specified ADAS file.')
+
+    # decode file and write out rates
+    rate = parse_adf11(receiver_element, path)
+
+    # reshape rate dictionary to match cherab convention
+    ccd_rate = RecursiveDict()
+    ccd_rate[donor_element][donor_charge] = rate
+
+    repository.update_thermal_cx_rates(ccd_rate, repository_path)
+
+
 def install_adf11plt(element, file_path, download=False, repository_path=None, adas_path=None):
     """
     Adds the line radiated power rates defined in an ADF11 file to the repository.

cherab/openadas/openadas.py

@@ -86,6 +86,25 @@ def recombination_rate(self, ion, charge):
 
         return RecombinationRate(data, extrapolate=self._permit_extrapolation)
 
+    def thermal_cx_rate(self, donor_element, donor_charge, receiver_element, receiver_charge):
+
+        if isinstance(donor_element, Isotope):
+            donor_element = donor_element.element
+
+        if isinstance(receiver_element, Isotope):
+            receiver_element = receiver_element.element
+
+        try:
+            data = repository.get_thermal_cx_rate(donor_element, donor_charge,
+                                                  receiver_element, receiver_charge)
+
+        except RuntimeError:
+            if self._missing_rates_return_null:
+                return NullRecombinationRate()
+            raise
+
+        return ThermalCXRate(data, extrapolate=self._permit_extrapolation)
+
     def beam_cx_pec(self, donor_ion, receiver_ion, receiver_charge, transition):
         """
 

cherab/openadas/rates/atomic.pxd

@@ -18,6 +18,7 @@
 
 from cherab.core cimport IonisationRate as CoreIonisationRate
 from cherab.core cimport RecombinationRate as CoreRecombinationRate
+from cherab.core cimport ThermalCXRate as CoreThermalCXRate
 from cherab.core.math cimport Interpolate2DCubic
 
 
@@ -45,3 +46,16 @@ cdef class RecombinationRate(CoreRecombinationRate):
 
 cdef class NullRecombinationRate(CoreRecombinationRate):
     pass
+
+
+cdef class ThermalCXRate(CoreThermalCXRate):
+
+    cdef:
+        readonly dict raw_data
+        readonly double wavelength
+        readonly tuple density_range, temperature_range
+        Interpolate2DCubic _rate
+
+
+cdef class NullThermalCXRate(CoreThermalCXRate):
+    pass

cherab/openadas/rates/atomic.pyx

@@ -97,3 +97,43 @@ cdef class NullRecombinationRate(CoreRecombinationRate):
 
     cpdef double evaluate(self, double density, double temperature) except? -1e999:
         return 0.0
+
+
+cdef class ThermalCXRate(CoreThermalCXRate):
+
+    def __init__(self, dict data, extrapolate=False):
+        """
+        :param data: Dictionary containing rate data.
+        :param extrapolate: Enable extrapolation (default=False).
+        """
+
+        self.raw_data = data
+
+        # unpack
+        ne = data['ne']
+        te = data['te']
+        rate =  data['rate']
+
+        # store limits of data
+        self.density_range = ne.min(), ne.max()
+        self.temperature_range = te.min(), te.max()
+
+        # interpolate rate
+        self._rate = Interpolate2DCubic(
+            ne, te, rate, extrapolate=extrapolate, extrapolation_type="quadratic"
+        )
+
+    cpdef double evaluate(self, double density, double temperature) except? -1e999:
+
+        # prevent -ve values (possible if extrapolation enabled)
+        return max(0, self._rate.evaluate(density, temperature))
+
+
+cdef class NullThermalCXRate(CoreThermalCXRate):
+    """
+    A PEC rate that always returns zero.
+    Needed for use cases where the required atomic data is missing.
+    """
+
+    cpdef double evaluate(self, double density, double temperature) except? -1e999:
+        return 0.0

cherab/openadas/repository/atomic.py

@@ -108,6 +108,55 @@ def update_recombination_rates(rates, repository_path=None):
         _update_and_write_adf11(species, rate_data, path)
 
 
+def add_thermal_cx_rate(donor_element, donor_charge, receiver_element, rate, repository_path=None):
+
+    """
+    Adds a single thermal charge exchange rate to the repository.
+
+    If adding multiple rates, consider using the update_recombination_rates()
+    function instead. The update function avoids repeatedly opening and closing
+    the rate files.
+
+    :param donor_element: Element donating the electron.
+    :param donor_charge: Charge of the donating atom/ion
+    :param receiver_element: Element receiving the electron
+    :param rate: rates
+    :param repository_path:
+    :return:
+    """
+
+    rates2update = RecursiveDict()
+    rates2update[donor_element][donor_charge][receiver_element] = rate
+
+    update_thermal_cx_rates(rates2update, repository_path)
+
+
+def update_thermal_cx_rates(rates, repository_path=None):
+    """
+    Thermal charge exchange rate file structure
+
+    /thermal_cx/<donor_element>/<donor_charge>/<receiver_element>.json
+
+    File contains multiple rates, indexed by the ion charge state.
+    """
+
+    repository_path = repository_path or DEFAULT_REPOSITORY_PATH
+
+    for donor_element in rates.keys():
+        for donor_charge in rates[donor_element].keys():
+            for receiver_element, rate_data in rates[donor_element][donor_charge].items():
+
+                # sanitise and validate arguments
+                if not isinstance(receiver_element, Element):
+                    raise TypeError('The receiver_element must be an Element object.')
+
+                rate_path = 'thermal_cx/{0}/{1}/{2}.json'.format(donor_element.symbol.lower(),
+                                                                 donor_charge, receiver_element.symbol.lower())
+                path = os.path.join(repository_path, rate_path)
+
+                _update_and_write_adf11(receiver_element, rate_data, path)
+
+
 def _update_and_write_adf11(species, rate_data, path):
 
         # read in any existing rates
@@ -195,3 +244,27 @@ def get_recombination_rate(element, charge, repository_path=None):
     d['rate'] = np.array(d['rate'], np.float64)
 
     return d
+
+
+def get_thermal_cx_rate(donor_element, donor_charge, receiver_element, receiver_charge, repository_path=None):
+
+    repository_path = repository_path or DEFAULT_REPOSITORY_PATH
+
+    rate_path = 'thermal_cx/{0}/{1}/{2}.json'.format(donor_element.symbol.lower(), donor_charge,
+                                                     receiver_element.symbol.lower())
+    path = os.path.join(repository_path, rate_path)
+    try:
+        with open(path, 'r') as f:
+            content = json.load(f)
+        d = content[str(receiver_charge)]
+    except (FileNotFoundError, KeyError):
+        raise RuntimeError('Requested thermal charge-exchange rate (donor={}, donor charge={}, receiver={})'
+                           ' is not available.'
+                           ''.format(donor_element.symbol, donor_charge, receiver_element.symbol, receiver_charge))
+
+    # convert to numpy arrays
+    d['ne'] = np.array(d['ne'], np.float64)
+    d['te'] = np.array(d['te'], np.float64)
+    d['rate'] = np.array(d['rate'], np.float64)
+
+    return d

cherab/openadas/repository/create.py

@@ -65,6 +65,21 @@ def populate(download=True, repository_path=None, adas_path=None):
             (krypton, 'adf11/acd89/acd89_kr.dat'),
             (xenon, 'adf11/acd89/acd89_xe.dat'),
         ),
+        'adf11ccd': (
+            # (donor_element, donor_charge, receiver_element, file_path)
+            (hydrogen, 0, hydrogen, 'adf11/ccd96/ccd96_h.dat'),
+            (hydrogen, 0, helium, 'adf11/ccd96/ccd96_he.dat'),
+            (hydrogen, 0, lithium, 'adf11/ccd89/ccd89_li.dat'),
+            (hydrogen, 0, beryllium, 'adf11/ccd89/ccd89_be.dat'),
+            (hydrogen, 0, boron, 'adf11/ccd89/ccd89_b.dat'),
+            (hydrogen, 0, carbon, 'adf11/ccd96/ccd96_c.dat'),
+            (hydrogen, 0, nitrogen, 'adf11/ccd89/ccd89_n.dat'),
+            (hydrogen, 0, oxygen, 'adf11/ccd89/ccd89_o.dat'),
+            (hydrogen, 0, neon, 'adf11/ccd89/ccd89_ne.dat'),
+            (hydrogen, 0, argon, 'adf11/ccd89/ccd89_ar.dat'),
+            (hydrogen, 0, krypton, 'adf11/ccd89/ccd89_kr.dat'),
+            (hydrogen, 0, xenon, 'adf11/ccd89/ccd89_xe.dat'),
+        ),
         'adf11plt': (
             (hydrogen, 'adf11/plt12/plt12_h.dat'),
             (helium, 'adf11/plt96/plt96_he.dat'),

demos/plot_thermalxcrates.py

@@ -0,0 +1,46 @@
+import numpy as np
+import matplotlib.pyplot as plt
+from cherab.core.atomic import neon, carbon, helium, hydrogen
+from cherab.openadas import OpenADAS
+adas = OpenADAS(permit_extrapolation=True)
+
+
+electron_temperatures = [10**x for x in np.linspace(np.log10(1), np.log10(10000), num=100)]
+electron_density = 1e19
+
+elem = neon
+
+numstates = elem.atomic_number + 1
+
+# Collect rate coefficients
+coef_tcx = {}
+for i in np.arange(1, elem.atomic_number+1):
+    coef_tcx[i] = adas.thermal_cx_rate(hydrogen, 0, neon, int(i))
+
+    # test correctness of available charge numbers
+    try:
+        adas.thermal_cx_rate(hydrogen, 0, neon, i)
+    except RuntimeError:
+        print("Check that thermal charge exchange between a neutral element and neutral hydrogen "
+              "is not allowed.")
+
+
+fig_tcxrates = plt.subplots()
+ax = fig_tcxrates[1]
+for i in range(1, elem.atomic_number+1):
+    tcx_rate = [coef_tcx[i](electron_density, x) for x in electron_temperatures]
+    ax.loglog(electron_temperatures, tcx_rate, "-x", label = "Ne{}+".format(i))
+
+    plt.xlabel("Electron Temperature (eV)")
+    ax.set_ylabel("rate [m^3s^-1]")
+    plt.title("Thermal Charge Exchange Rates")
+
+
+plt.show()
+
+
+# test loading rates for CX between neutrals is not allowed
+try:
+    coef_notallowed = adas.thermal_cx_rate(hydrogen, 0, neon, 0)
+except RuntimeError:
+    print("All correct")