src/Model/Ensembles/Ensemble.py

@@ -31,7 +31,10 @@ class Ensemble(SQLSubModel):
 
     def __init__(self, id: int = -1,
                  name: str = "",
-                 function=-1, range=[], data_pid=-1,
+                 prob_dist="generic",
+                 function="generic",
+                 range=[],
+                 data_pid=-1,
                  status=None, owner_scenario=-1):
         super(Ensemble, self).__init__(
             id=id, status=status,
@@ -40,6 +43,7 @@ class Ensemble(SQLSubModel):
 
         self._name = name
         self._type = "generic"
+        self._prob_dist = prob_dist
         self._function = function
         self._range = range
         self._data_pid = data_pid
@@ -52,8 +56,8 @@ class Ensemble(SQLSubModel):
             deleted BOOLEAN NOT NULL DEFAULT FALSE,
             name TEXT NOT NULL,
             type VARCHAR(8) NOT NULL,
-            function INTEGER NOT NULL,
-            function_name VARCHAR(32) NOT NULL,
+            prob_dist VARCHAR(8) NOT NULL,
+            function VARCHAR(8) NOT NULL,
             range BLOB NOT NULL,
             range_len INTEGER NOT NULL,
             data_pid INTEGER NOT NULL,
@@ -83,12 +87,11 @@ class Ensemble(SQLSubModel):
     def _db_load(cls, execute, data=None):
         new = []
         scenario = data["scenario"]
-        functions = data["ens_functions"]
         loaded = data['loaded_pid']
 
         table = execute(
             "SELECT pamhyr_id, deleted, name, type, " +
-            "function, function_name, range, range_len, data_pid, " +
+            "prob_dist, function, range, range_len, data_pid, " +
             "scenario " +
             "FROM ensemble " +
             f"WHERE scenario = {scenario.id} " +
@@ -102,25 +105,18 @@ class Ensemble(SQLSubModel):
             deleted = (next(it) == 1)
             name = next(it)
             type = next(it)
-            fid = next(it)
-            fname = next(it)
+            prob_dist = next(it)
+            function = next(it)
             range = next(it)
             range_len = next(it)
             data_pid = next(it)
             owner_scenario = next(it)
 
-            def fn(f): return f._name == fname
-            if fid != -1:
-                def fn(f): return f._pamhyr_id == fid
-
-            function = next(
-                filter(fn, functions.lst),
-            )
-
             rdata = cls._decode_range(range, range_len)
 
             new_ensemble = cls(
                 id, name=name,
+                prob_dist=prob_dist,
                 function=function,
                 range=rdata,
                 status=data["status"],
@@ -146,24 +142,27 @@ class Ensemble(SQLSubModel):
 
         brange, length = self._encode_range()
 
-        fid = self._function._pamhyr_id
-        if self._function._type == "generic":
-            fid = -1
-
         execute(
             "INSERT INTO " +
             "ensemble(pamhyr_id, deleted, name, type, " +
-            "  function, function_name, range, range_len, " +
+            "  prob_dist, function, range, range_len, " +
             "  data_pid, scenario) " +
             "VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",
             self.pamhyr_id, self.is_deleted(),
             self._name, self._type,
-            fid, self._function._name, brange, length,
+            self._prob_dist, self._function, brange, length,
             self._data_pid, self._status.scenario_id
         )
 
         return True
 
+    def _data_traversal(self,
+                        predicate=lambda obj, data: True,
+                        modifier=lambda obj, data: None,
+                        data={}):
+        if predicate(self, data):
+            modifier(self, data)
+
     def _encode_range(self):
         length = len(self._range)
         data_format = ">" + ''.join(itertools.repeat("d", length))

src/Model/Ensembles/EnsembleList.py

@@ -41,6 +41,9 @@ class EnsembleList(PamhyrModelList):
             f"WHERE scenario = {self._status.scenario_id}"
         )
 
+        if data is None:
+            data = {}
+
         for ens in self._lst:
             ens._db_save(execute, data=data)
 

src/Model/Ensembles/Function.py

@@ -19,7 +19,6 @@
 import struct
 import logging
 import itertools
-import numpy as np
 
 from tools import logger_exception
 
@@ -150,14 +149,8 @@ class Uniform(Function):
             name=name, script=script
         )
 
-    def get_sample(self, ens_range):
-        b = ens_range[0]
-        t = ens_range[1]
-        n = ens_range[2]
-
-        step = (t - b) / n
-
-        return np.arange(b, t, step)
+    def get_sample(self, dist, ens_range):
+        return range(ens_range[0], ens_range[1])
 
 
 class Custom(Function):
@@ -173,11 +166,11 @@ class Custom(Function):
 
         self._type = "custom"
 
-    def get_sample(self, ens_range):
+    def get_sample(self, dist, ens_range):
         try:
             # Run script to (re)define the sample function
             f = eval(self._script)
-            return f(ens_range)
+            return f(dist, ens_range)
         except Exception as e:
             logger_exception(e)
             return None         # TODO: Raise helpful exception

src/Model/Ensembles/FunctionList.py

@@ -29,12 +29,7 @@ class FunctionList(PamhyrModelList):
     def _db_load(cls, execute, data=None):
         new = cls(status=data['status'])
 
-        # Default functions
-        new._lst = [
-            Uniform(status=data['status']),
-        ]
-
-        # DB custom functions
+        new._lst = [Uniform(status=data['status'])]
         new._lst += Function._db_load(
             execute, data
         )

src/Model/Study.py

@@ -87,7 +87,7 @@ class Study(SQLModel):
             self.scenarios[0] = s0
             self.status.scenario = s0
 
-            self._ens_functions = FunctionList(status=self.status)
+            self._ens_function = FunctionList(status=self.status)
 
             self._river = River(status=self.status)
         else:
@@ -400,12 +400,10 @@ class Study(SQLModel):
         data["loaded_pid"] = set()
 
         # Get ensemble function
-        new._ens_functions = FunctionList._db_load(
+        new._ens_function = FunctionList._db_load(
             sql_exec, data=data
         )
 
-        data["ens_functions"] = new._ens_functions
-
         # Load river data
         new._river = River._db_load(
             sql_exec, data=data
@@ -468,7 +466,7 @@ class Study(SQLModel):
 
         self._save_submodel(
             [
-                self.scenarios, self._ens_functions,
+                self.scenarios, self._ens_function,
                 self._river
             ],
             data=progress
@@ -588,8 +586,7 @@ class Study(SQLModel):
         data = {
             "status": self.status,
             "loaded_pid": set(),
-            "scenario": scenario,
-            "ens_functions": self._ens_functions,
+            "scenario": scenario
         }
 
         self.status.start_loading()