Skip to content

Commit

Permalink
Add simple Python bindings and an example
Browse files Browse the repository at this point in the history
  • Loading branch information
franzpoeschel committed Aug 1, 2023
1 parent e77a13c commit ba10099
Show file tree
Hide file tree
Showing 6 changed files with 101 additions and 15 deletions.
2 changes: 2 additions & 0 deletions CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -602,6 +602,7 @@ if(openPMD_HAVE_PYTHON)
src/binding/python/BaseRecordComponent.cpp
src/binding/python/ChunkInfo.cpp
src/binding/python/Container.cpp
src/binding/python/CustomHierarchy.cpp
src/binding/python/Dataset.cpp
src/binding/python/Datatype.cpp
src/binding/python/Error.cpp
Expand Down Expand Up @@ -773,6 +774,7 @@ set(openPMD_PYTHON_EXAMPLE_NAMES
11_particle_dataframe
12_span_write
13_write_dynamic_configuration
14_custom_hierarchy
)

if(openPMD_USE_INVASIVE_TESTS)
Expand Down
48 changes: 48 additions & 0 deletions examples/14_custom_hierarchy.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,48 @@
import numpy as np
import openpmd_api as io


def main():
if "bp" in io.file_extensions:
filename = "../samples/custom_hierarchy.bp"
else:
filename = "../samples/custom_hierarchy.json"
s = io.Series(filename, io.Access.create)
it = s.write_iterations()[100]

# write openPMD part
temp = it.meshes["temperature"]
temp.axis_labels = ["x", "y"]
temp.unit_dimension = {io.Unit_Dimension.T: 1}
temp.position = [0.5, 0.5]
temp.grid_spacing = [1, 1]
temp.grid_global_offset = [0, 0]
temp.reset_dataset(io.Dataset(np.dtype("double"), [5, 5]))
temp[()] = np.zeros((5, 5))

# write NeXus part
nxentry = it["Scan"]
nxentry.set_attribute("NX_class", "NXentry")
nxentry.set_attribute("default", "data")

data = nxentry["data"]
data.set_attribute("NX_class", "NXdata")
data.set_attribute("signal", "counts")
data.set_attribute("axes", ["two_theta"])
data.set_attribute("two_theta_indices", [0])

counts = data.as_container_of_datasets()["counts"]
counts.set_attribute("units", "counts")
counts.set_attribute("long_name", "photodiode counts")
counts.reset_dataset(io.Dataset(np.dtype("int"), [15]))
counts[()] = np.zeros(15, dtype=np.dtype("int"))

two_theta = data.as_container_of_datasets()["two_theta"]
two_theta.set_attribute("units", "degrees")
two_theta.set_attribute("long_name", "two_theta (degrees)")
two_theta.reset_dataset(io.Dataset(np.dtype("double"), [15]))
two_theta[()] = np.zeros(15)


if __name__ == "__main__":
main()
6 changes: 6 additions & 0 deletions src/binding/python/Container.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@

#include <pybind11/pybind11.h>

#include "openPMD/CustomHierarchy.hpp"
#include "openPMD/Iteration.hpp"
#include "openPMD/Mesh.hpp"
#include "openPMD/ParticlePatches.hpp"
Expand All @@ -52,6 +53,7 @@ using PyPatchRecordContainer = Container<PatchRecord>;
using PyRecordComponentContainer = Container<RecordComponent>;
using PyMeshRecordComponentContainer = Container<MeshRecordComponent>;
using PyPatchRecordComponentContainer = Container<PatchRecordComponent>;
using PyCustomHierarchyContainer = Container<CustomHierarchy>;
PYBIND11_MAKE_OPAQUE(PyIterationContainer)
PYBIND11_MAKE_OPAQUE(PyMeshContainer)
PYBIND11_MAKE_OPAQUE(PyPartContainer)
Expand All @@ -61,6 +63,7 @@ PYBIND11_MAKE_OPAQUE(PyPatchRecordContainer)
PYBIND11_MAKE_OPAQUE(PyRecordComponentContainer)
PYBIND11_MAKE_OPAQUE(PyMeshRecordComponentContainer)
PYBIND11_MAKE_OPAQUE(PyPatchRecordComponentContainer)
PYBIND11_MAKE_OPAQUE(PyCustomHierarchyContainer)

void init_Container(py::module &m)
{
Expand All @@ -85,4 +88,7 @@ void init_Container(py::module &m)
::detail::create_and_bind_container<
PyPatchRecordComponentContainer,
Attributable>(m, "Patch_Record_Component_Container");
::detail::
create_and_bind_container<PyCustomHierarchyContainer, Attributable>(
m, "Custom_Hierarchy_Container");
}
36 changes: 36 additions & 0 deletions src/binding/python/CustomHierarchy.cpp
Original file line number Diff line number Diff line change
@@ -0,0 +1,36 @@


#include "openPMD/CustomHierarchy.hpp"
#include "openPMD/ParticleSpecies.hpp"
#include "openPMD/RecordComponent.hpp"
#include "openPMD/backend/Attributable.hpp"
#include <pybind11/pybind11.h>

namespace py = pybind11;
using namespace openPMD;

void init_CustomHierarchy(py::module &m)
{
py::class_<CustomHierarchy, Container<CustomHierarchy>, Attributable>(
m, "CustomHierarchy")
.def(
"as_container_of_datasets",
&CustomHierarchy::asContainerOf<RecordComponent>)
.def("as_container_of_meshes", &CustomHierarchy::asContainerOf<Mesh>)
.def(
"as_container_of_particles",
&CustomHierarchy::asContainerOf<ParticleSpecies>)

.def_readwrite(
"meshes",
&CustomHierarchy::meshes,
py::return_value_policy::copy,
// garbage collection: return value must be freed before Iteration
py::keep_alive<1, 0>())
.def_readwrite(
"particles",
&CustomHierarchy::particles,
py::return_value_policy::copy,
// garbage collection: return value must be freed before Iteration
py::keep_alive<1, 0>());
}
22 changes: 7 additions & 15 deletions src/binding/python/Iteration.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@
#include <pybind11/pybind11.h>
#include <pybind11/stl.h>

#include "openPMD/CustomHierarchy.hpp"
#include "openPMD/Iteration.hpp"

#include <ios>
Expand All @@ -32,7 +33,11 @@ using namespace openPMD;

void init_Iteration(py::module &m)
{
py::class_<Iteration, Attributable>(m, "Iteration")
py::class_<
Iteration,
CustomHierarchy,
Container<CustomHierarchy>,
Attributable>(m, "Iteration")
.def(py::init<Iteration const &>())

.def(
Expand Down Expand Up @@ -83,18 +88,5 @@ void init_Iteration(py::module &m)
// TODO remove in future versions (deprecated)
.def("set_time", &Iteration::setTime<double>)
.def("set_dt", &Iteration::setDt<double>)
.def("set_time_unit_SI", &Iteration::setTimeUnitSI)

.def_readwrite(
"meshes",
&Iteration::meshes,
py::return_value_policy::copy,
// garbage collection: return value must be freed before Iteration
py::keep_alive<1, 0>())
.def_readwrite(
"particles",
&Iteration::particles,
py::return_value_policy::copy,
// garbage collection: return value must be freed before Iteration
py::keep_alive<1, 0>());
.def("set_time_unit_SI", &Iteration::setTimeUnitSI);
}
2 changes: 2 additions & 0 deletions src/binding/python/openPMD.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -41,6 +41,7 @@ void init_Dataset(py::module &);
void init_Datatype(py::module &);
void init_Error(py::module &);
void init_Helper(py::module &);
void init_CustomHierarchy(py::module &);
void init_Iteration(py::module &);
void init_IterationEncoding(py::module &);
void init_Mesh(py::module &);
Expand Down Expand Up @@ -94,6 +95,7 @@ PYBIND11_MODULE(openpmd_api_cxx, m)
init_Dataset(m);
init_Datatype(m);
init_Helper(m);
init_CustomHierarchy(m);
init_Iteration(m);
init_IterationEncoding(m);
init_BaseRecordComponent(m);
Expand Down

0 comments on commit ba10099

Please sign in to comment.