From 54f5c7c54817005354216d31bc1e74c6a94c88bc Mon Sep 17 00:00:00 2001
From: Pierre-Antoine Rouby <contact@parouby.fr>
Date: Thu, 30 Jul 2026 11:13:53 +0200
Subject: [PATCH] Ensemble: Switch data pid to data ref, and minor change.

---
 src/Model/Ensembles/Ensemble.py | 26 +++++++++++++++++++++++---
 src/Model/Ensembles/Function.py |  9 ++++-----
 2 files changed, 27 insertions(+), 8 deletions(-)

diff --git a/src/Model/Ensembles/Ensemble.py b/src/Model/Ensembles/Ensemble.py
index 25f0fc80..1b6b7e6d 100644
--- a/src/Model/Ensembles/Ensemble.py
+++ b/src/Model/Ensembles/Ensemble.py
@@ -31,7 +31,7 @@ class Ensemble(SQLSubModel):
 
     def __init__(self, id: int = -1,
                  name: str = "",
-                 function=-1, range=[], data_pid=-1,
+                 function=-1, range=[], target_data=None,
                  status=None, owner_scenario=-1):
         super(Ensemble, self).__init__(
             id=id, status=status,
@@ -42,7 +42,7 @@ class Ensemble(SQLSubModel):
         self._type = "generic"
         self._function = function
         self._range = range
-        self._data_pid = data_pid
+        self._target_data = target_data
 
     @classmethod
     def _db_create(cls, execute, ext=""):
@@ -84,6 +84,7 @@ class Ensemble(SQLSubModel):
         new = []
         scenario = data["scenario"]
         functions = data["ens_functions"]
+        stricklers = data["stricklers"]
         loaded = data['loaded_pid']
 
         table = execute(
@@ -115,6 +116,15 @@ class Ensemble(SQLSubModel):
 
             function = next(
                 filter(fn, functions.lst),
+                None
+            )
+
+            target_data = next(
+                filter(
+                    lambda s: s._pamhyr_id == data_pid,
+                    stricklers.lst
+                ),
+                None
             )
 
             rdata = cls._decode_range(range, range_len)
@@ -123,6 +133,7 @@ class Ensemble(SQLSubModel):
                 id, name=name,
                 function=function,
                 range=rdata,
+                target_data=target_data,
                 status=data["status"],
                 owner_scenario=owner_scenario
             )
@@ -150,6 +161,10 @@ class Ensemble(SQLSubModel):
         if self._function._type == "generic":
             fid = -1
 
+        data_pid = -1
+        if self._target_data is not None:
+            data_pid = self._target_data._pamhyr_id
+
         execute(
             "INSERT INTO " +
             "ensemble(pamhyr_id, deleted, name, type, " +
@@ -159,7 +174,7 @@ class Ensemble(SQLSubModel):
             self.pamhyr_id, self.is_deleted(),
             self._name, self._type,
             fid, self._function._name, brange, length,
-            self._data_pid, self._status.scenario_id
+            data_pid, self._status.scenario_id
         )
 
         return True
@@ -175,3 +190,8 @@ class Ensemble(SQLSubModel):
         data_format = ">" + ''.join(itertools.repeat("d", length))
 
         return struct.unpack(data_format, brange)
+
+    def get_sample(self, nb):
+        return self._function.get_sample(
+            self._range, nb
+        )
diff --git a/src/Model/Ensembles/Function.py b/src/Model/Ensembles/Function.py
index 38ad7f57..81250bdf 100644
--- a/src/Model/Ensembles/Function.py
+++ b/src/Model/Ensembles/Function.py
@@ -150,12 +150,11 @@ class Uniform(Function):
             name=name, script=script
         )
 
-    def get_sample(self, ens_range):
+    def get_sample(self, ens_range, nb):
         b = ens_range[0]
         t = ens_range[1]
-        n = ens_range[2]
 
-        step = (t - b) / n
+        step = (t - b) / nb
 
         return np.arange(b, t, step)
 
@@ -173,11 +172,11 @@ class Custom(Function):
 
         self._type = "custom"
 
-    def get_sample(self, ens_range):
+    def get_sample(self, ens_range, nb):
         try:
             # Run script to (re)define the sample function
             f = eval(self._script)
-            return f(ens_range)
+            return f(ens_range, nb)
         except Exception as e:
             logger_exception(e)
             return None         # TODO: Raise helpful exception