Fortran API: computing stencils for StructuredGrid

doc/CMakeLists.txt

@@ -6,6 +6,10 @@
 # granted to it by virtue of its status as an intergovernmental organisation nor
 # does it submit to any jurisdiction.
 
+if(atlas_HAVE_FORTRAN)
+  add_subdirectory(example-fortran)
+endif()
+
 if(atlas_HAVE_DOCS)
 
 add_subdirectory(user-guide)

doc/example-fortran/CMakeLists.txt

@@ -0,0 +1,10 @@
+# (C) Copyright 2024- ECMWF.
+#
+# This software is licensed under the terms of the Apache Licence Version 2.0
+# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+# In applying this licence, ECMWF does not waive the privileges and immunities
+# granted to it by virtue of its status as an intergovernmental organisation nor
+# does it submit to any jurisdiction.
+
+add_subdirectory(structured-grid-stencils)
+

doc/example-fortran/structured-grid-stencils/CMakeLists.txt

@@ -0,0 +1,12 @@
+# (C) Copyright 2024- ECMWF.
+#
+# This software is licensed under the terms of the Apache Licence Version 2.0
+# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+# In applying this licence, ECMWF does not waive the privileges and immunities
+# granted to it by virtue of its status as an intergovernmental organisation nor
+# does it submit to any jurisdiction.
+
+ecbuild_add_executable( TARGET example-fortran-structured-grid-stencils
+                        OUTPUT_NAME program NOINSTALL
+                        SOURCES program.F90
+                        LIBS atlas_f )

doc/example-fortran/structured-grid-stencils/program.F90

@@ -0,0 +1,208 @@
+! (C) Copyright 2023 ECMWF.
+!
+! This software is licensed under the terms of the Apache Licence Version 2.0
+! which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+! In applying this licence, ECMWF does not waive the privileges and immunities
+! granted to it by virtue of its status as an intergovernmental organisation nor
+! does it submit to any jurisdiction.
+
+! ------------------------------------------------------------------------------------------------
+! Example program on how to access and use StructuredColumns functionspace
+! Also includes computation of stencils around a given coordinate
+! ------------------------------------------------------------------------------------------------
+
+subroutine run
+
+use, intrinsic :: iso_fortran_env, only : real64, real32
+use atlas_module, only :   &
+  & atlas_StructuredGrid  ,&
+  & atlas_functionspace_StructuredColumns ,&
+  & atlas_Field, &
+  & ATLAS_KIND_IDX, &
+  & atlas_StructuredGrid_ComputeNorth ,&
+  & atlas_StructuredGrid_ComputeWest ,&
+  & atlas_StructuredGrid_ComputeStencil ,&
+  & atlas_StructuredGrid_Stencil, &
+  & atlas_real
+
+use fckit_mpi_module, only : fckit_mpi
+
+! ------------------------------------------------
+implicit none
+
+character(len=:), allocatable :: gridname
+type(atlas_StructuredGrid) :: grid
+type(atlas_functionspace_StructuredColumns) :: functionspace
+real(real64) :: zlon, zlat
+integer(ATLAS_KIND_IDX) :: jlon,jlat, jgp
+real(real64), pointer :: xy(:,:)
+integer :: log_rank
+
+! ------------------------------------------------
+
+log_rank = 1 ! First partition has rank=0 !!!
+if (fckit_mpi%size() == 1) then
+  log_rank = 0
+endif
+
+gridname = "O320"
+
+! Create grid and write some properties
+grid          = atlas_StructuredGrid(gridname)
+
+if (fckit_mpi%rank() == log_rank) then
+    write(0,*) "grid%name()  : ", grid%name()
+    write(0,*) "grid%size()  : ", grid%size()
+    write(0,*) "grid%ny()    : ", grid%ny()
+    write(0,*) "grid%nxmax() : ", grid%nxmax()
+endif
+
+! Create functionspace of grid, distributed across MPI tasks with default partitioner, and with halo of 2 wide
+!   The default partitioner is "equal_regions"
+functionspace = atlas_functionspace_StructuredColumns(grid, halo=2)
+
+! Latitude bounds without and with halos of this partition
+if (fckit_mpi%rank() == log_rank) then
+    write(0,*) "functionspace%size_owned() : ", functionspace%size_owned()
+    write(0,*) "functionspace%size()       : ", functionspace%size()
+    write(0,'(A,I0,A,I0)') " halo points between index ", functionspace%size_owned()+1 , " and " , functionspace%size()
+    write(0,*) "functionspace%j_begin(),      functionspace%j_end()      : ", &
+      & functionspace%j_begin(), functionspace%j_end()
+    write(0,*) "functionspace%j_begin_halo(), functionspace%j_end_halo() : ", &
+      & functionspace%j_begin_halo(), functionspace%j_end_halo()
+endif
+
+! Longitude bounds without and with halos of this partition for the first latitude
+jlat = functionspace%j_begin()
+if (fckit_mpi%rank() == log_rank) then
+  write(0,*) "functionspace%i_begin(jlat),      functionspace%i_end(jlat)      : ", functionspace%i_begin(jlat),     &
+   & functionspace%i_end(jlat)
+  write(0,*) "functionspace%i_begin_halo(jlat), functionspace%i_end_halo(jlat) : ", functionspace%i_begin_halo(jlat),&
+   & functionspace%i_end_halo(jlat)
+endif
+
+! Access i,j indices of grid points
+block
+  type(atlas_Field) :: field_index_i, field_index_j
+  integer(ATLAS_KIND_IDX), pointer :: index_i(:), index_j(:)
+  field_index_i = functionspace%index_i()
+  field_index_j = functionspace%index_j()
+  call field_index_i%data(index_i)
+  call field_index_j%data(index_j)
+  if (fckit_mpi%rank() == log_rank) then
+    write(0,*) "i,j of first partition point :", index_i(1), index_j(1)
+  endif
+  call field_index_i%final()
+  call field_index_j%final()
+end block
+
+! Access to xy coordinates
+block
+  type(atlas_Field) :: field_xy
+  field_xy = functionspace%xy()
+  call field_xy%data(xy)
+  call field_xy%final()
+end block
+
+! Creating a horizontal field and perform halo exchange
+block
+  type(atlas_Field) :: field
+  real(real32), pointer :: view(:)
+  field = functionspace%create_field(name="myfield", kind=atlas_real(real32))
+  call field%data(view)
+  if (fckit_mpi%rank() == log_rank) then
+    write(0,*) "shape( horizontal field )", shape(view)
+  endif
+  view(1:functionspace%size_owned()) = 1.
+  view(functionspace%size_owned()+1:functionspace%size()) = 0
+  call field%set_dirty() ! This marks that halos are in "dirty" state
+  call field%halo_exchange() ! Only exchanges halos when halos are marked as "dirty"
+  if (fckit_mpi%rank() == log_rank) then
+    write(0,*) "halo exhange success : ", all(view == 1.)
+  endif
+  call field%final()
+end block
+
+! Creating a horizontal/vertical field
+block
+  type(atlas_Field) :: field
+  real(real32), pointer :: view(:,:)
+  field = functionspace%create_field(name="myfield", kind=atlas_real(real32), levels=10)
+  call field%data(view)
+  if (fckit_mpi%rank() == log_rank) then
+    write(0,*) "shape( horizontal/vertical field )", shape(view)
+  endif
+  call field%final()
+end block
+
+! Set a coordinate somewhere south-east of the first grid point, but still within this partition
+jgp = 1
+zlon = xy(1,jgp) + 0.1
+zlat = xy(2,jgp) - 0.1
+
+! Compute nearest points to the north-west of a coordinate
+block
+  type(atlas_StructuredGrid_ComputeNorth) :: compute_north
+  type(atlas_StructuredGrid_ComputeWest) :: compute_west
+  compute_north = atlas_StructuredGrid_ComputeNorth(grid, halo=2)
+  compute_west  = atlas_StructuredGrid_ComputeWest(grid,  halo=2)
+
+  jlat = compute_north%execute(zlat)
+  jlon = compute_west%execute(zlon, jlat)
+  jgp = functionspace%index(jlon,jlat) ! gridpoint index in 1-dimensional array
+
+  if (fckit_mpi%rank() == log_rank) then
+  write(0,'(A,F6.2,A,I0)') "compute_north%execute(y=",zlat,") : ", jlat
+  write(0,'(A,F6.2,A,I0,A,I0)') "compute_west%execute(x=",zlon,", j=", jlat,") : ", jlon
+  write(0,'(A,I0,A,F6.2,A,F6.2,A)') "gridpoint north-west: jpg=",jgp,"    (lon,lat)=(", xy(1,jgp), ",", xy(2,jgp), ")"
+  endif
+  call compute_west%final()
+  call compute_north%final()
+end block
+
+! Stencil computer of 4x4 stencil around a coordinate
+block
+  type(atlas_StructuredGrid_ComputeStencil) :: compute_stencil
+  type(atlas_StructuredGrid_Stencil) :: stencil
+  call compute_stencil%setup(grid, stencil_width=4)
+  call compute_stencil%execute(zlon,zlat,stencil)
+  if (fckit_mpi%rank() == log_rank) then
+    write(0,'(A,A,dt)') 'stencil:', new_line('a'), stencil
+  endif
+  if (fckit_mpi%rank() == log_rank) then
+    write(0,'(A,A,dt)') 'stencil gridpoints:'
+    do jlat=1,stencil%width
+      do jlon=1,stencil%width
+        write(0,'(I8)',advance='no') functionspace%index(stencil%i(jlon,jlat),stencil%j(jlat))
+      enddo
+      write(0,'(A)') ""
+    enddo
+
+    write(0,'(A,A,dt)') 'stencil coordinates:'
+    do jlat=1,stencil%width
+      jgp = functionspace%index(stencil%i(1,jlat),stencil%j(jlat))
+      write(0,'(A, F6.2,A)', advance='no') " lat : ", xy(2,jgp), "  ---  lon : "
+      do jlon=1,stencil%width
+        jgp = functionspace%index(stencil%i(jlon,jlat),stencil%j(jlat))
+        write(0,'(F8.2)',advance='no') xy(1,jgp)
+      enddo
+      write(0,'(A)') ""
+    enddo
+
+  endif
+  call compute_stencil%final()
+end block
+
+call functionspace%final()
+call grid%final()
+end subroutine
+
+! ------------------------------------------------------------------------------------------------
+
+program main
+use atlas_module, only : atlas_library
+implicit none
+call atlas_library%initialize()
+call run()
+call atlas_library%finalize()
+end program

src/atlas/CMakeLists.txt

@@ -173,6 +173,8 @@ grid/detail/grid/RegionalVariableResolution.h
 grid/detail/grid/RegionalVariableResolution.cc
 grid/detail/grid/Healpix.h
 grid/detail/grid/Healpix.cc
+grid/detail/grid/StencilComputerInterface.h
+grid/detail/grid/StencilComputerInterface.cc
 
 grid/detail/tiles/Tiles.cc
 grid/detail/tiles/Tiles.h

src/atlas/grid/detail/grid/StencilComputerInterface.cc

@@ -0,0 +1,48 @@
+/*
+ * (C) Copyright 2023 ECMWF.
+ *
+ * This software is licensed under the terms of the Apache Licence Version 2.0
+ * which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+ * In applying this licence, ECMWF does not waive the privileges and immunities
+ * granted to it by virtue of its status as an intergovernmental organisation
+ * nor does it submit to any jurisdiction.
+ */
+
+#include "StencilComputerInterface.h"
+
+extern "C" {
+    atlas::grid::ComputeNorth* atlas__grid__ComputeNorth__new(const atlas::StructuredGrid::Implementation* grid, int halo) {
+        return new atlas::grid::ComputeNorth{atlas::StructuredGrid{grid}, halo};
+    }
+
+    void atlas__grid__ComputeNorth__delete(atlas::grid::ComputeNorth* This) {
+        delete This;
+    }
+
+    atlas::idx_t atlas__grid__ComputeNorth__execute_real32(const atlas::grid::ComputeNorth* This, float y) {
+        return This->operator()(y);
+    }
+
+    atlas::idx_t atlas__grid__ComputeNorth__execute_real64(const atlas::grid::ComputeNorth* This, double y) {
+        return This->operator()(y);
+    }
+
+
+    atlas::grid::ComputeWest* atlas__grid__ComputeWest__new(const atlas::StructuredGrid::Implementation* grid, int halo) {
+        return new atlas::grid::ComputeWest{atlas::StructuredGrid{grid}, halo};
+    }
+
+    void atlas__grid__ComputeWest__delete(atlas::grid::ComputeWest* This) {
+        delete This;
+    }
+
+    atlas::idx_t atlas__grid__ComputeWest__execute_real32(const atlas::grid::ComputeWest* This, float x, atlas::idx_t j) {
+        return This->operator()(x, j);
+    }
+
+    atlas::idx_t atlas__grid__ComputeWest__execute_real64(const atlas::grid::ComputeWest* This, double x, atlas::idx_t j) {
+        return This->operator()(x, j);
+    }
+
+
+}

src/atlas/grid/detail/grid/StencilComputerInterface.h

@@ -0,0 +1,27 @@
+/*
+ * (C) Copyright 2023 ECMWF.
+ *
+ * This software is licensed under the terms of the Apache Licence Version 2.0
+ * which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+ * In applying this licence, ECMWF does not waive the privileges and immunities
+ * granted to it by virtue of its status as an intergovernmental organisation
+ * nor does it submit to any jurisdiction.
+ */
+
+#pragma once
+
+#include "atlas/grid/StencilComputer.h"
+#include "atlas/grid/StructuredGrid.h"
+
+extern "C" {
+    atlas::grid::ComputeNorth* atlas__grid__ComputeNorth__new(const atlas::StructuredGrid::Implementation* grid, int halo);
+    void atlas__grid__ComputeNorth__delete(atlas::grid::ComputeNorth* This);
+    atlas::idx_t atlas__grid__ComputeNorth__execute_real32(const atlas::grid::ComputeNorth* This, float y);
+    atlas::idx_t atlas__grid__ComputeNorth__execute_real64(const atlas::grid::ComputeNorth* This, double y);
+
+    atlas::grid::ComputeWest* atlas__grid__ComputeWest__new(const atlas::StructuredGrid::Implementation* grid, int halo);
+    void atlas__grid__ComputeWest__delete(atlas::grid::ComputeWest* This);
+    atlas::idx_t atlas__grid__ComputeWest__execute_real32(const atlas::grid::ComputeWest* This, float x, atlas::idx_t j);
+    atlas::idx_t atlas__grid__ComputeWest__execute_real64(const atlas::grid::ComputeWest* This, double x, atlas::idx_t j);
+
+}

src/atlas_f/CMakeLists.txt

@@ -49,7 +49,7 @@ function(generate_fortran_bindings output filename)
     OUTPUT ${outfile}
     COMMAND ${PYTHON_EXECUTABLE} ${PROJECT_SOURCE_DIR}/tools/c2f.py ${CMAKE_CURRENT_SOURCE_DIR}/${filename}
             -o ${outfile} -m ${_PAR_MODULE}
-            -t '{"idx_t":"${F_IDX}","gidx_t":"${F_GIDX}"}'
+            -t '{"idx_t":"${F_IDX}","gidx_t":"${F_GIDX}","atlas::idx_t":"${F_IDX}","atlas::gidx_t":"${F_GIDX}"}'
     DEPENDS ${filename} )
   set_source_files_properties(${outfile} PROPERTIES GENERATED TRUE)
 endfunction()
@@ -137,6 +137,10 @@ generate_fortran_bindings(FORTRAN_BINDINGS ../atlas/redistribution/detail/Redist
     MODULE    atlas_redistribution_c_binding
     OUTPUT    redistribution_c_binding.f90)
 
+generate_fortran_bindings(FORTRAN_BINDINGS ../atlas/grid/detail/grid/StencilComputerInterface.h
+    MODULE    atlas_grid_StencilComputer_c_binding
+    OUTPUT    grid_StencilComputer_c_binding.f90 )
+
 if( atlas_HAVE_ATLAS_NUMERICS )
     generate_fortran_bindings(FORTRAN_BINDINGS ../atlas/numerics/Nabla.h)
     generate_fortran_bindings(FORTRAN_BINDINGS ../atlas/numerics/Nabla.h)
@@ -230,6 +234,7 @@ ecbuild_add_library( TARGET atlas_f
         grid/atlas_GridDistribution_module.F90
         grid/atlas_Vertical_module.F90
         grid/atlas_Partitioner_module.F90
+        grid/atlas_StencilComputer_module.F90
         mesh/atlas_MeshBuilder_module.F90
         mesh/atlas_MeshGenerator_module.F90
         mesh/atlas_Mesh_module.F90

src/atlas_f/atlas_module.F90

@@ -89,6 +89,11 @@ module atlas_module
     & atlas_RegionalGrid
 use atlas_Vertical_module, only :&
     & atlas_Vertical
+use atlas_StencilComputer_module, only: &
+    & atlas_StructuredGrid_ComputeNorth, &
+    & atlas_StructuredGrid_ComputeWest, &
+    & atlas_StructuredGrid_ComputeStencil, &
+    & atlas_StructuredGrid_Stencil
 use atlas_functionspace_EdgeColumns_module, only: &
     & atlas_functionspace_EdgeColumns
 use atlas_functionspace_CellColumns_module, only: &