factoriax.engine.renderer

factoriax.engine.renderer#

Draw an EnvState as a pixel image.

JAX can trace render_map(), so it runs under jax.jit and vmaps over a batch of states. It draws three layers, from back to front:

The terrain, gathered from a block atlas by state.map.
The machines, gathered by (direction, machine_type) and blended onto the terrain through the alpha channel of the atlas.
The players, one cell for each (player_idx, direction), blended onto the tile of each player.

The sprites come from factoriax/assets/atlas.png, a grid of 32x32 RGBA cells. factoriax/assets/atlas.layout.md fixes the layout, and tests/test_atlas_fresh.py tests it. A block cell is opaque, so the terrain always paints. A machine cell and a player cell hold transparent regions, so a player on a belt does not hide the belt.

CAUTION: The row and column constants below are written into this file. They do not come from the atlas.json file next to the PNG. Copy every layout change into this file by hand.

Each atlas is built for one tile_px. The builder slices a row and resamples every cell from 32x32 with nearest-neighbour. A conveyor belt, a miner, an arm, a splitter, and a crossing each carry one cell for each Direction. Every other machine repeats its one sprite across the four direction rows, so the gather needs no branch.

The module serves two readers. JaxRenderer holds the atlases on the device for the JAX path. The *_rgba helpers return numpy arrays on the host for the pygame surfaces in the editor and the play UI. Both come from the same PNG, so one change to the art changes both.

The editor and the play UI draw the HUD. This module does not.

Usage:

renderer = JaxRenderer(tile_px=8)
img = renderer.jit_render_map(state)              # single state
imgs = renderer.vmap_render_map(batched_states)   # batched

Module Attributes

`DEFAULT_TILE_PX`	Tile side length in pixels when a caller gives no other size.
`DIGIT_H`	Rows in a digit glyph from `build_digit_atlas()`.
`DIGIT_W`	Columns in a digit glyph from `build_digit_atlas()`.

factoriax.engine.renderer.DEFAULT_TILE_PX: int = 8#: Tile side length in pixels when a caller gives no other size. A map of H by W tiles renders to (H * 8, W * 8, 3).

factoriax.engine.renderer.DIGIT_H: int = 5#: Rows in a digit glyph from build_digit_atlas().

factoriax.engine.renderer.DIGIT_W: int = 3#: Columns in a digit glyph from build_digit_atlas().

class factoriax.engine.renderer.JaxRenderer(tile_px: int = 8)#

Hold the sprite atlases for one tile size and render states with them.

An atlas build reads and resamples the PNG. A caller that renders every step must therefore keep one renderer alive, and must not build one for each frame. The atlases move to the device at construction and stay there. The compile happens later. The jit wrapper and the vmap wrapper build on first use, and each one traces again when the shape of a state changes.

One instance serves one tile size. A second size needs a second renderer.

Example:

renderer = JaxRenderer(tile_px=8)
img = renderer.jit_render_map(state)              # single state
imgs = renderer.vmap_render_map(batched_states)   # batched

Parameters:: tile_px – Tile side length in pixels. Any size except 32 resamples the atlas by nearest neighbour. A size that does not divide 32 exactly therefore drops or repeats rows of pixels.

jit_render_map(state: EnvState) → Array#

Render one state through the compiled renderer.

The first call for a map size and player count pays for the compile, and a later call with that shape reuses it. A play loop and a video recorder must use this method.

Parameters:: state – One unbatched state.
Returns:: The image that render_map() describes. It lives on the device, so a pygame or imageio reader must call np.asarray first. That call is where the render blocks.
Return type:: jnp.ndarray

render_map_single(state: EnvState) → Array#

Render one state with no compile first.

Each primitive dispatches on its own. This is slower for each call than jit_render_map(), but it pays for no compile. Use it for a single frame, or to step through the layers in a debugger.

Parameters:: state – One unbatched state.
Returns:: Image as described in render_map().
Return type:: jnp.ndarray

vmap_render_map(batched_state: EnvState) → Array#

Render a batch of states in one call.

Parameters:: batched_state – A state in which every field carries a leading batch axis, as a vmapped environment produces. An unbatched state raises during the trace and renders no image.
Returns:: Shape (batch, H * tile_px, W * tile_px, 3), uint8. The device holds every image at the same time. A large batch at a large tile size can therefore fill the device memory long before the states themselves do.
Return type:: jnp.ndarray

factoriax.engine.renderer.block_textures_rgba(size: int) → ndarray#

Return every terrain texture as one stack that the editor can index.

Parameters:: size – Tile side length in pixels.
Returns:: Shape (num_block_types, size, size, 4), uint8. A BlockType value indexes it. The one stack lets the editor canvas blit a whole viewport with one advanced index, and not with a loop over the tiles. The alpha channel is 255 everywhere, because the terrain is the layer that every other layer paints over. A cache holds this array and every caller shares it, so do not write to it.
Return type:: np.ndarray

factoriax.engine.renderer.build_block_atlas(tile_px: int) → Array#

Build the terrain layer atlas that render_map() gathers from.

Parameters:: tile_px – Tile side length in pixels.
Returns:: Shape (num_block_types, tile_px, tile_px, 4), uint8. A BlockType value indexes it. Every block cell is opaque, so the alpha channel is 255 everywhere and render_map() drops it. The row length is the highest enum value plus one. A gap in the numbering therefore leaves a magenta cell, which the atlas builder wrote as a marker for a missing sprite.
Return type:: jnp.ndarray

factoriax.engine.renderer.build_digit_atlas() → Array#

Build the bitmap font for the digits 0 to 9.

Nothing in this module draws digits. factoriax/assets/build_atlas.py is the one caller. It writes these glyphs into the digits row of the atlas at build time.

Returns:: Shape (10, DIGIT_H, DIGIT_W), bool. The digit itself indexes it. True marks an inked pixel and carries no colour, so the caller selects one.
Return type:: jnp.ndarray

factoriax.engine.renderer.build_machine_atlas(tile_px: int) → Array#

Build the machine layer atlas, with one sprite set for each facing.

Parameters:: tile_px – Tile side length in pixels.
Returns:: Shape (4, num_machine_types, tile_px, tile_px, 4), uint8. The pair (direction - 1, machine_type) indexes it, and the direction axis is in value order LEFT, RIGHT, UP, DOWN. A machine that looks the same in every direction repeats its one sprite across the four rows, so the gather at render time needs no test for a directional kind. The alpha channel carries meaning here. Machine.NONE is fully transparent, and that lets render_map() blend an empty tile away instead of a mask.
Return type:: jnp.ndarray

factoriax.engine.renderer.build_player_sprite(tile_px: int) → Array#

Build the player layer atlas, with one sprite for each slot and facing.

Parameters:: tile_px – Tile side length in pixels.
Returns:: Shape (8, 4, tile_px, tile_px, 4), uint8. The pair (player_slot, direction - 1) indexes it. There are eight slots in total. A scenario with more players than eight reuses the earlier colours, because the atlas holds no ninth palette. Each cell carries a transparent margin, so the tile under a player stays visible.
Return type:: jnp.ndarray

factoriax.engine.renderer.machine_icon_rgba(machine_type: int, size: int, direction: int) → ndarray#

Return one machine sprite for a pygame surface to blit.

The editor canvas calls this function instead of factoriax.playground.ui.icons.render_item_icon(). The map that the editor draws and the map that render_map() draws therefore come from the same art.

Parameters:

machine_type – Machine value. A value past the end of the row returns an empty array and raises nothing.
size – Side length of the returned sprite, in pixels.
direction – Direction value. A value outside 1 to 4 clamps to the LEFT sprite. This includes the 0 of an entity that is not on the map, and of an entity with no direction.

Returns:

Shape (size, size, 4), uint8. Machine.NONE gives a fully transparent cell. A cache holds this array and every caller shares it, so do not write to it.

Return type:

np.ndarray

factoriax.engine.renderer.player_icon_rgba(player_idx: int, size: int, direction: int) → ndarray#

Return one player sprite for a pygame surface to blit.

Parameters:

player_idx – Player index. The atlas holds eight sprites, so an index past the eighth returns to the first colour. render_map() uses the same count, so a sprite here matches the sprite on the map.
size – Side length of the returned sprite, in pixels.
direction – Direction value. A value outside 1 to 4 clamps to the LEFT sprite.

Returns:

Shape (size, size, 4), uint8, with a transparent margin. A cache holds this array and every caller shares it, so do not write to it.

Return type:

np.ndarray

factoriax.engine.renderer.render_map(state: EnvState, block_atlas: Array, machine_atlas: Array, player_sprite: Array) → Array#

Draw terrain, machines, and players into one image.

The function reads only the state fields that it draws. It allocates a new image and changes nothing, so it traces under jax.jit and vmaps over a batch of states. The tile size comes from block_atlas, not from an argument.

The function clamps input outside the valid range and does not refuse it, because a traced gather cannot raise. A block id or a machine id past the end of its atlas draws the last cell. A facing outside 1 to 4 draws the LEFT cell. A player outside the map appears at the nearest edge.

CAUTION: The function reports none of these cases. A scenario with a bad id therefore renders the wrong sprite and gives no error.

Parameters:

state – One unbatched state. The machines come from machine_types and tile_entity, which both index by [row, column]. player_positions holds (x, y), in the other order.
block_atlas – Terrain sprites from build_block_atlas().
machine_atlas – Machine sprites from build_machine_atlas().
player_sprite – Player sprites from build_player_sprite(). The length of its first axis limits the number of different player colours.

Returns:

Shape (H * tile_px, W * tile_px, 3), uint8, with H and W from state.map. Row 0 is the top of the map. The result is opaque. The blend consumes the alpha channel and does not return it, so this is a finished picture and not a layer for a further blend.

Return type:

jnp.ndarray

factoriax.engine.renderer

Contents

factoriax.engine.renderer#