Repository: PMEAL/OpenPNM
Pull request #1289 (2019-05-30 02:37:02)

Added set_label method to simplify manipulation of labels
This addresses is #1284 


Patch:
diff --git a/openpnm/__init__.py b/openpnm/__init__.py
index 7dc98efe3..22adfe21f 100644
--- a/openpnm/__init__.py
+++ b/openpnm/__init__.py
@@ -52,7 +52,7 @@ It consists of the following submodules:
 
 """
 
-__version__ = '2.1.0b1'
+__version__ = '2.1.0-beta'
 
 from . import utils
 from .utils import Workspace, Project
diff --git a/openpnm/core/Base.py b/openpnm/core/Base.py
index b9af6f7ca..cb38ad898 100644
--- a/openpnm/core/Base.py
+++ b/openpnm/core/Base.py
@@ -666,6 +666,66 @@ class Base(dict):
                                       mode=mode)
         return labels
 
+    def set_label(self, label, pores=None, throats=None, mode='add'):
+        r"""
+        Creates or updates a label array
+
+        Parameters
+        ----------
+        label : string
+                The label to apply to the specified locations
+        pores : array_like
+            A list of pore indices or a boolean mask of where given label
+            should be added or removed (see ``mode``)
+        throats : array_like
+            A list of throat indices or a boolean mask of where given label
+            should be added or removed (see ``mode``)
+        mode : string
+            Controls how the labels are handled.  Options are:
+
+            *'add'* - Adds the given label to the specified locations while
+            keeping existing labels (default)
+
+            *'overwrite'* - Removes existing label from all locations before
+            adding the label in the specified locations
+
+            *'remove'* - Removes the  given label from the specified locations
+            leaving the remainder intact.
+
+            *'purge'* - Removes the specified label from the object
+
+        """
+        if mode == 'purge':
+            if label.split('.')[0] in ['pore', 'throat']:
+                if label in self.labels():
+                    del self[label]
+                else:
+                    logger.warning(label + ' is not a label, skpping')
+            else:
+                self.set_label(label='pore.'+label, mode='purge')
+                self.set_label(label='throat.'+label, mode='purge')
+        else:
+            if label.split('.')[0] in ['pore', 'throat']:
+                label = label.split('.', 1)[1]
+            if pores is not None:
+                pores = self._parse_indices(pores)
+                if (mode == 'overwrite') or ('pore.'+label not in self.labels()):
+                    self['pore.' + label] = False
+                if mode in ['remove']:
+                    self['pore.' + label][pores] = False
+                else:
+                    self['pore.' + label][pores] = True
+            if throats is not None:
+                throats = self._parse_indices(throats)
+                if (mode == 'overwrite') or ('throat.'+label not in self.labels()):
+                    self['throat.' + label] = False
+                if mode in ['remove']:
+                    self['throat.' + label][throats] = False
+                else:
+                    self['throat.' + label][throats] = True
+            if pores is None and throats is None:
+                del self
+
     def _get_indices(self, element, labels='all', mode='or'):
         r"""
         This is the actual method for getting indices, but should not be called


Test patch:
diff --git a/tests/unit/core/BaseTest.py b/tests/unit/core/BaseTest.py
index 498b4e283..c50d8830c 100644
--- a/tests/unit/core/BaseTest.py
+++ b/tests/unit/core/BaseTest.py
@@ -958,6 +958,61 @@ class BaseTest:
         with pytest.raises(KeyError):
             pn['pore.fo']
 
+    def test_set_label_add_to_pores(self):
+        pn = op.network.Cubic(shape=[5, 5, 5])
+        pn.set_label(label='tester', pores=[1, 2])
+        assert pn['pore.tester'].sum() == 2
+
+    def test_set_label_add_to_throats(self):
+        pn = op.network.Cubic(shape=[5, 5, 5])
+        pn.set_label(label='tester', throats=[1, 2])
+        assert pn['throat.tester'].sum() == 2
+
+    def test_set_label_overwrite_on_pores(self):
+        pn = op.network.Cubic(shape=[5, 5, 5])
+        pn.set_label(label='tester', pores=[1, 2])
+        pn.set_label(label='tester', pores=[2, 3, 4], mode='overwrite')
+        assert pn['pore.tester'].sum() == 3
+
+    def test_set_label_overwrite_on_throats(self):
+        pn = op.network.Cubic(shape=[5, 5, 5])
+        pn.set_label(label='tester', throats=[1, 2])
+        pn.set_label(label='tester', throats=[2, 3, 4], mode='overwrite')
+        assert pn['throat.tester'].sum() == 3
+
+    def test_set_label_remove_from_pores(self):
+        pn = op.network.Cubic(shape=[5, 5, 5])
+        pn.set_label(label='tester', pores=[1, 2])
+        assert pn['pore.tester'].sum() == 2
+        pn.set_label(label='tester', pores=[1, 2, 3], mode='remove')
+        assert pn['pore.tester'].sum() == 0
+
+    def test_set_label_remove_from_throats(self):
+        pn = op.network.Cubic(shape=[5, 5, 5])
+        pn.set_label(label='tester', throats=[1, 2])
+        assert pn['throat.tester'].sum() == 2
+        pn.set_label(label='tester', throats=[1, 2, 3], mode='remove')
+        assert pn['throat.tester'].sum() == 0
+
+    def test_set_label_purge_from_pores(self):
+        pn = op.network.Cubic(shape=[5, 5, 5])
+        pn.set_label(label='tester', pores=[1, 2])
+        assert pn['pore.tester'].sum() == 2
+        pn.set_label(label='tester', mode='purge')
+        assert 'pore.tester' not in pn.keys()
+
+    def test_set_label_purge_from_throats(self):
+        pn = op.network.Cubic(shape=[5, 5, 5])
+        pn.set_label(label='tester', throats=[1, 2])
+        assert pn['throat.tester'].sum() == 2
+        pn.set_label(label='tester', mode='purge')
+        assert 'throat.tester' not in pn.keys()
+
+    def test_set_label_purge_nonexistent_label(self):
+        pn = op.network.Cubic(shape=[5, 5, 5])
+        pn.set_label(label='tester', mode='purge')
+        # Should only issue warning
+
 
 if __name__ == '__main__':