Merge pull request #92 from chrhansk/feature-integration-solver

Add integration solver

docs/conf.py

@@ -14,7 +14,7 @@
 project = "pygradflow"
 copyright = "2023, Christoph Hansknecht"
 author = "Christoph Hansknecht"
-release = "0.4.22"
+release = "0.5.0"
 
 # -- General configuration ---------------------------------------------------
 # https://www.sphinx-doc.org/en/master/usage/configuration.html#general-configuration

pygradflow/display.py

@@ -1,10 +1,27 @@
 from abc import ABC, abstractmethod
 from typing import List, Literal
 
+import numpy as np
 from termcolor import colored
 
+from pygradflow.log import logger
 from pygradflow.params import Params
 from pygradflow.problem import Problem
+from pygradflow.timer import SimpleTimer
+
+
+class StateData:
+    def __init__(self):
+        self._entries = dict()
+
+    def __setitem__(self, key, value):
+        self._entries[key] = value
+
+    def __getitem__(self, key):
+        entry = self._entries[key]
+        if callable(entry):
+            return entry()
+        return entry
 
 
 class Format:
@@ -65,7 +82,7 @@ def header(self) -> str:
         return "{:^{}s}".format(self.name, self.width)
 
     @abstractmethod
-    def content(self, state) -> str:
+    def content(self, state, last_state) -> str:
         raise NotImplementedError()
 
 
@@ -80,20 +97,38 @@ def __init__(self, name: str, width: int, format, attr):
 
         self.attr = attr
 
-    def content(self, state) -> str:
+    def content(self, state, _) -> str:
         return self.format(self.attr(state))
 
 
 class Display:
-    def __init__(self, cols):
+    def __init__(self, cols, interval=None):
         self.cols = cols
+        self.interval = interval
+
+        self.timer = None
+        if self.interval is not None:
+            assert self.interval > 0
+            self.timer = SimpleTimer()
+        self.last_state = None
+
+    def should_display(self):
+        if self.timer is None:
+            return True
+
+        return self.timer.elapsed() >= self.interval
 
     @property
     def header(self) -> str:
         return " ".join([col.header for col in self.cols])
 
     def row(self, state) -> str:
-        return " ".join([col.content(state) for col in self.cols])
+        if self.timer is not None:
+            self.timer.reset()
+
+        row = " ".join([col.content(state, self.last_state) for col in self.cols])
+        self.last_state = state
+        return row
 
 
 class StateAttr:
@@ -102,7 +137,7 @@ def __init__(self, name: str):
 
     def __call__(self, state):
         value = state[self.name]
-        return value()
+        return value
 
 
 class IterateAttr:
@@ -121,9 +156,12 @@ def __init__(self):
     def empty(self):
         return "{:^{}s}".format("--", self.width)
 
-    def content(self, state):
-        curr_active_set = state["curr_active_set"]()
-        last_active_set = state["last_active_set"]()
+    def content(self, state, last_state):
+        curr_active_set = state["active_set"]
+        last_active_set = None
+
+        if (last_state is not None) and (state["iter"] == (last_state["iter"] + 1)):
+            last_active_set = last_state["active_set"]
 
         if curr_active_set is None:
             return self.empty()
@@ -142,12 +180,9 @@ def content(self, state):
             return "{:^{}s}".format("--", self.width)
 
 
-def problem_display(problem: Problem, params: Params):
+def iter_cols(problem):
     is_bounded = problem.var_bounded
-
-    cols: List[Column] = []
-
-    cols.append(AttrColumn("Iter", 6, BoldFormatter("{:6d}"), StateAttr("iter")))
+    cols = []
     cols.append(AttrColumn("Aug Lag", 16, "{:16.8e}", StateAttr("aug_lag")))
     cols.append(AttrColumn("Objective", 16, "{:16.8e}", IterateAttr("obj")))
 
@@ -158,6 +193,17 @@ def problem_display(problem: Problem, params: Params):
 
     cols.append(AttrColumn("Cons inf", 16, "{:16.8e}", IterateAttr("cons_violation")))
     cols.append(AttrColumn("Dual inf", 16, "{:16.8e}", IterateAttr("stat_res")))
+
+    return cols
+
+
+def solver_display(problem: Problem, params: Params):
+    cols: List[Column] = []
+
+    cols.append(AttrColumn("Iter", 6, BoldFormatter("{:6d}"), StateAttr("iter")))
+
+    cols += iter_cols(problem)
+
     cols.append(
         AttrColumn("Primal step", 16, "{:16.8e}", StateAttr("primal_step_norm"))
     )
@@ -177,6 +223,65 @@ def problem_display(problem: Problem, params: Params):
 
     cols.append(AttrColumn("Type", 8, StepFormatter(), StateAttr("step_accept")))
 
+    return Display(cols, interval=params.display_interval)
+
+
+class FilterColumn(Column):
+    def __init__(self):
+        super().__init__("Filter", 10)
+
+    def empty(self):
+        return "{:^{}s}".format("--", self.width)
+
+    def content(self, state, last_state):
+        curr_filter = state["filter"]
+        last_filter = None
+
+        if (last_state is not None) and (state["iter"] == (last_state["iter"] + 1)):
+            last_filter = last_state["filter"]
+
+        if curr_filter is None:
+            return self.empty()
+
+        display_curr = False
+
+        if last_filter is None:
+            display_curr = True
+        elif (curr_filter != last_filter).any():
+            display_curr = True
+
+        if display_curr:
+            num_active = curr_filter.sum()
+            return "{:^{}d}".format(num_active, self.width)
+        else:
+            return "{:^{}s}".format("--", self.width)
+
+
+class StepTypeColumn(Column):
+    def __init__(self):
+        super().__init__("Step Type", 10)
+        self.format = BoldFormatter("{:^10s}")
+
+    def content(self, state, last_state):
+        step_type = state["step_type"]
+        return self.format(step_type.name())
+
+
+def integrator_display(problem: Problem, params: Params):
+    cols: List[Column] = []
+
+    cols.append(AttrColumn("Iter", 6, BoldFormatter("{:6d}"), StateAttr("iter")))
+    cols += iter_cols(problem)
+
+    cols.append(FilterColumn())
+
+    cols.append(AttrColumn("Func evals", 10, "{:10d}", StateAttr("num_func_evals")))
+    cols.append(AttrColumn("Jac evals", 10, "{:10d}", StateAttr("num_jac_evals")))
+    cols.append(AttrColumn("Steps", 10, "{:10d}", StateAttr("num_steps")))
+    cols.append(AttrColumn("dt", 12, "{:6e}", StateAttr("dt")))
+
+    cols.append(StepTypeColumn())
+
     return Display(cols)
 
 
@@ -189,3 +294,57 @@ def inner_display(problem: Problem, params: Params):
     cols.append(AttrColumn("Active set", 10, "{:10d}", StateAttr("active_set_size")))
 
     return Display(cols)
+
+
+def print_problem_stats(problem, iterate):
+    num_vars = problem.num_vars
+    num_cons = problem.num_cons
+
+    logger.info("Solving problem with %s variables, %s constraints", num_vars, num_cons)
+
+    cons_jac = iterate.cons_jac
+    lag_hess = iterate.lag_hess(iterate.y)
+
+    var_lb = problem.var_lb
+    var_ub = problem.var_ub
+
+    inf_lb = var_lb == -np.inf
+    inf_ub = var_ub == np.inf
+
+    unbounded = np.logical_and(inf_lb, inf_ub)
+    ranged = np.logical_and(np.logical_not(inf_lb), np.logical_not(inf_ub))
+    bounded_below = np.logical_and(np.logical_not(inf_lb), inf_ub)
+    bounded_above = np.logical_and(inf_lb, np.logical_not(inf_ub))
+    bounded = np.logical_or(bounded_below, bounded_above)
+
+    num_unbounded = np.sum(unbounded)
+    num_ranged = np.sum(ranged)
+    num_bounded = np.sum(bounded)
+
+    has_cons = num_cons > 0
+
+    logger.info("           Hessian nnz: %s", lag_hess.nnz)
+    if has_cons:
+        logger.info("          Jacobian nnz: %s", cons_jac.nnz)
+
+    logger.info("        Free variables: %s", num_unbounded)
+    logger.info("     Bounded variables: %s", num_bounded)
+    logger.info("      Ranged variables: %s", num_ranged)
+
+    if not has_cons:
+        return
+
+    cons_lb = problem.cons_lb
+    cons_ub = problem.cons_ub
+
+    equation = cons_lb == cons_ub
+    ranged = np.logical_and(np.isfinite(cons_lb), np.isfinite(cons_ub))
+    ranged = np.logical_and(ranged, np.logical_not(equation))
+
+    num_equations = equation.sum()
+    num_ranged = ranged.sum()
+    num_inequalities = num_cons - num_equations - num_ranged
+
+    logger.info("  Equality constraints: %s", num_equations)
+    logger.info("Inequality constraints: %s", num_inequalities)
+    logger.info("    Ranged constraints: %s", num_ranged)

pygradflow/implicit_func.py

@@ -6,6 +6,7 @@
 
 from pygradflow.iterate import Iterate
 from pygradflow.problem import Problem
+from pygradflow.util import keep_rows
 
 
 class StepFunc(abc.ABC):
@@ -88,24 +89,8 @@ def apply_project_deriv(self, mat: sp.sparse.spmatrix, active_set: np.ndarray):
 
         assert (num_rows,) == lb.shape
 
-        mat = mat.tocoo()
-
         inactive_set = np.logical_not(active_set)
-        inactive_indices = np.where(inactive_set)[0]
-
-        alive_indices = np.isin(mat.row, inactive_indices)
-
-        assert inactive_set[mat.row[alive_indices]].all()
-
-        next_rows = mat.row[alive_indices]
-        next_cols = mat.col[alive_indices]
-        next_entries = mat.data[alive_indices]
-
-        proj_mat = sp.sparse.coo_matrix(
-            (next_entries, (next_rows, next_cols)), mat.shape
-        )
-
-        assert inactive_set[proj_mat.row].all()
+        proj_mat = keep_rows(mat, inactive_set)
 
         return proj_mat
 

pygradflow/integration/events.py

@@ -0,0 +1,47 @@
+from enum import Enum, auto
+
+import numpy as np
+
+
+class EventResultType(Enum):
+    CONVERGED = auto()
+    UNBOUNDED = auto()
+    FILTER_CHANGED = auto()
+    FREE_GRAD_ZERO = auto()
+    PENALTY = auto()
+
+
+class EventResult:
+    def __init__(self, t: float, z: np.ndarray):
+        self.t = t
+        self.z = z
+
+
+class ConvergedResult(EventResult):
+    def __init__(self, t: float, z: np.ndarray):
+        super().__init__(t, z)
+        self.type = EventResultType.CONVERGED
+
+
+class UnboundedResult(EventResult):
+    def __init__(self, t: float, z: np.ndarray):
+        super().__init__(t, z)
+        self.type = EventResultType.UNBOUNDED
+
+
+class FilterChangedResult(EventResult):
+    def __init__(self, t: float, z: np.ndarray, filter: np.ndarray, j: int):
+        super().__init__(t, z)
+        self.type = EventResultType.FILTER_CHANGED
+
+        next_filter = np.copy(filter)
+        (size,) = filter.shape
+        assert 0 <= j < size
+        next_filter[j] = not (filter[j])
+        self.filter = next_filter
+
+
+class PenaltyResult(EventResult):
+    def __init__(self, t: float, z: np.ndarray):
+        super().__init__(t, z)
+        self.type = EventResultType.PENALTY