Get groups

python/sunbeam/sunbeam.cpp

@@ -25,7 +25,7 @@ struct ptime_to_python_datetime {
 };
 
 template< typename T >
-py::list vector_to_pylist( const std::vector< T >& v ) {
+py::list iterable_to_pylist( const T& v ) {
     py::list l;
     for( const auto& x : v ) l.append( x );
     return l;
@@ -39,6 +39,10 @@ std::vector< Well > get_wells( const Schedule& sch ) {
     return wells;
 }
 
+py::list group_wellnames( const Group& g, size_t timestep ) {
+    return iterable_to_pylist( g.getWells( timestep )  );
+}
+
 /* alias some of boost's long names and operations */
 using ref = py::return_internal_reference<>;
 using copy = py::return_value_policy< py::copy_const_reference >;
@@ -81,7 +85,7 @@ py::list get_timesteps( const Schedule* s ) {
 
     for( size_t i = 0; i < tm.size(); ++i ) v.push_back( tm[ i ] );
 
-    return vector_to_pylist( v );
+    return iterable_to_pylist( v );
 }
 
 }
@@ -131,12 +135,18 @@ py::class_< std::vector< Well > >( "WellList", py::no_init )
     .def( py::vector_indexing_suite< std::vector< Well > >() )
     ;
 
+py::class_< Group >( "Group", py::no_init )
+    .def( "name",       mkcopy( &Group::name ) )
+    .def( "_wellnames", group_wellnames )
+    ;
+
 py::class_< Schedule >( "Schedule", py::no_init )
     .add_property( "_wells", get_wells )
     .add_property( "start",  get_start_time )
     .add_property( "end",    get_end_time )
     .add_property( "timesteps", get_timesteps )
     .def( "__contains__", &Schedule::hasWell )
+    .def( "_group", &Schedule::getGroup, ref() )
     ;
 
 void (ParseContext::*ctx_update)(const std::string&, InputError::Action) = &ParseContext::update;

python/sunbeam/sunbeam.py

@@ -46,6 +46,9 @@ class Schedule(object):
     def wells(self):
         return map(Well, self._wells)
 
+    def group(self, name):
+        return Group(self._group(name), self)
+
 @delegate(lib.EclipseState)
 class EclipseState(object):
     @property
@@ -75,6 +78,17 @@ class Well(object):
         def fn(well): return well.isproducer(timestep)
         return fn
 
+@delegate(lib.Group)
+class Group(object):
+    def __init__(self, _, schedule):
+        self._schedule = schedule
+
+    def wells(self, timestep):
+        names = self._wellnames(timestep)
+        wells = { well.name: well for well in self._schedule.wells }
+        return map(wells.__getitem__, filter(wells.__contains__, names))
+
+
 def _parse_context(actions):
     ctx = lib.ParseContext()