solve for ne instead of using atmospheric value

src/statmech.jl

@@ -77,6 +77,13 @@ arguments:
 - a Dict of log molecular equilibrium constants, `log_equilibrium_constants`, in partial pressure form. 
   The keys of `equilibrium_constants` act as a list of all molecules.
 
+keyword arguments:
+- `x0` (default: `nothing`) is an initial guess for the solution (in the format internal to 
+  `chemical_equilibrium`). If not supplied, a good guess is computed by neglecting molecules.
+- `electron_number_density_warn_threshold` (default: `0.25`) is the fractional difference between 
+  the calculated electron number density and the model atmosphere electron number density at which 
+  a warning is issued.
+
 The system of equations is specified with the number densities of the neutral atoms as free 
 parameters.  Each equation specifies the conservation of a particular species, e.g. (simplified)
 
@@ -95,27 +102,40 @@ Equilibrium constants are defined in terms of partial pressures, so e.g.
 
     K(OH)  ==  (p(O) p(H)) / p(OH)  ==  (n(O) n(H)) / n(OH)) kT
 """
-function chemical_equilibrium(T, nₜ, nₑ, absolute_abundances, ionization_energies, 
-                              partition_fns, log_equilibrium_constants; x0=nothing)
+function chemical_equilibrium(T, nₜ, model_atm_nₑ, absolute_abundances, ionization_energies, 
+                              partition_fns, log_equilibrium_constants;  x0=nothing,
+                              electron_number_density_warn_threshold=0.25)
     if x0 === nothing
         #compute good first guess by neglecting molecules
-        x0 = map(1:MAX_ATOMIC_NUMBER) do atom
-            wII, wIII =  saha_ion_weights(T, nₑ, atom, ionization_energies, partition_fns)
-            nₜ*absolute_abundances[atom] / (1 + wII + wIII)
+        x0 = map(1:MAX_ATOMIC_NUMBER) do Z
+            wII, wIII =  saha_ion_weights(T, model_atm_nₑ, Z, ionization_energies, partition_fns)
+            1 / (1 + wII + wIII) 
         end
+        # this wacky maneuver ensures that x0 has the approprate dual number type for autodiff
+        # if that is going on.  I'm sure there's a better way...
+        x0 = x0 .* (absolute_abundances[1] / absolute_abundances[1])
+        push!(x0, model_atm_nₑ / nₜ * 1e5)
     end
 
     #numerically solve for equlibrium.
-    sol = nlsolve(chemical_equilibrium_equations(T, nₜ, nₑ, absolute_abundances, ionization_energies,
-                                                 partition_fns, log_equilibrium_constants),
-                  x0; iterations=1_000, store_trace=true, ftol=nₜ*1e-12, autodiff=:forward)
+    residuals! = chemical_equilibrium_equations(T, nₜ, absolute_abundances, ionization_energies, 
+                                                partition_fns, log_equilibrium_constants)
+    sol = nlsolve(residuals!, x0; method=:newton, iterations=1_000, store_trace=true, ftol=1e-8, autodiff=:forward)
     if !sol.f_converged
         error("Molecular equlibrium unconverged. \n", sol)
+    elseif !all(isfinite, sol.zero)
+        error("Molecular equlibrium solution contains non-finite values.")
     end
 
-    # start with the neutral atomic species.  Only the absolute value of sol.zero is
-    # necessarilly correct.
-    number_densities = Dict(Species.(Formula.(1:MAX_ATOMIC_NUMBER), 0) .=> abs.(sol.zero))
+    #  Only the absolute value of sol.zero is necessarilly correct.
+    nₑ = abs(sol.zero[end]) * nₜ * 1e-5
+    if abs((nₑ - model_atm_nₑ) / model_atm_nₑ) > electron_number_density_warn_threshold
+        @warn "Electron number density differs from model atmosphere by a factor greater than $electron_number_density_warn_threshold. (calculated nₑ = $nₑ, model atmosphere nₑ = $model_atm_nₑ)"
+    end 
+
+    # start with the neutral atomic species.
+    number_densities = Dict(Species.(Formula.(1:MAX_ATOMIC_NUMBER), 0) 
+                            .=> nₜ .* absolute_abundances .* abs.(sol.zero[1:end-1]))
     #now the ionized atomic species
     for a in 1:MAX_ATOMIC_NUMBER
         wII, wIII = saha_ion_weights(T, nₑ, a, ionization_energies, partition_fns)
@@ -129,39 +149,64 @@ function chemical_equilibrium(T, nₜ, nₑ, absolute_abundances, ionization_ene
         number_densities[mol] = 10^(sum(element_log_ns) - log_nK)
     end
 
-    number_densities
+    nₑ, number_densities
 end
 
-function chemical_equilibrium_equations(T, nₜ, nₑ, absolute_abundances, ionization_energies, 
+function chemical_equilibrium_equations(T, nₜ, absolute_abundances, ionization_energies, 
                                         partition_fns, log_equilibrium_constants)
     molecules = collect(keys(log_equilibrium_constants))
     atom_number_densities = nₜ .* absolute_abundances
 
-    # ion_factors is a vector of ( n(X I) + n(X II)+ n(X III) ) / n(X I) for each element X
-    ion_factors = map(1:MAX_ATOMIC_NUMBER) do elem
-        wII, wIII = saha_ion_weights(T, nₑ, elem, ionization_energies, partition_fns)
-        (1 + wII + wIII)
-    end
     # precalculate equilibrium coefficients. Here, K is in terms of number density, not partial
     # pressure, unlike those in equilibrium_constants.
     log_nKs = get_log_nK.(molecules, T, Ref(log_equilibrium_constants))
-
-    #`residuals!` puts the residuals the system of molecular equilibrium equations in `F`
-    #`x` is a vector containing the number density of the neutral species of each element
-    function residuals!(F, x)
-        # Don't allow negative number densities.  This is a trick to bound the possible values 
-        # of x. Taking the log was less performant in tests.
-        x = abs.(x) 
-
-        # LHS: total number of atoms, RHS: first through third ionization states
-        F .= atom_number_densities .- ion_factors .* x
-        for (m, log_nK) in zip(molecules, log_nKs)
-            els = get_atoms(m.formula)
-            n_mol = 10^(sum(log10(x[el]) for el in els) - log_nK)
-            # RHS: atoms which are part of molecules
-            for el in els
-                F[el] -= n_mol
+
+    # precompute the ratio of singly and doubly ionized to neutral atoms with factors of nₑ^-1 and 
+    # nₑ^-2 divided out
+    pairs = map(1:MAX_ATOMIC_NUMBER) do Z
+        saha_ion_weights(T, 1, Z, ionization_energies, partition_fns)
+    end
+    wII_ne, wIII_ne2 = first.(pairs), last.(pairs)
+
+    let nₜ=nₜ, log_nKs=log_nKs, molecules=molecules, atom_number_densities=atom_number_densities, wII_ne=wII_ne, wIII_ne2=wIII_ne2
+        #`residuals!` puts the residuals the system of molecular equilibrium equations in `F`
+        #`x` is a vector containing the number density of the neutral species of each element
+        function residuals!(F, x)
+            # the last is the electron number density in units of n_tot/10^5. The scaling allows us to 
+            # better specify the tolerance for the solver.
+            # Don't allow negative number densities.  This is a trick to bound the possible values 
+            # of x. Taking the log was less performant in tests.
+            nₑ = abs(x[end]) * nₜ * 1e-5 
+
+            # the first 92 elements of x are the fraction of each element in it's neutral atomic form
+            neutral_number_densities = atom_number_densities .* abs.(view(x, 1:MAX_ATOMIC_NUMBER))
+
+            F[end] = 0
+
+            # ion_factors is a vector of ( n(X I) + n(X II)+ n(X III) ) / n(X I) for each element X
+            for Z in 1:MAX_ATOMIC_NUMBER
+                wII = wII_ne[Z] / nₑ
+                wIII = wIII_ne2[Z] / nₑ^2
+                # LHS: total number of atoms, RHS: first through third ionization states
+                F[Z] = atom_number_densities[Z] - (1 + wII + wIII) * neutral_number_densities[Z]
+                # RHS: electrons freed from each ion
+                F[end] += (wII + 2wIII) * neutral_number_densities[Z]
             end
+            F[end] -= nₑ #LHS: total electron number density
+
+            # now the first 92 elements of x are log10(neutral number densities)
+            neutral_number_densities .= log10.(neutral_number_densities)
+            for (m, log_nK) in zip(molecules, log_nKs)
+                els = get_atoms(m.formula)
+                n_mol = 10^(sum(neutral_number_densities[el] for el in els) - log_nK)
+                # RHS: atoms which are part of molecules
+                for el in els
+                    F[el] -= n_mol
+                end
+            end
+
+            F[1:end-1] ./= atom_number_densities
+            F[end] /= nₑ * 1e-5
         end
     end
 end

src/synthesize.jl

@@ -28,6 +28,7 @@ A named tuple with keys:
    size (layers x wavelengths)
 - `number_densities`: A dictionary mapping `Species` to vectors of number densities at each 
    atmospheric layer
+- `electron_number_density`: the electron number density at each atmospheric layer
 - `wavelengths`: The vacuum wavelenths (in Å) over which the synthesis was performed.  If 
   `air_wavelengths=true` this will not be the same as the input wavelenths.
 - `subspectra`: A vector of ranges which can be used to index into `flux` to extract the spectrum 
@@ -69,6 +70,9 @@ solution = synthesize(atm, linelist, A_X, 5000, 5100)
    at which line profiles are truncated.  This has major performance impacts, since line absorption
    calculations dominate more syntheses.  Turn it down for more precision at the expense of runtime.
    The default value should effect final spectra below the 10^-3 level.
+- `electron_number_density_warn_threshold` (default: `0.25`): if the relative difference between the 
+   calculated electron number density and the input electron number density is greater than this value,
+   a warning is printed.
 - `ionization_energies`, a `Dict` mapping `Species` to their first three ionization energies, 
    defaults to `Korg.ionization_energies`.
 - `partition_funcs`, a `Dict` mapping `Species` to partition functions (in terms of ln(T)). Defaults 
@@ -89,6 +93,7 @@ function synthesize(atm::ModelAtmosphere, linelist, A_X::Vector{<:Real},
                     air_wavelengths=false, wavelength_conversion_warn_threshold=1e-4,
                     hydrogen_lines=true, use_MHD_for_hydrogen_lines=true, 
                     hydrogen_line_window_size=150, n_mu_points=20, line_cutoff_threshold=3e-4, 
+                    electron_number_density_warn_threshold=0.25,
                     bezier_radiative_transfer=false, ionization_energies=ionization_energies, 
                     partition_funcs=default_partition_funcs, 
                     log_equilibrium_constants=default_log_equilibrium_constants)
@@ -149,7 +154,7 @@ function synthesize(atm::ModelAtmosphere, linelist, A_X::Vector{<:Real},
     end
 
     abs_abundances = @. 10^(A_X - 12) # n(X) / n_tot
-    abs_abundances ./= sum(abs_abundances) #normalize so that sum(N_x/N_total) = 1
+    abs_abundances ./= sum(abs_abundances) #normalize so that sum(n(X)/n_tot) = 1
 
     #float-like type general to handle dual numbers
     α_type = typeof(promote(atm.layers[1].temp, length(linelist) > 0 ? linelist[1].wl : 1.0, 
@@ -159,44 +164,42 @@ function synthesize(atm::ModelAtmosphere, linelist, A_X::Vector{<:Real},
     # each layer's absorption at reference λ (5000 Å)
     # This isn't used with bezier radiative transfer.
     α5 = Vector{α_type}(undef, length(atm.layers)) 
-    pairs = map(enumerate(atm.layers)) do (i, layer)
-        n_dict = chemical_equilibrium(layer.temp, layer.number_density, layer.electron_number_density, 
-                                      abs_abundances, ionization_energies, 
-                                      partition_funcs, log_equilibrium_constants)
-
-        α_cntm_vals = reverse(total_continuum_absorption(sorted_cntmνs, layer.temp,
-                                                         layer.electron_number_density,
-                                                         n_dict, partition_funcs))
+    triples = map(enumerate(atm.layers)) do (i, layer)
+        nₑ, n_dict = chemical_equilibrium(layer.temp, layer.number_density, 
+                                          layer.electron_number_density, 
+                                          abs_abundances, ionization_energies, 
+                                          partition_funcs, log_equilibrium_constants; 
+                                          electron_number_density_warn_threshold=electron_number_density_warn_threshold)
+
+        α_cntm_vals = reverse(total_continuum_absorption(sorted_cntmνs, layer.temp, nₑ, n_dict, 
+                                                         partition_funcs))
         α_cntm_layer = LinearInterpolation(cntmλs, α_cntm_vals)
         α[i, :] .= α_cntm_layer.(all_λs)
 
         if ! bezier_radiative_transfer
-            α5[i] = total_continuum_absorption([c_cgs/5e-5], layer.temp, 
-                                               layer.electron_number_density, n_dict,
-                                               partition_funcs)[1]
+            α5[i] = total_continuum_absorption([c_cgs/5e-5], layer.temp, nₑ, n_dict, partition_funcs)[1]
         end
 
         if hydrogen_lines
-            hydrogen_line_absorption!(view(α, i, :), wl_ranges, layer.temp, layer.electron_number_density, 
+            hydrogen_line_absorption!(view(α, i, :), wl_ranges, layer.temp, nₑ,
                                       n_dict[species"H_I"],  n_dict[species"He I"],
                                       partition_funcs[species"H_I"](log(layer.temp)), vmic*1e5, 
                                       hydrogen_line_window_size*1e-8; 
                                       use_MHD=use_MHD_for_hydrogen_lines)
         end
 
-        n_dict, α_cntm_layer
+        nₑ, n_dict, α_cntm_layer
     end
+    nₑs = first.(triples)
     #put number densities in a dict of vectors, rather than a vector of dicts.
-    n_dicts = first.(pairs)
+    n_dicts = getindex.(triples, 2)
     number_densities = Dict([spec=>[n[spec] for n in n_dicts] for spec in keys(n_dicts[1]) 
                              if spec != species"H III"])
     #vector of continuum-absorption interpolators
-    α_cntm = last.(pairs) 
-
-    line_absorption!(α, linelist, wl_ranges, [layer.temp for layer in atm.layers], 
-                    [layer.electron_number_density for layer in atm.layers], number_densities,
-                    partition_funcs, vmic*1e5, α_cntm, cutoff_threshold=line_cutoff_threshold)
+    α_cntm = last.(triples) 
 
+    line_absorption!(α, linelist, wl_ranges, [layer.temp for layer in atm.layers], nₑs,
+        number_densities, partition_funcs, vmic*1e5, α_cntm, cutoff_threshold=line_cutoff_threshold)
 
     source_fn = blackbody.((l->l.temp).(atm.layers), all_λs')
     flux, intensity = if bezier_radiative_transfer
@@ -215,7 +218,7 @@ function synthesize(atm::ModelAtmosphere, linelist, A_X::Vector{<:Real},
     end
 
     (flux=flux, intensity=intensity, alpha=α, number_densities=number_densities, 
-     wavelengths=all_λs.*1e8, subspectra=subspectra)
+    electron_number_density=nₑs, wavelengths=all_λs.*1e8, subspectra=subspectra)
 end
 
 """

test/runtests.jl

@@ -285,11 +285,14 @@ end
 
     linelist = read_linelist("data/linelists/5000-5005.vald")
     wls = [6564:0.01:6565]
-    for atm_file in ["data/sun.mod",
-             "data/s6000_g+1.0_m0.5_t05_st_z+0.00_a+0.00_c+0.00_n+0.00_o+0.00_r+0.00_s+0.00.mod"]
+    for (atm_file, threshold) in [("data/sun.mod", 0.1),
+             ("data/s6000_g+1.0_m0.5_t05_st_z+0.00_a+0.00_c+0.00_n+0.00_o+0.00_r+0.00_s+0.00.mod", 100)]
         atm = read_model_atmosphere(atm_file)
-        flux(p) = synthesize(atm, linelist, format_A_X(p[1], Dict("Ni"=>p[2])), 
-                             wls; vmic=p[3]).flux
+        # the second model atmosphere happens to be in a weird (probably unphysical) part of 
+        # parameter space where the electron number densities calculated doesn't match the marcs
+        # numbers.  We suppress the warnings in this case by setting the threshold to 100.
+        flux(p) = synthesize(atm, linelist, format_A_X(p[1], Dict("Ni"=>p[2])), wls; 
+                             vmic=p[3], electron_number_density_warn_threshold=threshold).flux
         #make sure this works.
         J = ForwardDiff.jacobian(flux, [0.0, 0.0, 1.5])
         @test .! any(isnan.(J))

test/statmech.jl

@@ -52,27 +52,34 @@ end
         nX_ntot ./= sum(nX_ntot)
 
         nₜ = 1e15 
-        nₑ = 1e-3 * nₜ #arbitrary
+        nₑ_initial_guess = 1e12
         T = 5700
-        n = Korg.chemical_equilibrium(T, nₜ, nₑ, nX_ntot, Korg.ionization_energies, 
-                                      Korg.default_partition_funcs, 
-                                      Korg.default_log_equilibrium_constants; x0=nothing)
+        nₑ, n_dict = Korg.chemical_equilibrium(T, nₜ, nₑ_initial_guess, nX_ntot, 
+                                               Korg.ionization_energies, Korg.default_partition_funcs, 
+                                               Korg.default_log_equilibrium_constants)
+
+        @test_logs (:warn, r"Electron number density differs") Korg.chemical_equilibrium(
+                                            T, nₜ, 1.0, nX_ntot, Korg.ionization_energies, 
+                                            Korg.default_partition_funcs, 
+                                            Korg.default_log_equilibrium_constants)
+
+
+        # plasma is net-neutral
+        positive_charge_density =  mapreduce(+, pairs(n_dict)) do (species, n)
+            n * species.charge
+        end
+        @test nₑ ≈ positive_charge_density rtol=1e-5
 
         #make sure number densities are sensible
-        @test (n[Korg.species"C_III"] < n[Korg.species"C_II"] < n[Korg.species"C_I"] < 
-               n[Korg.species"H_II"] < n[Korg.species"H_I"])
+        @test (n_dict[Korg.species"C_III"] < n_dict[Korg.species"C_II"] < n_dict[Korg.species"C_I"] < 
+               n_dict[Korg.species"H_II"] < n_dict[Korg.species"H_I"])
 
-        #total number of carbons is correct
-        total_C = map(collect(keys(n))) do species
-            if species == Korg.species"C2"
-                n[species] * 2
-            elseif 0x06 in Korg.get_atoms(species.formula)
-                n[species]
-            else
-                0.0
+        @testset "conservation of nuclei: $(Korg.atomic_symbols[Z])" for Z in 1:Korg.MAX_ATOMIC_NUMBER
+            total_n = mapreduce(+, collect(keys(n_dict))) do species
+                sum(Korg.get_atoms(species.formula) .== Z) * n_dict[species]
             end
-        end |> sum
-        @test total_C ≈ nX_ntot[Korg.atomic_numbers["C"]] * nₜ
+            @test total_n ≈ nX_ntot[Z] * nₜ
+        end
     end
 
     @testset "compare to Barklem and Collet partiion functions" begin

test/transfer.jl

@@ -23,7 +23,7 @@ end
     atm = read_model_atmosphere("data/sun.mod")
     bezier_sol = synthesize(atm, [], format_A_X(), 5000, 5001; bezier_radiative_transfer=true, n_mu_points=50)
     sol = synthesize(atm, [], format_A_X(), 5000, 5001 )
-    @test assert_allclose_grid(bezier_sol.flux, sol.flux, [("λ" , sol.wavelengths, "Å")]; rtol=0.025)
+    @test assert_allclose_grid(bezier_sol.flux, sol.flux, [("λ" , sol.wavelengths, "Å")]; rtol=0.03)
 end
 
 end