Repository: PMEAL/OpenPNM

Issue:
Add remove_source to Algorithm objects?

Discussion:

Patch:
diff --git a/openpnm/algorithms/ReactiveTransport.py b/openpnm/algorithms/ReactiveTransport.py
index a5f422105..15ea1402a 100644
--- a/openpnm/algorithms/ReactiveTransport.py
+++ b/openpnm/algorithms/ReactiveTransport.py
@@ -185,7 +185,7 @@ class ReactiveTransport(GenericTransport):
         if variable_props:
             self.settings['variable_props'] = []
 
-    def set_source(self, propname, pores):
+    def set_source(self, propname, pores, mode='overwrite'):
         r"""
         Applies a given source term to the specified pores
 
@@ -193,9 +193,14 @@ class ReactiveTransport(GenericTransport):
         ----------
         propname : string
             The property name of the source term model to be applied
-
         pores : array_like
             The pore indices where the source term should be applied
+        mode : str
+            Controls how the sources are applied. Options are:
+                
+            'merge' - Adds supplied source term to already existing ones.
+            'overwrite' - (default) Deletes all existing source terms of the 
+            given ``propname`` then adds the specified new ones
 
         Notes
         -----
@@ -209,12 +214,34 @@ class ReactiveTransport(GenericTransport):
         if (locs & locs_BC).any():
             raise Exception('Boundary conditions already present in given '
                             + 'pores, cannot also assign source terms')
-        # Set source term
-        self[propname] = locs
+        if propname not in self.keys():
+            self[propname] = False
+        if mode == 'merge':
+            self[propname][locs] = True
+        elif mode == 'overwrite':
+            self[propname] = locs
         # Check if propname already in source term list
         if propname not in self.settings['sources']:
             self.settings['sources'].append(propname)
 
+    def remove_source(self, propname, pores=None):
+        r"""
+        Removes source terms from specified pores
+
+        Parameters
+        ----------
+        propname : str
+            The property name of the source term model to be removed
+        pores : array_like
+            The pore indices where the source term should be applied
+        """
+        if pores is None:
+            pores = self.Ps
+        locs = self.tomask(pores=pores)
+        if propname not in self.keys():
+            self[propname] = False
+        self[propname][locs] = False
+
     def _set_variable_props(self, propnames):
         r"""
         Inform the algorithm which properties are variable, so those on which


Test patch:
diff --git a/tests/unit/algorithms/ReactiveTransportTest.py b/tests/unit/algorithms/ReactiveTransportTest.py
index 4b70b92be..a01bc2b07 100644
--- a/tests/unit/algorithms/ReactiveTransportTest.py
+++ b/tests/unit/algorithms/ReactiveTransportTest.py
@@ -149,6 +149,34 @@ class ReactiveTransportTest:
         c_mean_relaxed = self.alg['pore.concentration'].mean()
         assert_allclose(c_mean_base, c_mean_relaxed, rtol=1e-6)
 
+    def test_set_source_with_modes(self):
+        self.alg.reset(bcs=True, source_terms=True)
+        self.alg.set_source(pores=self.net.pores('left'),
+                            propname='pore.reaction',
+                            mode='overwrite')
+        assert self.alg['pore.reaction'].sum() == self.net.num_pores('left')
+        self.alg.set_source(pores=[0, 1, 2],
+                            propname='pore.reaction',
+                            mode='overwrite')
+        assert self.alg['pore.reaction'].sum() == 3
+        self.alg.set_source(pores=[2, 3, 4],
+                            propname='pore.reaction',
+                            mode='merge')
+        assert self.alg['pore.reaction'].sum() == 5
+
+    def test_remove_source(self):
+        self.alg.reset(bcs=True, source_terms=True)
+        self.alg.set_source(pores=[2, 3, 4],
+                            propname='pore.reaction',
+                            mode='overwrite')
+        assert self.alg['pore.reaction'].sum() == 3
+        self.alg.remove_source(propname='pore.reaction', pores=[0, 1])
+        assert self.alg['pore.reaction'].sum() == 3
+        self.alg.remove_source(propname='pore.reaction', pores=[0, 2])
+        assert self.alg['pore.reaction'].sum() == 2
+        self.alg.remove_source(propname='pore.reaction')
+        assert self.alg['pore.reaction'].sum() == 0
+
     def test_source_relaxation_consistency_w_base_solution(self):
         self.alg.reset(bcs=True, source_terms=True)
         self.alg.set_source(pores=self.net.pores('bottom'), propname='pore.reaction')