Merge pull request #1048 from jdebacker/dask_graph

rickecon · web-flow · commit ef6695e09760 · 2025-08-21T17:12:25.000-06:00
Merging
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -6,6 +6,12 @@ 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.14.7] - 2025-08-21 17:00:00
+
+### Added
+
+- Refactor calls to dask in `SS.py` and `TPI.py`.  See PR [#1048](https://github.com/PSLmodels/OG-Core/pull/1048)
+
 ## [0.14.6] - 2025-08-15 14:00:00
 
 ### Added
@@ -403,6 +409,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Any earlier versions of OG-USA can be found in the [`OG-Core`](https://github.com/PSLmodels/OG-Core) repository [release history](https://github.com/PSLmodels/OG-Core/releases) from [v.0.6.4](https://github.com/PSLmodels/OG-Core/releases/tag/v0.6.4) (Jul. 20, 2021) or earlier.
 
 
+[0.14.7]: https://github.com/PSLmodels/OG-Core/compare/v0.14.6...v0.14.7
 [0.14.6]: https://github.com/PSLmodels/OG-Core/compare/v0.14.5...v0.14.6
 [0.14.5]: https://github.com/PSLmodels/OG-Core/compare/v0.14.4...v0.14.5
 [0.14.4]: https://github.com/PSLmodels/OG-Core/compare/v0.14.3...v0.14.4
diff --git a/environment.yml b/environment.yml
@@ -13,9 +13,8 @@ dependencies:
 - dask>=2.30.0
 - dask-core>=2.30.0
 - distributed>=2.30.1
-- paramtools>=0.15.0
+- paramtools>=0.20.0
 - sphinx>=3.5.4
-- marshmallow<4.0.0
 - sphinx-argparse
 - sphinxcontrib-bibtex>=2.0.0
 - sphinx-math-dollar
diff --git a/ogcore/SS.py b/ogcore/SS.py
@@ -169,6 +169,60 @@ def euler_equation_solver(guesses, *args):
     return errors
 
 
+def solve_for_j(
+    guesses,
+    r_p,
+    w,
+    p_tilde,
+    bq_j,
+    rm_j,
+    tr_j,
+    ubi_j,
+    factor,
+    j,
+    p_future,
+):
+    """
+    Solves the household's optimization problem for a given type j.
+
+    Args:
+        guesses (Numpy array): initial guesses for b and n, length 2S
+        r_p (scalar): return on household investment portfolio
+        w (scalar): real wage rate
+        p_tilde (scalar): composite good price
+        bq_j (Numpy array): bequest amounts by age, length S
+        rm_j (Numpy array): remittance amounts by age, length S
+        tr_j (Numpy array): government transfer amount by age, length S
+        ubi_j (vector): universal basic income (UBI) payment, length S
+        factor (scalar): scaling factor converting model units to dollars
+        j (int): household type index
+        p_future (OG-Core Specifications object): future model parameters
+
+    Returns:
+        root (OptimizeResult): the optimization result
+    """
+    # scattered_p is either the original object (serial case)
+    # or a Future pointing to it (distributed case)
+    return opt.root(
+        euler_equation_solver,
+        guesses * 0.9,
+        args=(
+            r_p,
+            w,
+            p_tilde,
+            bq_j,
+            rm_j,
+            tr_j,
+            ubi_j,
+            factor,
+            j,
+            p_future,
+        ),
+        method=p_future.FOC_root_method,
+        tol=MINIMIZER_TOL,
+    )
+
+
 def inner_loop(outer_loop_vars, p, client):
     """
     This function solves for the inner loop of the SS.  That is, given
@@ -235,50 +289,19 @@ def inner_loop(outer_loop_vars, p, client):
     tr = household.get_tr(TR, None, p, "SS")
     ubi = p.ubi_nom_array[-1, :, :] / factor
 
-    scattered_p = client.scatter(p, broadcast=True) if client else p
-
-    lazy_values = []
-    for j in range(p.J):
-        guesses = np.append(bssmat[:, j], nssmat[:, j])
+    results = []
+    # from dask.base import dask_sizeof
 
-        # Create a delayed function that will access the scattered_p
-        @delayed
-        def solve_for_j(
-            guesses,
-            r_p,
-            w,
-            p_tilde,
-            bq_j,
-            rm_j,
-            tr_j,
-            ubi_j,
-            factor,
-            j,
-            scattered_p_future,
-        ):
-            # This function will be executed on workers with access to scattered_p
-            return opt.root(
-                euler_equation_solver,
-                guesses * 0.9,
-                args=(
-                    r_p,
-                    w,
-                    p_tilde,
-                    bq_j,
-                    rm_j,
-                    tr_j,
-                    ubi_j,
-                    factor,
-                    j,
-                    scattered_p_future,
-                ),
-                method=p.FOC_root_method,
-                tol=MINIMIZER_TOL,
-            )
-
-        # Add the delayed computation to our list
-        lazy_values.append(
-            solve_for_j(
+    if client:
+        # Scatter p only once and only if client not equal None
+        scattered_p_future = client.scatter(p, broadcast=True)
+
+        # Launch in parallel with submit (or map)
+        futures = []
+        for j in range(p.J):
+            guesses = np.append(bssmat[:, j], nssmat[:, j])
+            f = client.submit(
+                solve_for_j,
                 guesses,
                 r_p,
                 w,
@@ -289,21 +312,30 @@ def solve_for_j(
                 ubi[:, j],
                 factor,
                 j,
-                scattered_p,
+                scattered_p_future,
             )
-        )
+            futures.append(f)
 
-    if client:
-        # Compute all the values
-        futures = client.compute(lazy_values)
-        # Later, gather the results when needed
         results = client.gather(futures)
+
     else:
-        results = compute(
-            *lazy_values,
-            scheduler=dask.multiprocessing.get,
-            num_workers=p.num_workers,
-        )
+        # Serial fallback (no dask client)
+        for j in range(p.J):
+            guesses = np.append(bssmat[:, j], nssmat[:, j])
+            res = solve_for_j(
+                guesses,
+                r_p,
+                w,
+                p_tilde,
+                bq[:, j],
+                rm[:, j],
+                tr[:, j],
+                ubi[:, j],
+                factor,
+                j,
+                p,  # pass the raw object directly
+            )
+            results.append(res)
 
     for j, result in enumerate(results):
         euler_errors[:, j] = result.fun
@@ -1221,7 +1253,6 @@ def run_SS(p, client=None):
             results
 
     """
-
     # Create list of deviation factors for initial guesses of r and TR
     dev_factor_list = [
         [1.00, 1.0],
diff --git a/ogcore/TPI.py b/ogcore/TPI.py
@@ -756,6 +756,10 @@ def run_TPI(p, client=None):
     euler_errors = np.zeros((p.T, 2 * p.S, p.J))
     TPIdist_vec = np.zeros(p.maxiter)
 
+    # scatter parameters to workers
+    if client:
+        scattered_p_future = client.scatter(p, broadcast=True)
+
     # TPI loop
     while (TPIiter < p.maxiter) and (TPIdist >= p.mindist_TPI):
         outer_loop_vars = (r_p, r, w, p_m, BQ, RM, TR, theta)
@@ -766,38 +770,40 @@ def run_TPI(p, client=None):
         ).sum(axis=2)
         p_tilde = aggr.get_ptilde(p_i[:, :], p.tau_c[:, :], p.alpha_c, "TPI")
 
-        # scatter parameters to workers
-        scattered_p = client.scatter(p, broadcast=True) if client else p
-
+        # Initialize Euler errors
         euler_errors = np.zeros((p.T, 2 * p.S, p.J))
-        lazy_values = []
-        for j in range(p.J):
-            guesses = (guesses_b[:, :, j], guesses_n[:, :, j])
-
-            # Add the delayed computation to our list
-            lazy_values.append(
-                delayed(inner_loop)(
+        # Solve for household decisions
+        results = []
+        if client:
+            futures = []
+            for j in range(p.J):
+                guesses = (guesses_b[:, :, j], guesses_n[:, :, j])
+                f = client.submit(
+                    inner_loop,
                     guesses,
                     outer_loop_vars,
                     initial_values,
                     ubi,
                     j,
                     ind,
-                    scattered_p,
+                    scattered_p_future,
                 )
-            )
-        if client:
-            # Compute all the values
-            futures = client.compute(lazy_values)
-            # Later, gather the results when needed
+                futures.append(f)
             results = client.gather(futures)
         else:
-            results = compute(
-                *lazy_values,
-                scheduler=dask.multiprocessing.get,
-                num_workers=p.num_workers,
-            )
-
+            # Serial fallback (no dask client)
+            for j in range(p.J):
+                guesses = (guesses_b[:, :, j], guesses_n[:, :, j])
+                res = inner_loop(
+                    guesses,
+                    outer_loop_vars,
+                    initial_values,
+                    ubi,
+                    j,
+                    ind,
+                    p,  # pass the raw object directly
+                )
+                results.append(res)
         for j, result in enumerate(results):
             euler_errors[:, :, j], b_mat[:, :, j], n_mat[:, :, j] = result
 
diff --git a/ogcore/__init__.py b/ogcore/__init__.py
@@ -20,4 +20,4 @@
 from ogcore.txfunc import *
 from ogcore.utils import *
 
-__version__ = "0.14.6"
+__version__ = "0.14.7"
diff --git a/setup.py b/setup.py
@@ -5,7 +5,7 @@
 
 setuptools.setup(
     name="ogcore",
-    version="0.14.6",
+    version="0.14.7",
     author="Jason DeBacker and Richard W. Evans",
     license="CC0 1.0 Universal (CC0 1.0) Public Domain Dedication",
     description="A general equilibrium overlapping generations model for fiscal policy analysis",
