Extrapolating nested lists

ogcore/parameters.py

@@ -4,7 +4,11 @@
 import paramtools
 import ogcore
 from ogcore import elliptical_u_est
-from ogcore.utils import rate_conversion, extrapolate_arrays
+from ogcore.utils import (
+    rate_conversion,
+    extrapolate_array,
+    extrapolate_nested_list,
+)
 from ogcore.constants import BASELINE_DIR
 
 CURRENT_PATH = os.path.abspath(os.path.dirname(__file__))
@@ -170,7 +174,7 @@ def compute_default_params(self):
         ]
         for item in tp_param_list:
             param_in = getattr(self, item)
-            param_out = extrapolate_arrays(
+            param_out = extrapolate_array(
                 param_in, dims=(self.T + self.S,), item=item
             )
             setattr(self, item, param_out)
@@ -183,15 +187,15 @@ def compute_default_params(self):
         ]
         for item in tp_param_list2:
             param_in = getattr(self, item)
-            param_out = extrapolate_arrays(
+            param_out = extrapolate_array(
                 param_in, dims=(self.T + self.S, self.M), item=item
             )
             setattr(self, item, param_out)
         # Deal with parameters that vary across consumption good and over time
         tp_param_list3 = ["tau_c"]
         for item in tp_param_list3:
             param_in = getattr(self, item)
-            param_out = extrapolate_arrays(
+            param_out = extrapolate_array(
                 param_in, dims=(self.T + self.S, self.I), item=item
             )
             setattr(self, item, param_out)
@@ -202,7 +206,7 @@ def compute_default_params(self):
         ]
         for item in tp_param_list3:
             param_in = getattr(self, item)
-            param_out = extrapolate_arrays(
+            param_out = extrapolate_array(
                 param_in, dims=(self.T + self.S, self.J), item=item
             )
             setattr(self, item, param_out)
@@ -212,7 +216,7 @@ def compute_default_params(self):
         ]
         for item in tp_param_list4:
             param_in = getattr(self, item)
-            param_out = extrapolate_arrays(
+            param_out = extrapolate_array(
                 param_in, dims=(self.T + self.S, self.S), item=item
             )
             setattr(self, item, param_out)
@@ -223,70 +227,38 @@ def compute_default_params(self):
             "mtry_params",
         ]
         for item in tax_params_to_TP:
-            tax_to_set = getattr(self, item)
+            tax_to_set_in = getattr(self, item)
             try:
-                tax_to_set = (
-                    tax_to_set.tolist()
-                )  # in case parameters are numpy arrays
-            except AttributeError:  # catches if they are lists already
-                pass
-            if len(tax_to_set) == 1 and isinstance(tax_to_set[0], float):
-                setattr(
-                    self,
-                    item,
-                    [
-                        [[tax_to_set] for i in range(self.S)]
-                        for t in range(self.T)
-                    ],
-                )
-            elif any(
-                [
-                    isinstance(tax_to_set[i][j], list)
-                    for i, v in enumerate(tax_to_set)
-                    for j, vv in enumerate(tax_to_set[i])
-                ]
-            ):
-                if len(tax_to_set) > self.T + self.S:
-                    tax_to_set = tax_to_set[: self.T + self.S]
-                if len(tax_to_set) < self.T + self.S:
-                    tax_params_to_add = [tax_to_set[-1]] * (
-                        self.T + self.S - len(tax_to_set)
-                    )
-                    tax_to_set.extend(tax_params_to_add)
-                if len(tax_to_set[0]) > self.S:
-                    for t, v in enumerate(tax_to_set):
-                        tax_to_set[t] = tax_to_set[t][: self.S]
-                if len(tax_to_set[0]) < self.S:
-                    tax_params_to_add = [tax_to_set[:][-1]] * (
-                        self.S - len(tax_to_set[0])
-                    )
-                    tax_to_set[0].extend(tax_params_to_add)
-                setattr(self, item, tax_to_set)
-            else:
+                len(tax_to_set_in[0][0])
+            except TypeError:
                 print(
                     "please give a "
                     + item
-                    + " that is a single element or nested lists of"
+                    + " that is a nested lists of"
                     + " lists that is three lists deep"
                 )
                 assert False
+            tax_to_set_out = extrapolate_nested_list(
+                tax_to_set_in, dims=(self.T, self.S, len(tax_to_set_in[0][0]))
+            )
+            setattr(self, item, tax_to_set_out)
 
         # Try to deal with size of eta.  It may vary by S, J, T, but
         # want to allow user to enter one that varies by only S, S and J,
         # S and T, or T and S and J.
         param_in = getattr(self, "eta")
-        param_out = extrapolate_arrays(
+        param_out = extrapolate_array(
             param_in, dims=(self.T + self.S, self.S, self.J), item="eta"
         )
         setattr(self, "eta", param_out)
         param_in = getattr(self, "e")
-        param_out = extrapolate_arrays(
+        param_out = extrapolate_array(
             param_in, dims=(self.T, self.S, self.J), item="e"
         )
         setattr(self, "e", param_out)
         # Extrapolate chi_n over T + S
         param_in = getattr(self, "chi_n")
-        param_out = extrapolate_arrays(
+        param_out = extrapolate_array(
             param_in, dims=(self.T + self.S, self.S), item="chi_n"
         )
         setattr(self, "chi_n", param_out)

ogcore/utils.py

@@ -869,7 +869,7 @@ def avg_by_bin(x, y, weights=None, bins=10, eql_pctl=True):
     return x_binned, y_binned, weights_binned
 
 
-def extrapolate_arrays(param_in, dims=None, item="Parameter Name"):
+def extrapolate_array(param_in, dims=None, item="Parameter Name"):
     """
     Extrapolates input values to fit model dimensions. Using this allows
     users to input smaller dimensional arrays and have the model infer
@@ -1022,6 +1022,53 @@ def extrapolate_arrays(param_in, dims=None, item="Parameter Name"):
     return param_out
 
 
+def extrapolate_nested_list(list_in, dims=(400, 80, 1)):
+    """
+    Function to extrapolate a nested list to a specified size.
+
+    Currently only set up for 3 deep nested lists, but could be
+    generalized to deeper or shallower lists.
+
+    Args:
+        list_in (list): list to extrapolate
+        dims (tuple): dimensions of the output list
+
+    Returns:
+        list_out (list): extrapolated list
+    """
+    T, S, num_params = dims
+    try:
+        list_in = list_in.tolist()  # in case parameters are numpy arrays
+    except AttributeError:  # catches if they are lists already
+        pass
+    assert isinstance(list_in, list), "please give a list"
+
+    def depth(L):
+        return isinstance(L, list) and max(map(depth, L)) + 1
+
+    # for now, just have this work for 3 deep lists since
+    # the only OG-Core use case is for tax function parameters
+    assert depth(list_in) == 3, "please give a list that is three lists deep"
+    assert depth(list_in) == len(
+        dims
+    ), "please make sure the depth of nested list is equal to the length of dims to extrapolate"
+    # Extrapolate along the first dimension
+    if len(list_in) > T + S:
+        list_in = list_in[: T + S]
+    if len(list_in) < T + S:
+        params_to_add = [list_in[-1]] * (T + S - len(list_in))
+        list_in.extend(params_to_add)
+    # Extrapolate along the second dimension
+    for t in range(len(list_in)):
+        if len(list_in[t]) > S:
+            list_in[t] = list_in[t][:S]
+        if len(list_in[t]) < S:
+            params_to_add = [list_in[t][-1]] * (S - len(list_in[t]))
+            list_in[t].extend(params_to_add)
+
+    return list_in
+
+
 class CustomHttpAdapter(requests.adapters.HTTPAdapter):
     """
     The UN Data Portal server doesn't support "RFC 5746 secure renegotiation". This causes and error when the client is using OpenSSL 3, which enforces that standard by default.

tests/test_utils.py

@@ -790,15 +790,63 @@ def test_avg_by_bin():
     test_data,
     ids=["2D, scalar in", "2D, 1D in", "2D in", "1D out", "3D out"],
 )
-def test_extrapolate_arrays(param_in, dims, expected):
+def test_extrapolate_array(param_in, dims, expected):
     """
-    Test of the utils.extrapolate_arrays function
+    Test of the utils.extrapolate_array function
     """
-    test_value = utils.extrapolate_arrays(param_in, dims=dims)
+    test_value = utils.extrapolate_array(param_in, dims=dims)
 
     assert np.allclose(test_value, expected)
 
 
+T_L = 20
+S_L = 4
+list1 = [[[3]]] * (T_L + 10)
+list2 = [[[3]]] * (T_L - 6)
+list3 = [[[3]] * (S_L + 2)]
+list4 = [[[3]] * (S_L - 2)]
+list5 = np.ones((18, 2, 3))
+test_data = [
+    (list1, (T_L, S_L, 1)),
+    (list2, (T_L, S_L, 1)),
+    (list3, (T_L, S_L, 1)),
+    (list4, (T_L, S_L, 1)),
+    (list5, (T_L, S_L, 3)),
+]
+
+
+@pytest.mark.parametrize(
+    "list_in,dims",
+    test_data,
+    ids=[" > T + S", "<T", ">S", "<S", "Numpy array"],
+)
+def test_extrapolate_nested_list(list_in, dims):
+    """
+    Test of the utils.extrapolate_nested_list function
+    """
+
+    test_value = utils.extrapolate_nested_list(list_in, dims=dims)
+
+    min_S = dims[1]
+    max_S = dims[1]
+    min_p = dims[2]
+    max_p = dims[2]
+    for t in range(len(test_value)):
+        min_S = min(min_S, len(test_value[t]))
+        max_S = max(max_S, len(test_value[t]))
+        for s in range(len(test_value[t])):
+            min_p = min(len(test_value[t][s]), min_p)
+            max_p = max(len(test_value[t][s]), max_p)
+
+    assert len(test_value) == dims[0] + dims[1]
+    assert len(test_value[0]) == dims[1]
+    assert len(test_value[0][0]) == dims[2]
+    assert min_S == dims[1]
+    assert max_S == dims[1]
+    assert min_p == dims[2]
+    assert max_p == dims[2]
+
+
 arr1 = np.array([1.0, 2.0, 2.0, 3.0])
 expected_array1 = np.tile(arr1.reshape(1, 4), (20, 1))
 expected_array2 = expected_array1.copy()