Unify the e-ph loop argument lists, drop the batched dummy padding, and name LoopContext construction

README_GPU.md

@@ -196,7 +196,9 @@ payload, named for which momentum is the outer loop and which is batched on the
   `eph_batched_bytes_per_point(calc, EPDataKBatched)` for the loop's memory-adaptive batch
   sizing. The k side is streamed per k-batch (host-staged, no whole-grid device stack), and the
   payload is trimmed to the batch's actual width — a consumer reads its own size from any field
-  (e.g. `size(eps, 4)`) and never sees a padded tail (the outer-k convention).
+  (e.g. `size(eps, 4)`) and never sees a padded tail (the outer-k convention). Both loops consume
+  their max-width staging buffers through contiguous trailing-prefix views, so a partial final batch
+  computes only its own columns; nothing is padded with duplicated data.
 - Both loops fold their device-buffer byte accounting into `src/calculator/eph_device_staging.jl`:
   `_outer_{k,q}_staging_bytes(…)` return the loop's `(per_point, committed)` device-byte counts, and
   `plan_batch(backend, per_point, committed, cap; …)` turns those into the memory-adaptive batch
@@ -254,10 +256,6 @@ unchanged). No window handling is needed in the calculator beyond addressing its
   benchmarked and validated bit-identical, but the gain is small on the GPU (CUDA's pool already
   recycles device buffers), so it is deferred. Best done together with the calculator loop, where
   one workspace allocated at loop setup is reused across all (k, q).
-- **View-instead-of-fill for the outer-q staging buffers.** `run_eph_over_q_and_k`'s per-batch
-  staging fills padded per-k buffers (`Uk_batch`, …). Passing `nk_batch` and taking width-`nk` views
-  into the padded buffers (instead of filling) would avoid the copy. A hot-path change with no
-  correctness component, so benchmark it on its own before adopting.
 - **Energy window and long-range/polar on the GPU** — left on the CPU per-k path for now.
 - **MPI / multi-GPU** for the GPU loop — not in this foundation.
 - **Backend as a type parameter instead of a backend object (future).** The `use_gpu` keyword is
@@ -278,7 +276,7 @@ unchanged). No window handling is needed in the calculator beyond addressing its
 - **A QR-based batched eigensolve (future).** The batched eigensolve uses `CUSOLVER.heevjBatched!`
   (Jacobi). A QR-based `HEEV` (e.g. via cuSolverDx) may be faster for the small matrices here;
   worth evaluating, but not in this PR. Accuracy is not a motivation — Jacobi is already at
-  machine precision. Note that `cusolverDnXsyevBatched` is *not* the answer: it is 2–3× faster per
+  machine precision. Note that `cusolverDnXsyevBatched` is *not* the answer: it is 2-3× faster per
   matrix but needs ~1.05 MB of workspace per matrix (65× heevj), which caps one call at ~72k
   matrices on an 80 GB A100 and competes with the device-resident e-ph tiles.
 - **Parametrize `Model` over its `WannierObject` array type (future).** Widening `WannierObject`

TODO.md

@@ -14,22 +14,24 @@
   https://github.com/JuliaMolSim/DFTK.jl/blob/master/src/architecture.jl and follow the pattern in
   https://docs.dftk.org/stable/developer/gpu_computations/ .
 
-- [ ] Remove all the `nk_batch+1:nk_batch_max` dummy padding fills (e.g.
-  `src/calculator/run_eph_over_k_and_kq.jl` ~L702) by passing the actual `nk_batch` through to the
-  batched interpolation/kernels instead of padding partial tails with duplicated valid data to run on
-  dense `nk_batch_max`-sized arrays.
+- [x] Remove all the `nk_batch+1:nk_batch_max` dummy padding fills — done at both sites
+  (`run_eph_over_k_and_kq.jl`, `run_eph_over_q_and_k.jl`). The batched interpolation/eigensolve/kernels
+  are handed contiguous width-`nk_batch` trailing-prefix views of the max-width staging buffers, so a
+  partial final batch computes only its own columns. `get_eph_Rq_to_kq_batched!` now asserts its
+  workspace sizes with `>=` and prefix-views them internally, as its sibling
+  `get_eph_kR_to_kq_batched!` already did.
 
 - [ ] JML should review `TiledDeviceOutput` (the Sᵢ tiling machinery: `tile_begin!` / `tile_download!` /
   `tile_offset` / `tile_length`, used by `BoltzmannCalculator`'s batched path).
 
-- [ ] Unify the CPU/GPU e-ph loop signatures (`_loop_eph_over_k_and_kq` vs
-  `_loop_eph_over_k_and_kq_batched`, `src/calculator/run_eph_over_k_and_kq.jl`). Decided direction (JML,
-  PR #9): keep **two functions with different names**, but give them the **identical positional
-  argument list** and push the path-specific data into **keyword arguments** — CPU kwargs
-  `(epstates, ep_ekpRs, epmat, ep_ekpR_obj, dyn_threads, epmat_R, epobj_ekpR_R, ep_ekpR_Rs)`, GPU
-  kwargs `(epmat_dev, backend)`. Constraint (`_setup`/`_loop` Core.Box rule): the CPU `_loop` must
-  destructure any NamedTuple into locals at the top before `@threads`, never index it inside the
-  threaded closure.
+- [x] Unify the per-point/batched e-ph loop signatures — done for both driver families. Each pair
+  takes the identical positional list of what BOTH need (outer-k: `(model, kpts, qpts, kqpts,
+  el_k_save, el_kq_save, ph_save, precompute_ph, backend)`; outer-q: `(model, kpts, qpts, el_k_save,
+  ph_save, eph_buffers, backend)`) and its own path-specific data as individual keyword arguments
+  (per-point: the host interpolators/channels, resp. the precomputed k+q states; batched:
+  `epmat_dev`, resp. `el_ham_dev`). `backend` is the last shared positional rather than a batched-only
+  kwarg, so no `_loop_*` hardcodes `CPUBackend()` any more (`run_eph_over_k_and_q` too, which has no
+  batched twin). The convention is stated once, above `_loop_eph_over_k_and_kq`.
 
 - [x] ~~Reconsider whether `backend` should be built inside `_loop_eph_over_k_and_kq_batched` rather
   than in `_setup_eph_over_k_and_kq`.~~ Subsumed: `backend` is now a user-facing driver keyword, so it
@@ -43,11 +45,8 @@
   which is exactly right now that batched-on-`CPUBackend` is a supported configuration. The real
   defect was the flag: `calc.on_gpu = backend isa GPUBackend` inferred the loop shape from the
   backend, which builds the wrong buffers under CPU+batched. It is now `calc.batched = mode isa
-  BatchedMode`, from the `mode::LoopMode` keyword the drivers pass to `setup_calculator!`.
-
-- [ ] Clean up the `LoopContext` construction at the batch/per-k scope (deferred out of the
-  `backend`/`batched` PR as an independent follow-up)
-  (`src/calculator/run_eph_over_k_and_kq.jl` ~L713/L728). The batch-scope `ctx_batch` and per-k
-  `ctx_k` are built from positional constructors that are disambiguated by whether the argument is an
-  `Integer` outer index (`ik`) or a `UnitRange` batch (`iks_batch`) — flagged as flaky in review.
-  Consider a clearer, explicitly-named construction API for the two scopes.
+  BatchedMode`, from the positional `mode::LoopMode` the drivers pass to `setup_calculator!`.
+
+- [x] Clean up the `LoopContext` construction at the batch/per-k scope — done. The `BatchedMode`
+  convenience constructor is keyword-only (`batch` / `outer_index` / required `n_batch_max`), so
+  argument 3 is never positional there, and the per-k context is `with_outer_index(ctx_batch, ik)`.

src/EPData.jl

@@ -86,10 +86,13 @@ batch of matrices takes the plural name `eps`; `ek`/`ekq`/`uk`/`ukq`/`wtk` are n
 they are (the batch is just their trailing axis).
 
 All fields are trimmed to the batch's actual width `nk` (the outer-k convention): the final partial
-batch has `nk < n_batch_max`, and the loop hands width-`nk` views into its internally padded staging
-buffers (a trailing-prefix device view, so it stays contiguous). A consumer therefore reads its own
-size from any field (e.g. `size(eps, 4)`) and never sees a padded tail. (The loop still zeros the
-internal `wtk` padding as defense-in-depth, but that padding is not exposed here.)
+batch has `nk < n_batch_max`, and the loop hands width-`nk` views into its max-width staging buffers
+(a trailing-prefix device view, so it stays contiguous). A consumer therefore reads its own size from
+any field (e.g. `size(eps, 4)`) and never sees a padded tail. There is no padded tail internally
+either — the loop computes only the batch's own columns — so the payload's width is the *only*
+guarantee: a consumer that sized itself from `ctx.n_batch_max` instead of `size(eps, 4)` would read
+stale data from the previous batch (uninitialised if the run has only one, partial batch), which is
+worse than uninitialised because it is finite and plausible.
 
 Fields (`m` = k+q band, `n` = k band, `k` = batch column):
 - `eps` :: `(nw, nw, nmodes, nk)` — eigenbasis e-ph matrices. Out-of-window bands are already zeroed

src/calculator/AbstractCalculator.jl

@@ -93,8 +93,21 @@ Fields:
 - `mode`        :: `SingleMode()` (per-(k, q) host loop) or `BatchedMode()` (device-batched loop).
 - `outer_index` :: current outer index (`ik` for outer-k loops, `iq` for the outer-q loop); `0` at
   batch scope.
-- `batch`       :: outer-iteration range of the current batch (`1:0` on the CPU paths).
-- `n_batch_max` :: loop batch cap, for device-buffer sizing.
+- `batch`       :: outer-iteration range of the current batch (`1:0` when the loop does not batch its
+  outer axis — the per-point loops, and the batched outer-q loop, which batches the INNER k axis).
+- `n_batch_max` :: loop batch cap, for device-buffer sizing. Which axis it caps follows the loop: the
+  OUTER-k width in the batched outer-k loop, the INNER-k width in the batched outer-q loop. A
+  consumer that reads it as an outer-axis width (as `TiledDeviceOutput` does, pairing it with
+  `ctx.batch`) is therefore outer-k-only.
+
+The two batched forms and the per-k derivation are constructed by name, never positionally:
+
+```julia
+LoopContext(backend, SingleMode(), ik)                                    # per-point
+LoopContext(backend, BatchedMode(); batch = iks, n_batch_max = nk_max)    # outer-k, batch scope
+LoopContext(backend, BatchedMode(); outer_index = iq, n_batch_max = nk_max)  # outer-q, one q
+with_outer_index(ctx_batch, ik)                                           # one k of a batch
+```
 """
 struct LoopContext{BT <: AbstractBackend, MT <: LoopMode}
     backend     :: BT
@@ -109,10 +122,24 @@ end
 LoopContext(backend::AbstractBackend, ::SingleMode, outer_index::Integer) =
     LoopContext(backend, SingleMode(), outer_index, 1:0, 0)
 
-# BatchedMode context at batch scope (device loops): there is no single outer index spanning the whole
-# batch, so `outer_index = 0` is the "no single outer index — use `batch`" sentinel.
-LoopContext(backend::AbstractBackend, ::BatchedMode, batch::UnitRange, n_batch_max::Integer) =
-    LoopContext(backend, BatchedMode(), 0, batch, n_batch_max)
+# BatchedMode context, keyword-only so the two scopes a batched loop needs are told apart by NAME
+# rather than by whether argument 3 is an `Integer` or a `UnitRange`:
+#   * batch scope (outer-k loop): pass `batch`; `outer_index = 0` is the "no single outer index — use
+#     `batch`" sentinel.
+#   * one outer iteration (outer-q loop): pass `outer_index`; `batch = 1:0` means "the outer axis is
+#     not batched here" (that loop batches its inner k axis).
+LoopContext(backend::AbstractBackend, ::BatchedMode; outer_index::Integer = 0,
+            batch::UnitRange = 1:0, n_batch_max::Integer) =
+    LoopContext(backend, BatchedMode(), outer_index, batch, n_batch_max)
+
+"""
+    with_outer_index(ctx::LoopContext, outer_index::Integer) -> LoopContext
+
+The same context at one outer index, all other fields kept. The batched outer-k loop uses it to
+derive its per-k context from the batch-scope one, which is exactly what the two contexts differ by.
+"""
+with_outer_index(ctx::LoopContext, outer_index::Integer) =
+    LoopContext(ctx.backend, ctx.mode, outer_index, ctx.batch, ctx.n_batch_max)
 
 
 # =============================================================================
@@ -246,7 +273,7 @@ if VERSION >= v"1.11.0-DEV.469"
         "postprocess_calculator!, calculator_begin!, calculator_end!, " *
         "OuterKLoop, OuterQLoop, OuterIteration, OuterIterationBatch, " *
         "AbstractElPhPayload, EPData, EPDataQBatched, EPDataKBatched, " *
-        "LoopContext, SingleMode, BatchedMode, LoopMode, " *
+        "LoopContext, with_outer_index, SingleMode, BatchedMode, LoopMode, " *
         "AbstractBackend, CPUBackend, GPUBackend, gpu_backend, alloc, free_bytes, synchronize, " *
         "batched_gemm!, eph_window_scatter!, bte_window_accumulate!, " *
         "eph_batched_bytes_per_point, allowed_eph_phonon_basis, " *