get rid of kargs argument in the callback function

grainlearning/bayesian_calibration.py

@@ -96,7 +96,7 @@ def __init__(
 
         self.num_iter = num_iter
 
-        self.curr_iter = curr_iter
+        self.set_curr_iter(curr_iter)
 
         self.save_fig = save_fig
 
@@ -114,7 +114,7 @@ def run(self):
 
         # Bayesian calibration continue until curr_iter = num_iter or sigma_max < tolerance
         for _ in range(self.num_iter - 1):
-            self.curr_iter += 1
+            self.increase_curr_iter()
             print(f"Bayesian calibration iter No. {self.curr_iter}")
             self.run_one_iteration()
             if self.system.sigma_max < self.system.sigma_tol:
@@ -134,7 +134,7 @@ def run_one_iteration(self, index: int = -1):
         self.system.num_samples = self.system.param_data.shape[0]
 
         # Run the model realizations
-        self.system.run(curr_iter=self.curr_iter)
+        self.system.run()
 
         # Load model data from disk
         self.load_system()
@@ -150,8 +150,8 @@ def load_system(self):
         """Load existing simulation data from disk into the dynamic system
         """
         if isinstance(self.system, IODynamicSystem):
-            self.system.load_param_data(self.curr_iter)
-            self.system.get_sim_data_files(self.curr_iter)
+            self.system.load_param_data()
+            self.system.get_sim_data_files()
             self.system.load_sim_data()
         else:
             if self.system.param_data is None or self.system.sim_data is None:
@@ -165,7 +165,7 @@ def load_and_run_one_iteration(self):
         self.load_system()
         self.calibration.add_curr_param_data_to_list(self.system.param_data)
         self.calibration.solve(self.system)
-        self.system.write_params_to_table(self.curr_iter + 1)
+        self.system.write_params_to_table()
         self.calibration.sigma_list.append(self.system.sigma_max)
         self.plot_uq_in_time()
 
@@ -176,7 +176,7 @@ def load_and_process(self, sigma: float = 0.1):
         """
         self.load_system()
         self.calibration.add_curr_param_data_to_list(self.system.param_data)
-        self.calibration.load_proposal_from_file(self.system, self.curr_iter)
+        self.calibration.load_proposal_from_file(self.system)
         self.calibration.inference.data_assimilation_loop(sigma, self.system)
         self.system.compute_estimated_params(self.calibration.inference.posteriors)
 
@@ -189,7 +189,7 @@ def resample(self):
         self.calibration.posterior = self.calibration.inference.get_posterior_at_time()
         self.calibration.run_sampling(self.system, )
         resampled_param_data = self.calibration.param_data_list[-1]
-        self.system.write_params_to_table(self.curr_iter + 1)
+        self.system.write_params_to_table()
         return resampled_param_data
 
     def plot_uq_in_time(self):
@@ -247,6 +247,20 @@ def get_most_prob_params(self):
         most_prob = argmax(self.calibration.posterior)
         return self.system.param_data[most_prob]
 
+    def set_curr_iter(self, curr_iter: int):
+        """Set the current iteration step
+
+        param curr_iter: Current iteration step
+        """
+        self.system.curr_iter = curr_iter
+        self.curr_iter = self.system.curr_iter
+
+    def increase_curr_iter(self):
+        """Increase the current iteration step by one
+        """
+        self.system.curr_iter += 1
+        self.curr_iter += 1
+
     @classmethod
     def from_dict(
         cls: Type["BayesianCalibration"],

grainlearning/dynamic_systems.py

@@ -90,6 +90,7 @@ class DynamicSystem:
     :param param_min: List of parameter lower bounds
     :param param_max: List of parameter Upper bounds
     :param callback: Callback function, defaults to None
+    :param curr_iter: current iteration ID, defaults to 0
     :param ctrl_data: control data (e.g, time), defaults to None, optional
     :param obs_names: Column names of the observation data, defaults to None, optional
     :param ctrl_name: Column name of the control data, defaults to None, optional
@@ -162,6 +163,9 @@ class DynamicSystem:
     #: Callback function to run the forward predictions
     callback: Callable
 
+    #: Current iteration ID
+    curr_iter: int
+
     ##### Uncertainty #####
 
     #: Minimum value of the uncertainty
@@ -195,6 +199,7 @@ def __init__(
         sim_name: str = None,
         sim_data: np.ndarray = None,
         callback: Callable = None,
+        curr_iter: int = 0,
         param_data: np.ndarray = None,
         param_names: List[str] = None,
         sigma_max: float = 1.0e6,
@@ -232,6 +237,8 @@ def __init__(
 
         self.callback = callback
 
+        self.curr_iter = curr_iter
+
         self.param_min = param_min
 
         self.param_max = param_max
@@ -284,7 +291,7 @@ def from_dict(cls: Type["DynamicSystem"], obj: dict):
             sigma_tol=obj.get("sigma_tol", 0.001),
         )
 
-    def run(self, **kwargs):
+    def run(self):
         """Run the callback function
 
         TODO design a better wrapper to avoid kwargs?
@@ -294,7 +301,7 @@ def run(self, **kwargs):
         if self.callback is None:
             raise ValueError("No callback function defined")
 
-        self.callback(self, **kwargs)
+        self.callback(self)
 
     def set_sim_data(self, data: list):
         """Set the simulation data
@@ -346,19 +353,19 @@ def compute_estimated_params(self, posteriors: np.array):
                 self.estimated_params_cv[stp_id, :]) / self.estimated_params[stp_id, :]
 
     @classmethod
-    def load_param_data(cls: Type["DynamicSystem"], curr_iter: int):
+    def load_param_data(cls: Type["DynamicSystem"]):
         """Virtual function to load param data from disk"""
 
     @classmethod
-    def get_sim_data_files(cls: Type["DynamicSystem"], curr_iter: int):
+    def get_sim_data_files(cls: Type["DynamicSystem"]):
         """Virtual function to get simulation data files from disk"""
 
     @classmethod
     def load_sim_data(cls: Type["DynamicSystem"]):
         """Virtual function to load simulation data"""
 
     @classmethod
-    def write_params_to_table(cls, curr_iter):
+    def write_params_to_table(cls):
         """Write the parameter data into a text file"""
 
     @classmethod
@@ -425,6 +432,7 @@ class IODynamicSystem(DynamicSystem):
     :param sim_data_dir: Simulation data directory, defaults to './sim_data'
     :param sim_data_file_ext: Simulation data file extension, defaults to '.npy'
     :param callback: Callback function, defaults to None
+    :param curr_iter: Current iteration ID, defaults to 0
     :param param_data_file: Parameter data file, defaults to None, optional
     :param obs_data: observation or reference data, optional
     :param ctrl_data: Control data (e.g, time), defaults to None, optional
@@ -470,6 +478,7 @@ def __init__(
         inv_obs_weight: List[float] = None,
         sim_data: np.ndarray = None,
         callback: Callable = None,
+        curr_iter: int = 0,
         param_data_file: str = '',
         param_data: np.ndarray = None,
         param_names: List[str] = None,
@@ -490,6 +499,7 @@ def __init__(
             sim_name,
             sim_data,
             callback,
+            curr_iter,
             param_data,
             param_names
         )
@@ -584,10 +594,9 @@ def get_obs_data(self):
             if len(self.obs_data) == 1:
                 self.obs_data = self.obs_data.reshape([1, self.obs_data.shape[0]])
 
-    def get_sim_data_files(self, curr_iter: int = 0):
-        """Get the simulation data files from the simulation data directory.
-
-        :param curr_iter: Current iteration number, default to 0.
+    def get_sim_data_files(self):
+        """
+        Get the simulation data files from the simulation data directory.
         """
         mag = floor(log(self.num_samples, 10)) + 1
         self.sim_data_files = []
@@ -597,7 +606,7 @@ def get_sim_data_files(self, curr_iter: int = 0):
                 sim_data_file_ext = '_sim' + self.sim_data_file_ext
             else:
                 sim_data_file_ext = self.sim_data_file_ext
-            file_name = self.sim_data_dir.rstrip('/') + f'/iter{curr_iter}/{self.sim_name}*Iter{curr_iter}*' \
+            file_name = self.sim_data_dir.rstrip('/') + f'/iter{self.curr_iter}/{self.sim_name}*Iter{self.curr_iter}*' \
                         + str(i).zfill(mag) + '*' + sim_data_file_ext
             files = glob(file_name)
 
@@ -630,19 +639,17 @@ def load_sim_data(self):
             params = np.array([data[key] for key in self.param_names])
             np.testing.assert_allclose(params, self.param_data[i, :], rtol=1e-5)
 
-    def load_param_data(self, curr_iter: int = 0):
+    def load_param_data(self):
         """
         Load parameter data from a table written in a text file.
-
-        :param curr_iter: Current iteration number, default to 0.
         """
         if os.path.exists(self.param_data_file):
             # we assume parameter data are always in the last columns.
             self.param_data = np.genfromtxt(self.param_data_file, comments='!')[:, -self.num_params:]
             self.num_samples = self.param_data.shape[0]
         else:
             # if param_data_file does not exit, get parameter data from text files
-            files = glob(self.sim_data_dir + f'/iter{curr_iter}/{self.sim_name}*_param*{self.sim_data_file_ext}')
+            files = glob(self.sim_data_dir + f'/iter{self.curr_iter}/{self.sim_name}*_param*{self.sim_data_file_ext}')
             self.num_samples = len(files)
             self.sim_data_files = sorted(files)
             self.param_data = np.zeros([self.num_samples, self.num_params])
@@ -654,48 +661,44 @@ def load_param_data(self, curr_iter: int = 0):
                 params = [data[key][0] for key in self.param_names]
                 self.param_data[i, :] = params
 
-    def run(self, **kwargs):
-        """Run the callback function
-
-        TODO design a better wrapper to avoid kwargs?
-        :param kwargs: keyword arguments to pass to the callback function
+    def run(self):
+        """
+        Run the callback function
         """
 
         if self.callback is None:
             raise ValueError("No callback function defined")
 
         # create a directory to store simulation data
-        curr_iter = kwargs['curr_iter']
         sim_data_dir = self.sim_data_dir.rstrip('/')
-        sim_data_sub_dir = f'{sim_data_dir}/iter{curr_iter}'
+        sim_data_sub_dir = f'{sim_data_dir}/iter{self.curr_iter}'
         os.makedirs(sim_data_sub_dir)
 
         # write the parameter data into a text file
-        self.write_params_to_table(curr_iter)
+        self.write_params_to_table()
 
         # run the callback function
-        self.callback(self, **kwargs)
+        self.callback(self)
 
         # move simulation data files and corresponding parameter table into the directory defined per iteration
-        files = glob(f'{os.getcwd()}/{self.sim_name}_Iter{curr_iter}*{self.sim_data_file_ext}')
+        files = glob(f'{os.getcwd()}/{self.sim_name}_Iter{self.curr_iter}*{self.sim_data_file_ext}')
         for file in files:
             f_name = os.path.relpath(file, os.getcwd())
             os.replace(f'{file}', f'{sim_data_sub_dir}/{f_name}')
 
         # redefine the parameter data file since its location is changed
         self.param_data_file = f'{sim_data_sub_dir}/' + os.path.relpath(self.param_data_file, os.getcwd())
 
-    def write_params_to_table(self, curr_iter: int):
+    def write_params_to_table(self):
         """Write the parameter data into a text file.
 
-        :param curr_iter: Current iteration number, default to 0.
         :return param_data_file: The name of the parameter data file
         """
         self.param_data_file = write_to_table(
             self.sim_name,
             self.param_data,
             self.param_names,
-            curr_iter)
+            self.curr_iter)
 
     def backup_sim_data(self):
         """Backup simulation data files to a backup directory."""

grainlearning/iterative_bayesian_filter.py

@@ -149,15 +149,14 @@ def initialize(self, system: Type["DynamicSystem"]):
         system.param_data = generate_params_qmc(system, system.num_samples, self.initial_sampling)
         self.param_data_list.append(system.param_data)
 
-    def run_inference(self, system: Type["DynamicSystem"], curr_iter: int = 0):
+    def run_inference(self, system: Type["DynamicSystem"]):
         """Compute the posterior distribution of model states such that the target effective sample size is reached.
 
         :param system: Dynamic system class
-        :param curr_iter: Current iteration number
         """
         # if the name of proposal data file is given, make use of the proposal density during Bayesian updating
         if self.proposal_data_file is not None and self.proposal is None:
-            self.load_proposal_from_file(system, curr_iter=curr_iter)
+            self.load_proposal_from_file(system)
 
         result = optimize.minimize_scalar(
             self.inference.data_assimilation_loop,
@@ -208,11 +207,10 @@ def add_curr_param_data_to_list(self, param_data: np.ndarray):
         """
         self.param_data_list.append(param_data)
 
-    def load_proposal_from_file(self, system: Type["IODynamicSystem"], curr_iter: int):
+    def load_proposal_from_file(self, system: Type["IODynamicSystem"]):
         """Load the proposal density from a file.
 
         :param system: Dynamic system class
-        :param curr_iter: Current iteration number
         """
         if system.param_data is None:
             raise RuntimeError("parameter samples not yet loaded...")
@@ -221,7 +219,7 @@ def load_proposal_from_file(self, system: Type["IODynamicSystem"], curr_iter: in
             return
 
         # load the proposal density from a file
-        self.sampling.load_gmm_from_file(f'{system.sim_data_dir}/iter{curr_iter-1}/{self.proposal_data_file}')
+        self.sampling.load_gmm_from_file(f'{system.sim_data_dir}/iter{system.curr_iter-1}/{self.proposal_data_file}')
 
         samples = np.copy(system.param_data)
         samples /= self.sampling.max_params
@@ -235,10 +233,9 @@ def load_proposal_from_file(self, system: Type["IODynamicSystem"], curr_iter: in
             proposal[np.where(proposal < 0.0)] = min(proposal[np.where(proposal > 0.0)])
             self.proposal = proposal / sum(proposal)
 
-    def save_proposal_to_file(self, system: Type["IODynamicSystem"], curr_iter: int):
+    def save_proposal_to_file(self, system: Type["IODynamicSystem"]):
         """Save the proposal density to a file.
 
         :param system: Dynamic system class
-        :param curr_iter: Current iteration number
         """
-        self.sampling.save_gmm_to_file(f'{system.sim_data_dir}/iter{curr_iter-1}/{self.proposal_data_file}')
+        self.sampling.save_gmm_to_file(f'{system.sim_data_dir}/iter{system.curr_iter-1}/{self.proposal_data_file}')

tests/integration/test_gmm.py

@@ -57,7 +57,7 @@ def test_gmm():
 
     # %%
     # write new parameter table to the simulation directory
-    calibration.system.write_params_to_table(calibration.curr_iter + 1)
+    calibration.system.write_params_to_table()
 
     # %%
     check_list = np.isclose(resampled_param_data_ref, resampled_param_data)

tests/integration/test_lenreg_IO.py

@@ -12,18 +12,17 @@
 executable = f'python {PATH}/linear_model.py'
 
 
-def run_sim(model, **kwargs):
+def run_sim(model):
     """
     Run the external executable and passes the parameter sample to generate the output file.
     """
     # keep the naming convention consistent between iterations
     mag = floor(log(model.num_samples, 10)) + 1
-    curr_iter = kwargs['curr_iter']
     # check the software name and version
     print("*** Running external software... ***\n")
     # loop over and pass parameter samples to the executable
     for i, params in enumerate(model.param_data):
-        description = 'Iter' + str(curr_iter) + '_Sample' + str(i).zfill(mag)
+        description = 'Iter' + str(model.curr_iter) + '_Sample' + str(i).zfill(mag)
         print(" ".join([executable, "%.8e %.8e" % tuple(params), description]))
         os.system(' '.join([executable, "%.8e %.8e" % tuple(params), description]))