Tree utilities#

Tree rewriting accepts a single pybamm.Symbol root. To rewrite several expressions with the same mapping, iterate explicitly and share a cache:

cache = {}
equations = {
    variable: pybamm.replace(equation, mapping, cache=cache)
    for variable, equation in model.rhs.items()
}

The cache belongs to this mapping; use a fresh cache for a different mapping. Model parameter processing remains available through parameter_values.process_model(model).

class pybamm.TreeDef(node_type, n_children, children_container, leaf_spec, state)[source]#

The static structure of a symbol: everything needed to rebuild it from its leaves. Returned by tree_flatten() and consumed by tree_unflatten().

Parameters:

  • node_type (type) – The symbol class.

  • n_children (int) – How many leading leaves are the node’s children.

  • children_container (type) – list or tuple, whichever the node stores its children in.

  • leaf_spec (tuple) – (field, kind, count, container) per extra leaf field, in order.

  • state (dict) – The node’s remaining (non-leaf, non-derived) state.

pybamm.tree_flatten(symbol: Symbol) → tuple[tuple[Symbol, ...], TreeDef][source]#

Split a symbol into its leaves (children and the other symbols it depends on) and a TreeDef describing how to put it back together.

pybamm.tree_unflatten(treedef: TreeDef, leaves: Sequence[Symbol]) → Symbol[source]#

Rebuild a symbol from a TreeDef and its leaves, without simplification or validation. The result has no id yet.

pybamm.rebuild(node: Symbol, new_leaves: Sequence[Symbol], simplify: bool = True) → Symbol[source]#

Return node with its leaves swapped for new_leaves.

node itself is returned when no leaf changed, so untouched subtrees are shared rather than copied. When only children changed the node’s constructor is used (so derived names, domains and, if simplify, simplifications are applied); otherwise the structure is rebuilt as-is with tree_unflatten().

pybamm.tree_map(fn: Callable[[Symbol, tuple[Symbol, ...]], Symbol], tree: Symbol, cache: MutableMapping[Symbol, Symbol] | None = None, is_leaf: Callable[[Symbol], bool] | None = None) → Symbol[source]#

Rewrite an expression tree bottom-up, out of place.

Parameters:

  • fn (callable) – fn(node, new_leaves) returns the node to use in place of node, given the already-rewritten leaves. Use rebuild() for the default “same node with new leaves” behaviour.

  • tree (pybamm.Symbol) – The root of the expression tree to rewrite.

  • cache (dict, optional) – Memo of {original node: rewritten node}. Shared nodes are rewritten once; pass the same dict across calls to share work between trees.

  • is_leaf (callable, optional) – is_leaf(node) returning True stops the descent at node: fn is then called as fn(node, ()) with nothing below it rewritten.

Return type:

The rewritten symbol. tree is not modified.

pybamm.replace(tree: Symbol, mapping: Mapping[Symbol, Any], cache: MutableMapping[Symbol, Symbol] | None = None, simplify: bool = True) → Symbol[source]#

Substitute symbols throughout an expression tree, out of place.

Every node equal (by id) to a key of mapping is replaced by the corresponding value, which is not itself traversed. All other nodes are rebuilt with rebuild() only if something below them changed, so the result shares every untouched subtree with tree.

Parameters:

  • tree (pybamm.Symbol) – The root of the expression tree to substitute into.

  • mapping (dict) – {symbol: replacement}; replacements may be numbers.

  • cache (dict, optional) – Memo shared between calls made with the same mapping.

  • simplify (bool, optional) – Whether rebuilt operators apply their usual simplifications (default True).

Migrating in-place updates#

Public symbol setters and domain/list edits remain available but emit SymbolMutationDeprecationWarning. Prefer constructing replacements:

updated = variable.create_copy(scale=2, reference=0, bounds=(0, 1))
expression = pybamm.replace(expression, {variable: updated})
symbol = symbol.with_domains({"primary": ["negative electrode"]})
symbol = symbol.without_domains()
symbol = symbol.with_mesh(mesh)

For names, scalar values and function inputs, construct a new symbol of the appropriate class. IDs update automatically; callers need not call set_id.

Deprecated updates invalidate cached IDs and derived values across live symbols. This deliberately slow compatibility path adds no per-symbol tracking to normal replacement. Public domain and child-list views support deprecated edits while internal storage remains shared. Mutating symbols already used as dictionary keys cannot repair those dictionaries: use a fresh replacement memo and rebuild any parameter-processed or discretised model after legacy edits. Changes to a child’s domains do not propagate into constructor-derived parent domains; reconstruct those expressions as well. Legacy mutation is rejected while gc.freeze() is active because frozen ancestors cannot be discovered for cache invalidation; use out-of-place replacement or unfreeze the collector first.

The test suite forbids deprecated mutation even with warnings ignored. Only isolated compatibility subprocesses exercise the legacy path; notebook kernels inherit the same guard as other tests.