Bump-on-tail problem; arbitrary additional ion/electron populations; separate diagnostics(...) and simulation(...)

examples/Landau_damping.py

@@ -1,7 +1,7 @@
 ## Landau_damping.py
 # Example of electric field damping in a plasma
 from jaxincell import plot
-from jaxincell import simulation
+from jaxincell import simulation, diagnostics
 import jax.numpy as jnp
 from jax import block_until_ready
 
@@ -19,7 +19,7 @@
 }
 
 solver_parameters = {
-    "field_solver"           : 0,    # Algorithm to solve E and B fields - 0: Curl_EB, 1: Gauss_1D_FFT, 2: Gauss_1D_Cartesian, 3: Poisson_1D_FFT, 
+    "field_solver"           : 0,    # Algorithm to solve E and B fields - 0: Curl_EB, 1: Gauss_1D_FFT, 2: Gauss_1D_Cartesian, 3: Poisson_1D_FFT,
     "number_grid_points"     : 81,  # Number of grid points
     "number_pseudoelectrons" : 3000, # Number of pseudoelectrons
     "total_steps"            : 200, # Total number of time steps

examples/Langmuir_wave.py

@@ -1,7 +1,7 @@
 ## Langmuir_wave.py
 # Example of plasma oscillations of electrons
 from jaxincell import plot
-from jaxincell import simulation
+from jaxincell import simulation, diagnostics
 import jax.numpy as jnp
 from jax import block_until_ready
 
@@ -18,7 +18,7 @@
 }
 
 solver_parameters = {
-    "field_solver"           : 0,    # Algorithm to solve E and B fields - 0: Curl_EB, 1: Gauss_1D_FFT, 2: Gauss_1D_Cartesian, 3: Poisson_1D_FFT, 
+    "field_solver"           : 0,    # Algorithm to solve E and B fields - 0: Curl_EB, 1: Gauss_1D_FFT, 2: Gauss_1D_Cartesian, 3: Poisson_1D_FFT,
     "number_grid_points"     : 33,  # Number of grid points
     "number_pseudoelectrons" : 3000, # Number of pseudoelectrons
     "total_steps"            : 1000, # Total number of time steps

examples/Weibel_instability.py

@@ -1,7 +1,7 @@
 ## Weibel_instability.py
 # Example of plasma oscillations of electrons
 from jaxincell import plot
-from jaxincell import simulation
+from jaxincell import simulation, diagnostics
 from jax import block_until_ready
 
 input_parameters = {
@@ -25,7 +25,7 @@
 }
 
 solver_parameters = {
-    "field_solver"           : 0,    # Algorithm to solve E and B fields - 0: Curl_EB, 1: Gauss_1D_FFT, 2: Gauss_1D_Cartesian, 3: Poisson_1D_FFT, 
+    "field_solver"           : 0,    # Algorithm to solve E and B fields - 0: Curl_EB, 1: Gauss_1D_FFT, 2: Gauss_1D_Cartesian, 3: Poisson_1D_FFT,
     "number_grid_points"     : 201,  # Number of grid points
     "number_pseudoelectrons" : 3000, # Number of pseudoelectrons
     "total_steps"            : 6000, # Total number of time steps
@@ -34,4 +34,7 @@
 
 output = block_until_ready(simulation(input_parameters, **solver_parameters))
 
+# Post-process: segregate ions/electrons, compute energies, compute FFT
+diagnostics(output)
+
 plot(output, direction="xz")  # Plot the results in x and z direction

examples/auto-differentiability.py

@@ -6,7 +6,7 @@
 import jax.numpy as jnp
 import matplotlib.pyplot as plt
 from jax import jit, grad, lax, block_until_ready, debug
-from jaxincell import simulation, load_parameters
+from jaxincell import simulation, load_parameters, diagnostics
 
 # Read from input.toml
 input_parameters, solver_parameters = load_parameters('input.toml')
@@ -15,6 +15,8 @@
 def mean_electric_field(electron_drift_speed):
     input_parameters["electron_drift_speed_x"] = electron_drift_speed
     output = block_until_ready(simulation(input_parameters, **solver_parameters))
+    # Post-process: segregate ions/electrons, compute energies, compute FFT
+    diagnostics(output)
     electric_field = jnp.mean(output['electric_field_x'][:, :, 0], axis=1)
     mean_E = jnp.mean(lax.slice(electric_field, [solver_parameters["total_steps"]//2], [solver_parameters["total_steps"]]))
     return mean_E

examples/bump-on-tail.py

@@ -0,0 +1,32 @@
+#!/usr/bin/env python
+"""
+bump-on-tail.py
+Weak electron beam on tail of bulk Maxwellian electron distribution drives
+slowly-growing Langmuir waves with Im(omega) << omega_pe.
+"""
+
+import numpy as np
+from datetime import datetime
+
+from jax import block_until_ready
+from jaxincell import plot, simulation, load_parameters, diagnostics
+
+input_parameters, solver_parameters = load_parameters('bump-on-tail.toml')
+
+# Run the simulation
+started = datetime.now()
+output = block_until_ready(simulation(input_parameters, **solver_parameters))
+print("Simulation done, elapsed:", datetime.now()-started)
+
+# Post-process: segregate ions/electrons, compute energies, compute FFT
+diagnostics(output)
+
+# Plot the results
+plot(output)
+
+# Save the output to a file
+np.savez("simulation_output.npz", **output)
+
+# # Load the output from the file
+# data = np.load("simulation_output.npz", allow_pickle=True)
+# output2 = dict(data)

examples/bump-on-tail.toml

@@ -0,0 +1,160 @@
+# ---------------------------------------------------
+# bump-on-tail.toml
+# ---------------------------------------------------
+# For the following plasma:
+#     nbeam/n0 = 0.03
+#     vbeam = 0.25*c
+#     sqrt(Te/me) = 0.05 for both bulk and beam populations
+# Fastest-growing mode predicted by non-relativistic electrostatic dispersion
+# has parameters:
+#     Re(omega) = 1.0*omega_pe
+#     Im(omega) = 0.075*omega_pe
+#     k = 4.9*omega_pe/c
+# where omega_pe is the electron plasma frequency, k = 2*pi/wavelength is the
+# wavenumber, c is speed of light, vbeam is (non-relativistic) three-velocity,
+# nbeam is beam density, n0 = background (bulk) density.
+# ---------------------------------------------------
+
+[input_parameters]
+# -----------------------------
+# species intrinsic properties
+# -----------------------------
+ion_mass_over_proton_mass = 1e9
+ion_charge_over_elementary_charge = 1
+electron_charge_over_elementary_charge = -1
+# -----------------------------
+# particle spatial distribution
+# -----------------------------
+random_positions_x = true
+random_positions_y = false
+random_positions_z = false
+amplitude_perturbation_x = 0  # Amplitude of sinusoidal perturbation
+amplitude_perturbation_y = 0
+amplitude_perturbation_z = 0
+wavenumber_electrons_x = 0  # Wavenumber of sinusoidal density perturbation (factor of 2pi/length)
+wavenumber_electrons_y = 0
+wavenumber_electrons_z = 0
+wavenumber_ions_x = 0
+wavenumber_ions_y = 0
+wavenumber_ions_z = 0
+# ------------------------------
+# electron velocity distribution
+# ------------------------------
+vth_electrons_over_c_x = 0.07071067812  # sqrt(2*T/m)/c = sqrt(2)*0.05
+vth_electrons_over_c_y = 0
+vth_electrons_over_c_z = 0
+electron_drift_speed_x = -2.25e6 # -0.25c * 3e-2  # drift in m/s; compensate bump drift so that j(t)=0
+electron_drift_speed_y = 0
+electron_drift_speed_z = 0
+velocity_plus_minus_electrons_x = false # create two groups of electrons moving in opposite directions in the x direction
+velocity_plus_minus_electrons_y = false
+velocity_plus_minus_electrons_z = false
+# ------------------------------
+# ion velocity distribution
+# ------------------------------
+ion_drift_speed_x = 0  # Drift speed of ions in the x direction
+ion_drift_speed_y = 0
+ion_drift_speed_z = 0
+ion_temperature_over_electron_temperature_x = 1e-9  # Temperature ratio of ions to electrons in the x direction
+ion_temperature_over_electron_temperature_y = 1e-9
+ion_temperature_over_electron_temperature_z = 1e-9
+velocity_plus_minus_ions_x = false  # create two groups of ions moving in opposite directions in the x direction
+velocity_plus_minus_ions_y = false
+velocity_plus_minus_ions_z = false
+# ------------------------------
+# domain, spacetime gridding
+# ------------------------------
+length = 1  # Simulation box size in meters (sets dimensionful normalization, but does not change dimensionless ratios?)
+grid_points_per_Debye_length = 2.565  # dx over Debye length (ignore the conflicting variable name)
+timestep_over_spatialstep_times_c = 1.0  # dt * speed_of_light / dx
+particle_BC_left  = 0  # 0: periodic, 1: reflective, 2: absorbing
+particle_BC_right = 0
+field_BC_left  = 0  # 0: periodic, 1: reflective, 2: absorbing
+field_BC_right = 0
+# ------------------------------
+# other controls
+# ------------------------------
+print_info = true
+relativistic = false  # Use relativistic Boris pusher
+seed = 250724  # Random seed for reproducibility
+tolerance_Picard_iterations_implicit_CN = 1e-6  # Tolerance for Picard iterations in implicit Crank-Nicholson method
+
+# ------------------------------
+# Additional Maxwellian particle populations with bulk drift.
+# To add new populations, copy this block and modify parameter values,
+# but don't change the header "[[species]]" or parameter names.
+# To only use default ion/electron plasma, comment out all [[species]] blocks.
+# All parameters must be specified; internal code doesn't set default values.
+# ------------------------------
+[[species]]
+mass_over_proton_mass = 0.0005446170199382714  # me/mp using values from jaxincell/_constants.py
+charge_over_elementary_charge = -1
+weight_ratio = 3e-2  # sets density ratio between bulk/beam, in combination with "number_pseudoparticles_species"
+# spatial distribution
+random_positions_x = true
+random_positions_y = false
+random_positions_z = false
+amplitude_perturbation_x = 0  # Amplitude of sinusoidal perturbation
+amplitude_perturbation_y = 0
+amplitude_perturbation_z = 0
+wavenumber_perturbation_x = 0  # Wavenumber of sinusoidal density perturbation (factor of 2pi/length)
+wavenumber_perturbation_y = 0
+wavenumber_perturbation_z = 0
+# velocity distribution
+vth_over_c_x = 0.07071067812  # sqrt(2*T/m)/c = sqrt(2)*0.05
+vth_over_c_y = 0
+vth_over_c_z = 0
+drift_speed_x = 7.5e7  # drift speed 0.25c converted to m/s
+drift_speed_y = 0
+drift_speed_z = 0
+# RNG control
+# "seed_position", "seed_position_override" parameters allow you to put
+# different particle species at identical positions, so as to ensure exact
+# charge neutrality at t=0 for multiple ion/electron populations.
+# WARNING: to avoid unphysically correlated coordinates, choose the position
+# seed to not coincide with RNG seeds used elsewhere in the program.
+seed_position_override = true
+seed_position = 10
+
+# enforce charge neutrality for the bump distribution
+[[species]]
+mass_over_proton_mass = 1e9
+charge_over_elementary_charge = 1
+weight_ratio = 3e-2
+# spatial distribution
+random_positions_x = true
+random_positions_y = false
+random_positions_z = false
+amplitude_perturbation_x = 0  # Amplitude of sinusoidal perturbation
+amplitude_perturbation_y = 0
+amplitude_perturbation_z = 0
+wavenumber_perturbation_x = 0  # Wavenumber of sinusoidal density perturbation (factor of 2pi/length)
+wavenumber_perturbation_y = 0
+wavenumber_perturbation_z = 0
+# velocity distribution
+vth_over_c_x = 1e-100
+vth_over_c_y = 1e-100
+vth_over_c_z = 1e-100
+drift_speed_x = 0
+drift_speed_y = 0
+drift_speed_z = 0
+# RNG control
+# "seed_position", "seed_position_override" parameters allow you to put
+# different particle species at identical positions, so as to ensure exact
+# charge neutrality at t=0 for multiple ion/electron populations.
+# WARNING: to avoid unphysically correlated coordinates, choose the position
+# seed to not coincide with RNG seeds used elsewhere in the program.
+seed_position_override = true
+seed_position = 10
+
+[solver_parameters]
+time_evolution_algorithm = 0  # 0: Boris solver, 1: implicit Crank-Nicholson
+max_number_of_Picard_iterations_implicit_CN = 30
+number_of_particle_substeps_implicit_CN = 2
+number_grid_points = 40  # number of grid CELLS, not edges/vertices
+field_solver = 0  # 0: Curl_EB, 1: Gauss_1D_FFT, 2: Gauss_1D_Cartesian, 3: Poisson_1D_FFT, 
+number_pseudoelectrons = 40000  # Total in entire domain
+total_steps = 2000  # Total number of time steps
+# number of particles for each additional species
+# provide one list value per [[species]] block
+number_pseudoparticles_species = [40000,40000,]

examples/optimize_two_stream_saturation.py

@@ -4,7 +4,7 @@
 from jax import jit, grad
 import matplotlib.pyplot as plt
 from scipy.optimize import least_squares
-from jaxincell import simulation, epsilon_0, load_parameters
+from jaxincell import simulation, epsilon_0, load_parameters, diagnostics
 
 # Read from input.toml
 input_parameters, solver_parameters = load_parameters('input.toml')
@@ -24,6 +24,8 @@ def objective_function(Ti):
     Ti = jnp.atleast_1d(Ti)[0]
     params["ion_temperature_over_electron_temperature_x"] = Ti
     output = simulation(params, **solver_parameters)
+    # Post-process: segregate ions/electrons, compute energies, compute FFT
+    diagnostics(output)
     abs_E_squared              = jnp.sum(output['electric_field']**2, axis=-1)
     integral_E_squared         = jnp.trapezoid(abs_E_squared, dx=output['dx'], axis=-1)
     energy = (epsilon_0/2) * integral_E_squared
@@ -33,6 +35,8 @@ def objective_function(Ti):
 print(f'Perform a first run to see one objective function')
 input_parameters["ion_temperature_over_electron_temperature_x"] = x0_optimization
 output = simulation(input_parameters, **solver_parameters)
+# Post-process: segregate ions/electrons, compute energies, compute FFT
+diagnostics(output)
 objective = objective_function(x0_optimization)
 plt.figure(figsize=(8,6))
 plt.plot(output['time_array']*output['plasma_frequency'],output['electric_field_energy'], label='Electric Field Energy')

examples/scaling_energy_time.py

@@ -28,7 +28,7 @@
 }
 
 solver_parameters = {
-    "field_solver"           : 1,    # Algorithm to solve E and B fields - 0: Curl_EB, 1: Gauss_1D_FFT, 2: Gauss_1D_Cartesian, 3: Poisson_1D_FFT, 
+    "field_solver"           : 1,    # Algorithm to solve E and B fields - 0: Curl_EB, 1: Gauss_1D_FFT, 2: Gauss_1D_Cartesian, 3: Poisson_1D_FFT,
     "number_grid_points"     : 50,   # Number of grid points
     "number_pseudoelectrons" : 1500, # Number of pseudoelectrons
     "total_steps"            : 2000, # Total number of time steps
@@ -46,7 +46,7 @@
 
 # Function to compute the maximum relative energy error
 def max_relative_energy_error(output):
-    diagnostics(output)
+    #diagnostics(output)  # moved to outer caller level to be more consistent with other jaxincell example problems
     relative_energy_error = jnp.abs((output["total_energy"] - output["total_energy"][0]) / output["total_energy"][0])
     return jnp.max(relative_energy_error)
 
@@ -78,6 +78,8 @@ def measure_time_and_error(parameter_list, param_name):
             #     solver_parameters['number_grid_points'] = old_grid_points
         elapsed_time = time.time() - start
         times.append(elapsed_time)
+        # Post-process: segregate ions/electrons, compute energies, compute FFT
+        diagnostics(output)
         max_relative_errors.append(max_relative_energy_error(output))
         print(f"{param_name.capitalize()}: {param}, Time: {elapsed_time}s")
     return times, max_relative_errors
@@ -147,4 +149,4 @@ def plot_and_save_results(x_data_list, y_data_list, x_labels, y_labels, file_nam
 # Plotting the relative energy error
 plot_and_save_results(x_data_list, error_y_data_list, x_labels, [error_y_label] * 4, 'scaling_energy_error.pdf')
 
-plt.show()
+plt.show()

examples/two-stream_instability.py

@@ -2,9 +2,6 @@
 import time
 from jax import block_until_ready
 from jaxincell import plot, simulation, load_parameters
-import numpy as np
-import pickle
-import json
 
 # Read from input.toml
 input_parameters, solver_parameters = load_parameters('input.toml')
@@ -22,8 +19,10 @@
 plot(output)
 
 # # Save the output to a file
+# import numpy as np
 # np.savez("simulation_output.npz", **output)
 
 # # Load the output from the file
+# import numpy as np
 # data = np.load("simulation_output.npz", allow_pickle=True)
 # output2 = dict(data)

jaxincell/__init__.py

@@ -5,4 +5,4 @@
 from ._particles import *
 from ._plot import *
 from ._simulation import *
-from ._sources import *
+from ._sources import *

jaxincell/__main__.py

@@ -19,4 +19,4 @@ def main(cl_args=sys.argv[1:]):
     plot(output)
 
 if __name__ == "__main__":
-    main(sys.argv[1:])
+    main(sys.argv[1:])