Tiled splineTransform! for RLR/SLR and multi-tile stride fix for 3D z-major layouts

src/tiling.jl

@@ -421,10 +421,14 @@ end
 """
     num_columns(grid::SpringsteelGrid{CartesianGeometry, <:SplineBasisArray, <:SplineBasisArray, <:SplineBasisArray}) -> Int
 
-Return `b_jDim * b_kDim` for 3-D Cartesian Spline×Spline×Spline (RRR) grids.
+Return `iDim * jDim` for 3-D Cartesian Spline×Spline×Spline (RRR) grids: the
+number of physical vertical columns (one per horizontal gridpoint), matching
+the RZ/RiRk and RLZ/RLR/SLZ/SLR convention that model drivers advance the
+state one z-fastest vertical column at a time. (Previously returned the
+spectral count `b_jDim * b_kDim`, which no caller used.)
 """
 function num_columns(grid::SpringsteelGrid{CartesianGeometry, <:SplineBasisArray, <:SplineBasisArray, <:SplineBasisArray})
-    return grid.params.b_jDim * grid.params.b_kDim
+    return grid.params.iDim * grid.params.jDim
 end
 
 # ────────────────────────────────────────────────────────────────────────────
@@ -3630,13 +3634,75 @@ function splineTransform!(sharedSpectral::SharedArray{real},
     b_kDim       = tile.params.b_kDim
     kDim_wn      = tile.params.iDim + tile.params.patchOffsetL
     nvars        = length(tile.params.vars)
+    # Sub-blocks are aligned by wavenumber, but an inner tile carries FEWER
+    # wavenumbers than the patch (kDim_wn follows the tile's outermost ring), so
+    # the z-level strides differ between the tile and patch layouts.
+    patch_kDim   = patch.params.iDim + patch.params.patchOffsetL
     nblocks      = 1 + kDim_wn * 2
+    nblocks_p    = 1 + patch_kDim * 2
 
     for v in 1:nvars
         for z_b in 1:b_kDim
             # z-level base indices (tile-local and patch-level)
             tr_base = (z_b - 1) * b_iDim_tile * nblocks
-            pr_base = (z_b - 1) * b_iDim_patch * nblocks
+            pr_base = (z_b - 1) * b_iDim_patch * nblocks_p
+
+            # k = 0 block
+            pr1 = pr_base + 1
+            tr1 = tr_base + 1
+            patch.spectral[pr1:pr1+b_iDim_patch-1, v] .=
+                SAtransform(patch.ibasis.data[1, v],
+                            view(sharedSpectral, pr1:pr1+b_iDim_patch-1, v))
+            tile.spectral[tr1:tr1+b_iDim_tile-1, v] .=
+                patch.spectral[pr1+siL-1:pr1+siL+b_iDim_tile-2, v]
+
+            # k >= 1: real and imaginary
+            for k in 1:kDim_wn
+                p = (k - 1) * 2
+
+                # Real part
+                pp1 = pr_base + b_iDim_patch + p * b_iDim_patch + 1
+                tp1 = tr_base + b_iDim_tile + p * b_iDim_tile + 1
+                patch.spectral[pp1:pp1+b_iDim_patch-1, v] .=
+                    SAtransform(patch.ibasis.data[2, v],
+                                view(sharedSpectral, pp1:pp1+b_iDim_patch-1, v))
+                tile.spectral[tp1:tp1+b_iDim_tile-1, v] .=
+                    patch.spectral[pp1+siL-1:pp1+siL+b_iDim_tile-2, v]
+
+                # Imaginary part
+                pp1 = pp1 + b_iDim_patch
+                tp1 = tp1 + b_iDim_tile
+                patch.spectral[pp1:pp1+b_iDim_patch-1, v] .=
+                    SAtransform(patch.ibasis.data[3, v],
+                                view(sharedSpectral, pp1:pp1+b_iDim_patch-1, v))
+                tile.spectral[tp1:tp1+b_iDim_tile-1, v] .=
+                    patch.spectral[pp1+siL-1:pp1+siL+b_iDim_tile-2, v]
+            end
+        end
+    end
+    return nothing
+end
+
+function splineTransform!(sharedSpectral::SharedArray{real},
+                            patch::_RLRGrid,
+                            tile::_RLRGrid)
+    # Same per-z_b block layout as the RLZ method above (k = 0 block, then
+    # real/imag pairs per wavenumber), with the spline vertical's b_kDim blocks.
+    b_iDim_tile  = tile.params.b_iDim
+    b_iDim_patch = patch.params.b_iDim
+    siL          = tile.params.spectralIndexL
+    b_kDim       = tile.params.b_kDim
+    kDim_wn      = tile.params.iDim + tile.params.patchOffsetL
+    nvars        = length(tile.params.vars)
+    patch_kDim   = patch.params.iDim + patch.params.patchOffsetL
+    nblocks      = 1 + kDim_wn * 2
+    nblocks_p    = 1 + patch_kDim * 2
+
+    for v in 1:nvars
+        for z_b in 1:b_kDim
+            # z-level base indices (tile-local and patch-level)
+            tr_base = (z_b - 1) * b_iDim_tile * nblocks
+            pr_base = (z_b - 1) * b_iDim_patch * nblocks_p
 
             # k = 0 block
             pr1 = pr_base + 1
@@ -3683,12 +3749,14 @@ function splineTransform!(sharedSpectral::SharedArray{real},
     b_kDim       = tile.params.b_kDim
     kDim_wn      = tile.params.iDim + tile.params.patchOffsetL
     nvars        = length(tile.params.vars)
+    patch_kDim   = patch.params.iDim + patch.params.patchOffsetL
     nblocks      = 1 + kDim_wn * 2
+    nblocks_p    = 1 + patch_kDim * 2
 
     for v in 1:nvars
         for z_b in 1:b_kDim
             tr_base = (z_b - 1) * b_iDim_tile * nblocks
-            pr_base = (z_b - 1) * b_iDim_patch * nblocks
+            pr_base = (z_b - 1) * b_iDim_patch * nblocks_p
 
             pr1 = pr_base + 1
             tr1 = tr_base + 1
@@ -3721,3 +3789,60 @@ function splineTransform!(sharedSpectral::SharedArray{real},
     end
     return nothing
 end
+
+function splineTransform!(sharedSpectral::SharedArray{real},
+                            patch::_SLRGrid,
+                            tile::_SLRGrid)
+    # Same per-z_b block layout as the RLR method (k = 0 block, then real/imag
+    # pairs per wavenumber), with the spline vertical's b_kDim blocks.
+    b_iDim_tile  = tile.params.b_iDim
+    b_iDim_patch = patch.params.b_iDim
+    siL          = tile.params.spectralIndexL
+    b_kDim       = tile.params.b_kDim
+    kDim_wn      = tile.params.iDim + tile.params.patchOffsetL
+    nvars        = length(tile.params.vars)
+    patch_kDim   = patch.params.iDim + patch.params.patchOffsetL
+    nblocks      = 1 + kDim_wn * 2
+    nblocks_p    = 1 + patch_kDim * 2
+
+    for v in 1:nvars
+        for z_b in 1:b_kDim
+            # z-level base indices (tile-local and patch-level)
+            tr_base = (z_b - 1) * b_iDim_tile * nblocks
+            pr_base = (z_b - 1) * b_iDim_patch * nblocks_p
+
+            # k = 0 block
+            pr1 = pr_base + 1
+            tr1 = tr_base + 1
+            patch.spectral[pr1:pr1+b_iDim_patch-1, v] .=
+                SAtransform(patch.ibasis.data[1, v],
+                            view(sharedSpectral, pr1:pr1+b_iDim_patch-1, v))
+            tile.spectral[tr1:tr1+b_iDim_tile-1, v] .=
+                patch.spectral[pr1+siL-1:pr1+siL+b_iDim_tile-2, v]
+
+            # k >= 1: real and imaginary
+            for k in 1:kDim_wn
+                p = (k - 1) * 2
+
+                # Real part
+                pp1 = pr_base + b_iDim_patch + p * b_iDim_patch + 1
+                tp1 = tr_base + b_iDim_tile + p * b_iDim_tile + 1
+                patch.spectral[pp1:pp1+b_iDim_patch-1, v] .=
+                    SAtransform(patch.ibasis.data[2, v],
+                                view(sharedSpectral, pp1:pp1+b_iDim_patch-1, v))
+                tile.spectral[tp1:tp1+b_iDim_tile-1, v] .=
+                    patch.spectral[pp1+siL-1:pp1+siL+b_iDim_tile-2, v]
+
+                # Imaginary part
+                pp1 = pp1 + b_iDim_patch
+                tp1 = tp1 + b_iDim_tile
+                patch.spectral[pp1:pp1+b_iDim_patch-1, v] .=
+                    SAtransform(patch.ibasis.data[3, v],
+                                view(sharedSpectral, pp1:pp1+b_iDim_patch-1, v))
+                tile.spectral[tp1:tp1+b_iDim_tile-1, v] .=
+                    patch.spectral[pp1+siL-1:pp1+siL+b_iDim_tile-2, v]
+            end
+        end
+    end
+    return nothing
+end

test/tiling.jl

@@ -617,6 +617,66 @@
             @test length(single_tile) == 1
         end
 
+        @testset "3D multi-tile splineTransform! (RLZ / RLR / SLR)" begin
+            # A tiled B → A round trip must reproduce the physical field on
+            # EVERY tile. The regression this guards: sub-blocks are aligned by
+            # wavenumber but an inner tile carries fewer wavenumbers than the
+            # patch, so the patch z-level stride (b_iDim_p * (1 + 2*patch_kDim))
+            # differs from the tile's — using the tile's stride for the patch
+            # base index reads the wrong shared rows on every z level past the
+            # first. The seeded field carries wavenumber-1 content so the k >= 1
+            # blocks are exercised, not just the k = 0 spline.
+            function tiled_roundtrip_3d(geometry)
+                cheb_vert = geometry in ("RLZ", "SLZ")
+                gp = SpringsteelGridParameters(;
+                    geometry = geometry,
+                    num_cells = 9,
+                    iMin = 0.0,
+                    iMax = startswith(geometry, "S") ? Float64(π) : 75.0,
+                    kMin = 0.0, kMax = 10.0,
+                    (cheb_vert ? (kDim = 6,) : (num_cells_k = 4,))...,
+                    vars = Dict("u" => 1),
+                    BCL = Dict("u" => CubicBSpline.R0),
+                    BCR = Dict("u" => CubicBSpline.R0),
+                    BCB = Dict("u" => cheb_vert ? Chebyshev.R0 : CubicBSpline.R0),
+                    BCT = Dict("u" => cheb_vert ? Chebyshev.R0 : CubicBSpline.R0))
+                patch = createGrid(gp)
+                pts = getGridpoints(patch)
+                rMax = patch.params.iMax
+                zMax = patch.params.kMax
+                f(r, l, z) = sin(π * r / rMax) * (0.5 + z / zMax) *
+                             (1.0 + 0.5 * (r / rMax) * cos(l))
+                for i in 1:size(pts, 1)
+                    patch.physical[i, 1, 1] = f(pts[i, 1], pts[i, 2], pts[i, 3])
+                end
+                spectralTransform!(patch)
+
+                shared = SharedArray{Float64}(size(patch.spectral))
+                shared[:, :] .= patch.spectral
+
+                tiles = calcTileSizes(patch, 3)
+                worst_val = 0.0
+                for tile in tiles
+                    splineTransform!(shared, patch, tile)
+                    tile_phys = zeros(Float64, size(tile.physical))
+                    tileTransform!(shared, tile, tile_phys, tile.spectral)
+                    tpts = getGridpoints(tile)
+                    for i in 1:size(tpts, 1)
+                        worst_val = max(worst_val,
+                            abs(tile_phys[i, 1, 1] - f(tpts[i, 1], tpts[i, 2], tpts[i, 3])))
+                    end
+                    @test all(isfinite, tile_phys)
+                end
+                return worst_val
+            end
+
+            # Smooth field over 9 cells: cubic-spline accuracy. An inner-tile
+            # stride bug produces O(1) garbage, far outside this tolerance.
+            @test tiled_roundtrip_3d("RLZ") < 5e-3
+            @test tiled_roundtrip_3d("RLR") < 5e-3
+            @test tiled_roundtrip_3d("SLR") < 5e-3
+        end
+
         @testset "Multi-dim tiling" begin
 
             # ── 2D tiling on RR_Grid (3 i-tiles × 2 j-tiles = 6 tiles) ──────