pressure observable toplevel api

wrappers/python/src/cxx/common/CommonModule.cpp

@@ -32,6 +32,7 @@
 #include <readdy/common/ReaDDyVec3.h>
 #include <readdy/common/Timer.h>
 #include <readdy/io/BloscFilter.h>
+#include <pybind11/numpy.h>
 #include "SpdlogPythonSink.h"
 
 namespace py = pybind11;
@@ -46,6 +47,8 @@ void exportIO(py::module &);
 
 void exportUtils(py::module& m);
 
+using np_array = py::array_t<readdy::scalar, py::array::c_style>;
+
 void exportCommon(py::module& common) {
     using namespace pybind11::literals;
     common.def("set_logging_level", [](const std::string &level, bool python_console_out) -> void {
@@ -187,8 +190,21 @@ void exportCommon(py::module& common) {
                 return self[i];
             });
 
-    py::class_<readdy::Matrix33>(common, "Matrix33")
-            .def(py::self + py::self)
+    py::class_<readdy::Matrix33>(common, "Matrix33", py::buffer_protocol())
+            .def_buffer([](readdy::Matrix33 &m) -> py::buffer_info {
+                return py::buffer_info(
+                        m.data().data(),
+                        sizeof(readdy::Matrix33::data_arr::value_type),
+                        py::format_descriptor<readdy::Matrix33::data_arr::value_type>::format(),
+                        2,
+                        {readdy::Matrix33::n(), readdy::Matrix33::m()},
+                        { sizeof(readdy::Matrix33::data_arr::value_type) * readdy::Matrix33::n(),
+                          sizeof(readdy::Matrix33::data_arr::value_type)}
+                );
+            })
+            ;
+    /**
+     * .def(py::self + py::self)
             //.def(py::self - py::self)
             .def(readdy::scalar() * py::self)
             //.def(py::self / readdy::scalar())
@@ -197,12 +213,13 @@ void exportCommon(py::module& common) {
             .def(py::self == py::self)
             .def(py::self != py::self)
             .def("toarray", [](const readdy::Matrix33 &self) {
-                std::array<std::array<readdy::scalar, 3>, 3> arr {};
+                np_array arr {3, 3};
                 for(readdy::Matrix33::size_type i = 0; i < readdy::Matrix33::n(); ++i) {
                     for(readdy::Matrix33::size_type j = 0; j < readdy::Matrix33::m(); ++j) {
-                        arr[i][j] = self.at(i, j);
+                        arr.mutable_at(i, j) = self.at(i, j);
                     }
                 }
                 return arr;
-            });
+            })
+     */
 }

wrappers/python/src/python/readdy/api/registry/observables.py

@@ -25,7 +25,10 @@
 @author: clonker
 """
 
+import numpy as _np
+
 from typing import Optional as _Optional, Dict as _Dict, Union as _Union
+from readdy.util.observable_utils import calculate_pressure as _calculate_pressure
 
 def _parse_save_args(save_args):
     assert save_args is None or isinstance(save_args, (dict, bool)), \
@@ -201,18 +204,22 @@ def virial(self, stride, callback=None, save: _Optional[_Union[_Dict, str]]='def
         if isinstance(save, str) and save == 'default':
             save = {"name": "virial", "chunk_size": 100}
 
-        handle = self._sim.register_observable_virial(stride, lambda x: callback(x.toarray()))
+        handle = self._sim.register_observable_virial(stride, lambda x: callback(_np.ndarray((3,3), buffer=x)))
         self._add_observable_handle(*_parse_save_args(save), handle)
 
     def pressure(self, stride, physical_particles=None, callback=None, save: _Optional[_Union[_Dict, str]]='default'):
         """
-
-        :param stride:
-        :param physical_particles:
-        :param save:
+        This observable will report back the pressure of the system. As the pressure can be computed from the number of
+        physical particles and the virial tensor, this observable actually delegates to the n_particles and virial
+        observables. The default save behavior is to save virial and n_particles under
+        "virial_pressure" and "n_particles_pressure", respectively.
+        :param stride: skip `stride` time steps before evaluating the observable again
+        :param physical_particles: a list of physical particles, if None, all particles are considered physical
+        :param callback: callback function that takes the current pressure as argument
+        :param save: dictionary containing `name` and `chunk_size` or None to not save the observable to file
         """
         if isinstance(save, str) and save == 'default':
-            save = {"name": "_pressure", "chunk_size": 100}
+            save = {"name": "_pressure", "chunk_size": 500}
         save_name, save_chunk_size = _parse_save_args(save)
         save_n_particles = None
         save_virial = None
@@ -222,11 +229,13 @@ def pressure(self, stride, physical_particles=None, callback=None, save: _Option
 
         class PressureCallback(object):
 
-            def __init__(self, user_callback):
+            def __init__(self, user_callback, kbt, volume):
 
                 self._user_callback = user_callback
                 self._n = None
                 self._v = None
+                self._kbt = kbt
+                self._volume = volume
 
                 def pressure_callback_n_particles(n):
                     self._n = n
@@ -241,11 +250,12 @@ def pressure_callback_virial(virial):
 
             def _eval_user_callback(self):
                 if self._n is not None and self._v is not None:
-                    # todo calc
+                    self._user_callback(_calculate_pressure(box_volume=self._volume, kbt=self._kbt,
+                                                            n_particles=self._n, virial=self._v))
                     self._n = None
                     self._v = None
 
-        pressure_callback = PressureCallback(callback)
+        pressure_callback = PressureCallback(callback, self._sim.context.kbt, self._sim.context.box_volume())
         self.number_of_particles(stride, types=physical_particles, callback=pressure_callback.n_particles_callback,
                                  save=save_n_particles)
         self.virial(stride, callback=pressure_callback.virial_callback, save=save_virial)

wrappers/python/src/python/readdy/api/trajectory.py

@@ -32,6 +32,7 @@
 
 from readdy._internal.readdybinding.common.util import read_reaction_observable as _read_reaction_observable
 from readdy._internal.readdybinding.common.util import read_trajectory as _read_trajectory
+from readdy.util.observable_utils import calculate_pressure as _calculate_pressure
 
 import readdy.util.io_utils as _io_utils
 
@@ -427,3 +428,42 @@ def read_observable_forces(self, data_set_name="forces"):
             time = f[group_path]["time"][:]
             forces = f[group_path]["data"][:]
             return time, forces
+
+    def read_observable_virial(self, data_set_name="virial"):
+        """
+        Reads back the output of the "virial" observable.
+        :param data_set_name: The data set name as given in the simulation setup
+        :return: a tuple which contains an array corresponding to the time as first entry and an array containing the
+                 corresponding virial per time step
+        """
+        group_path = "readdy/observables/{}".format(data_set_name)
+        with _h5py.File(self._filename, "r") as f:
+            if not group_path in f:
+                raise ValueError("The virial observable was not recorded in the file or recorded under a "
+                                 "different name!")
+            time = f[group_path]["time"][:]
+            virial = f[group_path]["data"][:]
+            return time, virial
+
+    def read_observable_pressure(self, data_set_name="_pressure"):
+        """
+        Reads back the output of the "pressure" observable. As the pressure can be computed from the number of particles
+        and the virial, this actually reads back the n_particles and virial observables. The data_set_name serves as a
+        postfix, where the default value corresponds to the data sets as they are created when using the default
+        settings of the observable.
+
+        :param data_set_name: the data set name postfix, default="_pressure",
+               yielding "n_particles_pressure" and "virial_pressure", respectively
+        :return: a tuple which contains an array corresponding to the time as first entry and an array containing the
+                 corresponding pressure per time step
+        """
+        time_n_particles, n_particles = self.read_observable_number_of_particles("n_particles{}".format(data_set_name))
+        time_virial, virial = self.read_observable_virial("virial{}".format(data_set_name))
+
+        if not _np.array_equal(time_n_particles, time_virial):
+            raise RuntimeError("For Pressure it is required to know the number of particles and the virial. "
+                               "However, these two observables were recorded with different strides.")
+
+        pressure = _np.array([_calculate_pressure(self.box_volume, self.kbt, n_particles[i], virial[i])
+                              for i in range(len(time_n_particles))])
+        return time_virial, pressure

wrappers/python/src/python/readdy/tests/test_observables_io.py

@@ -99,7 +99,7 @@ def test_virial_observable_CPU(self):
 
         virials = []
         def virial_callback(virial):
-            virials.append(np.array(virial.toarray()))
+            virials.append(np.ndarray((3,3), buffer=virial))
 
         handle = sim.register_observable_virial(1, virial_callback)
         with closing(io.File.create(fname)) as f:
@@ -126,7 +126,7 @@ def test_virial_observable_SCPU(self):
 
         virials = []
         def virial_callback(virial):
-            virials.append(np.array(virial.toarray()))
+            virials.append(np.ndarray((3,3), buffer=virial))
 
         handle = sim.register_observable_virial(1, virial_callback)
         with closing(io.File.create(fname)) as f:

wrappers/python/src/python/readdy/tests/test_toplevel_api.py

@@ -276,8 +276,25 @@ def _run_readwrite_test_for(self, kernel, reaction_handler):
         sim.observe.reaction_counts(1)
         sim.observe.forces(1)
         sim.observe.energy(1)
+        pressures = []
+        pressures_inactive = []
+
+        class PressureCallback(object):
+
+            def __init__(self):
+                self.active = True
+
+            def __call__(self, p):
+                if self.active:
+                    pressures.append(p)
+                else:
+                    pressures_inactive.append(p)
+        pressure_callback = PressureCallback()
+
+        sim.observe.pressure(1, callback=pressure_callback)
         sim.run(50, 1e-3, False)
 
+        pressure_callback.active = False
         sim.output_file = traj_fname2
         sim.run(50, 1e-3, False)
 
@@ -293,6 +310,16 @@ def _run_readwrite_test_for(self, kernel, reaction_handler):
             np.testing.assert_equal(len(time), 51)
             np.testing.assert_equal(len(positions), 51)
 
+            time, pressure = traj.read_observable_pressure()
+            np.testing.assert_equal(len(time), 51)
+            np.testing.assert_equal(len(pressure), 51)
+            np.testing.assert_equal(len(pressures), 51)
+            np.testing.assert_equal(len(pressures_inactive), 51)
+            if fname == traj_fname:
+                np.testing.assert_array_almost_equal(pressure, np.array(pressures))
+            else:
+                np.testing.assert_array_almost_equal(pressure, np.array(pressures_inactive))
+
             time, types, ids, positions = traj.read_observable_particles()
             np.testing.assert_equal(len(time), 51)
             np.testing.assert_equal(len(types), 51)

wrappers/python/src/python/readdy/util/observable_utils.py

@@ -23,4 +23,7 @@
 
 
 def calculate_pressure(box_volume, kbt, n_particles, virial):
-    return (n_particles * kbt * _np.ones(shape=(3, 3), dtype=virial.dtype) + virial) / (3. * box_volume)
+    if isinstance(n_particles, (list, tuple, _np.ndarray)):
+        n_particles = _np.sum(_np.array(n_particles).squeeze())
+    VP = (n_particles * kbt * _np.identity(3, dtype=virial.dtype) + virial) / 3.
+    return _np.mean(VP.diagonal()) / box_volume