Jet in astronomix

Imports

# ==== GPU selection ====
from autocvd import autocvd
autocvd(num_gpus = 1)
# ruff: noqa: E402
# =======================

# general
import jax.numpy as jnp

# astronomix constants
from astronomix import (
    PERIODIC_BOUNDARY,
)

# astronomix containers
from astronomix import (
    SimulationConfig,
    SimulationParams,
    BoundarySettings,
    BoundarySettings1D,
    SnapshotSettings,
)

# astronomix functions
from astronomix import (
    get_registered_variables,
    construct_primitive_state,
    time_integration,
    finalize_config,
    setup_magnetic_fields_from_vector_potential,
)

# plotting
import matplotlib.pyplot as plt
import matplotlib.animation as animation
from IPython.display import HTML

Setup simulation

gamma = 5/3

box_size = 24.0
num_cells = 64
grid_spacing = box_size / num_cells
x_center = box_size / 2.0
y_center = box_size / 2.0
z_center = box_size / 2.0

config = SimulationConfig(
    grid_spacing = grid_spacing,
    mhd = True,
    progress_bar = True,
    dimensionality = 3,
    box_size = box_size,
    num_cells = num_cells,
    boundary_settings = BoundarySettings(
        BoundarySettings1D(
            left_boundary = PERIODIC_BOUNDARY,
            right_boundary = PERIODIC_BOUNDARY
        ),
        BoundarySettings1D(
            left_boundary = PERIODIC_BOUNDARY,
            right_boundary = PERIODIC_BOUNDARY
        ),
        BoundarySettings1D(
            left_boundary = PERIODIC_BOUNDARY,
            right_boundary = PERIODIC_BOUNDARY
        )
    ),
    return_snapshots = True,
    num_snapshots = 100,
    snapshot_settings = SnapshotSettings(
      return_states=True,
      return_final_state=True,
      return_magnetic_divergence=True
    ),
)

registered_variables = get_registered_variables(config)

C_CFL = 1.5

rho_0 = 1.0
p_0 = 1.0

# define the vector potential function
def jet_vector_potential(X, Y, Z):
    r = jnp.sqrt((X - x_center)**2 + (Y - y_center)**2 + (Z - z_center)**2)
    A0 = 20.0

    A_x = -jnp.exp(-r ** 2) * (Y - y_center)
    A_y = jnp.exp(-r ** 2) * (X - x_center)
    A_z = 0.5 * A0 * jnp.exp(-r ** 2)

    return A_x, A_y, A_z

# init b fields from vector potential
B_x, B_y, B_z, bxb, byb, bzb = setup_magnetic_fields_from_vector_potential(
    config=config,
    vector_potential_func=jet_vector_potential
)

# Set up primitive hydro arrays
rho = jnp.ones((config.num_cells, config.num_cells, config.num_cells)) * rho_0
u_x = jnp.zeros((config.num_cells, config.num_cells, config.num_cells))
u_y = jnp.zeros((config.num_cells, config.num_cells, config.num_cells))
u_z = jnp.zeros((config.num_cells, config.num_cells, config.num_cells))
p = jnp.ones((config.num_cells, config.num_cells, config.num_cells)) * p_0

# simulation params
params = SimulationParams(
    C_cfl = C_CFL,
    t_end = 5.0,
    gamma = gamma,
)

# construct primitive state
initial_state = construct_primitive_state(
    config = config,
    registered_variables=registered_variables,
    density = rho,
    velocity_x = u_x,
    velocity_y = u_y,
    velocity_z = u_z,
    gas_pressure = p,
    magnetic_field_x = B_x,
    magnetic_field_y = B_y,
    magnetic_field_z = B_z,
    interface_magnetic_field_x = bxb,
    interface_magnetic_field_y = byb,
    interface_magnetic_field_z = bzb,
)

config = finalize_config(config, initial_state.shape)

OPTIMAL_BACKEND: using the PALLAS backend (GPU compute capability >= 8.0).
For 3D simulations with periodic boundaries, setting boundary handling to PERIODIC_ROLL and num_ghost_cells to 0 for better performance.
Setting time integrator to RK4_SSP for finite difference solver mode.

Run the simulation

result = time_integration(initial_state, config, params, registered_variables)

 |████████████████████████████████████████████████████████████████████| 100.0%  

Animate the result

# what we want to plot
y_index = num_cells // 2
density_slices = result.states[:, registered_variables.density_index, :, y_index, :]

fig_anim, ax_anim = plt.subplots(1, 1, figsize=(6, 6))

# init image
current_slice = density_slices[0, :, :].T
img = ax_anim.imshow(current_slice, cmap="YlOrRd", origin='lower')
ax_anim.set_title("Density Slice Animation")
ax_anim.set_xlabel("X-axis")
ax_anim.set_ylabel("Z-axis")

# update image
def update(frame):
    current_slice = density_slices[frame, :, :].T
    img.set_array(current_slice)
    img.set_clim(current_slice.min(), current_slice.max())
    ax_anim.set_title(f"Density Slice Animation")
    return [img]

# create anim
ani = animation.FuncAnimation(
    fig_anim,
    update,
    frames=density_slices.shape[0],
    blit=True,
    interval=50 # mills between frames
)

fig_anim.tight_layout()
plt.close(fig_anim)

HTML(ani.to_jshtml())

Once Loop Reflect