ELEX-3031: Potential Voterflow Solvers

.github/workflows/test.yml

@@ -4,10 +4,10 @@ jobs:
   test:
     name: Run unit tests
     runs-on: ubuntu-latest
-    timeout-minutes: 5
+    timeout-minutes: 15
     strategy:
       matrix:
-        python-version: ['3.10']
+        python-version: ['3.11']
     steps:
       - uses: actions/checkout@v2
       - name: Setup Python

README.md

@@ -16,8 +16,8 @@ We have our own implementation of ordinary least squares in Python because this
 ## Quantile Regression
 Since we did not find any implementations of quantile regression in Python that fit our needs, we decided to write one ourselves. At the moment this uses two libraries, the version that solves the non-regularized problem uses `numpy`and solves the dual based on [this](https://arxiv.org/pdf/2305.12616.pdf) paper. The version that solves the regularized problem uses [`cvxpy`](https://www.cvxpy.org/#) and sets up the problem as a normal optimization problem. Eventually, we are planning on replacing the regularized version with the dual also.
 
-## Transition matrices
-We also have a solver for transition matrices. While this works arbitrarily, we have used this in the past for our primary election night model. We can still use this to create the sankey diagram coefficients.
+## Transition Matrices
+We also have a matrix regression solver built with `cvxpy`.  We've used this for our primary election model and analysis.  The transitions it generates form the transitions displayed in our sankey diagrams.
 
 ## Development
 We welcome contributions to this repo. Please open a Github issue for any issues or comments you have.

setup.cfg

@@ -2,4 +2,8 @@
 universal = 1
 
 [pycodestyle]
-max-line-length = 160
+max-line-length = 160
+
+[pylint]
+max-line-length = 160
+disable = invalid-name, duplicate-code, missing-function-docstring, too-many-instance-attributes, too-many-arguments, too-many-locals, missing-module-docstring

setup.py

@@ -3,7 +3,7 @@
 
 from setuptools import find_packages, setup
 
-INSTALL_REQUIRES = ["cvxpy~=1.4", "numpy~=1.26", "scipy~=1.11"]
+INSTALL_REQUIRES = ["cvxpy~=1.4", "numpy~=1.26", "scipy~=1.12"]
 
 THIS_FILE_DIR = os.path.dirname(__file__)
 
@@ -29,7 +29,7 @@
         "Intended Audience :: Developers",
         "License :: OSI Approved :: MIT License",
         "Programming Language :: Python",
-        "Programming Language :: Python :: 3.10",
+        "Programming Language :: Python :: 3.11",
     ],
     description="A package for optimization solvers",
     long_description=LONG_DESCRIPTION,

src/elexsolver/TransitionMatrixSolver.py

@@ -1,28 +1,107 @@
+import logging
+import warnings
+
 import cvxpy as cp
+import numpy as np
+
+from elexsolver.logging import initialize_logging
+from elexsolver.TransitionSolver import TransitionSolver
+
+initialize_logging()
+
+LOG = logging.getLogger(__name__)
+
 
+class TransitionMatrixSolver(TransitionSolver):
+    """
+    Matrix regression transition solver using CVXPY.
+    """
 
-class TransitionMatrixSolver:
-    def __init__(self):
-        self.transition_matrix = None
+    def __init__(self, strict: bool = True, lam: float | None = None):
+        """
+        Parameters
+        ----------
+        `strict` : bool, default True
+            If True, solution will be constrainted so that all coefficients are >= 0,
+            <= 1, and the sum of each row equals 1.
+        `lam` : float, optional
+            `lam` != 0 will enable L2 regularization (Ridge).
+        """
+        super().__init__()
+        self._strict = strict
+        self._lambda = lam
 
     @staticmethod
-    def __get_constraint(X, strict):
+    def __get_constraints(coef: np.ndarray, strict: bool) -> list:
         if strict:
-            return [cp.sum(X, axis=1) == 1]
-        return [cp.sum(X, axis=1) <= 1.1, cp.sum(X, axis=1) >= 0.9]
-
-    def __solve(self, A, B, strict):
-        transition_matrix = cp.Variable((A.shape[1], B.shape[1]))
-        loss_function = cp.norm(A @ transition_matrix - B, "fro")
-        objective = cp.Minimize(loss_function)
-        constraint = TransitionMatrixSolver.__get_constraint(transition_matrix, strict)
-        problem = cp.Problem(objective, constraint)
-        problem.solve()
+            return [0 <= coef, coef <= 1, cp.sum(coef, axis=1) == 1]
+        return [cp.sum(coef, axis=1) <= 1.1, cp.sum(coef, axis=1) >= 0.9]
+
+    def __standard_objective(self, A: np.ndarray, B: np.ndarray, beta: np.ndarray) -> cp.Minimize:
+        loss_function = cp.norm(A @ beta - B, "fro")
+        return cp.Minimize(loss_function)
+
+    def __ridge_objective(self, A: np.ndarray, B: np.ndarray, beta: np.ndarray) -> cp.Minimize:
+        # Based on https://www.cvxpy.org/examples/machine_learning/ridge_regression.html
+        lam = cp.Parameter(nonneg=True, value=self._lambda)
+        loss_function = cp.pnorm(A @ beta - B, p=2) ** 2
+        regularizer = cp.pnorm(beta, p=2) ** 2
+        return cp.Minimize(loss_function + lam * regularizer)
+
+    def __solve(self, A: np.ndarray, B: np.ndarray, weights: np.ndarray) -> np.ndarray:
+        transition_matrix = cp.Variable((A.shape[1], B.shape[1]), pos=True)
+        Aw = np.dot(weights, A)
+        Bw = np.dot(weights, B)
+
+        if self._lambda is None or self._lambda == 0:
+            objective = self.__standard_objective(Aw, Bw, transition_matrix)
+        else:
+            objective = self.__ridge_objective(Aw, Bw, transition_matrix)
+
+        constraints = TransitionMatrixSolver.__get_constraints(transition_matrix, self._strict)
+        problem = cp.Problem(objective, constraints)
+
+        with warnings.catch_warnings():
+            warnings.simplefilter("error")
+            try:
+                problem.solve(solver=cp.CLARABEL)
+            except (UserWarning, cp.error.SolverError) as e:
+                raise RuntimeError(e) from e
+
         return transition_matrix.value
 
-    def fit(self, A, B, strict=False):
-        transition_matrix = self.__solve(A, B, strict)
-        self.transition_matrix = transition_matrix
+    def fit_predict(self, X: np.ndarray, Y: np.ndarray, weights: np.ndarray | None = None) -> np.ndarray:
+        self._check_data_type(X)
+        self._check_data_type(Y)
+        self._check_any_element_nan_or_inf(X)
+        self._check_any_element_nan_or_inf(Y)
+
+        # matrices should be (units x things), where the number of units is > the number of things
+        if X.shape[1] > X.shape[0]:
+            X = X.T
+        if Y.shape[1] > Y.shape[0]:
+            Y = Y.T
+
+        if X.shape[0] != Y.shape[0]:
+            raise ValueError(f"Number of units in X ({X.shape[0]}) != number of units in Y ({Y.shape[0]}).")
+
+        self._check_dimensions(X)
+        self._check_dimensions(Y)
+        self._check_for_zero_units(X)
+        self._check_for_zero_units(Y)
+
+        if not isinstance(X, np.ndarray):
+            X = X.to_numpy()
+        if not isinstance(Y, np.ndarray):
+            Y = Y.to_numpy()
+
+        X_expected_totals = X.sum(axis=0) / X.sum(axis=0).sum()
+
+        X = self._rescale(X)
+        Y = self._rescale(Y)
+
+        weights = self._check_and_prepare_weights(X, Y, weights)
 
-    def predict(self, A):
-        return A @ self.transition_matrix
+        percentages = self.__solve(X, Y, weights)
+        self._transitions = np.diag(X_expected_totals) @ percentages
+        return percentages

src/elexsolver/TransitionSolver.py

@@ -0,0 +1,111 @@
+import logging
+import warnings
+from abc import ABC
+
+import numpy as np
+
+from elexsolver.logging import initialize_logging
+
+initialize_logging()
+
+LOG = logging.getLogger(__name__)
+
+
+class TransitionSolver(ABC):
+    """
+    Abstract class for transition solvers.
+    """
+
+    def __init__(self):
+        self._transitions = None
+
+    def fit_predict(self, X: np.ndarray, Y: np.ndarray, weights: np.ndarray | None = None):
+        """
+        After this method finishes, transitions will be available in the `transitions` class member.
+
+        Parameters
+        ----------
+        `X` : np.ndarray matrix or pandas.DataFrame of int
+            Must have the same number of rows as `Y` but can have any number of columns greater than the number of rows.
+        `Y` : np.ndarray matrix or pandas.DataFrame of int
+            Must have the same number of rows as `X` but can have any number of columns greater than the number of rows.
+        `weights` : list, np.ndarray, or pandas.Series of int, optional
+            Must have the same length (number of rows) as both `X` and `Y`.
+
+        Returns
+        -------
+        np.ndarray matrix of float of shape (number of columns in `X`) x (number of columns in `Y`).
+        Each float represents the percent of how much of row x is part of column y.
+        """
+        raise NotImplementedError
+
+    def get_prediction_interval(self, pi: float):
+        raise NotImplementedError
+
+    @property
+    def transitions(self) -> np.ndarray:
+        return self._transitions
+
+    def _check_any_element_nan_or_inf(self, A: np.ndarray):
+        """
+        Check whether any element in a matrix or vector is NaN or infinity
+        """
+        if np.any(np.isnan(A)) or np.any(np.isinf(A)):
+            raise ValueError("Matrix contains NaN or Infinity.")
+
+    def _check_data_type(self, A: np.ndarray):
+        if not np.all(A.astype("int64") == A):
+            raise ValueError("Matrix must contain integers.")
+
+    def _check_dimensions(self, A: np.ndarray):
+        """
+        Ensure that in our (units x things) matrix, the number of units is
+        at least twice as large as the number of things.
+        """
+        if A.shape[0] <= A.shape[1] or (A.shape[0] // 2) <= A.shape[1]:
+            raise ValueError(f"Not enough units ({A.shape[0]}) relative to the number of things ({A.shape[1]}).")
+
+    def _check_for_zero_units(self, A: np.ndarray):
+        """
+        If we have at least one unit whose columns are all zero, we can't continue.
+        """
+        if np.any(np.sum(A, axis=1) == 0):
+            raise ValueError("Matrix cannot contain any rows (units) where all columns (things) are zero.")
+
+    def _rescale(self, A: np.ndarray) -> np.ndarray:
+        """
+        Rescale rows (units) to ensure they sum to 1 (100%).
+        """
+        A = A.copy().astype(float)
+
+        with warnings.catch_warnings():
+            warnings.filterwarnings("ignore", category=RuntimeWarning, message="invalid value encountered in divide")
+            A = (A.T / A.sum(axis=1)).T
+
+        return np.nan_to_num(A, nan=0, posinf=0, neginf=0)
+
+    def _check_and_prepare_weights(self, X: np.ndarray, Y: np.ndarray, weights: np.ndarray | None) -> np.ndarray:
+        """
+        If `weights` is not None, and `weights` has the same number of rows in both matrices `X` and `Y`,
+        we'll rescale the weights by taking the square root after dividing them by their sum,
+        then return a diagonal matrix containing these now-normalized weights.
+        If `weights` is None, return a diagonal matrix of ones.
+
+        Parameters
+        ----------
+        `X` : np.ndarray matrix of int (same number of rows as `Y`)
+        `Y` : np.ndarray matrix of int (same number of rows as `X`)
+        `weights` : np.ndarray of int of the shape (number of rows in `X` and `Y`, 1), optional
+        """
+
+        if weights is not None:
+            if len(weights) != X.shape[0] and len(weights) != Y.shape[0]:
+                raise ValueError("weights must be the same length as the number of rows in X and Y.")
+            if isinstance(weights, list):
+                weights = np.array(weights).copy()
+            elif not isinstance(weights, np.ndarray):
+                # pandas.Series
+                weights = weights.values.copy()
+            return np.diag(np.sqrt(weights.flatten() / weights.sum()))
+
+        return np.diag(np.ones((Y.shape[0],)))