python - 使用并行 NetCDF 保存分布式 3D 复数数组
问题描述
我有一个用 Fortran 编写的基于 MPI 的程序，它在每个节点（2D 时间序列的部分）处生成复杂数据的 3D 数组。我想使用并行 I/O 将这些数组写入单个文件，该文件可以相对容易地在 python 中打开以进行进一步分析/可视化。理想情况下，我希望解决方案具有内存效率（即避免创建中间临时数组）。
使用 NetCDF，我设法调整了一个子例程，该子例程实现了实数的 3D 数组。然而，当涉及到复杂的数组时，我遇到了一个绊脚石。
在下面的代码中，我试图通过创建由两个实数组成的复合数据类型并假设 Fortran 复数数据类型的实部和虚部连续存储在 3D 数组的第一维中，将子例程从实数扩展到复数。
module IO
use NetCDF
use, intrinsic :: iso_fortran_env, only: dp => real64
implicit none
contains
subroutine output_3D(dataname, starts, ends, global_data_dims, &
local_data, MPI_communicator)
character(len=*), intent(in) :: dataname
integer, dimension(3), intent(in) :: starts
integer, dimension(3), intent(in) :: ends
integer, dimension(3), intent(in) :: global_data_dims
complex(dp), intent(in) :: local_data( 1:(ends(1) - starts(1)+ 1), &
1:(ends(2) - starts(2) + 1), &
1:(ends(3) - starts(3) + 1))
integer, dimension(3) :: expanded_starts
integer, intent(in) :: MPI_communicator
integer :: ncid, varid, dimid(3)
integer :: counts(3)
integer :: typeid
expanded_starts(1) = (starts(1))* 2 + 1
expanded_starts = starts(2)
expanded_starts(3) = starts(3)
call check(nf90_create( trim(dataname)//'.cdf', &
IOR(NF90_NETCDF4, NF90_MPIIO), &
ncid, &
comm = MPI_communicator, &
info = MPI_INFO_NULL))
call check(nf90_def_dim(ncid, "x", global_data_dims(1), dimid(1)))
call check(nf90_def_dim(ncid, "y", global_data_dims(2) * 2, dimid(2)))
call check(nf90_def_dim(ncid, "z", global_data_dims(3), dimid(3)))
! define a complex data type consisting of two real(8)
call check(nf90_def_compound(ncid, 16, "COMPLEX", typeid))
call check(nf90_insert_compound(ncid, typeid, "REAL", 0, NF90_DOUBLE))
call check(nf90_insert_compound(ncid, typeid, "IMAG", 8, NF90_DOUBLE))
! define a 3D variable of type "complex"
call check(nf90_def_var(ncid, dataname, typeid, dimid, varid))
! exit define mode
call check(nf90_enddef(ncid))
! Now in NETCDF data mode
! set to use MPI/PnetCDF collective I/O
call check(nf90_var_par_access(ncid, varid, NF90_COLLECTIVE))
counts(1) = (ends(1) - starts(1) + 1) * 2
counts(2) = (ends(2) - starts(2) + 1)
counts(3) = (ends(3) - starts(3) + 1)
call check(nf90_put_var(ncid, &
varid, &
local_data, &
start = expanded_starts,&
count = counts))
! close the file
call check(nf90_close(ncid))
return
end subroutine output_3D
subroutine check(status)
integer, intent ( in) :: status
if(status /= nf90_noerr) then
print *, trim(nf90_strerror(status))
stop 2
end if
end subroutine check
end module IO
program test_write
use IO
use MPI
complex(dp) :: data(2,2,3)
integer :: flock
integer :: rank
integer :: ierr
integer :: i, j, k
call MPI_init(ierr)
call MPI_comm_size(MPI_comm_world, flock, ierr)
call MPI_comm_rank(MPI_comm_world, rank, ierr)
do k = 1, 3
do j = 1, 2
do i = 1, 2
data(i,j,k) = cmplx(i, j, 8)
enddo
enddo
enddo
if (rank == 0) then
call output_3D_hdf5('out', [1,1,1], [2,2,3], [2,2,6], &
data, MPI_comm_world)
else
call output_3D_hdf5('out', [1,1,4], [2,2,6], [2,2,6], &
data, MPI_comm_world)
endif
call MPI_finalize(ierr)
end program test_write
上面的代码导致编译时出现“There is no specific function for nf90_put_var”错误。这表明该函数对输入数组的数据类型不满意，因此显然我在使用复合数据类型方面缺少一些东西。
编辑：一种简单的解决方法是将复杂数组分配给真实指针，如本文所述。然后可以使用 numpy 对数组进行重新整形/重铸，以达到 python 中的复杂数组。它有点笨重，而且有点不满意——但对于我现在的目的来说可能已经足够了。
解决方案
由于您将在下面看到的原因，这只是部分答案 - 但评论太长了。希望我能够找到丢失的信息并“升级”它，但这是我目前所拥有的。
要在复合类型中写入数据，首先使用
nf90_def_compound创建类型，多次调用nf90_insert_compound添加到复合类型，然后使用适当nf90_put_var1的 、nf90_put_vara、nf90_put_vars或写入数据nf90_put_varm call。
请注意，它根本没有提及nf90_put_var，而是 4 个不同的功能。这有某种意义，nf90_put_var大概是很好地重载来处理所有 NetCDF 支持的内在类型（它完全是废话，它不支持复杂），所以对于非内在类型，大概有一些类似 C 的接口，类似于void *，我猜这就是上面提到的四个函数实现的。
到目前为止一切顺利，您应该调用nf90_put_var1, nf90_put_vara,nf90_put_vars或nf90_put_varm而不是nf90_put_var. 现在是坏消息 - 我找不到这 4 个函数的任何文档。等效的 C 函数位于https://www.unidata.ucar.edu/software/netcdf/docs/group__variables.html，因此您可以从那里找出所需的内容，但这不是很好 - 我会至少向 Unidata 提交一个错误报告，但这对我来说缺乏对复杂的内在支持足以让我在其他地方寻找我的 I/O 解决方案......
虽然我在这里，但您确实不应该对变量的种类使用显式数字，但我可以向您展示编译器complex(8)将无法编译的地方。而是使用Selected_real_kind或类似的，或使用内部模块 iso_fortran_env 中的常量，或者可能使用 iso_c_binding 中的常量，并且复数的种类与组成它的实数的种类相同。