style: black the files left unformatted by the dyno-importer branch

`pixi run black . --check` (CI "Code Quality") failed on 10 files added or
reworked on the merged feature branch.  Reformat with the pinned black
(26.5.1); `black . --check` is now clean (91 files).

No functional change.

Author: albop

src/dolo/algos/newton.py

@@ -344,7 +344,12 @@ def newton(
         if linsolve_cfg["method"] == "iti":
             linsolve_cfg["method"] = "neumann"
     else:
-        linsolve_cfg = {"method": "gmres", "maxit": 50, "tol": 1e-8, "precondition": True}
+        linsolve_cfg = {
+            "method": "gmres",
+            "maxit": 50,
+            "tol": 1e-8,
+            "precondition": True,
+        }
 
     if invmethod is not None:
         linsolve_cfg["method"] = "neumann" if invmethod == "iti" else invmethod

src/dolo/algos/time_iteration.py

@@ -11,7 +11,10 @@
 from dolo.numeric.decision_rule import DecisionRule
 from dolo.numeric.grids import CartesianGrid
 from dolo.numeric.optimize.ncpsolve import ncpsolve
-from dolo.numeric.optimize.newton import SerialDifferentiableFunction, newton as serial_newton
+from dolo.numeric.optimize.newton import (
+    SerialDifferentiableFunction,
+    newton as serial_newton,
+)
 from dolo.numeric.processes import DiscretizedIIDProcess
 
 from dolo.numeric import DiscretizedModel, GArray
@@ -27,7 +30,9 @@
 from .results import AlgoResult, TimeIterationResult
 
 
-def euler_residuals_tau(dmodel: DiscretizedModel, controls: GArray, dr: DecisionRule) -> GArray:
+def euler_residuals_tau(
+    dmodel: DiscretizedModel, controls: GArray, dr: DecisionRule
+) -> GArray:
     """Universal Euler residuals written entirely in terms of GArray and Tau."""
     res = GArray.zeros_like(controls)
     for i_m, m_val, s_endo in dmodel.regimes():
@@ -175,7 +180,9 @@ def time_iteration(
         dprocess_options = kwargs.pop("dprocess_options", {})
         if dprocess is not None:
             dprocess_options["dprocess"] = dprocess
-        dmodel = DiscretizedModel(model, grid_options=grid_options, dprocess_options=dprocess_options)
+        dmodel = DiscretizedModel(
+            model, grid_options=grid_options, dprocess_options=dprocess_options
+        )
 
     grid = dmodel.grid
     dprocess = dmodel.dprocess
@@ -273,9 +280,11 @@ def time_iteration(
         if trace:
             trace_details.append({"dr": copy.deepcopy(mdr)})
 
-        fn = lambda x_val: euler_residuals_tau(
-            dmodel, controls_ga.unravel(x_val), mdr
-        ).ravel().reshape((-1, n_x))
+        fn = (
+            lambda x_val: euler_residuals_tau(dmodel, controls_ga.unravel(x_val), mdr)
+            .ravel()
+            .reshape((-1, n_x))
+        )
         dfn = SerialDifferentiableFunction(fn)
 
         res_0 = fn(controls_0)
@@ -305,7 +314,6 @@ def time_iteration(
         err_SA = err_step / err_0 if not numpy.isnan(err_0) and err_0 > 0 else numpy.nan
         err_0 = err_step
 
-
         # Policy acceleration step
         if acc_cfg is not None and it > acc_cfg["wait"]:
             d = (controls - controls_0).reshape(sh_c)

src/dolo/misc/itprinter.py

@@ -3,6 +3,7 @@
 
 try:
     from rich.console import Console
+
     HAS_RICH = True
 except ImportError:
     HAS_RICH = False
@@ -82,9 +83,7 @@ def print_header(self, msg=None):
             if msg:
                 self.console.print(f"[bold cyan]{msg:^{self.table_width}}[/bold cyan]")
             self.console.print(self.box_top)
-            hdr_cells = [
-                f" {h:>{w}} " for h, w in zip(self.names, self.col_widths)
-            ]
+            hdr_cells = [f" {h:>{w}} " for h, w in zip(self.names, self.col_widths)]
             self.console.print("[bold]│" + "│".join(hdr_cells) + "│[/bold]")
             self.console.print(self.box_sep)
             self._header_printed = True
@@ -153,7 +152,9 @@ def print_finished(self):
         elapsed = time.time() - self.t_start
 
         if self.use_rich:
-            if getattr(self, "_header_printed", False) and not getattr(self, "_closed", False):
+            if getattr(self, "_header_printed", False) and not getattr(
+                self, "_closed", False
+            ):
                 self.console.print(self.box_bot)
                 self._closed = True
             self.console.print(

src/dolo/numeric/__init__.py

@@ -20,4 +20,3 @@
     "trembling_hand",
     "DiscretizedModel",
 ]
-

src/dolo/numeric/decision_rule.py

@@ -58,14 +58,18 @@ def __call__(self, *args):
 
         if len(args) == 1 and isinstance(args[0], SliceQP):
             sqp = args[0]
-            is_empty_exo = hasattr(self, "exo_grid") and isinstance(self.exo_grid, EmptyGrid)
+            is_empty_exo = hasattr(self, "exo_grid") and isinstance(
+                self.exo_grid, EmptyGrid
+            )
             if sqp.loc is not None and not is_empty_exo:
                 return self.eval_is(sqp.loc, sqp.endo)
             return self.eval_s(sqp.endo)
 
         if len(args) == 1 and isinstance(args[0], QP):
             qp = args[0]
-            is_empty_exo = hasattr(self, "exo_grid") and isinstance(self.exo_grid, EmptyGrid)
+            is_empty_exo = hasattr(self, "exo_grid") and isinstance(
+                self.exo_grid, EmptyGrid
+            )
             if isinstance(qp.loc, tuple) and len(qp.loc) >= 2 and not is_empty_exo:
                 i, s_endo = qp.loc[0], qp.loc[1]
                 s_arr = np.asarray(s_endo)
@@ -139,6 +143,7 @@ def __init__(
 
     def set_values(self, x):
         from dolo.numeric.garray import GArray
+
         if isinstance(x, GArray):
             x = x.data
         self.coefficients = self.__get_coefficients__(

src/dolo/numeric/discretized_model.py

@@ -89,7 +89,9 @@ def lower_bounds(self) -> Optional[GArray]:
         if "arbitrage_lb" not in self.model.functions:
             return None
         lb_fun = self.model.functions["arbitrage_lb"]
-        controls_lb = GArray.zeros(self.grid, len(self.model.symbols["controls"]), self.n_regimes)
+        controls_lb = GArray.zeros(
+            self.grid, len(self.model.symbols["controls"]), self.n_regimes
+        )
         N = self.N_endo
         p = np.repeat(self.parms[None, :], N, axis=0)
 
@@ -103,7 +105,9 @@ def upper_bounds(self) -> Optional[GArray]:
         if "arbitrage_ub" not in self.model.functions:
             return None
         ub_fun = self.model.functions["arbitrage_ub"]
-        controls_ub = GArray.zeros(self.grid, len(self.model.symbols["controls"]), self.n_regimes)
+        controls_ub = GArray.zeros(
+            self.grid, len(self.model.symbols["controls"]), self.n_regimes
+        )
         N = self.N_endo
         p = np.repeat(self.parms[None, :], N, axis=0)
 
@@ -158,7 +162,7 @@ def ergodic_distribution(self, controls: GArray) -> GArray:
         """
         P = self.transition_matrix(controls)
         # Solve mu P = mu  <=>  (P^T - I) mu = 0 with sum(mu) = 1
-        A = (P.T - np.eye(self.N_total))
+        A = P.T - np.eye(self.N_total)
         A[-1, :] = 1.0
         b = np.zeros(self.N_total)
         b[-1] = 1.0
@@ -168,4 +172,3 @@ def ergodic_distribution(self, controls: GArray) -> GArray:
 
         mu = GArray.zeros(self.grid, 1, self.n_regimes)
         return mu.unravel(mu_flat)
-

src/dolo/numeric/garray.py

@@ -37,6 +37,7 @@ class QualifiedPoint:
     val : np.ndarray
         Evaluated economic coordinates / values.
     """
+
     loc: Any
     val: np.ndarray
 
@@ -73,6 +74,7 @@ class QualifiedBatch:
     endo : np.ndarray
         Endogenous state matrix, shape (N, n_s).
     """
+
     loc: Any
     exo: np.ndarray
     endo: np.ndarray
@@ -201,4 +203,3 @@ def __neg__(self) -> GArray:
     def __repr__(self) -> str:
         grid_name = type(self.grid).__name__
         return f"GArray(grid={grid_name}, shape={self.data.shape})"
-

src/dolo/numeric/tau.py

@@ -130,9 +130,7 @@ def slice_transitions(
             target_regime = 0 if self.dprocess.n_nodes == 1 else I_m
             yield prob, SliceQP(loc=target_regime, exo=M_val, endo=S_endo)
 
-    def point_transitions(
-        self, s: QP, x: np.ndarray
-    ) -> Iterator[Tuple[float, QP]]:
+    def point_transitions(self, s: QP, x: np.ndarray) -> Iterator[Tuple[float, QP]]:
         """Point-level transition (pure Julia-style).
 
         Given a single state point QP(loc, val) and control vector x,

tests/test_tau_garray.py

@@ -8,7 +8,12 @@
 from dolo.compiler.model_v2 import ModelV2, yaml_compose
 from dolo.numeric.decision_rule import DecisionRule
 from dolo.numeric.garray import GArray, QP, SliceQP
-from dolo.numeric.grids import UniformCartesianGrid, ProductGrid, UnstructuredGrid, EmptyGrid
+from dolo.numeric.grids import (
+    UniformCartesianGrid,
+    ProductGrid,
+    UnstructuredGrid,
+    EmptyGrid,
+)
 from dolo.numeric.tau import Tau, trembling_hand
 from dolo.numeric.discretized_model import DiscretizedModel
 from dolo.algos.time_iteration import residuals_simple, time_iteration
@@ -25,6 +30,7 @@ def _load_v2(path: str) -> ModelV2:
 # 1. GArray and QP Unit Tests
 # ---------------------------------------------------------------------------
 
+
 def test_garray_cartesian():
     grid = UniformCartesianGrid(min=[0.1, 1.0], max=[0.5, 5.0], n=[10, 20])
     ga = GArray.zeros(grid, n_vars=3)
@@ -79,7 +85,9 @@ def test_qualified_point_and_batch():
     assert "QualifiedPoint" in repr(qp)
 
     # Construct QualifiedBatch
-    batch = QualifiedBatch(loc=1, exo=np.array([0.02]), endo=np.array([[1.0, 2.0], [1.5, 2.5]]))
+    batch = QualifiedBatch(
+        loc=1, exo=np.array([0.02]), endo=np.array([[1.0, 2.0], [1.5, 2.5]])
+    )
     assert batch.val.shape == (2, 3)
     assert np.allclose(batch.val[0], [0.02, 1.0, 2.0])
     assert np.allclose(batch.val[1], [0.02, 1.5, 2.5])
@@ -106,7 +114,10 @@ def test_trembling_hand_1d():
 # 2. Universal Euler Residuals & Parity
 # ---------------------------------------------------------------------------
 
-def euler_residuals(dmodel: DiscretizedModel, controls: GArray, dr: DecisionRule) -> GArray:
+
+def euler_residuals(
+    dmodel: DiscretizedModel, controls: GArray, dr: DecisionRule
+) -> GArray:
     """Universal Euler residuals using Tau, GArray, and direct dr(SliceQP) calls."""
     res = GArray.zeros_like(controls)
     for i, m_val, s_endo in dmodel.regimes():
@@ -119,10 +130,13 @@ def euler_residuals(dmodel: DiscretizedModel, controls: GArray, dr: DecisionRule
     return res
 
 
-@pytest.mark.parametrize("model_path", [
-    "examples/models/rbc_mc.yaml",
-    "examples/models/rbc_iid.yaml",
-])
+@pytest.mark.parametrize(
+    "model_path",
+    [
+        "examples/models/rbc_mc.yaml",
+        "examples/models/rbc_iid.yaml",
+    ],
+)
 def test_tau_euler_parity(model_path: str):
     model = _load_v2(model_path)
     dmodel = DiscretizedModel(model)
@@ -156,6 +170,7 @@ def test_tau_euler_parity(model_path: str):
 # 3. DecisionRule with QP and SliceQP
 # ---------------------------------------------------------------------------
 
+
 def test_decision_rule_qp_dispatch():
     model = _load_v2("examples/models/rbc_mc.yaml")
     dmodel = DiscretizedModel(model)
@@ -185,10 +200,14 @@ def test_decision_rule_qp_dispatch():
 # 4. Transition Matrix & Ergodic Distribution
 # ---------------------------------------------------------------------------
 
-@pytest.mark.parametrize("model_path", [
-    "examples/models/rbc_mc.yaml",
-    "examples/models/rbc_iid.yaml",
-])
+
+@pytest.mark.parametrize(
+    "model_path",
+    [
+        "examples/models/rbc_mc.yaml",
+        "examples/models/rbc_iid.yaml",
+    ],
+)
 def test_transition_matrix_and_ergodic(model_path: str):
     model = _load_v2(model_path)
     dmodel = DiscretizedModel(model)
@@ -205,18 +224,24 @@ def test_transition_matrix_and_ergodic(model_path: str):
     # Ergodic distribution
     mu = dmodel.ergodic_distribution(controls)
     assert isinstance(mu, GArray)
-    assert np.isclose(mu.data.sum(), 1.0), f"Ergodic distribution does not sum to 1: {mu.data.sum()}"
+    assert np.isclose(
+        mu.data.sum(), 1.0
+    ), f"Ergodic distribution does not sum to 1: {mu.data.sum()}"
     assert np.all(mu.data >= 0.0), "Negative stationary probabilities!"
 
 
 # ---------------------------------------------------------------------------
 # 5. Full Time Iteration Solver Parity
 # ---------------------------------------------------------------------------
 
-@pytest.mark.parametrize("model_path", [
-    "examples/models/rbc_mc.yaml",
-    "examples/models/rbc_iid.yaml",
-])
+
+@pytest.mark.parametrize(
+    "model_path",
+    [
+        "examples/models/rbc_mc.yaml",
+        "examples/models/rbc_iid.yaml",
+    ],
+)
 def test_full_solver_parity(model_path: str):
     model = _load_v2(model_path)
     dmodel = DiscretizedModel(model)
@@ -243,22 +268,30 @@ def test_full_solver_parity(model_path: str):
         it += 1
         dr.set_values(controls)
 
-        fn = lambda flat_x: euler_residuals(
-            dmodel, controls.unravel(flat_x), dr
-        ).ravel().reshape((-1, n_x))
+        fn = (
+            lambda flat_x: euler_residuals(dmodel, controls.unravel(flat_x), dr)
+            .ravel()
+            .reshape((-1, n_x))
+        )
 
         dfn = SerialDifferentiableFunction(fn)
-        flat_controls, _ = newton(dfn, controls.ravel().reshape((-1, n_x)), verbose=False, maxit=10)
+        flat_controls, _ = newton(
+            dfn, controls.ravel().reshape((-1, n_x)), verbose=False, maxit=10
+        )
         new_controls = controls.unravel(flat_controls)
 
         err = (new_controls - controls).norm()
         controls = new_controls
 
     # Parity check on final decision rule coefficients
     if isinstance(sol_legacy.dr.coefficients, list):
-        dr_diff = max(float(np.max(np.abs(c1 - c2))) for c1, c2 in zip(sol_legacy.dr.coefficients, dr.coefficients))
+        dr_diff = max(
+            float(np.max(np.abs(c1 - c2)))
+            for c1, c2 in zip(sol_legacy.dr.coefficients, dr.coefficients)
+        )
     else:
         dr_diff = float(np.max(np.abs(sol_legacy.dr.coefficients - dr.coefficients)))
 
-    assert dr_diff < 1e-5, f"Decision rule parity failed on {model_path}! diff={dr_diff}"
-
+    assert (
+        dr_diff < 1e-5
+    ), f"Decision rule parity failed on {model_path}! diff={dr_diff}"

tests/test_time_iteration_newton.py

@@ -42,7 +42,9 @@ def test_time_iteration_inner_options(rbc_model):
 
 
 def test_newton_gmres_backtracking(rbc_model):
-    sol = dolo.newton(rbc_model, linsolve="gmres", linesearch="backtracking", verbose=False)
+    sol = dolo.newton(
+        rbc_model, linsolve="gmres", linesearch="backtracking", verbose=False
+    )
     assert sol.converged
     assert sol.iterations < 15
     assert sol.error_res < 1e-7