SFI.statefunc.structexpr module

StructuredExpr — declarative structured tensor expressions (inner world).

Defines StructuredExpr, the user-facing expression type for the inner physics world, and the lightweight _StructNode IR that encodes the computation graph. Expressions are symbolic and cannot be evaluated directly; use Layout.embed() to compile them into outer-world StateExpr objects.

class SFI.statefunc.structexpr.StructuredExpr(sdims, n_features, param_suite, labels, _layout_id, _node)[source]

Bases: object

Declarative structured tensor expression (inner world).

Users build these via Layout methods and algebra. They are symbolic and cannot be evaluated directly — use Layout.embed() to compile to an outer-world StateExpr.

Variables:

  • sdims (tuple[int, ...]) – Per-axis sizes. () = scalar, (2,) = 2-vector, etc.

  • n_features (int) – Number of independent regression channels (last axis).

  • param_suite (ParamSuite | None) – Learnable parameters (None = pure basis).

  • labels (tuple[str, ...]) – Human-readable feature labels.

Parameters:

  • sdims (tuple[int, ...])

  • n_features (int)

  • param_suite (ParamSuite | None)

  • labels (tuple[str, ...])

  • _layout_id (int)

  • _node (_StructNode)

property T: StructuredExpr

Swap last two rank axes. Requires srank >= 2.

abs()[source]
Return type:

StructuredExpr

cos()[source]
Return type:

StructuredExpr

dense(n_out, *, name='dense')[source]

Affine projection of features. Adds a learnable weight matrix.

Parameters:

  • n_out (int)

  • name (str)

Return type:

StructuredExpr

dot(other)[source]

Contract last axis of self with last axis of other.

Parameters:

other (StructuredExpr)

Return type:

StructuredExpr

static einsum(spec, *operands)[source]

Einstein summation over rank axes.

Example:

Q = StructuredExpr.einsum('i,j->ij', n, n)
Parameters:
Return type:

StructuredExpr

elementwisemap(fn, *, name=None)[source]

Apply a JAX-traceable function elementwise.

Records a _UnaryOp("ew", …) node in the IR tree. At embed time this compiles to StateExpr.elementwisemap(fn).

Parameters:

  • name (str, optional) – Override the function name used in auto-generated labels. Defaults to fn.__name__.

  • fn (Callable)

Return type:

StructuredExpr

exp()[source]
Return type:

StructuredExpr

classmethod eye(sdim, *, layout_id=-1)[source]

Identity matrix with sdims=(sdim, sdim), n_features=1.

Parameters:

  • sdim (int)

  • layout_id (int)

Return type:

StructuredExpr

features_to_rank(target_sdims)[source]

Promote features into rank axes. Requires srank == 0.

sdims=(), n_features=12 → features_to_rank((3,4)) → sdims=(3,4), n_features=1.

Parameters:

target_sdims (tuple[int, ...])

Return type:

StructuredExpr

labels: tuple[str, ...]
log()[source]
Return type:

StructuredExpr

n_features: int
param_suite: ParamSuite | None
rank_to_features()[source]

Flatten all rank axes into the features axis.

sdims=(2,3), n_features=5 → sdims=(), n_features=30.

Return type:

StructuredExpr

sdims: tuple[int, ...]
sin()[source]
Return type:

StructuredExpr

sqrt()[source]
Return type:

StructuredExpr

property srank: int

Number of structured dimensions (rank axes).

classmethod stack(exprs, *, sdim=None)[source]

Stack expressions along a new leading rank axis.

All inputs must share the same sdims and n_features. sdim defaults to len(exprs).

Parameters:
Return type:

StructuredExpr

tanh()[source]
Return type:

StructuredExpr

with_label(label)[source]

Return a copy of this expression with a single human-readable label.

Useful for annotating derived quantities (arithmetic, einsum, …) so that print_report shows a meaningful term name instead of a generic fallback.

Parameters:

label (str) – Human-readable name for this term, e.g. "|Q|²Q".

Returns:

Identical expression with labels=(label,).

Return type:

StructuredExpr