Merge pull request #13 from utiasASRL/new_plotting_feature

New plotting feature and cleanup

.github/workflows/python-app.yml

@@ -36,4 +36,5 @@ jobs:
         flake8 . --count --exit-zero --max-complexity=10 --max-line-length=127 --statistics
     - name: Test with pytest
       run: |
+        pip install -e .
         pytest

_test/test_ls.py

@@ -1,12 +1,6 @@
 import itertools
-import sys
-from os.path import dirname
 
 import numpy as np
-import pytest
-
-sys.path.append(dirname(__file__) + "/../")
-print("appended:", sys.path[-1])
 
 from poly_matrix.least_squares_problem import LeastSquaresProblem
 

_test/test_poly_matrix.py

@@ -1,16 +1,10 @@
-import sys
-from os.path import dirname
-
 import numpy as np
 import pytest
 
-sys.path.append(dirname(__file__) + "/../")
-print("appended:", sys.path[-1])
-
 from poly_matrix import PolyMatrix, sorted_dict
 
 
-def get_Ai(test=False):
+def get_Ai():
     """
     Creates the following matrices:
 
@@ -202,10 +196,6 @@ def test_Q():
     get_Q(test=True)
 
 
-def test_Ai():
-    get_Ai(test=True)
-
-
 def test_operations_simple():
     # mat1 =
     # 1 0
@@ -405,7 +395,7 @@ def test_scaling():
 
         print("getting matrix...", end="")
 
-        mat.get_matrix(verbose=True)
+        mat.get_matrix()
         print("...done")
 
 
@@ -450,21 +440,19 @@ def test_multiply():
 
     AT = A.transpose()
     S = A.multiply(B.multiply(AT))
-    np.testing.assert_allclose(
-        S["x1", "x1"], np.c_[np.r_[3 * 7 * 7, 3 * 7 * 8], np.r_[3 * 7 * 8, 3 * 8 * 8]]
-    )
-    np.testing.assert_allclose(
-        S["x2", "x2"],
-        np.c_[np.r_[2 * 9 * 9, 2 * 9 * 10], np.r_[2 * 9 * 10, 2 * 10 * 10]],
+    S_test = (
+        A.get_matrix_dense((["x1", "x2"], ["z1", "z2"]))
+        @ B.get_matrix_dense(["z1", "z2"])
+        @ AT.get_matrix_dense((["z1", "z2"], ["x1", "x2"]))
     )
+    np.testing.assert_allclose(S.get_matrix_dense(["x1", "x2"]), S_test)
 
 
 if __name__ == "__main__":
     test_multiply()
 
     test_get_empty()
 
-    test_Ai()
     test_Q()
 
     test_join_dicts()

environment.yml

@@ -1,12 +1,20 @@
-name: polymat 
+name: poly_matrix
 channels:
     - defaults
     - conda-forge
     - anaconda
+
 dependencies:
     - python=3.10
     - pip=22.3
+    - flake8
+
+    - pandas>=1.5.3
+    - numpy>=1.23.5
+    - scipy>=1.9.3
+    - pytest>=7.2.1
+    - matplotlib>=3.6.2
 
     - pip:
       - -r requirements.txt
-      - -e . 
+      - -e .  # build local package

poly_matrix/__init__.py

@@ -1,3 +1,3 @@
 from .poly_matrix import *
 
-__version__ = "0.0.2"
+__version__ = "0.1.1"

poly_matrix/poly_matrix.py

@@ -1,3 +1,4 @@
+import itertools
 from copy import deepcopy
 import warnings
 
@@ -106,11 +107,7 @@ def __init__(self, symmetric=True):
         self.variable_dict_i = {}
         self.variable_dict_j = {}
 
-        # TODO(FD) technically, adjacency_i has redundant information, since
-        # self.matrix.keys() could be used. Consider removing it (for now,
-        # it is kept for analogy with adjacency_j).
         # adjacency_j allows for fast starting of the adjacency variables of j.
-        self.adjacency_i = {}
         self.adjacency_j = {}
 
         self.shape = (0, 0)
@@ -155,17 +152,14 @@ def init_from_sparse(A, var_dict, unfold=False):
             return self, new_var_dict
         return self, var_dict
 
-    def drop(self, variables_i):
-        for v in variables_i:
+    def drop(self, variables):
+        for v in variables:
             if v in self.matrix:
                 self.matrix.pop(v)
-            if v in self.adjacency_i:
-                j_list = self.adjacency_i[v]
-                for j in j_list:
-                    self.adjacency_j[j].remove(v)
-                self.adjacency_i.pop(v)
             if v in self.variable_dict_i:
                 self.variable_dict_i.pop(v)
+            if v in self.variable_dict_j:
+                self.variable_dict_j.pop(v)
 
     def __getitem__(self, key):
         key_i, key_j = key
@@ -206,13 +200,7 @@ def add_variable_j(self, key, size):
         self.last_var_j_index += size
 
     def add_key_pair(self, key_i, key_j):
-        if key_i in self.adjacency_i.keys():
-            if not (key_j in self.adjacency_i[key_i]):
-                self.adjacency_i[key_i].append(key_j)
-        else:
-            self.adjacency_i[key_i] = [key_j]
-
-            assert key_i not in self.matrix.keys()
+        if key_i not in self.matrix.keys():
             self.matrix[key_i] = {}
 
         if key_j in self.adjacency_j.keys():
@@ -252,7 +240,7 @@ def __setitem__(self, key_pair, val, symmetric=None):
 
         # make sure the dimensions of new block are consistent with
         # previously inserted blocks.
-        if key_i in self.adjacency_i:
+        if key_i in self.matrix.keys():
             assert (
                 val.shape[0] == self.variable_dict_i[key_i]
             ), f"mismatch in height of filled value for key_i {key_i}: got {val.shape[0]} but expected {self.variable_dict_i[key_i]}"
@@ -317,7 +305,7 @@ def generate_variable_dict(self, variables=None, key="i"):
     def generate_variable_dict_i(self, variables=None):
         """Regenerate last_var_index using new ordering."""
         if variables is None:
-            variables = list(self.adjacency_i.keys())
+            variables = list(self.matrix.keys())
         return self._generate_variable_dict(variables, self.variable_dict_i)
 
     def generate_variable_dict_j(self, variables=None):
@@ -330,7 +318,9 @@ def _generate_variable_dict(self, variables, variable_dict):
         if isinstance(variables, dict):
             return {key: variable_dict.get(key, variables[key]) for key in variables}
         else:
-            return {key: variable_dict[key] for key in variables}
+            return {
+                key: variable_dict[key] for key in variables if key in variable_dict
+            }
 
     def get_variables(self, key=None):
         """Return variable names starting with key.
@@ -350,6 +340,7 @@ def get_variables(self, key=None):
     # TODO(FD): there is probably a much cleaner way of doing this -- basically,
     # keeping track of N somewhere else. For now, this is a quick hack.
     def get_max_index(self):
+        print("Warning: get_max_index is inefficient.")
         max_ = 0
         for key in self.variable_dict.keys():
             if isinstance(key, str):
@@ -439,12 +430,13 @@ def get_matrix_poly(self, variables, verbose=False):
         out_matrix.variable_dict_j = variable_dict_j
         return out_matrix
 
-    def get_matrix_dense(self, variables, verbose=False):
+    def get_matrix_dense(self, variables=None, verbose=False):
         """Return a small submatrix in dense format
 
         :param variables: same as in self.get_matrix, but None is not allowed
         """
-        assert variables is not None
+        if variables is None:
+            variables = self.get_variables()
         if isinstance(variables, list):
             variable_dict_i = self.generate_variable_dict_i(variables)
             variable_dict_j = self.generate_variable_dict_j(variables)
@@ -478,6 +470,14 @@ def get_matrix_dense(self, variables, verbose=False):
             index_i += size_i
         return matrix
 
+    def get_start_indices(self, axis=0):
+        if axis in [0, "i"]:
+            return generate_indices(self.variable_dict_i)
+        elif axis in [1, "j"]:
+            return generate_indices(self.variable_dict_j)
+        else:
+            raise ValueError(f"Invalid axis {axis}")
+
     def get_matrix_sparse(self, variables=None, output_type="coo", verbose=False):
         """Return a sparse matrix in desired format.
 
@@ -531,20 +531,22 @@ def get_matrix_sparse(self, variables=None, output_type="coo", verbose=False):
         data_list = []
 
         # Loop through blocks of stored matrices
-        for key_i in self.matrix:
-            for key_j in self.matrix[key_i]:
-                # Check if blocks appear in variable dictionary
-                if key_i in variable_dict["i"] and key_j in variable_dict["j"]:
+        for key_i in variable_dict["i"]:
+            for key_j in variable_dict["j"]:
+                try:
                     values = self.matrix[key_i][key_j]
-                    assert values.shape == (
-                        variable_dict["i"][key_i],
-                        variable_dict["j"][key_j],
-                    ), f"Variable size does not match input matrix size, variables: {(variable_dict['i'][key_i], variable_dict['j'][key_j])}, matrix: {values.shape}"
-                    # generate list of indices for sparse mat input
-                    rows, cols = np.nonzero(values)
-                    i_list = np.append(i_list, rows + indices_i[key_i])
-                    j_list = np.append(j_list, cols + indices_j[key_j])
-                    data_list = np.append(data_list, values[rows, cols])
+                except KeyError:
+                    continue
+                # Check if blocks appear in variable dictionary
+                assert values.shape == (
+                    variable_dict["i"][key_i],
+                    variable_dict["j"][key_j],
+                ), f"Variable size does not match input matrix size, variables: {(variable_dict['i'][key_i], variable_dict['j'][key_j])}, matrix: {values.shape}"
+                # generate list of indices for sparse mat input
+                rows, cols = np.nonzero(values)
+                i_list += list(rows + indices_i[key_i])
+                j_list += list(cols + indices_j[key_j])
+                data_list += list(values[rows, cols])
 
         shape = get_shape(variable_dict["i"], variable_dict["j"])
 
@@ -798,6 +800,24 @@ def matshow(
         )
         return fig, ax, im
 
+    def plot_box(self, ax, clique_keys, symmetric=True, **kwargs):
+        delta = 0.499
+        indices = self.get_start_indices(axis=1)
+        for key_i, key_j in itertools.combinations_with_replacement(clique_keys, 2):
+            i1 = indices[key_i] - delta
+            i2 = i1 + self.variable_dict_j[key_i]
+            j1 = indices[key_j] - delta
+            j2 = j1 + self.variable_dict_j[key_j]
+            ax.plot([j1, j2], [i1, i1], **kwargs)
+            ax.plot([j1, j2], [i2, i2], **kwargs)
+            ax.plot([j1, j1], [i1, i2], **kwargs)
+            ax.plot([j2, j2], [i1, i2], **kwargs)
+            if symmetric:
+                ax.plot([i1, i2], [j1, j1], **kwargs)
+                ax.plot([i1, i2], [j2, j2], **kwargs)
+                ax.plot([i1, i1], [j1, j2], **kwargs)
+                ax.plot([i2, i2], [j1, j2], **kwargs)
+
     def __repr__(self, variables=None, binary=False):
         """Called by the print() function"""
         if self.shape is None:
@@ -840,8 +860,8 @@ def __add__(self, other, inplace=False):
                 "Both matrices must be symmetric or non-symmetric to add."
             )
             # Loop through second matrix elements
-            for key_i in other.adjacency_i:
-                for key_j in other.adjacency_i[key_i]:
+            for key_i in other.matrix:
+                for key_j in other.matrix[key_i]:
                     # Check if element exists in first matrix
                     if key_i in res.matrix and key_j in res.matrix[key_i]:
                         # Check shapes of matrices to be added
@@ -863,8 +883,8 @@ def __add__(self, other, inplace=False):
                         )
         else:
             # simply add constant to all non-zero elements
-            for key_i in res.adjacency_i.keys():
-                for key_j in res.adjacency_i[key_i]:
+            for key_i in res.matrix:
+                for key_j in res.matrix[key_i]:
                     if other[key_i, key_j] is not None:
                         res.matrix[key_i][key_j] += other
         return res
@@ -887,27 +907,28 @@ def __mul__(self, scalar, inplace=False):
             res = self
         else:
             res = self.copy()
-        for key_i in res.adjacency_i.keys():
-            for key_j in res.adjacency_i[key_i]:
+        for key_i in res.matrix:
+            for key_j in res.matrix[key_i]:
                 res.matrix[key_i][key_j] *= scalar
         return res
 
     def transpose(self):
         res = deepcopy(self)
 
+        res.adjacency_j = {}
         matrix_tmp = {}
-        for key_i, key_j_list in res.adjacency_i.items():
+        for key_i, key_j_list in res.matrix.items():
             for key_j in key_j_list:
                 if key_j in matrix_tmp.keys():
                     matrix_tmp[key_j][key_i] = res.matrix[key_i][key_j].T
                 else:
                     matrix_tmp[key_j] = {key_i: res.matrix[key_i][key_j].T}
+                if key_i in res.adjacency_j:
+                    res.adjacency_j[key_i].append(key_j)
+                else:
+                    res.adjacency_j[key_i] = [key_j]
         res.matrix = matrix_tmp
 
-        tmp = deepcopy(res.adjacency_i)
-        res.adjacency_i = res.adjacency_j
-        res.adjacency_j = tmp
-
         tmp = deepcopy(res.variable_dict_i)
         res.variable_dict_i = res.variable_dict_j
         res.variable_dict_j = tmp
@@ -924,11 +945,11 @@ def multiply(self, other_mat):
         """
         output_mat = PolyMatrix(symmetric=False)
 
-        rows = self.adjacency_i.keys()
+        rows = self.matrix.keys()
         cols = other_mat.adjacency_j.keys()
         for key_i in rows:
             for key_j in cols:
-                common_elements = set(self.adjacency_i[key_i]).intersection(
+                common_elements = set(self.matrix[key_i].keys()).intersection(
                     other_mat.adjacency_j[key_j]
                 )
                 for key_mul in common_elements:

requirements.txt

@@ -1,6 +1,5 @@
 pandas>=1.5.3
 numpy>=1.24.1
+scipy==1.9.1
 pytest>=7.2.1
 matplotlib>=3.6.2
-scipy
-pytest-cov

setup.cfg

@@ -1,6 +1,6 @@
 [metadata]
 name = poly_matrix
-version = 0.0.2
+version = 0.1.1
 authors = [
     {name = "Frederike Dümbgen", email = "[email protected]" },
     {name = "Connor Holmes", email = "[email protected]" }]
@@ -16,7 +16,7 @@ license = { file="LICENSE" }
 [options]
 packages = find:
 install_requires=
-  scipy
+  scipy==1.9.1
 
 [options.packages.find] # do not mistake tests/ for a package directory
 exclude=_test*