start functions for inequality constraints

src/raking/inequality_constraints.py

@@ -0,0 +1,71 @@
+"""Module with methods to compute the inequality constraint matrix"""
+
+import numpy as np
+
+
+def inequality_infant_mortality(
+    n1: np.ndarray,
+    n2: np.ndarray,
+    t1: float,
+    t2: float
+) -> tuple[np.ndarray, np.ndarray]:
+    """Compute the constraints matrix C and the inequality vector c for the infant mortality problem.
+
+    We need to define the inequality constraints C beta < c
+    for the infant mortality problem of the Population, Fertility, and Mortality team:
+    Probability of death for 0-to-1-month-olds must be lower than probability of death for 0-to-1-year-olds.
+
+    Parameters
+    ----------
+    n1: np.ndarray
+        Population number for 0-to-1-month-olds
+    n2: np.ndarray
+        Population number for 0-to-1-year-olds
+    t1: float
+        Time interval between 0 and 1 month
+    t2: float
+        Time interval between 0 and 1 year
+
+    Returns
+    -------
+    C : np.ndarray
+        2I * 2I constraints matrix
+    c : np.ndarray
+        length 2I inequality vector vector
+    """
+    assert isinstance(
+        n1, np.ndarray
+    ), "The population number for 0-to-1-month-olds must be a Numpy array."
+    assert (
+        len(n1.shape) == 1
+    ), "The population number for 0-to-1-month-olds must be a 1D Numpy array."
+    assert isinstance(
+        n2, np.ndarray
+    ), "The population number for 0-to-1-month-olds must be a Numpy array."
+    assert (
+        len(n2.shape) == 1
+    ), "The population number for 0-to-1-month-olds must be a 1D Numpy array."
+    assert (
+        len(n1) == len(n2)
+    ), "The two population number arrays must have the same length."
+    assert isinstance(
+        t1, float
+    ), "The time interval between 0 and 1 month must be a float."
+     assert isinstance(
+        t2, float
+    ), "The time interval between 0 and 1 year must be a float."
+   assert (
+        t2 > t1
+    ), "The time interval between 0 and 1 year must be larger than the time interval between 0 and a month."
+
+    I = len(n1)
+    C = np.concatenate(
+        (
+            np.concatenate((np.diag(t1 / n1), np.zeros((I, I))), axis=1),
+            np.concatenate((np.zeros((I, I)), np.diag(- t2 / n2)), axis=1),
+        ),
+        axis=0,
+    )
+    c = np.zeros(2 * I)
+    return (C, c)
+

src/raking/raking_inequality.py

@@ -0,0 +1,29 @@
+"""Module with methods to solve the raking problem with inequality constraints"""
+
+import numpy as np
+
+from scipy.optimize import LinearConstraint
+from scipy.optimize import minimize
+
+def raking_chi2_inequality(
+    y: np.ndarray,
+    A: np.ndarray,
+    s: np.ndarray,
+    C: np.ndarray,
+    c: np.ndarray
+) -> tuple[np.ndarray, np.ndarray]:
+
+    equality = LinearConstraint(A, s, s)
+    inequality = LinearConstraint(C, np.repeat(-np.inf, len(c)), c)
+
+    def distance(x):
+        return np.sum(np.square(x - y) / (2.0 * y))
+    def jacobian(x):
+        return x / y - 1.0
+    def hessian(x):
+        return np.diag(1.0 / y)
+
+    res = minimize(distance, y, method='trust-constr', \
+        jac=jacobian, hess=hessian, constraints=[equality, inequality])
+    return res
+

src/raking/set_inequality_problems.py

@@ -0,0 +1,57 @@
+"""Module with methods to set up the problems with enequality constraints"""
+
+import numpy as np
+
+from raking.compute_constraints import constraints_1D
+
+from inequality_constraints import inequality_infant_mortality
+
+def set_infant_mortality(
+    n1: np.ndarray,
+    n2: np.ndarray,
+    t1: float,
+    t2: float,
+    y1: np.ndarray,
+    y2: np.ndarray,
+    s1: float,
+    s2: float
+) -> tuple[np.ndarray, np.ndarray]:
+    """Set up the optimization problem for the infant mortality problem.
+
+    We need to define the problem:
+        min_beta f(beta, y) s.t. A beta = s and C beta <= c
+
+    Parameters
+    ----------
+    n1: np.ndarray
+        Population number for 0-to-1-month-olds
+    n2: np.ndarray
+        Population number for 0-to-1-year-olds
+    t1: float
+        Time interval between 0 and 1 month
+    t2: float
+        Time interval between 0 and 1 year
+
+    Returns
+    -------
+    y
+    A
+    s
+    C
+    c
+    """
+    I = len(y1
+    y = np.concatenate((y1, y2))
+    (A1, s1) = constraints_1D(s1, I)
+    (A2, s2) = constraints_1D(s2, I)
+    A = np.concatenate(
+        (
+            np.concatenate((A1, np.zeros(I)), axis=1),
+            np.concatenate((np.zeros(I), A2), axis=1),
+        ),
+        axis=0,
+    )
+    s = np.array([s1, s2])
+    (C, c) = inequality_infant_mortality(n1, n2, t1, t2)
+    return (y, A, s, C, c)
+