Internals redesign proposal

[archermarx]

Been thinking about how we might extend the code in the future. One promising possible use for the code is in investigating more chemically-complex propellants, and incorporating those into the current framework would be a headache. Additionally, the present architecture of the code derives from when we used DifferentialEquations.jl as our ODE driver. As such, we had to squeeze everything in to fit their structure, namely a state matrix U and an arbitrary collection of params. Since we no longer rely in DiffEq, I've been slowly moving away from that structure and reorganizing things in ways that make sense to me. With the PR I submitted last night (https://github.com/UM-PEPL/HallThruster.jl/pull/129), the electron solver is a bit less dependent on these details, but the ion solver is still basically the same as it was when we used DiffEq. I would like to split things up a bit differently as follows:

1. Instead of a single state vector, each species gets a struct like this

struct Species
    name::Symbol
    charge_state::Int
    density::Vector{Float64}
    momentum::Vector{Float64}
    energy::Vector{Float64}
    ...
end

This could also be parameterized based on the conservation law type (continuity only, isothermal euler, euler) to dispatch to the right update methods. Each species would be updated independently, which would enable different species to be updated at different rates (neutrals could be supercycled for example), as well as potentially improve performance (less logic in the internals related to tracking indices and possibly better cache locality). The logic here is that during the convection phase of the update, each species is totally independent from the others.

2. Reactions no longer have reactant_inds and product_inds but reactants and products, something like

   struct Reaction
   reactants::Vector{Species}
   product::Species
   ...
   end

   Applying reactions would then just become iterating over all reactions and applying them one at a time to their respective reactants and product, and would require less chasing things around through the index struct.

3. Handle source terms separately from the fluid update

When we add more reactions, we will want the flexibility to timestep the reactions (which could be stiff) separately from the rest of the ion/neutral update. On the flip side, we also don't want to perform reactions more often than is needed, as they're one of the more expensive steps in the code,

[archermarx]

Here's a prototype of the new reaction design, where sol is just a Hall thruster solution I had previously generated. The main advantage of the new approach is that we can vectorize a ton of the computations

using ProtoStructs

using HallThruster: HallThruster, Species, Xenon

struct SpeciesData
    species::Species
    density::Vector{Float64}
    momentum::Vector{Float64}
    energy::Vector{Float64}
end

function SpeciesData(species::Species, num_cells::Int)
    density = zeros(num_cells)
    momentum = zeros(num_cells)
    energy = zeros(num_cells)
    return SpeciesData(species, density, momentum, energy)
end

# new-style ionization reaction
struct IonizationRxn{X, Y}
    reactant::SpeciesData
    product::SpeciesData
    energy::Float64
    rate_coeff::HallThruster.LinearInterpolation{X, Y}
end

function apply_reaction!(rxn, rate_coeff, delta_n, delta_p, electron_density, electron_energy, dt)
    (;reactant, product) = rxn

    # compute rate coefficient
    rate_coeff .= rxn.rate_coeff.(electron_energy)

    # change in reactant density
    @. delta_n = HallThruster.reaction_rate(rate_coeff, electron_density, reactant.density) * dt

    # change in reactant momentum
    @. delta_p = delta_n * rxn.reactant.momentum / rxn.reactant.density

    # apply changes to reactants and products
    @. rxn.reactant.density -= delta_n
    @. rxn.product.density  += delta_n
    @. rxn.reactant.momentum -= delta_p
    @. rxn.product.momentum += delta_p

    # apply changes to electrons
    electrons_created = product.species.Z - reactant.species.Z
    @. electron_density += electrons_created * delta_n  # increase/decrease in number of electrons due to ionization
    @. electron_energy -= rxn.energy * delta_n          # energy loss due to ionization
end

function apply_reactions!(rxns, rate_coeff, delta_n, delta_p, electron_density, electron_energy, dt)
    @inbounds for rxn in rxns
        apply_reaction!(rxn, rate_coeff, delta_n, delta_p, electron_density, electron_energy, dt)
    end
end

function apply_reactions_old!(sol)
    for i in 2:sol.params.ncells-1
        HallThruster.apply_reactions!(sol.params.cache.k, sol.u[end], sol.params, i)
    end
end

using BenchmarkTools

sol = sols[1]

# given an existing Hall thruster solution, create the new required data structures
function create_data_newstyle(sol)
    (;ncells, fluids) = sol.params
    species = [f.species for f in fluids]
    species_data = [SpeciesData(s, ncells) for s in species]
    reactions = HallThruster.load_reactions(HallThruster.IonizationLookup(), species)
    species_data_map = Dict(zip([s.symbol for s in species], species_data))

    rxns = [
        IonizationRxn(
            species_data_map[r.reactant.symbol],
            species_data_map[r.product.symbol],
            r.energy, r.rate_coeff
        ) for r in reactions
    ]

    # fill species structs
    electron_density = copy(sol.params.cache.ne)
    electron_energy = @. 1.5 * sol.params.cache.Tev + sol.params.cache.K

    for s in species_data
        Z = s.species.Z
        if (Z == 0)
            @. s.density  = sol.params.cache.nn
            @. s.momentum = sol.params.config.neutral_velocity * s.density
        else
            @. s.density  = sol.params.cache.ni[Z, :]
            @. s.momentum = sol.params.cache.niui[Z, :]
        end
    end

    return rxns, species, electron_density, electron_energy
end

rxns, species, electron_density, electron_energy = create_data_newstyle(sols[1])

begin
    println("Old reaction method")
    @btime apply_reactions_old!($sol)
end

begin
    println("New reaction method")
    rate_coeff = zeros(sol.params.ncells)
    delta_n = zeros(sol.params.ncells)
    delta_p = zeros(sol.params.ncells)
    @btime apply_reactions!(rxns, rate_coeff, delta_n, delta_p, electron_density, electron_energy, 1e-8)
end

This runs about twice as fast as the old-style version

Old reaction method
  32.300 μs (0 allocations: 0 bytes)

New reaction method
  17.100 μs (0 allocations: 0 bytes)