1#  ___________________________________________________________________________
2#
3#  Pyomo: Python Optimization Modeling Objects
4#  Copyright 2017 National Technology and Engineering Solutions of Sandia, LLC
5#  Under the terms of Contract DE-NA0003525 with National Technology and
6#  Engineering Solutions of Sandia, LLC, the U.S. Government retains certain
7#  rights in this software.
8#  This software is distributed under the 3-clause BSD License.
9#  ___________________________________________________________________________
10
11from collections import namedtuple
12from pyomo.common.dependencies import networkx as nx
13from pyomo.contrib.incidence_analysis.common.dulmage_mendelsohn import (
14    dulmage_mendelsohn as dm_nx,
15    )
16"""
17This module imports the general Dulmage-Mendelsohn-on-a-graph function
18from "common" and implements an interface for coo_matrix-like objects.
19"""
20
21RowPartition = namedtuple(
22        "RowPartition",
23        ["unmatched", "overconstrained", "underconstrained", "square"],
24        )
25ColPartition = namedtuple(
26        "ColPartition",
27        ["unmatched", "underconstrained", "overconstrained", "square"],
28        )
29
30def dulmage_mendelsohn(matrix_or_graph, top_nodes=None, matching=None):
31    """
32    COO matrix or NetworkX graph interface to the coarse Dulmage Mendelsohn
33    partition. The matrix or graph should correspond to a Pyomo model.
34    top_nodes must be provided if a NetworkX graph is used, and should
35    correspond to Pyomo constraints.
36
37    """
38    if isinstance(matrix_or_graph, nx.Graph):
39        # The purpose of handling graphs here is that if we construct NX graphs
40        # directly from Pyomo expressions, we can eliminate the overhead of
41        # convering expressions to a matrix, then the matrix to a graph.
42        #
43        # In this case, top_nodes should correspond to constraints.
44        graph = matrix_or_graph
45        if top_nodes is None:
46            raise ValueError(
47                    "top_nodes must be specified if a graph is provided,"
48                    "\notherwise the result is ambiguous."
49                    )
50        partition = dm_nx(graph, top_nodes=top_nodes, matching=matching)
51        # RowPartition and ColPartition do not make sense for a general graph.
52        # However, here we assume that this graph comes from a Pyomo model,
53        # and that "top nodes" are constraints.
54        partition = (RowPartition(*partition[0]), ColPartition(*partition[1]))
55    else:
56        # Assume matrix_or_graph is a scipy coo_matrix
57        matrix = matrix_or_graph
58        M, N = matrix.shape
59        nxb = nx.algorithms.bipartite
60        from_biadjacency_matrix = nxb.matrix.from_biadjacency_matrix
61
62        if matching is not None:
63            # If a matching was provided for a COO matrix, we assume it
64            # maps row indices to column indices, for compatibility with
65            # output of our maximum_matching function.
66
67            # NetworkX graph has column nodes offset by M
68            matching = {i: j + M for i, j in matching.items()}
69            inv_matching = {j: i for i, j in matching.items()}
70            # DM function requires matching map to contain inverse matching too
71            matching.update(inv_matching)
72
73        # Matrix rows have bipartite=0, columns have bipartite=1
74        bg = from_biadjacency_matrix(matrix)
75        row_partition, col_partition = dm_nx(
76                bg,
77                top_nodes=list(range(M)),
78                matching=matching,
79                )
80
81        partition = (
82                row_partition,
83                tuple([n - M for n in subset] for subset in col_partition)
84                # Column nodes have values in [M, M+N-1]. Apply the offset
85                # to get values corresponding to indices in user's matrix.
86                )
87
88    partition = (RowPartition(*partition[0]), ColPartition(*partition[1]))
89    return partition
90