graphnet.models.coarsening module

Class(es) for coarsening operations (i.e., clustering, or local pooling).

graphnet.models.coarsening.unbatch_edge_index(edge_index, batch)[source]

Splits the edge_index according to a batch vector.

Parameters:
  • edge_index (Tensor) – The edge_index tensor. Must be ordered.

  • batch (LongTensor) – The batch vector \(\mathbf{b} \in {\{ 0, \ldots, B-1\}}^N\), which assigns each node to a specific example. Must be ordered.

Return type:

List[Tensor]

class graphnet.models.coarsening.Coarsening(*args, **kwargs)[source]

Bases: Model

Base class for coarsening operations.

Construct Coarsening.

Parameters:
  • args (Any)

  • kwargs (Any)

Return type:

object

reduce_options = {'avg': (<function avg_pool>, <function avg_pool_x>), 'max': (<function max_pool>, <function max_pool_x>), 'min': (<function min_pool>, <function min_pool_x>), 'sum': (<function sum_pool>, <function sum_pool_x>)}
forward(data)[source]

Perform coarsening operation.

Return type:

Union[Data, Batch]

Parameters:

data (Data | Batch)

class graphnet.models.coarsening.AttributeCoarsening(*args, **kwargs)[source]

Bases: Coarsening

Coarsen pulses based on specified attributes.

Construct SimpleCoarsening.

Parameters:
  • args (Any)

  • kwargs (Any)

Return type:

object

class graphnet.models.coarsening.DOMCoarsening(*args, **kwargs)[source]

Bases: Coarsening

Coarsen pulses to DOM-level.

Cluster pulses on the same DOM.

Parameters:
  • args (Any)

  • kwargs (Any)

Return type:

object

class graphnet.models.coarsening.CustomDOMCoarsening(*args, **kwargs)[source]

Bases: DOMCoarsening

Coarsen pulses to DOM-level with additional attributes.

Cluster pulses on the same DOM.

Parameters:
  • args (Any)

  • kwargs (Any)

Return type:

object

class graphnet.models.coarsening.DOMAndTimeWindowCoarsening(*args, **kwargs)[source]

Bases: Coarsening

Coarsen pulses to DOM-level, with additional time-window clustering.

Cluster pulses on the same DOM within time_window.

Parameters:
  • args (Any)

  • kwargs (Any)

Return type:

object