diff --git a/new_python/genericio.cpp b/new_python/genericio.cpp
index 8719477a9dbf27d70df43db99cf430fb2eba6951..a2b1e0c6ef4e46ab37354c1c27e076d6d57bf3fd 100644
--- a/new_python/genericio.cpp
+++ b/new_python/genericio.cpp
@@ -17,14 +17,18 @@ namespace py = pybind11;
class PyGenericIO : public gio::GenericIO {
public:
- PyGenericIO(const std::string& filename, gio::GenericIO::FileIO method=gio::GenericIO::FileIOPOSIX, gio::GenericIO::MismatchBehavior redistribute=gio::GenericIO::MismatchRedistribute)
+ PyGenericIO(
+ const std::string& filename,
+ gio::GenericIO::FileIO method=gio::GenericIO::FileIOPOSIX,
+ gio::GenericIO::MismatchBehavior redistribute=gio::GenericIO::MismatchRedistribute,
+ int eff_rank = -1)
#ifdef GENERICIO_NO_MPI
: gio::GenericIO(filename, method), num_ranks(0) {
#else
: gio::GenericIO(MPI_COMM_WORLD, filename, method), num_ranks(0) {
#endif
// open headers and rank info
- openAndReadHeader(redistribute);
+ openAndReadHeader(redistribute, eff_rank);
num_ranks = readNRanks();
// read variable info
getVariableInfo(variables);
@@ -62,13 +66,21 @@ public:
}
std::map<std::string, py::array> read(
- std::optional<std::vector<std::string>> var_names,
- bool print_stats=true,
- bool collective_stats=true
+ std::optional<std::vector<std::string>> var_names,
+ bool print_stats=true,
+ bool collective_stats=true,
+ int eff_rank=-1
) {
+ int rank;
+ #ifdef GENERICIO_NO_MPI
+ rank = 0;
+ #else
+ MPI_Comm_rank(MPI_COMM_WORLD, &rank);
+ #endif
+
// read number of elements
- int64_t num_elem = readNumElems();
-
+ int64_t num_elem = readNumElems(eff_rank);
+
// if no argument, read all
if(!var_names.has_value()) {
var_names.emplace(std::vector<std::string>());
@@ -82,8 +94,8 @@ public:
for(const std::string& var_name: *var_names) {
auto varp = std::find_if(
- variables.begin(),
- variables.end(),
+ variables.begin(),
+ variables.end(),
[&var_name](const auto& v){ return v.Name == var_name; }
);
if (varp != variables.end()) {
@@ -107,8 +119,10 @@ public:
}
}
}
-
- readData(-1, print_stats, collective_stats);
+
+ readData(eff_rank, print_stats, collective_stats);
+ clearVariables();
+
#ifndef GENERICIO_NO_MPI
MPI_Barrier(MPI_COMM_WORLD);
#endif
@@ -135,31 +149,38 @@ public:
return scale;
}
+ std::vector<int> get_source_ranks() {
+ std::vector<int> sr;
+ getSourceRanks(sr);
+ return sr;
+ }
+
private:
int num_ranks;
std::vector<gio::GenericIO::VariableInfo> variables;
};
std::map<std::string, py::array> read_genericio(
- std::string filename,
- std::optional<std::vector<std::string>> var_names,
- PyGenericIO::FileIO method=PyGenericIO::FileIO::FileIOPOSIX,
- PyGenericIO::MismatchBehavior redistribute=PyGenericIO::MismatchBehavior::MismatchRedistribute,
+ std::string filename,
+ std::optional<std::vector<std::string>> var_names,
+ PyGenericIO::FileIO method=PyGenericIO::FileIO::FileIOPOSIX,
+ PyGenericIO::MismatchBehavior redistribute=PyGenericIO::MismatchBehavior::MismatchRedistribute,
bool print_stats=true,
bool collective_stats=true,
- bool rebalance_source_ranks=false
+ bool rebalance_source_ranks=false,
+ int eff_rank=-1
) {
- PyGenericIO reader(filename, method, redistribute);
+ PyGenericIO reader(filename, method, redistribute, eff_rank);
#ifndef GENERICIO_NO_MPI
if(rebalance_source_ranks)
reader.rebalanceSourceRanks();
#endif
- return reader.read(var_names, print_stats, collective_stats);
+ return reader.read(var_names, print_stats, collective_stats, eff_rank);
}
void inspect_genericio(
- std::string filename,
- PyGenericIO::FileIO method=PyGenericIO::FileIO::FileIOPOSIX,
+ std::string filename,
+ PyGenericIO::FileIO method=PyGenericIO::FileIO::FileIOPOSIX,
PyGenericIO::MismatchBehavior redistribute=PyGenericIO::MismatchBehavior::MismatchRedistribute
) {
PyGenericIO reader(filename, method, redistribute);
@@ -168,9 +189,9 @@ void inspect_genericio(
#ifndef GENERICIO_NO_MPI
void write_genericio(
- std::string filename,
- std::map<std::string, py::array> variables,
- std::array<double, 3> phys_scale, std::array<double, 3> phys_origin,
+ std::string filename,
+ std::map<std::string, py::array> variables,
+ std::array<double, 3> phys_scale, std::array<double, 3> phys_origin,
PyGenericIO::FileIO method=PyGenericIO::FileIO::FileIOPOSIX
) {
// check data integrity, find particle count
@@ -197,15 +218,15 @@ void write_genericio(
}
for(auto& [name, data]: variables) {
- if(py::isinstance<py::array_t<float>>(data))
+ if(py::isinstance<py::array_t<float>>(data))
writer.addVariable(name.c_str(), reinterpret_cast<float*>(data.mutable_data()));
- else if(py::isinstance<py::array_t<double>>(data))
+ else if(py::isinstance<py::array_t<double>>(data))
writer.addVariable(name.c_str(), reinterpret_cast<double*>(data.mutable_data()));
- else if(py::isinstance<py::array_t<int32_t>>(data))
+ else if(py::isinstance<py::array_t<int32_t>>(data))
writer.addVariable(name.c_str(), reinterpret_cast<int32_t*>(data.mutable_data()));
- else if(py::isinstance<py::array_t<int64_t>>(data))
+ else if(py::isinstance<py::array_t<int64_t>>(data))
writer.addVariable(name.c_str(), reinterpret_cast<int64_t*>(data.mutable_data()));
- else if(py::isinstance<py::array_t<uint16_t>>(data))
+ else if(py::isinstance<py::array_t<uint16_t>>(data))
writer.addVariable(name.c_str(), reinterpret_cast<uint16_t*>(data.mutable_data()));
else
throw std::runtime_error("array dtype not supported for " + name);
@@ -228,7 +249,7 @@ PYBIND11_MODULE(pygio, m) {
MPI_Initialized(&initialized);
if(!initialized) {
int level_provided;
- MPI_Init_thread(nullptr, nullptr, MPI_THREAD_SINGLE, &level_provided);
+ MPI_Init_thread(nullptr, nullptr, MPI_THREAD_SINGLE, &level_provided);
}
});
#endif
@@ -246,9 +267,9 @@ PYBIND11_MODULE(pygio, m) {
.value("MismatchRedistribute", PyGenericIO::MismatchBehavior::MismatchRedistribute);
pyGenericIO.def(
- py::init<std::string, PyGenericIO::FileIO, PyGenericIO::MismatchBehavior>(),
- py::arg("filename"),
- py::arg("method")=PyGenericIO::FileIO::FileIOPOSIX,
+ py::init<std::string, PyGenericIO::FileIO, PyGenericIO::MismatchBehavior>(),
+ py::arg("filename"),
+ py::arg("method")=PyGenericIO::FileIO::FileIOPOSIX,
py::arg("redistribute")=PyGenericIO::MismatchBehavior::MismatchRedistribute)
.def("inspect", &PyGenericIO::inspect, "Print variable infos and size of GenericIO file")
.def("get_variables", &PyGenericIO::get_variables, "Get a list of VariableInformations defined in the GenericIO file")
@@ -256,8 +277,14 @@ PYBIND11_MODULE(pygio, m) {
.def("read_total_num_elems", (uint64_t (PyGenericIO::*)(void))(&PyGenericIO::readTotalNumElems))
.def("read_phys_origin", &PyGenericIO::read_phys_origin)
.def("read_phys_scale", &PyGenericIO::read_phys_scale)
- .def("read", &PyGenericIO::read, py::arg("variables")=nullptr, py::arg("print_stats")=true, py::arg("collective_stats")=true)
- .def("get_source_ranks", &PyGenericIO::getSourceRanks)
+ .def("read", &PyGenericIO::read,
+ py::arg("variables")=nullptr,
+ py::kw_only(),
+ py::arg("print_stats")=true,
+ py::arg("collective_stats")=true,
+ py::arg("eff_rank")=-1)
+ .def("get_source_ranks", &PyGenericIO::get_source_ranks)
+ .def("read_nranks", (int (PyGenericIO::*)(void))(&PyGenericIO::readNRanks))
#ifndef GENERICIO_NO_MPI
.def("rebalance_source_ranks", &PyGenericIO::rebalanceSourceRanks)
#endif
@@ -273,29 +300,30 @@ PYBIND11_MODULE(pygio, m) {
(vi.IsFloat ? "float" : "int") + " name='" + vi.Name + "'>";
});
- m.def("read_genericio", &read_genericio,
- py::arg("filename"),
- py::arg("variables")=nullptr,
+ m.def("read_genericio", &read_genericio,
+ py::arg("filename"),
+ py::arg("variables")=nullptr,
py::kw_only(),
- py::arg("method")=PyGenericIO::FileIO::FileIOPOSIX,
- py::arg("redistribute")=PyGenericIO::MismatchBehavior::MismatchRedistribute,
+ py::arg("method")=PyGenericIO::FileIO::FileIOPOSIX,
+ py::arg("redistribute")=PyGenericIO::MismatchBehavior::MismatchRedistribute,
py::arg("print_stats")=true,
py::arg("collective_stats")=true,
- py::arg("rebalance_sourceranks")=false,
+ py::arg("rebalance_sourceranks")=false,
+ py::arg("eff_rank")=-1,
py::return_value_policy::move);
- m.def("inspect_genericio", &inspect_genericio,
- py::arg("filename"),
+ m.def("inspect_genericio", &inspect_genericio,
+ py::arg("filename"),
py::kw_only(),
- py::arg("method")=PyGenericIO::FileIO::FileIOPOSIX,
+ py::arg("method")=PyGenericIO::FileIO::FileIOPOSIX,
py::arg("redistribute")=PyGenericIO::MismatchBehavior::MismatchRedistribute);
#ifndef GENERICIO_NO_MPI
- m.def("write_genericio", &write_genericio,
- py::arg("filename"),
- py::arg("variables"),
- py::arg("phys_scale"),
- py::arg("phys_origin") = std::array<double, 3>({0., 0., 0.}),
+ m.def("write_genericio", &write_genericio,
+ py::arg("filename"),
+ py::arg("variables"),
+ py::arg("phys_scale"),
+ py::arg("phys_origin") = std::array<double, 3>({0., 0., 0.}),
py::kw_only(),
py::arg("method")=PyGenericIO::FileIO::FileIOPOSIX);
#endif