Harmonize interface of methods and create unified interface

.envs/testenv.yml

@@ -1,10 +1,8 @@
 ---
 name: upper-envelope
-
 channels:
   - conda-forge
   - nodefaults
-
 dependencies:
   - pip
   - setuptools_scm
@@ -24,4 +22,4 @@ dependencies:
 
   # Install locally
   - pip:
-      - -e ../
+      - -e ../

.github/workflows/main.yml

@@ -45,4 +45,4 @@ jobs:
         if: runner.os == 'Linux' && matrix.python-version == '3.10'
         uses: codecov/codecov-action@v3
         with:
-          token: ${{ secrets.CODECOV_TOKEN }}
+          token: ${{ secrets.CODECOV_TOKEN }}

.yamllint.yml

@@ -33,4 +33,4 @@ rules:
   quoted-strings: disable
   trailing-spaces: enable
   truthy:
-    level: warning
+    level: warning

README.md

@@ -4,15 +4,15 @@
 [![Codecov](https://codecov.io/gh/OpenSourceEconomics/upper-envelope/branch/main/graph/badge.svg)](https://app.codecov.io/gh/OpenSourceEconomics/upper-envelope)
 [![Black](https://img.shields.io/badge/code%20style-black-000000.svg)](https://github.com/psf/black)
 
-Extension of the Fast Upper-Envelope Scan (FUES) for solving discrete-continuous dynamic programming problems based on Dobrescu & Shanker (2022). Both
-`jax` and `numba` versions are available.
+Extension of the Fast Upper-Envelope Scan (FUES) for solving discrete-continuous dynamic
+programming problems based on Dobrescu & Shanker (2022). Both `jax` and `numba` versions
+are available.
 
 ## References
 
 1. Iskhakov, Jorgensen, Rust, & Schjerning (2017).
    [The Endogenous Grid Method for Discrete-Continuous Dynamic Choice Models with (or without) Taste Shocks](http://onlinelibrary.wiley.com/doi/10.3982/QE643/full).
    *Quantitative Economics*
 
-
 1. Loretti I. Dobrescu & Akshay Shanker (2022).
-   [Fast Upper-Envelope Scan for Discrete-Continuous Dynamic Programming](https://dx.doi.org/10.2139/ssrn.4181302).
+   [Fast Upper-Envelope Scan for Discrete-Continuous Dynamic Programming](https://dx.doi.org/10.2139/ssrn.4181302).

codecov.yml

@@ -9,7 +9,7 @@ coverage:
   status:
     patch:
       default:
-        target: 70%
+        target: 80%
     project:
       default:
         target: 90%

docs/tutorials/getting_started.ipynb

@@ -0,0 +1,192 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "outputs": [],
+   "source": [
+    "import numpy as np\n",
+    "from upper_envelope.fues_numba.fues_numba import fast_upper_envelope_wrapper\n",
+    "import numba as nb\n",
+    "from collections import namedtuple\n"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "ExecuteTime": {
+     "end_time": "2024-05-24T06:42:04.589484405Z",
+     "start_time": "2024-05-24T06:42:04.167227740Z"
+    }
+   },
+   "id": "c71bb4b54fd1da58",
+   "execution_count": 1
+  },
+  {
+   "cell_type": "code",
+   "outputs": [],
+   "source": [
+    "@nb.njit\n",
+    "def utility_crra(consumption: np.array, params_ntuple) -> np.array:\n",
+    "    \"\"\"Computes the agent's current utility based on a CRRA utility function.\n",
+    "\n",
+    "    Args:\n",
+    "        consumption (jnp.array): Level of the agent's consumption.\n",
+    "            Array of shape (i) (n_quad_stochastic * n_grid_wealth,)\n",
+    "            when called by :func:`~dcgm.call_egm_step.map_exog_to_endog_grid`\n",
+    "            and :func:`~dcgm.call_egm_step.get_next_period_value`, or\n",
+    "            (ii) of shape (n_grid_wealth,) when called by\n",
+    "            :func:`~dcgm.call_egm_step.get_current_period_value`.\n",
+    "        choice (int): Choice of the agent, e.g. 0 = \"retirement\", 1 = \"working\".\n",
+    "        params_dict (dict): Dictionary containing model parameters.\n",
+    "            Relevant here is the CRRA coefficient theta.\n",
+    "\n",
+    "    Returns:\n",
+    "        utility (jnp.array): Agent's utility . Array of shape\n",
+    "            (n_quad_stochastic * n_grid_wealth,) or (n_grid_wealth,).\n",
+    "\n",
+    "    \"\"\"\n",
+    "    utility_consumption = (consumption ** (1 - params_ntuple.rho) - 1) / (\n",
+    "        1 - params_ntuple.rho\n",
+    "    )\n",
+    "\n",
+    "    utility = utility_consumption - (1 - params_ntuple.choice) * params_ntuple.delta\n",
+    "\n",
+    "    return utility"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "ExecuteTime": {
+     "end_time": "2024-05-24T06:42:04.611236211Z",
+     "start_time": "2024-05-24T06:42:04.609614665Z"
+    }
+   },
+   "id": "ab4b0dcb970dac41",
+   "execution_count": 2
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "initial_id",
+   "metadata": {
+    "collapsed": true,
+    "ExecuteTime": {
+     "end_time": "2024-05-24T06:42:04.613370376Z",
+     "start_time": "2024-05-24T06:42:04.609954700Z"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "max_wealth = 50\n",
+    "n_grid_wealth = 500\n",
+    "exog_savings_grid = np.linspace(0, max_wealth, n_grid_wealth)\n",
+    "\n",
+    "beta = 0.95  # discount_factor\n",
+    "\n",
+    "utility_kwargs = {\n",
+    "    \"choice\": 0,\n",
+    "    \"rho\": 1.95,\n",
+    "    \"delta\": 0.35,\n",
+    "}"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "outputs": [],
+   "source": [
+    "resource_dir = \"../../tests/resources/\"\n",
+    "value_egm = np.genfromtxt(\n",
+    "    resource_dir + \"upper_envelope_period_tests/val2.csv\",\n",
+    "    delimiter=\",\",\n",
+    ")\n",
+    "policy_egm = np.genfromtxt(\n",
+    "    resource_dir + \"upper_envelope_period_tests/pol2.csv\",\n",
+    "    delimiter=\",\",\n",
+    ")\n",
+    "n_constrained_points_to_add = int(0.1 * len(policy_egm[0]))\n",
+    "n_final_wealth_grid = int(1.2 * (len(policy_egm[0])))\n",
+    "tuning_params = {\n",
+    "    \"n_final_wealth_grid\": n_final_wealth_grid,\n",
+    "    \"jump_thresh\": 2,\n",
+    "    \"n_constrained_points_to_add\": n_constrained_points_to_add,\n",
+    "    \"n_points_to_scan\": 10,\n",
+    "}"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "ExecuteTime": {
+     "end_time": "2024-05-24T06:42:04.614201145Z",
+     "start_time": "2024-05-24T06:42:04.610074325Z"
+    }
+   },
+   "id": "61b7b971c6f81237",
+   "execution_count": 4
+  },
+  {
+   "cell_type": "code",
+   "outputs": [],
+   "source": [
+    "\n",
+    "utility_ntuple = namedtuple(\"utility_params\", utility_kwargs.keys())(\n",
+    "    *utility_kwargs.values()\n",
+    ")\n",
+    "tuning_params_tuple = namedtuple(\"tunings\", tuning_params.keys())(\n",
+    "    *tuning_params.values()\n",
+    ")"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "ExecuteTime": {
+     "end_time": "2024-05-24T06:42:04.614607376Z",
+     "start_time": "2024-05-24T06:42:04.610206856Z"
+    }
+   },
+   "id": "415a872b98d65631",
+   "execution_count": 5
+  },
+  {
+   "cell_type": "code",
+   "outputs": [],
+   "source": [
+    "endog_grid_refined, policy_refined, value_refined = fast_upper_envelope_wrapper(\n",
+    "    endog_grid=policy_egm[0, 1:],\n",
+    "    policy=policy_egm[1, 1:],\n",
+    "    value=value_egm[1, 1:],\n",
+    "    expected_value_zero_savings=value_egm[1, 0],\n",
+    "    exog_grid=exog_savings_grid,\n",
+    "    utility_function=utility_crra,\n",
+    "    utility_kwargs=utility_ntuple,\n",
+    "    discount_factor=beta,\n",
+    "    tuning_params=tuning_params_tuple,\n",
+    ")"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "ExecuteTime": {
+     "end_time": "2024-05-24T06:42:14.783923576Z",
+     "start_time": "2024-05-24T06:42:04.611563430Z"
+    }
+   },
+   "id": "df76ba014c57934d",
+   "execution_count": 6
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 2
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython2",
+   "version": "2.7.6"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

environment.yml

@@ -1,10 +1,8 @@
 ---
 name: upper-envelope
-
 channels:
   - conda-forge
   - defaults
-
 dependencies:
   - python=3.10
   - pip
@@ -37,8 +35,7 @@ dependencies:
   - conda-build
   - conda-verify
   - tox-conda
-
   - pip:
       - blackcellmagic
       - furo
-      - -e .  # Install locally
+      - -e .  # Install locally

pyproject.toml

@@ -7,6 +7,7 @@ write_to = "src/upper_envelope/_version.py"
 [tool.ruff]
 target-version = "py310"
 fix = true
+ignore = ["F401"]
 
 [tool.yamlfix]
 line_length = 88

setup.py

@@ -1,4 +1,4 @@
 from setuptools import setup
 
 if __name__ == "__main__":
-    setup()
+    setup()

src/upper_envelope/__init__.py

@@ -0,0 +1,4 @@
+from upper_envelope.fues_jax.fues_jax import fues_jax
+from upper_envelope.fues_jax.fues_jax import fues_jax_unconstrained
+from upper_envelope.fues_numba.fues_numba import fues_numba
+from upper_envelope.fues_numba.fues_numba import fues_numba_unconstrained