netket.utils.walk_tree_with_path

netket.utils.walk_tree_with_path(tree, root, *, visit_leaf, enter_node=None, expand_node=None, path_fn=<function _append_path>, **kwargs)

Traverse a nested tree while threading a path to each leaf.

Parameters:

• tree (Any) – the tree to traverse.

• root (Any) – the current path accumulator.

• visit_leaf (Callable[..., Any]) – function applied to leaves as visit_leaf(root, leaf, **kwargs).

• enter_node (Callable[..., dict[str, Any] | None] | None) – optional hook run before descending into a node. It can return a dict of keyword arguments to override for that subtree. The expected signature is enter_node(root, node, **kwargs) -> dict[str, Any] | None.

• expand_node (Callable[..., Iterable[tuple[Any, Any]] | None] | None) – optional hook returning custom children for a node as (key, child) pairs. Return None to use the default traversal. The expected signature is expand_node(root, node, **kwargs) -> Iterable[tuple[key, child]] | None.

• path_fn (Callable[[Any, Any], Any]) – function used to build child paths from (root, key).

• **kwargs – additional keyword arguments forwarded to hooks and leaves.

Return type:

Any

Returns:

The original root after the traversal completes.

Example

Collect all leaves into a list of (path, value) pairs:

from netket.utils.tree_walk import walk_tree_with_path

def visit_leaf(path, leaf, *, out, prefix=""):
    full_path = f"{prefix}/{path}" if prefix else path
    out.append((full_path, leaf))

def enter_node(path, node, *, prefix=""):
    if path == "":
        return {"prefix": "data"}
    return None

def path_fn(path, key):
    return f"{path}/{key}" if path else str(key)

out = []
walk_tree_with_path(
    {"energy": [1.0, 2.0], "stats": {"mean": 0.5}},
    "",
    visit_leaf=visit_leaf,
    enter_node=enter_node,
    path_fn=path_fn,
    out=out,
)

In practice:

• visit_leaf does the work at each leaf.

• enter_node can inject or override state for a whole subtree.

• expand_node is only needed when a node should expose custom children, such as splitting a complex number into (“re”, value.real) and (“im”, value.imag).