Update of python bindings for EclFile

-> two extra overloads for getting (array name + array name, occurrence)
   -> extra member function count(std::string), count occurrences of arrays with a specific name.

python/cxx/eclipse_io.cpp

@@ -13,26 +13,70 @@ namespace py = pybind11;
 
 namespace {
 
-py::array get_vector(Opm::EclIO::EclFile * file_ptr, std::size_t array_index) {
+using npArray = std::tuple<py::array, Opm::EclIO::eclArrType>;
+using EclEntry = std::tuple<std::string, Opm::EclIO::eclArrType, long int>;
+
+npArray get_vector_index(Opm::EclIO::EclFile * file_ptr, std::size_t array_index)
+{
     auto array_type = std::get<1>(file_ptr->getList()[array_index]);
+
     if (array_type == Opm::EclIO::INTE)
-        return convert::numpy_array( file_ptr->get<int>(array_index) );
+        return std::make_tuple (convert::numpy_array( file_ptr->get<int>(array_index)), array_type);
 
     if (array_type == Opm::EclIO::REAL)
-        return convert::numpy_array( file_ptr->get<float>(array_index) );
+        return std::make_tuple (convert::numpy_array( file_ptr->get<float>(array_index)), array_type);
 
     if (array_type == Opm::EclIO::DOUB)
-        return convert::numpy_array( file_ptr->get<double>(array_index) );
+        return std::make_tuple (convert::numpy_array( file_ptr->get<double>(array_index)), array_type);
 
     if (array_type == Opm::EclIO::LOGI)
-        return convert::numpy_array( file_ptr->get<bool>(array_index) );
+        return std::make_tuple (convert::numpy_array( file_ptr->get<bool>(array_index)), array_type);
 
     if (array_type == Opm::EclIO::CHAR)
-        return convert::numpy_string_array( file_ptr->get<std::string>(array_index) );
+        return std::make_tuple (convert::numpy_string_array( file_ptr->get<std::string>(array_index)), array_type);
 
     throw std::logic_error("Data type not supported");
 }
 
+size_t get_array_index(const std::vector<EclEntry>& array_list, const std::string& array_name, size_t occurence)
+{
+    size_t cidx = 0;
+
+    auto it = std::find_if(array_list.begin(), array_list.end(),
+                           [&cidx, &array_name, occurence](const EclEntry& entry)
+                           {
+                              if (std::get<0>(entry) == array_name)
+                                  ++cidx;
+
+                              return cidx == occurence + 1;
+                           });
+
+    return std::distance(array_list.begin(), it);
+}
+
+npArray get_vector_name(Opm::EclIO::EclFile * file_ptr, const std::string& array_name)
+{
+    if (file_ptr->hasKey(array_name) == false)
+        throw std::logic_error("Array " + array_name + " not found in EclFile");
+
+    auto array_list = file_ptr->getList();
+    size_t array_index = get_array_index(array_list, array_name, 0);
+
+    return get_vector_index(file_ptr, array_index);
+}
+
+npArray get_vector_occurrence(Opm::EclIO::EclFile * file_ptr, const std::string& array_name, size_t occurrence)
+{
+    if (occurrence >= file_ptr->count(array_name) )
+        throw std::logic_error("Occurrence " + std::to_string(occurrence) + " not found in EclFile");
+
+    auto array_list = file_ptr->getList();
+    size_t array_index = get_array_index(array_list, array_name, occurrence);
+
+    return get_vector_index(file_ptr, array_index);
+}
+
+
 }
 
 
@@ -51,10 +95,11 @@ void python::common::export_IO(py::module& m) {
 
     py::class_<Opm::EclIO::EclFile>(m, "EclFile")
         .def(py::init<const std::string &, bool>(), py::arg("filename"), py::arg("preload") = false)
-        .def_property_readonly("arrays", &Opm::EclIO::EclFile::getList)
-        .def("getListOfArrays", &Opm::EclIO::EclFile::getList)
+        .def("__get_list_of_arrays", &Opm::EclIO::EclFile::getList)
         .def("__contains__", &Opm::EclIO::EclFile::hasKey)
         .def("__len__", &Opm::EclIO::EclFile::size)
-        .def("__getitem__", &get_vector)
-        .def("get", &get_vector);
+        .def("count", &Opm::EclIO::EclFile::count)
+        .def("__get_data", &get_vector_index)
+        .def("__get_data", &get_vector_name)
+        .def("__get_data", &get_vector_occurrence);
 }

python/python/opm/io/ecl/__init__.py

@@ -1,24 +1,40 @@
 from opm._common import eclArrType
 from opm._common import EclFile
 
+import sys
+import datetime
+import numpy as np
 
 # When extracting the strings from CHAR keywords we get a character array, in
 # Python this becomes a list of bytes. This desperate monkey-patching is to
 # ensure the EclFile class returns normal Python strings in the case of CHAR
 # arrays. The return value is normal Python list of strings.
 
-def getitem(self, index):
-    data = self.__getitem(index)
-    array_type = self.arrays[index][1]
+@property
+def eclfile_get_list_of_arrays(self):
+
+    if sys.version_info.major == 2:
+        rawData = self.__get_list_of_arrays()
+        return [ ( x[0].encode("utf-8"), x[1], x[2] ) for x in rawData ]
+    else:
+        return self.__get_list_of_arrays()
+
+
+def getitem_eclfile(self, arg):
+
+    if isinstance(arg, tuple):
+        data, array_type = self.__get_data(str(arg[0]), int(arg[1]))
+    else:
+        data, array_type = self.__get_data(arg)
+
     if array_type == eclArrType.CHAR:
         return [ x.decode("utf-8") for x in data ]
 
     return data
 
 
-def make_getitem(cls):
-    setattr(cls, "__getitem", cls.__getitem__)
-    setattr(cls, "__getitem__", getitem)
 
 
-make_getitem(EclFile)
+setattr(EclFile, "__getitem__", getitem_eclfile)
+setattr(EclFile, "arrays", eclfile_get_list_of_arrays)
+

python/tests/test_eclfile.py

@@ -2,7 +2,7 @@ import unittest
 import sys
 import numpy as np
 
-from opm.io.ecl import EclFile
+from opm.io.ecl import EclFile, eclArrType
 try:
     from tests.utils import test_path
 except ImportError:
@@ -21,6 +21,7 @@ def array_index(ecl_file, target_kw):
 
 class TestEclFile(unittest.TestCase):
 
+
     def test_arrays(self):
 
         refList=["INTEHEAD","LOGIHEAD","DOUBHEAD","PORV","DEPTH","DX","DY","DZ","PORO",
@@ -32,12 +33,22 @@ class TestEclFile(unittest.TestCase):
         file2uf = EclFile(test_path("data/SPE9.INIT"), preload=False)
         self.assertEqual(len(file2uf), 24)
 
-        arrList = [ x[0] for x in file2uf.arrays ]
-        self.assertEqual(arrList, refList)
+        arr_string_list = [ x[0] for x in file2uf.arrays ]
+        self.assertEqual(arr_string_list, refList)
 
         file2f = EclFile(test_path("data/SPE9.FINIT"))
         self.assertEqual(len(file2f), 24)
+        self.assertTrue(isinstance(file2uf.arrays, list))
+        self.assertEqual(len(file2uf.arrays), len(refList))
 
+        for str1, str2 in zip(file2uf.arrays, refList):
+            self.assertEqual(str1[0], str2)
+
+        self.assertEqual( file2uf.arrays[3] , ("PORV", eclArrType.REAL, 9000) )
+        self.assertEqual( file2uf.arrays[16] , ("TABDIMS", eclArrType.INTE,100) )
+        self.assertEqual( file2uf.arrays[17] , ("TAB", eclArrType.DOUB, 885) )
+
+        
     def test_get_function(self):
 
         file1 = EclFile(test_path("data/SPE9.INIT"), preload=True)
@@ -125,7 +136,7 @@ class TestEclFile(unittest.TestCase):
         self.assertEqual(len(logih), 121)
         self.assertEqual(logih[0], True)
         self.assertEqual(logih[2], False)
-        self.assertEqual(logih[8], True)
+        self.assertEqual(logih[8], False)
 
     def test_get_function_char(self):
 
@@ -139,6 +150,43 @@ class TestEclFile(unittest.TestCase):
         self.assertEqual(keyw[0], "TIME")
         self.assertEqual(keyw[16], "FWCT")
 
+
+    def test_get_occurence(self):
+
+        file1 = EclFile(test_path("data/SPE9.UNRST"))
+
+        self.assertTrue("PRESSURE" in file1)
+
+        with self.assertRaises(RuntimeError):
+            test = file1["PRESSURE", int(10)]
+
+        #first occurence of pressure
+        pres = file1["PRESSURE"]
+        pres0 = file1["PRESSURE", 0]
+
+        self.assertEqual(len(pres), len(pres0))
+
+        for v1, v2 in zip(pres, pres0):
+            self.assertEqual(v1, v2)
+
+        #occurence number 2 of pressure
+        pres2 = file1["PRESSURE", 1]
+
+        self.assertTrue(isinstance(pres2, np.ndarray))
+        self.assertEqual(pres2.dtype, "float32")
+
+        self.assertEqual(len(pres2), 9000)
+
+        seqn0 = file1["SEQNUM", 0]
+        self.assertEqual(seqn0[0], 37)
+
+        seqn1 = file1["SEQNUM", 1]
+        self.assertEqual(seqn1[0], 74)
+
+        self.assertEqual(file1.count("PRESSURE"), 2)
+        self.assertEqual(file1.count("XXXX"), 0)
+
+
 if __name__ == "__main__":
 
     unittest.main()