Improve constraint handling, typing, and doc

.github/copilot-instructions.md

@@ -0,0 +1,20 @@
+# Copilot Instructions
+
+YOU MUST UBSOLUTELY FOLLOW THESE INSTRUCTIONS AT ALL TIMES.
+
+## General guidelines to be used throughout:
+- Use spaces, not tabs: 4 spaces per indentation level 
+- Variables are snake_case, Classes are CamelCase.
+
+## When I ask you to overhaul a module, follow these instructions:
+- Add type hints to all functions and methods. If a variable is initialized to None, make sure to use | None in the type hint.
+- ABSOLUTELY do not document standard functions such as __init__, __str__, __repr__, unless there is something special to document.
+- ABSOLUTELY do not document property funcitons (those that use @property decorator), unless there is something special to document.
+- Never remove inline comments, those starting with #. They are important for understanding the code. You may correct or complete them.
+- Add at least one-line docstrings to each function, unless the function starts with an underscore, and unless it is clear from the functon name what it does.
+- ABSOLUTELY if the function or class already has a docstring, do not rewrite it, only correct or complete it.
+- Fix typos in documentation. Do not make any unnecessary changes to language, only if something is grammatically incorrect. 
+- Make code more efficient without making it less readable. 
+- If the module contains mostly one class, just provide one very brief module-level docstring. All the main information should be in the class top-level docstring. 
+- If this module contains many helper functions, then add a top-level docstring if it does not exist yet. If there is one, make sure that it features all important aspects of the code. Do not add information such as “Requires:, Author:, Date:”, only the technical content of the file. 
+- Do not introduce dummy variables, like d=self.d inside a function. Always use self.d for better readability.

docs/source/conf.py

@@ -120,6 +120,9 @@
 # Do not show "View Page Source"
 html_show_sourcelink = False
 
+# Make sure that inheriting classes do not show the docstring of the parent class
+autodoc_inherit_docstrings = False
+
 
 def setup(app):
     app.add_css_file("custom_css")

popcor/auto_template.py

@@ -267,25 +267,30 @@ def is_tight(self, verbose=False, data_dict={}, tightness=None):
         elif tightness == "cost":
             return cost_tight
 
-    def get_A_list(self, var_dict=None):
+    def get_A_b_list(self, var_dict=None):
+        A_0, b_0 = self.lifter.get_A0()
+        A_b_list = list(zip(A_0, b_0))
         if var_dict is None:
-            A_known = []
             if self.use_known:
-                A_known += [constraint.A_sparse_ for constraint in self.templates_known]
-            return A_known + [constraint.A_sparse_ for constraint in self.constraints]
+                A_b_list += [
+                    (constraint.A_sparse_, constraint.rhs)
+                    for constraint in self.templates_known
+                ]
+            return A_b_list + [
+                (constraint.A_sparse_, constraint.rhs)
+                for constraint in self.constraints
+            ]
         else:
-            A_known = []
+            A_b_list = []
             if self.use_known:
-                A_known += [constraint.A_poly_ for constraint in self.templates_known]
-            A_list_poly = A_known + [
-                constraint.A_poly_ for constraint in self.constraints
+                A_b_list += [
+                    (constraint.A_poly_, constraint.rhs)
+                    for constraint in self.templates_known
+                ]
+            A_list_poly = A_b_list + [
+                (constraint.A_poly_, constraint.rhs) for constraint in self.constraints
             ]
-            return [A.get_matrix(var_dict) for A in A_list_poly]
-
-    def get_A_b_list(self):
-        A_list = self.get_A_list()
-        A_b_list_all = self.lifter.get_A_b_list(A_list)
-        return A_b_list_all
+            return [(A.get_matrix(var_dict), b) for A, b in A_list_poly]
 
     def generate_minimal_subset(
         self,
@@ -449,11 +454,11 @@ def find_local_solution(self, n_inits=None, verbose=False, plot=False):
             for key, qcqp_that_local in info.items():
                 if key.startswith("local solution"):
                     solution_idx = key.strip("local solution ")
-                    error_dict = self.lifter.get_error(qcqp_that_local)
+                    error_dict = self.lifter.get_error(np.asarray(qcqp_that_local))
                     self.solver_vars.update(
                         {
                             f"local {solution_idx} {error_name}": err
-                            for error_name, err in error_dict.items()
+                            for error_name, err in error_dict.items()  # type: ignore
                         }
                     )
 
@@ -517,7 +522,7 @@ def find_global_solution(self, data_dict={}):
 
         if x is not None:
             theta = self.lifter.get_theta(x)
-            cost = self.lifter.get_cost(theta, self.solver_vars["y"])
+            cost = self.lifter.get_cost(theta, self.solver_vars["y"])  # type: ignore
             data_dict["global theta"] = theta
             data_dict["global cost"] = cost
             return True
@@ -584,11 +589,13 @@ def learn_templates(self, plot=False, data_dict=None):
         if self.use_incremental:
             for c in self.templates:
                 ai = get_vec(c.A_poly_.get_matrix(mat_var_dict))
+                assert isinstance(ai, np.ndarray)
                 bi = self.lifter.augment_using_zero_padding(ai, param_dict)
                 a_vectors.append(bi)
         if self.use_known:
             for c in self.templates_known_sub:
                 ai = get_vec(c.A_poly_.get_matrix(mat_var_dict))
+                assert isinstance(ai, np.ndarray)
                 bi = self.lifter.augment_using_zero_padding(ai, param_dict)
                 a_vectors.append(bi)
         Y = np.vstack([Y] + a_vectors)
@@ -732,8 +739,8 @@ def get_known_templates(self, unroll=False):
 
         # TODO(FD) we should not always recompute from scratch, but it's not very expensive so it's okay for now.
         target_dict = self.lifter.get_var_dict(unroll_keys=unroll)
-        for i, Ai in enumerate(
-            self.lifter.get_A_known(var_dict=target_dict, output_poly=True)
+        for i, (Ai, bi) in enumerate(
+            zip(*self.lifter.get_A_known(var_dict=target_dict, output_poly=True))
         ):
             template = Constraint.init_from_A_poly(
                 lifter=self.lifter,
@@ -743,6 +750,7 @@ def get_known_templates(self, unroll=False):
                 template_idx=self.constraint_index,
                 mat_var_dict=self.lifter.var_dict,
                 compute_polyrow_b=True,
+                rhs=bi,
             )
             self.constraint_index += 1
             templates_known.append(template)
@@ -825,7 +833,7 @@ def sort_fun_sparsity(series):
                 )
             )
         df = pd.DataFrame(
-            series, dtype="Sparse[float]", index=templates_poly.variable_dict_i
+            series, dtype="Sparse[float]", index=templates_poly.variable_dict_i  # type: ignore
         )
         df.dropna(axis=1, how="all", inplace=True)
 
@@ -1033,7 +1041,7 @@ def save_matrices_sparsity(self, A_matrices=None, fname_root="", title=""):
 
         vmin = min(-np.max(Q), np.min(Q))
         vmax = max(np.max(Q), -np.min(Q))
-        norm = matplotlib.colors.SymLogNorm(10**-5, vmin=vmin, vmax=vmax)
+        norm = matplotlib.colors.SymLogNorm(10**-5, vmin=vmin, vmax=vmax)  # type: ignore
         im1 = axs[1].matshow(Q, norm=norm)
 
         for ax in axs:
@@ -1228,10 +1236,14 @@ def apply(self, lifter: StateLifter, use_known: bool = False) -> list:
         """Apply the learned templates to a new lifter."""
         constraints = lifter.apply_templates(self.templates)
 
+        A_0, b_0 = lifter.get_A0()
         if use_known:
             # if we set use_known=True in running AutoTemplate, then we learned only
             # constraints that were not already known, so we need to add them to the
             # overall set of constraints.
-            A_known = lifter.get_A_known()
+            A_known, b_known = lifter.get_A_known()
             assert isinstance(A_known, list)
-        return A_known + [c.A_sparse_ for c in constraints]  # type: ignore
+            A_b_list = list(zip(A_0 + A_known, b_0 + b_known))
+        else:
+            A_b_list = list(zip(A_0, b_0))
+        return A_b_list + [(c.A_sparse_, c.rhs) for c in constraints]  # type: ignore