Encoder¶

AlhenaEncoder combines the UNI2-h image encoder with the ALHENA low-rank adaptation and projection heads.

Three methods cover the main encoding levels:

Method	Level	Returns
`encode_tiles()`	tile	`TileEmbeddings` with embeddings, identifiers, spatial bounds when available, and provenance
`encode_cells()`	cell	`CellEmbeddings` with pre-projection and projected cell representations, positions when available, and provenance
`encode()`	low-level tile operation	embedding matrix only

encode_tiles() and encode_cells() are the high-level analysis workflows. encode() is the lower-level operation for in-memory arrays when only the embedding matrix is required.

Two sets of tiles, told apart

The tiles encode_tiles reads are the tile-level result: one embedding per tile of a sampling of the image. The tiles encode_cells reads provide context for the cells, so by default they cover only the tiles those cells require rather than forming a sampling of the image. Both come back as TileEmbeddings, and purpose says which is which, in memory and in the table write_h5ad writes. See Dedicated tiles and cell-context tiles.

Image geometry is recorded with the result

Tile padding and cropping can affect the representations produced by the encoder. High-level result objects retain pad_fraction and crop_fraction, and the relevant encoding methods emit geometry warnings when appropriate.

See Geometry for interpretation.

alhena.encoder.AlhenaEncoder ¶

AlhenaEncoder(backbone, heads, info, device='cpu')

The adapted backbone together with the projection heads that read out of it.

load_encoder and load_bundle build it. Every method runs in float32 under inference mode, with autocast disabled on CPU and CUDA, and no output is normalized.

Parameters:

Name	Type	Description	Default
`backbone`	`Module`	UNI2-h from `build_backbone`, with adapters inserted by `insert_adapters`.	required
`heads`	`Module`	The projection heads, as `build_heads` makes them.	required
`info`	`BundleInfo`	The identity and geometry of the bundle, a `BundleInfo`.	required
`device`	`device \| str`	Device to run on, such as `"cpu"`, `"cuda"` or `"mps"`.	`'cpu'`

Attributes:

Name	Type	Description
`info`		Identity and geometry of the bundle.
`device`		Device the backbone and the heads sit on.

Raises:

Type	Description
`AlhenaValueError`	If `device` names no device this machine has.
`BundleError`	If the backbone is not as wide as the bundle declares.

encode ¶

encode(tiles, mpp=None, batch_size=DEFAULT_BATCH_SIZE, progress=None)

Embed tiles with the backbone.

Any white pad on the canvas raises a GeometryWarning, and so does a center crop discarding more than CROP_WARNING_FRACTION of a field; prepared tiles raise neither.

Parameters:

Name	Type	Description	Default
`tiles`	`ndarray \| Sequence[ndarray] \| Tensor`	Raw uint8 RGB tiles, or a float tensor already prepared for the model.	required
`mpp`	`float \| None`	Microns per pixel of the tiles. Refused for prepared tiles.	`None`
`batch_size`	`int`	Tiles per forward pass.	`DEFAULT_BATCH_SIZE`
`progress`	`bool \| None`	Whether to report progress on standard error. None follows `ALHENA_DISABLE_PROGRESS_BARS`.	`None`

Returns:

Type	Description
`ndarray`	An `(n_tiles, embed_dim)` float32 array, not normalized.

Raises:

Type	Description
`AlhenaValueError`	If `mpp` or `batch_size` is out of range, or no tile is given.
`GeometryError`	If the tiles cannot reach the geometry the model reads.

encode_tiles ¶

encode_tiles(tiles, mpp=None, batch_size=DEFAULT_BATCH_SIZE, progress=None)

Embed the tiles of a source, with where each one was read from.

The entry point for tile-level analysis: the names, bounds and frame of the source travel with the result, and arrays come back named by position and without bounds. Tiles are read one batch at a time. The warnings on white pad and the center crop come once: with the first batch where every tile has one size, after the pass where sizes differ.

Parameters:

Name	Type	Description	Default
`tiles`	`TileSource \| TileInput`	A tile source, raw uint8 RGB tiles as one array per tile or one stacked array, or a float tensor already prepared for the model.	required
`mpp`	`float \| None`	Microns per pixel of the tiles. Defaults to the resolution the source carries. Refused for prepared tiles.	`None`
`batch_size`	`int`	Tiles per forward pass, which bounds the memory the pass holds.	`DEFAULT_BATCH_SIZE`
`progress`	`bool \| None`	Whether to report progress on standard error. None follows `ALHENA_DISABLE_PROGRESS_BARS`.	`None`

Returns:

Type	Description
`TileEmbeddings`	One embedding per tile of the source, in its order, with the tile bounds and the
`TileEmbeddings`	provenance of the pass; `purpose` reads `tiles`.

Raises:

Type	Description
`AlhenaValueError`	If `mpp` or `batch_size` is out of range, or the tiles are unreadable.
`GeometryError`	If the tiles cannot reach the geometry the model reads.

encode_with_tokens ¶

encode_with_tokens(tiles, mpp=None, batch_size=DEFAULT_BATCH_SIZE, progress=None)

Embed tiles and return their patch tokens as well.

The tokens follow the final layer normalization, without the class and register tokens, in row-major order of the grid; pooling tokens read before that normalization gives another cell embedding. No warning is raised here about white pad or the center crop: the cell path keeps the pad out of pooling and warns about the crop itself.

Parameters:

Name	Type	Description	Default
`tiles`	`ndarray \| Sequence[ndarray] \| Tensor`	Raw uint8 RGB tiles, or a float tensor already prepared for the model.	required
`mpp`	`float \| None`	Microns per pixel of the tiles. Refused for prepared tiles.	`None`
`batch_size`	`int`	Tiles per forward pass.	`DEFAULT_BATCH_SIZE`
`progress`	`bool \| None`	Whether to report progress on standard error. None follows `ALHENA_DISABLE_PROGRESS_BARS`.	`None`

Returns:

Type	Description
`ndarray`	An `(n_tiles, embed_dim)` array of tile embeddings and an
`ndarray`	`(n_tiles, n_tokens, embed_dim)` array of patch tokens, neither normalized.

Raises:

Type	Description
`AlhenaValueError`	If `mpp` or `batch_size` is out of range, or no tile is given.
`BundleError`	If the backbone yields another token grid than the bundle reads.
`GeometryError`	If the tiles cannot reach the geometry the model reads.

encode_cells ¶

encode_cells(tiles, nuclei, mpp=None, *, cell_id=None, all_tiles=False, batch_size=DEFAULT_BATCH_SIZE, progress=None)

Embed the cells of a set of tiles in one pass, each tile embedded once.

The tiles and the nuclei are checked, and every nucleus resolved to the tile it is read from, before the tiles are embedded one batch at a time. By default a tile no cell is read from is skipped, and the embeddings of the tiles read come back on CellEmbeddings.tiles. The cells equal those encode_with_tokens and pool_cells give, to float32 rounding.

Parameters:

Name	Type	Description	Default
`tiles`	`TileInput`	Raw uint8 RGB tiles, one array per tile or one stacked array, a tile source, or a float tensor already prepared for the model.	required
`nuclei`	`NucleiInput`	Label maps, the mapping an `alhena.io` reader of nuclei returns, or a membership table.	required
`mpp`	`float \| None`	Microns per pixel of the tiles, declared once for the images and the nuclei. Refused for prepared tiles.	`None`
`cell_id`	`Sequence[Any] \| ndarray \| None`	Unique identifier of each nucleus, for the builders that take positions.	`None`
`all_tiles`	`bool`	Read every tile, including those no cell is read from.	`False`
`batch_size`	`int`	Tiles per forward pass, which bounds the memory the pass holds.	`DEFAULT_BATCH_SIZE`
`progress`	`bool \| None`	Whether to report progress on standard error. None follows `ALHENA_DISABLE_PROGRESS_BARS`.	`None`

Returns:

Type	Description
`CellEmbeddings`	One pre-projection and one cell-space embedding per cell that reached a token, ordered
`CellEmbeddings`	by identifier, with their provenance and the embeddings of the tiles read.

Raises:

Type	Description
`AlhenaValueError`	If `mpp` is out of range, or the tiles and the nuclei do not match.
`BundleError`	If the bundle carries no cell head.
`GeometryError`	If the tiles cannot reach the canvas, or the nuclei lie on another grid.
`MembershipError`	If the nuclei and the tiles cannot be reconciled.

project_tile ¶

project_tile(embeddings)

Project tile embeddings with the tile head.

Parameters:

Name	Type	Description	Default
`embeddings`	`ndarray \| Tensor`	A (n_tiles, embed_dim) array of pre-projection tile embeddings.	required

Returns:

Type	Description
`ndarray`	A (n_tiles, cell_dim) float32 array, not normalized.

Raises:

Type	Description
`AlhenaValueError`	If the array is not two-dimensional.
`BundleError`	If the bundle carries no tile head, or the array is the wrong width.

project_cell ¶

project_cell(pooled_tokens)

Project pooled patch tokens with the cell head, the image side of the retrieval space.

Parameters:

Name	Type	Description	Default
`pooled_tokens`	`ndarray \| Tensor`	An `(n_cells, embed_dim)` array of nucleus-pooled patch tokens.	required

Returns:

Type	Description
`ndarray`	An `(n_cells, cell_dim)` float32 array, not normalized.

Raises:

Type	Description
`AlhenaValueError`	If the array is not two-dimensional.
`BundleError`	If the bundle carries no cell head, or if the array is the wrong width.

project_reference ¶

project_reference(scconcept_embeddings)

Project single-cell embeddings with the reference head, the reference side of retrieval.

Parameters:

Name	Type	Description	Default
`scconcept_embeddings`	`ndarray \| Tensor`	An `(n_cells, scconcept_dim)` array of single-cell embeddings.	required

Returns:

Type	Description
`ndarray`	An `(n_cells, cell_dim)` float32 array, not normalized.

Raises:

Type	Description
`AlhenaValueError`	If the array is not two-dimensional.
`BundleError`	If the bundle carries no reference head, or the array is the wrong width.

as_module ¶

as_module()

Return a module mapping tiles on the unit scale to pre-projection tile embeddings.

The module takes a (B, 3, 224, 224) float tensor in [0, 1], normalizes it, and is in evaluation mode, as the tile pipelines that call it expect. It shares the backbone with the encoder.

alhena.encoder.load_bundle ¶

load_bundle(bundle, backbone, *, device=None, merge=False, progress=None)

Build the encoder from a bundle and the user's UNI2-h weights.

Parameters:

Name	Type	Description	Default
`bundle`	`str \| Path`	Directory holding the bundle files, or the repository publishing them.	required
`backbone`	`str \| Path`	Path of the UNI2-h weight file.	required
`device`	`device \| str \| None`	Device to run on, such as `"cpu"`, `"cuda"`, or `"mps"` for the GPU of an Apple silicon Mac. A CUDA device is chosen when one is available, otherwise the processor.	`None`
`merge`	`bool`	Fold the adapters into the fused weights instead of inserting them, which agrees with the published forward pass to float32 rounding.	`False`
`progress`	`bool \| None`	Whether to announce on standard error the reading of the UNI2-h weights. None follows `ALHENA_DISABLE_PROGRESS_BARS`.	`None`

Returns:

Type	Description
`AlhenaEncoder`	The encoder, on the resolved device.

Raises:

Type	Description
`AlhenaValueError`	If `device` names no device this machine has.
`BackboneNotFoundError`	If the UNI2-h weight file does not exist.
`BundleError`	If a file is missing or does not match what the bundle declares.

alhena.encoder.LoRAQKV ¶

LoRAQKV(linear_layer, rank=LORA_RANK, alpha=LORA_ALPHA)

Bases: Module

Low-rank adaptation of a fused query-key-value projection.

The adapters add a low-rank update to the query and value thirds of the output, and none to the key third. Every parameter starts at zero, so the module reproduces the wrapped layer until weights are loaded.

Parameters:

Name	Type	Description	Default
`linear_layer`	`Linear`	The fused `qkv` layer of a transformer block, whose output is three times as wide as its input.	required
`rank`	`int`	Rank of the update.	`LORA_RANK`
`alpha`	`float`	Scaling of the update, which is multiplied by `alpha / rank`.	`LORA_ALPHA`

Raises:

Type	Description
`BundleError`	If the layer is not a fused `qkv` linear layer, or `rank` is below one.

delta ¶

delta(x)

Return the update the adapters add to the wrapped layer's output.

Parameters:

Name	Type	Description	Default
`x`	`Tensor`	Input of the wrapped layer, `(..., in)`.	required

Returns:

Type	Description
`Tensor`	The update, `(..., 3 * in)`, zero over the key third.

delta_weight ¶

delta_weight()

Return the update the adapters add to the fused weight, shape (3 * in, in).

forward ¶

forward(x)

Apply the wrapped layer and add the adapters' update.

Parameters:

Name	Type	Description	Default
`x`	`Tensor`	Input of the wrapped layer, `(..., in)`.	required

Returns:

Type	Description
`Tensor`	The fused query, key and value, `(..., 3 * in)`.

alhena.encoder.insert_adapters ¶

insert_adapters(vit, rank=LORA_RANK, alpha=LORA_ALPHA)

Wrap the fused qkv projection of every block in a low-rank adapter.

Parameters:

Name	Type	Description	Default
`vit`	`Module`	A vision transformer exposing `blocks`.	required
`rank`	`int`	Rank of the adapters.	`LORA_RANK`
`alpha`	`float`	Scaling numerator; the update is scaled by `alpha / rank`.	`LORA_ALPHA`

Returns:

Type	Description
`int`	The number of blocks adapted.

Raises:

Type	Description
`BundleError`	If the model exposes no blocks, or a block is already adapted.

alhena.encoder.load_adapter_weights ¶

load_adapter_weights(vit, tensors)

Copy adapter weights into the adapters inserted in a transformer.

Keys are matched from blocks. onward, so the prefixes of a training checkpoint are accepted. Every adapted block must receive its four tensors and nothing else.

Parameters:

Name	Type	Description	Default
`vit`	`Module`	A transformer whose blocks already carry adapters.	required
`tensors`	`Mapping[str, Tensor]`	Adapter weights, keyed as in the bundle.	required

Raises:

Type	Description
`BundleError`	If the tensors do not match the adapted blocks, in names or in shapes.

alhena.encoder.merge_adapters ¶

merge_adapters(vit)

Fold the inserted adapters into the fused weights and restore plain linear layers.

The merged and inserted forms agree to float32 rounding: on the released bundle, at most 2.3e-05 in absolute terms and 2.9e-06 in relative terms.

Parameters:

Name	Type	Description	Default
`vit`	`Module`	A transformer whose blocks carry adapters holding their final weights.	required

Raises:

Type	Description
`BundleError`	If a block carries no adapter.

alhena.encoder.build_backbone ¶

build_backbone(backbone, architecture=BACKBONE_ARCHITECTURE, timm_kwargs=None)

Instantiate the vision transformer and load the UNI2-h weights into it.

The architecture is built locally and the weights read from the file named, with no network.

Parameters:

Name	Type	Description	Default
`backbone`	`str \| Path`	Path of the UNI2-h weight file, `pytorch_model.bin` or its safetensors form.	required
`architecture`	`str`	timm architecture name.	`BACKBONE_ARCHITECTURE`
`timm_kwargs`	`Mapping[str, Any] \| None`	Constructor arguments; those of the released model by default.	`None`

Returns:

Type	Description
`Module`	The transformer in float32, in evaluation mode, with gradients disabled.

Raises:

Type	Description
`BackboneNotFoundError`	If the weight file does not exist.
`BundleError`	If the weights do not match the architecture.

alhena.encoder.build_heads ¶

build_heads(tensors, embed_dim=BACKBONE_DIM, cell_dim=CELL_SPACE_DIM, scconcept_dim=SCCONCEPT_DIM)

Build the projection heads of a bundle from their weights.

The heads are bias-free linear maps: image_proj and img_tok_proj read backbone features but were trained against different teachers, so they are not interchangeable, and st_node_proj reads single-cell embeddings. A head whose weight is absent is not built.

Parameters:

Name	Type	Description	Default
`tensors`	`Mapping[str, Tensor]`	Head weights, keyed as in the bundle.	required
`embed_dim`	`int`	Width the two image-side heads read.	`BACKBONE_DIM`
`cell_dim`	`int`	Width all three heads write.	`CELL_SPACE_DIM`
`scconcept_dim`	`int`	Width the reference head reads.	`SCCONCEPT_DIM`

Returns:

Type	Description
`Module`	A module holding the heads, in evaluation mode, with gradients disabled.

Raises:

Type	Description
`BundleError`	If a tensor is not one of the three head weights, or has the wrong shape.

alhena.encoder.bundle_identity ¶

bundle_identity(adapter_bytes, head_tensors, scconcept_model=SCCONCEPT_MODEL, scconcept_version=SCCONCEPT_VERSION)

Return the encoder identity and the cell space identity of a bundle.

Each is the first sixteen hexadecimal digits of a SHA-256 digest: of the adapter file, and of alhena-cell-space-v1, the adapter digest, the digests of the cell and reference head weights, the single-cell model name and its version, joined by |.

Parameters:

Name	Type	Description	Default
`adapter_bytes`	`bytes`	Contents of the adapter file.	required
`head_tensors`	`Mapping[str, Tensor]`	Head weights, keyed as in the bundle.	required
`scconcept_model`	`str`	Single-cell model the reference embeddings come from.	`SCCONCEPT_MODEL`
`scconcept_version`	`str`	Version of that model.	`SCCONCEPT_VERSION`

Returns:

Type	Description
`tuple[str, str]`	The encoder identity and the cell space identity.

alhena.encoder.adapter_key ¶

adapter_key(block_index, name)

Return the key one adapter tensor is stored under, relative to the transformer.

Parameters:

Name	Type	Description	Default
`block_index`	`int`	Index of the transformer block.	required
`name`	`str`	Name of the tensor, one of `ADAPTER_NAMES`.	required

Returns:

Type	Description
`str`	The key, such as `blocks.0.attn.qkv.adapter_Q_downsample`.

alhena.encoder.ADAPTER_NAMES `module-attribute` ¶

ADAPTER_NAMES = ('adapter_Q_downsample', 'adapter_Q_upsample', 'adapter_V_downsample', 'adapter_V_upsample')

Names of the four tensors one adapter carries, as the bundle stores them.

alhena.encoder.DEFAULT_BATCH_SIZE `module-attribute` ¶

DEFAULT_BATCH_SIZE = 32

Tiles per forward pass where a call does not say.

alhena.encoder.TILE_EMBEDDING_CONTRACT `module-attribute` ¶

TILE_EMBEDDING_CONTRACT = 'A tile embedding is the class token over every token of the canvas, taken unmasked and unweighted, so it moves with the geometry the tiles reach. Tile embeddings computed at different geometries are not comparable and must not be pooled.'

What a tile embedding is, as the warnings on white pad and on patch resolution state it.

Next: Cell typer

Encoder¶

alhena.encoder.AlhenaEncoder ¶

encode ¶

encode_tiles ¶

encode_with_tokens ¶

encode_cells ¶

project_tile ¶

project_cell ¶

project_reference ¶

as_module ¶

alhena.encoder.load_bundle ¶

alhena.encoder.LoRAQKV ¶

delta ¶

delta_weight ¶

forward ¶

alhena.encoder.insert_adapters ¶

alhena.encoder.load_adapter_weights ¶

alhena.encoder.merge_adapters ¶

alhena.encoder.build_backbone ¶

alhena.encoder.build_heads ¶

alhena.encoder.bundle_identity ¶

alhena.encoder.adapter_key ¶

alhena.encoder.ADAPTER_NAMES module-attribute ¶

alhena.encoder.DEFAULT_BATCH_SIZE module-attribute ¶

alhena.encoder.TILE_EMBEDDING_CONTRACT module-attribute ¶

alhena.encoder.ADAPTER_NAMES `module-attribute` ¶

alhena.encoder.DEFAULT_BATCH_SIZE `module-attribute` ¶

alhena.encoder.TILE_EMBEDDING_CONTRACT `module-attribute` ¶