Source code for jf1uids.data_classes.simulation_helper_data

from typing import NamedTuple
import jax.numpy as jnp
from jf1uids._geometry.geometry import _center_of_volume, _r_hat_alpha
from jf1uids.option_classes.simulation_config import CARTESIAN, CYLINDRICAL, SPHERICAL, SimulationConfig

# Helper data like the radii and cell volumes 
# in the simulation or cooling tables etc.

[docs] class HelperData(NamedTuple): """Helper data used throughout the simulation.""" #: The geometric centers of the cells. geometric_centers: jnp.ndarray = None #: The volumetric centers of the cells. #: Same as the geometric centers for Cartesian geometry. volumetric_centers: jnp.ndarray = None #: cell center to box center distances #: only for config.dimensionality > 1 r: jnp.ndarray = None #: A helper variable, defined as #: \hat{r}^\alpha = V_j / (2 * \alpha * \pi * \Delta r) #: with V_j the volume of cell j, \alpha the geometry factor #: and \Delta r the cell width. r_hat_alpha: jnp.ndarray = None #: The cell volumes. cell_volumes: jnp.ndarray = None #: Coordinates of the inner cell boundaries. inner_cell_boundaries: jnp.ndarray = None #: Coordinates of the outer cell boundaries. outer_cell_boundaries: jnp.ndarray = None
[docs] def get_helper_data(config: SimulationConfig) -> HelperData: """Generate the helper data for the simulation from the configuration.""" if config.dimensionality > 1: x = jnp.linspace(0, config.box_size, config.num_cells) y = jnp.linspace(0, config.box_size, config.num_cells) if config.dimensionality == 3: z = jnp.linspace(0, config.box_size, config.num_cells) geometric_centers = jnp.meshgrid(x, y, z) else: geometric_centers = jnp.meshgrid(x, y) # calculate the distances from the cell centers to the box center box_center = jnp.zeros(config.dimensionality) + config.box_size / 2 geometric_centers = jnp.array(geometric_centers) geometric_centers = jnp.moveaxis(geometric_centers, 0, -1) volumetric_centers = geometric_centers r = jnp.linalg.norm(geometric_centers - box_center, axis = -1) return HelperData(geometric_centers = geometric_centers, volumetric_centers = volumetric_centers, r = r) grid_spacing = config.box_size / (config.num_cells - 1) if config.geometry == CARTESIAN: r = jnp.linspace(0, config.box_size, config.num_cells) r_hat = grid_spacing * jnp.ones_like(r) # not really cell_volumes = grid_spacing * jnp.ones_like(r) inner_cell_boundaries = r - grid_spacing / 2 outer_cell_boundaries = r + grid_spacing / 2 return HelperData(geometric_centers = r, r_hat_alpha = r_hat, cell_volumes = cell_volumes, inner_cell_boundaries = inner_cell_boundaries, outer_cell_boundaries = outer_cell_boundaries, volumetric_centers = r) elif config.geometry == SPHERICAL or config.geometry == CYLINDRICAL: # r = jnp.linspace(- 3 * grid_spacing/2, config.box_size - 3 * grid_spacing / 2, config.num_cells) r = jnp.linspace(grid_spacing / 2, config.box_size + grid_spacing / 2, config.num_cells) inner_cell_boundaries = r - grid_spacing / 2 outer_cell_boundaries = r + grid_spacing / 2 volumetric_centers = _center_of_volume(r, grid_spacing, config.geometry) r_hat = _r_hat_alpha(r, grid_spacing, config.geometry) cell_volumes = 2 * config.geometry * jnp.pi * grid_spacing * r_hat return HelperData(geometric_centers = r, volumetric_centers = volumetric_centers, r_hat_alpha = r_hat, cell_volumes = cell_volumes, inner_cell_boundaries = inner_cell_boundaries, outer_cell_boundaries = outer_cell_boundaries)