From a4373b54fd9ed35e5c2e09063f0b84ab8ce169ba Mon Sep 17 00:00:00 2001
From: Tomas Stolker <stolker@strw.leidenuniv.nl>
Date: Tue, 9 Mar 2021 12:32:51 +0100
Subject: [PATCH] Updated kzz parameter to log_kzz

---
 species/analysis/retrieval.py | 10 ++++++++--
 species/read/read_radtrans.py | 17 ++++++++++++-----
 species/util/data_util.py     | 19 ++++++++++++++-----
 3 files changed, 34 insertions(+), 12 deletions(-)

diff --git a/species/analysis/retrieval.py b/species/analysis/retrieval.py
index 23f2665b..6cdb0ada 100644
--- a/species/analysis/retrieval.py
+++ b/species/analysis/retrieval.py
@@ -303,7 +303,10 @@ def set_parameters(self,
             The chemistry type: 'equilibrium' for equilibrium chemistry or 'free' for retrieval
             of free abundances (but constant with altitude).
         quenching : str, None
-            TODO Fitting a quenching pressure.
+            Quenching type for CO/CH4/H2O abundances. Either the quenching pressure (bar) is a free
+            parameter (``quenching='pressure'``) or the quenching pressure is calculated from the
+            mixing and chemical timescales (``quenching='diffusion'``). The quenching is not
+            applied if the argument is set to ``None``.
         pt_profile : str
             The parametrization for the pressure-temperature profile ('molliere', 'free', or
             'monotonic').
@@ -510,7 +513,10 @@ def run_multinest(self,
             The chemistry type: 'equilibrium' for equilibrium chemistry or 'free' for retrieval
             of free abundances (but constant with altitude).
         quenching : str, None
-            TODO Fitting a quenching pressure.
+            Quenching type for CO/CH4/H2O abundances. Either the quenching pressure (bar) is a free
+            parameter (``quenching='pressure'``) or the quenching pressure is calculated from the
+            mixing and chemical timescales (``quenching='diffusion'``). The quenching is not
+            applied if the argument is set to ``None``.
         pt_profile : str
             The parametrization for the pressure-temperature profile ('molliere', 'free', or
             'monotonic').
diff --git a/species/read/read_radtrans.py b/species/read/read_radtrans.py
index 1a6b2258..b08ef4c4 100644
--- a/species/read/read_radtrans.py
+++ b/species/read/read_radtrans.py
@@ -173,13 +173,13 @@ def get_model(self,
 
         # Determine chemistry type
 
-        check_free = True
+        check_free_abund = True
 
-        for item in self.cloud_species:
+        for item in self.line_species:
             if item not in model_param:
-                check_free = False
+                check_free_abund = False
 
-        if check_free:
+        if check_free_abund:
             chemistry = 'free'
 
         elif 'metallicity' in model_param and 'c_o_ratio' in model_param:
@@ -322,10 +322,17 @@ def get_model(self,
 
             # Calculate the petitRADTRANS spectrum for a cloudy atmosphere
 
+            if 'kzz' in model_param:
+                # Backward compatibility
+                log_kzz = model_param['kzz']
+
+            elif 'log_kzz' in model_param:
+                log_kzz = model_param['log_kzz']
+
             wavelength, flux, emission_contr = retrieval_util.calc_spectrum_clouds(
                 self.rt_object, self.pressure, temp, c_o_ratio, metallicity,
                 log_p_quench, log_x_abund, log_x_base, model_param['fsed'],
-                model_param['kzz'], model_param['logg'], model_param['sigma_lnorm'],
+                log_kzz, model_param['logg'], model_param['sigma_lnorm'],
                 chemistry=chemistry, pressure_grid=self.pressure_grid,
                 plotting=False, contribution=True, tau_cloud=tau_cloud)
 
diff --git a/species/util/data_util.py b/species/util/data_util.py
index 11af8d90..5b757954 100644
--- a/species/util/data_util.py
+++ b/species/util/data_util.py
@@ -703,7 +703,7 @@ def retrieval_spectrum(indices: Dict[str, np.int64],
                        pt_profile: str,
                        line_species: List[str],
                        cloud_species: List[str],
-                       quenching: np.bool_,
+                       quenching: Optional[str],
                        spec_res: float,
                        distance: Optional[float],
                        pt_smooth: Optional[float],
@@ -724,8 +724,11 @@ def retrieval_spectrum(indices: Dict[str, np.int64],
         List with the line species.
     cloud_species : list(str)
         List with the cloud species.
-    quenching : bool
-        Use a quenching pressure for CH4/CO.
+    quenching : str, None
+        Quenching type for CO/CH4/H2O abundances. Either the quenching pressure (bar) is a free
+        parameter (``quenching='pressure'``) or the quenching pressure is calculated from the
+        mixing and chemical timescales (``quenching='diffusion'``). The quenching is not applied
+        if the argument is set to ``None``.
     spec_res : float
         Spectral resolution.
     distance : float, None
@@ -786,16 +789,22 @@ def retrieval_spectrum(indices: Dict[str, np.int64],
         for species_item in line_species:
             model_param[species_item] = sample[indices[species_item]]
 
-    if quenching:
+    if quenching == 'pressure':
         model_param['log_p_quench'] = sample[indices['log_p_quench']]
 
     # Add cloud parameters
 
     if len(cloud_species) > 0:
         model_param['fsed'] = sample[indices['fsed']]
-        model_param['kzz'] = sample[indices['kzz']]
         model_param['sigma_lnorm'] = sample[indices['sigma_lnorm']]
 
+        if 'kzz' in indices:
+            # Backward compatibility
+            model_param['kzz'] = sample[indices['kzz']]
+
+        elif 'log_kzz' in indices:
+            model_param['log_kzz'] = sample[indices['log_kzz']]
+
         for cloud_item in cloud_species:
             cloud_param = f'{cloud_item[:-3].lower()}_fraction'