Repository: PMEAL/OpenPNM
Pull request #1204 (2019-01-20 03:31:30)

Return sub dicts
Since we allow subdicts to be written, we should probably support reading them too.  This PR allows the following:

```
pn['pore.foo.bar'] = 1
pn['pore.foo.baz'] = 2
d = pn['pore.foo']
d.keys()  # will be ['pore.foo.bar', 'pore.foo.baz']
```

This works with all the interleaving and multiple subdomains


Patch:
diff --git a/openpnm/core/Base.py b/openpnm/core/Base.py
index fce2359ee..656799f25 100644
--- a/openpnm/core/Base.py
+++ b/openpnm/core/Base.py
@@ -1,4 +1,5 @@
 import unyt
+from flatdict import FlatDict
 from collections import namedtuple
 import matplotlib.pyplot as plt
 from openpnm.utils import Workspace, logging
@@ -244,6 +245,29 @@ class Base(dict):
             else:
                 raise Exception('Cannot write array, wrong length: '+key)
 
+    def __getitem__(self, key):
+        element, prop = key.split('.', 1)
+        if key in self.keys():
+            # Get values if present on self
+            vals = self.get(key)
+        elif key in self.keys(mode='all', deep=True):
+            # Interleave values from geom if found there
+            vals = self.interleave_data(key)
+        elif any([k.startswith(key + '.') for k in self.keys()]):
+            # Create a subdict of values present on self
+            vals = {}
+            keys = self.keys()
+            vals.update({k: self.get(k) for k in keys if k.startswith(key + '.')})
+        elif any([k.startswith(key + '.') for k in self.keys(mode='all', deep=True)]):
+            # Create a subdict of values in subdomains by interleaving
+            vals = {}
+            keys = self.keys(mode='all', deep=True)
+            vals.update({k: self.interleave_data(k) for k in keys
+                         if k.startswith(key + '.')})
+        else:
+            raise KeyError(key)
+        return vals
+
     def _set_name(self, name, validate=True):
         if not hasattr(self, '_name'):
             self._name = None
@@ -276,6 +300,14 @@ class Base(dict):
 
     project = property(fget=_get_project)
 
+    @property
+    def network(self):
+        r"""
+        A shortcut to get a handle to the associated network
+        There can only be one so this works
+        """
+        return self.project.network
+
     def clear(self, element=None, mode='all'):
         r"""
         A subclassed version of the standard dict's clear method.  This can be
@@ -437,62 +469,6 @@ class Base(dict):
 
         return temp
 
-    def get(self, keys, default=None):
-        r"""
-        This subclassed method can be used to obtain a dictionary containing
-        subset of data on the object
-
-        Parameters
-        ----------
-        keys : string or list of strings
-            The item or items to retrieve.
-
-        default : any object
-            The value to return in the event that the requested key(s) is not
-            found.  The default is ``None``.
-
-        Returns
-        -------
-        If a single string is given in ``keys``, this method behaves exactly
-        as the ``dict's`` native ``get`` method and returns just the item
-        requested (or the ``default`` if not found).  If, however, a list of
-        strings is received, then a dictionary containing each of the
-        requested items is returned.
-
-        Notes
-        -----
-        This is useful for creating Pandas Dataframes of a specific subset of
-        data.  Note that a Dataframe can be initialized with a ``dict``, but
-        all columns must be the same length.  (e.g. ``df = pd.Dataframe(d)``)
-
-        Examples
-        --------
-        >>> import openpnm as op
-        >>> pn = op.network.Cubic(shape=[5, 5, 5])
-        >>> pore_props = pn.props(element='pore')
-        >>> subset = pn.get(keys=pore_props)
-        >>> print(len(subset))  # Only pore.coords, so gives dict with 1 array
-        1
-        >>> subset = pn.get(['pore.top', 'pore.bottom'])
-        >>> print(len(subset))  # Returns a dict with the 2 requested array
-        2
-
-        It behaves exactly as normal with a dict key string is supplied:
-
-        >>> array = pn.get('pore.coords')
-        >>> print(array.shape)  # Returns requested array
-        (125, 3)
-
-        """
-        # If a list of several keys is passed, then create a subdict
-        if isinstance(keys, list):
-            ret = {}
-            for k in keys:
-                ret[k] = super().get(k, default)
-        else:  # Otherwise return numpy array
-            ret = super().get(keys, default)
-        return ret
-
     # -------------------------------------------------------------------------
     """Data Query Methods"""
     # -------------------------------------------------------------------------
diff --git a/openpnm/core/Subdomain.py b/openpnm/core/Subdomain.py
index 2358c62c8..7db7fc911 100644
--- a/openpnm/core/Subdomain.py
+++ b/openpnm/core/Subdomain.py
@@ -43,12 +43,19 @@ class Subdomain(Base):
         element = key.split('.')[0]
         # Find boss object (either phase or network)
         boss = self.project.find_full_domain(self)
-        # Get values if present, or regenerate them
+        # Try to get vals directly first
         vals = self.get(key)
-        # If still not found, check with boss object
-        if vals is None:
+        if vals is None:  # Otherwise invoke search
             inds = boss._get_indices(element=element, labels=self.name)
-            vals = boss[key][inds]
+            try:  # Will invoke interleave data if necessary
+                vals = boss[key]  # Will return nested dict if present
+                if type(vals) is dict:  # Index into each array in nested dict
+                    for item in vals:
+                        vals[item] = vals[item][inds]
+                else:  # Otherwise index into single array
+                    vals = vals[inds]
+            except KeyError:
+                vals = super().__getitem__(key)
         return vals
 
     def __setitem__(self, key, value):
diff --git a/openpnm/geometry/GenericGeometry.py b/openpnm/geometry/GenericGeometry.py
index 2086d8638..96017364a 100644
--- a/openpnm/geometry/GenericGeometry.py
+++ b/openpnm/geometry/GenericGeometry.py
@@ -89,11 +89,3 @@ class GenericGeometry(Subdomain, ModelsMixin):
             network['pore.'+self.name] = False
             network['throat.'+self.name] = False
             self.add_locations(pores=pores, throats=throats)
-
-    @property
-    def network(self):
-        r"""
-        A shortcut to get a handle to the associated network
-        There can only be one so this works
-        """
-        return self.project.network
diff --git a/openpnm/network/GenericNetwork.py b/openpnm/network/GenericNetwork.py
index a20fa7c1c..99a2842c3 100644
--- a/openpnm/network/GenericNetwork.py
+++ b/openpnm/network/GenericNetwork.py
@@ -141,24 +141,23 @@ class GenericNetwork(Base, ModelsMixin):
         super().__setitem__(key, value)
 
     def __getitem__(self, key):
+        element, prop = key.split('.', 1)
         # Deal with special keys first
         if key.split('.')[-1] == self.name:
             element = key.split('.')[0]
             return self[element+'.all']
         if key.split('.')[-1] == '_id':
             self._gen_ids()
-        # Now get values if present, or regenerate them
-        vals = self.get(key)
-        if vals is None:  # Invoke interleave data
-            vals = self.interleave_data(key)
+            return self.get(element+'._id')
+        vals = super().__getitem__(key)
         return vals
 
     def _gen_ids(self):
-        IDs = super().get('pore._id', sp.array([], ndmin=1, dtype=sp.int64))
+        IDs = self.get('pore._id', sp.array([], ndmin=1, dtype=sp.int64))
         if len(IDs) < self.Np:
             temp = ws._gen_ids(size=self.Np - len(IDs))
             self['pore._id'] = sp.concatenate((IDs, temp))
-        IDs = super().get('throat._id', sp.array([], ndmin=1, dtype=sp.int64))
+        IDs = self.get('throat._id', sp.array([], ndmin=1, dtype=sp.int64))
         if len(IDs) < self.Nt:
             temp = ws._gen_ids(size=self.Nt - len(IDs))
             self['throat._id'] = sp.concatenate((IDs, temp))
diff --git a/openpnm/phases/GenericPhase.py b/openpnm/phases/GenericPhase.py
index c54a15e25..72a9fe5d2 100644
--- a/openpnm/phases/GenericPhase.py
+++ b/openpnm/phases/GenericPhase.py
@@ -86,17 +86,5 @@ class GenericPhase(Base, ModelsMixin):
             return net[element+'._id']
         if prop == self.name:
             return self[element+'.all']
-        # Now get values if present, or regenerate them
-        vals = self.get(key)
-        if vals is None:
-            if element == 'throat' and 'pore.'+prop in self.keys():
-                vals = self.interpolate_data(propname='pore.'+prop)
-                logger.info(key + ', not found, interpolating from ' +
-                            'pore.' + prop)
-            elif element == 'pore' and 'throat.'+prop in self.keys():
-                vals = self.interpolate_data(propname='throat.'+prop)
-                logger.info(key + ', not found, interpolating from ' +
-                            'throat.' + prop)
-            else:
-                vals = self.interleave_data(key)
+        vals = super().__getitem__(key)
         return vals


Test patch:
diff --git a/tests/unit/core/BaseTest.py b/tests/unit/core/BaseTest.py
index ad09f09aa..498b4e283 100644
--- a/tests/unit/core/BaseTest.py
+++ b/tests/unit/core/BaseTest.py
@@ -808,10 +808,6 @@ class BaseTest:
         geo2['pore.test'] = sp.rand(geo2.Np, ) * unyt.m
         assert hasattr(pn['pore.test'], 'units')
 
-    def test_get_list_of_strings(self):
-        a = self.net.get(['pore.coords', 'throat.conns'])
-        assert len(a) == 2
-
     def test_interpolate_date_with_unyts(self):
         import unyt
         pn = op.network.Cubic(shape=[10, 1, 1])
@@ -827,6 +823,141 @@ class BaseTest:
         b = geo.interpolate_data('pore.test')
         assert not hasattr(b, 'units')
 
+    def test_subdict_getitem_on_network_from_network(self):
+        pn = op.network.Cubic(shape=[5, 5, 5])
+        pn['pore.foo.bar'] = 1
+        pn['pore.foo.baz'] = 2
+        d = pn['pore.foo']
+        assert len(d) == 2
+        assert 'pore.foo.bar' in d.keys()
+        assert 'pore.foo.baz' in d.keys()
+
+    def test_subdict_getitem_on_network_from_one_geometry(self):
+        pn = op.network.Cubic(shape=[5, 5, 5])
+        pn['pore.foo.bar'] = 1
+        pn['pore.foo.baz'] = 2
+        geo = op.geometry.GenericGeometry(network=pn, pores=pn.Ps,
+                                          throats=pn.Ts)
+        d = geo['pore.foo']
+        assert len(d) == 2
+        assert 'pore.foo.bar' in d.keys()
+        assert 'pore.foo.baz' in d.keys()
+
+    def test_subdict_getitem_on_network_from_two_geometries(self):
+        pn = op.network.Cubic(shape=[5, 5, 5])
+        pn['pore.foo.bar'] = 1
+        pn['pore.foo.baz'] = 2
+        geo1 = op.geometry.GenericGeometry(network=pn,
+                                           pores=pn.Ps[:75],
+                                           throats=pn.Ts[:75])
+        geo2 = op.geometry.GenericGeometry(network=pn,
+                                           pores=pn.Ps[75:],
+                                           throats=pn.Ts[75:])
+        d = geo1['pore.foo']
+        assert len(d) == 2
+        assert 'pore.foo.bar' in d.keys()
+        assert 'pore.foo.baz' in d.keys()
+        assert d['pore.foo.bar'].size == geo1.Np
+        assert d['pore.foo.baz'].size == geo1.Np
+        d = geo2['pore.foo']
+        assert len(d) == 2
+        assert 'pore.foo.bar' in d.keys()
+        assert 'pore.foo.baz' in d.keys()
+        assert d['pore.foo.bar'].size == geo2.Np
+        assert d['pore.foo.baz'].size == geo2.Np
+
+    def test_subdict_getitem_on_phase_from_phase(self):
+        pn = op.network.Cubic(shape=[5, 5, 5])
+        air = op.phases.GenericPhase(network=pn)
+        air['pore.foo.bar'] = 1
+        air['pore.foo.baz'] = 2
+        d = air['pore.foo']
+        assert len(d) == 2
+        assert 'pore.foo.bar' in d.keys()
+        assert 'pore.foo.baz' in d.keys()
+
+    def test_subdict_getitem_on_phase_from_one_physics(self):
+        pn = op.network.Cubic(shape=[5, 5, 5])
+        geo = op.geometry.GenericGeometry(network=pn, pores=pn.Ps,
+                                          throats=pn.Ts)
+        air = op.phases.GenericPhase(network=pn)
+        phys = op.physics.GenericPhysics(network=pn, phase=air, geometry=geo)
+        air['pore.foo.bar'] = 1
+        air['pore.foo.baz'] = 2
+        d = phys['pore.foo']
+        assert len(d) == 2
+        assert 'pore.foo.bar' in d.keys()
+        assert 'pore.foo.baz' in d.keys()
+
+    def test_subdict_getitem_on_phase_from_two_physics(self):
+        pn = op.network.Cubic(shape=[5, 5, 5])
+        geo1 = op.geometry.GenericGeometry(network=pn,
+                                           pores=pn.Ps[:75],
+                                           throats=pn.Ts[:75])
+        geo2 = op.geometry.GenericGeometry(network=pn,
+                                           pores=pn.Ps[75:],
+                                           throats=pn.Ts[75:])
+        air = op.phases.GenericPhase(network=pn)
+        phys1 = op.physics.GenericPhysics(network=pn, phase=air, geometry=geo1)
+        phys2 = op.physics.GenericPhysics(network=pn, phase=air, geometry=geo2)
+        air['pore.foo.bar'] = 1
+        air['pore.foo.baz'] = 2
+
+        d = phys1['pore.foo']
+        assert len(d) == 2
+        assert 'pore.foo.bar' in d.keys()
+        assert 'pore.foo.baz' in d.keys()
+        assert d['pore.foo.bar'].size == geo1.Np
+        assert d['pore.foo.baz'].size == geo1.Np
+        d = phys2['pore.foo']
+        assert len(d) == 2
+        assert 'pore.foo.bar' in d.keys()
+        assert 'pore.foo.baz' in d.keys()
+        assert d['pore.foo.bar'].size == geo2.Np
+        assert d['pore.foo.baz'].size == geo2.Np
+
+    def test_subdict_getitem_on_one_geometry(self):
+        pn = op.network.Cubic(shape=[5, 5, 5])
+        geo = op.geometry.GenericGeometry(network=pn, pores=pn.Ps,
+                                          throats=pn.Ts)
+        geo['pore.foo.bar'] = 1
+        geo['pore.foo.baz'] = 2
+        d = pn['pore.foo']
+        assert len(d) == 2
+        assert 'pore.foo.bar' in d.keys()
+        assert 'pore.foo.baz' in d.keys()
+
+    def test_subdict_getitem_on_two_geometries(self):
+        pn = op.network.Cubic(shape=[5, 5, 5])
+        geo1 = op.geometry.GenericGeometry(network=pn,
+                                           pores=pn.Ps[:75],
+                                           throats=pn.Ts[:75])
+        geo2 = op.geometry.GenericGeometry(network=pn,
+                                           pores=pn.Ps[75:],
+                                           throats=pn.Ts[75:])
+        geo1['pore.foo.bar'] = 1
+        geo1['pore.foo.baz'] = 2
+        d = pn['pore.foo']
+        assert len(d) == 2
+        assert 'pore.foo.bar' in d.keys()
+        assert 'pore.foo.baz' in d.keys()
+        assert sp.any(sp.isnan(d['pore.foo.bar']))
+        assert sp.any(sp.isnan(d['pore.foo.baz']))
+        geo2['pore.foo.bar'] = 1
+        geo2['pore.foo.baz'] = 2
+        d = pn['pore.foo']
+        assert sp.all(~sp.isnan(d['pore.foo.bar']))
+        assert sp.all(~sp.isnan(d['pore.foo.baz']))
+
+    def test_subdict_lookup_errors(self):
+        pn = op.network.Cubic(shape=[5, 5, 5])
+        pn['pore.foo.bar'] = 1
+        pn['pore.foo.baz'] = 2
+        with pytest.raises(KeyError):
+            pn['pore.foo.b']
+        with pytest.raises(KeyError):
+            pn['pore.fo']
+
 
 if __name__ == '__main__':