Expose grouptree ultimate (#47)

python/sunbeam/schedule.py

@@ -13,14 +13,11 @@ class Schedule(object):
     def wells(self):
         return map(Well, self._wells)
 
-    def group(self, name):
-        return Group(self._group(name), self)
+    def group(self, timestep=0):
+        return {grp.name: grp for grp in self.groups(timestep)}
 
-    @property
-    def groups(self):
-        def mk(x): return Group(x, self)
-        def not_field(x): return x.name != 'FIELD'
-        return map(mk, filter(not_field, self._groups))
+    def groups(self, timestep=0):
+        return [Group(x, self, timestep) for x in self._groups if x.name != 'FIELD']
 
 
 @delegate(lib.Well)
@@ -68,12 +65,43 @@ class Well(object):
         def fn(well): return well.status(timestep) == 'AUTO'
         return fn
 
+
 @delegate(lib.Group)
 class Group(object):
-    def __init__(self, _, schedule):
+    def __init__(self, _, schedule, timestep):
+
+        try:
+            if not timestep == int(timestep):
+                raise ValueError
+        except ValueError:
+            raise ValueError('timestep must be int, not {}'.format(type(timestep)))
+
+        if not 0 <= timestep < len(schedule.timesteps):
+            raise IndexError('Timestep out of range')
+
         self._schedule = schedule
+        self.timestep = timestep
+
+    def __getitem__(self, name):
+        return Group(self._schedule._group(name), self._schedule, self.timestep)
 
     def wells(self, timestep):
         names = self._wellnames(timestep)
-        wells = { well.name: well for well in self._schedule.wells }
+        wells = {well.name: well for well in self._schedule.wells}
         return map(wells.__getitem__, filter(wells.__contains__, names))
+
+    @property
+    def parent(self):
+        par = self._schedule._group_tree(self.timestep)._parent(self.name)
+        if self.name == 'FIELD':
+            return None
+        else:
+            return Group(self._schedule._group(par), self._schedule, self.timestep)
+
+    @property
+    def children(self):
+        l = []
+        chl = self._schedule._group_tree(self.timestep)._children(self.name)
+        for elem in chl:
+            l.append(Group(self._schedule._group(elem), self._schedule, self.timestep))
+        return l

sunbeam/CMakeLists.txt

@@ -17,7 +17,9 @@ add_library( sunbeam SHARED sunbeam.cpp
                             parser.cpp
                             schedule.cpp
                             table_manager.cpp
-                            well.cpp )
+                            well.cpp
+                            group_tree.cpp )
+
 set_target_properties( sunbeam PROPERTIES LIBRARY_OUTPUT_DIRECTORY ${PROJECT_BINARY_DIR}/python/sunbeam )
 target_link_libraries( sunbeam ${Boost_LIBRARIES} ${opm-parser_LIBRARIES} )
 

sunbeam/group_tree.cpp

@@ -0,0 +1,25 @@
+#include <opm/parser/eclipse/EclipseState/Schedule/GroupTree.hpp>
+#include <opm/parser/eclipse/EclipseState/Schedule/Schedule.hpp>
+
+#include "sunbeam.hpp"
+
+namespace {
+
+    std::string parent( const GroupTree& gt, const std::string& name ) {
+        return gt.parent(name);
+    }
+
+    py::list children( const GroupTree& gt, const std::string& name ) {
+        return iterable_to_pylist(gt.children(name)) ;
+    }
+}
+
+void sunbeam::export_GroupTree() {
+
+    py::class_< GroupTree >( "GroupTree", py::no_init)
+
+        .def( "_parent", &parent,                       "parent function returning parent of a group")
+        .def( "_children", &children,                   "children function returning python"
+                                                       " list containing children of a group")
+        ;
+}

sunbeam/schedule.cpp

@@ -20,6 +20,9 @@ namespace {
         throw key_error( name );
     }
 
+    const GroupTree& get_grouptree ( const Schedule& sch, const size_t& timestep) {
+        return sch.getGroupTree(timestep);
+    }
 
     boost::posix_time::ptime get_start_time( const Schedule& s ) {
         return boost::posix_time::from_time_t( s.posixStartTime() );
@@ -62,6 +65,7 @@ void sunbeam::export_Schedule() {
         .def( "__contains__", &Schedule::hasWell )
         .def( "__getitem__", &get_well, ref() )
         .def( "_group", &Schedule::getGroup, ref() )
+        .def( "_group_tree", &get_grouptree, ref() )
         ;
 
 }

sunbeam/sunbeam.cpp

@@ -29,6 +29,6 @@ BOOST_PYTHON_MODULE(libsunbeam) {
     sunbeam::export_Schedule();
     sunbeam::export_TableManager();
     sunbeam::export_Well();
-
+    sunbeam::export_GroupTree();
 
 }

sunbeam/sunbeam.hpp

@@ -26,6 +26,7 @@ void export_Parser();
 void export_Schedule();
 void export_TableManager();
 void export_Well();
+void export_GroupTree();
 
 }
 

tests/CMakeLists.txt

@@ -2,6 +2,7 @@ configure_file(spe3/SPE3CASE1.DATA spe3/SPE3CASE1.DATA COPYONLY)
 configure_file(data/CORNERPOINT_ACTNUM.DATA data/CORNERPOINT_ACTNUM.DATA COPYONLY)
 configure_file(data/JFUNC.DATA data/JFUNC.DATA COPYONLY)
 
-foreach(prog deck parse_deck parse state props schedule wells)
+
+foreach(prog deck group_tree parse_deck parse state props schedule wells)
     add_python_test(${prog} ${prog}.py)
 endforeach()

tests/group_tree.py

@@ -0,0 +1,32 @@
+import unittest
+import sunbeam
+
+
+class TestGroupTree(unittest.TestCase):
+    def setUp(self):
+        norne = '../../examples/data/norne/NORNE_ATW2013.DATA'
+        self.es = sunbeam.parse(norne, ('PARSE_RANDOM_SLASH', sunbeam.action.ignore))
+
+    def test_group(self):
+        gr = self.es.schedule.group(timestep=2)['PROD']
+
+        self.assertEqual('PROD', gr.name)
+        self.assertEqual('MANI-B1', gr.children[0].name)
+        self.assertEqual(6, len(gr.children))
+        self.assertEqual('FIELD', gr.parent.name)
+        self.assertEqual(2, gr.timestep)
+        self.assertEqual(None, gr.parent.parent)
+
+    def test_timestep_groups(self):
+        for group in self.es.schedule.groups(timestep=3):
+            for child in group.children:
+                self.assertIsNotNone(child.name)
+
+        group = self.es.schedule.group(timestep=3)['PROD']
+        children = ['MANI-B1', 'MANI-B2', 'MANI-D1', 'MANI-D2', 'MANI-E1', 'MANI-E2']
+        for child in group.children:
+            self.assertIn(child.name, children)
+
+
+if __name__ == '__main__':
+    unittest.main()

tests/schedule.py

@@ -34,7 +34,7 @@ class TestSchedule(unittest.TestCase):
         self.assertEqual(dt.date(2015, 12, 31), timesteps[7])
 
     def testGroups(self):
-        g1 = self.sch.group('G1').wells(0)
+        g1 = self.sch.group()['G1'].wells(0)
         self.assertEqual(2, len(g1))
 
         def head(xs): return next(iter(xs))
@@ -45,5 +45,5 @@ class TestSchedule(unittest.TestCase):
         self.assertEqual(self.sch['INJ'],  inje)
         self.assertEqual(self.sch['PROD'], prod)
 
-        with self.assertRaises(ValueError):
-            self.sch.group('foo')
+        with self.assertRaises(KeyError):
+            self.sch.group()['foo']