/*******************************************************************************
 * This file is part of SWIFT.
 * Copyright (c) 2017 Matthieu Schaller (matthieu.schaller@durham.ac.uk)
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Lesser General Public License as published
 * by the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU Lesser General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 *
 ******************************************************************************/
/**
 * @file src/cooling/CHIMES/cooling.c
 * @brief CHIMES cooling functions
 */

/* Config parameters. */
#include <config.h>

/* Some standard headers. */
#include <float.h>
#include <hdf5.h>
#include <math.h>
#include <time.h>

/* Local includes. */
#include "adiabatic_index.h"
#include "chemistry.h"
#include "chimes/src/chimes_interpol.h"
#include "colibre_subgrid.h"
#include "cooling_properties.h"
#include "dust.h"
#include "entropy_floor.h"
#include "error.h"
#include "exp10.h"
#include "hydro.h"
#include "io_properties.h"
#include "parser.h"
#include "part.h"
#include "physical_constants.h"
#include "space.h"
#include "star_formation.h"
#include "units.h"

const char *CHIMES_tables_version_string = "14Oct2022";

/* Pre-declarations to simplify the code */

double chimes_mu(const struct cooling_function_data *cooling,
                 const struct hydro_props *hydro_props, const struct part *p,
                 const struct xpart *xp);

void chimes_update_element_abundances(
    const struct phys_const *phys_const, const struct unit_system *us,
    const struct cosmology *cosmo, const struct cooling_function_data *cooling,
    const struct hydro_props *hydro_props, const struct part *p,
    const struct xpart *xp, struct gasVariables *ChimesGasVars,
    const double u_0, const int mode);

void chimes_update_gas_vars(
    const double u_cgs, const struct phys_const *phys_const,
    const struct unit_system *us, const struct cosmology *cosmo,
    const struct hydro_props *hydro_properties,
    const struct entropy_floor_properties *floor_props,
    const struct cooling_function_data *cooling, const struct dustevo_props *dp,
    const struct part *p, const struct xpart *xp,
    struct gasVariables *ChimesGasVars, const float dt_cgs);

float cooling_get_temperature_on_entropy_floor(
    const struct phys_const *phys_const, const struct hydro_props *hydro_props,
    const struct unit_system *us, const struct cosmology *cosmo,
    const struct cooling_function_data *cooling, const struct part *p,
    const struct xpart *xp);

double calculate_colibre_N_ref(const struct phys_const *phys_const,
                               const struct unit_system *us,
                               const struct cosmology *cosmo,
                               const struct cooling_function_data *cooling,
                               const struct hydro_props *hydro_properties,
                               const struct part *p, const struct xpart *xp,
                               const double temperature);

void cooling_set_HIIregion_chimes_abundances(
    struct gasVariables *ChimesGasVars,
    const struct cooling_function_data *cooling);

void cooling_set_FB_particle_chimes_abundances(
    struct gasVariables *ChimesGasVars,
    const struct cooling_function_data *cooling);

/**
 * @brief Returns whether the gas particle is in an HII region or not.
 */
int is_HII_region(const struct xpart *xp) {
  return xp->tracers_data.HIIregion_timer_gas > 0.;
}

/**
 * @brief Computes the mean molecular weight for an HII region at a given H
 * abundance.
 *
 * @param cooling The properties of the cooling scheme.
 * @param XH The Hydrogen abundance of the gas.
 */
double get_mu_HII_region(const struct cooling_function_data *cooling,
                         const double XH) {

  return 4.0 / ((1.0 + cooling->HIIregion_fion) * (1.0 + (3.0 * XH)));
}

/**
 * @brief Compute the temperature from mu and the internal energy in CGS.
 */
double get_T_from_u_cgs(const double u_cgs, const double mu,
                        const struct cooling_function_data *cooling) {
  return u_cgs * hydro_gamma_minus_one * mu *
         cooling->proton_mass_over_boltzmann_k_cgs;
}

/**
 * @brief Compute the temperature from mu and the internal energy in internal
 * units.
 */
double get_T_from_u_internal(const double u, const double mu,
                             const struct phys_const *phys_const) {
  return u * hydro_gamma_minus_one * mu * phys_const->const_proton_mass /
         phys_const->const_boltzmann_k;
}

/**
 * @brief Compute the the internal energy in CGS from mu and the temperature.
 */
double get_u_cgs_from_T(const double T, const double mu,
                        const struct cooling_function_data *cooling) {
  return T * hydro_one_over_gamma_minus_one /
         (cooling->proton_mass_over_boltzmann_k_cgs * mu);
}

/**
 * @brief Compute the the internal energy in internal units from mu and the
 * temperature.
 */
double get_u_internal_from_T(const double T, const double mu,
                             const struct phys_const *phys_const) {
  return T * hydro_one_over_gamma_minus_one * phys_const->const_boltzmann_k /
         (phys_const->const_proton_mass * mu);
}

/**
 * @brief Return the hydrogen mass fraction of a particle.
 *
 * Deals with the case where we don't have a chemistry model.
 */
double get_hydrogen_mass_fraction(const struct part *p,
                                  const struct hydro_props *hydro_props) {

#if defined(CHEMISTRY_COLIBRE) || defined(CHEMISTRY_EAGLE)

  const float *const metal_fraction =
      chemistry_get_metal_mass_fraction_for_cooling(p);

  return metal_fraction[chemistry_element_H];
#else

  /* Without COLIBRE or EAGLE chemistry, the metal abundances are unavailable.
   * --> Set to primordial abundances. */
  return hydro_props->hydrogen_mass_fraction;

#endif /* CHEMISTRY_COLIBRE || CHEMISTRY_EAGLE */
}

/**
 * @brief Get the dust ratio for a given particle
 * and for the different dust modes
 */
double get_dust_ratio(const struct part *p,
                      const struct cooling_function_data *cooling,
                      const struct dustevo_props *dp, const double Z) {

#if defined(DUST_NONE)
  return Z;
#else
  if (dp->pair_to_cooling) {
    float dfrac = 0.;

    for (int grain = 0; grain < dust_grain_species_count; grain++) {
      double grainfrac = p->dust_data.grain_mass_fraction[grain];

#if defined(DUST_T20MS)
      if (grain > (dust_grain_species_count / 2)) {
        /* the second half of the grain array has a smaller radius.
         * Grain processes scale with cross section, so we apply
         * a boost factor based on the size ratios. */
        grainfrac *= dust_get_large_to_small_ratio();
      }
#endif

      dfrac += grainfrac;
    }
    return dfrac / dust_get_local_dust_to_gas_ratio();
  } else {
    return Z;
  }
#endif
}

double get_dust_boost_factor(const double nH_cgs,
                             const struct cooling_function_data *cooling) {

  if (nH_cgs <= cooling->dust_boost_nH_min_cgs) {

    /* Below the min density for boost --> no boost */
    return 1.0f;

  } else if (nH_cgs >= cooling->dust_boost_nH_max_cgs) {

    /* Above the max density for boost --> max boost */
    return cooling->max_dust_boost_factor;

  } else {

    /* In the transition region --> Compute the boost factor */
    const float log_nH = log10f(nH_cgs);
    const float log_nH_min = log10f(cooling->dust_boost_nH_min_cgs);
    const float log_nH_max = log10f(cooling->dust_boost_nH_max_cgs);
    const float log_max_boost = log10f(cooling->max_dust_boost_factor);
    const float delta_log_nH = chimes_max(log_nH_max - log_nH_min, FLT_MIN);

    /* Linear ramp (in log) between min and max */
    const float log_boost =
        ((log_nH - log_nH_min) / delta_log_nH) * log_max_boost;

    return exp10f(log_boost);
  }
}

double get_CR_rate(const double N_ref_cgs, const double nH_cgs,
                   const struct cosmology *cosmo,
                   const struct phys_const *phys_const,
                   const struct cooling_function_data *cooling) {

  /* Set the cosmic ray rate */
  if (cooling->UV_field_choice != UV_field_COLIBRE) {

    /* Constant cosmic ray rate */
    return cooling->cosmic_ray_rate;

  } else {

    /* The cosmic ray rate scales with the reference column density N_ref.
     * Based on Ploeckinger & Schaye (2020) */
    const double cr_coldensratio =
        N_ref_cgs / cooling->cr_transition_column_density_cgs;

    double cosmic_ray_rate;

    if (cr_coldensratio < 1.) {
      /* N_ref < N_trans: power law index 1 */
      cosmic_ray_rate = pow(cr_coldensratio, cooling->cr_powerlaw_slope_1);
    } else {
      /* N_ref >=N_trans: power law index 2 */
      cosmic_ray_rate = pow(cr_coldensratio, cooling->cr_powerlaw_slope_2);
    }

    cosmic_ray_rate *= cooling->cr_rate_at_transition_cgs;
    const double log_cosmic_ray_rate_cgs = log10(cosmic_ray_rate);

    /* Now, apply some saturation */

    const double log_cr_rate_min =
        log10(cooling->cr_rate_at_transition_cgs *
              pow(cooling->nref_column_density_min_cgs /
                      cooling->cr_transition_column_density_cgs,
                  cooling->cr_powerlaw_slope_1));

    /* low density cutoff between cr_nmax and cr_nmin
     * with k regulating the steepness */
    const double cr_nmax_over_dens =
        cosmo->critical_density * cooling->cr_cutoff_over_density *
        cooling->nref_XH_const / phys_const->const_proton_mass;

    const double cr_nmax_over_dens_cgs =
        cr_nmax_over_dens * cooling->number_density_to_cgs;

    /* Max density */
    const double cr_nmax =
        fmax(cr_nmax_over_dens_cgs, cooling->cr_cutoff_phys_density_cgs);

    /* Min density: (1. / 10^(cutoff_width_dex)) * cr_max */
    const double cr_nmin = exp10(-cooling->cr_cutoff_width_dex) * cr_nmax;

    const double numerator = (log_cosmic_ray_rate_cgs - log_cr_rate_min);
    const double argument = sqrt(cr_nmin / cr_nmax) * cr_nmax / nH_cgs;
    const double denominator =
        1. + pow(argument, -cooling->cr_cutoff_steepness_k);

    const double log_cosmic_ray_rate =
        log_cosmic_ray_rate_cgs - numerator / denominator;

    return exp10(log_cosmic_ray_rate);
  }
}

/**
 * @brief Find the index of the current redshift along the redshift dimension
 * of the cooling tables.
 *
 * Since the redshift table is not evenly spaced, compare z with each
 * table value in decreasing order starting with the previous redshift index
 *
 * The returned difference is expressed in units of the table separation. This
 * means dx = (x - table[i]) / (table[i+1] - table[i]). It is always between
 * 0 and 1.
 *
 * @param z Redshift we are searching for.
 * @param z_index (return) Index of the redshift in the table.
 * @param dz (return) Difference in redshift between z and table[z_index].
 * @param cooling #cooling_function_data structure containing redshift table.
 */
__attribute__((always_inline)) INLINE void get_redshift_index(
    const float z, int *z_index, float *dz,
    struct colibre_cooling_tables *table) {

  /* If z is larger than the maximum table redshift, use the last value */
  if (z >= table->Redshifts[table->number_of_redshifts - 1]) {
    *z_index = table->number_of_redshifts - 2;
    *dz = 1.0;
  }
  /* At the end, just use the last value */
  else if (z <= table->Redshifts[0]) {
    *z_index = 0;
    *dz = 0.0;
  }

  /* Normal case: search... */
  else {

    /* start at the previous index and search */
    for (int i = table->previous_z_index; i >= 0; i--) {
      if (z > table->Redshifts[i]) {

        *z_index = i;
        table->previous_z_index = i;

        *dz = (z - table->Redshifts[i]) /
              (table->Redshifts[i + 1] - table->Redshifts[i]);
        break;
      }
    }
  }
}

/**
 * @brief Checks CHIMES eqm tables headers.
 *
 * If we are using the redshift-dependent
 * CHIMES eqm abundances with the COLIBRE
 * radiation field and shielding models,
 * we need to check that the parameters
 * used in the tables match the
 * hard-coded values.
 *
 * @param fname_chimes filename of the eqm table.
 * @param cooling #cooling_function_data struct.
 */
void chimes_check_colibre_eqm_tables_header(
    const char *fname_chimes, const struct cooling_function_data *cooling) {

  const hid_t file_id = H5Fopen(fname_chimes, H5F_ACC_RDONLY, H5P_DEFAULT);

  if (file_id < 0)
    error("Unable to open equilibrium abundance data file '%s'!", fname_chimes);

  /* Check table version */
  const hid_t attr = H5Aopen_by_name(file_id, "/Header", "version_date",
                                     H5P_DEFAULT, H5P_DEFAULT);
  if (attr < 0)
    error(
        "Error reading the version string. You are likely using old tables "
        "that do not match this code version.");

  const hid_t type = H5Aget_type(attr);
  char *version_string = NULL;
  const hid_t status = H5Aread(attr, type, &version_string);
  if (status < 0)
    error("Error reading version_date in CHIMES eqm table header!");

  if (strcmp(version_string, CHIMES_tables_version_string) != 0)
    error(
        "The table and code version do not match, please use table version "
        "'%s' (you are using version '%s').",
        CHIMES_tables_version_string, version_string);

  H5free_memory(version_string);
  H5Tclose(type);
  H5Aclose(attr);
  H5Fclose(file_id);
}

/**
 * @brief Initialises properties stored in the cooling_function_data struct
 *
 * @param parameter_file The parsed parameter file
 * @param us Internal system of units data structure
 * @param phys_const #phys_const data structure
 * @param hydro_props The properties of the hydro scheme.
 * @param cooling #cooling_function_data struct to initialize
 */
void cooling_init_backend(struct swift_params *parameter_file,
                          const struct unit_system *us,
                          const struct phys_const *phys_const,
                          const struct hydro_props *hydro_props,
                          struct cooling_function_data *cooling,
                          struct dustevo_props *dp) {

  char chimes_data_dir[256];
  char string_buffer[196];

  /* The following CHIMES variables control the options involved with coupling
   * the thermo-chemistry solver in CHIMES to a Radiative Transfer (RT) solver.
   * However, we have not yet implemented  the coupling between CHIMES and the
   * RT modules in Swift. For now, we therefore simply switch off the RT
   * coupling options in CHIMES. */
  cooling->ChimesGlobalVars.rt_update_flux = 0;
  cooling->ChimesGlobalVars.rt_use_on_the_spot_approx = 0;
  cooling->ChimesGlobalVars.rt_density_absoluteTolerance = 1.0e-10f;
  cooling->ChimesGlobalVars.rt_flux_absoluteTolerance = 1.0e-10f;

  /* read the parameters */

  /* Set paths to CHIMES data files. */
  parser_get_param_string(parameter_file, "CHIMESCooling:data_path",
                          chimes_data_dir);
  sprintf(cooling->ChimesGlobalVars.MainDataTablePath,
          "%s/chimes_main_data.hdf5", chimes_data_dir);

  parser_get_param_string(parameter_file, "CHIMESCooling:EqmAbundanceTable",
                          string_buffer);
  if ((strcmp("None", string_buffer) == 0) ||
      (strcmp("none", string_buffer) == 0))
    sprintf(cooling->ChimesGlobalVars.EqAbundanceTablePath, "%s",
            string_buffer);
  else
    sprintf(cooling->ChimesGlobalVars.EqAbundanceTablePath,
            "%s/EqAbundancesTables/%s", chimes_data_dir, string_buffer);

  parser_get_param_string(parameter_file, "CHIMESCooling:UV_field",
                          string_buffer);
  if (strcmp(string_buffer, "none") == 0) {
    cooling->UV_field_choice = UV_field_no_UV;
  } else if (strcmp(string_buffer, "user-defined") == 0) {
    cooling->UV_field_choice = UV_field_user_defined;
  } else if (strcmp(string_buffer, "COLIBRE") == 0) {
    cooling->UV_field_choice = UV_field_COLIBRE;
  } else {
    error(
        "Invalid choice of 'CHIMESCooling:UV_field'. Must be 'none', "
        "'user-defined', or 'COLIBRE'");
  }

  parser_get_param_string(parameter_file, "CHIMESCooling:UVB_z_dependence",
                          string_buffer);
  if (strcmp(string_buffer, "constant") == 0) {
    cooling->UVB_dependence_choice = UVB_constant;
  } else if (strcmp(string_buffer, "cross sections") == 0) {
    cooling->UVB_dependence_choice = UVB_zdep_cross_secs;
  } else if (strcmp(string_buffer, "COLIBRE") == 0) {
    cooling->UVB_dependence_choice = UVB_zdep_cross_secs_and_eqm;
  } else {
    error(
        "Invalid choice of 'CHIMESCooling:UVB_z_dependence'. Must be "
        "'constant', 'cross sections', or 'COLIBRE'");
  }

  parser_get_param_string(parameter_file, "CHIMESCooling:Shielding_model",
                          string_buffer);
  if (strcmp(string_buffer, "none") == 0) {
    cooling->Shielding_choice = Shielding_none;
  } else if (strcmp(string_buffer, "Jeans length") == 0) {
    cooling->Shielding_choice = Shielding_Jeans_length;
  } else if (strcmp(string_buffer, "COLIBRE") == 0) {
    cooling->Shielding_choice = Shielding_COLIBRE;
  } else {
    error(
        "Invalid choice of 'CHIMESCooling:Shielding_model'. Must be 'none', "
        "'Jeans length', or 'COLIBRE'");
  }

  cooling->rad_field_norm_factor = 0.0;

  if (cooling->UVB_dependence_choice == UVB_constant) {

    cooling->ChimesGlobalVars.redshift_dependent_UVB_index = -1;
    cooling->ChimesGlobalVars.use_redshift_dependent_eqm_tables = 0;

  } else {

    if (cooling->UV_field_choice == UV_field_no_UV)
      error(
          "CHIMES ERROR: redshift-dependent UVB has been switched on, but "
          "UV_field_choice is 'no_UV'. These options are incompatible.");

    cooling->ChimesGlobalVars.redshift_dependent_UVB_index = 0;

    if (cooling->UVB_dependence_choice == UVB_zdep_cross_secs)
      cooling->ChimesGlobalVars.use_redshift_dependent_eqm_tables = 0;
    else if (cooling->UVB_dependence_choice == UVB_zdep_cross_secs_and_eqm)
      cooling->ChimesGlobalVars.use_redshift_dependent_eqm_tables = 1;
    else
      error("Invalid choice of 'CHIMESCooling:UVB_z_dependence");
  }

  if (cooling->UV_field_choice == UV_field_no_UV) {

    cooling->ChimesGlobalVars.N_spectra = 0;

  } else if (cooling->UV_field_choice == UV_field_user_defined) {

    cooling->rad_field_norm_factor = (ChimesFloat)parser_get_opt_param_double(
        parameter_file, "CHIMESCooling:rad_field_norm_factor", 1.0);

    cooling->ChimesGlobalVars.N_spectra = 1;

    parser_get_param_string(parameter_file, "CHIMESCooling:PhotoIonTable",
                            string_buffer);
    sprintf(cooling->ChimesGlobalVars.PhotoIonTablePath[0], "%s/%s",
            chimes_data_dir, string_buffer);

  } else if (cooling->UV_field_choice == UV_field_COLIBRE) {

    cooling->ChimesGlobalVars.N_spectra = 2;

    parser_get_param_string(parameter_file, "CHIMESCooling:PhotoIonTable_UVB",
                            string_buffer);
    sprintf(cooling->ChimesGlobalVars.PhotoIonTablePath[0], "%s/%s",
            chimes_data_dir, string_buffer);

    parser_get_param_string(parameter_file, "CHIMESCooling:PhotoIonTable_ISRF",
                            string_buffer);
    sprintf(cooling->ChimesGlobalVars.PhotoIonTablePath[1], "%s/%s",
            chimes_data_dir, string_buffer);
  } else {
    error("Invalid 'UV_field_choice'");
  }

  if (cooling->Shielding_choice == Shielding_none)
    cooling->ChimesGlobalVars.cellSelfShieldingOn = 0;
  else if (cooling->Shielding_choice == Shielding_Jeans_length)
    cooling->ChimesGlobalVars.cellSelfShieldingOn = 1;
  else if (cooling->Shielding_choice == Shielding_COLIBRE)
    cooling->ChimesGlobalVars.cellSelfShieldingOn = 1;

  if (cooling->Shielding_choice == Shielding_Jeans_length) {

    /* Maximum shielding length (in kpc).
     *If negative, do not impose a maximum. */
    cooling->max_shielding_length_kpc = parser_get_opt_param_double(
        parameter_file, "CHIMESCooling:max_shielding_length_kpc", -1.0);

    /* Convert the maximum shielding length to cm. */
    cooling->max_shielding_length_cm =
        cooling->max_shielding_length_kpc * 1.0e3 *
        units_cgs_conversion_factor(us, UNIT_CONV_LENGTH) *
        phys_const->const_parsec;
  }

  if (cooling->Shielding_choice != Shielding_none) {

    /* This parameter is needed if shielding is switched on */
    /* Factor to re-scale shielding length. */
    cooling->shielding_length_factor = parser_get_opt_param_double(
        parameter_file, "CHIMESCooling:shielding_length_factor", 1.0);
  }

  if (cooling->Shielding_choice == Shielding_COLIBRE) {
    /* The maximum shielding length can depend on redshift */
    /* in the COLIBRE shielding model */
    cooling->lshmax_lowz_kpc = parser_get_opt_param_float(
        parameter_file, "CHIMESCooling:LSHmax_lowz_kpc", 50.0);
    cooling->lshmax_highz_kpc = parser_get_opt_param_float(
        parameter_file, "CHIMESCooling:LSHmax_highz_kpc", 10.0);
    cooling->lshmax_transition_redshift = parser_get_opt_param_float(
        parameter_file, "CHIMESCooling:LSHmax_transition_redshift", 7.0);
    cooling->lshmax_transition_width_dz = parser_get_opt_param_float(
        parameter_file, "CHIMESCooling:LSHmax_transition_width_dz", 0.2);
    cooling->lshmax_transition_steepness_k = parser_get_opt_param_float(
        parameter_file, "CHIMESCooling:LSHmax_transition__steepness_k", 100.);
  }

  if ((cooling->Shielding_choice == Shielding_COLIBRE) ||
      (cooling->UV_field_choice == UV_field_COLIBRE)) {

    /* Read in the parameters needed for calculating Nref;
     * The reference column density is used for the shielding length
     * if Shielding_flag == COLIBRE
     * and for the ISRF if UV_field_flag == COLIBRE */

    /* Constant hydrogen mass fraction to calculate Nref */
    cooling->nref_XH_const = parser_get_opt_param_double(
        parameter_file, "CHIMESCooling:NREF_XH_const", 0.75);

    /* Constant mean particle mass to calculate Nref */
    cooling->nref_mu_const = parser_get_opt_param_double(
        parameter_file, "CHIMESCooling:NREF_mu_const", 1.24);

    /* Minimum column density for Nref */
    cooling->nref_column_density_min_cgs = parser_get_opt_param_double(
        parameter_file, "CHIMESCooling:NREF_column_density_min_cgs", 3.086e15);

    /* Maximum column density for Nref */
    cooling->nref_column_density_max_cgs = parser_get_opt_param_double(
        parameter_file, "CHIMESCooling:NREF_column_density_max_cgs", 1.0e24);

    /* Minimum temperature for transition to nref_column_density_min_cgs */
    cooling->nref_trans_temperature_min_K = parser_get_opt_param_float(
        parameter_file, "CHIMESCooling:NREF_trans_temperature_min_K", 1.e4);

    /* Maximum temperature for transition to nref_column_density_min_cgs */
    cooling->nref_trans_temperature_max_K = parser_get_opt_param_float(
        parameter_file, "CHIMESCooling:NREF_trans_temperature_max_K", 1.e5);

    /* Steepness of the transition to nref_column_density_min_cgs */
    cooling->nref_trans_steepness_k = parser_get_opt_param_float(
        parameter_file, "CHIMESCooling:NREF_trans_steepness_k", 2.);

    /* Maximum length of reference column density */
    cooling->nref_max_length_kpc = parser_get_opt_param_double(
        parameter_file, "CHIMESCooling:NREF_column_length_max_kpc", 50.);

    /* Convert the maximum length of reference column density to cm */
    cooling->nref_max_length_cgs =
        cooling->nref_max_length_kpc * 1.0e3 *
        units_cgs_conversion_factor(us, UNIT_CONV_LENGTH) *
        phys_const->const_parsec;
  }

  if (cooling->UV_field_choice == UV_field_COLIBRE) {

    /* cosmic ray rate in s-1 at the transition column density */
    cooling->cr_rate_at_transition_cgs = parser_get_opt_param_double(
        parameter_file, "CHIMESCooling:CR_rate_at_transition_cgs", 2.e-16);

    /* cosmic ray transition column density from power law 1 to 2 in cm-2 */
    cooling->cr_transition_column_density_cgs = parser_get_opt_param_double(
        parameter_file, "CHIMESCooling:CR_transition_column_density_cgs",
        1.e21);

    /* cosmic rays: slope of power law 1 */
    cooling->cr_powerlaw_slope_1 = parser_get_opt_param_float(
        parameter_file, "CHIMESCooling:CR_powerlaw_slope_1", 1.4);

    /* cosmic rays: slope of power law 2 */
    cooling->cr_powerlaw_slope_2 = parser_get_opt_param_float(
        parameter_file, "CHIMESCooling:CR_powerlaw_slope_2", 0.0);

    /* cosmic rays: overdensity (rel to rho_crit) cutoff */
    cooling->cr_cutoff_over_density = parser_get_opt_param_float(
        parameter_file, "CHIMESCooling:CR_cutoff_over_density", 100.);

    /* cosmic rays: physical density cutoff in cm-3 */
    cooling->cr_cutoff_phys_density_cgs = parser_get_opt_param_float(
        parameter_file, "CHIMESCooling:CR_cutoff_phys_density_cgs", 1.e-2);

    /* cosmic rays: width of cutoff in dex (gas density) */
    cooling->cr_cutoff_width_dex = parser_get_opt_param_float(
        parameter_file, "CHIMESCooling:CR_cutoff_width_dex", 1.);

    /* cosmic rays: steepness of the cutoff */
    cooling->cr_cutoff_steepness_k = parser_get_opt_param_float(
        parameter_file, "CHIMESCooling:CR_cutoff_steepness_k", 2.);

    /* Ionising photon density relative to that in the cross sections file
     * at the transition column density */
    cooling->isrf_norm_at_transition = parser_get_opt_param_float(
        parameter_file, "CHIMESCooling:ISRF_norm_at_transition", 19.);

    /* ISRF transition column density from power law 1 to 2 in cm-2 */
    cooling->isrf_transition_column_density_cgs = parser_get_opt_param_double(
        parameter_file, "CHIMESCooling:ISRF_transition_column_density_cgs",
        1.e22);

    /* ISRF: slope of power law 1 */
    cooling->isrf_powerlaw_slope_1 = parser_get_opt_param_float(
        parameter_file, "CHIMESCooling:ISRF_powerlaw_slope_1", 1.4);

    /* ISRF: slope of power law 2 */
    cooling->isrf_powerlaw_slope_2 = parser_get_opt_param_float(
        parameter_file, "CHIMESCooling:ISRF_powerlaw_slope_2", 0.0);

    /* ISRF: overdensity (rel to rho_crit) cutoff */
    cooling->isrf_cutoff_over_density = parser_get_opt_param_float(
        parameter_file, "CHIMESCooling:ISRF_cutoff_over_density", 100.);

    /* ISRF: physical cutoff in cm-3 */
    cooling->isrf_cutoff_phys_density_cgs = parser_get_opt_param_float(
        parameter_file, "CHIMESCooling:ISRF_cutoff_phys_density_cgs", 1.e-2);

    /* ISRF: width of cutoff in dex (gas density) */
    cooling->isrf_cutoff_width_dex = parser_get_opt_param_float(
        parameter_file, "CHIMESCooling:ISRF_cutoff_width_dex", 1.);

    /* ISRF: steepness of the cutoff */
    cooling->isrf_cutoff_steepness_k = parser_get_opt_param_float(
        parameter_file, "CHIMESCooling:ISRF_cutoff_steepness_k", 2.);

  } else {

    /* Cosmic ray ionisation rate of HI. */
    cooling->cosmic_ray_rate = (ChimesFloat)parser_get_param_double(
        parameter_file, "CHIMESCooling:cosmic_ray_rate_cgs");
  }

  /* In CHIMES, the UVB is switched off above
   * the redshift of H-reionisation. */
  cooling->ChimesGlobalVars.reionisation_redshift =
      (ChimesFloat)parser_get_param_float(parameter_file,
                                          "CHIMESCooling:UVB_cutoff_z");

  /* Parameters used for the COLIBRE ISRF and shielding length. These have
   * just been hard-coded for now
   * The values are taken from Ploeckinger & Schaye (2020). */
  cooling->N_H0 = 3.63e20; /* cgs */

  /* Parameter to control use of turbulent Jeans length.
   * Only used with the Colibre UV field and shielding options. */
  cooling->colibre_use_turbulent_jeans_length = parser_get_param_int(
      parameter_file, "CHIMESCooling:colibre_use_turbulent_jeans_length");

  /* Turbulent 1D velocity dispersion, in km/s. */
  cooling->sigma_turb_km_p_s = (ChimesFloat)parser_get_param_float(
      parameter_file, "CHIMESCooling:turbulent_velocity_dispersion_km_p_s");

  parser_get_param_string(parameter_file, "CHIMESCooling:init_abundance_mode",
                          string_buffer);
  if (strcmp(string_buffer, "single") == 0) {
    cooling->init_abundance_choice = init_abundance_single;
  } else if (strcmp(string_buffer, "read") == 0) {
    cooling->init_abundance_choice = init_abundance_read;
  } else if (strcmp(string_buffer, "compute") == 0) {
    cooling->init_abundance_choice = init_abundance_compute;
  } else {
    error(
        "Invalid choice of 'CHIMESCooling:init_abundance_mode'. Must be "
        "'single', 'read', or 'compute'");
  }

  if (cooling->init_abundance_choice == init_abundance_single)
    cooling->InitIonState =
        parser_get_param_int(parameter_file, "CHIMESCooling:InitIonState");
  else
    cooling->InitIonState = 0;

  /* CHIMES tolerance parameters */
  cooling->ChimesGlobalVars.relativeTolerance =
      (ChimesFloat)parser_get_param_double(parameter_file,
                                           "CHIMESCooling:relativeTolerance");
  cooling->ChimesGlobalVars.absoluteTolerance =
      (ChimesFloat)parser_get_param_double(parameter_file,
                                           "CHIMESCooling:absoluteTolerance");
  cooling->ChimesGlobalVars.explicitTolerance =
      (ChimesFloat)parser_get_param_double(parameter_file,
                                           "CHIMESCooling:explicitTolerance");
  cooling->ChimesGlobalVars.scale_metal_tolerances =
      (ChimesFloat)parser_get_param_double(
          parameter_file, "CHIMESCooling:scale_metal_tolerances");

  /* Maximum temperature for the molecular network */
  cooling->ChimesGlobalVars.Tmol_K = (ChimesFloat)parser_get_param_double(
      parameter_file, "CHIMESCooling:Tmol_K");

  /* For now, the CMB temperature in CHIMES is just set
   * to the present day value. For cosmological runs, this
   * will need to be updated for the current redshift. */

  /* CMB temperature at redshift zero. */
  cooling->T_CMB_0 = phys_const->const_T_CMB_0 *
                     units_cgs_conversion_factor(us, UNIT_CONV_TEMPERATURE);

  /* Initially set the CMB temperature in CHIMES to the redshift zero value.
   * This will be updated to the current redshift in cooling_update(). */
  cooling->ChimesGlobalVars.cmb_temperature = (ChimesFloat)cooling->T_CMB_0;

  /* Equilibrium mode:
   * 0 --> Evolve in non-equilibrium.
   * 1 --> Evolve with equilibrium abundances.
   */
  cooling->ChemistryEqmMode =
      parser_get_param_int(parameter_file, "CHIMESCooling:ChemistryEqmMode");

  /* Flag switch thermal evolution on/off:
   * 0 --> Fixed temperature (but chemical abundances can still be evolved).
   * 1 --> Enable thermal evolution.
   */
  cooling->ThermEvolOn =
      parser_get_param_int(parameter_file, "CHIMESCooling:ThermEvolOn");

  /* Flag to switch on extra debug print
   * statements in CHIMES.
   * 0 --> No extra debug prints.
   * 1 --> If CVode returns a non-zero flag (i.e. it produces a
   *       CVODE error/warning), then we print out everything in
   *       the ChimesGasVars struct.
   */
  cooling->ChimesGlobalVars.chimes_debug =
      parser_get_param_int(parameter_file, "CHIMESCooling:chimes_debug");

  /* The following CHIMES parameters do not need to be modified
   * by the user. These are just hard-coded for now. */
  cooling->ChimesGlobalVars.grain_temperature = parser_get_opt_param_float(
      parameter_file, "CHIMESCooling:dust_grain_temperature_K", 10.f);

  /* Physical velocity divergence isn't easily accessible, so just
   * run with static molecular cooling for now. */
  cooling->ChimesGlobalVars.StaticMolCooling = 1;

  /* Optional parameters to define S and Ca
   * relative to Si. */
  cooling->S_over_Si_ratio_in_solar = parser_get_opt_param_float(
      parameter_file, "CHIMESCooling:S_over_Si_in_solar", 1.f);
  cooling->Ca_over_Si_ratio_in_solar = parser_get_opt_param_float(
      parameter_file, "CHIMESCooling:Ca_over_Si_in_solar", 1.f);

  /* Solar mass fractions of Si, S and Ca are hard-coded here,
   * because we might not have access to the COLIBRE cooling tables to
   * get these. */
  cooling->Si_solar_mass_fraction = 6.64948e-4;
  cooling->S_solar_mass_fraction = 3.09171e-4;
  cooling->Ca_solar_mass_fraction = 6.41451e-5;

  /* CHIMES uses a solar metallicity of 0.0129. */
  cooling->Zsol = 0.0129;

  /* Dust depletion factors in the solar neighbourhood.
   * Taken from Ploeckinger & Schaye (2020) */
  cooling->f_dust0_C = 0.34385;
  cooling->f_dust0_O = 0.31766;
  cooling->f_dust0_Mg = 0.94338;
  cooling->f_dust0_Si = 0.94492;
  cooling->f_dust0_Ca = 0.9999;
  cooling->f_dust0_Fe = 0.99363;

  /* Parameter that controls whether to reduce gas-phase metal abundances
   * due to dust depletion. */
  cooling->colibre_metal_depletion = parser_get_param_int(
      parameter_file, "CHIMESCooling:colibre_metal_depletion");

  /* Distance from EOS to use thermal equilibrium temperature for subgrid
   * props, and to evolve the chemistry in equilibrium. */
  cooling->dlogT_EOS = parser_get_param_float(
      parameter_file, "CHIMESCooling:delta_logTEOS_subgrid_properties");

  /* Flag to use the subgrid density and
   * temperature from the COLIBRE cooling tables. */
  cooling->use_colibre_subgrid_EOS = parser_get_param_int(
      parameter_file, "CHIMESCooling:use_colibre_subgrid_EOS");

  /* Flag to set particles to chemical equilibrium
   * if they have been heated by feedback. */
  cooling->set_FB_particles_to_eqm = parser_get_param_int(
      parameter_file, "CHIMESCooling:set_FB_particles_to_eqm");

  /* Threshold in dt / t_cool above which we are in the rapid cooling regime.
   * If negative, we never use this scheme (i.e. always drift the internal
   * energies). */
  cooling->rapid_cooling_threshold = parser_get_param_double(
      parameter_file, "CHIMESCooling:rapid_cooling_threshold");

  /* Properties of the HII region model */
  cooling->HIIregion_fion = parser_get_opt_param_float(
      parameter_file, "COLIBREFeedback:HIIregion_ionization_fraction", 1.0);
  cooling->HIIregion_temp = parser_get_opt_param_float(
      parameter_file, "COLIBREFeedback:HIIregion_temperature_K", 1.0e4);

  /* Maximum boost factor for reactions on the surfaces
   * of dust grains. */
  cooling->max_dust_boost_factor = parser_get_param_float(
      parameter_file, "CHIMESCooling:max_dust_boost_factor");

  /* Density below which the dust boost factor is
   * equal to unity, in units of cm^-3. */
  cooling->dust_boost_nH_min_cgs = parser_get_param_float(
      parameter_file, "CHIMESCooling:dust_boost_nH_min_cgs");

  /* Density above which the dust boost factor is
   * equal to the maximum, in units of cm^-3. */
  cooling->dust_boost_nH_max_cgs = parser_get_param_float(
      parameter_file, "CHIMESCooling:dust_boost_nH_max_cgs");

  /* Flag to indicate whether or not we redistribute FB heated
   * particle's dust in the cooling step */
  cooling->destroy_FB_heated_dust = parser_get_param_int(
      parameter_file, "CHIMESCooling:destroy_FB_heated_dust");

  /* Check that the density thresholds for the dust boost factor are sensible.
   * They must both be positive (as we will be taking logs), and nH_min must be
   * less than nH_max. */
  if ((cooling->dust_boost_nH_min_cgs <= 0.0f) ||
      (cooling->dust_boost_nH_max_cgs <= 0.0f) ||
      (cooling->dust_boost_nH_min_cgs >= cooling->dust_boost_nH_max_cgs))
    error(
        "CHIMES Error: dust_boost_nH_min_cgs = %.2e, dust_boost_nH_max_cgs = "
        "%.2e. Both must be positive, and dust_boost_nH_min_cgs must be less "
        "than dust_boost_nH_max_cgs.",
        cooling->dust_boost_nH_min_cgs, cooling->dust_boost_nH_max_cgs);

  /* Switch for Hybrid cooling */
  cooling->ChimesGlobalVars.hybrid_cooling_mode =
      parser_get_param_int(parameter_file, "CHIMESCooling:use_hybrid_cooling");

  /* Determine which metal elements to include in the CHIMES network.
   * Note that H and He are always included. */
  cooling->ChimesGlobalVars.element_included[0] =
      parser_get_param_int(parameter_file, "CHIMESCooling:IncludeCarbon");
  cooling->ChimesGlobalVars.element_included[1] =
      parser_get_param_int(parameter_file, "CHIMESCooling:IncludeNitrogen");
  cooling->ChimesGlobalVars.element_included[2] =
      parser_get_param_int(parameter_file, "CHIMESCooling:IncludeOxygen");
  cooling->ChimesGlobalVars.element_included[3] =
      parser_get_param_int(parameter_file, "CHIMESCooling:IncludeNeon");
  cooling->ChimesGlobalVars.element_included[4] =
      parser_get_param_int(parameter_file, "CHIMESCooling:IncludeMagnesium");
  cooling->ChimesGlobalVars.element_included[5] =
      parser_get_param_int(parameter_file, "CHIMESCooling:IncludeSilicon");
  cooling->ChimesGlobalVars.element_included[6] =
      parser_get_param_int(parameter_file, "CHIMESCooling:IncludeSulphur");
  cooling->ChimesGlobalVars.element_included[7] =
      parser_get_param_int(parameter_file, "CHIMESCooling:IncludeCalcium");
  cooling->ChimesGlobalVars.element_included[8] =
      parser_get_param_int(parameter_file, "CHIMESCooling:IncludeIron");

  /* The COLIBRE cooling tables are needed if we are using hybrid cooling or
   * if we are using the subgrid density and temperature on the EOS. */
  if ((cooling->ChimesGlobalVars.hybrid_cooling_mode == 1) ||
      (cooling->use_colibre_subgrid_EOS == 1)) {

    /* Path to colibre cooling table */
    parser_get_param_string(parameter_file, "CHIMESCooling:colibre_table_path",
                            cooling->colibre_table.cooling_table_path);

    /* Filebase for split cooling tables */
    parser_get_param_string(parameter_file,
                            "CHIMESCooling:colibre_table_filebase",
                            cooling->colibre_table.cooling_table_filebase);

    /* Set the S/Si and Ca/Si ratios to
     * the values already read in for CHIMES. */
    cooling->colibre_table.S_over_Si_ratio_in_solar =
        cooling->S_over_Si_ratio_in_solar;
    cooling->colibre_table.Ca_over_Si_ratio_in_solar =
        cooling->Ca_over_Si_ratio_in_solar;

    /* Physical constants that we will need, in cgs units */
    const float dimension_k[5] = {1, 2, -2, 0, -1};
    const float dimension_G[5] = {-1, 3, -2, 0, 0};
    cooling->boltzmann_k_cgs =
        phys_const->const_boltzmann_k *
        units_general_cgs_conversion_factor(us, dimension_k);
    cooling->newton_G_cgs =
        phys_const->const_newton_G *
        units_general_cgs_conversion_factor(us, dimension_G);
    cooling->proton_mass_cgs = phys_const->const_proton_mass *
                               units_cgs_conversion_factor(us, UNIT_CONV_MASS);

    cooling->proton_mass_over_boltzmann_k_cgs =
        cooling->proton_mass_cgs / cooling->boltzmann_k_cgs;

    cooling->colibre_table.log10_kB_cgs = log10(cooling->boltzmann_k_cgs);
    cooling->colibre_table.inv_proton_mass_cgs = 1. / cooling->proton_mass_cgs;
    cooling->colibre_table.proton_mass_cgs = cooling->proton_mass_cgs;

    /* Compute conversion factors */
    cooling->internal_energy_to_cgs =
        units_cgs_conversion_factor(us, UNIT_CONV_ENERGY_PER_UNIT_MASS);
    cooling->density_to_cgs =
        units_cgs_conversion_factor(us, UNIT_CONV_DENSITY);
    cooling->number_density_to_cgs =
        units_cgs_conversion_factor(us, UNIT_CONV_NUMBER_DENSITY);
    cooling->time_to_cgs = units_cgs_conversion_factor(us, UNIT_CONV_TIME);
    cooling->internal_energy_over_time_to_cgs = units_cgs_conversion_factor(
        us, UNIT_CONV_ENERGY_PER_UNIT_MASS_PER_UNIT_TIME);

    cooling->internal_energy_from_cgs = 1. / cooling->internal_energy_to_cgs;
    cooling->internal_energy_over_time_from_cgs =
        1. / cooling->internal_energy_over_time_to_cgs;

    /* Same for the COLIBRE tables */
    cooling->colibre_table.pressure_to_cgs =
        units_cgs_conversion_factor(us, UNIT_CONV_PRESSURE);
    cooling->colibre_table.internal_energy_to_cgs =
        units_cgs_conversion_factor(us, UNIT_CONV_ENERGY_PER_UNIT_MASS);
    cooling->colibre_table.internal_energy_from_cgs =
        1. / cooling->colibre_table.internal_energy_to_cgs;
    cooling->colibre_table.number_density_to_cgs =
        units_cgs_conversion_factor(us, UNIT_CONV_NUMBER_DENSITY);
    cooling->colibre_table.number_density_from_cgs =
        1. / cooling->colibre_table.number_density_to_cgs;
    cooling->colibre_table.density_to_cgs =
        units_cgs_conversion_factor(us, UNIT_CONV_DENSITY);
    cooling->colibre_table.density_from_cgs =
        1. / cooling->colibre_table.density_to_cgs;

    /* Pre-factor for the Compton y terms */
    cooling->y_compton_factor =
        phys_const->const_thomson_cross_section *
        phys_const->const_boltzmann_k /
        (phys_const->const_speed_light_c * phys_const->const_speed_light_c *
         phys_const->const_electron_mass);

    /* Recover the minimal energy allowed (in internal units) */
    const double u_min = hydro_props->minimal_internal_energy;

    /* Convert to CGS units */
    cooling->colibre_table.umin_cgs =
        u_min * cooling->colibre_table.internal_energy_to_cgs;

    /* If running with dust, read if pairing to cooling? */
#if defined(DUST_NONE)
    dp->pair_to_cooling = 0;
#else
    dp->pair_to_cooling = parser_get_opt_param_int(
        parameter_file, "DustEvolution:pair_to_cooling", 1);

    if (cooling->colibre_metal_depletion && dp->pair_to_cooling) {
      error(
          "May not run with DustEvolution:pair_to_cooling:1 and "
          "CHIMESCooling:colibre_metal_depletion:1 simultaneously, as "
          "this double-counts dust depletion effects. Please try switching "
          "one of these flags off in the configuration file.");
    }
#endif

    /* Read the Colibre table. */
    message("Reading Colibre cooling table.");

    get_cooling_redshifts(&(cooling->colibre_table));

    /* Read in cooling table header */
    char fname[colibre_table_path_name_length + colibre_table_filebase_length +
               12];
    sprintf(fname, "%s/%s_z0.00.hdf5",
            cooling->colibre_table.cooling_table_path,
            cooling->colibre_table.cooling_table_filebase);

    read_cooling_header(fname, &(cooling->colibre_table));

    /* Allocate memory for the tables */
    allocate_cooling_tables(&(cooling->colibre_table));

    /* Set the redshift indices to invalid values */
    cooling->colibre_table.z_index = -10;

    /* set previous_z_index and to last value of redshift table*/
    cooling->colibre_table.previous_z_index =
        cooling->colibre_table.number_of_redshifts - 2;
  }

  /* Set redshift to a very high value, just
   * while we initialise the CHIMES module.
   * It will be set to the correct redshift in
   * the cooling_update() routine. */
  cooling->ChimesGlobalVars.redshift = 1000.0;

  if ((cooling->UVB_dependence_choice == UVB_zdep_cross_secs_and_eqm) &&
      (cooling->UV_field_choice == UV_field_COLIBRE) &&
      (cooling->Shielding_choice == Shielding_COLIBRE)) {

    /* If we are using the redshift-dependent CHIMES eqm abundances with the
     * COLIBRE radiation field and shielding models, we need to check that the
     * parameters used in the tables match the hard-coded values. */
    char fname_chimes[520];
    sprintf(fname_chimes, "%s/z%.3f_eqm.hdf5",
            cooling->ChimesGlobalVars.EqAbundanceTablePath, 0.0);

    /* Implement better version control check but not while we are
     * still testing */
    /* chimes_check_colibre_eqm_tables_header(fname_chimes, cooling); */
  }

  /* Initialise the CHIMES module. */
  message("Initialising CHIMES cooling module.");
  init_chimes(&cooling->ChimesGlobalVars);

  /* Check that the size of the CHIMES network
   * determined by the various Include<Element>
   * parameters matches the size specified
   * through the configure script. */
  if (cooling->ChimesGlobalVars.totalNumberOfSpecies != CHIMES_NETWORK_SIZE)
    error(
        "The size of the CHIMES network based on the Include<Element> "
        "parameters is %d. However, Swift was configured and compiled using "
        "--with-cooling=CHIMES_%d. If you are sure that you have "
        "correctly set which metals to include in the network in the parameter "
        "file, then you will need to re-compile Swift using ./configure "
        "--with-cooling=CHIMES_%d.",
        cooling->ChimesGlobalVars.totalNumberOfSpecies, CHIMES_NETWORK_SIZE,
        cooling->ChimesGlobalVars.totalNumberOfSpecies);

  /* Construct the array of species names
   * for the reduced network. */
  for (int idx = 0; idx < CHIMES_TOTSIZE; idx++) {
    if (cooling->ChimesGlobalVars.speciesIndices[idx] >= 0)
      strcpy(cooling->chimes_species_names_reduced[cooling->ChimesGlobalVars
                                                       .speciesIndices[idx]],
             chimes_species_names[idx]);
  }

  if (cooling->ChimesGlobalVars.hybrid_cooling_mode == 1) {

    /* Create data structure for hybrid cooling,
     * and store pointer to the Colibre table. */
    cooling->ChimesGlobalVars.hybrid_data =
        (void *)malloc(sizeof(struct global_hybrid_data_struct));
    struct global_hybrid_data_struct *myData;
    myData = (struct global_hybrid_data_struct *)
                 cooling->ChimesGlobalVars.hybrid_data;
    myData->table = &(cooling->colibre_table);
    myData->Zsol = cooling->Zsol;

    /* Set the hybrid cooling function pointers. */
    cooling->ChimesGlobalVars.hybrid_cooling_fn =
        &colibre_metal_cooling_rate_temperature;
    cooling->ChimesGlobalVars.allocate_gas_hybrid_data_fn =
        &chimes_allocate_gas_hybrid_data;
    cooling->ChimesGlobalVars.free_gas_hybrid_data_fn =
        &chimes_free_gas_hybrid_data;
  }
}

/**
 * @brief Prints the properties of the cooling model to stdout.
 *
 * @param cooling #cooling_function_data struct.
 */
void cooling_print_backend(const struct cooling_function_data *cooling) {
  message("Cooling function is 'SPLITCHIMES'.");
}

/**
 * @brief Common operations performed on the cooling function at a
 * given time-step or redshift.
 *
 * @param cosmo The current cosmological model.
 * @param pressure_floor Properties of the pressure floor.
 * @param cooling The #cooling_function_data used in the run.
 * @param s The #space containing all the particles.
 * @param dustevo The dust evolution properties
 * @param time The current system time
 */
void cooling_update(const struct cosmology *cosmo,
                    const struct pressure_floor_props *pressure_floor,
                    struct cooling_function_data *cooling, struct space *s,
                    struct dustevo_props *dustevo, const double time) {

  /* Update redshift stored in ChimesGlobalVars. */
  cooling->ChimesGlobalVars.redshift = cosmo->z;

  /* Update T_CMB in CHIMES to current redshift. */
  cooling->ChimesGlobalVars.cmb_temperature =
      (ChimesFloat)cooling->T_CMB_0 * (1.0 + cosmo->z);

  /* Update redshift-dependent UVB. */
  if (cooling->UVB_dependence_choice == UVB_zdep_cross_secs ||
      cooling->UVB_dependence_choice == UVB_zdep_cross_secs_and_eqm)
    interpolate_redshift_dependent_UVB(&cooling->ChimesGlobalVars);

  /* The COLIBRE cooling tables are needed
   * if we are using hybrid cooling or
   * if we are using the subgrid density
   * and temperature on the EOS. */
  if ((cooling->ChimesGlobalVars.hybrid_cooling_mode == 1) ||
      (cooling->use_colibre_subgrid_EOS == 1)) {

    /* Current redshift */
    const float redshift = cosmo->z;

    /* What is the current table index along the redshift axis? */
    int z_index = -1;
    float dz = 0.f;
    get_redshift_index(redshift, &z_index, &dz, &(cooling->colibre_table));
    cooling->colibre_table.dz = dz;

    /* Do we already have the correct tables loaded? */
    if (cooling->colibre_table.z_index == z_index) return;

    /* Normal case: two tables bracketing the current z */
    const int low_z_index = z_index;
    const int high_z_index = z_index + 1;

    get_cooling_table(&(cooling->colibre_table), low_z_index, high_z_index,
                      dustevo->logfD, dustevo->pair_to_cooling);

    /* Store the currently loaded index */
    cooling->colibre_table.z_index = z_index;
  }
}

/**
 * @brief Update ChimesGasVars structure.
 *
 * Updates the ChimesGasVars structure with the various
 * thermodynamics quantities of the gas particle that
 * will be needed for the CHIMES chemistry solver.
 *
 * @param u_cgs The internal energy, in cgs units.
 * @param phys_const The physical constants in internal units.
 * @param us The internal system of units.
 * @param cosmo The current cosmological model.
 * @param hydro_properties the hydro_props struct
 * @param floor_props Properties of the entropy floor.
 * @param cooling The #cooling_function_data used in the run.
 * @param p Pointer to the particle data.
 * @param xp Pointer to the extended particle data.
 * @param ChimesGasVars CHIMES gasVariables structure.
 * @param dt_cgs The cooling time-step of this particle.
 */
void chimes_update_gas_vars(
    const double u_cgs, const struct phys_const *phys_const,
    const struct unit_system *us, const struct cosmology *cosmo,
    const struct hydro_props *hydro_properties,
    const struct entropy_floor_properties *floor_props,
    const struct cooling_function_data *cooling, const struct dustevo_props *dp,
    const struct part *p, const struct xpart *xp,
    struct gasVariables *ChimesGasVars, const float dt_cgs) {

  /* Get the mean molecular weight */
  const double mu = chimes_mu(cooling, hydro_properties, p, xp);

  /* Get the Hydrogen mass fraction */
  const ChimesFloat XH = get_hydrogen_mass_fraction(p, hydro_properties);

  /* Get the metallicity (metal mass fraction / solar metal mass frac.) */
#if defined(CHEMISTRY_COLIBRE) || defined(CHEMISTRY_EAGLE)

  ChimesGasVars->metallicity =
      (ChimesFloat)chemistry_get_total_metal_mass_fraction_for_cooling(p);
  ChimesGasVars->metallicity /= cooling->Zsol;
  ChimesGasVars->metallicity = max(ChimesGasVars->metallicity, FLT_MIN);

#else
  /* Without COLIBRE or EAGLE chemistry, the metal abundances are unavailable.
   * --> Set to primordial abundances. */
  ChimesGasVars->metallicity = 0.0;
#endif /* CHEMISTRY_COLIBRE || CHEMISTRY_EAGLE */

  dust_check_for_NaNs(p)

      /* Get the dust ratio */
      ChimesGasVars->dust_ratio =
      get_dust_ratio(p, cooling, dp, ChimesGasVars->metallicity);

  /* Limit imposed by the entropy floor */
  const double A_floor = entropy_floor(p, cosmo, floor_props);
  const double rho = hydro_get_physical_density(p, cosmo);
  const double u_floor = gas_internal_energy_from_entropy(rho, A_floor);
  const double u_floor_cgs = u_floor * cooling->internal_energy_to_cgs;

  /* Set the temperature floor */
  if (xp->tracers_data.HIIregion_timer_gas > 0.) {

    /* Particles is an HII region. */

    /* Temperature on the entropy floor, but using  the HII region mean
     * molecular weight. */
    const double mu_HII = get_mu_HII_region(cooling, XH);
    const double T_floor = get_T_from_u_cgs(u_floor_cgs, mu_HII, cooling);

    /* Set the temperature floor in CHIMES to be  either T_floor or
     * HIIregion_temp, whichever is higher. */
    ChimesGasVars->TempFloor =
        (ChimesFloat)chimes_max(T_floor, cooling->HIIregion_temp);

    /* Check TempFloor doesn't go below minimal_temperature. */
    ChimesGasVars->TempFloor =
        chimes_max(ChimesGasVars->TempFloor,
                   (ChimesFloat)hydro_properties->minimal_temperature);

  } else {

    /* Particle is *not* an HII region. */

    /* Temperature on the entropy floor using current mean molecular weight. */
    const double T_floor = get_T_from_u_cgs(u_floor_cgs, mu, cooling);

    /* If using the entropy floor, passing it as a negative number means
     * CHIMES will interpret it as a minimum thermal energy, in erg/cm^3,
     * rather than a minimum temperature. */
    if (T_floor < hydro_properties->minimal_temperature) {
      ChimesGasVars->TempFloor =
          (ChimesFloat)hydro_properties->minimal_temperature;
    } else {
      ChimesGasVars->TempFloor =
          (ChimesFloat)(-u_floor_cgs * rho * cooling->density_to_cgs);
    }
  }

  ChimesGasVars->temp_floor_mode = 1;

  /* Select energy to use based on where we sit w.r.t. the entropy floor */
  double u_actual_cgs;
  if (u_cgs < u_floor_cgs) {

    /* Particle is below the entropy floor.
     *
     * Set internal energy to the floor and chemistry will be evolved in
     * equilibrium. */
    u_actual_cgs = u_floor_cgs;
    ChimesGasVars->ForceEqOn = 1;

  } else if (u_cgs < exp10(cooling->dlogT_EOS) * u_floor_cgs) {

    /* Particle is above the entropy floor, but close enough that we will need
     * to evolve the chemistry in equilibrium. */
    u_actual_cgs = u_cgs;
    ChimesGasVars->ForceEqOn = 1;

  } else {

    /* Particle is well above the entropy floor.
     *
     * Evolve chemistry as usual, according to the user-provided parameter. */
    u_actual_cgs = u_cgs;
    ChimesGasVars->ForceEqOn = cooling->ChemistryEqmMode;
  }

  /* Now convert this to a temperature */
  ChimesGasVars->temperature =
      (ChimesFloat)get_T_from_u_cgs(u_actual_cgs, mu, cooling);

  /* Recover the Hydrogen number density */
  const double nH =
      hydro_get_physical_density(p, cosmo) * XH / phys_const->const_proton_mass;
  const double nH_cgs = (nH * cooling->number_density_to_cgs);
  ChimesGasVars->nH_tot = (ChimesFloat)nH_cgs;

  /* Recover the dust boost factor to use based on the H number density */
  ChimesGasVars->dust_boost_factor = get_dust_boost_factor(nH_cgs, cooling);

  ChimesGasVars->hydro_timestep = (ChimesFloat)dt_cgs;
  ChimesGasVars->constant_heating_rate = 0.0;
  ChimesGasVars->ThermEvolOn = cooling->ThermEvolOn;
  ChimesGasVars->divVel = 0.0;

  /* N_ref is used by both the COLIBRE ISRF and the COLIBRE shielding length,
   * and to scale metal depletion. Based on Ploeckinger & Schaye (2020). */
  double N_ref_cgs = 0.0;
  if ((cooling->UV_field_choice == UV_field_COLIBRE) ||
      (cooling->Shielding_choice == Shielding_COLIBRE) ||
      (cooling->colibre_metal_depletion)) {

    N_ref_cgs = calculate_colibre_N_ref(phys_const, us, cosmo, cooling,
                                        hydro_properties, p, xp,
                                        ChimesGasVars->temperature);
  }

  /* Get the CR rate we use */
  ChimesGasVars->cr_rate =
      get_CR_rate(N_ref_cgs, nH_cgs, cosmo, phys_const, cooling);

  /* Set the radiation field by copying over the sectrum parameters from the
   * chimes tables. */
  if (cooling->UV_field_choice == UV_field_user_defined) {

    /* Single, constant radiation field. */
    ChimesGasVars->G0_parameter[0] = chimes_table_spectra.G0_parameter[0];
    ChimesGasVars->H2_dissocJ[0] = chimes_table_spectra.H2_dissocJ[0];
    ChimesGasVars->isotropic_photon_density[0] =
        chimes_table_spectra.isotropic_photon_density[0];
    ChimesGasVars->isotropic_photon_density[0] *=
        cooling->rad_field_norm_factor;

  } else if (cooling->UV_field_choice == UV_field_COLIBRE) {

    /* Extra-galactic UVB */
    ChimesGasVars->G0_parameter[0] = chimes_table_spectra.G0_parameter[0];
    ChimesGasVars->H2_dissocJ[0] = chimes_table_spectra.H2_dissocJ[0];
    ChimesGasVars->isotropic_photon_density[0] =
        chimes_table_spectra.isotropic_photon_density[0];

    /* ISRF, scales with Nref */
    ChimesGasVars->G0_parameter[1] = chimes_table_spectra.G0_parameter[1];
    ChimesGasVars->H2_dissocJ[1] = chimes_table_spectra.H2_dissocJ[1];

    const double log_isrf_min =
        log10(chimes_table_spectra.isotropic_photon_density[1] *
              cooling->isrf_norm_at_transition *
              pow(cooling->nref_column_density_min_cgs /
                      cooling->isrf_transition_column_density_cgs,
                  cooling->isrf_powerlaw_slope_1));

    const double isrf_coldensratio =
        N_ref_cgs / cooling->isrf_transition_column_density_cgs;

    double isophotdens;

    if (isrf_coldensratio < 1.) {
      /* N_ref < N_trans: power law index 1 */
      isophotdens = chimes_table_spectra.isotropic_photon_density[1] *
                    cooling->isrf_norm_at_transition *
                    pow(isrf_coldensratio, cooling->isrf_powerlaw_slope_1);
    } else {
      /* N_ref >=N_trans: power law index 2 */
      isophotdens = chimes_table_spectra.isotropic_photon_density[1] *
                    cooling->isrf_norm_at_transition *
                    pow(isrf_coldensratio, cooling->isrf_powerlaw_slope_2);
    }

    const double log_isophotdens_cgs = log10(isophotdens);

    /* Now, apply some saturation */

    /* low density cutoff between isrf_nmax and isrf_nmin
     * with k regulating the steepness */
    const double isrf_nmax_over_dens =
        cosmo->critical_density * cooling->isrf_cutoff_over_density *
        cooling->nref_XH_const / phys_const->const_proton_mass;

    const double isrf_nmax_over_dens_cgs =
        isrf_nmax_over_dens * cooling->number_density_to_cgs;

    /* Max density */
    const double isrf_nmax =
        max(isrf_nmax_over_dens_cgs, cooling->isrf_cutoff_phys_density_cgs);

    /* Min density: (1. / 10^(cutoff_width_dex)) * isrf_max */
    const double isrf_nmin = exp10(-cooling->isrf_cutoff_width_dex) * isrf_nmax;

    const double numerator = log_isophotdens_cgs - log_isrf_min;
    const double argument =
        sqrt(isrf_nmin / isrf_nmax) * isrf_nmax / ChimesGasVars->nH_tot;
    const double denominator =
        1. + pow(argument, -cooling->isrf_cutoff_steepness_k);

    const double log_isophotdens =
        log_isophotdens_cgs - numerator / denominator;

    ChimesGasVars->isotropic_photon_density[1] = exp10(log_isophotdens);
  }

  double N_scale_rad_cgs = N_ref_cgs;

  if (cooling->colibre_metal_depletion) {

    /* Scale dust_ratio by N_scale_rad, but only if N_scale_rad < N_H0 */
    if (N_scale_rad_cgs < cooling->N_H0)
      ChimesGasVars->dust_ratio *=
          (ChimesFloat)pow(N_scale_rad_cgs / cooling->N_H0, 1.4);
  }

  /* Set shielding length */
  if (cooling->Shielding_choice == Shielding_none) {

    ChimesGasVars->cell_size = 0;

  } else if (cooling->Shielding_choice == Shielding_Jeans_length) {

    /* Jeans length */
    ChimesGasVars->cell_size = (ChimesFloat)sqrt(
        M_PI * hydro_gamma * cooling->boltzmann_k_cgs *
        ((double)ChimesGasVars->temperature) /
        (mu * cooling->newton_G_cgs *
         (cooling->proton_mass_cgs * ((double)ChimesGasVars->nH_tot) / XH) *
         cooling->proton_mass_cgs));
    ChimesGasVars->cell_size *= cooling->shielding_length_factor;

    if ((cooling->max_shielding_length_cm > 0.0) &&
        (ChimesGasVars->cell_size > cooling->max_shielding_length_cm))
      ChimesGasVars->cell_size = (ChimesFloat)cooling->max_shielding_length_cm;

  } else if (cooling->Shielding_choice == Shielding_COLIBRE) {

    double shielding_length_cgs = cooling->shielding_length_factor * N_ref_cgs /
                                  ((double)ChimesGasVars->nH_tot);

    /* limit shielding length by redshift dependent maximum value */
    const double numerator =
        (cooling->lshmax_lowz_kpc - cooling->lshmax_highz_kpc);
    const double argument = sqrt((cooling->lshmax_transition_redshift -
                                  cooling->lshmax_transition_width_dz) *
                                 (cooling->lshmax_transition_redshift)) /
                            fmax(cosmo->z, 1.e-10);
    const double denominator =
        1. + pow(fmin(argument, 10.), cooling->lshmax_transition_steepness_k);
    const double shielding_length_max_kpc =
        cooling->lshmax_lowz_kpc - numerator / denominator;

    double shielding_length_max_cgs =
        shielding_length_max_kpc * 1.0e3 *
        units_cgs_conversion_factor(us, UNIT_CONV_LENGTH) *
        phys_const->const_parsec;

    ChimesGasVars->cell_size =
        (ChimesFloat)(fmin(shielding_length_cgs, shielding_length_max_cgs));
  }

  /* Doppler broadening parameter, for H2 self-shielding.
   * Value is set in the parameter file. */
  ChimesGasVars->doppler_broad = M_SQRT2 * cooling->sigma_turb_km_p_s;

  ChimesGasVars->InitIonState = cooling->InitIonState;

  /* If using hybrid cooling, we need to set the abundance_ratio array using
   * the corresponding routine from COLIBRE. */
  if (cooling->ChimesGlobalVars.hybrid_cooling_mode == 1) {

    struct gas_hybrid_data_struct *myData;
    myData = (struct gas_hybrid_data_struct *)ChimesGasVars->hybrid_data;
    abundance_ratio_to_solar(p, &(cooling->colibre_table),
                             myData->abundance_ratio);
  }
}

/**
 * @brief Update CHIMES element abundances.
 *
 * Updates the element abundances based on the particle's current mass
 * fractions. If the element abundances have changed since the last time this
 * routine was called, for example due to enrichment or metal diffusion, then
 * the ion and molecule abundances will now be inconsistent with their
 * corresponding total element abundances. Therefore, if mode == 1,  we call the
 * check_constraint_equations() routine from the CHIMES module, which adds all
 * species of each element and compares to the total abundance of that element.
 * If they are inconsistent, the species abundances are adjusted, preserving
 * their relative fractions. If a metal is now present that was not present
 * before (e.g. if a particle was primordial but has recently been enriched),
 * then that metal is introduced as fully neutral.
 *
 * We do not call check_constraint_equations() if mode == 0 (for example, when
 * we initialise the abundance arrays for the first time).
 *
 * @param phys_const #phys_const data structure.
 * @param us The internal system of units.
 * @param cosmo #cosmology data structure.
 * @param cooling #cooling_function_data struct.
 * @param hydro_properties The properties of the hydro scheme
 * @param p #part data.
 * @param xp Pointer to the #xpart data.
 * @param ChimesGasVars CHIMES gasVariables structure.
 * @param u_0 internal energy (in code units).
 * @param mode Set to zero if particle not fully initialised.
 */
void chimes_update_element_abundances(
    const struct phys_const *phys_const, const struct unit_system *us,
    const struct cosmology *cosmo, const struct cooling_function_data *cooling,
    const struct hydro_props *hydro_props, const struct part *p,
    const struct xpart *xp, struct gasVariables *ChimesGasVars,
    const double u_0, const int mode) {

  const struct globalVariables *ChimesGlobalVars = &cooling->ChimesGlobalVars;

  /* Special case when there is no chemistry */
#if !defined(CHEMISTRY_COLIBRE) && !defined(CHEMISTRY_EAGLE)

  /* Without COLIBRE or EAGLE chemistry,the metal abundances are unavailable.
   * --> Set to primordial abundances. */
  ChimesGasVars->element_abundances[0] = 0.0833; /* He */

  for (int i = 1; i < 10; i++) ChimesGasVars->element_abundances[i] = 0.0;

  if (mode == 1) check_constraint_equations(ChimesGasVars, ChimesGlobalVars);

  return;
#endif /* !CHEMISTRY_COLIBRE && !CHEMISTRY_EAGLE */

  /* Normal case */

  /* Get element mass fractions */
  const float *const metal_fraction =
      chemistry_get_metal_mass_fraction_for_cooling(p);

  const float XH = metal_fraction[chemistry_element_H];

  ChimesGasVars->element_abundances[0] =
      (ChimesFloat)metal_fraction[chemistry_element_He] / (4.0 * XH);
  ChimesGasVars->element_abundances[1] =
      (ChimesFloat)metal_fraction[chemistry_element_C] / (12.0 * XH);
  ChimesGasVars->element_abundances[2] =
      (ChimesFloat)metal_fraction[chemistry_element_N] / (14.0 * XH);
  ChimesGasVars->element_abundances[3] =
      (ChimesFloat)metal_fraction[chemistry_element_O] / (16.0 * XH);
  ChimesGasVars->element_abundances[4] =
      (ChimesFloat)metal_fraction[chemistry_element_Ne] / (20.0 * XH);
  ChimesGasVars->element_abundances[5] =
      (ChimesFloat)metal_fraction[chemistry_element_Mg] / (24.0 * XH);
  ChimesGasVars->element_abundances[6] =
      (ChimesFloat)metal_fraction[chemistry_element_Si] / (28.0 * XH);
  ChimesGasVars->element_abundances[9] =
      (ChimesFloat)metal_fraction[chemistry_element_Fe] / (56.0 * XH);

  ChimesGasVars->element_abundances[7] =
      (ChimesFloat)metal_fraction[chemistry_element_Si] *
      cooling->S_over_Si_ratio_in_solar *
      (cooling->S_solar_mass_fraction / cooling->Si_solar_mass_fraction) /
      (32.0 * XH);
  ChimesGasVars->element_abundances[8] =
      (ChimesFloat)metal_fraction[chemistry_element_Si] *
      cooling->Ca_over_Si_ratio_in_solar *
      (cooling->Ca_solar_mass_fraction / cooling->Si_solar_mass_fraction) /
      (40.0 * XH);

  /* When chimes_update_element_abundances() is called on a particle that is
   * being initialised for the first time, the chimes_abundance array has not
   * yet been initialised. We therefore set mu = 1. */
  const double mu = (mode == 0) ? 1.0 : chimes_mu(cooling, hydro_props, p, xp);

  if (cooling->colibre_metal_depletion) {

    const double u_cgs = u_0 * cooling->internal_energy_to_cgs;
    const double T = get_T_from_u_cgs(u_cgs, mu, cooling);

    /* Reference column density */
    const double N_ref = calculate_colibre_N_ref(phys_const, us, cosmo, cooling,
                                                 hydro_props, p, xp, T);

    /* Dust depletion factor */
    const double factor = fmin(pow(N_ref / cooling->N_H0, 1.4), 1.0);

    /* Reduce gas-phase metal abundances due to dust depletion. */
    ChimesGasVars->element_abundances[1] *=
        (ChimesFloat)(1.0 - (cooling->f_dust0_C * factor));
    ChimesGasVars->element_abundances[3] *=
        (ChimesFloat)(1.0 - (cooling->f_dust0_O * factor));
    ChimesGasVars->element_abundances[5] *=
        (ChimesFloat)(1.0 - (cooling->f_dust0_Mg * factor));
    ChimesGasVars->element_abundances[6] *=
        (ChimesFloat)(1.0 - (cooling->f_dust0_Si * factor));
    ChimesGasVars->element_abundances[8] *=
        (ChimesFloat)(1.0 - (cooling->f_dust0_Ca * factor));
    ChimesGasVars->element_abundances[9] *=
        (ChimesFloat)(1.0 - (cooling->f_dust0_Fe * factor));
  }

  /* Zero the abundances of any elements that not included in the network. */
  for (int i = 0; i < 9; i++)
    ChimesGasVars->element_abundances[i + 1] *=
        (ChimesFloat)ChimesGlobalVars->element_included[i];

  if (mode == 1) check_constraint_equations(ChimesGasVars, ChimesGlobalVars);
}

/**
 * @brief Apply the cooling function to a particle.
 *
 * We use the CHIMES module to integrate the cooling rate
 * and chemical abundances over the time-step.
 *
 * @param phys_const The physical constants in internal units.
 * @param us The internal system of units.
 * @param cosmo The current cosmological model.
 * @param hydro_properties the hydro_props struct
 * @param floor_props Properties of the entropy floor.
 * @param pressure_floor Properties of the pressure floor.
 * @param cooling The #cooling_function_data used in the run.
 * @param p Pointer to the particle data.
 * @param xp Pointer to the extended particle data.
 * @param dt The cooling time-step of this particle.
 * @param dt_therm The hydro time-step of this particle.
 * @param time Time since Big Bang
 */
void cooling_cool_part(const struct phys_const *phys_const,
                       const struct unit_system *us,
                       const struct cosmology *cosmo,
                       const struct hydro_props *hydro_properties,
                       const struct entropy_floor_properties *floor_props,
                       const struct pressure_floor_props *pressure_floor,
                       const struct cooling_function_data *cooling,
                       const struct dustevo_props *dp, struct part *p,
                       struct xpart *xp, const float dt, const float dt_therm,
                       const double time) {

  /* No cooling happens over zero time */
  if (dt == 0.) {

    /* But we still set the subgrid properties to a valid state */
    cooling_set_subgrid_properties(phys_const, us, cosmo, hydro_properties,
                                   floor_props, cooling, dp, p, xp);

    return;
  }

  ticks tic_chimes = getticks();

  /* Verify whether we are still in an HII region or not
   * (i.e. has the counter expired?) */
  if ((time > xp->tracers_data.HIIregion_timer_gas) &&
      (xp->tracers_data.HIIregion_timer_gas > 0.)) {

    /* Not anymore an HII region */
    xp->tracers_data.HIIregion_timer_gas = -1.;
    xp->tracers_data.HIIregion_starid = -1;
  }

  /* CHIMES structures. */
  const struct globalVariables *ChimesGlobalVars = &cooling->ChimesGlobalVars;

  /* Allocate memory to arrays within ChimesGasVars. */
  struct gasVariables ChimesGasVars;
  allocate_gas_abundances_memory(&ChimesGasVars, ChimesGlobalVars);

  /* Copy abundances over from xp to ChimesGasVars. */
  for (int i = 0; i < ChimesGlobalVars->totalNumberOfSpecies; i++)
    ChimesGasVars.abundances[i] = xp->cooling_data.chimes_abundances[i];

  /* Get internal energy at the last kick step */
  const float u_start = hydro_get_physical_internal_energy(p, xp, cosmo);

  /* Get the change in internal energy due to hydro forces */
  const float hydro_du_dt = hydro_get_physical_internal_energy_dt(p, cosmo);

  /* Get internal energy at the end of the next kick step (assuming dt does not
   * increase) */
  double u_0 = (u_start + hydro_du_dt * dt_therm);

  /* Prior to computing updated element abundances, do we destroy dust? */
  if (xp->cooling_data.heated_by_FB && cooling->destroy_FB_heated_dust) {
    dust_redistribute_masses_cooling_step(p, dp);
    dust_check_for_NaNs(p);
  }

  /* Update element abundances from metal mass fractions.
   *
   * We need to do this here, and not later on in chimes_update_gas_vars(),
   * because the element abundances need to be set in ChimesGasVars before we
   * can calculate the mean molecular weight. */
  chimes_update_element_abundances(phys_const, us, cosmo, cooling,
                                   hydro_properties, p, xp, &ChimesGasVars, u_0,
                                   /*mode=*/1);

  /* Check for minimal energy limit at the start of the step.
   *
   * Note that, to maintain consistency with the temperature
   * floor that is imposed within CHIMES, we compute the minimal energy from the
   * minimal temperature based on the particle's actual mean molecular weight
   * mu, rather than an assumed constant mu.
   */
  const double mu_start = chimes_mu(cooling, hydro_properties, p, xp);
  const double minimal_internal_energy_start = get_u_internal_from_T(
      hydro_properties->minimal_temperature, mu_start, phys_const);

  u_0 = max(u_0, minimal_internal_energy_start);

  /* Convert to CGS units */
  const double u_0_cgs = u_0 * cooling->internal_energy_to_cgs;
  const double dt_cgs = dt * cooling->time_to_cgs;

  /* Update the ChimesGasVars structure with the
   * particle's thermodynamic variables. */
  chimes_update_gas_vars(u_0_cgs, phys_const, us, cosmo, hydro_properties,
                         floor_props, cooling, dp, p, xp, &ChimesGasVars,
                         dt_cgs);

  /* Store the updated dust boost factor */
  xp->cooling_data.dust_boost_factor = ChimesGasVars.dust_boost_factor;

  /* Special treatment for particles recently heated by feedback */
  if (xp->cooling_data.heated_by_FB && cooling->set_FB_particles_to_eqm) {

    /* Set the particle abundances to equilibrium values */
    cooling_set_FB_particle_chimes_abundances(&ChimesGasVars, cooling);
  }

  if (xp->cooling_data.heated_by_FB &&
      (cooling->set_FB_particles_to_eqm || cooling->destroy_FB_heated_dust)) {
    /* Reset the heated by FB flag, Particles that were heated are
       treated as regular particles from now on.*/
    xp->cooling_data.heated_by_FB = 0;
  }

  /* If particle is an HII region, immediately heat it up to 1e4 K (if it was
   * colder), and force it to be evolved with equilibrium cooling. */
  if (is_HII_region(xp)) {
    ChimesGasVars.temperature = chimes_max(
        ChimesGasVars.temperature, (ChimesFloat)cooling->HIIregion_temp);
    ChimesGasVars.ForceEqOn = 1;
  }

  /* Check that the temperature passed to CHIMES is not unreasonably high. */
  if (ChimesGasVars.temperature > 1.0e12) {
    chimes_print_gas_vars(stderr, &ChimesGasVars, ChimesGlobalVars);
    error(
        "CHIMES ERROR: the temperature that is being passed to CHIMES is %.6e. "
        "This is too high!",
        ChimesGasVars.temperature);
  }

  /* Make a backup of the initial gas vars. */
  struct gasVariables ChimesGasVars_save;
  ChimesGasVars_save = ChimesGasVars;
  allocate_gas_abundances_memory(&ChimesGasVars_save, ChimesGlobalVars);
  memcpy(ChimesGasVars_save.isotropic_photon_density,
         ChimesGasVars.isotropic_photon_density,
         sizeof(ChimesFloat) * ChimesGlobalVars->N_spectra);
  memcpy(ChimesGasVars_save.G0_parameter, ChimesGasVars.G0_parameter,
         sizeof(ChimesFloat) * ChimesGlobalVars->N_spectra);
  memcpy(ChimesGasVars_save.H2_dissocJ, ChimesGasVars.H2_dissocJ,
         sizeof(ChimesFloat) * ChimesGlobalVars->N_spectra);
  memcpy(ChimesGasVars_save.abundances, ChimesGasVars.abundances,
         sizeof(ChimesFloat) * ChimesGlobalVars->totalNumberOfSpecies);

  /* Call CHIMES to integrate the chemistry and cooling over the time-step. */
  const int chimes_return_code =
      chimes_network(&ChimesGasVars, ChimesGlobalVars);

  /* Check whether the abundances became  unphysical before enforcing the
   * constraint equations within CHIMES.
   * If they did, repeat CHIMES with equilibrium cooling. */
  if (chimes_return_code == 1) {

    if ((strcmp(ChimesGlobalVars->EqAbundanceTablePath, "None") == 0) ||
        (strcmp(ChimesGlobalVars->EqAbundanceTablePath, "none") == 0)) {
      /* Special case if no tables were provided... */
      error(
          "CHIMES ERROR: chimes_network() returned an error code. This "
          "suggests we might need to use stricter tolerances in CHIMES.");
    }

    /* Be verbose about the failure! */
    chimes_print_gas_vars(stderr, &ChimesGasVars_save, ChimesGlobalVars);
    chimes_print_gas_vars(stderr, &ChimesGasVars, ChimesGlobalVars);

    fprintf(stderr,
            "CHIMES Warning: chimes_network() returned an error code. We will "
            "repeat the CHIMES integration with equilibrium cooling.");

    /* Copy over original CHIMES abundances and temperature. */
    memcpy(ChimesGasVars.abundances, ChimesGasVars_save.abundances,
           sizeof(ChimesFloat) * ChimesGlobalVars->totalNumberOfSpecies);
    ChimesGasVars.temperature = ChimesGasVars_save.temperature;

    /* Use equilibrium abundances. */
    ChimesGasVars.ForceEqOn = 1;

    /* Integrate CHIMES in equilibrium.
     * (Note we assume this does *not* fail) */
    chimes_network(&ChimesGasVars, ChimesGlobalVars);
  }

  /* Check that the temperature returned is not unreasonably high. */
  if (ChimesGasVars.temperature > 1.0e12) {
    chimes_print_gas_vars(stderr, &ChimesGasVars_save, ChimesGlobalVars);
    chimes_print_gas_vars(stderr, &ChimesGasVars, ChimesGlobalVars);
    error(
        "CHIMES ERROR: the temperature that is being returned from CHIMES is "
        "%.6e. This is too high!",
        ChimesGasVars.temperature);
  }

  /* Check that radiative heating alone has  not increased the temperature to
   * an unphysical level.
   *
   * Radiative heating alone should not increase the temperature
   * above 1e5 K, so above this we should have net cooling from CHIMES.
   * Although changes in mean molecular weight could
   * make the temperature increase by a factor of a few. */
  const float T_rad_heat_max = 1.0e5f;
  const float T_relative_change_limit = 10.0f;
  if ((ChimesGasVars.temperature > T_rad_heat_max) &&
      (ChimesGasVars.temperature / ChimesGasVars_save.temperature >
       T_relative_change_limit)) {
    chimes_print_gas_vars(stderr, &ChimesGasVars_save, ChimesGlobalVars);
    chimes_print_gas_vars(stderr, &ChimesGasVars, ChimesGlobalVars);
    error(
        "CHIMES ERROR: After the CHIMES integration the temperature has "
        "increased from %.6e to %.6e. Radiative heating alone should not "
        "increase the temperature above 1e5 K, so this suggests we might need "
        "to use stricter tolerances in CHIMES.",
        ChimesGasVars_save.temperature, ChimesGasVars.temperature);
  }

  /* If particle is in an HII region, manually set it to be ionised,
   * according to the HIIregion_fion parameter. */
  if (is_HII_region(xp)) {
    cooling_set_HIIregion_chimes_abundances(&ChimesGasVars, cooling);
  }

  /* Calculate net cooling rate at the end of the time-step.
   *
   * If the particle is in an HII region, set the cooling rate to zero.
   * If the particle is on the EOS, the cooling rate will be
   * overwritten when we compute the subgrid properties below. */
  const float rho_phys_cgs =
      hydro_get_physical_density(p, cosmo) * cooling->density_to_cgs;

  if (is_HII_region(xp)) {

    xp->cooling_data.net_cooling_rate = 0.0f;

  } else {

    /* cooling rate in erg/cm^3/s. */
    float net_cooling_rate_cgs = (float)chimes_cooling_rate_external(
        &ChimesGasVars, ChimesGlobalVars, /*mode=*/0);

    /* Convert to erg/g/s */
    net_cooling_rate_cgs /= rho_phys_cgs;

    /* Convert to code units. */
    xp->cooling_data.net_cooling_rate =
        net_cooling_rate_cgs * cooling->internal_energy_over_time_from_cgs;
  }

  /* Recover the Hydrogen mass fraction */
  const double XH = get_hydrogen_mass_fraction(p, hydro_properties);

  /* Copy abundances from ChimesGasVars back to the particle. */
  for (int i = 0; i < ChimesGlobalVars->totalNumberOfSpecies; i++)
    xp->cooling_data.chimes_abundances[i] = ChimesGasVars.abundances[i];

  /* Calculate mean molecular weight at the end of the time step. */
  const float mu_end = chimes_mu(cooling, hydro_properties, p, xp);

  /* Compute the internal energy at the end of the
   * step using the final temperature from CHIMES. */
  const double u_final_cgs =
      get_u_cgs_from_T((double)ChimesGasVars.temperature, mu_end, cooling);

  /* Convert back to internal units */
  double u_final = u_final_cgs * cooling->internal_energy_from_cgs;

  /* We now need to check that we are not going to go below any of the limits */

  /* Recompute new minimal internal energy
   * from the new mean molecular weight. */
  const double minimal_internal_energy_end = get_u_internal_from_T(
      hydro_properties->minimal_temperature, mu_end, phys_const);

  u_final = max(u_final, minimal_internal_energy_end);

  /* Limit imposed by the entropy floor */
  const double A_floor = entropy_floor(p, cosmo, floor_props);
  const double rho = hydro_get_physical_density(p, cosmo);
  const double u_floor = gas_internal_energy_from_entropy(rho, A_floor);

  u_final = max(u_final, u_floor);

  /* Expected change in energy over the next kick step
     (assuming no change in dt) */
  const double delta_u = u_final - max(u_start, u_floor);

  /* Determine if we are in the slow- or rapid-cooling regime,
   * by comparing dt / t_cool to the rapid_cooling_threshold.
   *
   * Note that dt / t_cool = fabs(delta_u) / u_start. */
  const double dt_over_t_cool = fabs(delta_u) / max(u_start, u_floor);

  /* If rapid_cooling_threshold < 0, always use the slow-cooling regime. */
  if ((cooling->rapid_cooling_threshold >= 0.0) &&
      (dt_over_t_cool >= cooling->rapid_cooling_threshold)) {

    /* Rapid-cooling regime. */

    /* Update the particle's u and du/dt */
    hydro_set_physical_internal_energy(p, xp, cosmo, u_final);
    hydro_set_drifted_physical_internal_energy(p, cosmo, pressure_floor,
                                               u_final);
    hydro_set_physical_internal_energy_dt(p, cosmo, 0.);

  } else {

    /* Slow-cooling regime. */

    /* Update du/dt so that we can subsequently drift internal energy. */
    const float cooling_du_dt = delta_u / dt_therm;

    /* Update the internal energy time derivative */
    hydro_set_physical_internal_energy_dt(p, cosmo, cooling_du_dt);
  }

  if (is_HII_region(xp)) {

    /* Check whether final internal energy has reached the HII region
     * temperature. */
    const double mu_HII = get_mu_HII_region(cooling, XH);
    const double u_HII_cgs =
        get_u_cgs_from_T(cooling->HIIregion_temp, mu_HII, cooling);
    const double u_HII = u_HII_cgs * cooling->internal_energy_from_cgs;

    /* The energy has reached, or is below, the HII target energy */
    if (u_final <= u_HII) {

      /* Immediately set the particle to u_HII */
      hydro_set_physical_internal_energy(p, xp, cosmo, u_HII);
      hydro_set_drifted_physical_internal_energy(p, cosmo, pressure_floor,
                                                 u_HII);

      /* Internal energy should stay constant for the timestep */
      hydro_set_physical_internal_energy_dt(p, cosmo, 0.f);
    }
  }

  ticks toc_chimes = getticks();

  if (clocks_from_ticks(toc_chimes - tic_chimes) > 1000.f) {
    warning("Particle %lld took more than 1000ms in CHIMES!", p->id);
    chimes_print_gas_vars(stderr, &ChimesGasVars_save, ChimesGlobalVars);
  }

  /* Store the radiated energy */
  xp->cooling_data.radiated_energy -= hydro_get_mass(p) * (u_final - u_0);

  /* Free CHIMES memory. */
  free_gas_abundances_memory(&ChimesGasVars, ChimesGlobalVars);
  free_gas_abundances_memory(&ChimesGasVars_save, ChimesGlobalVars);

  /* Set subgrid properties. If the particle is below logT_EOS_max, the CHIMES
   * abundances will be re-computed from the subgrid temperature and density. */
  cooling_set_subgrid_properties(phys_const, us, cosmo, hydro_properties,
                                 floor_props, cooling, dp, p, xp);
}

/**
 * @brief Returns the total radiated energy by this particle.
 *
 * @param xp The extended particle data
 */
float cooling_get_radiated_energy(const struct xpart *xp) {

  return xp->cooling_data.radiated_energy;
}

/**
 * @brief Split the cooling content of a particle into n pieces
 *
 * @param p The #part.
 * @param xp The #xpart.
 * @param n The number of pieces to split into.
 */
void cooling_split_part(struct part *p, struct xpart *xp, double n) {

  xp->cooling_data.radiated_energy /= n;
}

/**
 *
 * @brief Calculate mean molecular weight from CHIMES abundances and
 * also deals with special cases (HII regions, feedback).
 *
 * @param cooling The #cooling_function_data used in the run.
 * @param hydro_props The properties of the hydro scheme.
 * @param p Pointer to the particle data.
 * @param xp Pointer to the extended particle data.
 */
double chimes_mu(const struct cooling_function_data *cooling,
                 const struct hydro_props *hydro_props, const struct part *p,
                 const struct xpart *xp) {

  /* Special cases first */

  /* Particles heated by FB and we run with heated particles at equilibrium */
  if (xp->cooling_data.heated_by_FB && cooling->set_FB_particles_to_eqm) {
    return hydro_props->mu_ionised;
  }

  /* Particles in an HII region */
  if (is_HII_region(xp)) {

    /* Recover the Hydrogen mass fraction */
    const double XH = get_hydrogen_mass_fraction(p, hydro_props);

    return get_mu_HII_region(cooling, XH);
  }

  /* Now, normal case: Compute mu from the particle's metal content and CHIMES
   * abundances */
  const struct globalVariables *ChimesGlobalVars = &cooling->ChimesGlobalVars;
  double numerator = 1.0;

  /* Get element mass fractions */
#if defined(CHEMISTRY_COLIBRE) || defined(CHEMISTRY_EAGLE)
  const float *const metal_fraction =
      chemistry_get_metal_mass_fraction_for_cooling(p);
  const float XH = metal_fraction[chemistry_element_H];

  numerator += (double)metal_fraction[chemistry_element_He] / XH;

  if (ChimesGlobalVars->element_included[0] == 1)
    numerator += (double)metal_fraction[chemistry_element_C] / XH;

  if (ChimesGlobalVars->element_included[1] == 1)
    numerator += (double)metal_fraction[chemistry_element_N] / XH;

  if (ChimesGlobalVars->element_included[2] == 1)
    numerator += (double)metal_fraction[chemistry_element_O] / XH;

  if (ChimesGlobalVars->element_included[3] == 1)
    numerator += (double)metal_fraction[chemistry_element_Ne] / XH;

  if (ChimesGlobalVars->element_included[4] == 1)
    numerator += (double)metal_fraction[chemistry_element_Mg] / XH;

  if (ChimesGlobalVars->element_included[5] == 1)
    numerator += (double)metal_fraction[chemistry_element_Si] / XH;

  if (ChimesGlobalVars->element_included[8] == 1)
    numerator += (double)metal_fraction[chemistry_element_Fe] / XH;

  if (ChimesGlobalVars->element_included[6] == 1)
    numerator +=
        (double)metal_fraction[chemistry_element_Si] *
        cooling->S_over_Si_ratio_in_solar *
        (cooling->S_solar_mass_fraction / cooling->Si_solar_mass_fraction) / XH;

  if (ChimesGlobalVars->element_included[7] == 1)
    numerator +=
        (double)metal_fraction[chemistry_element_Si] *
        cooling->Ca_over_Si_ratio_in_solar *
        (cooling->Ca_solar_mass_fraction / cooling->Si_solar_mass_fraction) /
        XH;
#else

  /* Without COLIBRE or EAGLE chemistry, the metal abundances are unavailable.
   * -->Set to primordial abundances. */
  numerator += 0.0833 * 4.0;
#endif /* CHEMISTRY_COLIBRE || CHEMISTRY_EAGLE */

  double denominator = 0.0;
  for (int i = 0; i < ChimesGlobalVars->totalNumberOfSpecies; i++)
    denominator += (double)xp->cooling_data.chimes_abundances[i];

  return numerator / denominator;
}

/**
 * @brief Compute the temperature of a #part based on the cooling function.
 *
 * This uses the mean molecular weight computed from the
 * CHIMES abundance array.
 *
 * @param phys_const #phys_const data structure.
 * @param hydro_props The properties of the hydro scheme.
 * @param us The internal system of units.
 * @param cosmo #cosmology data structure.
 * @param cooling #cooling_function_data struct.
 * @param p #part data.
 * @param xp Pointer to the #xpart data.
 */
float cooling_get_temperature(const struct phys_const *phys_const,
                              const struct hydro_props *hydro_props,
                              const struct unit_system *us,
                              const struct cosmology *cosmo,
                              const struct cooling_function_data *cooling,
                              const struct part *p, const struct xpart *xp) {

  /* Mean molecular weight */
  const double mu = chimes_mu(cooling, hydro_props, p, xp);

  /* Internal energy, in code units. */
  const double u = hydro_get_physical_internal_energy(p, xp, cosmo);

  /* Convert to cgs. */
  const double u_cgs = u * cooling->internal_energy_to_cgs;

  /* Return particle temperature */
  return (float)get_T_from_u_cgs(u_cgs, mu, cooling);
}

/**
 * @brief Compute the temperature of a #part on the entropy floor.
 *
 * This uses the mean molecular weight computed from the
 * CHIMES equilibrium abundance tables, not the particle's
 * actual abundance array. This is useful for particles
 * on the entropy floor.
 *
 * @param phys_const #phys_const data structure.
 * @param hydro_props The properties of the hydro scheme.
 * @param us The internal system of units.
 * @param cosmo #cosmology data structure.
 * @param cooling #cooling_function_data struct.
 * @param p #part data.
 * @param xp Pointer to the #xpart data.
 */
float cooling_get_temperature_on_entropy_floor(
    const struct phys_const *phys_const, const struct hydro_props *hydro_props,
    const struct unit_system *us, const struct cosmology *cosmo,
    const struct cooling_function_data *cooling, const struct part *p,
    const struct xpart *xp) {

  /* Internal energy, in code units. */
  const double u = hydro_get_physical_internal_energy(p, xp, cosmo);

  /* Convert to cgs. */
  const double u_cgs = u * cooling->internal_energy_to_cgs;

  /* Recover the Hydrogen mass fraction */
  const double XH = get_hydrogen_mass_fraction(p, hydro_props);

  /* Helium abundance. */
#if defined(CHEMISTRY_COLIBRE) || defined(CHEMISTRY_EAGLE)
  const float *const metal_fraction =
      chemistry_get_metal_mass_fraction_for_cooling(p);

  float metallicity = chemistry_get_total_metal_mass_fraction_for_cooling(p);
  if (metallicity <= 0.0) metallicity = FLT_MIN;

  const ChimesFloat x_He =
      (ChimesFloat)metal_fraction[chemistry_element_He] / (4.0 * XH);
#else

  /* Without COLIBRE or EAGLE chemistry, the metal abundances are unavailable.
   * --> Set to primordial abundances. */
  const float metallicity = FLT_MIN;
  const ChimesFloat x_He = 0.0833;
#endif

  /* Interpolate equilibrium table to get H2 and electron fractions. */
  const ChimesFloat nH = (ChimesFloat)((hydro_get_physical_density(p, cosmo) *
                                        XH / phys_const->const_proton_mass) *
                                       cooling->number_density_to_cgs);
  const ChimesFloat Z_over_Zsol = (ChimesFloat)metallicity / cooling->Zsol;

  /* Initial guess for T, assuming mu = 1 */
  const ChimesFloat T = get_T_from_u_cgs(u_cgs, /*mu=*/1.0, cooling);

  /* Take log's, but protect against log(0) */
  const ChimesFloat log_nH = log10(chimes_max(nH, FLT_MIN));
  const ChimesFloat log_Z_over_Zsol = log10(chimes_max(Z_over_Zsol, FLT_MIN));
  const ChimesFloat log_T = log10(chimes_max(T, FLT_MIN));
  const ChimesFloat log_u = log10(chimes_max(u_cgs, FLT_MIN));

  int T_index, nH_index, Z_index;
  ChimesFloat dT, dnH, dZ;
  const int N_T = chimes_table_eqm_abundances.N_Temperatures;
  const int N_nH = chimes_table_eqm_abundances.N_Densities;
  const int N_Z = chimes_table_eqm_abundances.N_Metallicities;
  chimes_get_table_index(chimes_table_eqm_abundances.Temperatures, N_T, log_T,
                         &T_index, &dT);
  chimes_get_table_index(chimes_table_eqm_abundances.Densities, N_nH, log_nH,
                         &nH_index, &dnH);
  chimes_get_table_index(chimes_table_eqm_abundances.Metallicities, N_Z,
                         log_Z_over_Zsol, &Z_index, &dZ);

  /* Position of electrons and H2 in the CHIMES abundance arrays. */
  const int elec_ind = cooling->ChimesGlobalVars.speciesIndices[sp_elec];
  const int H2_ind = cooling->ChimesGlobalVars.speciesIndices[sp_H2];

  /* Find the temperature bins in the eqm tables that bracket u_cgs */
  dT = 0.0;

  ChimesFloat mu, x_H2, x_elec, u_low, u_hi, T_low, T_hi;

  x_elec = exp10(chimes_interpol_4d_fix_x(
      chimes_table_eqm_abundances.Abundances, elec_ind, T_index, nH_index,
      Z_index, dT, dnH, dZ, N_T, N_nH, N_Z));
  x_H2 = exp10(chimes_interpol_4d_fix_x(chimes_table_eqm_abundances.Abundances,
                                        H2_ind, T_index, nH_index, Z_index, dT,
                                        dnH, dZ, N_T, N_nH, N_Z));
  mu = (1.0 + (4.0 * x_He)) / (1.0 + x_He + x_elec - x_H2);
  T_low = exp10(chimes_table_eqm_abundances.Temperatures[T_index]);
  u_low = get_u_cgs_from_T(T_low, mu, cooling);

  x_elec = exp10(chimes_interpol_4d_fix_x(
      chimes_table_eqm_abundances.Abundances, elec_ind, T_index + 1, nH_index,
      Z_index, dT, dnH, dZ, N_T, N_nH, N_Z));
  x_H2 = exp10(chimes_interpol_4d_fix_x(chimes_table_eqm_abundances.Abundances,
                                        H2_ind, T_index + 1, nH_index, Z_index,
                                        dT, dnH, dZ, N_T, N_nH, N_Z));
  mu = (1.0 + (4.0 * x_He)) / (1.0 + x_He + x_elec - x_H2);
  T_hi = exp10(chimes_table_eqm_abundances.Temperatures[T_index + 1]);
  u_hi = get_u_cgs_from_T(T_hi, mu, cooling);

  if (u_cgs < u_low) {

    while (u_cgs < u_low) {
      T_index -= 1;
      if (T_index < 0)
        error(
            "CHIMES failed to convert internal energy to temperature for "
            "snapshot, as T_index < 0.");

      x_elec = exp10(chimes_interpol_4d_fix_x(
          chimes_table_eqm_abundances.Abundances, elec_ind, T_index, nH_index,
          Z_index, dT, dnH, dZ, N_T, N_nH, N_Z));
      x_H2 = exp10(chimes_interpol_4d_fix_x(
          chimes_table_eqm_abundances.Abundances, H2_ind, T_index, nH_index,
          Z_index, dT, dnH, dZ, N_T, N_nH, N_Z));
      mu = (1.0 + (4.0 * x_He)) / (1.0 + x_He + x_elec - x_H2);
      T_low = exp10(chimes_table_eqm_abundances.Temperatures[T_index]);
      u_low = get_u_cgs_from_T(T_low, mu, cooling);

      x_elec = exp10(chimes_interpol_4d_fix_x(
          chimes_table_eqm_abundances.Abundances, elec_ind, T_index + 1,
          nH_index, Z_index, dT, dnH, dZ, N_T, N_nH, N_Z));
      x_H2 = exp10(chimes_interpol_4d_fix_x(
          chimes_table_eqm_abundances.Abundances, H2_ind, T_index + 1, nH_index,
          Z_index, dT, dnH, dZ, N_T, N_nH, N_Z));
      mu = (1.0 + (4.0 * x_He)) / (1.0 + x_He + x_elec - x_H2);
      T_hi = exp10(chimes_table_eqm_abundances.Temperatures[T_index + 1]);
      u_hi = get_u_cgs_from_T(T_hi, mu, cooling);
    }
  } else if (u_cgs > u_hi) {

    while (u_cgs > u_hi) {
      T_index += 1;

      if (T_index > N_T - 2)
        error(
            "CHIMES failed to convert internal energy to temperature for "
            "snapshot, as T_index > N_T - 2.");

      x_elec = exp10(chimes_interpol_4d_fix_x(
          chimes_table_eqm_abundances.Abundances, elec_ind, T_index, nH_index,
          Z_index, dT, dnH, dZ, N_T, N_nH, N_Z));
      x_H2 = exp10(chimes_interpol_4d_fix_x(
          chimes_table_eqm_abundances.Abundances, H2_ind, T_index, nH_index,
          Z_index, dT, dnH, dZ, N_T, N_nH, N_Z));
      mu = (1.0 + (4.0 * x_He)) / (1.0 + x_He + x_elec - x_H2);
      T_low = exp10(chimes_table_eqm_abundances.Temperatures[T_index]);
      u_low = get_u_cgs_from_T(T_low, mu, cooling);

      x_elec = exp10(chimes_interpol_4d_fix_x(
          chimes_table_eqm_abundances.Abundances, elec_ind, T_index + 1,
          nH_index, Z_index, dT, dnH, dZ, N_T, N_nH, N_Z));
      x_H2 = exp10(chimes_interpol_4d_fix_x(
          chimes_table_eqm_abundances.Abundances, H2_ind, T_index + 1, nH_index,
          Z_index, dT, dnH, dZ, N_T, N_nH, N_Z));
      mu = (1.0 + (4.0 * x_He)) / (1.0 + x_He + x_elec - x_H2);
      T_hi = exp10(chimes_table_eqm_abundances.Temperatures[T_index + 1]);
      u_hi = get_u_cgs_from_T(T_hi, mu, cooling);
    }
  }

  /* Finally, interpolate mu between
   * u_low and u_hi. */
  if (!((u_cgs >= u_low) && (u_cgs <= u_hi)))
    error(
        "CHIMES failed to convert internal energy to temperature for snapshot, "
        "as u_low and u_hi do not bracket u_cgs.");

  dT = (log_u - log10(u_low)) / (log10(u_hi) - log10(u_low));

  x_elec = exp10(chimes_interpol_4d_fix_x(
      chimes_table_eqm_abundances.Abundances, elec_ind, T_index, nH_index,
      Z_index, dT, dnH, dZ, N_T, N_nH, N_Z));
  x_H2 = exp10(chimes_interpol_4d_fix_x(chimes_table_eqm_abundances.Abundances,
                                        H2_ind, T_index, nH_index, Z_index, dT,
                                        dnH, dZ, N_T, N_nH, N_Z));

  mu = (1.0 + (4.0 * x_He)) / (1.0 + x_He + x_elec - x_H2);

  /* Return particle temperature */
  return (float)get_T_from_u_cgs(u_cgs, mu, cooling);
}

/**
 * @brief Compute the temperature of a #part for the snapshot.
 *
 * This function returns the temperature that will be written
 * out to the snapshot. If it is above the entropy floor,
 * it uses the mean molecular weight calculated from the
 * particle's abundance array. On the entropy floor, it
 * uses the equilibrium abundance tables.
 *
 * @param phys_const #phys_const data structure.
 * @param hydro_props The properties of the hydro scheme.
 * @param us The internal system of units.
 * @param cosmo #cosmology data structure.
 * @param cooling #cooling_function_data struct.
 * @param p #part data.
 * @param xp Pointer to the #xpart data.
 */
float cooling_get_snapshot_temperature(
    const struct phys_const *phys_const, const struct hydro_props *hydro_props,
    const struct entropy_floor_properties *floor_props,
    const struct unit_system *us, const struct cosmology *cosmo,
    const struct cooling_function_data *cooling, const struct part *p,
    const struct xpart *xp) {

  /* Limit imposed by the entropy floor */
  const double A_EOS = entropy_floor(p, cosmo, floor_props);
  const double rho = hydro_get_physical_density(p, cosmo);
  const double u_EOS = gas_internal_energy_from_entropy(rho, A_EOS);
  const double u_EOS_max = u_EOS * exp10(cooling->dlogT_EOS);

  /* Particle's internal energy */
  const double u = hydro_get_physical_internal_energy(p, xp, cosmo);

  if ((u < u_EOS_max) && !is_HII_region(xp)) {

    /* If the particle is within dlogT_EOS of the entropy floor, and is not an
     * HII region, its CHIMES abundance array  would have been set according to
     * the subgrid properties.
     * We therefore  need to calculate the temperature based on the equilibrium
     * abundance tables  at the particle's actual temperature, not subgrid.
     */
    return cooling_get_temperature_on_entropy_floor(phys_const, hydro_props, us,
                                                    cosmo, cooling, p, xp);

  } else {

    /* Above u_EOS_max, we use the particle's actual abundance array; */
    return cooling_get_temperature(phys_const, hydro_props, us, cosmo, cooling,
                                   p, xp);
  }
}

/**
 * @brief Compute the electron number density of a #part based on the cooling
 * function.
 *
 * The electron density returned is in physical internal units.
 *
 * @param phys_const #phys_const data structure.
 * @param hydro_props The properties of the hydro scheme.
 * @param floor_props The properties of the entropy floor.
 * @param us The internal system of units.
 * @param cosmo #cosmology data structure.
 * @param cooling #cooling_function_data struct.
 * @param dust_props The properties of the dust model.
 * @param p #part data.
 * @param xp Pointer to the #xpart data.
 */
double cooling_get_electron_density(
    const struct phys_const *phys_const, const struct hydro_props *hydro_props,
    const struct entropy_floor_properties *floor_props,
    const struct unit_system *us, const struct cosmology *cosmo,
    const struct cooling_function_data *cooling,
    const struct dustevo_props *dust_props, const struct part *p,
    const struct xpart *xp) {

  const struct colibre_cooling_tables *table = &(cooling->colibre_table);
  const struct globalVariables *ChimesGlobalVars = &(cooling->ChimesGlobalVars);

  /* Create ChimesGasVars struct */
  struct gasVariables ChimesGasVars;
  allocate_gas_abundances_memory(&ChimesGasVars, ChimesGlobalVars);

  /* Copy abundances over from xp to ChimesGasVars */
  for (int i = 0; i < ChimesGlobalVars->totalNumberOfSpecies; i++)
    ChimesGasVars.abundances[i] = xp->cooling_data.chimes_abundances[i];

  /* Particle's internal energy */
  const double u = hydro_get_physical_internal_energy(p, xp, cosmo);
  const double u_cgs = u * cooling->internal_energy_to_cgs;

  /* Update element abundances */
  chimes_update_element_abundances(phys_const, us, cosmo, cooling, hydro_props,
                                   p, xp, &ChimesGasVars, u,
                                   /*mode=*/1);

  /* Special treatment for particles recently heated by feedback */
  if (xp->cooling_data.heated_by_FB && cooling->set_FB_particles_to_eqm) {

    /* Update the ChimesGasVars with the particle's thermo variables. */
    chimes_update_gas_vars(u_cgs, phys_const, us, cosmo, hydro_props,
                           floor_props, cooling, dust_props, p, xp,
                           &ChimesGasVars, 1.0);

    /* Set the particle abundances to equilibrium values */
    cooling_set_FB_particle_chimes_abundances(&ChimesGasVars, cooling);
  }
  /* Special treatment for particles that are part of HII regions */
  else if (is_HII_region(xp)) {

    /* Set HII region abundance array */
    cooling_set_HIIregion_chimes_abundances(&ChimesGasVars, cooling);
  }

  /* Recover the Hydrogen mass fraction */
  const double XH = get_hydrogen_mass_fraction(p, hydro_props);

  /* Calculate hydrogen number density nH
   * in internal units. */
  const double n_H =
      hydro_get_physical_density(p, cosmo) * XH / phys_const->const_proton_mass;
  const double n_H_cgs = (n_H * cooling->number_density_to_cgs);
  const double T_cgs = cooling_get_temperature(phys_const, hydro_props, us,
                                               cosmo, cooling, p, xp) *
                       units_cgs_conversion_factor(us, UNIT_CONV_TEMPERATURE);

  /* Calculate electron density in internal units
   * using CHIMES abundance array. */
  const int elec_ind = ChimesGlobalVars->speciesIndices[sp_elec];
  const double n_e_non_eq = n_H * ChimesGasVars.abundances[elec_ind];

  /* Free CHIMES memory. */
  free_gas_abundances_memory(&ChimesGasVars, ChimesGlobalVars);

  double n_e_metal;
  if (ChimesGlobalVars->hybrid_cooling_mode == 1) {
    /* When we use hybrid cooling, we need to add
     * the contribution of metals that are not included
     * in the non-eq network, which we read in from
     * the equilibrium cooling tables. */

    /* Get the total metallicity in units of solar */
    float abundance_ratio[colibre_cooling_N_elementtypes];
    const float logZZsol =
        abundance_ratio_to_solar(p, &(cooling->colibre_table), abundance_ratio);

    // Set weights for electron densities.
    float weights_electron[colibre_cooling_N_electrontypes];
    for (int i = 0; i < colibre_cooling_N_electrontypes; i++) {
      if (i < colibre_cooling_N_elementtypes - 1)
        weights_electron[i] = abundance_ratio[i];
      else
        weights_electron[i] = 1.f;
    }

    /* Get indices of T, nH, and metallicity*/
    int T_index, n_H_index, met_index, red_index;
    float d_T, d_n_H, d_met;

    cooling_get_index_1d(table->Temp, colibre_cooling_N_temperature,
                         log10(T_cgs), &T_index, &d_T);
    cooling_get_index_1d(table->Metallicity, colibre_cooling_N_metallicity,
                         logZZsol, &met_index, &d_met);
    cooling_get_index_1d(table->nH, colibre_cooling_N_density, log10(n_H_cgs),
                         &n_H_index, &d_n_H);

    if (table->dz > 0.5)
      red_index = 1;
    else
      red_index = 0;

    if (met_index == 0) {
      /* Primordial abundances; metal
       * contribution is zero. */
      n_e_metal = 0.0;
    } else {
      /* From metals, with actual metal ratios.
       * We also need to exclude any metals
       * that are already included in CHIMES. */
      for (int i = element_C; i < element_OA; i++) {
        if (ChimesGlobalVars->element_included[i - 2] == 1)
          weights_electron[i] = 0.f;
      }

      n_e_metal =
          n_H * interpolation3d_fix_x_plus_summation(
                    table->Telectron_fraction, weights_electron, element_C,
                    colibre_cooling_N_electrontypes - 4, red_index, T_index,
                    met_index, n_H_index, d_T, d_met, d_n_H,
                    colibre_cooling_N_loaded_redshifts,
                    colibre_cooling_N_temperature,
                    colibre_cooling_N_metallicity, colibre_cooling_N_density,
                    colibre_cooling_N_electrontypes);
    }
  } else {
    /* If we are not using hybrid cooling, we
     * only include the electrons that are in
     * the non-eq network. */
    n_e_metal = 0.0;
  }

  return n_e_non_eq + n_e_metal;
}

/**
 * @brief Compute the electron pressure of a #part based on the cooling
 * function.
 *
 * Returns the total electron pressure of the particle, P_e * V, in code units.
 *
 * @param phys_const #phys_const data structure.
 * @param hydro_props The properties of the hydro scheme.
 * @param floor_props The properties of the entropy floor.
 * @param us The internal system of units.
 * @param cosmo #cosmology data structure.
 * @param cooling #cooling_function_data struct.
 * @param dust_props The properties of the dust model.
 * @param p #part data.
 * @param xp Pointer to the #xpart data.
 */
double cooling_get_electron_pressure(
    const struct phys_const *phys_const, const struct hydro_props *hydro_props,
    const struct entropy_floor_properties *floor_props,
    const struct unit_system *us, const struct cosmology *cosmo,
    const struct cooling_function_data *cooling,
    const struct dustevo_props *dust_props, const struct part *p,
    const struct xpart *xp) {

  /* Get internal energy in physical frame */
  const float u_phys = hydro_get_physical_internal_energy(p, xp, cosmo);

  /* Calculate temperature in internal units. */
  const double mu = chimes_mu(cooling, hydro_props, p, xp);
  const double temperature = get_T_from_u_internal(u_phys, mu, phys_const);

  /* Calculate total electron density in internal units using
   * CHIMES abundance array and elements evolved in equilibrium. */
  const double n_e =
      cooling_get_electron_density(phys_const, hydro_props, floor_props, us,
                                   cosmo, cooling, dust_props, p, xp);

  return n_e * phys_const->const_boltzmann_k * temperature;
}

/**
 * @brief Compute the y-Compton contribution of a #part based on the cooling
 * function.
 *
 * This is eq. (2) of McCarthy et al. (2017).
 *
 * @param phys_const #phys_const data structure.
 * @param hydro_props The properties of the hydro scheme.
 * @param floor_props The properties of the entropy floor.
 * @param us The internal system of units.
 * @param cosmo #cosmology data structure.
 * @param cooling #cooling_function_data struct.
 * @param dust_props The properties of the dust model.
 * @param p #part data.
 * @param xp Pointer to the #xpart data.
 */
double cooling_get_ycompton(const struct phys_const *phys_const,
                            const struct hydro_props *hydro_props,
                            const struct entropy_floor_properties *floor_props,
                            const struct unit_system *us,
                            const struct cosmology *cosmo,
                            const struct cooling_function_data *cooling,
                            const struct dustevo_props *dust_props,
                            const struct part *p, const struct xpart *xp) {

  /* Get quantities in physical frame */
  const float u_phys = hydro_get_physical_internal_energy(p, xp, cosmo);
  const float rho_phys = hydro_get_physical_density(p, cosmo);
  const double m = hydro_get_mass(p);

  /* Calculate temperature in internal units. */
  const double mu = chimes_mu(cooling, hydro_props, p, xp);
  const double temperature = get_T_from_u_internal(u_phys, mu, phys_const);

  /* Calculate total electron density in internal units using
   * CHIMES abundance array and elements evolved in equilibrium. */
  const double n_e =
      cooling_get_electron_density(phys_const, hydro_props, floor_props, us,
                                   cosmo, cooling, dust_props, p, xp);

  return cooling->y_compton_factor * temperature * m * n_e / rho_phys;
}

/**
 * @brief Calculate the N_ref column density.
 *
 * This routine returns the column density N_ref
 * as defined in Ploeckinger et al. (in prep),
 * which is used to scale the ISRF, cosmic rays,
 * dust depletion and shielding column density
 * in COLIBRE.
 *
 * @param phys_const #phys_const data structure.
 * @param us The internal system of units.
 * @param cosmo #cosmology data structure.
 * @param cooling #cooling_function_data struct.
 * @param p #part data.
 * @param xp Pointer to the #xpart data.
 * @param temperature The gas temperature.
 * @param mu Mean molecular weight.
 */
double calculate_colibre_N_ref(const struct phys_const *phys_const,
                               const struct unit_system *us,
                               const struct cosmology *cosmo,
                               const struct cooling_function_data *cooling,
                               const struct hydro_props *hydro_props,
                               const struct part *p, const struct xpart *xp,
                               const double temperature) {

  /* Based on Ploeckinger & Schaye (2020). */
  const float N_min = cooling->nref_column_density_min_cgs;
  const float N_max = cooling->nref_column_density_max_cgs;
  const float T_min = cooling->nref_trans_temperature_min_K;
  const float T_max = cooling->nref_trans_temperature_max_K;
  const double l_max_cgs = cooling->nref_max_length_cgs;
  const float k = cooling->nref_trans_steepness_k;

  const double XH_const = cooling->nref_XH_const;
  const double mu_const = cooling->nref_mu_const;

  /* Recover the Hydrogen mass fraction */
  const double XH = get_hydrogen_mass_fraction(p, hydro_props);

  /* Density */
  const double nH =
      hydro_get_physical_density(p, cosmo) * XH / phys_const->const_proton_mass;
  const double nH_cgs = nH * cooling->number_density_to_cgs;

  /* Jeans column density */
  const double N_J_thermal =
      nH_cgs * sqrt(hydro_gamma * cooling->boltzmann_k_cgs * temperature /
                    (mu_const * cooling->newton_G_cgs *
                     (cooling->proton_mass_cgs * nH_cgs / XH_const) *
                     cooling->proton_mass_cgs));

  double N_J = N_J_thermal;

  /* Special case where we use the turbulent Jeans scale */
  if (cooling->colibre_use_turbulent_jeans_length) {

    const double N_J_turb =
        sqrt(hydro_gamma * XH_const * nH_cgs /
             (cooling->proton_mass_cgs * cooling->newton_G_cgs)) *
        (cooling->sigma_turb_km_p_s * 1.0e5);

    N_J = chimes_max(N_J_thermal, N_J_turb);
  }

  double N_ref_prime = chimes_min(N_J, N_max);

  if (l_max_cgs > 0.0)
    N_ref_prime = chimes_min(N_ref_prime, l_max_cgs * nH_cgs);

  /* Apply the sigmoid */
  double N_ref = exp10(log10(N_ref_prime) -
                       (log10(N_ref_prime) - log10(N_min)) /
                           (1.0 + pow(sqrt(T_min * T_max) / temperature, k)));
  return N_ref;
}

/**
 * @brief Write a cooling struct to the given FILE as a stream of bytes.
 *
 * @param cooling the struct
 * @param stream the file stream
 */
void cooling_struct_dump(const struct cooling_function_data *cooling,
                         struct dustevo_props *dp, FILE *stream) {

  /* Zero all pointers in the colibre_table within
   * the cooling_function_data struct. */
  struct cooling_function_data cooling_copy = *cooling;
  cooling_copy.colibre_table.Tcooling = NULL;
  cooling_copy.colibre_table.Tcooling_reduced = NULL;
  cooling_copy.colibre_table.Theating = NULL;
  cooling_copy.colibre_table.Theating_reduced = NULL;
  cooling_copy.colibre_table.Telectron_fraction = NULL;
  cooling_copy.colibre_table.Redshifts = NULL;
  cooling_copy.colibre_table.nH = NULL;
  cooling_copy.colibre_table.Temp = NULL;
  cooling_copy.colibre_table.Metallicity = NULL;
  cooling_copy.colibre_table.LogAbundances = NULL;
  cooling_copy.colibre_table.Abundances = NULL;
  cooling_copy.colibre_table.Abundances_inv = NULL;
  cooling_copy.colibre_table.atomicmass = NULL;
  cooling_copy.colibre_table.atomicmass_inv = NULL;
  cooling_copy.colibre_table.LogMassFractions = NULL;
  cooling_copy.colibre_table.MassFractions = NULL;
  cooling_copy.colibre_table.Zsol = NULL;
  cooling_copy.colibre_table.Zsol_inv = NULL;
  cooling_copy.ChimesGlobalVars.hybrid_data = NULL;

  restart_write_blocks((void *)&cooling_copy,
                       sizeof(struct cooling_function_data), 1, stream,
                       "cooling", "cooling function");
}

/**
 * @brief Restore a hydro_props struct from the given FILE as a stream of
 * bytes.
 *
 * @param cooling the struct
 * @param stream the file stream
 * @param cosmo #cosmology structure
 */
void cooling_struct_restore(struct cooling_function_data *cooling,
                            struct dustevo_props *dp, FILE *stream,
                            const struct cosmology *cosmo) {

  restart_read_blocks((void *)cooling, sizeof(struct cooling_function_data), 1,
                      stream, NULL, "cooling function");

  /* Initialise the CHIMES module. */
  if (engine_rank == 0) message("Initialising CHIMES cooling module.");
  init_chimes(&cooling->ChimesGlobalVars);

  if ((cooling->ChimesGlobalVars.hybrid_cooling_mode == 1) ||
      (cooling->use_colibre_subgrid_EOS == 1)) {

    /* Read the Colibre table. */
    if (engine_rank == 0) message("Reading Colibre cooling table.");

    get_cooling_redshifts(&(cooling->colibre_table));

    char fname[colibre_table_path_name_length + colibre_table_filebase_length +
               12];
    sprintf(fname, "%s/%s_z0.00.hdf5",
            cooling->colibre_table.cooling_table_path,
            cooling->colibre_table.cooling_table_filebase);

    read_cooling_header(fname, &(cooling->colibre_table));

    /* Allocate memory for the tables */
    allocate_cooling_tables(&(cooling->colibre_table));

    /* Force a re-read of the cooling tables */
    cooling->colibre_table.z_index = -10;
    cooling->colibre_table.previous_z_index =
        cooling->colibre_table.number_of_redshifts - 2;
    cooling_update(cosmo, /*pfloor=*/NULL, cooling, /*space=*/NULL, dp,
                   /*time=*/0);
  }

  /* Check that the size of the CHIMES network
   * determined by the various Include<Element>
   * parameters matches the size specified
   * through the configure script. */
  if (cooling->ChimesGlobalVars.totalNumberOfSpecies != CHIMES_NETWORK_SIZE)
    error(
        "The size of the CHIMES network based on the Include<Element> "
        "parameters is %d. However, Swift was configured and compiled using "
        "--with-chimes-network-size=%d. If you are sure that you have "
        "correctly set which metals to include in the network in the parameter "
        "file, then you will need to re-compile Swift using ./configure "
        "--with-cooling=CHIMES --with-chimes-network-size=%d.",
        cooling->ChimesGlobalVars.totalNumberOfSpecies, CHIMES_NETWORK_SIZE,
        cooling->ChimesGlobalVars.totalNumberOfSpecies);

  if (cooling->ChimesGlobalVars.hybrid_cooling_mode == 1) {
    /* Create data structure for hybrid cooling,
     * and store pointer to the Colibre table. */
    cooling->ChimesGlobalVars.hybrid_data =
        (void *)malloc(sizeof(struct global_hybrid_data_struct));
    struct global_hybrid_data_struct *myData;
    myData = (struct global_hybrid_data_struct *)
                 cooling->ChimesGlobalVars.hybrid_data;
    myData->table = &(cooling->colibre_table);
    myData->Zsol = cooling->Zsol;

    /* Set the hybrid cooling function pointers. */
    cooling->ChimesGlobalVars.hybrid_cooling_fn =
        &colibre_metal_cooling_rate_temperature;
    cooling->ChimesGlobalVars.allocate_gas_hybrid_data_fn =
        &chimes_allocate_gas_hybrid_data;
    cooling->ChimesGlobalVars.free_gas_hybrid_data_fn =
        &chimes_free_gas_hybrid_data;
  }
}

/**
 * @brief Converts cooling quantities of a particle at the start of a run
 *
 * This function is called once at the end of the engine_init_particle()
 * routine (at the start of a calculation) after the densities of
 * particles have been computed.
 *
 * For this cooling module, this routine is used to set the cooling
 * properties of the (x-)particles to a valid start state, in particular
 * the CHIMES abundance array.
 *
 * This is controlled by the cooling->init_abundance_mode as follows:
 * 0 -- Set each element to one ionisation state, determined by the
 *      ChimesGlobalVars.InitIonState parameter.
 * 1 -- Read abundances from eqm abundance tables.
 * 2 -- Compute initial equilibrium abundances.
 *
 * @param p The particle to act upon
 * @param xp The extended particle to act upon
 * @param cosmo The cosmological model.
 * @param hydro_props The constants used in the scheme.
 * @param phys_const #phys_const data structure.
 * @param us Internal system of units data structure.
 * @param floor_props Properties of the entropy floor.
 * @param cooling #cooling_function_data data structure.
 */
void cooling_convert_quantities(
    struct part *p, struct xpart *xp, const struct cosmology *cosmo,
    const struct hydro_props *hydro_props, const struct phys_const *phys_const,
    const struct unit_system *us,
    const struct entropy_floor_properties *floor_props,
    const struct cooling_function_data *cooling,
    const struct dustevo_props *dp) {

  const struct globalVariables *ChimesGlobalVars = &cooling->ChimesGlobalVars;

  /* Allocate memory to arrays within ChimesGasVars. */
  struct gasVariables ChimesGasVars;
  allocate_gas_abundances_memory(&ChimesGasVars, ChimesGlobalVars);

  /* Set element abundances from
   * metal mass fractions. */
  const double u_0 = hydro_get_physical_internal_energy(p, xp, cosmo);
  chimes_update_element_abundances(phys_const, us, cosmo, cooling, hydro_props,
                                   p, xp, &ChimesGasVars, u_0, /*mode=*/0);

  /* Zero particle's radiated energy. */
  xp->cooling_data.radiated_energy = 0.f;

  /* Set initial values for CHIMES abundance array. */
  ChimesGasVars.InitIonState = cooling->InitIonState;
  initialise_gas_abundances(&ChimesGasVars, ChimesGlobalVars);

  if (cooling->init_abundance_choice == init_abundance_single) {

    /* Copy abundances over to xp. */
    for (int i = 0; i < ChimesGlobalVars->totalNumberOfSpecies; i++)
      xp->cooling_data.chimes_abundances[i] = ChimesGasVars.abundances[i];

  } else if ((cooling->init_abundance_choice == init_abundance_read) ||
             (cooling->init_abundance_choice == init_abundance_compute)) {

    /* Copy initial abundances over to xp. */
    for (int i = 0; i < ChimesGlobalVars->totalNumberOfSpecies; i++)
      xp->cooling_data.chimes_abundances[i] = ChimesGasVars.abundances[i];

    /* Convert internal energy to cgs */
    const double u_0_cgs = u_0 * cooling->internal_energy_to_cgs;

    /* If computing the eqm (init_abundance_compute), we will integrate the
     * chemistry ten times for 1 Gyr per iteration. Multiple iterations are
     * required so that the shielding column densities can be updated between
     * each iteration. If reading from tables, we only need 1 iteration. */
    int n_iterations;
    if (cooling->init_abundance_choice == init_abundance_read)
      n_iterations = 1;
    else /* cooling->init_abundance_choice == init_abundance_compute */
      n_iterations = 10;

    /* 1 Gyr in cgs */
    const double dt_cgs = 1e9 * 365.25 * 24. * 3600.;

    for (int i = 0; i < n_iterations; i++) {
      /* Update element abundances. This
       * accounts for the dust depletion
       * factors. */
      chimes_update_element_abundances(phys_const, us, cosmo, cooling,
                                       hydro_props, p, xp, &ChimesGasVars, u_0,
                                       /*mode=*/1);

      /* Update ChimesGasVars with the particle's
       * thermodynamic variables. */
      chimes_update_gas_vars(u_0_cgs, phys_const, us, cosmo, hydro_props,
                             floor_props, cooling, dp, p, xp, &ChimesGasVars,
                             dt_cgs);

      /* Set temperature evolution off, so that we
       * compute equilibrium at fixed temperature. */
      ChimesGasVars.ThermEvolOn = 0;

      /* Determine whether to use eqm tables or compute the eqm state. */
      if (cooling->init_abundance_choice == init_abundance_read)
        ChimesGasVars.ForceEqOn = 1;
      else
        ChimesGasVars.ForceEqOn = 0;

      /* Integrate to chemical equilibrium. */
      const int chimes_return_code =
          chimes_network(&ChimesGasVars, ChimesGlobalVars);

      /* Check whether the abundances became
       * unphysical before enforcing the
       * constraint equations within CHIMES. */
      if (chimes_return_code == 1) {
        chimes_print_gas_vars(stderr, &ChimesGasVars, ChimesGlobalVars);
        error(
            "CHIMES ERROR: When computing the initial equilibrium abundances, "
            "chimes_network() returned an error code. This suggests we might "
            "need to use stricter tolerances in CHIMES.");
      }
    }

    /* Calculate net cooling rate. */
    const float rho_phys_cgs =
        hydro_get_physical_density(p, cosmo) * cooling->density_to_cgs;

    /* cooling rate in erg/cm^3/s. */
    float net_cooling_rate_cgs = (float)chimes_cooling_rate_external(
        &ChimesGasVars, ChimesGlobalVars, 0);

    /* Convert to erg/g/s */
    net_cooling_rate_cgs /= rho_phys_cgs;

    /* Convert to code units. */
    xp->cooling_data.net_cooling_rate =
        net_cooling_rate_cgs * cooling->internal_energy_over_time_from_cgs;

    /* Copy final abundances over to xp. */
    for (int i = 0; i < ChimesGlobalVars->totalNumberOfSpecies; i++)
      xp->cooling_data.chimes_abundances[i] = ChimesGasVars.abundances[i];
  }

  /* Free CHIMES memory. */
  free_gas_abundances_memory(&ChimesGasVars, ChimesGlobalVars);
}

/**
 * @brief Set the subgrid properties of the gas particle
 *
 * @param phys_const The physical constants in internal units.
 * @param us The internal system of units.
 * @param cosmo The current cosmological model.
 * @param hydro_props the hydro_props struct
 * @param floor_props Properties of the entropy floor.
 * @param cooling The #cooling_function_data used in the run.
 * @param p Pointer to the particle data.
 * @param xp Pointer to the extended particle data.
 */
void cooling_set_subgrid_properties(
    const struct phys_const *phys_const, const struct unit_system *us,
    const struct cosmology *cosmo, const struct hydro_props *hydro_props,
    const struct entropy_floor_properties *floor_props,
    const struct cooling_function_data *cooling, const struct dustevo_props *dp,
    struct part *p, struct xpart *xp) {

  /* Limit imposed by the entropy floor */
  const double A_EOS = entropy_floor(p, cosmo, floor_props);
  const double rho = hydro_get_physical_density(p, cosmo);
  const double u_EOS = gas_internal_energy_from_entropy(rho, A_EOS);
  const double u_EOS_max = u_EOS * exp10(cooling->dlogT_EOS);
  const float log10_u_EOS_max_cgs = log10f(
      u_EOS_max * cooling->colibre_table.internal_energy_to_cgs + FLT_MIN);

  /* Particle's internal energy */
  const double u = hydro_get_physical_internal_energy(p, xp, cosmo);

  /* Particle's temperature */
  const float T = cooling_get_snapshot_temperature(
      phys_const, hydro_props, floor_props, us, cosmo, cooling, p, xp);

  if (cooling->use_colibre_subgrid_EOS == 1) {
    const float log10_T = log10f(T);

    /* Recover the Hydrogen mass fraction */
    const float XH = get_hydrogen_mass_fraction(p, hydro_props);

    /* Get the particle pressure */
    const float P_phys = hydro_get_physical_pressure(p, cosmo);

    /* Are we in an HII region? */
    const int HII_region = xp->tracers_data.HIIregion_timer_gas > 0.;

    /* Get the total metallicity in units of solar */
    float abundance_ratio[colibre_cooling_N_elementtypes];
    const float logZZsol =
        abundance_ratio_to_solar(p, &(cooling->colibre_table), abundance_ratio);

    /* Compute the subgrid properties */
    p->cooling_data.subgrid_temp = compute_subgrid_property(
        cooling, phys_const, floor_props, cosmo, rho, logZZsol, XH, P_phys,
        log10_T, log10_u_EOS_max_cgs, HII_region, abundance_ratio,
        log10(u * cooling->colibre_table.internal_energy_to_cgs),
        colibre_compute_subgrid_temperature);
    p->cooling_data.subgrid_dens = compute_subgrid_property(
        cooling, phys_const, floor_props, cosmo, rho, logZZsol, XH, P_phys,
        log10_T, log10_u_EOS_max_cgs, HII_region, abundance_ratio,
        log10(u * cooling->colibre_table.internal_energy_to_cgs),
        colibre_compute_subgrid_density);

  } else {
    p->cooling_data.subgrid_temp = T;
    p->cooling_data.subgrid_dens = rho;
  }

  if (xp->tracers_data.HIIregion_timer_gas > 0.) {
    /* If the particle is flagged as an HII region,
     * we need to set its abundance array accordingly. */

    const struct globalVariables *ChimesGlobalVars = &cooling->ChimesGlobalVars;

    /* Create ChimesGasVars struct */
    struct gasVariables ChimesGasVars;
    allocate_gas_abundances_memory(&ChimesGasVars, ChimesGlobalVars);

    /* Copy abundances over from xp to ChimesGasVars */
    for (int i = 0; i < ChimesGlobalVars->totalNumberOfSpecies; i++)
      ChimesGasVars.abundances[i] = xp->cooling_data.chimes_abundances[i];

    /* Update element abundances */
    chimes_update_element_abundances(phys_const, us, cosmo, cooling,
                                     hydro_props, p, xp, &ChimesGasVars, u,
                                     /*mode=*/1);

    /* Set HII region abundance array */
    cooling_set_HIIregion_chimes_abundances(&ChimesGasVars, cooling);

    /* Copy abundances from ChimesGasVars to xp. */
    for (int i = 0; i < ChimesGlobalVars->totalNumberOfSpecies; i++)
      xp->cooling_data.chimes_abundances[i] = ChimesGasVars.abundances[i];

    /* Free CHIMES memory. */
    free_gas_abundances_memory(&ChimesGasVars, ChimesGlobalVars);

  } else if (u < u_EOS_max) {

    /* If the particle is within dlogT_EOS of the entropy floor, we need to
     * re-set its abundance array to be in equilibriu using the subgrid density
     * and temperature.
     *
     * Note that, if use_colibe_subgrid_EOS == 0, the subgrid density and
     * temperature have
     * simply been set to the particle density and temperature. */
    const struct globalVariables *ChimesGlobalVars = &cooling->ChimesGlobalVars;

    struct gasVariables ChimesGasVars;
    allocate_gas_abundances_memory(&ChimesGasVars, ChimesGlobalVars);

    for (int i = 0; i < ChimesGlobalVars->totalNumberOfSpecies; i++)
      ChimesGasVars.abundances[i] = xp->cooling_data.chimes_abundances[i];
    chimes_update_element_abundances(phys_const, us, cosmo, cooling,
                                     hydro_props, p, xp, &ChimesGasVars, u,
                                     /*mode=*/1);

    const double mu = chimes_mu(cooling, hydro_props, p, xp);
    const double T_subgrid = p->cooling_data.subgrid_temp;
    const double u_subgrid = get_u_internal_from_T(T_subgrid, mu, phys_const);
    const double u_subgrid_cgs = u_subgrid * cooling->internal_energy_to_cgs;

    chimes_update_gas_vars(u_subgrid_cgs, phys_const, us, cosmo, hydro_props,
                           floor_props, cooling, dp, p, xp, &ChimesGasVars,
                           1.0);

    /* Recover the Hydrogen mass fraction */
    const ChimesFloat XH =
        (ChimesFloat)get_hydrogen_mass_fraction(p, hydro_props);

    ChimesGasVars.nH_tot = XH * p->cooling_data.subgrid_dens *
                           cooling->number_density_to_cgs /
                           phys_const->const_proton_mass;
    ChimesGasVars.temperature = p->cooling_data.subgrid_temp;
    ChimesGasVars.ForceEqOn = 1;
    ChimesGasVars.ThermEvolOn = 0;

    chimes_network(&ChimesGasVars, ChimesGlobalVars);

    /* Copy abundances from ChimesGasVars back to xp. */
    for (int i = 0; i < ChimesGlobalVars->totalNumberOfSpecies; i++)
      xp->cooling_data.chimes_abundances[i] = ChimesGasVars.abundances[i];

    /* Set the net cooling rate to zero */
    xp->cooling_data.net_cooling_rate = 0.0f;

    /* Free CHIMES memory. */
    free_gas_abundances_memory(&ChimesGasVars, ChimesGlobalVars);

  } else if (xp->cooling_data.heated_by_FB &&
             (cooling->set_FB_particles_to_eqm ||
              cooling->destroy_FB_heated_dust)) {

    /* Prior to computing updated element abundances, destroy dust. */
    if (cooling->destroy_FB_heated_dust) {
      dust_redistribute_masses_cooling_step(p, dp);
    }

    if (cooling->set_FB_particles_to_eqm) {

      /* If the particle is not on, or near, the EOS but has been heated by
       * feedback since it last went through the cooling routines, then we need
       * to re-set its abundance array to be in equilibrium. */

      const struct globalVariables *ChimesGlobalVars =
          &cooling->ChimesGlobalVars;

      /* Create ChimesGasVars struct */
      struct gasVariables ChimesGasVars;
      allocate_gas_abundances_memory(&ChimesGasVars, ChimesGlobalVars);

      /* Copy abundances over from xp to ChimesGasVars */
      for (int i = 0; i < ChimesGlobalVars->totalNumberOfSpecies; i++)
        ChimesGasVars.abundances[i] = xp->cooling_data.chimes_abundances[i];

      /* Update element abundances */
      chimes_update_element_abundances(phys_const, us, cosmo, cooling,
                                       hydro_props, p, xp, &ChimesGasVars, u,
                                       /*mode=*/1);

      const double u_cgs = u * cooling->internal_energy_to_cgs;

      /* Update ChimesGasVars */
      chimes_update_gas_vars(u_cgs, phys_const, us, cosmo, hydro_props,
                             floor_props, cooling, dp, p, xp, &ChimesGasVars,
                             1.0);

      /* Set abundances to equilibrium */
      cooling_set_FB_particle_chimes_abundances(&ChimesGasVars, cooling);

      /* Copy abundances from ChimesGasVars to xp. */
      for (int i = 0; i < ChimesGlobalVars->totalNumberOfSpecies; i++)
        xp->cooling_data.chimes_abundances[i] = ChimesGasVars.abundances[i];

      /* Re-calculate the net cooling rate. */
      const float rho_phys_cgs =
          hydro_get_physical_density(p, cosmo) * cooling->density_to_cgs;

      /* cooling rate in erg/cm^3/s. */
      float net_cooling_rate_cgs = (float)chimes_cooling_rate_external(
          &ChimesGasVars, ChimesGlobalVars, /*mode=*/0);

      /* Convert to erg/g/s */
      net_cooling_rate_cgs /= rho_phys_cgs;

      /* Convert to code units. */
      xp->cooling_data.net_cooling_rate =
          net_cooling_rate_cgs * cooling->internal_energy_over_time_from_cgs;

      /* Free CHIMES memory. */
      free_gas_abundances_memory(&ChimesGasVars, ChimesGlobalVars);
    }

    /* Re-set heated_by_FB flag. */
    xp->cooling_data.heated_by_FB = 0;
  }
}

float cooling_get_particle_subgrid_HI_fraction(
    const struct unit_system *us, const struct phys_const *phys_const,
    const struct cosmology *cosmo, const struct hydro_props *hydro_props,
    const struct entropy_floor_properties *floor_props,
    const struct cooling_function_data *cooling, const struct part *p,
    const struct xpart *xp) {

  const struct globalVariables *ChimesGlobalVars = &cooling->ChimesGlobalVars;

  return xp->cooling_data
      .chimes_abundances[ChimesGlobalVars->speciesIndices[sp_HI]];
}

float cooling_get_particle_subgrid_HII_fraction(
    const struct unit_system *us, const struct phys_const *phys_const,
    const struct cosmology *cosmo, const struct hydro_props *hydro_props,
    const struct entropy_floor_properties *floor_props,
    const struct cooling_function_data *cooling, const struct part *p,
    const struct xpart *xp) {

  const struct globalVariables *ChimesGlobalVars = &cooling->ChimesGlobalVars;

  return xp->cooling_data
      .chimes_abundances[ChimesGlobalVars->speciesIndices[sp_HII]];
}

float cooling_get_particle_subgrid_H2_fraction(
    const struct unit_system *us, const struct phys_const *phys_const,
    const struct cosmology *cosmo, const struct hydro_props *hydro_props,
    const struct entropy_floor_properties *floor_props,
    const struct cooling_function_data *cooling, const struct part *p,
    const struct xpart *xp) {

  const struct globalVariables *ChimesGlobalVars = &cooling->ChimesGlobalVars;

  return xp->cooling_data
      .chimes_abundances[ChimesGlobalVars->speciesIndices[sp_H2]];
}

/**
 * @brief Compute the temperature of the gas.
 *
 * For particles on the entropy floor, we use pressure equilibrium to
 * infer the properties of the particle.
 *
 * @param us The internal system of units.
 * @param phys_const The physical constants.
 * @param hydro_props The properties of the hydro scheme.
 * @param cosmo The cosmological model.
 * @param floor_props The properties of the entropy floor.
 * @param cooling The properties of the cooling scheme.
 * @param p The #part.
 * @param xp The #xpart.
 */
float cooling_get_subgrid_temperature(
    const struct unit_system *us, const struct phys_const *phys_const,
    const struct cosmology *cosmo, const struct hydro_props *hydro_props,
    const struct entropy_floor_properties *floor_props,
    const struct cooling_function_data *cooling, const struct part *p,
    const struct xpart *xp) {

  /* Get the particle's temperature */
  const float T = cooling_get_temperature(phys_const, hydro_props, us, cosmo,
                                          cooling, p, xp);

  if (cooling->use_colibre_subgrid_EOS == 1) {
    /* Particle's internal energy */
    const double u = hydro_get_physical_internal_energy(p, xp, cosmo);

    const float log10_T = log10f(T);

    /* Physical density of this particle */
    const float rho_phys = hydro_get_physical_density(p, cosmo);

    /* Recover the Hydrogen mass fraction */
    const float XH = get_hydrogen_mass_fraction(p, hydro_props);

    /* Get the particle pressure */
    const float P_phys = hydro_get_physical_pressure(p, cosmo);

    /* Are we in an HII region? */
    const int HII_region = xp->tracers_data.HIIregion_timer_gas > 0.;

    /* Get the total metallicity in units of solar */
    float abundance_ratio[colibre_cooling_N_elementtypes];
    const float logZZsol =
        abundance_ratio_to_solar(p, &(cooling->colibre_table), abundance_ratio);

    /* Limit imposed by the entropy floor */
    const double A_EOS = entropy_floor(p, cosmo, floor_props);
    const double u_EOS = gas_internal_energy_from_entropy(rho_phys, A_EOS);
    const double u_EOS_max = u_EOS * exp10(cooling->dlogT_EOS);

    const float log10_u_EOS_max_cgs = log10f(
        u_EOS_max * cooling->colibre_table.internal_energy_to_cgs + FLT_MIN);

    return compute_subgrid_property(
        cooling, phys_const, floor_props, cosmo, rho_phys, logZZsol, XH, P_phys,
        log10_T, log10_u_EOS_max_cgs, HII_region, abundance_ratio,
        log10(u * cooling->colibre_table.internal_energy_to_cgs),
        colibre_compute_subgrid_temperature);
  } else {
    return T;
  }
}

/**
 * @brief Compute the physical density of the gas.
 *
 * For particles on the entropy floor, we use pressure equilibrium to
 * infer the properties of the particle.
 *
 * Note that we return the density in physical coordinates.
 *
 * @param us The internal system of units.
 * @param phys_const The physical constants.
 * @param hydro_props The properties of the hydro scheme.
 * @param cosmo The cosmological model.
 * @param floor_props The properties of the entropy floor.
 * @param cooling The properties of the cooling scheme.
 * @param p The #part.
 * @param xp The #xpart.
 */
float cooling_get_subgrid_density(
    const struct unit_system *us, const struct phys_const *phys_const,
    const struct cosmology *cosmo, const struct hydro_props *hydro_props,
    const struct entropy_floor_properties *floor_props,
    const struct cooling_function_data *cooling, const struct part *p,
    const struct xpart *xp) {

  if (cooling->use_colibre_subgrid_EOS == 1) {

    /* Particle's internal energy */
    const double u = hydro_get_physical_internal_energy(p, xp, cosmo);
    /* Particle's temperature */
    const float T = cooling_get_temperature(phys_const, hydro_props, us, cosmo,
                                            cooling, p, xp);
    const float log10_T = log10f(T);

    /* Physical density of this particle */
    const float rho_phys = hydro_get_physical_density(p, cosmo);

    /* Recover the Hydrogen mass fraction */
    const float XH = get_hydrogen_mass_fraction(p, hydro_props);

    /* Get the particle pressure */
    const float P_phys = hydro_get_physical_pressure(p, cosmo);

    /* Are we in an HII region? */
    const int HII_region = xp->tracers_data.HIIregion_timer_gas > 0.;

    /* Get the total metallicity in units of solar */
    float abundance_ratio[colibre_cooling_N_elementtypes];
    const float logZZsol =
        abundance_ratio_to_solar(p, &(cooling->colibre_table), abundance_ratio);

    /* Limit imposed by the entropy floor */
    const double A_EOS = entropy_floor(p, cosmo, floor_props);
    const double u_EOS = gas_internal_energy_from_entropy(rho_phys, A_EOS);
    const double u_EOS_max = u_EOS * exp10(cooling->dlogT_EOS);

    const float log10_u_EOS_max_cgs = log10f(
        u_EOS_max * cooling->colibre_table.internal_energy_to_cgs + FLT_MIN);

    return compute_subgrid_property(
        cooling, phys_const, floor_props, cosmo, rho_phys, logZZsol, XH, P_phys,
        log10_T, log10_u_EOS_max_cgs, HII_region, abundance_ratio,
        log10(u * cooling->colibre_table.internal_energy_to_cgs),
        colibre_compute_subgrid_density);
  } else {
    return hydro_get_physical_density(p, cosmo);
  }
}

/**
 * @brief Set CHIMES abundances for HII regions.
 *
 * Sets the CHIMES abundance array according to the
 * HIIregion_fion parameter. This parameter gives the
 * fraction of each element that is singly ionised,
 * with the remainder set to neutral.
 *
 * @param ChimesGasVars CHIMES gasVariables structure.
 * @param cooling The #cooling_function_data used in the run.
 */
void cooling_set_HIIregion_chimes_abundances(
    struct gasVariables *ChimesGasVars,
    const struct cooling_function_data *cooling) {

  const struct globalVariables *ChimesGlobalVars = &cooling->ChimesGlobalVars;

  /* First zero all abundances */
  for (int i = 0; i < ChimesGlobalVars->totalNumberOfSpecies; i++)
    ChimesGasVars->abundances[i] = 0.0;

  /* Set H and He ionisation fractions. */
  ChimesGasVars->abundances[ChimesGlobalVars->speciesIndices[sp_HII]] =
      (ChimesFloat)cooling->HIIregion_fion;
  ChimesGasVars->abundances[ChimesGlobalVars->speciesIndices[sp_HI]] =
      (ChimesFloat)(1.0 -
                    ChimesGasVars
                        ->abundances[ChimesGlobalVars->speciesIndices[sp_HII]]);
  ChimesGasVars->abundances[ChimesGlobalVars->speciesIndices[sp_elec]] +=
      ChimesGasVars->abundances[ChimesGlobalVars->speciesIndices[sp_HII]];

  ChimesGasVars->abundances[ChimesGlobalVars->speciesIndices[sp_HeII]] =
      (ChimesFloat)ChimesGasVars->element_abundances[0] *
      cooling->HIIregion_fion;
  ChimesGasVars->abundances[ChimesGlobalVars->speciesIndices[sp_HeI]] =
      (ChimesFloat)ChimesGasVars->element_abundances[0] *
      (1.0 - cooling->HIIregion_fion);
  ChimesGasVars->abundances[ChimesGlobalVars->speciesIndices[sp_elec]] +=
      ChimesGasVars->abundances[ChimesGlobalVars->speciesIndices[sp_HeII]];

  /* Set metal ionisation fractions. */
  if ((ChimesGlobalVars->element_included[0] == 1) &&
      (ChimesGasVars->element_abundances[1] > 0.0)) {
    ChimesGasVars->abundances[ChimesGlobalVars->speciesIndices[sp_CII]] =
        (ChimesFloat)ChimesGasVars->element_abundances[1] *
        cooling->HIIregion_fion;
    ChimesGasVars->abundances[ChimesGlobalVars->speciesIndices[sp_CI]] =
        (ChimesFloat)ChimesGasVars->element_abundances[1] *
        (1.0 - cooling->HIIregion_fion);
    ChimesGasVars->abundances[ChimesGlobalVars->speciesIndices[sp_elec]] +=
        ChimesGasVars->abundances[ChimesGlobalVars->speciesIndices[sp_CII]];
  }

  if ((ChimesGlobalVars->element_included[1] == 1) &&
      (ChimesGasVars->element_abundances[2] > 0.0)) {
    ChimesGasVars->abundances[ChimesGlobalVars->speciesIndices[sp_NII]] =
        (ChimesFloat)ChimesGasVars->element_abundances[2] *
        cooling->HIIregion_fion;
    ChimesGasVars->abundances[ChimesGlobalVars->speciesIndices[sp_NI]] =
        (ChimesFloat)ChimesGasVars->element_abundances[2] *
        (1.0 - cooling->HIIregion_fion);
    ChimesGasVars->abundances[ChimesGlobalVars->speciesIndices[sp_elec]] +=
        ChimesGasVars->abundances[ChimesGlobalVars->speciesIndices[sp_NII]];
  }

  if ((ChimesGlobalVars->element_included[2] == 1) &&
      (ChimesGasVars->element_abundances[3] > 0.0)) {
    ChimesGasVars->abundances[ChimesGlobalVars->speciesIndices[sp_OII]] =
        (ChimesFloat)ChimesGasVars->element_abundances[3] *
        cooling->HIIregion_fion;
    ChimesGasVars->abundances[ChimesGlobalVars->speciesIndices[sp_OI]] =
        (ChimesFloat)ChimesGasVars->element_abundances[3] *
        (1.0 - cooling->HIIregion_fion);
    ChimesGasVars->abundances[ChimesGlobalVars->speciesIndices[sp_elec]] +=
        ChimesGasVars->abundances[ChimesGlobalVars->speciesIndices[sp_OII]];
  }

  if ((ChimesGlobalVars->element_included[3] == 1) &&
      (ChimesGasVars->element_abundances[4] > 0.0)) {
    ChimesGasVars->abundances[ChimesGlobalVars->speciesIndices[sp_NeII]] =
        (ChimesFloat)ChimesGasVars->element_abundances[4] *
        cooling->HIIregion_fion;
    ChimesGasVars->abundances[ChimesGlobalVars->speciesIndices[sp_NeI]] =
        (ChimesFloat)ChimesGasVars->element_abundances[4] *
        (1.0 - cooling->HIIregion_fion);
    ChimesGasVars->abundances[ChimesGlobalVars->speciesIndices[sp_elec]] +=
        ChimesGasVars->abundances[ChimesGlobalVars->speciesIndices[sp_NeII]];
  }

  if ((ChimesGlobalVars->element_included[4] == 1) &&
      (ChimesGasVars->element_abundances[5] > 0.0)) {
    ChimesGasVars->abundances[ChimesGlobalVars->speciesIndices[sp_MgII]] =
        (ChimesFloat)ChimesGasVars->element_abundances[5] *
        cooling->HIIregion_fion;
    ChimesGasVars->abundances[ChimesGlobalVars->speciesIndices[sp_MgI]] =
        (ChimesFloat)ChimesGasVars->element_abundances[5] *
        (1.0 - cooling->HIIregion_fion);
    ChimesGasVars->abundances[ChimesGlobalVars->speciesIndices[sp_elec]] +=
        ChimesGasVars->abundances[ChimesGlobalVars->speciesIndices[sp_MgII]];
  }

  if ((ChimesGlobalVars->element_included[5] == 1) &&
      (ChimesGasVars->element_abundances[6] > 0.0)) {
    ChimesGasVars->abundances[ChimesGlobalVars->speciesIndices[sp_SiII]] =
        (ChimesFloat)ChimesGasVars->element_abundances[6] *
        cooling->HIIregion_fion;
    ChimesGasVars->abundances[ChimesGlobalVars->speciesIndices[sp_SiI]] =
        (ChimesFloat)ChimesGasVars->element_abundances[6] *
        (1.0 - cooling->HIIregion_fion);
    ChimesGasVars->abundances[ChimesGlobalVars->speciesIndices[sp_elec]] +=
        ChimesGasVars->abundances[ChimesGlobalVars->speciesIndices[sp_SiII]];
  }

  if ((ChimesGlobalVars->element_included[6] == 1) &&
      (ChimesGasVars->element_abundances[7] > 0.0)) {
    ChimesGasVars->abundances[ChimesGlobalVars->speciesIndices[sp_SII]] =
        (ChimesFloat)ChimesGasVars->element_abundances[7] *
        cooling->HIIregion_fion;
    ChimesGasVars->abundances[ChimesGlobalVars->speciesIndices[sp_SI]] =
        (ChimesFloat)ChimesGasVars->element_abundances[7] *
        (1.0 - cooling->HIIregion_fion);
    ChimesGasVars->abundances[ChimesGlobalVars->speciesIndices[sp_elec]] +=
        ChimesGasVars->abundances[ChimesGlobalVars->speciesIndices[sp_SII]];
  }

  if ((ChimesGlobalVars->element_included[7] == 1) &&
      (ChimesGasVars->element_abundances[8] > 0.0)) {
    ChimesGasVars->abundances[ChimesGlobalVars->speciesIndices[sp_CaII]] =
        (ChimesFloat)ChimesGasVars->element_abundances[8] *
        cooling->HIIregion_fion;
    ChimesGasVars->abundances[ChimesGlobalVars->speciesIndices[sp_CaI]] =
        (ChimesFloat)ChimesGasVars->element_abundances[8] *
        (1.0 - cooling->HIIregion_fion);
    ChimesGasVars->abundances[ChimesGlobalVars->speciesIndices[sp_elec]] +=
        ChimesGasVars->abundances[ChimesGlobalVars->speciesIndices[sp_CaII]];
  }

  if ((ChimesGlobalVars->element_included[8] == 1) &&
      (ChimesGasVars->element_abundances[9] > 0.0)) {
    ChimesGasVars->abundances[ChimesGlobalVars->speciesIndices[sp_FeII]] =
        (ChimesFloat)ChimesGasVars->element_abundances[9] *
        cooling->HIIregion_fion;
    ChimesGasVars->abundances[ChimesGlobalVars->speciesIndices[sp_FeI]] =
        (ChimesFloat)ChimesGasVars->element_abundances[9] *
        (1.0 - cooling->HIIregion_fion);
    ChimesGasVars->abundances[ChimesGlobalVars->speciesIndices[sp_elec]] +=
        ChimesGasVars->abundances[ChimesGlobalVars->speciesIndices[sp_FeII]];
  }
}

/**
 * @brief Set CHIMES abundances for particles heated by feedback.
 *
 * Sets gas particle that has just been heated by feedback
 * to be in chemical equilibrium at its new temperature.
 *
 * @param ChimesGasVars CHIMES gasVariables structure.
 * @param cooling The #cooling_function_data used in the run.
 */
void cooling_set_FB_particle_chimes_abundances(
    struct gasVariables *ChimesGasVars,
    const struct cooling_function_data *cooling) {

  const struct globalVariables *ChimesGlobalVars = &cooling->ChimesGlobalVars;

  /* Save equilibrium and ThermEvol flags. */
  int ForceEqOn_save = ChimesGasVars->ForceEqOn;
  int ThermEvolOn_save = ChimesGasVars->ThermEvolOn;

  /* Use equilibrium abundances. */
  ChimesGasVars->ForceEqOn = 1;

  /* Disable temperature evolution. */
  ChimesGasVars->ThermEvolOn = 0;

  /* Set abundance array to equilibrium. */
  chimes_network(ChimesGasVars, ChimesGlobalVars);

  /* Revert flags to saved values. */
  ChimesGasVars->ForceEqOn = ForceEqOn_save;
  ChimesGasVars->ThermEvolOn = ThermEvolOn_save;
}
