Revise and optimise PMSL calculation

interfaces/physics_schemes_interface/source/algorithm/pmsl_alg_mod.x90

@@ -16,13 +16,15 @@ module pmsl_alg_mod
   use level_heights_mod,       only: eta_theta_levels
   use sci_geometric_constants_mod,   &
                                only: get_height_fv, get_panel_id, get_dx_at_w2
-  use fs_continuity_mod,       only: Wtheta, W3
+  use fs_continuity_mod,       only: Wtheta, W3, W2
   use function_space_collection_mod, only: function_space_collection
+  use function_space_mod,            only: function_space_type
   use mesh_mod,                only: mesh_type
   use initialise_diagnostics_mod, only: init_diag => init_diagnostic_field
   use planet_constants_mod, only: recip_kappa_def
   use planet_config_mod, only: p_zero
   use um_domain_init_mod, only: n_iter_pmsl
+  use transport_config_mod, only: halo_size => dep_pt_stencil_extent
 
   implicit none
   private
@@ -49,6 +51,7 @@ contains
     use pmsl_fgeo_kernel_mod, only: pmsl_fgeo_kernel_type
     use pmsl_ffunc_kernel_mod, only: pmsl_ffunc_kernel_type
     use pmsl_solve_kernel_mod, only: pmsl_solve_kernel_type
+    use pmsl_copy_kernel_mod, only: pmsl_copy_kernel_type
 
     implicit none
 
@@ -62,6 +65,7 @@ contains
     ! Define PMSL field
     type( field_type ) :: pmsl, pmsl_unsmooth, t_at_mean_sea_level
     type( field_type ) :: f_vg, f_ug, exner_m, f_func, exner_sm
+    type( field_type ) :: height_wth_halo, dx_at_w2_halo
 
     ! Internal variables
 
@@ -71,7 +75,10 @@ contains
     type( field_type ), pointer     :: dx_at_w2 => null()
     type( field_type ), pointer     :: panel_id => null()
 
-    integer(kind=i_def) :: k, levelupper, n
+    type(function_space_type), pointer :: vector_space
+    type(mesh_type), pointer :: mesh
+
+    integer(kind=i_def) :: k, levelupper, n, m
 
     logical :: pmsl_flag, pmsl_unsmooth_flag
 
@@ -80,6 +87,7 @@ contains
 
     real(kind=r_def), parameter :: upperheight=2000.0_r_def
     integer(kind=i_def), parameter :: sten_size = 1
+    integer(kind=i_def) :: sten_size_big
 
     if ( subroutine_timers ) call timer('pmsl_alg')
 
@@ -118,9 +126,12 @@ contains
       ! Calculate the geostrohpic forcing functions used for the smoothed PMSL
       call pmsl%copy_field_properties(f_vg)
       call pmsl%copy_field_properties(f_ug)
-      call pmsl%copy_field_properties(exner_m)
-      call pmsl%copy_field_properties(f_func)
-      call pmsl%copy_field_properties(exner_sm)
+
+      vector_space => function_space_collection % get_fs(twod_mesh, 0, 0, W3)
+      call exner_m%initialise(vector_space, halo_depth=(halo_size+1))
+      call f_func%initialise(vector_space, halo_depth=(halo_size+1))
+      call exner_sm%initialise(vector_space, halo_depth=(halo_size+1))
+
       dx_at_w2 => get_dx_at_w2(theta_wth%get_mesh_id())
       panel_id => get_panel_id(theta_wth%get_mesh_id())
 
@@ -133,13 +144,25 @@ contains
                                          f_ug, sten_size, panel_id, sten_size, &
                                          f_func) )
 
-      do n = 1, n_iter_pmsl
-        ! Solve for new PMSL value
-        call invoke(pmsl_solve_kernel_type(dx_at_w2, f_func, exner_m, &
-                                           sten_size, height_wth, exner_sm) )
-
-        ! Set old value to new value
-        call invoke(setval_X(exner_m, exner_sm))
+      mesh => exner_w3%get_mesh()
+      vector_space => function_space_collection % get_fs(mesh, 0, 0, Wtheta)
+      call height_wth_halo%initialise(vector_space, halo_depth=(halo_size+1))
+      vector_space => function_space_collection % get_fs(mesh, 0, 0, W2)
+      call dx_at_w2_halo%initialise(vector_space, halo_depth=(halo_size+1))
+      call invoke(setval_X(height_wth_halo, height_wth), &
+                  setval_X(dx_at_w2_halo, dx_at_w2) )
+
+      do n = 1, n_iter_pmsl/(halo_size+1)
+        do m = halo_size,0,-1
+          sten_size_big = m+sten_size
+          ! Solve for new PMSL value
+          call invoke(pmsl_solve_kernel_type(dx_at_w2_halo, f_func, exner_m, &
+                                             sten_size_big, height_wth_halo, &
+                                             exner_sm, m) )
+
+          ! Set old value to new value
+          call invoke(pmsl_copy_kernel_type(exner_m, exner_sm, m))
+        end do
       end do
 
       ! Convert exner to PMSL

interfaces/physics_schemes_interface/source/kernel/pmsl_copy_kernel_mod.F90

@@ -0,0 +1,67 @@
+!-------------------------------------------------------------------------------
+! (c) Crown copyright 2026 Met Office. All rights reserved.
+! The file LICENCE, distributed with this code, contains details of the terms
+! under which the code may be used.
+!-------------------------------------------------------------------------------
+!> @brief Holds code to copy PMSL in halo
+
+module pmsl_copy_kernel_mod
+
+  use argument_mod,         only: arg_type,                  &
+                                  GH_FIELD, GH_SCALAR,       &
+                                  GH_READ, GH_WRITE,         &
+                                  GH_REAL, OWNED_AND_HALO_CELL_COLUMN,      &
+                                  ANY_DISCONTINUOUS_SPACE_1, &
+                                  STENCIL, CROSS2D
+  use fs_continuity_mod,    only: WTHETA, W2
+  use constants_mod,        only: r_def, i_def
+  use kernel_mod,           only: kernel_type
+
+  implicit none
+
+  private
+
+  !> Kernel metadata for Psyclone
+  type, public, extends(kernel_type) :: pmsl_copy_kernel_type
+    private
+    type(arg_type) :: meta_args(2) = (/                                    &
+         arg_type(GH_FIELD, GH_REAL, GH_READ,  ANY_DISCONTINUOUS_SPACE_1), & ! exner_in
+         arg_type(GH_FIELD, GH_REAL, GH_WRITE, ANY_DISCONTINUOUS_SPACE_1)  &
+         /) ! exner_out
+    integer :: operates_on = OWNED_AND_HALO_CELL_COLUMN
+  contains
+    procedure, nopass :: pmsl_copy_code
+  end type pmsl_copy_kernel_type
+
+  public :: pmsl_copy_code
+
+contains
+
+  !> @brief Copy field including a depth in the halo
+  !> @param[in]     nlayers       The number of layers
+  !> @param[in]     exner_in      Initial guess for exner at MSL
+  !> @param[in,out] exner_out     Next guess for exner at MSL
+  !> @param[in]     ndf_2d        Number of degrees of freedom per cell for 2d fields
+  !> @param[in]     undf_2d       Number of total degrees of freedom for 2d fields
+  !> @param[in]     map_2d        Dofmap for the cell at the base of the column for 2d fields
+  subroutine pmsl_copy_code(nlayers,                         &
+                            exner_in,                        &
+                            exner_out,                       &
+                            ndf_2d, undf_2d, map_2d)
+
+    implicit none
+
+    ! Arguments added automatically in call to kernel
+    integer(kind=i_def), intent(in) :: nlayers
+    integer(kind=i_def), intent(in) :: ndf_2d, undf_2d
+    integer(kind=i_def), intent(in), dimension(ndf_2d)  :: map_2d
+
+    ! Arguments passed explicitly from algorithm
+    real(kind=r_def),    intent(in), dimension(undf_2d) :: exner_in
+    real(kind=r_def),    intent(inout), dimension(undf_2d) :: exner_out
+
+    exner_out(map_2d(1)) = exner_in(map_2d(1))
+
+  end subroutine pmsl_copy_code
+
+end module pmsl_copy_kernel_mod

interfaces/physics_schemes_interface/source/kernel/pmsl_solve_kernel_mod.F90

@@ -10,7 +10,7 @@ module pmsl_solve_kernel_mod
   use argument_mod,         only: arg_type,                  &
                                   GH_FIELD, GH_SCALAR,       &
                                   GH_READ, GH_WRITE,         &
-                                  GH_REAL, CELL_COLUMN,      &
+                                  GH_REAL, OWNED_AND_HALO_CELL_COLUMN,      &
                                   ANY_DISCONTINUOUS_SPACE_1, &
                                   STENCIL, CROSS2D
   use fs_continuity_mod,    only: WTHETA, W2
@@ -31,7 +31,7 @@ module pmsl_solve_kernel_mod
          arg_type(GH_FIELD, GH_REAL, GH_READ,  WTHETA),                    & ! height_wth
          arg_type(GH_FIELD, GH_REAL, GH_WRITE, ANY_DISCONTINUOUS_SPACE_1)  &
          /) ! exner_out
-    integer :: operates_on = CELL_COLUMN
+    integer :: operates_on = OWNED_AND_HALO_CELL_COLUMN
   contains
     procedure, nopass :: pmsl_solve_code
   end type pmsl_solve_kernel_type
@@ -107,13 +107,13 @@ subroutine pmsl_solve_code(nlayers,                         &
     ! Calculate which cell in the x branch of the stencil to use
     ! This sets the point to use to be the stencil point (2) if it exists,
     ! or the centre point (1) if it doesn't (i.e. we are at a domain edge)
-    xp1 = smap_2d_size(3)
-    xm1 = smap_2d_size(1)
+    xp1 = min(2,smap_2d_size(3))
+    xm1 = min(2,smap_2d_size(1))
     ! Calculate which cell in the y branch of the stencil to use
     ! This sets the point to use to be the stencil point (2) if it exists,
     ! or the centre point (1) if it doesn't (i.e. we are at a domain edge)
-    yp1 = smap_2d_size(4)
-    ym1 = smap_2d_size(2)
+    yp1 = min(2,smap_2d_size(4))
+    ym1 = min(2,smap_2d_size(2))
 
     ! Only calculated above a certain height or if stencil point exists
     if (height_wth(map_wth(1)) > pmsl_smooth_height .and. &