Combinators

Combinators for building complex manifolds from simpler ones.

Provides product manifolds (Pair, Triple), homogeneous products (Replicated), and the zero-dimensional Null. Each combinator stores coordinates as a flat concatenation and provides split_coords / join_coords for component access.

Class Hierarchy

Null Manifold

class Null

Bases: Manifold

A zero-dimensional manifold with no coordinates.

property dim: int

The dimension of the manifold.

Product Manifolds

class Tuple

Bases: Manifold, ABC

Abstract Cartesian product of manifolds, with coordinates stored as a flat concatenation.

Mathematically, the Cartesian product \(\mathcal M_1 \times \cdots \times \mathcal M_k\) has \(\dim = \sum_i \dim(\mathcal M_i)\). Subclasses (Pair, Triple) fix the arity and provide typed split_coords / join_coords.

abstractmethod split_coords(coords: Array) → tuple[Array, ...]

Split coordinates into tuple components.

abstractmethod join_coords(*components: Array) → Array

Join tuple components into a single array.

Note: Subclasses implement this with specific numbers of coordinates (e.g., Pair takes exactly 2, Triple takes exactly 3).

class Pair

Bases: Tuple, ABC, Generic

Binary Cartesian product, with coordinates stored as [fst | snd].

abstract property fst_man: First

First component manifold.

abstract property snd_man: Second

Second component manifold.

property dim: int

The dimension of the manifold.

split_coords(coords: Array) → tuple[Array, Array]

Split into (fst, snd) components.

join_coords(fst_coords: Array, snd_coords: Array) → Array

Concatenate component coordinates.

class Triple

Bases: Tuple, ABC, Generic

Product of three manifolds, with coordinates stored as [fst | snd | trd].

abstract property fst_man: First

First component manifold.

abstract property snd_man: Second

Second component manifold.

abstract property trd_man: Third

Third component manifold.

property dim: int

Total dimension is the sum of component dimensions.

split_coords(coords: Array) → tuple[Array, Array, Array]

Split into (fst, snd, trd) components.

join_coords(fst_coords: Array, snd_coords: Array, trd_coords: Array) → Array

Concatenate component coordinates.

class Replicated(rep_man: M, n_reps: int)

Bases: Manifold, Generic

Homogeneous product of \(n\) copies of the same manifold, stored flat as [n_reps * rep_man.dim].

Used for collections where every element lives on the same manifold (e.g. mixture components, time series states). The map method applies a function across copies via vmap.

Mathematically, the \(n\)-fold product \(\mathcal M^n = \mathcal M \times \cdots \times \mathcal M\) with \(\dim = n \cdot \dim(\mathcal M)\). Unlike Tuple, the homogeneity allows vmap-based operations over the copies.

rep_man: M

The base manifold being replicated.

n_reps: int

Number of copies of the base manifold.

property dim: int

Total dimension is product of base dimension and number of copies.

get_replicate(coords: Array, idx: int) → Array

Extract the idx-th replicate from flat coordinates.

to_2d(coords: Array) → Array

Convert flat coordinates to 2D array with shape [n_reps, rep_man.dim].

to_1d(array: Array) → Array

Convert 2D array back to flat coordinates.

map(f: Callable[[Array], Array], coords: Array, flatten: bool = False) → Array

Map a function across replicates.

By default, returns stacked 2D results for easier indexing and inspection. Use flatten=True when the result should be flat coordinates on another manifold.

Parameters:

• f – Function that takes coordinates for one replicate (shape [rep_man.dim])

• coords – Flat array of replicated coordinates (shape [n_reps * rep_man.dim])

• flatten – If True, return flat array [n_reps * f_result_dim]. If False (default), return stacked array [n_reps, *f_result_shape] for easier indexing and inspection.

Returns:

Stacked 2D array by default, or flat 1D array if flatten=True