jax_sn.tools.parametrisation module

Module to handle flat parametrised modules in a relatively simple way. In essense, a parametrised module is a graph (a sweep data depends on face normals which depend on vertices), and so we use some graph tools to create the graph, check that what you are modifying makes sense, etc. Used for parametrizing the jax_sn modules: the primitives are the parameters, everything else follows.

[Claude generated most of this!]

exception MissingQuantityError

Bases: AttributeError

Raised when a quantity cannot be derived because a primitive it depends on is None.

class Graph(quantities: FrozenSet[str], primitives: FrozenSet[str], stored: FrozenSet[str], projected: FrozenSet[str], derived: FrozenSet[str], frozen: FrozenSet[str], free: FrozenSet[str], conditionally_free: FrozenSet[str], deps: Mapping[str, FrozenSet[str]], primitive_deps: Mapping[str, FrozenSet[str]], required_deps: Mapping[str, FrozenSet[str]], projections: Mapping[str, Tuple[str, Callable]])

Bases: NamedTuple

What a class’s storage policy implies. All derived from the registry plus which quantities the class declared as fields; see ParametrisedModule.graph().

quantities: FrozenSet[str]

Alias for field number 0

primitives: FrozenSet[str]

Alias for field number 1

stored: FrozenSet[str]

Alias for field number 2

projected: FrozenSet[str]

Alias for field number 3

derived: FrozenSet[str]

Alias for field number 4

frozen: FrozenSet[str]

Alias for field number 5

free: FrozenSet[str]

Alias for field number 6

conditionally_free: FrozenSet[str]

Alias for field number 7

deps: Mapping[str, FrozenSet[str]]

Alias for field number 8

primitive_deps: Mapping[str, FrozenSet[str]]

Alias for field number 9

required_deps: Mapping[str, FrozenSet[str]]

Alias for field number 10

projections: Mapping[str, Tuple[str, Callable]]

Alias for field number 11

class ParametrisedModule

Bases: Module

Base for modules built from a derivation registry.

Subclasses set DERIVATIONS to a {quantity: function} mapping whose functions take their dependencies as named parameters. A concrete subclass declares as fields the primitives plus whichever derived quantities it wants to store; accessors for the rest are injected automatically.

Declare intermediate classes abstract=True:

class Interface(ParametrisedModule, abstract=True): ...

which suppresses that injection. Since it is an injected accessor that makes a class unsafe to inherit from – a field cannot override one – an abstract class is safe to inherit from, so shared fields and methods can live there and be declared once.

DERIVATIONS

The derivation registry. Inherited by subclasses unless overridden.

Type:

ClassVar[Mapping[str, Callable]]

PROJECTIONS

quantity -> {source quantity: extractor}, for the derivations that repackage rather than compute: how to read a quantity back out of a stored one that already holds it.

Type:

ClassVar[Mapping[str, Mapping[str, Callable]]]

NON_DIFFERENTIABLE

Quantities to stop_gradient. Marked here rather than inside the derivation, so no derivation can forget it and it holds on the stored and projected paths too.

Type:

ClassVar[FrozenSet[str]]

RETAINED_UNDER

quantity -> {primitive: predicate(held_value, new_primitive) -> bool}: when a held value of quantity stays valid across a replacement of primitive.

A dependency edge states only that a quantity is a function of a primitive. Some edges are much coarser than that – the harmonics depend on the cross sections only through their l_max – and without the predicate the graph must treat the two as incompatible: hold the quantity and the primitive is frozen, keep the primitive free and the quantity is recomputed. Declaring when the held value still serves relaxes both halves at once:

• updated() replaces such a primitive even though it is frozen, when the predicate holds, and refuses loudly when it does not – never a silent rebuild of what the policy promised to keep. These primitives are Graph.conditionally_free. Here the whole walk counts: a retention anywhere on the path shields what is stored above it, provided its own value is stored or projected, so there is something to check it against.

• frontier() stores a quantity only such a primitive held back. Here only a quantity’s own retentions count: one declared further down is on a value nothing will hold.

Two obligations. The predicate must be sound – it asserts the held value is a correct value for the new primitive, not merely a convenient one. And it runs in Python, so it may read only static structure (types, shapes, capabilities), never a traced leaf.

Type:

ClassVar[Mapping[str, Mapping[str, Callable[[Any, Any], bool]]]]

DERIVATIONS: ClassVar[Mapping[str, Callable]] = {}

PROJECTIONS: ClassVar[Mapping[str, Mapping[str, Callable]]] = {}

NON_DIFFERENTIABLE: ClassVar[FrozenSet[str]] = frozenset({})

RETAINED_UNDER: ClassVar[Mapping[str, Mapping[str, Callable[[Any, Any], bool]]]] = {}

classmethod graph()

This class’s storage policy and what it implies. Memoised on the class; it cannot be built in __init_subclass__ because dataclasses.fields is only populated after equinox’s metaclass has finished.

Return type:

Graph

classmethod frontier(*frozen: str)

The quantities worth storing if frozen never changes – plus those held back only by a primitive they have a retention for, which lands them in graph().conditionally_free.

Return type:

FrozenSet[str]

classmethod freezing(*frozen: str)

A subclass storing everything that frozen determines. Freezing a primitive no quantity is fully determined by is a no-op; the result’s graph().frozen says what was achieved. See frozen() for the instance-side spelling.

Return type:

type

classmethod storing(*quantities: str)

A subclass storing exactly quantities. Unlike freezing(), this states the storage directly rather than deriving it, so it can hold something a frontier would not pick – use it to seed an instance from data you already have, then materialise() into the policy you want. See stored() for the instance-side spelling.

Return type:

type

classmethod varying(*free: str)

A subclass that keeps free updatable and stores as much of the rest as it can. See varied() for the instance-side spelling.

Return type:

type

frozen(*frozen: str)

This instance under freezing().

Return type:

ParametrisedModule

stored(*quantities: str)

This instance under storing().

Return type:

ParametrisedModule

varied(*free: str)

This instance under varying().

Return type:

ParametrisedModule

available_quantities()

The quantities this instance can actually produce: a quantity is unavailable when a primitive it needs is None (e.g. angular quantities on a builder with no quadrature set).

Return type:

FrozenSet[str]

updated(**primitives: Any)

Replace primitives, keeping everything else. Only graph().free primitives may be replaced – plus those in graph().conditionally_free whose retention predicate accepts the new value (see RETAINED_UNDER); everything derived from them follows automatically.

Return type:

ParametrisedModule

materialise(cls: type)

Convert to another class over the same registry, storing whatever that class stores. Promoting pays a repeated derivation once; demoting frees more primitives to update. Raises MissingQuantityError if the target stores something unavailable here.

Return type:

ParametrisedModule