Repository: PMEAL/OpenPNM

Issue:
Change solution objects to be nested inside a dict
All solutions should be returned inside a dict, so that ``sol = alg.run()`` and ``sol['pore.concentration']`` points to the actual Solution object. This way (a) the solution knows the quantity it contains (pore.concentration) and (b) multiphysics can be combined into one object, like ``sol['pore.concentration]`` and ``sol['pore.temperature']``.  It also allows us to add functionality to the sol object higher up than the subclassed numpy array, so the namespace is not so messy.

Discussion:

Patch:
diff --git a/openpnm/algorithms/_generic_transport.py b/openpnm/algorithms/_generic_transport.py
index 7dcf7a3d1..272d8ea2d 100644
--- a/openpnm/algorithms/_generic_transport.py
+++ b/openpnm/algorithms/_generic_transport.py
@@ -7,7 +7,7 @@ from openpnm.algorithms import BCsMixin
 from openpnm.utils import logging, prettify_logger_message
 from openpnm.utils import Docorator, SettingsAttr
 from openpnm.solvers import PardisoSpsolve
-from ._solution import SteadyStateSolution
+from ._solution import SteadyStateSolution, SolutionContainer
 docstr = Docorator()
 logger = logging.getLogger(__name__)
 
@@ -228,7 +228,8 @@ class GenericTransport(GenericAlgorithm, BCsMixin):
         self["pore.initial_guess"] = x0
         self._validate_x0()
         # Initialize the solution object
-        self.soln = SteadyStateSolution(x0)
+        self.soln = SolutionContainer()
+        self.soln[self.settings['quantity']] = SteadyStateSolution(x0)
         self.soln.is_converged = False
         # Build A and b, then solve the system of equations
         self._update_A_and_b()
@@ -246,7 +247,7 @@ class GenericTransport(GenericAlgorithm, BCsMixin):
         # are outdated
         self._update_A_and_b()
         # Update SteadyStateSolution object on algorithm
-        self.soln[:] = self.x
+        self.soln[self.settings['quantity']][:] = self.x
         self.soln.is_converged = not bool(exit_code)
 
     def _update_A_and_b(self):
diff --git a/openpnm/algorithms/_solution.py b/openpnm/algorithms/_solution.py
index 2fa0bcf77..232d17552 100644
--- a/openpnm/algorithms/_solution.py
+++ b/openpnm/algorithms/_solution.py
@@ -7,6 +7,11 @@ class Solution(np.ndarray):
     pass
 
 
+class SolutionContainer(dict):
+    """Brief description of 'SolutionContainer'"""
+    pass
+
+
 class SteadyStateSolution(np.ndarray):
     r"""Brief description of 'SteadyStateSolution'"""
 
diff --git a/openpnm/algorithms/_transient_reactive_transport.py b/openpnm/algorithms/_transient_reactive_transport.py
index 5d595f5a2..c3c59855d 100644
--- a/openpnm/algorithms/_transient_reactive_transport.py
+++ b/openpnm/algorithms/_transient_reactive_transport.py
@@ -2,6 +2,7 @@ import numpy as np
 from openpnm.algorithms import ReactiveTransport
 from openpnm.utils import logging, Docorator, SettingsAttr
 from openpnm.integrators import ScipyRK45
+from openpnm.algorithms._solution import SolutionContainer
 docstr = Docorator()
 logger = logging.getLogger(__name__)
 
@@ -87,7 +88,10 @@ class TransientReactiveTransport(ReactiveTransport):
         # Build RHS (dx/dt = RHS), then integrate the system of ODEs
         rhs = self._build_rhs()
         # Integrate RHS using the given solver
-        self.soln = integrator.solve(rhs, x0, tspan, saveat)
+        soln = integrator.solve(rhs, x0, tspan, saveat)
+        # Return solution as dictionary
+        self.soln = SolutionContainer()
+        self.soln[self.settings['quantity']] = soln
         return self.soln
 
     def _run_special(self, x0): ...
diff --git a/openpnm/contrib/_transient_multiphysics.py b/openpnm/contrib/_transient_multiphysics.py
index de40b0d74..3abbeac14 100644
--- a/openpnm/contrib/_transient_multiphysics.py
+++ b/openpnm/contrib/_transient_multiphysics.py
@@ -2,6 +2,7 @@ import numpy as np
 from openpnm.utils import logging, SettingsAttr, Docorator
 from openpnm.integrators import ScipyRK45
 from openpnm.algorithms import GenericAlgorithm
+from openpnm.algorithms._solution import SolutionContainer, TransientSolution
 logger = logging.getLogger(__name__)
 docstr = Docorator()
 
@@ -80,12 +81,16 @@ class TransientMultiPhysics(GenericAlgorithm):
         # Build RHS (dx/dt = RHS), then integrate the system of ODEs
         rhs = self._build_rhs()
         # Integrate RHS using the given solver
-        self.soln = integrator.solve(rhs, x0, tspan, saveat)
-        # Attach solution to each algorithm
+        soln = integrator.solve(rhs, x0, tspan, saveat)
+        # Return dictionary containing solution
+        self.soln = SolutionContainer()
         for i, alg in enumerate(self._algs):
-            # FIXME: alg.soln needs to be a proper TransientSolution object
-            # not just a slice of the composite TransientSolution.
-            alg.soln = self.soln[i*alg.Np:(i+1)*alg.Np, :]
+            # Slice soln and attach as TransientSolution object to each alg
+            t = soln.t
+            x = soln[i*alg.Np:(i+1)*alg.Np, :]
+            alg.soln = TransientSolution(t, x)
+            # Add solution of each alg to solution dictionary
+            self.soln[alg.settings['quantity']] = alg.soln
         return self.soln
 
     def _run_special(self, x0): ...
diff --git a/scripts/example_multiphysics_solver.py b/scripts/example_multiphysics_solver.py
index 0fbf7ed7b..52627b43b 100644
--- a/scripts/example_multiphysics_solver.py
+++ b/scripts/example_multiphysics_solver.py
@@ -95,11 +95,11 @@ tmp = op.contrib.TransientMultiPhysics(algorithms=[tfc, tfd], network=net)
 sol = tmp.run(y0, tspan, saveat=t)
 
 # plot
-T = sol[0:net.Np, -1]
-C = sol[net.Np:, -1]
+T = sol['pore.temperature']
+C = sol['pore.concentration']
 
 not_BC_pores = net.pores("left", mode='nor')
-C_avg = sol[net.Np:, 0:][not_BC_pores].mean(axis=0)
+C_avg = C[not_BC_pores].mean(axis=0)
 
 # Load COMSOL data from .npz file
 time = np.linspace(0, 2000, 401)
@@ -109,8 +109,8 @@ C_comsol = data['a']
 # Plot concentration and temperature profiles
 plt.figure(1)
 fig, ax = plt.subplots(ncols=2)
-im_1 = ax[0].imshow(T.reshape((Nx, Nx)))
-im_2 = ax[1].imshow(C.reshape((Nx, Nx)))
+im_1 = ax[0].imshow(T[:, -1].reshape((Nx, Nx)))
+im_2 = ax[1].imshow(C[:, -1].reshape((Nx, Nx)))
 fig.colorbar(im_1, ax=ax[0], fraction=0.046, pad=0.04)
 fig.colorbar(im_2, ax=ax[1], fraction=0.046, pad=0.04)
 ax[0].title.set_text('Temperature (K)')


Test patch:
diff --git a/tests/unit/algorithms/GenericTransportTest.py b/tests/unit/algorithms/GenericTransportTest.py
index ead10a746..ee5788abf 100644
--- a/tests/unit/algorithms/GenericTransportTest.py
+++ b/tests/unit/algorithms/GenericTransportTest.py
@@ -57,9 +57,10 @@ class GenericTransportTest:
         soln = alg.run()
         # Check returned data type
         from openpnm.algorithms._solution import SteadyStateSolution
-        assert isinstance(soln, SteadyStateSolution)
+        quantity = alg.settings['quantity']
+        assert isinstance(soln[quantity], SteadyStateSolution)
         # Ensure solution object is attached to the algorithm
-        assert isinstance(alg.soln, SteadyStateSolution)
+        assert isinstance(alg.soln[quantity], SteadyStateSolution)
 
     def test_two_value_conditions(self):
         alg = op.algorithms.GenericTransport(network=self.net,
diff --git a/tests/unit/algorithms/TransientMultiPhysicsTest.py b/tests/unit/algorithms/TransientMultiPhysicsTest.py
index 974273b82..6c6e766d6 100644
--- a/tests/unit/algorithms/TransientMultiPhysicsTest.py
+++ b/tests/unit/algorithms/TransientMultiPhysicsTest.py
@@ -113,7 +113,7 @@ class TransientMultiPhysicsTest:
         self.tmp.run(y0, tspan, saveat=t)
 
     def test_concentration(self):
-        C = self.tmp.soln[self.net.Np:2*self.net.Np, :]
+        C = self.tmp.soln[self.tfd.settings['quantity']]
         C_avg = C[self.net.pores("left", mode='nor')].mean(axis=0)
         actual = C_avg
         desired = [50.00000000, 50.46652183, 51.00059257, 51.56467322,
diff --git a/tests/unit/algorithms/TransientReactiveTransportTest.py b/tests/unit/algorithms/TransientReactiveTransportTest.py
index 853d180d2..37e9a89b0 100644
--- a/tests/unit/algorithms/TransientReactiveTransportTest.py
+++ b/tests/unit/algorithms/TransientReactiveTransportTest.py
@@ -47,18 +47,19 @@ class TransientReactiveTransportTest:
         out = self.alg.run(x0=0, tspan=(0, 1), saveat=0.1)
         # Test datatype
         from openpnm.algorithms._solution import TransientSolution
-        assert isinstance(out, TransientSolution)
+        quantity = self.alg.settings['quantity']
+        assert isinstance(out[quantity], TransientSolution)
         # Ensure solution object is attached to the algorithm
-        assert isinstance(self.alg.soln, TransientSolution)
+        assert isinstance(self.alg.soln[quantity], TransientSolution)
         # Test shape
-        nt.assert_array_equal(self.alg.soln.shape, (self.alg.Np, 11))
+        nt.assert_array_equal(self.alg.soln[quantity].shape, (self.alg.Np, 11))
         # Test stored time points
-        nt.assert_array_equal(self.alg.soln.t, np.arange(0, 1.1, 0.1))
+        nt.assert_array_equal(self.alg.soln[quantity].t, np.arange(0, 1.1, 0.1))
         # Ensure solution is callable (i.e., interpolates intermediate times)
-        assert hasattr(out, "__call__")
+        assert hasattr(out[quantity], "__call__")
         # Test solution interpolation
-        f = interp1d(out.t, out)
-        nt.assert_allclose(f(0.05), out(0.05))
+        f = interp1d(out[quantity].t, out[quantity])
+        nt.assert_allclose(f(0.05), out[quantity](0.05))
         # Ensure no extrapolation
         with nt.assert_raises(Exception):
             out(1.01)
@@ -70,7 +71,8 @@ class TransientReactiveTransportTest:
         # Integrate from 0 to 0.3 in one go
         out2 = self.alg.run(x0=0, tspan=(0, 0.3))
         # Ensure the results match
-        nt.assert_allclose(out1(0.3), out2(0.3), rtol=1e-5)
+        quantity = self.alg.settings['quantity']
+        nt.assert_allclose(out1[quantity](0.3), out2[quantity](0.3), rtol=1e-5)
 
 
     def test_adding_bc_over_sources(self):