diff --git a/CHANGELOG.md b/CHANGELOG.md index bbbd74e62..d3314628b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,12 @@ All notable changes to this project will be documented in this file. The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). +## [0.11.12] - 2024-07-26 01:00:00 + +### Bug Fix + +- Fixes extrapolation of nested lists of tax function parameters. + ## [0.11.11] - 2024-06-24 01:00:00 ### Added diff --git a/ogcore/__init__.py b/ogcore/__init__.py index a7a60577e..7039682e2 100644 --- a/ogcore/__init__.py +++ b/ogcore/__init__.py @@ -20,4 +20,4 @@ from ogcore.txfunc import * from ogcore.utils import * -__version__ = "0.11.11" +__version__ = "0.11.12" diff --git a/ogcore/parameters.py b/ogcore/parameters.py index db9a59698..f900d5bd9 100644 --- a/ogcore/parameters.py +++ b/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,7 +187,7 @@ 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) @@ -191,7 +195,7 @@ def compute_default_params(self): 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) diff --git a/ogcore/utils.py b/ogcore/utils.py index a88003a80..1bde6c585 100644 --- a/ogcore/utils.py +++ b/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. diff --git a/setup.py b/setup.py index ad1c98168..70e57cab1 100755 --- a/setup.py +++ b/setup.py @@ -5,7 +5,7 @@ setuptools.setup( name="ogcore", - version="0.11.11", + version="0.11.12", author="Jason DeBacker and Richard W. Evans", license="CC0 1.0 Universal (CC0 1.0) Public Domain Dedication", description="A general equilibribum overlapping generations model for fiscal policy analysis", diff --git a/tests/test_utils.py b/tests/test_utils.py index d183ef122..b73284ba2 100644 --- a/tests/test_utils.py +++ b/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", "S", "