moving some functions for general model use into a utils module

salamanca/models/calibrate_ineq.py

@@ -9,44 +9,7 @@
 
 
 from salamanca import ineq
-
-#
-# Pyomo-enabled helper functions
-#
-
-
-def below_threshold(x, i, t, mean=True):
-    """Compute the CDF of the lognormal distribution at x using an approximation of
-    the error function:
-
-    .. math::
-
-        erf(x) \approx \tanh(\sqrt \pi \log x)
-
-    Parameters
-    ----------
-    x : numeric
-       threshold income
-    i : numeric
-       mean income (per capita)
-    t : numeric or Pyomo variable
-       theil coefficient
-    mean : bool, default: True
-       treat income as mean income
-    """
-    # see
-    # https://math.stackexchange.com/questions/321569/approximating-the-error-function-erf-by-analytical-functions
-    sigma2 = 2 * t  # t is var
-    # f(var), adjust for mean income vs. median
-    mu = math.log(i)
-    if mean:
-        mu -= sigma2 / 2
-    # f(var), argument for error function
-    arg = (math.log(x) - mu) / mo.sqrt(2 * sigma2)
-    # coefficient for erf approximation
-    k = math.pi ** 0.5 * math.log(2)
-    # definition of cdf with tanh(kx) approximating erf(x)
-    return 0.5 + 0.5 * mo.tanh(k * arg)
+from salamanca.models.utils import below_threshold, model_T_w
 
 
 #
@@ -111,9 +74,8 @@ def l2_norm_obj(m):
 
 
 def theil_sum_obj(m):
-    _t_w = lambda m, idx: m.data['i'][idx] * \
-        m.data['n'][idx] / m.data['G'] * m.t[idx]
-    return (m.data['T_w'] - sum(_t_w(m, idx) for idx in m.idxs)) ** 2
+    T_w = model_T_w(m, income_from_data=True)
+    return (m.data['T_w'] - T_w) ** 2
 
 
 def threshold_obj(m, factors=[1.0], weights=[1.0], relative=True):

salamanca/models/project.py

@@ -0,0 +1,28 @@
+from __future__ import division
+
+import itertools
+import math
+
+import numpy as np
+import pandas as pd
+import pyomo.environ as mo
+
+
+from salamanca import ineq
+from salamanca.models.utils import below_threshold
+
+#
+# Constraints
+#
+
+
+#
+# Objectives
+#
+
+def theil_total_sum_obj(m):
+    _t_w = lambda m, idx: m.i[idx] * \
+        m.data['n'][idx] / m.data['G'] * m.t[idx]
+    T_w = sum(_t_w(m, idx) for idx in m.idxs)
+    T_b = sum(_t_b(m, idx) for idx in m.idxs)
+    return (m.data['T'] - T_w - T_b) ** 2

salamanca/models/utils.py

@@ -0,0 +1,50 @@
+import math
+
+import pyomo.environ as mo
+
+
+#
+# Pyomo-enabled helper functions
+#
+
+
+def below_threshold(x, i, t, mean=True):
+    """Compute the CDF of the lognormal distribution at x using an approximation of
+    the error function:
+
+    .. math::
+
+        erf(x) \approx \tanh(\sqrt \pi \log x)
+
+    Parameters
+    ----------
+    x : numeric
+       threshold income
+    i : numeric
+       mean income (per capita)
+    t : numeric or Pyomo variable
+       theil coefficient
+    mean : bool, default: True
+       treat income as mean income
+    """
+    # see
+    # https://math.stackexchange.com/questions/321569/approximating-the-error-function-erf-by-analytical-functions
+    sigma2 = 2 * t  # t is var
+    # f(var), adjust for mean income vs. median
+    mu = math.log(i)
+    if mean:
+        mu -= sigma2 / 2
+    # f(var), argument for error function
+    arg = (math.log(x) - mu) / mo.sqrt(2 * sigma2)
+    # coefficient for erf approximation
+    k = math.pi ** 0.5 * math.log(2)
+    # definition of cdf with tanh(kx) approximating erf(x)
+    return 0.5 + 0.5 * mo.tanh(k * arg)
+
+
+def model_T_w(m, income_from_data=True):
+    def _t_w(m, idx):
+        i = m.data['i'] if income_from_data else m.i
+        return i[idx] * m.data['n'][idx] / m.data['G'] * m.t[idx]
+    T_w = sum(_t_w(m, idx) for idx in m.idxs)
+    return T_w