Grid

RoomsGrid

Bases: PyTreeNode

A grid of rooms. Each room is represented as a 2D grid of shape (room_height, room_width), with walls set to -1 and empty tiles set to 0. The grid of rooms is represented as a 2D grid of shape (rows * (room_height + 1), cols * (room_width + 1)), with walls set to -1 and empty tiles set to 0. The grid of rooms is represented as a 2D grid of shape (rows * (room_height + 1), cols * (room_width + 1)), with walls set to -1 and empty tiles set to 0.

create(num_rows, num_cols, room_size) classmethod

Creates a grid of rooms with the given number of rows and columns, and the given room size.

Parameters:

Name	Type	Description	Default
num_rows	int	The number of rows.	required
num_cols	int	The number of columns.	required
room_size	Tuple[int, int]	The size of each room (height, width).	required

Returns:

Name	Type	Description
RoomsGrid	RoomsGrid	A grid of rooms.

get_grid(occupied_positions=None)

Computes the array representation of the grid of rooms, with walls set to -1 and empty tiles set to 0.

Parameters:

Name	Type	Description	Default
occupied_positions	Array	A batch of extra occupied positions of shape (n, 2). Defaults to None.	None

Returns:

Name	Type	Description
Array	Array	A 2D grid of shape (rows * (room_height + 1), cols * (room_width + 1)).

position_in_room(row, col, *, key)

Generates a random position in a given room.

Parameters:

Name	Type	Description	Default
row	Array	The row index of the room.	required
col	Array	The column index of the room.	required
key	Array	A random key.	required

Returns:

Name	Type	Description
Array	Array	A random position in the given room.

position_on_border(row, col, side, *, key)

Generates a random position on the border of a given room. Side is 0: west, 1: east, 2: north, 3: south (like padding)

Parameters:

Name	Type	Description	Default
row	Array	The row index of the room.	required
col	Array	The column index of the room.	required
side	int	The side of the room.	required
key	Array	A random key.	required

Returns:

Name	Type	Description
Array	Array	A random position on the border of the given room.

align(patch, current_direction, desired_direction)

Aligns a patch of the grid from the current direction to the desired direction.

Parameters:

Name	Type	Description	Default
patch	Array	A patch of the grid.	required
current_direction	Array	The current direction in the range [0, 1, 2, 3] representing the cardinal directions [east, south, west, north].	required
desired_direction	Array	The desired direction in the range [0, 1, 2, 3] representing the cardinal directions [east, south, west, north].	required

Returns:

Name	Type	Description
Array	Array	A patch of the grid aligned to the desired direction.

coordinates(grid)

Returns a tuple of 2D coordinates [(col, row), ...] for each cell in the grid. A grid array of shape i32[height, width] will return a tuple of length (height * width), containing two arrays, each of shape i32[2].

Parameters:

Name	Type	Description	Default
grid	Array	A 2D grid of shape (height, width).	required

Returns:

Type	Description
Coordinates	Tuple[Array, Array]: A tuple of two arrays containing the 2D coordinates of each cell in the grid.

coordinates_from_idx(grid, idx)

Converts a flat index of shape i32[] into a 2D coordinate i32[2] containing (col, row) data. The index is calculated as idx = row * width + col.

Parameters:

Name	Type	Description	Default
grid	Array	A 2D grid of shape (height, width).	required
idx	Array	A flat index of shape i32[].	required

Returns:

Name	Type	Description
Array	Array	A 2D coordinate of shape i32[2] containing the (col, row) data.

crop(grid, origin, direction, radius, padding_value=0)

Crops a grid around a given origin, facing a given direction, with a given radius.

Parameters:

Name	Type	Description	Default
grid	Array	A 2D grid of shape (height, width).	required
origin	Array	The origin of the crop.	required
direction	Array	The direction the crop is facing.	required
radius	int	The radius of the crop.	required
padding_value	int	The padding value. Defaults to 0.	0

Returns:

Name	Type	Description
Array	Array	A cropped grid.

from_ascii_map(ascii_map, mapping={})

Converts an ASCII map into a 2D grid. The ASCII map is a string where each character represents a tile in the grid. The mapping dictionary can be used to map ASCII characters to integer values. By default, the mapping is as follows: - # is mapped to -1 - . is mapped to 0

Parameters:

Name	Type	Description	Default
ascii_map	str	The ASCII map.	required
mapping	Dict[str, int]	A dictionary mapping ASCII characters to integer values. Defaults to {}.	{}

Returns:

Name	Type	Description
Array	Array	A 2D grid representing the ASCII map.

horizontal_wall(grid, col_idx, opening_row_idx=None)

Creates a horizontal wall in the grid at the given column index, with an opening at the given row index.

Parameters:

Name	Type	Description	Default
grid	Array	A 2D grid of shape (height, width).	required
col_idx	int	The column index where the wall is placed.	required
opening_row_idx	Array	The row index where the opening is placed. Defaults to None.	None

Returns:

Name	Type	Description
Array	Array	A 2D grid of shape (height, width) with a horizontal wall.

idx_from_coordinates(grid, coordinates)

Converts a batch of 2D coordinates [(col, row), ...] into a flat index

Parameters:

Name	Type	Description	Default
grid	Array	A 2D grid of shape (height, width).	required
coordinates	Array	A batch of 2D coordinates of shape (batch_size, 2).	required

Returns:

Name	Type	Description
Array	Array	A flat index of shape i32[batch_size] for each coordinate in the batch.

mask_by_coordinates(grid, address, comparison_fn=jnp.greater_equal)

This is a workaround to compute dynamicly-sized masks in XLA, which would not be possible otherwise. Returns a mask of the same shape as grid where the value is 1 if the corresponding element in grid satisfies the comparison_fn with the corresponding element in address (col, row) and 0 otherwise.

Parameters:

Name	Type	Description	Default
grid	Array	A 2D grid of shape (height, width).	required
address	Coordinates	A tuple of 2D coordinates (col, row).	required
comparison_fn	Callable[[Array, Array], Array]	A comparison function. Defaults to jnp.greater_equal.	greater_equal

Returns:

Name	Type	Description
Array	Array	A boolean mask of the same shape as grid.

positions_equal(a, b)

Checks if two points are equal.

Parameters:

Name	Type	Description	Default
a	Array	A 2D coordinate of shape i32[2] containing the (col, row) data.	required
b	Array	A 2D coordinate of shape i32[2] containing the (col, row) data.	required

Returns:

random_colour(key, n=1)

Generates n random colours in the range [0, 1, 2, 3, 4, 5].

Parameters:

Name	Type	Description	Default
key	Array	A random key.	required
n	int	The number of random colours to generate. Defaults to 1.	1

Returns:

Name	Type	Description
Array	Array	A batch of random colours of shape u8[n].

random_directions(key, n=1)

Generates n random directions in the range [0, 1, 2, 3] representing the cardinal directions [east, south, west, north].

Parameters:

Name	Type	Description	Default
key	Array	A random key.	required
n	int	The number of random directions to generate. Defaults to 1.	1

Returns:

Name	Type	Description
Array	Array	A batch of random directions of shape i32[n].

random_positions(key, grid, n=1, exclude=jnp.asarray((-1, -1)))

Generates n random positions in the grid, excluding the exclude position.

Parameters:

Name	Type	Description	Default
key	Array	A random key.	required
grid	Array	A 2D grid of shape (height, width).	required
n	int	The number of random positions to generate. Defaults to 1.	1
exclude	Array	The position to exclude. Defaults to jnp.asarray((-1, -1)).	asarray((-1, -1))

Returns:

Name	Type	Description
Array	Array	A batch of random positions of shape i32[n, 2].

room(height, width)

Creates an array representing a room of size height x width, including a set of walls around the room. The room is represented as a 2D grid of shape (height, width), including walls, with walls set to -1 and empty tiles set to 0.

Parameters:

Name	Type	Description	Default
height	int	The height of the room.	required
width	int	The width of the room.	required

Returns:

Name	Type	Description
Array	Array	A 2D grid of shape (height, width) representing a room.

rotate(direction, spin)

Changes a direction vectory by a given number of spins.

Parameters:

Name	Type	Description	Default
direction	Array	A direction vector of shape i32[] in the range [0, 3] representing the cardinal directions [east, south, west, north].	required
spin	int	The number of spins to apply.	required

Returns:

Name	Type	Description
Array	Array	A direction vector of shape i32[] in the range [0, 3] representing the cardinal directions [east, south, west, north].

translate(position, direction, modulus=jnp.asarray(1))

Translates a point in a grid by a given direction and modulus.

Parameters:

Name	Type	Description	Default
position	Array	A 2D coordinate of shape i32[2] containing the (col, row) data.	required
direction	Array	A direction in the range [0, 1, 2, 3] representing the cardinal directions [east, south, west, north].	required
modulus	Array	The modulus of the translation. Defaults to jnp.asarray(1).	asarray(1)

Returns:

Name	Type	Description
Array	Array	A 2D coordinate of shape i32[2] containing the (col, row) data.

translate_forward(position, forward_direction, modulus)

Translates a point in a grid by a given forward direction and modulus.

Parameters:

Name	Type	Description	Default
position	Array	A 2D coordinate of shape i32[2] containing the (col, row) data.	required
forward_direction	Array	A direction in the range [0, 1, 2, 3] representing the cardinal directions [east, south, west, north].	required
modulus	Array	The modulus of the translation.	required

Returns:

Name	Type	Description
Array		A 2D coordinate of shape i32[2] containing the (col, row) data.

translate_left(position, forward_direction, modulus)

Translates a point in a grid by a given left direction and modulus.

Parameters:

Name	Type	Description	Default
position	Array	A 2D coordinate of shape i32[2] containing the (col, row) data.	required
forward_direction	Array	A direction in the range [0, 1, 2, 3] representing the cardinal directions [east, south, west, north].	required
modulus	Array	The modulus of the translation.	required

Returns:

Name	Type	Description
Array		A 2D coordinate of shape i32[2] containing the (col, row) data.

translate_right(position, forward_direction, modulus)

Translates a point in a grid by a given right direction and modulus.

Parameters:

Name	Type	Description	Default
position	Array	A 2D coordinate of shape i32[2] containing the (col, row) data.	required
forward_direction	Array	A direction in the range [0, 1, 2, 3] representing the cardinal directions [east, south, west, north].	required
modulus	Array	The modulus of the translation.	required

Returns:

Name	Type	Description
Array		A 2D coordinate of shape i32[2] containing the (col, row) data.

two_rooms(height, width, key)

Creates a 2D grid representing two rooms of size height x width, separated by a wall. The rooms are represented as a 2D grid of shape (height, width), including walls, with walls set to -1 and empty tiles set to 0.

Parameters:

Name	Type	Description	Default
height	int	The height of the rooms.	required
width	int	The width of the rooms.	required
key	Array	A random key, determining the position of the wall separating the rooms.	required

Returns:

Type	Description
Tuple[Array, Array]	Tuple[Array, Array]: A tuple containing the 2D grid representing the rooms and the column index of the wall separating the rooms.

vertical_wall(grid, row_idx, opening_col_idx=None)

Creates a vertical wall in the grid at the given row index, with an opening at the given column index.

Parameters:

Name	Type	Description	Default
grid	Array	A 2D grid of shape (height, width).	required
row_idx	int	The row index where the wall is placed.	required
opening_col_idx	Array	The column index where the opening is placed. Defaults to None.	None

Returns:

Name	Type	Description
Array	Array	A 2D grid of shape (height, width) with a vertical wall.

view_cone(transparency_map, origin, radius)

Computes the view cone of a given origin in a grid with a given radius. The view cone is a boolean map of transparent (1) and opaque (0) tiles, indicating whether a tile is visible from the origin or not.

Parameters:

Name	Type	Description	Default
transparency_map	Array	A boolean map of transparent (1) and opaque (0) tiles.	required
origin	Array	The origin of the view cone.	required
radius	int	The radius of the view cone.	required

Returns:

Name	Type	Description
Array	Array	The view cone of the given origin in the grid with the given radius.