"""Plot the cell geometry given a set of cell properties. This script generates a 2D grid diagram of cells with a nested region. The script takes parameters from a YAML file to configure the grid structure. Example: python plot_zoom_geometry.py --params params.yaml """ import argparse from dataclasses import dataclass import h5py import matplotlib.pyplot as plt import numpy as np import yaml from matplotlib.lines import Line2D # Declare the task diff grav constant zoom_bkg_subdepth_diff_grav_default = 0 zoom_bkg_subdepth_diff_grav = zoom_bkg_subdepth_diff_grav_default class StoreDictKeyPair(argparse.Action): """A class for storing key-value pairs from the command line.""" def __call__(self, parser, namespace, values, option_string=None): """ Store the key-value pair in the namespace. Args: parser (argparse.ArgumentParser): The parser instance. namespace (argparse.Namespace): The namespace to store the key-value pair in. values (str): The key-value pair to store. option_string (str): The option string. """ key = option_string.lstrip("--") if not hasattr(namespace, "swaps"): setattr(namespace, "swaps", {}) getattr(namespace, "swaps")[key] = values @dataclass class Space: """A dataclass to represent the space properties.""" box_size: float bkg_top_level_cells: int zoom_cell_depth: int = 2 bkg_cdim: np.ndarray = None buffer_cell_depth: int = 0 region_pad_factor: float = 1.1 zoom_shift: list = None com: list = None void_lower_bounds: list = None void_upper_bounds: list = None void_dim: list = None zoom_props: dict = None bkg_cells: list = None zoom_cells: list = None buffer_cells: list = None cdim: np.ndarray = np.zeros(3, dtype=int) width: np.ndarray = np.zeros(3) dim: np.ndarray = np.zeros(3) iwidth: np.ndarray = np.zeros(3) def __post_init__(self): """Post-initialization method to set default values.""" if self.bkg_cdim is None: self.bkg_cdim = np.full(3, self.bkg_top_level_cells) if self.zoom_shift is None: self.zoom_shift = [0.0, 0.0, 0.0] if self.com is None: self.com = [0.0, 0.0, 0.0] if self.void_lower_bounds is None: self.void_lower_bounds = [0.0, 0.0, 0.0] if self.void_upper_bounds is None: self.void_upper_bounds = [0.0, 0.0, 0.0] if self.void_dim is None: self.void_dim = [0.0, 0.0, 0.0] if self.zoom_props is None: self.zoom_props = { "region_pad_factor": self.region_pad_factor, "neighbour_max_tree_depth": -1, "bkg_subdepth_diff_grav": zoom_bkg_subdepth_diff_grav_default, "bkg_cdim": self.bkg_cdim, "buffer_cell_depth": self.buffer_cell_depth, "zoom_cell_depth": self.zoom_cell_depth, "zoom_shift": self.zoom_shift, "com": self.com, "void_lower_bounds": self.void_lower_bounds, "void_upper_bounds": self.void_upper_bounds, "void_dim": self.void_dim, "buffer_lower_bounds": [0.0, 0.0, 0.0], "buffer_upper_bounds": [0.0, 0.0, 0.0], "buffer_dim": [0.0, 0.0, 0.0], "buffer_cdim": [0, 0, 0], "buffer_width": [0.0, 0.0, 0.0], "buffer_iwidth": [0.0, 0.0, 0.0], "region_lower_bounds": [0.0, 0.0, 0.0], "region_upper_bounds": [0.0, 0.0, 0.0], "dim": [0.0, 0.0, 0.0], "cdim": [0, 0, 0], "width": [0.0, 0.0, 0.0], "iwidth": [0.0, 0.0, 0.0], "cell_min": 0.0, "nr_bkg_cells": 0, "nr_buffer_cells": 0, "nr_zoom_cells": 0, "bkg_cell_offset": 0, "buffer_cell_offset": 0, } self.dim = np.array([self.box_size] * 3) if self.bkg_cells is None: self.bkg_cells = [] if self.zoom_cells is None: self.zoom_cells = [] if self.buffer_cells is None: self.buffer_cells = [] class Cell: """A class to represent a cell in the grid.""" def __init__(self, x, y, width, ctype="background"): """ Initialize the cell with the given x, y coordinates and width. Args: x (int): The x coordinate of the lower left corner of the cell. y (int): The y coordinate of the lower left corner of the cell. width (int): The width of the cell. """ self.x = x self.y = y self.width = width self.type = ctype def draw_cell(self, ax): """ Draw the cell on the given axis. Args: ax (matplotlib.axes.Axes): The axis to draw the cell on. """ # Define the color based on the cell type if self.type == "background": color = "#bbe6e4" order = 0 elif self.type == "buffer": color = "#42bfdd" order = 1 elif self.type == "zoom": color = "#084b83" order = 2 else: raise ValueError(f"Unknown cell type: {self.type}") ax.add_patch( plt.Rectangle( (self.x, self.y), self.width[0], self.width[1], fill=True, facecolor=color, edgecolor="black", zorder=order, alpha=0.9, ) ) def read_params(params_file): """ Read the parameters from the given YAML file. Args: params_file (str): The path to the YAML file containing the parameters. Returns: dict: A dictionary containing the parameters read from the file. """ # Read the parameters from the YAML file with open(params_file, "r") as file: params = yaml.safe_load(file) return params def zoom_parse_params(params): """ Parse parameters for ZoomRegion properties and update Space instance. Args: params (dict): Dictionary containing the parameters. Returns: Space: An instance of Space populated with parameters. """ global zoom_bkg_subdepth_diff_grav # Extract the initial conditions filename ic_filename = params["InitialConditions"]["file_name"] # Set the zoom cell depth zoom_cell_depth = params["ZoomRegion"].get("zoom_top_level_depth", 2) # Set the target background cdim bkg_cdim = params["ZoomRegion"].get("bkg_top_level_cells", 16) # Get the buffer cell depth buffer_cell_depth = params["ZoomRegion"].get("buffer_top_level_depth", 0) # Ensure the buffer cell depth is less than the zoom cell depth if buffer_cell_depth > zoom_cell_depth: raise ValueError("Buffer cell depth must be less than the zoom cell depth.") # Extract the zoom width boost factor region_pad_factor = params["ZoomRegion"].get("region_pad_factor", 1.1) # Extract the minimum difference for background cells zoom_bkg_subdepth_diff_grav = params["ZoomRegion"].get( "bkg_subdepth_diff_grav", zoom_bkg_subdepth_diff_grav_default ) # Open the HDF5 file and read the box size with h5py.File(ic_filename, "r") as f: if params["InitialConditions"].get("cleanup_h_factors", False): h = params["Cosmology"].get("h", 0.6777) else: h = 1.0 box_size = f["Header"].attrs["BoxSize"] / h # Create and return the Space instance return Space( box_size=box_size, bkg_top_level_cells=bkg_cdim, zoom_cell_depth=zoom_cell_depth, buffer_cell_depth=buffer_cell_depth, region_pad_factor=region_pad_factor, ) def zoom_get_cdim_at_depth(region_dim, parent_width, child_depth): """ Compute the number of child cells in a single parent cell at a given depth. Args: region_dim (float): The dimension of the region. parent_width (float): The width of the parent cell. child_depth (int): The depth of the child cell within the parent. Returns: int: The number of child cells in a single parent cell. """ # How many parent_widths are in the region? (ensure correct rounding) region_parent_cdim = np.floor((region_dim + (0.1 * parent_width)) / parent_width) # Calculate cdim as the number of parents times the number of children per parent return int(region_parent_cdim * (2 ** child_depth)) def zoom_get_region_dim_and_shift(space, params): """ Compute the zoom region centre and boundaries. Args: space (Space): The space object containing the properties. ic_filename (str): The path to the initial conditions HDF5 file. Returns: float: The maximum side length of the zoom region. """ # Are we cleaning up h factors? if params["InitialConditions"].get("cleanup_h_factors", False): h = params["Cosmology"].get("h", 0.6777) else: h = 1.0 # Open the HDF5 file and read particle data with h5py.File(params["InitialConditions"]["file_name"], "r") as f: particles = f["PartType1/Coordinates"][:] / h masses = f["PartType1/Masses"][:] / h # Initialize boundaries and center of mass (COM) calculations min_bounds = np.min(particles, axis=0) max_bounds = np.max(particles, axis=0) com = np.average(particles, axis=0, weights=masses) # Calculate initial dimensions and midpoint ini_dims = max_bounds - min_bounds midpoint = min_bounds + (ini_dims / 2.0) box_mid = np.array([space.box_size / 2.0] * 3) # Calculate the shift needed to place the midpoint at the center # of the box zoom_shift = box_mid - midpoint # Adjust shift if incremental is too small for i in range(3): if abs(zoom_shift[i]) < 0.01 * space.box_size: zoom_shift[i] = 0.0 # If the volume isn't periodic, we cannot shift for i in range(3): if abs(zoom_shift[i]) > 0.01 * space.box_size: raise ValueError( "Cannot shift the zoom region to the centre of the box " "when the box is not periodic." ) # Update space with computed shift and COM space.zoom_shift = zoom_shift.tolist() space.zoom_props["zoom_shift"] = zoom_shift.tolist() space.com = (com + zoom_shift).tolist() space.zoom_props["com"] = (com + zoom_shift).tolist() # Compute the maximum side length of the zoom region ini_dim = max(ini_dims) return ini_dim def zoom_get_void_geometry(space, region_dim): """ Compute the void region geometry. The void region is the region covered by background cells above the zoom region. If the void region is sufficiently close to the zoom region size, then the two will be made equivalent later on. Otherwise, the void region is equivalent to the buffer region. Args: space (Space): The space object containing the properties. region_dim (float): The dimension of the zoom region before tessellating the volume. Returns: int: The number of zoom regions that tessellate the void region. """ # Get the lower and upper bounds of the zoom region based on initial dimensions and padding factor lower_bounds = [(space.box_size / 2) - (region_dim / 2.0) for _ in range(3)] upper_bounds = [(space.box_size / 2) + (region_dim / 2.0) for _ in range(3)] # Assign these bounds to the zoom region properties temporarily space.zoom_props["region_lower_bounds"] = lower_bounds space.zoom_props["region_upper_bounds"] = upper_bounds # Find the background cell edges that contain these bounds print(lower_bounds, space.box_size, space.bkg_cdim) void_lower_bounds = [ np.floor(lower_bounds[i] * space.bkg_cdim[i] / space.box_size) * (space.box_size / space.bkg_cdim[i]) for i in range(3) ] void_upper_bounds = [ (np.floor(upper_bounds[i] * space.bkg_cdim[i] / space.box_size) + 1) * (space.box_size / space.bkg_cdim[i]) for i in range(3) ] # Assign the void bounds space.void_lower_bounds = void_lower_bounds space.void_upper_bounds = void_upper_bounds # Compute the void region dimensions space.void_dim = [void_upper_bounds[i] - void_lower_bounds[i] for i in range(3)] # Compute the number of zoom regions that tessellate the void region nr_zoom_regions = int(np.ceil(space.void_dim[0] / region_dim)) return nr_zoom_regions def zoom_get_geometry_no_buffer_cells(space): """ Compute the geometry when no buffer cells are needed. Args: space (Space): The space object containing the properties. """ # If we have a buffer cell depth, warn that we will ignore it if space.buffer_cell_depth > 0: print("No buffer cells are needed, ignoring buffer cell depth.") space.buffer_cell_depth = 0 # Zero the buffer region properties explicitly space.zoom_props["buffer_lower_bounds"] = [0.0, 0.0, 0.0] space.zoom_props["buffer_upper_bounds"] = [0.0, 0.0, 0.0] space.zoom_props["buffer_dim"] = [0.0, 0.0, 0.0] space.zoom_props["buffer_cdim"] = [0, 0, 0] space.zoom_props["buffer_width"] = [0.0, 0.0, 0.0] # Match the zoom region bounds to the void region bounds space.zoom_props["region_lower_bounds"] = space.void_lower_bounds space.zoom_props["region_upper_bounds"] = space.void_upper_bounds # Compute the zoom region dimensions space.zoom_props["dim"] = [ space.zoom_props["region_upper_bounds"][i] - space.zoom_props["region_lower_bounds"][i] for i in range(3) ] # Compute the number of zoom cells in the void region cdim = zoom_get_cdim_at_depth( space.zoom_props["dim"][0], space.box_size / space.bkg_top_level_cells, space.zoom_cell_depth, ) # Compute the zoom cdim and cell width space.zoom_props["cdim"] = [cdim] * 3 space.zoom_props["width"] = [space.zoom_props["dim"][i] / cdim for i in range(3)] space.zoom_props["iwidth"] = [1.0 / space.zoom_props["width"][i] for i in range(3)] def zoom_get_geometry_with_buffer_cells(space): """ Compute the geometry when buffer cells are needed. Args: space (Space): The space object containing the properties. """ # Ensure we have a buffer cell depth if space.buffer_cell_depth == 0: zoom_report_cell_properties(space) raise ValueError( "Current cell structure requires buffer cells but no buffer cell depth has " "been given. ZoomRegion:buffer_top_level_depth must be greater than 0." ) # Match the buffer region bounds to the void region bounds space.zoom_props["buffer_lower_bounds"] = space.void_lower_bounds[:] space.zoom_props["buffer_upper_bounds"] = space.void_upper_bounds[:] # Compute the buffer region dimensions space.zoom_props["buffer_dim"] = [ space.zoom_props["buffer_upper_bounds"][i] - space.zoom_props["buffer_lower_bounds"][i] for i in range(3) ] # Compute the number of buffer cells in the void region buffer_cdim = zoom_get_cdim_at_depth( space.zoom_props["buffer_dim"][0], space.box_size / space.bkg_top_level_cells, space.buffer_cell_depth, ) # Compute the buffer cdim and cell width space.zoom_props["buffer_cdim"] = [buffer_cdim] * 3 space.zoom_props["buffer_width"] = [ space.zoom_props["buffer_dim"][i] / buffer_cdim for i in range(3) ] space.zoom_props["buffer_iwidth"] = [ 1.0 / space.zoom_props["buffer_width"][i] for i in range(3) ] # Find the buffer cell edges that contain the zoom region bounds region_lower_bounds = [] region_upper_bounds = [] for i in range(3): lower = int( np.floor( ( space.zoom_props["region_lower_bounds"][i] - space.zoom_props["buffer_lower_bounds"][i] ) * space.zoom_props["buffer_iwidth"][i] ) ) upper = int( np.floor( ( space.zoom_props["region_upper_bounds"][i] - space.zoom_props["buffer_lower_bounds"][i] ) * space.zoom_props["buffer_iwidth"][i] ) ) region_lower_bounds.append( lower * space.zoom_props["buffer_width"][i] + space.zoom_props["buffer_lower_bounds"][i] ) region_upper_bounds.append( (upper + 1) * space.zoom_props["buffer_width"][i] + space.zoom_props["buffer_lower_bounds"][i] ) # Assign the new aligned zoom bounds space.zoom_props["region_lower_bounds"] = region_lower_bounds space.zoom_props["region_upper_bounds"] = region_upper_bounds # Compute the zoom region dimensions space.zoom_props["dim"] = [ region_upper_bounds[i] - region_lower_bounds[i] for i in range(3) ] # Compute the number of zoom cells in the zoom region cdim = zoom_get_cdim_at_depth( space.zoom_props["dim"][0], space.zoom_props["buffer_width"][0], space.zoom_cell_depth - space.buffer_cell_depth, ) # Compute the zoom cdim and cell width space.zoom_props["cdim"] = [cdim] * 3 space.zoom_props["width"] = [space.zoom_props["dim"][i] / cdim for i in range(3)] space.zoom_props["iwidth"] = [1.0 / space.zoom_props["width"][i] for i in range(3)] def zoom_report_cell_properties(space): """ Report Zoom Region Properties. This function prints out a table containing the properties of the zoom region, if it is enabled. The table includes information such as dimensions, center, CDIM, background CDIM, buffer CDIM, region buffer ratio, zoom boost factor, minimum zoom cell width, background cell width, buffer width, and the number of wanderers. Args: space (Space): The space object containing the properties. """ zoom_props = space.zoom_props # Cdims print( f"{'Background cdim':>28} = [{space.bkg_cdim[0]}, {space.bkg_cdim[1]}, {space.bkg_cdim[2]}]" ) if zoom_props.get("with_buffer_cells"): print( f"{'Buffer cdim':>28} = [{zoom_props['buffer_cdim'][0]}, {zoom_props['buffer_cdim'][1]}, {zoom_props['buffer_cdim'][2]}]" ) print( f"{'Zoom cdim':>28} = [{zoom_props['cdim'][0]}, {zoom_props['cdim'][1]}, {zoom_props['cdim'][2]}]" ) # Dimensions print( f"{'Background Dimensions':>28} = [{space.box_size}, {space.box_size}, {space.box_size}]" ) if zoom_props.get("with_buffer_cells"): print( f"{'Buffer Region Dimensions':>28} = [{zoom_props['buffer_dim'][0]}, {zoom_props['buffer_dim'][1]}, {zoom_props['buffer_dim'][2]}]" ) print( f"{'Zoom Region Dimensions':>28} = [{zoom_props['dim'][0]}, {zoom_props['dim'][1]}, {zoom_props['dim'][2]}]" ) # Cell Widths print( f"{'Background Cell Width':>28} = [{space.box_size / space.bkg_top_level_cells}, {space.box_size / space.bkg_top_level_cells}, {space.box_size / space.bkg_top_level_cells}]" ) if zoom_props.get("with_buffer_cells"): print( f"{'Buffer Cell Width':>28} = [{zoom_props['buffer_width'][0]}, {zoom_props['buffer_width'][1]}, {zoom_props['buffer_width'][2]}]" ) print( f"{'Zoom Cell Width':>28} = [{zoom_props['width'][0]}, {zoom_props['width'][1]}, {zoom_props['width'][2]}]" ) # Number of Cells print(f"{'Number of Background Cells':>28} = {zoom_props.get('nr_bkg_cells', 0)}") if zoom_props.get("with_buffer_cells"): print( f"{'Number of Buffer Cells':>28} = {zoom_props.get('nr_buffer_cells', 0)}" ) print(f"{'Number of Zoom Cells':>28} = {zoom_props.get('nr_zoom_cells', 0)}") # Bounds if zoom_props.get("with_buffer_cells"): print( f"{'Buffer Bounds':>28} = [{zoom_props['buffer_lower_bounds'][0]}-{zoom_props['buffer_upper_bounds'][0]}, {zoom_props['buffer_lower_bounds'][1]}-{zoom_props['buffer_upper_bounds'][1]}, {zoom_props['buffer_lower_bounds'][2]}-{zoom_props['buffer_upper_bounds'][2]}]" ) print( f"{'Zoom Region Bounds':>28} = [{zoom_props['region_lower_bounds'][0]}-{zoom_props['region_upper_bounds'][0]}, {zoom_props['region_lower_bounds'][1]}-{zoom_props['region_upper_bounds'][1]}, {zoom_props['region_lower_bounds'][2]}-{zoom_props['region_upper_bounds'][2]}]" ) # Depths if zoom_props.get("with_buffer_cells"): print(f"{'Buffer Top Level Depth':>28} = {zoom_props['buffer_cell_depth']}") print(f"{'Zoom Top Level Depth':>28} = {zoom_props['zoom_cell_depth']}") # Assorted extra zoom properties print(f"{'Zoom Region Pad Factor':>28} = {zoom_props['region_pad_factor']}") print( f"{'Zoom Region Shift':>28} = [{zoom_props['zoom_shift'][0]}, {zoom_props['zoom_shift'][1]}, {zoom_props['zoom_shift'][2]}]" ) print( f"{'Zoom Region Center':>28} = [{zoom_props['region_lower_bounds'][0] + (zoom_props['dim'][0] / 2.0)}, {zoom_props['region_lower_bounds'][1] + (zoom_props['dim'][1] / 2.0)}, {zoom_props['region_lower_bounds'][2] + (zoom_props['dim'][2] / 2.0)}]" ) def zoom_report_cell_properties_text(space): """ Report Zoom Region Properties. Args: space (Space): The space object containing the properties. Returns: str: The formatted string containing the zoom region properties. """ zoom_props = space.zoom_props report_lines = [] # Cdims report_lines.append( f"{'Background cdim':>28} = [{space.bkg_cdim[0]}, {space.bkg_cdim[1]}, {space.bkg_cdim[2]}]" ) if zoom_props.get("with_buffer_cells"): report_lines.append( f"{'Buffer cdim':>28} = [{zoom_props['buffer_cdim'][0]}, {zoom_props['buffer_cdim'][1]}, {zoom_props['buffer_cdim'][2]}]" ) report_lines.append( f"{'Zoom cdim':>28} = [{zoom_props['cdim'][0]}, {zoom_props['cdim'][1]}, {zoom_props['cdim'][2]}]" ) # Dimensions report_lines.append( f"{'Background Dimensions':>28} = [{space.box_size}, {space.box_size}, {space.box_size}]" ) if zoom_props.get("with_buffer_cells"): report_lines.append( f"{'Buffer Region Dimensions':>28} = [{zoom_props['buffer_dim'][0]}, {zoom_props['buffer_dim'][1]}, {zoom_props['buffer_dim'][2]}]" ) report_lines.append( f"{'Zoom Region Dimensions':>28} = [{zoom_props['dim'][0]}, {zoom_props['dim'][1]}, {zoom_props['dim'][2]}]" ) # Cell Widths report_lines.append( f"{'Background Cell Width':>28} = [{space.box_size / space.bkg_top_level_cells}, {space.box_size / space.bkg_top_level_cells}, {space.box_size / space.bkg_top_level_cells}]" ) if zoom_props.get("with_buffer_cells"): report_lines.append( f"{'Buffer Cell Width':>28} = [{zoom_props['buffer_width'][0]}, {zoom_props['buffer_width'][1]}, {zoom_props['buffer_width'][2]}]" ) report_lines.append( f"{'Zoom Cell Width':>28} = [{zoom_props['width'][0]}, {zoom_props['width'][1]}, {zoom_props['width'][2]}]" ) # Number of Cells report_lines.append( f"{'Number of Background Cells':>28} = {zoom_props.get('nr_bkg_cells', 0)}" ) if zoom_props.get("with_buffer_cells"): report_lines.append( f"{'Number of Buffer Cells':>28} = {zoom_props.get('nr_buffer_cells', 0)}" ) report_lines.append( f"{'Number of Zoom Cells':>28} = {zoom_props.get('nr_zoom_cells', 0)}" ) # Bounds if zoom_props.get("with_buffer_cells"): report_lines.append( f"{'Buffer Lower Bounds':>28} = [{zoom_props['buffer_lower_bounds'][0]}, {zoom_props['buffer_lower_bounds'][1]}, {zoom_props['buffer_lower_bounds'][2]}]" ) report_lines.append( f"{'Buffer Upper Bounds':>28} = [{zoom_props['buffer_upper_bounds'][0]}, {zoom_props['buffer_upper_bounds'][1]}, {zoom_props['buffer_upper_bounds'][2]}]" ) report_lines.append( f"{'Zoom Region Lower Bounds':>28} = [{zoom_props['region_lower_bounds'][0]}, {zoom_props['region_lower_bounds'][1]}, {zoom_props['region_lower_bounds'][2]}]" ) report_lines.append( f"{'Zoom Region Upper Bounds':>28} = [{zoom_props['region_upper_bounds'][0]}, {zoom_props['region_upper_bounds'][1]}, {zoom_props['region_upper_bounds'][2]}]" ) # Depths if zoom_props.get("with_buffer_cells"): report_lines.append( f"{'Buffer Top Level Depth':>28} = {zoom_props['buffer_cell_depth']}" ) report_lines.append( f"{'Zoom Top Level Depth':>28} = {zoom_props['zoom_cell_depth']}" ) # Assorted extra zoom properties report_lines.append( f"{'Zoom Region Pad Factor':>28} = {zoom_props['region_pad_factor']}" ) report_lines.append( f"{'Zoom Region Shift':>28} = [{zoom_props['zoom_shift'][0]}, {zoom_props['zoom_shift'][1]}, {zoom_props['zoom_shift'][2]}]" ) report_lines.append( f"{'Zoom Region Center':>28} = [{zoom_props['region_lower_bounds'][0] + (zoom_props['dim'][0] / 2.0)}, {zoom_props['region_lower_bounds'][1] + (zoom_props['dim'][1] / 2.0)}, {zoom_props['region_lower_bounds'][2] + (zoom_props['dim'][2] / 2.0)}]" ) return "\n".join(report_lines) def zoom_region_init(space, params, verbose=False): """ Initialise the zoom region geometry. This will compute the cell grid properties ready for cell construction when zoom_construct_tl_cells. Args: space (Space): The space object to update. params (dict): The parameters to use for the zoom region. verbose (bool): Whether to print verbose output. """ # Compute the extent of the zoom region and calculate the shift to move # the zoom region to the centre of the box ini_dim = zoom_get_region_dim_and_shift(space, params) # Apply the shift to the particles for particle_group in ["parts", "gparts", "sparts", "bparts", "sinks"]: for particle in getattr(space, particle_group, []): particle["x"][0] += space.zoom_props["zoom_shift"][0] particle["x"][1] += space.zoom_props["zoom_shift"][1] particle["x"][2] += space.zoom_props["zoom_shift"][2] # Include the requested padding around the high-resolution particles max_dim = ini_dim * space.zoom_props["region_pad_factor"] # Define the background grid (we'll treat this as gospel) for i in range(3): space.cdim[i] = space.zoom_props["bkg_cdim"][i] space.width[i] = space.dim[i] / space.cdim[i] space.iwidth[i] = 1.0 / space.width[i] # Compute the void region bounds and number of zoom regions that tessellate it nr_zoom_regions = zoom_get_void_geometry(space, max_dim) # Check the user gave a sensible background cdim if nr_zoom_regions >= 64: zoom_report_cell_properties(space) raise ValueError( "Background cell size is too large relative to the zoom region! " "Increase ZoomRegion:bkg_top_level_cells (would have needed {} zoom " "cells in the void region).".format(nr_zoom_regions) ) # Warn the user if the number of zoom regions is large but not unreasonable if nr_zoom_regions >= 16: if verbose: print( "Warning: Background cell size is large relative to the zoom region! " "(we'll need at least {} buffer cells which may be slow).".format( nr_zoom_regions ) ) # Decide whether to use buffer cells based on the padding due to background cells if nr_zoom_regions <= 2: space.zoom_props["with_buffer_cells"] = 0 zoom_get_geometry_no_buffer_cells(space) else: space.zoom_props["with_buffer_cells"] = 1 zoom_get_geometry_with_buffer_cells(space) # Store what the true boost factor ended up being input_pad_factor = space.zoom_props["region_pad_factor"] space.zoom_props["region_pad_factor"] = space.zoom_props["dim"][0] / ini_dim # Ensure the zoom region is not smaller than the high-resolution particle distribution if space.zoom_props["dim"][0] < ini_dim: zoom_report_cell_properties(space) raise ValueError( "Found a zoom region smaller than the high-resolution particle " "distribution! Adjust the cell structure " "(ZoomRegion:bkg_top_level_cells, ZoomRegion:zoom_top_level_cells)" ) # Warn if the size of the requested padding region has drastically changed if (space.zoom_props["region_pad_factor"] / input_pad_factor) >= 2: if verbose: print( "Warning: The pad region has to be {} times larger than requested. " "Either increase ZoomRegion:region_pad_factor, increase the number of " "background cells, or increase the depths of the zoom cells (and buffer cells if using).".format( int(space.zoom_props["region_pad_factor"] / input_pad_factor) ) ) # Set the neighbour max tree depth if not explicitly provided if space.zoom_props["neighbour_max_tree_depth"] < 0: space.zoom_props["neighbour_max_tree_depth"] = space.zoom_props[ "zoom_cell_depth" ] # Set the minimum allowed zoom cell width zoom_dmax = max(space.zoom_props["dim"]) space.zoom_props["cell_min"] = 0.99 * zoom_dmax / space.zoom_props["cdim"][0] # Set the minimum background cell size dmax = max(space.dim) space.cell_min = 0.99 * dmax / space.cdim[0] # Store cell numbers and offsets space.zoom_props["bkg_cell_offset"] = ( space.zoom_props["cdim"][0] * space.zoom_props["cdim"][1] * space.zoom_props["cdim"][2] ) space.zoom_props["nr_zoom_cells"] = space.zoom_props["bkg_cell_offset"] space.zoom_props["nr_bkg_cells"] = space.cdim[0] * space.cdim[1] * space.cdim[2] space.zoom_props["buffer_cell_offset"] = ( space.zoom_props["bkg_cell_offset"] + space.zoom_props["nr_bkg_cells"] ) space.zoom_props["nr_buffer_cells"] = ( space.zoom_props["buffer_cdim"][0] * space.zoom_props["buffer_cdim"][1] * space.zoom_props["buffer_cdim"][2] ) def construct_tl_cells(space): """ Construct the top level cells. Args: space (Space): The space object containing the properties. """ # Create the background cells space.bkg_cells = [] for i in range(space.cdim[0]): for j in range(space.cdim[1]): cell = Cell( i * space.width[0], j * space.width[1], space.width, ctype="background" ) space.bkg_cells.append(cell) # Create the zoom cells space.zoom_cells = [] for i in range(space.zoom_props["cdim"][0]): for j in range(space.zoom_props["cdim"][1]): cell = Cell( i * space.zoom_props["width"][0] + space.zoom_props["region_lower_bounds"][0], j * space.zoom_props["width"][1] + space.zoom_props["region_lower_bounds"][1], space.zoom_props["width"], ctype="zoom", ) space.zoom_cells.append(cell) # Create the buffer cells if space.zoom_props.get("with_buffer_cells"): space.buffer_cells = [] for i in range(space.zoom_props["buffer_cdim"][0]): for j in range(space.zoom_props["buffer_cdim"][1]): cell = Cell( i * space.zoom_props["buffer_width"][0] + space.zoom_props["buffer_lower_bounds"][0], j * space.zoom_props["buffer_width"][1] + space.zoom_props["buffer_lower_bounds"][1], space.zoom_props["buffer_width"], ctype="buffer", ) space.buffer_cells.append(cell) def draw_grid(space, params): """ Draw the grid diagram of the cells. Args: space (Space): The space object containing the properties. """ # Setup the figure fig, ax = plt.subplots(figsize=(6, 6)) ax.set_aspect("equal") ax.set_xlim(0, space.box_size) ax.set_ylim(0, space.box_size) # Draw each cell for cell in space.bkg_cells: cell.draw_cell(ax) for cell in space.buffer_cells: cell.draw_cell(ax) for cell in space.zoom_cells: cell.draw_cell(ax) # Define handles for the legend as square markers handles = [ Line2D( [0], [0], marker="s", # Use a square marker color="w", # Set the color of the line to be invisible markerfacecolor="#bbe6e4", markeredgecolor="black", markersize=15, label="Background", alpha=0.9, ) ] if len(space.buffer_cells) > 0: handles.append( Line2D( [0], [0], marker="s", # Use a square marker color="w", # Set the color of the line to be invisible markerfacecolor="#42bfdd", markeredgecolor="black", markersize=15, label="Buffer", alpha=0.9, ) ) handles.append( Line2D( [0], [0], marker="s", # Use a square marker color="w", # Set the color of the line to be invisible markerfacecolor="#084b83", markeredgecolor="black", markersize=15, label="Zoom", alpha=0.9, ) ) # Place a legend below the x axis with the cell types ax.legend( handles=handles, loc="upper center", bbox_to_anchor=(0.5, -0.05), ncol=len(handles), frameon=False, ) # Save the figure including relevant information in the filename ic_name = params["InitialConditions"]["file_name"].split("/")[-1].split(".")[0] fig.savefig(f"zoom_geometry_{ic_name}.png", dpi=300) plt.show() plt.close(fig) def main(): """ Draw a 2D slice of the cell grid geometry. Example: python plot_zoom_geometry.py --params params.yaml """ parser = argparse.ArgumentParser( description="Draw a 2D grid diagram of cells with a nested region." ) parser.add_argument( dest="params", type=str, help="Path to the YAML file containing the parameters." ) # Add arbitrary arguments parser.add_argument( "--", dest="overrides", action=StoreDictKeyPair, nargs=1, metavar="VALUE", help="Override the parameter file with the given key-value pair.", ) # Unpack the arguments args, unknown_args = parser.parse_known_args() args.overrides = {} # Process unknown_args manually while unknown_args: if unknown_args[0].startswith("--"): key = unknown_args.pop(0).lstrip("--") if ( unknown_args and not unknown_args[0].startswith("-") and not unknown_args[0].startswith("--") ): value = unknown_args.pop(0) args.overrides[key] = value else: args.overrides[key] = None else: unknown_args.pop(0) # Read the parameters from the YAML file params = read_params(args.params) # Override the parameters with the provided key-value pairs for key, value in args.overrides.items(): if key not in params["ZoomRegion"]: raise ValueError(f"Invalid parameter: {key}") else: print(f"Overriding parameter: {key} = {value}") # Handle whether this is an integer or float if "." in value: params["ZoomRegion"][key] = float(value) else: params["ZoomRegion"][key] = int(value) # Parse the parameter file and populate the Space instance space = zoom_parse_params(params) # Initialize the zoom region geometry zoom_region_init(space, params, verbose=True) zoom_report_cell_properties(space) # Create the top level cells construct_tl_cells(space) # Draw the grid draw_grid(space, params) if __name__ == "__main__": main()