diff --git a/docs/src/collisions.md b/docs/src/collisions.md
index cd641389..9670ebfd 100644
--- a/docs/src/collisions.md
+++ b/docs/src/collisions.md
@@ -212,13 +212,15 @@ This method returns an integer corresponding to the maximum allowed charge state
 ```julia
 import HallThruster: maximum_charge_state
 
-maximum_charge_state(::MyIonizationModel) = 1
+HallThruster.maximum_charge_state(::MyIonizationModel) = 1
 ```
 
 ### `load_reactions(model::ReactionModel, species::Vector{Species})`
 
-This is the most important method to define. `load_reactions` takes our model along with a vector of `Species` (generated using the `propellant` and `ncharge` fields in the user-provided `Config`) and returns a vector of `IonizationReaction{I}` or `ExcitationReaction{I}`, where `I` is the type of the rate coefficient function. It is important that the `Reactions` all have the same type so that the returned vector will have a concrete type, which will prevent dynamic dispatch during runtime and prevent unnecessary slowdowns. This means that if you have multiple reactions with different rate coefficient functions, you should use something like [`FunctionWrappers.jl`](https://github.com/yuyichao/FunctionWrappers.jl).  Let's implement a simple ionization curve with the form
+`load_reactions` takes our model along with a vector of `Species` (generated using the `propellant` and `ncharge` fields in the user-provided `Config`) and returns a vector of `IonizationReaction` or `ExcitationReaction`. If you are not using a lookup table for this purpose, the `rate_coeffs` field
+of these structs can be left empty. In that case you also need to define the `rate_coeff` function, as by default that looks for a lookup table in the reaction struct.
 
+Let's say our ionization rate coefficient is
 ```math
 k_{iz}(\epsilon) = 4\times 10^{-20} \exp\left(\frac{-7.3}{\frac{2}{3} \epsilon}\right)\sqrt{\frac{8 (\frac{2}{3}\epsilon)}{\pi m_e}}
 ```
@@ -226,12 +228,14 @@ k_{iz}(\epsilon) = 4\times 10^{-20} \exp\left(\frac{-7.3}{\frac{2}{3} \epsilon}\
 We would implement this as so:
 
 ```julia
-using FunctionWrappers
-import HallThruster: e, me, load_reactions
+import HallThruster: e, me, load_reactions, rate_coefficient
 
-kiz(ϵ) = 4e-20 * exp(-7.3 / (2/3 * ϵ)) * sqrt(8 * 2/3 * ϵ / pi / me)
+function rate_coeff(::MyIonizationModel, ::Reaction, energy::Float64)
+    Tev = 2/3 * energy
+    return 4e-20 * exp(-7.3 / Tev) * sqrt(8 * Tev / pi / me)
+end
 
-function load_reactions(model::MyIonizationModel, species)
+function load_reactions(::MyIonizationModel, species)
     rxn = IonizationReaction(
     	energy = -7.3,
         #=
@@ -241,8 +245,8 @@ function load_reactions(model::MyIonizationModel, species)
         =#
         reactant = species[1],
         product = species[2],
-        # Use a function wrapper here, though not necessary with only one reaction
-        rate_coeff = FunctionWrapper{Float64, Tuple{Float64}}(kiz)
+        # Empty lookup table, as we will not be using it
+        rate_coeff = Float64[]
     )
     return rxn
 end
diff --git a/src/collisions/collision_frequencies.jl b/src/collisions/collision_frequencies.jl
index 1d56fa11..06ad5a74 100644
--- a/src/collisions/collision_frequencies.jl
+++ b/src/collisions/collision_frequencies.jl
@@ -2,18 +2,18 @@
     freq_electron_neutral(model::ElectronNeutralModel, nn, Tev)
 Effective frequency of electron scattering caused by collisions with neutrals
 """
-@inline function freq_electron_neutral(collisions::Vector{ElasticCollision{T}}, nn::Number, Tev::Number) where T
+@inline function freq_electron_neutral(model::ElectronNeutralModel, collisions::Vector{ElasticCollision}, nn::Number, Tev::Number)
     νen = 0.0
     @inbounds for c in collisions
-        νen += c.rate_coeff(3/2 * Tev) * nn
+        νen += rate_coeff(model, c, 3/2 * Tev) * nn
     end
     return νen
 end
 
-function freq_electron_neutral(params::NamedTuple, i::Int)
+function freq_electron_neutral(params, i)
     nn = params.cache.nn[i]
     Tev = params.cache.Tev[i]
-    return freq_electron_neutral(params.electron_neutral_collisions, nn, Tev)
+    return freq_electron_neutral(params.config.electron_neutral_model, params.electron_neutral_collisions, nn, Tev)
 end
 
 """
diff --git a/src/collisions/elastic.jl b/src/collisions/elastic.jl
index b4e439e7..4308482e 100644
--- a/src/collisions/elastic.jl
+++ b/src/collisions/elastic.jl
@@ -1,48 +1,39 @@
-Base.@kwdef struct ElasticCollision{I} <: Reaction
+Base.@kwdef struct ElasticCollision <: Reaction
     species::Species
-    rate_coeff::I
+    rate_coeffs::Vector{Float64}
 end
 
 abstract type ElectronNeutralModel <: ReactionModel end
 
+# Definitions for NoElectronNeutral
 struct NoElectronNeutral <: ElectronNeutralModel end
-
 supported_species(::NoElectronNeutral) = Gas[]
+load_reactions(::NoElectronNeutral, species) = ElasticCollision[]
+rate_coeff(::NoElectronNeutral, ::Reaction, ::Float64) = 0.0
 
-load_reactions(::NoElectronNeutral, species) = ElasticCollision{Nothing}[]
-
-
+# Definitions for LandmarkElectronNeutral
 struct LandmarkElectronNeutral <: ElectronNeutralModel end
-
 supported_species(::LandmarkElectronNeutral) = [Xenon]
+load_reactions(::LandmarkElectronNeutral, species) = [ElasticCollision(Xenon(0), Float64[])]
+rate_coeff(::LandmarkElectronNeutral, ::Reaction, ::Float64) = 2.5e-13
 
-function load_reactions(::LandmarkElectronNeutral, species)
-    landmark_electron_neutral = Returns(2.5e-13)
-    return [ElasticCollision(Xenon(0), landmark_electron_neutral)]
-end
-
-
-struct GKElectronNeutral <: ElectronNeutralModel end
-
-supported_species(::GKElectronNeutral) = [Xenon]
-
+# Definitions for GKElectronNeutral
 """
-    σ_en(Tev)
-
-Electron neutral collision cross section in m² as a function of electron temperature in eV.
+Electron neutral collision model as defined by
 Eq. 3.6-13, from Fundamentals of Electric Propulsion, Goebel and Katz, 2008.
-
 """
-@inline function σ_en(Tev)
-    return max(0.0, 6.6e-19 * ((Tev / 4 - 0.1) / (1 + (Tev / 4)^1.6)))
-end
+struct GKElectronNeutral <: ElectronNeutralModel end
+supported_species(::GKElectronNeutral) = [Xenon]
+load_reactions(::GKElectronNeutral, species) = [ElasticCollision(Xenon(0), Float64[])]
 
-function load_reactions(::GKElectronNeutral, species)
-    rate_coeff(ϵ) = σ_en(2/3 * ϵ) * electron_sound_speed(2/3 * ϵ)
-    return [ElasticCollision(Xenon(0), rate_coeff)]
-end
+@inline σ_en(Tev) = max(0.0, 6.6e-19 * ((Tev / 4 - 0.1) / (1 + (Tev / 4)^1.6)))
 
+@inline function rate_coeff(::GKElectronNeutral, ::Reaction, energy::Float64)
+    Tev = 2/3 * energy
+    return σ_en(Tev) * electron_sound_speed(Tev)
+end
 
+# Definitions for ElectronNeutralLookup
 Base.@kwdef struct ElectronNeutralLookup <: ElectronNeutralModel
     directories::Vector{String} = String[]
 end
@@ -59,15 +50,14 @@ function load_reactions(model::ElectronNeutralLookup, species)
         reactant = species_sorted[i]
         if reactant.Z > 0
             break
-        else
-            for folder in folders
-                filename = rate_coeff_filename(reactant, product, collision_type, folder)
-                if ispath(filename)
-                    _, rate_coeff = load_rate_coeffs(reactant, product, collision_type,folder)
-                    reaction = HallThruster.ElasticCollision(reactant, rate_coeff)
-                    push!(reactions, reaction)
-                    break
-                end
+        end
+        for folder in folders
+            filename = rate_coeff_filename(reactant, product, collision_type, folder)
+            if ispath(filename)
+                _, rate_coeff = load_rate_coeffs(reactant, product, collision_type,folder)
+                reaction = HallThruster.ElasticCollision(reactant, rate_coeff)
+                push!(reactions, reaction)
+                break
             end
         end
     end
diff --git a/src/collisions/excitation.jl b/src/collisions/excitation.jl
index fd1e12be..926dc06c 100644
--- a/src/collisions/excitation.jl
+++ b/src/collisions/excitation.jl
@@ -1,7 +1,7 @@
-Base.@kwdef struct ExcitationReaction{I} <: Reaction
+Base.@kwdef struct ExcitationReaction <: Reaction
     energy::Float64
     reactant::Species
-    rate_coeff::I
+    rate_coeffs::Vector{Float64}
 end
 
 function Base.show(io::IO, i::ExcitationReaction)
@@ -64,7 +64,7 @@ Supports only singly-charged Xenon.
 """
 struct LandmarkExcitationLookup <: ExcitationModel end
 
-function load_reactions(model::LandmarkExcitationLookup, species)
+function load_reactions(::LandmarkExcitationLookup, species)
     rates = readdlm(LANDMARK_RATES_FILE, ',', skipstart = 1)
     ϵ = rates[:, 1]
     k_iz = rates[:, 2]
@@ -75,6 +75,8 @@ function load_reactions(model::LandmarkExcitationLookup, species)
     # We infer the excitation loss coefficient by subtracting the contribution of the ionization
     # loss coefficient from the inelastic loss term
     k_excitation = @. (k_loss - ionization_energy * k_iz) / excitation_energy
-    rate_coeff = LinearInterpolation(ϵ, k_excitation)
-    return [ExcitationReaction(excitation_energy, Xenon(0), rate_coeff)]
+    itp = LinearInterpolation(ϵ, k_excitation)
+    xs = 0:1.0:255
+    rate_coeffs = itp.(xs)
+    return [ExcitationReaction(excitation_energy, Xenon(0), rate_coeffs)]
 end
diff --git a/src/collisions/ionization.jl b/src/collisions/ionization.jl
index 0092271e..8d04d7cc 100644
--- a/src/collisions/ionization.jl
+++ b/src/collisions/ionization.jl
@@ -1,8 +1,8 @@
-Base.@kwdef struct IonizationReaction{I} <: Reaction
+Base.@kwdef struct IonizationReaction <: Reaction
     energy::Float64
     reactant::Species
     product::Species
-    rate_coeff::I
+    rate_coeffs::Vector{Float64}
 end
 
 function Base.show(io::IO, i::IonizationReaction)
@@ -74,6 +74,8 @@ function load_reactions(::LandmarkIonizationLookup, species)
     rates = readdlm(LANDMARK_RATES_FILE, ',', skipstart = 1)
     ϵ = rates[:, 1]
     k = rates[:, 2]
-    rate_coeff = LinearInterpolation(ϵ, k)
-    return [IonizationReaction(12.12, Xenon(0), Xenon(1), rate_coeff)]
+    itp = LinearInterpolation(ϵ, k)
+    xs = 0:1.0:255
+    rate_coeffs = itp.(xs)
+    return [IonizationReaction(12.12, Xenon(0), Xenon(1), rate_coeffs)]
 end
diff --git a/src/collisions/reactions.jl b/src/collisions/reactions.jl
index 0e54f85d..973ca1e0 100644
--- a/src/collisions/reactions.jl
+++ b/src/collisions/reactions.jl
@@ -1,5 +1,6 @@
 abstract type Reaction end
 
+
 function rate_coeff_filename(reactant, product, reaction_type, folder = REACTION_FOLDER)
     fname = if product === nothing
         joinpath(folder, join([reaction_type, repr(reactant)], "_") * ".dat")
@@ -56,12 +57,27 @@ function load_rate_coeffs(reactant, product, reaction_type, folder = REACTION_FO
     end
     ϵ = rates[:, 1]
     k = rates[:, 2]
-    rate_coeff = LinearInterpolation(ϵ, k, resample_uniform=true)
-    return energy, rate_coeff
+    itp = LinearInterpolation(ϵ, k)
+    xs = 0:1.0:255
+    rate_coeffs = itp.(xs)
+    return energy, rate_coeffs
 end
 
 abstract type ReactionModel end
 
+"""
+By default, rate_coeff looks for a lookup table stored in the reaction struct
+"""
+function rate_coeff(::ReactionModel, rxn::Reaction, energy)
+    ind = Base.trunc(Int64, energy)
+    N = length(rxn.rate_coeffs) - 2
+    ind = ind > N ? N : ind
+    r1 = rxn.rate_coeffs[ind+1]
+    r2 = rxn.rate_coeffs[ind+2]
+    t = energy - ind
+    return (1 - t) * r1 + t * r2
+end
+
 """
     load_reactions(model::ReactionModel, species)::Vector{IonizationReaction}
 Load ionization reactions for the provided species and ionization model
@@ -109,13 +125,11 @@ function _load_reactions(model::ReactionModel, species)
 end
 
 function _indices(symbol, reactions, species_range_dict)
-    indices = [Int[] for _ in reactions]
+    indices = zeros(Int, length(reactions))
     for (i, reaction) in enumerate(reactions)
         species = getfield(reaction, symbol).symbol
-        ranges = species_range_dict[species]
-        for r in ranges
-            push!(indices[i], r[1])
-        end
+        range = species_range_dict[species]
+        indices[i] = range[1]
     end
     return indices
 end
diff --git a/src/simulation/configuration.jl b/src/simulation/configuration.jl
index 306d8958..8b86a93e 100644
--- a/src/simulation/configuration.jl
+++ b/src/simulation/configuration.jl
@@ -111,7 +111,7 @@ function Config(;
 
     anode_Te = convert_to_float64(anode_Te, u"eV")
     cathode_Te = convert_to_float64(cathode_Te, u"eV")
-    
+
     default_neutral_velocity = 150.0 # m/s
     default_neutral_temp = 500.0 # K
     if isnothing(neutral_velocity) && isnothing(neutral_temperature)
@@ -217,10 +217,10 @@ function configure_fluids(config)
     species = [f.species for f in fluids]
 
     fluid_ranges = ranges(fluids)
-    species_range_dict = Dict(Symbol(fluid.species) => UnitRange{Int64}[] for fluid in fluids)
+    species_range_dict = Dict(Symbol(fluid.species) => 0:0 for fluid in fluids)
 
     for (fluid, fluid_range) in zip(fluids, fluid_ranges)
-        push!(species_range_dict[Symbol(fluid.species)], fluid_range)
+        species_range_dict[Symbol(fluid.species)] = fluid_range
     end
 
     last_fluid_index = fluid_ranges[end][end]
diff --git a/src/simulation/electronenergy.jl b/src/simulation/electronenergy.jl
index c5debf08..1af5a4b2 100644
--- a/src/simulation/electronenergy.jl
+++ b/src/simulation/electronenergy.jl
@@ -22,10 +22,13 @@ function update_electron_energy!(params, dt)
 
     bϵ[end] = 1.5 * Te_R * ne[end]
 
+    # Compute energy source terms
+    Q = cache.cell_cache_1
+    source_electron_energy!(Q, params)
+
     @inbounds for i in 2:ncells-1
-        Q = source_electron_energy(params, i)
         # User-provided source term
-        Q += config.source_energy(params, i)
+        Q[i] += config.source_energy(params, i)
 
         neL = ne[i-1]
         ne0 = ne[i]
@@ -117,7 +120,7 @@ function update_electron_energy!(params, dt)
         flux = 5/3 * nϵ0 * ue0
 
         # Explicit right-hand-side
-        bϵ[i] = nϵ[i] + dt * (Q - explicit * F_explicit)
+        bϵ[i] = nϵ[i] + dt * (Q[i] - explicit * F_explicit)
         bϵ[i] -= dt * flux * dlnA_dz
     end
 
diff --git a/src/simulation/restart.jl b/src/simulation/restart.jl
index fb9416de..fc12b2b7 100644
--- a/src/simulation/restart.jl
+++ b/src/simulation/restart.jl
@@ -34,14 +34,14 @@ function read_restart(path::AbstractString)
     )
 end
 
-function load_restart(grid, fluids, config, path::AbstractString)
+function load_restart(grid, config, path::AbstractString)
     sol = read_restart(path)
-    U, cache = load_restart(grid, fluids, config, sol)
+    U, cache = load_restart(grid, config, sol)
     return U, cache
 end
 
-function load_restart(grid, fluids, config, sol::Solution)
-    U, cache = HallThruster.allocate_arrays(grid, fluids)
+function load_restart(grid, config, sol::Solution)
+    U, cache = HallThruster.allocate_arrays(grid, config)
 
     z_cell = sol.params.z_cell
 
diff --git a/src/simulation/simulation.jl b/src/simulation/simulation.jl
index 3dae86be..e6942927 100644
--- a/src/simulation/simulation.jl
+++ b/src/simulation/simulation.jl
@@ -1,23 +1,20 @@
-function allocate_arrays(grid, fluids, anom_model = HallThruster.NoAnom())
-    # Number of variables in the state vector U
-    nvariables = 0
-    for fluid in fluids
-        if fluid.type == _ContinuityOnly
-            nvariables += 1
-        elseif fluid.type == _IsothermalEuler
-            nvariables += 2
-        elseif fluid.type == _EulerEquations
-            nvariables += 3
-        end
-    end
+function allocate_arrays(grid, config)
+    (;ncharge, anom_model) = config
+
+    # made less general to handle common use cases as part of fluid refactor
+    nvariables = 1 + 2 * ncharge    # 1 variable for neutrals and 2 for each ion fluid
 
     ncells = grid.ncells + 2
     nedges = grid.ncells + 1
 
+    # Main state vector
     U = zeros(nvariables, ncells)
-    B = zeros(ncells)
-    Aϵ = Tridiagonal(ones(ncells-1), ones(ncells), ones(ncells-1)) #for energy
-    bϵ = zeros(ncells) #for energy
+
+    # Caches for energy solve
+    Aϵ = Tridiagonal(ones(ncells-1), ones(ncells), ones(ncells-1))
+    bϵ = zeros(ncells)
+
+    # Collision frequencies
     νan = zeros(ncells)
     νc = zeros(ncells)
     νei = zeros(ncells)
@@ -25,35 +22,58 @@ function allocate_arrays(grid, fluids, anom_model = HallThruster.NoAnom())
     radial_loss_frequency = zeros(ncells)
     νew_momentum = zeros(ncells)
     νiw = zeros(ncells)
+    νe = zeros(ncells)
+    νiz = zeros(ncells)
+    νex = zeros(ncells)
+
+    # Magnetic field
+    B = zeros(ncells)
+
+    # Conductivity and mobility
     κ = zeros(ncells)
     μ = zeros(ncells)
+
+    # Potential and electric field
     ϕ = zeros(ncells)
     ∇ϕ = zeros(ncells)
+
+    # Electron number density
     ne = zeros(ncells)
+
+    # Electron energy density
     nϵ = zeros(ncells)
+
+    # Electron temperature and energy [eV]
     Tev = zeros(ncells)
+    ϵ = zeros(ncells)
+
+    # Electron pressure and pressure gradient
     pe = zeros(ncells)
     ∇pe = zeros(ncells)
+
+    # Electron axial velocity and kinetic energy
     ue = zeros(ncells)
-    F = zeros(nvariables, nedges)
-    UL = zeros(nvariables, nedges)
-    UR = zeros(nvariables, nedges)
-    Z_eff = zeros(ncells)
-    λ_global = zeros(length(fluids))
-    νe = zeros(ncells)
-    νiz = zeros(ncells)
-    νex = zeros(ncells)
-    K   = zeros(ncells)
-    ji  = zeros(ncells)
+    K = zeros(ncells)
+
+    λ_global = zeros(ncharge + 1)
 
+    # Electron source terms
     ohmic_heating = zeros(ncells)
     wall_losses = zeros(ncells)
     inelastic_losses = zeros(ncells)
 
-    ncharge = maximum(f.species.Z for f in fluids)
+    # Effective charge number
+    Z_eff = zeros(ncells)
+
+    # Ion density, velocity, and number flux
     ni = zeros(ncharge, ncells)
     ui = zeros(ncharge, ncells)
     niui = zeros(ncharge, ncells)
+
+    # ion current
+    ji  = zeros(ncells)
+
+    # Neutral density
     nn = zeros(ncells)
     γ_SEE = zeros(ncells)
     Id  = [0.0]
@@ -65,6 +85,11 @@ function allocate_arrays(grid, fluids, anom_model = HallThruster.NoAnom())
     errors = [0.0, 0.0, 0.0]
     dcoeffs = [0.0, 0.0, 0.0, 0.0]
 
+    # Edge state caches
+    F = zeros(nvariables, nedges)
+    UL = zeros(nvariables, nedges)
+    UR = zeros(nvariables, nedges)
+
     # timestepping caches
     k = copy(U)
     u1 = copy(U)
@@ -84,21 +109,20 @@ function allocate_arrays(grid, fluids, anom_model = HallThruster.NoAnom())
     anom_variables = allocate_anom_variables(anom_model, size(U, 2))
 
     # Timesteps
-    dt_iz = zeros(ncells)
-    dt_E = zeros(ncells)
-    dt_u = zeros(nedges)
-    dt_cell = zeros(ncells)
+    dt_iz = zeros(1)
     dt = zeros(1)
+    dt_E = zeros(1)
+    dt_u = zeros(nedges)
 
     cache = (;
-        Aϵ, bϵ, nϵ, B, νan, νc, μ, ϕ, ∇ϕ, ne, Tev, pe, ue, ∇pe,
+        Aϵ, bϵ, nϵ, B, νan, νc, μ, ϕ, ∇ϕ, ne, ϵ, Tev, pe, ue, ∇pe,
         νen, νei, radial_loss_frequency, νew_momentum, νiw, νe, κ, F, UL, UR, Z_eff, λ_global, νiz, νex, K, Id, ji,
         ni, ui, Vs, niui, nn, k, u1, γ_SEE, cell_cache_1,
         error_integral, Id_smoothed, anom_multiplier, smoothing_time_constant,
         errors, dcoeffs,
         ohmic_heating, wall_losses, inelastic_losses,
         channel_area, dA_dz, channel_height, inner_radius, outer_radius, tanδ,
-        anom_variables, dt_iz, dt_E, dt_u, dt, dt_cell
+        anom_variables, dt_iz, dt_E, dt_u, dt
     )
 
     return U, cache
@@ -114,7 +138,7 @@ function setup_simulation(
         Kp = 0.0, Ti = Inf, Td = 0.0, time_constant = 5e-4,
         dtmin = 0.0, dtmax = 1e-7,
     )
-    
+
     #check that Landmark uses the correct thermal conductivity
     if config.LANDMARK && !(config.conductivity_model isa LANDMARK_conductivity)
         error("LANDMARK configuration needs to use the LANDMARK thermal conductivity model.")
@@ -147,9 +171,9 @@ function setup_simulation(
     use_restart = restart !== nothing
 
     if use_restart
-        U, cache = load_restart(grid1d, fluids, config, restart)
+        U, cache = load_restart(grid1d, config, restart)
     else
-        U, cache = allocate_arrays(grid1d, fluids, config.anom_model)
+        U, cache = allocate_arrays(grid1d, config)
     end
 
     z_cell = grid1d.cell_centers
@@ -180,7 +204,7 @@ function setup_simulation(
 
     cache.smoothing_time_constant[] = time_constant
     cache.dt .= dt
-    cache.dt_cell .= dt
+
 
     # Simulation parameters
     params = (;
diff --git a/src/simulation/sourceterms.jl b/src/simulation/sourceterms.jl
index 32d9e49c..5917b9e2 100644
--- a/src/simulation/sourceterms.jl
+++ b/src/simulation/sourceterms.jl
@@ -1,31 +1,38 @@
-function apply_reactions!(dU::AbstractArray{T}, U::AbstractArray{T}, params, i::Int64) where T
-    (;index, ionization_reactions, index, ionization_reactant_indices, ionization_product_indices, cache) = params
-    (;inelastic_losses, νiz) = cache
-
-    ne = electron_density(U, params, i)
-    K = if params.config.LANDMARK
-        0.0
-    else
-        params.cache.K[i]
+function apply_reactions!(dU, U, params)
+    (;config, index, ionization_reactions, index, ionization_reactant_indices, ionization_product_indices, cache, ncells) = params
+    (;inelastic_losses, νiz, ϵ, ne, K) = cache
+    model = config.ionization_model
+
+    inv_m = inv(config.propellant.m)
+
+    νiz .= 0.0
+    inelastic_losses .= 0.0
+    @inbounds for i in 1:ncells
+        ne[i] = 0.0
+        for Z in 1:config.ncharge
+            ne[i] += Z * U[index.ρi[Z], i]
+        end
+        ne[i] *= inv_m
+    end
+    inv_ne = cache.cell_cache_1
+    @. inv_ne = inv(cache.ne)
+    @. ϵ = cache.nϵ * inv_ne
+    if !config.LANDMARK
+        @. ϵ += K
     end
-
-    ϵ  = cache.nϵ[i] / ne + K
 
     dt_max = Inf
-    νiz[i] = 0.0
-    inelastic_losses[i] = 0.0
-
-    @inbounds for (rxn, reactant_inds, product_inds) in zip(ionization_reactions, ionization_reactant_indices, ionization_product_indices)
-        product_index = product_inds[]
-        rate_coeff = rxn.rate_coeff(ϵ)
 
-        for reactant_index in reactant_inds
+    @inbounds for (rxn, reactant_index, product_index) in zip(ionization_reactions, ionization_reactant_indices, ionization_product_indices)
+        for i in 2:ncells-1
+            r = rate_coeff(model, rxn, ϵ[i])
             ρ_reactant = U[reactant_index, i]
-            ρdot = reaction_rate(rate_coeff, ne, ρ_reactant)
-            ndot = ρdot / params.config.propellant.m
-            νiz[i] += ndot / ne
-            inelastic_losses[i] += ndot * rxn.energy
+            ρdot = reaction_rate(r, ne[i], ρ_reactant)
             dt_max = min(dt_max, ρ_reactant / ρdot)
+            ndot = ρdot * inv_m
+            νiz[i] += ndot * inv_ne[i]
+
+            inelastic_losses[i] += ndot * rxn.energy
 
             # Change in density due to ionization
             dU[reactant_index, i] -= ρdot
@@ -36,7 +43,7 @@ function apply_reactions!(dU::AbstractArray{T}, U::AbstractArray{T}, params, i::
                 if reactant_index == index.ρn
                     reactant_velocity = params.config.neutral_velocity
                 else
-                    reactant_velocity = U[reactant_index + 1, i] / U[reactant_index, i]
+                    reactant_velocity = U[reactant_index + 1, i] / ρ_reactant
                     dU[reactant_index + 1, i] -= ρdot * reactant_velocity
                 end
 
@@ -45,129 +52,121 @@ function apply_reactions!(dU::AbstractArray{T}, U::AbstractArray{T}, params, i::
         end
     end
 
-    params.cache.dt_iz[i] = dt_max
+    cache.dt_iz[] = dt_max
 end
 
 @inline reaction_rate(rate_coeff, ne, n_reactant) = rate_coeff * ne * n_reactant
-@inline reaction_rate(rxn, ne, n_reactant, ϵ) = rxn.rate_coeff(ϵ) * n_reactant * ne
 
-function apply_ion_acceleration!(dU, U, params, i)
-    (;cache, config, index, z_edge) = params
-    E = -cache.∇ϕ[i]
+function apply_ion_acceleration!(dU, U, params)
+    (;cache, config, index, z_edge, ncells) = params
+
     mi = params.config.propellant.m
+    inv_m = inv(mi)
+    inv_e = inv(e)
     dt_max = Inf
 
-    Δx = z_edge[right_edge(i)] - z_edge[left_edge(i)]
+    @inbounds for i in 2:ncells-1
+        E = -cache.∇ϕ[i]
+        Δz = z_edge[right_edge(i)] - z_edge[left_edge(i)]
+        inv_E = inv(abs(E))
 
-    @inbounds for Z in 1:config.ncharge
-        Q_accel = Z * e * U[index.ρi[Z], i] / mi * E
-        dt_max = min(dt_max, sqrt(mi * Δx / Z / e / abs(E)))
-        dU[index.ρiui[Z], i] += Q_accel
+        @inbounds for Z in 1:config.ncharge
+            Q_accel = Z * e * U[index.ρi[Z], i] * inv_m * E
+            dt_max = min(dt_max, sqrt(mi * Δz * inv_e * inv_E / Z))
+            dU[index.ρiui[Z], i] += Q_accel
+        end
     end
 
-    params.cache.dt_E[i] = dt_max
+    params.cache.dt_E[] = dt_max
 end
 
-function apply_ion_wall_losses!(dU, U, params, i)
-    (;index, config, z_cell, L_ch) = params
+function apply_ion_wall_losses!(dU, U, params)
+    (;index, config, z_cell, L_ch, ncells, cache, ncells) = params
     (;ncharge, propellant, wall_loss_model, thruster) = config
 
-    geometry = thruster.geometry
-    Δr = geometry.outer_radius - geometry.inner_radius
+    if wall_loss_model isa NoWallLosses
+        return
+    end
+
+    inv_Δr = inv(thruster.geometry.outer_radius - thruster.geometry.inner_radius)
+    e_inv_m = e / propellant.m
 
-    mi = propellant.m
     if wall_loss_model isa WallSheath
         α = wall_loss_model.α
-    elseif wall_loss_model isa NoWallLosses
-        return
     else
         α = 1.0
     end
 
-    @inbounds for Z in 1:ncharge
-        in_channel = linear_transition(z_cell[i], L_ch, params.config.transition_length, 1.0, 0.0)
-        u_bohm = sqrt(Z * e * params.cache.Tev[i] / mi)
-        h = edge_to_center_density_ratio()
-        νiw = α * in_channel * u_bohm / Δr * h
+    h = α * edge_to_center_density_ratio()
 
-        density_loss  = U[index.ρi[Z], i] * νiw
-        momentum_loss = U[index.ρiui[Z], i] * νiw
+    @inbounds for i in 2:ncells-1
+        u_bohm = sqrt(e_inv_m * cache.Tev[i])
+        in_channel = linear_transition(z_cell[i], L_ch, params.config.transition_length, 1.0, 0.0)
+        νiw_base = in_channel * u_bohm * inv_Δr * h
 
-        dU[index.ρi[Z],   i] -= density_loss
-        dU[index.ρiui[Z], i] -= momentum_loss
+        for Z in 1:ncharge
+            νiw = sqrt(Z) * νiw_base
+            density_loss  = U[index.ρi[Z], i] * νiw
+            momentum_loss = U[index.ρiui[Z], i] * νiw
 
-        # Neutrals gain density due to ion recombination at the walls
-        dU[index.ρn, i] += density_loss
+            # Neutrals gain density due to ion recombination at the walls
+            dU[index.ρi[Z], i] -= density_loss
+            dU[index.ρn,    i] += density_loss
+            dU[index.ρiui[Z], i] -= momentum_loss
+        end
     end
 end
 
-function excitation_losses!(params, i)
-    (;excitation_reactions, cache, excitation_reactant_indices) = params
-    (;νex, Tev, inelastic_losses, nn) = cache
-    mi = params.config.propellant.m
-    ne = cache.ne[i]
-
-    K = if params.config.LANDMARK
-        # Neglect electron kinetic energy if LANDMARK
-        0.0
-    else
-        # Include kinetic energy contribution
-        cache.K[i]
-    end
-
-    # Total electron energy
-    ϵ = 3/2 * Tev[i] + K
-
-    νex[i] = 0.0
-    @inbounds for (reactant_inds, rxn) in zip(excitation_reactant_indices, excitation_reactions)
-        rate_coeff = rxn.rate_coeff(ϵ)
-        for _ in reactant_inds
-            n_reactant = nn[i]
-            ndot = reaction_rate(rate_coeff, ne, n_reactant)
-            inelastic_losses[i] += ndot * (rxn.energy - K)
-            νex[i] += ndot / ne
+function excitation_losses!(Q, params)
+    (;excitation_reactions, cache, ncells, config) = params
+    (;νex, ϵ, nn, ne, K) = cache
+    model = config.excitation_model
+
+    @. νex = 0.0
+    @inbounds for rxn in excitation_reactions
+        for i in 2:ncells-1
+            r = rate_coeff(model, rxn, ϵ[i])
+            ndot = reaction_rate(r, ne[i], nn[i])
+            νex[i] += ndot / ne[i]
+            Q[i] += ndot * (rxn.energy - !config.LANDMARK * K[i])
         end
     end
 
-    return inelastic_losses[i]
-end
-
-function electron_kinetic_energy(params, i)
-    νe = params.cache.νe[i]
-    B  = params.cache.B[i]
-    ue = params.cache.ue[i]
-    Ωe = e * B / me / νe
-    return 0.5 *  me * (1 + Ωe^2) * ue^2 / e
+    return nothing
 end
 
-function source_electron_energy(params, i)
-    ne = params.cache.ne[i]
-    ue = params.cache.ue[i]
-    ∇ϕ = params.cache.∇ϕ[i]
-
+function ohmic_heating!(Q, params)
+    (;cache, config) = params
+    (;ne, ue, ∇ϕ, K, νe, ue, ∇pe) = cache
     # Compute ohmic heating term, which is the rate at which energy is transferred out of the electron
-    # drift (kinetic energy) into thermal inergy
-    if params.config.LANDMARK
+    # drift (kinetic energy) into thermal energy
+    if (config.LANDMARK)
         # Neglect kinetic energy, so the rate of energy transfer into thermal energy is equivalent to
         # the total input power into the electrons (j⃗ ⋅ E⃗ = -mₑnₑ|uₑ|²νₑ)
-        ohmic_heating  = ne * ue * ∇ϕ
+        @. Q = ne * ue * ∇ϕ
     else
         # Do not neglect kinetic energy, so ohmic heating term is mₑnₑ|uₑ|²νₑ + ue ∇pe = 2nₑKνₑ + ue ∇pe
         # where K is the electron bulk kinetic energy, 1/2 * mₑ|uₑ|²
-        K = params.cache.K[i]
-        νe = params.cache.νe[i]
-        ∇pe = params.cache.∇pe[i]
-        ohmic_heating = 2 * ne * K * νe + ue * ∇pe
+        @. Q = 2 * ne * K * νe + ue * ∇pe
     end
+    return nothing
+end
+
+function source_electron_energy!(Q, params)
+    (;cache, config) = params
+    (;ne, ohmic_heating, wall_losses, inelastic_losses) = cache
+
+    # compute ohmic heating
+    ohmic_heating!(ohmic_heating, params)
 
     # add excitation losses to total inelastic losses
-    excitation_losses!(params, i)
+    excitation_losses!(inelastic_losses, params)
 
     # compute wall losses
-    wall_loss = ne * wall_power_loss(params.config.wall_loss_model, params, i)
+    wall_power_loss!(wall_losses, config.wall_loss_model, params)
 
-    params.cache.wall_losses[i] = wall_loss
-    params.cache.ohmic_heating[i] = ohmic_heating
+    # Compute net energy source, i.e heating minus losses
+    @. Q = ohmic_heating - ne * wall_losses - inelastic_losses
 
-    return ohmic_heating - wall_loss - params.cache.inelastic_losses[i]
+    return Q
 end
diff --git a/src/simulation/update_electrons.jl b/src/simulation/update_electrons.jl
index bba5bef2..7a9a4d41 100644
--- a/src/simulation/update_electrons.jl
+++ b/src/simulation/update_electrons.jl
@@ -4,7 +4,7 @@ function update_electrons!(params, t = 0)
     (;control_current, target_current, Kp, Ti, pe_factor, ncells) = params
     (;
         B, ue, Tev, ∇ϕ, ϕ, pe, ne, nϵ, μ, ∇pe, νan, νc, νen, νei, radial_loss_frequency,
-        Z_eff, νiz, νex, νe, ji, Id, νew_momentum, κ, Vs, nn,
+        Z_eff, νiz, νex, νe, ji, Id, νew_momentum, κ, Vs, nn, K,
         Id_smoothed, smoothing_time_constant, anom_multiplier,
         errors, channel_area
     ) = params.cache
@@ -27,7 +27,7 @@ function update_electrons!(params, t = 0)
     # Update other collisions
     @inbounds for i in 1:ncells
         # Compute electron-neutral and electron-ion collision frequencies
-        νen[i] = freq_electron_neutral(params.electron_neutral_collisions, nn[i], Tev[i])
+        νen[i] = freq_electron_neutral(params, i)
 
         # Compute total classical collision frequency
         # If we're not running the LANDMARK benchmark, include momentum transfer due to inelastic collisions
@@ -65,11 +65,11 @@ function update_electrons!(params, t = 0)
     @inbounds for i in 1:ncells
         # je + ji = Id / A
         ue[i] = (ji[i] - Id[] / channel_area[i]) / e / ne[i]
-
-        # Kinetic energy in both axial and azimuthal directions is accounted for
-        params.cache.K[i] = electron_kinetic_energy(params, i)
     end
 
+    # Kinetic energy in both axial and azimuthal directions is accounted for
+    electron_kinetic_energy!(K, params)
+
     # Compute potential gradient and pressure gradient
     compute_pressure_gradient!(∇pe, params)
 
@@ -121,7 +121,7 @@ function integrate_discharge_current(params)
 
     apply_drag = false & !params.config.LANDMARK & (iteration[] > 5)
 
-    if (apply_drag) 
+    if (apply_drag)
         (;νei, νen, νan, ui) = cache
     end
 
@@ -131,7 +131,7 @@ function integrate_discharge_current(params)
 
         int1_1 = (ji[i]   / e / μ[i]   + ∇pe[i])   / ne[i]
         int1_2 = (ji[i+1] / e / μ[i+1] + ∇pe[i+1]) / ne[i+1]
-        
+
         if (apply_drag)
             ion_drag_1 = ui[1, i  ] * (νei[i  ] + νan[i  ]) * me / e
             ion_drag_2 = ui[1, i+1] * (νei[i+1] + νan[i+1]) * me / e
@@ -161,14 +161,13 @@ function compute_electric_field!(∇ϕ, params)
     (;ji, Id, ne, μ, ∇pe, channel_area, ui, νei, νen, νan) = cache
 
     apply_drag = false & !params.config.LANDMARK & (iteration[] > 5)
-    
-    if (apply_drag) 
+
+    if (apply_drag)
         (;νei, νen, νan, ui) = cache
     end
 
 
     for i in eachindex(∇ϕ)
-
         E = ((Id[] / channel_area[i] - ji[i]) / e / μ[i] - ∇pe[i]) / ne[i]
 
         if (apply_drag)
@@ -183,6 +182,17 @@ function compute_electric_field!(∇ϕ, params)
     return ∇ϕ
 end
 
+
+function electron_kinetic_energy!(K, params)
+    (;νe, B, ue) = params.cache
+    # K = 1/2 me ue^2
+    #   = 1/2 me (ue^2 + ue_θ^2)
+    #   = 1/2 me (ue^2 + Ωe^2 ue^2)
+    #   = 1/2 me (1 + Ωe^2) ue^2
+    #   divide by e to get units of eV
+    @. K = 0.5 *  me * (1 + (e * B / me / νe)^2) * ue^2 / e
+end
+
 function compute_pressure_gradient!(∇pe, params)
     (; pe) = params.cache
     (;z_cell, ncells) = params
diff --git a/src/simulation/update_heavy_species.jl b/src/simulation/update_heavy_species.jl
index 6d37b413..feb38ce3 100644
--- a/src/simulation/update_heavy_species.jl
+++ b/src/simulation/update_heavy_species.jl
@@ -1,5 +1,5 @@
 function iterate_heavy_species!(dU, U, params; apply_boundary_conditions = true)
-    (;index, Δz_cell, config, cache, ncells) = params
+    (;index, Δz_cell, config, cache, ncells, CFL) = params
     (;
         source_neutrals, source_ion_continuity, source_ion_momentum,
         ncharge, ion_wall_losses
@@ -9,8 +9,6 @@ function iterate_heavy_species!(dU, U, params; apply_boundary_conditions = true)
 
     compute_fluxes!(F, UL, UR, U, params; apply_boundary_conditions)
 
-    params.cache.dt[] = Inf
-
     @inbounds for i in 2:ncells-1
         left = left_edge(i)
         right = right_edge(i)
@@ -38,26 +36,22 @@ function iterate_heavy_species!(dU, U, params; apply_boundary_conditions = true)
             dU[index.ρi[Z],   i] += source_ion_continuity[Z](U, params, i)
             dU[index.ρiui[Z], i] += source_ion_momentum[Z  ](U, params, i)
         end
+    end
 
-        apply_ion_acceleration!(dU, U, params, i)
-        apply_reactions!(dU, U, params, i)
-
-        if ion_wall_losses
-            apply_ion_wall_losses!(dU, U, params, i)
-        end
+    apply_reactions!(dU, U, params)
 
-        params.cache.dt_cell[i] = min(
-            sqrt(params.CFL) * params.cache.dt_E[i],
-            params.CFL * params.cache.dt_iz[i],
-            params.CFL * params.cache.dt_u[left],
-            params.CFL * params.cache.dt_u[right]
-        )
+    apply_ion_acceleration!(dU, U, params)
 
-        params.cache.dt[] = min(params.cache.dt_cell[i], params.cache.dt[])
+    if ion_wall_losses
+        apply_ion_wall_losses!(dU, U, params)
     end
 
-    params.cache.dt_cell[1] = params.cache.dt_cell[2]
-    params.cache.dt_cell[end] = params.cache.dt_cell[end-1]
+    # Compute maximum allowable timestep
+    params.cache.dt[] = min(
+        CFL * cache.dt_iz[],                          # max ionization timestep
+        sqrt(CFL) * cache.dt_E[],                     # max acceleration timestep
+        CFL * minimum(@views cache.dt_u[1:ncells-1]), # max fluid timestep
+    )
 
     @. @views dU[:, 1] = 0.0
     @. @views dU[:, end] = 0.0
@@ -83,26 +77,27 @@ function integrate_heavy_species!(U, params, dt, apply_boundary_conditions = tru
 end
 
 function update_heavy_species!(U, params)
-    (;index, ncells, cache) = params
-    (;nn, ne, ni, ui, niui, Z_eff, ji) = cache
+    (;index, ncells, cache, config) = params
+    (;nn, ne, ni, ui, niui, Z_eff, ji, K, ϵ, nϵ) = cache
     mi = params.config.propellant.m
+    inv_m = inv(mi)
 
     # Apply fluid boundary conditions
     @views left_boundary_state!(U[:, 1], U, params)
     @views right_boundary_state!(U[:, end], U, params)
 
+    # Compute neutral number density
+    @. @views nn = U[index.ρn, :] * inv_m
+
     # Update plasma quantities
     @inbounds for i in 1:ncells
-        # Compute number density for each neutral fluid
-        nn[i] = U[index.ρn, i] / params.config.propellant.m
-
         # Compute ion derived quantities
         ne[i] = 0.0
         Z_eff[i] = 0.0
         ji[i] = 0.0
         @inbounds for Z in 1:params.config.ncharge
-            _ni = U[index.ρi[Z], i] / mi
-            _niui = U[index.ρiui[Z], i] / mi
+            _ni = U[index.ρi[Z], i] * inv_m
+            _niui = U[index.ρiui[Z], i] * inv_m
             ni[Z, i] = _ni
             niui[Z, i] = _niui
             ui[Z, i] = _niui / _ni
@@ -117,4 +112,6 @@ function update_heavy_species!(U, params)
         # Effective ion charge state (density-weighted average charge state)
         Z_eff[i] = max(1.0, ne[i] / Z_eff[i])
     end
+
+    @. ϵ = nϵ / ne + config.LANDMARK * K
 end
diff --git a/src/visualization/recipes.jl b/src/visualization/recipes.jl
index cca497d2..34579917 100644
--- a/src/visualization/recipes.jl
+++ b/src/visualization/recipes.jl
@@ -1,5 +1,6 @@
 @recipe function f(sol::Solution, frame::Int = length(sol.u))
     (;z_cell, ionization_reactions, ncharge, L_ch) = sol.params
+    model = params.config.ionization_model
 
     subplot_width = 500
     subplot_height = 500
@@ -105,7 +106,7 @@
                     else
                         reactant_density = sol[:ni, rxn.reactant.Z][frame]
                     end
-                    ionization_rate .+= reactant_density .* ne .* rxn.rate_coeff.(3/2 .* Tev)
+                    ionization_rate .+= reactant_density .* ne .* rate_coeff.(model, rxn, 3/2 .* Tev)
                 end
             end
         end
diff --git a/src/wall_loss_models/constant_sheath_potential.jl b/src/wall_loss_models/constant_sheath_potential.jl
index b3fdee8b..ac84618b 100644
--- a/src/wall_loss_models/constant_sheath_potential.jl
+++ b/src/wall_loss_models/constant_sheath_potential.jl
@@ -6,17 +6,15 @@ end
 
 freq_electron_wall(::ConstantSheathPotential, params, i) = 1e7
 
-function wall_power_loss(model::ConstantSheathPotential, params, i)
-    (;z_cell, L_ch, config, cache) = params
-
-    ne = cache.ne[i]
-    nϵ = cache.nϵ[i]
-    ϵ = nϵ / ne
-
+function wall_power_loss!(Q, model::ConstantSheathPotential, params)
+    (;z_cell, L_ch, config, cache, ncells) = params
+    (;ϵ) = cache
     (;sheath_potential, inner_loss_coeff, outer_loss_coeff) = model
-    αϵ = linear_transition(z_cell[i], L_ch, config.transition_length, inner_loss_coeff, outer_loss_coeff)
 
-    W = 1e7 * αϵ * ϵ * exp(-sheath_potential / ϵ)
+    @inbounds for i in 2:ncells-1
+        αϵ = linear_transition(z_cell[i], L_ch, config.transition_length, inner_loss_coeff, outer_loss_coeff)
+        Q[i] = 1e7 * αϵ * ϵ[i] * exp(-sheath_potential / ϵ[i])
+    end
 
-    return W
+    return nothing
 end
diff --git a/src/wall_loss_models/no_wall_losses.jl b/src/wall_loss_models/no_wall_losses.jl
index beeb0a60..0f560d4e 100644
--- a/src/wall_loss_models/no_wall_losses.jl
+++ b/src/wall_loss_models/no_wall_losses.jl
@@ -2,4 +2,6 @@ struct NoWallLosses <: WallLossModel end
 
 freq_electron_wall(model::NoWallLosses, params, i) = 0.0
 wall_electron_current(model::NoWallLosses, params, i) = 0.0
-wall_power_loss(model::NoWallLosses, params, i) = 0.0
+function wall_power_loss!(Q, ::NoWallLosses, _)
+  Q .= 0.0
+end
diff --git a/src/wall_loss_models/wall_losses.jl b/src/wall_loss_models/wall_losses.jl
index 590390af..14e8d148 100644
--- a/src/wall_loss_models/wall_losses.jl
+++ b/src/wall_loss_models/wall_losses.jl
@@ -13,7 +13,7 @@ Abstract type for wall loss models in the electron energy equation. Types includ
 
 Users implementing their own `WallLossModel` will need to implement at least three methods
     1) `freq_electron_wall(model, params, i)`: Compute the electron-wall momentum transfer collision frequency in cell `i`
-    2) `wall_power_loss(model, params, i)`: Compute the electron power lost to the walls
+    2) `wall_power_loss!(Q, model, params)`: Compute the electron power lost to the walls in array Q
 
 A third method, `wall_electron_current(model, params, i)`, will compute the electron current to the walls in cell `i`. If left unimplemented,
 it defaults to Ie,w = e ne νew V_cell where V_cell is the cell volume.
@@ -27,7 +27,7 @@ function freq_electron_wall(model::WallLossModel, params, i)
     error("freq_electron_wall not implemented for wall loss model of type $(typeof(model)). See documentation for WallLossModel for a list of required methods")
 end
 
-function wall_power_loss(model::WallLossModel, params, i)
+function wall_power_loss(Q, model::WallLossModel, params)
     error("wall_power_loss not implemented for wall loss model of type $(typeof(model)). See documentation for WallLossModel for a list of required methods")
 end
 
diff --git a/src/wall_loss_models/wall_sheath.jl b/src/wall_loss_models/wall_sheath.jl
index 201450bb..3c1d9307 100644
--- a/src/wall_loss_models/wall_sheath.jl
+++ b/src/wall_loss_models/wall_sheath.jl
@@ -81,25 +81,27 @@ function freq_electron_wall(model::WallSheath, params, i)
     return νew
 end
 
-function wall_power_loss(model::WallSheath, params, i)
-    (;config, cache, z_cell, L_ch) = params
+function wall_power_loss!(Q, ::WallSheath, params)
+    (;config, cache, z_cell, L_ch, ncells) = params
     mi = config.propellant.m
 
-    Tev = wall_electron_temperature(params, i)
+    @inbounds for i in 2:ncells-1
+        Tev = wall_electron_temperature(params, i)
 
-    # space charge limited SEE coefficient
-    γ = params.cache.γ_SEE[i]
+        # space charge limited SEE coefficient
+        γ = params.cache.γ_SEE[i]
 
-    # Space charge-limited sheath potential
-    ϕ_s = sheath_potential(Tev, γ, mi)
+        # Space charge-limited sheath potential
+        ϕ_s = sheath_potential(Tev, γ, mi)
 
-    # Compute electron wall collision frequency with transition function for energy wall collisions in plume
-    νew = cache.radial_loss_frequency[i] * linear_transition(
-        z_cell[i], L_ch, config.transition_length, 1.0, config.electron_plume_loss_scale
-    )
+        # Compute electron wall collision frequency with transition function for energy wall collisions in plume
+        νew = cache.radial_loss_frequency[i] * linear_transition(
+            z_cell[i], L_ch, config.transition_length, 1.0, config.electron_plume_loss_scale
+        )
 
-    # Compute wall power loss rate
-    W = νew * (2 * Tev +  (1 - γ) * ϕ_s)
+        # Compute wall power loss rate
+        Q[i] = νew * (2 * Tev +  (1 - γ) * ϕ_s)
+    end
 
-    return W
+    return nothing
 end
diff --git a/test/order_verification/ovs_energy.jl b/test/order_verification/ovs_energy.jl
index 1fb1d550..52a6257f 100644
--- a/test/order_verification/ovs_energy.jl
+++ b/test/order_verification/ovs_energy.jl
@@ -2,8 +2,11 @@ module OVS_Energy
 
 include("ovs_funcs.jl")
 
+
 using Symbolics, HallThruster, LinearAlgebra
 
+struct R <: HallThruster.Reaction end
+
 @variables x t
 
 Dt = Differential(t)
@@ -34,7 +37,7 @@ ue_func = eval(build_function(expand_derivatives(ue), [x]))
 nn_func = eval(build_function(nn, [x]))
 ϵ_func = eval(build_function(ϵ, [x]))
 
-k(ϵ) = 8.32 * OVS_rate_coeff_ex(ϵ)
+k(ϵ) = 8.32 * HallThruster.rate_coeff(OVS_Excitation(), R(), ϵ)
 W(ϵ) = 1e7 * ϵ * exp(-20 / ϵ)
 energy_eq = Dt(nϵ) + Dx(5/3 * nϵ * ue - 10/9 * μ * nϵ * Dx(nϵ/ne)) + ne * (-ue * Dx(ϕ) + nn * k(ϵ) + W(ϵ))
 source_energy = eval(build_function(expand_derivatives(energy_eq), [x]))
@@ -65,59 +68,40 @@ end
 
 function verify_energy(ncells; niters = 20000)
     grid = HallThruster.generate_grid(HallThruster.SPT_100.geometry, (0.0, 0.05), UnevenGrid(ncells))
-
     z_cell = grid.cell_centers
     z_edge = grid.edges
     ncells = length(z_cell)
 
-    μ = μ_func.(z_cell)
-    κ = κ_func.(z_cell)
-    ne = ne_func.(z_cell)
-    ϕ = zeros(ncells)
-    ue = ue_func.(z_cell)
-    ∇ϕ = ∇ϕ_func.(z_cell)
-    nn = nn_func.(z_cell)
-    niui = niui_func.(z_cell)
-    Tev = ϵ_func.(z_cell) * 2/3
-    νex = zeros(ncells)
-    νiz = zeros(ncells)
-    channel_area = ones(ncells)
-    dA_dz = zeros(ncells)
+    # fill cache values
+    _, cache = HallThruster.allocate_arrays(grid, (;ncharge=1, anom_model=HallThruster.NoAnom()))
+    @. cache.μ = μ_func(z_cell)
+    @. cache.κ = κ_func(z_cell)
+    @. cache.ne = ne_func(z_cell)
+    @. cache.ue = ue_func(z_cell)
+    @. cache.∇ϕ = ∇ϕ_func(z_cell)
+    @. cache.nn = nn_func(z_cell)
+    @. cache.niui = niui_func(z_cell)'
+    @. cache.Tev = 2/3 * ϵ_func(z_cell)
+    @. cache.channel_area = 1.0
+    @. cache.ni = cache.ne'
 
     nϵ_exact = nϵ_func.(z_cell)
-    pe = copy(nϵ_exact)
-
-    Te_L = nϵ_exact[1] / ne[1]
-    Te_R = nϵ_exact[end] / ne[end]
-
-    nϵ = Te_L * ne # Set initial temp to 3 eV
-
-    Aϵ = Tridiagonal(ones(ncells-1), ones(ncells), ones(ncells-1))
-    bϵ = zeros(ncells)
-
-    min_electron_temperature = 0.1 * min(Te_L, Te_R)
+    @. cache.pe = copy(nϵ_exact)
 
-    source_func = (params, i) -> source_energy(params.z_cell[i])
-
-    # Test backward difference implicit solve
-    dt = 1e-6
-
-    transition_length = 0.0
-
-    excitation_model = OVS_Excitation()
-    ionization_model = OVS_Ionization()
-
-    wall_loss_model = HallThruster.ConstantSheathPotential(20.0, 1.0, 1.0)
-    L_ch = 0.025
-    propellant = HallThruster.Xenon
-    LANDMARK = true
-
-    geometry = (;channel_length = L_ch)
+    Te_L = nϵ_exact[1] / cache.ne[1]
+    Te_R = nϵ_exact[end] / cache.ne[end]
+    @. cache.nϵ = Te_L * cache.ne
 
     config = (;
-        ncharge = 1, source_energy = source_func, implicit_energy = 1.0,
-        min_electron_temperature, transition_length, LANDMARK, propellant,
-        ionization_model, excitation_model, wall_loss_model, geometry,
+        ncharge = 1,
+        source_energy = (params, i) -> source_energy(params.z_cell[i]),
+        min_electron_temperature = 0.1 * min(Te_L, Te_R),
+        transition_length = 0.0,
+        LANDMARK = true,
+        propellant = Xenon,
+        ionization_model = OVS_Ionization(), excitation_model =  OVS_Excitation(),
+        wall_loss_model = HallThruster.ConstantSheathPotential(20.0, 1.0, 1.0),
+        geometry = (;channel_length = 0.025),
         anode_boundary_condition = :dirichlet,
         conductivity_model = HallThruster.LANDMARK_conductivity(),
     )
@@ -132,59 +116,35 @@ function verify_energy(ncells; niters = 20000)
     excitation_reactions = HallThruster._load_reactions(config.excitation_model, species)
     excitation_reactant_indices = HallThruster.reactant_indices(excitation_reactions, species_range_dict)
 
-    wall_losses = zeros(ncells)
-    inelastic_losses = zeros(ncells)
-    ohmic_heating = zeros(ncells)
-    cache = (;
-        Aϵ, bϵ, μ, ϕ, ne, ue, ∇ϕ, Tev, pe, νex, νiz,
-        wall_losses, ohmic_heating, inelastic_losses,
-        channel_area, dA_dz, κ, nϵ, nn, ni = ne, niui
-    )
-
     Δz_cell, Δz_edge = HallThruster.grid_spacing(grid)
 
-    params = (;
-        z_cell, z_edge, Te_L = 2/3 * Te_L, Te_R = 2/3 * Te_R, cache, config,
-        dt, L_ch, propellant,
+    dt = 8 / maximum(abs.(cache.ue)) * (z_cell[2] - z_cell[1])
+    params_base = (;
+        dt,
+        z_cell, z_edge,
+        Te_L = 2/3 * Te_L, Te_R = 2/3 * Te_R,
+        cache = deepcopy(cache),
+        L_ch = config.geometry.channel_length,
         ionization_reactions,
         ionization_reactant_indices,
         ionization_product_indices,
         excitation_reactions,
         excitation_reactant_indices,
         Δz_cell, Δz_edge,
-        ncells
+        ncells,
     )
 
-    solve_energy!(params, niters, dt)
-    results_implicit = (;z = z_cell, exact = nϵ_exact, sim = params.cache.nϵ[:])
+    # Test backward euler implicit solve
+    params_implicit = (;params_base..., config = (;config..., implicit_energy = 1.0))
+    solve_energy!(params_implicit, niters, dt)
+    results_implicit = (;z = z_cell, exact = nϵ_exact, sim = params_implicit.cache.nϵ[:])
 
     # Test crank-nicholson implicit solve
-    params.cache.nϵ .= ne * Te_L
-
-    config = (;
-        ncharge = 1, source_energy = source_func, implicit_energy = 0.5,
-        min_electron_temperature, transition_length, LANDMARK, propellant,
-        ionization_model, excitation_model, wall_loss_model, geometry,
-        anode_boundary_condition = :dirichlet,
-        conductivity_model = HallThruster.LANDMARK_conductivity(),
-    )
-
-    dt = 8 / maximum(abs.(ue)) * (z_cell[2] - z_cell[1])
-    params = (;
-        z_cell, z_edge, Te_L = 2/3 * Te_L, Te_R = 2/3 * Te_R, cache, config,
-        dt, L_ch, propellant,
-        ionization_reactions,
-        ionization_reactant_indices,
-        ionization_product_indices,
-        excitation_reactions,
-        excitation_reactant_indices,
-        Δz_cell, Δz_edge, ncells
-    )
-
-    solve_energy!(params, niters, dt)
-    results_crank_nicholson = (;z = z_cell, exact = nϵ_exact, sim = params.cache.nϵ[:])
+    params_cn = (;params_base..., config = (;config..., implicit_energy = 0.5))
+    solve_energy!(params_cn, niters, dt)
+    results_cn = (;z = z_cell, exact = nϵ_exact, sim = params_cn.cache.nϵ[:])
 
-    return (results_implicit, results_crank_nicholson)
+    return (results_implicit, results_cn)
 end
 
 end
diff --git a/test/order_verification/ovs_funcs.jl b/test/order_verification/ovs_funcs.jl
index c29c7cb0..8b8b9f26 100644
--- a/test/order_verification/ovs_funcs.jl
+++ b/test/order_verification/ovs_funcs.jl
@@ -62,23 +62,18 @@ function refines(num_refinements, initial, factor)
     ]
 end
 
-struct OVS_Ionization <: HallThruster.IonizationModel end
-struct OVS_Excitation <: HallThruster.ExcitationModel end
+import HallThruster: load_reactions, rate_coeff, IonizationModel, ExcitationModel, IonizationReaction, ExcitationReaction, Reaction
 
-import HallThruster.load_reactions
+struct OVS_Ionization <: IonizationModel end
+struct OVS_Excitation <: ExcitationModel end
 
-function OVS_rate_coeff_iz(ϵ)
-    return 1e-12 * exp(-12.12/ϵ)
-end
-
-function OVS_rate_coeff_ex(ϵ)
-    return 1e-12 * exp(-8.32/ϵ)
-end
+HallThruster.rate_coeff(::OVS_Ionization, ::Reaction, ϵ) = 1e-12 * exp(-12.12/ϵ)
+HallThruster.rate_coeff(::OVS_Excitation, ::Reaction, ϵ) = 1e-12 * exp(-8.32/ϵ)
 
 function HallThruster.load_reactions(::OVS_Ionization, species)
-    return [HallThruster.IonizationReaction(12.12, HallThruster.Xenon(0), HallThruster.Xenon(1), OVS_rate_coeff_iz)]
+    return [IonizationReaction(12.12, Xenon(0), Xenon(1), Float64[])]
 end
 
 function HallThruster.load_reactions(::OVS_Excitation, species)
-    return [HallThruster.ExcitationReaction(8.32, HallThruster.Xenon(0), OVS_rate_coeff_ex)]
+    return [ExcitationReaction(8.32, Xenon(0), Float64[])]
 end
diff --git a/test/order_verification/ovs_ions.jl b/test/order_verification/ovs_ions.jl
index 90740dbd..b1045042 100644
--- a/test/order_verification/ovs_ions.jl
+++ b/test/order_verification/ovs_ions.jl
@@ -6,12 +6,15 @@ using HallThruster
 using LinearAlgebra
 using Symbolics
 
+struct R <: HallThruster.Reaction end
+
 @variables x t
 
 Dt = Differential(t)
 Dx = Differential(x)
 
-const k_ionization = OVS_rate_coeff_iz
+k_ionization(ϵ) = rate_coeff(OVS_Ionization(), R(), ϵ)
+
 const un = 1000
 const mi = HallThruster.Xenon.m
 const e = HallThruster.e
@@ -50,28 +53,17 @@ source_ρi = eval(build_function(expand_derivatives(continuity_ions), [x]))
 source_ρiui = eval(build_function(expand_derivatives(momentum_ions), [x]))
 
 function solve_ions(ncells, scheme; t_end = 1e-4)
-    grid = HallThruster.generate_grid(HallThruster.SPT_100.geometry, (0.0, 0.05), UnevenGrid(ncells))
-
-    Δz_cell, Δz_edge = HallThruster.grid_spacing(grid)
-    propellant = HallThruster.Xenon
-
+    # Create config struct
     thruster = HallThruster.SPT_100
-
     A_ch = HallThruster.channel_area(thruster)
     anode_mass_flow_rate = un * (ρn_func(0.0) + ui_func(0.0) / un * ρi_func(0.0))* A_ch
 
     config = (;
         thruster,
-        source_neutrals = (
-            (U, p, i) -> source_ρn(p.z_cell[i]),
-        ),
-        source_ion_continuity = (
-            (U, p, i) -> source_ρi(p.z_cell[i]),
-        ),
-        source_ion_momentum = (
-            (U, p, i) -> source_ρiui(p.z_cell[i]),
-        ),
-        propellant,
+        source_neutrals = ((_, p, i) -> source_ρn(p.z_cell[i]),),
+        source_ion_continuity = ((_, p, i) -> source_ρi(p.z_cell[i]),),
+        source_ion_momentum = ((_, p, i) -> source_ρiui(p.z_cell[i]),),
+        propellant = HallThruster.Xenon,
         ncharge = 1,
         min_electron_temperature = 1.0,
         neutral_velocity = un,
@@ -87,58 +79,45 @@ function solve_ions(ncells, scheme; t_end = 1e-4)
         conductivity_model = HallThruster.LANDMARK_conductivity(),
         ion_wall_losses = false,
         anode_boundary_condition = :dirichlet,
+        anom_model = HallThruster.NoAnom(),
     )
 
+    # Construct grid
+    grid = HallThruster.generate_grid(HallThruster.SPT_100.geometry, (0.0, 0.05), UnevenGrid(ncells))
+    Δz_cell, Δz_edge = HallThruster.grid_spacing(grid)
     z_edge = grid.edges
     z_cell = grid.cell_centers
 
-    nedges = length(z_edge)
-
-    nϵ = nϵ_func.(z_cell)
-    ρn_exact = ρn_func.(z_cell)
-    ρi_exact = ρi_func.(z_cell)
-    ui_exact = ui_func.(z_cell)
-    ∇ϕ = ∇ϕ_func.(z_cell)
-
+    # Create fluids
     fluids, fluid_ranges, species, species_range_dict, is_velocity_index = HallThruster.configure_fluids(config)
     ionization_reactions = HallThruster._load_reactions(config.ionization_model, species)
     index = HallThruster.configure_index(fluids, fluid_ranges)
 
-    nvars = fluid_ranges[end][end]
-
-    F = zeros(nvars, nedges)
-    UL = zeros(nvars, nedges)
-    UR = zeros(nvars, nedges)
-    λ_global = zeros(2)
-
-    channel_area = A_ch .* ones(ncells+2)
-    dA_dz = zeros(ncells+2)
-
-    dt = zeros(1)
-    dt_cell = zeros(ncells+2)
-    dt_u = zeros(nedges)
-    dt_iz = zeros(ncells+2)
-    dt_E = zeros(ncells+2)
+    # Allocate arrays and fill variables
+    U, cache = HallThruster.allocate_arrays(grid, config)
+    ρn_exact = ρn_func.(z_cell)
+    ρi_exact = ρi_func.(z_cell)
+    ui_exact = ui_func.(z_cell)
+    @. cache.nϵ = nϵ_func.(z_cell)
+    @. cache.ϵ = ϵ_func.(z_cell)
+    @. cache.∇ϕ = ∇ϕ_func.(z_cell)
+    @. cache.channel_area = A_ch
 
-    U = zeros(nvars, ncells+2)
-    z_end = z_cell[end]
+    # Fill initial condition
     z_start = z_cell[1]
+    z_end = z_cell[end]
     line(v0, v1, z) = v0 + (v1 - v0) * (z - z_start) / (z_end - z_start)
     U[index.ρn, :] = [line(ρn_func(z_start), ρn_func(z_end), z) for z in z_cell]
     U[index.ρi[1], :] = [line(ρi_func(z_start), ρi_func(z_end), z) for z in z_cell]
     U[index.ρiui[1], :] = U[index.ρi[1], :] * ui_func(0.0)
 
+    # Compute timestep
     amax = maximum(abs.(ui_exact) .+ sqrt.(2/3 * e * ϵ_func.(z_cell) / mi))
-    dt .= 0.1 * minimum(Δz_cell) / amax
-    dt_cell .= dt[]
-    k = zeros(size(U))
-    u1 = zeros(size(U))
-    inelastic_losses = zeros(ncells+2)
-    νiz = zeros(ncells+2)
-    cache = (;k, u1, F, UL, UR, ∇ϕ, λ_global, channel_area, dA_dz, dt_cell, dt, dt_u, dt_iz, dt_E, nϵ, νiz, inelastic_losses)
+    cache.dt[] = 0.1 * minimum(Δz_cell) / amax
 
+    # Create params struct
     params = (;
-        ncells = ncells+2,
+        ncells = length(z_cell),
         index,
         config,
         cache,
@@ -160,8 +139,6 @@ function solve_ions(ncells, scheme; t_end = 1e-4)
         CFL = 0.9,
     )
 
-    dU = copy(U)
-
     t = 0.0
     while t < t_end
         @views U[:, end] = U[:, end-1]
diff --git a/test/runtests.jl b/test/runtests.jl
index b6ae47e5..884f1358 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -27,9 +27,8 @@ include("unit_tests/test_json.jl")
 
 using Symbolics
 include("order_verification/ovs_funcs.jl")
-include("order_verification/ovs_energy.jl")
 @testset "Order verification (electron energy)" begin
-
+    include("order_verification/ovs_energy.jl")
     vfunc = x -> OVS_Energy.verify_energy(x)
     refinements = refines(6, 20, 2)
 
@@ -49,8 +48,8 @@ include("order_verification/ovs_energy.jl")
     end
 end
 
-include("order_verification/ovs_ions.jl")
 @testset "Order verification (neutrals and ions)" begin
+    include("order_verification/ovs_ions.jl")
     refinements = refines(5, 10, 2)
 
     limiter = HallThruster.van_leer
diff --git a/test/unit_tests/test_electrons.jl b/test/unit_tests/test_electrons.jl
index b85990cf..7a1dd9ad 100644
--- a/test/unit_tests/test_electrons.jl
+++ b/test/unit_tests/test_electrons.jl
@@ -62,7 +62,7 @@
 
     U = [mi * nn; mi * ne; ne * 3/2 * Tev ;;]
     @test HallThruster.freq_electron_neutral(params_landmark, 1) == 2.5e-13 * nn
-    @test HallThruster.freq_electron_neutral(params_gk, 1) == HallThruster.σ_en(Tev) * nn * sqrt(8 * e * Tev / π / me)
+    @test isapprox(HallThruster.freq_electron_neutral(params_gk, 1), HallThruster.σ_en(Tev) * nn * sqrt(8 * e * Tev / π / me), rtol = 0.01)
     @test HallThruster.freq_electron_neutral(params_none, 1) == 0.0
 
     Z = 1
diff --git a/test/unit_tests/test_initialization.jl b/test/unit_tests/test_initialization.jl
index 40a93b7c..36ae006b 100644
--- a/test/unit_tests/test_initialization.jl
+++ b/test/unit_tests/test_initialization.jl
@@ -27,7 +27,7 @@
     ncells = 100
     fluids, fluid_ranges, species, species_range_dict, is_velocity_index = HallThruster.configure_fluids(config)
     grid = HallThruster.generate_grid(config.thruster.geometry, domain, EvenGrid(ncells))
-    U, cache = HallThruster.allocate_arrays(grid, fluids, config.anom_model)
+    U, cache = HallThruster.allocate_arrays(grid, config)
     index = HallThruster.configure_index(fluids, fluid_ranges)
 
     params = (;
@@ -113,10 +113,10 @@ end
     @test fluid_ranges == [1:1, 2:3, 4:5, 6:7]
     @test species == [Xenon(0), Xenon(1), Xenon(2), Xenon(3)]
     @test species_range_dict == Dict(
-        Symbol("Xe") => [1:1],
-        Symbol("Xe+") => [2:3],
-        Symbol("Xe2+") => [4:5],
-        Symbol("Xe3+") => [6:7],
+        Symbol("Xe") => 1:1,
+        Symbol("Xe+") => 2:3,
+        Symbol("Xe2+") => 4:5,
+        Symbol("Xe3+") => 6:7,
     )
 
     @test fluids[1] == HallThruster.ContinuityOnly(species[1], config.neutral_velocity, config.neutral_temperature)
@@ -133,14 +133,14 @@ end
     # load collisions and reactions
     ionization_reactions = HallThruster._load_reactions(config.ionization_model, unique(species))
     ionization_reactant_indices = HallThruster.reactant_indices(ionization_reactions, species_range_dict)
-    @test ionization_reactant_indices == [[1], [1], [1], [2], [2], [4]]
+    @test ionization_reactant_indices == [1, 1, 1, 2, 2, 4]
 
     ionization_product_indices = HallThruster.product_indices(ionization_reactions, species_range_dict)
-    @test ionization_product_indices == [[2], [4], [6], [4], [6], [6]]
+    @test ionization_product_indices == [2, 4, 6, 4, 6, 6]
 
     excitation_reactions = HallThruster._load_reactions(config.excitation_model, unique(species))
     excitation_reactant_indices = HallThruster.reactant_indices(excitation_reactions, species_range_dict)
-    @test excitation_reactant_indices == [[1]]
+    @test excitation_reactant_indices == [1]
 
     # Test that initialization and configuration works properly when background neutrals are included
 
@@ -157,10 +157,10 @@ end
     @test fluid_ranges == [1:1, 2:3, 4:5, 6:7]
     @test species == [Xenon(0), Xenon(1), Xenon(2), Xenon(3)]
     @test species_range_dict == Dict(
-        Symbol("Xe") => [1:1],
-        Symbol("Xe+") => [2:3],
-        Symbol("Xe2+") => [4:5],
-        Symbol("Xe3+") => [6:7],
+        Symbol("Xe") => 1:1,
+        Symbol("Xe+") => 2:3,
+        Symbol("Xe2+") => 4:5,
+        Symbol("Xe3+") => 6:7,
     )
 
     @test fluids[1] == HallThruster.ContinuityOnly(species[1], config.neutral_velocity, config.neutral_temperature)
@@ -177,12 +177,12 @@ end
     # load collisions and reactions
     ionization_reactions = HallThruster._load_reactions(config.ionization_model, unique(species))
     ionization_reactant_indices = HallThruster.reactant_indices(ionization_reactions, species_range_dict)
-    @test ionization_reactant_indices == [[1], [1], [1], [2], [2], [4]]
+    @test ionization_reactant_indices == [1, 1, 1, 2, 2, 4]
 
     ionization_product_indices = HallThruster.product_indices(ionization_reactions, species_range_dict)
-    @test ionization_product_indices == [[2], [4], [6], [4], [6], [6]]
+    @test ionization_product_indices == [2, 4, 6, 4, 6, 6]
 
     excitation_reactions = HallThruster._load_reactions(config.excitation_model, unique(species))
     excitation_reactant_indices = HallThruster.reactant_indices(excitation_reactions, species_range_dict)
-    @test excitation_reactant_indices == [[1]]
+    @test excitation_reactant_indices == [1]
 end
diff --git a/test/unit_tests/test_misc.jl b/test/unit_tests/test_misc.jl
index b1d37b9f..c46a05d6 100644
--- a/test/unit_tests/test_misc.jl
+++ b/test/unit_tests/test_misc.jl
@@ -92,21 +92,20 @@ end
     @test length(grid.edges) == ncells+1
 
     Xe_0 = HallThruster.Fluid(HallThruster.Xenon(0), 100., 100.)
-    Xe_0_background = HallThruster.Fluid(HallThruster.Xenon(0), 100., 100.)
     Xe_I = HallThruster.Fluid(HallThruster.Xenon(1), T = 100.)
     Xe_II = HallThruster.Fluid(HallThruster.Xenon(2), T = 100.)
-    Xe_III = HallThruster.Fluid(HallThruster.Xenon(3))
+    Xe_III = HallThruster.Fluid(HallThruster.Xenon(3), T = 300.0)
     fluids = [
-        Xe_0,
         Xe_0,
         Xe_I,
         Xe_II,
         Xe_III
     ]
 
-    nvars = 2 + 2 + 2 + 3
+    nvars = 1 + 2 + 2 + 2
+    config = (;ncharge = 3, anom_model = HallThruster.NoAnom())
 
-    U, cache = HallThruster.allocate_arrays(grid, fluids)
+    U, cache = HallThruster.allocate_arrays(grid, config)
 
     @test size(U) == (nvars, ncells+2)
 
diff --git a/test/unit_tests/test_reactions.jl b/test/unit_tests/test_reactions.jl
index 7687d626..b3e5883d 100644
--- a/test/unit_tests/test_reactions.jl
+++ b/test/unit_tests/test_reactions.jl
@@ -5,10 +5,11 @@
     Xe_III  = HallThruster.Xenon(3)
     Xe_IV   = HallThruster.Xenon(4)
 
-    rxn_0_I = HallThruster.IonizationReaction(0.0, Xe_0, Xe_I, Te -> 0.0)
-    rxn_0_II = HallThruster.IonizationReaction(0.0, Xe_0, Xe_II, Te -> 0.0)
-    rxn_0_III = HallThruster.IonizationReaction(0.0, Xe_0, Xe_III, Te -> 0.0)
-    rxn_I_III = HallThruster.IonizationReaction(0.0, Xe_I, Xe_III, Te -> 0.0)
+    r = zeros(256)
+    rxn_0_I = HallThruster.IonizationReaction(0.0, Xe_0, Xe_I, r)
+    rxn_0_II = HallThruster.IonizationReaction(0.0, Xe_0, Xe_II, r)
+    rxn_0_III = HallThruster.IonizationReaction(0.0, Xe_0, Xe_III, r)
+    rxn_I_III = HallThruster.IonizationReaction(0.0, Xe_I, Xe_III, r)
     @test repr(rxn_0_I) == "e- + Xe -> 2e- + Xe+"
     @test repr(rxn_0_II) == "e- + Xe -> 3e- + Xe2+"
     @test repr(rxn_0_III) == "e- + Xe -> 4e- + Xe3+"
@@ -48,7 +49,7 @@
     @test_throws ArgumentError HallThruster._load_reactions(landmark_lut, [Xe_0, Xe_I, Xe_II, Xe_III])
     landmark_rxns = HallThruster._load_reactions(landmark_lut, [Xe_0, Xe_I])
     @test length(landmark_rxns) == 1
-    @test landmark_rxns[1].rate_coeff(19) ≈ 5.690E-14
+    @test HallThruster.rate_coeff(landmark_lut, landmark_rxns[1], 19.0) ≈ 5.690E-14
 
     # Test behavior of general lookup
     @test_throws ArgumentError HallThruster.load_reactions(lookup, [Ar_0, Ar_I])
@@ -68,16 +69,16 @@
     lookup_2 = HallThruster.IonizationLookup([joinpath(HallThruster.PACKAGE_ROOT, "test", "unit_tests", "reaction_tests")])
     lookup_2_rxns = HallThruster._load_reactions(lookup_2, [Ar_0, Ar_I])
     @test length(lookup_2_rxns) == 1
-    @test lookup_2_rxns[1].rate_coeff(0.3878e-01) ≈ 0.0
+    @test HallThruster.rate_coeff(lookup_2, lookup_2_rxns[1], 0.3878e-01) |> abs < eps(Float64)
     @test lookup_2_rxns[1].energy == 13.0
     lookup_2_rxns_Xe = HallThruster._load_reactions(lookup_2, [Xe_0, Xe_I, Xe_II, Xe_III])
     @test length(lookup_2_rxns_Xe) == 6
-    @test lookup_2_rxns_Xe[1].rate_coeff(1.0) ≈ 0.0
-    @test lookup_2_rxns_Xe[1].rate_coeff(1.0) != lookup_rxns[1].rate_coeff(1.0)
+    @test HallThruster.rate_coeff(lookup_2, lookup_2_rxns_Xe[1], 1.0) |> abs < eps(Float64)
+    @test HallThruster.rate_coeff(lookup_2, lookup_2_rxns_Xe[1], 1.0) != HallThruster.rate_coeff(lookup, lookup_rxns[1], 1.0)
     @test lookup_2_rxns_Xe[1].energy == 15.0
 
     # Excitation reactions
-    ex_1 = HallThruster.ExcitationReaction(0.0, Xe_0, Te -> 0.0)
+    ex_1 = HallThruster.ExcitationReaction(0.0, Xe_0, r)
     @test repr(ex_1) == "e- + Xe -> e- + Xe*"
 
     ex_lookup = HallThruster.ExcitationLookup()
@@ -85,13 +86,13 @@
     ex_rxns = HallThruster._load_reactions(ex_lookup, [Xe_0])
     @test length(ex_rxns) == 1
     @test ex_rxns[1].energy == 8.32
-    @test ex_rxns[1].rate_coeff(10) ≈ 2.3579448453671528e-14
+    @test HallThruster.rate_coeff(ex_lookup, ex_rxns[1], 10.0) ≈ 2.358251483e-14
 
     ex_lookup_2 = HallThruster.ExcitationLookup([joinpath(HallThruster.PACKAGE_ROOT, "test", "unit_tests", "reaction_tests")])
     @test !isempty(HallThruster._load_reactions(ex_lookup_2, [Ar_0]))
     ex_rxns = HallThruster._load_reactions(ex_lookup_2, [Ar_0])
     @test ex_rxns[1].energy == 10.0
-    @test ex_rxns[1].rate_coeff(1.0) ≈ 0.0
+    @test HallThruster.rate_coeff(ex_lookup_2, ex_rxns[1], 1.0) |> abs < eps(Float64)
 
     ex_lookup_landmark = HallThruster.LandmarkExcitationLookup()
     ex_landmark_rxn = HallThruster.load_reactions(ex_lookup_landmark, [Xe_0])[1]
@@ -101,5 +102,5 @@
     loss_itp = HallThruster.LinearInterpolation(landmark_table[:, 1], landmark_table[:, 3])
     iz_landmark_rxn = landmark_rxns[1]
 
-    @test 8.32 * ex_landmark_rxn.rate_coeff(14.32) + 12.12 * iz_landmark_rxn.rate_coeff(14.32) ≈ loss_itp(14.32)
+    @test 8.32 * HallThruster.rate_coeff(ex_lookup_landmark, ex_landmark_rxn, 14.32) + 12.12 * HallThruster.rate_coeff(landmark_lut, iz_landmark_rxn, 14.32) ≈ loss_itp(14.32)
 end
diff --git a/test/unit_tests/test_restarts.jl b/test/unit_tests/test_restarts.jl
index 9eb3fd86..820ba1e8 100644
--- a/test/unit_tests/test_restarts.jl
+++ b/test/unit_tests/test_restarts.jl
@@ -31,8 +31,8 @@
     @test sol.t == restart.t
 
     # test loading restart, generating cache, U, etc, without interpolation, changing number of charges, or changing domain
-    U, cache = HallThruster.load_restart(grid, fluids, config, sol)
-    U2, cache2 = HallThruster.load_restart(grid, fluids, config, restart_path)
+    U, cache = HallThruster.load_restart(grid, config, sol)
+    U2, cache2 = HallThruster.load_restart(grid, config, restart_path)
 
     # Check that we get the same result by loading solution directly or by loading filename
     @test U == U2
diff --git a/test/unit_tests/test_walls.jl b/test/unit_tests/test_walls.jl
index 7f102644..b7497ada 100644
--- a/test/unit_tests/test_walls.jl
+++ b/test/unit_tests/test_walls.jl
@@ -19,6 +19,7 @@
         ne = [ne, ne, ne, ne],
         ni = [ne ne ne ne],
         Tev = [Tev, Tev, Tev, Tev],
+        ϵ = 1.5 .* [Tev, Tev, Tev, Tev],
         nϵ = [nϵ, nϵ, nϵ, nϵ],
         Z_eff = [1.0, 1.0, 1.0, 1.0],
         γ_SEE = [0.0, 0.0, 0.0, 0.0],
@@ -34,6 +35,7 @@
     )
     index = (;ρi = [1], nϵ = 2)
 
+    ncells = 2
     grid = HallThruster.generate_grid(HallThruster.geometry_SPT_100, (0, 2 * L_ch), EvenGrid(2))
 
     z_cell = grid.cell_centers
@@ -63,12 +65,16 @@
     params = (;
         cache, config, index, z_cell = grid.cell_centers,
         z_edge = grid.edges, L_ch = L_ch, A_ch = A_ch,
-        Δz_edge, Δz_cell, γ_SEE_max = γmax
+        Δz_edge, Δz_cell, γ_SEE_max = γmax, ncells = ncells+2
     )
 
-    @test HallThruster.wall_power_loss(no_losses, params, 2) == 0.0
-    @test HallThruster.wall_power_loss(landmark_losses, params, 2) ≈ αin *  6.0 * 1e7 * exp(-20 / 6.0)
-    @test HallThruster.wall_power_loss(landmark_losses, params, 3) ≈ αout *  6.0 * 1e7 * exp(-20 / 6.0)
+    arr = zeros(ncells+2)
+    HallThruster.wall_power_loss!(arr, no_losses, params)
+    @test arr[2] == 0.0
+
+    HallThruster.wall_power_loss!(arr, landmark_losses, params)
+    @test arr[2] ≈ αin *  6.0 * 1e7 * exp(-20 / 6.0)
+    @test arr[3] ≈ αout *  6.0 * 1e7 * exp(-20 / 6.0)
 
     γ1 = 0.5
     γ2 = 1.0
@@ -121,8 +127,9 @@
     @test HallThruster.freq_electron_wall(sheath_model, params, 2) * HallThruster.linear_transition(params.z_cell[2], L_ch, params.config.transition_length, 1.0, 0.0) ≈ νew
     @test HallThruster.freq_electron_wall(sheath_model, params, 3) * HallThruster.linear_transition(params.z_cell[3], L_ch, params.config.transition_length, 1.0, 0.0) ≈ 0.0
 
-    @test HallThruster.wall_power_loss(sheath_model, params, 2) ≈ νew * (2 * Tev + (1 - γ) * Vs)
-    @test HallThruster.wall_power_loss(sheath_model, params, 4) ≈ 0.0
+    HallThruster.wall_power_loss!(arr, sheath_model, params)
+    @test arr[2] ≈ νew * (2 * Tev + (1 - γ) * Vs)
+    @test arr[4] ≈ 0.0
 end
 
 @testset "Ion wall losses" begin
@@ -214,7 +221,7 @@ end
     z_edge = grid.edges
     z_cell = grid.cell_centers
     γ_SEE_max = 1 - 8.3 * sqrt(HallThruster.me/mi)
-    base_params = (;cache, L_ch, A_ch, fluids, z_cell, z_edge, index, Δz_cell, Δz_edge, γ_SEE_max)
+    base_params = (;cache, L_ch, A_ch, fluids, z_cell, z_edge, index, Δz_cell, Δz_edge, γ_SEE_max, ncells=4)
 
     config_no_losses = (;config..., wall_loss_model = HallThruster.NoWallLosses())
     params_no_losses = (;base_params..., config = config_no_losses)
@@ -227,8 +234,7 @@ end
     u_bohm_2 = sqrt(2) * u_bohm
 
     # Test 1: no wall losses
-    HallThruster.apply_ion_wall_losses!(dU, U, params_no_losses, 2)
-    HallThruster.apply_ion_wall_losses!(dU, U, params_no_losses, 3)
+    HallThruster.apply_ion_wall_losses!(dU, U, params_no_losses)
 
     @test all(dU .≈ 0.0)
 
@@ -249,7 +255,7 @@ end
 
     dU .= 0.0
     # Check that wall losses work correctly
-    HallThruster.apply_ion_wall_losses!(dU, U, params_constant_sheath, i)
+    HallThruster.apply_ion_wall_losses!(dU, U, params_constant_sheath)
 
     # Neutrals should recombine at walls
     @test dU[index.ρn, i] ≈ -(dU[index.ρi[1], i] + dU[index.ρi[2], i])
@@ -274,7 +280,7 @@ end
     @test HallThruster.wall_ion_current(constant_sheath, params_constant_sheath, i, 2) == 0.0
     @test HallThruster.wall_electron_current(constant_sheath, params_constant_sheath, i) == 0.0
 
-    HallThruster.apply_ion_wall_losses!(dU, U, params_constant_sheath, 3)
+    HallThruster.apply_ion_wall_losses!(dU, U, params_constant_sheath)
     @test all(dU[:, 3] .≈ 0.0)
 
     # Test 3: Self-consistent wall sheath
@@ -299,7 +305,7 @@ end
     @test Iiw_2 ≈ 2 * Iew * ni_2 / ne * (1 - γ)
     @test Iew ≈ inv(1 - γ) * (Iiw_1 + Iiw_2)
 
-    HallThruster.apply_ion_wall_losses!(dU, U, params_wall_sheath, i)
+    HallThruster.apply_ion_wall_losses!(dU, U, params_wall_sheath)
 
     # Neutrals should recombine at walls
     @test dU[index.ρn, i] ≈ -(dU[index.ρi[1], i] + dU[index.ρi[2], i])
@@ -318,6 +324,6 @@ end
     @test HallThruster.wall_ion_current(wall_sheath, params_wall_sheath, i, 2) == 0.0
     @test HallThruster.wall_electron_current(wall_sheath, params_wall_sheath, i) == 0.0
 
-    HallThruster.apply_ion_wall_losses!(dU, U, params_wall_sheath, 3)
+    HallThruster.apply_ion_wall_losses!(dU, U, params_wall_sheath)
     @test all(dU[:, 3] .≈ 0.0)
 end