jaxincell._simulation

Classes

Simulation

Calling Simulation(parameters) will create a Simulation object with the provided parameters.

Functions

load_parameters(input_file)

Load parameters from a given .toml input file given the path to the file.

Module Contents

jaxincell._simulation.load_parameters(input_file)

Load parameters from a given .toml input file given the path to the file.

class jaxincell._simulation.Simulation(parameters=None)

Calling Simulation(parameters) will create a Simulation object with the provided parameters. The parameters should be provided as a dictionary, but can also be provided as a path to a .toml file containing the parameters. The parameters will be cleaned and initialized using the appropriate cleaner functions for each parameter section, and the simulation state will be initialized based on the cleaned parameters.

The Simulation object will expose any differentiable parameters that were provided in input_parameters through Simulation_object.input_parameters. These input_parameters can then be passed to the simulation() or run() bound functions to run a simulation with these input_parameters exposed such that grads can be taken with respect to them. If simulation() or run() is called without input_parameters, it will use the parameters provided at initialization (which may include user-provided parameters and defaults) and will not overwrite any parameters with the input_parameters.

Example without input parameters:

sim = Simulation(parameters) simulation_output = sim.simulation() or simulation_output = sim.run()

Example with input parameters:

sim = Simulation(parameters) input_parameters = sim.input_parameters simulation_output = sim.simulation(input_parameters) or simulation_output = sim.run(input_parameters)

Additionally,

def scalar_objective(input_parameters):

simulation_output = sim.run(input_parameters) return some_scalar_function_of(simulation_output)

grad(scalar_objective)(input_parameters)

will provide the gradient of some scalar objective function of the simulation output with respect to the input_parameters.

simulation(input_parameters=None)

Exposed simulation call which doesn’t expose the hash values for each section to prevent unintentionally forcing recompiles or not recompiling when necessary.

input_parameters is a dictionary of differentiable parameters If simulation is called without input parameters, it will use the user input parameters or default parameters previously provided. input_parameters will overwrite any differentiable parameters previously provided. This is meant to make it simple to expose grads derivatives with respect to the input parameters.

run(input_parameters=None)

_simulation(input_parameters=None, domain_hash='', species_hash='', external_field_hash='', source_hash='', solver_hash='')

Run a plasma physics simulation using a Particle-In-Cell (PIC) method in JAX.

This function simulates the evolution of a plasma system by solving for particle motion (electrons and ions) and self-consistent electromagnetic fields on a grid. It uses the Boris algorithm for particle updates and a leapfrog scheme for field updates.

Parameters:

user_parametersdict

User-defined parameters for the simulation. These can include: - Physical parameters: box size, number of particles, thermal velocities. - Numerical parameters: grid resolution, time step size. - Boundary conditions for particles and fields. - Random seed for reproducibility.

Returns:

output : dict

assemble_output(simulation_output, input_parameters)

clean_and_initialize_parameters(parameters)

classify_and_sort_input_parameters(parameters)

Sort through input parameters to move parameters into their respective dictionaries to overwrite defaults and move differentiable parameters into a separate input_parameters dictionary. This input_parameters dictionary can then be accessed to use them as inputs to the simulation function without having to write multiple toml files for the differentiable inputs and the non-differentiable parameters.

build_hash_values()

Build the hashes for each of the parameter sections to help with determining when Jax needs to recompile the _simulation() function due to new parameters being passed.

reinitialize_simulation_state()

Reinitialize the simulation state based on the current parameter sections. This should be called whenever parameters are updated after initialization to ensure that the simulation state is consistent with the new parameters.

clean_runtime_input_parameters(input_parameters=None)

Clean the input_parameters provided to the simulation(…) or run(…) functions at runtime.

current_domain_state()

build_domain()

initialize_particles()

initialize_fields()

set_parameter_section(section_name, new_parameters)

Helper which is used by setters for the parameter sections to automatically clean parameters and reinitialize the state of the simulation including creating new hashes.

property domain_parameters

property species_parameters

property external_field_parameters

property source_parameters

property solver_parameters

property input_parameters