{ "cells": [ { "cell_type": "markdown", "metadata": { "id": "SxheI9KkPF3D" }, "source": [ "# Jet in astronomix" ] }, { "cell_type": "markdown", "metadata": { "id": "fwKNzezpT1Yb" }, "source": [ "## Imports" ] }, { "cell_type": "code", "execution_count": 1, "metadata": { "id": "bZQxDqwjKcpu" }, "outputs": [], "source": [ "# ==== GPU selection ====\n", "from autocvd import autocvd\n", "autocvd(num_gpus = 1)\n", "# ruff: noqa: E402\n", "# =======================\n", "\n", "# general\n", "import jax.numpy as jnp\n", "\n", "# astronomix constants\n", "from astronomix import (\n", " PERIODIC_BOUNDARY,\n", ")\n", "\n", "# astronomix containers\n", "from astronomix import (\n", " SimulationConfig,\n", " SimulationParams,\n", " BoundarySettings,\n", " BoundarySettings1D,\n", " SnapshotSettings,\n", ")\n", "\n", "# astronomix functions\n", "from astronomix import (\n", " get_registered_variables,\n", " construct_primitive_state,\n", " time_integration,\n", " finalize_config,\n", " setup_magnetic_fields_from_vector_potential,\n", ")\n", "\n", "# plotting\n", "import matplotlib.pyplot as plt\n", "import matplotlib.animation as animation\n", "from IPython.display import HTML" ] }, { "cell_type": "markdown", "metadata": { "id": "kYaOj5N2T9SN" }, "source": [ "## Setup simulation" ] }, { "cell_type": "code", "execution_count": 2, "metadata": { "colab": { "base_uri": "https://localhost:8080/" }, "id": "iyx0rUp7KZ5W", "outputId": "6b419755-066e-4d15-fba1-f2431bec8dc7" }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "OPTIMAL_BACKEND: using the PALLAS backend (GPU compute capability >= 8.0).\n", "For 3D simulations with periodic boundaries, setting boundary handling to PERIODIC_ROLL and num_ghost_cells to 0 for better performance.\n", "Setting time integrator to RK4_SSP for finite difference solver mode.\n" ] } ], "source": [ "gamma = 5/3\n", "\n", "box_size = 24.0\n", "num_cells = 64\n", "grid_spacing = box_size / num_cells\n", "x_center = box_size / 2.0\n", "y_center = box_size / 2.0\n", "z_center = box_size / 2.0\n", "\n", "config = SimulationConfig(\n", " grid_spacing = grid_spacing,\n", " mhd = True,\n", " progress_bar = True,\n", " dimensionality = 3,\n", " box_size = box_size,\n", " num_cells = num_cells,\n", " boundary_settings = BoundarySettings(\n", " BoundarySettings1D(\n", " left_boundary = PERIODIC_BOUNDARY,\n", " right_boundary = PERIODIC_BOUNDARY\n", " ),\n", " BoundarySettings1D(\n", " left_boundary = PERIODIC_BOUNDARY,\n", " right_boundary = PERIODIC_BOUNDARY\n", " ),\n", " BoundarySettings1D(\n", " left_boundary = PERIODIC_BOUNDARY,\n", " right_boundary = PERIODIC_BOUNDARY\n", " )\n", " ),\n", " return_snapshots = True,\n", " num_snapshots = 100,\n", " snapshot_settings = SnapshotSettings(\n", " return_states=True,\n", " return_final_state=True,\n", " return_magnetic_divergence=True\n", " ),\n", ")\n", "\n", "registered_variables = get_registered_variables(config)\n", "\n", "C_CFL = 1.5\n", "\n", "rho_0 = 1.0\n", "p_0 = 1.0\n", "\n", "# define the vector potential function\n", "def jet_vector_potential(X, Y, Z):\n", " r = jnp.sqrt((X - x_center)**2 + (Y - y_center)**2 + (Z - z_center)**2)\n", " A0 = 20.0\n", "\n", " A_x = -jnp.exp(-r ** 2) * (Y - y_center)\n", " A_y = jnp.exp(-r ** 2) * (X - x_center)\n", " A_z = 0.5 * A0 * jnp.exp(-r ** 2)\n", "\n", " return A_x, A_y, A_z\n", "\n", "# init b fields from vector potential\n", "B_x, B_y, B_z, bxb, byb, bzb = setup_magnetic_fields_from_vector_potential(\n", " config=config,\n", " vector_potential_func=jet_vector_potential\n", ")\n", "\n", "# Set up primitive hydro arrays\n", "rho = jnp.ones((config.num_cells, config.num_cells, config.num_cells)) * rho_0\n", "u_x = jnp.zeros((config.num_cells, config.num_cells, config.num_cells))\n", "u_y = jnp.zeros((config.num_cells, config.num_cells, config.num_cells))\n", "u_z = jnp.zeros((config.num_cells, config.num_cells, config.num_cells))\n", "p = jnp.ones((config.num_cells, config.num_cells, config.num_cells)) * p_0\n", "\n", "# simulation params\n", "params = SimulationParams(\n", " C_cfl = C_CFL,\n", " t_end = 5.0,\n", " gamma = gamma,\n", ")\n", "\n", "# construct primitive state\n", "initial_state = construct_primitive_state(\n", " config = config,\n", " registered_variables=registered_variables,\n", " density = rho,\n", " velocity_x = u_x,\n", " velocity_y = u_y,\n", " velocity_z = u_z,\n", " gas_pressure = p,\n", " magnetic_field_x = B_x,\n", " magnetic_field_y = B_y,\n", " magnetic_field_z = B_z,\n", " interface_magnetic_field_x = bxb,\n", " interface_magnetic_field_y = byb,\n", " interface_magnetic_field_z = bzb,\n", ")\n", "\n", "config = finalize_config(config, initial_state.shape)" ] }, { "cell_type": "markdown", "metadata": { "id": "cQp7wdrrUAEn" }, "source": [ "## Run the simulation" ] }, { "cell_type": "code", "execution_count": 3, "metadata": { "colab": { "background_save": true }, "id": "ZFDl60IgSr5K", "outputId": "c7f1933b-10d8-463e-e5db-60dcf076989e" }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ " |████████████████████████████████████████████████████████████████████| 100.0% \n" ] } ], "source": [ "result = time_integration(initial_state, config, params, registered_variables)" ] }, { "cell_type": "markdown", "metadata": { "id": "BObkRwpMUOtD" }, "source": [ "## Animate the result" ] }, { "cell_type": "code", "execution_count": 4, "metadata": { "colab": { "background_save": true }, "id": "V8c3Y-o5Taru" }, "outputs": [ { "data": { "text/html": [ "\n", "\n", "\n", "\n", "\n", "\n", "