diff --git a/include/openmc/capi.h b/include/openmc/capi.h index 5aba98453..43f5f9869 100644 --- a/include/openmc/capi.h +++ b/include/openmc/capi.h @@ -3,6 +3,7 @@ #include #include +#include #ifdef __cplusplus extern "C" { @@ -106,7 +107,7 @@ extern "C" { int openmc_tally_get_scores(int32_t index, int** scores, int* n); int openmc_tally_get_type(int32_t index, int32_t* type); int openmc_tally_reset(int32_t index); - int openmc_tally_results(int32_t index, double** ptr, int shape_[3]); + int openmc_tally_results(int32_t index, double** ptr, size_t shape_[3]); int openmc_tally_set_active(int32_t index, bool active); int openmc_tally_set_estimator(int32_t index, const char* estimator); int openmc_tally_set_filters(int32_t index, int n, const int32_t* indices); diff --git a/include/openmc/hdf5_interface.h b/include/openmc/hdf5_interface.h index cf7fb732f..cce7ca164 100644 --- a/include/openmc/hdf5_interface.h +++ b/include/openmc/hdf5_interface.h @@ -8,6 +8,7 @@ #include // for strlen #include #include +#include #include #include "hdf5.h" @@ -333,7 +334,9 @@ write_attribute(hid_t obj_id, const char* name, const std::vector& buffer) // Templates/overloads for write_dataset //============================================================================== -template inline void +// Template for scalars (ensured by SFINAE) +template inline +std::enable_if_t>::value> write_dataset(hid_t obj_id, const char* name, T buffer) { write_dataset(obj_id, 0, nullptr, name, H5TypeMap::type_id, &buffer, false); @@ -359,18 +362,11 @@ write_dataset(hid_t obj_id, const char* name, const std::vector& buffer) write_dataset(obj_id, 1, dims, name, H5TypeMap::type_id, buffer.data(), false); } -template inline void -write_dataset(hid_t obj_id, const char* name, const xt::xarray& arr) -{ - auto s = arr.shape(); - std::vector dims {s.cbegin(), s.cend()}; - write_dataset(obj_id, dims.size(), dims.data(), name, H5TypeMap::type_id, - arr.data(), false); -} - -template inline void -write_dataset(hid_t obj_id, const char* name, const xt::xtensor& arr) +// Template for xarray, xtensor, etc. +template inline void +write_dataset(hid_t obj_id, const char* name, const xt::xcontainer& arr) { + using T = typename D::value_type; auto s = arr.shape(); std::vector dims {s.cbegin(), s.cend()}; write_dataset(obj_id, dims.size(), dims.data(), name, H5TypeMap::type_id, diff --git a/include/openmc/tallies/tally.h b/include/openmc/tallies/tally.h index a68278bb1..a07cd6bd1 100644 --- a/include/openmc/tallies/tally.h +++ b/include/openmc/tallies/tally.h @@ -9,9 +9,11 @@ namespace openmc { extern "C" double total_weight; -xt::xtensor tally_results(int idx); +template +using adaptor_type = xt::xtensor_adaptor, N>; -xt::xtensor global_tallies(); +adaptor_type<2> global_tallies(); +adaptor_type<3> tally_results(int idx); #ifdef OPENMC_MPI //! Collect all tally results onto master process diff --git a/openmc/capi/tally.py b/openmc/capi/tally.py index 56c4799aa..a414ead45 100644 --- a/openmc/capi/tally.py +++ b/openmc/capi/tally.py @@ -1,5 +1,5 @@ from collections.abc import Mapping -from ctypes import c_int, c_int32, c_double, c_char_p, c_bool, POINTER +from ctypes import c_int, c_int32, c_size_t, c_double, c_char_p, c_bool, POINTER from weakref import WeakValueDictionary import numpy as np @@ -60,7 +60,7 @@ _dll.openmc_tally_reset.argtypes = [c_int32] _dll.openmc_tally_reset.restype = c_int _dll.openmc_tally_reset.errcheck = _error_handler _dll.openmc_tally_results.argtypes = [ - c_int32, POINTER(POINTER(c_double)), POINTER(c_int*3)] + c_int32, POINTER(POINTER(c_double)), POINTER(c_size_t*3)] _dll.openmc_tally_results.restype = c_int _dll.openmc_tally_results.errcheck = _error_handler _dll.openmc_tally_set_active.argtypes = [c_int32, c_bool] @@ -296,7 +296,7 @@ class Tally(_FortranObjectWithID): @property def results(self): data = POINTER(c_double)() - shape = (c_int*3)() + shape = (c_size_t*3)() _dll.openmc_tally_results(self._index, data, shape) return as_array(data, tuple(shape)) diff --git a/src/simulation.cpp b/src/simulation.cpp index 249dd090a..386e00510 100644 --- a/src/simulation.cpp +++ b/src/simulation.cpp @@ -3,10 +3,7 @@ #include "openmc/capi.h" #include "openmc/message_passing.h" #include "openmc/settings.h" - -#ifdef OPENMC_MPI -#include -#endif +#include "openmc/tallies/tally.h" #include @@ -106,29 +103,25 @@ void calculate_work() #ifdef OPENMC_MPI void broadcast_results() { - // Broadcast tally results so that each process has access to results - double* buffer; - if (n_tallies > 0) { - for (int i=1; i <= n_tallies; ++i) { - // Create a new datatype that consists of all values for a given filter - // bin and then use that to broadcast. This is done to minimize the - // chance of the 'count' argument of MPI_BCAST exceeding 2**31 - int shape[3]; - openmc_tally_results(i, &buffer, shape); + // Broadcast tally results so that each process has access to results + for (int i = 1; i <= n_tallies; ++i) { + // Create a new datatype that consists of all values for a given filter + // bin and then use that to broadcast. This is done to minimize the + // chance of the 'count' argument of MPI_BCAST exceeding 2**31 + auto results = tally_results(i); - int n = shape[0]; - int count_per_filter = shape[1] * shape[2]; - MPI_Datatype result_block; - MPI_Type_contiguous(count_per_filter, MPI_DOUBLE, &result_block); - MPI_Type_commit(&result_block); - MPI_Bcast(buffer, n, result_block, 0, mpi::intracomm); - MPI_Type_free(&result_block); - } + auto shape = results.shape(); + int count_per_filter = shape[1] * shape[2]; + MPI_Datatype result_block; + MPI_Type_contiguous(count_per_filter, MPI_DOUBLE, &result_block); + MPI_Type_commit(&result_block); + MPI_Bcast(results.data(), shape[0], result_block, 0, mpi::intracomm); + MPI_Type_free(&result_block); } // Also broadcast global tally results - openmc_global_tallies(&buffer); - MPI_Bcast(buffer, 4, MPI_DOUBLE, 0, mpi::intracomm); + auto gt = global_tallies(); + MPI_Bcast(gt.data(), gt.size(), MPI_DOUBLE, 0, mpi::intracomm); // These guys are needed so that non-master processes can calculate the // combined estimate of k-effective diff --git a/src/tallies/tally.cpp b/src/tallies/tally.cpp index c98442b2c..98783f449 100644 --- a/src/tallies/tally.cpp +++ b/src/tallies/tally.cpp @@ -8,31 +8,33 @@ #include "xtensor/xbuilder.hpp" // for empty_like #include "xtensor/xview.hpp" +#include #include namespace openmc { -xt::xtensor global_tallies() +adaptor_type<2> global_tallies() { // Get pointer to global tallies double* buffer; openmc_global_tallies(&buffer); // Adapt into xtensor - std::vector shape {N_GLOBAL_TALLIES, 3}; - int size {3*N_GLOBAL_TALLIES}; + std::array shape = {N_GLOBAL_TALLIES, 3}; + std::size_t size {3*N_GLOBAL_TALLIES}; + return xt::adapt(buffer, size, xt::no_ownership(), shape); } -xt::xtensor tally_results(int idx) +adaptor_type<3> tally_results(int idx) { // Get pointer to tally results double* results; - std::vector shape(3); + std::array shape; openmc_tally_results(idx, &results, shape.data()); // Adapt array into xtensor with no ownership - int size {shape[0] * shape[1] * shape[2]}; + std::size_t size {shape[0] * shape[1] * shape[2]}; return xt::adapt(results, size, xt::no_ownership(), shape); } @@ -59,7 +61,7 @@ void reduce_tally_results() // Transfer values on master and reset on other ranks if (mpi::master) { - values_view = mpi::master; + values_view = values_reduced; } else { values_view = 0.0; } diff --git a/src/tallies/tally_header.F90 b/src/tallies/tally_header.F90 index 85cc5133b..06eb6d2a2 100644 --- a/src/tallies/tally_header.F90 +++ b/src/tallies/tally_header.F90 @@ -681,7 +681,7 @@ contains ! allows a user to obtain in-memory tally results from Python directly. integer(C_INT32_T), intent(in), value :: index type(C_PTR), intent(out) :: ptr - integer(C_INT), intent(out) :: shape_(3) + integer(C_SIZE_T), intent(out) :: shape_(3) integer(C_INT) :: err if (index >= 1 .and. index <= size(tallies)) then