easy_output.h

Go to the documentation of this file.

#pragma once
 
#include <exception>
#include <netcdf.h>
#include "dg/topology/grid.h"
#ifdef MPI_VERSION
#include "dg/backend/mpi_vector.h"
#include "dg/topology/mpi_grid.h"
#include "dg/backend/mpi_init.h"
#endif //MPI_VERSION
#include "nc_error.h"
#include "nc_hyperslab.h"
 
namespace dg
{
namespace file
{
namespace detail
{
 
template<class T>
inline int put_var_T( int ncid, int varid, const T* data)
{
    assert( false && "Type not supported!\n" );
    return NC_EBADTYPE;
}
template<>
inline int put_var_T<float>( int ncid, int varid, const float* data){
    return nc_put_var_float( ncid, varid, data);
}
template<>
inline int put_var_T<double>( int ncid, int varid, const double* data){
    return nc_put_var_double( ncid, varid, data);
}
template<>
inline int put_var_T<int>( int ncid, int varid, const int* data){
    return nc_put_var_int( ncid, varid, data);
}
template<>
inline int put_var_T<unsigned>( int ncid, int varid, const unsigned* data){
    return nc_put_var_uint( ncid, varid, data);
}
template<class T>
inline int put_vara_T( int ncid, int varid, const size_t* startp, const size_t* countp, const T* data)
{
    assert( false && "Type not supported!\n" );
    return NC_EBADTYPE;
}
template<>
inline int put_vara_T<float>( int ncid, int varid, const size_t* startp, const size_t* countp, const float* data){
    return nc_put_vara_float( ncid, varid, startp, countp, data);
}
template<>
inline int put_vara_T<double>( int ncid, int varid, const size_t* startp, const size_t* countp, const double* data){
    return nc_put_vara_double( ncid, varid, startp, countp, data);
}
template<>
inline int put_vara_T<int>( int ncid, int varid, const size_t* startp, const size_t* countp, const int* data){
    return nc_put_vara_int( ncid, varid, startp, countp, data);
}
template<>
inline int put_vara_T<unsigned>( int ncid, int varid, const size_t* startp, const size_t* countp, const unsigned* data){
    return nc_put_vara_uint( ncid, varid, startp, countp, data);
}
#ifdef MPI_VERSION
template<class host_vector>
void put_vara_detail(int ncid, int varid,
        const MPINcHyperslab& slab,
        const host_vector& data,
        thrust::host_vector<dg::get_value_type<host_vector>>& receive,
        MPI_Comm global_comm = MPI_COMM_WORLD
        )
{
    MPI_Comm comm = slab.communicator();
    // we need to identify the global root rank within the groups and mark the
    // entire group
    int local_root_rank = dg::file::detail::mpi_comm_global2local_rank(comm, 0,
        global_comm);
    if (local_root_rank == MPI_UNDEFINED)
        return;
    unsigned ndim = slab.ndim(); // same on all processes
    int rank, size;
    MPI_Comm_rank( comm, &rank);
    MPI_Comm_size( comm, &size);
 
    // Send start and count vectors to root
    std::vector<size_t> r_start( rank == local_root_rank ? size * ndim : 0);
    std::vector<size_t> r_count( rank == local_root_rank ? size * ndim : 0);
    MPI_Gather( slab.startp(), ndim, dg::getMPIDataType<size_t>(),
                &r_start[0], ndim, dg::getMPIDataType<size_t>(),
                local_root_rank, comm);
    MPI_Gather( slab.countp(), ndim, dg::getMPIDataType<size_t>(),
                &r_count[0], ndim, dg::getMPIDataType<size_t>(),
                local_root_rank, comm);
 
    MPI_Datatype mpitype = dg::getMPIDataType<get_value_type<host_vector>>();
    int err = NC_NOERR; // store error flag
    if( rank == local_root_rank )
    {
        // Sanity check
        int ndims;
        int e = nc_inq_varndims( ncid, varid, &ndims);
        if( not e)
            err = e;
        if( (unsigned)ndims != slab.ndim())
            err = 1001; // Our own error code
 
        std::vector<size_t> sizes( size, 1);
        for( int r = 0 ; r < size; r++)
            for( unsigned u=0; u<ndim; u++)
                sizes[r]*= r_count[r*ndim + u];
 
        // host_vector could be a View
        unsigned max_size = *std::max_element( sizes.begin(), sizes.end());
        receive.resize( max_size);
        for( int r=0; r<size; r++)
        {
            if(r!=rank)
            {
                MPI_Status status;
                MPI_Recv( receive.data(), (int)sizes[r], mpitype,
                      r, r, comm, &status);
                int e = detail::put_vara_T( ncid, varid, &r_start[r*ndim],
                        &r_count[r*ndim], receive.data()); // write received
                if( not e) err = e;
 
            }
            else // write own data
            {
                int e = detail::put_vara_T( ncid, varid, slab.startp(),
                        slab.countp(), thrust::raw_pointer_cast(data.data()));
                if( not e) err = e;
            }
        }
    }
    else
    {
        size_t num = 1;
        for( unsigned u=0; u<ndim; u++)
            num*= slab.count()[u];
        MPI_Send( thrust::raw_pointer_cast(data.data()), num, mpitype,
            local_root_rank, rank, comm);
    }
    MPI_Bcast( &err, 1, dg::getMPIDataType<int>(), local_root_rank, comm);
    if( err)
        throw NC_Error( err);
    return;
}
 
// all comms must be same size
template<class host_vector, class MPITopology>
void put_vara_detail(int ncid, int varid, unsigned slice,
        const MPITopology& grid, const MPI_Vector<host_vector>& data,
        bool vara, bool parallel = false)
{
    MPINcHyperslab slab( grid);
    if( vara)
        slab = MPINcHyperslab( slice, grid);
    if( parallel)
    {
        file::NC_Error_Handle err;
        err = detail::put_vara_T( ncid, varid,
                slab.startp(), slab.countp(), data.data().data());
    }
    else
    {
        thrust::host_vector<dg::get_value_type<host_vector>> receive;
        put_vara_detail( ncid, varid, slab, data.data(), receive);
    }
}
#endif // MPI_VERSION
} // namespace detail
//
 
template<class host_vector, class Topology>
void put_var( int ncid, int varid, const Topology& grid,
    const host_vector& data, bool parallel = false)
{
    file::NC_Error_Handle err;
    err = detail::put_var_T( ncid, varid, data.data());
}
 
template<class host_vector, class Topology>
void put_vara( int ncid, int varid, unsigned slice, const Topology& grid,
    const host_vector& data, bool parallel = false)
{
    file::NC_Error_Handle err;
    NcHyperslab slab( slice, grid);
    err = detail::put_vara_T( ncid, varid, slab.startp(), slab.countp(),
            data.data());
}
// scalar data
 
template<class T, class real_type>
void put_var( int ncid, int varid, const RealGrid0d<real_type>& grid,
    T data, bool parallel = false)
{
    file::NC_Error_Handle err;
    err = detail::put_var_T( ncid, varid, &data);
}
template<class T, class real_type>
void put_vara( int ncid, int varid, unsigned slice, const RealGrid0d<real_type>& grid,
    T data, bool parallel = false)
{
    file::NC_Error_Handle err;
    size_t count = 1;
    size_t start = slice; // convert to size_t
    err = detail::put_vara_T( ncid, varid, &start, &count, &data);
}
 
 
#ifdef MPI_VERSION
 
// array data
template<class host_vector, class MPITopology>
void put_var(int ncid, int varid, const MPITopology& grid,
    const dg::MPI_Vector<host_vector>& data, bool parallel = false)
{
    detail::put_vara_detail( ncid, varid, 0, grid, data, false, parallel);
}
 
template<class host_vector, class MPITopology>
void put_vara(int ncid, int varid, unsigned slice,
    const MPITopology& grid, const dg::MPI_Vector<host_vector>& data,
    bool parallel = false)
{
    detail::put_vara_detail( ncid, varid, slice, grid, data, true, parallel);
}
// scalar data
 
template<class T, class real_type>
void put_var( int ncid, int varid, const RealMPIGrid0d<real_type>& grid,
    T data, bool parallel = false)
{
    int rank;
    MPI_Comm_rank( MPI_COMM_WORLD, &rank);
    if( rank == 0)
        put_var( ncid, varid, dg::RealGrid0d<real_type>(), data, parallel);
}
template<class T, class real_type>
void put_vara( int ncid, int varid, unsigned slice, const RealMPIGrid0d<real_type>& grid,
    T data, bool parallel = false)
{
    int rank;
    MPI_Comm_rank( MPI_COMM_WORLD, &rank);
    if( rank == 0)
        put_vara( ncid, varid, slice, dg::RealGrid0d<real_type>(), data, parallel);
}
 
#endif //MPI_VERSION
 
template<class host_vector, class Topology>
void put_var_double(int ncid, int varid, const Topology& grid,
    const host_vector& data, bool parallel = false)
{
    put_var( ncid, varid, grid, data, parallel);
}
 
template<class host_vector, class Topology>
void put_vara_double(int ncid, int varid, unsigned slice, const Topology& grid,
    const host_vector& data, bool parallel = false)
{
    put_vara( ncid, varid, slice, grid, data, parallel);
}
 
}//namespace file
}//namespace dg