Add updated Dimant-Oppenheim files

src/metablate/dimant_oppenheim_plottest.py

@@ -0,0 +1,145 @@
+import numpy as np
+import matplotlib.pyplot as plt
+
+# import
+from metablate.models.dimant_oppenheim_2017 import (
+    plasma_distribution_morphology_array,
+    plasma_distribution,
+    collisional_cross_section_bronshten_1983,
+    ionization_probability_Na_vondrak_2008,
+    mean_free_path,
+    ablated_thermal_speed_bronshten_1983
+)
+
+
+def main():
+    # 1) Build (r, theta) grid in units of mean free path (lambda) 
+    # r is dimensionless distance in units of lambda
+    r = np.linspace(0.001, 0.1, 100)
+    # theta = np.linspace(0.5*np.pi, 1.5*np.pi, 50)
+    theta = np.array([3.36]) 
+    R, TH = np.meshgrid(r, theta)
+
+    # 2) Compute morphology
+    Xlam, Ylam, Rlam, ne_morph, t1, t2, t3 = plasma_distribution_morphology_array(
+        R, TH, pbar=True
+    )
+
+    Xlam, Ylam, Rlam, ne_morph1, t1, t2, t3 = plasma_distribution_morphology_array(
+        R, TH, pbar=True, quad_abs_tol = 1e-3
+    )
+
+    Xlam, Ylam, Rlam, ne_morph2, t1, t2, t3 = plasma_distribution_morphology_array(
+        R, TH, pbar=True, quad_abs_tol = 1e-8
+    )
+
+
+    # 3) Scale to physical electron density using plasma_distribution
+    # Put in example-ish parameters 
+    total_atm_n = 1e19  # atmospheric number density [m^-3] 
+    v = 40000.0  # meteroid velocity [m/s] 
+    r_m = 5e-5  # meteoroid radius [m]
+    plasma_source_density = 1e18  # example
+    R_lambda = 1.0  # dimensionless scaling in DO paper
+    ablated_thermal_speed = ablated_thermal_speed_bronshten_1983(meteoroid_surface_temperature = 2000, meteoroid_molecular_mass = 3.8e-26)
+    sigma = collisional_cross_section_bronshten_1983(v)  # m^2
+    lamb = mean_free_path(
+        ablated_thermal_speed,
+        total_atm_n,
+        sigma,
+        v)
+    Pion = ionization_probability_Na_vondrak_2008(v)  # dimensionless
+
+    ne = plasma_distribution(
+        total_atmospheric_number_density=total_atm_n,
+        meteoroid_velocity=v,
+        meteoroid_radius=r_m,
+        plasma_source_density=plasma_source_density,
+        collisional_cross_section=sigma,
+        ionization_probability=Pion,
+        base_plasma_distribution_function=ne_morph,
+        R_lambda=R_lambda,
+    )
+
+    # 4) Plot density curves (1D slices)
+    # a) along the meteoroid path axis: theta = 0
+    # b) perpendicular cut: theta = pi/2
+    # both morphology and scaled density
+
+    i0 = np.argmin(np.abs(theta)) # directly behind meteoroid
+    i90 = np.argmin(np.abs(theta - (np.pi / 2))) # perpendicular to the meteoroid
+
+
+
+    # figure -1
+    plt.figure()
+    plt.plot(r, ne_morph1[0, :] - ne_morph2[0, :], label = "ne_morph1")
+    plt.title("Components")
+    plt.legend()
+    plt.show() 
+
+    # figure 0
+    plt.figure()
+    plt.plot(r, t1[0, :], label = "t1")
+    plt.plot(r, t2[0, :], label = "t2")
+    plt.plot(r, t3[0, :], label = "t3")
+    plt.plot(r, ne_morph[0, :], label = "ne_morph")
+    plt.title("Components")
+    plt.legend()
+    plt.show()
+
+
+
+    # figure 1
+    plt.figure()
+    plt.plot(r, ne_morph[i0, :], label="morphology, θ=0")
+    plt.plot(r, ne_morph[i90, :], label="morphology, θ=π/2")
+    plt.yscale("log")
+    plt.xlabel("r / λ")
+    plt.ylabel("ne_morph (dimensionless)")
+    plt.title("Plasma morphology curves")
+    plt.legend()
+    plt.grid(True, which="both")
+    plt.show()
+
+    # figure 2: with dips
+    plt.figure()
+    for index in range(len(theta)):
+        plt.plot(r, ne_morph[index, :], label=f"{theta[index]}")
+    plt.axvline(r_m/lamb, color = "red")
+    plt.yscale("log")
+    plt.xlabel("r / λ")
+    plt.ylabel("ne_morph (dimensionless)")
+    plt.title("Plasma morphology curves")
+    plt.legend()
+    plt.grid(True, which="both")
+    plt.show()
+
+    # figure 3
+    plt.figure()
+    plt.plot(r, ne[i0, :], label="density, θ=0")
+    plt.plot(r, ne[i90, :], label="density, θ=π/2")
+    plt.yscale("log")
+    plt.xlabel("r / λ")
+    plt.ylabel(r"ne (m$^{-3}$)")
+    plt.title("Scaled electron density curves")
+    plt.legend()
+    plt.grid(True, which="both")
+    plt.show()
+
+    # 5) 2D map (sanity check)
+    plt.figure()
+    plt.pcolormesh(Xlam, Ylam, ne_morph, shading="auto")
+    plt.gca().set_aspect("equal", "box")
+    plt.colorbar(label="ne_morph")
+    plt.xlabel("X / λ")
+    plt.ylabel("Y / λ")
+    plt.title("Plasma morphology (2D)")
+    plt.show()
+
+
+if __name__ == "__main__":
+    main()
+
+#note: ne_morph dip at r/lambda = 0.06 (ne_morph = 9e-5)
+

src/metablate/models/dimant_oppenheim_2017.py

@@ -6,7 +6,7 @@
 from tqdm import tqdm
 import multiprocessing as mp
 
-from ..types import NDArray_N
+from metablate.array_types import NDArray_N
 
 
 def ablated_thermal_speed_bronshten_1983(meteoroid_surface_temperature, meteoroid_molecular_mass):
@@ -125,12 +125,228 @@ def do_int_2_quad(a, b, rlam23, cth, sth, quad_abs_tol):
 def plasma_distribution_morphology_array(
     r: NDArray_N,
     theta: NDArray_N,
-    quad_abs_tol: float = 1e-7,
-    upper_limit_delta: float = 1e-6,
-    pbar: bool = True,
-):
+    quad_abs_tol: float = 1e-5,
+    # upper_limit_delta: float = 1e-6,
+    pbar: bool = True):
+
     # Generate density distribution
     # Solve for a 2D plane
+
+    ### bookmark 1 ###
+
+    ## r and theta as arrays, broadcasting used to achieve same shape of arrays
+    r = np.asarray(r, dtype=float)
+    theta = np.asarray(theta, dtype=float)
+    r, theta = np.broadcast_arrays(r, theta)
+
+    ## cos and sin theta
+    costh = np.cos(theta)
+    sinth = np.sin(theta)
+
+    ## mean free path x and y components
+    Xlam = r * costh
+    Ylam = r * sinth
+
+    ## r as Rlam 
+    Rlam = r.copy()
+    Rlam[Rlam == 0] = np.nan
+    Rlam23 = Rlam ** (2/3)
+
+    ## Func 1 (same as before)
+    Func1 = np.sqrt(2 * np.pi / 3) / Rlam * scs.erf(
+        np.sqrt(3 / 2) * Rlam ** (1 / 3) * np.abs(costh) ** (1 / 3)
+    ) - np.exp(-3 / 2 * Rlam23 * np.abs(costh) ** (2 / 3)) * (
+        (4 - np.pi)
+        * np.abs(costh)
+        / (2 * np.sqrt((1 + (4 - np.pi) ** 2 * Rlam23 * np.abs(costh) ** (2 / 3)) / (2 * np.pi)))
+        + 2 * np.abs(costh) ** (1 / 3) / Rlam23
+    )
+
+    ## Func 2 (same as before)
+    Func2 = np.sqrt(2 * np.pi / 3) / Rlam * scs.erf(np.sqrt(3 * Rlam23 / 2)) - (
+        1 + 2 / Rlam23
+    ) * np.exp(-3 * Rlam23 / 2)
+
+    ## flatten but ravel for faster computation of quad loop
+    orig_shape = costh.shape
+    costh_f = costh.ravel()
+    sinth_f = sinth.ravel()
+    Rlam23_f = Rlam23.ravel()
+
+    low_lim = np.abs(costh_f)
+    up_lim = np.full_like(costh_f, 1.0)
+
+
+    # ## limiting the limits of the integral
+    # bad = low_lim >= up_lim
+    # low_lim[bad] = up_lim[bad]
+
+    ## creating empty arrays with same size and shape as flattened costh
+    Func3_1 = np.empty_like(costh_f)
+    Func3_2 = np.empty_like(costh_f)
+
+    ## older sin(theta) -> (sth) is super small but not zero, and we need it to be not zero when dividing
+    sth_eps = 1e-14
+
+    ## progress bar
+    if pbar:
+        tqdm_bar = tqdm(desc = "Computing electron density", total = costh_f.size)
+
+    ## quad loop for int1 and int2 of equation 5 or third term (Sugar et al)
+    for i in range(costh_f.size):
+        # filter NaNs
+        if not np.isfinite(Rlam23_f[i]) or not np.isfinite(costh_f[i]):
+            Func3_1[i] = 0.0
+            Func3_2[i] = 0.0
+        else:
+            Func3_1[i] = do_int_1_quad(low_lim[i], up_lim[i], Rlam23_f[i], costh_f[i], quad_abs_tol)
+
+            # integral 2 becomes unstable when sintheta ~ theta (we use sth_eps here)
+            if abs(sinth_f[i]) < sth_eps:
+                Func3_2[i] = 0.0
+            else:
+                Func3_2[i] = do_int_2_quad(
+                    low_lim[i], up_lim[i], Rlam23_f[i], costh_f[i], sinth_f[i], quad_abs_tol
+                )
+
+        if pbar:
+            tqdm_bar.update(1)
+
+    if pbar:
+        tqdm_bar.close()
+
+
+    # reshaping func3_1 and func3_2 after flattening
+    Func3_1 = Func3_1.reshape(orig_shape)
+    Func3_2 = Func3_2.reshape(orig_shape)
+
+    # combining terms and returning values
+    t1 = np.abs(costh) * Func1
+    t2 = costh * Func2
+    t3 = np.abs(costh) * Func3_2 + Func3_1
+    t3[np.isnan(t3)] = 0
+
+    ne_morph = np.abs(t1 + t2 + t3)
+
+
+
+    return Xlam, Ylam, Rlam, ne_morph, t1, t2, t3, quad_abs_tol
+
+
+# Sample values to test
+theta_sample = 1  # radians
+r_sample = np.array([0.001, 0.01, 0.1, 1.0, 5.0, 10.0]) #Rlam
+theta_grid = np.full_like(r_sample, theta_sample)
+
+_, _, Rlam, ne_morph, *_ = plasma_distribution_morphology_array(
+    r_sample,
+    theta_grid,
+    pbar=False,
+    quad_abs_tol=1e-5
+)
+
+print("Sample Rlam and ne_morph values:")
+for r_val, ne_val in zip(Rlam, ne_morph):
+    print("Rlam:", r_val, "ne_morph:", ne_val)
+
+### Solve r from ne_morph and theta ###
+
+# """
+# This is a solver that uses root finder to find r, 
+# such that theta and ne_morph are known. 
+# Note that we are not directly finding the root of ne_morph. 
+# Instead, we are finding the difference between the ne_morph with
+# trial r and user input ne_morph. If the difference is 0,
+# trial r gives the target ne_morph.
+
+# Returns:
+# - best_r: r found by the root finder (based on user input theta and ne_morph)
+# - best_ne_morph: ne_morph evaluated at best_r
+# - error: best_ne_morph - input_ne_morph
+
+# DISCLAIMER: r here is actually Rlam - in mean-free-path units.
+# """
+
+
+from scipy.optimize import root_scalar # solves one-dimensional equations
+
+
+def solve_r(
+        theta,
+        input_ne_morph,
+        r_min = 0.001, # to be specified later
+        r_max = 10.0, # to be specified later
+        quad_abs_tol = 1e-5,
+        xtol = 1e-6 # root finder tolerance
+): 
+
+    def ne_morph_at_r(
+            r # solver provides trial r
+    ):
+        _, _, Rlam, ne_morph, *_ = plasma_distribution_morphology_array( 
+            np.array([r]),
+            np.array([theta]), 
+            pbar = True, 
+            quad_abs_tol=quad_abs_tol
+    )
+
+        return float(ne_morph[0])
+
+    def objective(r):
+        return ne_morph_at_r(r) - input_ne_morph
+
+    f_min = objective(r_min) # ne_morph difference
+    f_max = objective(r_max)
+
+    if f_min == 0:
+        best_r = r_min
+    elif f_max == 0:
+        best_r = r_max
+    elif f_min * f_max > 0: # checking bounds
+        raise ValueError(
+            "Root is not bracketed. Values at r_min and r_max do not have opposite signs. Try changing r_min and/or r_max."
+        )
+    else:
+        result = root_scalar( # actual root finder function from scipy
+            objective,
+            bracket = [r_min, r_max],
+            method = "brentq",
+            xtol = xtol
+        )
+
+        if not result.converged:
+            raise RuntimeError("Root finder did not converge.")
+
+        best_r = result.root
+
+    best_ne_morph = ne_morph_at_r(best_r)
+    error = best_ne_morph - input_ne_morph
+
+    return best_r, best_ne_morph, error
+
+
+theta = float(input("Theta? [rad]"))
+input_ne_morph = float(input("Target ne_morph?"))
+
+best_r, best_ne_morph, error = solve_r(
+    theta = theta,
+    input_ne_morph = input_ne_morph,
+    r_min = 0.001, # to be specified later
+    r_max = 10.0, # to be specified later
+    quad_abs_tol = 1e-5
+)
+
+print("best r (in mean-free-path units, Rlam):", best_r)
+print("best ne_morph:", best_ne_morph)
+print("percent error:", error)
+
+
+
+
+
+
+'''
+    ### bookmark 2 - old code
     [Xlam, Ylam] = np.meshgrid(x_grid, y_grid)
 
     Rlam = np.sqrt(Xlam**2 + Ylam**2)
@@ -323,7 +539,7 @@ def plasma_distribution_morphology_grid(
     #     ne_morph[zero_point, :] = (ne_morph[zero_point - 1, :] + ne_morph[zero_point + 1, :]) / 2
 
     return Xlam, Ylam, Rlam, ne_morph
-
+'''
 
 def plasma_distribution(
     total_atmospheric_number_density,