caskade package

Contents

caskade package#

Submodules#

caskade.backend module#

Backend abstraction for array operations.

Provides a unified Backend class that delegates array creation and manipulation to one of three libraries: torch, jax, or numpy. A module-level backend instance is created on import and serves as the primary interface for users.

class caskade.backend.Backend(backend=None)[source]#

Bases: object

Unified interface for array operations across torch, jax, and numpy.

Provides a single API for creating and manipulating arrays regardless of the underlying library. Methods such as make_array, concatenate, to, sigmoid, and logit are dynamically bound when the backend is set, delegating to the appropriate library-specific implementation.

Parameters:

backend (str, optional) – Backend name: "torch", "jax", or "numpy". If None, reads from the CASKADE_BACKEND environment variable, defaulting to "torch".

Examples

Use the module-level backend instance to switch backends:

from caskade import backend
backend.backend = "numpy"
arr = backend.make_array([1.0, 2.0, 3.0])
all(array)[source]#

Test whether all elements evaluate to True.

Parameters:

array (ArrayLike) – Input array.

Returns:

Scalar result; True if every element is non-zero.

Return type:

ArrayLike

any(array)[source]#

Test whether any element evaluates to True.

Parameters:

array (ArrayLike) – Input array.

Returns:

Scalar result; True if any element is non-zero.

Return type:

ArrayLike

property array_type#

The array class for the active backend.

Returns torch.Tensor, jax.numpy.ndarray, or numpy.ndarray depending on the current backend. Useful for isinstance checks.

Returns:

The array class used by the active backend.

Return type:

type

Examples

isinstance(my_array, backend.array_type)
Type:

type

property backend#

Name of the active backend ("torch", "jax", or "numpy").

Type:

str

exp(array)[source]#

Compute the exponential element-wise.

Parameters:

array (ArrayLike) – Input array.

Returns:

Element-wise exponential of the input.

Return type:

ArrayLike

log(array)[source]#

Compute the natural logarithm element-wise.

Parameters:

array (ArrayLike) – Input array.

Returns:

Element-wise natural logarithm of the input.

Return type:

ArrayLike

setup_jax()[source]#
setup_numpy()[source]#
setup_torch()[source]#
sum(array, axis=None)[source]#

Sum array elements over a given axis.

Parameters:

  • array (ArrayLike) – Input array.

  • axis (int or None, optional) – Axis along which to sum. If None, sums all elements.

Returns:

Sum of elements.

Return type:

ArrayLike

caskade.backend.backend = <caskade.backend.Backend object>#

Module-level Backend instance used as the default entry point. Import and configure this object to switch backends globally:

from caskade import backend
backend.backend = "numpy"

caskade.base module#

class caskade.base.Node(name: str | None = None, link: Node | tuple[Node] | None = None, description: str = '')[source]#

Bases: object

Base graph node class for caskade objects.

The Node object is the base class for all caskade objects. It is used to construct the directed acyclic graph (DAG). The primary function of the Node object is to manage the parent-child relationships between nodes in the graph. There is limited functionality for the Node object, though it implements the base versions of the active state and to / update_graph methods. The active state is used to communicate through the graph that the simulator is currently running. The to method is used to move and/or cast the values of the parameter. The update_graph method is used signal all parents that the graph below them has changed.

Examples

Example making some Node objects and then linking/unlinking them:

n1 = Node()
n2 = Node()
n1.link("subnode", n2) # link n2 as a child of n1, may use any str as the key
n1.unlink("subnode") # alternately n1.unlink(n2) to unlink by object
property active: bool#

True if the node is currently in an active simulation run.

Type:

bool

add_memo(memo)[source]#

Add a memo string and propagate it to all children.

Parameters:

memo (str) – The memo message to add. Children in subgraphs receive the memo with the child name appended (memo|child_name).

append_state(saveto: str | File)[source]#

Append the current state to an existing HDF5 file.

The file must have been previously created by save_state with appendable=True. The graph structure in the file is verified before appending.

Parameters:

saveto (str or File) – Path to an HDF5 file ('.h5' or '.hdf5') or an open HDF5 File object.

Raises:

  • GraphError – If the graph structure no longer matches the file.

  • NotImplementedError – If the file path does not end with a supported extension.

property children: dict[str, Node]#

Mapping of link keys to child nodes.

Type:

dict[str, Node]

graph_dict() → dict[str, dict][source]#

Return a nested dictionary representation of the graph.

Each key is a string of the form "name|node_type" and the value is a dict containing the same structure for that node’s children.

Returns:

Nested dictionary mirroring the DAG hierarchy.

Return type:

dict[str, dict]

graph_print(dag: dict, depth: int = 0, indent: int = 4, result: str = '') → str[source]#

Recursively render a graph dictionary as an indented string.

Parameters:

  • dag (dict[str, dict]) – A nested dictionary as returned by graph_dict.

  • depth (int, optional) – Current indentation depth (used during recursion). Defaults to 0.

  • indent (int, optional) – Number of spaces per indentation level. Defaults to 4.

  • result (str, optional) – Accumulator string (used during recursion). Defaults to "".

Returns:

A human-readable, indented representation of the graph.

Return type:

str

graphviz(saveto: str | None = None) → graphviz.Digraph[source]#

Return a graphviz Digraph representing the DAG below this node.

Parameters:

saveto (str, optional) – If provided, save the rendered graph to this file path. The file extension determines the output format (e.g. '.pdf', '.png'). Defaults to None.

Returns:

The constructed directed-graph object.

Return type:

graphviz.Digraph

property graphviz_style#

Link the current Node object to another Node object as a child in a hierarchical manner. See link for more detail on linking. A hierarchical link will allow batching internally to the simulator.

Parameters:

  • key (str) – The key to link the child node with.

  • child (Node) – The child Node object to link to.

Examples

parent = Node(name="parent")
child = Node(name="child")
parent.hierarchical_link("child", child)

Link the current Node object to another Node object as a child.

Parameters:

  • key ((Union[str, Node])) – The key to link the child node with. This will also become the attribute to access the child node. After linking you will have node.key == child

  • child ((Optional[Node], optional)) – The child Node object to link to. Defaults to None in which case the key is used as the child and the child.name is used as the key.

Examples

Example making some Node objects and then linking/unlinking them, demonstrating multiple ways to link/unlink:

n1 = Node()
n2 = Node()

n1.link("subnode", n2)  # may use any str as the key
n1.unlink("subnode")

# Alternatively, link by object
n1.link(n2)
n1.unlink(n2)
load_state(loadfrom: str | File, index: int = -1, **kwargs)[source]#

Load node state (and children) from an HDF5 file.

Parameters:

  • loadfrom (str or File) – Path to an HDF5 file ('.h5' or '.hdf5') or an open HDF5 File object.

  • index (int, optional) – Sample index to load when the file was saved in appendable mode. Defaults to -1 (last sample).

  • **kwargs – Additional keyword arguments forwarded to h5py.File (e.g. driver).

Raises:

  • GraphError – If the graph structure no longer matches the file.

  • NotImplementedError – If the file path does not end with a supported extension.

property memos: set[str]#

Current set of memo strings held by this node.

Type:

set[str]

property name: str#

The name of this node.

Type:

str

property node_str#
property online: bool#

True if the node is online within a hierarchical sub-graph.

Type:

bool

property parents: set[Node]#

Set of parent nodes that link to this node.

Type:

set[Node]

remove_memo(memo)[source]#

Remove a memo string and propagate removal to all children.

Parameters:

memo (str) – The memo message to remove. The same propagation rules as add_memo apply.

save_state(saveto: str | File, appendable: bool = False)[source]#

Save the state of the node and its children, currently only works for HDF5 file types (.h5 and .hdf5).

The “state” of a node is considered to be the value of its params, however it is also possible to save other attributes of the node by adding them to the Node.saveattrs set. Simply call Node.saveattrs.add(‘attribute’) and then Node.attribute will be saved if possible. The HDF5 file will be created with the same structure as the graph, even if there are multiple paths to the same node. For example if N1 has children N2 and N3, and both N2 and N3 have the child N4, the HDF5 file will reflect this. It will be possible to find the N4 params under both ‘N1/N2/N4’ and ‘N1/N3/N4’ if inspecting the HDF5 file manually. Specifically, if N4 has the param P1 then you could access its value like this:

If the save had been set as appendable, then the value will have an extra dimension for the number of samples, this will always be the first dimension. If appendable was false then the value will simply equal the param value.

Note

You need the optional h5py package installed to use this method.

Parameters:

  • saveto ((Union[str, File])) – The file to save the state to. If a string, it should be the path to an HDF5 file (ending in ‘.h5’ or ‘.hdf5’). If a File object, it should be an open HDF5 file.

  • appendable ((bool, optional)) – Whether to save the state in an appendable format. If True, the values will have an extra dimension for the number of samples. Defaults to False.

property subgraphs: set[Node]#

Subset of children linked hierarchically.

Type:

set[Node]

to(device=None, dtype=None)[source]#

Moves and/or casts the values of the Node to a particular device and/or dtype.

Parameters:

  • device ((optional)) – The device to move the values to. Defaults to None.

  • dtype ((optional)) – The desired data type. Defaults to None.

topological_ordering() → tuple[Node][source]#

Return a topological ordering of the graph below the current node.

Performs a recursive depth-first search with post-order traversal to resolve dependencies. The result starts with this node and proceeds to its descendants in dependency order.

Returns:

All nodes reachable from (and including) this node, ordered so that every parent appears before its children.

Return type:

tuple[Node]

Unlink one or more Node objects from this Node.

Parameters:

key ((str, Node, list, tuple, or None, optional)) – The key, Node object, or collection of keys/nodes to unlink. If a string, the child with that key is unlinked. If a Node object, the matching child is located and unlinked. If a list or tuple, each element is unlinked in turn. If None (the default), all children are unlinked.

Raises:

GraphError – If the graph is currently active.

update_graph()[source]#

Triggers a call to all parents that the graph below them has been updated. The base Node object does nothing with this information, but other node types may use this to update internal state.

caskade.collection module#

class caskade.collection.NodeCollection(name: str | None = None, link: Node | tuple[Node] | None = None, description: str = '')[source]#

Bases: Node, GetSetValues

Base mixin for collections of nodes that track parameters.

Provides shared functionality for traversing, querying, and converting parameters within a graph of nodes. Subclasses such as NodeTuple and NodeList combine this mixin with a standard Python sequence type.

copy()[source]#
deepcopy()[source]#
property dynamic#

Whether any node in this collection has dynamic parameters.

Returns:

True if at least one contained node is dynamic.

Return type:

bool

property dynamic_param_groups: tuple[int]#

Sorted unique group identifiers of all dynamic parameters.

Returns:

Sorted group indices present among the dynamic parameters.

Return type:

tuple of int

property dynamic_params: tuple[Param]#

All dynamic parameters in the graph below this node.

Returns:

Dynamic (non-static, non-pointer) parameters found via topological ordering.

Return type:

tuple of Param

property pointer_params: tuple[Param]#

All pointer parameters in the graph below this node.

Returns:

Parameters that act as pointers to other parameters, found via topological ordering.

Return type:

tuple of Param

property static#

Whether all nodes in this collection are static.

Returns:

True if no contained node is dynamic.

Return type:

bool

property static_params: tuple[Param]#

All static parameters in the graph below this node.

Returns:

Static (non-dynamic, non-pointer) parameters found via topological ordering.

Return type:

tuple of Param

to_dynamic(children_only=True)[source]#

Change all parameters to dynamic parameters.

Parameters:

children_only ((bool, optional)) – If True, only convert the children of this module to dynamic. If False, convert all parameters in the graph below this module. Defaults to True.

to_static(children_only=True)[source]#

Change all parameters to static parameters.

Parameters:

children_only ((bool, optional)) – If True, only convert children of this module. If False, convert all parameters in the graph below this module. Defaults to True.

class caskade.collection.NodeDict(mapping=None, name=None)[source]#

Bases: NodeCollection, dict

Mutable, keyed collection of nodes.

Behaves like a standard dict but also participates in the caskade node graph. All elements must be Node instances. Graph links are automatically updated whenever the dict is modified.

Parameters:

  • mapping (mapping of str to Node, optional) – Nodes to include in the dict. Defaults to an empty dict.

  • name (str, optional) – Human-readable name for this collection of nodes.

clear()[source]#

Remove all nodes from the dict and update graph links.

property dynamic#

Whether any node in this collection has dynamic parameters.

Returns:

True if at least one contained node is dynamic.

Return type:

bool

property graphviz_style#
pop(key, *args)[source]#

Remove and return a node from the dict and update graph links.

popitem()[source]#

Remove and return an arbitrary (key, node) pair from the dict (the last one inserted) and update graph links.

setdefault(key, default)[source]#

If key is in the dictionary, return its value. If not, insert key with a value of default and return default. Update graph links.

update(mapping=None, **kwargs)[source]#

Update the dict with another mapping (i.e. dict) and update graph links.

class caskade.collection.NodeList(iterable=(), name=None)[source]#

Bases: NodeCollection, list

Mutable, ordered collection of nodes.

Behaves like a standard list but also participates in the caskade node graph. All elements must be Node instances. Graph links are automatically updated whenever the list is modified.

Parameters:

  • iterable (iterable of Node, optional) – Nodes to include in the list. Defaults to an empty iterable.

  • name (str, optional) – Human-readable name for this collection node.

append(node)[source]#

Append a node to the list and update graph links.

clear()[source]#

Remove all nodes from the list and update graph links.

extend(iterable)[source]#

Extend the list with nodes from an iterable and update graph links.

property graphviz_style#
insert(index, node)[source]#

Insert a node at the given index and update graph links.

pop(index=-1)[source]#

Remove and return a node at the given index, updating graph links.

remove(value)[source]#

Remove the first occurrence of a node and update graph links.

class caskade.collection.NodeTuple(iterable=None, name=None)[source]#

Bases: NodeCollection, tuple

Immutable, ordered collection of nodes.

Behaves like a standard tuple but also participates in the caskade node graph. All elements must be Node instances and are automatically linked as children upon construction.

Parameters:

  • iterable (iterable of Node, optional) – Nodes to include in the tuple.

  • name (str, optional) – Human-readable name for this collection node.

property graphviz_style#

caskade.context module#

class caskade.context.ActiveContext(module: Module)[source]#

Bases: object

Context manager to activate a module for a simulation.

Only inside an ActiveContext is it possible to fill or clear the dynamic and live parameters. On entry, the module is marked as active (or its current parameter state is saved if already active). On exit, the state is restored.

Parameters:

module (Module) – The module to activate for the duration of the context.

Raises:

ActiveStateError – If the module is already running a simulation (module.online is True).

Examples

Activate a module, fill parameters, and run a forward pass:

with ActiveContext(my_module):
    my_module.fill_params(params)
    result = my_module.my_forward(x)
class caskade.context.OverrideParam(param: Param, value)[source]#

Bases: object

Context manager to override a parameter value.

Only inside an OverrideParam will the parameter be set to the new value. The original value (and the values of any parent pointer parameters) are saved on entry and restored on exit.

Parameters:

  • param (Param) – The parameter whose value should be temporarily overridden.

  • value (object) – The temporary value to assign to param.

Examples

Override a parameter inside a @forward method so that it uses new_value regardless of what was passed via params:

class MySim(Module):
    def __init__(self):
        super().__init__()
        self.a = Param("a", None)
        self.b = Param("b", None)

    @forward
    def __call__(self, x, a=None, b=None):
        with OverrideParam(self.b, 5.0):
            # b will always be 5.0 here, ignoring params
            return x + a + self.b.value
class caskade.context.ValidContext(module: Module)[source]#

Bases: object

Context manager that transforms parameter values to an unconstrained space.

Inside a ValidContext, all parameter values are automatically mapped into the range (-inf, inf) via each parameter’s to_valid / from_valid transformations. This is useful when interfacing with samplers or optimizers that expect unconstrained parameters—any value they propose will be mapped back into the parameter’s original valid range on exit.

Parameters:

module (Module) – The module whose parameters should be transformed.

Examples

Get unconstrained parameter values for use with an optimizer:

with ValidContext(my_module):
    unconstrained_params = my_module.get_values()
    # unconstrained_params live in (-inf, inf)

caskade.decorators module#

class caskade.decorators.active_cache(func)[source]#

Bases: object

Caches the first evaluated result of a Module method for the duration of a simulation.

This decorator ensures that an expensive method is executed exactly once per active simulation run. Once calculated, subsequent calls to the decorated method will return the stored value, ignoring any arguments passed to it.

Warning

If the method is called multiple times with different arguments in one simulation, the cached result will still be returned, which may lead to unexpected behavior. Use with caution!

Notes

If you are stacking multiple decorators on a method (such as @forward or @jax.jit), @active_cache MUST be the outermost (top) decorator.

Examples

class FluxModel(Module):
    def __init__(self, nodes, x, M):
        super().__init__()
        self.nodes = nodes
        self.x = Param("x", x)
        self.M = Param("M", M)

    @active_cache
    @jax.jit  # Notice active_cache is placed at the top
    @forward
    def compute_intrinsic_sed(self, w, x, M):
        print("Computing SED...")
        return jnp.interp(w, self.nodes, x * M)

    @forward
    def compute_flux(self, wavelengths):
        sed = self.compute_intrinsic_sed(wavelengths)  # Cached after first call
        flux = jnp.sum(sed)
        sed = self.compute_intrinsic_sed(wavelengths)  # Returns cached result, no print
        peak = jnp.max(sed)
        return flux, peak

model = FluxModel(np.linspace(400, 700, 10), x=1.0, M=np.random.rand(10))

# Compute flux only calls compute_intrinsic_sed once due to caching
flux, peak = model.compute_flux(wavelengths)
caskade.decorators.forward(method)[source]#

Decorator to define a forward method for a module.

Manages parameter passing and activation for the decorated method. When called, it automatically fills keyword arguments from the module’s Param children and handles parameter overrides and active context.

Parameters:

method ((Callable)) – The forward method to be decorated.

Returns:

The decorated forward method.

Return type:

Callable

Examples

Standard usage of the forward decorator:

class ExampleSim(Module):
    def __init__(self, a, b, c):
        super().__init__("example_sim")
        self.a = a
        self.b = Param("b", b)
        self.c = Param("c", c)

    @forward
    def example_func(self, x, b=None):
        return x + self.a + b

E = ExampleSim(a=1, b=None, c=3)
print(E.example_func(4, params=[5]))
# Output: 10

caskade.errors module#

exception caskade.errors.ActiveStateError[source]#

Bases: CaskadeException

Exception for active-state errors in caskade.

Raised when an operation requires a node to be in a particular active state (enabled or disabled) and that condition is not met.

exception caskade.errors.BackendError[source]#

Bases: CaskadeException

Exception for backend-related errors in caskade.

Raised when the selected numerical backend encounters an unsupported operation or configuration issue.

exception caskade.errors.CaskadeException[source]#

Bases: Exception

Base class for all exceptions in caskade.

All custom exceptions raised by caskade inherit from this class, allowing users to catch any caskade-specific error with a single except clause.

exception caskade.errors.FillParamsArrayError(name, input_params, params)[source]#

Bases: FillParamsError

Exception raised when filling parameters with an array fails.

Raised when the shape of the input array does not match the total number of flattened parameters registered on a node.

Parameters:

  • name (str) – Name of the node whose parameters are being filled.

  • input_params (ArrayLike) – The input array that was provided.

  • params (tuple of Param) – Registered parameters whose shapes are compared against the input.

exception caskade.errors.FillParamsError[source]#

Bases: CaskadeException

Base exception for errors when filling parameters in caskade.

Raised when the input data provided to fill node parameters is invalid. Subclasses handle specific input types (array, sequence, mapping).

exception caskade.errors.FillParamsMappingError(name, children, missing_key=None)[source]#

Bases: FillParamsError

Exception raised when filling parameters with a mapping fails.

Raised when a key in the input dictionary does not correspond to any registered child node.

Parameters:

  • name (str) – Name of the node whose parameters are being filled.

  • children (dict) – Dictionary of registered child nodes.

  • missing_key (str, optional) – The key from the input mapping that was not found among the node’s children.

exception caskade.errors.FillParamsSequenceError(name, input_params, dynamic_params)[source]#

Bases: FillParamsError

Exception raised when filling parameters with a sequence fails.

Raised when the length of the input sequence does not match the number of dynamic parameters registered on a node.

Parameters:

  • name (str) – Name of the node whose parameters are being filled.

  • input_params (sequence) – The input sequence (list, tuple, etc.) that was provided.

  • dynamic_params (tuple of Param) – Registered dynamic parameters expected by the node.

exception caskade.errors.GraphError[source]#

Bases: CaskadeException

Exception for graph-related errors in caskade.

Raised when an operation on the computational graph is invalid, such as creating cycles or referencing nonexistent nodes.

exception caskade.errors.LinkToAttributeError[source]#

Bases: GraphError

Exception raised when linking to an attribute fails.

Raised when an attempt is made to create a link to a node attribute that does not exist or is not a valid link target.

exception caskade.errors.NodeConfigurationError[source]#

Bases: CaskadeException

Exception for node configuration errors in caskade.

Raised when a node is configured with invalid or incompatible settings.

exception caskade.errors.ParamConfigurationError[source]#

Bases: NodeConfigurationError

Exception for parameter configuration errors in caskade.

Raised when a parameter is defined with an invalid shape, type, or constraint.

exception caskade.errors.ParamTypeError[source]#

Bases: CaskadeException

Exception for parameter type errors in caskade.

Raised when a value assigned to a parameter does not match its expected type.

caskade.mixins module#

class caskade.mixins.GetSetValues[source]#

Bases: object

Mixin providing methods for getting and setting parameter values.

Provides array, list, and dict interfaces for reading and writing dynamic (or static) parameter values on a Module or NodeCollection. Also includes helpers for locating parameters by index and for converting between raw and valid (transformed) parameter spaces.

find_index(param: Param | tuple[Param] | Module, scheme: str = 'array') → int | slice[source]#

Identify what index is associated with a param in the dynamic params array.

Parameters:

  • param (Union[Param, tuple[Param], Module]) – The param for which to find the associated index.

  • scheme (str) – Whether to search the array (default) params or list version of params. dict is currently unsupported.

Returns:

param_info – A int giving the index associated with the provided Param object. If the param is multi-dimensional then the result will be a slice over all indices associated with that param.

Return type:

Union[int, slice]

find_param(idx: int | tuple[int], group: int | None = None, scheme: str = 'array') → tuple[Param, tuple[int]][source]#

Identify which param is associated with the provided index in the dynamic params array.

Parameters:

  • idx (Union[int, tuple[int]]) – The index in the params array at which we wish to find the associated param.

  • group (Optional[int]) – If the dynamic params have multiple group values, then this argument specifies which group to check.

  • scheme (str) – Whether to search the array (default) params or list version of params. dict is currently unsupported.

Returns:

param_info – A tuple with the Param object and the index within the Param value associated with idx (empty tuple if scalar). If idx is a tuple then the result is a tuple of these results.

Return type:

tuple[Param, tuple[int]]

from_valid(valid_params: ArrayLike | Sequence | Mapping, param_list: tuple[Param] | None = None, group: int | None = None) → ArrayLike | Sequence | Mapping[source]#

Map parameter values from the unconstrained space back to their natural range.

Takes values in the unconstrained domain (-inf, inf) (as produced by to_valid() or proposed by an optimizer/sampler) and maps them back into each parameter’s original valid range (e.g. 0-1 for an axis ratio).

Parameters:

  • valid_params (Union[ArrayLike, Sequence, Mapping]) – Parameter values in the unconstrained (-inf, inf) domain, in any supported format (array, sequence, or mapping).

  • param_list (tuple of Param or None, optional) – Subset of parameters to transform. Defaults to all dynamic parameters.

  • group (int or None, optional) – Restrict to a specific parameter group. When None and multiple groups exist, all groups are transformed.

Returns:

Inverse-transformed values in the same format as valid_params.

Return type:

Union[ArrayLike, Sequence, Mapping]

get_values(scheme: str = 'array', dynamic: bool = True, attribute: str | Callable = 'value', group: int | None = None, respect_valid: bool = True) → ArrayLike | list[ArrayLike] | dict[str, dict | ArrayLike][source]#

Retrieve parameter values from the module.

Parameters:

  • scheme (str, optional) –

    Output format, one of "array" (default), "list", or "dict".

    • "array" / "tensor" - returns a single flat array with all parameter values concatenated along the last axis.

    • "list" - returns a list of raw parameter values.

    • "dict" - returns a nested dictionary mirroring the graph structure.

  • dynamic (bool, optional) – If True (default), retrieves dynamic parameters; otherwise retrieves static parameters.

  • attribute (Union[str, Callable], optional) – The Param attribute to read from, by default "value".

  • group (int or None, optional) – Restrict to a specific parameter group. When None (default) and multiple groups exist, returns a list of per-group results.

  • respect_valid (bool, optional) – If True (default), will apply to_valid transformation.

Returns:

Parameter values in the format specified by scheme. When multiple groups exist and group is None, a list of per-group results is returned.

Return type:

Union[ArrayLike, list[ArrayLike], dict[str, Union[dict, ArrayLike]]]

set_values(params: ArrayLike | Sequence | Mapping, dynamic: bool = True, attribute: str = 'value', respect_valid: bool = True)[source]#

Fill parameter values of the module from the provided data.

Parameters:

  • params (Union[ArrayLike, Sequence, Mapping]) –

    Values to assign to the parameters. Accepted formats:

    • ArrayLike - a flat (or batched) array whose last dimension is concatenated parameter values in topological order.

    • Sequence - one element per parameter, matched by position.

    • Mapping - keys matching child names, values being the parameter data (may be nested).

    When multiple dynamic parameter groups exist, params should be a sequence of per-group containers.

  • dynamic (bool, optional) – If True (default), sets dynamic parameters; otherwise sets static parameters.

  • attribute (str, optional) – The Param attribute to write to, by default "value".

  • respect_valid (bool, optional) – If True (default), will apply from_valid transformations before assignment.

Raises:

ActiveStateError – If the module is currently in an active (tracing) state.

to_valid(params: ArrayLike | Sequence | Mapping, param_list: tuple[Param] | None = None, group: int | None = None) → ArrayLike | Sequence | Mapping[source]#

Map parameter values from their natural range to an unconstrained space.

Takes parameter values that lie within each parameter’s valid range (e.g. 0-1 for an axis ratio) and maps them into the unconstrained domain (-inf, inf). The inverse mapping from_valid() will map any value in (-inf, inf) back into the original valid range. This is useful for interfacing with samplers and optimizers that require unconstrained parameters.

Parameters:

  • params (Union[ArrayLike, Sequence, Mapping]) – Raw parameter values in any supported format (array, sequence, or mapping).

  • param_list (tuple of Param or None, optional) – Subset of parameters to transform. Defaults to all dynamic parameters.

  • group (int or None, optional) – Restrict to a specific parameter group. When None and multiple groups exist, all groups are transformed.

Returns:

Transformed values in the same format as params.

Return type:

Union[ArrayLike, Sequence, Mapping]

property valid_context: bool#

Return True if the module is in a valid context.

caskade.module module#

class caskade.module.Module(name: str | None = None, **kwargs)[source]#

Bases: Node, GetSetValues

Node to represent a simulation module in the graph.

The Module object is used to represent a simulation module in the graph. These are python objects that contain the calculations for a simulation, they also hold the Param objects that are used in the calculations. The Module object has additional functionality to manage the Param objects below it in the graph, it keeps track of all dynamic Param objects so that at runtime their values may be filled. The Module object manages its links to other nodes through attributes of the class.

Examples

Example of a nested pair of Module objects and how their @forward methods are called:

class MySim(Module):
    def __init__(self, a, b=None):
        super().__init__()
        self.a = a
        self.b = Param("b", b)

    @forward
    def myfunc(self, x, b=None):
        return x * self.a.otherfun(x) + b

class OtherSim(Module):
    def __init__(self, c=None):
        super().__init__()
        self.c = Param("c", c)

    @forward
    def otherfun(self, x, c = None):
        return x + c

othersim = OtherSim()
mysim = MySim(a=othersim)
#                       b                         c
params = [torch.tensor([1.0, 2.0]), torch.tensor([3.0, 4.0])]
result = mysim.myfunc(3.0, params=params)
# result is tensor([19.0, 23.0])
property all_params#

All parameters below this module in the DAG.

Returns:

Concatenation of static, dynamic, and pointer parameters.

Return type:

tuple of Param

clear_state()[source]#

Clear the active state _value for all params below this Module in the DAG. This should not be used by a user under normal circumstances.

property dynamic: bool#

Return True if the module has dynamic parameters as direct children.

Returns:

True if any direct children are dynamic parameters.

Return type:

bool

fill_kwargs(keys: tuple[str]) → dict[str, ArrayLike][source]#

Fill the kwargs for an @forward method with the values of the dynamic parameters. The requested keys are matched to names of Param objects owned by the Module. This should not be used by the user under normal circumstances.

fill_params(params: ArrayLike | Sequence | Mapping, dynamic=True)[source]#

Fill the dynamic/static parameters of the module with the input values from params.

Parameters:

  • params ((Union[ArrayLike, Sequence, Mapping])) – The input values to fill the dynamic parameters with. The input can be an ArrayLike, a Sequence, or a Mapping.

  • dynamic (bool) – Operate on dynamic parameters (True, default) or static parameters (False).

property graphviz_style#
property node_str: str#

Returns a string representation of the node for graph visualization.

param_order()[source]#

Return a human-readable string of dynamic parameter ordering.

Each line corresponds to a parameter group and lists the parameters in the format parent_name: param_name.

Returns:

Multi-line string describing the dynamic parameter order.

Return type:

str

remove_memo(memo)[source]#

Remove a memo string and propagate removal to all children.

Parameters:

memo (str) – The memo message to remove. The same propagation rules as add_memo apply.

property static: bool#

Return True if the module has no dynamic parameters as direct children.

Returns:

True if none of the direct children are dynamic parameters.

Return type:

bool

to_dynamic(children_only=True)[source]#

Change all parameters to dynamic parameters.

Parameters:

children_only ((bool, optional)) – If True, only convert the children of this module to dynamic. If False, convert all parameters in the graph below this module. Defaults to True.

to_static(children_only=True)[source]#

Change all parameters to static parameters.

Parameters:

children_only ((bool, optional)) – If True, only convert children of this module to static. If False, convert all parameters in the graph below this module. Defaults to True.

update_graph()[source]#

Maintain a tuple of dynamic, static, and pointer parameters at all points lower in the DAG.

caskade.param module#

class caskade.param.Param(name: str, value: ArrayLike | float | int | None = None, shape: tuple[int, ...] | None = None, cyclic: bool = False, valid: tuple[ArrayLike | float | int | None] | None = None, units: str | None = None, dynamic: bool | None = None, group: int = 0, batch_shape: tuple[int] | None = None, dtype: Any | None = None, device: Any | None = None, **kwargs)[source]#

Bases: Node

Node to represent a parameter in the graph.

The Param object is used to represent a parameter in the graph. During runtime this will represent a value which can be used in various calculations. The Param object can be set to a constant value (static); None meaning the value is to be provided at runtime (dynamic); another Param object meaning it will take on that value at runtime (pointer); or a function of other Param objects to be computed at runtime (also pointer, see user guides). These options allow users to flexibly set the behavior of the simulator.

Examples

Example making some Param objects:

p1 = Param("test", (1.0, 2.0)) # constant value, length 2 vector
p2 =Param("p2", None, (2,2)) # dynamic 2x2 matrix value
p3 = Param("p3", p1) # pointer to another parameter
p4 = Param("p4", lambda p: p.children["other"].value * 2) # arbitrary function of another parameter
p5 = Param("p5", valid=(0.0,2*pi), units="radians", cyclic=True) # parameter with metadata
Parameters:

  • name ((str)) – The name of the parameter.

  • value ((Optional[Union[ArrayLike, float, int]], optional)) – The value of the parameter. Defaults to None meaning dynamic.

  • shape ((Optional[tuple[int, ...]], optional)) – The shape of the parameter. Defaults to () meaning scalar.

  • cyclic ((bool, optional)) – Whether the parameter is cyclic, imposing periodic boundary conditions. Such as a rotation from 0 to 2pi. Defaults to False.

  • valid ((Optional[tuple[Union[ArrayLike, float, int, None]]], optional)) – The valid range of the parameter. Defaults to None meaning all of -inf to inf is valid.

  • units ((Optional[str], optional)) – The units of the parameter. Defaults to None.

  • dynamic ((bool, optional)) – Force param to be dynamic if True. If a value is provided and param is dynamic then it has a default value at call time.

  • (bool (batched) – If True, the param is assumed batched and the shape may now take the form (*B, *D) where *D is the shape of the value.

  • optional) – If True, the param is assumed batched and the shape may now take the form (*B, *D) where *D is the shape of the value.

  • dtype ((Optional[Any], optional)) – The data type of the parameter. Defaults to None meaning the data type will be inferred from the value.

  • device ((Optional[Any], optional)) – The device of the parameter. Defaults to None meaning the device will be inferred from the value.

property batch_shape: tuple[int, ...]#

The batch dimensions of the parameter value.

Batch dimensions are the leading dimensions of the value that precede the event shape. If an explicit batch shape was set it is returned directly; otherwise it is inferred from the value.

Returns:

The batch shape, or () if the parameter is not batched.

Return type:

tuple of int

property batched: bool#

Whether this parameter carries batch dimensions.

Returns:

True if batch_shape is non-empty.

Return type:

bool

property cyclic: bool#

Whether the parameter has cyclic (periodic) boundary conditions.

When True, values wrap around the valid range (e.g. an angle from 0 to 2π).

Returns:

True if the parameter is cyclic.

Return type:

bool

property device: str | None#

The device on which the parameter value resides.

If no explicit device was set, the device is inferred from the current value.

Returns:

The device, or None if unknown.

Return type:

device or None

property dtype: str | None#

The data type of the parameter value.

If no explicit dtype was set, the dtype is inferred from the current value.

Returns:

The data type, or None if unknown.

Return type:

dtype or None

property dynamic: bool#

Whether this parameter is dynamic.

Returns:

True if the parameter’s value is provided at runtime.

Return type:

bool

property graphviz_style#
property group: int#

The group index of this parameter.

Parameters that share the same group index are collected together into a single params object when calling a simulator’s @forward method, as well as when using get_values or set_values.

Returns:

The group index (default 0).

Return type:

int

is_valid(value=None) → bool[source]#

Check whether a value lies within the allowed range.

Parameters:

value (ArrayLike or None, optional) – The value to check. If None (default), the parameter’s current value is used.

Returns:

True if the value is within the valid range or if no constraints are set. False otherwise; a warning is also emitted.

Return type:

bool

property node_str: str#

Returns a string representation of the node for graph visualization.

property node_type#

The current type of this parameter node.

Returns:

One of "static", "dynamic", or "pointer".

Return type:

str

property npvalue: ndarray#

The current value converted to a NumPy array.

Returns:

The value as a NumPy ndarray.

Return type:

numpy.ndarray

property pointer: bool#

Whether this parameter is a pointer.

Returns:

True if the parameter points to another Param or a callable that is evaluated at runtime.

Return type:

bool

property shape: tuple[int, ...]#

The event (non-batch) shape of the parameter value.

Wildcard dimensions (None) in the declared shape are resolved using the current value. If no shape was declared, the shape of the current value is returned directly.

Returns:

The resolved shape of the parameter.

Return type:

tuple of int

property static: bool#

Whether this parameter is static.

Returns:

True if the parameter holds a fixed value that does not change at runtime.

Return type:

bool

to(device=None, dtype=None) → Param[source]#

Moves and/or casts the values of the parameter.

Parameters:

  • device ((optional)) – The device to move the values to. Defaults to None.

  • dtype ((optional)) – The desired data type. Defaults to None.

to_dynamic(value=<object object>)[source]#

Change this parameter to a dynamic parameter.

If a value is provided, it is stored as the default dynamic value. When called without arguments the existing value (if any) is kept.

Parameters:

value (ArrayLike, float, int, None, or sentinel, optional) – The default value for the dynamic parameter. Must not be a Param or callable. By default the current value is retained.

Raises:
to_pointer(value, link=())[source]#

Change this parameter to a pointer parameter.

The parameter’s value will be computed at runtime by dereferencing another Param or by calling a user-supplied function.

Parameters:

  • value (Param or callable) – A Param whose value will be mirrored, or a callable f(param) -> ArrayLike evaluated at runtime.

  • link (Node or tuple of Node, optional) – Additional nodes to link into the graph when creating the pointer. Defaults to an empty tuple.

Raises:
to_static(value=<object object>)[source]#

Change this parameter to a static parameter.

If a value is provided, it is stored as the fixed static value. When called without arguments the existing value (if any) is kept.

Parameters:

value (ArrayLike, float, int, None, or sentinel, optional) – The constant value for the static parameter. Must not be a Param or callable. By default the current value is retained.

Raises:
property valid: tuple[ArrayLike | None, ArrayLike | None]#

The valid range of the parameter value.

Returns:

(lower_bound, upper_bound). Either bound may be None indicating no constraint on that side.

Return type:

tuple of (ArrayLike or None, ArrayLike or None)

property value: ArrayLike | None#

The current value of the parameter.

For static and dynamic parameters the stored value is returned. For pointer parameters the linked callable is evaluated. During an active simulation the result is cached.

Returns:

The parameter value, or None if no value has been set.

Return type:

ArrayLike or None

caskade.param.valid_shape(batch_shape, shape, value_shape)[source]#

Check whether a value’s shape is compatible with a parameter’s shape.

Validates that value_shape is consistent with the declared shape and optional batch_shape. Dimensions set to None in shape act as wildcards and match any size.

Parameters:

  • batch_shape (tuple of int or None) – Leading batch dimensions, or None if the parameter is not batched.

  • shape (tuple of int or None, or None) – Expected event dimensions. Individual entries may be None (wildcard). If the entire argument is None, any shape is accepted.

  • value_shape (tuple of int) – The actual shape of the value to validate.

Returns:

True if the shapes are compatible, False otherwise.

Return type:

bool

caskade.tests module#

caskade.tests.test()[source]#

Run a basic integration test to verify caskade is installed and working.

Exercises core functionality including Module and Param creation, parameter linking, and forward method execution.

Examples

import caskade
caskade.test()
# Output: Success!

caskade.utils module#

caskade.utils.broadcast_cat_jax(arrays, dim=-1)[source]#

Concatenates JAX arrays with broadcasting.

Behaves like jnp.concatenate, but first broadcasts the arrays to match on all dimensions EXCEPT the concatenation dimension.

Parameters:

  • arrays (sequence of jnp.ndarray) – Arrays to concatenate.

  • dim (int) – The dimension along which to concatenate.

Returns:

The concatenated array.

Return type:

jnp.ndarray

caskade.utils.broadcast_cat_numpy(arrays, dim=-1)[source]#

Concatenates NumPy arrays with broadcasting.

Behaves like np.concatenate, but first broadcasts the arrays to match on all dimensions EXCEPT the concatenation dimension.

Parameters:

  • arrays (sequence of np.ndarray) – Arrays to concatenate.

  • dim (int) – The dimension along which to concatenate.

Returns:

The concatenated array.

Return type:

np.ndarray

caskade.utils.broadcast_cat_torch(tensors, dim=-1)[source]#

Concatenates tensors with broadcasting.

It behaves like torch.cat, but first broadcasts the tensors to match on all dimensions EXCEPT the concatenation dimension.

Parameters:

  • tensors (sequence of Tensors) – Tensors to concatenate.

  • dim (int) – The dimension along which to concatenate. Must be a negative index to ensure consistency across tensors of different ranks (e.g., -1 for the last dimension).

Returns:

The concatenated tensor.

Return type:

Tensor

caskade.warnings module#

exception caskade.warnings.CaskadeWarning[source]#

Bases: Warning

Base warning class for caskade.

All custom warnings issued by caskade inherit from this class, allowing users to filter or catch any caskade-specific warning.

exception caskade.warnings.InvalidValueWarning(name, value, valid)[source]#

Bases: CaskadeWarning

Warning issued when a parameter value is outside its valid range.

Indicates that the assigned value may cause errors or unexpected behavior during computation.

Parameters:

  • name (str) – Name of the parameter with the out-of-range value.

  • value (ArrayLike) – The value that was assigned.

  • valid (tuple) – A (lower, upper) tuple defining the valid range, where None represents negative or positive infinity.

exception caskade.warnings.SaveStateWarning[source]#

Bases: CaskadeWarning

Warning issued when saving state encounters a non-fatal problem.

Issued when the state serialization completes but with potential data loss or format issues that the user should be aware of.

Module contents#

caskade - Build scientific simulators as directed acyclic graphs.

Caskade provides a framework for constructing modular scientific simulators by composing computational steps into a directed acyclic graph (DAG). It handles parameter management, caching, and context-dependent evaluation.

Main Public API#

Node : Base class for building computational graph nodes. Module : High-level container for assembling simulator components. Param : Declare and manage parameters within nodes. forward : Decorator to define the forward computation of a node. active_cache : Decorator for caching intermediate results. NodeCollection, NodeList, NodeTuple : Collections of nodes. ActiveContext, ValidContext, OverrideParam : Context managers for evaluation. backend : Array backend abstraction (NumPy, PyTorch, etc.). utils : Utility functions.

class caskade.ActiveContext(module: Module)[source]#

Bases: object

Context manager to activate a module for a simulation.

Only inside an ActiveContext is it possible to fill or clear the dynamic and live parameters. On entry, the module is marked as active (or its current parameter state is saved if already active). On exit, the state is restored.

Parameters:

module (Module) – The module to activate for the duration of the context.

Raises:

ActiveStateError – If the module is already running a simulation (module.online is True).

Examples

Activate a module, fill parameters, and run a forward pass:

with ActiveContext(my_module):
    my_module.fill_params(params)
    result = my_module.my_forward(x)
exception caskade.ActiveStateError[source]#

Bases: CaskadeException

Exception for active-state errors in caskade.

Raised when an operation requires a node to be in a particular active state (enabled or disabled) and that condition is not met.

exception caskade.BackendError[source]#

Bases: CaskadeException

Exception for backend-related errors in caskade.

Raised when the selected numerical backend encounters an unsupported operation or configuration issue.

exception caskade.CaskadeException[source]#

Bases: Exception

Base class for all exceptions in caskade.

All custom exceptions raised by caskade inherit from this class, allowing users to catch any caskade-specific error with a single except clause.

exception caskade.CaskadeWarning[source]#

Bases: Warning

Base warning class for caskade.

All custom warnings issued by caskade inherit from this class, allowing users to filter or catch any caskade-specific warning.

exception caskade.FillParamsArrayError(name, input_params, params)[source]#

Bases: FillParamsError

Exception raised when filling parameters with an array fails.

Raised when the shape of the input array does not match the total number of flattened parameters registered on a node.

Parameters:

  • name (str) – Name of the node whose parameters are being filled.

  • input_params (ArrayLike) – The input array that was provided.

  • params (tuple of Param) – Registered parameters whose shapes are compared against the input.

exception caskade.FillParamsError[source]#

Bases: CaskadeException

Base exception for errors when filling parameters in caskade.

Raised when the input data provided to fill node parameters is invalid. Subclasses handle specific input types (array, sequence, mapping).

exception caskade.FillParamsMappingError(name, children, missing_key=None)[source]#

Bases: FillParamsError

Exception raised when filling parameters with a mapping fails.

Raised when a key in the input dictionary does not correspond to any registered child node.

Parameters:

  • name (str) – Name of the node whose parameters are being filled.

  • children (dict) – Dictionary of registered child nodes.

  • missing_key (str, optional) – The key from the input mapping that was not found among the node’s children.

exception caskade.FillParamsSequenceError(name, input_params, dynamic_params)[source]#

Bases: FillParamsError

Exception raised when filling parameters with a sequence fails.

Raised when the length of the input sequence does not match the number of dynamic parameters registered on a node.

Parameters:

  • name (str) – Name of the node whose parameters are being filled.

  • input_params (sequence) – The input sequence (list, tuple, etc.) that was provided.

  • dynamic_params (tuple of Param) – Registered dynamic parameters expected by the node.

exception caskade.GraphError[source]#

Bases: CaskadeException

Exception for graph-related errors in caskade.

Raised when an operation on the computational graph is invalid, such as creating cycles or referencing nonexistent nodes.

exception caskade.InvalidValueWarning(name, value, valid)[source]#

Bases: CaskadeWarning

Warning issued when a parameter value is outside its valid range.

Indicates that the assigned value may cause errors or unexpected behavior during computation.

Parameters:

  • name (str) – Name of the parameter with the out-of-range value.

  • value (ArrayLike) – The value that was assigned.

  • valid (tuple) – A (lower, upper) tuple defining the valid range, where None represents negative or positive infinity.

exception caskade.LinkToAttributeError[source]#

Bases: GraphError

Exception raised when linking to an attribute fails.

Raised when an attempt is made to create a link to a node attribute that does not exist or is not a valid link target.

class caskade.Memo(module: Node, memo: str)[source]#

Bases: object

Sends a “memo” (a small message) to all nodes below the current one in the graph. This can be used to communicate state changes in the graph with all lower nodes. By default, the message will skip any subgraphs (hierarchical graphs) but this can be changed to ensure all nodes hear the message.

Note that memos are stored as a python set, so duplicates will be merged. Depending on your use case, it may be wise to ensure that your memo is unique.

Parameters:

  • module (Module) – The caskade Module object that will propogate the memo

  • memo (str) – The message to send down the graph

  • skip_subgraphs (bool) – If True (default) any subgraphs, otherwise known as hierarchical graphs, will not get the memo.

class caskade.Module(name: str | None = None, **kwargs)[source]#

Bases: Node, GetSetValues

Node to represent a simulation module in the graph.

The Module object is used to represent a simulation module in the graph. These are python objects that contain the calculations for a simulation, they also hold the Param objects that are used in the calculations. The Module object has additional functionality to manage the Param objects below it in the graph, it keeps track of all dynamic Param objects so that at runtime their values may be filled. The Module object manages its links to other nodes through attributes of the class.

Examples

Example of a nested pair of Module objects and how their @forward methods are called:

class MySim(Module):
    def __init__(self, a, b=None):
        super().__init__()
        self.a = a
        self.b = Param("b", b)

    @forward
    def myfunc(self, x, b=None):
        return x * self.a.otherfun(x) + b

class OtherSim(Module):
    def __init__(self, c=None):
        super().__init__()
        self.c = Param("c", c)

    @forward
    def otherfun(self, x, c = None):
        return x + c

othersim = OtherSim()
mysim = MySim(a=othersim)
#                       b                         c
params = [torch.tensor([1.0, 2.0]), torch.tensor([3.0, 4.0])]
result = mysim.myfunc(3.0, params=params)
# result is tensor([19.0, 23.0])
property all_params#

All parameters below this module in the DAG.

Returns:

Concatenation of static, dynamic, and pointer parameters.

Return type:

tuple of Param

clear_state()[source]#

Clear the active state _value for all params below this Module in the DAG. This should not be used by a user under normal circumstances.

property dynamic: bool#

Return True if the module has dynamic parameters as direct children.

Returns:

True if any direct children are dynamic parameters.

Return type:

bool

fill_kwargs(keys: tuple[str]) → dict[str, ArrayLike][source]#

Fill the kwargs for an @forward method with the values of the dynamic parameters. The requested keys are matched to names of Param objects owned by the Module. This should not be used by the user under normal circumstances.

fill_params(params: ArrayLike | Sequence | Mapping, dynamic=True)[source]#

Fill the dynamic/static parameters of the module with the input values from params.

Parameters:

  • params ((Union[ArrayLike, Sequence, Mapping])) – The input values to fill the dynamic parameters with. The input can be an ArrayLike, a Sequence, or a Mapping.

  • dynamic (bool) – Operate on dynamic parameters (True, default) or static parameters (False).

property graphviz_style#
property node_str: str#

Returns a string representation of the node for graph visualization.

param_order()[source]#

Return a human-readable string of dynamic parameter ordering.

Each line corresponds to a parameter group and lists the parameters in the format parent_name: param_name.

Returns:

Multi-line string describing the dynamic parameter order.

Return type:

str

remove_memo(memo)[source]#

Remove a memo string and propagate removal to all children.

Parameters:

memo (str) – The memo message to remove. The same propagation rules as add_memo apply.

property static: bool#

Return True if the module has no dynamic parameters as direct children.

Returns:

True if none of the direct children are dynamic parameters.

Return type:

bool

to_dynamic(children_only=True)[source]#

Change all parameters to dynamic parameters.

Parameters:

children_only ((bool, optional)) – If True, only convert the children of this module to dynamic. If False, convert all parameters in the graph below this module. Defaults to True.

to_static(children_only=True)[source]#

Change all parameters to static parameters.

Parameters:

children_only ((bool, optional)) – If True, only convert children of this module to static. If False, convert all parameters in the graph below this module. Defaults to True.

update_graph()[source]#

Maintain a tuple of dynamic, static, and pointer parameters at all points lower in the DAG.

class caskade.Node(name: str | None = None, link: Node | tuple[Node] | None = None, description: str = '')[source]#

Bases: object

Base graph node class for caskade objects.

The Node object is the base class for all caskade objects. It is used to construct the directed acyclic graph (DAG). The primary function of the Node object is to manage the parent-child relationships between nodes in the graph. There is limited functionality for the Node object, though it implements the base versions of the active state and to / update_graph methods. The active state is used to communicate through the graph that the simulator is currently running. The to method is used to move and/or cast the values of the parameter. The update_graph method is used signal all parents that the graph below them has changed.

Examples

Example making some Node objects and then linking/unlinking them:

n1 = Node()
n2 = Node()
n1.link("subnode", n2) # link n2 as a child of n1, may use any str as the key
n1.unlink("subnode") # alternately n1.unlink(n2) to unlink by object
property active: bool#

True if the node is currently in an active simulation run.

Type:

bool

add_memo(memo)[source]#

Add a memo string and propagate it to all children.

Parameters:

memo (str) – The memo message to add. Children in subgraphs receive the memo with the child name appended (memo|child_name).

append_state(saveto: str | File)[source]#

Append the current state to an existing HDF5 file.

The file must have been previously created by save_state with appendable=True. The graph structure in the file is verified before appending.

Parameters:

saveto (str or File) – Path to an HDF5 file ('.h5' or '.hdf5') or an open HDF5 File object.

Raises:

  • GraphError – If the graph structure no longer matches the file.

  • NotImplementedError – If the file path does not end with a supported extension.

property children: dict[str, Node]#

Mapping of link keys to child nodes.

Type:

dict[str, Node]

graph_dict() → dict[str, dict][source]#

Return a nested dictionary representation of the graph.

Each key is a string of the form "name|node_type" and the value is a dict containing the same structure for that node’s children.

Returns:

Nested dictionary mirroring the DAG hierarchy.

Return type:

dict[str, dict]

graph_print(dag: dict, depth: int = 0, indent: int = 4, result: str = '') → str[source]#

Recursively render a graph dictionary as an indented string.

Parameters:

  • dag (dict[str, dict]) – A nested dictionary as returned by graph_dict.

  • depth (int, optional) – Current indentation depth (used during recursion). Defaults to 0.

  • indent (int, optional) – Number of spaces per indentation level. Defaults to 4.

  • result (str, optional) – Accumulator string (used during recursion). Defaults to "".

Returns:

A human-readable, indented representation of the graph.

Return type:

str

graphviz(saveto: str | None = None) → graphviz.Digraph[source]#

Return a graphviz Digraph representing the DAG below this node.

Parameters:

saveto (str, optional) – If provided, save the rendered graph to this file path. The file extension determines the output format (e.g. '.pdf', '.png'). Defaults to None.

Returns:

The constructed directed-graph object.

Return type:

graphviz.Digraph

property graphviz_style#

Link the current Node object to another Node object as a child in a hierarchical manner. See link for more detail on linking. A hierarchical link will allow batching internally to the simulator.

Parameters:

  • key (str) – The key to link the child node with.

  • child (Node) – The child Node object to link to.

Examples

parent = Node(name="parent")
child = Node(name="child")
parent.hierarchical_link("child", child)

Link the current Node object to another Node object as a child.

Parameters:

  • key ((Union[str, Node])) – The key to link the child node with. This will also become the attribute to access the child node. After linking you will have node.key == child

  • child ((Optional[Node], optional)) – The child Node object to link to. Defaults to None in which case the key is used as the child and the child.name is used as the key.

Examples

Example making some Node objects and then linking/unlinking them, demonstrating multiple ways to link/unlink:

n1 = Node()
n2 = Node()

n1.link("subnode", n2)  # may use any str as the key
n1.unlink("subnode")

# Alternatively, link by object
n1.link(n2)
n1.unlink(n2)
load_state(loadfrom: str | File, index: int = -1, **kwargs)[source]#

Load node state (and children) from an HDF5 file.

Parameters:

  • loadfrom (str or File) – Path to an HDF5 file ('.h5' or '.hdf5') or an open HDF5 File object.

  • index (int, optional) – Sample index to load when the file was saved in appendable mode. Defaults to -1 (last sample).

  • **kwargs – Additional keyword arguments forwarded to h5py.File (e.g. driver).

Raises:

  • GraphError – If the graph structure no longer matches the file.

  • NotImplementedError – If the file path does not end with a supported extension.

property memos: set[str]#

Current set of memo strings held by this node.

Type:

set[str]

property name: str#

The name of this node.

Type:

str

property node_str#
property online: bool#

True if the node is online within a hierarchical sub-graph.

Type:

bool

property parents: set[Node]#

Set of parent nodes that link to this node.

Type:

set[Node]

remove_memo(memo)[source]#

Remove a memo string and propagate removal to all children.

Parameters:

memo (str) – The memo message to remove. The same propagation rules as add_memo apply.

save_state(saveto: str | File, appendable: bool = False)[source]#

Save the state of the node and its children, currently only works for HDF5 file types (.h5 and .hdf5).

The “state” of a node is considered to be the value of its params, however it is also possible to save other attributes of the node by adding them to the Node.saveattrs set. Simply call Node.saveattrs.add(‘attribute’) and then Node.attribute will be saved if possible. The HDF5 file will be created with the same structure as the graph, even if there are multiple paths to the same node. For example if N1 has children N2 and N3, and both N2 and N3 have the child N4, the HDF5 file will reflect this. It will be possible to find the N4 params under both ‘N1/N2/N4’ and ‘N1/N3/N4’ if inspecting the HDF5 file manually. Specifically, if N4 has the param P1 then you could access its value like this:

If the save had been set as appendable, then the value will have an extra dimension for the number of samples, this will always be the first dimension. If appendable was false then the value will simply equal the param value.

Note

You need the optional h5py package installed to use this method.

Parameters:

  • saveto ((Union[str, File])) – The file to save the state to. If a string, it should be the path to an HDF5 file (ending in ‘.h5’ or ‘.hdf5’). If a File object, it should be an open HDF5 file.

  • appendable ((bool, optional)) – Whether to save the state in an appendable format. If True, the values will have an extra dimension for the number of samples. Defaults to False.

property subgraphs: set[Node]#

Subset of children linked hierarchically.

Type:

set[Node]

to(device=None, dtype=None)[source]#

Moves and/or casts the values of the Node to a particular device and/or dtype.

Parameters:

  • device ((optional)) – The device to move the values to. Defaults to None.

  • dtype ((optional)) – The desired data type. Defaults to None.

topological_ordering() → tuple[Node][source]#

Return a topological ordering of the graph below the current node.

Performs a recursive depth-first search with post-order traversal to resolve dependencies. The result starts with this node and proceeds to its descendants in dependency order.

Returns:

All nodes reachable from (and including) this node, ordered so that every parent appears before its children.

Return type:

tuple[Node]

Unlink one or more Node objects from this Node.

Parameters:

key ((str, Node, list, tuple, or None, optional)) – The key, Node object, or collection of keys/nodes to unlink. If a string, the child with that key is unlinked. If a Node object, the matching child is located and unlinked. If a list or tuple, each element is unlinked in turn. If None (the default), all children are unlinked.

Raises:

GraphError – If the graph is currently active.

update_graph()[source]#

Triggers a call to all parents that the graph below them has been updated. The base Node object does nothing with this information, but other node types may use this to update internal state.

class caskade.NodeCollection(name: str | None = None, link: Node | tuple[Node] | None = None, description: str = '')[source]#

Bases: Node, GetSetValues

Base mixin for collections of nodes that track parameters.

Provides shared functionality for traversing, querying, and converting parameters within a graph of nodes. Subclasses such as NodeTuple and NodeList combine this mixin with a standard Python sequence type.

copy()[source]#
deepcopy()[source]#
property dynamic#

Whether any node in this collection has dynamic parameters.

Returns:

True if at least one contained node is dynamic.

Return type:

bool

property dynamic_param_groups: tuple[int]#

Sorted unique group identifiers of all dynamic parameters.

Returns:

Sorted group indices present among the dynamic parameters.

Return type:

tuple of int

property dynamic_params: tuple[Param]#

All dynamic parameters in the graph below this node.

Returns:

Dynamic (non-static, non-pointer) parameters found via topological ordering.

Return type:

tuple of Param

property pointer_params: tuple[Param]#

All pointer parameters in the graph below this node.

Returns:

Parameters that act as pointers to other parameters, found via topological ordering.

Return type:

tuple of Param

property static#

Whether all nodes in this collection are static.

Returns:

True if no contained node is dynamic.

Return type:

bool

property static_params: tuple[Param]#

All static parameters in the graph below this node.

Returns:

Static (non-dynamic, non-pointer) parameters found via topological ordering.

Return type:

tuple of Param

to_dynamic(children_only=True)[source]#

Change all parameters to dynamic parameters.

Parameters:

children_only ((bool, optional)) – If True, only convert the children of this module to dynamic. If False, convert all parameters in the graph below this module. Defaults to True.

to_static(children_only=True)[source]#

Change all parameters to static parameters.

Parameters:

children_only ((bool, optional)) – If True, only convert children of this module. If False, convert all parameters in the graph below this module. Defaults to True.

exception caskade.NodeConfigurationError[source]#

Bases: CaskadeException

Exception for node configuration errors in caskade.

Raised when a node is configured with invalid or incompatible settings.

class caskade.NodeDict(mapping=None, name=None)[source]#

Bases: NodeCollection, dict

Mutable, keyed collection of nodes.

Behaves like a standard dict but also participates in the caskade node graph. All elements must be Node instances. Graph links are automatically updated whenever the dict is modified.

Parameters:

  • mapping (mapping of str to Node, optional) – Nodes to include in the dict. Defaults to an empty dict.

  • name (str, optional) – Human-readable name for this collection of nodes.

clear()[source]#

Remove all nodes from the dict and update graph links.

property dynamic#

Whether any node in this collection has dynamic parameters.

Returns:

True if at least one contained node is dynamic.

Return type:

bool

property graphviz_style#
pop(key, *args)[source]#

Remove and return a node from the dict and update graph links.

popitem()[source]#

Remove and return an arbitrary (key, node) pair from the dict (the last one inserted) and update graph links.

setdefault(key, default)[source]#

If key is in the dictionary, return its value. If not, insert key with a value of default and return default. Update graph links.

update(mapping=None, **kwargs)[source]#

Update the dict with another mapping (i.e. dict) and update graph links.

class caskade.NodeList(iterable=(), name=None)[source]#

Bases: NodeCollection, list

Mutable, ordered collection of nodes.

Behaves like a standard list but also participates in the caskade node graph. All elements must be Node instances. Graph links are automatically updated whenever the list is modified.

Parameters:

  • iterable (iterable of Node, optional) – Nodes to include in the list. Defaults to an empty iterable.

  • name (str, optional) – Human-readable name for this collection node.

append(node)[source]#

Append a node to the list and update graph links.

clear()[source]#

Remove all nodes from the list and update graph links.

extend(iterable)[source]#

Extend the list with nodes from an iterable and update graph links.

property graphviz_style#
insert(index, node)[source]#

Insert a node at the given index and update graph links.

pop(index=-1)[source]#

Remove and return a node at the given index, updating graph links.

remove(value)[source]#

Remove the first occurrence of a node and update graph links.

class caskade.NodeTuple(iterable=None, name=None)[source]#

Bases: NodeCollection, tuple

Immutable, ordered collection of nodes.

Behaves like a standard tuple but also participates in the caskade node graph. All elements must be Node instances and are automatically linked as children upon construction.

Parameters:

  • iterable (iterable of Node, optional) – Nodes to include in the tuple.

  • name (str, optional) – Human-readable name for this collection node.

property graphviz_style#
class caskade.OverrideParam(param: Param, value)[source]#

Bases: object

Context manager to override a parameter value.

Only inside an OverrideParam will the parameter be set to the new value. The original value (and the values of any parent pointer parameters) are saved on entry and restored on exit.

Parameters:

  • param (Param) – The parameter whose value should be temporarily overridden.

  • value (object) – The temporary value to assign to param.

Examples

Override a parameter inside a @forward method so that it uses new_value regardless of what was passed via params:

class MySim(Module):
    def __init__(self):
        super().__init__()
        self.a = Param("a", None)
        self.b = Param("b", None)

    @forward
    def __call__(self, x, a=None, b=None):
        with OverrideParam(self.b, 5.0):
            # b will always be 5.0 here, ignoring params
            return x + a + self.b.value
class caskade.Param(name: str, value: ArrayLike | float | int | None = None, shape: tuple[int, ...] | None = None, cyclic: bool = False, valid: tuple[ArrayLike | float | int | None] | None = None, units: str | None = None, dynamic: bool | None = None, group: int = 0, batch_shape: tuple[int] | None = None, dtype: Any | None = None, device: Any | None = None, **kwargs)[source]#

Bases: Node

Node to represent a parameter in the graph.

The Param object is used to represent a parameter in the graph. During runtime this will represent a value which can be used in various calculations. The Param object can be set to a constant value (static); None meaning the value is to be provided at runtime (dynamic); another Param object meaning it will take on that value at runtime (pointer); or a function of other Param objects to be computed at runtime (also pointer, see user guides). These options allow users to flexibly set the behavior of the simulator.

Examples

Example making some Param objects:

p1 = Param("test", (1.0, 2.0)) # constant value, length 2 vector
p2 =Param("p2", None, (2,2)) # dynamic 2x2 matrix value
p3 = Param("p3", p1) # pointer to another parameter
p4 = Param("p4", lambda p: p.children["other"].value * 2) # arbitrary function of another parameter
p5 = Param("p5", valid=(0.0,2*pi), units="radians", cyclic=True) # parameter with metadata
Parameters:

  • name ((str)) – The name of the parameter.

  • value ((Optional[Union[ArrayLike, float, int]], optional)) – The value of the parameter. Defaults to None meaning dynamic.

  • shape ((Optional[tuple[int, ...]], optional)) – The shape of the parameter. Defaults to () meaning scalar.

  • cyclic ((bool, optional)) – Whether the parameter is cyclic, imposing periodic boundary conditions. Such as a rotation from 0 to 2pi. Defaults to False.

  • valid ((Optional[tuple[Union[ArrayLike, float, int, None]]], optional)) – The valid range of the parameter. Defaults to None meaning all of -inf to inf is valid.

  • units ((Optional[str], optional)) – The units of the parameter. Defaults to None.

  • dynamic ((bool, optional)) – Force param to be dynamic if True. If a value is provided and param is dynamic then it has a default value at call time.

  • (bool (batched) – If True, the param is assumed batched and the shape may now take the form (*B, *D) where *D is the shape of the value.

  • optional) – If True, the param is assumed batched and the shape may now take the form (*B, *D) where *D is the shape of the value.

  • dtype ((Optional[Any], optional)) – The data type of the parameter. Defaults to None meaning the data type will be inferred from the value.

  • device ((Optional[Any], optional)) – The device of the parameter. Defaults to None meaning the device will be inferred from the value.

property batch_shape: tuple[int, ...]#

The batch dimensions of the parameter value.

Batch dimensions are the leading dimensions of the value that precede the event shape. If an explicit batch shape was set it is returned directly; otherwise it is inferred from the value.

Returns:

The batch shape, or () if the parameter is not batched.

Return type:

tuple of int

property batched: bool#

Whether this parameter carries batch dimensions.

Returns:

True if batch_shape is non-empty.

Return type:

bool

property cyclic: bool#

Whether the parameter has cyclic (periodic) boundary conditions.

When True, values wrap around the valid range (e.g. an angle from 0 to 2π).

Returns:

True if the parameter is cyclic.

Return type:

bool

property device: str | None#

The device on which the parameter value resides.

If no explicit device was set, the device is inferred from the current value.

Returns:

The device, or None if unknown.

Return type:

device or None

property dtype: str | None#

The data type of the parameter value.

If no explicit dtype was set, the dtype is inferred from the current value.

Returns:

The data type, or None if unknown.

Return type:

dtype or None

property dynamic: bool#

Whether this parameter is dynamic.

Returns:

True if the parameter’s value is provided at runtime.

Return type:

bool

property graphviz_style#
property group: int#

The group index of this parameter.

Parameters that share the same group index are collected together into a single params object when calling a simulator’s @forward method, as well as when using get_values or set_values.

Returns:

The group index (default 0).

Return type:

int

is_valid(value=None) → bool[source]#

Check whether a value lies within the allowed range.

Parameters:

value (ArrayLike or None, optional) – The value to check. If None (default), the parameter’s current value is used.

Returns:

True if the value is within the valid range or if no constraints are set. False otherwise; a warning is also emitted.

Return type:

bool

property node_str: str#

Returns a string representation of the node for graph visualization.

property node_type#

The current type of this parameter node.

Returns:

One of "static", "dynamic", or "pointer".

Return type:

str

property npvalue: ndarray#

The current value converted to a NumPy array.

Returns:

The value as a NumPy ndarray.

Return type:

numpy.ndarray

property pointer: bool#

Whether this parameter is a pointer.

Returns:

True if the parameter points to another Param or a callable that is evaluated at runtime.

Return type:

bool

property shape: tuple[int, ...]#

The event (non-batch) shape of the parameter value.

Wildcard dimensions (None) in the declared shape are resolved using the current value. If no shape was declared, the shape of the current value is returned directly.

Returns:

The resolved shape of the parameter.

Return type:

tuple of int

property static: bool#

Whether this parameter is static.

Returns:

True if the parameter holds a fixed value that does not change at runtime.

Return type:

bool

to(device=None, dtype=None) → Param[source]#

Moves and/or casts the values of the parameter.

Parameters:

  • device ((optional)) – The device to move the values to. Defaults to None.

  • dtype ((optional)) – The desired data type. Defaults to None.

to_dynamic(value=<object object>)[source]#

Change this parameter to a dynamic parameter.

If a value is provided, it is stored as the default dynamic value. When called without arguments the existing value (if any) is kept.

Parameters:

value (ArrayLike, float, int, None, or sentinel, optional) – The default value for the dynamic parameter. Must not be a Param or callable. By default the current value is retained.

Raises:
to_pointer(value, link=())[source]#

Change this parameter to a pointer parameter.

The parameter’s value will be computed at runtime by dereferencing another Param or by calling a user-supplied function.

Parameters:

  • value (Param or callable) – A Param whose value will be mirrored, or a callable f(param) -> ArrayLike evaluated at runtime.

  • link (Node or tuple of Node, optional) – Additional nodes to link into the graph when creating the pointer. Defaults to an empty tuple.

Raises:
to_static(value=<object object>)[source]#

Change this parameter to a static parameter.

If a value is provided, it is stored as the fixed static value. When called without arguments the existing value (if any) is kept.

Parameters:

value (ArrayLike, float, int, None, or sentinel, optional) – The constant value for the static parameter. Must not be a Param or callable. By default the current value is retained.

Raises:
property valid: tuple[ArrayLike | None, ArrayLike | None]#

The valid range of the parameter value.

Returns:

(lower_bound, upper_bound). Either bound may be None indicating no constraint on that side.

Return type:

tuple of (ArrayLike or None, ArrayLike or None)

property value: ArrayLike | None#

The current value of the parameter.

For static and dynamic parameters the stored value is returned. For pointer parameters the linked callable is evaluated. During an active simulation the result is cached.

Returns:

The parameter value, or None if no value has been set.

Return type:

ArrayLike or None

exception caskade.ParamConfigurationError[source]#

Bases: NodeConfigurationError

Exception for parameter configuration errors in caskade.

Raised when a parameter is defined with an invalid shape, type, or constraint.

exception caskade.ParamTypeError[source]#

Bases: CaskadeException

Exception for parameter type errors in caskade.

Raised when a value assigned to a parameter does not match its expected type.

exception caskade.SaveStateWarning[source]#

Bases: CaskadeWarning

Warning issued when saving state encounters a non-fatal problem.

Issued when the state serialization completes but with potential data loss or format issues that the user should be aware of.

class caskade.ValidContext(module: Module)[source]#

Bases: object

Context manager that transforms parameter values to an unconstrained space.

Inside a ValidContext, all parameter values are automatically mapped into the range (-inf, inf) via each parameter’s to_valid / from_valid transformations. This is useful when interfacing with samplers or optimizers that expect unconstrained parameters—any value they propose will be mapped back into the parameter’s original valid range on exit.

Parameters:

module (Module) – The module whose parameters should be transformed.

Examples

Get unconstrained parameter values for use with an optimizer:

with ValidContext(my_module):
    unconstrained_params = my_module.get_values()
    # unconstrained_params live in (-inf, inf)
class caskade.active_cache(func)[source]#

Bases: object

Caches the first evaluated result of a Module method for the duration of a simulation.

This decorator ensures that an expensive method is executed exactly once per active simulation run. Once calculated, subsequent calls to the decorated method will return the stored value, ignoring any arguments passed to it.

Warning

If the method is called multiple times with different arguments in one simulation, the cached result will still be returned, which may lead to unexpected behavior. Use with caution!

Notes

If you are stacking multiple decorators on a method (such as @forward or @jax.jit), @active_cache MUST be the outermost (top) decorator.

Examples

class FluxModel(Module):
    def __init__(self, nodes, x, M):
        super().__init__()
        self.nodes = nodes
        self.x = Param("x", x)
        self.M = Param("M", M)

    @active_cache
    @jax.jit  # Notice active_cache is placed at the top
    @forward
    def compute_intrinsic_sed(self, w, x, M):
        print("Computing SED...")
        return jnp.interp(w, self.nodes, x * M)

    @forward
    def compute_flux(self, wavelengths):
        sed = self.compute_intrinsic_sed(wavelengths)  # Cached after first call
        flux = jnp.sum(sed)
        sed = self.compute_intrinsic_sed(wavelengths)  # Returns cached result, no print
        peak = jnp.max(sed)
        return flux, peak

model = FluxModel(np.linspace(400, 700, 10), x=1.0, M=np.random.rand(10))

# Compute flux only calls compute_intrinsic_sed once due to caching
flux, peak = model.compute_flux(wavelengths)
caskade.forward(method)[source]#

Decorator to define a forward method for a module.

Manages parameter passing and activation for the decorated method. When called, it automatically fills keyword arguments from the module’s Param children and handles parameter overrides and active context.

Parameters:

method ((Callable)) – The forward method to be decorated.

Returns:

The decorated forward method.

Return type:

Callable

Examples

Standard usage of the forward decorator:

class ExampleSim(Module):
    def __init__(self, a, b, c):
        super().__init__("example_sim")
        self.a = a
        self.b = Param("b", b)
        self.c = Param("c", c)

    @forward
    def example_func(self, x, b=None):
        return x + self.a + b

E = ExampleSim(a=1, b=None, c=3)
print(E.example_func(4, params=[5]))
# Output: 10
caskade.test()[source]#

Run a basic integration test to verify caskade is installed and working.

Exercises core functionality including Module and Param creation, parameter linking, and forward method execution.

Examples

import caskade
caskade.test()
# Output: Success!
Contents