Utility Functions

PyOCN includes several utility functions for working with channel networks, optimization schedules, and network analysis.

Network Generation

PyOCN.utils.net_type_to_dag(net_type: Literal['I', 'H', 'V', 'E'], dims: tuple, pbar: bool = False) → DiGraph[source]

Create a predefined OCN initialization network as a NetworkX DiGraph.

Parameters:

  • net_type ({"I", "H", "V", "E"}) –

    The type of network to create. Descriptions of allowed types:

    • ”I”:

      O--O--O--O--O
      |
O--O--O--O--O
      |
O--O--O--O--O
      |
O--O--X--O--O

    • ”V”:

      O  O  O  O  O
 \  \ | /  /
O  O  O  O  O
 \  \ | /  /
O  O  O  O  O
 \  \ | /  /
O--O--X--O--O

    • ”H”:

      O  O  O  O
|  |  | /
O  O  O--O
|  | /
O  O--O--O
| /
X--O--O--O

    • ”E”: A network where every node on the edge of the grid is a root. Initial flow moves away from center towards edges.

  • dims (tuple) – The network dimensions as (rows, cols). Both must be positive even integers.

  • pbar (bool, default False) – If True, display a progress bar while constructing the graph.

Returns:

A directed acyclic graph representing a valid initial OCN configuration.

Return type:

networkx.DiGraph

Raises:

ValueError – If net_type is invalid or dims are not two positive even integers.

Note

The returned graph assigns each grid cell exactly one node with a pos attribute equal to (row, col).

Optimization

PyOCN.utils.simulated_annealing_schedule(dims: tuple[int, int], E0: float, constant_phase: float, n_iterations: int, cooling_rate: float) → Callable[[int], float | ndarray][source]

Create a simulated-annealing cooling schedule for OCN optimization.

This returns a callable schedule(i) that returns the temperature at iteration i. The schedule consists of a constant-temperature phase followed by an exponentially decaying phase.

Parameters:

  • dims (tuple[int, int]) – The dimensions of the grid as (rows, cols).

  • E0 (float) – Initial energy value.

  • constant_phase (float) – Fraction of iterations (0 <= fraction <= 1) during which the temperature remains constant at Energy[0].

  • n_iterations (int) – Total number of optimization iterations.

  • cooling_rate (float) – Positive decay rate controlling the exponential temperature decrease after the constant phase.

Returns:

A function mapping an iteration index i to a temperature value. If vectorized evaluation is used, may return a NumPy array of temperatures.

Return type:

Callable[[int], Union[float, np.ndarray]]

Note

The exponential phase follows the form

\[T_i = E_0 \exp\left(-\frac{r\cdot(i - n_0)}{N}\right),\]

where E0 is the initial energy, n0 is the number of iterations in the constant phase, r is the cooling rate,and N = rows * cols.

Example usage:

import PyOCN as po

ocn = po.OCN.from_net_type("I", dims=(64, 64))
n_iterations = 50000
# Create a custom cooling schedule
schedule = po.utils.simulated_annealing_schedule(
    dims=(64, 64),
    E0=ocn.energy,
    constant_phase=0.1,
    n_iterations=n_iterations,
    cooling_rate=1.0
)

# Use in optimization
ocn.fit_custom_cooling(cooling_func=schedule, n_iterations=n_iterations)

Network Analysis

PyOCN.utils.unwrap_digraph(dag: DiGraph, dims: tuple[int, int]) → DiGraph[source]

“unwrap” gridcell coordinate attributes in a directed acyclic graph to place connected nodes adjacent to each other, removing periodic boundary conditions.

Parameters:

  • dag (nx.DiGraph) – The input directed acyclic graph with periodic boundary conditions. Each node must have a ‘pos’ attribute indicating its (row, col) position.

  • dims (tuple[int, int]) – The dimensions of the grid as (rows, cols). Both must be positive integers.

Returns:

A new directed acyclic graph with unwrapped grid coordinates. May not be consistent with a grid structure.

Return type:

nx.DiGraph

Raises:

ValueError – If any node in the input graph lacks a ‘pos’ attribute or if the dimensions are not positive integers.

Important

The function assumes that the input graph is a valid DAG and that the ‘pos’ attributes are correctly assigned. The output graph will no longer span a toroidal topology and will no longer cover a dense grid of nodes.

PyOCN.utils.assign_subwatersheds(dag: DiGraph) → None[source]

Assign a ‘watershed_id’ attribute to each node in the DAG. The resulting watershed_ids will be of the highest order possible, meaning that ids are assigned based on watersheds that drain directly into the root nodes of the graph. To assign ids to lower order watersheds, consider first partitioning the graph into smaller subgraphs using the get_subwatersheds function.

Parameters:

dag (nx.DiGraph) – The input directed acyclic graph. Each node must have a ‘pos’ attribute indicating its (row, col) position.

Returns:

The function modifies the input graph in place by adding a ‘watershed_id’ attribute to each node.

Return type:

None

Raises:

ValueError – If any node in the input graph lacks a ‘pos’ attribute.

Note

A subwatershed is defined as the set of nodes that drain to a common outlet, where an outlet is a node with out-degree zero. Each subwatershed is assigned a unique integer ID, starting from 0. Nodes that are outlets themselves are assigned a watershed ID of -1.

PyOCN.utils.get_subwatersheds(dag: DiGraph, node: Any) → set[DiGraph][source]

Extract subwatershed subgraphs from the main DAG. Each subwatershed drains to a common outlet node node. Node node is not included in the returned subwatershed graphs.

Parameters:

  • dag (nx.DiGraph) – The input directed acyclic graph. Each node must have a ‘pos’ attribute indicating its (row, col) position.

  • node (Any) – A node in the graph representing the outlet of a subwatershed.

Returns:

A set of directed acyclic graphs, each representing a subwatershed.

Return type:

set of nx.DiGraph

Danger

The returned subwatersheds are subgraph views of the input graph and share node and edge data with the original graph. Unless copied, any changes to node or edge attributes in the subwatersheds will affect the original graph.

Example watershed analysis:

import PyOCN as po

# Create and optimize network
ocn = po.OCN.from_net_type("I", dims=(32, 32))
ocn.fit(n_iterations=10000)

# Get the network as a graph
dag = ocn.to_dag()

# Analyze subwatersheds
subwatersheds = po.utils.get_subwatersheds(dag, node=5)
print(f"Found {len(subwatersheds)} outlet nodes")