Momentum dynamics validation on the CubedSphere

src/Fields/field.jl

@@ -245,6 +245,10 @@ function offset_windowed_data(data, Loc, grid, indices)
         windowed_parent = parent(data)
     else
         parent_indices = map(parent_index_range, indices, loc, topo, halo)
+        if size(data)[3] == 1 && first.(axes(data))[3] == grid.Nz + 1 # take ssh into consideration
+            parent_indices = collect(parent_indices)
+            parent_indices[3] = 1:1
+        end
         windowed_parent = view(parent(data), parent_indices...)
     end
 

src/Models/HydrostaticFreeSurfaceModels/barotropic_pressure_correction.jl

@@ -11,14 +11,15 @@ const ExplicitFreeSurfaceHFSM      = HFSM{<:Any, <:Any, <:Any, <:ExplicitFreeSur
 const ImplicitFreeSurfaceHFSM      = HFSM{<:Any, <:Any, <:Any, <:ImplicitFreeSurface}
 const SplitExplicitFreeSurfaceHFSM = HFSM{<:Any, <:Any, <:Any, <:SplitExplicitFreeSurface}
 
+#=
 pressure_correct_velocities!(model::ExplicitFreeSurfaceHFSM, Δt; kwargs...) = nothing
+=#
 
 #####
 ##### Barotropic pressure correction for models with a free surface
 #####
 
-function pressure_correct_velocities!(model::ImplicitFreeSurfaceHFSM, Δt)
-
+function pressure_correct_velocities!(model::Union{ImplicitFreeSurfaceHFSM, ExplicitFreeSurfaceHFSM}, Δt)
     launch!(model.architecture, model.grid, :xyz,
             _barotropic_pressure_correction,
             model.velocities,

src/Models/HydrostaticFreeSurfaceModels/explicit_free_surface.jl

@@ -69,6 +69,9 @@ explicit_ab2_step_free_surface!(free_surface, model, Δt, χ) =
     i, j = @index(Global, NTuple)
 
     @inbounds begin
-        η[i, j, Nz+1] += Δt * ((FT(1.5) + χ) * Gηⁿ[i, j, Nz+1] - (FT(0.5) + χ) * Gη⁻[i, j, Nz+1])
+        # AB2 for η
+        # η[i, j, Nz+1] += Δt * ((FT(1.5) + χ) * Gηⁿ[i, j, Nz+1] - (FT(0.5) + χ) * Gη⁻[i, j, Nz+1])
+        # FW for η
+        η[i, j, Nz+1] += Δt * Gηⁿ[i, j, Nz+1]
     end
 end

src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_tendency_kernel_functions.jl

@@ -44,7 +44,9 @@ implicitly during time-stepping.
     model_fields = merge(hydrostatic_fields(velocities, free_surface, tracers), auxiliary_fields)
 
     return ( - U_dot_∇u(i, j, k, grid, advection, velocities)
+             #=
              - explicit_barotropic_pressure_x_gradient(i, j, k, grid, free_surface)
+             =#
              - x_f_cross_U(i, j, k, grid, coriolis, velocities)
              - ∂xᶠᶜᶜ(i, j, k, grid, hydrostatic_pressure_anomaly)
              - ∂ⱼ_τ₁ⱼ(i, j, k, grid, closure, diffusivities, clock, model_fields, buoyancy)
@@ -83,7 +85,9 @@ implicitly during time-stepping.
     model_fields = merge(hydrostatic_fields(velocities, free_surface, tracers), auxiliary_fields)
 
     return ( - U_dot_∇v(i, j, k, grid, advection, velocities)
+             #=
              - explicit_barotropic_pressure_y_gradient(i, j, k, grid, free_surface)
+             =#
              - y_f_cross_U(i, j, k, grid, coriolis, velocities)
              - ∂yᶜᶠᶜ(i, j, k, grid, hydrostatic_pressure_anomaly)
              - ∂ⱼ_τ₂ⱼ(i, j, k, grid, closure, diffusivities, clock, model_fields, buoyancy)

src/Models/HydrostaticFreeSurfaceModels/pcg_implicit_free_surface_solver.jl

@@ -5,7 +5,9 @@ using Oceananigans.Architectures
 using Oceananigans.Grids: with_halo, isrectilinear, halo_size
 using Oceananigans.Architectures: device
 
-import Oceananigans.Solvers: solve!, precondition!
+using Oceananigans.Fields: replace_horizontal_vector_halos!
+
+import Oceananigans.Solvers: solve!, precondition!, auxiliary_actions!
 import Oceananigans.Architectures: architecture
 
 """
@@ -50,7 +52,14 @@ function PCGImplicitFreeSurfaceSolver(grid::AbstractGrid, settings, gravitationa
     vertically_integrated_lateral_areas = (xᶠᶜᶜ = ∫ᶻ_Axᶠᶜᶜ, yᶜᶠᶜ = ∫ᶻ_Ayᶜᶠᶜ)
 
     @apply_regionally compute_vertically_integrated_lateral_areas!(vertically_integrated_lateral_areas)
-    fill_halo_regions!(vertically_integrated_lateral_areas)
+
+    u = vertically_integrated_lateral_areas.xᶠᶜᶜ
+    v = vertically_integrated_lateral_areas.yᶜᶠᶜ
+
+    for _ in 1:3
+        fill_halo_regions!(vertically_integrated_lateral_areas)
+        @apply_regionally replace_horizontal_vector_halos!((; u, v, w = nothing), grid, signed = false)
+    end
 
     # Set some defaults
     settings = Dict{Symbol, Any}(settings)
@@ -136,14 +145,17 @@ function implicit_free_surface_linear_operation!(L_ηⁿ⁺¹, ηⁿ⁺¹, ∫
 
     # Note: because of `fill_halo_regions!` below, we cannot use `PCGImplicitFreeSurface` on a
     # multi-region grid; `fill_halo_regions!` cannot be used within `@apply_regionally`
-    fill_halo_regions!(ηⁿ⁺¹)
 
     launch!(arch, grid, :xy, _implicit_free_surface_linear_operation!,
             L_ηⁿ⁺¹, grid,  ηⁿ⁺¹, ∫ᶻ_Axᶠᶜᶜ, ∫ᶻ_Ayᶜᶠᶜ, g, Δt)
 
     return nothing
 end
 
+ImplicitFreeSurfaceOperation = typeof(implicit_free_surface_linear_operation!)
+
+auxiliary_actions!(::ImplicitFreeSurfaceOperation, L_ηⁿ⁺¹, ηⁿ⁺¹, args...) = fill_halo_regions!(ηⁿ⁺¹)
+
 # Kernels that act on vertically integrated / surface quantities
 @inline ∫ᶻ_Ax_∂x_ηᶠᶜᶜ(i, j, k, grid, ∫ᶻ_Axᶠᶜᶜ, η) = @inbounds ∫ᶻ_Axᶠᶜᶜ[i, j, k] * ∂xᶠᶜᶠ(i, j, k, grid, η)
 @inline ∫ᶻ_Ay_∂y_ηᶜᶠᶜ(i, j, k, grid, ∫ᶻ_Ayᶜᶠᶜ, η) = @inbounds ∫ᶻ_Ayᶜᶠᶜ[i, j, k] * ∂yᶜᶠᶠ(i, j, k, grid, η)

src/Models/HydrostaticFreeSurfaceModels/update_hydrostatic_free_surface_model_state.jl

@@ -34,7 +34,10 @@ function update_state!(model::HydrostaticFreeSurfaceModel, grid, callbacks; comp
     @apply_regionally update_model_field_time_series!(model, model.clock)
 
     fill_halo_regions!(prognostic_fields(model), model.clock, fields(model); async = true)
+    #=
     @apply_regionally replace_horizontal_vector_halos!(model.velocities, model.grid)
+    =#
+    grid isa ConformalCubedSphereGrid ? fill_halo_regions!((model.velocities.u, model.velocities.v)) : nothing
     @apply_regionally compute_auxiliaries!(model)
 
     fill_halo_regions!(model.diffusivity_fields; only_local_halos = true)

src/MultiRegion/cubed_sphere_boundary_conditions.jl

@@ -23,7 +23,7 @@ end
 function fill_halo_regions!(field::CubedSphereField{<:Center, <:Center})
     grid = field.grid
 
-    Nx, Ny, Nz = size(grid)
+    Nx, Ny, Nz = size(field)
     Hx, Hy, Hz = halo_size(grid)
 
     Nx == Ny || error("horizontal grid size Nx and Ny must be the same")
@@ -32,14 +32,20 @@ function fill_halo_regions!(field::CubedSphereField{<:Center, <:Center})
     Hx == Hy || error("horizontal halo size Hx and Hy must be the same")
     Hc = Hx
 
+    if size(field[1].data)[3] == 1 && first.(axes(field[1].data))[3] == grid.Nz + 1 # take ssh into consideration
+        k_min, k_max = Nz+1, Nz+1
+    else
+        k_min, k_max = 1, Nz
+    end
+
     #-- one pass: only use interior-point values:
     for region in 1:6
 
         region_E, region_N, region_W, region_S = find_neighboring_panels(grid, region)
 
         if isodd(region)
             #- odd face number (1, 3, 5):
-            for k in -Hz+1:Nz+Hz
+            for k in k_min:k_max
                 #- E + W Halo for field:
                 field[region][Nc+1:Nc+Hc, 1:Nc, k] .=         field[region_E][1:Hc, 1:Nc, k]
                 field[region][1-Hc:0, 1:Nc, k]     .= reverse(field[region_W][1:Nc, Nc+1-Hc:Nc, k], dims=1)'
@@ -49,7 +55,7 @@ function fill_halo_regions!(field::CubedSphereField{<:Center, <:Center})
             end
         elseif iseven(region)
             #- even face number (2, 4, 6):
-            for k in -Hz+1:Nz+Hz
+            for k in k_min:k_max
                 #- E + W Halo for field:
                 field[region][Nc+1:Nc+Hc, 1:Nc, k] .= reverse(field[region_E][1:Nc, 1:Hc, k], dims=1)'
                 field[region][1-Hc:0, 1:Nc, k]     .=         field[region_W][Nc+1-Hc:Nc, 1:Nc, k]
@@ -63,10 +69,16 @@ function fill_halo_regions!(field::CubedSphereField{<:Center, <:Center})
     return nothing
 end
 
+# Halos have to be filled concurrently for pairs of XFaced and YFaced Fields, like `u` and `v`.
+# Therefore, e.g., fill_halo_regions!(u) should do nothing and the halos for `u` and `v` are filed
+# via fill_halo_regions!((u, v))
+fill_halo_regions!(field::CubedSphereField{<:Face, <:Center}) = nothing
+fill_halo_regions!(field::CubedSphereField{<:Center, <:Face}) = nothing
+
 function fill_halo_regions!(field::CubedSphereField{<:Face, <:Face})
     grid = field.grid
 
-    Nx, Ny, Nz = size(grid)
+    Nx, Ny, Nz = size(field)
     Hx, Hy, Hz = halo_size(grid)
 
     Nx == Ny || error("horizontal grid size Nx and Ny must be the same")
@@ -82,7 +94,7 @@ function fill_halo_regions!(field::CubedSphereField{<:Face, <:Face})
 
         if isodd(region)
             #- odd face number (1, 3, 5):
-            for k in -Hz+1:Nz+Hz
+            for k in 1:Nz
                 #- E + W Halo for field:
                 field[region][Nc+1:Nc+Hc, 1:Nc, k]   .=         field[region_E][1:Hc, 1:Nc, k]
                 field[region][1-Hc:0, 2:Nc+1, k]     .= reverse(field[region_W][1:Nc, Nc+1-Hc:Nc, k], dims=1)'
@@ -99,7 +111,7 @@ function fill_halo_regions!(field::CubedSphereField{<:Face, <:Face})
             end
         elseif iseven(region)
             #- even face number (2, 4, 6):
-            for k in -Hz+1:Nz+Hz
+            for k in 1:Nz
                 #- E + W Halo for field:
                 field[region][Nc+1:Nc+Hc, 2:Nc, k]   .= reverse(field[region_E][2:Nc, 1:Hc, k], dims=1)'
                 if Hc > 1
@@ -128,7 +140,7 @@ function fill_halo_regions!(field_1::CubedSphereField{<:Center, <:Center},
     field_1.grid == field_2.grid || error("fields must be on the same grid")
     grid = field_1.grid
 
-    Nx, Ny, Nz = size(grid)
+    Nx, Ny, Nz = size(field_1)
     Hx, Hy, Hz = halo_size(grid)
     signed ? plmn = -1 : plmn = 1
 
@@ -145,7 +157,7 @@ function fill_halo_regions!(field_1::CubedSphereField{<:Center, <:Center},
 
         if isodd(region)
             #- odd face number (1, 3, 5):
-            for k in -Hz+1:Nz+Hz
+            for k in 1:Nz
                 #- E Halo:
                 field_1[region][Nc+1:Nc+Hc, 1:Nc, k] .=         field_1[region_E][1:Hc, 1:Nc, k]
                 field_2[region][Nc+1:Nc+Hc, 1:Nc, k] .=         field_2[region_E][1:Hc, 1:Nc, k]
@@ -161,7 +173,7 @@ function fill_halo_regions!(field_1::CubedSphereField{<:Center, <:Center},
             end
         elseif iseven(region)
             #- even face number (2, 4, 6):
-            for k in -Hz+1:Nz+Hz
+            for k in 1:Nz
                 #- E Halo:
                 field_1[region][Nc+1:Nc+Hc, 1:Nc, k] .= reverse(field_2[region_E][1:Nc, 1:Hc, k], dims=1)'
                 field_2[region][Nc+1:Nc+Hc, 1:Nc, k] .= reverse(field_1[region_E][1:Nc, 1:Hc, k], dims=1)' * plmn
@@ -187,7 +199,7 @@ function fill_halo_regions!(field_1::CubedSphereField{<:Face, <:Center},
     field_1.grid == field_2.grid || error("fields must be on the same grid")
     grid = field_1.grid
 
-    Nx, Ny, Nz = size(grid)
+    Nx, Ny, Nz = size(field_1)
     Hx, Hy, Hz = halo_size(grid)
     signed ? plmn = -1 : plmn = 1
 
@@ -204,7 +216,7 @@ function fill_halo_regions!(field_1::CubedSphereField{<:Face, <:Center},
 
         if isodd(region)
             #- odd face number (1, 3, 5):
-            for k in -Hz+1:Nz+Hz
+            for k in 1:Nz
                 #- E + W Halo for field_1:
                 field_1[region][Nc+1:Nc+Hc, 1:Nc, k]   .=         field_1[region_E][1:Hc, 1:Nc, k]
                 field_1[region][1-Hc:0, 1:Nc, k]       .= reverse(field_2[region_W][1:Nc, Nc+1-Hc:Nc, k], dims=1)'
@@ -224,7 +236,7 @@ function fill_halo_regions!(field_1::CubedSphereField{<:Face, <:Center},
             end
         elseif iseven(region)
             #- even face number (2, 4, 6):
-            for k in -Hz+1:Nz+Hz
+            for k in 1:Nz
                 #- E + W Halo for field_1:
                 field_1[region][Nc+1:Nc+Hc, 1:Nc, k]   .= reverse(field_2[region_E][1:Nc, 1:Hc, k], dims=1)'
                 field_1[region][1-Hc:0, 1:Nc, k]       .=         field_1[region_W][Nc+1-Hc:Nc, 1:Nc, k]
@@ -249,7 +261,7 @@ function fill_halo_regions!(field_1::CubedSphereField{<:Face, <:Center},
     # (e.g., vort3(0,1) & (1,0)).
     if Hc > 1
         for region in 1:6
-            for k in -Hz+1:Nz+Hz
+            for k in 1:Nz
                 #- SW corner:
                 field_1[region][1-Hc:0, 0, k] .= field_2[region][1, 1-Hc:0, k]
                 field_2[region][0, 1-Hc:0, k] .= field_1[region][1-Hc:0, 1, k]'
@@ -275,7 +287,7 @@ function fill_halo_regions!(field_1::CubedSphereField{<:Face, <:Face},
     field_1.grid == field_2.grid || error("fields must be on the same grid")
     grid = field_1.grid
 
-    Nx, Ny, Nz = size(grid)
+    Nx, Ny, Nz = size(field_1)
     Hx, Hy, Hz = halo_size(grid)
     signed ? plmn = -1 : plmn = 1
 
@@ -292,7 +304,7 @@ function fill_halo_regions!(field_1::CubedSphereField{<:Face, <:Face},
 
         if isodd(region)
             #- odd face number (1, 3, 5):
-            for k in -Hz+1:Nz+Hz
+            for k in 1:Nz
                 #- E Halo:
                 field_1[region][Nc+1:Nc+Hc, 1:Nc, k]   .=         field_1[region_E][1:Hc, 1:Nc, k]
                 field_2[region][Nc+1:Nc+Hc, 1:Nc, k]   .=         field_2[region_E][1:Hc, 1:Nc, k]
@@ -318,7 +330,7 @@ function fill_halo_regions!(field_1::CubedSphereField{<:Face, <:Face},
             end
         else
             #- even face number (2, 4, 6):
-            for k in -Hz+1:Nz+Hz
+            for k in 1:Nz
                 #- E Halo:
                 field_1[region][Nc+1:Nc+Hc, 2:Nc, k]   .= reverse(field_2[region_E][2:Nc, 1:Hc, k], dims=1)'
                 field_2[region][Nc+1:Nc+Hc, 2:Nc+1, k] .= reverse(field_1[region_E][1:Nc, 1:Hc, k], dims=1)' * plmn

src/MultiRegion/multi_region_field.jl

@@ -24,7 +24,7 @@ const GriddedMultiRegionFieldTuple{N, T} = NTuple{N, T} where {N, T<:GriddedMult
 const GriddedMultiRegionFieldNamedTuple{S, N} = NamedTuple{S, N} where {S, N<:GriddedMultiRegionFieldTuple}
 
 # Utils
-Base.size(f::GriddedMultiRegionField) = size(getregion(f.grid, 1))
+Base.size(f::GriddedMultiRegionField) = size(getregion(f, 1))
 
 @inline isregional(f::GriddedMultiRegionField) = true
 @inline devices(f::GriddedMultiRegionField) = devices(f.grid)

src/MultiRegion/unified_implicit_free_surface_solver.jl

@@ -66,6 +66,8 @@ build_implicit_step_solver(::Val{:PreconditionedConjugateGradient}, grid::MultiR
     throw(ArgumentError("Cannot use PCG solver with Multi-region grids!! Select :Default or :HeptadiagonalIterativeSolver as solver_method"))
 build_implicit_step_solver(::Val{:Default}, grid::ConformalCubedSphereGrid, settings, gravitational_acceleration) =
     PCGImplicitFreeSurfaceSolver(grid, settings, gravitational_acceleration)
+build_implicit_step_solver(::Val{:PreconditionedConjugateGradient}, grid::ConformalCubedSphereGrid, settings, gravitational_acceleration) =
+    PCGImplicitFreeSurfaceSolver(grid, settings, gravitational_acceleration)
 build_implicit_step_solver(::Val{:HeptadiagonalIterativeSolver}, grid::ConformalCubedSphereGrid, settings, gravitational_acceleration) =
     throw(ArgumentError("Cannot use Matrix solvers with ConformalCubedSphereGrid!! Select :Default or :PreconditionedConjugateGradient as solver_method"))
 

src/Operators/vorticity_operators.jl

@@ -5,9 +5,23 @@
      ζ₃ᶠᶠᶜ(i, j, k, grid, u, v)
 
  The vertical vorticity associated with horizontal velocities ``u`` and ``v``.
- """
+"""
 @inline ζ₃ᶠᶠᶜ(i, j, k, grid, u, v) = Γᶠᶠᶜ(i, j, k, grid, u, v) / Azᶠᶠᶜ(i, j, k, grid)
 
+@inline function ζ₃ᶠᶠᶜ(i, j, k, grid::OrthogonalSphericalShellGrid{FT}, u, v) where FT
+    scaling = ifelse(on_corner(i, j, grid), convert(FT, 1), 1)
+    return scaling * Γᶠᶠᶜ(i, j, k, grid, u, v) / Azᶠᶠᶜ(i, j, k, grid)
+end
+
+# Corner
+@inline on_corner(i, j, grid) = on_south_west_corner(i, j, grid) | on_south_east_corner(i, j, grid) | on_north_west_corner(i, j, grid) | on_north_east_corner(i, j, grid)
+
+# South-west, south-east, north-west, north-east corners
+@inline on_south_west_corner(i, j, grid) = (i == 1 && j == 1)
+@inline on_south_east_corner(i, j, grid) = (i == grid.Nx+1 && j == 1)
+@inline on_north_east_corner(i, j, grid) = (i == grid.Nx+1 && j == grid.Ny+1)
+@inline on_north_west_corner(i, j, grid) = (i == 1 && j == grid.Ny+1)
+
 #####
 ##### Vertical circulation at the corners of the cubed sphere needs to treated in a special manner.
 ##### See: https://github.com/CliMA/Oceananigans.jl/issues/1584
@@ -17,26 +31,26 @@
      Γᶠᶠᶜ(i, j, k, grid, u, v)
 
  The vertical circulation associated with horizontal velocities ``u`` and ``v``.
- """
-@inline function Γᶠᶠᶜ(i, j, k, grid::OrthogonalSphericalShellGrid, u, v)
-    # South-west corner
-    if i == 1 && j == 1
-        return Δy_qᶜᶠᶜ(i, j, k, grid, v) - Δx_qᶠᶜᶜ(i, j, k, grid, u) + Δx_qᶠᶜᶜ(i, j-1, k, grid, u)
-
-    # South-east corner
-    elseif i == grid.Nx+1 && j == 1
-        return - Δy_qᶜᶠᶜ(i-1, j, k, grid, v) - Δx_qᶠᶜᶜ(i, j, k, grid, u) + Δx_qᶠᶜᶜ(i, j-1, k, grid, u)
-
-    # North-west corner
-    elseif i == 1 && j == grid.Ny+1
-        return Δy_qᶜᶠᶜ(i, j, k, grid, v) - Δx_qᶠᶜᶜ(i, j, k, grid, u) + Δx_qᶠᶜᶜ(i, j-1, k, grid, u)
-
-    # North-east corner
-    elseif i == grid.Nx+1 && j == grid.Ny+1
-        return - Δy_qᶜᶠᶜ(i-1, j, k, grid, v) - Δx_qᶠᶜᶜ(i, j, k, grid, u) + Δx_qᶠᶜᶜ(i, j-1, k, grid, u)
-
-    # Not a corner
+"""
+#=
+@inline Γᶠᶠᶜ(i, j, k, grid::OrthogonalSphericalShellGrid, u, v) = 
+    ifelse(on_south_west_corner(i, j, grid) | on_north_west_corner(i, j, grid),
+           Δy_qᶜᶠᶜ(i, j, k, grid, v) - Δx_qᶠᶜᶜ(i, j, k, grid, u) + Δx_qᶠᶜᶜ(i, j-1, k, grid, u), 
+           ifelse(on_south_east_corner(i, j, grid) | on_north_east_corner(i, j, grid),
+                  - Δy_qᶜᶠᶜ(i-1, j, k, grid, v) - Δx_qᶠᶜᶜ(i, j, k, grid, u) + Δx_qᶠᶜᶜ(i, j-1, k, grid, u), 
+                  δxᶠᶜᶜ(i, j, k, grid, Δy_qᶜᶠᶜ, v) - δyᶜᶠᶜ(i, j, k, grid, Δx_qᶠᶜᶜ, u)))
+=#
+
+function Γᶠᶠᶜ(i, j, k, grid::OrthogonalSphericalShellGrid, u, v) 
+    Hx, Hy = grid.Hx, grid.Hy
+    if i == 1 - Hx || j == 1 - Hy
+        Γ = NaN
+    elseif on_south_west_corner(i, j, grid) || on_north_west_corner(i, j, grid)
+        Γ = Δy_qᶜᶠᶜ(i, j, k, grid, v) - Δx_qᶠᶜᶜ(i, j, k, grid, u) + Δx_qᶠᶜᶜ(i, j-1, k, grid, u)
+    elseif on_south_east_corner(i, j, grid) || on_north_east_corner(i, j, grid)
+        Γ = - Δy_qᶜᶠᶜ(i-1, j, k, grid, v) + Δx_qᶠᶜᶜ(i, j-1, k, grid, u) - Δx_qᶠᶜᶜ(i, j, k, grid, u)
     else
-        return δxᶠᵃᵃ(i, j, k, grid, Δy_qᶜᶠᶜ, v) - δyᵃᶠᵃ(i, j, k, grid, Δx_qᶠᶜᶜ, u)
+        Γ = δxᶠᶠᶜ(i, j, k, grid, Δy_qᶜᶠᶜ, v) - δyᶠᶠᶜ(i, j, k, grid, Δx_qᶠᶜᶜ, u)
     end
-end
+    return Γ
+end

src/Solvers/preconditioned_conjugate_gradient_solver.jl

@@ -195,6 +195,10 @@ function iterate!(x, solver, b, args...)
 
     @apply_regionally perform_iteration!(q, p, ρ, z, solver, args...)
 
+    auxiliary_actions!(solver.linear_operation!, q, p, args...)
+
+    @apply_regionally solver.linear_operation!(q, p, args...)
+
     α = ρ / dot(p, q)
 
     @debug "PreconditionedConjugateGradientSolver $(solver.iteration), |q|: $(norm(q))"
@@ -208,6 +212,8 @@ function iterate!(x, solver, b, args...)
     return nothing
 end
 
+auxiliary_actions!(args...) = nothing
+
 """ first iteration of the PCG """
 function initialize_solution!(q, x, b, solver, args...)
     solver.linear_operation!(q, x, args...)