diff --git a/GEOSlandpert_GridComp/CMakeLists.txt b/GEOSlandpert_GridComp/CMakeLists.txt index 9e20b363..8bd7bf0c 100644 --- a/GEOSlandpert_GridComp/CMakeLists.txt +++ b/GEOSlandpert_GridComp/CMakeLists.txt @@ -2,16 +2,21 @@ esma_set_this () set (SRCS nr_ran2_gasdev.F90 nr_jacobi.F90 nr_fft.F90 - random_fields.F90 land_pert.F90 force_and_cat_progn_pert_types.F90 LDAS_PertRoutines.F90 + abstract_random_fields.F90 rectangle_random_fields.F90 sphere_random_fields.F90 + land_pert.F90 force_and_cat_progn_pert_types.F90 LDAS_PertRoutines.F90 GEOS_LandPertGridComp.F90 ) +esma_add_subdirectory(@smerfs) + esma_add_library (${this} SRCS ${SRCS} - DEPENDENCIES GEOS_LdasShared GEOSensavg_GridComp GEOSland_GridComp MAPL ${MKL_LIBRARIES} + DEPENDENCIES GEOS_LdasShared GEOSensavg_GridComp GEOSland_GridComp MAPL smerfs_f ${MKL_LIBRARIES} INCLUDES ${INC_ESMF} ${MKL_INCLUDE_DIRS}) # Enable MKL paths only when MKL was detected and the compiler supports it. if (MKL_FOUND AND NOT CMAKE_Fortran_COMPILER_ID MATCHES "NAG") target_compile_definitions(${this} PRIVATE MKL_AVAILABLE) endif () + +esma_add_subdirectory(compare_pert) diff --git a/GEOSlandpert_GridComp/LDAS_PertRoutines.F90 b/GEOSlandpert_GridComp/LDAS_PertRoutines.F90 index a351dde8..4161a421 100644 --- a/GEOSlandpert_GridComp/LDAS_PertRoutines.F90 +++ b/GEOSlandpert_GridComp/LDAS_PertRoutines.F90 @@ -163,7 +163,8 @@ subroutine read_ens_prop_inputs( & kw_xcorr_force_pert, & kw_ycorr_force_pert, & kw_tcorr_force_pert, & - kw_ccorr_force_pert & + kw_ccorr_force_pert, & + kw_use_sphere_pert & ) ! read ensemble propagation inputs from namelist file @@ -243,6 +244,9 @@ subroutine read_ens_prop_inputs( & type(force_pert_ccor_type), intent(out), optional :: & kw_ccorr_force_pert + ! Sphere-pert optional keyword output + logical, intent(out), optional :: kw_use_sphere_pert + ! ------------------------ ! locals @@ -286,6 +290,9 @@ subroutine read_ens_prop_inputs( & type(force_pert_ccor_type) :: ccorr_force_pert + ! Sphere-pert local variable + logical :: use_sphere_pert + ! Errlong variables character(len=*), parameter :: Iam = 'read_ens_prop_inputs' @@ -322,11 +329,15 @@ subroutine read_ens_prop_inputs( & xcorr_force_pert, & ycorr_force_pert, & tcorr_force_pert, & - ccorr_force_pert + ccorr_force_pert, & + use_sphere_pert root_proc = (myid ==0) + ! Initialize sphere-pert default (safe: flat-FFT path is the default) + use_sphere_pert = .false. + ! --------------------------------------------------------------------- ! ! initialize selected name list inputs @@ -488,6 +499,9 @@ subroutine read_ens_prop_inputs( & if (present(kw_ccorr_force_pert)) kw_ccorr_force_pert = ccorr_force_pert + ! sphere-pert keyword + if (present(kw_use_sphere_pert)) kw_use_sphere_pert = use_sphere_pert + ! ------------------------------------------------------------------ ! ! save driver inputs into *ens_prop_inputs.nml file @@ -804,13 +818,17 @@ subroutine get_progn_pert_param( pert_grid_l, N_progn_pert, & integer, dimension(N_progn_pert_max) :: progn_pert_select + ! Sphere-pert parameter for prognostic perturbations + logical :: use_sphere_pert + ! --------------------------------------------------------------- call get_progn_pert_inputs( pert_grid_l, & descr_progn_pert, zeromean_progn_pert, coarsen_progn_pert, & std_normal_max_progn_pert, std_progn_pert, & xcorr_progn_pert, ycorr_progn_pert, & - tcorr_progn_pert, typ_progn_pert, ccorr_progn_pert ) + tcorr_progn_pert, typ_progn_pert, ccorr_progn_pert, & + use_sphere_pert ) call get_pert_select( N_progn_pert_max, pert_grid_l, std_progn_pert, & N_progn_pert, progn_pert_select ) @@ -827,7 +845,8 @@ subroutine get_progn_pert_param( pert_grid_l, N_progn_pert, & std_normal_max_progn_pert, & std_progn_pert, xcorr_progn_pert, ycorr_progn_pert, & tcorr_progn_pert, typ_progn_pert, ccorr_progn_pert, & - progn_pert_select, progn_pert_param ) + progn_pert_select, progn_pert_param, & + use_sphere_pert=use_sphere_pert ) end if @@ -872,13 +891,17 @@ subroutine get_force_pert_param( pert_grid_l, N_force_pert, force_pert_param ) integer, dimension(N_force_pert_max) :: force_pert_select + ! Shared sphere-pert parameter for forcing perturbations + logical :: use_sphere_pert + ! --------------------------------------------------------------- call get_force_pert_inputs( pert_grid_l, & descr_force_pert, zeromean_force_pert, coarsen_force_pert, & std_normal_max_force_pert, std_force_pert, & xcorr_force_pert, ycorr_force_pert, & - tcorr_force_pert, typ_force_pert, ccorr_force_pert ) + tcorr_force_pert, typ_force_pert, ccorr_force_pert, & + use_sphere_pert ) call get_pert_select( N_force_pert_max, pert_grid_l, std_force_pert, & N_force_pert, force_pert_select ) @@ -894,7 +917,8 @@ subroutine get_force_pert_param( pert_grid_l, N_force_pert, force_pert_param ) std_normal_max_force_pert, & std_force_pert, xcorr_force_pert, ycorr_force_pert, & tcorr_force_pert, typ_force_pert, ccorr_force_pert, & - force_pert_select, force_pert_param ) + force_pert_select, force_pert_param, & + use_sphere_pert=use_sphere_pert ) end if @@ -906,7 +930,8 @@ subroutine get_force_pert_inputs( pert_grid_l, & descr_force_pert, zeromean_force_pert, coarsen_force_pert, & std_normal_max_force_pert, std_force_pert, & xcorr_force_pert, ycorr_force_pert, & - tcorr_force_pert, typ_force_pert, ccorr_force_pert ) + tcorr_force_pert, typ_force_pert, ccorr_force_pert, & + use_sphere_pert ) ! get inputs for forcing perturbations for ALL forcing ! variables (including zero standard deviations) on a grid (typically, @@ -959,7 +984,9 @@ subroutine get_force_pert_inputs( pert_grid_l, & typ_force_pert real, dimension(N_force_pert_max,N_force_pert_max), intent(out) :: & - ccorr_force_pert + ccorr_force_pert + + logical, intent(out) :: use_sphere_pert ! --------------------------------------------------------------- ! @@ -1004,6 +1031,13 @@ subroutine get_force_pert_inputs( pert_grid_l, & root_proc = (myid==0) + ! Read the shared sphere-pert option for forcing perturbations. + if (root_proc) then + call read_ens_prop_inputs(kw_echo=.false., & + kw_use_sphere_pert=use_sphere_pert) + end if + call MPI_Bcast(use_sphere_pert, 1, MPI_LOGICAL, 0, mpicomm, mpierr) + ! --------- ! ! DESCR @@ -1297,7 +1331,8 @@ subroutine get_progn_pert_inputs( pert_grid_l, & descr_progn_pert, zeromean_progn_pert, coarsen_progn_pert, & std_normal_max_progn_pert, std_progn_pert, & xcorr_progn_pert, ycorr_progn_pert, & - tcorr_progn_pert, typ_progn_pert, ccorr_progn_pert ) + tcorr_progn_pert, typ_progn_pert, ccorr_progn_pert, & + use_sphere_pert ) ! get inputs for perturbations to prognostic variables for ALL prognostic ! variables (including zero standard deviations) on a grid (typically, @@ -1348,7 +1383,9 @@ subroutine get_progn_pert_inputs( pert_grid_l, & typ_progn_pert real, dimension(N_progn_pert_max,N_progn_pert_max), intent(out) :: & - ccorr_progn_pert + ccorr_progn_pert + + logical, intent(out) :: use_sphere_pert ! --------------------------------------------------------------- ! @@ -1393,6 +1430,13 @@ subroutine get_progn_pert_inputs( pert_grid_l, & root_proc = (myid==0) + ! Read the shared sphere-pert option for prognostic perturbations. + if (root_proc) then + call read_ens_prop_inputs(kw_echo=.false., & + kw_use_sphere_pert=use_sphere_pert) + end if + call MPI_Bcast(use_sphere_pert, 1, MPI_LOGICAL, 0, mpicomm, mpierr) + ! ------- ! ! DESCR @@ -1686,13 +1730,15 @@ subroutine assemble_pert_param( & std_normal_max_pert, & std_pert, xcorr_pert, ycorr_pert, tcorr_pert, typ_pert, & ccorr_pert, & - pert_select, pert_param ) + pert_select, pert_param, & + use_sphere_pert ) ! assemble perturbation parameters for the prognostics or forcings ! that have been selected to be subject to perturbations ! ! reichle, 30 Nov 2001 ! reichle, 27 May 2005 - redesign + ! wjiang, 2026-07 - added optional sphere-pert arguments ! ! ----------------------------------------------------------------- @@ -1725,6 +1771,10 @@ subroutine assemble_pert_param( & type(pert_param_type), dimension(:), pointer :: pert_param ! output + ! Optional sphere-pert arguments: + ! use_sphere_pert : .true. => use SMERFS for all selected fields in this call + logical, intent(in), optional :: use_sphere_pert + ! -------------------------- ! local variables @@ -1738,8 +1788,14 @@ subroutine assemble_pert_param( & character(len=400) :: err_msg character(len=*), parameter :: Iam = 'assemble_pert_param' + logical :: do_sphere_pert + ! ----------------------------------------------------------------- + ! Determine whether sphere-pert is requested for this call + do_sphere_pert = .false. + if (present(use_sphere_pert)) do_sphere_pert = use_sphere_pert + ! extract info into pert_param of length N_pert ! (only those states listed in pert_select are affected by perturbations) @@ -1765,8 +1821,10 @@ subroutine assemble_pert_param( & pert_param(m)%tcorr = tcorr_pert(k) pert_param(m)%typ = nint(typ_pert(k)) + ! Sphere-pert fields (safe defaults ensure flat-FFT when not requested) + pert_param(m)%use_sphere_pert = do_sphere_pert - n = 0 + n = 0 do l=1,N_pert_max diff --git a/GEOSlandpert_GridComp/abstract_random_fields.F90 b/GEOSlandpert_GridComp/abstract_random_fields.F90 new file mode 100644 index 00000000..3e8f8ae0 --- /dev/null +++ b/GEOSlandpert_GridComp/abstract_random_fields.F90 @@ -0,0 +1,68 @@ +module abstract_random_fieldsMod + use nr_ran2_gasdev, only: NRANDSEED + implicit none + private + public :: abstract_random_fields + + type, abstract :: abstract_random_fields + character(len=:), allocatable :: ID_string + contains + procedure(generate_2d_random_field_interface), deferred :: generate_2d_Random_field + procedure(generate_white_field_interface), deferred :: generate_white_field + procedure(finalize_random_field_interface), deferred :: finalize + end type abstract_random_fields + + abstract interface + subroutine generate_2d_random_field_interface(this, rseed, rfield, rfield2, lx, ly, dx, dy) + import :: abstract_random_fields, NRANDSEED + class(abstract_random_fields), intent(inout) :: this + integer, intent(inout) :: rseed(NRANDSEED) + real, intent(out) :: rfield(:,:), rfield2(:,:) + real, intent(in) :: lx, ly, dx, dy + end subroutine generate_2d_random_field_interface + + subroutine generate_white_field_interface(this, rseed, rfield) + import :: abstract_random_fields, NRANDSEED + class(abstract_random_fields), intent(inout) :: this + integer, intent(inout) :: rseed(NRANDSEED) + real, target,intent(out) :: rfield(:,:) + end subroutine generate_white_field_interface + + subroutine finalize_random_field_interface(this, rc) + import :: abstract_random_fields + class(abstract_random_fields), intent(inout) :: this + integer, optional, intent(out) :: rc + end subroutine finalize_random_field_interface + end interface +end module abstract_random_fieldsMod + +module StringAbstractRandom_fieldsMapMod + use abstract_random_fieldsMod + +#include "types/key_deferredLengthString.inc" +#define _value class (abstract_random_fields) +#define _value_allocatable +#define _value_equal_defined + +! Work around Intel assignment handling for polymorphic map components. +#define _ASSIGN(dest,src) allocate(dest%key,source=src%key); if(allocated(src%value)) allocate(dest%value,source=src%value) +#define _MOVE(dest,src) call move_alloc(from=src%key,to=dest%key); if (allocated(src%value)) call move_alloc(from=src%value,to=dest%value) +#define _FREE(x) deallocate(x%key,x%value) + +#define _map StringAbstractRandom_fieldsMap +#define _iterator StringAbstractRandom_fieldsMapIterator +#define _alt + +#include "templates/map.inc" + +#undef _alt +#undef _iterator +#undef _map +#undef _value +#undef _value_allocatable +#undef _key +#undef _value_equal_defined +#undef _ASSIGN +#undef _MOVE +#undef _FREE +end module StringAbstractRandom_fieldsMapMod diff --git a/GEOSlandpert_GridComp/compare_pert/CMakeLists.txt b/GEOSlandpert_GridComp/compare_pert/CMakeLists.txt new file mode 100644 index 00000000..9356e902 --- /dev/null +++ b/GEOSlandpert_GridComp/compare_pert/CMakeLists.txt @@ -0,0 +1,33 @@ +# ----------------------------------------------------------------------- +# compare_pert/CMakeLists.txt +# +# Builds the MPI comparison program for flat lat-lon FFT and sphere-native +# SMERFS perturbations on a 360x180 grid. +# +# Executable: +# compare_pert.x - runs both perturbation paths and writes binary output +# +# Python script (installed to ${CMAKE_INSTALL_PREFIX}/bin): +# compare_pert.py - reads both .bin outputs, produces comparison figure +# +# Both executables link against GEOSlandpert_GridComp, which transitively +# carries MAPL, ESMF, MKL, smerfs_f, MPI, netCDF, and HDF5. +# The MKL_AVAILABLE compile definition must be repeated here because the +# parent library applies it only as PRIVATE. +# ----------------------------------------------------------------------- + +# ---- Combined comparison executable ------------------------------------ +ecbuild_add_executable( + TARGET compare_pert.x + SOURCES compare_pert.F90 + LIBS GEOSlandpert_GridComp MAPL) + +# ---- MKL compile definition (matches parent library treatment) --------- +if (MKL_FOUND AND NOT CMAKE_Fortran_COMPILER_ID MATCHES "NAG") + target_compile_definitions(compare_pert.x PRIVATE MKL_AVAILABLE) +endif () + +# ---- Install Python comparison script ---------------------------------- +install( + PROGRAMS compare_pert.py + DESTINATION bin) diff --git a/GEOSlandpert_GridComp/compare_pert/compare_pert.F90 b/GEOSlandpert_GridComp/compare_pert/compare_pert.F90 new file mode 100644 index 00000000..2b6148e1 --- /dev/null +++ b/GEOSlandpert_GridComp/compare_pert/compare_pert.F90 @@ -0,0 +1,117 @@ +program compare_pert + use mpi + use, intrinsic :: iso_c_binding, only: c_double + use nr_ran2_gasdev, only: NRANDSEED, init_randseed + use rectangle_random_fieldsMod, only: rectangle_random_fields + use sphere_random_fields_mod, only: sphere_random_fields + use LDAS_PertTypes, only: pert_param_type + use LDAS_TileCoordType, only: grid_def_type + implicit none + + integer, parameter :: N_LON = 720, N_LAT = 360 + real, parameter :: DLON = 360.0 / real(N_LON), DLAT = 180.0 / real(N_LAT) + real, parameter :: XCORR = 10.0, YCORR = 10.0 + integer, parameter :: RSEED0 = -777 + integer :: mpierr, rank, npes + integer :: seed_template(NRANDSEED) + + call MPI_Init(mpierr) + call MPI_Comm_rank(MPI_COMM_WORLD, rank, mpierr) + call MPI_Comm_size(MPI_COMM_WORLD, npes, mpierr) + if (rank == 0) write(*,'(a,i0)') '=== compare_pert, MPI ranks: ', npes + + ! Initialize one RNG state and give an identical copy to each case. The + ! generators consume different numbers of variates internally, so they + ! must not share the same mutable state while running. + seed_template = 0 + seed_template(1) = RSEED0 + call init_randseed(seed_template) + + call run_latlon(MPI_COMM_WORLD, rank, seed_template) + call run_sphere(MPI_COMM_WORLD, rank, seed_template) + call MPI_Finalize(mpierr) + +contains + + subroutine run_latlon(comm, rank, seed_template) + integer, intent(in) :: comm, rank + integer, intent(in) :: seed_template(NRANDSEED) + integer :: mpierr, unit + integer :: rseed(NRANDSEED) + real :: rfield(N_LON,N_LAT), rfield2(N_LON,N_LAT) + real :: t0, t_init, t_draw + type(rectangle_random_fields) :: rf + type(pert_param_type) :: pp + type(grid_def_type) :: grid + + pp%xcorr = XCORR + pp%ycorr = YCORR + pp%coarsen = .false. + grid%N_lon = N_LON; grid%N_lat = N_LAT + grid%dlon = DLON; grid%dlat = DLAT + rseed = seed_template + + call MPI_Barrier(comm, mpierr); t0 = MPI_Wtime() +#ifdef MKL_AVAILABLE + rf = rectangle_random_fields(pp, grid, comm=comm) +#else + rf = rectangle_random_fields(pp, grid) +#endif + call MPI_Barrier(comm, mpierr); t_init = MPI_Wtime() - t0 + call MPI_Barrier(comm, mpierr); t0 = MPI_Wtime() + call rf%generate_2d_Random_field(rseed, rfield, rfield2, XCORR, YCORR, DLON, DLAT) + call MPI_Barrier(comm, mpierr); t_draw = MPI_Wtime() - t0 + + if (rank == 0) then + open(newunit=unit, file='latlon_pert.bin', status='replace', access='stream', form='unformatted') + write(unit) rfield; close(unit) + call print_stats('latlon', rfield) + write(*,'(a,f9.4,a,f9.4)') ' init=', t_init, ' draw=', t_draw + end if + call rf%finalize() + end subroutine run_latlon + + subroutine run_sphere(comm, rank, seed_template) + integer, intent(in) :: comm, rank + integer, intent(in) :: seed_template(NRANDSEED) + integer :: mpierr, unit + integer :: rseed(NRANDSEED) + real :: rfield(N_LON,N_LAT), rfield2(N_LON,N_LAT) + real :: t0, t_build, t_draw + type(sphere_random_fields) :: sf + type(pert_param_type) :: pp + type(grid_def_type) :: grid + + pp%xcorr = XCORR + pp%ycorr = YCORR + grid%N_lon = N_LON; grid%N_lat = N_LAT + grid%dlon = DLON; grid%dlat = DLAT + rseed = seed_template + + call MPI_Barrier(comm, mpierr); t0 = MPI_Wtime() + sf = sphere_random_fields(pp, grid, comm=comm) + call MPI_Barrier(comm, mpierr); t_build = MPI_Wtime() - t0 + call MPI_Barrier(comm, mpierr); t0 = MPI_Wtime() + call sf%generate_2d_Random_field(rseed, rfield, rfield2, XCORR, YCORR, DLON, DLAT) + call MPI_Barrier(comm, mpierr); t_draw = MPI_Wtime() - t0 + + if (rank == 0) then + open(newunit=unit, file='sphere_pert.bin', status='replace', access='stream', form='unformatted') + write(unit) rfield; close(unit) + call print_stats('sphere', rfield) + write(*,'(a,f9.4,a,f9.4)') ' build=', t_build, ' draw=', t_draw + end if + call sf%finalize() + end subroutine run_sphere + + subroutine print_stats(tag, field) + character(*), intent(in) :: tag + real, intent(in) :: field(N_LON,N_LAT) + real :: mean, std + mean = sum(field) / real(N_LON*N_LAT) + std = sqrt(sum((field-mean)**2) / real(N_LON*N_LAT-1)) + write(*,'(a,a,a,f9.5,a,f9.5,a,f9.5,a,f9.5)') ' [',tag,'] mean=',mean, & + ' std=',std,' min=',minval(field),' max=',maxval(field) + end subroutine print_stats + +end program compare_pert diff --git a/GEOSlandpert_GridComp/compare_pert/compare_pert.png b/GEOSlandpert_GridComp/compare_pert/compare_pert.png new file mode 100644 index 00000000..f1e58ca0 Binary files /dev/null and b/GEOSlandpert_GridComp/compare_pert/compare_pert.png differ diff --git a/GEOSlandpert_GridComp/compare_pert/compare_pert.py b/GEOSlandpert_GridComp/compare_pert/compare_pert.py new file mode 100644 index 00000000..3c2c2c8e --- /dev/null +++ b/GEOSlandpert_GridComp/compare_pert/compare_pert.py @@ -0,0 +1,226 @@ +#!/usr/bin/env python3 +""" +compare_pert.py + +Read latlon_pert.bin and sphere_pert.bin (both raw Fortran binary streams +of real(4), shape 720 x 360, column-major = lon-fastest), and produce: + + 1. Summary statistics table + 2. Zonal-mean profiles (mean and std vs latitude) + 3. Empirical covariance as a function of angular separation for + a sample of reference points + 4. Side-by-side global maps (lat-lon pert and sphere pert) + 5. Difference map + +Output: compare_pert.png (single figure with 5 panels) + compare_pert.txt (plain-text statistics) + +Usage: + python3 compare_pert.py + +Requires: numpy, matplotlib (cartopy optional for nicer maps) +""" + +import numpy as np +import os +import sys + +# ── parameters matching the Fortran programs ───────────────────────────────── +N_LON = 720 +N_LAT = 360 +XCORR = 10.0 # [degrees] + +# ── load binary fields ──────────────────────────────────────────────────────── +def load_field(fname): + """Load a flat Fortran stream binary of N_LON*N_LAT real(4) values.""" + data = np.fromfile(fname, dtype=np.float32) + if data.size != N_LON * N_LAT: + raise ValueError(f"{fname}: expected {N_LON*N_LAT} values, got {data.size}") + # Fortran column-major: first index (lon) varies fastest + return data.reshape(N_LON, N_LAT, order='F') # shape (lon, lat) + +for fname in ('latlon_pert.bin', 'sphere_pert.bin'): + if not os.path.exists(fname): + sys.exit(f"ERROR: {fname} not found. Run both MPI programs first.") + +ll = load_field('latlon_pert.bin') # (720, 360) +sp = load_field('sphere_pert.bin') # (720, 360) + +lons = np.linspace(-180 + 0.5, 180 - 0.5, N_LON) +lats = np.linspace(-90 + 0.5, 90 - 0.5, N_LAT) +LON2D, LAT2D = np.meshgrid(lons, lats, indexing='ij') # (360,180) + +# ── statistics ──────────────────────────────────────────────────────────────── +def stats(f, name): + return dict(name=name, + mean=float(f.mean()), std=float(f.std()), + min=float(f.min()), max=float(f.max())) + +st = {k: stats(f, k) for k, f in [('latlon', ll), ('sphere', sp)]} + +lines = [] +lines.append(f"{'':12s} {'mean':>10s} {'std':>10s} {'min':>10s} {'max':>10s}") +lines.append('-' * 60) +for k, s in st.items(): + lines.append(f"{s['name']:12s} {s['mean']:10.5f} {s['std']:10.5f}" + f" {s['min']:10.5f} {s['max']:10.5f}") +diff = ll - sp +d_st = stats(diff, 'diff (ll-sp)') +lines.append(f"\n{d_st['name']:12s} {d_st['mean']:10.5f} {d_st['std']:10.5f}" + f" {d_st['min']:10.5f} {d_st['max']:10.5f}") + +with open('compare_pert.txt', 'w') as fh: + fh.write('\n'.join(lines) + '\n') +print('\n'.join(lines)) + +# ── empirical zonal-mean and zonal-std vs latitude ─────────────────────────── +ll_zmean = ll.mean(axis=0) # (180,) +sp_zmean = sp.mean(axis=0) +ll_zstd = ll.std(axis=0) +sp_zstd = sp.std(axis=0) + +# ── empirical covariance vs angular distance ────────────────────────────────── +# Sample a few reference points; compute covariance along a meridional transect. +def empirical_cov_meridional(field, ref_lat_idx, n_lons=36): + """ + Average the covariance between a reference row (ref_lat_idx) and all + other rows, using n_lons equally spaced longitudes. + Returns (delta_lat_deg, cov) arrays. + """ + lon_step = N_LON // n_lons + cov = np.zeros(N_LAT) + ref = field[::lon_step, ref_lat_idx] # (n_lons,) + for j in range(N_LAT): + cov[j] = np.mean(ref * field[::lon_step, j]) + # normalise by cov[0] + if cov[ref_lat_idx] > 0: + cov /= cov[ref_lat_idx] + dlat = lats - lats[ref_lat_idx] + return dlat, cov + +ref_idx = N_LAT // 2 # equator +dlat_ll, cov_ll = empirical_cov_meridional(ll, ref_idx) +dlat_sp, cov_sp = empirical_cov_meridional(sp, ref_idx) + +# ── plotting ────────────────────────────────────────────────────────────────── +try: + import matplotlib + matplotlib.use('Agg') + import matplotlib.pyplot as plt + import matplotlib.colors as mcolors +except ImportError: + sys.exit("matplotlib not available; statistics written to compare_pert.txt") + +fig = plt.figure(figsize=(18, 14)) +fig.suptitle(f'Perturbation comparison: flat lat-lon FFT vs sphere SMERFS\n' + f'Grid {N_LON}x{N_LAT}, xcorr = ycorr = {XCORR}°', + fontsize=13, fontweight='bold') + +cmap = 'RdBu_r' +vmax = max(abs(ll).max(), abs(sp).max()) * 0.8 +levels = np.linspace(-vmax, vmax, 41) + +# -- Panel 1: lat-lon FFT map ------------------------------------------------ +ax1 = fig.add_subplot(3, 3, 1) +c1 = ax1.contourf(lons, lats, ll.T, levels=levels, cmap=cmap, extend='both') +fig.colorbar(c1, ax=ax1, shrink=0.8) +ax1.set_title('Flat lat-lon FFT') +ax1.set_xlabel('Longitude [°]') +ax1.set_ylabel('Latitude [°]') + +# -- Panel 2: sphere SMERFS map ----------------------------------------------- +ax2 = fig.add_subplot(3, 3, 2) +c2 = ax2.contourf(lons, lats, sp.T, levels=levels, cmap=cmap, extend='both') +fig.colorbar(c2, ax=ax2, shrink=0.8) +ax2.set_title('Sphere SMERFS') +ax2.set_xlabel('Longitude [°]') + +# -- Panel 3: difference ------------------------------------------------------- +ax3 = fig.add_subplot(3, 3, 3) +dlev = np.linspace(-vmax*0.5, vmax*0.5, 41) +c3 = ax3.contourf(lons, lats, (ll - sp).T, levels=dlev, cmap='bwr', extend='both') +fig.colorbar(c3, ax=ax3, shrink=0.8) +ax3.set_title('Difference (FFT − Sphere)') +ax3.set_xlabel('Longitude [°]') + +# -- Panel 4: zonal-mean profile ---------------------------------------------- +ax4 = fig.add_subplot(3, 3, 4) +ax4.plot(ll_zmean, lats, 'b-', label='FFT zonal mean') +ax4.plot(sp_zmean, lats, 'r--', label='Sphere zonal mean') +ax4.axvline(0, color='k', lw=0.5) +ax4.set_xlabel('Zonal mean') +ax4.set_ylabel('Latitude [°]') +ax4.set_title('Zonal mean vs latitude') +ax4.legend(fontsize=8) +ax4.grid(True, lw=0.3) + +# -- Panel 5: zonal-std profile ----------------------------------------------- +ax5 = fig.add_subplot(3, 3, 5) +ax5.plot(ll_zstd, lats, 'b-', label='FFT zonal std') +ax5.plot(sp_zstd, lats, 'r--', label='Sphere zonal std') +ax5.set_xlabel('Zonal std') +ax5.set_ylabel('Latitude [°]') +ax5.set_title('Zonal std vs latitude') +ax5.legend(fontsize=8) +ax5.grid(True, lw=0.3) + +# -- Panel 6: empirical covariance vs angular distance ------------------------ +ax6 = fig.add_subplot(3, 3, 6) +efold_line = np.exp(-np.abs(dlat_ll) / XCORR) +ax6.plot(dlat_ll, cov_ll, 'b-', label='FFT (meridional)') +ax6.plot(dlat_sp, cov_sp, 'r--', label='Sphere (meridional)') +ax6.plot(dlat_ll, efold_line, 'k:', lw=1, label=f'exp(-|Δlat|/{XCORR}°)') +ax6.axhline(np.exp(-1), color='gray', lw=0.7, ls='--', label='e⁻¹ level') +ax6.axhline(0, color='k', lw=0.5) +ax6.set_xlabel('Δ latitude [°]') +ax6.set_ylabel('Normalised covariance') +ax6.set_title(f'Meridional covariance (ref: equator)') +ax6.legend(fontsize=8) +ax6.set_xlim(-90, 90) +ax6.grid(True, lw=0.3) + +# -- Panel 7: power spectrum (variance per latitude band) -------------------- +ax7 = fig.add_subplot(3, 3, 7) +ax7.plot(ll_zstd**2, lats, 'b-', label='FFT variance') +ax7.plot(sp_zstd**2, lats, 'r--', label='Sphere variance') +ax7.set_xlabel('Zonal variance') +ax7.set_ylabel('Latitude [°]') +ax7.set_title('Variance vs latitude (isotropy test)') +ax7.legend(fontsize=8) +ax7.grid(True, lw=0.3) + +# -- Panel 8: histogram comparison ------------------------------------------- +ax8 = fig.add_subplot(3, 3, 8) +bins = np.linspace(-4, 4, 60) +ax8.hist(ll.ravel(), bins=bins, alpha=0.6, color='blue', + density=True, label='FFT') +ax8.hist(sp.ravel(), bins=bins, alpha=0.6, color='red', + density=True, label='Sphere') +xg = np.linspace(-4, 4, 200) +from math import pi, exp +ax8.plot(xg, [1/((2*pi)**0.5)*exp(-x**2/2) for x in xg], + 'k-', lw=1.5, label='N(0,1)') +ax8.set_xlabel('Value') +ax8.set_ylabel('PDF') +ax8.set_title('Marginal distribution') +ax8.legend(fontsize=8) +ax8.grid(True, lw=0.3) + +# -- Panel 9: scatter plot --------------------------------------------------- +ax9 = fig.add_subplot(3, 3, 9) +skip = 4 # downsample for readability +ax9.scatter(ll[::skip, ::skip].ravel(), sp[::skip, ::skip].ravel(), + s=1, alpha=0.2, color='purple') +lim = vmax +ax9.plot([-lim, lim], [-lim, lim], 'k-', lw=0.8) +ax9.set_xlabel('FFT value') +ax9.set_ylabel('Sphere value') +ax9.set_title('FFT vs Sphere (scatter, decimated)') +ax9.set_xlim(-lim, lim) +ax9.set_ylim(-lim, lim) +ax9.set_aspect('equal') +ax9.grid(True, lw=0.3) + +plt.tight_layout(rect=[0, 0, 1, 0.95]) +plt.savefig('compare_pert.png', dpi=130, bbox_inches='tight') +print('\nFigure saved to compare_pert.png') diff --git a/GEOSlandpert_GridComp/land_pert.F90 b/GEOSlandpert_GridComp/land_pert.F90 index ee948e0e..ff354a0b 100644 --- a/GEOSlandpert_GridComp/land_pert.F90 +++ b/GEOSlandpert_GridComp/land_pert.F90 @@ -12,7 +12,7 @@ ! (otherwise set initialize=.false.) ! ! compile line for test program: -! cpp -P -C nr_ran2_gasdev.f90 nr_ran2_gasdev.cpp.f90; cpp -P -C random_fields.f90 random_fields.cpp.f90; cpp -P -C land_pert.f90 land_pert.cpp.f90; f90 nr_ran2_gasdev.cpp.f90 random_fields.cpp.f90 land_pert.cpp.f90 +! cpp -P -C nr_ran2_gasdev.f90 nr_ran2_gasdev.cpp.f90; cpp -P -C rectangle_random_fields.f90 rectangle_random_fields.cpp.f90; cpp -P -C land_pert.f90 land_pert.cpp.f90 ! ! reichle, 24 Jan 2005 ! reichle, 11 Feb 2005 @@ -37,8 +37,10 @@ module land_pert_routines NRANDSEED, & init_randseed - use Random_FieldsMod - use StringRandom_fieldsMapMod + use rectangle_random_fieldsMod + use sphere_random_fields_mod, only: sphere_random_fields, sphere_random_fields_id + use StringAbstractRandom_fieldsMapMod + use abstract_random_fieldsMod, only: abstract_random_fields use nr_jacobi, ONLY: & jacobi @@ -68,7 +70,7 @@ module land_pert_routines ! ********************************************************************** - type(StringRandom_fieldsMap) :: random_fieldsMap + type(StringAbstractRandom_fieldsMap) :: random_fieldsMap contains @@ -455,7 +457,7 @@ subroutine propagate_pert( & ! ! ---------------------- - integer :: i, j, m, n, Nx, Ny, Nx_fft, Ny_fft, xStride, yStride, imax, jmax + integer :: i, j, m, n, xStride, yStride, imax, jmax real :: cc, dd, xCorr, yCorr, tCorr, tmpReal, rdlon, rdlat @@ -467,7 +469,7 @@ subroutine propagate_pert( & integer :: tmpInt, xstart, xend, ystart, yend - type(random_fields), pointer :: rf + class(abstract_random_fields), pointer :: rf type(ESMF_VM) :: vm integer :: mpicomm, status @@ -504,10 +506,28 @@ subroutine propagate_pert( & end if - ! get grid parameters for generation of new random fields on - ! possibly coarsened grid + ! initialize instance rf and get its grid parameters +#ifdef MKL_AVAILABLE + rf => find_rf(pert_param(m), pert_grid_f, comm=mpicomm) +#else + rf => find_rf(pert_param(m), pert_grid_f) +#endif - call calc_fft_grid(pert_param(m), pert_grid_f, Nx, Ny, Nx_fft, Ny_fft, xStride, yStride, rdlon, rdlat) + if (pert_param(m)%use_sphere_pert) then + ! Sphere fields are generated directly on the fine lat-lon grid. + xStride = 1 + yStride = 1 + rdlon = pert_grid_f%dlon + rdlat = pert_grid_f%dlat + else + select type (rectangle => rf) + class is (rectangle_random_fields) + xStride = rectangle%xStride + yStride = rectangle%yStride + rdlon = rectangle%rdlon + rdlat = rectangle%rdlat + end select + end if ptr2rfield => rfield( 1:pert_grid_f%N_lon:xStride,1:pert_grid_f%N_lat:yStride) ptr2rfield2 => rfield2(1:pert_grid_f%N_lon:xStride,1:pert_grid_f%N_lat:yStride) @@ -523,15 +543,6 @@ subroutine propagate_pert( & stored_field = .false. - ! initialize instance rf of class random_fields - ! this needs to be done for each pert field -#ifdef MKL_AVAILABLE - ! W.J Note: hardcoded comm = mpicomm to activate parallel fft - rf => find_rf(Nx, Ny, Nx_fft, Ny_fft, comm=mpicomm ) -#else - rf => find_rf(Nx, Ny, Nx_fft, Ny_fft) -#endif - do n=1,N_ens ! generate a random field @@ -542,11 +553,11 @@ subroutine propagate_pert( & else ! spatially correlated random fields - ! NOTE: rfg2d_fft() relies on CXML math library (22 Feb 05) - ! rfg2d_fft() now relies on Intel MKL (19 Jun 13) + ! The rectangle generator uses the configured FFT backend. if (.not. stored_field) then - call rf%rfg2d_fft(Pert_rseed(:,n), ptr2rfield, ptr2rfield2, xCorr, yCorr, rdlon, rdlat) + call rf%generate_2d_Random_field(Pert_rseed(:,n), ptr2rfield, ptr2rfield2, & + xCorr, yCorr, rdlon, rdlat) stored_field = .true. else rfield = rfield2 @@ -619,57 +630,6 @@ subroutine propagate_pert( & end subroutine propagate_pert - ! ****************************************************************************************** - - subroutine calc_fft_grid(pert_param, pert_grid_f, Nx, Ny, N_x_fft, N_y_fft, xStride, yStride, rdlon, rdlat) - - type(pert_param_type), intent(in) :: pert_param - type(grid_def_type), intent(in) :: pert_grid_f - - integer, intent(out) :: Nx, Ny, N_x_fft, N_y_fft, xStride, yStride - real, intent(out) :: rdlon, rdlat - - ! local variables - - integer :: Nx_fft, Ny_fft - real, parameter :: mult_of_xcorr = 2. - real, parameter :: mult_of_ycorr = 2. - real, parameter :: coarsen_param = 0.8 - real :: xCorr, yCorr - - xCorr = pert_param%xcorr - yCorr = pert_param%ycorr - - xStride = 1 - yStride = 1 - if (pert_param%coarsen) then - xStride = max( 1, floor(coarsen_param * xCorr / pert_grid_f%dlon) ) - yStride = max( 1, floor(coarsen_param * yCorr / pert_grid_f%dlat) ) - endif - rdlon = real(xStride)*pert_grid_f%dlon - rdlat = real(yStride)*pert_grid_f%dlat - - ! NOTE: number of grid cells of coarsened grid might not evenly divide - ! that of pert_grid_f - - Nx = pert_grid_f%N_lon / xStride - Ny = pert_grid_f%N_lat / yStride - - if (mod(pert_grid_f%N_lon,xStride)>0) Nx = Nx + 1 - if (mod(pert_grid_f%N_lat,yStride)>0) Ny = Ny + 1 - - ! add minimum required correlation lengths - Nx_fft = Nx + ceiling(mult_of_xcorr*xCorr/rdlon) - Ny_fft = Ny + ceiling(mult_of_ycorr*yCorr/rdlat) - - ! ensure N_x_fft, N_y_fft are powers of two - N_x_fft = 2**ceiling(log(real(Nx_fft))/log(2.)) - N_y_fft = 2**ceiling(log(real(Ny_fft))/log(2.)) - - end subroutine - - ! ****************************************************************** - subroutine truncate_std_normal( N_x, N_y, std_normal_max, grid_data ) ! truncate a realization of standard normal variables @@ -1269,24 +1229,43 @@ end subroutine get_sqrt_corr_matrix ! ************************************************************************ - function find_rf(Nx, Ny, Nx_fft, Ny_fft, comm) result (rf) - - type(random_fields), pointer :: rf - integer, intent(in) :: Nx, Ny, Nx_fft, Ny_fft + function find_rf(pert_param, pert_grid_f, comm) result (rf) + + class(abstract_random_fields), pointer :: rf + type(pert_param_type), intent(in) :: pert_param + type(grid_def_type), intent(in) :: pert_grid_f integer, optional, intent(in) :: comm ! local variables - type(StringRandom_fieldsMapIterator) :: iter - Character(len=:), allocatable :: id_string - type(random_fields) :: rf_tmp - - id_string = i_to_string(Nx)//":"//i_to_string(Ny)//":"//i_to_string(Nx_fft)//":"//i_to_string(Ny_fft) + type(StringAbstractRandom_fieldsMapIterator) :: iter + character(len=:), allocatable :: id_string + type(rectangle_random_fields) :: rf_tmp + type(sphere_random_fields) :: sf_tmp + + if (pert_param%use_sphere_pert) then + id_string = sphere_random_fields_id(pert_param, pert_grid_f) + else + id_string = rectangle_random_fields_id(pert_param, pert_grid_f) + end if iter = random_fieldsMap%find(id_string) if (iter == random_fieldsMap%end() ) then - rf_tmp = random_fields(Nx, Ny, Nx_fft, Ny_fft, comm=comm) - call random_fieldsMap%insert(id_string, rf_tmp) - iter = random_fieldsMap%find(id_string) + if (pert_param%use_sphere_pert) then + if (present(comm)) then + sf_tmp = sphere_random_fields(pert_param, pert_grid_f, comm=comm) + else + sf_tmp = sphere_random_fields(pert_param, pert_grid_f) + end if + call random_fieldsMap%insert(id_string, sf_tmp) + else + if (present(comm)) then + rf_tmp = rectangle_random_fields(pert_param, pert_grid_f, comm=comm) + else + rf_tmp = rectangle_random_fields(pert_param, pert_grid_f) + end if + call random_fieldsMap%insert(id_string, rf_tmp) + end if + iter = random_fieldsMap%find(id_string) endif rf => iter%value() @@ -1296,8 +1275,8 @@ end function find_rf subroutine clear_rf() - type(StringRandom_fieldsMapIterator) :: iter - type(random_fields), pointer :: rf_ptr + type(StringAbstractRandom_fieldsMapIterator) :: iter + class(abstract_random_fields), pointer :: rf_ptr iter = random_fieldsMap%begin() do while (iter /= random_fieldsMap%end()) diff --git a/GEOSlandpert_GridComp/random_fields.F90 b/GEOSlandpert_GridComp/rectangle_random_fields.F90 similarity index 79% rename from GEOSlandpert_GridComp/random_fields.F90 rename to GEOSlandpert_GridComp/rectangle_random_fields.F90 index 0b37d6a7..189a5675 100644 --- a/GEOSlandpert_GridComp/random_fields.F90 +++ b/GEOSlandpert_GridComp/rectangle_random_fields.F90 @@ -1,17 +1,17 @@ ! ========================================================================= ! -! random_fields.f90 +! rectangle_random_fields.f90 ! ! random field generator in 2d: ! generate a pair of random fields in 2d with zero mean ! -! subroutines rfg2d_fft() and sqrt_gauss_spectrum are translated from +! subroutines generate_2d_Random_field() and sqrt_gauss_spectrum are translated from ! C++ code rfg2d.C written for MIT EnKF work by reichle ! (see janus:~reichle/nasa/EnKF) ! ! covariance is specified through its spectrum, so far only Gaussian ! -! IMPORTANT: read comments for function rfg2d_fft() +! IMPORTANT: read comments for generate_2d_Random_field() ! ! written for NSIPP - EnKF ! Type: f90 @@ -33,7 +33,10 @@ #include "MAPL_ErrLog.h" #include "unused_dummy.H" -module random_fieldsMod +module rectangle_random_fieldsMod + use abstract_random_fieldsMod, only: abstract_random_fields + use LDAS_PertTypes, only: pert_param_type + use LDAS_TileCoordType, only: grid_def_type #ifdef MKL_AVAILABLE use, intrinsic :: iso_c_binding, only: c_loc, c_f_pointer, c_ptr, c_sizeof, C_NULL_PTR @@ -53,15 +56,18 @@ module random_fieldsMod implicit none - private + private + public :: rectangle_random_fields_id - real, parameter :: TWO_PI = 2.*3.14159265 - real, parameter :: SQRT2 = sqrt(2.0) + real, parameter :: TWO_PI = 2.*3.14159265 + real, parameter :: SQRT2 = sqrt(2.0) - type, public :: random_fields + type, public, extends(abstract_random_fields) :: rectangle_random_fields private - integer :: N_x, N_y - integer :: N_x_fft, N_y_fft ! computed by calc_fft_grid + integer, public :: N_x, N_y + integer, public :: N_x_fft, N_y_fft + integer, public :: xStride, yStride + real, public :: rdlon, rdlat, xCorr, yCorr real, allocatable :: field1_fft(:,:), field2_fft(:,:) integer :: fft_lens(2) ! length of each dim for 2D transform #ifdef MKL_AVAILABLE @@ -81,43 +87,93 @@ module random_fieldsMod ! procedure, public :: initialize procedure, public :: finalize - procedure, public :: rfg2d_fft + procedure, public :: generate_2d_Random_field procedure, public :: generate_white_field procedure, private :: sqrt_gauss_spectrum_2d #ifdef MKL_AVAILABLE procedure, private :: win_allocate procedure, private :: win_deallocate #endif - end type random_fields + end type rectangle_random_fields - interface random_fields - module procedure new_random_fields - end interface random_fields + interface rectangle_random_fields + module procedure new_rectangle_random_fields + end interface rectangle_random_fields contains - + + function rectangle_random_fields_id(pert_param, pert_grid_f) result(id_string) + type(pert_param_type), intent(in) :: pert_param + type(grid_def_type), intent(in) :: pert_grid_f + character(len=:), allocatable :: id_string + integer :: nx, ny, nx_fft, ny_fft, xstride, ystride + real :: rdlon, rdlat + character(len=256) :: id + real, parameter :: mult_of_xcorr = 2. + real, parameter :: mult_of_ycorr = 2. + real, parameter :: coarsen_param = 0.8 + + xstride = 1 + ystride = 1 + if (pert_param%coarsen) then + xstride = max(1, floor(coarsen_param * pert_param%xcorr / pert_grid_f%dlon)) + ystride = max(1, floor(coarsen_param * pert_param%ycorr / pert_grid_f%dlat)) + endif + rdlon = real(xstride) * pert_grid_f%dlon + rdlat = real(ystride) * pert_grid_f%dlat + nx = (pert_grid_f%N_lon + xstride - 1) / xstride + ny = (pert_grid_f%N_lat + ystride - 1) / ystride + nx_fft = 2**ceiling(log(real(nx + ceiling(mult_of_xcorr*pert_param%xcorr/rdlon)))/log(2.)) + ny_fft = 2**ceiling(log(real(ny + ceiling(mult_of_ycorr*pert_param%ycorr/rdlat)))/log(2.)) + write(id,'(i0,":",i0,":",i0,":",i0,":",i0,":",i0,":",es16.8,":",es16.8)') & + nx, ny, nx_fft, ny_fft, xstride, ystride, pert_param%xcorr, pert_param%ycorr + id_string = trim(id) + end function rectangle_random_fields_id + ! constructor (set parameter values), allocate memory - function new_random_fields(Nx, Ny, Nx_fft, Ny_fft, comm, rc) result (rf) + function new_rectangle_random_fields(pert_param, pert_grid_f, comm, rc) result (rf) - ! input/output variables [NEED class(random_fields) - ! instead of type(random_fields)] - F2003 quirk?!? - type(random_fields) :: rf - integer, intent(in) :: Nx, Ny, Nx_fft, Ny_fft + ! input/output variables [NEED class(rectangle_random_fields) + ! instead of type(rectangle_random_fields)] - F2003 quirk?!? + type(rectangle_random_fields) :: rf + type(pert_param_type), intent(in) :: pert_param + type(grid_def_type), intent(in) :: pert_grid_f integer, optional, intent(in) :: comm integer, optional, intent(out) :: rc ! local variables integer :: status, ierror integer :: rank, npes, local_dim1, local_dim2, remainder - integer :: Stride(2) + integer :: Stride(2), Nx_fft, Ny_fft + real, parameter :: mult_of_xcorr = 2. + real, parameter :: mult_of_ycorr = 2. + real, parameter :: coarsen_param = 0.8 + character(len=256) :: id_string ! set obj param vals - rf%N_x = Nx - rf%N_y = Ny - - ! ensure N_x_fft, N_y_fft are powers of two - rf%N_x_fft = Nx_fft - rf%N_y_fft = Ny_fft + rf%xCorr = pert_param%xcorr + rf%yCorr = pert_param%ycorr + rf%xStride = 1 + rf%yStride = 1 + if (pert_param%coarsen) then + rf%xStride = max(1, floor(coarsen_param * rf%xCorr / pert_grid_f%dlon)) + rf%yStride = max(1, floor(coarsen_param * rf%yCorr / pert_grid_f%dlat)) + endif + rf%rdlon = real(rf%xStride) * pert_grid_f%dlon + rf%rdlat = real(rf%yStride) * pert_grid_f%dlat + rf%N_x = pert_grid_f%N_lon / rf%xStride + rf%N_y = pert_grid_f%N_lat / rf%yStride + if (mod(pert_grid_f%N_lon,rf%xStride)>0) rf%N_x = rf%N_x + 1 + if (mod(pert_grid_f%N_lat,rf%yStride)>0) rf%N_y = rf%N_y + 1 + + ! ensure N_x_fft, N_y_fft are powers of two + Nx_fft = rf%N_x + ceiling(mult_of_xcorr*rf%xCorr/rf%rdlon) + Ny_fft = rf%N_y + ceiling(mult_of_ycorr*rf%yCorr/rf%rdlat) + rf%N_x_fft = 2**ceiling(log(real(Nx_fft))/log(2.)) + rf%N_y_fft = 2**ceiling(log(real(Ny_fft))/log(2.)) + write(id_string,'(i0,":",i0,":",i0,":",i0,":",i0,":",i0,":",es16.8,":",es16.8)') & + rf%N_x, rf%N_y, rf%N_x_fft, rf%N_y_fft, rf%xStride, rf%yStride, rf%xCorr, rf%yCorr + rf%ID_string = trim(id_string) ! allocate memory @@ -138,36 +194,36 @@ function new_random_fields(Nx, Ny, Nx_fft, Ny_fft, comm, rc) result (rf) _FAIL('Parallel FFT failed') endif - call rf%win_allocate(Nx_fft, Ny_fft, _RC) + call rf%win_allocate(rf%N_x_fft, rf%N_y_fft, _RC) ! distribution of the grid for fft allocate(rf%dim1_counts(npes),rf%dim2_counts(npes)) - local_dim1 = Nx_fft/npes + local_dim1 = rf%N_x_fft/npes rf%dim1_counts = local_dim1 - remainder = mod(Nx_fft, npes) + remainder = mod(rf%N_x_fft, npes) rf%dim1_counts(1:remainder) = local_dim1 + 1 local_dim1 = rf%dim1_counts(rank+1) - local_dim2 = Ny_fft/npes + local_dim2 = rf%N_y_fft/npes rf%dim2_counts = local_dim2 - remainder = mod(Ny_fft, npes) + remainder = mod(rf%N_y_fft, npes) rf%dim2_counts(1:remainder) = local_dim2 + 1 local_dim2 = rf%dim2_counts(rank+1) status = DftiCreateDescriptor(rf%Desc_Handle_Dim1, DFTI_SINGLE,& - DFTI_COMPLEX, 1, Nx_fft ) + DFTI_COMPLEX, 1, rf%N_x_fft ) _VERIFY(status) status = DftiCreateDescriptor(rf%Desc_Handle_Dim2, DFTI_SINGLE,& - DFTI_COMPLEX, 1, Ny_fft ) + DFTI_COMPLEX, 1, rf%N_y_fft ) _VERIFY(status) ! perform local_dim2 one-dimensional transforms along 1st dimension status = DftiSetValue( rf%Desc_Handle_Dim1, DFTI_NUMBER_OF_TRANSFORMS, local_dim2 ) _VERIFY(status) - status = DftiSetValue( rf%Desc_Handle_Dim1, DFTI_INPUT_DISTANCE, Nx_fft ) + status = DftiSetValue( rf%Desc_Handle_Dim1, DFTI_INPUT_DISTANCE, rf%N_x_fft ) _VERIFY(status) - status = DftiSetValue( rf%Desc_Handle_Dim1, DFTI_OUTPUT_DISTANCE, Nx_fft ) + status = DftiSetValue( rf%Desc_Handle_Dim1, DFTI_OUTPUT_DISTANCE, rf%N_x_fft ) _VERIFY(status) status = DftiCommitDescriptor( rf%Desc_Handle_Dim1 ) _VERIFY(status) @@ -190,7 +246,7 @@ function new_random_fields(Nx, Ny, Nx_fft, Ny_fft, comm, rc) result (rf) else rf%comm = MPI_COMM_NULL ! allocate mem and init mkl dft - status = DftiCreateDescriptor(rf%Desc_Handle, DFTI_SINGLE, DFTI_COMPLEX, 2, [Nx_fft, Ny_fft]) + status = DftiCreateDescriptor(rf%Desc_Handle, DFTI_SINGLE, DFTI_COMPLEX, 2, [rf%N_x_fft, rf%N_y_fft]) _VERIFY(status) ! initialize for actual dft computation @@ -199,7 +255,7 @@ function new_random_fields(Nx, Ny, Nx_fft, Ny_fft, comm, rc) result (rf) endif #endif _RETURN(_SUCCESS) - end function new_random_fields + end function new_rectangle_random_fields ! ************************************************************************** @@ -207,7 +263,7 @@ end function new_random_fields subroutine finalize(this, rc) ! input/output variables - class(random_fields), intent(inout) :: this + class(rectangle_random_fields), intent(inout) :: this integer, optional, intent(out) :: rc ! local variable integer :: status @@ -271,7 +327,7 @@ end subroutine finalize subroutine sqrt_gauss_spectrum_2d(this, lx, ly, dx, dy) ! input/output variables - class(random_fields), intent(inout) :: this + class(rectangle_random_fields), intent(inout) :: this real, intent(in) :: lx, ly, dx, dy ! local variables @@ -331,7 +387,7 @@ end subroutine sqrt_gauss_spectrum_2d ! ---------------------------------------------------------------------- ! - ! subroutine rfg2d_fft() + ! subroutine generate_2d_Random_field() ! ! generate a pair of 2d zero-mean random fields using FFT method ! (so far only Gaussian covariance implemented) @@ -353,7 +409,7 @@ end subroutine sqrt_gauss_spectrum_2d ! generated on a grid that is two correlation lenghts bigger ! than the field on which the grid is desired. Then cut out ! fields of the necessary size. - ! This procedure is included in rfg2d_fft(). + ! This procedure is included in generate_2d_Random_field(). ! ! NOTE: The variance specified as input is the theoretical variance of ! the complex field that is obtained from the inverse fft of the @@ -382,12 +438,12 @@ end subroutine sqrt_gauss_spectrum_2d ! The individual sample variances within each pair vary from ! realization to realization. - subroutine rfg2d_fft(this, rseed, rfield, rfield2, lx, ly, dx, dy) + subroutine generate_2d_Random_field(this, rseed, rfield, rfield2, lx, ly, dx, dy) ! input/output variables - class(random_fields), intent(inout) :: this ! ffield*_fft is modified + class(rectangle_random_fields), intent(inout) :: this ! ffield*_fft is modified integer, intent(inout) :: rseed(NRANDSEED) ! nr_ran2 modifies rseed - real, dimension(this%N_x,this%N_y), intent(out) :: rfield, rfield2 + real, dimension(:,:), intent(out) :: rfield, rfield2 real, intent(in) :: lx, ly, dx, dy ! local variables @@ -566,7 +622,7 @@ subroutine rfg2d_fft(this, rseed, rfield, rfield2, lx, ly, dx, dy) rfield = SQRT2*N_xy_fft_real*this%field1_fft(1:this%N_x,1:this%N_y) rfield2 = SQRT2*N_xy_fft_real*this%field2_fft(1:this%N_x,1:this%N_y) - end subroutine rfg2d_fft + end subroutine generate_2d_Random_field @@ -586,9 +642,9 @@ subroutine generate_white_field(this, rseed, rfield) implicit none ! input/output variables - class(random_fields), intent(in) :: this + class(rectangle_random_fields), intent(inout) :: this integer, intent(inout) :: rseed(NRANDSEED) ! nr_gasdev modifies rseed - real, dimension(this%N_x,this%N_y), intent(out), target :: rfield + real, dimension(:,:), intent(out), target :: rfield ! local variables integer :: Nxy, index @@ -629,7 +685,7 @@ end subroutine quit #ifdef MKL_AVAILABLE subroutine win_allocate(this, nx, ny, rc) - class(random_fields), intent(inout) :: this + class(rectangle_random_fields), intent(inout) :: this integer, intent(in) :: nx, ny integer, optional, intent(out) :: rc complex :: dummy @@ -655,7 +711,7 @@ subroutine win_allocate(this, nx, ny, rc) end subroutine win_allocate subroutine win_deallocate(this, rc) - class(random_fields), intent(inout) :: this + class(rectangle_random_fields), intent(inout) :: this integer, optional, intent(out) :: rc integer :: status call MPI_Win_fence(0, this%win, status) @@ -668,36 +724,13 @@ subroutine win_deallocate(this, rc) end subroutine win_deallocate #endif -end module Random_fieldsMod - -module StringRandom_fieldsMapMod - use Random_fieldsMod - -#include "types/key_deferredLengthString.inc" -#define _value type (random_fields) -#define _value_equal_defined - -#define _map StringRandom_fieldsMap -#define _iterator StringRandom_fieldsMapIterator - -#define _alt - -#include "templates/map.inc" - -#undef _alt -#undef _iterator -#undef _map -#undef _value -#undef _key -#undef _value_equal_defined -end module StringRandom_fieldsMapMod - +end module rectangle_random_fieldsMod #ifdef TEST_RFG2D !program test_rfg2d ! -! use Random_fieldsMod +! use rectangle_random_fieldsMod ! use nr_ran2_gasdev ! ! implicit none @@ -716,8 +749,8 @@ end module StringRandom_fieldsMapMod ! ! character(5) :: fft_tag ! -! ! instance of random_fields -! type(random_fields) :: rf +! ! instance of rectangle_random_fields +! type(rectangle_random_fields) :: rf ! ! ! start ! RSEEDCONST = -777 @@ -747,8 +780,8 @@ end module StringRandom_fieldsMapMod ! N_e_tot = 10 ! do n_e=1,N_e_tot,2 ! -! rf = random_fields(N_x, N_y, Nx_fft, Ny_fft) -! call rf%rfg2d_fft(rseed, field1, field2, lx, ly, dx, dy) +! rf = rectangle_random_fields(N_x, N_y, Nx_fft, Ny_fft) +! call rf%generate_2d_Random_field(rseed, field1, field2, lx, ly, dx, dy) ! !call rf%generate_white_field(rseed, field1) ! call rf%finalize ! @@ -784,4 +817,3 @@ end module StringRandom_fieldsMapMod ! ======= EOF ================================================== - diff --git a/GEOSlandpert_GridComp/sphere_random_fields.F90 b/GEOSlandpert_GridComp/sphere_random_fields.F90 new file mode 100644 index 00000000..d694d868 --- /dev/null +++ b/GEOSlandpert_GridComp/sphere_random_fields.F90 @@ -0,0 +1,1414 @@ +! +! sphere_random_fields.F90 +! +! Module for generating isotropic Gaussian random fields directly on the +! sphere using the SMERFS library (Creasey & Lang 2018). +! +! Parallel strategy — Option A (m-mode distribution): +! +! build_sphere_filter (comm optional): +! Rank 0 computes the full all_innov / all_trans tables, then +! MPI_Bcast to all other ranks on the node. Every rank therefore +! holds a copy of the complete filter. The broadcast is cheap +! relative to the hypergeometric evaluation cost. +! +! draw_sphere_realisation (comm optional): +! When comm is supplied the ranks within a node split the n_m +! Fourier m-modes evenly (Option A). Each rank walks the Kalman +! filter for its strip j_lo:j_hi of m-modes independently (no +! communication needed during the walk). The resulting partial +! half-spectrum is written into a shared MPI window +! res_cplx(n_m, nz), then a fence ensures all ranks see the +! complete spectrum before the IRFFT. The IRFFT output rows are +! also distributed: each rank computes rfield_out(:, row_lo:row_hi) +! into the shared rfield window. After a final fence every rank +! returns with a complete rfield_out(nphi, nz). +! +! Without comm the subroutine falls back to the serial path. +! +! Key public interface: +! sphere_filter_type - derived type holding pre-computed filter +! build_sphere_filter - one-time setup (rank 0 computes, bcast) +! draw_sphere_realisation - per-call random field generation (parallel) +! finalize_sphere_filter - deallocate allocatables inside sf +! sphere_to_latlon - remap sphere field to equidistant lat-lon +! xcorr_deg_to_c2 - derive c2 from e-folding distance in degrees +! +! Memory layout note (inherited from SMERFS example1_f.F90): +! Arrays passed to C kernels use reversed Fortran dimension order so +! column-major Fortran byte layout matches C row-major layout. +! +! wjiang, 2026-07 +! +! ----------------------------------------------------------------------- + +module sphere_random_fields_mod + use abstract_random_fieldsMod, only: abstract_random_fields + use LDAS_PertTypes, only: pert_param_type + use LDAS_TileCoordType, only: grid_def_type + use MAPL_MathConstantsMod, only: PI => MAPL_PI_R8, MAPL_DEGREES_TO_RADIANS_R8 + + use, intrinsic :: iso_c_binding, only : c_double, c_float, c_int32_t, c_int64_t, & + c_double_complex, c_loc, & + c_f_pointer, c_ptr, & + c_sizeof + use mpi + use nr_ran2_gasdev, only: NRANDSEED, init_randseed, nr_ran2_2d, nr_gasdev + use smerfs_interface +#ifdef MKL_AVAILABLE + use MKL_DFTI +#endif + + implicit none + private + + public :: sphere_random_fields + public :: sphere_random_fields_id + + integer, parameter :: MORD_SPHERE = 2 + + + ! ----------------------------------------------------------------------- + ! + ! sphere_filter_type + ! + ! Holds all pre-computed state-space filter coefficients for one + ! choice of (nz, nphi, c0, c2), plus the MPI context needed by + ! draw_sphere_realisation so it does not need to be re-established + ! on every call. + + type :: sphere_filter_type + + ! Grid dimensions + integer :: nz = 0 ! total iso-latitude rings (S.pole to N.pole) + integer :: nphi = 0 ! phi points per ring + integer :: n_m = 0 ! nphi/2 + 1 (half-spectrum width) + integer :: uhalf = 0 ! nz/2 + 1 (upper hemisphere ring count) + integer :: Mord = MORD_SPHERE + + ! Power-spectrum coefficients: C_l = 1 / (c0 + c2*(l*(l+1))^2) + real(c_double) :: c0 = 1.0d0 + real(c_double) :: c2 = 0.0d0 + + ! Pre-computed state-space filter arrays (full set, all m-modes): + ! all_innov(Mord, Mord, uhalf, n_m) + ! all_trans(Mord, Mord, uhalf-1, n_m) + real(c_double), allocatable :: all_innov(:,:,:,:) + real(c_double), allocatable :: all_trans(:,:,:,:) + + ! cos(theta) at each upper-hemisphere ring (uhalf values, equator first) + real(c_double), allocatable :: z_pts(:) + + ! MPI context stored at build time for reuse in draw_sphere_realisation. + ! node_comm is MPI_COMM_NULL when running serially (no comm supplied). + integer :: node_comm = MPI_COMM_NULL + integer :: node_rank = 0 + integer :: node_npes = 1 + ! m-mode strip owned by this rank: [j_lo, j_hi] (1-based, inclusive) + integer :: j_lo = 1 + integer :: j_hi = 0 ! set to n_m in serial path + + ! Ring strip for IRFFT (set at build time, reused in draw) + integer :: row_lo = 1 + integer :: row_hi = 0 + integer :: n_row_local = 0 + +#ifdef MKL_AVAILABLE + ! MKL DFTI descriptor for the batched 1-D real backward FFT. + ! One descriptor per rank, created once in build_sphere_filter. + ! Transform: complex CCE half-spectrum (n_m) -> real (nphi), + ! batched over n_row_local rings. + type(DFTI_DESCRIPTOR), pointer :: dfti_irfft => null() +#endif + + logical :: initialized = .false. + + contains + end type sphere_filter_type + + type, extends(abstract_random_fields) :: sphere_random_fields + private + type(sphere_filter_type) :: filter + real(c_double), public :: c0 = 1.0d0 + real(c_double), public :: c2 = 0.0d0 + integer :: output_nlon = 0 + integer :: output_nlat = 0 + integer, public :: nz_sphere = 0 + integer, public :: nphi_sphere = 0 + contains + procedure, public :: generate_2d_Random_field + procedure, public :: generate_white_field + procedure, public :: generate_2d_Random_field_seeded + procedure, public :: sphere_to_latlon + procedure, public :: finalize + procedure, public :: initialized + end type sphere_random_fields + + ! for example1_f in smerfs, so the random numbers are the same as those generated by python + type :: mt19937_state + integer(c_int32_t) :: mt(624) + integer :: index = 625 + end type mt19937_state + + interface sphere_random_fields + module procedure new_sphere_random_fields + end interface sphere_random_fields + +contains + + function sphere_random_fields_id(pert_param, pert_grid_f) result(id_string) + type(pert_param_type), intent(in) :: pert_param + type(grid_def_type), intent(in) :: pert_grid_f + character(len=:), allocatable :: id_string + character(len=256) :: id + + write(id,'(a,":",i0,":",i0,":",es16.8)') 'sphere', & + pert_grid_f%N_lat, pert_grid_f%N_lon, pert_param%xcorr + id_string = trim(id) + end function sphere_random_fields_id + + function new_sphere_random_fields(pert_param, pert_grid_f, comm, rc, c2) result(rf) + type(sphere_random_fields) :: rf + type(pert_param_type), intent(in) :: pert_param + type(grid_def_type), intent(in) :: pert_grid_f + integer, optional, intent(in) :: comm + integer, optional, intent(out) :: rc + real(c_double), optional, intent(in) :: c2 + real(c_double) :: xcorr_deg + real(c_double) :: c2_use + character(len=:), allocatable :: id_string + + xcorr_deg = real(pert_param%xcorr, c_double) + ! Sphere perturbations always use the fine grid; no coarsening is applied. + rf%nphi_sphere = pert_grid_f%N_lon + rf%nz_sphere = pert_grid_f%N_lat + rf%output_nlon = pert_grid_f%N_lon + rf%output_nlat = pert_grid_f%N_lat + id_string = sphere_random_fields_id(pert_param, pert_grid_f) + rf%ID_string = id_string + c2_use = -1.0d0 + if (present(c2)) c2_use = c2 + if (c2_use <= 0.0d0) then + if (xcorr_deg > 0.0d0) then + call xcorr_deg_to_c2(xcorr_deg, rf%nz_sphere, c2_use) + else + c2_use = 1.0d-12 + end if + end if + rf%c2 = c2_use + if (present(comm)) then + call build_sphere_filter(rf%filter, rf%nz_sphere, rf%nphi_sphere, rf%c0, rf%c2, comm) + else + call build_sphere_filter(rf%filter, rf%nz_sphere, rf%nphi_sphere, rf%c0, rf%c2) + end if + end function new_sphere_random_fields + + subroutine generate_2d_Random_field(this, rseed, rfield, rfield2, lx, ly, dx, dy) + class(sphere_random_fields), intent(inout) :: this + integer, intent(inout) :: rseed(NRANDSEED) + real, intent(out) :: rfield(:,:), rfield2(:,:) + real, intent(in) :: lx, ly, dx, dy + real(c_double), allocatable :: field_d(:,:), field2_d(:,:) + real, allocatable :: output_field(:,:), output_field2(:,:) + real :: field_mean, field_norm + allocate(field_d(this%filter%nphi, this%filter%nz)) + allocate(field2_d(this%filter%nphi, this%filter%nz)) + call generate_2d_random_field_sphere(this%filter, rseed, field_d, field2_d, lx, ly, dx, dy) + allocate(output_field(size(rfield,1), size(rfield,2))) + allocate(output_field2(size(rfield2,1), size(rfield2,2))) + call sphere_to_latlon_field(field_d, this%filter, size(rfield,1), size(rfield,2), output_field) + call sphere_to_latlon_field(field2_d, this%filter, size(rfield2,1), size(rfield2,2), output_field2) + rfield = output_field + rfield2 = output_field2 + field_mean = sum(rfield) / real(size(rfield)) + rfield = rfield - field_mean + field_norm = sqrt(sum(rfield**2) / real(size(rfield))) + if (field_norm > 0.0) rfield = rfield / field_norm + field_mean = sum(rfield2) / real(size(rfield2)) + rfield2 = rfield2 - field_mean + field_norm = sqrt(sum(rfield2**2) / real(size(rfield2))) + if (field_norm > 0.0) rfield2 = rfield2 / field_norm + deallocate(output_field, output_field2) + deallocate(field_d, field2_d) + end subroutine generate_2d_Random_field + + subroutine generate_white_field(this, rseed, rfield) + class(sphere_random_fields), intent(inout) :: this + integer, intent(inout) :: rseed(NRANDSEED) + real, target, intent(out) :: rfield(:,:) + real :: values(2) + integer :: i, j + + do j = 1, size(rfield, 2) + do i = 1, size(rfield, 1), 2 + call nr_gasdev(rseed, values) + rfield(i,j) = values(1) + if (i + 1 <= size(rfield, 1)) then + rfield(i + 1,j) = values(2) + end if + end do + end do + end subroutine generate_white_field + + subroutine generate_2d_Random_field_seeded(this, seed, rfield, rfield2, lx, ly, dx, dy) + class(sphere_random_fields), intent(inout) :: this + integer, intent(in) :: seed + real, intent(out) :: rfield(:,:), rfield2(:,:) + real, intent(in) :: lx, ly, dx, dy + ! The Python reference uses numpy.random.RandomState(seed), whose + ! generator is MT19937. Use the same raw 32-bit stream here instead of + ! the project L'Ecuyer stream used by the general perturbation interface. + call generate_2d_random_field_mt(this, seed, rfield, rfield2, lx, ly, dx, dy) + end subroutine generate_2d_Random_field_seeded + + subroutine finalize(this, rc) + class(sphere_random_fields), intent(inout) :: this + integer, optional, intent(out) :: rc + call finalize_sphere_filter(this%filter) + if (present(rc)) rc = 0 + end subroutine finalize + + logical function initialized(this) + class(sphere_random_fields), intent(in) :: this + initialized = this%filter%initialized + end function initialized + + subroutine sphere_to_latlon(this, sphere_field, nlon_out, nlat_out, latlon_field) + class(sphere_random_fields), intent(in) :: this + real, intent(in) :: sphere_field(this%filter%nphi, this%filter%nz) + integer, intent(in) :: nlon_out, nlat_out + real, intent(out) :: latlon_field(nlon_out, nlat_out) + real(c_double), allocatable :: sphere_d(:,:) + real, allocatable :: latlon_d(:,:) + allocate(sphere_d(this%filter%nphi, this%filter%nz)) + allocate(latlon_d(nlon_out, nlat_out)) + sphere_d = real(sphere_field, c_double) + call sphere_to_latlon_field(sphere_d, this%filter, nlon_out, nlat_out, latlon_d) + latlon_field = latlon_d + deallocate(sphere_d, latlon_d) + end subroutine sphere_to_latlon + + subroutine draw_sphere_realisation_seeded(sf, seed, rfield_out) + type(sphere_filter_type), intent(in) :: sf + integer(c_int32_t), intent(in) :: seed + real(c_double), intent(out) :: rfield_out(sf%nphi, sf%nz) + integer(c_int32_t), allocatable :: rand_ints(:) + integer :: n_needed, n_got, i + integer(c_int32_t) :: state + + n_needed = sf%nz * sf%Mord * (sf%j_hi - sf%j_lo + 1) * 2 + n_got = n_needed + n_needed/32 + 64 + allocate(rand_ints(n_got)) + state = seed + do i = 1, n_got + state = 1664525_c_int32_t * state + 1013904223_c_int32_t + rand_ints(i) = state + end do + call draw_sphere_realisation(sf, rand_ints, n_got, rfield_out) + deallocate(rand_ints) + end subroutine draw_sphere_realisation_seeded + + ! ====================================================================== + ! + ! build_sphere_filter + ! + ! Pre-compute state-space filter coefficients for the isotropic GRF + ! with power spectrum C_l = 1 / (c0 + c2*(l*(l+1))^2). + ! + ! Parallel strategy: + ! When comm is supplied, rank 0 of the global communicator runs the + ! full computation (hyp_llp1_f loop + state_space_f loop) and then + ! broadcasts all_innov and all_trans to all other ranks. This is + ! cheaper than having every rank redo the expensive hypergeometric + ! evaluations. The node communicator is also split here and stored + ! in sf for reuse in draw_sphere_realisation. + ! + ! Arguments: + ! sf (out) : sphere_filter_type to be initialised + ! nz (in) : number of iso-latitude rings (even) + ! nphi (in) : number of phi points per ring (even) + ! c0 (in) : constant term of power spectrum denominator + ! c2 (in) : quartic term of power spectrum denominator + ! comm (in, optional) : MPI communicator (global, e.g. mpicomm) + + subroutine build_sphere_filter(sf, nz, nphi, c0, c2, comm) + + type(sphere_filter_type), intent(out) :: sf + integer, intent(in) :: nz, nphi + real(c_double), intent(in) :: c0, c2 + integer, optional, intent(in) :: comm + + ! --- local variables ------------------------------------------------ + + integer :: uhalf, n_m, m_max + integer :: i, iroot, imode, im_hyp, rc + integer :: global_rank, mpierr + integer :: innov_size, trans_size + + ! Partial fraction decomposition (Mord=2 poles) + complex(c_double_complex) :: roots_llp1(MORD_SPHERE) + complex(c_double_complex) :: norms_pf(MORD_SPHERE) + complex(c_double_complex) :: norm, llp1_val, lam_val, sin_pi_lam + real(c_double) :: llp1_re, llp1_im + real(c_double) :: disc_re, tmp + + ! z-points and twiddle factors + real(c_double), allocatable :: tau_p(:,:) ! (uhalf, 2*MORD_SPHERE-1) + real(c_double), allocatable :: eta_ratio(:) ! (uhalf-1) + real(c_double), allocatable :: xvals(:) ! (uhalf) + real(c_double), allocatable :: yvals(:) ! (uhalf) + real(c_double) :: theta_val + + ! Hypergeometric tables + complex(c_double_complex), allocatable :: Fmat(:,:) ! (uhalf, m_max+MORD_SPHERE+1) + complex(c_double_complex), allocatable :: Hmat(:,:) ! (uhalf, m_max+MORD_SPHERE+1) + + ! Covariance arrays + real(c_double), allocatable :: cov(:,:,:,:) ! (MORD_SPHERE,MORD_SPHERE,uhalf,0:m_max) + real(c_double), allocatable :: cross_cov(:,:,:,:) ! (MORD_SPHERE,MORD_SPHERE,uhalf-1,0:m_max) + + ! Temporary per-mode state-space arrays + real(c_double), allocatable :: innov_m(:,:,:) ! (MORD_SPHERE,MORD_SPHERE,uhalf) + real(c_double), allocatable :: trans_m(:,:,:) ! (MORD_SPHERE,MORD_SPHERE,uhalf-1) + + ! -------------------------------------------------------------------- + + uhalf = nz/2 + 1 + n_m = nphi/2 + 1 + m_max = n_m - 1 + + ! Store grid info in sf + sf%nz = nz + sf%nphi = nphi + sf%n_m = n_m + sf%uhalf = uhalf + sf%Mord = MORD_SPHERE + sf%c0 = c0 + sf%c2 = c2 + + ! --- MPI context ---------------------------------------------------- + ! Split the supplied communicator into node-local communicators. + ! Store node_comm in sf so draw_sphere_realisation can reuse it. + ! Compute the m-mode strip owned by this rank. + + if (present(comm)) then + call MPI_Comm_split_type(comm, MPI_COMM_TYPE_SHARED, 0, & + MPI_INFO_NULL, sf%node_comm, mpierr) + call MPI_Comm_rank(sf%node_comm, sf%node_rank, mpierr) + call MPI_Comm_size(sf%node_comm, sf%node_npes, mpierr) + ! Also get global rank to decide who does the filter computation + call MPI_Comm_rank(comm, global_rank, mpierr) + else + sf%node_comm = MPI_COMM_NULL + sf%node_rank = 0 + sf%node_npes = 1 + global_rank = 0 + end if + + ! m-mode strip: distribute n_m modes as evenly as possible + call compute_strip(n_m, sf%node_npes, sf%node_rank, sf%j_lo, sf%j_hi) + + ! --- Allocate filter arrays (every rank needs the full tables) ------ + allocate(sf%all_innov(MORD_SPHERE, MORD_SPHERE, uhalf, n_m)) + allocate(sf%all_trans(MORD_SPHERE, MORD_SPHERE, uhalf-1, n_m)) + + ! --- z-points ------------------------------------------------------- + allocate(sf%z_pts(uhalf)) + do i = 1, uhalf + theta_val = (real(uhalf - i, c_double) + 0.5d0) * & + (PI / real(nz, c_double)) + sf%z_pts(i) = cos(theta_val) + end do + + ! --- Only rank 0 (globally) performs the expensive computation ------ + ! All other ranks wait for the MPI_Bcast below. + + if (global_rank == 0) then + + ! Twiddle factors + allocate(tau_p (uhalf, 2*MORD_SPHERE-1)) + allocate(eta_ratio(uhalf-1)) + tau_p(:, MORD_SPHERE) = 1.0d0 + tau_p(:, MORD_SPHERE+1) = sqrt((1.0d0 - sf%z_pts) / (1.0d0 + sf%z_pts)) + tau_p(:, MORD_SPHERE-1) = 1.0d0 / tau_p(:, MORD_SPHERE+1) + do i = 2, MORD_SPHERE-1 + tau_p(:, MORD_SPHERE+i) = tau_p(:, MORD_SPHERE+i-1) * tau_p(:, MORD_SPHERE+1) + tau_p(:, MORD_SPHERE-i) = tau_p(:, MORD_SPHERE-i+1) * tau_p(:, MORD_SPHERE-1) + end do + do i = 1, uhalf-1 + eta_ratio(i) = tau_p(i+1, MORD_SPHERE+1) * tau_p(i, MORD_SPHERE-1) + end do + + allocate(xvals(uhalf)) + allocate(yvals(uhalf)) + xvals = 0.5d0 * (1.0d0 - sf%z_pts) + yvals = 0.5d0 * (1.0d0 + sf%z_pts) + + ! Partial fraction decomposition of 1/(c0 + c2*k^2), c1=0 + disc_re = -4.0d0 * c0 * c2 + if (disc_re >= 0.0d0) then + roots_llp1(1) = cmplx( sqrt(disc_re)/(2.0d0*c2), 0.0d0, c_double_complex) + roots_llp1(2) = cmplx(-sqrt(disc_re)/(2.0d0*c2), 0.0d0, c_double_complex) + else + tmp = sqrt(-disc_re) / (2.0d0*c2) + roots_llp1(1) = cmplx(0.0d0, tmp, c_double_complex) + roots_llp1(2) = cmplx(0.0d0, -tmp, c_double_complex) + end if + do iroot = 1, MORD_SPHERE + norms_pf(iroot) = 1.0d0 / & + (2.0d0 * cmplx(c2, 0.0d0, c_double_complex) * roots_llp1(iroot)) + end do + + ! Covariance accumulators + allocate(cov (MORD_SPHERE, MORD_SPHERE, uhalf, 0:m_max)) + allocate(cross_cov(MORD_SPHERE, MORD_SPHERE, uhalf-1, 0:m_max)) + allocate(Fmat(uhalf, m_max+MORD_SPHERE+1)) + allocate(Hmat(uhalf, m_max+MORD_SPHERE+1)) + allocate(innov_m(MORD_SPHERE, MORD_SPHERE, uhalf)) + allocate(trans_m(MORD_SPHERE, MORD_SPHERE, uhalf-1)) + cov = 0.0d0 + cross_cov = 0.0d0 + + ! Hypergeometric evaluation + covariance accumulation + do iroot = 1, MORD_SPHERE + llp1_val = roots_llp1(iroot) + llp1_re = real(llp1_val, c_double) + llp1_im = aimag(llp1_val) + norm = norms_pf(iroot) + + if (llp1_re < -0.25d0) then + lam_val = cmplx(-0.5d0, 0.0d0, c_double_complex) + & + cmplx(0.0d0, 1.0d0, c_double_complex) * sqrt(-0.25d0 - llp1_val) + else + lam_val = cmplx(-0.5d0, 0.0d0, c_double_complex) - & + sqrt(cmplx(0.25d0, 0.0d0, c_double_complex) + llp1_val) + end if + sin_pi_lam = sin(PI * lam_val) + norm = -(0.25d0 / PI) * norm * PI / sin_pi_lam + + do im_hyp = 0, m_max + MORD_SPHERE + call hyp_llp1_f(llp1_re, llp1_im, im_hyp, uhalf, xvals, Fmat(:,im_hyp+1), rc) + if (rc /= 0) then + write(*,'(a,i0)') 'build_sphere_filter: hyp_llp1_f F failed at im_hyp=', im_hyp + stop 1 + end if + call hyp_llp1_f(llp1_re, llp1_im, im_hyp, uhalf, yvals, Hmat(:,im_hyp+1), rc) + if (rc /= 0) then + write(*,'(a,i0)') 'build_sphere_filter: hyp_llp1_f H failed at im_hyp=', im_hyp + stop 1 + end if + end do + + call update_cov_f(m_max, uhalf, MORD_SPHERE, & + real(norm,c_double), aimag(norm), & + llp1_re, llp1_im, & + Fmat, Hmat, tau_p, eta_ratio, & + cov, cross_cov, rc) + if (rc /= 0) then + write(*,'(a,i0)') 'build_sphere_filter: update_cov_f failed, rc=', rc + stop 1 + end if + end do ! iroot + + ! State-space decomposition for each m-mode + do imode = 0, n_m-1 + call state_space_f(uhalf, MORD_SPHERE, & + cross_cov(:,:,:,imode), & + cov(:,:,:,imode), & + innov_m, trans_m, rc) + if (rc /= 0) then + write(*,'(a,i0,a,i0)') & + 'build_sphere_filter: state_space_f failed at imode=', imode, ', rc=', rc + stop 1 + end if + sf%all_innov(:,:,:,imode+1) = innov_m + sf%all_trans(:,:,:,imode+1) = trans_m + end do + + deallocate(tau_p, eta_ratio, xvals, yvals) + deallocate(Fmat, Hmat, cov, cross_cov, innov_m, trans_m) + + end if ! global_rank == 0 + + ! --- Broadcast filter tables to all ranks --------------------------- + ! Use the global communicator (comm if present, else no-op since + ! global_rank == 0 is the only rank). + if (present(comm)) then + innov_size = MORD_SPHERE * MORD_SPHERE * uhalf * n_m + trans_size = MORD_SPHERE * MORD_SPHERE * (uhalf-1) * n_m + call MPI_Bcast(sf%all_innov, innov_size, MPI_DOUBLE_PRECISION, 0, comm, mpierr) + call MPI_Bcast(sf%all_trans, trans_size, MPI_DOUBLE_PRECISION, 0, comm, mpierr) + call MPI_Bcast(sf%z_pts, uhalf, MPI_DOUBLE_PRECISION, 0, comm, mpierr) + end if + + ! --- Pre-compute the ring strip owned by this rank for IRFFT -------- + ! Store in sf so draw_sphere_realisation can reuse it without + ! recomputing compute_strip every call. + call compute_strip(nz, sf%node_npes, sf%node_rank, sf%row_lo, sf%row_hi) + sf%n_row_local = sf%row_hi - sf%row_lo + 1 + +#ifdef MKL_AVAILABLE + ! --- Create MKL DFTI descriptor for batched 1-D real IRFFT ---------- + ! Each rank transforms sf%n_row_local rings of length nphi (single prec). + ! + ! Transform: DFTI_REAL domain, DFTI_BACKWARD (complex half-spec -> real) + ! Storage: DFTI_NOT_INPLACE, separate input/output buffers. + ! Input = complex(c_float) half-spectrum, length n_m per ring. + ! Output = real(c_float), length nphi per ring. + ! + ! Batching over n_row_local rings: + ! DFTI_NUMBER_OF_TRANSFORMS = n_row_local + ! DFTI_INPUT_DISTANCE = n_m (complex elements between ring starts) + ! DFTI_OUTPUT_DISTANCE = nphi (real elements between ring starts) + ! + ! DFTI_COMPLEX_STORAGE defaults to DFTI_COMPLEX_COMPLEX for real domain + ! backward transforms, meaning the input is treated as n_m complex values + ! (the CCE half-spectrum). No explicit storage-format call is needed. + block + integer :: dfti_stat + if (sf%n_row_local > 0) then + dfti_stat = DftiCreateDescriptor(sf%dfti_irfft, & + DFTI_SINGLE, DFTI_REAL, 1, nphi) + if (dfti_stat /= DFTI_NO_ERROR) then + write(*,'(a,i0)') & + 'build_sphere_filter: DftiCreateDescriptor failed, stat=', dfti_stat + stop 1 + end if + ! In-place real backward FFT: input and output share the same buffer. + ! work(:) is real(c_float), size 2*n_m*n_row_local. + ! MKL CCE in-place format: + ! INPUT_DISTANCE = distance between ring starts in complex elements = n_m + ! OUTPUT_DISTANCE = distance between ring starts in real elements = 2*n_m + dfti_stat = DftiSetValue(sf%dfti_irfft, & + DFTI_NUMBER_OF_TRANSFORMS, sf%n_row_local) + dfti_stat = DftiSetValue(sf%dfti_irfft, DFTI_INPUT_DISTANCE, n_m) + dfti_stat = DftiSetValue(sf%dfti_irfft, DFTI_OUTPUT_DISTANCE, 2*n_m) + dfti_stat = DftiCommitDescriptor(sf%dfti_irfft) + if (dfti_stat /= DFTI_NO_ERROR) then + write(*,'(a,i0)') & + 'build_sphere_filter: DftiCommitDescriptor failed, stat=', dfti_stat + stop 1 + end if + end if + end block +#endif + + sf%initialized = .true. + + end subroutine build_sphere_filter + + ! ====================================================================== + ! + ! draw_sphere_realisation (Option A parallel implementation) + ! + ! Generate one realisation of the isotropic GRF on the sphere. + ! + ! Parallel strategy (Option A — m-mode distribution): + ! + ! The n_m Fourier m-modes are split evenly across the node_npes ranks + ! of the node communicator stored in sf (established at build time). + ! Each rank: + ! 1. Generates noise only for its m-mode strip j_lo:j_hi. + ! 2. Runs the Kalman walk independently for those m-modes. + ! 3. Writes res_cplx(j_lo:j_hi, :) into a shared MPI window. + ! After MPI_Win_fence, all ranks see the complete res_cplx(n_m, nz). + ! The IRFFT rows (nz rings) are then distributed across ranks: + ! Each rank transforms rfield_out(:, row_lo:row_hi) into a shared + ! rfield window. After a second fence every rank holds the full + ! rfield_out(nphi, nz). + ! + ! When sf%node_comm == MPI_COMM_NULL (serial build / no comm), the + ! subroutine runs the original serial path without any MPI calls. + ! + ! Arguments: + ! sf (in) : initialised sphere_filter_type + ! rand_ints (in) : array of random 32-bit integers + ! Serial: size >= n_rand_got (all modes) + ! Parallel: size >= strip_rand_got (local strip) + ! n_rand_got (in) : length of rand_ints supplied + ! rfield_out (out) : real(c_double) field (sf%nphi, sf%nz) + ! On return every rank holds the complete field. + + subroutine draw_sphere_realisation(sf, rand_ints, n_rand_got, rfield_out) + + type(sphere_filter_type), intent(in) :: sf + integer, intent(in) :: n_rand_got + integer(c_int32_t), intent(in) :: rand_ints(n_rand_got) + real(c_double), intent(out) :: rfield_out(sf%nphi, sf%nz) + + ! --- local variables ------------------------------------------------ + + integer :: uhalf, n_m, nz, nphi, Mord + integer :: j_lo, j_hi, n_m_local ! m-mode strip for this rank + integer :: row_lo, row_hi, n_row_local ! ring strip for IRFFT (from sf) + integer :: n_rand_needed_local + integer :: i, j, ip, iq, nn, rc + integer :: skip + logical :: parallel + + real(c_float), allocatable :: noise_real(:) ! standard normals + complex(c_float), allocatable :: noise(:,:,:) ! (Mord, n_m_local, nz) + complex(c_float), allocatable :: fp_f(:,:) ! (Mord, n_m_local) + complex(c_float), allocatable :: fp_fstart(:,:) ! (Mord, n_m_local) + complex(c_float), allocatable :: fp_fprev(:,:) ! evolving state + + ! Half-spectrum and real-space field (shared windows in parallel, + ! private allocatables in serial) + complex(c_float), allocatable, target :: res_cplx_priv(:,:) ! (n_m, nz) serial + real(c_double), allocatable, target :: rfield_priv(:,:) ! (nphi, nz) serial + + complex(c_float), pointer :: res_cplx(:,:) ! points into window or priv + real(c_double), pointer :: rfield_win(:,:) ! points into window or priv + + ! MPI shared-window variables + integer :: win_res, win_rf + integer :: mpierr + integer(MPI_ADDRESS_KIND) :: win_size + integer :: disp_unit + type(c_ptr) :: base_res, base_rf + complex(c_float) :: dummy_c + real(c_double) :: dummy_r + + ! -------------------------------------------------------------------- + + uhalf = sf%uhalf + n_m = sf%n_m + nz = sf%nz + nphi = sf%nphi + Mord = sf%Mord + j_lo = sf%j_lo + j_hi = sf%j_hi + + parallel = (sf%node_comm /= MPI_COMM_NULL) + + n_m_local = j_hi - j_lo + 1 + if (n_m_local <= 0) n_m_local = 0 ! edge case: more ranks than modes + + ! Ring strip pre-computed at build time; reuse from sf + row_lo = sf%row_lo + row_hi = sf%row_hi + n_row_local = sf%n_row_local + + ! --- 1. Allocate or map shared windows for res_cplx and rfield_out -- + + if (parallel) then + ! --- res_cplx shared window: shape (n_m, nz), complex(c_float) --- + disp_unit = 4 + win_size = 0_MPI_ADDRESS_KIND + if (sf%node_rank == 0) & + win_size = int(n_m, MPI_ADDRESS_KIND) * int(nz, MPI_ADDRESS_KIND) & + * int(c_sizeof(dummy_c), MPI_ADDRESS_KIND) + call MPI_Win_allocate_shared(win_size, disp_unit, MPI_INFO_NULL, & + sf%node_comm, base_res, win_res, mpierr) + if (sf%node_rank /= 0) & + call MPI_Win_shared_query(win_res, 0, win_size, disp_unit, base_res, mpierr) + call MPI_Win_fence(0, win_res, mpierr) + call c_f_pointer(base_res, res_cplx, [n_m, nz]) + + ! --- rfield_out shared window: shape (nphi, nz), real(c_double) --- + win_size = 0_MPI_ADDRESS_KIND + if (sf%node_rank == 0) & + win_size = int(nphi, MPI_ADDRESS_KIND) * int(nz, MPI_ADDRESS_KIND) & + * int(c_sizeof(dummy_r), MPI_ADDRESS_KIND) + call MPI_Win_allocate_shared(win_size, disp_unit, MPI_INFO_NULL, & + sf%node_comm, base_rf, win_rf, mpierr) + if (sf%node_rank /= 0) & + call MPI_Win_shared_query(win_rf, 0, win_size, disp_unit, base_rf, mpierr) + call MPI_Win_fence(0, win_rf, mpierr) + call c_f_pointer(base_rf, rfield_win, [nphi, nz]) + else + ! Serial: use private allocatables and point to them + allocate(res_cplx_priv(n_m, nz)) + allocate(rfield_priv (nphi, nz)) + res_cplx => res_cplx_priv + rfield_win => rfield_priv + n_m_local = n_m ! serial owns all modes + j_lo = 1 ; j_hi = n_m + row_lo = 1 ; row_hi = nz ; n_row_local = nz + end if + + ! --- 2. Noise generation for the local m-mode strip ----------------- + ! n_rand_needed for this rank: nz * Mord * n_m_local * 2 + ! The south-hemisphere walk reuses noise at indices uhalf+1..nz, so + ! the full nz rings of noise are required (not just uhalf). + n_rand_needed_local = nz * Mord * n_m_local * 2 + + allocate(noise_real(n_rand_needed_local)) + allocate(noise (Mord, n_m_local, nz)) + allocate(fp_f (Mord, n_m_local)) + allocate(fp_fstart(Mord, n_m_local)) + allocate(fp_fprev (Mord, n_m_local)) + + if (n_m_local > 0) then + call zigg_f(n_rand_needed_local, n_rand_got, rand_ints, noise_real, rc) + if (rc /= 0) then + write(*,'(a,i0,a)') & + 'WARNING draw_sphere_realisation: zigg exhausted, ', rc, ' samples missing' + end if + + ! Pack into complex noise: noise(ip, jlocal, i) = (re + i*im) / sqrt(2) + ! Loop over all nz rings so south-hemisphere indices (uhalf+1..nz) + ! are filled from the same noise stream as the reference Python code. + nn = 0 + do i = 1, nz + do ip = 1, Mord + do j = 1, n_m_local + nn = nn + 1 + noise(ip, j, i) = cmplx(noise_real(2*nn-1), noise_real(2*nn), c_float) & + * real(0.5d0**0.5d0, c_float) + end do + end do + end do + + ! --- 3. Kalman walk for local m-modes j_lo:j_hi ------------------- + ! Upward walk: equator (ring 1 in sf%z_pts) -> pole (ring uhalf) + + ! Initialise at ring 1 + do j = 1, n_m_local + do ip = 1, Mord + fp_f(ip, j) = cmplx(0.0, 0.0, c_float) + do iq = 1, Mord + fp_f(ip, j) = fp_f(ip, j) & + + real(sf%all_innov(iq, ip, 1, j_lo+j-1), c_float) * noise(iq, j, 1) + end do + end do + end do + fp_fstart = fp_f + fp_fprev = fp_f + res_cplx(j_lo:j_hi, uhalf) = fp_f(1, :) + + do i = 1, uhalf-1 + do j = 1, n_m_local + do ip = 1, Mord + fp_f(ip, j) = cmplx(0.0, 0.0, c_float) + do iq = 1, Mord + fp_f(ip, j) = fp_f(ip, j) & + + real(sf%all_trans(iq, ip, i, j_lo+j-1), c_float) * fp_fprev(iq, j) & + + real(sf%all_innov(iq, ip, i+1, j_lo+j-1), c_float) * noise(iq, j, i+1) + end do + end do + end do + fp_fprev = fp_f + res_cplx(j_lo:j_hi, uhalf - i) = fp_f(1, :) + end do + + ! Reflect across the equator + do ip = 1, Mord + if (mod(ip-1, 2) == 0) then + fp_f(ip, :) = fp_fstart(ip, :) + else + fp_f(ip, :) = -fp_fstart(ip, :) + end if + end do + fp_fprev = fp_f + + ! For even nz, the equatorial transition was already represented by + ! the northern walk, so skip that transition on the southward walk. + skip = 1 - mod(nz, 2) + + do i = 0, nz - uhalf - 1 + do j = 1, n_m_local + do ip = 1, Mord + fp_f(ip, j) = cmplx(0.0, 0.0, c_float) + do iq = 1, Mord + fp_f(ip, j) = fp_f(ip, j) & + + real(sf%all_trans(iq, ip, i+skip+1, j_lo+j-1), c_float) * fp_fprev(iq, j) & + + real(sf%all_innov(iq, ip, i+skip+2, j_lo+j-1), c_float) * noise(iq, j, uhalf+i+1) + end do + end do + end do + fp_fprev = fp_f + res_cplx(j_lo:j_hi, uhalf + i + 1) = fp_f(1, :) + end do + + ! Scale DC mode (m=0, j_lo=1 only on the rank that owns mode 0) + if (j_lo == 1) then + do i = 1, nz + res_cplx(1, i) = cmplx(real(res_cplx(1, i)) * sqrt(2.0), 0.0, c_float) + end do + end if + + end if ! n_m_local > 0 + + ! --- 4. Fence: all ranks have written their m-mode strip ------------ + if (parallel) then + call MPI_Win_fence(0, win_res, mpierr) + end if + + ! --- 5. Distributed IRFFT: each rank transforms its ring strip ------ + ! MKL path: batched 1-D real backward FFT using the pre-built descriptor. + ! Fallback: pure-Fortran O(nphi^2) DFT (NR path or n_row_local==0). +#ifdef MKL_AVAILABLE + if (n_row_local > 0) then + block + ! In-place real backward FFT using the pre-built descriptor. + ! The work buffer holds n_row_local rings, each padded to 2*n_m + ! real(c_float) elements (= n_m complex elements). + ! MKL reads the complex half-spectrum from the first n_m complex + ! positions (= 2*n_m reals) and writes the real output (nphi reals) + ! starting at the same address. + real(c_float), allocatable, target :: work(:) ! (2*n_m * n_row_local) + type(c_ptr) :: cptr_w + complex(c_float), pointer :: cplx_view(:,:) ! (n_m, n_row_local) + real(c_float), pointer :: real_view(:,:) ! (2*n_m, n_row_local) + integer :: dfti_stat, row + + allocate(work(2 * n_m * n_row_local)) + + ! Map complex view onto work buffer, copy half-spectrum in + cptr_w = c_loc(work(1)) + call c_f_pointer(cptr_w, cplx_view, [n_m, n_row_local]) + cplx_view = res_cplx(:, row_lo:row_hi) + + ! In-place backward FFT + dfti_stat = DftiComputeBackward(sf%dfti_irfft, work) + if (dfti_stat /= DFTI_NO_ERROR) then + write(*,'(a,i0)') & + 'draw_sphere_realisation: DftiComputeBackward failed, stat=', dfti_stat + stop 1 + end if + + ! Real output occupies first nphi reals of each padded row (stride 2*n_m) + call c_f_pointer(cptr_w, real_view, [2*n_m, n_row_local]) + do row = 1, n_row_local + rfield_win(:, row_lo + row - 1) = real(real_view(1:nphi, row), c_double) + end do + + deallocate(work) + end block + end if +#else + if (n_row_local > 0) then + call irfft_rows(res_cplx(:, row_lo:row_hi), n_row_local, n_m, nphi, & + rfield_win(:, row_lo:row_hi)) + end if +#endif + + ! --- 6. Fence: all ranks have written their ring strip --------------- + if (parallel) call MPI_Win_fence(0, win_rf, mpierr) + + ! --- 7. Normalise to unit variance (global empirical) ---------------- + ! The IRFFT output has a variance that depends on nz, nphi, Mord, and + ! the power spectrum. We normalise the full field to zero mean and unit + ! variance empirically so that the std passed from land_pert.F90 can be + ! applied directly. + ! + ! IMPORTANT: this normalisation is computed and applied on the FULL + ! rfield_win (all nz rings) AFTER the fence, so all rings are scaled by + ! the same factor and the relative ring-to-ring structure is preserved. + ! Each rank computes the sum-of-squares over its own IRFFT ring strip + ! (row_lo:row_hi); MPI_Allreduce gathers the global total. Because + ! every rank holds the full rfield_win after the fence, each rank then + ! divides its strip by the global norm_fac. A second fence ensures that + ! rfield_win is fully normalised before being copied to rfield_out. + block + real(c_double) :: local_sum, global_sum, global_mean + real(c_double) :: local_sum2, global_sum2, global_var, norm_fac + integer :: ntot, mpierr2 + + if (n_row_local > 0) then + local_sum = sum(rfield_win(:, row_lo:row_hi)) + else + local_sum = 0.0d0 + end if + ntot = nphi * nz + + if (parallel) then + call MPI_Allreduce(local_sum, global_sum, 1, & + MPI_DOUBLE_PRECISION, MPI_SUM, sf%node_comm, mpierr2) + else + global_sum = local_sum + end if + global_mean = global_sum / real(ntot, c_double) + + if (n_row_local > 0) & + rfield_win(:, row_lo:row_hi) = rfield_win(:, row_lo:row_hi) - global_mean + + if (n_row_local > 0) then + local_sum2 = sum(rfield_win(:, row_lo:row_hi)**2) + else + local_sum2 = 0.0d0 + end if + + if (parallel) then + call MPI_Allreduce(local_sum2, global_sum2, 1, & + MPI_DOUBLE_PRECISION, MPI_SUM, sf%node_comm, mpierr2) + else + global_sum2 = local_sum2 + end if + + global_var = global_sum2 / real(ntot, c_double) + norm_fac = sqrt(global_var) + + if (n_row_local > 0 .and. norm_fac > 0.0d0) & + rfield_win(:, row_lo:row_hi) = rfield_win(:, row_lo:row_hi) / norm_fac + end block + + ! --- 8. Final fence + copy to caller -------------------------------- + if (parallel) then + call MPI_Win_fence(0, win_rf, mpierr) + rfield_out = rfield_win + call MPI_Win_free(win_res, mpierr) + call MPI_Win_free(win_rf, mpierr) + else + rfield_out = rfield_win + deallocate(res_cplx_priv, rfield_priv) + nullify(res_cplx, rfield_win) + end if + + deallocate(noise_real, noise, fp_f, fp_fstart, fp_fprev) + + end subroutine draw_sphere_realisation + + ! ====================================================================== + ! + ! finalize_sphere_filter + + subroutine finalize_sphere_filter(sf) + type(sphere_filter_type), intent(inout) :: sf + integer :: mpierr + + if (allocated(sf%all_innov)) deallocate(sf%all_innov) + if (allocated(sf%all_trans)) deallocate(sf%all_trans) + if (allocated(sf%z_pts)) deallocate(sf%z_pts) +#ifdef MKL_AVAILABLE + if (associated(sf%dfti_irfft)) then + mpierr = DftiFreeDescriptor(sf%dfti_irfft) + nullify(sf%dfti_irfft) + end if +#endif + if (sf%node_comm /= MPI_COMM_NULL) then + call MPI_Comm_free(sf%node_comm, mpierr) + sf%node_comm = MPI_COMM_NULL + end if + sf%initialized = .false. + end subroutine finalize_sphere_filter + + ! ====================================================================== + ! + ! sphere_to_latlon + ! + ! Remap sphere-native field rfield_sphere(nphi, nz) onto the standard + ! equidistant lat-lon grid rfield_latlon(N_lon, N_lat). + ! Bilinear interpolation in z=sin(latitude) and periodic longitude. + + subroutine sphere_to_latlon_field(rfield_sphere, sf, N_lon, N_lat, rfield_latlon) + + type(sphere_filter_type), intent(in) :: sf + real(c_double), intent(in) :: rfield_sphere(sf%nphi, sf%nz) + integer, intent(in) :: N_lon, N_lat + real, intent(out) :: rfield_latlon(N_lon, N_lat) + + integer :: i_lon, j_lat, k1, k2, p1, p2, k + real(c_double) :: lat_deg, lat_rad, z_target, z1, z2, wz + real(c_double) :: phi_pos, wp + real(c_double) :: v1, v2 + do j_lat = 1, N_lat + lat_deg = -90.0d0 + (real(j_lat, c_double) - 0.5d0) * 180.0d0 / real(N_lat, c_double) + lat_rad = lat_deg * MAPL_DEGREES_TO_RADIANS_R8 + z_target = sin(lat_rad) + + if (z_target >= ring_z(1, sf)) then + k1 = 1; k2 = 1; wz = 0.0d0 + else if (z_target <= ring_z(sf%nz, sf)) then + k1 = sf%nz; k2 = sf%nz; wz = 0.0d0 + else + k1 = 1 + do k = 1, sf%nz - 1 + if (ring_z(k, sf) >= z_target .and. z_target >= ring_z(k+1, sf)) then + k1 = k + exit + end if + end do + k2 = k1 + 1 + z1 = ring_z(k1, sf) + z2 = ring_z(k2, sf) + wz = (z1 - z_target) / (z1 - z2) + end if + + do i_lon = 1, N_lon + phi_pos = 0.5d0 + (real(i_lon, c_double) - 0.5d0) * real(sf%nphi, c_double) / real(N_lon, c_double) + p1 = floor(phi_pos) + wp = phi_pos - real(p1, c_double) + p1 = modulo(p1 - 1, sf%nphi) + 1 + p2 = modulo(p1, sf%nphi) + 1 + v1 = (1.0d0 - wz) * rfield_sphere(p1,k1) + wz * rfield_sphere(p1,k2) + v2 = (1.0d0 - wz) * rfield_sphere(p2,k1) + wz * rfield_sphere(p2,k2) + rfield_latlon(i_lon, j_lat) = real((1.0d0-wp)*v1 + wp*v2) + end do + end do + + end subroutine sphere_to_latlon_field + + ! ====================================================================== + ! + ! xcorr_deg_to_c2 + ! + ! Derive SMERFS power-spectrum coefficient c2 from a user-supplied + ! horizontal correlation length xcorr_deg [degrees]. + ! + ! xcorr_deg is interpreted as the Gaussian sigma of the covariance, i.e. + ! the angle at which C(theta) = C(0) * exp(-0.5). This matches the + ! flat-FFT path in random_fields.F90, where the power spectrum amplitude + ! is exp(-0.25*(lx*kx)^2), giving covariance C(r) = exp(-r^2/(2*lx^2)) + ! and Gaussian sigma = lx = xcorr [in whatever units lx and dx share]. + ! + ! The true e-folding distance (where C = C(0)/e) is xcorr_deg * sqrt(2). + ! + ! Uses bisection on the analytic Legendre-series covariance. + + subroutine xcorr_deg_to_c2(xcorr_deg, nz, c2_out) + + real(c_double), intent(in) :: xcorr_deg + integer, intent(in) :: nz + real(c_double), intent(out) :: c2_out + + integer, parameter :: lmax_mult = 3 + real(c_double), parameter :: tol_deg = 0.05d0 + integer, parameter :: max_iter = 80 + + integer :: iter, lmax + real(c_double) :: c2_lo, c2_hi, c2_mid, theta_mid, xcorr_rad + + xcorr_rad = xcorr_deg * MAPL_DEGREES_TO_RADIANS_R8 + lmax = lmax_mult * nz + + ! Bracket: larger c2 -> smoother spectrum (more low-l power) -> LONGER + ! correlation. So c2_lo gives SHORT correlation, c2_hi gives LONG. + ! c2_lo = 1e-10 -> Gaussian sigma ~ 3.5 deg (very short) + ! c2_hi = 1e+2 -> Gaussian sigma ~ 127 deg (nearly uniform) + c2_lo = 1.0d-10 + c2_hi = 1.0d2 + + if (xcorr_rad < compute_efold_theta(c2_lo, lmax)) then + write(*,'(a,f8.3,a)') 'xcorr_deg_to_c2: xcorr=', xcorr_deg, & + ' deg is very small (< 5 deg); using c2_lo' + c2_out = c2_lo ; return + end if + if (xcorr_rad > compute_efold_theta(c2_hi, lmax)) then + write(*,'(a,f8.3,a)') 'xcorr_deg_to_c2: xcorr=', xcorr_deg, & + ' deg is very large; using c2_hi' + c2_out = c2_hi ; return + end if + + ! Bisection: theta_mid > xcorr_rad means c2_mid gives too long a scale, + ! so move the upper bracket down. + c2_mid = c2_lo + do iter = 1, max_iter + c2_mid = 0.5d0 * (c2_lo + c2_hi) + theta_mid = compute_efold_theta(c2_mid, lmax) + if (abs(theta_mid - xcorr_rad) < tol_deg * MAPL_DEGREES_TO_RADIANS_R8) exit + if (theta_mid > xcorr_rad) then + c2_hi = c2_mid ! c2_mid gives too long a scale; reduce c2 + else + c2_lo = c2_mid ! c2_mid gives too short a scale; increase c2 + end if + end do + c2_out = c2_mid + + contains + + real(c_double) function compute_efold_theta(c2_val, lmax_in) + real(c_double), intent(in) :: c2_val + integer, intent(in) :: lmax_in + + real(c_double) :: C0_z1, C_z, z, dz, C_l_val, llp1, P0, P1, P2 + integer :: l, iz, nz_scan + real(c_double), parameter :: inv4pi = 0.25d0 / PI + + nz_scan = 500 + dz = 2.0d0 / real(nz_scan, c_double) + + C0_z1 = 0.0d0 ; P0 = 1.0d0 ; P1 = 1.0d0 + do l = 0, lmax_in + llp1 = real(l,c_double) * real(l+1,c_double) + C_l_val = 1.0d0 / (1.0d0 + c2_val * llp1**2) + if (l == 0) then + C0_z1 = C0_z1 + C_l_val * inv4pi + else if (l == 1) then + C0_z1 = C0_z1 + C_l_val * 3.0d0 * inv4pi + else + P2 = (real(2*l-1,c_double)*P1 - real(l-1,c_double)*P0) / real(l,c_double) + C0_z1 = C0_z1 + C_l_val * real(2*l+1,c_double) * inv4pi + P0 = P1 ; P1 = P2 + end if + end do + + compute_efold_theta = PI + do iz = 0, nz_scan + z = 1.0d0 - real(iz,c_double) * dz + C_z = 0.0d0 ; P0 = 1.0d0 ; P1 = z + do l = 0, lmax_in + llp1 = real(l,c_double) * real(l+1,c_double) + C_l_val = 1.0d0 / (1.0d0 + c2_val * llp1**2) + if (l == 0) then + C_z = C_z + C_l_val * inv4pi + else if (l == 1) then + C_z = C_z + C_l_val * 3.0d0 * z * inv4pi + else + P2 = (real(2*l-1,c_double)*z*P1 - real(l-1,c_double)*P0) / real(l,c_double) + C_z = C_z + C_l_val * real(2*l+1,c_double) * P2 * inv4pi + P0 = P1 ; P1 = P2 + end if + end do + ! Use exp(-0.5) as the target so that xcorr_deg is the Gaussian sigma + ! (C(sigma) = C(0)*exp(-0.5)), matching the flat-FFT definition where + ! the covariance is exp(-r^2/(2*xcorr^2)). + if (C_z <= exp(-0.5d0) * C0_z1) then + compute_efold_theta = acos(max(-1.0d0, min(1.0d0, z))) + return + end if + end do + end function compute_efold_theta + + end subroutine xcorr_deg_to_c2 + + subroutine generate_2d_random_field_sphere(this, rseed, rfield, rfield2, lx, ly, dx, dy) + class(sphere_filter_type), intent(inout) :: this + integer, intent(inout) :: rseed(NRANDSEED) + real(c_double), intent(out) :: rfield(this%nphi, this%nz) + real(c_double), intent(out) :: rfield2(this%nphi, this%nz) + real, intent(in) :: lx, ly, dx, dy + integer(c_int32_t), allocatable :: ints(:) + real, allocatable :: uniforms(:) + integer :: n_needed, n_got + + n_needed = this%nz * this%Mord * (this%j_hi - this%j_lo + 1) * 2 + n_got = n_needed + n_needed/32 + 64 + allocate(ints(n_got), uniforms(n_got)) + call nr_ran2_2d(1, n_got, rseed, uniforms) + ints = int(uniforms * 2147483647.0, c_int32_t) + call draw_sphere_realisation(this, ints, n_got, rfield) + call nr_ran2_2d(1, n_got, rseed, uniforms) + ints = int(uniforms * 2147483647.0, c_int32_t) + call draw_sphere_realisation(this, ints, n_got, rfield2) + deallocate(ints, uniforms) + end subroutine generate_2d_random_field_sphere + + ! Generate a field using the raw MT19937 stream used by numpy's + ! RandomState(seed). The returned integers have the same bit pattern as + ! numpy's uint32 values and are consumed by the same C Ziggurat routine. + subroutine generate_2d_random_field_mt(this, seed, rfield, rfield2, lx, ly, dx, dy) + class(sphere_random_fields), intent(inout) :: this + integer, intent(in) :: seed + real, intent(out) :: rfield(:,:) + real, intent(out) :: rfield2(:,:) + real, intent(in) :: lx, ly, dx, dy + integer(c_int32_t), allocatable :: ints(:) + integer :: n_needed, n_got + type(mt19937_state) :: state + real(c_double), allocatable :: field_d(:,:), field2_d(:,:) + real :: field_mean, field_norm + + n_needed = this%filter%nz * this%filter%Mord * & + (this%filter%j_hi - this%filter%j_lo + 1) * 2 + n_got = n_needed + n_needed / 32 + 64 + allocate(ints(n_got)) + allocate(field_d(this%filter%nphi, this%filter%nz)) + allocate(field2_d(this%filter%nphi, this%filter%nz)) + call mt19937_init(state, seed) + call mt19937_fill(state, ints) + call draw_sphere_realisation(this%filter, ints, n_got, field_d) + call mt19937_fill(state, ints) + call draw_sphere_realisation(this%filter, ints, n_got, field2_d) + rfield = real(field_d) + rfield2 = real(field2_d) + field_mean = sum(rfield) / real(size(rfield)) + rfield = rfield - field_mean + field_norm = sqrt(sum(rfield**2) / real(size(rfield))) + if (field_norm > 0.0) rfield = rfield / field_norm + field_mean = sum(rfield2) / real(size(rfield2)) + rfield2 = rfield2 - field_mean + field_norm = sqrt(sum(rfield2**2) / real(size(rfield2))) + if (field_norm > 0.0) rfield2 = rfield2 / field_norm + deallocate(ints, field_d, field2_d) + end subroutine generate_2d_random_field_mt + + subroutine mt19937_init(state, seed) + type(mt19937_state), intent(out) :: state + integer, intent(in) :: seed + integer :: i + integer(c_int64_t) :: value + integer(c_int64_t), parameter :: modulus = 4294967296_c_int64_t + + state%mt(1) = int(seed, c_int32_t) + do i = 2, 624 + value = 1812433253_c_int64_t * & + iand(int(ieor(state%mt(i-1), ishft(state%mt(i-1), -30)), c_int64_t), & + modulus - 1_c_int64_t) + i - 1 + state%mt(i) = int(modulo(value, modulus), c_int32_t) + end do + state%index = 625 + end subroutine mt19937_init + + subroutine mt19937_twist(state) + type(mt19937_state), intent(inout) :: state + integer :: i, j + integer(c_int32_t) :: y + integer(c_int32_t), parameter :: upper_mask = int(z'80000000', c_int32_t) + integer(c_int32_t), parameter :: lower_mask = int(z'7fffffff', c_int32_t) + integer(c_int32_t), parameter :: matrix_a = int(z'9908b0df', c_int32_t) + + do i = 1, 624 + j = i + 1 + if (j > 624) j = 1 + y = ior(iand(state%mt(i), upper_mask), iand(state%mt(j), lower_mask)) + j = i + 397 + if (j > 624) j = j - 624 + state%mt(i) = ieor(state%mt(j), ishft(y, -1)) + if (iand(y, 1_c_int32_t) /= 0_c_int32_t) & + state%mt(i) = ieor(state%mt(i), matrix_a) + end do + state%index = 1 + end subroutine mt19937_twist + + function mt19937_next(state) result(value) + type(mt19937_state), intent(inout) :: state + integer(c_int32_t) :: value + integer(c_int32_t) :: y + + if (state%index > 624) call mt19937_twist(state) + y = state%mt(state%index) + state%index = state%index + 1 + y = ieor(y, iand(ishft(y, -11), int(z'ffffffff', c_int32_t))) + y = ieor(y, iand(ishft(y, 7), int(z'9d2c5680', c_int32_t))) + y = ieor(y, iand(ishft(y, 15), int(z'efc60000', c_int32_t))) + y = ieor(y, ishft(y, -18)) + value = y + end function mt19937_next + + subroutine mt19937_fill(state, values) + type(mt19937_state), intent(inout) :: state + integer(c_int32_t), intent(out) :: values(:) + integer :: i + do i = 1, size(values) + values(i) = mt19937_next(state) + end do + end subroutine mt19937_fill + + ! ====================================================================== + ! Private: analytic_var + ! + ! Compute the analytic lag-0 variance of the GRF: + ! C(0) = sum_{l=0}^{lmax} (2l+1) * C_l / (4*pi) + ! C_l = 1 / (c0 + c2*(l*(l+1))^2) + + real(c_double) function analytic_var(c0, c2, lmax) + real(c_double), intent(in) :: c0, c2 + integer, intent(in) :: lmax + integer :: l + real(c_double) :: llp1, Cl + real(c_double), parameter :: inv4pi = 0.25d0 / PI + analytic_var = 0.0d0 + do l = 0, lmax + llp1 = real(l,c_double) * real(l+1,c_double) + Cl = 1.0d0 / (c0 + c2 * llp1**2) + analytic_var = analytic_var + real(2*l+1, c_double) * Cl * inv4pi + end do + end function analytic_var + + ! ====================================================================== + ! Private: compute_strip + ! + ! Distribute n total items across npes ranks; return the 1-based + ! inclusive range [lo, hi] owned by rank r (0-based). + ! Remainder items go to the first (mod) ranks. + + subroutine compute_strip(n, npes, r, lo, hi) + integer, intent(in) :: n, npes, r + integer, intent(out) :: lo, hi + integer :: base, rem + base = n / npes + rem = mod(n, npes) + ! ranks 0..rem-1 get (base+1) items; ranks rem..npes-1 get base items + if (r < rem) then + lo = r * (base + 1) + 1 + hi = lo + base + else + lo = rem * (base + 1) + (r - rem) * base + 1 + hi = lo + base - 1 + end if + end subroutine compute_strip + + ! ====================================================================== + ! Private: native ring z coordinate + + real(c_double) function ring_z(k, sf) + integer, intent(in) :: k + type(sphere_filter_type), intent(in) :: sf + + if (k <= sf%uhalf) then + ring_z = sf%z_pts(sf%uhalf + 1 - k) + else + ring_z = -sf%z_pts(k - sf%uhalf + 1) + end if + end function ring_z + + ! Private: find_nearest_ring + + integer function find_nearest_ring(z_target, sf) + real(c_double), intent(in) :: z_target + type(sphere_filter_type), intent(in) :: sf + + integer :: k, k_best + real(c_double) :: dz, dz_best, z_k + + ! Ring k layout: + ! k = 1..uhalf : north walk result, z = +z_pts(uhalf+1-k) [pole->equator] + ! k = uhalf+1..nz: south walk result, z = -z_pts(k-uhalf+1) [equator->S-pole] + ! + ! The south-walk step i_s (0-based) stores in ring k = uhalf+i_s+1 the field + ! produced by the filter transition from z_pts(i_s+skip+1) to z_pts(i_s+skip+2), + ! i.e. the field lives at z = -z_pts(i_s+skip+2) = -z_pts(k-uhalf+1) (skip=1 for + ! even nz). The previous formula used -z_pts(k-uhalf), which was off by one. + dz_best = 1.0d10 ; k_best = 1 + do k = 1, sf%nz + z_k = ring_z(k, sf) + dz = abs(z_k - z_target) + if (dz < dz_best) then + dz_best = dz ; k_best = k + end if + end do + find_nearest_ring = k_best + end function find_nearest_ring + + ! ====================================================================== + ! Private: irfft_rows + ! + ! Inverse real FFT of each column (ring) of the half-spectrum. + ! inp(nm, nrows) complex -> out(nph, nrows) real + ! Pure Fortran O(nph^2) DFT. + + subroutine irfft_rows(inp, nrows, nm, nph, out) + + integer, intent(in) :: nrows, nm, nph + complex(c_float), intent(in) :: inp(nm, nrows) + real(c_double), intent(out) :: out(nph, nrows) + + integer :: row, j, k + real(c_double) :: angle, re_sum, two_pi_over_n + complex(c_double_complex) :: Xk + ! The SMERFS Kalman filter is calibrated to produce a unit-variance + ! field when using the DFT convention without the 1/N normalisation + ! factor (matching example1_f.F90 and the Python `res *= p` step). + two_pi_over_n = 2.0d0 * PI / real(nph, c_double) + + do row = 1, nrows + do j = 0, nph-1 + re_sum = real(inp(1, row), c_double) + do k = 1, nm-2 + angle = two_pi_over_n * real(j*k, c_double) + Xk = cmplx(real(inp(k+1, row), c_double), & + real(aimag(inp(k+1, row)), c_double), c_double_complex) + re_sum = re_sum + 2.0d0*(real(Xk)*cos(angle) - aimag(Xk)*sin(angle)) + end do + k = nm - 1 + angle = two_pi_over_n * real(j*k, c_double) + re_sum = re_sum + real(inp(nm, row), c_double) * cos(angle) + out(j+1, row) = re_sum + end do + end do + + end subroutine irfft_rows + +end module sphere_random_fields_mod diff --git a/GEOSldas_App/LDASsa_DEFAULT_inputs_ensprop.nml b/GEOSldas_App/LDASsa_DEFAULT_inputs_ensprop.nml index 425b8c91..373cce2e 100644 --- a/GEOSldas_App/LDASsa_DEFAULT_inputs_ensprop.nml +++ b/GEOSldas_App/LDASsa_DEFAULT_inputs_ensprop.nml @@ -103,6 +103,9 @@ progn_pert_dtstep = 10800 ! time step for generation of AR(1) prognostics perts force_pert_dtstep = 10800 ! time step for generation of AR(1) forcing perts [s] +! use spherical random fields for both forcing and prognostic perturbations? +use_sphere_pert = .false. + ! --------------------------------------------------------------------- ! ! forcing error (or forcing perturbation) parameters diff --git a/LDAS_Shared/LDAS_PertTypes.F90 b/LDAS_Shared/LDAS_PertTypes.F90 index c0da62c1..20d0240b 100644 --- a/LDAS_Shared/LDAS_PertTypes.F90 +++ b/LDAS_Shared/LDAS_PertTypes.F90 @@ -83,8 +83,20 @@ module LDAS_PertTypes real :: xcorr ! correlation length along latitudes [deg] real :: ycorr ! correlation length along longitudes [deg] real :: tcorr ! temporal correlation length [s] - - end type pert_param_type + + ! ------------------------------------------------------------------ + ! Sphere-native perturbation parameters (SMERFS-based). + ! + ! When use_sphere_pert = .true., perturbations are generated directly + ! on the sphere using the SMERFS + ! Power spectrum: C_l = 1 / (c0 + c2 * (l*(l+1))^2), with c0=1 + ! and c2 derived from xcorr by sphere_random_fields. + ! Native sphere dimensions are automatically calculated from xcorr and + ! the fine-grid spacing. + ! ------------------------------------------------------------------ + logical :: use_sphere_pert = .false. ! .true. => SMERFS; .false. => flat FFT + + end type pert_param_type ! **********************************************************************