使用并行 NetCDF 保存分布式 3D 复杂数组

我有一个用 Fortran 编写的基于 MPI 的程序，它在每个节点（2D 时间序列的部分）生成复杂数据的 3D 数组。我想使用并行 I/O 将这些数组写入单个文件，该文件可以相对轻松地在 python 中打开以进行进一步分析/可视化。理想情况下，我希望解决方案具有内存效率（即避免创建中间临时数组）。

使用 NetCDF，我成功地适应了子程序 https://www.archer.ac.uk/documentation/white-papers/fortanIO_netCDF/fortranIO_netCDF.pdf它针对 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 certain function for nf90_put_var”错误。这表明该函数对输入数组的数据类型不满意，因此显然我在复合数据类型的使用方面缺少一些东西。

编辑：一种简单的解决方法是将复杂数组分配给真实指针，如中所述this https://stackoverflow.com/questions/54817724/how-to-assign-pointers-to-a-complex-3d-array-in-fortran邮政。然后可以使用 numpy 对该数组进行重塑/重新转换，以得到 python 中的复杂数组。它有点笨重，而且有些令人不满意——但对于我现在的目的来说可能已经足够了。

这只是部分答案，原因您将在下面看到 - 但对于评论来说太长了。希望我能够找到丢失的信息并“升级”它，但这就是我到目前为止所拥有的。

如果您查看“复合类型介绍”下的 NetCDF4 文档，网址为https://www.unidata.ucar.edu/software/netcdf/fortran/docs/f90-user-define-data-types.html#f90-compound-types-introduction https://www.unidata.ucar.edu/software/netcdf/fortran/docs/f90-user-defined-data-types.html#f90-compound-types-introduction你会看见：

要以复合类型写入数据，首先使用nf90_def_compound要创建类型，请多次调用nf90_insert_compound添加到复合类型，然后用适当的方式写入数据nf90_put_var1, nf90_put_vara, nf90_put_vars, or nf90_put_varm call.

注意它没有提到nf90_put_var根本没有，但是有4个不同的功能。这有一定的道理，nf90_put_var大概可以很好地重载来处理 NetCDF 支持的所有内在类型（而且它完全是废话，它不支持复杂），因此对于非内在类型，大概有一些类似 C 的接口，类似于void *，我猜上面提到的四个函数就是这个实现的。

到目前为止一切顺利，您应该致电其中之一nf90_put_var1, nf90_put_vara, nf90_put_vars, or nf90_put_varm而不是nf90_put_var。现在是坏消息 - 我找不到这 4 个函数的任何文档。等效的 C 函数位于https://www.unidata.ucar.edu/software/netcdf/docs/group__variables.html 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 中的常量，并且复数的类型与组成它的实数的类型相同。