OCN Class

The OCN class is the main interface for creating and optimizing channel networks.

Basic Usage

import PyOCN as po

# Create from a network type
ocn = po.OCN.from_net_type("I", dims=(64, 64), gamma=0.5)

# Optimize the network
ocn.fit()

# Access properties
print(f"Final energy: {ocn.energy}")

Class Reference

class PyOCN.OCN(dag: DiGraph, resolution: float = 1.0, gamma: float = 0.5, random_state=None, verbosity: int = 0, validate: bool = True, wrap: bool = False, vertical_exaggeration: float = 1.0)[source]

Bases: object

The main class for interacting with Optimized Channel Networks.

Use OCN.from_net_type() or OCN.from_digraph() to construct an instance.

from_net_type()[source]

Create an OCN from a predefined network type and dimensions.

from_digraph()[source]

Create an OCN from an existing NetworkX DiGraph.

to_digraph()[source]

Export the current grid to a NetworkX DiGraph.

to_numpy()[source]

Export raster arrays (energy, drained area, elevation) as numpy arrays.

to_xarray()[source]

Export raster arrays as an xarray Dataset (requires xarray).

to_gtiff()[source]

Export raster arrays to a GeoTIFF file (requires rasterio).

copy()[source]

Create a deep copy of the OCN.

single_erosion_event()

Perform a single erosion event at a given temperature.

fit()[source]

Optimize the network using the simulated annealing method from Carraro et al (2020).

fit_custom_cooling()[source]

Optimize the network using a custom cooling function.

compute_energy()[source]

Compute the current energy of the network.

copy()[source]

Create a deep copy of the OCN.

energy

Current energy of the network (read-only property).

Type:

float

dims

Grid dimensions (rows, cols) of the FlowGrid (read-only property).

Type:

tuple[int, int]

resolution

The side length of each grid cell (read-only property).

Type:

float

nroots

Number of root nodes in the current OCN grid (read-only property).

Type:

int

gamma

Exponent in the energy model.

Type:

float

verbosity

Verbosity level (0-2). 0 is silent, 1 prints high-level progress messages, 2 prints detailed logging messages. Can be overridden in methods that support it.

Type:

int

wrap

If true, enables periodic boundary conditions on the FlowGrid (read-only property).

Type:

bool

history

numpy array of shape (n_iterations, 3) recording the iteration index, energy, and temperature at each iteration during optimization. Updated each time an optimization method is called.

Type:

np.ndarray

rng

the current random state of the internal RNG

Type:

int

vertical_exaggeration

Vertical exaggeration factor for the network. Used when computing elevation.

Type:

float

Examples

The following is a simple example of creating, optimizing, and plotting an OCN using PyOCN and Matplotlib. More examples are available in the demo.ipynb notebook in the repository (https://github.com/alextsfox/PyOCN).

>>> # Fit an OCN from a "Hip-roof" initial network shape and periodic boundary conditions
>>> import matplotlib.pyplot as plt
>>> import matplotlib as mpl
>>> import PyOCN as po
>>> ocn = po.OCN.from_net_type("H", dims=(64, 64), wrap=True, random_state=8472, verbosity=0)
>>> ocn.fit()
>>> po.plot_ocn_raster(ocn, norm=mpl.colors.PowerNorm(gamma=0.5), attribute='drained_area')
>>> plt.show()
__init__(dag: DiGraph, resolution: float = 1.0, gamma: float = 0.5, random_state=None, verbosity: int = 0, validate: bool = True, wrap: bool = False, vertical_exaggeration: float = 1.0)[source]

Construct an OCN from a valid NetworkX DiGraph.

Attention

Please use the classmethods OCN.from_net_type() or OCN.from_digraph() to instantiate an OCN.

compute_energy() → float[source]

Compute the current energy of the network.

Returns:

The computed energy value.

Return type:

float

copy() → OCN[source]

Create a deep copy of the OCN, including the underlying FlowGrid_C. Also copies the current RNG state. The new copy and the original will be independent from each other and behave identically statistically.

property dims: tuple[int, int]
property energy: float
fit(cooling_rate: float = 1.0, constant_phase: float = 0.0, n_iterations: int = None, pbar: bool = False, array_reports: int = 0, tol: float | None = None, max_iterations_per_loop=10000, unwrap: bool = True, calculate_full_energy: bool = False) → xr.Dataset | None[source]

Convenience function to optimize the OCN using the simulated annealing algorithm from Carraro et al (2020). For finer control over the optimization process, use fit_custom_cooling() or use single_erosion_event() in a loop.

This performs n_iterations erosion events, accepting or rejecting proposals according to a temperature schedule defined by the annealing algorithm. A proposal consists of changing the outflow direction of a randomly selected vertex. The new outflow direction is chosen uniformly from the valid neighbors. A proposal is valid if it maintains a well-formed graph structure.

The history attribute is updated in-place after optimization finishes.

Parameters:

  • cooling_rate (float, default 1.0) – Cooling rate parameter in the annealing algorithm. Typical range is 0.5-1.5. Higher values result in faster cooling and a greedier search. Lower values result in slower cooling and more thorough exploration of the solution space, but slower convergence and lower stability.

  • constant_phase (float, default 0.0) – Amount of time to hold temeprature constant at the start of the optimization. This is a fraction of n_iterations, and must be in the range [0, 1]. A value of 0.0 (default) means the temperature starts cooling immediately from the initial temperature. A value of 1.0 means the temperature is held constant for the entire optimization.

  • n_iterations (int, optional) – defaults to int((63 * exp(-0.448 * cooling_rate)) * rows * cols * (1 + constant_phase)), which was empirically found to work well across a range of cooling rates. When constant_phase = 1 and constant_phase = 0, this reduces to 40 * rows * cols. Clamped to at least energy_reports * 10 (this should only matter for extremely small grids, where rows * cols < 256).

  • pbar (bool, default True) – Whether to display a progress bar.

  • array_reports (int, default 0) – Number of timepoints (approximately) at which to save the state of the FlowGrid. If 0 (default), returns None. If >0, returns an xarray.Dataset containing the state of the FlowGrid at approximately evenly spaced intervals throughout the optimization, including the initial and final states. Requires xarray to be installed. See notes on xarray output for details.

  • tol (float | None, optional) – If provided, optimization will stop early if the average relative energy reduction per iteration is less than tol for two consecutive checks. Must be positive. If None, no early stopping is performed. Recommended values are in the range 1e-9 to 1e-6 per iteration. A good default is 3.2e-8.

  • max_iterations_per_loop (int, optional) – If provided, the number of iterations steps to perform in each “chunk” of optimization. Energy and output arrays can be reported no more often than this. Recommended values are 1_000-1_000_000. Default is 10_000.

  • unwrap (bool, default True) – If True and the current OCN has periodic boundaries, the resulting rasters will be transformed so connected grid cells are adjacent in the output. This will result in a larger raster with some nan values. If False or the current OCN does not have periodic boundaries, then no transformation is applied and the resulting raster will have the same dimensions as the current OCN grid.

  • calculate_full_energy (bool, default False) – If True, the full energy of the graph is recalculated when considering the proposed change. If False, a more efficient incremental update is used. Small numerical differences may arise between the two methods due to floating point precision. If precision is of the utmost importance, set this to True, but note that this comes with a significant performance penalty.

Returns:

ds – If array_reports > 0, an xarray.Dataset containing the state of the FlowGrid at approximately evenly spaced intervals throughout the optimization, including the initial and final states. See notes on xarray output for details. If array_reports == 0, returns None.

Return type:

xr.Dataset | None

Note

The returned xarray.Dataset will have coordinates:

  • y (float) representing the northing coordinate of each row

  • x (float) representing the easting coordinate of each column

  • iteration (int) representing the iteration index at which the data was recorded

and data variables:

  • energy (np.float64) representing energy at each grid cell

  • area (np.float64) representing drained area at each grid cell

  • elevation (np.float64) representing elevation at each grid cell

If the OCN has a periodic boundary condition, the following changes apply:

  • The (0,0) coordinate will be set to the position of the “main” root node, defined as the root node with the smallest row*cols + col value

  • The rasters will be unwrapped to a non-periodic representation, which may result in larger rasters.

  • The size of the final rasters are the maximum extent of the unwrapped grid, taken across all iterations.

Generating reports requires additional memory and computation time.

Note

At iteration i, the outflow of a random grid cell if proposed to be rerouted. The proposal is accepted with the probability

\[P(\text{accept}) = e^{-\Delta E / T},\]

where \(\Delta E\) is the change in energy the change would cause and \(T\) is the temperature of the network.

The total energy of the system is computed from the drained areas of each grid cell \(k\) as

\[E = \sum_k A_k^\gamma\]

The temperature of the network is governed by a cooling schedule, which is a function of iteration index.

Note that when \(\Delta E < 0\), the move is always accepted.

The cooling schedule used by this method is a piecewise function of iteration index:

\[\begin{split}T(i) = \begin{cases} E_0 & i < C N \\ E_0 \cdot e^{i - C N} & i \ge C N \end{cases}\end{split}\]

where \(E_0\) is the initial energy, \(N\) is the total number of iterations, and \(C\) is constant_phase. Decreasing-energy moves (\(\Delta E < 0\)) are always accepted.

Alternative cooling schedules can be implemented using fit_custom_cooling().

fit_custom_cooling(cooling_func: Callable[[ndarray], ndarray], n_iterations: int = None, iteration_start: int = 0, pbar: bool = False, array_reports: int = 0, tol: float | None = None, max_iterations_per_loop=10000, unwrap: bool = True, calculate_full_energy: bool = False) → xr.Dataset | None[source]

Optimize the OCN using the a custom cooling schedule. This allows for multi-stage optimizations or other custom cooling schedules not covered by the default simulated annealing schedule from Carraro et al (2020).

See fit() for additional details on the optimization algorithm and parameters.

Parameters:

  • cooling_func (Callable[[np.ndarray], np.ndarray]) – A function that takes an array of iteration indices and returns an array of temperatures. This function defines the cooling schedule for the optimization. Note that the function should return temperatures that are appropriate for the current energy of the OCN.

  • n_iterations (int, optional)

  • iteration_start (int, default 0) – The starting iteration index. This is useful for continuing an optimization from a previous run. Must be a non-negative integer. If provided, n_iterations is the number of additional iterations to perform. The iteration number passed to cooling_func will be iteration_start + i where i is the current iteration index in the range [0, n_iterations-1].

  • pbar (bool, default True)

  • array_reports (int, default 0)

  • tol (float, optional)

  • max_iterations_per_loop (int, optional)

  • unwrap (bool, default True)

  • calculate_full_energy (bool, default False)

Return type:

xr.Dataset | None

classmethod from_digraph(dag: DiGraph, resolution: float = 1, gamma=0.5, random_state=None, verbosity: int = 0, wrap: bool = False, vertical_exaggeration: float = 1.0)[source]

Create an OCN from an existing NetworkX DiGraph.

Parameters:

  • dag (nx.DiGraph) – Directed acyclic graph (DAG) representing the stream network.

  • resolution (int, optional) – The side length of each grid cell.

  • gamma (float, default 0.5) – Exponent in the energy model.

  • random_state (int | numpy.random.Generator | None, optional) – Seed or generator for RNG seeding.

  • verbosity (int) – Verbosity level (0-2). 0 is silent, 1 prints high-level progress messages, 2 prints detailed logging messages. Can be overridden in methods that support it.

  • wrap (bool, default False) – If true, allows wrapping around the edges of the grid (toroidal). If false, no wrapping is applied.

  • vertical_exaggeration (float, default 1.0) – Vertical exaggeration factor for the network. Used when computing elevation.

Returns:

A newly constructed instance encapsulating the provided graph.

Return type:

OCN

Important

The input graph must satisfy all of the following:

  • It is a directed acyclic graph (DAG).

  • Each node has attribute pos=(row:int, col:int) specifying its grid position with non-negative coordinates. Any other attributes are ignored.

  • The graph can be partitioned into one or more spanning trees over a dense grid of shape (m, n): each grid cell corresponds to exactly one node; each non-root node has out_degree == 1; the roots have out_degree == 0.

  • Edges connect only to one of the 8 neighbors (cardinal or diagonal), i.e., no jumps over rows or columns. If wrap=True, edges may connect across the grid boundaries (i.e. row 0 can connect to row m-1 and col 0 can connect to col n-1).

  • Edges do not cross in the row-column plane.

  • Both m and n are positive integers, and there are at least four vertices.

Examples

>>> # creating a "zig-zag" network
>>> # O O O O
>>> # |/|/|/|
>>> # O O O X
>>> import networkx as nx
>>> import PyOCN as po
>>> import matplotlib as mpl
>>> dag = nx.DiGraph()
>>> for i in range(8):
...     dag.add_node(i, pos=(i % 4, i // 4))
...     if i < 7: dag.add_edge(i, i + 1)
>>> ocn = OCN.from_digraph(dag, random_state=8472)
>>> ocn.fit()
>>> po.plot_ocn_raster(ocn, norm=mpl.colors.PowerNorm(gamma=0.5), attribute='drained_area')
>>> plt.show()
classmethod from_net_type(net_type: str, dims: tuple[int, int], resolution: float = 1, gamma: float = 0.5, random_state=None, verbosity: int = 0, wrap: bool = False, vertical_exaggeration: float = 1.0)[source]

Create an OCN from a predefined network type and dimensions.

Parameters:

  • net_type (str) – Predefined network type to instantiate from. Valid types are “H”, “I”, “E”, and “V”. See net_type_to_dag() for more information.

  • dims (tuple[int, int]) – Grid dimensions (rows, cols). Both must be positive even integers.

  • resolution (int, optional) – The side length of each grid cell.

  • gamma (float, default 0.5) – Exponent in the energy model.

  • random_state (int | numpy.random.Generator | None, optional) – Seed or generator for RNG seeding.

  • verbosity (int) – Verbosity level (0-2). 0 is silent, 1 prints high-level progress messages, 2 prints detailed logging messages. Can be overridden in methods that support it.

  • wrap (bool, default False) – If true, allows wrapping around the edges of the grid (toroidal). If false, no wrapping is applied.

  • vertical_exaggeration (float, default 1.0) – Vertical exaggeration factor for the network. Used when computing elevation.

Returns:

A newly constructed instance initialized from the specified network type and dimensions.

Return type:

OCN

property history: ndarray

The optimization history as a numpy array of shape (n_iterations, 3). Each row corresponds to an iteration. The columns are iteration index, energy, and temperature.

property nroots: int
property resolution: float
property rng: int
single_iteration(temperature: float, array_report: bool = False, unwrap: bool = True, calculate_full_energy: bool = False) → xr.Dataset | None[source]

Perform a single iteration of the optimization algorithm at a given temperature. Updates the internal history attribute. See fit() for details on the algorithm.

Parameters:

  • temperature (float) – Temperature parameter governing acceptance probability. Typical range is a fraction of ocn.energy.

  • array_report (bool, default False) – If True (default), the returned result will be an xarray.Dataset. See fit() for details. The returned object will have an iteration dimension of size 1. Requires xarray to be installed.

  • unwrap (bool, default True) – If True and the current OCN has periodic boundaries, the resulting rasters will be transformed so connected grid cells are adjacent in the output. This will result in a larger raster with some nan values. If False or the current OCN does not have periodic boundaries, then no transformation is applied and the resulting raster will have the same dimensions as the current OCN grid.

  • calculate_full_energy (bool, default False) – If True, the full energy of the graph is recalculated when considering the proposed change. If False, a more efficient incremental update is used. Small numerical differences may arise between the two methods due to floating point precision. If precision is of the utmost importance, set this to True, but note that this comes with a significant performance penalty.

Returns:

If array_report == True, an xarray.Dataset containing the state of the FlowGrid after the iteration. See to_xarray() for details. If array_report == False, returns None.

Return type:

xr.Dataset | None

Raises:

LibOCNError – If the underlying C routine reports an error status.

to_digraph() → DiGraph[source]

Create a NetworkX DiGraph view of the current grid.

Returns:

  • nx.DiGraph

  • # TODO (implement elevation, length, slope, q) – A DAG with the following node attributes per node:

    • pos: (row, col) grid position

    • drained_area: drained area value

    • energy: cumulative energy at the node

    • elevation: elevation at the node

    and the following edge attributes per edge: - length: length of the edge/link - slope: slope of the edge/link - q: the landscape-forming discharge through the edge/link

to_gtiff(west: float, north: float, crs: Any, path: str | PathLike | rasterio.io.MemoryFile, unwrap: bool = True)[source]

Export a raster of the current FlowGrid to a GeoTIFF file using rasterio. The resulting raster has 3 bands: energy, drained_area, and elevation.

Caution

This uses the resolution attribute to set pixel size in the raster and for computing elevation. This function does not check for unit compatibility, so it is up to the user to ensure the resolution and CRS units match. By default, the resolution is set to 1.0 units. Using a lat-long (geographic) CRS may give undesireable results. It is recommended to use a linear (projected) CRS with units of length (e.g., meters) for best results.

Parameters:

  • west (float) – The western border of the raster in crs units, corresponding to column 0.

  • north (float) – The northern border of the raster in crs units, corresponding to row 0.

  • crs (Any) – The crs for the resulting gtiff, passed to rasterio.open

  • path (str or Pathlike or rasterio.io.MemoryFile) – The output path for the resulting gtiff file. Can also be rasterio.io.MemoryFile.

  • unwrap (bool, default True) – If True and the current OCN has periodic boundaries, the resulting raster will be transformed so connected grid cells are adjacent in the raster. This will result in a larger raster with some nan values. If False or the current OCN does not have periodic boundaries, then no transformation is applied and the raster will have the same dimensions as the current OCN grid.

to_numpy(unwrap: bool = True) → ndarray[source]

Export the current FlowGrid to a numpy array with shape (3, rows, cols). Has three channels: 0=energy, 1=drained_area, 2=elevation.

Parameters:

unwrap (bool, default True) – If True and the current OCN has periodic boundaries, the resulting array will be transformed so connected grid cells are adjacent in the array. This will result in a larger array with some nan values. If False or the current OCN does not have periodic boundaries, then no transformation is applied and the resulting array will have the same dimensions as the current OCN grid.

to_xarray(unwrap: bool = True) → xr.Dataset[source]

Export the current FlowGrid to an xarray Dataset

Parameters:

unwrap (bool, default True) –

If True and the current OCN has periodic boundaries, the resulting rasters will be transformed so connected grid cells are adjacent in the output. This will result in a larger raster with some nan values. If False or the current OCN does not have periodic boundaries, then no transformation is applied and the resulting raster will have the same dimensions as the current OCN grid.

When unwrapping, the (0,0) coordinate will be set to the position of the “main” root node, defined as the root node with the smallest row*cols + col value. Otherwise, (0,0) will be the top-left corner of the grid.

Returns:

  • xr.Dataset –

    an xarray Dataset with data variables:

    • energy (np.float64) representing energy at each grid cell

    • area (np.float64) representing drained area at each grid cell

    • elevation (np.float64) representing elevation at each grid cell

  • and coordinates –

    • y (float) representing the northing coordinate of each row.

    • x (float) representing the easting coordinate of each column.

property wrap: bool