sgd hot loop compiled

Author: simutisernestas

src/pixwake/optim/sgd.py

@@ -780,7 +780,6 @@ def topfarm_sgd_solve(
     boundary: Boundary,
     min_spacing: float,
     settings: SGDSettings | None = None,
-    jit: bool = True,
     progress: Literal["none", "print", "bar"] = "none",
     progress_callback: Callable[[SGDRecord], None] | None = None,
     record: bool = False,
@@ -810,15 +809,14 @@ def topfarm_sgd_solve(
             multi-polygon boundaries.
         min_spacing: Minimum inter-turbine distance.
         settings: SGD configuration. Uses defaults if None.
-        jit: Whether to JIT compile the gradient computations.
         progress: Progress rendering mode: none, print, or bar.
         progress_callback: Optional host callback receiving per-iteration progress.
         record: Whether to record and return structured optimization diagnostics.
         wind: Optional :class:`~pixwake.wind_resource.SampledWind` sampler.
-            When provided, :meth:`~pixwake.wind_resource.SampledWind.inc_rng` is
-            called outside JIT before each gradient step so that every iteration
-            uses a fresh, independent sample batch.  ``objective_fn`` must then
-            accept ``(x, y, wind) -> scalar``.
+            When provided, a fresh PRNG subkey is split from the carry key on
+            every gradient step so that every iteration uses an independent
+            sample batch.  ``objective_fn`` must then accept
+            ``(x, y, wind) -> scalar``.
 
     Returns:
         Tuple of (optimized_x, optimized_y), optionally with per-iteration `SGDRecord`.
@@ -854,9 +852,13 @@ def _constraint_penalty(x: jnp.ndarray, y: jnp.ndarray) -> jnp.ndarray:
             x, y, min_spacing
         )
 
-    jit_wrap = jax.jit if jit else (lambda f: f)
-    grad_obj_fn = jit_wrap(jax.grad(objective_fn, argnums=(0, 1)))
-    grad_con_fn = jit_wrap(jax.grad(_constraint_penalty, argnums=(0, 1)))
+    def _pen_per_turbine(x: jnp.ndarray, y: jnp.ndarray) -> jnp.ndarray:
+        return _boundary_penalty_per_turbine(
+            x, y, boundary_polys
+        ) + _spacing_penalty_per_turbine(x, y, min_spacing)
+
+    grad_obj_fn = jax.jit(jax.grad(objective_fn, argnums=(0, 1)))
+    grad_con_fn = jax.jit(jax.grad(_constraint_penalty, argnums=(0, 1)))
 
     if progress not in {"none", "print", "bar"}:
         raise ValueError(
@@ -865,74 +867,89 @@ def _constraint_penalty(x: jnp.ndarray, y: jnp.ndarray) -> jnp.ndarray:
     progress_enabled = progress != "none" or progress_callback is not None
     emit_progress = _make_emit_progress(progress, progress_callback, settings.max_iter)
 
-    x, y = init_x, init_y
-    state = _init_sgd_state(settings)
-    current_wind = wind
     record_list: list[SGDRecord] = []
 
+    def _side_effects(rec: SGDRecord) -> None:
+        if record:
+            record_list.append(rec)
+        if progress_enabled:
+            emit_progress(rec)
+
     def _compute_layout_change(
         x: jnp.ndarray, y: jnp.ndarray, prev_x: jnp.ndarray, prev_y: jnp.ndarray
     ) -> jnp.ndarray:
         return jnp.max(jnp.abs(x - prev_x)) + jnp.max(jnp.abs(y - prev_y))
 
-    def _is_done(change: jnp.ndarray, it: int | jnp.ndarray) -> bool:
-        return float(change) <= settings.tol or int(it) >= settings.max_iter
+    def _is_done(change: jnp.ndarray, it: jnp.ndarray) -> jnp.ndarray:
+        return (change <= settings.tol) | (it >= settings.max_iter)
+
+    has_wind = wind is not None
+    init_key = wind._key if wind is not None else jax.random.PRNGKey(0)
 
-    done = _is_done(_compute_layout_change(x, y, x - 1.0, y - 1.0), state.iteration)
-    while not done:
-        if current_wind is not None:
-            current_wind = current_wind.inc_rng()
-            grad_obj_x, grad_obj_y = grad_obj_fn(x, y, current_wind)
+    def body_fn(
+        carry: tuple[jnp.ndarray, jnp.ndarray, SGDState, jnp.ndarray, jnp.ndarray],
+    ) -> tuple[jnp.ndarray, jnp.ndarray, SGDState, jnp.ndarray, jnp.ndarray]:
+        x, y, state, rng_key, _prev_change = carry
+
+        if has_wind:
+            rng_key, subkey = jax.random.split(rng_key)
+            assert wind is not None
+            grad_obj_x, grad_obj_y = grad_obj_fn(x, y, wind.with_key(subkey))
         else:
             grad_obj_x, grad_obj_y = grad_obj_fn(x, y)
-        grad_con_x, grad_con_y = grad_con_fn(x, y)
 
+        grad_con_x, grad_con_y = grad_con_fn(x, y)
         d_x, d_y = _project_gradient(grad_obj_x, grad_obj_y, grad_con_x, grad_con_y)
         delta_x, delta_y, new_state = _sgd_step(state, d_x, d_y, settings)
 
-        # Feasibility-preserving step: per-turbine bisection.
-        # Each turbine keeps its own scale s_i, initialised to 1.  On each
-        # bisection round only the turbines that are still infeasible have
-        # their scale halved, so a turbine far from any constraint is never
-        # held back by one that is close to the boundary.
-        # _project_gradient() already makes the gradient direction tangential
-        # to any active constraint at the current position.
-        step_x, step_y = delta_x, delta_y
-        s = jnp.ones(x.shape[0])
-        x_try = x - s * step_x
-        y_try = y - s * step_y
-        pen_i = _boundary_penalty_per_turbine(
-            x_try, y_try, boundary_polys
-        ) + _spacing_penalty_per_turbine(x_try, y_try, min_spacing)
-        while (
-            float(jnp.max(pen_i)) > settings.bisect_tol
-            and float(jnp.min(s)) > settings.bisect_s_min
-        ):
-            s = jnp.where(pen_i > settings.bisect_tol, s / 2.0, s)
-            x_try = x - s * step_x
-            y_try = y - s * step_y
-            pen_i = _boundary_penalty_per_turbine(
-                x_try, y_try, boundary_polys
-            ) + _spacing_penalty_per_turbine(x_try, y_try, min_spacing)
-
-        x_new, y_new = x_try, y_try
-
-        iter_rec = SGDRecord(
-            x=x_new,
-            y=y_new,
-            iteration=new_state.iteration,
-            penalty=_constraint_penalty(x_new, y_new),
-            change=(change := _compute_layout_change(x_new, y_new, x, y)),
-            learning_rate=new_state.learning_rate,
-            converged=float(change <= settings.tol),
-            is_final=(done := _is_done(change, new_state.iteration)),
+        # Feasibility-preserving per-turbine bisection inside compiled loop.
+        def bisect_body(
+            bc: tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray, jnp.ndarray],
+        ) -> tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray, jnp.ndarray]:
+            s, xt, yt, pen_i = bc
+            s = jnp.where(pen_i > settings.bisect_tol, s / 2, s)
+            xt = x - s * delta_x
+            yt = y - s * delta_y
+            return s, xt, yt, _pen_per_turbine(xt, yt)
+
+        s0 = jnp.ones(x.shape[0])
+        x0 = x - delta_x
+        y0 = y - delta_y
+        _, x_new, y_new, _ = while_loop(
+            lambda bc: (jnp.max(bc[3]) > settings.bisect_tol)
+            & (jnp.min(bc[0]) > settings.bisect_s_min),
+            bisect_body,
+            (s0, x0, y0, _pen_per_turbine(x0, y0)),
         )
-        if record:
-            record_list.append(iter_rec)
-        if progress_enabled:
-            emit_progress(iter_rec)
 
-        x, y, state = x_new, y_new, new_state
+        change = _compute_layout_change(x_new, y_new, x, y)
+        is_final = _is_done(change, jnp.asarray(new_state.iteration))
+        if record or progress_enabled:
+            iter_rec = SGDRecord(
+                x=x_new,
+                y=y_new,
+                iteration=new_state.iteration,
+                penalty=_constraint_penalty(x_new, y_new),
+                change=change,
+                learning_rate=new_state.learning_rate,
+                converged=change <= settings.tol,
+                is_final=is_final,
+            )
+            jax.debug.callback(_side_effects, iter_rec)
+        return x_new, y_new, new_state, rng_key, change
+
+    def cond_fn(
+        carry: tuple[jnp.ndarray, jnp.ndarray, SGDState, jnp.ndarray, jnp.ndarray],
+    ) -> jnp.ndarray:
+        _, _, state, _, change = carry
+        return ~_is_done(change, jnp.asarray(state.iteration))
+
+    init_state = _init_sgd_state(settings)
+    # Initial change > tol ensures the first iteration always runs.
+    init_change = jnp.asarray(settings.tol + 1.0)
+    x, y, _, _, _ = while_loop(
+        cond_fn, body_fn, (init_x, init_y, init_state, init_key, init_change)
+    )
 
     if record:
         return x, y, record_list
@@ -1150,13 +1167,13 @@ def create_layout_optimizer(
         wind: Wind conditions — a :class:`~pixwake.wind_resource.WeibullWindRose`,
             :class:`~pixwake.wind_resource.TimeSeriesWind`, or
             :class:`~pixwake.wind_resource.SampledWind`.  When a
-            :class:`~pixwake.wind_resource.SampledWind` is provided the solver
-            draws a fresh sample batch on every gradient step via
-            :meth:`~pixwake.wind_resource.SampledWind.inc_rng`.
+            :class:`~pixwake.wind_resource.SampledWind` is provided a fresh PRNG
+            subkey is split on every gradient step, drawing an independent sample
+            batch each iteration.
         settings: SGD configuration.
         **kwargs: Additional keyword arguments forwarded to
             :func:`topfarm_sgd_solve` (e.g. ``progress``, ``progress_callback``,
-            ``record``, ``jit``).
+            ``record``).
 
     Returns:
         Function ``(init_x, init_y) -> (opt_x, opt_y)``.  When ``record=True``