jax_sn.tools

assert_pytree_allclose(a, b, *, rtol=1e-12, atol=1e-12)[source]
class ParametrisedModule[source]

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()[source]

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)[source]

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)[source]

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)[source]

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)[source]

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)[source]

This instance under freezing().

Return type:

ParametrisedModule

stored(*quantities: str)[source]

This instance under storing().

Return type:

ParametrisedModule

varied(*free: str)[source]

This instance under varying().

Return type:

ParametrisedModule

available_quantities()[source]

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)[source]

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)[source]

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

Submodules