nonthermal.cc

#include "nonthermal.h"

#include <algorithm>
#include <array>
#include <cmath>
#include <cstddef>
#include <cstdio>
#include <cstdlib>
#include <filesystem>
#include <functional>
#include <ios>
#include <numeric>
#include <ranges>
#include <span>
#include <sstream>
#include <string>
#include <tuple>
#include <vector>

#pragma clang unsafe_buffer_usage begin
#ifdef EIGEN_OFF
#include <gsl/gsl_blas.h>
#include <gsl/gsl_cblas.h>
#include <gsl/gsl_linalg.h>
#include <gsl/gsl_matrix_double.h>
#include <gsl/gsl_permutation.h>
#include <gsl/gsl_vector_double.h>
#else
#include <Eigen/Core>
#include <Eigen/Dense>
#endif
#include <mpi.h>
#pragma clang unsafe_buffer_usage end

#include "artisoptions.h"
#include "atomic.h"
#include "constants.h"
#include "decay.h"
#include "globals.h"
#include "grid.h"
#include "input.h"
#include "ltepop.h"
#include "macroatom.h"
#include "packet.h"
#include "random.h"
#include "sn3d.h"
#include "stats.h"
#include "thermalbalance.h"

namespace nonthermal {

namespace {

// minimum number fraction of the total population to include in SF solution
constexpr double MIN_ION_OVER_NNTOT = 1.e-8;

// minimum deposition rate density (eV/s/cm^3) to solve SF equation
constexpr double MINDEPRATE = 0.;

// Bohr radius squared in cm^2
constexpr double A_naught_squared = 2.800285203e-17;

constexpr std::array<std::string, 28> shellnames{"K ", "L1", "L2", "L3", "M1", "M2", "M3", "M4", "M5", "N1",
                                                 "N2", "N3", "N4", "N5", "N6", "N7", "O1", "O2", "O3", "O4",
                                                 "O5", "O6", "O7", "P1", "P2", "P3", "P4", "Q1"};

std::vector<std::vector<double>> elements_electron_binding;
std::vector<std::vector<int>> allions_shell_occupancies;

struct ShellParams {
  int Z{-1};
  int ionstage{-1};
  int n{-1};
  int l{-1};
  double ionpot_ev{NAN};
  double A{NAN};
  double B{NAN};
  double C{NAN};
  double D{NAN};
  // track the statistical weight represented by the values below, so they can be updated with new g-weighted averaged
  // values
  double auger_g_accumulated = 0.;

  // probability of 0, 1, ..., NT_MAX_AUGER_ELECTRONS Auger electrons being ejected when the shell is ionised
  std::array<double, NT_MAX_AUGER_ELECTRONS + 1> prob_num_auger{};

  // the average kinetic energy released in Auger electrons after making a hole in this shell
  float en_auger_ev{NAN};
  float n_auger_elec_avg{NAN};

  ShellParams() {
    std::ranges::fill(prob_num_auger, 0.);
    prob_num_auger[0] = 1.;
  }
};

std::vector<ShellParams> colliondata;

static_assert(SF_EMIN > 0.);
constexpr double DELTA_E = (SF_EMAX - SF_EMIN) / (SFPTS - 1);

// energy grid on which solution is sampled [eV]
constexpr auto engrid(int index) -> double { return SF_EMIN + (index * DELTA_E); }

const auto logengrid = []() {
  std::vector<double> _logengrid(SFPTS);
  for (int i = 0; i < SFPTS; i++) {
    _logengrid[i] = std::log(engrid(i));
  }
  return _logengrid;
}();

// samples of the source function (energy distribution of deposited energy)
constexpr auto sourcevec(const int index) {
  assert_testmodeonly(index >= 0 && index < SFPTS);

  // spread the source over some energy width
  constexpr int source_spread_pts = static_cast<int>(SFPTS * 0.03333) + 1;
  constexpr double source_spread_en = source_spread_pts * DELTA_E;
  constexpr int sourcestartindex = SFPTS - source_spread_pts;

  return (index < sourcestartindex) ? 0. : 1. / source_spread_en;

  // or put all of the source into one point at SF_EMAX
  // return (index < SFPTS - 1) ? 0. : 1. / DELTA_E;
  // so that E_init_ev = SF_EMAX;
}

// the energy injection rate density (integral of E * S(e) dE) in eV/s/cm3 that the Spencer-Fano equation is solved for.
// This is arbitrary and and the solution will be scaled to match the actual energy deposition rate density.
constexpr double E_init_ev = []() {
  double integral = 0.;
  for (int s = 0; s < SFPTS; s++) {
    integral += sourcevec(s) * DELTA_E * engrid(s);
  }
  return integral;
}();

// rhs is the constant term (not dependent on y func) in each equation
constexpr auto rhsvec = []() {
  std::array<double, SFPTS> _rhsvec{};
  double source_integral_to_SF_EMAX = 0.;
  for (int i = SFPTS - 1; i >= 0; i--) {
    _rhsvec[i] = source_integral_to_SF_EMAX * DELTA_E;
    source_integral_to_SF_EMAX += sourcevec(i);
  }
  return _rhsvec;
}();

// Monte Carlo result - compare to analytical expectation
double nt_energy_deposited = 0;

struct NonThermalExcitation {
  double frac_deposition;  // the fraction of the non-thermal deposition energy going to the excitation transition
  double ratecoeffperdeposition;  // the excitation rate coefficient divided by the deposition rate density
  int alltransindex;
};

// pointer to either local or node-shared memory excitation list of all cells
std::span<NonThermalExcitation> excitations_list_all_cells{};

// the minimum of MAX_NT_EXCITATIONS_STORED and the number of included excitation transitions in the atomic dataset
int nt_excitations_stored = 0;

struct NonThermalSolutionIon {
  float eff_ionpot{0.};  // these are used to calculate the non-thermal ionisation rate
  double fracdep_ionisation_ion{0.};  // the fraction of the non-thermal deposition energy going to ionizing each ion

  // probability that one ionisation of this ion will produce n Auger electrons.
  // items sum to 1.0 for a given ion
  std::array<float, NT_MAX_AUGER_ELECTRONS + 1> prob_num_auger{};
  // like prob_num_auger, but energy weighted. items sum to 1.0 for an ion
  std::array<float, NT_MAX_AUGER_ELECTRONS + 1> ionenfrac_num_auger{};
};

std::span<NonThermalSolutionIon> ion_data_all_cells{};

struct NonThermalCellSolution {
  float frac_heating = 1.;  // energy fractions should add up to 1.0 if the solution is good
  float frac_ionisation = 0.;  // fraction of deposition energy going to ionisation
  float frac_excitation = 0.;  // fraction of deposition energy going to excitation

  int frac_excitations_list_size = 0;

  int timestep_last_solved = -1;  // the quantities above were calculated for this timestep
  float nneperion_when_solved{NAN};  // the nne when the solver was last run
};

std::span<NonThermalCellSolution> nt_solution;

std::span<double> deposition_rate_density_all_cells;

constexpr auto uppertriangular(const int i, const int j) -> int {
  assert_testmodeonly(i >= 0);
  assert_testmodeonly(i < SFPTS);
  // sometimes you might want to get an offset for a row using j = 0 < i, so that j can be added to it.
  // assert_testmodeonly(j >= i);
  assert_testmodeonly(j < SFPTS);
  return (SFPTS * i) - (i * (i + 1) / 2) + j;
}

constexpr void decompactify_triangular_matrix(std::span<double> matrix) {
  // the indices work in decreasing order because we're working in-place and
  // don't want to overwrite data that we still need to read
  for (int i = SFPTS - 1; i > 0; i--) {
    const int rowoffset = uppertriangular(i, 0);
    for (int j = SFPTS - 1; j >= i; j--) {
      matrix[(i * SFPTS) + j] = matrix[rowoffset + j];
    }
    for (int j = i - 1; j >= 0; j--) {
      matrix[(i * SFPTS) + j] = 0.;
    }
  }
}

auto get_approx_shell_occupancies(const int atomic_number, const int nbound) {
  assert_always(nbound > 0);
  const int ioncharge = atomic_number - nbound;
  assert_always(ioncharge >= 0);
  std::vector<int> q;
  q.resize(std::max(10UZ, elements_electron_binding[atomic_number - 1].size()), 0);

  for (int electron_loop = 0; electron_loop < nbound; electron_loop++) {
    if (q[0] < 2) {
      q[0]++;  // K 1s
    } else if (q[1] < 2) {
      q[1]++;  // L1 2s
    } else if (q[2] < 2) {
      q[2]++;  // L2 2p[1/2]
    } else if (q[3] < 4) {
      q[3]++;  // L3 2p[3/2]
    } else if (q[4] < 2) {
      q[4]++;  // M1 3s
    } else if (q[5] < 2) {
      q[5]++;  // M2 3p[1/2]
    } else if (q[6] < 4) {
      q[6]++;  // M3 3p[3/2]
    } else if (ioncharge == 0) {
      if (q[9] < 2) {
        q[9]++;  // N1 4s
      } else if (q[7] < 4) {
        q[7]++;  // M4 3d[3/2]
      } else if (q[8] < 6) {
        q[8]++;  // M5 3d[5/2]
      } else {
        assert_always(false);  // Going beyond the 4s shell
      }
    } else if (ioncharge == 1) {
      if (q[9] < 1) {
        q[9]++;  // N1 4s
      } else if (q[7] < 4) {
        q[7]++;  // M4 3d[3/2]
      } else if (q[8] < 6) {
        q[8]++;  // M5 3d[5/2]
      } else {
        assert_always(false);  // Going beyond the 4s shell
      }
    } else if (ioncharge > 1) {
      if (q[7] < 4) {
        q[7]++;  // M4 3d[3/2]
      } else if (q[8] < 6) {
        q[8]++;  // M5 3d[5/2]
      } else {
        assert_always(false);  // Going beyond the 4s shell
      }
    }
  }
  assert_always(nbound == std::accumulate(q.begin(), q.end(), 0));
  return q;
}

auto calculate_ion_shell_occupancies(const int atomic_number, const int nbound,
                                     const std::vector<int>& element_shells_q_neutral) {
  assert_testmodeonly(nbound >= 0);

  const auto shellcount =
      std::min(element_shells_q_neutral.size(), elements_electron_binding[atomic_number - 1].size());
  std::vector<int> element_shells_q;
  resize_exactly(element_shells_q, shellcount);

  int electron_count = 0;
  for (auto shellindex = 0ZU; shellindex < shellcount; shellindex++) {
    const int electronsinshell_neutral = element_shells_q_neutral.at(shellindex);

    int electronsinshell = 0;
    if ((electron_count + electronsinshell_neutral) <= nbound) {
      electronsinshell = electronsinshell_neutral;
    } else {
      electronsinshell = nbound - electron_count;
    }
    assert_always(electronsinshell <= electronsinshell_neutral);
    element_shells_q[shellindex] = electronsinshell;
    electron_count += electronsinshell;
    assert_always(electron_count <= nbound);
  }

  return element_shells_q;
}

auto read_shell_configs() {
  assert_always(NT_WORKFUNCTION_USE_SHELL_OCCUPANCY_FILE);
  auto shells_file = fstream_required("electron_shell_occupancy.txt", std::ios::in);

  int nshells = 0;  // number of shell in binding energy file
  int n_z_binding = 0;  // number of elements in file

  std::string line;
  assert_always(get_noncommentline(shells_file, line));
  std::istringstream{line} >> nshells >> n_z_binding;
  printlnlog("Reading electron_shell_occupancy.txt with {} elements and {} shells", n_z_binding, nshells);

  std::vector<std::vector<int>> elements_shells_q;

  elements_shells_q.resize(n_z_binding, std::vector<int>(nshells, 0.));

  assert_always(elements_shells_q.size() == elements_electron_binding.size());

  int zminusone = 0;
  while (get_noncommentline(shells_file, line)) {
    std::istringstream ssline(line);

    int z_element = 0;
    assert_always(ssline >> z_element);

    assert_always(elements_shells_q[zminusone].size() == elements_electron_binding[zminusone].size());
    for (int shell = 0; shell < nshells; shell++) {
      int q = 0;
      assert_always(ssline >> q);
      elements_shells_q.at(zminusone).at(shell) = q;
    }
    zminusone++;
  }
  return elements_shells_q;
}

void read_binding_energies() {
  const bool binding_en_newformat_local = std::filesystem::exists("binding_energies_lotz_tab1and2.txt") ||
                                          std::filesystem::exists("data/binding_energies_lotz_tab1and2.txt");
  bool binding_en_newformat = binding_en_newformat_local;
  // just in case the file system was faulty and the ranks disagree on the existence of the files
  MPI_Allreduce_safe(binding_en_newformat, MPI_LOR, MPI_COMM_WORLD);

  int nshells = 0;  // number of shell in binding energy file
  int n_z_binding = 0;  // number of elements in binding energy file

  const auto* filename = binding_en_newformat ? "binding_energies_lotz_tab1and2.txt" : "binding_energies.txt";
  auto binding_energies_file = fstream_required(filename, std::ios::in);

  std::string line;
  assert_always(get_noncommentline(binding_energies_file, line));
  std::istringstream{line} >> nshells >> n_z_binding;
  printlnlog("Reading binding energies file '{}' with {} elements and {} shells", filename, n_z_binding, nshells);

  elements_electron_binding.resize(n_z_binding, std::vector<double>(nshells, 0.));

  for (int zminusone = 0; zminusone < n_z_binding; zminusone++) {
    assert_always(get_noncommentline(binding_energies_file, line));
    std::istringstream ssline(line);
    // new file as an atomic number column
    if (binding_en_newformat) {
      int z_element{-1};
      ssline >> z_element;
      assert_always(z_element == (zminusone + 1));
    }
    for (int shell = 0; shell < nshells; shell++) {
      float bindingenergy = 0.;
      assert_always(ssline >> bindingenergy);
      elements_electron_binding.at(zminusone).at(shell) = bindingenergy * EV;
    }
  }

  std::vector<std::vector<int>> elements_neutral_shells_q;
  if constexpr (NT_WORKFUNCTION_USE_SHELL_OCCUPANCY_FILE) {
    assert_always(binding_en_newformat);
    elements_neutral_shells_q = read_shell_configs();
  }

  resize_exactly(allions_shell_occupancies, get_includedions());
  for (int element = 0; element < get_nelements(); element++) {
    for (int ion = 0; ion < get_nions(element); ion++) {
      const int ioncharge = get_ionstage(element, ion) - 1;
      const int atomic_number = get_atomicnumber(element);
      const int nbound = atomic_number - ioncharge;
      if (nbound <= 0) {
        continue;
      }
      allions_shell_occupancies[get_uniqueionindex(element, ion)] =
          NT_WORKFUNCTION_USE_SHELL_OCCUPANCY_FILE
              ? calculate_ion_shell_occupancies(atomic_number, nbound, elements_neutral_shells_q.at(atomic_number - 1))
              : get_approx_shell_occupancies(atomic_number, nbound);
    }
  }
}

[[nodiscard]] auto get_cell_ntexcitations(const ptrdiff_t nonemptymgi) {
  return excitations_list_all_cells.subspan(nonemptymgi * nt_excitations_stored,
                                            nt_solution[nonemptymgi].frac_excitations_list_size);
}

[[nodiscard]] auto get_cell_allions_data(const ptrdiff_t nonemptymgi) {
  return ion_data_all_cells.subspan(nonemptymgi * get_includedions(), get_includedions());
}

auto get_auger_probability(const ptrdiff_t nonemptymgi, const int element, const int ion, const int naugerelec) {
  assert_always(naugerelec <= NT_MAX_AUGER_ELECTRONS);
  const int uniqueionindex = get_uniqueionindex(element, ion);
  return get_cell_allions_data(nonemptymgi)[uniqueionindex].prob_num_auger[naugerelec];
}

auto get_ion_auger_enfrac(const ptrdiff_t nonemptymgi, const int element, const int ion, const int naugerelec) {
  assert_always(naugerelec <= NT_MAX_AUGER_ELECTRONS);
  const int uniqueionindex = get_uniqueionindex(element, ion);
  return get_cell_allions_data(nonemptymgi)[uniqueionindex].ionenfrac_num_auger[naugerelec];
}

void check_auger_probabilities(const ptrdiff_t nonemptymgi) {
  bool problem_found = false;

  for (int element = 0; element < get_nelements(); element++) {
    for (int ion = 0; ion < get_nions(element) - 1; ion++) {
      double prob_sum = 0.;
      double ionenfrac_sum = 0.;
      for (int a = 0; a <= NT_MAX_AUGER_ELECTRONS; a++) {
        prob_sum += get_auger_probability(nonemptymgi, element, ion, a);
        ionenfrac_sum += get_ion_auger_enfrac(nonemptymgi, element, ion, a);
      }

      if (fabs(prob_sum - 1.0) > 0.001) {
        printlnlog("Problem with Auger probabilities for cell {} Z={} ionstage {} prob_sum {:g}",
                   grid::get_mgi_of_nonemptymgi(nonemptymgi), get_atomicnumber(element), get_ionstage(element, ion),
                   prob_sum);
        for (int a = 0; a <= NT_MAX_AUGER_ELECTRONS; a++) {
          printlnlog("{}: {:g}", a, get_auger_probability(nonemptymgi, element, ion, a));
        }
        problem_found = true;
      }

      if (fabs(ionenfrac_sum - 1.0) > 0.001) {
        printlnlog("Problem with Auger energy frac sum for cell {} Z={} ionstage {} ionenfrac_sum {:g}",
                   grid::get_mgi_of_nonemptymgi(nonemptymgi), get_atomicnumber(element), get_ionstage(element, ion),
                   ionenfrac_sum);
        for (int a = 0; a <= NT_MAX_AUGER_ELECTRONS; a++) {
          printlnlog("{}: {:g}", a, get_ion_auger_enfrac(nonemptymgi, element, ion, a));
        }
        problem_found = true;
      }
    }
  }

  assert_always(!problem_found);
}

void read_auger_data() {
  printlnlog("Reading Auger effect data...");
  auto augerfile = fstream_required("auger-km1993-table2.txt", std::ios::in);

  // map x-ray notation shells K L1 L2 L3 M1 M2 M3 to quantum numbers n and l
  constexpr std::array<int, 7> xrayn = {1, 2, 2, 2, 3, 3, 3};
  constexpr std::array<int, 7> xrayl = {0, 0, 1, 1, 0, 1, 1};
  constexpr std::array<int, 7> xrayg = {2, 2, 2, 4, 2, 2, 4};  // g statistical weight = 2j + 1

  std::string strline;
  while (get_noncommentline(augerfile, strline)) {
    int Z = -1;
    int ionstage = -1;
    int shellnum = -1;

    int linepos = 0;
    int offset = 0;

    assert_always(sscanf(strline.c_str(), "%d %d%n", &Z, &ionstage, &offset) == 2);
    assert_always(offset == 5);
    linepos += offset;

    const int element = get_elementindex(Z);

    if (element >= 0 && get_ionstage(element, 0) <= ionstage &&
        ionstage < (get_ionstage(element, 0) + get_nions(element))) {
      float ionpot_ev = -1;
      float en_auger_ev_total_nocorrection = -1;
      int epsilon_e3 = -1;

      assert_always(sscanf(strline.substr(linepos).c_str(), "%d %g %g %d%n", &shellnum, &ionpot_ev,
                           &en_auger_ev_total_nocorrection, &epsilon_e3, &offset) == 4);
      assert_always(offset == 20);

      float n_auger_elec_avg = 0;
      std::array<double, (NT_MAX_AUGER_ELECTRONS + 1)> prob_num_auger{};
      for (int a = 0; a < 9; a++) {
        linepos = 26 + (a * 5);
        // have to read out exactly 5 characters at a time because the columns are sometimes not separated by a space
        std::array<char, 6> strprob{"00000"};
        assert_always(sscanf(strline.substr(linepos).c_str(), "%5c%n", strprob.data(), &offset) == 1);
        assert_always(offset == 5);
        strprob[5] = '\0';

        int probnaugerelece4 = -1;
        assert_always(sscanf(strprob.data(), "%d", &probnaugerelece4) == 1);

        const double probnaugerelec = probnaugerelece4 / 10000.;

        assert_always(probnaugerelec <= 1.0);

        n_auger_elec_avg += static_cast<float>(a * probnaugerelec);

        if (a <= NT_MAX_AUGER_ELECTRONS) {
          prob_num_auger.at(a) = probnaugerelec;
        } else {
          // add the rates of all higher ionisations to the top one
          prob_num_auger.at(NT_MAX_AUGER_ELECTRONS) += probnaugerelec;
        }
      }

      // use the epsilon correction factor as in equation 7 of Kaastra & Mewe (1993)
      auto en_auger_ev = static_cast<float>(en_auger_ev_total_nocorrection - (epsilon_e3 / 1000. * ionpot_ev));

      const int n = xrayn[shellnum - 1];
      const int l = xrayl[shellnum - 1];
      const int g = xrayg[shellnum - 1];

      if (!std::isfinite(en_auger_ev) || en_auger_ev < 0) {
        printlnlog("  WARNING: Z={:2} ionstage {:2} shellnum {} en_auger_ev is {:g}. Setting to zero.", Z, ionstage,
                   shellnum, en_auger_ev);
        en_auger_ev = 0.;
      }

      // now loop through shells with impact ionisation cross sections and apply Auger data that matches n, l values
      for (auto& collionrow : colliondata) {
        if (collionrow.Z == Z && collionrow.ionstage == ionstage && collionrow.n == n && collionrow.l == l) {
          printlog(
              "Z={:2} ionstage {:2} shellnum {} n {} l {} ionpot {:7.2f} E_A {:8.1f} E_A' {:8.1f} epsilon {:6} "
              "<n_Auger> {:5.1f} P(n_Auger)",
              Z, ionstage, shellnum, n, l, ionpot_ev, en_auger_ev_total_nocorrection, en_auger_ev, epsilon_e3,
              n_auger_elec_avg);

          const double prob_sum = std::ranges::fold_left(prob_num_auger, 0.0, std::plus<>());
          for (int a = 0; a <= NT_MAX_AUGER_ELECTRONS; a++) {
            printlog(" {}: {:4.2f}", a, prob_num_auger[a]);
          }
          assert_always(fabs(prob_sum - 1.0) < 0.001);

          printlnlog("");
          const bool found_existing_data = (collionrow.auger_g_accumulated > 0.);

          // keep existing data but update according to statistical weight represented by existing and new data
          const double oldweight = collionrow.auger_g_accumulated / (g + collionrow.auger_g_accumulated);
          const double newweight = g / (g + collionrow.auger_g_accumulated);
          collionrow.auger_g_accumulated += g;

          // update the statistical-weight averaged values
          collionrow.en_auger_ev = static_cast<float>((oldweight * collionrow.en_auger_ev) + (newweight * en_auger_ev));
          collionrow.n_auger_elec_avg =
              static_cast<float>((oldweight * collionrow.n_auger_elec_avg) + (newweight * n_auger_elec_avg));

          for (int a = 0; a <= NT_MAX_AUGER_ELECTRONS; a++) {
            collionrow.prob_num_auger[a] = (oldweight * collionrow.prob_num_auger[a]) + (newweight * prob_num_auger[a]);
          }
          const double prob_sum_final = std::ranges::fold_left(collionrow.prob_num_auger, 0.0, std::plus<>());
          assert_always(fabs(prob_sum_final - 1.0) < 0.001);

          if (found_existing_data) {
            printlog("  same NL shell already has data from another X-ray shell. New g-weighted values: P(n_Auger)");

            for (int a = 0; a <= NT_MAX_AUGER_ELECTRONS; a++) {
              printlog(" {}: {:4.2f}", a, collionrow.prob_num_auger[a]);
            }
            printlnlog("");
          }
        }
      }
    }
  }
}

auto get_sum_q_over_binding_energy(const int element, const int ion) -> double {
  const int ioncharge = get_ionstage(element, ion) - 1;
  const int nbound = get_atomicnumber(element) - ioncharge;  // number of bound electrons

  if (nbound <= 0) {
    return 0.;
  }

  // get the approximate shell occupancy if we don't have the data file
  const auto& shells_q = allions_shell_occupancies[get_uniqueionindex(element, ion)];
  const auto& binding_energies = elements_electron_binding.at(get_atomicnumber(element) - 1);

  double total = 0.;
  const auto num_shells = std::ssize(shells_q);
  for (int shellindex = 0; shellindex < num_shells; shellindex++) {
    const int electronsinshell = shells_q[shellindex];

    if (electronsinshell <= 0) {
      continue;
    }
    double enbinding = binding_energies.at(shellindex);
    if (enbinding <= 0) {
      // if we don't have the shell's binding energy, use the previous one
      enbinding = binding_energies.at(shellindex - 1);
      assert_always(enbinding > 0);
    }
    total += electronsinshell / std::max(get_ionpot(element, ion), enbinding);
  }

  return total;
}

void read_collion_data() {
  printlnlog("Reading collisional ionisation data from collion.txt...");

  auto cifile = fstream_required("collion.txt", std::ios::in);
  std::string line;
  get_noncommentline(cifile, line);
  std::istringstream ssline(line);
  int colliondatacount = 0;
  assert_always(ssline >> colliondatacount);
  printlnlog("Reading {} collisional transition rows", colliondatacount);
  assert_always(colliondatacount >= 0);

  for (int i = 0; i < colliondatacount; i++) {
    ShellParams collionrow{};
    int nelec = -1;
    get_noncommentline(cifile, line);
    assert_always(std::istringstream{line} >> collionrow.Z >> nelec >> collionrow.n >> collionrow.l >>
                  collionrow.ionpot_ev >> collionrow.A >> collionrow.B >> collionrow.C >> collionrow.D);

    assert_always(nelec > 0);
    collionrow.ionstage = collionrow.Z - nelec + 1;

    const int element = get_elementindex(collionrow.Z);
    if (element < 0 || collionrow.ionstage < get_ionstage(element, 0) ||
        collionrow.ionstage > get_ionstage(element, get_nions(element) - 1)) {
      continue;
    }

    std::ranges::fill(collionrow.prob_num_auger, 0.);
    collionrow.prob_num_auger[0] = 1.;

    collionrow.auger_g_accumulated = 0.;
    collionrow.en_auger_ev = 0.;
    collionrow.n_auger_elec_avg = 0.;

    colliondata.push_back(collionrow);
  }
  printlnlog("Stored {} of {} input shell cross sections", colliondata.size(), colliondatacount);
  for (int element = 0; element < get_nelements(); element++) {
    const int Z = get_atomicnumber(element);
    for (int ion = 0; ion < get_nions(element); ion++) {
      const int ionstage = get_ionstage(element, ion);
      const int ioncharge = ionstage - 1;
      const int nbound = Z - ioncharge;  // number of bound electrons
      if (nbound <= 0) {
        continue;
      }
      const bool any_data_matched = std::ranges::any_of(colliondata, [Z, ionstage](const ShellParams& collionrow) {
        return collionrow.Z == Z && collionrow.ionstage == ionstage;
      });
      if (!any_data_matched) {
        const double ionpot_ev = get_ionpot(element, ion) / EV;
        printlnlog(
            "No collisional ionisation data for Z={} ionstage {}. Using Lotz approximation with ionpot = {:g} eV", Z,
            ionstage, ionpot_ev);

        // get the approximate shell occupancy if we don't have the data file
        const auto& shells_q = allions_shell_occupancies[get_uniqueionindex(element, ion)];
        int electron_count = 0;
        const auto num_shells = std::ssize(shells_q);
        for (int shellindex = 0; shellindex < num_shells; shellindex++) {
          const int electronsinshell = shells_q.at(shellindex);

          electron_count += electronsinshell;

          if (electronsinshell <= 0) {
            continue;
          }
          double enbinding = elements_electron_binding.at(Z - 1).at(shellindex);
          const double ionpot = ionpot_ev * EV;
          if (enbinding <= 0) {
            // if we don't have the shell's binding energy, use the previous one
            enbinding = elements_electron_binding.at(Z - 1).at(shellindex - 1);
            assert_always(enbinding > 0);
          }
          const double p = std::max(ionpot, enbinding);
          ShellParams collionrow{};
          collionrow.Z = Z;
          collionrow.ionstage = ionstage;
          collionrow.n = -1;
          collionrow.l = -shellindex;
          collionrow.ionpot_ev = p / EV;
          collionrow.A = -1.;
          collionrow.B = -1.;
          collionrow.C = -1.;
          collionrow.D = -1.;
          std::ranges::fill(collionrow.prob_num_auger, 0.);
          collionrow.prob_num_auger[0] = 1.;
          collionrow.auger_g_accumulated = 0.;
          collionrow.en_auger_ev = 0.;
          collionrow.n_auger_elec_avg = 0.;

          colliondata.push_back(collionrow);
          if (electron_count >= nbound) {
            break;
          }
        }
      }
    }
  }
  colliondata.shrink_to_fit();
  std::ranges::stable_sort(colliondata, [](const ShellParams& a, const ShellParams& b) {
    return std::tie(a.Z, a.ionstage, a.ionpot_ev, a.n, a.l) < std::tie(b.Z, b.ionstage, b.ionpot_ev, b.n, b.l);
  });

  if (NT_MAX_AUGER_ELECTRONS > 0) {
    read_auger_data();
  }
}

auto get_possible_nt_excitation_count() -> int {
  // count the number of excitation transitions that pass the MAXNLEVELS_LOWER and MAXNLEVELS_UPPER conditions
  // this count might be higher than the number of stored ratecoeffs due to the MAX_NT_EXCITATIONS_STORED limit
  int ntexcitationcount = 0;
  for (int element = 0; element < get_nelements(); element++) {
    for (int ion = 0; ion < get_nions(element); ion++) {
      const int lower_nlevels = std::min(NTEXCITATION_MAXNLEVELS_LOWER, get_nlevels(element, ion));
      for (int lower = 0; lower < lower_nlevels; lower++) {
        const int nuptrans = get_nuptrans(element, ion, lower);
        const auto alltrans_startup = get_alltrans_startup(element, ion, lower);
        for (int t = 0; t < nuptrans; t++) {
          const int upper = globals::alltrans.targetlevelindex[alltrans_startup + t];
          if (upper < NTEXCITATION_MAXNLEVELS_UPPER) {
            ntexcitationcount++;
          }
        }
      }
    }
  }
  return ntexcitationcount;
}

void zero_all_effionpot(const ptrdiff_t nonemptymgi) {
  for (int uniqueionindex = 0; uniqueionindex < get_includedions(); uniqueionindex++) {
    auto& celliondata = get_cell_allions_data(nonemptymgi)[uniqueionindex];
    celliondata.eff_ionpot = 0.;

    std::ranges::fill(celliondata.prob_num_auger, 0.);
    celliondata.prob_num_auger[0] = 1.;
    std::ranges::fill(celliondata.ionenfrac_num_auger, 0.);
    celliondata.ionenfrac_num_auger[0] = 1.;

    const auto [element, ion] = get_ionfromuniqueionindex(uniqueionindex);
    assert_always(fabs(get_auger_probability(nonemptymgi, element, ion, 0) - 1.0) < 1e-3);
    assert_always(fabs(get_ion_auger_enfrac(nonemptymgi, element, ion, 0) - 1.0) < 1e-3);
  }
  check_auger_probabilities(nonemptymgi);
}

[[nodiscard]] constexpr auto get_energyindex_ev_lteq(const double energy_ev) -> int
// finds the highest energy point <= energy_ev
{
  const int index = std::floor((energy_ev - SF_EMIN) / DELTA_E);

  return std::clamp(index, 0, SFPTS - 1);
}

[[nodiscard]] constexpr auto get_energyindex_ev_gteq(const double energy_ev) -> int
// finds the highest energy point <= energy_ev
{
  const int index = std::ceil((energy_ev - SF_EMIN) / DELTA_E);

  return std::clamp(index, 0, SFPTS - 1);
}

// interpolate the y flux values to get the value at a given energy
// y has units of particles / cm2 / s / eV
[[nodiscard]] constexpr auto get_y(const std::array<double, SFPTS>& yfunc, const double energy_ev) -> double {
  if (energy_ev <= 0) {
    return 0.;
  }

  const int index = static_cast<int>((energy_ev - SF_EMIN) / DELTA_E);

  // assert_always(index > 0);
  if (index < 0) {
    // return 0.;
    assert_always(std::isfinite(yfunc[0]));
    return yfunc[0];
  }
  if (index >= SFPTS - 1) {
    return 0.;
  }
  const double enbelow = engrid(index);
  const double enabove = engrid(index + 1);
  const double ybelow = yfunc[index];
  const double yabove = yfunc[index + 1];
  const double x = (energy_ev - enbelow) / (enabove - enbelow);
  return ((1 - x) * ybelow) + (x * yabove);

  // or return the nearest neighbour
  // return yfunc[index];
}

constexpr auto xs_ionisation_lotz(const double en_erg, const ShellParams& colliondata_ion, const int electronsinshell)
    -> double {
  const double ionpot = colliondata_ion.ionpot_ev * EV;
  if (en_erg < ionpot) {
    return 0.;
  }
  // const double gamma = (en_erg / (ME * std::pow(CLIGHT, 2))) + 1;
  // const double beta = std::sqrt(1.0 - (1.0 / (std::pow(gamma, 2))));
  const double beta = std::sqrt(2 * en_erg / ME) / CLIGHT;

  const int ioncharge = colliondata_ion.ionstage - 1;
  const int nbound = colliondata_ion.Z - ioncharge;  // number of bound electrons

  if (nbound <= 0) {
    return 0.;
  }

  if (en_erg > ionpot) {
    // Equation 3.38 of Axelrod (1980) attributed to Lotz (1967)
    const double part_sigma_shell =
        (electronsinshell / ionpot *
         (std::log(pow2(beta) * ME * pow2(CLIGHT) / 2.0 / ionpot) - std::log10(1 - pow2(beta)) - pow2(beta)));
    if (part_sigma_shell > 0.) {
      constexpr double Aconst = 1.33e-14 * EV * EV;
      const double sigma = 2 * Aconst / pow2(beta) / ME / pow2(CLIGHT) * part_sigma_shell;
      assert_always(sigma >= 0);
      return sigma;
    }
  }

  return 0.;
}

auto get_xs_ionisation_vector_lotz(std::array<double, SFPTS>& xs_vec, const ShellParams& colliondata_ion) -> int {
  const int element = get_elementindex(colliondata_ion.Z);
  const int ion = colliondata_ion.ionstage - get_ionstage(element, 0);
  const int shellindex = -colliondata_ion.l;
  const int electronsinshell = allions_shell_occupancies[get_uniqueionindex(element, ion)][shellindex];

  const double ionpot_ev = colliondata_ion.ionpot_ev;
  const int startindex = get_energyindex_ev_gteq(ionpot_ev);
  std::fill_n(xs_vec.begin(), startindex, 0.);

  for (int i = startindex; i < SFPTS; i++) {
    xs_vec[i] = xs_ionisation_lotz(engrid(i) * EV, colliondata_ion, electronsinshell);
  }

  return startindex;
}

// xs_vec will be set with impact ionisation cross sections [cm2] for E > ionpot_ev (and zeros below this energy)
// returns the index of the first energy point >= ionpot_ev
auto get_xs_ionisation_vector(std::array<double, SFPTS>& xs_vec, const ShellParams& colliondata_ion) -> int {
  const double A = colliondata_ion.A;
  if (A < 0) {
    return get_xs_ionisation_vector_lotz(xs_vec, colliondata_ion);
  }

  const double ionpot_ev = colliondata_ion.ionpot_ev;
  const int startindex = get_energyindex_ev_gteq(ionpot_ev);

  std::fill_n(xs_vec.begin(), startindex, 0.);

  const double B = colliondata_ion.B;
  const double C = colliondata_ion.C;
  const double D = colliondata_ion.D;

  for (int i = startindex; i < SFPTS; i++) {
    const double u = engrid(i) / ionpot_ev;
    const double xs_ioniz =
        1e-14 * ((A * (1 - (1 / u))) + (B * std::pow((1 - (1 / u)), 2)) + (C * std::log(u)) + (D * std::log(u) / u)) /
        (u * std::pow(ionpot_ev, 2));
    xs_vec[i] = xs_ioniz;
  }

  return startindex;
}

// distribution of secondary electron energies for primary electron with energy e_p
// Opal, Peterson, & Beaty (1971)
[[nodiscard]] constexpr auto Psecondary(const double e_p, const double epsilon, const double I, const double J)
    -> double {
  const double e_s = epsilon - I;

  if (e_p <= I || e_s < 0.) {
    return 0.;
  }
  assert_testmodeonly(J > 0);
  assert_testmodeonly(e_p >= I);
  assert_testmodeonly(e_s >= 0);
  assert_testmodeonly(std::isfinite(std::atan((e_p - I) / 2 / J)));
  return 1 / (J * std::atan((e_p - I) / 2 / J) * (1 + pow2(e_s / J)));
}

[[nodiscard]] constexpr auto get_J(const int Z, const int ionstage, const double ionpot_ev) -> double {
  // returns an energy in eV
  // values from Opal et al. 1971 as applied by Kozma & Fransson 1992
  if (ionstage == 1) {
    if (Z == 2) {  // He I
      return 15.8;
    }
    if (Z == 10) {  // Ne I
      return 24.2;
    }
    if (Z == 18) {  // Ar I
      return 10.;
    }
  }

  return 0.6 * ionpot_ev;
}

// collisional excitation cross section in cm^2
// energies are in erg
constexpr auto xs_excitation(const int element, const int ion, const int lower, const int uptransindex,
                             const double epsilon_trans, const double lowerstatweight, const double energy) -> double {
  if (energy < epsilon_trans) {
    return 0.;
  }

  const auto alltrans_startup = get_alltrans_startup(element, ion, lower);
  const auto alltransindex = alltrans_startup + uptransindex;
  if (globals::alltrans.coll_str[alltransindex] >= 0) {
    // collision strength is available, so use it
    // Li et al. 2012 equation 11
    return pow2(H_ionpot / energy) / lowerstatweight * globals::alltrans.coll_str[alltransindex] * PI *
           A_naught_squared;
  }
  if (!globals::alltrans.forbidden[alltransindex]) {
    // permitted E1 electric dipole transitions
    const double U = energy / epsilon_trans;

    // constexpr double g_bar = 0.2;
    constexpr double A = 0.28;
    constexpr double B = 0.15;
    const double g_bar = (A * std::log(U)) + B;

    constexpr double prefactor = 45.585750051;  // 8 * pi^2/sqrt(3)
    // Eq 4 of Mewe 1972, possibly from Seaton 1962?
    return prefactor * A_naught_squared * pow2(H_ionpot / epsilon_trans) *
           globals::alltrans.osc_strength[alltransindex] * g_bar / U;
  }
  return 0.;
}

// -dE / dx for fast electrons
// energy is in ergs
// nne is the thermal electron density [cm^-3]
// return value has units of erg/cm
constexpr auto electron_loss_rate(const double energy, const double nne) -> double {
  if (energy <= 0.) {
    return 0;
  }

  // normally set to 1.0, but Shingles et al. (2021) boosted this to increase heating
  constexpr double boostfactor = 1.;

  const double omegap = std::sqrt(4 * PI * nne * pow2(QE) / ME);
  const double zetae = H * omegap / 2 / PI;
  if (energy > 14 * EV) {
    return boostfactor * nne * 2 * PI * pow4(QE) / energy * std::log(2 * energy / zetae);
  }
  const double v = std::sqrt(2 * energy / ME);
  return boostfactor * nne * 2 * PI * pow4(QE) / energy * std::log(ME * pow3(v) / (EULERGAMMA * pow2(QE) * omegap));
}

// impact ionisation cross section in cm^2
// energy and ionisation_potential should be in eV
// fitting formula of Younger 1981
// called Q_i(E) in KF92 equation 7
constexpr auto xs_impactionisation(const double energy_ev, const ShellParams& colliondata_ion) -> double {
  const double ionpot_ev = colliondata_ion.ionpot_ev;
  const double u = energy_ev / ionpot_ev;

  if (u <= 1.) {
    return 0;
  }
  const double A = colliondata_ion.A;
  if (A < 0) {
    const int element = get_elementindex(colliondata_ion.Z);
    const int ion = colliondata_ion.ionstage - get_ionstage(element, 0);
    const int shellindex = -colliondata_ion.l;
    const int electronsinshell = allions_shell_occupancies[get_uniqueionindex(element, ion)][shellindex];
    return xs_ionisation_lotz(energy_ev / EV, colliondata_ion, electronsinshell);
  }

  const double B = colliondata_ion.B;
  const double C = colliondata_ion.C;
  const double D = colliondata_ion.D;

  return 1e-14 * ((A * (1 - (1 / u))) + (B * pow2((1 - (1 / u)))) + (C * std::log(u)) + (D * std::log(u) / u)) /
         (u * pow2(ionpot_ev));
}

// Kozma & Fransson equation 6.
// Something related to a number of electrons, needed to calculate the heating fraction in equation 3
// not valid for energy > SF_EMIN
auto N_e(const int nonemptymgi, const double energy, const std::array<double, SFPTS>& yfunc) -> double {
  const double energy_ev = energy / EV;
  const double nnion_tot = get_nnion_tot(nonemptymgi);
  double N_e = 0.;

  for (int element = 0; element < get_nelements(); element++) {
    const int Z = get_atomicnumber(element);
    const int nions = get_nions(element);

    for (int ion = 0; ion < nions; ion++) {
      double N_e_ion = 0.;
      const int ionstage = get_ionstage(element, ion);
      const double nnion = get_nnion(nonemptymgi, element, ion);

      if (nnion < MIN_ION_OVER_NNTOT * nnion_tot) {  // skip negligible ions
        continue;
      }

      // excitation terms

      const int nlevels_all = get_nlevels(element, ion);
      const int nlevels = std::min(NTEXCITATION_MAXNLEVELS_LOWER, nlevels_all);

      for (int lower = 0; lower < nlevels; lower++) {
        const auto uniquelevelindex = get_uniquelevelindex(element, ion, lower);
        const double nnlevel = calculate_levelpop(nonemptymgi, element, ion, lower);
        const double epsilon_lower = epsilon(uniquelevelindex);
        const auto statweight_lower = stat_weight(uniquelevelindex);
        const int nuptrans = get_nuptrans(uniquelevelindex);
        const auto alltrans_startup = get_alltrans_startup(uniquelevelindex);
        for (int t = 0; t < nuptrans; t++) {
          const int upper = globals::alltrans.targetlevelindex[alltrans_startup + t];
          if (upper >= NTEXCITATION_MAXNLEVELS_UPPER) {
            continue;
          }
          const double epsilon_trans = epsilon(element, ion, upper) - epsilon_lower;
          const double epsilon_trans_ev = epsilon_trans / EV;
          N_e_ion += (nnlevel / nnion) * get_y(yfunc, energy_ev + epsilon_trans_ev) *
                     xs_excitation(element, ion, lower, t, epsilon_trans, statweight_lower, energy + epsilon_trans);
        }
      }

      // ionisation terms
      for (const auto& collionrow : colliondata) {
        if (collionrow.Z == Z && collionrow.ionstage == ionstage) {
          const double ionpot_ev = collionrow.ionpot_ev;
          const double J = get_J(Z, ionstage, ionpot_ev);
          const double lambda = std::min(SF_EMAX - energy_ev, energy_ev + ionpot_ev);

          const int integral1startindex = get_energyindex_ev_lteq(ionpot_ev);
          const int integral1stopindex = get_energyindex_ev_lteq(lambda);

          // integral from ionpot up to lambda
          for (int i = integral1startindex; i <= integral1stopindex; i++) {
            const double endash = engrid(i);

            N_e_ion += get_y(yfunc, energy_ev + endash) * xs_impactionisation(energy_ev + endash, collionrow) *
                       Psecondary(energy_ev + endash, endash, ionpot_ev, J) * DELTA_E;
          }

          // integral from 2E + I up to E_max
          const int integral2startindex = get_energyindex_ev_lteq((2 * energy_ev) + ionpot_ev);
          for (int i = integral2startindex; i < SFPTS; i++) {
            const double endash = engrid(i);
            N_e_ion += yfunc[i] * xs_impactionisation(endash, collionrow) *
                       Psecondary(endash, energy_ev + ionpot_ev, ionpot_ev, J) * DELTA_E;
          }
        }
      }

      N_e += nnion * N_e_ion;
    }
  }

  // source term, should be zero at the low end anyway
  N_e += sourcevec(get_energyindex_ev_lteq(energy_ev));

  assert_always(std::isfinite(N_e));
  return N_e;
}

// fraction of deposited energy that goes into heating the thermal electrons
// Kozma & Fransson equation 3
auto calculate_frac_heating(const int nonemptymgi, const std::array<double, SFPTS>& yfunc) -> float {
  // frac_heating multiplied by E_init, which will be divided out at the end
  double frac_heating_Einit = 0.;
  const float nne = grid::get_nne(nonemptymgi);

  for (int i = 0; i < SFPTS; i++) {
    const double endash = engrid(i);

    // first term
    frac_heating_Einit += yfunc[i] * (electron_loss_rate(endash * EV, nne) / EV) * DELTA_E;
  }

  // second term
  frac_heating_Einit += SF_EMIN * get_y(yfunc, SF_EMIN) * (electron_loss_rate(SF_EMIN * EV, nne) / EV);

  double N_e_contrib_Einit = 0.;
  // third term (integral from zero to SF_EMIN)
  constexpr int nsteps = (static_cast<int>(SF_EMIN / DELTA_E) + 1) * 10;
  static_assert(nsteps > 0);
  constexpr double delta_endash = SF_EMIN / nsteps;
  for (int j = 1; j < nsteps; j++) {
    const double endash = delta_endash * j;
    N_e_contrib_Einit += N_e(nonemptymgi, endash * EV, yfunc) * endash * delta_endash;
  }
  // there should not be more heating for E< EMIN than in the whole solution
  assert_always(N_e_contrib_Einit < frac_heating_Einit);
  frac_heating_Einit += N_e_contrib_Einit;

  const auto frac_heating = static_cast<float>(frac_heating_Einit / E_init_ev);

  assert_always(std::isfinite(frac_heating));
  assert_always(frac_heating >= 0.);
  assert_always(frac_heating <= 1.);

  return frac_heating;
}

// fraction of deposited energy that goes into ionisation
auto get_nt_frac_ionisation(const int nonemptymgi) -> float {
  if (!NT_ON) {
    return 0.;
  }
  if (!NT_SOLVE_SPENCERFANO) {
    return 0.03;  // Axelrod 1980 approximation
  }

  assert_always(nt_solution[nonemptymgi].frac_ionisation >= 0.);

  return nt_solution[nonemptymgi].frac_ionisation;
}

// fraction of deposited energy that goes into collisional excitation
auto get_nt_frac_excitation(const int nonemptymgi) -> float {
  if (!NT_ON || !NT_SOLVE_SPENCERFANO) {
    return 0.;
  }

  const float frac_excitation = nt_solution[nonemptymgi].frac_excitation;

  if (frac_excitation < 0 || !std::isfinite(frac_excitation)) {
    printlnlog("ERROR: get_nt_frac_excitation called with no valid solution stored for cell {}. frac_excitation = {:g}",
               grid::get_mgi_of_nonemptymgi(nonemptymgi), frac_excitation);
    std::abort();
  }

  return frac_excitation;
}

// compute the approximate work per ion pair without using the Spencer-Fano equation
// Makes use of EXTREMELY SIMPLE approximations - high energy limits only (can be used as an alternative to the
// Spencer-Fano solver)
// Luke (2026-02-20): This is so far from the Spencer-Fano eff_ionpot that I think there is a mistake in the calculation
auto get_oneoverw_approx_axelrod(const int element, const int ion, const int nonemptymgi) -> double {
  // Work in terms of 1/W since this is actually what we want. It is given by sigma/(Latom + Lelec).
  // We are going to start by taking all the high energy limits and ignoring Lelec, so that the
  // denominator is extremely simplified. Need to get the mean Z value.

  double Zbar = 0.;  // mass-weighted average atomic number
  for (int ielement = 0; ielement < get_nelements(); ielement++) {
    Zbar += grid::get_elem_abundance(nonemptymgi, ielement) * get_atomicnumber(ielement);
  }

  const double binding = get_sum_q_over_binding_energy(element, ion);
  constexpr double Aconst = 1.33e-14 * EV * EV;

  return Aconst * binding / Zbar / (2 * PI * std::pow(QE, 4));
}

// the fraction of deposited energy that goes into ionising electrons in a particular shell
auto calculate_nt_frac_ionisation_shell(const int nonemptymgi, const int element, const int ion,
                                        const ShellParams& collionrow, const std::array<double, SFPTS>& yfunc)
    -> double {
  const double nnion = get_nnion(nonemptymgi, element, ion);
  const double ionpot_ev = collionrow.ionpot_ev;

  std::array<double, SFPTS> cross_section_vec{};
  get_xs_ionisation_vector(cross_section_vec, collionrow);

  const double y_dot_crosssection_de =
      std::inner_product(yfunc.begin(), yfunc.end(), cross_section_vec.begin(), 0.0) * DELTA_E;

  return nnion * ionpot_ev * y_dot_crosssection_de / E_init_ev;
}

// non-thermal ionisation rate coefficient (multiply by population to get rate)
auto nt_ionisation_ratecoeff_wfapprox(const int nonemptymgi, const int element, const int ion) -> double {
  const double deposition_rate_density = get_deposition_rate_density(nonemptymgi);
  // to get the non-thermal ionisation rate we need to divide the energy deposited
  // per unit volume per unit time in the grid cell (sum of terms above)
  // by the total ion number density and the "work per ion pair"
  return deposition_rate_density / get_nnion_tot(nonemptymgi) * get_oneoverw_approx_axelrod(element, ion, nonemptymgi);
}

// Integrate the ionisation cross section over the electron degradation function to get the ionisation rate
// coefficient i.e. multiply this by ion population to get a rate of ionisations per second Do not call during packet
// propagation, as the y vector may not be in memory! IMPORTANT: we are dividing by the shell potential, not the
// valence potential here! To change this set assumeshellpotentialisvalence to true
auto calculate_nt_ionisation_ratecoeff(const int nonemptymgi, const int element, const int ion,
                                       const bool assumeshellpotentialisvalence, const std::array<double, SFPTS>& yfunc)
    -> double {
  std::array<double, SFPTS> cross_section_vec{};
  std::array<double, SFPTS> cross_section_vec_allshells{};

  const int Z = get_atomicnumber(element);
  const int ionstage = get_ionstage(element, ion);
  double ionpot_valence = -1;

  for (const auto& collionrow : colliondata) {
    if (collionrow.Z == Z && collionrow.ionstage == ionstage) {
      get_xs_ionisation_vector(cross_section_vec, collionrow);

      if (assumeshellpotentialisvalence) {
        const double ionpot_shell = collionrow.ionpot_ev * EV;
        if (ionpot_valence < 0) {
          ionpot_valence = ionpot_shell;
        }

        // ensure that the first shell really was the valence shell (we assumed ascending energy order)
        assert_always(ionpot_shell >= ionpot_valence);

        // boost the ionisation rate by assuming shell vacancy energy is used to eject valence electrons
        for (int i = 0; i < SFPTS; i++) {
          cross_section_vec[i] *= ionpot_shell / ionpot_valence;
        }
      }

      for (int i = 0; i < SFPTS; i++) {
        cross_section_vec_allshells[i] += cross_section_vec[i];
      }
    }
  }

  const double y_xs_de =
      std::inner_product(yfunc.begin(), yfunc.end(), cross_section_vec_allshells.begin(), 0.0) * DELTA_E;

  const double deposition_rate_density_ev = get_deposition_rate_density(nonemptymgi) / EV;
  const double yscalefactor = deposition_rate_density_ev / E_init_ev;

  return yscalefactor * y_xs_de;
}

// Kozma & Fransson 1992 equation 12, except modified to be a sum over all shells of an ion.
// the result is in [erg]
void calculate_eff_ionpot_auger_rates(const int nonemptymgi, const int element, const int ion,
                                      const std::array<double, SFPTS>& yfunc) {
  const int Z = get_atomicnumber(element);
  const int ionstage = get_ionstage(element, ion);
  const int uniqueionindex = get_uniqueionindex(element, ion);
  const double nnion = get_nnion(nonemptymgi, element, ion);  // ions/cm^3
  const double tot_nion = get_nnion_tot(nonemptymgi);
  const double X_ion = nnion / tot_nion;  // molar fraction of this ion

  // The ionisation rates of all shells of an ion add to make the ion's total ionisation rate,
  // i.e., Gamma_ion = Gamma_shell_a + Gamma_shell_b + ...
  // And since the ionisation rate is inversely proportional to the effective ion potential,
  // we solve:
  // (eta_ion / ionpot_ion) = (eta_shell_a / ionpot_shell_a) + (eta_shell_b / ionpot_shell_b) + ...
  // where eta is the fraction of the deposition energy going into ionisation of the ion or shell

  std::array<double, NT_MAX_AUGER_ELECTRONS + 1> eta_nauger_ionise_over_ionpot_sum{};
  std::array<double, NT_MAX_AUGER_ELECTRONS + 1> eta_nauger_ionise_sum{};
  auto& celliondata = get_cell_allions_data(nonemptymgi)[uniqueionindex];
  std::ranges::fill(celliondata.prob_num_auger, 0.);
  std::ranges::fill(celliondata.ionenfrac_num_auger, 0.);

  double eta_over_ionpot_sum = 0.;
  double eta_sum = 0.;
  double ionpot_valence = -1;
  int matching_nlsubshell_count = 0;
  for (const auto& collionrow : colliondata) {
    if (collionrow.Z == Z && collionrow.ionstage == ionstage) {
      matching_nlsubshell_count++;
      const double frac_ionisation_shell =
          calculate_nt_frac_ionisation_shell(nonemptymgi, element, ion, collionrow, yfunc);
      eta_sum += frac_ionisation_shell;
      const double ionpot_shell = collionrow.ionpot_ev * EV;

      if (ionpot_valence < 0) {
        ionpot_valence = ionpot_shell;
      }

      // ensure that the first shell really was the valence shell (we assumed ascending energy order)
      assert_always(ionpot_shell >= ionpot_valence);

      const double ionpot = NT_USE_VALENCE_IONPOTENTIAL ? ionpot_valence : ionpot_shell;
      const double eta_over_ionpot = frac_ionisation_shell / ionpot;  // this is proportional to rate

      eta_over_ionpot_sum += eta_over_ionpot;

      for (auto a = 0Z; a < std::ssize(eta_nauger_ionise_sum); a++) {
        eta_nauger_ionise_over_ionpot_sum[a] += eta_over_ionpot * collionrow.prob_num_auger[a];
        eta_nauger_ionise_sum[a] += frac_ionisation_shell * collionrow.prob_num_auger[a];
      }
    }
  }

  if (NT_MAX_AUGER_ELECTRONS > 0 && matching_nlsubshell_count > 0) {
    const int nions = get_nions(element);
    const int topion = nions - 1;
    if (ion < topion)  // don't try to ionise the top ion
    {
      for (int a = 0; a <= NT_MAX_AUGER_ELECTRONS; a++) {
        const int upperion = ion + 1 + a;
        if (upperion <= topion)  // not too many Auger electrons to exceed the top ion of this element
        {
          celliondata.prob_num_auger[a] =
              static_cast<float>(eta_nauger_ionise_over_ionpot_sum[a] / eta_over_ionpot_sum);
          celliondata.ionenfrac_num_auger[a] = static_cast<float>(eta_nauger_ionise_sum[a] / eta_sum);
        } else {
          // the following ensures that multiple ionisations can't send you to an ion stage that is not in
          // the model. Send it to the highest ion stage instead
          const int a_replace = topion - ion - 1;
          celliondata.prob_num_auger.at(a_replace) = static_cast<float>(
              celliondata.prob_num_auger.at(a_replace) + (eta_nauger_ionise_over_ionpot_sum[a] / eta_over_ionpot_sum));
          celliondata.ionenfrac_num_auger.at(a_replace) =
              static_cast<float>(celliondata.ionenfrac_num_auger.at(a_replace) + (eta_nauger_ionise_sum[a] / eta_sum));

          celliondata.prob_num_auger[a] = 0;
          celliondata.ionenfrac_num_auger[a] = 0.;
        }
      }
    }
  } else {
    const int a = 0;
    celliondata.prob_num_auger[a] = 1.;
    celliondata.ionenfrac_num_auger[a] = 1.;
  }

  if (matching_nlsubshell_count > 0) {
    double eff_ionpot = X_ion / eta_over_ionpot_sum;
    if (!std::isfinite(eff_ionpot)) {
      eff_ionpot = 0.;
    }
    celliondata.eff_ionpot = static_cast<float>(eff_ionpot);
  } else {
    printlnlog("WARNING! No matching subshells in NT impact ionisation cross section data for Z={} ionstage {}.",
               get_atomicnumber(element), get_ionstage(element, ion));
    printlnlog(
        "-> Defaulting to work function approximation and ionisation energy is not accounted for in Spencer-Fano "
        "solution.");

    celliondata.eff_ionpot = static_cast<float>(1. / get_oneoverw_approx_axelrod(element, ion, nonemptymgi));
  }
}

// get the effective ion potential from the stored value
// a value of 0. should be treated as invalid
auto get_eff_ionpot(const int nonemptymgi, const int element, const int ion) {
  return get_cell_allions_data(nonemptymgi)[get_uniqueionindex(element, ion)].eff_ionpot;
  // OR
  // return calculate_eff_ionpot(modelgridindex, element, ion);
}

// Kozma & Fransson 1992 equation 13
// returns the rate coefficient in s^-1
auto nt_ionisation_ratecoeff_sf(const int nonemptymgi, const int element, const int ion) -> double {
  const double deposition_rate_density = get_deposition_rate_density(nonemptymgi);
  if (deposition_rate_density > 0.) {
    return deposition_rate_density / get_nnion_tot(nonemptymgi) / get_eff_ionpot(nonemptymgi, element, ion);
    // alternatively, if the y vector is still in memory:
    // return calculate_nt_ionisation_ratecoeff(nonemptymgi, element, ion);
  }

  return 0.;
}

// vector of collisional excitation cross sections in cm^2
// epsilon_trans is in erg
// returns the index of the first valid cross section point (en >= epsilon_trans)
// all elements below this index are invalid and should not be used
auto get_xs_excitation_vector(const int alltransindex, const double statweight_lower, const double epsilon_trans)
    -> std::tuple<std::array<double, SFPTS>, int> {
  std::array<double, SFPTS> xs_excitation_vec{};
  if (globals::alltrans.coll_str[alltransindex] >= 0) {
    // collision strength is available, so use it
    // Li et al. 2012 equation 11
    const double constantfactor =
        std::pow(H_ionpot, 2) / statweight_lower * globals::alltrans.coll_str[alltransindex] * PI * A_naught_squared;

    const int en_startindex = get_energyindex_ev_gteq(epsilon_trans / EV);

    std::fill_n(xs_excitation_vec.begin(), en_startindex, 0.);

    for (int j = en_startindex; j < SFPTS; j++) {
      const double energy = engrid(j) * EV;
      xs_excitation_vec[j] = constantfactor * std::pow(energy, -2);
    }
    return {xs_excitation_vec, en_startindex};
  }
  if (!globals::alltrans.forbidden[alltransindex]) {
    const double trans_osc_strength = globals::alltrans.osc_strength[alltransindex];
    // permitted E1 electric dipole transitions

    // constexpr double g_bar = 0.2;
    constexpr double A = 0.28;
    constexpr double B = 0.15;

    constexpr double prefactor = 45.585750051;  // 8 * pi^2/sqrt(3)
    const double epsilon_trans_ev = epsilon_trans / EV;

    // Eq 4 of Mewe 1972, possibly from Seaton 1962?
    const double constantfactor =
        epsilon_trans_ev * prefactor * A_naught_squared * std::pow(H_ionpot / epsilon_trans, 2) * trans_osc_strength;

    const int en_startindex = get_energyindex_ev_gteq(epsilon_trans_ev);

    std::fill_n(xs_excitation_vec.begin(), en_startindex, 0.);

    // U = en / epsilon
    // g_bar = A * std::log(U) + b
    // xs[j] = constantfactor * g_bar / engrid(j)
    const double logepsilon = std::log(epsilon_trans_ev);
    for (int j = en_startindex; j < SFPTS; j++) {
      const double logU = logengrid[j] - logepsilon;
      const double g_bar = (A * logU) + B;
      xs_excitation_vec[j] = constantfactor * g_bar / engrid(j);
    }

    return {xs_excitation_vec, en_startindex};
  }

  return {xs_excitation_vec, -1};
}

// Kozma & Fransson equation 9 divided by level population and epsilon_trans
// returns the rate coefficient in s^-1 divided by deposition rate density in erg/cm^3/s
auto calculate_nt_excitation_ratecoeff_perdeposition(const std::array<double, SFPTS>& yvec, const int alltransindex,
                                                     const double statweight_lower, const double epsilon_trans)
    -> double {
  const auto [xs_excitation_vec, xsstartindex] =
      get_xs_excitation_vector(alltransindex, statweight_lower, epsilon_trans);

  if (xsstartindex >= 0) {
    double y_xs_de = 0.;
    for (int i = xsstartindex; i < SFPTS; i++) {
      y_xs_de += yvec[i] * xs_excitation_vec[i];
    }
    // multiply by DELTA_E to get the integral over the energy grid
    y_xs_de *= DELTA_E;

    return y_xs_de / E_init_ev / EV;
  }

  return 0.;
}

// returns the energy rate [erg/cm3/s] going toward non-thermal ionisation of lowerion
auto ion_ntion_energyrate(const int nonemptymgi, const int element, const int lowerion) -> double {
  const double nnlowerion = get_nnion(nonemptymgi, element, lowerion);
  double enrate = 0.;
  const auto maxupperion = nt_ionisation_maxupperion(element, lowerion);
  for (int upperion = lowerion + 1; upperion <= maxupperion; upperion++) {
    const double upperionprobfrac = nt_ionisation_upperion_probability(nonemptymgi, element, lowerion, upperion, false);
    // for (int lower = 0; lower < get_nlevels(element, lowerion); lower++)
    // {
    //   const double epsilon_trans = epsilon(element, upperion, 0) - epsilon(element, lowerion, lower);
    //   const double nnlower = get_cellcache_levelpop(nonemptymgi, element, lowerion, lower);
    //   enrate += nnlower * upperionprobfrac * epsilon_trans;
    // }
    const double epsilon_trans = epsilon(element, upperion, 0) - epsilon(element, lowerion, 0);
    enrate += nnlowerion * upperionprobfrac * epsilon_trans;
  }

  const double gamma_nt = nt_ionisation_ratecoeff(nonemptymgi, element, lowerion);
  return gamma_nt * enrate;
}

// returns the energy rate [erg/s] going toward non-thermal ionisation in a modelgrid cell
auto get_ntion_energyrate(const int nonemptymgi) -> double {
  double ratetotal = 0.;
  for (int ielement = 0; ielement < get_nelements(); ielement++) {
    const int nions = get_nions(ielement);
    for (int ilowerion = 0; ilowerion < nions - 1; ilowerion++) {
      ratetotal += ion_ntion_energyrate(nonemptymgi, ielement, ilowerion);
    }
  }
  return ratetotal;
}

auto select_nt_ionisation(const int nonemptymgi) -> std::tuple<int, int> {
  const double zrand = rng_uniform();

  // // select based on stored frac_deposition for each ion
  // double frac_deposition_ion_sum = 0.;
  // // zrand is between zero and frac_ionisation
  // // keep subtracting off deposition fractions of ionisations transitions until we hit the right one
  // // e.g. if zrand was less than frac_dep_trans1, then use the first transition
  // // e.g. if zrand was between frac_dep_trans1 and frac_dep_trans2 then use the second transition, etc
  // for (int allionindex = 0; allionindex < get_includedions(); allionindex++) {
  //   frac_deposition_ion_sum += nt_solution[nonemptymgi].fracdep_ionisation_ion[allionindex];
  //   if (frac_deposition_ion_sum >= zrand) {
  //     get_ionfromuniqueionindex(allionindex, element, lowerion);

  //     return;
  //   }
  // }
  // assert_always(false);  // should not reach here

  const double ratetotal = get_ntion_energyrate(nonemptymgi);

  // select based on the calculated energy going to ionisation for each ion
  double ratesum = 0.;
  for (int ielement = 0; ielement < get_nelements(); ielement++) {
    const int nions = get_nions(ielement);
    for (int ilowerion = 0; ilowerion < nions - 1; ilowerion++) {
      ratesum += ion_ntion_energyrate(nonemptymgi, ielement, ilowerion);
      if (ratesum >= zrand * ratetotal) {
        return {ielement, ilowerion};
      }
    }
  }
  assert_always(false);
  return {-1, -1};
}

void analyse_sf_solution(const int nonemptymgi, const int timestep, const bool enable_sfexcitation,
                         const std::array<double, SFPTS>& yfunc) {
  const auto nne = grid::get_nne(nonemptymgi);
  const auto nntot = get_nnion_tot(nonemptymgi);
  const auto nnetot = grid::get_nnetot(nonemptymgi);

  double frac_excitation_total = 0.;
  double frac_ionisation_total = 0.;

  // temporary storage of the full excitation list for current cell before possible truncation and copying to
  // node-shared memory
  THREADLOCALONHOST std::vector<NonThermalExcitation> tmp_excitation_list;
  tmp_excitation_list.clear();

  for (int element = 0; element < get_nelements(); element++) {
    const int Z = get_atomicnumber(element);
    const int nions = get_nions(element);
    for (int ion = 0; ion < nions; ion++) {
      const int uniqueionindex = get_uniqueionindex(element, ion);

      const int ionstage = get_ionstage(element, ion);
      const int ioncharge = ionstage - 1;
      const int nbound = Z - ioncharge;  // number of bound electrons
      if (nbound <= 0) {
        continue;
      }
      const double nnion = get_nnion(nonemptymgi, element, ion);

      // if (nnion < minionfraction * get_nnion_tot(nonemptymgi)) // skip negligible ions
      if (nnion <= 0.) {  // skip zero-abundance ions
        continue;
      }

      double frac_ionisation_ion = 0.;
      double frac_excitation_ion = 0.;
      printlnlog("  Z={} ionstage {}:", Z, ionstage);
      printlnlog("    nnion/nntot: {:g}", nnion / nntot);

      calculate_eff_ionpot_auger_rates(nonemptymgi, element, ion, yfunc);

      int matching_subshell_count = 0;
      for (const auto& collionrow : colliondata) {
        if (collionrow.Z == Z && collionrow.ionstage == ionstage) {
          const double frac_ionisation_ion_shell =
              calculate_nt_frac_ionisation_shell(nonemptymgi, element, ion, collionrow, yfunc);
          frac_ionisation_ion += frac_ionisation_ion_shell;
          matching_subshell_count++;
          printlog("      shell ");
          if (collionrow.n >= 0) {
            printlog("n {}, l {}", collionrow.n, collionrow.l);
          } else {
            printlog("{} (Lotz)", shellnames.at(-collionrow.l));
          }
          printlog(" I {:5.1f} eV: frac_ionisation {:10.4e}", collionrow.ionpot_ev, frac_ionisation_ion_shell);

          if (NT_MAX_AUGER_ELECTRONS > 0) {
            printlog("  prob(n Auger elec):");
            for (int a = 0; a <= NT_MAX_AUGER_ELECTRONS; a++) {
              printlog(" {}: {:.2f}", a, collionrow.prob_num_auger[a]);
            }
          }
          printlnlog("");
        }
      }

      // do not ionise the top ion
      if (ion < nions - 1) {
        get_cell_allions_data(nonemptymgi)[uniqueionindex].fracdep_ionisation_ion = frac_ionisation_ion;

        frac_ionisation_total += frac_ionisation_ion;
      } else {
        get_cell_allions_data(nonemptymgi)[uniqueionindex].fracdep_ionisation_ion = 0.;
      }
      printlnlog("    frac_ionisation: {:g} ({} subshells)", frac_ionisation_ion, matching_subshell_count);

      // excitation from all levels is very SLOW
      const int nlevels_all = get_nlevels(element, ion);
      // So limit the lower levels to improve performance
      int nlevels = (nlevels_all > NTEXCITATION_MAXNLEVELS_LOWER) ? NTEXCITATION_MAXNLEVELS_LOWER : nlevels_all;
      if (!enable_sfexcitation) {
        nlevels = -1;  // disable all excitations
      }
      const bool above_minionfraction = (nnion >= MIN_ION_OVER_NNTOT * get_nnion_tot(nonemptymgi));

      for (int lower = 0; lower < nlevels; lower++) {
        const auto uniquelevelindex = get_uniquelevelindex(element, ion, lower);
        const double statweight_lower = stat_weight(uniquelevelindex);
        const auto alltrans_startup = get_alltrans_startup(uniquelevelindex);
        const int nuptrans = get_nuptrans(uniquelevelindex);
        const double nnlevel = calculate_levelpop(nonemptymgi, element, ion, lower);
        const double epsilon_lower = epsilon(uniquelevelindex);

        for (int t = 0; t < nuptrans; t++) {
          const int alltransindex = alltrans_startup + t;
          const int upper = globals::alltrans.targetlevelindex[alltransindex];
          if (upper >= NTEXCITATION_MAXNLEVELS_UPPER) {
            continue;
          }

          const double epsilon_trans = epsilon(element, ion, upper) - epsilon_lower;
          const double ratecoeffperdeposition =
              calculate_nt_excitation_ratecoeff_perdeposition(yfunc, alltransindex, statweight_lower, epsilon_trans);
          const double frac_excitation_thistrans = nnlevel * epsilon_trans * ratecoeffperdeposition;
          frac_excitation_ion += frac_excitation_thistrans;

          if constexpr (NT_EXCITATION_ON) {
            assert_always(std::isfinite(ratecoeffperdeposition));
            // the atomic data set was limited for Fe V, which caused the ground multiplet to be massively
            // depleted, and then almost no recombination happened!
            if (above_minionfraction && ratecoeffperdeposition > 0 && (Z != 26 || ionstage != 5)) {
              // if (get_coll_str(lineindex) < 0) // if collision strength is not defined, the rate coefficient is
              // unreliable
              //   ratecoeffperdeposition = 0.;

              tmp_excitation_list.push_back({
                  .frac_deposition = frac_excitation_thistrans,
                  .ratecoeffperdeposition = ratecoeffperdeposition,
                  .alltransindex = alltransindex,
              });
            }
          }  // NT_EXCITATION_ON
        }  // for t
      }  // for lower

      printlnlog("    frac_excitation: {:g}", frac_excitation_ion);
      if (frac_excitation_ion > 1. || !std::isfinite(frac_excitation_ion)) {
        printlnlog("      WARNING: invalid frac_excitation. Replacing with zero");
        frac_excitation_ion = 0.;
      }
      frac_excitation_total += frac_excitation_ion;
      printlnlog("    approxworkfn: {:9.2f} eV  (without using the Spencer-Fano solution)",
                 (1. / get_oneoverw_approx_axelrod(element, ion, nonemptymgi)) / EV);
      printlnlog("    eff_ionpot:   {:9.2f} eV  (always use valence potential is {})",
                 get_eff_ionpot(nonemptymgi, element, ion) / EV, (NT_USE_VALENCE_IONPOTENTIAL ? "true" : "false"));

      printlnlog("    approxworkfn Gamma:      {:9.3e}", nt_ionisation_ratecoeff_wfapprox(nonemptymgi, element, ion));

      if constexpr (NT_USE_VALENCE_IONPOTENTIAL) {
        printlnlog("    SF integral Gamma:       {:9.3e} (alternative if NT_USE_VALENCE_IONPOTENTIAL was disabled)",
                   calculate_nt_ionisation_ratecoeff(nonemptymgi, element, ion, false, yfunc));
      } else {
        printlnlog("    SF integral(I=Iv) Gamma: {:9.3e}  (alternative if NT_USE_VALENCE_IONPOTENTIAL was enabled)",
                   calculate_nt_ionisation_ratecoeff(nonemptymgi, element, ion, true, yfunc));
      }

      printlnlog("    ARTIS using Gamma:       {:9.3e}", nt_ionisation_ratecoeff(nonemptymgi, element, ion));

      // the ion values (unlike shell ones) have been collapsed down to ensure that upperion < nions
      if (ion < nions - 1) {
        printlog("    probability to ionstage:");
        double prob_sum = 0.;
        for (int upperion = ion + 1; upperion <= nt_ionisation_maxupperion(element, ion); upperion++) {
          const double probability = nt_ionisation_upperion_probability(nonemptymgi, element, ion, upperion, false);
          prob_sum += probability;
          if (probability > 0.) {
            printlog(" {}: {:.3f}", get_ionstage(element, upperion), probability);
          }
        }
        printlnlog("");
        assert_always((fabs(prob_sum - 1.0) <= 1e-2) ||
                      (nt_ionisation_ratecoeff_sf(nonemptymgi, element, ion) < 1e-20));

        printlog("         enfrac to ionstage:");
        double enfrac_sum = 0.;
        for (int upperion = ion + 1; upperion <= nt_ionisation_maxupperion(element, ion); upperion++) {
          const double probability = nt_ionisation_upperion_probability(nonemptymgi, element, ion, upperion, true);
          enfrac_sum += probability;
          if (probability > 0.) {
            printlog(" {}: {:.3f}", get_ionstage(element, upperion), probability);
          }
        }
        printlnlog("");
        assert_always(fabs(enfrac_sum - 1.0) <= 1e-2 ||
                      (nt_ionisation_ratecoeff_sf(nonemptymgi, element, ion) < 1e-20));
      }
    }
  }

  if (nt_excitations_stored > 0) {
    // sort by descending frac_deposition
    std::ranges::SORT_OR_STABLE_SORT(tmp_excitation_list, std::ranges::greater{},
                                     &NonThermalExcitation::frac_deposition);

    // the excitation list is now sorted by frac_deposition descending
    const double deposition_rate_density = get_deposition_rate_density(nonemptymgi);

    if (std::ssize(tmp_excitation_list) > nt_excitations_stored) {
      // truncate the sorted list to save memory
      printlnlog("  Truncating non-thermal excitation list from {} to {} transitions.", tmp_excitation_list.size(),
                 nt_excitations_stored);
      tmp_excitation_list.resize(nt_excitations_stored);
    }

    nt_solution[nonemptymgi].frac_excitations_list_size = static_cast<int>(std::ssize(tmp_excitation_list));
    std::ranges::copy(tmp_excitation_list, get_cell_ntexcitations(nonemptymgi).begin());

    const auto T_e = grid::get_Te(nonemptymgi);
    printlnlog("  Top non-thermal excitation fractions (total excitations = {}):",
               nt_solution[nonemptymgi].frac_excitations_list_size);
    const int ntransdisplayed = std::min(10, nt_solution[nonemptymgi].frac_excitations_list_size);

    for (int excitationindex = 0; excitationindex < ntransdisplayed; excitationindex++) {
      const auto& ntexc = tmp_excitation_list[excitationindex];
      if (ntexc.frac_deposition > 0.) {
        const auto alltransindex = ntexc.alltransindex;
        const auto lineindex = globals::alltrans.lineindex[alltransindex];
        const int element = globals::linelist.elementindex[lineindex];
        const int ion = globals::linelist.ionindex[lineindex];
        const int lower = globals::linelist.lowerlevelindex[lineindex];
        const int upper = globals::linelist.upperlevelindex[lineindex];
        const auto lower_uniquelevelindex = get_uniquelevelindex(element, ion, lower);
        const auto upper_uniquelevelindex = get_uniquelevelindex(element, ion, upper);
        const auto nnlevel_lower = calculate_levelpop(nonemptymgi, element, ion, lower);
        const auto nnlevel_upper = calculate_levelpop(nonemptymgi, element, ion, upper);
        const auto statweight_lower = stat_weight(lower_uniquelevelindex);

        const double epsilon_trans = epsilon(upper_uniquelevelindex) - epsilon(lower_uniquelevelindex);

        const double ntcollexc_ratecoeff = ntexc.ratecoeffperdeposition * deposition_rate_density;
        const auto statweight_upper = stat_weight(upper_uniquelevelindex);

        const double t_mid = globals::timesteps[timestep].mid;
        const double radexc_ratecoeff = rad_excitation_ratecoeff(
            nonemptymgi, statweight_upper, globals::alltrans.einstein_A[alltransindex], epsilon_trans, nnlevel_lower,
            nnlevel_upper, statweight_lower, alltransindex, t_mid);

        const double collexc_ratecoeff =
            col_excitation_ratecoeff(T_e, nne, statweight_upper, alltransindex, epsilon_trans, statweight_lower);

        const double exc_ratecoeff = radexc_ratecoeff + collexc_ratecoeff + ntcollexc_ratecoeff;
        const auto coll_str = globals::alltrans.coll_str[alltransindex];

        printlnlog(
            "    frac_deposition {:.3e} Z={:2} ionstage {} lower {:4} upper {:4} rad_exc {:.1e} coll_exc {:.1e} nt_exc "
            "{:.1e} nt/tot {:.1e} collstr {:.1e} lineindex {}",
            ntexc.frac_deposition, get_atomicnumber(element), get_ionstage(element, ion), lower, upper,
            radexc_ratecoeff, collexc_ratecoeff, ntcollexc_ratecoeff, ntcollexc_ratecoeff / exc_ratecoeff, coll_str,
            lineindex);
      }
    }

    // sort the excitation list by ascending lineindex for fast lookup with a binary search
    std::ranges::SORT_OR_STABLE_SORT(get_cell_ntexcitations(nonemptymgi), std::ranges::less{},
                                     &NonThermalExcitation::alltransindex);

  }  // NT_EXCITATION_ON

  // calculate number density of non-thermal electrons
  const double deposition_rate_density_ev = get_deposition_rate_density(nonemptymgi) / EV;
  const double yscalefactor = deposition_rate_density_ev / E_init_ev;

  double nne_nt_max = 0.;
  for (int i = 0; i < SFPTS; i++) {
    const double endash = engrid(i);
    const double delta_endash = DELTA_E;
    const double oneovervelocity = std::sqrt(9.10938e-31 / 2 / endash / 1.60218e-19) / 100;  // in sec/cm
    nne_nt_max += yscalefactor * yfunc[i] * oneovervelocity * delta_endash;
  }

  nt_solution[nonemptymgi].frac_excitation = static_cast<float>(frac_excitation_total);
  nt_solution[nonemptymgi].frac_ionisation = static_cast<float>(frac_ionisation_total);

  printlnlog("  deposition:  {:9.2f} eV/s/cm^3", deposition_rate_density_ev);
  printlnlog("  nne:         {:9.3e} e-/cm^3", nne);
  printlnlog("  nnetot:      {:9.3e} e-/cm^3", nnetot);
  printlnlog("  nne_nt     < {:9.3e} e-/cm^3", nne_nt_max);
  printlnlog("  nne_nt/nne < {:9.3e}", nne_nt_max / nne);

  // store the solution properties now while the NT spectrum is in memory (in case we free before packet prop)
  nt_solution[nonemptymgi].frac_heating = calculate_frac_heating(nonemptymgi, yfunc);

  printlnlog("  frac_heating_tot:    {:g}", nt_solution[nonemptymgi].frac_heating);
  printlnlog("  frac_excitation_tot: {:g}", frac_excitation_total);
  printlnlog("  frac_ionisation_tot: {:g}", frac_ionisation_total);
  const double frac_sum = nt_solution[nonemptymgi].frac_heating + frac_excitation_total + frac_ionisation_total;
  printlnlog("  frac_sum:            {:g} (should be close to 1.0)", frac_sum);

  nt_solution[nonemptymgi].frac_heating = static_cast<float>(1. - frac_excitation_total - frac_ionisation_total);
  printlnlog("  (replacing calculated frac_heating_tot with {:g} to make frac_sum = 1.0)",
             nt_solution[nonemptymgi].frac_heating);
}

void sfmatrix_add_excitation(std::span<double> sfmatrixuppertri, const int nonemptymgi, const int element,
                             const int ion) {
  // excitation terms

  const int nlevels_all = get_nlevels(element, ion);
  const int nlevels = (nlevels_all > NTEXCITATION_MAXNLEVELS_LOWER) ? NTEXCITATION_MAXNLEVELS_LOWER : nlevels_all;

  const auto lowers = std::ranges::iota_view{0, nlevels};
  std::for_each(lowers.begin(), lowers.end(), [&](const int lower) {
    const auto uniquelevelindex = get_uniquelevelindex(element, ion, lower);
    const double statweight_lower = stat_weight(uniquelevelindex);
    const double nnlevel = calculate_levelpop(nonemptymgi, element, ion, lower);
    const double epsilon_lower = epsilon(uniquelevelindex);
    const auto alltrans_startup = get_alltrans_startup(uniquelevelindex);
    const int nuptrans = get_nuptrans(uniquelevelindex);
    for (int t = 0; t < nuptrans; t++) {
      const int alltransindex = alltrans_startup + t;
      const int upper = globals::alltrans.targetlevelindex[alltransindex];
      if (upper >= NTEXCITATION_MAXNLEVELS_UPPER) {
        continue;
      }
      const double epsilon_trans = epsilon(element, ion, upper) - epsilon_lower;
      const double epsilon_trans_ev = epsilon_trans / EV;
      if (epsilon_trans_ev < SF_EMIN) {
        continue;
      }

      auto [vec_xs_excitation_deltae, xsstartindex] =
          get_xs_excitation_vector(alltransindex, statweight_lower, epsilon_trans);
      if (xsstartindex >= 0) {
        for (int j = xsstartindex; j < SFPTS; j++) {
          vec_xs_excitation_deltae[j] *= DELTA_E;
        }

        for (int i = 0; i < SFPTS; i++) {
          const int rowoffset = uppertriangular(i, 0);
          const double en = engrid(i);
          const int stopindex = get_energyindex_ev_lteq(en + epsilon_trans_ev);

          const int startindex = std::max(i, xsstartindex);
          for (int j = startindex; j < stopindex; j++) {
            atomicadd(sfmatrixuppertri[rowoffset + j], nnlevel * vec_xs_excitation_deltae[j]);
          }

          // do the last bit separately because we're not using the full delta_e interval
          const double delta_en_actual = (en + epsilon_trans_ev - engrid(stopindex));
          atomicadd(sfmatrixuppertri[rowoffset + stopindex],
                    nnlevel * vec_xs_excitation_deltae[stopindex] * delta_en_actual / DELTA_E);
        }
      }
    }
  });
}

void sfmatrix_add_ionisation(std::span<double> sfmatrixuppertri, const int Z, const int ionstage, const double nnion)
// add the ionisation terms to the Spencer-Fano matrix
{
  std::array<double, SFPTS> vec_xs_ionisation{};
  for (const auto& collionrow : colliondata) {
    if (collionrow.Z == Z && collionrow.ionstage == ionstage) {
      const double ionpot_ev = collionrow.ionpot_ev;
      const double en_auger_ev = collionrow.en_auger_ev;
      // const double n_auger_elec_avg = colliondata[n].n_auger_elec_avg;
      const double J = get_J(Z, ionstage, ionpot_ev);

      assert_always(ionpot_ev >= SF_EMIN);

      const int xsstartindex = get_xs_ionisation_vector(vec_xs_ionisation, collionrow);
      assert_always(xsstartindex >= 0);
      // Luke Shingles: the use of min and max on the epsilon limits keeps energies
      // from becoming unphysical. This insight came from reading the
      // CMFGEN Fortran source code (Li, Dessart, Hillier 2012, doi:10.1111/j.1365-2966.2012.21198.x)
      // I had neglected this, so the limits of integration were incorrect. The fix didn't massively affect
      // ionisation rates or spectra, but it was a source of error that led to energy fractions not adding up to
      // 100%
      std::array<double, SFPTS> int_eps_upper = {0};
      std::array<double, SFPTS> prefactors = {0};
      for (int j = xsstartindex; j < SFPTS; j++) {
        const double endash = engrid(j);
        const double epsilon_upper = std::min((endash + ionpot_ev) / 2, endash);
        int_eps_upper[j] = std::atan((epsilon_upper - ionpot_ev) / J);
        prefactors[j] = vec_xs_ionisation[j] * nnion / std::atan((endash - ionpot_ev) / 2 / J);
      }

      for (int i = 0; i < SFPTS; i++) {
        // i is the matrix row index, which corresponds to an energy E at which we are solve from y(E)
        const double en = engrid(i);
        const int rowoffset = uppertriangular(i, 0);

        // endash ranges from en to SF_EMAX, but skip over the zero-cross section points
        const int jstart = std::max(i, xsstartindex);
        for (int j = jstart; j < SFPTS; j++) {
          // j is the matrix column index which corresponds to the piece of the integral at y(E') where E' >= E and E'
          // = engrid(j)
          const double endash = engrid(j);

          // J * atan[(epsilon - ionpot_ev) / J] is the indefinite integral of 1/[1 + (epsilon - ionpot_ev)^2/ J^2]
          // in Kozma & Fransson 1992 equation 4

          const double epsilon_lower =
              std::max(endash - en, ionpot_ev);  // and epsilon_upper = (endash + ionpot_ev) / 2;
          const double int_eps_lower = std::atan((epsilon_lower - ionpot_ev) / J);
          if (int_eps_lower <= int_eps_upper[j]) {
            sfmatrixuppertri[rowoffset + j] += prefactors[j] * (int_eps_upper[j] - int_eps_lower) * DELTA_E;
          }
        }

        // below is std::atan((epsilon_lower - ionpot_ev) / J) where epsilon_lower = en + ionpot_ev;
        const double int_eps_lower2 = std::atan(en / J);

        // endash ranges from 2 * en + ionpot_ev to SF_EMAX
        if ((2 * en) + ionpot_ev <= SF_EMAX) {
          const int secondintegralstartindex = std::max(xsstartindex, get_energyindex_ev_lteq((2 * en) + ionpot_ev));
          for (int j = secondintegralstartindex; j < SFPTS; j++) {
            // epsilon_lower = en + ionpot_ev;
            // epsilon_upper = (endash + ionpot_ev) / 2;
            if (int_eps_lower2 <= int_eps_upper[j]) {
              sfmatrixuppertri[rowoffset + j] -= prefactors[j] * (int_eps_upper[j] - int_eps_lower2) * DELTA_E;
            }
          }
        }
      }

      if constexpr (SF_AUGER_CONTRIBUTION_ON) {
        int augerstopindex = 0;
        if constexpr (SF_AUGER_CONTRIBUTION_DISTRIBUTE_EN) {
          // en_auger_ev is (if LJS understands it correctly) averaged to include some probability of zero Auger
          // electrons so we need a boost to get the average energy of Auger electrons given that there are one or
          // more
          const double en_boost = 1 / (1. - collionrow.prob_num_auger[0]);

          augerstopindex = get_energyindex_ev_gteq(en_auger_ev * en_boost);
        } else {
          augerstopindex = get_energyindex_ev_gteq(en_auger_ev);
        }

        for (int i = 0; i < augerstopindex; i++) {
          const int rowoffset = uppertriangular(i, 0);
          const double en = engrid(i);
          const int jstart = std::max(i, xsstartindex);
          for (int j = jstart; j < SFPTS; j++) {
            const double xs = vec_xs_ionisation[j];
            if constexpr (SF_AUGER_CONTRIBUTION_DISTRIBUTE_EN) {
              const double en_boost = 1 / (1. - collionrow.prob_num_auger[0]);
              for (int a = 1; a <= NT_MAX_AUGER_ELECTRONS; a++) {
                if (en < (en_auger_ev * en_boost / a)) {
                  sfmatrixuppertri[rowoffset + j] -= nnion * xs * collionrow.prob_num_auger[a] * a;
                }
              }
            } else {
              assert_always(en < en_auger_ev);
              sfmatrixuppertri[rowoffset + j] -= nnion * xs;  // * n_auger_elec_avg; // * en_auger_ev???
            }
          }
        }
      }
    }
  }
}

// solve the Spencer-Fano matrix equation and return the y vector (samples of the Spencer-Fano solution function).
// Multiply y by energy interval [eV] to get non-thermal electron number flux. y(E) * dE is the flux of electrons with
// energy in the range (E, E + dE) in units of particles/cm2/s. y has units of particles/cm2/s/eV
auto sfmatrix_solve(const std::span<const double> sfmatrix) -> std::array<double, SFPTS> {
  // solve the matrix-vector equation sfmatrix * yvec = rhsvec for yvec

  THREADLOCALONHOST std::array<double, SFPTS> yvec_arr{};

#ifdef EIGEN_OFF

  auto gsl_yvec = gsl_vector_view_array(yvec_arr.data(), SFPTS).vector;
  const auto gsl_rhsvec = gsl_vector_const_view_array(rhsvec.data(), SFPTS).vector;

  THREADLOCALONHOST std::array<size_t, SFPTS> vec_permutation{};
  gsl_permutation p{.size = SFPTS, .data = vec_permutation.data()};
  gsl_permutation_init(&p);

  const auto gsl_sfmatrix = gsl_matrix_const_view_array(sfmatrix.data(), SFPTS, SFPTS).matrix;

  // sfmatrix is already in upper triangular form
  const auto& gsl_sfmatrix_LU = gsl_sfmatrix;

  // solve matrix equation: sf_matrix * y_vec = rhsvec for yvec
  gsl_linalg_LU_solve(&gsl_sfmatrix_LU, &p, &gsl_rhsvec, &gsl_yvec);

#else

  const Eigen::Map<const Eigen::Vector<double, SFPTS>> eigen_rhsvec{rhsvec.data()};
  const Eigen::Map<const Eigen::Matrix<double, SFPTS, SFPTS, Eigen::RowMajor>> eigen_sfmatrix{sfmatrix.data()};
  assert_testmodeonly(eigen_sfmatrix.isUpperTriangular());

  auto eigen_sfmatrix_LU = eigen_sfmatrix.triangularView<Eigen::Upper>();
  Eigen::Map<Eigen::Vector<double, SFPTS>> eigen_yvec{yvec_arr.data()};
  eigen_yvec = eigen_sfmatrix_LU.solve(eigen_rhsvec);

#endif

  // refine the solution

  THREADLOCALONHOST std::array<double, SFPTS> yvec_best{};
#ifdef EIGEN_OFF
  THREADLOCALONHOST std::array<double, SFPTS> residual_vec{};
  auto gsl_residual_vec = gsl_vector_view_array(residual_vec.data(), SFPTS).vector;
#else
  THREADLOCALONHOST Eigen::Vector<double, SFPTS> eigen_residual_vec{};
#endif

  double error_best = -1.;
  int iteration = 0;
  for (iteration = 0; iteration < 10; iteration++) {
#ifdef EIGEN_OFF
    if (iteration > 0) {
      // use gsl_vec_residual as the temporary work vector
      gsl_linalg_LU_refine(&gsl_sfmatrix, &gsl_sfmatrix_LU, &p, &gsl_rhsvec, &gsl_yvec, &gsl_residual_vec);
    }
    // residual = sfmatrix * yvec - rhsvec
    std::ranges::copy(rhsvec, residual_vec.begin());
    gsl_blas_dgemv(CblasNoTrans, 1.0, &gsl_sfmatrix, &gsl_yvec, -1.0, &gsl_residual_vec);
    // error is the largest absolute residual element
    const double error = fabs(gsl_vector_get(&gsl_residual_vec, gsl_blas_idamax(&gsl_residual_vec)));
#else
    if (iteration > 0) {
      eigen_yvec += eigen_sfmatrix_LU.solve(eigen_residual_vec);
    }
    eigen_residual_vec = eigen_rhsvec - eigen_sfmatrix * eigen_yvec;
    const double error = eigen_residual_vec.cwiseAbs().maxCoeff();
#endif

    if (error < error_best || error_best < 0.) {
      std::ranges::copy(yvec_arr, yvec_best.begin());
      error_best = error;
    }
    // printlnlog("  SF solver LU_refine: After {} iterations, solution vector has a max residual of {:g}", iteration,
    //            error);
  }

  assert_always(error_best >= 0.);
  std::ranges::copy(yvec_best, yvec_arr.begin());

  if (error_best > 1e-10) {
    printlnlog("  SF solver LU_refine: After {} iterations, best solution vector has a max residual of {:g}", iteration,
               error_best);
  }

  assert_always(std::ranges::all_of(yvec_arr, [](double y) { return y >= 0.; }));

  return yvec_arr;
}

}  // anonymous namespace

void init() {
  const ptrdiff_t nonempty_npts_model = grid::get_nonempty_npts_model();

  deposition_rate_density_all_cells = MPI_shared_malloc_span<double>(nonempty_npts_model);

  if (globals::rank_in_node == 0) {
    std::ranges::fill(deposition_rate_density_all_cells, -1.);
  }

  if (!NT_ON) {
    return;
  }

  read_binding_energies();

  if (!NT_SOLVE_SPENCERFANO) {
    return;
  }

  printlnlog("Initializing non-thermal solver with:");
  printlnlog("  NT_EXCITATION {}", NT_EXCITATION_ON ? "on" : "off");
  printlnlog("  MAX_NT_EXCITATIONS_STORED {}", MAX_NT_EXCITATIONS_STORED);
  printlnlog("  NTEXCITATION_MAXNLEVELS_LOWER {}", NTEXCITATION_MAXNLEVELS_LOWER);
  printlnlog("  NTEXCITATION_MAXNLEVELS_UPPER {}", NTEXCITATION_MAXNLEVELS_UPPER);
  printlnlog("  SFPTS {}", SFPTS);
  printlnlog("  SF_EMIN {:g} eV", SF_EMIN);
  printlnlog("  SF_EMAX {:g} eV", SF_EMAX);
  printlnlog("  NT_USE_VALENCE_IONPOTENTIAL {}", NT_USE_VALENCE_IONPOTENTIAL ? "on" : "off");
  printlnlog("  NT_MAX_AUGER_ELECTRONS {}", NT_MAX_AUGER_ELECTRONS);
  printlnlog("  SF_AUGER_CONTRIBUTION {}", SF_AUGER_CONTRIBUTION_ON ? "on" : "off");
  printlnlog("  SF_AUGER_CONTRIBUTION_DISTRIBUTE_EN {}", SF_AUGER_CONTRIBUTION_DISTRIBUTE_EN ? "on" : "off");

  if (NT_EXCITATION_ON) {
    nt_excitations_stored = std::min(MAX_NT_EXCITATIONS_STORED, get_possible_nt_excitation_count());
    printlnlog("[info] mem_usage: storing {} non-thermal excitations for non-empty cells occupies {:.3f} MB",
               nt_excitations_stored,
               nonempty_npts_model * sizeof(NonThermalExcitation) * nt_excitations_stored / 1024. / 1024.);

    excitations_list_all_cells =
        MPI_shared_malloc_span<NonThermalExcitation>(nonempty_npts_model * nt_excitations_stored);
  }

  ion_data_all_cells = MPI_shared_malloc_span<NonThermalSolutionIon>(nonempty_npts_model * get_includedions());

  nt_solution = MPI_shared_malloc_span<NonThermalCellSolution>(nonempty_npts_model);

  if (globals::rank_in_node == 0) {
    for (auto nonemptymgi = 0Z; nonemptymgi < nonempty_npts_model; nonemptymgi++) {
      // should make these negative?
      nt_solution[nonemptymgi].frac_heating = 0.97;
      nt_solution[nonemptymgi].frac_ionisation = 0.03;
      nt_solution[nonemptymgi].frac_excitation = 0.;

      nt_solution[nonemptymgi].nneperion_when_solved = -1.;
      nt_solution[nonemptymgi].timestep_last_solved = -1;

      zero_all_effionpot(nonemptymgi);

      nt_solution[nonemptymgi].frac_excitations_list_size = 0;
    }
  }
  MPI_Barrier(globals::mpi_comm_node);

  const double sourceintegral = std::ranges::fold_left(
      std::views::iota(0, SFPTS), 0.0, [](double sum, int s) { return sum + (sourcevec(s) * DELTA_E); });

  printlnlog("E_init: {:14.7e} eV/s/cm3", E_init_ev);
  printlnlog("source function integral: {:14.7e}", sourceintegral);

  read_collion_data();

  printlnlog("Finished initializing non-thermal solver");
}

// set total non-thermal deposition rate from individual gamma/positron/electron/alpha rates. This should be called
// after packet propagation is finished for this timestep and normalise_deposition_estimators() has been done
void calculate_deposition_rate_density(const int nonemptymgi, const int timestep,
                                       HeatingCoolingRates& heatingcoolingrates) {
  heatingcoolingrates.dep_gamma = globals::dep_estimator_gamma[nonemptymgi];

  const double tmid = globals::timesteps[timestep].mid;
  const double rho = grid::get_rho(nonemptymgi);

  // if PARTICLE_THERMALISATION_SCHEME == ThermalisationScheme::INSTANT, use the analytic rate at t_mid since it will
  // have no Monte Carlo noise (although strictly, it should be an integral from the timestep start to the end)
  // with time-dependent deposition, we don't have an analytic rate, so we use the Monte Carlo rate

  heatingcoolingrates.eps_gamma_ana = rho * decay::get_gamma_emission_rate(nonemptymgi, tmid);

  heatingcoolingrates.eps_positron_ana =
      rho * decay::get_particle_injection_rate(nonemptymgi, tmid, decay::DECAYTYPE_BETAPLUS);

  heatingcoolingrates.eps_electron_ana =
      (rho * decay::get_particle_injection_rate(nonemptymgi, tmid, decay::DECAYTYPE_BETAMINUS));

  heatingcoolingrates.eps_alpha_ana =
      rho * decay::get_particle_injection_rate(nonemptymgi, tmid, decay::DECAYTYPE_ALPHA);

  heatingcoolingrates.eps_spfission_ana =
      rho * decay::get_particle_injection_rate(nonemptymgi, tmid, decay::DECAYTYPE_SPONTFISSION);

  if (PARTICLE_THERMALISATION_SCHEME == ThermalisationScheme::INSTANT) {
    heatingcoolingrates.dep_positron = heatingcoolingrates.eps_positron_ana;
    heatingcoolingrates.dep_electron = heatingcoolingrates.eps_electron_ana;
    heatingcoolingrates.dep_alpha = heatingcoolingrates.eps_alpha_ana;
  } else {
    heatingcoolingrates.dep_positron = globals::dep_estimator_positron[nonemptymgi];
    heatingcoolingrates.dep_electron = globals::dep_estimator_electron[nonemptymgi];
    heatingcoolingrates.dep_alpha = globals::dep_estimator_alpha[nonemptymgi];
  }

  // spontaneous fission contribution is always treated as instant deposition
  heatingcoolingrates.dep_spfission = heatingcoolingrates.eps_spfission_ana;

  deposition_rate_density_all_cells[nonemptymgi] =
      (heatingcoolingrates.dep_gamma + heatingcoolingrates.dep_positron + heatingcoolingrates.dep_electron);
}

// get non-thermal deposition rate density in erg / s / cm^3 previously stored by calculate_deposition_rate_density()
DEVICE_FUNC auto get_deposition_rate_density(const int nonemptymgi) -> double {
  assert_always(deposition_rate_density_all_cells[nonemptymgi] >= 0);
  return deposition_rate_density_all_cells[nonemptymgi];
}

void close_file() {
  if (!NT_ON || !NT_SOLVE_SPENCERFANO) {
    return;
  }

  colliondata.clear();
}

auto get_nt_frac_heating(const int nonemptymgi) -> float {
  if (!NT_ON) {
    return 1.;
  }
  if (!NT_SOLVE_SPENCERFANO) {
    return 0.97;
  }
  const float frac_heating = nt_solution[nonemptymgi].frac_heating;
  return frac_heating;
}

DEVICE_FUNC auto nt_ionisation_upperion_probability(const int nonemptymgi, const int element, const int lowerion,
                                                    const int upperion, const bool energyweighted) -> double {
  assert_always(upperion > lowerion);
  assert_always(upperion < get_nions(element));
  assert_always(upperion <= nt_ionisation_maxupperion(element, lowerion));
  if (NT_SOLVE_SPENCERFANO && NT_MAX_AUGER_ELECTRONS > 0) {
    const int numaugerelec = upperion - lowerion - 1;  // number of Auger electrons to go from lowerin to upper ion
    const int uniqueionindex = get_uniqueionindex(element, lowerion);
    const auto& cell_ion_data = get_cell_allions_data(nonemptymgi)[uniqueionindex];
    if (numaugerelec < NT_MAX_AUGER_ELECTRONS) {
      if (energyweighted) {
        return cell_ion_data.ionenfrac_num_auger[numaugerelec];
      }
      return cell_ion_data.prob_num_auger[numaugerelec];
    }
    if (numaugerelec == NT_MAX_AUGER_ELECTRONS) {
      double prob_remaining = 1.;
      for (int a = 0; a < NT_MAX_AUGER_ELECTRONS; a++) {
        if (energyweighted) {
          prob_remaining -= cell_ion_data.ionenfrac_num_auger[a];
        } else {
          prob_remaining -= cell_ion_data.prob_num_auger[a];
        }
      }
      if (energyweighted) {
        assert_always(fabs(prob_remaining - cell_ion_data.ionenfrac_num_auger[numaugerelec]) < 0.001);
      } else {
        assert_always(fabs(prob_remaining - cell_ion_data.prob_num_auger[numaugerelec]) < 0.001);
      }
      return prob_remaining;
    }

    return 0.;
  }
  return (upperion == lowerion + 1) ? 1.0 : 0.;
}

DEVICE_FUNC auto nt_ionisation_maxupperion(const int element, const int lowerion) -> int {
  const int nions = get_nions(element);
  assert_always(lowerion < nions - 1);
  int maxupper = lowerion + 1;

  if (NT_SOLVE_SPENCERFANO) {
    maxupper += NT_MAX_AUGER_ELECTRONS;
  }

  maxupper = std::min(maxupper, nions - 1);

  return maxupper;
}

DEVICE_FUNC auto nt_random_upperion(const int nonemptymgi, const int element, const int lowerion,
                                    const bool energyweighted) -> int {
  assert_testmodeonly(lowerion < get_nions(element) - 1);
  if (NT_SOLVE_SPENCERFANO && NT_MAX_AUGER_ELECTRONS > 0) {
    const double zrand = rng_uniform();

    double prob_sum = 0.;
    for (int upperion = lowerion + 1; upperion <= nt_ionisation_maxupperion(element, lowerion); upperion++) {
      prob_sum += nt_ionisation_upperion_probability(nonemptymgi, element, lowerion, upperion, energyweighted);

      if (zrand <= prob_sum) {
        return upperion;
      }
    }

    assert_always(false);  // should never get here
    return -1;
  }
  return lowerion + 1;
}

DEVICE_FUNC auto nt_ionisation_ratecoeff(const int nonemptymgi, const int element, const int ion) -> double {
  assert_always(NT_ON);

  if (NT_SOLVE_SPENCERFANO) {
    const double Y_nt = nt_ionisation_ratecoeff_sf(nonemptymgi, element, ion);
    if (!std::isfinite(Y_nt)) {
      // probably because eff_ionpot = 0 because the solver hasn't been run yet, or no impact ionisation cross sections
      // exist
      const double Y_nt_wfapprox = nt_ionisation_ratecoeff_wfapprox(nonemptymgi, element, ion);
      return Y_nt_wfapprox;
    }
    assert_always(Y_nt >= 0.);

    return Y_nt;
  }
  return nt_ionisation_ratecoeff_wfapprox(nonemptymgi, element, ion);
}

DEVICE_FUNC auto nt_excitation_ratecoeff(const int nonemptymgi, const int lowerlevel, const int upperlevel,
                                         const int alltransindex) -> double {
  if constexpr (!NT_EXCITATION_ON) {
    return 0.;
  }
  if (lowerlevel >= NTEXCITATION_MAXNLEVELS_LOWER) {
    return 0.;
  }
  if (upperlevel >= NTEXCITATION_MAXNLEVELS_UPPER) {
    return 0.;
  }

  // binary search, assuming the excitation list is sorted by lineindex ascending
  const auto ntexclist = get_cell_ntexcitations(nonemptymgi);
  const auto ntexcitation =
      std::ranges::lower_bound(ntexclist, alltransindex, {}, &NonThermalExcitation::alltransindex);
  if (ntexcitation == ntexclist.end() || ntexcitation->alltransindex != alltransindex) {
    return 0.;
  }

  const double deposition_rate_density = get_deposition_rate_density(nonemptymgi);
  const double ratecoeffperdeposition = ntexcitation->ratecoeffperdeposition;

  return ratecoeffperdeposition * deposition_rate_density;
}

DEVICE_FUNC void do_ntalpha_fisprod_deposit(Packet& pkt) {
  // if ionisation by alpha particles is found to be important for the ionisation state, we could do a separate
  // Spencer-Fano solution. For now, just treat alpha deposition as pure heating (even though the alpha deposition rate
  // was calculated from the sum of ionisation and plasma heating)
  atomicadd(nt_energy_deposited, pkt.e_cmf);
  pkt.type = TYPE_KPKT;
  stats::increment(stats::Counter::NT_STAT_TO_KPKT);
}

DEVICE_FUNC void do_ntlepton_deposit(Packet& pkt) {
  atomicadd(nt_energy_deposited, pkt.e_cmf);

  const int modelgridindex = grid::get_propcell_modelgridindex(pkt.where);
  const auto nonemptymgi = grid::get_nonemptymgi_of_mgi(modelgridindex);

  // macroatom should not be activated in thick cells
  if (NT_ON && NT_SOLVE_SPENCERFANO && grid::thick_allcells[nonemptymgi] != 1) {
    // here there is some probability to cause ionisation or excitation to a macroatom packet
    // instead of converting directly to k-packet (unless the heating channel is selected)

    double zrand = rng_uniform();
    // zrand is initially between [0, 1), but we will subtract off each
    // component of the deposition fractions
    // until we end and select transition_ij when zrand < dep_frac_transition_ij

    // const double frac_ionisation = get_nt_frac_ionisation(nonemptymgi);
    const double frac_ionisation = get_ntion_energyrate(nonemptymgi) / get_deposition_rate_density(nonemptymgi);
    // printlnlog("frac_ionisation compare {:g} and {:g}", frac_ionisation, get_nt_frac_ionisation(nonemptymgi));
    // const double frac_ionisation = 0.;

    if (zrand < frac_ionisation) {
      const auto [element, lowerion] = select_nt_ionisation(nonemptymgi);
      const int upperion = nt_random_upperion(nonemptymgi, element, lowerion, true);
      // const int upperion = lowerion + 1;

      stats::increment(stats::Counter::MA_STAT_ACTIVATION_NTCOLLION);
      stats::increment(stats::Counter::INTERACTIONS);
      pkt.trueemissiontype = EMTYPE_NOTSET;
      pkt.trueem_pos = {NAN, NAN, NAN};

      stats::increment(stats::Counter::NT_STAT_TO_IONISATION);

      do_macroatom(pkt, {.element = element, .ion = upperion, .level = 0, .activatingline = -99});
      return;
    }

    // const double frac_excitation = NT_EXCITATION_ON ? get_nt_frac_excitation(nonemptymgi) : 0;
    const double frac_excitation = 0.;
    if (zrand < (frac_ionisation + frac_excitation)) {
      zrand -= frac_ionisation;
      // now zrand is between zero and frac_excitation
      // the selection algorithm is the same as for the ionisation transitions
      for (const auto& ntexcitation : get_cell_ntexcitations(nonemptymgi)) {
        const double frac_deposition_exc = ntexcitation.frac_deposition;
        if (zrand < frac_deposition_exc) {
          const auto lineindex = globals::alltrans.lineindex[ntexcitation.alltransindex];
          const int element = globals::linelist.elementindex[lineindex];
          const int ion = globals::linelist.ionindex[lineindex];
          const int upper = globals::linelist.upperlevelindex[lineindex];

          stats::increment(stats::Counter::MA_STAT_ACTIVATION_NTCOLLEXC);
          stats::increment(stats::Counter::INTERACTIONS);
          pkt.trueemissiontype = EMTYPE_NOTSET;
          pkt.trueem_pos = {NAN, NAN, NAN};

          stats::increment(stats::Counter::NT_STAT_TO_EXCITATION);

          do_macroatom(pkt, {.element = element, .ion = ion, .level = upper, .activatingline = -99});
          return;
        }
        zrand -= frac_deposition_exc;
      }
      // in case we reached here because the excitation reactions that were stored didn't add up to frac_excitation_ion
      // then just convert it to a kpkt
    }
  }

  pkt.type = TYPE_KPKT;
  stats::increment(stats::Counter::NT_STAT_TO_KPKT);
}

// solve the Spencer-Fano equation to get the non-thermal electron flux energy distribution
// based on Equation (2) of Li et al. (2012)
void solve_spencerfano(const int nonemptymgi, const int timestep, const int iteration) {
  const auto modelgridindex = grid::get_mgi_of_nonemptymgi(nonemptymgi);
  bool skip_solution = false;
  if (timestep < globals::num_lte_timesteps + 1) {
    printlnlog("Skipping Spencer-Fano solution for first NLTE timestep");
    skip_solution = true;
  } else if (get_deposition_rate_density(nonemptymgi) / EV < MINDEPRATE) {
    printlnlog(
        "Non-thermal deposition rate of {:g} eV/cm/s/cm^3 below  MINDEPRATE {:g} in cell {} at timestep {}. Skipping "
        "Spencer-Fano solution.",
        get_deposition_rate_density(nonemptymgi) / EV, MINDEPRATE, modelgridindex, timestep);

    skip_solution = true;
  }

  if (skip_solution) {
    // Axelrod values
    nt_solution[nonemptymgi].frac_heating = 0.97;
    nt_solution[nonemptymgi].frac_ionisation = 0.03;
    nt_solution[nonemptymgi].frac_excitation = 0.;

    nt_solution[nonemptymgi].nneperion_when_solved = -1.;
    nt_solution[nonemptymgi].timestep_last_solved = -1;

    nt_solution[nonemptymgi].frac_excitations_list_size = 0;

    zero_all_effionpot(nonemptymgi);
    return;
  }

  const auto nne = grid::get_nne(nonemptymgi);  // electrons per cm^3
  const double nne_per_ion = nne / get_nnion_tot(nonemptymgi);
  const double nne_per_ion_last = nt_solution[nonemptymgi].nneperion_when_solved;
  const double nne_per_ion_fracdiff = fabs((nne_per_ion_last / nne_per_ion) - 1.);
  const int timestep_last_solved = nt_solution[nonemptymgi].timestep_last_solved;

  printlnlog(
      "Spencer-Fano solver at timestep {} (last solution was at timestep {}) nne/niontot = {:g}, at last solution was "
      "{:g} fracdiff {:g}",
      timestep, timestep_last_solved, nne_per_ion, nne_per_ion_last, nne_per_ion_fracdiff);

  if ((nne_per_ion_fracdiff < NT_MAX_FRACDIFF_NNEPERION_BETWEEN_SOLUTIONS) &&
      (timestep - timestep_last_solved <= SF_MAX_TIMESTEPS_BETWEEN_SOLUTIONS) &&
      timestep_last_solved > globals::num_lte_timesteps) {
    printlnlog(
        "Keeping Spencer-Fano solution from timestep {} because x_e fracdiff {:g} < {:g} and because timestep {} - {} "
        "< {}",
        timestep_last_solved, nne_per_ion_fracdiff, NT_MAX_FRACDIFF_NNEPERION_BETWEEN_SOLUTIONS, timestep,
        timestep_last_solved, SF_MAX_TIMESTEPS_BETWEEN_SOLUTIONS);

    return;
  }
  printlnlog(
      "Setting up Spencer-Fano equation with {} energy points from {:g} eV to {:g} eV in cell {} at timestep {} "
      "iteration {} (nne={:g} e-/cm^3)",
      SFPTS, SF_EMIN, SF_EMAX, modelgridindex, timestep, iteration, nne);

  nt_solution[nonemptymgi].nneperion_when_solved = static_cast<float>(nne_per_ion);
  nt_solution[nonemptymgi].timestep_last_solved = timestep;

  const bool enable_sfexcitation = true;
  const bool enable_sfionisation = true;
  // if (timestep <= globals::num_lte_timesteps) {
  //   // for the first run of the solver at the first NLTE timestep (which usually requires many iterations),
  //   // do a fast initial solution but mark it has an invalid nne per ion so it gets replaced at the next timestep
  //   nt_solution[nonemptymgi].nneperion_when_solved = -1.;
  //   enable_sfexcitation = false;
  //   enable_sfionisation = false;

  //   printlnlog(
  //       "Doing a fast initial solution without ionisation or excitation in the SF equation for the first NLTE "
  //       "timestep.");
  // }
  // if (timestep <= globals::num_lte_timesteps + 2) {
  //   // run the solver in a faster mode for the first couple of NLTE timesteps
  //   nt_solution[nonemptymgi].nneperion_when_solved = -1.;
  //   enable_sfexcitation = false;
  //   // enable_sfionisation = false;
  //   printlnlog("Doing a faster solution without excitation in the SF equation for the first couple of NLTE
  //   timesteps.");
  // }

  // sfmatrix will be a compacted upper triangular matrix during construction and then expanded into a full matrix (with
  // lots of zeros) just before the solver is called
  THREADLOCALONHOST std::vector<double> sfmatrix(SFPTS * SFPTS);
  // while we are filling the matrix, we only need to fill the upper triangular part in a compacted form (to reduce
  // cache misses). Later when we go to solve it, we will expand it back to a full square matrix with zeros in the lower
  // triangle
  std::fill_n(sfmatrix.begin(), SFPTS * (SFPTS + 1) / 2, 0.);
  // loss terms and source terms
  for (int i = 0; i < SFPTS; i++) {
    sfmatrix[uppertriangular(i, i)] += electron_loss_rate(engrid(i) * EV, nne) / EV;
  }

  if (enable_sfexcitation || enable_sfionisation) {
    for (int element = 0; element < get_nelements(); element++) {
      const int Z = get_atomicnumber(element);
      const int nions = get_nions(element);
      bool first_included_ion_of_element = true;
      for (int ion = 0; ion < nions; ion++) {
        const double nnion = get_nnion(nonemptymgi, element, ion);

        // skip negligible ions
        if (nnion < MIN_ION_OVER_NNTOT * get_nnion_tot(nonemptymgi)) {
          continue;
        }

        const int ionstage = get_ionstage(element, ion);
        if (first_included_ion_of_element) {
          printlog("  including Z={:2} ionstages: ", Z);
          for (int i = 1; i < get_ionstage(element, ion); i++) {
            printlog("  ");
          }
          first_included_ion_of_element = false;
        }

        printlog("{} ", ionstage);

        if (enable_sfexcitation) {
          sfmatrix_add_excitation(sfmatrix, nonemptymgi, element, ion);
        }

        if (enable_sfionisation && (ion < nions - 1)) {
          sfmatrix_add_ionisation(sfmatrix, Z, ionstage, nnion);
        }
      }
      if (!first_included_ion_of_element) {
        printlnlog("");
      }
    }
  }

  decompactify_triangular_matrix(sfmatrix);
  const auto yfunc = sfmatrix_solve(sfmatrix);

  analyse_sf_solution(nonemptymgi, timestep, enable_sfexcitation, yfunc);
}

void write_restart_data(FILE* gridsave_file) {
  printlog("non-thermal solver, ");

  fprintf(gridsave_file, "%d\n", 24724518);  // special number marking the beginning of NT data
  fprintf(gridsave_file, "%d %la %la\n", SFPTS, SF_EMIN, SF_EMAX);

  for (int nonemptymgi = 0; nonemptymgi < grid::get_nonempty_npts_model(); nonemptymgi++) {
    fprintf(gridsave_file, "%d %la ", nonemptymgi, deposition_rate_density_all_cells[nonemptymgi]);

    if (NT_ON && NT_SOLVE_SPENCERFANO) {
      check_auger_probabilities(nonemptymgi);

      fprintf(gridsave_file, "%a %a %a %a\n", nt_solution[nonemptymgi].nneperion_when_solved,
              nt_solution[nonemptymgi].frac_heating, nt_solution[nonemptymgi].frac_ionisation,
              nt_solution[nonemptymgi].frac_excitation);

      for (int uniqueionindex = 0; uniqueionindex < get_includedions(); uniqueionindex++) {
        const auto& celliondata = get_cell_allions_data(nonemptymgi)[uniqueionindex];
        fprintf(gridsave_file, "%la ", celliondata.fracdep_ionisation_ion);
        fprintf(gridsave_file, "%a ", celliondata.eff_ionpot);

        for (int a = 0; a <= NT_MAX_AUGER_ELECTRONS; a++) {
          fprintf(gridsave_file, "%a %a ", celliondata.prob_num_auger[a], celliondata.ionenfrac_num_auger[a]);
        }
      }

      // write NT excitations
      fprintf(gridsave_file, "%d\n", nt_solution[nonemptymgi].frac_excitations_list_size);

      for (const auto& excitation : get_cell_ntexcitations(nonemptymgi)) {
        fprintf(gridsave_file, "%la %la %d\n", excitation.frac_deposition, excitation.ratecoeffperdeposition,
                excitation.alltransindex);
      }
    }
  }
}

void read_restart_data(FILE* gridsave_file) {
  printlnlog("Reading restart data for non-thermal solver");

  int code_check = 0;
  assert_always(fscanf(gridsave_file, "%d\n", &code_check) == 1);
  assert_always(code_check == 24724518);  // special number marking the beginning of NT data

  int sfpts_in = 0;
  double SF_EMIN_in{NAN};
  double SF_EMAX_in{NAN};
  assert_always(fscanf(gridsave_file, "%d %la %la\n", &sfpts_in, &SF_EMIN_in, &SF_EMAX_in) == 3);

  if (sfpts_in != SFPTS || SF_EMIN_in != SF_EMIN || SF_EMAX_in != SF_EMAX) {
    printlnlog("ERROR: gridsave file specifies {} Spencer-Fano samples, SF_EMIN {:g} SF_EMAX {:g}", sfpts_in,
               SF_EMIN_in, SF_EMAX_in);
    printlnlog("ERROR: This simulation has {} Spencer-Fano samples, SF_EMIN {:g} SF_EMAX {:g}", SFPTS, SF_EMIN,
               SF_EMAX);
    std::abort();
  }

  for (int nonemptymgi = 0; nonemptymgi < grid::get_nonempty_npts_model(); nonemptymgi++) {
    int nonemptymgi_in = 0;
    assert_always(fscanf(gridsave_file, "%d %la ", &nonemptymgi_in, &deposition_rate_density_all_cells[nonemptymgi]) ==
                  2);
    assert_always(nonemptymgi_in == nonemptymgi);

    if (NT_ON && NT_SOLVE_SPENCERFANO) {
      assert_always(fscanf(gridsave_file, "%a %a %a %a\n", &nt_solution[nonemptymgi].nneperion_when_solved,
                           &nt_solution[nonemptymgi].frac_heating, &nt_solution[nonemptymgi].frac_ionisation,
                           &nt_solution[nonemptymgi].frac_excitation) == 4);

      for (int uniqueionindex = 0; uniqueionindex < get_includedions(); uniqueionindex++) {
        auto& celliondata = get_cell_allions_data(nonemptymgi)[uniqueionindex];
        assert_always(fscanf(gridsave_file, "%la ", &celliondata.fracdep_ionisation_ion) == 1);
        assert_always(fscanf(gridsave_file, "%a ", &celliondata.eff_ionpot) == 1);

        for (int a = 0; a <= NT_MAX_AUGER_ELECTRONS; a++) {
          assert_always(fscanf(gridsave_file, "%a %a ", &celliondata.prob_num_auger[a],
                               &celliondata.ionenfrac_num_auger[a]) == 2);
        }
      }

      check_auger_probabilities(nonemptymgi);

      // read NT excitations
      int frac_excitations_list_size_in = 0;
      assert_always(fscanf(gridsave_file, "%d\n", &frac_excitations_list_size_in) == 1);

      nt_solution[nonemptymgi].frac_excitations_list_size = frac_excitations_list_size_in;

      for (auto& excitation : get_cell_ntexcitations(nonemptymgi)) {
        assert_always(fscanf(gridsave_file, "%la %la %d\n", &excitation.frac_deposition,
                             &excitation.ratecoeffperdeposition, &excitation.alltransindex) == 3);
      }
    }
  }
}

void nt_MPI_Bcast(const ptrdiff_t nonemptymgi, const int root_node_id) {
  MPI_Bcast_safe(deposition_rate_density_all_cells[nonemptymgi], root_node_id, globals::mpi_comm_internode);

  if (NT_ON && NT_SOLVE_SPENCERFANO) {
    if (globals::rank_in_node == 0) {
      MPI_Bcast_safe(nt_solution[nonemptymgi].nneperion_when_solved, root_node_id, globals::mpi_comm_internode);
      MPI_Bcast_safe(nt_solution[nonemptymgi].timestep_last_solved, root_node_id, globals::mpi_comm_internode);
      MPI_Bcast_safe(nt_solution[nonemptymgi].frac_heating, root_node_id, globals::mpi_comm_internode);
      MPI_Bcast_safe(nt_solution[nonemptymgi].frac_ionisation, root_node_id, globals::mpi_comm_internode);
      MPI_Bcast_safe(nt_solution[nonemptymgi].frac_excitation, root_node_id, globals::mpi_comm_internode);

      MPI_Bcast_safe(nt_solution[nonemptymgi].frac_excitations_list_size, root_node_id, globals::mpi_comm_internode);

      MPI_Bcast_safe(get_cell_ntexcitations(nonemptymgi), root_node_id, globals::mpi_comm_internode);

      MPI_Bcast_safe(get_cell_allions_data(nonemptymgi), root_node_id, globals::mpi_comm_internode);
    }

    MPI_Barrier(MPI_COMM_WORLD);

    check_auger_probabilities(nonemptymgi);
  }
}

void reset_stats() { nt_energy_deposited = 0.; }

void print_stats(const double modelvolume, const double deltat) {
  const double deposition_rate_density_montecarlo = nt_energy_deposited / EV / modelvolume / deltat;

  printlnlog("nt_energy_deposited = {:g} [eV/s/cm^3]", deposition_rate_density_montecarlo);
}

}  // namespace nonthermal