Export shared rng to Python

equinor · Feb 15, 2021 · c1cc5f5 · c1cc5f5
1 parent d5737fc
commit c1cc5f5
Show file tree

Hide file tree

Showing 4 changed files with 19 additions and 0 deletions.
diff --git a/lib/enkf/enkf_main.cpp b/lib/enkf/enkf_main.cpp
@@ -1700,6 +1700,10 @@ rng_config_type * enkf_main_get_rng_config( const enkf_main_type * enkf_main ) {
   return res_config_get_rng_config(enkf_main->res_config);
 }
 
+rng_type * enkf_main_get_shared_rng( enkf_main_type * enkf_main ) {
+  return enkf_main->shared_rng;
+}
+
 
 void enkf_main_rng_init( enkf_main_type * enkf_main) {
   enkf_main->rng_manager = rng_config_alloc_rng_manager( enkf_main_get_rng_config(enkf_main) );

diff --git a/lib/include/ert/enkf/enkf_main.hpp b/lib/include/ert/enkf/enkf_main.hpp
@@ -129,6 +129,7 @@ extern "C" {
 
   rng_config_type     * enkf_main_get_rng_config( const enkf_main_type * enkf_main );
   void                  enkf_main_rng_init( enkf_main_type * enkf_main);
+  rng_type            * enkf_main_get_shared_rng( enkf_main_type * enkf_main );
 
   bool enkf_main_export_field(const enkf_main_type * enkf_main,
                               const char * kw,

diff --git a/python/res/enkf/enkf_main.py b/python/res/enkf/enkf_main.py
@@ -35,6 +35,7 @@
 from res.enkf.key_manager import KeyManager
 from res.util import Log
 from res.util.substitution_list import SubstitutionList
+from ecl.util.util import rng
 
 
 class EnKFMain(BaseCClass):
@@ -252,6 +253,7 @@ class _RealEnKFMain(BaseCClass):
     _add_node = ResPrototype("void enkf_main_add_node(enkf_main, enkf_config_node)")
     _get_res_config = ResPrototype("res_config_ref enkf_main_get_res_config(enkf_main)")
     _init_run = ResPrototype("void enkf_main_init_run(enkf_main, ert_run_context)")
+    _get_shared_rng = ResPrototype("rng_ref enkf_main_get_shared_rng(enkf_main)")
 
     def __init__(self, config, strict=True, verbose=False):
         """ Please don't use this class directly. See EnKFMain instead """
@@ -454,3 +456,7 @@ def getRunpathList(self):
 
     def addNode(self, enkf_config_node):
         self._add_node(enkf_config_node)
+
+    def rng(self):
+        "Will return the random number generator used for updates."
+        return self._get_shared_rng()
diff --git a/tests/res/enkf/test_enkf.py b/tests/res/enkf/test_enkf.py
@@ -280,6 +280,14 @@ def test_run_context(self):
             run_arg = run_context[2]
             self.assertTrue(isinstance(run_arg, RunArg))
 
+            rng1 = main.rng()
+            rng1.setState("ABCDEFGHIJ012345")
+            d1 = rng1.getDouble()
+            rng1.setState("ABCDEFGHIJ012345")
+            rng2 = main.rng()
+            d2 = rng2.getDouble()
+            self.assertEqual(d1, d2)
+
     @tmpdir()
     def test_run_context_from_external_folder(self):
         with TestAreaContext("enkf_test") as work_area: