Merge pull request #7 from jokva/wells.py

Rudimentary Wells support.

python/sunbeam/sunbeam.cpp

@@ -1,11 +1,33 @@
 #include <boost/python.hpp>
+#include <boost/python/suite/indexing/vector_indexing_suite.hpp>
 
 #include <opm/parser/eclipse/EclipseState/EclipseState.hpp>
+#include <opm/parser/eclipse/EclipseState/Schedule/Well.hpp>
 #include <opm/parser/eclipse/Parser/Parser.hpp>
 
 namespace py = boost::python;
 using namespace Opm;
 
+namespace {
+
+std::vector< Well > get_wells( const Schedule& sch ) {
+    std::vector< Well > wells;
+    for( const auto& w : sch.getWells() )
+        wells.push_back( *w );
+
+    return wells;
+}
+
+/* alias some of boost's long names and operations */
+using ref = py::return_internal_reference<>;
+
+template< typename F >
+auto mkref( F f ) -> decltype( py::make_function( f, ref() ) ) {
+    return py::make_function( f, ref() );
+}
+
+}
+
 BOOST_PYTHON_MODULE(libsunbeam) {
 
 EclipseState (*parse_file)( const std::string&, const ParseContext& ) = &Parser::parse;
@@ -13,6 +35,32 @@ py::def( "parse", parse_file );
 
 py::class_< EclipseState >( "EclipseState", py::no_init )
     .add_property( "title", &EclipseState::getTitle )
+    .def( "_schedule", &EclipseState::getSchedule, ref() )
+    ;
+
+int    (Well::*headI)() const = &Well::getHeadI;
+int    (Well::*headJ)() const = &Well::getHeadI;
+double (Well::*refD)()  const = &Well::getRefDepth;
+
+int    (Well::*headIts)(size_t) const = &Well::getHeadI;
+int    (Well::*headJts)(size_t) const = &Well::getHeadI;
+double (Well::*refDts)(size_t)  const = &Well::getRefDepth;
+
+py::class_< Well >( "Well", py::no_init )
+    .def( "I",   headI )
+    .def( "I",   headIts )
+    .def( "J",   headJ )
+    .def( "J",   headJts )
+    .def( "ref", refD )
+    .def( "ref", refDts )
+    ;
+
+py::class_< std::vector< Well > >( "WellList", py::no_init )
+    .def( py::vector_indexing_suite< std::vector< Well > >() )
+    ;
+
+py::class_< Schedule >( "Schedule", py::no_init )
+    .add_property( "wells", get_wells )
     ;
 
 void (ParseContext::*ctx_update)(const std::string&, InputError::Action) = &ParseContext::update;

python/sunbeam/sunbeam.py

@@ -1,7 +1,50 @@
 import libsunbeam as lib
 
-class EclipseState(lib.EclipseState):
-    pass
+class _delegate(object):
+    def __init__(self, name, attr):
+        self._name = name
+        self._attr = attr
+
+    def __get__(self, instance, _):
+        if instance is None: return self
+        return getattr(self.delegate(instance), self._attr)
+
+    def __set__(self, instance, value):
+        setattr(self.delegate(instance), self._attr, value)
+
+    def delegate(self, instance):
+        return getattr(instance, self._name)
+
+    def __repr__(self):
+        return '_delegate(' + repr(self._name) + ", " + repr(self._attr) + ")"
+
+def delegate(delegate_cls, to = '_sun'):
+    attributes = set(delegate_cls.__dict__.keys())
+
+    def inner(cls):
+        class _property(object):
+            pass
+
+        setattr(cls, to, _property())
+        for attr in attributes - set(cls.__dict__.keys()):
+            setattr(cls, attr, _delegate(to, attr))
+
+        # inject __init__
+        def default_init(self, this):
+            setattr(self, to, this)
+
+        setattr(cls, '__init__', default_init)
+
+        return cls
+
+    return inner
+
+@delegate(lib.EclipseState)
+class EclipseState(object):
+
+    @property
+    def schedule(self):
+        return self._schedule()
 
 def _parse_context(actions):
     ctx = lib.ParseContext()
@@ -24,6 +67,5 @@ def _parse_context(actions):
 
     return ctx
 
-
 def parse(path, actions = None):
-    return lib.parse(path, _parse_context(actions))
+    return EclipseState(lib.parse(path, _parse_context(actions)))

tests/CMakeLists.txt

@@ -1,5 +1,5 @@
 configure_file(spe3/SPE3CASE1.DATA spe3/SPE3CASE1.DATA COPYONLY)
 
-foreach(prog parse)
+foreach(prog parse wells)
     add_python_test(${prog} ${prog}.py)
 endforeach()

tests/parse.py

@@ -6,9 +6,6 @@ class TestParse(unittest.TestCase):
     def setUp(self):
         self.spe3 = 'spe3/SPE3CASE1.DATA'
 
-    def parse(self):
-        sunbeam.parse(self.spe3)
-
     def testParse(self):
         spe3 = sunbeam.parse(self.spe3)
         self.assertEqual('SPE 3 - CASE 1', spe3.title)

tests/wells.py

@@ -0,0 +1,34 @@
+import unittest
+import sunbeam
+
+spe3 = sunbeam.parse('spe3/SPE3CASE1.DATA')
+
+class TestWells(unittest.TestCase):
+    def setUp(self):
+        self.wells = spe3.schedule.wells
+
+    def testWells(self):
+        self.assertEqual(2, len(self.wells))
+
+        well = wells
+        well.I
+        well.I( timestep )
+        (I, J, refDepth ) = well.pos( timestep )
+
+        well.at( timestep ).I
+
+    def testWellPos0(self):
+        well = self.wells[0]
+        i, j, refdepth = well.pos(0)
+
+        self.assertEqual(6, i)
+        self.assertEqual(6, j)
+        self.assertEqual(2247.9, refdepth)
+
+    def testWellProperty(self):
+        well = self.wells[0]
+        i, j, refdepth = well.pos
+
+        self.assertEqual(6, i)
+        self.assertEqual(6, j)
+        self.assertEqual(2247.9, refdepth)