JAx and numba new (#21)

Co-authored-by: Sebastian Gsell <sebastian.gsell93@gmail.com>

Author: MaxBlesch

.pre-commit-config.yaml

@@ -6,7 +6,7 @@ repos:
       - id: check-useless-excludes
     # -   id: identity  # Prints all files passed to pre-commits. Debugging.
   - repo: https://github.com/adrienverge/yamllint.git
-    rev: v1.37.1
+    rev: v1.38.0
     hooks:
       - id: yamllint
   - repo: https://github.com/lyz-code/yamlfix
@@ -42,7 +42,7 @@ repos:
       - id: python-use-type-annotations
       - id: text-unicode-replacement-char
   - repo: https://github.com/pycqa/isort
-    rev: 7.0.0
+    rev: 8.0.1
     hooks:
       - id: isort
         name: isort
@@ -55,13 +55,13 @@ repos:
   #       args:
   #         - --py37-plus
   - repo: https://github.com/psf/black-pre-commit-mirror
-    rev: 25.12.0
+    rev: 26.1.0
     hooks:
       - id: black
         language_version: python3.12
         exclude: tests/utils/fast_upper_envelope_org.py
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.14.10
+    rev: v0.15.4
     hooks:
       - id: ruff
         # exclude: |

docs/time_period2_ops.py

@@ -8,7 +8,8 @@
 import numpy as np
 from numba import njit
 
-import upper_envelope as upenv
+import upper_envelope.jax as ue_jax
+import upper_envelope.numba as ue_numba
 
 jax.config.update("jax_enable_x64", True)
 
@@ -60,7 +61,7 @@ def value_func_jax(consumption, choice, params):
 
 
 def fues_jax_partial(endog, pol, val, exp_val_zero):
-    return upenv.fues_jax(
+    return ue_jax.fues_jax(
         endog_grid=jnp.asarray(endog),
         policy=jnp.asarray(pol),
         value=jnp.asarray(val),
@@ -103,7 +104,7 @@ def fues_jax_partial(endog, pol, val, exp_val_zero):
 
 
 def drued_jorg_jax_partial(endog, pol, val, m_grid, exp_val_zero):
-    return upenv.drued_jorg_jax(
+    return ue_jax.drued_jorg_jax(
         endog_grid=endog,
         policy=pol,
         value=val,
@@ -168,7 +169,7 @@ def drued_jorg_jax_partial(endog, pol, val, m_grid, exp_val_zero):
 # Numba FUES
 start = time.time()
 jax.block_until_ready(
-    upenv.fues_numba(
+    ue_numba.fues_numba(
         endog_grid=policy_egm[0, 1:],
         policy=policy_egm[1, 1:],
         value=value_egm[1, 1:],
@@ -184,7 +185,7 @@ def drued_jorg_jax_partial(endog, pol, val, m_grid, exp_val_zero):
 for _ in range(n_runs):
     start = time.time()
     jax.block_until_ready(
-        upenv.fues_numba(
+        ue_numba.fues_numba(
             endog_grid=policy_egm[0, 1:],
             policy=policy_egm[1, 1:],
             value=value_egm[1, 1:],
@@ -201,7 +202,7 @@ def drued_jorg_jax_partial(endog, pol, val, m_grid, exp_val_zero):
 # Numba DRUED-JORG
 start = time.time()
 jax.block_until_ready(
-    upenv.drued_jorg_numba(
+    ue_numba.drued_jorg_numba(
         endog_grid=policy_egm[0, 1:],
         policy=policy_egm[1, 1:],
         value=value_egm[1, 1:],
@@ -218,7 +219,7 @@ def drued_jorg_jax_partial(endog, pol, val, m_grid, exp_val_zero):
 for _ in range(n_runs):
     start = time.time()
     jax.block_until_ready(
-        upenv.drued_jorg_numba(
+        ue_numba.drued_jorg_numba(
             endog_grid=policy_egm[0, 1:],
             policy=policy_egm[1, 1:],
             value=value_egm[1, 1:],

docs/tutorials/ue_drued_jorg.ipynb

@@ -2,15 +2,12 @@
 # Project metadata
 # ======================================================================================
 [project]
-name               = "upper_envelope"
-description        = "Upper envelope scan for dynamic discrete-continuous life cycle models."
-version         = "0.1.3"
-requires-python    = ">=3.10"
+name = "upper_envelope"
+description = "Upper envelope scan for dynamic discrete-continuous life cycle models."
+dynamic = ["version"]
+requires-python = ">=3.10"
 dependencies       = [
     "numpy",
-    "pandas",
-    "scipy",
-    "jax"
 ]
 keywords = [
     "Dynamic programming",
@@ -32,7 +29,7 @@ classifiers = [
     "Topic :: Scientific/Engineering",
 ]
 authors            = [
-  { name="Max Blesch", email="maximilian.blesch@hu-berlin.de" },
+  { name="Max Blesch", email="maxblesch@gmail.com" },
   { name="Sebastian Gsell", email="gsell.sebastian@gmail.com" },
 ]
 maintainers = [
@@ -55,7 +52,7 @@ Github = "https://github.com/OpensourceEconomics/upper-envelope"
 # ======================================================================================
 
 [build-system]
-requires = ["hatchling", "hatch_vcs"]
+requires = ["hatchling", "hatch-vcs"]
 build-backend = "hatchling.build"
 
 [tool.hatch.build.hooks.vcs]

pyproject.toml

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

src/upper_envelope/__init__.py

@@ -0,0 +1,2 @@
+from upper_envelope.jax.drued_jorg_jax import drued_jorg_jax
+from upper_envelope.jax.fues_jax.fues_jax import fues_jax, fues_jax_unconstrained

src/upper_envelope/jax/__init__.py

@@ -13,7 +13,7 @@
 import jax.numpy as jnp
 from jax import vmap
 
-from upper_envelope.fues_jax.check_and_scan_funcs import (
+from upper_envelope.jax.fues_jax.check_and_scan_funcs import (
     determine_cases_and_conduct_necessary_scans,
 )
 from upper_envelope.math_funcs import calc_intersection_and_extrapolate_policy