Graph¶

In this package, the class Graph is the core data representation. It is used to represent contexts \(x\) (i.e input data), decisions \(y\) (i.e. output data), and gradients \(\nabla_y f\). A Graph is composed of multiple HyperEdgeSet objects, each defined by a series of ports and features.

The class Graph can represent both a single graph instance or a batch of graphs. The backend (NumPy or JAX) is controlled by passing a NumpyBackend or JaxBackend instance at construction time.

Note

Here is a typical instance of Graph.

>>> print(graph)
Mass
          ports      features
            node_id    weight         x         y         z
object_id
0               0.0  5.322265  0.202435  0.202435  0.242032
1               1.0  3.496568  0.962326  0.962326  0.306690
2               2.0  3.535864  0.060886  0.060886  0.094170
3               3.0  7.213709  0.984766  0.984766  0.068853
Spring
          ports              features
           node1_id node2_id         k
object_id
0               0.0      1.0  0.020424
1               1.0      2.0  0.037591
2               2.0      3.0  0.045405
Registry
[0. 1. 2. 3.]

Graph¶

class Graph(*, backend=None, hyper_edge_sets, true_shape, current_shape, non_fictitious_addresses)[source]¶

Bases: dict

Hyper Heterogeneous Multi-Graph (H2MG) container.

Stores hyper-edge sets, shapes, and address masks for single or batched graphs. All array operations are delegated to the provided backend.

Parameters:

backend (Backend | None) – Array backend (NumpyBackend or JaxBackend).
hyper_edge_sets (dict[str, HyperEdgeSet]) – Dict of hyper-edge sets.
true_shape (GraphShape) – True shape, unaffected by padding.
current_shape (GraphShape) – Current shape, consistent with padding.
non_fictitious_addresses – 1 for real addresses, 0 otherwise.

`Graph.from_dict`	Build a Graph from a dict of HyperEdgeSet and a number of addresses.
`Graph.to_pickle`	Save this graph to a pickle file.
`Graph.from_pickle`	Load a graph from a pickle file.
`Graph.is_batch`	True if all hyper-edge sets are batched and the address mask is 2-D.
`Graph.is_single`	True if all hyper-edge sets are single and the address mask is 1-D.
`Graph.feature_flat_array`	Concatenated flat features of all hyper-edge sets.
`Graph.pad`	Pad hyper-edge sets and address mask to `target_shape`.
`Graph.unpad`	Remove padding to restore the true shape.
`Graph.count_connected_components`	Count connected components and return per-address component labels.
`Graph.offset_addresses`	Add `offset` to all port addresses; used before graph concatenation.
`Graph.quantiles`	Compute quantiles of all hyper-edge set features.

HyperEdgeSet¶

class HyperEdgeSet(*, backend=None, port_dict, feature_array, feature_names, non_fictitious)[source]¶

Bases: dict

A collection of hyper-edges of the same class, optionally batched.

Internally this is a dict storing four entries. All array operations are delegated to the provided backend, making instances transparent to both NumPy and JAX pipelines.

Parameters:

backend (Backend | None) – Array backend (NumpyBackend or JaxBackend).
port_dict (dict | None) – Mapping from a port name to an integer address array of shape (n_edges,) or (batch, n_edges).
feature_array – Array that contains all hyper-edge features.
feature_names (dict | None) – Dictionary from feature names to index in feature_array.
non_fictitious – Mask array set to 1 for real objects and 0 for fictitious ones.

`HyperEdgeSet.from_dict`	Build a HyperEdgeSet from raw dicts of ports and features.
`HyperEdgeSet.array`	Concatenate (features, ports) along the last axis.
`HyperEdgeSet.is_batch`	True if `array` is 3-D: `(batch, n_obj, features+ports)`.
`HyperEdgeSet.is_single`	True if `array` is 2-D: `(n_obj, features+ports)`.
`HyperEdgeSet.n_obj`	Number of hyper-edges per instance.
`HyperEdgeSet.n_batch`	Number of batches; valid only when `is_batch` is True.
`HyperEdgeSet.port_array`	Stacked port array of shape `(n_obj, n_ports)` or `(batch, n_obj, n_ports)`.
`HyperEdgeSet.port_names`	Maps port name to column index in `port_array`.
`HyperEdgeSet.feature_dict`	Unstack `feature_array` into a dict: feature_name → array slice.
`HyperEdgeSet.feature_flat_array`	Flatten all features into one long vector per `(batch,)` in Fortran order.
`HyperEdgeSet.pad`	Pad a single HyperEdgeSet to `target_shape` objects with zeros/zeros.
`HyperEdgeSet.unpad`	Remove padding to restore `target_shape` objects in a single HyperEdgeSet.
`HyperEdgeSet.offset_addresses`	Add `offset` to every port address; used before graph concatenation.

GraphShape¶

class GraphShape(*, backend=None, hyper_edge_sets, addresses)[source]¶

Bases: dict

Represents the shape of a graph: per-class object counts and registry size.

Parameters:

backend (Backend | None) – Array backend.
hyper_edge_sets (dict) – Dict mapping hyper-edge class name to count array.
addresses – Number of addresses in the graph.

`GraphShape.from_dict`	Build a GraphShape from a hyper-edge set dictionary and a non-fictitious mask.
`GraphShape.to_jsonable_dict`	Serialize GraphShape to a JSON-friendly dict.
`GraphShape.from_jsonable_dict`	Deserialize GraphShape from a JSON-friendly dictionary.
`GraphShape.max`	Return the element-wise maximum of two GraphShape objects.
`GraphShape.sum`	Return the element-wise sum of two GraphShape objects.
`GraphShape.array`	Concatenated hyper-edge-set shape values as a single array.
`GraphShape.is_single`	True if the array is 1-D.
`GraphShape.is_batch`	True if the array is 2-D.
`GraphShape.n_batch`	Return the batch size; raises if not batched.

Graph, hyper-edge set, and shape manipulation functions¶

The following functions help to manipulate graphs, hyper-edge sets, shapes objects and to proceed operations on them.

`collate_graphs`	Collate a list of Graphs into a single batched Graph (all must share `current_shape`).
`concatenate_graphs`	Concatenate multiple single Graphs into one single Graph (no new batch dim).
`get_statistics`	Extract summary statistics from each feature array in the graph.
`separate_graphs`	Split a batched Graph into a list of single Graphs (reverses `collate_graphs()`).
`check_hyper_edge_set_dict_type`	Validate that the mapping is a dict of HyperEdgeSet instances.
`collate_hyper_edge_sets`	Collate a list of HyperEdgeSet into a single batched HyperEdgeSet.
`concatenate_hyper_edge_sets`	Concatenate several single HyperEdgeSet into one single HyperEdgeSet (no new batch dim).
`separate_hyper_edge_sets`	Split a batched HyperEdgeSet into its constituent HyperEdgeSet instances.
`check_dict_shape`	Ensure all arrays in a dict share the same last-axis size.
`build_hyper_edge_set_shape`	Return a scalar int32 NumPy array with the number of hyper-edges.
`dict2array`	Stack a sorted dict of NumPy arrays into a single array along the last axis.
`check_dict_or_none`	Validate that the input is a dict or None.
`check_no_nan`	Ensure there are no NaN values in port or feature arrays.
`collate_shapes`	Batch a list of GraphShape into one batched GraphShape.
`max_shape`	Return the element-wise maximum over a list of GraphShape objects.
`separate_shapes`	Split a batched GraphShape into individual GraphShape instances.
`sum_shapes`	Return the element-wise sum over a list of GraphShape objects.
`to_numpy`	Converts a NumPy array, JAX array, or tuple of values into a NumPy array preserving its dtype, or converts the values in a dictionary accordingly.
`np_to_jnp`	Convert NumPy arrays or dictionary of NumPy arrays to JAX arrays.
`jnp_to_np`	Convert JAX arrays or mappings of JAX arrays back to NumPy arrays.