From 3fc091e22d6e855ec86207730ff927ad338c814a Mon Sep 17 00:00:00 2001
From: MatejTomes <tomes@ipp.cas.cz>
Date: Thu, 21 Mar 2019 11:11:22 +0100
Subject: [PATCH 1/5] Adding OpenADAS ADF11 ccd rates

ccd rates are thermal charge exchange rates between element and neutral hydrogen.
---
 cherab/openadas/install.py           | 23 +++++++++++
 cherab/openadas/openadas.py          | 15 +++++++
 cherab/openadas/rates/atomic.pyx     | 41 ++++++++++++++++++
 cherab/openadas/repository/atomic.py | 62 ++++++++++++++++++++++++++++
 cherab/openadas/repository/create.py | 14 +++++++
 5 files changed, 155 insertions(+)

diff --git a/cherab/openadas/install.py b/cherab/openadas/install.py
index 83fffd1..d411cc7 100644
--- a/cherab/openadas/install.py
+++ b/cherab/openadas/install.py
@@ -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,26 @@ def install_adf11acd(element, file_path, download=False, repository_path=None, a
     repository.update_recombination_rates(rate, repository_path)
 
 
+def install_adf11ccd(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 element: The element described by the rate file.
+    :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(element, path)
+    repository.update_thermalchargeexchange_rates(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.
diff --git a/cherab/openadas/openadas.py b/cherab/openadas/openadas.py
index 0847bfe..ac4d39f 100644
--- a/cherab/openadas/openadas.py
+++ b/cherab/openadas/openadas.py
@@ -86,6 +86,21 @@ def recombination_rate(self, ion, charge):
 
         return RecombinationRate(data, extrapolate=self._permit_extrapolation)
 
+    def thermalchargeexchange_rate(self, ion, charge):
+
+        if isinstance(ion, Isotope):
+            ion = ion.element
+
+        try:
+            data = repository.get_recombination_rate(ion, charge)
+
+        except RuntimeError:
+            if self._missing_rates_return_null:
+                return NullRecombinationRate()
+            raise
+
+        return RecombinationRate(data, extrapolate=self._permit_extrapolation)
+
     def beam_cx_pec(self, donor_ion, receiver_ion, receiver_charge, transition):
         """
 
diff --git a/cherab/openadas/rates/atomic.pyx b/cherab/openadas/rates/atomic.pyx
index bbdf33c..066cf39 100644
--- a/cherab/openadas/rates/atomic.pyx
+++ b/cherab/openadas/rates/atomic.pyx
@@ -97,3 +97,44 @@ cdef class NullRecombinationRate(CoreRecombinationRate):
 
     cpdef double evaluate(self, double density, double temperature) except? -1e999:
         return 0.0
+
+
+cdef class ThermalChargeExchangeRate(CoreRecombinationRate):
+
+    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 NullThermalChargeExchangeRate(CoreRecombinationRate):
+    """
+    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
\ No newline at end of file
diff --git a/cherab/openadas/repository/atomic.py b/cherab/openadas/repository/atomic.py
index 2b54bab..0245427 100644
--- a/cherab/openadas/repository/atomic.py
+++ b/cherab/openadas/repository/atomic.py
@@ -108,6 +108,47 @@ def update_recombination_rates(rates, repository_path=None):
         _update_and_write_adf11(species, rate_data, path)
 
 
+def add_thermalchargeexchange_rate(species, charge, 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 repository_path:
+    :return:
+    """
+
+    update_thermalchargeexchange_rates({
+        species: {
+            charge: rate
+        }
+    }, repository_path)
+
+
+def update_thermalchargeexchange_rates(rates, repository_path=None):
+    """
+    Thermal charge exchange rate file structure
+
+    /recombination/<species>.json
+
+    File contains multiple rates, indexed by the ion charge state.
+    """
+
+    repository_path = repository_path or DEFAULT_REPOSITORY_PATH
+
+    for species, rate_data in rates.items():
+
+        # sanitise and validate arguments
+        if not isinstance(species, Element):
+            raise TypeError('The species must be an Element object.')
+
+        path = os.path.join(repository_path, 'thermalchargeexchange/{}.json'.format(species.symbol.lower()))
+
+        _update_and_write_adf11(species, rate_data, path)
+
+
 def _update_and_write_adf11(species, rate_data, path):
 
         # read in any existing rates
@@ -195,3 +236,24 @@ def get_recombination_rate(element, charge, repository_path=None):
     d['rate'] = np.array(d['rate'], np.float64)
 
     return d
+
+
+def get_thermalchargeexchange_rate(element, charge, repository_path=None):
+
+    repository_path = repository_path or DEFAULT_REPOSITORY_PATH
+
+    path = os.path.join(repository_path, 'thermalchargeexchange/{}.json'.format(element.symbol.lower()))
+    try:
+        with open(path, 'r') as f:
+            content = json.load(f)
+        d = content[str(charge)]
+    except (FileNotFoundError, KeyError):
+        raise RuntimeError('Requested recombination rate (element={}, charge={})'
+                           ' is not available.'.format(element.symbol, 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
diff --git a/cherab/openadas/repository/create.py b/cherab/openadas/repository/create.py
index 2ceaf09..95d327c 100644
--- a/cherab/openadas/repository/create.py
+++ b/cherab/openadas/repository/create.py
@@ -65,6 +65,20 @@ def populate(download=True, repository_path=None, adas_path=None):
             (krypton, 'adf11/acd89/acd89_kr.dat'),
             (xenon, 'adf11/acd89/acd89_xe.dat'),
         ),
+        'adf11ccd': (
+            (hydrogen, 'adf11/ccd96/ccd96_h.dat'),
+            (helium, 'adf11/ccd96/ccd96_he.dat'),
+            (lithium, 'adf11/ccd89/ccd89_li.dat'),
+            (beryllium, 'adf11/ccd89/ccd89_be.dat'),
+            (boron, 'adf11/ccd89/ccd89_b.dat'),
+            (carbon, 'adf11/ccd96/ccd96_c.dat'),
+            (nitrogen, 'adf11/ccd89/ccd89_n.dat'),
+            (oxygen, 'adf11/ccd89/ccd89_o.dat'),
+            (neon, 'adf11/ccd89/ccd89_ne.dat'),
+            (argon, 'adf11/ccd89/ccd89_ar.dat'),
+            (krypton, 'adf11/ccd89/ccd89_kr.dat'),
+            (xenon, 'adf11/ccd89/ccd89_xe.dat'),
+        ),
         'adf11plt': (
             (hydrogen, 'adf11/plt12/plt12_h.dat'),
             (helium, 'adf11/plt96/plt96_he.dat'),

From 8edbdb25d291d8463421b61b7b18a4d26ea3461d Mon Sep 17 00:00:00 2001
From: MatejTomes <tomes@ipp.cas.cz>
Date: Mon, 25 Mar 2019 12:21:55 +0100
Subject: [PATCH 2/5] Reaction to PR #26 review

Names of functions changed and pointed out bugs fixed.
---
 cherab/openadas/install.py           | 13 ++++++---
 cherab/openadas/openadas.py          | 14 +++++----
 cherab/openadas/rates/atomic.pxd     | 14 +++++++++
 cherab/openadas/rates/atomic.pyx     |  4 +--
 cherab/openadas/repository/atomic.py | 32 ++++++++++++---------
 cherab/openadas/repository/create.py | 25 ++++++++--------
 demos/plot_thermalxcrates.py         | 43 ++++++++++++++++++++++++++++
 7 files changed, 108 insertions(+), 37 deletions(-)
 create mode 100644 demos/plot_thermalxcrates.py

diff --git a/cherab/openadas/install.py b/cherab/openadas/install.py
index d411cc7..a56e73f 100644
--- a/cherab/openadas/install.py
+++ b/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/'
@@ -107,7 +107,8 @@ def install_adf11acd(element, file_path, download=False, repository_path=None, a
     repository.update_recombination_rates(rate, repository_path)
 
 
-def install_adf11ccd(element, file_path, download=False, repository_path=None, adas_path=None):
+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.
 
@@ -124,8 +125,12 @@ def install_adf11ccd(element, file_path, download=False, repository_path=None, a
         raise ValueError('Could not locate the specified ADAS file.')
 
     # decode file and write out rates
-    rate = parse_adf11(element, path)
-    repository.update_thermalchargeexchange_rates(rate, repository_path)
+    rate = parse_adf11(receiver_element, path)
+
+    ccd_rate = RecursiveDict()
+    ccd_rate[donor_element][donor_charge] = rate #reshape rate dictionary to match cherab convention
+
+    repository.update_thermal_cx_rates(ccd_rate, repository_path)
 
 def install_adf11plt(element, file_path, download=False, repository_path=None, adas_path=None):
     """
diff --git a/cherab/openadas/openadas.py b/cherab/openadas/openadas.py
index ac4d39f..a355cbd 100644
--- a/cherab/openadas/openadas.py
+++ b/cherab/openadas/openadas.py
@@ -86,20 +86,24 @@ def recombination_rate(self, ion, charge):
 
         return RecombinationRate(data, extrapolate=self._permit_extrapolation)
 
-    def thermalchargeexchange_rate(self, ion, charge):
+    def thermal_cx_rate(self, donor_element, donor_charge, receiver_element, receiver_charge):
 
-        if isinstance(ion, Isotope):
-            ion = ion.element
+        if isinstance(donor_element, Isotope):
+            donor_element = donor_element.element
+
+        if isinstance(receiver_element, Isotope):
+            receiver_element = receiver_element.element
 
         try:
-            data = repository.get_recombination_rate(ion, charge)
+            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 RecombinationRate(data, extrapolate=self._permit_extrapolation)
+        return ThermalCXRate(data, extrapolate=self._permit_extrapolation)
 
     def beam_cx_pec(self, donor_ion, receiver_ion, receiver_charge, transition):
         """
diff --git a/cherab/openadas/rates/atomic.pxd b/cherab/openadas/rates/atomic.pxd
index 8767ad7..cd59574 100644
--- a/cherab/openadas/rates/atomic.pxd
+++ b/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
diff --git a/cherab/openadas/rates/atomic.pyx b/cherab/openadas/rates/atomic.pyx
index 066cf39..e647402 100644
--- a/cherab/openadas/rates/atomic.pyx
+++ b/cherab/openadas/rates/atomic.pyx
@@ -99,7 +99,7 @@ cdef class NullRecombinationRate(CoreRecombinationRate):
         return 0.0
 
 
-cdef class ThermalChargeExchangeRate(CoreRecombinationRate):
+cdef class ThermalCXRate(CoreThermalCXRate):
 
     def __init__(self, dict data, extrapolate=False):
         """
@@ -130,7 +130,7 @@ cdef class ThermalChargeExchangeRate(CoreRecombinationRate):
         return max(0, self._rate.evaluate(density, temperature))
 
 
-cdef class NullThermalChargeExchangeRate(CoreRecombinationRate):
+cdef class NullThermalCXRate(CoreThermalCXRate):
     """
     A PEC rate that always returns zero.
     Needed for use cases where the required atomic data is missing.
diff --git a/cherab/openadas/repository/atomic.py b/cherab/openadas/repository/atomic.py
index 0245427..64f5903 100644
--- a/cherab/openadas/repository/atomic.py
+++ b/cherab/openadas/repository/atomic.py
@@ -108,7 +108,7 @@ def update_recombination_rates(rates, repository_path=None):
         _update_and_write_adf11(species, rate_data, path)
 
 
-def add_thermalchargeexchange_rate(species, charge, rate, repository_path=None):
+def add_thermal_cx_rate(species, charge, rate, repository_path=None):
     """
     Adds a single thermal charge exchange rate to the repository.
 
@@ -120,14 +120,14 @@ def add_thermalchargeexchange_rate(species, charge, rate, repository_path=None):
     :return:
     """
 
-    update_thermalchargeexchange_rates({
+    update_thermal_cx_rates({
         species: {
             charge: rate
         }
     }, repository_path)
 
 
-def update_thermalchargeexchange_rates(rates, repository_path=None):
+def update_thermal_cx_rates(rates, repository_path=None):
     """
     Thermal charge exchange rate file structure
 
@@ -138,15 +138,18 @@ def update_thermalchargeexchange_rates(rates, repository_path=None):
 
     repository_path = repository_path or DEFAULT_REPOSITORY_PATH
 
-    for species, rate_data in rates.items():
+    for donor_element in rates.keys():
+        for donor_charge in rates[donor_element].keys():
+            for species, rate_data in rates[donor_element][donor_charge].items():
 
-        # sanitise and validate arguments
-        if not isinstance(species, Element):
-            raise TypeError('The species must be an Element object.')
+                # sanitise and validate arguments
+                if not isinstance(species, Element):
+                    raise TypeError('The species must be an Element object.')
 
-        path = os.path.join(repository_path, 'thermalchargeexchange/{}.json'.format(species.symbol.lower()))
+                path = os.path.join(repository_path, 'thermal_cx/{0}/{1}/{2}.json'.format(donor_element.symbol.lower(),
+                                                                                          donor_charge, species.symbol.lower()))
 
-        _update_and_write_adf11(species, rate_data, path)
+                _update_and_write_adf11(species, rate_data, path)
 
 
 def _update_and_write_adf11(species, rate_data, path):
@@ -238,18 +241,19 @@ def get_recombination_rate(element, charge, repository_path=None):
     return d
 
 
-def get_thermalchargeexchange_rate(element, charge, repository_path=None):
+def get_thermal_cx_rate(donor_element, donor_charge, receiver_element, receiver_charge, repository_path=None):
 
     repository_path = repository_path or DEFAULT_REPOSITORY_PATH
 
-    path = os.path.join(repository_path, 'thermalchargeexchange/{}.json'.format(element.symbol.lower()))
+    path = os.path.join(repository_path, 'thermal_cx/{0}/{1}/{2}.json'.format(donor_element.symbol.lower(),
+                                                                              donor_charge, receiver_element.symbol.lower()))
     try:
         with open(path, 'r') as f:
             content = json.load(f)
-        d = content[str(charge)]
+        d = content[str(receiver_charge)]
     except (FileNotFoundError, KeyError):
-        raise RuntimeError('Requested recombination rate (element={}, charge={})'
-                           ' is not available.'.format(element.symbol, charge))
+        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)
diff --git a/cherab/openadas/repository/create.py b/cherab/openadas/repository/create.py
index 95d327c..e168bce 100644
--- a/cherab/openadas/repository/create.py
+++ b/cherab/openadas/repository/create.py
@@ -66,18 +66,19 @@ def populate(download=True, repository_path=None, adas_path=None):
             (xenon, 'adf11/acd89/acd89_xe.dat'),
         ),
         'adf11ccd': (
-            (hydrogen, 'adf11/ccd96/ccd96_h.dat'),
-            (helium, 'adf11/ccd96/ccd96_he.dat'),
-            (lithium, 'adf11/ccd89/ccd89_li.dat'),
-            (beryllium, 'adf11/ccd89/ccd89_be.dat'),
-            (boron, 'adf11/ccd89/ccd89_b.dat'),
-            (carbon, 'adf11/ccd96/ccd96_c.dat'),
-            (nitrogen, 'adf11/ccd89/ccd89_n.dat'),
-            (oxygen, 'adf11/ccd89/ccd89_o.dat'),
-            (neon, 'adf11/ccd89/ccd89_ne.dat'),
-            (argon, 'adf11/ccd89/ccd89_ar.dat'),
-            (krypton, 'adf11/ccd89/ccd89_kr.dat'),
-            (xenon, 'adf11/ccd89/ccd89_xe.dat'),
+            #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'),
diff --git a/demos/plot_thermalxcrates.py b/demos/plot_thermalxcrates.py
new file mode 100644
index 0000000..c3f9316
--- /dev/null
+++ b/demos/plot_thermalxcrates.py
@@ -0,0 +1,43 @@
+import numpy as np
+import matplotlib.pyplot as plt
+from cherab.core.atomic import neon, carbon, helium, hydrogen
+from cherab.openadas import OpenADAS
+from cherab.openadas.repository.atomic import get_thermal_cx_rate
+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 for thermal charge exchange between neutral element and neutral hydrogen should not indeed be 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")
+
+#test loading rates fro not allowed CX between neutrals
+try:
+    coef_notallowed = adas.thermal_cx_rate(hydrogen, 0, neon, 0)
+except RuntimeError:
+    print("All correct")

From e9458e956af735110764c7f8d3c7403889e7d9de Mon Sep 17 00:00:00 2001
From: MatejTomes <tomes@ipp.cas.cz>
Date: Mon, 25 Mar 2019 12:27:27 +0100
Subject: [PATCH 3/5] Updating function documentation

---
 cherab/openadas/install.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/cherab/openadas/install.py b/cherab/openadas/install.py
index a56e73f..6e21025 100644
--- a/cherab/openadas/install.py
+++ b/cherab/openadas/install.py
@@ -112,7 +112,9 @@ def install_adf11ccd(donor_element, donor_charge, receiver_element, file_path, d
     """
     Adds the thermal charge exchange rate defined in an ADF11 file to the repository.
 
-    :param element: The element described by the rate file.
+    :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).

From 2ff1cfd8fe2fc8878b398c421c32f93c2a408d8c Mon Sep 17 00:00:00 2001
From: MatejTomes <tomes@ipp.cas.cz>
Date: Mon, 25 Mar 2019 12:59:01 +0100
Subject: [PATCH 4/5] updating add_thermal_cx_rate

accommodating new structure of the thermal cx repository/data
---
 cherab/openadas/repository/atomic.py | 16 ++++++++++------
 1 file changed, 10 insertions(+), 6 deletions(-)

diff --git a/cherab/openadas/repository/atomic.py b/cherab/openadas/repository/atomic.py
index 64f5903..0353229 100644
--- a/cherab/openadas/repository/atomic.py
+++ b/cherab/openadas/repository/atomic.py
@@ -108,7 +108,8 @@ def update_recombination_rates(rates, repository_path=None):
         _update_and_write_adf11(species, rate_data, path)
 
 
-def add_thermal_cx_rate(species, charge, rate, repository_path=None):
+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.
 
@@ -116,15 +117,18 @@ def add_thermal_cx_rate(species, charge, rate, repository_path=None):
     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:
     """
 
-    update_thermal_cx_rates({
-        species: {
-            charge: rate
-        }
-    }, repository_path)
+    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):

From 18c381c6fda46b84e1288f9e8f86520d5cce76e9 Mon Sep 17 00:00:00 2001
From: mattngc <matthewncarr@gmail.com>
Date: Mon, 25 Mar 2019 13:04:22 +0000
Subject: [PATCH 5/5] Minor tidy up before merge.

---
 cherab/openadas/install.py           |  7 +++++--
 cherab/openadas/openadas.py          |  4 ++--
 cherab/openadas/rates/atomic.pyx     |  1 -
 cherab/openadas/repository/atomic.py | 23 +++++++++++++----------
 cherab/openadas/repository/create.py |  2 +-
 demos/plot_thermalxcrates.py         | 15 +++++++++------
 6 files changed, 30 insertions(+), 22 deletions(-)

diff --git a/cherab/openadas/install.py b/cherab/openadas/install.py
index 6e21025..cfcab77 100644
--- a/cherab/openadas/install.py
+++ b/cherab/openadas/install.py
@@ -112,7 +112,8 @@ def install_adf11ccd(donor_element, donor_charge, receiver_element, file_path, d
     """
     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_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.
@@ -129,11 +130,13 @@ def install_adf11ccd(donor_element, donor_charge, receiver_element, file_path, d
     # 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 #reshape rate dictionary to match cherab convention
+    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.
diff --git a/cherab/openadas/openadas.py b/cherab/openadas/openadas.py
index a355cbd..510edd8 100644
--- a/cherab/openadas/openadas.py
+++ b/cherab/openadas/openadas.py
@@ -95,8 +95,8 @@ def thermal_cx_rate(self, donor_element, donor_charge, receiver_element, receive
             receiver_element = receiver_element.element
 
         try:
-            data = repository.get_thermal_cx_rate(donor_element, donor_charge, receiver_element,
-                                                             receiver_charge)
+            data = repository.get_thermal_cx_rate(donor_element, donor_charge,
+                                                  receiver_element, receiver_charge)
 
         except RuntimeError:
             if self._missing_rates_return_null:
diff --git a/cherab/openadas/rates/atomic.pyx b/cherab/openadas/rates/atomic.pyx
index e647402..bd8a131 100644
--- a/cherab/openadas/rates/atomic.pyx
+++ b/cherab/openadas/rates/atomic.pyx
@@ -123,7 +123,6 @@ cdef class ThermalCXRate(CoreThermalCXRate):
             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)
diff --git a/cherab/openadas/repository/atomic.py b/cherab/openadas/repository/atomic.py
index 0353229..aa800c7 100644
--- a/cherab/openadas/repository/atomic.py
+++ b/cherab/openadas/repository/atomic.py
@@ -135,7 +135,7 @@ def update_thermal_cx_rates(rates, repository_path=None):
     """
     Thermal charge exchange rate file structure
 
-    /recombination/<species>.json
+    /thermal_cx/<donor_element>/<donor_charge>/<receiver_element>.json
 
     File contains multiple rates, indexed by the ion charge state.
     """
@@ -144,16 +144,17 @@ def update_thermal_cx_rates(rates, repository_path=None):
 
     for donor_element in rates.keys():
         for donor_charge in rates[donor_element].keys():
-            for species, rate_data in rates[donor_element][donor_charge].items():
+            for receiver_element, rate_data in rates[donor_element][donor_charge].items():
 
                 # sanitise and validate arguments
-                if not isinstance(species, Element):
-                    raise TypeError('The species must be an Element object.')
+                if not isinstance(receiver_element, Element):
+                    raise TypeError('The receiver_element must be an Element object.')
 
-                path = os.path.join(repository_path, 'thermal_cx/{0}/{1}/{2}.json'.format(donor_element.symbol.lower(),
-                                                                                          donor_charge, species.symbol.lower()))
+                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(species, rate_data, path)
+                _update_and_write_adf11(receiver_element, rate_data, path)
 
 
 def _update_and_write_adf11(species, rate_data, path):
@@ -249,15 +250,17 @@ def get_thermal_cx_rate(donor_element, donor_charge, receiver_element, receiver_
 
     repository_path = repository_path or DEFAULT_REPOSITORY_PATH
 
-    path = os.path.join(repository_path, 'thermal_cx/{0}/{1}/{2}.json'.format(donor_element.symbol.lower(),
-                                                                              donor_charge, receiver_element.symbol.lower()))
+    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))
+                           ' 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)
diff --git a/cherab/openadas/repository/create.py b/cherab/openadas/repository/create.py
index e168bce..7997aae 100644
--- a/cherab/openadas/repository/create.py
+++ b/cherab/openadas/repository/create.py
@@ -66,7 +66,7 @@ def populate(download=True, repository_path=None, adas_path=None):
             (xenon, 'adf11/acd89/acd89_xe.dat'),
         ),
         'adf11ccd': (
-            #donor_element, donor_charge, receiver_element file_path,
+            # (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'),
diff --git a/demos/plot_thermalxcrates.py b/demos/plot_thermalxcrates.py
index c3f9316..b9dec7d 100644
--- a/demos/plot_thermalxcrates.py
+++ b/demos/plot_thermalxcrates.py
@@ -2,7 +2,6 @@
 import matplotlib.pyplot as plt
 from cherab.core.atomic import neon, carbon, helium, hydrogen
 from cherab.openadas import OpenADAS
-from cherab.openadas.repository.atomic import get_thermal_cx_rate
 adas = OpenADAS(permit_extrapolation=True)
 
 
@@ -18,12 +17,12 @@
 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
+    # test correctness of available charge numbers
     try:
-        adas.thermal_cx_rate(hydrogen,0, neon, i)
+        adas.thermal_cx_rate(hydrogen, 0, neon, i)
     except RuntimeError:
-        print("check for thermal charge exchange between neutral element and neutral hydrogen should not indeed be allowed")
+        print("Check that thermal charge exchange between a neutral element and neutral hydrogen "
+              "is not allowed.")
 
 
 fig_tcxrates = plt.subplots()
@@ -36,7 +35,11 @@
     ax.set_ylabel("rate [m^3s^-1]")
     plt.title("Thermal Charge Exchange Rates")
 
-#test loading rates fro not allowed CX between neutrals
+
+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: