Add callback mechanism to TransientReactiveTransport

openpnm/algorithms/_transient_reactive_transport.py

@@ -47,6 +47,7 @@ def __init__(self, phase, name='trans_react_?', **kwargs):
         self.settings._update(TransientReactiveTransportSettings())
         self.settings['phase'] = phase.name
         self["pore.ic"] = np.nan
+        self._callbacks = []
 
     def run(self, x0, tspan, saveat=None, integrator=None):
         """
@@ -101,6 +102,37 @@ def run(self, x0, tspan, saveat=None, integrator=None):
     def _run_special(self, x0):
         pass
 
+    def _set_callback(self, func):
+        """
+        Stores the given function handle as a callback.
+
+        Callback functions are called at the beginning of every time step.
+
+        Parameters
+        ----------
+        func : function
+            Function to be called as a callback. Must be a function of
+            t and y.
+
+        Returns
+        -------
+        None
+
+        Examples
+        --------
+        >>> import openpnm as op
+        >>> net = op.network.Demo([1, 2, 3])
+        >>> air = op.phase.Air(network=net)
+        >>> air.add_model_collection(op.models.collections.physics.standard)
+        >>> trt = op.algorithms.TransientReactiveTransport(network=net, phase=air)
+        >>> func = lambda t, y: print(t)
+        >>> trt._set_callback(func)
+
+        """
+        if not callable(func):
+            raise Exception("'func' must be a function. See Examples section.")
+        self._callbacks.append(func)
+
     def _build_rhs(self):
         """
         Returns a function handle, which calculates dy/dt = rhs(y, t).
@@ -113,6 +145,7 @@ def _build_rhs(self):
         """
         def ode_func(t, y):
             # TODO: add a cache mechanism
+            self._apply_callbacks(t, y)
             self.x = y
             self._update_A_and_b()
             A = self.A.tocsc()
@@ -128,3 +161,7 @@ def _merge_inital_and_boundary_values(self):
         x0[bc_pores] = self['pore.bc.value'][bc_pores]
         quantity = self.settings['quantity']
         self[quantity] = x0
+
+    def _apply_callbacks(self, t, y):
+        for func in self._callbacks:
+            func(t, y)

tests/integration/Callback_Functionality.py

@@ -0,0 +1,42 @@
+import numpy as np
+import openpnm as op
+import matplotlib.pyplot as plt
+
+
+def test_callback_functionality():
+
+    Nx = 101
+    shape = [Nx, 1, 1]
+    spacing = 1/Nx * 5
+    net = op.network.Cubic(shape=shape, spacing=spacing)
+    air = op.phase.Air(network=net)
+    air["throat.diffusive_conductance"] = spacing * np.ones(net.Nt)
+    net["pore.volume"] = spacing**3
+
+    # Set up transient Fickian diffusion algorithm
+    tfd = op.algorithms.TransientFickianDiffusion(network=net, phase=air)
+    tfd.set_value_BC(net.pores("left"), 0)
+    tfd.set_value_BC(net.pores("right"), 0)
+
+    # Define a pulse signal
+    def pulse(t, y):
+        if 0 <= t <= 0.05:
+            y[net.Np//2] = 1.0
+
+    # Add the pulse signal to the algorithm as a callback
+    tfd._set_callback(pulse)
+
+    # Solve the transient algorithm
+    c0 = np.zeros(tfd.Np)
+    tspan = [0, 0.4]
+    tfd.run(x0=c0, tspan=tspan)
+
+    # Plot c vs. time
+    tout = np.linspace(tspan[0], tspan[1], 10)
+    fig, ax = plt.subplots()
+    for i, t in enumerate(tout):
+        ax.plot(tfd.soln['pore.concentration'](t), label=f"{t:.2f} (s)")
+    ax.legend()
+    ax.set_title("Dissipation of a pulse signal , c(x=0,L) = 0")
+    ax.set_xlabel("distance (m)")
+    ax.set_ylabel("concentration (mol/m$^3$)")

tests/integration/PoreSpyIO_on_Berea.py

@@ -1,19 +1,18 @@
-if __name__ == "__main__":
+import openpnm as op
+import porespy as ps
+import numpy as np
+import os
+from pathlib import Path
 
-    import openpnm as op
-    import porespy as ps
-    import numpy as np
-    import os
-    from pathlib import Path
 
+def test_porespy_io_on_berea():
 
     # %% Read image from file in fixtures
     path = Path(os.path.realpath(__file__),
                 '../../../tests/fixtures/berea_100_to_300.npz')
     data = np.load(path.resolve())
     im = data['im']
 
-
     # %% Note meta data for this image
     data = {
         'shape': {
@@ -63,7 +62,6 @@
     gas.add_model_collection(op.models.collections.physics.basic)
     gas.regenerate_models()
 
-
     # %% Perform Fickian Diffusion to find formation factor
     fd = op.algorithms.FickianDiffusion(network=pn, phase=gas)
     fd.set_value_BC(pores=pn.pores('xmin'), values=1.0)