From 8da6669f59df533b5a8f44d193706e42f59c50d1 Mon Sep 17 00:00:00 2001
From: chyalexcheng <chyalexcheng@gmail.com>
Date: Thu, 1 Feb 2024 23:08:39 +0100
Subject: [PATCH 01/21] Create a new preprocessor to handle data directly
 coming from GL. Change naming from load_sequences, contact_parameters,
 outputs to inputs, params, outputs

---
 grainlearning/rnn/models.py            |   6 +-
 grainlearning/rnn/predict.py           |  22 +--
 grainlearning/rnn/preprocessor.py      | 200 ++++++++++++++++++++++---
 grainlearning/rnn/windows.py           |  20 +--
 tutorials/data_driven/LSTM/test_rnn.py | 112 ++++++++++++++
 5 files changed, 313 insertions(+), 47 deletions(-)
 create mode 100644 tutorials/data_driven/LSTM/test_rnn.py

diff --git a/grainlearning/rnn/models.py b/grainlearning/rnn/models.py
index e08c8fe7..391d0916 100644
--- a/grainlearning/rnn/models.py
+++ b/grainlearning/rnn/models.py
@@ -20,7 +20,7 @@ def rnn_model(
     Takes in a load sequence and contact parameters, and outputs the macroscopic responses.
     The contact parameters are used to initialize the hidden state of the LSTM.
 
-    :param input_shapes: Dictionary containing `'num_load_features'`, `'num_contact_params'`,
+    :param input_shapes: Dictionary containing `'num_input_features'`, `'num_params'`,
         `'num_labels'`. It can contain other keys but hese are the ones used here.
     :param window_size: Length of time window.
     :param lstm_units: Number of units of the hidden state of the LSTM.
@@ -34,8 +34,8 @@ def rnn_model(
 
     sequence_length = window_size
     load_sequence = layers.Input(
-            shape=(sequence_length, input_shapes['num_load_features']), name='load_sequence')
-    contact_params = layers.Input(shape=(input_shapes['num_contact_params'],), name='contact_parameters')
+            shape=(sequence_length, input_shapes['num_input_features']), name='inputs')
+    contact_params = layers.Input(shape=(input_shapes['num_params'],), name='params')
 
     # compute hidden state of LSTM based on contact parameters
     state_h = layers.Dense(lstm_units, activation='tanh', name='state_h')(contact_params)
diff --git a/grainlearning/rnn/predict.py b/grainlearning/rnn/predict.py
index ab7807d7..f42ee89e 100644
--- a/grainlearning/rnn/predict.py
+++ b/grainlearning/rnn/predict.py
@@ -68,7 +68,7 @@ def get_pretrained_model(path_to_model: str):
     :return:
         - model: keras model ready to use.
         - train_stats: Array containing the values used to standardize the data (if config.standardize_outputs = True),
-          and lenghts of sequences, load_features, contact_params, labels, window_size and window_step.
+          and lengths of sequences, input_features, params, labels, window_size and window_step.
         - config: dictionary with the model configuration
     """
     path_to_model = Path(path_to_model)
@@ -166,7 +166,7 @@ def predict_macroscopics(
     If 'standardize_outputs' in config, rescale the predictions to their original units.
 
     :param model: Keras RNN model
-    :param data: Tensorflow dataset containing inputs: 'load_sequence' and 'contact_parameters', and outputs.
+    :param data: Tensorflow dataset containing inputs: 'inputs' and 'params', and outputs.
     :param train_stats: Dictionary containing statistics of the training set.
     :param config: Dictionary containing the configuration with which the model was trained.
     :param batch_size: Size of batches to use.
@@ -177,18 +177,18 @@ def predict_macroscopics(
     inputs = list(data)[0][0]
 
     # Check that input sizes of data correspond to those of the pre-trained model
-    if inputs['load_sequence'].shape[2] != train_stats['num_load_features']:
-        raise ValueError(f"Number of elements in load_sequence of data does not match the model load_sequence shape. \
-            Got {inputs['load_sequence'].shape[2]}, expected {train_stats['num_load_features']}.")
+    if inputs['inputs'].shape[2] != train_stats['num_input_features']:
+        raise ValueError(f"Number of elements in inputs of data does not match the model inputs shape. \
+            Got {inputs['inputs'].shape[2]}, expected {train_stats['num_input_features']}.")
 
-    if inputs['contact_parameters'].shape[1] != train_stats['num_contact_params']:
-        raise ValueError(f"Number of elements in contact_parameters of data does not match the model \
-            contact_parameters shape. \
-            Got {inputs['contact_parameters'].shape[2]}, expected {train_stats['num_contact_params']}.")
+    if inputs['params'].shape[1] != train_stats['num_params']:
+        raise ValueError(f"Number of elements in params of data does not match the model \
+            params shape. \
+            Got {inputs['params'].shape[2]}, expected {train_stats['num_params']}.")
 
-    if inputs['load_sequence'].shape[1] != train_stats['sequence_length']:
+    if inputs['inputs'].shape[1] != train_stats['sequence_length']:
         raise ValueError(f"Sequence length of the train_stats {train_stats['sequence_length']} does not match \
-            that of the data {inputs['load_sequence'].shape[1]}. If the train_stats are does of the model, \
+            that of the data {inputs['inputs'].shape[1]}. If the train_stats are does of the model, \
             check pad_length. Can be that the trained model is not compatible.")
 
     predictions = predict_over_windows(inputs, model, config['window_size'], train_stats['sequence_length'])
diff --git a/grainlearning/rnn/preprocessor.py b/grainlearning/rnn/preprocessor.py
index 079b80ae..56a65a9b 100644
--- a/grainlearning/rnn/preprocessor.py
+++ b/grainlearning/rnn/preprocessor.py
@@ -2,6 +2,7 @@
 import numpy as np
 import tensorflow as tf
 import warnings
+from typing import Type
 
 from abc import ABC, abstractmethod
 
@@ -12,6 +13,15 @@ class Preprocessor(ABC):
     def __init__(self):
         pass
 
+    @classmethod
+    def from_dict(cls: Type["Preprocessor"], obj: dict):
+        """ Initialize the class using a dictionary style
+
+        :param obj: Dictionary object
+        :return: Preprocessor: Preprocessor object
+        """
+        raise NotImplementedError
+
     @abstractmethod
     def prepare_datasets(self, **kwargs):
         """
@@ -20,7 +30,7 @@ def prepare_datasets(self, **kwargs):
 
           * ``data_split``: Dictionary with keys ``'train'``, ``'val'``, ``'test'``, and values of the corresponding tensorflow Datasets.
           * ``train_stats``: Dictionary containing the shape of the data:
-                  ``sequence_length``, ``num_input_features``, ``num_contact_params``, ``num_labels``,
+                  ``sequence_length``, ``num_input_features``, ``num_params``, ``num_labels``,
                   and ``'mean'`` and ``'std'`` of the training set, in case ``standardize_outputs`` is True.
         """
         pass
@@ -31,8 +41,8 @@ def prepare_single_dataset(self):
         Abstract class to be implemented in each of the child classes.
         Must return a Tensorflow dataset containing a Tuple (inputs, outputs)
 
-          * ``inputs``: Dictionary with keys ``'load_sequence'``, ``'contact_parameters'``, that are the
-            corresponding tensorflow Datasets to input to an rnn model.
+          * ``inputs``: Dictionary with keys ``'input_sequence'``, ``'parameters'``, that are the
+            corresponding tensorflow Datasets to input to a rnn model.
 
           * ``outputs``: tensorflow Dataset containing the outputs or labels.
         """
@@ -48,7 +58,7 @@ def warning_config_field(cls, key: str, config: dict, default, add_default_to_co
         :param key: string of the key in dictionary ``config`` to look for.
         :param config: Dictionary containing keys and values
         :param default: default value associated to the key in config.
-        :param add_default_to_config: Wether to add the default value and the key to config in case it does not exists.
+        :param add_default_to_config: Whether to add the default value and the key to config in case it does not exist.
 
         :return: Updated ``config`` dictionary.
         """
@@ -68,8 +78,7 @@ def _custom_format_warning(msg, *_):
     def make_splits(cls, dataset: tuple, train_frac: float, val_frac: float, seed: int):
         """
         Split data into training, validation, and test sets.
-        The split is done on a sample by sample basis, so sequences are not broken up.
-        It is done randomly across all pressures and experiment types present.
+        The split is done on a sample-by-sample basis, so sequences are not broken up.
 
         :param dataset: Full dataset to split on, in form of a tuple (inputs, outputs),
                 where inputs is a dictionary and its keys and the outputs are numpy arrays.
@@ -92,8 +101,8 @@ def make_splits(cls, dataset: tuple, train_frac: float, val_frac: float, seed: i
                              f"lead to have {n_tot - n_train - n_val} samples in test dataset.")
 
         np.random.seed(seed=seed)
-        inds = np.random.permutation(np.arange(n_tot))
-        i_train, i_val, i_test = inds[:n_train], inds[n_train:n_train + n_val], inds[n_train + n_val:]
+        ids = np.random.permutation(np.arange(n_tot))
+        i_train, i_val, i_test = ids[:n_train], ids[n_train:n_train + n_val], ids[n_train + n_val:]
         split_data = {
                      'train': cls.get_split(dataset, i_train),
                      'val': cls.get_split(dataset, i_val),
@@ -102,21 +111,21 @@ def make_splits(cls, dataset: tuple, train_frac: float, val_frac: float, seed: i
         return split_data
 
     @classmethod
-    def get_split(cls, dataset:tuple, inds: np.array):
+    def get_split(cls, dataset:tuple, ids: np.array):
         """
-        Given a dataset and the indexes, return  the sub-dataset containing only the elements of indexes in ``inds``.
+        Given a dataset and the indexes, return  the sub-dataset containing only the elements of indexes in ``ids``.
         The returned dataset respects the format of ``inputs`` and ``outputs`` (see more info in return).
 
         :param dataset: Full dataset to split on, in form of a tuple ``(inputs, outputs)``,
           where ``inputs`` is a dictionary and its keys and ``outputs`` are numpy arrays.
-        :param inds: Indexes of the elements in ``dataset`` that are going to be gathered in this specific split.
+        :param ids: Indexes of the elements in ``dataset`` that are going to be gathered in this specific split.
 
         :return: tuple of the split, 2 dimensions:
           * inputs: dictionary containing ``'load_sequence'`` and ``'contact_params'``.
           * outputs: Labels
         """
-        X = {key: tf.gather(val, inds) for key, val in dataset[0].items()}
-        y = tf.gather(dataset[1], inds)
+        X = {key: tf.gather(val, ids) for key, val in dataset[0].items()}
+        y = tf.gather(dataset[1], ids)
 
         return X, y
 
@@ -128,7 +137,7 @@ def pad_initial(cls, array: np.array, pad_length: int, axis=1):
         of the same size.
 
         :param array: Array that is going to be modified
-        :param pad_lenght: number of copies of the initial state to be added at the beggining of ``array``.
+        :param pad_length: number of copies of the initial state to be added at the beginning of ``array``.
 
         :return: Modified array
         """
@@ -142,7 +151,7 @@ def pad_initial(cls, array: np.array, pad_length: int, axis=1):
     def standardize_outputs(cls, split_data: dict):
         """
         Standardize outputs of ``split_data`` using the mean and std of the training data
-        taken over both the samples and the timesteps.
+        taken over both the samples and the time steps.
         The 3 datasets ``'train'``, ``'val'``, ``'test'`` will be standardized.
 
         :param split_data: dictionary containing ``'train'``, ``'val'``, ``'test'`` keys and the respective datasets.
@@ -175,19 +184,163 @@ def get_dimensions(cls, data: tf.data.Dataset):
         :param data: The dataset to extract from.
 
         :return: Dictionary containing:
-                sequence_length, num_load_features, num_contact_params, num_labels
+                sequence_length, num_input_features, num_params, num_labels
         """
         train_sample = next(iter(data))  # just to extract a single batch
-        sequence_length, num_load_features = train_sample[0]['load_sequence'].shape
-        num_contact_params = train_sample[0]['contact_parameters'].shape[0]
+        sequence_length, num_input_features = train_sample[0]['inputs'].shape
+        num_params = train_sample[0]['params'].shape[0]
         num_labels = train_sample[1].shape[-1]
 
         return {'sequence_length': sequence_length,
-                'num_load_features': num_load_features,
-                'num_contact_params': num_contact_params,
+                'num_input_features': num_input_features,
+                'num_params': num_params,
                 'num_labels': num_labels,
                 }
 
+    def get_default_config(self):
+        pass
+
+
+class PreprocessorLSTM(Preprocessor):
+    """
+    Class to preprocess data for LSTM models which requires certain windowizing, inheriting from abstract class :class:`Preprocessor`
+    """
+    def __init__(
+        self,
+        input_data: np.ndarray,
+        param_data: np.ndarray,
+        output_data: np.ndarray,
+        train_frac: float,
+        val_frac: float,
+        pad_length: int,
+        window_size: int,
+        window_step: int = 1,
+        standardize_outputs: bool = True,
+    ):
+        """
+        Constructor of the preprocessor for LSTM models.
+        Initializes the values of required attributes to preprocess the data, from kwargs:
+
+        :param input_data: input sequence
+        :param param_data: ndarray of parameters
+        :param output_data: ndarray of output data
+        :param train_frac: Fraction of data used in the training set.
+        :param val_frac: Fraction of the data used in the validation set.
+        :param pad_length: Amount by which to pad the sequences from the start.
+        :param window_size: Number of time steps to include in a window.
+        :param window_step: Offset between subsequent windows.
+        :param standardize_outputs: Whether to transform the training set labels to have zero mean and unit variance.
+        """
+        super().__init__()
+        # reshape input_data to be of shape (sequence_length, num_input_features)
+        input_data = np.reshape(input_data, (input_data.shape[0], 1))
+        self.input_data = input_data
+        self.param_data = param_data
+        # swap the last two axes of the output data
+        output_data = np.swapaxes(output_data, 1, 2)
+        self.output_data = output_data
+        self.train_frac = train_frac
+        self.val_frac = val_frac
+        self.pad_length = pad_length
+        self.window_size = window_size
+        self.window_step = window_step
+        self.standardize_outputs = standardize_outputs
+
+    @classmethod
+    def from_dict(cls: Type["PreprocessorLSTM"], obj: dict):
+        """ Initialize the class using a dictionary style
+
+        :param obj: Dictionary object
+        :return: PreprocessorLSTM: PreprocessorLSTM object
+        """
+
+        assert "input_data" in obj.keys(), "Error no raw_data key found in input"
+        assert "param_data" in obj.keys(), "Error no train_frac key found in input"
+        assert "output_data" in obj.keys(), "Error no val_frac key found in input"
+        assert "train_frac" in obj.keys(), "Error no train_frac key found in input"
+        assert "val_frac" in obj.keys(), "Error no val_frac key found in input"
+        assert "window_size" in obj.keys(), "Error no window_size key found in input"
+        assert "window_step" in obj.keys(), "Error no window_step key found in input"
+        assert "standardize_outputs" in obj.keys(), "Error no standardize_outputs key found in input"
+
+        return cls(
+            input_data=obj["input_data"],
+            param_data=obj["param_data"],
+            output_data=obj["output_data"],
+            train_frac=obj["train_frac"],
+            val_frac=obj["val_frac"],
+            window_size=obj["window_size"],
+            pad_length=obj.get("pad_length", 0),
+            window_step=obj.get("window_step", 1),
+            standardize_outputs=obj.get("standardize_outputs", True),
+        )
+
+    def prepare_datasets(self, seed: int = 42):
+        """
+        Convert raw data into preprocessed split datasets.
+        First split the data into `train`, `val` and `test` datasets
+        and then apply the `Sliding windows` transformation.
+        This is to avoid having some parts of a dataset in `train` and some in `val` and/or in `test` (i.e. data leak).
+
+        :param seed: Random seed used to split the datasets.
+
+        :return: Tuple (split_data, train_stats)
+
+          * ``split_data``: Dictionary with keys ``'train'``, ``'val'``, ``'test'``, and values the
+            corresponding tensorflow Datasets.
+
+          * ``train_stats``: Dictionary containing the shape of the data:
+            ``sequence_length``, ``num_load_features``, ``num_contact_params``, ``num_labels``,
+            and ``'mean'`` and ``'std'`` of the training set, in case ``standardize_outputs`` is True.
+        """
+        if self.pad_length > 0:
+            self.input_data = super().pad_initial(self.input_data, self.pad_length)
+            self.output_data = super().pad_initial(self.output_data, self.pad_length)
+
+        self.input_data = np.repeat(self.input_data[np.newaxis, :, :], self.output_data.shape[0], axis=0)
+
+        dataset = ({'inputs': self.input_data, 'params': self.param_data}, self.output_data)
+        split_data = super().make_splits(dataset, self.train_frac, self.val_frac, seed)
+
+        if self.standardize_outputs:
+            split_data, train_stats = super().standardize_outputs(split_data)
+        else:
+            train_stats = {}
+
+        split_data = {key: tf.data.Dataset.from_tensor_slices(val) for key, val in split_data.items()}
+        train_stats.update(self.get_dimensions(split_data['train']))
+        split_data = windows.windowize_train_val_test(split_data, self.window_size, self.window_step)
+
+        return split_data, train_stats
+
+    def prepare_single_dataset(self):
+        """
+        Convert raw data into a tensorflow dataset with compatible format to predict and evaluate a rnn model.
+        """
+        if self.pad_length > 0:
+            self.input_data = super().pad_initial(self.input_data, self.pad_length)
+            self.output_data = super().pad_initial(self.output_data, self.pad_length)
+
+        dataset = ({'inputs': self.input_data, 'params': self.param_data}, self.output_data)
+        return tf.data.Dataset.from_tensor_slices(dataset)
+
+    @classmethod
+    def get_default_config(cls):
+        """
+        :return: Dictionary containing default values of the arguments that the user can set.
+        """
+        return {
+            'input_data': None,
+            'param_data': None,
+            'output_data': None,
+            'train_frac': 0.7,
+            'val_frac': 0.3,
+            'pad_length': 0,
+            'window_size': 10,
+            'window_step': 1,
+            'standardize_outputs': True
+        }
+
 
 class PreprocessorTriaxialCompression(Preprocessor):
     """
@@ -220,6 +373,7 @@ def __init__(self, **kwargs):
 
         :param add_experiment_type: Wether to add the experiment type to contact parameters.
         """
+        super().__init__()
         config = self.check_config(kwargs)
 
         # creates an attribute per each element in config
@@ -254,7 +408,7 @@ def prepare_datasets(self, seed: int = 42):
             inputs = super().pad_initial(inputs, self.pad_length)
             outputs = super().pad_initial(outputs, self.pad_length)
 
-        dataset = ({'load_sequence': inputs, 'contact_parameters': contacts}, outputs)
+        dataset = ({'inputs': inputs, 'params': contacts}, outputs)
         split_data = super().make_splits(dataset, self.train_frac, self.val_frac, seed)
 
         if self.standardize_outputs:
@@ -263,7 +417,7 @@ def prepare_datasets(self, seed: int = 42):
             train_stats = {}
 
         split_data = {key: tf.data.Dataset.from_tensor_slices(val) for key, val in split_data.items()}
-        train_stats.update(super().get_dimensions(split_data['train']))
+        train_stats.update(self.get_dimensions(split_data['train']))
         split_data = windows.windowize_train_val_test(split_data, self.window_size, self.window_step)
 
         return split_data, train_stats
@@ -288,7 +442,7 @@ def prepare_single_dataset(self):
             inputs = super().pad_initial(inputs, self.pad_length)
             outputs = super().pad_initial(outputs, self.pad_length)
 
-        dataset = ({'load_sequence': inputs, 'contact_parameters': contacts}, outputs)
+        dataset = ({'inputs': inputs, 'params': contacts}, outputs)
         return tf.data.Dataset.from_tensor_slices(dataset)
 
     @classmethod
diff --git a/grainlearning/rnn/windows.py b/grainlearning/rnn/windows.py
index b21f2e03..1ace048a 100644
--- a/grainlearning/rnn/windows.py
+++ b/grainlearning/rnn/windows.py
@@ -27,7 +27,7 @@ def windowize_single_dataset(
         seed: int = 42,
         **_):
     """
-    Take a dataset ((inputs, contact_params),outputs) of sequences of shape N, S, L and output
+    Take a dataset ((inputs, params),outputs) of sequences of shape N, S, L and output
     another dataset of shorter sequences of size ``window_size``, taken at intervals ``window_step``.
     The resulting dataset will have ``M=((sequence_length - window_size)/window_step) + 1)`` samples,
     corresponding to independent windows in the given dataset.
@@ -48,18 +48,18 @@ def windowize_single_dataset(
       * ``inputs`` of shape: (M, window_size, ``num_load_features``), with M >> N.
       * ``outputs`` of shape: (M, L_outputs)
     """
-    load_sequences, contact_parameters, outputs = extract_tensors(data)
+    inputs, params, outputs = extract_tensors(data)
     _, sequence_length, _ = outputs.shape
 
     if window_size >= sequence_length:
         raise ValueError(f"window_size {window_size} >= sequence_length {sequence_length}.")
 
-    # For brevity denote load_sequence, contacts, outputs as x, c, y
+    # For brevity denote inputs, params, outputs as x, c, y
     xs, cs, ys = [], [], []
     for end in range(window_size, sequence_length + 1, window_step):
-        xs.append(load_sequences[:, end - window_size:end]) # input window
+        xs.append(inputs[:, end - window_size:end]) # input window
         ys.append(outputs[:, end - 1]) # final output
-        cs.append(contact_parameters)
+        cs.append(params)
 
     xs = np.array(xs)
     cs = np.array(cs)
@@ -75,7 +75,7 @@ def windowize_single_dataset(
     # finally shuffle the windows
     xs, cs, ys =  _shuffle(xs, cs, ys, seed)
     # and convert back into a tensorflow dataset
-    return tf.data.Dataset.from_tensor_slices(({'load_sequence': xs, 'contact_parameters': cs}, ys))
+    return tf.data.Dataset.from_tensor_slices(({'inputs': xs, 'params': cs}, ys))
 
 
 def _shuffle(xs, cs, ys, seed):
@@ -98,7 +98,7 @@ def predict_over_windows(
     Note the length of the output sequence will be shorter by the window_size than
     the input sequence.
 
-    :param inputs: dict containing inputs: `'load_sequence'` and `'contact_parameters'`, both being tensorflow.Tensor.
+    :param inputs: dict containing inputs: `'load_sequence'` and `'params'`, both being tensorflow.Tensor.
     :param model: The model to predict with.
     :param window_size: Number of timesteps in a single window.
     :param sequence_length: Number of timesteps in a full sequence.
@@ -106,7 +106,7 @@ def predict_over_windows(
     :return: tensorflow.Tensor of predicted sequences.
     """
     predictions = [
-        model([inputs['load_sequence'][:, end - window_size:end], inputs['contact_parameters']])
+        model([inputs['inputs'][:, end - window_size:end], inputs['params']])
         for end in range(window_size, sequence_length)
         ]
     predictions = tf.stack(predictions, axis=1)
@@ -123,8 +123,8 @@ def extract_tensors(data: tf.data.Dataset):
     """
     inputs, contacts, outputs = [], [], []
     for _inputs, _outputs in iter(data):
-        inputs.append(_inputs['load_sequence'])
-        contacts.append(_inputs['contact_parameters'])
+        inputs.append(_inputs['inputs'])
+        contacts.append(_inputs['params'])
         outputs.append(_outputs)
 
     return np.array(inputs), np.array(contacts), np.array(outputs)
diff --git a/tutorials/data_driven/LSTM/test_rnn.py b/tutorials/data_driven/LSTM/test_rnn.py
new file mode 100644
index 00000000..b0e48283
--- /dev/null
+++ b/tutorials/data_driven/LSTM/test_rnn.py
@@ -0,0 +1,112 @@
+import grainlearning.rnn.train as train_rnn
+import grainlearning.rnn.predict as predict_rnn
+from grainlearning.rnn import preprocessor
+import numpy as np
+from grainlearning import BayesianCalibration
+from matplotlib import pyplot as plt
+from grainlearning.rnn.evaluate_model import plot_predictions
+import tensorflow as tf
+
+x_obs = np.arange(120)
+y_obs = 0.2 * x_obs + 5.0
+
+
+def run_sim(calib):
+    """This is the callback function that runs different realizations of the same model.
+
+    :param calib: The calibration object.
+    """
+    data = []
+    for params in calib.system.param_data:
+        # Run the model
+        y_sim = linear(calib.system.ctrl_data, params)
+        data.append(np.array(y_sim, ndmin=2))
+    calib.system.set_sim_data(data)
+
+
+def linear(x, params):
+    return params[0] * x + params[1]
+
+
+calibration = BayesianCalibration.from_dict(
+    {
+        "num_iter": 1,
+        "callback": run_sim,
+        "system": {
+            "param_min": [0.001, 0.001],
+            "param_max": [1, 10],
+            "param_names": ['a', 'b'],
+            "num_samples": 100,
+            "obs_names": ['f'],
+            "ctrl_name": 'u',
+            "obs_data": y_obs,
+            "ctrl_data": x_obs,
+            "sim_name": 'linear',
+            "sigma_tol": 0.01,
+        },
+        "calibration": {
+            "inference": {"ess_target": 0.3},
+            "sampling": {
+                "max_num_components": 1,
+                "n_init": 1,
+                "random_state": 0,
+                "slice_sampling": True,
+            },
+            "initial_sampling": "halton",
+        },
+        "save_fig": -1,
+    }
+)
+
+calibration.run()
+
+# 1. Create my dictionary of configuration
+my_config = {
+    'input_data': calibration.system.ctrl_data,
+    'param_data': calibration.system.param_data,
+    'output_data': calibration.system.sim_data,
+    'train_frac': 0.7,
+    'val_frac': 0.2,
+    'window_size': 10,
+    'window_step': 1,
+    'patience': 25,
+    'epochs': 20,
+    'learning_rate': 1e-4,
+    'lstm_units': 128,
+    'dense_units': 128,
+    'batch_size': 64,
+    'standardize_outputs': True,
+    'save_weights_only': True
+}
+
+# 2. Create an object Preprocessor to pre-process my data
+preprocessor = preprocessor.PreprocessorLSTM.from_dict(my_config)
+
+# 3. Run the training using bare tensorflow
+history_simple = train_rnn.train_without_wandb(preprocessor, config=my_config)
+
+plt.plot(history_simple.history['loss'], label='training loss')
+plt.plot(history_simple.history['val_loss'], label='validation loss')
+plt.xlabel("epoch")
+plt.ylabel("MSE")
+plt.legend()
+plt.show()
+
+# 4. Make predictions with the trained model
+model, train_stats, config = predict_rnn.get_pretrained_model('outputs')
+
+data_inputs = ({'inputs': preprocessor.input_data, 'params': calibration.system.param_data}, preprocessor.input_data)
+data_inputs = tf.data.Dataset.from_tensor_slices(data_inputs)
+
+predictions = predict_rnn.predict_macroscopics(model, data_inputs, train_stats, config,
+                                               batch_size=calibration.system.num_samples)
+# converting the predictions to GL format (temporal dimension at the end)
+predictions = np.moveaxis(predictions, 1, -1)
+
+# compute the mean square error per sample
+error = np.mean((predictions - calibration.system.sim_data[:, :, config['window_size']:]) ** 2, axis=-1)
+
+# plot the error
+plt.plot(error)
+plt.xlabel("sample")
+plt.ylabel("MSE")

From e6644f47d563af55c169409be4147c8ea297e622 Mon Sep 17 00:00:00 2001
From: chyalexcheng <chyalexcheng@gmail.com>
Date: Sat, 3 Feb 2024 07:15:11 +0100
Subject: [PATCH 02/21] add a nonlinear regression tutorial

---
 tutorials/data_driven/LSTM/test_rnn.py        |  1 +
 .../python_hyperbola_regression_solve.py      | 88 +++++++++++++++++++
 2 files changed, 89 insertions(+)
 create mode 100644 tutorials/simple_regression/nonlinear_regression/python_hyperbola_regression_solve.py

diff --git a/tutorials/data_driven/LSTM/test_rnn.py b/tutorials/data_driven/LSTM/test_rnn.py
index b0e48283..c0391aca 100644
--- a/tutorials/data_driven/LSTM/test_rnn.py
+++ b/tutorials/data_driven/LSTM/test_rnn.py
@@ -110,3 +110,4 @@ def linear(x, params):
 plt.plot(error)
 plt.xlabel("sample")
 plt.ylabel("MSE")
+plt.show()
diff --git a/tutorials/simple_regression/nonlinear_regression/python_hyperbola_regression_solve.py b/tutorials/simple_regression/nonlinear_regression/python_hyperbola_regression_solve.py
new file mode 100644
index 00000000..0a910a67
--- /dev/null
+++ b/tutorials/simple_regression/nonlinear_regression/python_hyperbola_regression_solve.py
@@ -0,0 +1,88 @@
+"""
+This tutorial shows how to link a nonlinear regression model implemented in Python to GrainLearning.
+The only changes made to the linear regression case are the definition of
+the observed data `y_obs` and the model `nonlinear`.
+"""
+import numpy as np
+from grainlearning import BayesianCalibration
+
+x_obs = np.arange(100)
+# hyperbola in a form similar to the Duncan-Chang material model, q = \eps / (a * 100 + b * \eps)
+y_obs = x_obs / (0.2 * 100 + 5.0 * x_obs)
+
+
+def run_sim(calib):
+    """This is the callback function that runs different realizations of the same model.
+
+    :param calib: The calibration object.
+    """
+    data = []
+    for params in calib.system.param_data:
+        # Run the model
+        y_sim = nonlinear(calib.system.ctrl_data, params)
+        data.append(np.array(y_sim, ndmin=2))
+    calib.system.set_sim_data(data)
+
+
+def nonlinear(x, params):
+    """
+    A hyperbola in a form similar to the Duncan-Chang material model,
+
+    .. math::
+        sig = eps / (a * 100 + b * eps) where s is stress and e is strain.
+
+        y_t & = h(x_t) + r_t
+
+    where
+    :math:`sig` is stress,
+    :math:`eps` is strain in percent,
+    :math:`a` and `b` are the material parameters.
+    """
+    return x_obs / (params[0] * 100 + params[1] * x_obs)
+
+
+calibration = BayesianCalibration.from_dict(
+    {
+        "num_iter": 10,
+        "callback": run_sim,
+        "system": {
+            "param_min": [0.001, 0.001],
+            "param_max": [1, 10],
+            "param_names": ['a', 'b'],
+            "num_samples": 20,
+            "obs_names": ['f'],
+            "ctrl_name": 'u',
+            "obs_data": y_obs,
+            "ctrl_data": x_obs,
+            "sim_name": 'nonlinear',
+            "sigma_tol": 0.01,
+        },
+        "calibration": {
+            "inference": {
+                "ess_target": 0.3,
+                "scale_cov_with_max": True,
+            },
+            "sampling": {
+                "max_num_components": 1,
+                "n_init": 1,
+                "random_state": 0,
+                "slice_sampling": True,
+            },
+            "initial_sampling": "halton",
+        },
+        "save_fig": -1,
+    }
+)
+
+calibration.run()
+
+most_prob_params = calibration.get_most_prob_params()
+print(f'Most probable parameter values: {most_prob_params}')
+
+error_tolerance = 0.1
+
+error = most_prob_params - [0.2, 5.0]
+assert abs(
+    error[0]) / 0.2 < error_tolerance, f"Model parameters are not correct, expected 0.2 but got {most_prob_params[0]}"
+assert abs(
+    error[1]) / 5.0 < error_tolerance, f"Model parameters are not correct, expected 5.0 but got {most_prob_params[1]}"

From 8235eca9f0121558d9a8ac822324bc51a55a73f3 Mon Sep 17 00:00:00 2001
From: chyalexcheng <chyalexcheng@gmail.com>
Date: Sat, 3 Feb 2024 10:58:49 +0100
Subject: [PATCH 03/21] add a rnn example with nonlinear function

---
 tutorials/data_driven/LSTM/test_rnn.py | 16 ++++++++++------
 1 file changed, 10 insertions(+), 6 deletions(-)

diff --git a/tutorials/data_driven/LSTM/test_rnn.py b/tutorials/data_driven/LSTM/test_rnn.py
index c0391aca..6a45b51c 100644
--- a/tutorials/data_driven/LSTM/test_rnn.py
+++ b/tutorials/data_driven/LSTM/test_rnn.py
@@ -7,8 +7,9 @@
 from grainlearning.rnn.evaluate_model import plot_predictions
 import tensorflow as tf
 
-x_obs = np.arange(120)
-y_obs = 0.2 * x_obs + 5.0
+x_obs = np.arange(100)
+# hyperbola in a form similar to the Duncan-Chang material model, q = \eps / (a * 100 + b * \eps)
+y_obs = x_obs / (0.2 * 100 + 5.0 * x_obs)
 
 
 def run_sim(calib):
@@ -19,13 +20,13 @@ def run_sim(calib):
     data = []
     for params in calib.system.param_data:
         # Run the model
-        y_sim = linear(calib.system.ctrl_data, params)
+        y_sim = nonlinear(calib.system.ctrl_data, params)
         data.append(np.array(y_sim, ndmin=2))
     calib.system.set_sim_data(data)
 
 
-def linear(x, params):
-    return params[0] * x + params[1]
+def nonlinear(x, params):
+    return x / (params[0] * 100 + params[1] * x)
 
 
 calibration = BayesianCalibration.from_dict(
@@ -45,7 +46,10 @@ def linear(x, params):
             "sigma_tol": 0.01,
         },
         "calibration": {
-            "inference": {"ess_target": 0.3},
+            "inference": {
+                "ess_target": 0.3,
+                "scale_cov_with_max": True,
+            },
             "sampling": {
                 "max_num_components": 1,
                 "n_init": 1,

From 3da2496b327214162e2127e84f3be219077ac50f Mon Sep 17 00:00:00 2001
From: chyalexcheng <chyalexcheng@gmail.com>
Date: Wed, 7 Feb 2024 11:58:56 +0100
Subject: [PATCH 04/21] change some naming to make the parameters more general.
 Update the hyperbola calibration example to allow use of ML surrogate

---
 grainlearning/rnn/evaluate_model.py           |  3 +-
 grainlearning/rnn/models.py                   | 30 ++++++++++++++--
 grainlearning/rnn/predict.py                  |  2 +-
 grainlearning/rnn/preprocessor.py             |  4 +--
 .../{test_rnn.py => hyperbola_calibration.py} | 36 +++++++++++++++----
 5 files changed, 61 insertions(+), 14 deletions(-)
 rename tutorials/data_driven/LSTM/{test_rnn.py => hyperbola_calibration.py} (69%)

diff --git a/grainlearning/rnn/evaluate_model.py b/grainlearning/rnn/evaluate_model.py
index 01f444d5..aebba53c 100644
--- a/grainlearning/rnn/evaluate_model.py
+++ b/grainlearning/rnn/evaluate_model.py
@@ -29,8 +29,7 @@ def plot_predictions(model: tf.keras.Model, data: tf.data.Dataset, train_stats:
     plt.rcParams['axes.labelsize'] = 25
     plt.rcParams['font.family'] = 'sans-serif'
 
-    predictions = predict.predict_macroscopics(model, data, train_stats, config,
-                                       batch_size=batch_size)
+    predictions = predict.predict_batch(model, data, train_stats, config, batch_size=batch_size)
     # extract tensors from dataset
     test_inputs, labels = next(iter(data.batch(batch_size)))
 
diff --git a/grainlearning/rnn/models.py b/grainlearning/rnn/models.py
index 391d0916..174eadfd 100644
--- a/grainlearning/rnn/models.py
+++ b/grainlearning/rnn/models.py
@@ -17,11 +17,11 @@ def rnn_model(
     """
     Neural network with an LSTM layer.
 
-    Takes in a load sequence and contact parameters, and outputs the macroscopic responses.
-    The contact parameters are used to initialize the hidden state of the LSTM.
+    Takes in an input sequence and the parameters and produces an output sequence.
+    The parameters are used to initialize the hidden state of the LSTM.
 
     :param input_shapes: Dictionary containing `'num_input_features'`, `'num_params'`,
-        `'num_labels'`. It can contain other keys but hese are the ones used here.
+        `'num_labels'`. It can contain other keys but these are the ones used here.
     :param window_size: Length of time window.
     :param lstm_units: Number of units of the hidden state of the LSTM.
     :param dense_units: Number of units used in the dense layer after the LSTM.
@@ -52,3 +52,27 @@ def rnn_model(
     model = Model(inputs=[load_sequence, contact_params], outputs=outputs)
 
     return model
+
+
+def rnn_model_for_triax(
+        input_shapes: dict,
+        window_size: int = 20,
+        lstm_units: int = 50,
+        dense_units: int = 20,
+        seed: int = 42,
+        **_,
+        ):
+    """
+    A wrapper of neural network model with an LSTM layer for triaxial loading conditions.
+    :param input_shapes: Dictionary containing `'num_load_features'`, `'num_contact_params'`, and `'num_labels'`.
+    :param window_size: Length of time window.
+    :param lstm_units: Number of units of the hidden state of the LSTM.
+    :param dense_units: Number of units used in the dense layer after the LSTM.
+    :param seed: The random seed used to initialize the weights.
+
+    :return: A Keras model.
+    """
+    # change the name of the keys to match the original model
+    input_shapes['num_input_features'] = input_shapes.pop('num_load_features')
+    input_shapes['num_params'] = input_shapes.pop('num_contact_params')
+    return rnn_model(input_shapes, window_size, lstm_units, dense_units, seed)
diff --git a/grainlearning/rnn/predict.py b/grainlearning/rnn/predict.py
index f42ee89e..9853f23d 100644
--- a/grainlearning/rnn/predict.py
+++ b/grainlearning/rnn/predict.py
@@ -154,7 +154,7 @@ def load_model(path_to_model: Path, train_stats: dict, config: dict):
     return model
 
 
-def predict_macroscopics(
+def predict_batch(
         model: tf.keras.Model,
         data: tf.data.Dataset,
         train_stats: dict,
diff --git a/grainlearning/rnn/preprocessor.py b/grainlearning/rnn/preprocessor.py
index 56a65a9b..e71ccafb 100644
--- a/grainlearning/rnn/preprocessor.py
+++ b/grainlearning/rnn/preprocessor.py
@@ -293,12 +293,12 @@ def prepare_datasets(self, seed: int = 42):
             ``sequence_length``, ``num_load_features``, ``num_contact_params``, ``num_labels``,
             and ``'mean'`` and ``'std'`` of the training set, in case ``standardize_outputs`` is True.
         """
+        self.input_data = np.repeat(self.input_data[np.newaxis, :, :], self.output_data.shape[0], axis=0)
+
         if self.pad_length > 0:
             self.input_data = super().pad_initial(self.input_data, self.pad_length)
             self.output_data = super().pad_initial(self.output_data, self.pad_length)
 
-        self.input_data = np.repeat(self.input_data[np.newaxis, :, :], self.output_data.shape[0], axis=0)
-
         dataset = ({'inputs': self.input_data, 'params': self.param_data}, self.output_data)
         split_data = super().make_splits(dataset, self.train_frac, self.val_frac, seed)
 
diff --git a/tutorials/data_driven/LSTM/test_rnn.py b/tutorials/data_driven/LSTM/hyperbola_calibration.py
similarity index 69%
rename from tutorials/data_driven/LSTM/test_rnn.py
rename to tutorials/data_driven/LSTM/hyperbola_calibration.py
index 6a45b51c..3e7140b9 100644
--- a/tutorials/data_driven/LSTM/test_rnn.py
+++ b/tutorials/data_driven/LSTM/hyperbola_calibration.py
@@ -4,7 +4,6 @@
 import numpy as np
 from grainlearning import BayesianCalibration
 from matplotlib import pyplot as plt
-from grainlearning.rnn.evaluate_model import plot_predictions
 import tensorflow as tf
 
 x_obs = np.arange(100)
@@ -34,7 +33,7 @@ def nonlinear(x, params):
         "num_iter": 1,
         "callback": run_sim,
         "system": {
-            "param_min": [0.001, 0.001],
+            "param_min": [0.1, 1],
             "param_max": [1, 10],
             "param_names": ['a', 'b'],
             "num_samples": 100,
@@ -42,7 +41,7 @@ def nonlinear(x, params):
             "ctrl_name": 'u',
             "obs_data": y_obs,
             "ctrl_data": x_obs,
-            "sim_name": 'linear',
+            "sim_name": 'hyperbola',
             "sigma_tol": 0.01,
         },
         "calibration": {
@@ -71,6 +70,8 @@ def nonlinear(x, params):
     'output_data': calibration.system.sim_data,
     'train_frac': 0.7,
     'val_frac': 0.2,
+    # TODO: hide `pad_length` from the user
+    'pad_length': 10,
     'window_size': 10,
     'window_step': 1,
     'patience': 25,
@@ -102,16 +103,39 @@ def nonlinear(x, params):
 data_inputs = ({'inputs': preprocessor.input_data, 'params': calibration.system.param_data}, preprocessor.input_data)
 data_inputs = tf.data.Dataset.from_tensor_slices(data_inputs)
 
-predictions = predict_rnn.predict_macroscopics(model, data_inputs, train_stats, config,
-                                               batch_size=calibration.system.num_samples)
+predictions = predict_rnn.predict_batch(model, data_inputs, train_stats, config,
+                                        batch_size=calibration.system.num_samples)
 # converting the predictions to GL format (temporal dimension at the end)
 predictions = np.moveaxis(predictions, 1, -1)
 
 # compute the mean square error per sample
-error = np.mean((predictions - calibration.system.sim_data[:, :, config['window_size']:]) ** 2, axis=-1)
+error = np.mean((predictions - calibration.system.sim_data) ** 2, axis=-1)
 
 # plot the error
 plt.plot(error)
 plt.xlabel("sample")
 plt.ylabel("MSE")
 plt.show()
+
+
+# define the callback function using the ML surrogate
+def run_sim_surrogate(calib):
+    # extend the first dimension of the input data to the number of samples (TODO: hide this from the user)
+    input_data = preprocessor.input_data[0, :, :]
+    preprocessor.input_data = np.repeat(input_data[np.newaxis, :, :], calib.system.num_samples, axis=0)
+    data_inputs = ({'inputs': preprocessor.input_data, 'params': calib.system.param_data}, preprocessor.input_data)
+    data_inputs = tf.data.Dataset.from_tensor_slices(data_inputs)
+    # make predictions with the trained model
+    sim_data = predict_rnn.predict_batch(model, data_inputs, train_stats, config, batch_size=calib.system.num_samples)
+    # converting the predictions to GL format (temporal dimension at the end)
+    sim_data = np.moveaxis(sim_data, 1, -1)
+    # update sim_data in system
+    calib.system.set_sim_data(sim_data)
+
+
+calibration.callback = run_sim_surrogate
+
+# continue the calibration with the surrogate
+calibration.num_iter = 10
+calibration.save_fig = 0
+calibration.run()

From 1bb2cd96f6da57273e1e61b258d95d1106f03006 Mon Sep 17 00:00:00 2001
From: chyalexcheng <chyalexcheng@gmail.com>
Date: Wed, 7 Feb 2024 12:01:30 +0100
Subject: [PATCH 05/21] fix errors of rnn selftests

---
 .../rnn/plain_entire_model/train_stats.npy    | Bin 378 -> 371 bytes
 .../rnn/plain_only_weights/train_stats.npy    | Bin 378 -> 371 bytes
 .../rnn/wandb_entire_model/train_stats.npy    | Bin 378 -> 371 bytes
 .../rnn/wandb_only_weights/train_stats.npy    | Bin 378 -> 371 bytes
 tests/unit/conftest.py                        |  14 ++---
 tests/unit/test_rnn_model.py                  |  22 ++++----
 tests/unit/test_rnn_preprocessor.py           |  50 +++++++++---------
 7 files changed, 43 insertions(+), 43 deletions(-)

diff --git a/tests/data/rnn/plain_entire_model/train_stats.npy b/tests/data/rnn/plain_entire_model/train_stats.npy
index 96edf979c7b33e3ea608ba007c15b5ed532d0735..1dc822fa63cacc5c13ae37e68188de978c5ebdea 100644
GIT binary patch
delta 79
zcmeyx^qFac4x>940|P@|X>NQ@Vp3{OaUs7qON0<qC^N61v?M+)HL;|$D7Cmyz?&%o
Xrn(@pC^5IVP|%yTG_j;mC`k_h+eR3w

delta 86
zcmey&^owbO4x^tS0|P@|X>NQ@equ^|T54iRX;Er%A-^|Mgb-9FIX|x?F}WnZAh9Sh
ex42Nin>B(9rYSKgHK({x(3_<+v7}HaNe=+NT^n-%

diff --git a/tests/data/rnn/plain_only_weights/train_stats.npy b/tests/data/rnn/plain_only_weights/train_stats.npy
index 96edf979c7b33e3ea608ba007c15b5ed532d0735..1dc822fa63cacc5c13ae37e68188de978c5ebdea 100644
GIT binary patch
delta 79
zcmeyx^qFac4x>940|P@|X>NQ@Vp3{OaUs7qON0<qC^N61v?M+)HL;|$D7Cmyz?&%o
Xrn(@pC^5IVP|%yTG_j;mC`k_h+eR3w

delta 86
zcmey&^owbO4x^tS0|P@|X>NQ@equ^|T54iRX;Er%A-^|Mgb-9FIX|x?F}WnZAh9Sh
ex42Nin>B(9rYSKgHK({x(3_<+v7}HaNe=+NT^n-%

diff --git a/tests/data/rnn/wandb_entire_model/train_stats.npy b/tests/data/rnn/wandb_entire_model/train_stats.npy
index 96edf979c7b33e3ea608ba007c15b5ed532d0735..1dc822fa63cacc5c13ae37e68188de978c5ebdea 100644
GIT binary patch
delta 79
zcmeyx^qFac4x>940|P@|X>NQ@Vp3{OaUs7qON0<qC^N61v?M+)HL;|$D7Cmyz?&%o
Xrn(@pC^5IVP|%yTG_j;mC`k_h+eR3w

delta 86
zcmey&^owbO4x^tS0|P@|X>NQ@equ^|T54iRX;Er%A-^|Mgb-9FIX|x?F}WnZAh9Sh
ex42Nin>B(9rYSKgHK({x(3_<+v7}HaNe=+NT^n-%

diff --git a/tests/data/rnn/wandb_only_weights/train_stats.npy b/tests/data/rnn/wandb_only_weights/train_stats.npy
index 96edf979c7b33e3ea608ba007c15b5ed532d0735..1dc822fa63cacc5c13ae37e68188de978c5ebdea 100644
GIT binary patch
delta 79
zcmeyx^qFac4x>940|P@|X>NQ@Vp3{OaUs7qON0<qC^N61v?M+)HL;|$D7Cmyz?&%o
Xrn(@pC^5IVP|%yTG_j;mC`k_h+eR3w

delta 86
zcmey&^owbO4x^tS0|P@|X>NQ@equ^|T54iRX;Er%A-^|Mgb-9FIX|x?F}WnZAh9Sh
ex42Nin>B(9rYSKgHK({x(3_<+v7}HaNe=+NT^n-%

diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py
index 71111308..6b5e623f 100644
--- a/tests/unit/conftest.py
+++ b/tests/unit/conftest.py
@@ -20,21 +20,21 @@ def hdf5_test_file():
 
         # group 0.2e5
         inputs, outputs, contact_params = create_dataset(num_samples=2, sequence_length=3,
-                                            num_load_features=2, num_labels=4, num_contact_params=5)
+                                            num_input_features=2, num_labels=4, num_params=5)
         file['0.2e5/drained/inputs'] = inputs
         file['0.2e5/drained/outputs'] = outputs
         file['0.2e5/drained/contact_params'] = contact_params
 
         # group 1000000
         inputs, outputs, contact_params = create_dataset(num_samples=10, sequence_length=3,
-                                            num_load_features=2, num_labels=4, num_contact_params=5)
+                                            num_input_features=2, num_labels=4, num_params=5)
         file['1000000/undrained/inputs'] = inputs
         file['1000000/undrained/outputs'] = outputs
         file['1000000/undrained/contact_params'] = contact_params
 
         # group 2000000.6
         inputs, outputs, contact_params = create_dataset(num_samples=4, sequence_length=3,
-                                            num_load_features=2, num_labels=4, num_contact_params=5)
+                                            num_input_features=2, num_labels=4, num_params=5)
         file['2000000.6/undrained/inputs'] = inputs
         file['2000000.6/undrained/outputs'] = outputs
         file['2000000.6/undrained/contact_params'] = contact_params
@@ -42,15 +42,15 @@ def hdf5_test_file():
     return target
 
 
-def create_dataset(num_samples: int, sequence_length: int, num_load_features: int,
-                   num_labels: int, num_contact_params: int):
+def create_dataset(num_samples: int, sequence_length: int, num_input_features: int,
+                   num_labels: int, num_params: int):
     """
     Creates a dataset (numpy) with random numbers between 0-1, given the dimensions.
     We set the seed so that the creation of the dataset is deterministic.
     """
     np.random.seed(42)
-    inputs = np.random.rand(num_samples, sequence_length, num_load_features)
+    inputs = np.random.rand(num_samples, sequence_length, num_input_features)
     outputs = np.random.rand(num_samples, sequence_length, num_labels)
-    contact_params = np.random.rand(num_samples, num_contact_params)
+    contact_params = np.random.rand(num_samples, num_params)
 
     return inputs, outputs, contact_params
diff --git a/tests/unit/test_rnn_model.py b/tests/unit/test_rnn_model.py
index e0395d9c..66176ef6 100644
--- a/tests/unit/test_rnn_model.py
+++ b/tests/unit/test_rnn_model.py
@@ -43,8 +43,8 @@ def test_model_output_shape():
     """ Test if rnn model can be initialized and outputs the expected shape. """
     # Normally gotten from train_stats after data loading
     input_shapes = {
-            'num_contact_params': 6,
-            'num_load_features': 4,
+            'num_params': 6,
+            'num_input_features': 4,
             'num_labels': 5,
             'sequence_length': 200,
         }
@@ -53,11 +53,11 @@ def test_model_output_shape():
     model = rnn_model(input_shapes, window_size)
     assert len(model.layers) == 7 # 2 inputs, 2 hidden states lstm, lstm, dense, dense_output
 
-    test_input_sequence = np.random.normal(size=(batch_size, window_size, input_shapes['num_load_features']))
-    test_contacts = np.random.normal(size=(batch_size, input_shapes['num_contact_params']))
+    test_input_sequence = np.random.normal(size=(batch_size, window_size, input_shapes['num_input_features']))
+    test_contacts = np.random.normal(size=(batch_size, input_shapes['num_params']))
 
-    output = model({'load_sequence': test_input_sequence,
-                    'contact_parameters': test_contacts})
+    output = model({'inputs': test_input_sequence,
+                    'params': test_contacts})
 
     assert output.shape == (batch_size, input_shapes['num_labels'])
 
@@ -136,8 +136,8 @@ def test_get_pretrained_model(config_test):
         model.summary() # Will throw an exception if the model was not loaded correctly
 
         # test that train_stats has expected members and values
-        assert train_stats.keys() == {'sequence_length', 'num_load_features',
-                                      'num_contact_params', 'num_labels'}
+        assert train_stats.keys() == {'sequence_length', 'num_input_features',
+                                      'num_params', 'num_labels'}
 
         # test params in config matching original config
         if "only_weights" in path_to_model:
@@ -152,13 +152,13 @@ def test_predict_macroscopics():
     model, train_stats, config = predict.get_pretrained_model("./tests/data/rnn/wandb_only_weights/")
     preprocessor_TC = preprocessor.PreprocessorTriaxialCompression(**config)
     data, _ = preprocessor_TC.prepare_datasets()
-    predictions_1 = predict.predict_macroscopics(model, data['test'], train_stats, config, batch_size=1)
+    predictions_1 = predict.predict_batch(model, data['test'], train_stats, config, batch_size=1)
     config['pad_length'] = config['window_size']
     config['train_frac'] = 0.25
     config['val_frac'] = 0.25
     preprocessor_TC_2 = preprocessor.PreprocessorTriaxialCompression(**config)
     data_padded, train_stats_2 = preprocessor_TC_2.prepare_datasets()
-    predictions_2 = predict.predict_macroscopics(model, data_padded['test'], train_stats_2, config, batch_size=2)
+    predictions_2 = predict.predict_batch(model, data_padded['test'], train_stats_2, config, batch_size=2)
 
     assert isinstance(predictions_1, tf.Tensor)
     assert isinstance(predictions_2, tf.Tensor)
@@ -173,7 +173,7 @@ def test_predict_macroscopics():
     # model loaded: pad_length=0, config, pad_length=1. If using train_stats of the model -> incompatible.
     data_padded = preprocessor_TC_2.prepare_single_dataset()
     with pytest.raises(ValueError):
-        predict.predict_macroscopics(model, data_padded, train_stats, config, batch_size=2)
+        predict.predict_batch(model, data_padded, train_stats, config, batch_size=2)
 
     # check that standardize outputs has been correctly applied: cannot comprare.
 
diff --git a/tests/unit/test_rnn_preprocessor.py b/tests/unit/test_rnn_preprocessor.py
index af4e514e..aa3c7c45 100644
--- a/tests/unit/test_rnn_preprocessor.py
+++ b/tests/unit/test_rnn_preprocessor.py
@@ -9,8 +9,8 @@
 @pytest.fixture(scope="session")
 def dummy_dataset():
     inputs, outputs, contact_params = create_dataset(num_samples = 100, sequence_length = 10,
-                              num_load_features = 4, num_labels = 3, num_contact_params = 5)
-    return ({'load_sequence': inputs, 'contact_parameters': contact_params}, outputs)
+                              num_input_features = 4, num_labels = 3, num_params = 5)
+    return ({'inputs': inputs, 'params': contact_params}, outputs)
 
 @pytest.fixture(scope="function")
 def default_config(hdf5_test_file):
@@ -37,9 +37,9 @@ def dummy_preprocessor(default_config):
 def test_make_splits(dummy_dataset, dummy_preprocessor):
     num_samples = 100
     sequence_length = 10
-    num_load_features = 4
+    num_input_features = 4
     num_labels = 3
-    num_contact_params = 5
+    num_params = 5
 
     train_frac_list = [0.5, 0.7, 0.1]
     val_frac_list = [0.25, 0.15, 0.5]
@@ -61,8 +61,8 @@ def test_make_splits(dummy_dataset, dummy_preprocessor):
 
         # Check that the other dimensions of the arrays are still correct
         for key, n_samples in zip(split_data.keys(), [n_train, n_val, n_test]):
-            assert split_data[key][0]['load_sequence'].shape == (n_samples, sequence_length, num_load_features)
-            assert split_data[key][0]['contact_parameters'].shape == (n_samples, num_contact_params)
+            assert split_data[key][0]['inputs'].shape == (n_samples, sequence_length, num_input_features)
+            assert split_data[key][0]['params'].shape == (n_samples, num_params)
             assert split_data[key][1].shape == (n_samples, sequence_length, num_labels)
 
     with pytest.raises(ValueError): # check that it is not possible to have a 0 length validation dataset
@@ -75,12 +75,12 @@ def test_get_dimensions(dummy_dataset, dummy_preprocessor):
     for data_i in split_data.values():
         dimensions = dummy_preprocessor.get_dimensions(data_i)
 
-        assert dimensions.keys() == {'sequence_length', 'num_load_features', 'num_contact_params', 'num_labels'}
+        assert dimensions.keys() == {'sequence_length', 'num_input_features', 'num_params', 'num_labels'}
 
         # from dimensions of dummy_dataset
         assert dimensions['sequence_length'] == 10
-        assert dimensions['num_load_features'] == 4
-        assert dimensions['num_contact_params'] == 5
+        assert dimensions['num_input_features'] == 4
+        assert dimensions['num_params'] == 5
         assert dimensions['num_labels'] == 3
 
 def test_merge_datasets(hdf5_test_file, default_config):
@@ -89,22 +89,22 @@ def test_merge_datasets(hdf5_test_file, default_config):
         config_TC = default_config.copy()
         preprocessor_TC_1 = preprocessor.PreprocessorTriaxialCompression(**config_TC)
 
-        inputs, outputs, contact_parameters = preprocessor_TC_1._merge_datasets(datafile)
+        inputs, outputs, params = preprocessor_TC_1._merge_datasets(datafile)
         # check total (sum) num of samples, and that pressure and experiment type have been added to contact_params
         # Values to be modified if hdf5_test_file changes
         assert inputs.shape == (2 + 10 + 4, 3, 2) # num_samples, sequence_length, load_features
         assert outputs.shape == (2 + 10 + 4, 3, 4) # num_samples, sequence_length, num_labels
-        assert contact_parameters.shape == (2 + 10 + 4, 5 + 2) # num_samples, num_contact_params + 2 (pressure, experiment_type)
+        assert params.shape == (2 + 10 + 4, 5 + 2) # num_samples, num_params + 2 (pressure, experiment_type)
 
         # Case specific pressure and experiment
         config_TC["pressure"] = '0.2e5'
         config_TC["experiment_type"] = 'drained'
         preprocessor_TC_2 = preprocessor.PreprocessorTriaxialCompression(**config_TC)
 
-        inputs, outputs, contact_parameters = preprocessor_TC_2._merge_datasets(datafile)
+        inputs, outputs, params = preprocessor_TC_2._merge_datasets(datafile)
         assert inputs.shape == (2, 3, 2) # num_samples, sequence_length, load_features
         assert outputs.shape == (2, 3, 4) # num_samples, sequence_length, num_labels
-        assert contact_parameters.shape == (2, 5 + 2) # num_samples, num_contact_params + 2 (pressure, experiment_type)
+        assert params.shape == (2, 5 + 2) # num_samples, num_params + 2 (pressure, experiment_type)
 
 def test_prepare_datasets(default_config):
     config_TC = default_config.copy()
@@ -124,11 +124,11 @@ def test_prepare_datasets(default_config):
 
     # Test train_stats
     assert isinstance(train_stats, dict)
-    assert train_stats.keys() == {'mean', 'std', 'sequence_length', 'num_load_features',
-                                  'num_contact_params', 'num_labels'}
+    assert train_stats.keys() == {'mean', 'std', 'sequence_length', 'num_input_features',
+                                  'num_params', 'num_labels'}
     assert train_stats['sequence_length'] > 0
-    assert train_stats['num_load_features'] > 0
-    assert train_stats['num_contact_params'] > 0
+    assert train_stats['num_input_features'] > 0
+    assert train_stats['num_params'] > 0
     assert train_stats['num_labels'] > 0
 
     # Test length of contact parameters when adding e0, pressure, experiment type
@@ -157,9 +157,9 @@ def test_prepare_datasets(default_config):
     split_data_3, train_stats_3 = preprocessor_TC_3.prepare_datasets()
 
     # Comparison against train_stats: No additional contact parameters.
-    assert train_stats_1['num_contact_params'] == train_stats['num_contact_params'] + 1
-    assert train_stats_2['num_contact_params'] == train_stats['num_contact_params'] + 1
-    assert train_stats_3['num_contact_params'] == train_stats['num_contact_params'] + 1
+    assert train_stats_1['num_params'] == train_stats['num_params'] + 1
+    assert train_stats_2['num_params'] == train_stats['num_params'] + 1
+    assert train_stats_3['num_params'] == train_stats['num_params'] + 1
 
     # Test that standardize_outputs is applied correctly when standardize_outputs=False
     assert 'mean' not in train_stats_1
@@ -181,9 +181,9 @@ def test_prepare_datasets(default_config):
 def test_windowize_single_dataset(dummy_dataset, dummy_preprocessor):
     # dimensions with which dummy_dataset was generated.
     sequence_length = 10
-    num_load_features = 4
+    num_input_features = 4
     num_labels = 3
-    num_contact_params = 5
+    num_params = 5
 
     split_data = dummy_preprocessor.make_splits(dummy_dataset, train_frac=0.7, val_frac=0.15, seed=42)
     split_data = {key: tf.data.Dataset.from_tensor_slices(val) for key, val in split_data.items()}
@@ -202,9 +202,9 @@ def test_windowize_single_dataset(dummy_dataset, dummy_preprocessor):
         expected_num_indep_samples = 70 * int(((sequence_length - window_size)/window_step) + 1)
 
         assert len(dataset) == expected_num_indep_samples
-        assert len(inputs) == 2 # load_sequence and contact_params
-        assert inputs['load_sequence'].shape == (expected_num_indep_samples, window_size, num_load_features) # num_samples, window_size, num_load_features
-        assert inputs['contact_parameters'].shape == (expected_num_indep_samples, num_contact_params) # num_samples, window_size, num_contact_params
+        assert len(inputs) == 2 # inputs and contact_params
+        assert inputs['inputs'].shape == (expected_num_indep_samples, window_size, num_input_features) # num_samples, window_size, num_input_features
+        assert inputs['params'].shape == (expected_num_indep_samples, num_params) # num_samples, window_size, num_params
         assert outputs.shape == (expected_num_indep_samples, num_labels) # num_samples, num_labels
 
     with pytest.raises(ValueError): # expect "window_size bigger than sequence_length."

From 54ee30b188a85e32db5c1269d26110fc85bbe738 Mon Sep 17 00:00:00 2001
From: chyalexcheng <chyalexcheng@gmail.com>
Date: Wed, 7 Feb 2024 16:29:08 +0100
Subject: [PATCH 06/21] implement a function to create and update input data
 for rnn.predict_batch

---
 grainlearning/rnn/preprocessor.py             | 20 ++++++++++++++++---
 .../data_driven/LSTM/hyperbola_calibration.py | 13 +++---------
 2 files changed, 20 insertions(+), 13 deletions(-)

diff --git a/grainlearning/rnn/preprocessor.py b/grainlearning/rnn/preprocessor.py
index e71ccafb..6a567751 100644
--- a/grainlearning/rnn/preprocessor.py
+++ b/grainlearning/rnn/preprocessor.py
@@ -212,7 +212,6 @@ def __init__(
         output_data: np.ndarray,
         train_frac: float,
         val_frac: float,
-        pad_length: int,
         window_size: int,
         window_step: int = 1,
         standardize_outputs: bool = True,
@@ -241,7 +240,7 @@ def __init__(
         self.output_data = output_data
         self.train_frac = train_frac
         self.val_frac = val_frac
-        self.pad_length = pad_length
+        self.pad_length = window_size
         self.window_size = window_size
         self.window_step = window_step
         self.standardize_outputs = standardize_outputs
@@ -270,7 +269,6 @@ def from_dict(cls: Type["PreprocessorLSTM"], obj: dict):
             train_frac=obj["train_frac"],
             val_frac=obj["val_frac"],
             window_size=obj["window_size"],
-            pad_length=obj.get("pad_length", 0),
             window_step=obj.get("window_step", 1),
             standardize_outputs=obj.get("standardize_outputs", True),
         )
@@ -324,6 +322,22 @@ def prepare_single_dataset(self):
         dataset = ({'inputs': self.input_data, 'params': self.param_data}, self.output_data)
         return tf.data.Dataset.from_tensor_slices(dataset)
 
+    def prepare_input_data(self, param_data: np.ndarray):
+        """
+        Prepare the input data for the model.
+
+        :param param_data: ndarray of parameters
+
+        :return: Tuple of input data
+        """
+        # if the number of samples is different, update self.input_data
+        if param_data.shape[0] != self.input_data.shape[0]:
+            input_data = self.input_data[0, :, :]
+            self.input_data = np.repeat(input_data[np.newaxis, :, :], param_data.shape[0], axis=0)
+        data_inputs = ({'inputs': self.input_data, 'params': param_data}, self.input_data)
+        data_inputs = tf.data.Dataset.from_tensor_slices(data_inputs)
+        return data_inputs
+
     @classmethod
     def get_default_config(cls):
         """
diff --git a/tutorials/data_driven/LSTM/hyperbola_calibration.py b/tutorials/data_driven/LSTM/hyperbola_calibration.py
index 3e7140b9..c968c63f 100644
--- a/tutorials/data_driven/LSTM/hyperbola_calibration.py
+++ b/tutorials/data_driven/LSTM/hyperbola_calibration.py
@@ -4,7 +4,6 @@
 import numpy as np
 from grainlearning import BayesianCalibration
 from matplotlib import pyplot as plt
-import tensorflow as tf
 
 x_obs = np.arange(100)
 # hyperbola in a form similar to the Duncan-Chang material model, q = \eps / (a * 100 + b * \eps)
@@ -70,8 +69,6 @@ def nonlinear(x, params):
     'output_data': calibration.system.sim_data,
     'train_frac': 0.7,
     'val_frac': 0.2,
-    # TODO: hide `pad_length` from the user
-    'pad_length': 10,
     'window_size': 10,
     'window_step': 1,
     'patience': 25,
@@ -100,8 +97,7 @@ def nonlinear(x, params):
 # 4. Make predictions with the trained model
 model, train_stats, config = predict_rnn.get_pretrained_model('outputs')
 
-data_inputs = ({'inputs': preprocessor.input_data, 'params': calibration.system.param_data}, preprocessor.input_data)
-data_inputs = tf.data.Dataset.from_tensor_slices(data_inputs)
+data_inputs = preprocessor.prepare_input_data(calibration.system.param_data)
 
 predictions = predict_rnn.predict_batch(model, data_inputs, train_stats, config,
                                         batch_size=calibration.system.num_samples)
@@ -120,11 +116,7 @@ def nonlinear(x, params):
 
 # define the callback function using the ML surrogate
 def run_sim_surrogate(calib):
-    # extend the first dimension of the input data to the number of samples (TODO: hide this from the user)
-    input_data = preprocessor.input_data[0, :, :]
-    preprocessor.input_data = np.repeat(input_data[np.newaxis, :, :], calib.system.num_samples, axis=0)
-    data_inputs = ({'inputs': preprocessor.input_data, 'params': calib.system.param_data}, preprocessor.input_data)
-    data_inputs = tf.data.Dataset.from_tensor_slices(data_inputs)
+    data_inputs = preprocessor.prepare_input_data(calib.system.param_data)
     # make predictions with the trained model
     sim_data = predict_rnn.predict_batch(model, data_inputs, train_stats, config, batch_size=calib.system.num_samples)
     # converting the predictions to GL format (temporal dimension at the end)
@@ -133,6 +125,7 @@ def run_sim_surrogate(calib):
     calib.system.set_sim_data(sim_data)
 
 
+# set the callback function to the one that runs the ML surrogate
 calibration.callback = run_sim_surrogate
 
 # continue the calibration with the surrogate

From 076229e4fa646e026b48c1b3b2b072ef2cd2c08b Mon Sep 17 00:00:00 2001
From: chyalexcheng <chyalexcheng@gmail.com>
Date: Wed, 7 Feb 2024 21:30:28 +0100
Subject: [PATCH 07/21] first working example of mixed use of data-driven and
 equation-based models. 1. Generate sim_data for one iteration; 2. train ML;
 3. Split param samples into two subsets: one for ML surrogate, one for
 equation-based; 4. Continue for all remaining iterations

---
 grainlearning/rnn/preprocessor.py             |   1 +
 ...ation.py => hyperbola_calibration_lstm.py} |   0
 .../LSTM/hyperbola_calibration_mixed.py       | 145 ++++++++++++++++++
 3 files changed, 146 insertions(+)
 rename tutorials/data_driven/LSTM/{hyperbola_calibration.py => hyperbola_calibration_lstm.py} (100%)
 create mode 100644 tutorials/data_driven/LSTM/hyperbola_calibration_mixed.py

diff --git a/grainlearning/rnn/preprocessor.py b/grainlearning/rnn/preprocessor.py
index 6a567751..67df6d51 100644
--- a/grainlearning/rnn/preprocessor.py
+++ b/grainlearning/rnn/preprocessor.py
@@ -331,6 +331,7 @@ def prepare_input_data(self, param_data: np.ndarray):
         :return: Tuple of input data
         """
         # if the number of samples is different, update self.input_data
+        self.param_data = param_data
         if param_data.shape[0] != self.input_data.shape[0]:
             input_data = self.input_data[0, :, :]
             self.input_data = np.repeat(input_data[np.newaxis, :, :], param_data.shape[0], axis=0)
diff --git a/tutorials/data_driven/LSTM/hyperbola_calibration.py b/tutorials/data_driven/LSTM/hyperbola_calibration_lstm.py
similarity index 100%
rename from tutorials/data_driven/LSTM/hyperbola_calibration.py
rename to tutorials/data_driven/LSTM/hyperbola_calibration_lstm.py
diff --git a/tutorials/data_driven/LSTM/hyperbola_calibration_mixed.py b/tutorials/data_driven/LSTM/hyperbola_calibration_mixed.py
new file mode 100644
index 00000000..d343aee1
--- /dev/null
+++ b/tutorials/data_driven/LSTM/hyperbola_calibration_mixed.py
@@ -0,0 +1,145 @@
+import grainlearning.rnn.train as train_rnn
+import grainlearning.rnn.predict as predict_rnn
+from grainlearning.rnn import preprocessor
+import numpy as np
+from grainlearning import BayesianCalibration
+from matplotlib import pyplot as plt
+
+x_obs = np.arange(100)
+# hyperbola in a form similar to the Duncan-Chang material model, q = \eps / (a * 100 + b * \eps)
+y_obs = x_obs / (0.2 * 100 + 5.0 * x_obs)
+
+
+def run_sim(calib):
+    """This is the callback function that runs different realizations of the same model.
+
+    :param calib: The calibration object.
+    """
+    data = []
+    for params in calib.system.param_data:
+        # Run the model
+        y_sim = nonlinear(calib.system.ctrl_data, params)
+        data.append(np.array(y_sim, ndmin=2))
+    calib.system.set_sim_data(data)
+
+
+def nonlinear(x, params):
+    return x / (params[0] * 100 + params[1] * x)
+
+
+calibration = BayesianCalibration.from_dict(
+    {
+        "num_iter": 1,
+        "callback": run_sim,
+        "system": {
+            "param_min": [0.1, 1],
+            "param_max": [1, 10],
+            "param_names": ['a', 'b'],
+            "num_samples": 20,
+            "obs_names": ['f'],
+            "ctrl_name": 'u',
+            "obs_data": y_obs,
+            "ctrl_data": x_obs,
+            "sim_name": 'hyperbola',
+            "sigma_tol": 0.01,
+        },
+        "calibration": {
+            "inference": {
+                "ess_target": 0.3,
+                "scale_cov_with_max": True,
+            },
+            "sampling": {
+                "max_num_components": 1,
+                "n_init": 1,
+                "random_state": 0,
+                "slice_sampling": True,
+            },
+            "initial_sampling": "halton",
+        },
+        "save_fig": -1,
+    }
+)
+
+# 1. Run one iteration of the calibration to collect the simulation data for training the ML surrogate
+calibration.run()
+
+# 2.1. Create my dictionary of configuration
+my_config = {
+    'input_data': calibration.system.ctrl_data,
+    'param_data': calibration.system.param_data,
+    'output_data': calibration.system.sim_data,
+    'train_frac': 0.7,
+    'val_frac': 0.2,
+    'window_size': 10,
+    'window_step': 1,
+    'patience': 25,
+    'epochs': 20,
+    'learning_rate': 1e-4,
+    'lstm_units': 128,
+    'dense_units': 128,
+    'batch_size': 64,
+    'standardize_outputs': True,
+    'save_weights_only': True
+}
+
+# 2.2. Create an object Preprocessor to pre-process my data
+preprocessorLSTM = preprocessor.PreprocessorLSTM.from_dict(my_config)
+
+# 2.3. Run the training using bare tensorflow
+history_simple = train_rnn.train_without_wandb(preprocessorLSTM, config=my_config)
+
+# 2.4. Get the trained model
+model, train_stats, config = predict_rnn.get_pretrained_model('outputs')
+
+
+# 3. Define the callback function using the ML surrogate
+def run_sim_mixed(calib):
+    # split samples into two subsets to be used with the original function and the ML surrogate
+    np.random.seed()
+    ids = np.random.permutation(len(calib.system.param_data))
+    split_index = int(len(ids) * 0.5)
+    ids_origin, ids_surrogate = ids[:split_index], ids[split_index:]
+
+    # run the original function for the first half of the samples
+    param_data_origin = calib.system.param_data[ids_origin]
+    sim_data_origin = []
+    for params in param_data_origin:
+        # Run the model
+        y_sim = nonlinear(calib.system.ctrl_data, params)
+        sim_data_origin.append(np.array(y_sim, ndmin=2))
+    sim_data_origin = np.array(sim_data_origin)
+
+    # retrain the ML surrogate
+    my_config['input_data'] = calib.system.ctrl_data
+    my_config['param_data'] = np.vstack([my_config['param_data'], param_data_origin])
+    my_config['output_data'] = np.vstack([my_config['output_data'], sim_data_origin])
+
+    preprocessorLSTM = preprocessor.PreprocessorLSTM.from_dict(my_config)
+    history_simple = train_rnn.train_without_wandb(preprocessorLSTM, config=my_config)
+    model, train_stats, config = predict_rnn.get_pretrained_model('outputs')
+
+    # run the surrogate for the second half of the samples
+    param_data_surrogate = calib.system.param_data[ids_surrogate]
+    data_inputs = preprocessorLSTM.prepare_input_data(param_data_surrogate)
+    # make predictions with the trained model
+    sim_data_surrogate = predict_rnn.predict_batch(model, data_inputs, train_stats, config,
+                                                   batch_size=param_data_surrogate.shape[0])
+    # converting the predictions to GL format (temporal dimension at the end)
+    sim_data_surrogate = np.moveaxis(sim_data_surrogate, 1, -1)
+
+    # put the two subsets of simulation data together according to the original order
+    sim_data = np.zeros([calib.system.num_samples, calib.system.num_obs, calib.system.num_steps])
+    sim_data[ids_surrogate] = sim_data_surrogate
+    sim_data[ids_origin] = sim_data_origin
+
+    # set `sim_data` to system
+    calib.system.set_sim_data(sim_data)
+
+
+# set the callback function to the one that runs the ML surrogate
+calibration.callback = run_sim_mixed
+
+# continue the calibration with the surrogate
+calibration.num_iter = 10
+calibration.save_fig = 0
+calibration.run()

From 05be8a6ced22d902fcc0a620a947eed06a7ea111 Mon Sep 17 00:00:00 2001
From: chyalexcheng <chyalexcheng@gmail.com>
Date: Wed, 7 Feb 2024 22:17:45 +0100
Subject: [PATCH 08/21] rewrite a cleaner version of
 hyperbola_calibration_mixed.py

---
 .../LSTM/hyperbola_calibration_mixed.py       | 192 +++++++++---------
 1 file changed, 97 insertions(+), 95 deletions(-)

diff --git a/tutorials/data_driven/LSTM/hyperbola_calibration_mixed.py b/tutorials/data_driven/LSTM/hyperbola_calibration_mixed.py
index d343aee1..4fe74074 100644
--- a/tutorials/data_driven/LSTM/hyperbola_calibration_mixed.py
+++ b/tutorials/data_driven/LSTM/hyperbola_calibration_mixed.py
@@ -10,29 +10,113 @@
 y_obs = x_obs / (0.2 * 100 + 5.0 * x_obs)
 
 
-def run_sim(calib):
-    """This is the callback function that runs different realizations of the same model.
+def nonlinear(x, params):
+    return x / (params[0] * 100 + params[1] * x)
 
-    :param calib: The calibration object.
+
+# Define the configuration dictionary for the ML surrogate
+my_config = {
+    'input_data': None,
+    'param_data': None,
+    'output_data': None,
+    'train_frac': 0.7,
+    'val_frac': 0.2,
+    'window_size': 10,
+    'window_step': 1,
+    'patience': 25,
+    'epochs': 20,
+    'learning_rate': 1e-4,
+    'lstm_units': 128,
+    'dense_units': 128,
+    'batch_size': 64,
+    'standardize_outputs': True,
+    'save_weights_only': True
+}
+
+
+def run_sim_original(x, params):
+    """Run different realizations of the original model.
+
+    :param x: the input sequence
+    :param params: the parameters
     """
     data = []
-    for params in calib.system.param_data:
+    for params in params:
         # Run the model
-        y_sim = nonlinear(calib.system.ctrl_data, params)
-        data.append(np.array(y_sim, ndmin=2))
-    calib.system.set_sim_data(data)
+        y = nonlinear(x, params)
+        data.append(np.array(y, ndmin=2))
+    return np.array(data)
 
 
-def nonlinear(x, params):
-    return x / (params[0] * 100 + params[1] * x)
+def run_sim_surrogate(params_origin, output_origin, params_surrogate):
+    """Train the ML surrogate and evaluate model output with the ML surrogate.
+
+    :param params_origin: The parameter data used by the original model.
+    :param output_origin: The output data produced by the original model.
+    :param params_surrogate: The parameter data to be used by the ML surrogate.
+    """
+    # expend the parameter and output data
+    my_config['param_data'] = np.vstack([my_config['param_data'], params_origin])
+    my_config['output_data'] = np.vstack([my_config['output_data'], output_origin])
+
+    preprocessor_lstm = preprocessor.PreprocessorLSTM.from_dict(my_config)
+    _ = train_rnn.train_without_wandb(preprocessor_lstm, config=my_config)
+    model, train_stats, config = predict_rnn.get_pretrained_model('outputs')
+
+    # run the surrogate for the second half of the samples
+    data_inputs = preprocessor_lstm.prepare_input_data(params_surrogate)
+    # make predictions with the trained model
+    output_surrogate = predict_rnn.predict_batch(model, data_inputs, train_stats, config,
+                                                 batch_size=params_surrogate.shape[0])
+    # converting the predictions to GL format (temporal dimension at the end)
+    output_surrogate = np.moveaxis(output_surrogate, 1, -1)
+
+    return output_surrogate
+
+
+# 3. Define the callback function using the ML surrogate
+def run_sim_mixed(calib):
+    """This is the callback function that runs different realizations of the same model.
+
+    :param calib: The calibration object.
+    """
+    # if first iteration, run the original function
+    if calib.curr_iter == 0:
+        sim_data = run_sim_original(calib.system.ctrl_data, calib.system.param_data)
+        calib.system.set_sim_data(sim_data)
+        my_config['input_data'] = calib.system.ctrl_data
+        my_config['param_data'] = calib.system.param_data
+        my_config['output_data'] = calib.system.sim_data
+    else:
+        # split samples into two subsets to be used with the original function and the ML surrogate
+        np.random.seed()
+        ids = np.random.permutation(len(calib.system.param_data))
+        split_index = int(len(ids) * 0.5)
+        ids_origin, ids_surrogate = ids[:split_index], ids[split_index:]
+
+        # run the original function for the first half of the samples
+        param_data_origin = calib.system.param_data[ids_origin]
+        sim_data_origin = run_sim_original(calib.system.ctrl_data, param_data_origin)
+
+        # run the surrogate for the second half of the samples
+        param_data_surrogate = calib.system.param_data[ids_surrogate]
+        sim_data_surrogate = run_sim_surrogate(param_data_origin, sim_data_origin, param_data_surrogate)
+
+        # put the two subsets of simulation data together according to the original order
+        sim_data = np.zeros([calib.system.num_samples, calib.system.num_obs, calib.system.num_steps])
+        sim_data[ids_surrogate] = sim_data_surrogate
+        sim_data[ids_origin] = sim_data_origin
+
+        # set `sim_data` to system
+        calib.system.set_sim_data(sim_data)
 
 
 calibration = BayesianCalibration.from_dict(
     {
-        "num_iter": 1,
-        "callback": run_sim,
+        "num_iter": 5,
+        "callback": run_sim_mixed,
         "system": {
-            "param_min": [0.1, 1],
+            "param_min": [0.1, 0.1],
             "param_max": [1, 10],
             "param_names": ['a', 'b'],
             "num_samples": 20,
@@ -56,90 +140,8 @@ def nonlinear(x, params):
             },
             "initial_sampling": "halton",
         },
-        "save_fig": -1,
+        "save_fig": 0,
     }
 )
 
-# 1. Run one iteration of the calibration to collect the simulation data for training the ML surrogate
-calibration.run()
-
-# 2.1. Create my dictionary of configuration
-my_config = {
-    'input_data': calibration.system.ctrl_data,
-    'param_data': calibration.system.param_data,
-    'output_data': calibration.system.sim_data,
-    'train_frac': 0.7,
-    'val_frac': 0.2,
-    'window_size': 10,
-    'window_step': 1,
-    'patience': 25,
-    'epochs': 20,
-    'learning_rate': 1e-4,
-    'lstm_units': 128,
-    'dense_units': 128,
-    'batch_size': 64,
-    'standardize_outputs': True,
-    'save_weights_only': True
-}
-
-# 2.2. Create an object Preprocessor to pre-process my data
-preprocessorLSTM = preprocessor.PreprocessorLSTM.from_dict(my_config)
-
-# 2.3. Run the training using bare tensorflow
-history_simple = train_rnn.train_without_wandb(preprocessorLSTM, config=my_config)
-
-# 2.4. Get the trained model
-model, train_stats, config = predict_rnn.get_pretrained_model('outputs')
-
-
-# 3. Define the callback function using the ML surrogate
-def run_sim_mixed(calib):
-    # split samples into two subsets to be used with the original function and the ML surrogate
-    np.random.seed()
-    ids = np.random.permutation(len(calib.system.param_data))
-    split_index = int(len(ids) * 0.5)
-    ids_origin, ids_surrogate = ids[:split_index], ids[split_index:]
-
-    # run the original function for the first half of the samples
-    param_data_origin = calib.system.param_data[ids_origin]
-    sim_data_origin = []
-    for params in param_data_origin:
-        # Run the model
-        y_sim = nonlinear(calib.system.ctrl_data, params)
-        sim_data_origin.append(np.array(y_sim, ndmin=2))
-    sim_data_origin = np.array(sim_data_origin)
-
-    # retrain the ML surrogate
-    my_config['input_data'] = calib.system.ctrl_data
-    my_config['param_data'] = np.vstack([my_config['param_data'], param_data_origin])
-    my_config['output_data'] = np.vstack([my_config['output_data'], sim_data_origin])
-
-    preprocessorLSTM = preprocessor.PreprocessorLSTM.from_dict(my_config)
-    history_simple = train_rnn.train_without_wandb(preprocessorLSTM, config=my_config)
-    model, train_stats, config = predict_rnn.get_pretrained_model('outputs')
-
-    # run the surrogate for the second half of the samples
-    param_data_surrogate = calib.system.param_data[ids_surrogate]
-    data_inputs = preprocessorLSTM.prepare_input_data(param_data_surrogate)
-    # make predictions with the trained model
-    sim_data_surrogate = predict_rnn.predict_batch(model, data_inputs, train_stats, config,
-                                                   batch_size=param_data_surrogate.shape[0])
-    # converting the predictions to GL format (temporal dimension at the end)
-    sim_data_surrogate = np.moveaxis(sim_data_surrogate, 1, -1)
-
-    # put the two subsets of simulation data together according to the original order
-    sim_data = np.zeros([calib.system.num_samples, calib.system.num_obs, calib.system.num_steps])
-    sim_data[ids_surrogate] = sim_data_surrogate
-    sim_data[ids_origin] = sim_data_origin
-
-    # set `sim_data` to system
-    calib.system.set_sim_data(sim_data)
-
-
-# set the callback function to the one that runs the ML surrogate
-calibration.callback = run_sim_mixed
-
-# continue the calibration with the surrogate
-calibration.num_iter = 10
-calibration.save_fig = 0
 calibration.run()

From 0e4e0c4956a6e5369a242df8550e37d7c14ddb67 Mon Sep 17 00:00:00 2001
From: chyalexcheng <chyalexcheng@gmail.com>
Date: Mon, 12 Feb 2024 22:12:16 +0100
Subject: [PATCH 09/21] change name in models.py from 'contact_parameters' to
 'params'

---
 grainlearning/rnn/models.py                               | 8 ++++----
 tutorials/data_driven/LSTM/hyperbola_calibration_mixed.py | 1 -
 2 files changed, 4 insertions(+), 5 deletions(-)

diff --git a/grainlearning/rnn/models.py b/grainlearning/rnn/models.py
index 174eadfd..7f447926 100644
--- a/grainlearning/rnn/models.py
+++ b/grainlearning/rnn/models.py
@@ -35,11 +35,11 @@ def rnn_model(
     sequence_length = window_size
     load_sequence = layers.Input(
             shape=(sequence_length, input_shapes['num_input_features']), name='inputs')
-    contact_params = layers.Input(shape=(input_shapes['num_params'],), name='params')
+    params = layers.Input(shape=(input_shapes['num_params'],), name='params')
 
     # compute hidden state of LSTM based on contact parameters
-    state_h = layers.Dense(lstm_units, activation='tanh', name='state_h')(contact_params)
-    state_c = layers.Dense(lstm_units, activation='tanh', name='state_c')(contact_params)
+    state_h = layers.Dense(lstm_units, activation='tanh', name='state_h')(params)
+    state_c = layers.Dense(lstm_units, activation='tanh', name='state_c')(params)
     initial_state = [state_h, state_c]
 
     X = load_sequence
@@ -49,7 +49,7 @@ def rnn_model(
     X = layers.Dense(dense_units, activation='relu')(X)
     outputs = layers.Dense(input_shapes['num_labels'])(X)
 
-    model = Model(inputs=[load_sequence, contact_params], outputs=outputs)
+    model = Model(inputs=[load_sequence, params], outputs=outputs)
 
     return model
 
diff --git a/tutorials/data_driven/LSTM/hyperbola_calibration_mixed.py b/tutorials/data_driven/LSTM/hyperbola_calibration_mixed.py
index 4fe74074..4ae806f5 100644
--- a/tutorials/data_driven/LSTM/hyperbola_calibration_mixed.py
+++ b/tutorials/data_driven/LSTM/hyperbola_calibration_mixed.py
@@ -3,7 +3,6 @@
 from grainlearning.rnn import preprocessor
 import numpy as np
 from grainlearning import BayesianCalibration
-from matplotlib import pyplot as plt
 
 x_obs = np.arange(100)
 # hyperbola in a form similar to the Duncan-Chang material model, q = \eps / (a * 100 + b * \eps)

From 502afec9a14bb7181b25f50fc53c49c0a2e4a537 Mon Sep 17 00:00:00 2001
From: chyalexcheng <chyalexcheng@gmail.com>
Date: Thu, 15 Feb 2024 20:43:54 +0100
Subject: [PATCH 10/21] add a correct example of hyperparameter tuning using
 wandb

---
 pyproject.toml                                |   2 +-
 tutorials/data_driven/LSTM/train_rnn.py       |  32 ++++++
 .../data_driven/LSTM/train_rnn_hypertuning.py | 108 ++++++++++++++++++
 3 files changed, 141 insertions(+), 1 deletion(-)
 create mode 100644 tutorials/data_driven/LSTM/train_rnn.py
 create mode 100644 tutorials/data_driven/LSTM/train_rnn_hypertuning.py

diff --git a/pyproject.toml b/pyproject.toml
index e12fae4e..8745a398 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -39,7 +39,7 @@ ipykernel = {version = "*", optional = true}
 
 [tool.poetry.extras]
 docs = ["Sphinx", "sphinx-autodoc-typehints", "sphinx-mdinclude", "sphinx-rtd-theme"]
-dev = ["pytest", "pytest-cov", "prospector", "pyroma", "h5py", "tensorflow", "wandb"]
+dev = ["pytest", "pytest-cov", "prospector", "pyroma", "h5py", "tensorflow", "wandb", "ipykernel"]
 rnn = ["wandb", "tensorflow"]
 tutorials = ["ipykernel"]
 visuals = ["seaborn"]
diff --git a/tutorials/data_driven/LSTM/train_rnn.py b/tutorials/data_driven/LSTM/train_rnn.py
new file mode 100644
index 00000000..bb05f481
--- /dev/null
+++ b/tutorials/data_driven/LSTM/train_rnn.py
@@ -0,0 +1,32 @@
+import grainlearning.rnn.train as train_rnn
+from grainlearning.rnn import preprocessor
+
+# 1. Create my dictionary of configuration
+my_config = {
+     'raw_data': 'triaxial_compression_variable_input.hdf5',
+     'pressure': '0.2e6',
+     'experiment_type': 'drained',
+     'add_pressure': True,
+     'add_e0': True,
+     'train_frac': 0.1,
+     'val_frac': 0.1,
+     'window_size': 10,
+     'window_step': 1,
+     'patience': 25,
+     'epochs': 10,
+     'learning_rate': 0.006200000000000001,
+     'lstm_units': 122,
+     'dense_units': 224,
+     'batch_size': 216,
+     'standardize_outputs': True,
+     'save_weights_only': True
+ }
+
+# 2. Create an object Preprocessor to pre-process my data
+preprocessor_TC = preprocessor.PreprocessorTriaxialCompression(**my_config)
+
+# # 3. Run the training using bare tensorflow
+# history_simple = train_rnn.train_without_wandb(preprocessor_TC, config=my_config)
+
+# 3. Run the training using wandb
+history_wandb = train_rnn.train(preprocessor_TC, config=my_config)
diff --git a/tutorials/data_driven/LSTM/train_rnn_hypertuning.py b/tutorials/data_driven/LSTM/train_rnn_hypertuning.py
new file mode 100644
index 00000000..f93c07f2
--- /dev/null
+++ b/tutorials/data_driven/LSTM/train_rnn_hypertuning.py
@@ -0,0 +1,108 @@
+import wandb
+import grainlearning.rnn.train as train_rnn
+from grainlearning.rnn import preprocessor
+import os
+
+os.environ["WANDB_MODE"] = "offline"
+
+
+def my_training_function():
+    """ A function that wraps the training process"""
+    # update window_size of my_config from wandb
+    with wandb.init():
+        my_config['window_size'] = wandb.config['window_size']
+    preprocessor_TC = preprocessor.PreprocessorTriaxialCompression(**my_config)
+    train_rnn.train(preprocessor_TC)
+
+
+# 1. Create my dictionary of configuration
+my_config = {
+    'raw_data': 'triaxial_compression_variable_input.hdf5',
+    'pressure': '0.2e6',
+    'experiment_type': 'drained',
+    'add_experiment_type': False,
+    'add_pressure': True,
+    'add_e0': True,
+    'train_frac': 0.1,
+    'val_frac': 0.1,
+    'window_size': 20,
+    'pad_length': 10,
+    'window_step': 1,
+    'patience': 25,
+    'epochs': 10,
+    'learning_rate': 1e-4,
+    'lstm_units': 250,
+    'dense_units': 250,
+    'batch_size': 256,
+    'standardize_outputs': True,
+    'save_weights_only': True
+}
+
+
+# 2. A function to specify the sweep configuration, returning a sweep_id
+def get_sweep_id(method):
+    sweep_config = {
+        'method': 'random',
+        'name': method,
+        # TODO: @Luisa: how does user know what metric is defined in the model?
+        'metric': {'goal': 'minimize', 'name': 'mae'},
+        'parameters': {},
+        'early_terminate': {
+            'type': 'hyperband',
+            's': 2,
+            'eta': 3,
+            'max_iter': 27
+        }
+    }
+    # add default parameters from my_config into sweep_config, use 'values' as the key
+    for key, value in my_config.items():
+        sweep_config['parameters'].update({key: {'values': [value]}})
+    # update sweep_config with the parameters to be searched and their distributions
+    # check the documentation https://docs.wandb.ai/guides/sweeps/sweep-config-keys#distribution-options-for-random-and-bayesian-search on how to define the search space
+    search_space = {
+        'learning_rate': {
+            # a flat distribution between 1e-4 and 1e-2
+            'distribution': 'q_log_uniform_values',
+            'q': 1e-4,
+            'min': 1e-4,
+            'max': 1e-2
+        },
+        'lstm_units': {
+            'distribution': 'q_log_uniform_values',
+            'q': 1,
+            'min': 32,
+            'max': 256
+        },
+        'dense_units': {
+            'distribution': 'q_log_uniform_values',
+            'q': 1,
+            'min': 32,
+            'max': 256
+        },
+        'batch_size': {
+            'distribution': 'q_log_uniform_values',
+            'q': 1,
+            'min': 32,
+            'max': 256
+        },
+        'window_size': {
+            'distribution': 'q_uniform',
+            'q': 2,
+            'min': 4,
+            'max': 30
+        },
+    }
+    sweep_config['parameters'].update(search_space)
+    # create the sweep
+    sweep_id = wandb.sweep(sweep_config, project='triax_sweep')
+    return sweep_id
+
+
+# 3. Log in to wandb
+wandb.login()
+
+# 4. Configure the sweep
+sweep_id = get_sweep_id('random')
+
+# 5. Run an agent
+wandb.agent(sweep_id, function=my_training_function, count=100)

From ca9053d9578364d60476ee8800327f7922fab03e Mon Sep 17 00:00:00 2001
From: chyalexcheng <chyalexcheng@gmail.com>
Date: Thu, 15 Feb 2024 21:48:04 +0100
Subject: [PATCH 11/21] move hypertuning into the rnn.train module

---
 grainlearning/rnn/train.py                    |  82 +++++++++++++
 .../data_driven/LSTM/train_rnn_hypertuning.py | 113 ++++++++----------
 2 files changed, 130 insertions(+), 65 deletions(-)

diff --git a/grainlearning/rnn/train.py b/grainlearning/rnn/train.py
index a25f8de8..854ca788 100644
--- a/grainlearning/rnn/train.py
+++ b/grainlearning/rnn/train.py
@@ -270,3 +270,85 @@ def _get_optimizer_config(config):
             config_optimizer[key] = config[key]
 
     return config_optimizer
+
+
+class HyperTuning:
+    """
+    Class to run hyperparameter tuning with wandb.
+    """
+    def __init__(self, sweep_config, search_space, other_config = None, project_name = 'my_sweep'):
+        """
+        :param sweep_config: Dictionary containing the configuration of the sweep.
+
+            For example:
+
+            .. highlight:: python
+            .. code-block:: python
+
+            sweep_config = {
+                'method': 'random',
+                'name': my_sweep,
+                'metric': {'goal': 'minimize', 'name': 'mae'},
+                'early_terminate': {
+                    'type': 'hyperband',
+                    's': 2,
+                    'eta': 3,
+                    'max_iter': 27
+                }
+            }
+
+        :param search_space: Dictionary containing the search space of the sweep.
+
+            For example:
+
+            .. highlight:: python
+            .. code-block:: python
+
+            search_space = {
+                'learning_rate': {
+                    'distribution': 'q_log_uniform_values',
+                    'q': 1e-4,
+                    'min': 1e-4,
+                    'max': 1e-2
+                },
+                'lstm_units': {
+                    'distribution': 'q_log_uniform_values',
+                    'q': 1,
+                    'min': 32,
+                    'max': 256
+                },
+            }
+
+        :param other_config: Dictionary containing other relevant configuration parameters.
+        :param project_name: Name of the project in wandb.
+        """
+        # Log in to wandb
+        wandb.login()
+        # Set the configuration
+        self.sweep_config = sweep_config
+        self.search_space = search_space
+        self.other_config = other_config
+        self.project_name = project_name
+        self.sweep_id = None
+
+    def get_sweep_id(self):
+        """
+        Returns the sweep_id of a sweep created with the configuration specified in sweep_config and search_space.
+        """
+        # update sweep_config with the parameters to be searched and their distributions
+        self.sweep_config['parameters'] = self.search_space
+        # add default parameters from my_config into sweep_config, use 'values' as the key
+        for key, value in self.other_config.items():
+            self.sweep_config['parameters'].update({key: {'values': [value]}})
+        # create the sweep
+        sweep_id = wandb.sweep(self.sweep_config, project=self.project_name)
+        self.sweep_id = sweep_id
+
+    def run_sweep(self, training_func, count = 100):
+        """
+        Function to run hyperparameter tuning with `training_func`
+        """
+        # Configure the sweep
+        self.get_sweep_id()
+        # Run an agent
+        wandb.agent(self.sweep_id, function=training_func, count=count)
diff --git a/tutorials/data_driven/LSTM/train_rnn_hypertuning.py b/tutorials/data_driven/LSTM/train_rnn_hypertuning.py
index f93c07f2..0d56a5aa 100644
--- a/tutorials/data_driven/LSTM/train_rnn_hypertuning.py
+++ b/tutorials/data_driven/LSTM/train_rnn_hypertuning.py
@@ -1,6 +1,7 @@
-import wandb
 import grainlearning.rnn.train as train_rnn
+from grainlearning.rnn.train import HyperTuning
 from grainlearning.rnn import preprocessor
+import wandb
 import os
 
 os.environ["WANDB_MODE"] = "offline"
@@ -39,70 +40,52 @@ def my_training_function():
 }
 
 
-# 2. A function to specify the sweep configuration, returning a sweep_id
-def get_sweep_id(method):
-    sweep_config = {
-        'method': 'random',
-        'name': method,
-        # TODO: @Luisa: how does user know what metric is defined in the model?
-        'metric': {'goal': 'minimize', 'name': 'mae'},
-        'parameters': {},
-        'early_terminate': {
-            'type': 'hyperband',
-            's': 2,
-            'eta': 3,
-            'max_iter': 27
-        }
+# 2. Define the sweep configuration
+sweep_config = {
+    'method': 'random',
+    'metric': {'goal': 'minimize', 'name': 'mae'},
+    'early_terminate': {
+        'type': 'hyperband',
+        's': 2,
+        'eta': 3,
+        'max_iter': 27
     }
-    # add default parameters from my_config into sweep_config, use 'values' as the key
-    for key, value in my_config.items():
-        sweep_config['parameters'].update({key: {'values': [value]}})
-    # update sweep_config with the parameters to be searched and their distributions
-    # check the documentation https://docs.wandb.ai/guides/sweeps/sweep-config-keys#distribution-options-for-random-and-bayesian-search on how to define the search space
-    search_space = {
-        'learning_rate': {
-            # a flat distribution between 1e-4 and 1e-2
-            'distribution': 'q_log_uniform_values',
-            'q': 1e-4,
-            'min': 1e-4,
-            'max': 1e-2
-        },
-        'lstm_units': {
-            'distribution': 'q_log_uniform_values',
-            'q': 1,
-            'min': 32,
-            'max': 256
-        },
-        'dense_units': {
-            'distribution': 'q_log_uniform_values',
-            'q': 1,
-            'min': 32,
-            'max': 256
-        },
-        'batch_size': {
-            'distribution': 'q_log_uniform_values',
-            'q': 1,
-            'min': 32,
-            'max': 256
-        },
-        'window_size': {
-            'distribution': 'q_uniform',
-            'q': 2,
-            'min': 4,
-            'max': 30
-        },
-    }
-    sweep_config['parameters'].update(search_space)
-    # create the sweep
-    sweep_id = wandb.sweep(sweep_config, project='triax_sweep')
-    return sweep_id
-
-
-# 3. Log in to wandb
-wandb.login()
+}
 
-# 4. Configure the sweep
-sweep_id = get_sweep_id('random')
+search_space = {
+    'learning_rate': {
+        # a flat distribution between 1e-4 and 1e-2
+        'distribution': 'q_log_uniform_values',
+        'q': 1e-4,
+        'min': 1e-4,
+        'max': 1e-2
+    },
+    'lstm_units': {
+        'distribution': 'q_log_uniform_values',
+        'q': 1,
+        'min': 32,
+        'max': 256
+    },
+    # 'dense_units': {
+    #     'distribution': 'q_log_uniform_values',
+    #     'q': 1,
+    #     'min': 32,
+    #     'max': 256
+    # },
+    # 'batch_size': {
+    #     'distribution': 'q_log_uniform_values',
+    #     'q': 1,
+    #     'min': 32,
+    #     'max': 256
+    # },
+    'window_size': {
+        'distribution': 'q_uniform',
+        'q': 2,
+        'min': 4,
+        'max': 30
+    },
+}
 
-# 5. Run an agent
-wandb.agent(sweep_id, function=my_training_function, count=100)
+# 3. Run the sweep
+hyper_tuner = HyperTuning(sweep_config, search_space, my_config, project_name='my_sweep')
+hyper_tuner.run_sweep(my_training_function, count=100)

From e0f147476705f494fac2d8cb17abe8bfdf85e1b1 Mon Sep 17 00:00:00 2001
From: chyalexcheng <chyalexcheng@gmail.com>
Date: Thu, 15 Feb 2024 22:23:49 +0100
Subject: [PATCH 12/21] update the rnn tutorial

---
 docs/source/rnn.rst                           | 109 +++++++++++-------
 .../data_driven/LSTM/train_rnn_hypertuning.py |   1 +
 2 files changed, 69 insertions(+), 41 deletions(-)

diff --git a/docs/source/rnn.rst b/docs/source/rnn.rst
index f16dbd2c..11b92cdb 100644
--- a/docs/source/rnn.rst
+++ b/docs/source/rnn.rst
@@ -155,50 +155,77 @@ Create `my_sweep.py` where you would like to run the training. Configure the swe
 .. code-block:: python
    :caption: my_sweep.py
 
-   import wandb
    import grainlearning.rnn.train as train_rnn
+   from grainlearning.rnn.train import HyperTuning
    from grainlearning.rnn import preprocessor
-
+   import wandb
+   import os
+   
+   
    def my_training_function():
-     """ A function that wraps the training process"""
-     preprocessor_TC = preprocessor.PreprocessorTriaxialCompression(**wandb.config)
-     train_rnn.train(preprocessor_TC)
-
-   if __name__ == '__main__':
-      wandb.login()
-      sweep_configuration = {
-      'method': 'bayes',
-      'name': 'sweep',
-      'metric': {'goal': 'maximize', 'name': 'val_acc'},
-      'parameters':
-         {
-         'raw_data': 'my_path_to_dataset.hdf5',
-         'pressure': 'All',
-         'experiment_type': 'All',
-         'add_e0': False,
-         'add_pressure': True,
-         'add_experiment_type': True,
-         'train_frac': 0.7,
-         'val_frac': 0.15,
-         'window_size': 10,
-         'window_step': 1,
-         'pad_length': 0,
-         'lstm_units': 200,
-         'dense_units': 200,
-         'patience': 5,
-         'epochs': 100,
-         'learning_rate': 1e-3,
-         'batch_size': 256,
-         'standardize_outputs': True,
-         'save_weights_only': False
-         }
-      }
-      
-      # create a new sweep, here you can also configure your project and entity.
-      sweep_id = wandb.sweep(sweep=sweep_configuration)
-
-      # run an agent
-      wandb.agent(sweep_id, function=my_training_function, count=4)
+       """ A function that wraps the training process"""
+       # update window_size of my_config from wandb (only needed for the LSTM model)
+       with wandb.init():
+           my_config['window_size'] = wandb.config['window_size']
+       preprocessor_TC = preprocessor.PreprocessorTriaxialCompression(**my_config)
+       train_rnn.train(preprocessor_TC)
+  
+  
+   # 1. Create my dictionary of configuration
+   my_config = {
+       'raw_data': 'triaxial_compression_variable_input.hdf5',
+       'pressure': '0.2e6',
+       'experiment_type': 'drained',
+       'add_experiment_type': False,
+       'add_pressure': True,
+       'add_e0': True,
+       'train_frac': 0.1,
+       'val_frac': 0.1,
+       'window_size': 20,
+       'pad_length': 10,
+       'window_step': 1,
+       'patience': 25,
+       'epochs': 10,
+       'learning_rate': 1e-4,
+       'lstm_units': 250,
+       'dense_units': 250,
+       'batch_size': 256,
+       'standardize_outputs': True,
+       'save_weights_only': True
+   }
+  
+  
+   # 2. Define the sweep configuration
+   sweep_config = {
+       'method': 'random',
+       'metric': {'goal': 'minimize', 'name': 'mae'},
+       'early_terminate': {
+           'type': 'hyperband',
+           's': 2,
+           'eta': 3,
+           'max_iter': 27
+       }
+   }
+  
+   search_space = {
+       'learning_rate': {
+           # a flat distribution between 1e-4 and 1e-2
+           'distribution': 'q_log_uniform_values',
+           'q': 1e-4,
+           'min': 1e-4,
+           'max': 1e-2
+       },
+       'lstm_units': {
+           'distribution': 'q_log_uniform_values',
+           'q': 1,
+           'min': 32,
+           'max': 256
+       },
+   }
+
+   # 3. Run the sweep
+   hyper_tuner = HyperTuning(sweep_config, search_space, my_config, project_name='my_sweep')
+   hyper_tuner.run_sweep(my_training_function, count=100)
 
 Open a terminal where you have your file, activate the environment where grainLearning and rnn dependencies has been installed and run: ``python my_sweep.py``.
 
diff --git a/tutorials/data_driven/LSTM/train_rnn_hypertuning.py b/tutorials/data_driven/LSTM/train_rnn_hypertuning.py
index 0d56a5aa..c80faab6 100644
--- a/tutorials/data_driven/LSTM/train_rnn_hypertuning.py
+++ b/tutorials/data_driven/LSTM/train_rnn_hypertuning.py
@@ -43,6 +43,7 @@ def my_training_function():
 # 2. Define the sweep configuration
 sweep_config = {
     'method': 'random',
+    # TODO: how does the user know about the metric?
     'metric': {'goal': 'minimize', 'name': 'mae'},
     'early_terminate': {
         'type': 'hyperband',

From 73733a26707ee42c057c08f6209c66620067701f Mon Sep 17 00:00:00 2001
From: chyalexcheng <chyalexcheng@gmail.com>
Date: Thu, 15 Feb 2024 23:07:42 +0100
Subject: [PATCH 13/21] setting hypersearch range must be done after setting
 the default values of other config parameters.

---
 grainlearning/rnn/train.py                          | 9 +++++----
 tutorials/data_driven/LSTM/train_rnn_hypertuning.py | 3 ---
 2 files changed, 5 insertions(+), 7 deletions(-)

diff --git a/grainlearning/rnn/train.py b/grainlearning/rnn/train.py
index 854ca788..f0cf4c05 100644
--- a/grainlearning/rnn/train.py
+++ b/grainlearning/rnn/train.py
@@ -322,8 +322,6 @@ def __init__(self, sweep_config, search_space, other_config = None, project_name
         :param other_config: Dictionary containing other relevant configuration parameters.
         :param project_name: Name of the project in wandb.
         """
-        # Log in to wandb
-        wandb.login()
         # Set the configuration
         self.sweep_config = sweep_config
         self.search_space = search_space
@@ -335,11 +333,12 @@ def get_sweep_id(self):
         """
         Returns the sweep_id of a sweep created with the configuration specified in sweep_config and search_space.
         """
-        # update sweep_config with the parameters to be searched and their distributions
-        self.sweep_config['parameters'] = self.search_space
         # add default parameters from my_config into sweep_config, use 'values' as the key
+        self.sweep_config['parameters'] = {}
         for key, value in self.other_config.items():
             self.sweep_config['parameters'].update({key: {'values': [value]}})
+        # update sweep_config with the parameters to be searched and their distributions
+        self.sweep_config['parameters'].update(self.search_space)
         # create the sweep
         sweep_id = wandb.sweep(self.sweep_config, project=self.project_name)
         self.sweep_id = sweep_id
@@ -348,6 +347,8 @@ def run_sweep(self, training_func, count = 100):
         """
         Function to run hyperparameter tuning with `training_func`
         """
+        # Log in to wandb
+        wandb.login()
         # Configure the sweep
         self.get_sweep_id()
         # Run an agent
diff --git a/tutorials/data_driven/LSTM/train_rnn_hypertuning.py b/tutorials/data_driven/LSTM/train_rnn_hypertuning.py
index c80faab6..07264930 100644
--- a/tutorials/data_driven/LSTM/train_rnn_hypertuning.py
+++ b/tutorials/data_driven/LSTM/train_rnn_hypertuning.py
@@ -4,8 +4,6 @@
 import wandb
 import os
 
-os.environ["WANDB_MODE"] = "offline"
-
 
 def my_training_function():
     """ A function that wraps the training process"""
@@ -39,7 +37,6 @@ def my_training_function():
     'save_weights_only': True
 }
 
-
 # 2. Define the sweep configuration
 sweep_config = {
     'method': 'random',

From 5b92bd64b10ac12e0dea20fd471d9e6616f5d6d6 Mon Sep 17 00:00:00 2001
From: chyalexcheng <chyalexcheng@gmail.com>
Date: Mon, 19 Feb 2024 20:52:48 +0100
Subject: [PATCH 14/21] fix examples of parameter sweep using wandb

---
 grainlearning/rnn/evaluate_model.py           |  22 +++-
 grainlearning/rnn/predict.py                  |  28 +++--
 grainlearning/rnn/preprocessor.py             |   2 +-
 grainlearning/rnn/train.py                    | 103 ++++++++++--------
 .../{train_rnn.py => train_predict_rnn.py}    |  15 ++-
 ...ng.py => train_predict_rnn_hypertuning.py} |  23 +++-
 6 files changed, 125 insertions(+), 68 deletions(-)
 rename tutorials/data_driven/LSTM/{train_rnn.py => train_predict_rnn.py} (61%)
 rename tutorials/data_driven/LSTM/{train_rnn_hypertuning.py => train_predict_rnn_hypertuning.py} (72%)

diff --git a/grainlearning/rnn/evaluate_model.py b/grainlearning/rnn/evaluate_model.py
index aebba53c..cd2327ff 100644
--- a/grainlearning/rnn/evaluate_model.py
+++ b/grainlearning/rnn/evaluate_model.py
@@ -2,7 +2,7 @@
 import random
 import tensorflow as tf
 from matplotlib import pyplot as plt
-
+from sklearn.metrics import mean_absolute_error
 from grainlearning.rnn import predict
 
 
@@ -65,7 +65,7 @@ def _plot_sequence(i, j, y_key, i_s=0, x_key='steps', color='blue'):
     for i_s, color in zip(representative_idxs,
             ['blue', 'green', 'purple', 'darkgreen', 'navy', 'yellowgreen']):
 
-        p_label, e_label = _get_p_e_labels(config, test_inputs['contact_parameters'][i_s])
+        p_label, e_label = _get_p_e_labels(config, test_inputs['params'][i_s])
 
         _plot_sequence(0, 0, 'e', i_s=i_s, color=color)
         _plot_sequence(0, 1, 'f_0', i_s=i_s, color=color)
@@ -129,7 +129,7 @@ def _find_representatives(input_data, add_e0: bool, add_pressure: bool, add_expe
         global P_INDEX, E_INDEX
 
         representatives = []
-        contact_params = input_data['contact_parameters']
+        contact_params = input_data['params']
         if add_e0:
             P_INDEX -= 1
             E_INDEX -= 1
@@ -150,7 +150,7 @@ def _find_representatives(input_data, add_e0: bool, add_pressure: bool, add_expe
 
 
 def _find_random_samples(input_data, num_samples):
-    return np.random.choice(len(input_data['contact_parameters']), num_samples, replace=False)
+    return np.random.choice(len(input_data['params']), num_samples, replace=False)
 
 def _checks_extra_contact_params(config: dict):
     """
@@ -174,3 +174,17 @@ def _get_p_e_labels(config: dict, contact_params):
     if add_experiment_type: e_label = 'drained' if contact_params[E_INDEX]==1 else 'undrained'
     else: e_label = config['experiment_type']
     return p_label, e_label
+
+
+def plot_metric_distribution(data, predictions, config):
+    """
+    Plot the histogram of the scores.
+    """
+    test_inputs, labels = next(iter(data.batch(len(data))))
+    scores = [mean_absolute_error(labels[i, config['window_size']:, :].numpy(), predictions.numpy()[i, :, :])
+              for i in range(labels.numpy().shape[0])]
+    fig, ax = plt.subplots(1, 1)
+    ax.set_xlabel('Mean Absolute Error')
+    ax.set_ylabel('Frequency')
+    ax.hist(scores)
+    return fig
diff --git a/grainlearning/rnn/predict.py b/grainlearning/rnn/predict.py
index 9853f23d..a8db184d 100644
--- a/grainlearning/rnn/predict.py
+++ b/grainlearning/rnn/predict.py
@@ -1,6 +1,7 @@
 """
 Module containing functions to load a trained RNN model and make a prediction.
 """
+import glob
 import os
 from pathlib import Path
 
@@ -11,21 +12,21 @@
 import yaml
 
 from grainlearning.rnn.models import rnn_model
-#from grainlearning.rnn.preprocessing import prepare_datasets
 from grainlearning.rnn.windows import predict_over_windows
-from grainlearning.rnn.preprocessor import Preprocessor, PreprocessorTriaxialCompression
+from grainlearning.rnn.preprocessor import Preprocessor
 
 # Constants
 NAME_MODEL_H5 = 'model-best.h5'
 NAME_TRAIN_STATS = 'train_stats.npy'
 
 
-def get_best_run_from_sweep(entity_project_sweep_id: str, preprocessor: Preprocessor = None):
+def get_best_run_from_sweep(entity_project_sweep_id: str, order: str = None, preprocessor: Preprocessor = None):
     """
     Load the best performing model found with a weights and biases sweep.
     Also load the data splits it was trained on.
 
     :param entity_project_sweep_id: string of form <user>/<project>/<sweep_id>
+    :param order: string of the metric to use to select the best run.
     :param preprocessor: Preprocessor object to load and prepare the data.
       If None is given, then a PreprocessorTriaxialCompression will be considered
 
@@ -36,25 +37,28 @@ def get_best_run_from_sweep(entity_project_sweep_id: str, preprocessor: Preproce
         - config: The configuration dictionary used to train the model.
     """
     sweep = wandb.Api().sweep(entity_project_sweep_id)
-    best_run = sweep.best_run()
+    best_run = sweep.best_run(order=order)
     best_model = wandb.restore(
             NAME_MODEL_H5,  # this saves the model locally under this name
-            run_path=entity_project_sweep_id + best_run.id,
+            run_path=entity_project_sweep_id + '/' + best_run.id,
             replace=True,
         )
     config = best_run.config
-    if preprocessor is None:
-        preprocessor = PreprocessorTriaxialCompression(**config)
 
-    if os.path.exists(Path(best_model.dir)/NAME_TRAIN_STATS):
-        train_stats = np.load(Path(best_model.dir)/NAME_TRAIN_STATS, allow_pickle=True).item()
-        data, _ = preprocessor.prepare_datasets()
+    # get the training statistics file from local directories
+    # (best_model from wandb only gives the model not other files)
+    stat_files = glob.glob(str(Path(f'wandb/*{best_run.id}*/files/*{NAME_TRAIN_STATS}')))
+    if len(stat_files) == 1:
+        train_stats = np.load(stat_files[0], allow_pickle=True).item()
+    elif len(stat_files) > 1:
+        raise ValueError(f"Multiple files found with name {NAME_TRAIN_STATS} in the wandb directory. \
+            Please check the duplicates and try again.")
     else:
-        data, train_stats = preprocessor.prepare_datasets()
+        raise FileNotFoundError(f"File {NAME_TRAIN_STATS} was not found in the wandb directory.")
 
     model = rnn_model(train_stats, **config)
     model.load_weights(best_model.name)
-    return model, data, train_stats, config
+    return model, train_stats, config
 
 
 def get_pretrained_model(path_to_model: str):
diff --git a/grainlearning/rnn/preprocessor.py b/grainlearning/rnn/preprocessor.py
index 67df6d51..7083db9c 100644
--- a/grainlearning/rnn/preprocessor.py
+++ b/grainlearning/rnn/preprocessor.py
@@ -11,7 +11,7 @@
 class Preprocessor(ABC):
 
     def __init__(self):
-        pass
+        self.run_dir = None
 
     @classmethod
     def from_dict(cls: Type["Preprocessor"], obj: dict):
diff --git a/grainlearning/rnn/train.py b/grainlearning/rnn/train.py
index f0cf4c05..68efb180 100644
--- a/grainlearning/rnn/train.py
+++ b/grainlearning/rnn/train.py
@@ -15,7 +15,7 @@
 from grainlearning.rnn.windows import windowize_single_dataset
 
 
-def train(preprocessor: Preprocessor, config = None, model: tf.keras.Model = None):
+def train(preprocessor: Preprocessor, config=None, model: tf.keras.Model = None):
     """
     Train a model and report to weights and biases.
 
@@ -38,6 +38,7 @@ def train(preprocessor: Preprocessor, config = None, model: tf.keras.Model = Non
         config = wandb.config
         config = _check_config(config, preprocessor)
         config_optimizer = _get_optimizer_config(config)
+        preprocessor.run_dir = wandb.run.dir
 
         # preprocess data
         split_data, train_stats = preprocessor.prepare_datasets()
@@ -49,10 +50,10 @@ def train(preprocessor: Preprocessor, config = None, model: tf.keras.Model = Non
 
         optimizer = tf.keras.optimizers.Adam(**config_optimizer)
         model.compile(
-                optimizer=optimizer,
-                loss='mse',
-                metrics=['mae'],
-            )
+            optimizer=optimizer,
+            loss='mse',
+            metrics=['mae'],
+        )
 
         # create batches
         for split in ['train', 'val']:  # do not batch test set
@@ -60,25 +61,25 @@ def train(preprocessor: Preprocessor, config = None, model: tf.keras.Model = Non
 
         # set up training
         early_stopping = tf.keras.callbacks.EarlyStopping(
-                monitor='val_loss',
-                patience=config.patience,
-                restore_best_weights=True,
-            )
+            monitor='val_loss',
+            patience=config.patience,
+            restore_best_weights=True,
+        )
         wandb_callback = wandb.keras.WandbCallback(
-                monitor='val_loss',
-                save_model=True,
-                save_weights_only=config.save_weights_only,
-                validation_data=split_data['val'],
-            )
+            monitor='val_loss',
+            save_model=True,
+            save_weights_only=config.save_weights_only,
+            validation_data=split_data['val'],
+        )
         callbacks = [wandb_callback, early_stopping]
 
         # train
         history = model.fit(
-                split_data['train'],
-                epochs=config.epochs,
-                validation_data=split_data['val'],
-                callbacks=callbacks,
-            )
+            split_data['train'],
+            epochs=config.epochs,
+            validation_data=split_data['val'],
+            callbacks=callbacks,
+        )
 
         # Evaluate in test dataset and log to wandb the metrics
         test_data = windowize_single_dataset(split_data['test'], **config)
@@ -88,7 +89,8 @@ def train(preprocessor: Preprocessor, config = None, model: tf.keras.Model = Non
 
         return history
 
-def train_without_wandb(preprocessor: Preprocessor, config = None, model: tf.keras.Model = None):
+
+def train_without_wandb(preprocessor: Preprocessor, config=None, model: tf.keras.Model = None):
     """
     Train a model locally: no report to wandb.
     Saves either the model or its weight to folder outputs.
@@ -108,7 +110,8 @@ def train_without_wandb(preprocessor: Preprocessor, config = None, model: tf.ker
     if os.path.exists(path_save_data):
         delete_outputs = input(f"The contents of {path_save_data} will be permanently deleted,\
                                  do you want to proceed? [y/n]: ")
-        if delete_outputs == "y": shutil.rmtree(path_save_data)
+        if delete_outputs == "y":
+            shutil.rmtree(path_save_data)
         else:
             raise SystemExit("Cancelling training")
 
@@ -116,8 +119,8 @@ def train_without_wandb(preprocessor: Preprocessor, config = None, model: tf.ker
 
     # preprocess data
     split_data, train_stats = preprocessor.prepare_datasets()
-    np.save(path_save_data/'train_stats.npy', train_stats)
-    np.save(path_save_data/'config.npy', config)
+    np.save(path_save_data / 'train_stats.npy', train_stats)
+    np.save(path_save_data / 'config.npy', config)
 
     # set up the model
     if model is None:
@@ -125,36 +128,36 @@ def train_without_wandb(preprocessor: Preprocessor, config = None, model: tf.ker
 
     optimizer = tf.keras.optimizers.Adam(**config_optimizer)
     model.compile(
-            optimizer=optimizer,
-            loss='mse',
-            metrics=['mae'],
-        )
+        optimizer=optimizer,
+        loss='mse',
+        metrics=['mae'],
+    )
 
     # create batches
     for split in ['train', 'val']:  # do not batch test set
         split_data[split] = split_data[split].batch(config['batch_size'])
 
     # set up training
-    if config['save_weights_only'] : path_save_data = path_save_data/"weights.h5"
+    if config['save_weights_only']: path_save_data = path_save_data / "weights.h5"
     early_stopping = tf.keras.callbacks.EarlyStopping(
-            monitor='val_loss',
-            patience=config['patience'],
-            restore_best_weights=True,
-        )
+        monitor='val_loss',
+        patience=config['patience'],
+        restore_best_weights=True,
+    )
     checkpoint = tf.keras.callbacks.ModelCheckpoint(
-                str(path_save_data),
-                monitor='val_loss',
-                save_best_only=True,
-                save_weights_only=config['save_weights_only']
-            )
+        str(path_save_data),
+        monitor='val_loss',
+        save_best_only=True,
+        save_weights_only=config['save_weights_only']
+    )
 
     # train
     history = model.fit(
-            split_data['train'],
-            epochs=config['epochs'],
-            validation_data=split_data['val'],
-            callbacks=[early_stopping, checkpoint],
-        )
+        split_data['train'],
+        epochs=config['epochs'],
+        validation_data=split_data['val'],
+        callbacks=[early_stopping, checkpoint],
+    )
 
     # Evaluate in test dataset and print the metrics
     test_data = windowize_single_dataset(split_data['test'], **config)
@@ -238,15 +241,16 @@ def _check_config(config: dict, preprocessor: Preprocessor):
     return config
 
 
-def _warning_config_field(key, config, default, add_default_to_config = False):
+def _warning_config_field(key, config, default, add_default_to_config=False):
     """
     Raises a warning if key is not included in config dictionary.
     Also informs the default value that will be used.
     If add_default_to_config=True, then it adds the key and its default value to config.
     """
+
     # customized warning to print -only- the warning message
     def _custom_format_warning(msg, *_):
-        return str(msg) + '\n' # ignore everything except the message
+        return str(msg) + '\n'  # ignore everything except the message
 
     warnings.formatwarning = _custom_format_warning
 
@@ -276,7 +280,9 @@ class HyperTuning:
     """
     Class to run hyperparameter tuning with wandb.
     """
-    def __init__(self, sweep_config, search_space, other_config = None, project_name = 'my_sweep'):
+
+    def __init__(self, sweep_config, search_space, other_config=None, entity_name='grainlearning',
+                 project_name='my_sweep'):
         """
         :param sweep_config: Dictionary containing the configuration of the sweep.
 
@@ -327,6 +333,7 @@ def __init__(self, sweep_config, search_space, other_config = None, project_name
         self.search_space = search_space
         self.other_config = other_config
         self.project_name = project_name
+        self.entity_name = entity_name
         self.sweep_id = None
 
     def get_sweep_id(self):
@@ -343,7 +350,7 @@ def get_sweep_id(self):
         sweep_id = wandb.sweep(self.sweep_config, project=self.project_name)
         self.sweep_id = sweep_id
 
-    def run_sweep(self, training_func, count = 100):
+    def run_sweep(self, training_func, count=100):
         """
         Function to run hyperparameter tuning with `training_func`
         """
@@ -353,3 +360,7 @@ def run_sweep(self, training_func, count = 100):
         self.get_sweep_id()
         # Run an agent
         wandb.agent(self.sweep_id, function=training_func, count=count)
+        # Close the wandb session
+        wandb.finish()
+        # write the data of the HyperTuning object to a npy file
+        np.save(f'wandb/{self.entity_name}_{self.project_name}_{self.sweep_id}.npy', self.__dict__)
diff --git a/tutorials/data_driven/LSTM/train_rnn.py b/tutorials/data_driven/LSTM/train_predict_rnn.py
similarity index 61%
rename from tutorials/data_driven/LSTM/train_rnn.py
rename to tutorials/data_driven/LSTM/train_predict_rnn.py
index bb05f481..d07e2416 100644
--- a/tutorials/data_driven/LSTM/train_rnn.py
+++ b/tutorials/data_driven/LSTM/train_predict_rnn.py
@@ -1,5 +1,7 @@
 import grainlearning.rnn.train as train_rnn
 from grainlearning.rnn import preprocessor
+from grainlearning.rnn.predict import predict_batch, get_pretrained_model
+from grainlearning.rnn.evaluate_model import plot_metric_distribution
 
 # 1. Create my dictionary of configuration
 my_config = {
@@ -8,7 +10,7 @@
      'experiment_type': 'drained',
      'add_pressure': True,
      'add_e0': True,
-     'train_frac': 0.1,
+     'train_frac': 0.8,
      'val_frac': 0.1,
      'window_size': 10,
      'window_step': 1,
@@ -19,7 +21,7 @@
      'dense_units': 224,
      'batch_size': 216,
      'standardize_outputs': True,
-     'save_weights_only': True
+     'save_weights_only': False
  }
 
 # 2. Create an object Preprocessor to pre-process my data
@@ -30,3 +32,12 @@
 
 # 3. Run the training using wandb
 history_wandb = train_rnn.train(preprocessor_TC, config=my_config)
+
+# 5. Load input data to predict from
+model, train_stats, config = get_pretrained_model(preprocessor_TC.run_dir)
+data = preprocessor_TC.prepare_single_dataset()
+
+# 6. Make predictions and plot the histogram of errors
+predictions = predict_batch(history_wandb.model, data, train_stats, config, batch_size=len(data))
+fig = plot_metric_distribution(data, predictions, config)
+fig.show()
diff --git a/tutorials/data_driven/LSTM/train_rnn_hypertuning.py b/tutorials/data_driven/LSTM/train_predict_rnn_hypertuning.py
similarity index 72%
rename from tutorials/data_driven/LSTM/train_rnn_hypertuning.py
rename to tutorials/data_driven/LSTM/train_predict_rnn_hypertuning.py
index 07264930..a30184e1 100644
--- a/tutorials/data_driven/LSTM/train_rnn_hypertuning.py
+++ b/tutorials/data_driven/LSTM/train_predict_rnn_hypertuning.py
@@ -1,8 +1,9 @@
 import grainlearning.rnn.train as train_rnn
 from grainlearning.rnn.train import HyperTuning
 from grainlearning.rnn import preprocessor
+from grainlearning.rnn.predict import get_best_run_from_sweep, predict_batch
+from grainlearning.rnn.evaluate_model import plot_metric_distribution
 import wandb
-import os
 
 
 def my_training_function():
@@ -22,8 +23,8 @@ def my_training_function():
     'add_experiment_type': False,
     'add_pressure': True,
     'add_e0': True,
-    'train_frac': 0.1,
-    'val_frac': 0.1,
+    'train_frac': 0.7,
+    'val_frac': 0.15,
     'window_size': 20,
     'pad_length': 10,
     'window_step': 1,
@@ -87,3 +88,19 @@ def my_training_function():
 # 3. Run the sweep
 hyper_tuner = HyperTuning(sweep_config, search_space, my_config, project_name='my_sweep')
 hyper_tuner.run_sweep(my_training_function, count=100)
+
+# 4. Get the best model
+# change metric to validation loss
+hyper_tuner.sweep_config['metric']['name'] = 'val_loss'
+entity_project_sweep_id = f"{hyper_tuner.entity_name}/{hyper_tuner.project_name}/{hyper_tuner.sweep_id}"
+model, train_stats, config = get_best_run_from_sweep(entity_project_sweep_id)
+config['pressure'] = my_config['pressure']
+
+# 5. Load input data to predict from
+preprocessor_TC = preprocessor.PreprocessorTriaxialCompression(**config)
+data = preprocessor_TC.prepare_single_dataset()
+
+# 6. Make and plot predictions
+predictions = predict_batch(model, data, train_stats, config, batch_size=len(data))
+fig = plot_metric_distribution(data, predictions, config)
+fig.show()

From 291600d2d152fc5679a8245f57b9680e0bbc0d41 Mon Sep 17 00:00:00 2001
From: chyalexcheng <chyalexcheng@gmail.com>
Date: Mon, 19 Feb 2024 21:55:38 +0100
Subject: [PATCH 15/21] adapt rnn documentation

---
 docs/source/rnn.rst                           | 53 ++++++++++++-------
 .../LSTM/{predict.ipynb => predict_rnn.ipynb} |  0
 ...tion_GL.ipynb => rnn_GL_calibration.ipynb} |  0
 .../data_driven/LSTM/train_predict_rnn.py     | 16 +++---
 .../LSTM/train_predict_rnn_hypertuning.py     |  2 +-
 5 files changed, 44 insertions(+), 27 deletions(-)
 rename tutorials/data_driven/LSTM/{predict.ipynb => predict_rnn.ipynb} (100%)
 rename tutorials/data_driven/LSTM/{rnn_calibration_GL.ipynb => rnn_GL_calibration.ipynb} (100%)

diff --git a/docs/source/rnn.rst b/docs/source/rnn.rst
index 11b92cdb..1568ea63 100644
--- a/docs/source/rnn.rst
+++ b/docs/source/rnn.rst
@@ -4,12 +4,26 @@ RNN Module
 We implemented a `Recurrent Neural Network (RNN) <https://stanford.edu/~shervine/teaching/cs-230/cheatsheet-recurrent-neural-networks>`_ model
 in the tensorflow framework. For more information about the model go to section `The RNN model`_.
 
-There are four main usages of the RNN module:
+There are four main usages of the RNN module (click on the links to download the scripts):
 
 1. `Train a RNN with your own data`_.
 2. `Make a prediction with a pre-trained model`_.
 3. `Use a trained RNN model in grainLearning calibration process`_.
-   
+
+We have made one tutorial that demonstrates
+how to combine :download:`training and predicting <../../tutorials/data_driven/LSTM/train_predict_rnn.py>` of ML models within one script.
+and another one on performing click :download:`hyperparameter optimization <../../tutorials/data_driven/LSTM/train_predict_rnn_hypertuning.py>`
+before making predictions.
+
+We provide two simple tutorial to demonstrate how to:
+
+- use an LSTM model that mimic hyperbolic curves for GL Bayesian Calibration workflow 
+  (click :download:`here <../../tutorials/data_driven/LSTM/hyperbola_calibration_lstm.py>`)
+- and partly replace the close-form solution of the hyperbolic model (click :download:`here <../../tutorials/data_driven/LSTM/hyperbola_calibration_mixed.py>`) with the LSTM, 
+
+in addition to the example (click :download:`here <../../tutorials/data_driven/LSTM/rnn_GL_calibration.ipynb>`) where
+an LSTM is pretrained to predict the triaxial response of a granular material and then used in the calibration process.
+
 Train a RNN with your own data
 ------------------------------
 
@@ -129,7 +143,7 @@ Create `my_train.py` where you would like to run the training. Be aware to confi
    preprocessor_TC = preprocessor.PreprocessorTriaxialCompression(**my_config)
    
    # 3. Run the training Tensorflow and reporting to wandb
-   train_rnn.train(preprocessor_TC, config=my_config)
+   history_wandb = train_rnn.train(preprocessor_TC, config=my_config)
 
 Open a terminal where you have your file, activate the environment where grainLearning and rnn dependencies has been installed and run: ``python my_train.py``
 
@@ -173,14 +187,14 @@ Create `my_sweep.py` where you would like to run the training. Configure the swe
   
    # 1. Create my dictionary of configuration
    my_config = {
-       'raw_data': 'triaxial_compression_variable_input.hdf5',
+       'raw_data': 'my_path_to_dataset.hdf5',
        'pressure': '0.2e6',
        'experiment_type': 'drained',
        'add_experiment_type': False,
        'add_pressure': True,
        'add_e0': True,
-       'train_frac': 0.1,
-       'val_frac': 0.1,
+       'train_frac': 0.7,
+       'val_frac': 0.15,
        'window_size': 20,
        'pad_length': 10,
        'window_step': 1,
@@ -194,11 +208,10 @@ Create `my_sweep.py` where you would like to run the training. Configure the swe
        'save_weights_only': True
    }
   
-  
    # 2. Define the sweep configuration
    sweep_config = {
        'method': 'random',
-       'metric': {'goal': 'minimize', 'name': 'mae'},
+       'metric': {'goal': 'minimize', 'name': 'val_loss'},
        'early_terminate': {
            'type': 'hyperband',
            's': 2,
@@ -338,22 +351,24 @@ In this example, we are going to load the same dataset that we used for training
 
    from pathlib import Path
 
-   import grainlearning.rnn.predict as predict_rnn
+   from grainlearning.rnn.predict import predict_batch, plot_metric_distribution
    from grainlearning.rnn import preprocessor
 
    # 1. Define the location of the model to use
-   path_to_trained_model = Path('C:/trained_models/My_model_1')
+   path_to_trained_model = Path('my_path_to_run_directory')
 
    # 2. Get the model information
    model, train_stats, config = predict_rnn.get_pretrained_model(path_to_trained_model)
 
    # 3. Load input data to predict from
-   config['raw_data'] = '../train/data/my_database.hdf5'
+   config['raw_data'] = 'my_path_to_dataset.hdf5'
    preprocessor_TC = preprocessor.PreprocessorTriaxialCompression(**config)
    data, _ = preprocessor_TC.prepare_datasets()
 
-   #4. Make a prediction
-   predictions = predict_rnn.predict_macroscopics(model, data['test'], train_stats, config,batch_size=256, single_batch=True)
+   # 4. Make a prediction and plot the histogram of errors
+   predictions = predict_rnn.predict_batch(model, data['test'], train_stats, config, batch_size=len(data['test']))
+   fig = plot_metric_distribution(data, predictions, config)
+   fig.show()
 
 If the model was trained with ``standardize_outputs = True``, ``predictions`` are going to be unstandardized (i.e. no values between [0, 1] but with the original scale). 
 In our example, ``predictions`` is a tensorflow tensor of size ``(batch_size, length_sequences - window_size, num_labels)``.
@@ -372,18 +387,20 @@ Often this looks like `<entity>/<project>/<sweep_id>`.
    from grainlearning.rnn import preprocessor
 
    # 1. Define which sweep to look into
-   entity_project_sweep_id = 'grainlearning-escience/grainLearning-grainlearning_rnn/6zrc0vjb'
+   entity_project_sweep_id = 'grainlearning/project/sweep_id'
 
    # 2. Chose the best model from a sweep, and get the model information
-   model, data, train_stats, config = predict_rnn.get_best_run_from_sweep(entity_project_sweep_id)
+   model, train_stats, config = predict_rnn.get_best_run_from_sweep(entity_project_sweep_id)
 
    # 3. Load input data to predict from
-   config['raw_data'] = '../train/data/sequences.hdf5'
+   config['raw_data'] = 'my_path_to_dataset.hdf5'
    preprocessor_TC = preprocessor.PreprocessorTriaxialCompression(**config)
    data, _ = preprocessor_TC.prepare_datasets()
 
-   #4. Make a prediction
-   predictions = predict_rnn.predict_macroscopics(model, data['test'], train_stats, config,batch_size=256, single_batch=True)
+   # 4. Make a prediction and plot the histogram of errors
+   predictions = predict_rnn.predict_batch(model, data['test'], train_stats, config, batch_size=len(data['test']))
+   fig = plot_metric_distribution(data, predictions, config)
+   fig.show()
 
 This can fail if you have deleted some runs or if your wandb folder is not present in this folder. We advise to copy `config.yaml`, `train_stats.py` and `model_best.h5` from `wandb/runXXX/files` to another location and follow `Saved model`_ instructions. These files can also be downloaded from the wandb dashboard.
 
diff --git a/tutorials/data_driven/LSTM/predict.ipynb b/tutorials/data_driven/LSTM/predict_rnn.ipynb
similarity index 100%
rename from tutorials/data_driven/LSTM/predict.ipynb
rename to tutorials/data_driven/LSTM/predict_rnn.ipynb
diff --git a/tutorials/data_driven/LSTM/rnn_calibration_GL.ipynb b/tutorials/data_driven/LSTM/rnn_GL_calibration.ipynb
similarity index 100%
rename from tutorials/data_driven/LSTM/rnn_calibration_GL.ipynb
rename to tutorials/data_driven/LSTM/rnn_GL_calibration.ipynb
diff --git a/tutorials/data_driven/LSTM/train_predict_rnn.py b/tutorials/data_driven/LSTM/train_predict_rnn.py
index d07e2416..8363775a 100644
--- a/tutorials/data_driven/LSTM/train_predict_rnn.py
+++ b/tutorials/data_driven/LSTM/train_predict_rnn.py
@@ -10,18 +10,18 @@
      'experiment_type': 'drained',
      'add_pressure': True,
      'add_e0': True,
-     'train_frac': 0.8,
-     'val_frac': 0.1,
-     'window_size': 10,
+     'train_frac': 0.7,
+     'val_frac': 0.15,
+     'window_size': 20,
      'window_step': 1,
      'patience': 25,
      'epochs': 10,
-     'learning_rate': 0.006200000000000001,
-     'lstm_units': 122,
-     'dense_units': 224,
-     'batch_size': 216,
+     'learning_rate': 1e-4,
+     'lstm_units': 250,
+     'dense_units': 250,
+     'batch_size': 256,
      'standardize_outputs': True,
-     'save_weights_only': False
+     'save_weights_only': True
  }
 
 # 2. Create an object Preprocessor to pre-process my data
diff --git a/tutorials/data_driven/LSTM/train_predict_rnn_hypertuning.py b/tutorials/data_driven/LSTM/train_predict_rnn_hypertuning.py
index a30184e1..97c4e2c3 100644
--- a/tutorials/data_driven/LSTM/train_predict_rnn_hypertuning.py
+++ b/tutorials/data_driven/LSTM/train_predict_rnn_hypertuning.py
@@ -42,7 +42,7 @@ def my_training_function():
 sweep_config = {
     'method': 'random',
     # TODO: how does the user know about the metric?
-    'metric': {'goal': 'minimize', 'name': 'mae'},
+    'metric': {'goal': 'minimize', 'name': 'val_loss'},
     'early_terminate': {
         'type': 'hyperband',
         's': 2,

From 37e6ec1246ba3c2b0422dadb3bc19b6e80962f0c Mon Sep 17 00:00:00 2001
From: chyalexcheng <chyalexcheng@gmail.com>
Date: Mon, 19 Feb 2024 23:43:31 +0100
Subject: [PATCH 16/21] add hyperparameter tuning to GL workflow. Got error 400
 response executing GraphQL

---
 grainlearning/rnn/train.py                    |   2 +-
 ...hyperbola_calibration_mixed_hypertuning.py | 196 ++++++++++++++++++
 .../LSTM/train_predict_rnn_hypertuning.py     |   2 +-
 3 files changed, 198 insertions(+), 2 deletions(-)
 create mode 100644 tutorials/data_driven/LSTM/hyperbola_calibration_mixed_hypertuning.py

diff --git a/grainlearning/rnn/train.py b/grainlearning/rnn/train.py
index 68efb180..f50afb68 100644
--- a/grainlearning/rnn/train.py
+++ b/grainlearning/rnn/train.py
@@ -347,7 +347,7 @@ def get_sweep_id(self):
         # update sweep_config with the parameters to be searched and their distributions
         self.sweep_config['parameters'].update(self.search_space)
         # create the sweep
-        sweep_id = wandb.sweep(self.sweep_config, project=self.project_name)
+        sweep_id = wandb.sweep(self.sweep_config, entity=self.entity_name, project=self.project_name)
         self.sweep_id = sweep_id
 
     def run_sweep(self, training_func, count=100):
diff --git a/tutorials/data_driven/LSTM/hyperbola_calibration_mixed_hypertuning.py b/tutorials/data_driven/LSTM/hyperbola_calibration_mixed_hypertuning.py
new file mode 100644
index 00000000..12b6edbf
--- /dev/null
+++ b/tutorials/data_driven/LSTM/hyperbola_calibration_mixed_hypertuning.py
@@ -0,0 +1,196 @@
+import grainlearning.rnn.train as train_rnn
+import grainlearning.rnn.predict as predict_rnn
+from grainlearning.rnn import preprocessor
+import numpy as np
+from grainlearning import BayesianCalibration
+import wandb
+
+x_obs = np.arange(100)
+# hyperbola in a form similar to the Duncan-Chang material model, q = \eps / (a * 100 + b * \eps)
+y_obs = x_obs / (0.2 * 100 + 5.0 * x_obs)
+
+
+def nonlinear(x, params):
+    return x / (params[0] * 100 + params[1] * x)
+
+
+# Define the configuration dictionary for the ML surrogate
+my_config = {
+    'input_data': None,
+    'param_data': None,
+    'output_data': None,
+    'train_frac': 0.7,
+    'val_frac': 0.2,
+    'window_size': 10,
+    'window_step': 1,
+    'patience': 25,
+    'epochs': 20,
+    'learning_rate': 1e-4,
+    'lstm_units': 128,
+    'dense_units': 128,
+    'batch_size': 64,
+    'standardize_outputs': True,
+    'save_weights_only': True
+}
+
+# 2. Define the sweep configuration
+sweep_config = {
+    'method': 'random',
+    'metric': {'goal': 'minimize', 'name': 'mae'},
+    'early_terminate': {
+        'type': 'hyperband',
+        's': 2,
+        'eta': 3,
+        'max_iter': 27
+    }
+}
+
+search_space = {
+    'learning_rate': {
+        # a flat distribution between 1e-4 and 1e-2
+        'distribution': 'q_log_uniform_values',
+        'q': 1e-4,
+        'min': 1e-4,
+        'max': 1e-2
+    },
+    'lstm_units': {
+        'distribution': 'q_log_uniform_values',
+        'q': 1,
+        'min': 32,
+        'max': 256
+    },
+    'window_size': {
+        'distribution': 'q_uniform',
+        'q': 2,
+        'min': 4,
+        'max': 30
+    },
+}
+
+
+def run_sim_original(x, params):
+    """Run different realizations of the original model.
+
+    :param x: the input sequence
+    :param params: the parameters
+    """
+    data = []
+    for params in params:
+        # Run the model
+        y = nonlinear(x, params)
+        data.append(np.array(y, ndmin=2))
+    return np.array(data)
+
+
+def my_training_function():
+    preprocessor_lstm = preprocessor.PreprocessorLSTM.from_dict(my_config)
+    train_rnn.train(preprocessor_lstm, config=my_config)
+
+
+def hyper_train():
+    """Train the ML surrogate using hyperparameter tuning."""
+    hyper_tuner = train_rnn.HyperTuning(sweep_config, search_space, my_config, project_name='hyperbola_sweep')
+    hyper_tuner.run_sweep(my_training_function, count=10)
+
+    hyper_tuner.sweep_config['metric']['name'] = 'val_loss'
+    entity_project_sweep_id = f"{hyper_tuner.entity_name}/{hyper_tuner.project_name}/{hyper_tuner.sweep_id}"
+    model, train_stats, config = predict_rnn.get_best_run_from_sweep(entity_project_sweep_id)
+    return model, train_stats, config
+
+
+def run_sim_surrogate(params_origin, output_origin, params_surrogate):
+    """Train the ML surrogate and evaluate model output with the ML surrogate.
+
+    :param params_origin: The parameter data used by the original model.
+    :param output_origin: The output data produced by the original model.
+    :param params_surrogate: The parameter data to be used by the ML surrogate.
+    """
+    # expend the parameter and output data
+    my_config['param_data'] = np.vstack([my_config['param_data'], params_origin])
+    my_config['output_data'] = np.vstack([my_config['output_data'], output_origin])
+
+    model, train_stats, config = hyper_train()
+
+    # run the surrogate for the second half of the samples
+    preprocessor_lstm = preprocessor.PreprocessorLSTM.from_dict(config)
+    data_inputs = preprocessor_lstm.prepare_input_data(params_surrogate)
+    # make predictions with the trained model
+    output_surrogate = predict_rnn.predict_batch(model, data_inputs, train_stats, config,
+                                                 batch_size=params_surrogate.shape[0])
+    # converting the predictions to GL format (temporal dimension at the end)
+    output_surrogate = np.moveaxis(output_surrogate, 1, -1)
+
+    return output_surrogate
+
+
+# 3. Define the callback function using the ML surrogate
+def run_sim_mixed(calib):
+    """This is the callback function that runs different realizations of the same model.
+
+    :param calib: The calibration object.
+    """
+    # if first iteration, run the original function
+    if calib.curr_iter == 0:
+        sim_data = run_sim_original(calib.system.ctrl_data, calib.system.param_data)
+        calib.system.set_sim_data(sim_data)
+        my_config['input_data'] = calib.system.ctrl_data
+        my_config['param_data'] = calib.system.param_data
+        my_config['output_data'] = calib.system.sim_data
+    else:
+        # split samples into two subsets to be used with the original function and the ML surrogate
+        np.random.seed()
+        ids = np.random.permutation(len(calib.system.param_data))
+        split_index = int(len(ids) * 0.5)
+        ids_origin, ids_surrogate = ids[:split_index], ids[split_index:]
+
+        # run the original function for the first half of the samples
+        param_data_origin = calib.system.param_data[ids_origin]
+        sim_data_origin = run_sim_original(calib.system.ctrl_data, param_data_origin)
+
+        # run the surrogate for the second half of the samples
+        param_data_surrogate = calib.system.param_data[ids_surrogate]
+        sim_data_surrogate = run_sim_surrogate(param_data_origin, sim_data_origin, param_data_surrogate)
+
+        # put the two subsets of simulation data together according to the original order
+        sim_data = np.zeros([calib.system.num_samples, calib.system.num_obs, calib.system.num_steps])
+        sim_data[ids_surrogate] = sim_data_surrogate
+        sim_data[ids_origin] = sim_data_origin
+
+        # set `sim_data` to system
+        calib.system.set_sim_data(sim_data)
+
+
+calibration = BayesianCalibration.from_dict(
+    {
+        "num_iter": 5,
+        "callback": run_sim_mixed,
+        "system": {
+            "param_min": [0.1, 0.1],
+            "param_max": [1, 10],
+            "param_names": ['a', 'b'],
+            "num_samples": 20,
+            "obs_names": ['f'],
+            "ctrl_name": 'u',
+            "obs_data": y_obs,
+            "ctrl_data": x_obs,
+            "sim_name": 'hyperbola',
+            "sigma_tol": 0.01,
+        },
+        "calibration": {
+            "inference": {
+                "ess_target": 0.3,
+                "scale_cov_with_max": True,
+            },
+            "sampling": {
+                "max_num_components": 1,
+                "n_init": 1,
+                "random_state": 0,
+                "slice_sampling": True,
+            },
+            "initial_sampling": "halton",
+        },
+        "save_fig": 0,
+    }
+)
+
+calibration.run()
diff --git a/tutorials/data_driven/LSTM/train_predict_rnn_hypertuning.py b/tutorials/data_driven/LSTM/train_predict_rnn_hypertuning.py
index 97c4e2c3..105d7254 100644
--- a/tutorials/data_driven/LSTM/train_predict_rnn_hypertuning.py
+++ b/tutorials/data_driven/LSTM/train_predict_rnn_hypertuning.py
@@ -12,7 +12,7 @@ def my_training_function():
     with wandb.init():
         my_config['window_size'] = wandb.config['window_size']
     preprocessor_TC = preprocessor.PreprocessorTriaxialCompression(**my_config)
-    train_rnn.train(preprocessor_TC)
+    train_rnn.train(preprocessor_TC, config=my_config)
 
 
 # 1. Create my dictionary of configuration

From 9e85527e2192ea2d80ae5ad0296a1d000aef2247 Mon Sep 17 00:00:00 2001
From: chyalexcheng <chyalexcheng@gmail.com>
Date: Tue, 20 Feb 2024 22:14:20 +0100
Subject: [PATCH 17/21] resolve issue #71 by excluding input, parameter, and
 output arrays from the sweep_config of wandb.sweep. However, somehow the
 arrays get turned into string when reloaded from best_run.config

---
 grainlearning/rnn/train.py                         |  7 ++++---
 .../hyperbola_calibration_mixed_hypertuning.py     | 14 +++++++++-----
 .../LSTM/train_predict_rnn_hypertuning.py          |  2 +-
 3 files changed, 14 insertions(+), 9 deletions(-)

diff --git a/grainlearning/rnn/train.py b/grainlearning/rnn/train.py
index f50afb68..bbe57ca2 100644
--- a/grainlearning/rnn/train.py
+++ b/grainlearning/rnn/train.py
@@ -281,8 +281,7 @@ class HyperTuning:
     Class to run hyperparameter tuning with wandb.
     """
 
-    def __init__(self, sweep_config, search_space, other_config=None, entity_name='grainlearning',
-                 project_name='my_sweep'):
+    def __init__(self, sweep_config, search_space, other_config=None, entity_name='', project_name='my_sweep'):
         """
         :param sweep_config: Dictionary containing the configuration of the sweep.
 
@@ -343,7 +342,9 @@ def get_sweep_id(self):
         # add default parameters from my_config into sweep_config, use 'values' as the key
         self.sweep_config['parameters'] = {}
         for key, value in self.other_config.items():
-            self.sweep_config['parameters'].update({key: {'values': [value]}})
+            # if key is not listed in the following list which should be excluded
+            if key not in ['input_data', 'param_data', 'output_data']:
+                self.sweep_config['parameters'].update({key: {'values': [value]}})
         # update sweep_config with the parameters to be searched and their distributions
         self.sweep_config['parameters'].update(self.search_space)
         # create the sweep
diff --git a/tutorials/data_driven/LSTM/hyperbola_calibration_mixed_hypertuning.py b/tutorials/data_driven/LSTM/hyperbola_calibration_mixed_hypertuning.py
index 12b6edbf..266a240f 100644
--- a/tutorials/data_driven/LSTM/hyperbola_calibration_mixed_hypertuning.py
+++ b/tutorials/data_driven/LSTM/hyperbola_calibration_mixed_hypertuning.py
@@ -23,8 +23,8 @@ def nonlinear(x, params):
     'val_frac': 0.2,
     'window_size': 10,
     'window_step': 1,
-    'patience': 25,
-    'epochs': 20,
+    'patience': 5,
+    'epochs': 10,
     'learning_rate': 1e-4,
     'lstm_units': 128,
     'dense_units': 128,
@@ -83,14 +83,18 @@ def run_sim_original(x, params):
 
 
 def my_training_function():
+    # update window_size of my_config from wandb
+    with wandb.init():
+        my_config['window_size'] = wandb.config['window_size']
     preprocessor_lstm = preprocessor.PreprocessorLSTM.from_dict(my_config)
     train_rnn.train(preprocessor_lstm, config=my_config)
 
 
 def hyper_train():
     """Train the ML surrogate using hyperparameter tuning."""
-    hyper_tuner = train_rnn.HyperTuning(sweep_config, search_space, my_config, project_name='hyperbola_sweep')
-    hyper_tuner.run_sweep(my_training_function, count=10)
+    hyper_tuner = train_rnn.HyperTuning(sweep_config, search_space, my_config, entity_name='grainlearning',
+                                        project_name='hyperbola_sweep')
+    hyper_tuner.run_sweep(my_training_function, count=1)
 
     hyper_tuner.sweep_config['metric']['name'] = 'val_loss'
     entity_project_sweep_id = f"{hyper_tuner.entity_name}/{hyper_tuner.project_name}/{hyper_tuner.sweep_id}"
@@ -189,7 +193,7 @@ def run_sim_mixed(calib):
             },
             "initial_sampling": "halton",
         },
-        "save_fig": 0,
+        "save_fig": -1,
     }
 )
 
diff --git a/tutorials/data_driven/LSTM/train_predict_rnn_hypertuning.py b/tutorials/data_driven/LSTM/train_predict_rnn_hypertuning.py
index 105d7254..1f602c9f 100644
--- a/tutorials/data_driven/LSTM/train_predict_rnn_hypertuning.py
+++ b/tutorials/data_driven/LSTM/train_predict_rnn_hypertuning.py
@@ -86,7 +86,7 @@ def my_training_function():
 }
 
 # 3. Run the sweep
-hyper_tuner = HyperTuning(sweep_config, search_space, my_config, project_name='my_sweep')
+hyper_tuner = HyperTuning(sweep_config, search_space, my_config, entity_name='grainlearning', project_name='my_sweep')
 hyper_tuner.run_sweep(my_training_function, count=100)
 
 # 4. Get the best model

From 6710fe8c25c37e360c25261fdc2e65b7216e356c Mon Sep 17 00:00:00 2001
From: chyalexcheng <chyalexcheng@gmail.com>
Date: Fri, 2 Aug 2024 03:51:53 +0200
Subject: [PATCH 18/21] add tutorials for geomechanical models: norsand and
 dike

---
 grainlearning/tools.py                        |   38 +-
 .../LSTM/norsand_triax_calibration_mixed.py   |  266 +++
 .../2D_dike_stability/2D_dike.py              | 1732 +++++++++++++++++
 .../2D_dike_model_inference.py                |  121 ++
 .../nosand_triax/norsand_triax_calibration.py |  172 ++
 5 files changed, 2316 insertions(+), 13 deletions(-)
 create mode 100644 tutorials/data_driven/LSTM/norsand_triax_calibration_mixed.py
 create mode 100644 tutorials/physics_based/2D_dike_stability/2D_dike.py
 create mode 100644 tutorials/physics_based/2D_dike_stability/2D_dike_model_inference.py
 create mode 100644 tutorials/physics_based/nosand_triax/norsand_triax_calibration.py

diff --git a/grainlearning/tools.py b/grainlearning/tools.py
index a5b477b4..226ad4f1 100644
--- a/grainlearning/tools.py
+++ b/grainlearning/tools.py
@@ -411,7 +411,8 @@ def plot_param_stats(fig_name, param_names, means, covs, save_fig=0):
         plt.xlabel("'Time' step")
         plt.ylabel(f'Mean of {param_names[i]}')
         plt.grid(True)
-    plt.tight_layout()
+    mng = plt.get_current_fig_manager()
+    mng.full_screen_toggle()
     if save_fig:
         plt.savefig(f'{fig_name}_param_means.png')
 
@@ -422,7 +423,8 @@ def plot_param_stats(fig_name, param_names, means, covs, save_fig=0):
         plt.xlabel("'Time' step")
         plt.ylabel(f'Coefficient of variation of {param_names[i]}')
         plt.grid(True)
-    plt.tight_layout()
+    mng = plt.get_current_fig_manager()
+    mng.full_screen_toggle()
     if save_fig:
         plt.savefig(f'{fig_name}_param_covs.png')
 
@@ -451,7 +453,8 @@ def plot_posterior(fig_name, param_names, param_data, posterior, save_fig=0):
             plt.xlabel(r'$' + name + '$')
             plt.ylabel('Posterior probability mass')
             plt.grid(True)
-        plt.tight_layout()
+        mng = plt.get_current_fig_manager()
+        mng.full_screen_toggle()
         if save_fig:
             plt.savefig(f'{fig_name}_posterior_{name}.png')
 
@@ -473,8 +476,8 @@ def plot_param_data(fig_name, param_names, param_data_list, save_fig=0):
             plt.xlabel(r'$' + param_names[j] + '$')
             plt.ylabel(r'$' + param_names[j + 1] + '$')
             plt.legend()
-        plt.legend()
-        plt.tight_layout()
+    mng = plt.get_current_fig_manager()
+    mng.full_screen_toggle()
     if save_fig:
         plt.savefig(f'{fig_name}_param_space.png')
 
@@ -488,7 +491,7 @@ def plot_obs_and_sim(fig_name, ctrl_name, obs_names, ctrl_data, obs_data, sim_da
     :param ctrl_data: ndarray
     :param obs_data: ndarray
     :param sim_data: ndarray
-    :param posterior: ndarray
+    :param posteriors: ndarray
     :param save_fig: bool defaults to False
     """
     import matplotlib.pylab as plt
@@ -512,10 +515,7 @@ def plot_obs_and_sim(fig_name, ctrl_name, obs_names, ctrl_data, obs_data, sim_da
         for j in (-posteriors[-1, :]).argsort()[:3]:
             plt.plot(ctrl_data, sim_data[j, i, :], label=f'sim No. {j:d}')
 
-        if len(ctrl_data) < 20:
-            markevery = 1
-        else:
-            markevery = int(len(ctrl_data) / 10.)
+        markevery = int(np.ceil(len(ctrl_data) / 100))
 
         plt.plot(ctrl_data,
                  obs_data[i, :], 'ok',
@@ -525,15 +525,18 @@ def plot_obs_and_sim(fig_name, ctrl_name, obs_names, ctrl_data, obs_data, sim_da
 
         plt.xlabel(ctrl_name)
         plt.ylabel(obs_names[i])
-        plt.legend()
+        # Show legend in the first sub-plot
+        if i == 0:
+            plt.legend()
         plt.grid(True)
 
-    plt.tight_layout()
+    mng = plt.get_current_fig_manager()
+    mng.full_screen_toggle()
     if save_fig:
         plt.savefig(f'{fig_name}_obs_and_sim.png')
 
 
-def plot_pdf(fig_name, param_names, samples, save_fig=0):
+def plot_pdf(fig_name, param_names, samples, save_fig=0, true_params=None):
     """
     Plot the posterior density function of the parameter distribution
     :param fig_name: string
@@ -541,6 +544,7 @@ def plot_pdf(fig_name, param_names, samples, save_fig=0):
     :param param_names: list of strings containing parameter names
     :param samples: list of ndarray of shape (num_samples, self.num_params)
     :param save_fig: bool defaults to False
+    :param true_params: list of true parameter values, default=None
     """
     import seaborn as sns
     import pandas as pd
@@ -569,6 +573,14 @@ def plot_pdf(fig_name, param_names, samples, save_fig=0):
     fig.map_diag(sns.kdeplot, lw=2)
     fig.add_legend(loc='upper right')
     fig.fig.canvas.manager.set_window_title('Estimated posterior probability density function')
+    if true_params is not None:
+        for i in range(len(param_names)):
+            for j in range(len(param_names)):
+                # plot the true parameter values as a yellow start
+                if i != j:
+                    # keep the marker on top of other plots
+                    fig.axes[i, j].scatter(true_params[j], true_params[i], color='gold', edgecolors='k', marker='*',
+                                           s=200, zorder=10)
 
     fig.tight_layout()
     if save_fig:
diff --git a/tutorials/data_driven/LSTM/norsand_triax_calibration_mixed.py b/tutorials/data_driven/LSTM/norsand_triax_calibration_mixed.py
new file mode 100644
index 00000000..fa7510dc
--- /dev/null
+++ b/tutorials/data_driven/LSTM/norsand_triax_calibration_mixed.py
@@ -0,0 +1,266 @@
+import grainlearning.rnn.train as train_rnn
+import grainlearning.rnn.predict as predict_rnn
+from grainlearning.rnn import preprocessor
+import numpy as np
+from grainlearning import BayesianCalibration
+from math import log
+
+
+def norsand(x, params):
+    # norsand parameters
+    gamma = params[0]  # [0.9 - 1.4] 'altitude of CSL @ 1kPa'
+    lambda_val = params[1]  # [0.01 - 0.07] 'slope CSL defined on natural log'
+    M_tc = params[2]  # [1.2 - 1.5] 'critical state friction ratio triaxial compression'
+    N = params[3]  # [0.2 - 0.5] 'volumetric coupling coefficient'
+    H = params[4]  # [25 - 500] 'plastic hardening modulus for loading'
+    Xim_tc = params[5]  # [2 - 5] 'relates maximum dilatancy to state variable (psi)'
+    Ir = params[6]  # [100 - 600] 'shear rigidity'
+    nu = params[7]  # [0.1 0.3] 'Poissons ratio'
+
+    # Derived parameters
+    eqp_inc = eqp_tot / (1e2 * lst)  # [-] 'increment applied plastic deviatoric strain'
+    ratio = 1 / (1 - N) ** ((N - 1) / N)  # [-] 'ratio mean effective stress (p) and image stress (pim)'
+    Xim = Xim_tc / (
+        1 - Xim_tc * lambda_val / M_tc)  # [-] 'relationship between soil property Xi_tc to Norsand property Xi'
+
+    # Declarations
+    p = np.zeros(lst + 1)
+    q = np.zeros(lst + 1)
+    pim = np.zeros(lst + 1)
+    ev = np.zeros(lst + 1)
+    eq = np.zeros(lst + 1)
+    e = np.zeros(lst + 1)
+    psi = np.zeros(lst + 1)
+
+    # Initial conditions
+    p[0] = p0
+    pim[0] = ratio * p0
+    e[0] = e0
+    psi[0] = e0 - (gamma - lambda_val * np.log(pim[0])) + lambda_val * np.log(pim[0] / p[0])
+
+    # Loadstep cycle CD test
+    for i in range(lst):
+        # Update image state
+        e[i + 1] = e0 - (1 + e0) * ev[i]
+        psi[i + 1] = e[i + 1] - (gamma - lambda_val * np.log(pim[i])) + lambda_val * np.log(pim[i] / p[i])
+
+        # Update image friction ratio
+        Mim = M_tc + Xim * N * np.abs(psi[i + 1])
+
+        # Apply hardening increment
+        pim_max = p[i] * np.exp(-Xim * psi[i + 1] / M_tc)
+        pim_inc = H * (pim_max - pim[i]) * eqp_inc
+        pim[i + 1] = pim[i] + pim_inc
+
+        # Calculate plastic volumetric strain increment
+        Dp = Mim - q[i] / p[i]
+        evp_inc = Dp * eqp_inc
+
+        # Calculate bulk and shear modulus
+        mu = Ir * p[i]
+        K = mu * (2 * (1 + nu)) / (3 * (1 - 2 * nu))
+
+        # Apply consistency condition to calculate stress ratio increment (drained)
+        eta_inc = (pim_inc / pim[i]) / (
+            1 / M_tc + 1 / (3.0 - q[i] / p[i]))  # Note: 3.0 is used instead of 3 in the denominator
+        # Calculate mean effective stress
+        p_inc = p[i] * eta_inc / (3.0 - q[i] / p[i])
+        p[i + 1] = p[i] + p_inc
+
+        # Calculate new stress ratio and shear stress
+        eta = M_tc * (1 + np.log(pim[i + 1] / p[i + 1]))
+        q[i + 1] = eta * p[i + 1]
+
+        # Update volumetric and deviatoric strain increments
+        eve_inc = p_inc / K
+        eqe_inc = (q[i + 1] - q[i]) / (3 * mu)
+        ev[i + 1] = ev[i] + eve_inc + evp_inc
+        eq[i + 1] = eq[i] + eqe_inc + eqp_inc
+
+    return np.stack([(q / p)[::int(lst / 100)], e[::int(lst / 100)]], axis=0)
+
+
+# CD input
+p0 = 130.0  # [kPa] 'inital pressure'
+e0 = 0.60  # [-] 'init void ratio'
+eqp_tot = 15.0  # [%] 'total applied plastic deviatoric strain'
+lst = int(1e3)  # [-] 'loadsteps'
+
+# ground truth parameter
+gamma = 0.816  # [0.7 - 1.4] 'altitude of CSL @ 1kPa'
+lambda_val = 0.015  # [0.01 - 0.07] 'slope CSL defined on natural log'
+M_tc = 1.26  # [1.2 - 1.5] 'critical state friction ratio triaxial compression'
+N = 0.4  # [0.2 - 0.5] 'volumetric coupling coefficient'
+H = 200.0  # [25 - 500] 'plastic hardening modulus for loading'
+Xim_tc = 3.8  # [2 - 5] 'relates maximum dilatancy to state variable (psi)'
+Ir = 200.0  # [100 - 600] 'shear rigidity'
+nu = 0.20  # [0.1 0.3] 'Poissons ratio'
+
+# generate synthetic data
+x_obs = np.arange(lst + 1) / lst * eqp_tot
+x_obs = x_obs[::int(lst / 100)]
+y_obs = norsand(x_obs, [gamma, lambda_val, M_tc, N, H, Xim_tc, Ir, nu])
+
+# Define the configuration dictionary for the ML surrogate
+my_config = {
+    'input_data': None,
+    'param_data': None,
+    'output_data': None,
+    'train_frac': 0.7,
+    'val_frac': 0.2,
+    'window_size': 10,
+    'window_step': 1,
+    'patience': 25,
+    'epochs': 100,
+    'learning_rate': 1e-4,
+    'lstm_units': 128,
+    'dense_units': 128,
+    'batch_size': 64,
+    'standardize_outputs': True,
+    'save_weights_only': True
+}
+
+
+def run_sim_original(x, params):
+    """Run different realizations of the original model.
+
+    :param x: the input sequence
+    :param params: the parameters
+    """
+    data = []
+    for params in params:
+        # Run the model
+        y = norsand(x, params)
+        data.append(np.array(y, ndmin=2))
+    return np.array(data)
+
+
+def run_sim(calib):
+    """This is the callback function that runs different realizations of the same model.
+
+    :param calib: The calibration object.
+    """
+    sim_data = run_sim_original(calib.system.ctrl_data, calib.system.param_data)
+    calib.system.set_sim_data(sim_data)
+
+
+def run_sim_surrogate(params_origin, output_origin, params_surrogate):
+    """Train the ML surrogate and evaluate model output with the ML surrogate.
+
+    :param params_origin: The parameter data used by the original model.
+    :param output_origin: The output data produced by the original model.
+    :param params_surrogate: The parameter data to be used by the ML surrogate.
+    """
+    # expend the parameter and output data
+    my_config['param_data'] = np.vstack([my_config['param_data'], params_origin])
+    my_config['output_data'] = np.vstack([my_config['output_data'], output_origin])
+
+    preprocessor_lstm = preprocessor.PreprocessorLSTM.from_dict(my_config)
+    _ = train_rnn.train_without_wandb(preprocessor_lstm, model=calibration.model, config=my_config)
+    model, train_stats, config = predict_rnn.get_pretrained_model('outputs')
+    calibration.model = model
+
+    # run the surrogate for the second half of the samples
+    data_inputs = preprocessor_lstm.prepare_input_data(params_surrogate)
+    # make predictions with the trained model
+    output_surrogate = predict_rnn.predict_batch(model, data_inputs, train_stats, config,
+                                                 batch_size=params_surrogate.shape[0])
+    # converting the predictions to GL format (temporal dimension at the end)
+    output_surrogate = np.moveaxis(output_surrogate, 1, -1)
+
+    return output_surrogate
+
+
+# 3. Define the callback function using the ML surrogate
+def run_sim_mixed(calib):
+    """This is the callback function that runs different realizations of the same model.
+
+    :param calib: The calibration object.
+    """
+    calib.system.num_steps = 101
+    # if first iteration, run the original function
+    if calib.curr_iter == 0:
+        sim_data = run_sim_original(calib.system.ctrl_data, calib.system.param_data)
+        calib.system.set_sim_data(sim_data)
+        my_config['input_data'] = calib.system.ctrl_data
+        my_config['param_data'] = calib.system.param_data
+        my_config['output_data'] = calib.system.sim_data
+        calibration.model = None
+    else:
+        # split samples into two subsets to be used with the original function and the ML surrogate
+        np.random.seed()
+        ids = np.random.permutation(len(calib.system.param_data))
+        split_index = int(len(ids) * 0.5)
+        ids_origin, ids_surrogate = ids[:split_index], ids[split_index:]
+        calib.ids_origin, calib.ids_surrogate = ids_origin, ids_surrogate
+
+        # run the original function for the first half of the samples
+        param_data_origin = calib.system.param_data[ids_origin]
+        sim_data_origin = run_sim_original(calib.system.ctrl_data, param_data_origin)
+
+        # run the surrogate for the second half of the samples
+        param_data_surrogate = calib.system.param_data[ids_surrogate]
+        sim_data_surrogate = run_sim_surrogate(param_data_origin, sim_data_origin, param_data_surrogate)
+
+        # put the two subsets of simulation data together according to the original order
+        sim_data = np.zeros(
+            [calib.system.num_samples, calib.system.num_obs, calib.system.num_steps - my_config['window_size']])
+        sim_data[ids_surrogate] = sim_data_surrogate
+        sim_data[ids_origin] = sim_data_origin[:, :, :-my_config['window_size']]
+
+        # set `sim_data` to system
+        calib.system.set_sim_data(sim_data)
+        calib.system.num_steps = sim_data.shape[2]
+
+
+num_cores = 18
+param_names = ['gamma', 'lambda', 'M', 'N', 'H', 'Xim', 'Ir', 'nu']
+num_samples = int(10 * len(param_names) * log(len(param_names)))
+
+calibration = BayesianCalibration.from_dict(
+    {
+        "num_iter": 10,
+        "callback": run_sim_mixed,
+        "system": {
+            "param_min": [0.7, 0.01, 1.2, 0.2, 25, 2, 100, 0.1],
+            "param_max": [1.4, 0.07, 1.5, 0.5, 500, 5, 600, 0.3],
+            "param_names": param_names,
+            "num_samples": num_samples,
+            "obs_names": ['q/p', 'e'],
+            "ctrl_name": 'u',
+            "obs_data": y_obs,
+            "ctrl_data": x_obs,
+            "sim_name": 'triax',
+            "sigma_tol": 0.01,
+        },
+        "calibration": {
+            "inference": {
+                "ess_target": 0.3,
+                "scale_cov_with_max": True,
+            },
+            "sampling": {
+                "max_num_components": 10,
+                "n_init": 1,
+                "random_state": 0,
+                "slice_sampling": False,
+            },
+            "initial_sampling": "halton",
+        },
+        "save_fig": -1,
+    }
+)
+
+calibration.run()
+
+import matplotlib.pylab as plt
+from grainlearning.tools import plot_posterior, plot_param_data, plot_pdf
+
+plot_posterior('test', param_names, calibration.system.param_data[calibration.ids_surrogate],
+               calibration.calibration.inference.posteriors[:, calibration.ids_surrogate])
+plot_posterior('test', param_names, calibration.system.param_data[calibration.ids_origin],
+               calibration.calibration.inference.posteriors[:, calibration.ids_origin])
+plt.show()
+
+plot_pdf('test', param_names, [calibration.system.param_data[calibration.ids_surrogate],
+                               calibration.system.param_data[calibration.ids_origin]])
+plt.show()
diff --git a/tutorials/physics_based/2D_dike_stability/2D_dike.py b/tutorials/physics_based/2D_dike_stability/2D_dike.py
new file mode 100644
index 00000000..c28b6a66
--- /dev/null
+++ b/tutorials/physics_based/2D_dike_stability/2D_dike.py
@@ -0,0 +1,1732 @@
+import numpy as np
+import matplotlib.pyplot as plt
+import matplotlib.cm as cm
+import time
+import sys
+from scipy.interpolate import interp1d
+from scipy.linalg import cholesky
+from scipy.sparse.linalg import splu
+from scipy.sparse.linalg import spsolve
+from scipy.sparse import csc_matrix
+from scipy.sparse import csr_matrix
+from scipy.sparse import coo_matrix
+from multiprocessing import Pool
+
+
+# function list:
+
+# function 1: partitioning of the FEM domain
+# function 2: 8 node quadrilateral mesh generator
+# function 3: elevate mesh to form the free surface that was pre-defined
+# function 4: calculate the equivalent thickness of the pile
+# function 5: loop over surface's to find intersection points
+# function 6: find intersection point between the free surface and the water level
+# function 7: calculate cross product between two line in 2D
+# function 8: calculate additional stiffness due to anchor presence
+# function 9: find closest node on pile where anchor should be present
+# function 10: define element degree of freedom array
+# function 11: find dirichlet boundary node's
+# function 12: find Neumann boundary (due to water pressure)
+# function 13: find outward normal to a boundary edge
+# function 14: find integration point location
+# function 15: shape function 4 node quadrilateral
+# function 16: derivate shape function 4 node quadrilateral
+# function 17: shape function 8 node quadrilateral
+# function 18: derivate shape function 8 node quadrilateral
+# function 19: local coordinates integration points
+# function 20: generate random field, (Gaussian correlation fuction)
+# function 21: set Youngs modulus and Poisson's ratio throughout FEM domain
+# function 22: find interface integration points
+# function 23: find pile integration points
+# function 24: form global derivative matrix
+# function 25: calculate Lamé parameters
+# function 26: calculate strain-displacement matrix
+# function 27: form displacement stiffness matrix
+# function 28: set specific weight throughout FEM domain
+# function 29: form body load vector
+# function 30: form init head based on water surface wl
+# function 31: random field instance for friction angle, cohesion and hydraulic conductivity
+# function 32: log normal distribution
+# function 33: set hydraulic conductivity throughout FEM domain
+# function 34: set plasticity parameters throughout FEM domain
+# function 35: form groundwater stiffness matrix
+# function 36: FEM solver for groundwater flow
+# function 37: calculate Darcy flow
+# function 38: hydromechanical coupling porepressure to total stress
+# function 39: vectorized version of nodal coordinate array
+# function 40: vectorized version of array declarations displacement solver array
+# function 41: vectorized version of displacement vector array
+# function 42: vectorized version of strain calculation
+# function 43: vectorized version of stress calculation
+# function 44: vectorized version of indexing array
+# function 45: shear strength reduction routine
+# function 46: vectorized version stress invariant routine
+# function 47: vectorized version of Mohr Coulomb Yield function & derivatives
+# function 48: vectorized version of the global derivative array
+# function 49: select pile elements routine
+# function 50: vectorized solver for the displacement and stability calculation
+# function 51: find safety factor by executing displacement solvers with several factors of safety
+# function 52: display Monte Carlo results
+# function 53: display single run results
+# function 54: single Monte Carlo iteration (groundwater-displacement-factor of safety)
+# function 55: parallel task Monte Carlo Iterations
+# function 56: parallel execution Monte Carlo Iterations
+# function 57: non-vectorized solver for the displacement and stability calculation
+
+# function 1: partitioning of the FEM domain
+def divide_FEM_domain(free, disc, pile, equivalent_pile_thickness, intersection_point_list, active_pile):
+    # Calculate discretization in the y-direction
+    ey = np.max(free[:, 1]) / disc[1]
+    # Determine the Y-coordinates
+    if active_pile:
+        # If the pile is active, include specific ratio related to the pile position
+        Y = np.unique(np.concatenate(
+            (np.linspace(0, 1.0, int(ey) + 1), [(free[(pile[0] - 1), 1] - pile[1]) / free[(pile[0] - 1), 1]])))
+    else:
+        # If the pile is not active, use a linearly spaced array
+        Y = np.linspace(0, 1.0, int(ey) + 1)
+    # Initialize X as the x-coordinates in `free`
+    X = np.array(free[:, 0])
+    # Iteratively add x-coordinates, dividing each segment into equally spaced intervals
+    for i in range(free.shape[0] - 1):
+        L = free[i + 1, 0] - free[i, 0]
+        ex = int(np.ceil(L / disc[0]))
+        dx = L / ex
+        for j in range(ex):
+            X = np.append(X, free[i, 0] + j * dx)
+    # Add additional point for active pile considering equivalent pile thickness
+    if active_pile:
+        X = np.append(X, free[(pile[0] - 1), 0] + equivalent_pile_thickness)
+    # Include intersection points if they exist
+    if len(intersection_point_list) != 0:
+        X = np.concatenate([X, intersection_point_list[:, 0]])
+    # Ensure X contains only unique values
+    X = np.unique(X)
+    # Return the discretized coordinates in the FEM domain
+    return X, Y
+
+
+# function 2: 8 node quadrilateral mesh generator
+def mesh_generator(X, Y):
+    # Calculate the number of elements along the x and y axes
+    ex = len(X) - 1
+    ey = len(Y) - 1
+    # Calculate the total number of elements and nodes
+    ec = ex * ey
+    nx = ex + 1
+    ny = ey + 1
+    # Calculate the total number of nodes with 8 node quadrilateral mesh
+    nc = (2 * ex + 1) * (2 * ey + 1) - ec
+    # Initialize arrays for elements and nodes
+    el = np.zeros((ec, 8), dtype=int)
+    no = np.zeros((nc, 2))
+    # Loop through each element to define node positions and element connectivity
+    for i in range(ey):
+        for j in range(ex):
+            # Calculate the index of the current element
+            el_index = i * ex + j
+            # Define the node indices for the current element
+            el[el_index, :] = [
+                i * nx + j,
+                (i + 1) * nx + j,
+                (i + 1) * nx + j + 1,
+                i * nx + j + 1,
+                nx * ny + ex + i * (nx + ex) + j,
+                nx * ny + (i + 1) * (nx + ex) + j,
+                nx * ny + ex + i * (nx + ex) + j + 1,
+                nx * ny + i * (nx + ex) + j
+            ]
+            # Define the x and y coordinates for each node in the current element
+            el_nodes = el[el_index, :]
+            no[el_nodes, 0] = [X[j], X[j], X[j + 1], X[j + 1], X[j], (X[j] + X[j + 1]) / 2, X[j + 1],
+                               (X[j] + X[j + 1]) / 2]
+            no[el_nodes, 1] = [Y[i], Y[i + 1], Y[i + 1], Y[i], (Y[i] + Y[i + 1]) / 2, Y[i + 1], (Y[i] + Y[i + 1]) / 2,
+                               Y[i]]
+    # Return the node positions, the number of corners, the elements, and the number of elements
+    return no, nc, el, ec
+
+
+# function 3: elevate mesh to form the free surface that was pre-defined
+def elevate_mesh(free, X, no, el):
+    # Calculate the number of nodes of the 4 node quadrilateral mesh
+    nc = np.max(el[:, :4]) + 1
+    nx = len(X)
+    ex = nx - 1
+    # Interpolate the free surface using the interp1d function
+    interp_function = interp1d(free[:, 0], free[:, 1], fill_value="extrapolate")
+    free_surface = interp_function(X)
+    # Calculate the midpoints of the free surface segments
+    free_surface = np.concatenate((free_surface, (free_surface[:-1] + free_surface[1:]) / 2))
+    # Adjust the y-coordinates of nodes based on the free surface elevation (4 node quadrilateral)
+    for i in range(nx):
+        no[i:nc:nx, 1] *= free_surface[i]
+    # Adjust left and right mid node (8 node quadrilateral)
+    for i in range(nx):
+        no[nc + ex + i:no.shape[0]: ex + nx, 1] *= free_surface[i]
+    # Adjust up and down node (8 node quadrilateral)
+    for i in range(ex):
+        no[nc + i:no.shape[0]:ex + nx, 1] *= free_surface[nx + i]
+    # Identify nodes that are on the free surface
+    free_surface_node = np.concatenate([np.arange(nc - ex - 1, nc), np.arange(no.shape[0] - ex, no.shape[0])])
+    id_sorted = np.argsort(no[free_surface_node, 0])
+    free_surface_node = free_surface_node[id_sorted]
+    # Return the modified node positions and indices of nodes on the free surface
+    return no, free_surface_node
+
+
+# function 4: calculate the equivalent thickness of the pile
+def equivalent_thickness_rectangle(I):
+    equivalent_pile_thickness = (12 * I) ** (1 / 3)
+    return equivalent_pile_thickness
+
+
+# function 5: loop over surface's to find intersection points
+def find_intersection_points(free, wl):
+    # Initialize a list to store intersection points
+    intersection_point_list = []
+    # Iterate through each segment in 'free'
+    for i in range(free.shape[0] - 1):
+        # Iterate through each segment in 'wl' (water line)
+        for j in range(wl.shape[0] - 1):
+            # Check for intersection between a segment in 'free' and a segment in 'wl'
+            intersection_point = cross_check(free[i, :], free[i + 1, :], wl[j, :], wl[j + 1, :])
+            # If an intersection point is found, add it to the list
+            if intersection_point is not None:
+                intersection_point_list.append(intersection_point)
+    # Return the list of intersection points as a numpy array
+    return np.array(intersection_point_list)
+
+
+# function 6: find intersection point between the free surface and the water level
+def cross_check(p1, p2, p3, p4):
+    # Calculate the vectors for the two line segments
+    v1 = p2 - p1
+    v2 = p4 - p3
+    # Calculate the denominator using a 2D cross product function
+    denom = cross_product_2d(v1, v2)
+    # If the denominator is zero, the lines are parallel or coincident, so return None
+    if denom == 0:
+        return None
+    # Calculate a vector from one line's start point to the other's
+    v3 = p3 - p1
+    # Calculate the intersection parameters for both lines
+    t1 = cross_product_2d(v3, v2) / denom
+    t2 = cross_product_2d(v3, v1) / denom
+    # Check if the intersection point lies within both line segments
+    if 0 < t1 < 1 and 0 < t2 < 1:
+        # Calculate the intersection point
+        intersection_point = p1 + t1 * v1
+        return intersection_point
+    else:
+        # If there is no valid intersection within both line segments, return None
+        return None
+
+
+# function 7: calculate cross product between two line in 2D
+def cross_product_2d(v1, v2):
+    crossp = v1[0] * v2[1] - v1[1] * v2[0]
+    return crossp
+
+
+# function 8: calculate additional stiffness due to anchor presence
+def anchor_spring(free, pile, anchor, E_anchor, d_anchor, HoH_anchor, no, active_anchor):
+    # Check if the anchor is active
+    if active_anchor:
+        # Select the appropriate anchor node based on the provided criteria
+        anchor_id = select_anchor_node(free, pile, anchor, no)
+        # Calculate the cross-sectional area of the anchor
+        A_anchor = np.pi * (d_anchor / 2) ** 2
+        # Compute the distances in x and y directions from the anchor node to the anchor point
+        dx = no[anchor_id, 0] - anchor[1]
+        dy = no[anchor_id, 1] - anchor[2]
+        # Calculate the length of the anchor
+        L_anchor = np.sqrt(dx ** 2 + dy ** 2)
+        # Determine the angle of the anchor with respect to the horizontal
+        angle_anchor = np.degrees(np.arcsin(dy / L_anchor))
+        # Calculate the equivalent spring constant for the anchor
+        equivalent_anchor_spring = E_anchor * A_anchor / (L_anchor * HoH_anchor)
+        # Decompose the spring constant into x and y components
+        anchor_spring_x = equivalent_anchor_spring * np.abs(np.cos(np.radians(angle_anchor)))
+        anchor_spring_y = equivalent_anchor_spring * np.abs(np.sin(np.radians(angle_anchor)))
+    else:
+        # If the anchor is not active, set all outputs to empty arrays
+        anchor_id = np.array([])
+        anchor_spring_x = np.array([])
+        anchor_spring_y = np.array([])
+
+    # Return the anchor node ID and the spring constants in the x and y directions
+    return anchor_id, anchor_spring_x, anchor_spring_y
+
+
+# function 9: find closest node on pile with anchor point
+def select_anchor_node(free, pile, anchor, no):
+    distances = np.sqrt((no[:, 0] - free[(pile[0] - 1), 0]) ** 2 + (no[:, 1] - anchor[0]) ** 2)
+    id = np.argmin(distances)
+    return id
+
+
+# function 10: define element degree of freedom array
+def degree_freedom(el, ec):
+    df = np.zeros((ec, 16), dtype=int)
+    for i in range(8):
+        df[:, i * 2] = el[:, i] * 2
+        df[:, i * 2 + 1] = el[:, i] * 2 + 1
+    return df
+
+
+# function 11: find dirichlet boundary node's
+def boundary_dirichlet(wl, X, no, nc, el, free_surface_node):
+    # Find node indices at the minimum and maximum X positions and at Y=0
+    x0 = np.where(no[:, 0] == np.min(X))[0]
+    x1 = np.where(no[:, 0] == np.max(X))[0]
+    y0 = np.where(no[:, 1] == 0.0)[0]
+    # Create an array of restricted degrees of freedom (Dirichlet boundary conditions)
+    rd = np.unique(np.concatenate([x0 * 2, x1 * 2, y0 * 2, y0 * 2 + 1]))
+    # Create an array of active degrees of freedom
+    nd = np.arange(1, 2 * nc)
+    ad = nd[~np.isin(nd, rd)]
+    # Calculate pressure at free surface nodes based on water line (wl) interpolation
+    free_surface_pressure_node = free_surface_node[::2]
+    interp_function = interp1d(wl[:, 0], wl[:, 1], fill_value="extrapolate")
+    fsh = interp_function(no[free_surface_pressure_node, 0])
+    # Determine fixed head nodes where wl > free and left and right boundary's
+    el_max = np.max(el[:, :4]) + 1
+    fh1 = free_surface_pressure_node[fsh >= no[free_surface_pressure_node, 1]]
+    fh2 = np.where((no[:el_max, 0] == np.min(X)) & (no[:el_max, 1] <= fsh[0]))[0]
+    fh3 = np.where((no[:el_max, 0] == np.max(X)) & (no[:el_max, 1] <= fsh[-1]))[0]
+    # Determine possible seepage head nodes where wl < free and left and right boundary's
+    fs1 = free_surface_pressure_node[fsh < no[free_surface_pressure_node, 1]]
+    fs2 = np.where((no[:el_max, 0] == np.min(X)) & (no[:el_max, 1] > fsh[0]))[0]
+    fs3 = np.where((no[:el_max, 0] == np.max(X)) & (no[:el_max, 1] > fsh[-1]))[0]
+    # Concatenate node indices for hydrostatic and free surface
+    fh = np.concatenate([fh1, fh2, fh3])
+    fs = np.concatenate([fs1, fs2, fs3])
+    # Return the arrays of active degrees of freedom, hydrostatic and free surface nodes
+    return ad, fh, fs
+
+
+# function 12: find Neumann boundary (due to water pressure)
+def boundary_neumann(wl, no, nc, free_surface_node):
+    # Interpolate the water level at the free surface nodes
+    interp_function = interp1d(wl[:, 0], wl[:, 1], fill_value="extrapolate")
+    free_surface_water_level = interp_function(no[free_surface_node, 0])
+    # Calculate the depth of water at each free surface node
+    water_depth = free_surface_water_level - no[free_surface_node, 1]
+    # Compute the mean water depth between pairs of nodes
+    mean_water_depth = (water_depth[::2][:-1] + water_depth[2::2]) / 2
+    # Compute the length and normal vectors at each free surface segment
+    dL, nxx, nyy = normal_boundary(no[free_surface_node[::2][:-1]], no[free_surface_node[2::2]])
+    # Identify indices for distributing water pressure on nodes
+    id1 = np.arange(0, len(water_depth) - 2, 2)
+    id2 = np.arange(1, len(water_depth) - 1, 2)
+    id3 = np.arange(2, len(water_depth), 2)
+    # Calculate water pressure based on mean water depth
+    water_pressure = 10.0 * mean_water_depth
+    water_pressure[water_pressure <= 0.0] = 0.0
+    # Initialize an array for nodal water pressure
+    nodal_water_pressure = np.zeros(2 * nc)
+    # Distribute the water pressure to the nodes of the 8 node quadrilaterals
+    nodal_water_pressure[free_surface_node[id1] * 2] -= water_pressure / 6.0 * dL * nxx
+    nodal_water_pressure[free_surface_node[id1] * 2 + 1] -= water_pressure / 6.0 * dL * nyy
+    nodal_water_pressure[free_surface_node[id2] * 2] -= 2 * water_pressure / 3.0 * dL * nxx
+    nodal_water_pressure[free_surface_node[id2] * 2 + 1] -= 2 * water_pressure / 3.0 * dL * nyy
+    nodal_water_pressure[free_surface_node[id3] * 2] -= water_pressure / 6.0 * dL * nxx
+    nodal_water_pressure[free_surface_node[id3] * 2 + 1] -= water_pressure / 6.0 * dL * nyy
+    # Return the calculated nodal water pressure
+    return nodal_water_pressure
+
+
+# function 13: find outward normal to a boundary edge
+def normal_boundary(n1, n2):
+    # Calculate the difference in x and y coordinates between pairs of nodes
+    dx = n2[:, 0] - n1[:, 0]
+    dy = n2[:, 1] - n1[:, 1]
+    # Compute the length of the segment between each pair of nodes
+    dL = np.sqrt(dx ** 2 + dy ** 2)
+    # Calculate the normal vector components for each segment
+    # The normal vector is perpendicular to the segment
+    nxx = -dy / dL
+    nyy = dx / dL
+    # Return the length of each segment and the normal vector components
+    return dL, nxx, nyy
+
+
+# function 14: find integration point location
+def integration_point_location(no, el, ec):
+    # Obtain integration points in local coordinates (xi, yi)
+    xi, yi = integration_point()
+    # Initialize an array to store the locations of integration points
+    lip = np.zeros((4 * ec, 2))
+    # Iterate over each element in the mesh
+    for i in range(ec):
+        # Get the coordinates of the nodes of the current element
+        co = no[el[i, :], :]
+        # Iterate over each integration point (4 per element in this case)
+        for j in range(4):
+            # Compute the shape tensor at each integration point
+            N = shape_tensor8(xi[j], yi[j])
+            # Calculate the global coordinates of the integration point
+            lip[i * 4 + j, :] = np.dot(N, co)
+    # Return the global coordinates of all integration points
+    return lip
+
+
+# function 15: shape function 4 node quadrilateral
+def shape_tensor4(xi, yi):
+    # Calculate the modified coordinates for interpolation
+    xim = 1 - xi
+    xip = 1 + xi
+    yim = 1 - yi
+    yip = 1 + yi
+    # Compute the shape functions for a 4-node quadrilateral element
+    N1 = 0.25 * xim * yim
+    N2 = 0.25 * xim * yip
+    N3 = 0.25 * xip * yip
+    N4 = 0.25 * xip * yim
+    # Combine the shape functions into an array
+    N = np.array([N1, N2, N3, N4])
+    # Return the array of shape functions
+    return N
+
+
+# function 16: derivate shape function 4 node quadrilateral
+def local_derivative4(xi, yi):
+    # Calculate the modified coordinates for differentiation
+    xim = 1 - xi
+    xip = 1 + xi
+    yim = 1 - yi
+    yip = 1 + yi
+    # Compute the derivatives of the shape functions with respect to xi
+    dN1_1 = -0.25 * yim
+    dN1_2 = -0.25 * yip
+    dN1_3 = 0.25 * yip
+    dN1_4 = 0.25 * yim
+    # Compute the derivatives of the shape functions with respect to yi
+    dN2_1 = -0.25 * xim
+    dN2_2 = 0.25 * xim
+    dN2_3 = 0.25 * xip
+    dN2_4 = -0.25 * xip
+    # Combine the derivatives into a 2x4 array
+    dN = np.array([[dN1_1, dN1_2, dN1_3, dN1_4],
+                   [dN2_1, dN2_2, dN2_3, dN2_4]])
+    # Return the array of shape function derivatives
+    return dN
+
+
+# function 17: shape function 8 node quadrilateral
+def shape_tensor8(xi, yi):
+    # Calculate the modified coordinates for interpolation
+    xim = 1 - xi
+    xip = 1 + xi
+    yim = 1 - yi
+    yip = 1 + yi
+    # Compute the shape functions for an 8-node quadrilateral element
+    N1 = 0.25 * xim * yim * (-xi - yi - 1)
+    N2 = 0.25 * xim * yip * (-xi + yi - 1)
+    N3 = 0.25 * xip * yip * (xi + yi - 1)
+    N4 = 0.25 * xip * yim * (xi - yi - 1)
+    N5 = 0.5 * xim * yim * yip
+    N6 = 0.5 * xim * xip * yip
+    N7 = 0.5 * xip * yim * yip
+    N8 = 0.5 * xim * xip * yim
+    # Combine the shape functions into an array
+    N = np.array([N1, N2, N3, N4, N5, N6, N7, N8])
+    # Return the array of shape functions
+    return N
+
+
+# function 18: derivate shape function 8 node quadrilateral
+def local_derivative8(xi, yi):
+    # Calculate the modified coordinates for differentiation
+    xim = 1 - xi
+    xip = 1 + xi
+    yim = 1 - yi
+    yip = 1 + yi
+    # Compute the derivatives of the shape functions with respect to xi
+    dN1_1 = 0.25 * yim * (2 * xi + yi)
+    dN1_2 = 0.25 * yip * (2 * xi - yi)
+    dN1_3 = 0.25 * yip * (2 * xi + yi)
+    dN1_4 = 0.25 * yim * (2 * xi - yi)
+    dN1_5 = -0.5 * yim * yip
+    dN1_6 = -xi * yip
+    dN1_7 = 0.5 * yim * yip
+    dN1_8 = -xi * yim
+    # Compute the derivatives of the shape functions with respect to yi
+    dN2_1 = 0.25 * xim * (2 * yi + xi)
+    dN2_2 = 0.25 * xim * (2 * yi - xi)
+    dN2_3 = 0.25 * xip * (xi + 2 * yi)
+    dN2_4 = 0.25 * xip * (2 * yi - xi)
+    dN2_5 = -xim * yi
+    dN2_6 = 0.5 * xim * xip
+    dN2_7 = -xip * yi
+    dN2_8 = -0.5 * xim * xip
+    # Return all the derivatives
+    return dN1_1, dN1_2, dN1_3, dN1_4, dN1_5, dN1_6, dN1_7, dN1_8, dN2_1, dN2_2, dN2_3, dN2_4, dN2_5, dN2_6, dN2_7, dN2_8
+
+
+# function 19: local coordinates integration points
+def integration_point():
+    factor = 1 / np.sqrt(3)
+    xi = factor * np.array([-1, 1, -1, 1])
+    yi = factor * np.array([1, 1, -1, -1])
+    return xi, yi
+
+
+# function 20: generate random field, (Gaussian correlation fuction)
+def random_field(x, y, clx, cly):
+    # This function calculates the covariance based on the distance between points
+    cfun = lambda x1, y1, x2, y2: np.exp(-np.sqrt(((x1 - x2) ** 2 / clx ** 2) + ((y1 - y2) ** 2 / cly ** 2)))
+    # Create a meshgrid of x and y coordinates for calculating the covariance matrix
+    X, Y = np.meshgrid(x, y)
+    # The covariance is calculated for every pair of points in the meshgrid
+    covariance_matrix = cfun(X, Y, X.T, Y.T)
+    # Perform the Cholesky decomposition of the covariance matrix
+    L = cholesky(covariance_matrix, lower=True)
+    # Return the Cholesky factor (lower triangular matrix)
+    return L
+
+
+# function 21: set Youngs modulus and Poisson's ratio throughout FEM domain
+def elasticity_parameters_integration_point(free, pile, inter, E_soil, E_pile, E_inter, nu_soil, nu_pile, nu_inter,
+                                            equivalent_pile_thickness, ec, lip, active_pile, active_inter):
+    # Initialize the elasticity modulus (E) and Poisson's ratio (nu) arrays for each integration point
+    E = np.full(4 * ec, E_soil)
+    nu = np.full(4 * ec, nu_soil)
+    # If interfaces are active, update the elasticity parameters at relevant integration points
+    if active_inter:
+        # Select integration points associated with the interface
+        id_inter = select_inter_integration_points(free, inter, lip)
+        # Update the elasticity modulus and Poisson's ratio at these points to interface properties
+        E[id_inter] = E_inter
+        nu[id_inter] = nu_inter
+    # If piles are active, update the elasticity parameters at relevant integration points
+    if active_pile:
+        # Select integration points associated with the pile
+        id_pile = select_pile_integration_points(free, pile, equivalent_pile_thickness, lip)
+        # Update the elasticity modulus and Poisson's ratio at these points to pile properties
+        E[id_pile] = E_pile
+        nu[id_pile] = nu_pile
+    # Return the arrays containing the elasticity modulus and Poisson's ratio for each integration point
+    return E, nu
+
+
+# function 22: find interface integration points
+def select_inter_integration_points(free, inter, lip):
+    # Define a condition to select integration points that are within the interface region
+    condition = (
+        # Check if the x-coordinate of the integration point is between the x-coordinates of the interface nodes
+            (lip[:, 0] >= free[(inter[0] - 1), 0]) &
+            (lip[:, 0] <= free[(inter[1] - 1), 0]) &
+            # Check if the y-coordinate of the integration point is within a certain range defined by the interface
+            (lip[:, 1] >= (free[(inter[0] - 1), 1] + free[(inter[1] - 1), 1]) / 2 - inter[2])
+    )
+    # Find indices of integration points that satisfy the condition
+    id = np.where(condition)[0]
+    # Return the indices of the integration points within the interface region
+    return id
+
+
+# function 23: find pile integration points
+def select_pile_integration_points(free, pile, equivalent_pile_thickness, lip):
+    # Define a condition to select integration points that are within the pile region
+    condition = (
+        # Check if the x-coordinate of the integration point is within the x-range of the pile
+            (lip[:, 0] >= free[(pile[0] - 1), 0]) &
+            (lip[:, 0] <= free[(pile[0] - 1), 0] + equivalent_pile_thickness) &
+            # Check if the y-coordinate of the integration point is below a certain depth defined by the pile
+            (lip[:, 1] >= free[(pile[0] - 1), 1] - pile[1])
+    )
+    # Find indices of integration points that satisfy the condition
+    id = np.where(condition)[0]
+    # Return the indices of the integration points within the pile region
+    return id
+
+
+# function 24: form global derivative matrix
+def global_derivative(dN, co):
+    # Calculate the Jacobian matrix from the local derivatives and nodal coordinates
+    J = np.dot(dN, co)
+    # Compute the determinant of the Jacobian matrix
+    # This determinant is used for calculating the area (or volume in 3D) element in the global coordinate system
+    dA = np.linalg.det(J)
+    # Compute the global derivative of the shape functions
+    # This is done by solving the linear system J * D = dN, where D are the global derivatives
+    D = np.linalg.solve(J, dN)
+    # Return the global derivatives of the shape functions and the determinant of the Jacobian
+    return D, dA
+
+
+# function 25: calculate Lamé parameters
+def elasticity_parameters(E, nu):
+    # Lamé parameters 1
+    Lambda = E * nu / ((1 - 2 * nu) * (1 + nu))
+    # Lamé parameters 2 (shear modulus)
+    Mu = E / (2 * (1 + nu))
+    return Lambda, Mu
+
+
+# function 26: calculate strain-displacement matrix
+def strain_displacement_tensor(D):
+    # Initialize a zero matrix for the strain-displacement relationship
+    b = np.zeros((3, 16))
+    # Fill the matrix to relate displacements to strains
+    # For 2D problems, the strain-displacement matrix has 3 rows (for εx, εy, and γxy)
+    # and twice the number of nodes in columns (for x and y displacements at each node)
+    # The first row corresponds to strain in the x-direction (εx)
+    b[0, 0:16:2] = D[0, :]
+    # The second row corresponds to strain in the y-direction (εy)
+    b[1, 1:16:2] = D[1, :]
+    # The third row corresponds to shear strain (γxy)
+    b[2, 1:16:2] = D[0, :]
+    b[2, 0:16:2] = D[1, :]
+    # Return the strain-displacement matrix
+    return b
+
+
+# function 27: form displacement stiffness matrix
+def displacement_stiffness(no, nc, el, ec, df, ad, E, nu, anchorId, anchorSpringX, anchorSpringY):
+    # Calculate the Lame parameters (Lambda and Mu) from elasticity modulus and Poisson's ratio
+    Lambda, Mu = elasticity_parameters(E, nu)
+    # Obtain integration points in local coordinates
+    xi, yi = integration_point()
+    # Initialize arrays for sparse matrix construction
+    a1 = np.zeros(16 * 16 * ec)
+    a2 = np.zeros_like(a1)
+    a3 = np.zeros_like(a1)
+    # Iterate over each element in the mesh
+    for i in range(ec):
+        # Get the coordinates of the nodes of the current element
+        co = no[el[i, :], :]
+        # Initialize the element stiffness matrix
+        km = np.zeros((16, 16))
+        # Iterate over each integration point
+        for j in range(4):
+            # integration point index
+            id = i * 4 + j
+            # Get the local derivatives for the current integration point
+            dN1_1, dN1_2, dN1_3, dN1_4, dN1_5, dN1_6, dN1_7, dN1_8, dN2_1, dN2_2, dN2_3, dN2_4, dN2_5, dN2_6, dN2_7, dN2_8 = local_derivative8(
+                xi[j], yi[j])
+            dN = np.array([
+                [dN1_1, dN1_2, dN1_3, dN1_4, dN1_5, dN1_6, dN1_7, dN1_8],
+                [dN2_1, dN2_2, dN2_3, dN2_4, dN2_5, dN2_6, dN2_7, dN2_8]
+            ])
+            # Calculate the global derivatives and area element
+            D, dA = global_derivative(dN, co)
+            # Compute the strain-displacement matrix
+            b = strain_displacement_tensor(D)
+            # Create the elasticity matrix for the current integration point
+            Ce = np.array([
+                [Lambda[id] + 2 * Mu[id], Lambda[id], 0.0],
+                [Lambda[id], Lambda[id] + 2 * Mu[id], 0.0],
+                [0.0, 0.0, Mu[id]]
+            ])
+            # Accumulate the element stiffness matrix
+            km += np.dot(np.dot(b.T, Ce), b) * dA
+        # Store the element stiffness matrix in a linear array for sparse matrix construction
+        idx = slice(i * 256, (i + 1) * 256)
+        a1[idx] = np.repeat(df[i, :], 16)
+        a2[idx] = np.tile(df[i, :], 16)
+        a3[idx] = km.ravel()
+    # Add anchor stiffness to the global matrix
+    a1 = np.append(a1, anchorId * 2)
+    a1 = np.append(a1, anchorId * 2 + 1)
+    a2 = np.append(a2, anchorId * 2)
+    a2 = np.append(a2, anchorId * 2 + 1)
+    a3 = np.append(a3, anchorSpringX)
+    a3 = np.append(a3, anchorSpringY)
+    # Construct the global stiffness matrix as a sparse matrix
+    Km = csr_matrix((a3, (a1, a2)), shape=(2 * nc, 2 * nc))
+    # Convert to a Compressed Sparse Column matrix and use Sparse LU decomposition (on active degrees of freedom ad)
+    Km_csc = csc_matrix(Km)
+    Km_d = splu(Km_csc[ad, :][:, ad])
+    # Return the decomposed stiffness matrix
+    return Km_d
+
+
+# function 28: set specific weight throughout FEM domain
+def specificWeightIntegrationPoint(free, pile, inter, gamma_soil, gamma_pile, gamma_inter, equivalentPileThickness, ec,
+                                   lip, activePile, activeInter):
+    # Initialize the specific weight array for each integration point
+    gamma = np.full((4 * ec, 1), gamma_soil)
+    # If interfaces are active, update the specific weight at relevant integration points
+    if activeInter:
+        # Select integration points associated with the interface
+        id = select_inter_integration_points(free, inter, lip)
+        # Update the specific weight at these points to interface specific weight
+        gamma[id] = gamma_inter
+    # If piles are active, update the specific weight at relevant integration points
+    if activePile:
+        # Select integration points associated with the pile
+        id = select_pile_integration_points(free, pile, equivalentPileThickness, lip)
+        # Update the specific weight at these points to pile specific weight
+        gamma[id] = gamma_pile
+    # Return the array containing the specific weight for each integration point
+    return gamma
+
+
+# function 29: form body load vector
+def body_weight(no, nc, el, ec, gamma):
+    # Obtain integration points in local coordinates
+    xi, yi = integration_point()
+    # Initialize an array for nodal self-weight
+    nodal_self_weight = np.zeros(2 * nc)
+    # Iterate over each element in the mesh
+    for i in range(ec):
+        # Get the coordinates of the nodes of the current element
+        co = no[el[i, :], :]
+        # Iterate over each integration point
+        for j in range(4):
+            # Calculate the index for the specific weight at the integration point
+            id = (i * 4) + j
+            # Calculate the shape tensor for the integration point
+            N = shape_tensor8(xi[j], yi[j])
+            # Get the local derivatives for the current integration point
+            dN1_1, dN1_2, dN1_3, dN1_4, dN1_5, dN1_6, dN1_7, dN1_8, dN2_1, dN2_2, dN2_3, dN2_4, dN2_5, dN2_6, dN2_7, dN2_8 = local_derivative8(
+                xi[j], yi[j])
+            dN = np.array([
+                [dN1_1, dN1_2, dN1_3, dN1_4, dN1_5, dN1_6, dN1_7, dN1_8],
+                [dN2_1, dN2_2, dN2_3, dN2_4, dN2_5, dN2_6, dN2_7, dN2_8]
+            ])
+            # Calculate the global derivative and area element
+            D, dA = global_derivative(dN, co)
+            # Distribute the self-weight to the nodal points
+            # The y-component of the self-weight is considered (hence multiplying by 2 + 1)
+            nodal_self_weight[el[i, :] * 2 + 1] -= N.T * gamma[id] * dA
+    # Return the array of nodal self-weight
+    return nodal_self_weight
+
+
+# function 30: form init head based on water surface wl
+def initial_head(wl, X, no, el, free_surface_node):
+    # Interpolate the water level (head) at the free surface nodes using the provided water level data
+    fsh = np.interp(no[free_surface_node[0::2], 0], wl[:, 0], wl[:, 1])
+    # Get the number of unique x-coordinates in the mesh
+    nx = len(X)
+    # Initialize an array to store the initial head at each node
+    H = np.zeros(np.max(el[:, 0:4]) + 1)
+    # Distribute the interpolated head values to the corresponding nodes in the mesh
+    for i in range(nx):
+        H[i:np.max(el[:, 0:4]) + 1:nx] = fsh[i]
+    # Return the array of initial head values
+    return H
+
+
+# function 31: random field instance for friction angle, cohesion and hydraulic conductivity
+def random_field_instance(ks_soil, phi_soil, coh_soil, ec, L):
+    # Generate a random field instance using the Cholesky factor (L)
+    z = np.dot(L, np.random.randn(4 * ec))
+    # Calculate the hydraulic conductivity (ks) using a log-normal distribution
+    # The mean and variance of the log-normal distribution are based on ks_soil
+    ks = log_norm_distribution(ks_soil, z)
+    # Calculate the friction angle (phi) using a transformed distribution
+    # The transformation ensures that phi varies within the range defined by phi_soil
+    phi = phi_soil[0] + (phi_soil[1] - phi_soil[0]) * ((1 + np.tanh(z)) / 2)
+    # Calculate the cohesion (coh) using a log-normal distribution
+    # The mean and variance of the log-normal distribution are based on coh_soil
+    coh = log_norm_distribution(coh_soil, z)
+    # Return the generated values for hydraulic conductivity, friction angle, and cohesion
+    return ks, phi, coh
+
+
+def uniform_field_instance(ks_soil, phi_soil, coh_soil, ec):
+    z = np.ones(4 * ec)
+    ks = ks_soil[0] * z
+    phi = phi_soil[0] * z
+    coh = coh_soil[0] * z
+    # Return the generated values for hydraulic conductivity, friction angle, and cohesion
+    return ks, phi, coh
+
+
+# function 32: log normal distribution
+def log_norm_distribution(var, z):
+    log_norm_sig = np.sqrt(np.log(1 + (var[1] ** 2) / (var[0] ** 2)))
+    log_norm_mu = np.log(var[0]) - 0.5 * (log_norm_sig ** 2)
+    var_field = np.exp(log_norm_mu + log_norm_sig * z)
+    return var_field
+
+
+# function 33: set hydraulic conductivity throughout FEM domain
+def hydraulic_conductivity_integration_point(free, pile, ks, equivalent_pile_thickness, lip, active_pile):
+    # Check if the pile is active
+    if active_pile:
+        # Select integration points for the pile based on the free field, pile properties,
+        id = select_pile_integration_points(free, pile, equivalent_pile_thickness, lip)
+        # Set hydraulic conductivity to zero at the selected integration points.
+        ks[id] = 0.0
+    # Return the modified hydraulic conductivity array
+    return ks
+
+
+# function 34: set plasticity parameters throughout FEM domain
+def plasticity_parameters_integration_point(free, pile, inter, phi, phi_pile, phi_inter, dilation_factor, psi_pile,
+                                            psi_inter, coh, coh_pile, coh_inter, equivalent_pile_thickness, lip,
+                                            active_pile, active_inter):
+    # Calculate the initial dilation angle (psi) as a product of the dilation factor and the friction angle (phi)
+    psi = dilation_factor * phi
+    # Check if the integration points for the interface (inter) are active
+    if active_inter:
+        # Select integration points for the interface based on the free field, interface properties, and lip condition
+        id = select_inter_integration_points(free, inter, lip)
+        # Update the friction angle (phi), dilation angle (psi), and cohesion (coh) at the interface integration points
+        phi[id] = phi_inter
+        psi[id] = psi_inter
+        coh[id] = coh_inter
+    # Check if the integration points for the pile are active
+    if active_pile:
+        # Select integration points for the pile based on the free field, pile properties, equivalent pile thickness, and lip condition
+        id = select_pile_integration_points(free, pile, equivalent_pile_thickness, lip)
+        # Update the friction angle (phi), dilation angle (psi), and cohesion (coh) at the pile integration points
+        phi[id] = phi_pile
+        psi[id] = psi_pile
+        coh[id] = coh_pile
+    # Return the updated values of friction angle, dilation angle, and cohesion
+    return phi, psi, coh
+
+
+# function 35: form groundwater stiffness matrix
+def conductivity_stiffness(no, el, ec, ks):
+    # Obtain integration points in local coordinates
+    xi, yi = integration_point()
+    # Initialize arrays for sparse matrix construction
+    a1 = np.zeros(4 * 4 * ec)
+    a2 = np.zeros_like(a1)
+    a3 = np.zeros_like(a1)
+    # Iterate over each element in the mesh
+    for i in range(ec):
+        # Get the coordinates of the nodes of the current element (4-node elements)
+        co = no[el[i, 0:4], :]
+        # Initialize the element conductivity matrix
+        kc = np.zeros((4, 4))
+        # Iterate over each integration point
+        for j in range(4):
+            # integration point index
+            id = i * 4 + j
+            # Calculate the local derivative for the current integration point
+            dN = local_derivative4(xi[j], yi[j])
+            # Calculate the global derivative and area element
+            D, dA = global_derivative(dN, co)
+            # Create the conductivity matrix for the current integration point
+            ks_matrix = np.array([[ks[id], 0.0], [0.0, ks[id]]])
+            # Accumulate the element conductivity matrix
+            kc += D.T @ ks_matrix @ D * dA
+        # Store indices for the sparse matrix construction
+        indices = np.repeat(el[i, 0:4], 4)
+        a1[i * 16: (i + 1) * 16] = indices
+        a2[i * 16: (i + 1) * 16] = np.tile(el[i, 0:4], 4)
+        a3[i * 16: (i + 1) * 16] = kc.T.flatten()
+    # Get the total number of nodes
+    n = np.max(el[:, 0:4]) + 1
+    # Construct the global conductivity matrix as a sparse matrix
+    Kc = coo_matrix((a3, (a1, a2)), shape=(n, n))
+    # Return the global conductivity matrix
+    return Kc
+
+
+# function 36: FEM solver for groundwater flow
+def hydraulic_solver(H, Kc, no, el, fh, fs, lim, tol):
+    # Check if there are non-zero initial heads, zero init head's mean's no porepressure's
+    if np.sum(H) > 0:
+        # Define the set of all nodes
+        nh = np.arange(1, np.max(el[:, 0:4]) + 1)
+        # Initialize an array to store active free surface nodes
+        afs = np.array([], dtype=int)
+        # Identify active head nodes excluding fixed head and active free surface nodes
+        ah = np.setdiff1d(nh, np.union1d(fh, afs))
+        # Convert the conductivity matrix to Compressed Sparse Row format for efficient calculations
+        Kc_csr = Kc.tocsr()
+        # Iterative solver loop
+        for k in range(lim):
+            # Calculate internal flows
+            Qint = Kc_csr.dot(H)
+            # Copy the current head values for convergence checking
+            H0 = H.copy()
+            # Update the head values for active nodes
+            H[ah] -= spsolve(Kc_csr[np.ix_(ah, ah)], Qint[ah])
+            # Calculate the error for convergence check
+            er = np.max(np.abs(H - H0)) / np.max(np.abs(H0))
+            if er < tol:
+                break
+            elif k == lim - 1:
+                raise ValueError('no conv (hydraulic phase)')
+            # Check for seepage nodes that becoming active (total head > elevation)
+            Hs, fsId = np.max(H[fs] - no[fs, 1]), np.argmax(H[fs] - no[fs, 1])
+            if Hs > 0.0:
+                afs = np.append(afs, fs[fsId])
+            # Update the head for active seepage node
+            H[afs] = no[afs, 1]
+            # Update the list of active head nodes
+            ah = np.setdiff1d(nh, np.union1d(fh, afs))
+    # Return the updated head distribution
+    return H
+
+
+# function 37: calculate Darcy flow
+def flow_darcy(H, no, el, ec, ks):
+    # Obtain integration points in local coordinates
+    xi, yi = integration_point()
+    # Initialize arrays for Darcy velocity and flow rate
+    vee = np.zeros((2, 4 * ec))
+    flow_rate = np.zeros(4 * ec)
+    # Iterate over each element in the mesh
+    for i in range(ec):
+        # Get the coordinates of the nodes of the current element (4-node elements)
+        co = no[el[i, 0:4], :]
+        # Iterate over each integration point
+        for j in range(4):
+            # Calculate the index for the hydraulic conductivity at the integration point
+            id = (i * 4) + j
+            # Calculate the shape tensor for the integration point
+            N = shape_tensor4(xi[j], yi[j])
+            # Calculate the local derivative for the current integration point
+            dN = local_derivative4(xi[j], yi[j])
+            # Calculate the global derivative and area element
+            D, dA = global_derivative(dN, co)
+            # Compute the pore water pressure at the integration point
+            pw = N @ (H[el[i, 0:4]] - co[:, 1]) * 10.0
+            # Calculate the Darcy velocity if the pore water pressure is positive
+            if pw > 0.0:
+                vee[:, id] = -np.array([[ks[id], 0.0], [0.0, ks[id]]]) @ D @ H[el[i, 0:4]]
+            # Accumulate the flow rate at the nodal points
+            flow_rate[el[i, 0:4]] -= D.T @ vee[:, id] * dA
+    # Return the Darcy velocity and the norm of the flow rate
+    return vee, flow_rate
+
+
+# function 38: hydromechanical coupling porepressure to total stress
+def hydro_mechanical_coupling(H, no, nc, el, ec, df):
+    # Obtain integration points in local coordinates
+    xi, yi = integration_point()
+    # Calculate pore water pressure from the hydraulic head (H)
+    P = (H - no[:np.max(el[:, :4]) + 1, 1]) * 10.0  #
+    P[P < 0.0] = 0.0  # Negative pressures are not physically meaningful in this context
+    # Initialize an array for nodal pore pressure
+    nodal_pore_pressure = np.zeros(2 * nc)
+    # Iterate over each element in the mesh
+    for i in range(ec):
+        # Get the coordinates of the nodes of the current element
+        co = no[el[i, :], :]
+        # Iterate over each integration point
+        for j in range(4):
+            # Calculate the shape tensor for the integration point
+            N = shape_tensor4(xi[j], yi[j])
+            # Get the local derivatives for the current integration point
+            dN1_1, dN1_2, dN1_3, dN1_4, dN1_5, dN1_6, dN1_7, dN1_8, dN2_1, dN2_2, dN2_3, dN2_4, dN2_5, dN2_6, dN2_7, dN2_8 = local_derivative8(
+                xi[j], yi[j])
+            dN = np.array([
+                [dN1_1, dN1_2, dN1_3, dN1_4, dN1_5, dN1_6, dN1_7, dN1_8],
+                [dN2_1, dN2_2, dN2_3, dN2_4, dN2_5, dN2_6, dN2_7, dN2_8]
+            ])
+            # Calculate the global derivative and area element
+            D, dA = global_derivative(dN, co)
+            # Compute the strain-displacement matrix
+            b = strain_displacement_tensor(D)
+            # Distribute the pore pressure to the nodal points
+            nodal_pore_pressure[df[i, :]] += b.T @ np.array([1, 1, 0]) * (N @ P[el[i, :4]]) * dA
+    # Return the array of nodal pore pressure
+    return nodal_pore_pressure
+
+
+# function 39: vectorized version of nodal coordinate array
+def coordinate_array_vectorization(no, el):
+    # Extract x and y coordinates for each element node
+    c1_1 = no[el[:, 0], 0]
+    c1_2 = no[el[:, 0], 1]
+    c2_1 = no[el[:, 1], 0]
+    c2_2 = no[el[:, 1], 1]
+    c3_1 = no[el[:, 2], 0]
+    c3_2 = no[el[:, 2], 1]
+    c4_1 = no[el[:, 3], 0]
+    c4_2 = no[el[:, 3], 1]
+    c5_1 = no[el[:, 4], 0]
+    c5_2 = no[el[:, 4], 1]
+    c6_1 = no[el[:, 5], 0]
+    c6_2 = no[el[:, 5], 1]
+    c7_1 = no[el[:, 6], 0]
+    c7_2 = no[el[:, 6], 1]
+    c8_1 = no[el[:, 7], 0]
+    c8_2 = no[el[:, 7], 1]
+    # Return the extracted x and y coordinates as separate variables
+    return c1_1, c1_2, c2_1, c2_2, c3_1, c3_2, c4_1, c4_2, c5_1, c5_2, c6_1, c6_2, c7_1, c7_2, c8_1, c8_2
+
+
+# function 40: vectorized version of array declarations displacement solver array
+def array_declaration_vectorization(nc, ec):
+    # Create arrays with zeros and appropriate dimensions
+    U = np.zeros(2 * nc)
+    evpt1 = np.zeros(4 * ec)
+    evpt2 = np.zeros(4 * ec)
+    evpt3 = np.zeros(4 * ec)
+    evpt4 = np.zeros(4 * ec)
+    sigma = np.zeros(4 * ec)
+    f1 = np.zeros(ec)
+    f2 = np.zeros(ec)
+    f3 = np.zeros(ec)
+    f4 = np.zeros(ec)
+    f5 = np.zeros(ec)
+    f6 = np.zeros(ec)
+    f7 = np.zeros(ec)
+    f8 = np.zeros(ec)
+    f9 = np.zeros(ec)
+    f10 = np.zeros(ec)
+    f11 = np.zeros(ec)
+    f12 = np.zeros(ec)
+    f13 = np.zeros(ec)
+    f14 = np.zeros(ec)
+    f15 = np.zeros(ec)
+    f16 = np.zeros(ec)
+    # Return all the created arrays
+    return U, evpt1, evpt2, evpt3, evpt4, sigma, f1, f2, f3, f4, f5, f6, f7, f8, f9, f10, f11, f12, f13, f14, f15, f16
+
+
+# function 41: vectorized version of displacement vector array
+def displacement_vectorization(U, df):
+    # Extract displacement values from the U array based on the indices specified in the df array
+    u1 = U[df[:, 0]]
+    u2 = U[df[:, 1]]
+    u3 = U[df[:, 2]]
+    u4 = U[df[:, 3]]
+    u5 = U[df[:, 4]]
+    u6 = U[df[:, 5]]
+    u7 = U[df[:, 6]]
+    u8 = U[df[:, 7]]
+    u9 = U[df[:, 8]]
+    u10 = U[df[:, 9]]
+    u11 = U[df[:, 10]]
+    u12 = U[df[:, 11]]
+    u13 = U[df[:, 12]]
+    u14 = U[df[:, 13]]
+    u15 = U[df[:, 14]]
+    u16 = U[df[:, 15]]
+    # Return the extracted displacement values as separate variables
+    return u1, u2, u3, u4, u5, u6, u7, u8, u9, u10, u11, u12, u13, u14, u15, u16
+
+
+# function 42: vectorized version of strain calculation
+def strain_displacement_vectorized(u1, u2, u3, u4, u5, u6, u7, u8, u9, u10, u11, u12, u13, u14, u15, u16, evpt1, evpt2,
+                                   evpt3, evpt4, D1_1, D1_2, D1_3, D1_4, D1_5, D1_6, D1_7, D1_8, D2_1, D2_2, D2_3, D2_4,
+                                   D2_5, D2_6, D2_7, D2_8):
+    # Calculate the strain components based on the displacement field and global derivatives
+    # e1, e2, and e3 correspond to the normal strains in x, y, and shear strain respectively
+    e1 = D1_1 * u1 + D1_2 * u3 + D1_3 * u5 + D1_4 * u7 + D1_5 * u9 + D1_6 * u11 + D1_7 * u13 + D1_8 * u15
+    e2 = D2_1 * u2 + D2_2 * u4 + D2_3 * u6 + D2_4 * u8 + D2_5 * u10 + D2_6 * u12 + D2_7 * u14 + D2_8 * u16
+    e3 = D1_1 * u2 + D1_2 * u4 + D1_3 * u6 + D1_4 * u8 + D2_1 * u1 + D1_5 * u10 + D2_2 * u3 + D1_6 * u12 + D2_3 * u5 + D1_7 * u14 + D2_4 * u7 + D1_8 * u16 + D2_5 * u9 + D2_6 * u11 + D2_7 * u13 + D2_8 * u15
+    # Adjust the strain components by the total viscoplastic plastic strain components (evpt)
+    e1 -= evpt1
+    e2 -= evpt2
+    e3 -= evpt3
+    e4 = -evpt4  # plane strain conditions
+    # Return the calculated strain components
+    return e1, e2, e3, e4
+
+
+# function 43: vectorized version of stress calculation
+def stress_strain_vectorized(Lambda, Mu, e1, e2, e3, e4):
+    # Calculate the stress components based on the strain components and material properties (Lambda and Mu)
+    # s1, s2, s3, and s4 correspond to the stress components in the x, y, shear, and z directions respectively
+    s1 = (Lambda + 2 * Mu) * e1 + Lambda * (e2 + e4)
+    s2 = (Lambda + 2 * Mu) * e2 + Lambda * (e1 + e4)
+    s3 = Mu * e3
+    s4 = (Lambda + 2 * Mu) * e4 + Lambda * (e1 + e2)
+    # Return the calculated stress components
+    return s1, s2, s3, s4
+
+
+# function 44: vectorized version of indexing array
+def vector_degree_freedom(df):
+    vId = np.concatenate([df[:, 0], df[:, 1], df[:, 2], df[:, 3],
+                          df[:, 4], df[:, 5], df[:, 6], df[:, 7],
+                          df[:, 8], df[:, 9], df[:, 10], df[:, 11],
+                          df[:, 12], df[:, 13], df[:, 14], df[:, 15]])
+    return vId
+
+
+# function 45: shear strength reduction routine
+def shear_strength_reduction(FS, phi, psi, coh, E, nu):
+    # reduction friction angle
+    phir = np.arctan(np.tan(np.radians(phi)) / FS)
+    # reduction dilation angle
+    psir = np.arctan(np.tan(np.radians(psi)) / FS)
+    # recalculate Mohr-Coulomb shear strength parameters
+    snph = np.sin(phir)
+    csph = np.cos(phir)
+    snps = np.sin(psir)
+    # reduction cohesion
+    cohr = coh / FS
+    # calculate critical time step
+    dt = (4 * (1 + nu) * (1 - 2 * nu)) / (E * (1 - 2 * nu + snph ** 2))
+    # return critical time step and shear strength parameters
+    return snph, csph, snps, cohr, dt
+
+
+# function 46: vectorized version stress invariant routine
+def stress_invariants(s1, s2, s3, s4):
+    # Calculate the first stress invariant, p (mean stress)
+    p = (s1 + s2 + s4) / 3
+    # Calculate deviatoric stresses
+    d1 = (2 * s1 - s2 - s4) / 3
+    d2 = (2 * s2 - s1 - s4) / 3
+    d3 = (2 * s4 - s1 - s2) / 3
+    # Calculate the second stress invariant, J2 (magnitude of deviatoric stress tensor)
+    J2 = 1 / 2 * (d1 ** 2 + d2 ** 2 + 2 * s3 ** 2 + d3 ** 2)
+    # Initialize the Lode angle parameters
+    id = J2 > 0.0  # Only calculate where J2 is positive
+    q = np.zeros_like(s1)
+    q[id] = np.sqrt(3 * J2[id])  # q is the equivalent deviatoric stress
+    tp = np.zeros_like(s1)  # tp is an intermediate parameter for Lode angle calculation
+    tp[id] = -13.5 * (d1[id] * d2[id] * d3[id] - d3[id] * s3[id] ** 2) / q[id] ** 3
+    tp[tp >= 1] = 1  # compression limit lode angle
+    tp[tp <= -1] = -1  # extension limit lode angle
+    # Calculate the Lode angle, t
+    t = np.arcsin(tp) / 3
+    # Calculate partial derivatives for the consistency parameters (dp1, dp2, dp3, dp4)
+    dp1 = 1.0
+    dp2 = 1.0
+    dp3 = 0.0
+    dp4 = 1.0
+    # Calculate partial derivatives of J2 with respect to the stress components
+    dJ2_1 = d1
+    dJ2_2 = d2
+    dJ2_3 = 2 * s3
+    dJ2_4 = d3
+    # Calculate partial derivatives of J3 with respect to the stress components
+    dJ3_1 = d1 ** 2 + s3 ** 2 - 2 / 3 * J2
+    dJ3_2 = s3 ** 2 + d2 ** 2 - 2 / 3 * J2
+    dJ3_3 = 2 * (d1 * s3 + s3 * d2)
+    dJ3_4 = d3 ** 2 - 2 / 3 * J2
+    # Return the stress invariants and their derivatives
+    return p, q, t, dp1, dp2, dp3, dp4, dJ2_1, dJ2_2, dJ2_3, dJ2_4, dJ3_1, dJ3_2, dJ3_3, dJ3_4
+
+
+# function 47: vectorized version of Mohr Coulomb Yield function & derivatives
+def yield_mohr_coulomb(p, q, t, snph, csph, snps, coh):
+    # Calculate trigonometric functions and constants used in the Mohr-Coulomb material description
+    snth = np.sin(t)
+    csth = np.cos(t)
+    cs3th = np.cos(3.0 * t)
+    tn3th = np.tan(3.0 * t)
+    tnth = snth / csth
+    sq3 = np.sqrt(3.0)
+    # Calculate the Mohr-Coulomb yield function
+    f = p * snph - coh * csph + q * (csth / sq3 - (snph * snth) / 3.0)
+    f[f < 0.0] = 0.0
+    # Calculate the derivatives of the yield function with respect to p,q and t
+    dq1 = snps
+    dq2 = (csth * sq3 * (tn3th * tnth + (snps * sq3 * (tn3th - tnth)) / 3.0 + 1.0)) / (2.0 * q)
+    dq3 = (1.0 / q ** 2 * (3.0 * csth * snps / 2.0 + snth * 3.0 * sq3 / 2.0)) / cs3th
+    # Adjust dq2 and dq3 close to pure extension and compression to avoid singularities
+    mask1 = snth > 0.49
+    mask2 = snth < -0.49
+    dq2[mask1] = (sq3 / 2.0 - snps[mask1] / (2.0 * sq3)) * sq3 / (2.0 * q[mask1])
+    dq2[mask2] = (sq3 / 2.0 + snps[mask2] / (2.0 * sq3)) * sq3 / (2.0 * q[mask2])
+    dq3[mask1] = 0.0
+    dq3[mask2] = 0.0
+    # Calculate the shear strength, based on the Mohr-Coulomb criterion
+    sigma = (-p * snph + coh * csph) / (csth / sq3 - (snph * snth) / 3.0)
+    # Return the yield function, its derivatives, and the equivalent stress
+    return f, dq1, dq2, dq3, sigma
+
+
+# function 48: vectorized version of the global derivative array
+def global_derivative_vectorized(dN1_1, dN1_2, dN1_3, dN1_4, dN1_5, dN1_6, dN1_7, dN1_8, dN2_1, dN2_2, dN2_3, dN2_4,
+                                 dN2_5, dN2_6, dN2_7, dN2_8, c1_1, c1_2, c2_1, c2_2, c3_1, c3_2, c4_1, c4_2, c5_1, c5_2,
+                                 c6_1, c6_2, c7_1, c7_2, c8_1, c8_2):
+    # Calculate the Jacobian matrix components
+    J1_1 = c1_1 * dN1_1 + c2_1 * dN1_2 + c3_1 * dN1_3 + c4_1 * dN1_4 + c5_1 * dN1_5 + c6_1 * dN1_6 + c7_1 * dN1_7 + c8_1 * dN1_8
+    J1_2 = c1_2 * dN1_1 + c2_2 * dN1_2 + c3_2 * dN1_3 + c4_2 * dN1_4 + c5_2 * dN1_5 + c6_2 * dN1_6 + c7_2 * dN1_7 + c8_2 * dN1_8
+    J2_1 = c1_1 * dN2_1 + c2_1 * dN2_2 + c3_1 * dN2_3 + c4_1 * dN2_4 + c5_1 * dN2_5 + c6_1 * dN2_6 + c7_1 * dN2_7 + c8_1 * dN2_8
+    J2_2 = c1_2 * dN2_1 + c2_2 * dN2_2 + c3_2 * dN2_3 + c4_2 * dN2_4 + c5_2 * dN2_5 + c6_2 * dN2_6 + c7_2 * dN2_7 + c8_2 * dN2_8
+    # Compute the determinant of the Jacobian matrix
+    dA = J1_1 * J2_2 - J1_2 * J2_1
+    # Compute the inverse Jacobian matrix components
+    iJ1_1 = J2_2 / dA
+    iJ1_2 = -J1_2 / dA
+    iJ2_1 = -J2_1 / dA
+    iJ2_2 = J1_1 / dA
+    # Compute the global derivatives
+    D1_1 = dN1_1 * iJ1_1 + dN2_1 * iJ1_2
+    D1_2 = dN1_2 * iJ1_1 + dN2_2 * iJ1_2
+    D1_3 = dN1_3 * iJ1_1 + dN2_3 * iJ1_2
+    D1_4 = dN1_4 * iJ1_1 + dN2_4 * iJ1_2
+    D1_5 = dN1_5 * iJ1_1 + dN2_5 * iJ1_2
+    D1_6 = dN1_6 * iJ1_1 + dN2_6 * iJ1_2
+    D1_7 = dN1_7 * iJ1_1 + dN2_7 * iJ1_2
+    D1_8 = dN1_8 * iJ1_1 + dN2_8 * iJ1_2
+    D2_1 = dN1_1 * iJ2_1 + dN2_1 * iJ2_2
+    D2_2 = dN1_2 * iJ2_1 + dN2_2 * iJ2_2
+    D2_3 = dN1_3 * iJ2_1 + dN2_3 * iJ2_2
+    D2_4 = dN1_4 * iJ2_1 + dN2_4 * iJ2_2
+    D2_5 = dN1_5 * iJ2_1 + dN2_5 * iJ2_2
+    D2_6 = dN1_6 * iJ2_1 + dN2_6 * iJ2_2
+    D2_7 = dN1_7 * iJ2_1 + dN2_7 * iJ2_2
+    D2_8 = dN1_8 * iJ2_1 + dN2_8 * iJ2_2
+    # return derivatives
+    return D1_1, D1_2, D1_3, D1_4, D1_5, D1_6, D1_7, D1_8, D2_1, D2_2, D2_3, D2_4, D2_5, D2_6, D2_7, D2_8, dA
+
+
+# function 49: select pile elements routine
+def select_pile_elements(free, pile, no, el):
+    pn = np.where((no[:, 0] == free[(pile[0] - 1), 0]) & (no[:, 1] >= free[(pile[0] - 1), 1] - pile[1]))[0]
+    id = np.where(np.isin(el[:, 0], pn))[0]
+    return id
+
+
+# function 50: solver for the displacement and stability calculation
+def displacement_solver_vectorized(SF, no, nc, el, ec, df, phi, psi, coh, E, nu, Fext, Km_d, ad, lim, tol):
+    CONV = True  # Variable to check convergence
+    # Calculate the Lame parameters (Lambda and Mu) from elasticity modulus and Poisson's ratio
+    Lambda, Mu = elasticity_parameters(E, nu)
+    # Obtain integration points in local coordinates
+    xi, yi = integration_point()
+    # Vectorize the nodal coordinates for all elements
+    c1_1, c1_2, c2_1, c2_2, c3_1, c3_2, c4_1, c4_2, c5_1, c5_2, c6_1, c6_2, c7_1, c7_2, c8_1, c8_2 = coordinate_array_vectorization(
+        no, el)
+    # Create a vector of degrees of freedom
+    vId = vector_degree_freedom(df)
+    # Calculate shear strength parameters with a reduction factor equivalent to SF
+    snph, csph, snps, cohr, dt = shear_strength_reduction(SF, phi, psi, coh, E, nu)
+    # Initialize arrays for displacement, plastic strain, stress, yield function, and internal forces
+    U, evpt1, evpt2, evpt3, evpt4, sigma, f1, f2, f3, f4, f5, f6, f7, f8, f9, f10, f11, f12, f13, f14, f15, f16 = array_declaration_vectorization(
+        nc, ec)
+    # Solve the displacement field
+    U[ad] = Km_d.solve(Fext[ad])
+    k = 0
+    for k in range(lim):
+        # Vectorize the displacements for all nodes
+        u1, u2, u3, u4, u5, u6, u7, u8, u9, u10, u11, u12, u13, u14, u15, u16 = displacement_vectorization(U, df)
+        # Loop over each integration point
+        for j in range(4):
+            id = range(j, 4 * ec, 4)
+            # Get the local derivatives for the current integration point
+            dN1_1, dN1_2, dN1_3, dN1_4, dN1_5, dN1_6, dN1_7, dN1_8, dN2_1, dN2_2, dN2_3, dN2_4, dN2_5, dN2_6, dN2_7, dN2_8 = local_derivative8(
+                xi[j], yi[j])
+            # Calculate the global derivatives and area element
+            D1_1, D1_2, D1_3, D1_4, D1_5, D1_6, D1_7, D1_8, D2_1, D2_2, D2_3, D2_4, D2_5, D2_6, D2_7, D2_8, dA = global_derivative_vectorized(
+                dN1_1, dN1_2, dN1_3, dN1_4, dN1_5, dN1_6, dN1_7, dN1_8, dN2_1, dN2_2, dN2_3, dN2_4, dN2_5, dN2_6, dN2_7,
+                dN2_8, c1_1, c1_2, c2_1, c2_2, c3_1, c3_2, c4_1, c4_2, c5_1, c5_2, c6_1, c6_2, c7_1, c7_2, c8_1, c8_2)
+            # Calculate strains and stresses for each integration point
+            e1, e2, e3, e4 = strain_displacement_vectorized(u1, u2, u3, u4, u5, u6, u7, u8, u9, u10, u11, u12, u13, u14,
+                                                            u15, u16, evpt1[id], evpt2[id], evpt3[id], evpt4[id], D1_1,
+                                                            D1_2, D1_3, D1_4, D1_5, D1_6, D1_7, D1_8, D2_1, D2_2, D2_3,
+                                                            D2_4, D2_5, D2_6, D2_7, D2_8)
+            s1, s2, s3, s4 = stress_strain_vectorized(Lambda[id], Mu[id], e1, e2, e3, e4)
+            # Calculate stress invariant and derivatives
+            p, q, t, dp1, dp2, dp3, dp4, dJ2_1, dJ2_2, dJ2_3, dJ2_4, dJ3_1, dJ3_2, dJ3_3, dJ3_4 = stress_invariants(s1,
+                                                                                                                    s2,
+                                                                                                                    s3,
+                                                                                                                    s4)
+            # calculate Mohr Coulomb yield function and derivaives
+            F, dq1, dq2, dq3, sigma_id = yield_mohr_coulomb(p, q, t, snph[id], csph[id], snps[id], cohr[id])
+            sigma[id] = sigma_id
+            # caclulate viscoplastic strain increment
+            evp1 = dt[id] * F * (dq1 * dp1 + dq2 * dJ2_1 + dq3 * dJ3_1)
+            evp2 = dt[id] * F * (dq1 * dp2 + dq2 * dJ2_2 + dq3 * dJ3_2)
+            evp3 = dt[id] * F * (dq1 * dp3 + dq2 * dJ2_3 + dq3 * dJ3_3)
+            evp4 = dt[id] * F * (dq1 * dp4 + dq2 * dJ2_4 + dq3 * dJ3_4)
+            # add viscoplastic strain increment to total viscoplastic strain
+            evpt1[id] += evp1
+            evpt2[id] += evp2
+            evpt3[id] += evp3
+            evpt4[id] += evp4
+            # caclulate viscoplastic stress
+            sp1, sp2, sp3, sp4 = stress_strain_vectorized(Lambda[id], Mu[id], evp1, evp2, evp3, evp4)
+            # add viscoplastic stress to viscoplastic body loads
+            f1 += dA * (D1_1 * sp1 + D2_1 * sp3)
+            f2 += dA * (D1_1 * sp3 + D2_1 * sp2)
+            f3 += dA * (D1_2 * sp1 + D2_2 * sp3)
+            f4 += dA * (D1_2 * sp3 + D2_2 * sp2)
+            f5 += dA * (D1_3 * sp1 + D2_3 * sp3)
+            f6 += dA * (D1_3 * sp3 + D2_3 * sp2)
+            f7 += dA * (D1_4 * sp1 + D2_4 * sp3)
+            f8 += dA * (D1_4 * sp3 + D2_4 * sp2)
+            f9 += dA * (D1_5 * sp1 + D2_5 * sp3)
+            f10 += dA * (D1_5 * sp3 + D2_5 * sp2)
+            f11 += dA * (D1_6 * sp1 + D2_6 * sp3)
+            f12 += dA * (D1_6 * sp3 + D2_6 * sp2)
+            f13 += dA * (D1_7 * sp1 + D2_7 * sp3)
+            f14 += dA * (D1_7 * sp3 + D2_7 * sp2)
+            f15 += dA * (D1_8 * sp1 + D2_8 * sp3)
+            f16 += dA * (D1_8 * sp3 + D2_8 * sp2)
+        # assemble viscoplastic load
+        values = np.concatenate([f1, f2, f3, f4, f5, f6, f7, f8, f9, f10, f11, f12, f13, f14, f15, f16])
+        Fint = np.bincount(vId, weights=values, minlength=2 * nc)
+        # add viscoplastic load to external load vector
+        loads = Fext + Fint
+        # store previous displacement
+        U0 = U.copy()
+        # solve system of equations to derive new displacement (active degree of freedom only ad)
+        U[ad] = Km_d.solve(loads[ad])
+        # calculate the error
+        er = np.max(np.abs(U - U0)) / np.max(np.abs(U0))
+        if er < tol:
+            break
+    # check whether solver converged
+    if k == lim - 1:
+        CONV = False
+    # return displacement, viscoplastic strain, shear strength and convergence
+    return CONV, U, evpt1, evpt2, evpt3, evpt4, sigma
+
+
+# function 51: find safety factor by executing displacement solvers with several factors of safety
+def find_safety_factor(no, nc, el, ec, df, phi, psi, coh, E, nu, Fext, Km_d, ad, lim, tol):
+    # Create an array of safety factors from 0.1 to 5.0 with a step of 0.1
+    SF = [1]
+    i = 0
+    U = 0
+    evpt1 = 0
+    evpt2 = 0
+    evpt3 = 0
+    evpt4 = 0
+    sigma = 0
+    # Loop through each safety factor in the SF array
+    for i, sf in enumerate(SF):
+        # Call the displacement_solver function to calculate displacement and other values
+        CONV, U, evpt1, evpt2, evpt3, evpt4, sigma = displacement_solver_vectorized(sf, no, nc, el, ec, df, phi, psi,
+                                                                                    coh, E, nu, Fext, Km_d, ad, lim,
+                                                                                    tol)
+        # Check if the convergence flag is False (meaning the solver did not converge)
+        if not CONV:
+            break
+    # Determine the factor of safety that caused the solver to stop
+    factor_of_safety = SF[i]
+    # Return the calculated values and the factor of safety
+    return U, evpt1, evpt2, evpt3, evpt4, sigma, factor_of_safety
+
+
+# function 52: display Monte Carlo results
+def displayMonteCarlo(normal_flow_rate, FoS):
+    # Create a figure for displaying Monte Carlo results
+    plt.figure(figsize=(19.2, 10.8))
+    plt.suptitle('Monte Carlo results')
+    # Create the first subplot for the histogram of normal_flow_rate
+    plt.subplot(1, 2, 1)
+    plt.hist(normal_flow_rate, edgecolor='black')
+    plt.xlabel('Bin', fontsize=14)
+    plt.ylabel('Instances', fontsize=14)
+    plt.title('|Flow Rate| [m³/s]', fontsize=14)
+    # Create the second subplot for the histogram of FoS (Factor of Safety)
+    plt.subplot(1, 2, 2)
+    plt.hist(FoS, edgecolor='black')
+    plt.xlabel('Bin', fontsize=14)
+    plt.ylabel('Instances', fontsize=14)
+    plt.title('Factor of Safety [-]', fontsize=14)
+    # Display the entire figure with both subplots
+    plt.show()
+
+
+# function 53: display single run results
+def displayResults(free, wl, pile, def_scaling, equivalent_pile_thickness, anchor_Id, no, nc, el, ec,
+                   nodal_water_pressure, lip, FoS, ks, H, U, evpt1, evpt2, evpt3, evpt4, sigma, active_pile):
+    # Calculate the pore pressure P
+    P = (H - no[:np.max(el[:, :4]) + 1, 1]) * 10.0
+    P[P < 0.0] = 0.0
+    # Calculate the average pore pressure Pe for each element
+    Pe = (P[el[:, 0]] + P[el[:, 1]] + P[el[:, 2]] + P[el[:, 3]]) / 4
+    # Calculate a scaling factor for displacements
+    scaling_factor = np.round(def_scaling * np.max(no) / np.max(np.abs(U)))
+    # Scale the displacements U
+    U_scaling = scaling_factor * U
+    # Calculate the deformed node positions (no_d)
+    no_d = np.zeros_like(no)
+    no_d[:, 0] = no[:, 0] + U_scaling[0:2 * nc:2]
+    no_d[:, 1] = no[:, 1] + U_scaling[1:2 * nc:2]
+    # Calculate vee and normal_flow_rate using flow_darcy function
+    vee, flow_rate = flow_darcy(H, no, el, ec, ks)
+    # Calculate the norm of the flow rate vector for the entire mesh
+    normal_flow_rate = np.round(np.linalg.norm(flow_rate), 6)
+    # Modify sigma for active pile integration points
+    if active_pile == True:
+        id = select_pile_integration_points(free, pile, equivalent_pile_thickness, lip)
+        sigma[id] = 0.0
+        pileId = select_pile_elements(free, pile, no, el)
+    else:
+        pileId = []
+    # Create a colormap for visualization
+    cmap = cm.viridis
+    # Calculate kse (average hydraulic conductivity for elements)
+    kse = (ks[0::4] + ks[1::4] + ks[2::4] + ks[3::4]) / 4
+    # Normalize colors for hydraulic conductivity
+    norm = plt.Normalize(np.min(kse), np.max(kse))
+    colors0 = cmap(norm(kse))
+    sm1 = cm.ScalarMappable(norm=norm, cmap=cmap)
+    sm1.set_array([])
+    # Calculate sigmae (average shear strength for elements)
+    sigmae = (sigma[0::4] + sigma[1::4] + sigma[2::4] + sigma[3::4]) / 4
+    # Normalize colors for shear strength
+    norm = plt.Normalize(np.min(sigmae), np.max(sigmae))
+    colors1 = cmap(norm(sigmae))
+    sm2 = cm.ScalarMappable(norm=norm, cmap=cmap)
+    sm2.set_array([])
+    # Calculate nevpt (norm of evpt1, evpt2, evpt3, evpt4 for nodes)
+    nevpt = np.sqrt(evpt1 ** 2 + evpt2 ** 2 + evpt3 ** 2 + evpt4 ** 2)
+    # Calculate nevpte (average norm of evpt for elements)
+    nevpte = (nevpt[0::4] + nevpt[1::4] + nevpt[2::4] + nevpt[3::4]) / 4
+    # Normalize colors for evpt
+    norm = plt.Normalize(np.min(nevpte), np.max(nevpte))
+    colors2 = cmap(norm(nevpte))
+    # Normalize colors for Pe (average pore pressure for elements)
+    norm = plt.Normalize(np.min(Pe), np.max(Pe))
+    colors3 = cmap(norm(Pe))
+    sm3 = cm.ScalarMappable(norm=norm, cmap=cmap)
+    sm3.set_array([])
+    # figure 1: hydraulic conductivity
+    fig, ax = plt.subplots(figsize=(19.2, 10.8))
+    plt.title(f'Hydraulic Conductivity [m/s], |Flow Rate| {normal_flow_rate} [m3/d]', fontsize=14)
+    for i in range(ec):
+        indices = el[i, [0, 4, 1, 5, 2, 6, 3, 7]]
+        x = no[indices, 0]
+        y = no[indices, 1]
+        plt.fill(x, y, color=colors0[i], edgecolor='none', alpha=0.7)
+    for i in pileId:
+        x_coords = no[el[i, [0, 4, 1, 5, 2, 6, 3, 7]], 0]
+        y_coords = no[el[i, [0, 4, 1, 5, 2, 6, 3, 7]], 1]
+        plt.fill(x_coords, y_coords, 'red')
+    if anchor_Id.size == 1:
+        plt.scatter(no[anchor_Id, 0], no[anchor_Id, 1], 100, color='white', marker='s', facecolors='none')
+    for i in range(wl.shape[0] - 1):
+        plt.plot([wl[i, 0], wl[i + 1, 0]], [wl[i, 1], wl[i + 1, 1]], '-b', linewidth=2)
+    if np.sum(np.abs(nodal_water_pressure)) > 0:
+        plt.quiver(no[:, 0], no[:, 1], -nodal_water_pressure[::2], -nodal_water_pressure[1::2], scale=5,
+                   scale_units='xy', angles='xy')
+    plt.colorbar(sm1, ax=ax, orientation='vertical')
+    plt.xlabel('Width [m]', fontsize=14)
+    plt.ylabel('Height [m]', fontsize=14)
+    plt.axis('equal')
+    plt.tight_layout()
+    plt.show()
+    # figure 2: porepressure & Darcy flow
+    fig, ax = plt.subplots(figsize=(19.2, 10.8))
+    plt.title(f'Porepressure [kPa] & Darcy Flow [m/s], |Flow Rate| {normal_flow_rate} [m3/d]', fontsize=14)
+    for i in range(ec):
+        indices = el[i, [0, 4, 1, 5, 2, 6, 3, 7]]
+        x = no[indices, 0]
+        y = no[indices, 1]
+        plt.fill(x, y, color=colors3[i], edgecolor='none', alpha=0.7)
+    for i in pileId:
+        x_coords = no[el[i, [0, 4, 1, 5, 2, 6, 3, 7]], 0]
+        y_coords = no[el[i, [0, 4, 1, 5, 2, 6, 3, 7]], 1]
+        plt.fill(x_coords, y_coords, 'red')
+    if anchor_Id.size == 1:
+        plt.scatter(no[anchor_Id, 0], no[anchor_Id, 1], 100, color='white', marker='s', facecolors='none')
+    if np.sum(np.abs(vee)) > 0:
+        plt.quiver(lip[:, 0], lip[:, 1], vee[0, :], vee[1, :], scale_units='xy', angles='xy', alpha=0.3)
+    plt.colorbar(sm3, ax=ax, orientation='vertical')
+    plt.xlabel('Width [m]', fontsize=14)
+    plt.ylabel('Height [m]', fontsize=14)
+    plt.axis('equal')
+    plt.tight_layout()
+    plt.show()
+    # figure 3: shear strength
+    fig, ax = plt.subplots(figsize=(19.2, 10.8))
+    plt.title(f'Shear Strength [kPa], FS {FoS} [-]', fontsize=14)
+    for i in range(ec):
+        indices = el[i, [0, 4, 1, 5, 2, 6, 3, 7]]
+        x = no[indices, 0]
+        y = no[indices, 1]
+        plt.fill(x, y, color=colors1[i], edgecolor='none', alpha=0.7)
+    for i in pileId:
+        x_coords = no[el[i, [0, 4, 1, 5, 2, 6, 3, 7]], 0]
+        y_coords = no[el[i, [0, 4, 1, 5, 2, 6, 3, 7]], 1]
+        plt.fill(x_coords, y_coords, 'red')
+    if anchor_Id.size == 1:
+        plt.scatter(no[anchor_Id, 0], no[anchor_Id, 1], 100, color='white', marker='s', facecolors='none')
+    for i in range(wl.shape[0] - 1):
+        plt.plot([wl[i, 0], wl[i + 1, 0]], [wl[i, 1], wl[i + 1, 1]], '-b', linewidth=2)
+    if np.sum(np.abs(nodal_water_pressure)) > 0:
+        plt.quiver(no[:, 0], no[:, 1], -nodal_water_pressure[::2], -nodal_water_pressure[1::2], scale=5,
+                   scale_units='xy', angles='xy')
+    plt.colorbar(sm2, ax=ax, orientation='vertical')
+    plt.xlabel('Width [m]', fontsize=14)
+    plt.ylabel('Height [m]', fontsize=14)
+    plt.axis('equal')
+    plt.tight_layout()
+    plt.show()
+    # figure 4: Viscoplastic Strain
+    plt.figure(figsize=(19.2, 10.8))
+    plt.title(f'Viscoplastic Strain [-], FS {FoS} [-]', fontsize=14)
+    for i in range(ec):
+        indices = el[i, [0, 4, 1, 5, 2, 6, 3, 7]]
+        x = no[indices, 0]
+        y = no[indices, 1]
+        plt.fill(x, y, color=colors2[i], edgecolor='none', alpha=0.7)
+    plt.xlabel('Width [m]', fontsize=14)
+    plt.ylabel('Height [m]', fontsize=14)
+    plt.axis('equal')
+    plt.tight_layout()
+    plt.show()
+    # figure 5: Deformed Mesh
+    plt.figure(figsize=(19.2, 10.8))
+    plt.title(f'Deformed Mesh, scaled {scaling_factor} times, FS {FoS} [-]', fontsize=14)
+    for i in range(ec):
+        indices = el[i, [0, 4, 1, 5, 2, 6, 3, 7]]
+        x = no_d[indices, 0]
+        y = no_d[indices, 1]
+        plt.fill(x, y, 'green', edgecolor='black', alpha=0.3)
+    for i in pileId:
+        x_coords = no_d[el[i, [0, 4, 1, 5, 2, 6, 3, 7]], 0]
+        y_coords = no_d[el[i, [0, 4, 1, 5, 2, 6, 3, 7]], 1]
+        plt.fill(x_coords, y_coords, 'red')
+    if anchor_Id.size == 1:
+        plt.scatter(no_d[anchor_Id, 0], no_d[anchor_Id, 1], 100, color='white', marker='s', facecolors='none')
+    plt.xlabel('Width [m]', fontsize=14)
+    plt.ylabel('Height [m]', fontsize=14)
+    plt.axis('equal')
+    plt.tight_layout()
+    plt.show()
+
+
+# function 54: single Monte Carlo iteration (groundwater-displacement-factor of safety)
+def single_run(ks_soil, phi_soil, coh_soil, L, free, pile, inter, phi_pile, phi_inter, dilation_factor,
+               psi_pile, psi_inter, coh_pile, coh_inter, equivalent_pile_thickness, lip, active_pile,
+               active_inter, H0, fh, fs, no, nc, el, ec, df, E, nu, Km_d, ad, lim, tol):
+    # Generate random field instances
+    # ks, phi, coh = random_field_instance(ks_soil, phi_soil, coh_soil, ec, L)
+    ks, phi, coh = uniform_field_instance(ks_soil, phi_soil, coh_soil, ec)
+    # Update hydraulic conductivity at integration points
+    ks = hydraulic_conductivity_integration_point(free, pile, ks, equivalent_pile_thickness, lip, active_pile)
+    # Update hydraulic conductivity at integration points
+    phi, psi, coh = plasticity_parameters_integration_point(free, pile, inter, phi, phi_pile, phi_inter,
+                                                            dilation_factor, psi_pile, psi_inter, coh, coh_pile,
+                                                            coh_inter, equivalent_pile_thickness, lip, active_pile,
+                                                            active_inter)
+    # form conductivity matrix
+    Kc = conductivity_stiffness(no, el, ec, ks)
+    # hydraulic solver
+    H = hydraulic_solver(H0, Kc, no, el, fh, fs, lim, tol)
+    # Darcy flow
+    vee, flow_rate = flow_darcy(H, no, el, ec, ks)
+    # Calculate the norm of the flow rate vector for the entire mesh
+    normal_flow_rate = np.linalg.norm(flow_rate)
+    # hydromechanical coupling
+    nodal_pore_pressure = hydro_mechanical_coupling(H, no, nc, el, ec, df)
+    # add all external forces
+    Fext = nodal_self_weight + nodal_water_pressure + nodal_pore_pressure
+    # find safety factor
+    U, evpt1, evpt2, evpt3, evpt4, sigma, factor_of_safety = find_safety_factor(no, nc, el, ec, df, phi, psi, coh, E,
+                                                                                nu, Fext, Km_d, ad, lim, tol)
+    # Excluding restricted nodes and boundary nodes
+    ad_x = (ad[np.where(ad % 2 == 0)] / 2).astype(int)
+    ad_y = (ad[np.where(ad % 2 != 0)] / 2).astype(int)
+    # intersection between ad_x and ad_y
+    ad_x = np.intersect1d(ad_x, ad_y)
+    # take out the free surface nodes
+    mask = np.in1d(ad_x, free_surface_node)
+    ad_x = ad_x[~mask]
+    # collect input coordinates and input features
+    input_coords = no[ad_x]
+    input_features = np.array([E_soil_input, phi_soil_input, coh_soil_input, ks_soil_input]).reshape([1, 4])
+    # collect output features
+    displ = np.stack([U[0:2 * nc:2], U[1:2 * nc:2]], axis=1)
+    displ_ad = displ[ad_x]
+    # TODO: add pore water pressure or ground water flow velocity
+    output_features = np.stack([displ_ad[:, 0], displ_ad[:, 1]], axis=1)
+
+    # save input coordinates and input and output features into a npy file
+    np.save(f'{sim_name}_{description}_ml_input_coords.npy', input_coords)
+    np.save(f'{sim_name}_{description}_ml_input_feature.npy', input_features)
+    np.save(f'{sim_name}_{description}_ml_output_feature.npy', output_features)
+
+    # save simulation data and parameter data for the inference
+    # uniformly pick 1000 points from the input coordinates
+    n_iter = int(input_coords.shape[0] / num_obs)
+    data_file_name = f'{sim_name}_{description}_sim.txt'
+    sim_data = {'t': [1]}
+    for i, _ in enumerate(input_coords[::n_iter, 0]):
+        sim_data[f'displ_x_{i}'] = [displ_ad[i, 0]]
+        sim_data[f'displ_y_{i}'] = [displ_ad[i, 1]]
+        # sim_data[f'pwp_x_{i}'] = [pwp_ad[i, 1]]
+        # sim_data[f'pwp_y_{i}'] = [pwp_ad[i, 1]]
+
+    write_dict_to_file(sim_data, data_file_name)
+
+    data_param_name = f'{sim_name}_{description}_param.txt'
+    param_data = {'E': [E_soil_input], 'phi': [phi_soil_input], 'coh': [coh_soil_input], 'ks': [ks_soil_input]}
+    write_dict_to_file(param_data, data_param_name)
+
+    # return normalized flow, factor of safety, displacement, viscoplastic strain components, permeability and total head
+    return normal_flow_rate, factor_of_safety, nodal_pore_pressure, U, evpt1, evpt2, evpt3, evpt4, phi, psi, coh, sigma, ks, H
+
+
+# function 55: parallel task Monte Carlo Iterations
+def parallel_task(i):
+    print(f"working on iteration {i + 1}")
+    normal_flow_rate_add, factor_of_safety_add, nodal_pore_pressure, U, evpt1, evpt2, evpt3, evpt4, phi, psi, coh, sigma, ks, H = single_run(
+        ks_soil, phi_soil, coh_soil, L, free, pile, inter, phi_pile, phi_inter, dilation_factor, psi_pile, psi_inter,
+        coh_pile, coh_inter, equivalent_pile_thickness, lip, active_pile, active_inter, H0, fh, fs, no, nc, el, ec, df,
+        E, nu, Km_d, ad, lim, tol)
+    return normal_flow_rate_add, factor_of_safety_add
+
+
+# function 56: parallel execution Monte Carlo Iterations
+def execute_parallel(mc, num_cores):
+    # Create a Pool of workers to parallelize the task
+    with Pool(num_cores) as pool:
+        results = []
+        # Submit tasks asynchronously using apply_async
+        for i in range(mc):
+            result = pool.apply_async(parallel_task, args=(i,))
+            results.append(result)
+        # Retrieve results from async tasks
+        results = [res.get() for res in results]
+    # Extract normal_flow_rate and factor_of_safety from the results
+    normal_flow_rate = np.array([result[0] for result in results])
+    factor_of_safety = np.array([result[1] for result in results])
+    # Return the computed normal_flow_rate and factor_of_safety
+    return normal_flow_rate, factor_of_safety
+
+
+# def interpolate_from_ipts_to_nodes(field):
+#     # Obtain integration points in local coordinates
+#     xi, yi = integration_point()
+#     # Initialize arrays for nodal values
+#     interpolated_field = np.zeros([nc, 2])
+#     # Iterate over each element in the mesh
+#     for i in range(ec):
+#         # Iterate over each integration point
+#         for j in range(4):
+#             # Calculate the index for the hydraulic conductivity at the integration point
+#             id = (i * 4) + j
+#             # Calculate the shape tensor for the integration point
+#             N = shape_tensor4(xi[j], yi[j])
+#             interpolated_field[el[i, 0:4], 0] += N.T * field[id ,0]
+#             interpolated_field[el[i, 0:4], 1] += N.T * field[id ,1]
+#     # Return the interpolated field
+#     return interpolated_field
+
+
+def write_dict_to_file(data, file_name):
+    """
+    write a python dictionary data into a text file
+    """
+    with open(file_name, 'w') as f:
+        keys = data.keys()
+        f.write('# ' + ' '.join(keys) + '\n')
+        num = len(data[list(keys)[0]])
+        for i in range(num):
+            f.write(' '.join([str(data[key][i]) for key in keys]) + '\n')
+
+
+# %% input parameters from command line
+
+# Variables received from the user
+
+# Extract parameters from sys.argv
+# sys.argv[0] is the script name
+E_soil_input = float(sys.argv[1])  # [kPa] Young's modulus of soil
+phi_soil_input = float(sys.argv[2])  # [deg] friction angle of soil
+coh_soil_input = float(sys.argv[3])  # [kPa] cohesion of soil
+ks_soil_input = 0.01  # [m/d] hydraulic conductivity of soil
+sim_name = sys.argv[4]  # name of the simulation (dike2D)
+description = sys.argv[5]  # description of the simulation
+num_obs = int(sys.argv[6])  # number of observations for the inference or uncertainty quantification
+
+# %% Geometry Input
+free = np.array([
+    [0.0, 10.0],
+    [10.0, 10.0],
+    [14.0, 14.0],
+    [18.0, 14.0],
+    [30.0, 8.0],
+    [50.0, 8.0],
+])  # [m] free surface coordinates
+
+wl = np.array([
+    [0.0, 12.0],
+    [14.0, 12.0],
+    [30.0, 7.5],
+    [50.0, 7.5],
+])  # [m] water level surface coordinates
+
+disc = [1.0, 0.5]  # [m] FEM element size
+
+# %% Soil Input
+clx = 1e10  # [m] correlation length x-dir
+cly = 1e10  # [m] correlation length y-dir
+E_soil = E_soil_input  # [kPa] Young's modulus soil
+nu_soil = 0.30  # [-] Poison's ratio soil
+gamma_soil = 20.0  # [kN/m3] specific weight soil
+phi_soil = [phi_soil_input, phi_soil_input]  # [deg] friction angle [min|max]
+coh_soil = [coh_soil_input, 0]  # [kPa] cohesion [mean|standard deviation]
+ks_soil = [0.01, 0]  # [m/d] hydraulic conductivity [mean|standard deviation]
+dilation_factor = 0.1  # [psi/phi[-]] dilation angle = dilation_factor*friction angle
+
+# %% Pile Input
+active_pile = False  # [True|False] activate pile
+pile = [3, 6.0]  # [node|depth[m]] pile node and depth
+I_pile = 5e-5  # [m4] surface moment of inertia
+E_pile = 200e6  # [kPa] Young's modulus pile
+nu_pile = 0.20  # [-] Poison's ratio pile
+gamma_pile = 78.5  # [kN/m3] specific weight pile
+phi_pile = 0.0  # [deg] friction angle pile
+psi_pile = 0.0  # [deg] dilation angle pile
+coh_pile = 125e3  # [kPa] cohesion pile
+
+# %% Anchor Input
+active_anchor = False  # [True|False] activate anchor of the pile
+anchor = [12.0, 5.0, 8.0]  # [connection to pile y-axis [m]|base anchor x-axis[m]|base anchor y-axis[m]]
+E_anchor = 200e6  # [kPa] Young's modulus anchor
+d_anchor = 0.15  # [m] diameter anchor
+HoH_anchor = 2.5  # [m] Center on center distance anchor (out of plane direction)
+
+# %% Interface Input
+active_inter = False  # [True|False] activate interface
+inter = [3, 4, 3.0]  # [node 1|node 2|depth[m]] start interface node, end interface node, depth of the interface
+E_inter = 0.5e5  # [kPa] Young's modulus interface
+nu_inter = 0.28  # [-] Poison's ratio interface
+gamma_inter = 20.0  # [kN/m3] specific weight interface
+phi_inter = 18.0  # [deg] friction angle interface
+psi_inter = 0.0  # [deg] dilation angle interface
+coh_inter = 8.5  # [kPa] cohesion interface
+
+# %% Solver Input
+lim = 500  # [-] iteration limit
+tol = 1e-4  # [-] iteration tolerance
+def_scaling = 0.05  # [-] deformation scaling
+num_cores = 18  # [-] Number of Cores to be used
+
+# %% main script:
+
+# Calculate equivalent pile thickness
+equivalent_pile_thickness = equivalent_thickness_rectangle(I_pile)
+# Find intersection points
+intersection_point_list = find_intersection_points(free, wl)
+# Partition the domain
+X, Y = divide_FEM_domain(free, disc, pile, equivalent_pile_thickness, intersection_point_list, active_pile)
+# Generate the mesh
+no, nc, el, ec = mesh_generator(X, Y)
+# Elevate the mesh
+no, free_surface_node = elevate_mesh(free, X, no, el)
+# add anchor
+anchor_Id, anchorSpringX, anchorSpringY = anchor_spring(free, pile, anchor, E_anchor, d_anchor, HoH_anchor, no,
+                                                        active_anchor)
+# element degree of freedom array
+df = degree_freedom(el, ec)
+# Dirichlet boundary conditions
+ad, fh, fs = boundary_dirichlet(wl, X, no, nc, el, free_surface_node)
+# Neumann boundary conditions
+nodal_water_pressure = boundary_neumann(wl, no, nc, free_surface_node)
+# integration point location
+lip = integration_point_location(no, el, ec)
+# # generate random field base
+# L = random_field(lip[:, 0], lip[:, 1], clx, cly)
+L = None
+# elasticity parameters
+E, nu = elasticity_parameters_integration_point(free, pile, inter, E_soil, E_pile, E_inter, nu_soil, nu_pile, nu_inter,
+                                                equivalent_pile_thickness, ec, lip, active_pile, active_inter)
+# displacement stiffness matrix
+Km_d = displacement_stiffness(no, nc, el, ec, df, ad, E, nu, anchor_Id, anchorSpringX, anchorSpringY)
+# self weight integration point
+gamma = specificWeightIntegrationPoint(free, pile, inter, gamma_soil, gamma_pile, gamma_inter,
+                                       equivalent_pile_thickness, ec, lip, active_pile, active_inter)
+# self weight vector
+nodal_self_weight = body_weight(no, nc, el, ec, gamma)
+# initial total head
+H0 = initial_head(wl, X, no, el, free_surface_node)
+
+if __name__ == '__main__':
+    start_time = time.time()
+    normal_flow_rate, factor_of_safety, nodal_pore_pressure, U, evpt1, evpt2, evpt3, evpt4, phi, psi, coh, sigma, ks, H = single_run(
+        ks_soil, phi_soil, coh_soil, L, free, pile, inter, phi_pile, phi_inter, dilation_factor, psi_pile, psi_inter,
+        coh_pile, coh_inter, equivalent_pile_thickness, lip, active_pile, active_inter, H0, fh, fs, no, nc, el, ec, df,
+        E, nu, Km_d, ad, lim, tol)
+    end_time = time.time()
+    print(f"Single run completed: {end_time - start_time} seconds")
diff --git a/tutorials/physics_based/2D_dike_stability/2D_dike_model_inference.py b/tutorials/physics_based/2D_dike_stability/2D_dike_model_inference.py
new file mode 100644
index 00000000..4d6ab461
--- /dev/null
+++ b/tutorials/physics_based/2D_dike_stability/2D_dike_model_inference.py
@@ -0,0 +1,121 @@
+"""
+This tutorial shows how to perform iterative Bayesian calibration for a dike regression model
+ using GrainLearning.
+"""
+import os
+from math import floor, log
+
+import numpy as np
+
+from grainlearning import BayesianCalibration
+from grainlearning.dynamic_systems import IODynamicSystem
+from grainlearning.tools import plot_pdf
+import matplotlib.pylab as plt
+
+from joblib import Parallel, delayed
+
+PATH = os.path.abspath(os.path.dirname(__file__))
+executable = f'python {PATH}/2D_dike.py'
+
+
+def run_sim(calib):
+    """
+    Run the external executable and passes the parameter sample to generate the output file.
+    """
+    system = calib.system
+    # keep the naming convention consistent between iterations
+    mag = floor(log(system.num_samples, 10)) + 1
+    # check the software name and version
+    print("*** Running external software... ***\n")
+    # loop over and pass parameter samples to the executable
+    Parallel(n_jobs=num_cores)(
+        delayed(dike)(params, system.sim_name, 'Iter' + str(system.curr_iter) + '_Sample' + str(i).zfill(mag)) for
+        i, params in enumerate(system.param_data))
+
+
+def dike(params, sim_name, description):
+    print(" ".join([executable, "%.8e %.8e %.8e" % tuple(params), sim_name, description, str(num_obs)]))
+    os.system(' '.join([executable, "%.8e %.8e %.8e" % tuple(params), sim_name, description, str(num_obs)]))
+
+
+# overwrite a virtual function get_normalization_factor of the smc class
+def set_normalization_factor(self):
+    self.normalization_factor = 1. / abs(np.mean(self.obs_data, axis=0))
+
+
+#
+IODynamicSystem.set_normalization_factor = set_normalization_factor
+
+# %% Generate synthetic data
+sim_name = 'dike'
+description = 'synth_data'
+true_params = [1e5, 30, 5]
+coef_of_variation = 0.2
+num_obs = 100
+num_cores = 18
+param_mins = list((1 - 1 * coef_of_variation) * np.array(true_params))
+param_maxs = list((1 + 3 * coef_of_variation) * np.array(true_params))
+os.system(' '.join([executable, "%.8e %.8e %.8e" % tuple(true_params), sim_name, description, str(num_obs)]))
+
+# read key of the synthetic data from the first line
+with open(f'{sim_name}_{description}_sim.txt', 'r') as file:
+    # Read the first line
+    first_line = file.readline()
+ctrl_name = first_line.split()[1]
+obs_names = first_line.split()[2:]
+
+param_names = ['E', 'phi', 'coh']
+num_samples = int(5 * len(param_names) * log(len(param_names)))
+
+# %% Define the calibration
+calibration = BayesianCalibration.from_dict(
+    {
+        "num_iter": 4,
+        "callback": run_sim,
+        "system": {
+            "system_type": IODynamicSystem,
+            "param_min": param_mins,
+            "param_max": param_maxs,
+            "param_names": param_names,
+            "num_samples": num_samples,
+            "obs_data_file": PATH + '/dike_synth_data_sim.txt',
+            "obs_names": obs_names,
+            "ctrl_name": 't',
+            "sim_name": sim_name,
+            "sim_data_dir": PATH + '/sim_data/',
+            "sim_data_file_ext": '.txt',
+            "sigma_tol": 0.01,
+            "sigma_min": 0.01,
+            "sigma_max": 10,
+        },
+        "calibration": {
+            "inference": {"ess_target": 0.3},
+            "sampling": {
+                "max_num_components": 1,
+                "n_init": 1,
+                "random_state": 0,
+                "slice_sampling": True,
+            },
+            "initial_sampling": "halton",
+        },
+        "save_fig": 0,
+    }
+)
+
+# %% Run the calibration
+calibration.run()
+
+most_prob_params = calibration.get_most_prob_params()
+print(f'Most probable parameter values: {most_prob_params}')
+
+error_tolerance = 0.1
+errors = most_prob_params - true_params
+
+for error, true_param, prob_param in zip(errors, true_params, most_prob_params):
+    assert abs(
+        error) / true_param < error_tolerance, \
+        f"Model parameters are not correct, expected {true_param} but got {prob_param}"
+
+plot_pdf('2D_dike', param_names, calibration.calibration.param_data_list, save_fig=0,
+         true_params=true_params)
+plt.show()
diff --git a/tutorials/physics_based/nosand_triax/norsand_triax_calibration.py b/tutorials/physics_based/nosand_triax/norsand_triax_calibration.py
new file mode 100644
index 00000000..f4d8ac2c
--- /dev/null
+++ b/tutorials/physics_based/nosand_triax/norsand_triax_calibration.py
@@ -0,0 +1,172 @@
+import numpy as np
+from grainlearning import BayesianCalibration
+from math import log
+from grainlearning.tools import plot_pdf
+import matplotlib.pylab as plt
+
+
+def norsand(x, e0, params):
+    # norsand parameters
+    gamma = params[0]  # [0.9 - 1.4] 'altitude of CSL @ 1kPa'
+    lambda_val = params[1]  # [0.01 - 0.07] 'slope CSL defined on natural log'
+    M_tc = params[2]  # [1.2 - 1.5] 'critical state friction ratio triaxial compression'
+    N = params[3]  # [0.2 - 0.5] 'volumetric coupling coefficient'
+    H = params[4]  # [25 - 500] 'plastic hardening modulus for loading'
+    Xim_tc = params[5]  # [2 - 5] 'relates maximum dilatancy to state variable (psi)'
+    Ir = params[6]  # [100 - 600] 'shear rigidity'
+    nu = params[7]  # [0.1 0.3] 'Poissons ratio'
+
+    # Derived parameters
+    eqp_inc = eqp_tot / (1e2 * lst)  # [-] 'increment applied plastic deviatoric strain'
+    ratio = 1 / (1 - N) ** ((N - 1) / N)  # [-] 'ratio mean effective stress (p) and image stress (pim)'
+    Xim = Xim_tc / (
+        1 - Xim_tc * lambda_val / M_tc)  # [-] 'relationship between soil property Xi_tc to Norsand property Xi'
+
+    # Declarations
+    p = np.zeros(lst + 1)
+    q = np.zeros(lst + 1)
+    pim = np.zeros(lst + 1)
+    ev = np.zeros(lst + 1)
+    eq = np.zeros(lst + 1)
+    e = np.zeros(lst + 1)
+    psi = np.zeros(lst + 1)
+
+    # Initial conditions
+    p[0] = p0
+    pim[0] = ratio * p0
+    e[0] = e0
+    psi[0] = e0 - (gamma - lambda_val * np.log(pim[0])) + lambda_val * np.log(pim[0] / p[0])
+
+    # Loadstep cycle CD test
+    for i in range(lst):
+        # Update image state
+        e[i + 1] = e0 - (1 + e0) * ev[i]
+        psi[i + 1] = e[i + 1] - (gamma - lambda_val * np.log(pim[i])) + lambda_val * np.log(pim[i] / p[i])
+
+        # Update image friction ratio
+        Mim = M_tc + Xim * N * np.abs(psi[i + 1])
+
+        # Apply hardening increment
+        pim_max = p[i] * np.exp(-Xim * psi[i + 1] / M_tc)
+        pim_inc = H * (pim_max - pim[i]) * eqp_inc
+        pim[i + 1] = pim[i] + pim_inc
+
+        # Calculate plastic volumetric strain increment
+        Dp = Mim - q[i] / p[i]
+        evp_inc = Dp * eqp_inc
+
+        # Calculate bulk and shear modulus
+        mu = Ir * p[i]
+        K = mu * (2 * (1 + nu)) / (3 * (1 - 2 * nu))
+
+        # Apply consistency condition to calculate stress ratio increment (drained)
+        eta_inc = (pim_inc / pim[i]) / (
+            1 / M_tc + 1 / (3.0 - q[i] / p[i]))  # Note: 3.0 is used instead of 3 in the denominator
+        # Calculate mean effective stress
+        p_inc = p[i] * eta_inc / (3.0 - q[i] / p[i])
+        p[i + 1] = p[i] + p_inc
+
+        # Calculate new stress ratio and shear stress
+        eta = M_tc * (1 + np.log(pim[i + 1] / p[i + 1]))
+        q[i + 1] = eta * p[i + 1]
+
+        # Update volumetric and deviatoric strain increments
+        eve_inc = p_inc / K
+        eqe_inc = (q[i + 1] - q[i]) / (3 * mu)
+        ev[i + 1] = ev[i] + eve_inc + evp_inc
+        eq[i + 1] = eq[i] + eqe_inc + eqp_inc
+
+    return np.stack([(q / p)[::int(lst / 100)], e[::int(lst / 100)]], axis=0)
+
+
+# CD input
+p0 = 130.0  # [kPa] 'inital pressure'
+e0 = [0.60, 0.80]  # [-] 'init void ratio'
+eqp_tot = 15.0  # [%] 'total applied plastic deviatoric strain'
+lst = int(1e3)  # [-] 'loadsteps'
+
+# ground truth parameter
+gamma = 0.816  # [0.7 - 1.4] 'altitude of CSL @ 1kPa'
+lambda_val = 0.015  # [0.01 - 0.07] 'slope CSL defined on natural log'
+M_tc = 1.26  # [1.2 - 1.5] 'critical state friction ratio triaxial compression'
+N = 0.4  # [0.2 - 0.5] 'volumetric coupling coefficient'
+H = 200.0  # [25 - 500] 'plastic hardening modulus for loading'
+Xim_tc = 3.8  # [2 - 5] 'relates maximum dilatancy to state variable (psi)'
+Ir = 200.0  # [100 - 600] 'shear rigidity'
+nu = 0.20  # [0.1 0.3] 'Poissons ratio'
+
+# generate synthetic data
+x_obs = np.arange(lst + 1) / lst * eqp_tot
+x_obs = x_obs[::int(lst / 100)]
+y_obs_1 = norsand(x_obs, e0[0], [gamma, lambda_val, M_tc, N, H, Xim_tc, Ir, nu])
+y_obs_2 = norsand(x_obs, e0[1], [gamma, lambda_val, M_tc, N, H, Xim_tc, Ir, nu])
+y_obs = np.concatenate([y_obs_1, y_obs_2], axis=0)
+
+
+def run_sim_original(x, params):
+    """Run different realizations of the original model.
+
+    :param x: the input sequence
+    :param params: the parameters
+    """
+    data = []
+    for params in params:
+        y_i = []
+        # Run the model
+        for e in e0:
+            y_i.append(norsand(x, e, params))
+        data.append(np.concatenate(y_i, axis=0))
+    return np.array(data)
+
+
+def run_sim(calib):
+    """This is the callback function that runs different realizations of the same model.
+
+    :param calib: The calibration object.
+    """
+    sim_data = run_sim_original(calib.system.ctrl_data, calib.system.param_data)
+    calib.system.set_sim_data(sim_data)
+
+
+num_cores = 18
+param_names = ['gamma', 'lambda', 'M', 'N', 'H', 'Xim', 'Ir', 'nu']
+num_samples = int(5 * len(param_names) * log(len(param_names)))
+
+calibration = BayesianCalibration.from_dict(
+    {
+        "num_iter": 8,
+        "callback": run_sim,
+        "system": {
+            "param_min": [0.7, 0.01, 1.2, 0.2, 25, 2, 100, 0.1],
+            "param_max": [1.4, 0.07, 1.5, 0.5, 500, 5, 600, 0.3],
+            "param_names": param_names,
+            "num_samples": num_samples,
+            "obs_names": ['q/p (0.6)', 'e (0.6)', 'q/p (0.8)', 'e (0.8)'],
+            "ctrl_name": 'u',
+            "obs_data": y_obs,
+            "ctrl_data": x_obs,
+            "sim_name": 'triax',
+            "sigma_tol": 0.01,
+        },
+        "calibration": {
+            "inference": {
+                "ess_target": 0.3,
+                "scale_cov_with_max": True,
+            },
+            "sampling": {
+                "max_num_components": 10,
+                "n_init": 1,
+                "random_state": 0,
+                "slice_sampling": False,
+            },
+            "initial_sampling": "sobol",
+        },
+        "save_fig": 0,
+    }
+)
+
+calibration.run()
+
+true_params = [gamma, lambda_val, M_tc, N, H, Xim_tc, Ir, nu]
+plot_pdf('norsand_triax', param_names, calibration.calibration.param_data_list, save_fig=0, true_params=true_params)
+plt.show()

From 9c1a35eb3d6e31425458f81afad97747cf37ab21 Mon Sep 17 00:00:00 2001
From: chyalexcheng <chyalexcheng@gmail.com>
Date: Sat, 3 Aug 2024 00:54:32 +0200
Subject: [PATCH 19/21] add a 2D displacement field plot to the dike tutorial

---
 grainlearning/bayesian_calibration.py         |  2 +-
 grainlearning/tools.py                        |  2 +-
 .../2D_dike_stability/2D_dike.py              |  9 ++-
 .../2D_dike_model_inference.py                | 57 +++++++++++++++++--
 4 files changed, 59 insertions(+), 11 deletions(-)

diff --git a/grainlearning/bayesian_calibration.py b/grainlearning/bayesian_calibration.py
index fa4115fb..39c6a36d 100644
--- a/grainlearning/bayesian_calibration.py
+++ b/grainlearning/bayesian_calibration.py
@@ -121,7 +121,7 @@ def run(self):
             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:
+            if abs(self.system.sigma_max - self.system.sigma_tol) / self.system.sigma_tol < 1e-2:
                 self.num_iter = self.curr_iter + 1
                 break
 
diff --git a/grainlearning/tools.py b/grainlearning/tools.py
index 226ad4f1..02d6d29c 100644
--- a/grainlearning/tools.py
+++ b/grainlearning/tools.py
@@ -324,7 +324,7 @@ def stratified_resample(weights, expand_num=10):
 def systematic_resample(weights, expand_num=10):
     """ Performs the systemic resampling algorithm used by particle filters.
     This algorithm separates the sample space into N divisions. A single random
-    offset is used to to choose where to sample from for all divisions. This
+    offset is used to choose where to sample from for all divisions. This
     guarantees that every sample is exactly 1/N apart.
     Parameters
     ----------
diff --git a/tutorials/physics_based/2D_dike_stability/2D_dike.py b/tutorials/physics_based/2D_dike_stability/2D_dike.py
index c28b6a66..1c5ec592 100644
--- a/tutorials/physics_based/2D_dike_stability/2D_dike.py
+++ b/tutorials/physics_based/2D_dike_stability/2D_dike.py
@@ -1521,16 +1521,15 @@ def single_run(ks_soil, phi_soil, coh_soil, L, free, pile, inter, phi_pile, phi_
     output_features = np.stack([displ_ad[:, 0], displ_ad[:, 1]], axis=1)
 
     # save input coordinates and input and output features into a npy file
-    np.save(f'{sim_name}_{description}_ml_input_coords.npy', input_coords)
-    np.save(f'{sim_name}_{description}_ml_input_feature.npy', input_features)
-    np.save(f'{sim_name}_{description}_ml_output_feature.npy', output_features)
+    np.save(f'{sim_name}_{description}_input_coords.npy', input_coords)
+    np.save(f'{sim_name}_{description}_input_feature.npy', input_features)
+    np.save(f'{sim_name}_{description}_output_feature.npy', output_features)
 
     # save simulation data and parameter data for the inference
     # uniformly pick 1000 points from the input coordinates
-    n_iter = int(input_coords.shape[0] / num_obs)
     data_file_name = f'{sim_name}_{description}_sim.txt'
     sim_data = {'t': [1]}
-    for i, _ in enumerate(input_coords[::n_iter, 0]):
+    for i in np.arange(0, input_coords.shape[0], input_coords.shape[0]/num_obs, dtype=int):
         sim_data[f'displ_x_{i}'] = [displ_ad[i, 0]]
         sim_data[f'displ_y_{i}'] = [displ_ad[i, 1]]
         # sim_data[f'pwp_x_{i}'] = [pwp_ad[i, 1]]
diff --git a/tutorials/physics_based/2D_dike_stability/2D_dike_model_inference.py b/tutorials/physics_based/2D_dike_stability/2D_dike_model_inference.py
index 4d6ab461..b5a2211d 100644
--- a/tutorials/physics_based/2D_dike_stability/2D_dike_model_inference.py
+++ b/tutorials/physics_based/2D_dike_stability/2D_dike_model_inference.py
@@ -9,7 +9,7 @@
 
 from grainlearning import BayesianCalibration
 from grainlearning.dynamic_systems import IODynamicSystem
-from grainlearning.tools import plot_pdf
+from grainlearning.tools import plot_pdf, plot_param_data
 import matplotlib.pylab as plt
 
 from joblib import Parallel, delayed
@@ -31,6 +31,52 @@ def run_sim(calib):
     Parallel(n_jobs=num_cores)(
         delayed(dike)(params, system.sim_name, 'Iter' + str(system.curr_iter) + '_Sample' + str(i).zfill(mag)) for
         i, params in enumerate(system.param_data))
+    # plot the samples and the comparison between observation and prediction
+    if calib.curr_iter == 0:
+        return
+
+    # get the path to save the figures
+    path = f'{system.sim_data_dir}/iter{calib.curr_iter}' \
+        if isinstance(system, IODynamicSystem) \
+        else f'./{system.sim_name}/iter{calib.curr_iter}'
+    if not os.path.exists(path):
+        os.makedirs(path)
+    fig_name = f'{path}/{system.sim_name}'
+
+    # get the id of sample that has the highest probability
+    most_prob_params_id = np.argmax(calib.calibration.posterior)
+    sim_data = system.sim_data[most_prob_params_id, :, :]
+    obs_data = system.obs_data
+
+    # compute the error norm
+    sim_data_norm = np.sqrt(sim_data[0::2] ** 2 + sim_data[1::2] ** 2)
+    obs_data_norm = np.sqrt(obs_data[0::2] ** 2 + obs_data[1::2] ** 2)
+    error_norm = sim_data_norm - obs_data_norm
+
+    # plot ground truth, inferred field, and error
+    lims_truth = dict(cmap='RdBu_r', vmin=obs_data_norm.min(), vmax=obs_data_norm.max())
+    lims_error = dict(cmap='RdBu_r', vmin=error_norm.min(), vmax=error_norm.max())
+    fig, ax = plt.subplots(nrows=1, ncols=3, figsize=(12, 4))
+    im0 = ax[0].scatter(coords[:, 0], coords[:, 1], 100, obs_data_norm, edgecolor='w', lw=0.1, **lims_truth)
+    im1 = ax[1].scatter(coords[:, 0], coords[:, 1], 100, sim_data_norm, edgecolor='w', lw=0.1, **lims_truth)
+    fig.colorbar(im1, ax=ax[1])
+    im2 = ax[2].scatter(coords[:, 0], coords[:, 1], 100, error_norm, edgecolor='w', lw=0.1, **lims_error)
+    fig.colorbar(im2, ax=ax[2])
+    for j in range(3):
+        ax[j].set_xlabel('x')
+        ax[j].set_ylabel('y')
+        ax[0].set_title('Ground truth')
+        ax[1].set_title('Inferred field')
+        ax[2].set_title('Error')
+
+    # plot resampling in parameter space
+    plot_param_data(
+        fig_name,
+        system.param_names,
+        calib.calibration.param_data_list,
+        save_fig=0
+    )
+    plt.show()
 
 
 def dike(params, sim_name, description):
@@ -63,6 +109,9 @@ def set_normalization_factor(self):
     first_line = file.readline()
 ctrl_name = first_line.split()[1]
 obs_names = first_line.split()[2:]
+coords = np.load(f'{sim_name}_{description}_input_coords.npy')
+n_iter = int(coords.shape[0] / num_obs)
+coords = coords[np.arange(0, coords.shape[0], coords.shape[0] / num_obs, dtype=int)]
 
 param_names = ['E', 'phi', 'coh']
 num_samples = int(5 * len(param_names) * log(len(param_names)))
@@ -70,7 +119,7 @@ def set_normalization_factor(self):
 # %% Define the calibration
 calibration = BayesianCalibration.from_dict(
     {
-        "num_iter": 4,
+        "num_iter": 10,
         "callback": run_sim,
         "system": {
             "system_type": IODynamicSystem,
@@ -96,9 +145,9 @@ def set_normalization_factor(self):
                 "random_state": 0,
                 "slice_sampling": True,
             },
-            "initial_sampling": "halton",
+            "initial_sampling": "sobol",
         },
-        "save_fig": 0,
+        "save_fig": -1,
     }
 )
 

From 6778f4dc14ccfd6bd6722a56ad3cc25e918bfff0 Mon Sep 17 00:00:00 2001
From: chyalexcheng <chyalexcheng@gmail.com>
Date: Sun, 4 Aug 2024 15:26:40 +0200
Subject: [PATCH 20/21] improve efficiency of slice sampling

---
 grainlearning/dynamic_systems.py |  2 ++
 grainlearning/sampling.py        | 27 +++++++++++++++++++--------
 2 files changed, 21 insertions(+), 8 deletions(-)

diff --git a/grainlearning/dynamic_systems.py b/grainlearning/dynamic_systems.py
index 46a41733..acb83752 100644
--- a/grainlearning/dynamic_systems.py
+++ b/grainlearning/dynamic_systems.py
@@ -151,6 +151,8 @@ def __init__(
 
         self.num_samples = num_samples
 
+        self.num_samples_max = num_samples
+
         self.sim_name = sim_name
 
         self.sim_data = sim_data
diff --git a/grainlearning/sampling.py b/grainlearning/sampling.py
index 33351856..b353777c 100644
--- a/grainlearning/sampling.py
+++ b/grainlearning/sampling.py
@@ -6,6 +6,7 @@
 import numpy as np
 from sklearn.mixture import BayesianGaussianMixture
 from scipy.stats.qmc import Sobol, Halton, LatinHypercube
+from scipy.optimize import minimize_scalar
 from grainlearning.dynamic_systems import DynamicSystem
 
 
@@ -219,17 +220,27 @@ def regenerate_params(self, weight: np.ndarray, system: Type["DynamicSystem"]):
         :param system: Dynamic system class
         :return: Resampled parameter data
         """
+        # Train the model using the provided weight and parameters from the dynamic system
         self.train(weight, system)
 
+        # Minimum number of samples required
         minimum_num_samples = system.num_samples
 
-        new_params = self.draw_samples_within_bounds(system, system.num_samples)
+        def sample_count_obj(num_samples):
+            """Objective function to minimize the number of samples drawn."""
+            num_samples = int(np.ceil(num_samples))  # Ensure we have an integer number of samples
+            samples = self.draw_samples_within_bounds(system, num_samples)
+            valid_samples_count = samples.shape[0]
+            # Return the difference between required and obtained samples
+            return (minimum_num_samples - valid_samples_count) ** 2
 
-        # resample until all parameters are within the upper and lower bounds
-        test_num = system.num_samples
-        while system.param_min and system.param_max and new_params.shape[0] < minimum_num_samples:
-            test_num = int(np.ceil(1.01 * test_num))
-            new_params = self.draw_samples_within_bounds(system, test_num)
+        # Perform optimization to find the minimum number of samples needed
+        result = minimize_scalar(sample_count_obj, bounds=(system.num_samples_max, 100 * system.num_samples_max),
+                                 method='bounded')
+        system.num_samples_max = result.x
+
+        # Draw the required number of samples
+        new_params = self.draw_samples_within_bounds(system, int(np.ceil(result.x)))
 
         return new_params
 
@@ -323,12 +334,12 @@ def generate_params_qmc(system: Type["DynamicSystem"], num_samples: int, method:
     sampler = Halton(system.num_params, scramble=False)
 
     if method == "sobol":
-        sampler = Sobol(system.num_params,seed=seed)
+        sampler = Sobol(system.num_params, seed=seed)
         random_base = round(np.log2(num_samples))
         num_samples = 2 ** random_base
 
     elif method == "LH":
-        sampler = LatinHypercube(system.num_params,seed=seed)
+        sampler = LatinHypercube(system.num_params, seed=seed)
 
     param_table = sampler.random(n=num_samples)
 

From 5ad28e712770a934e93a277255d6276166ba4bdc Mon Sep 17 00:00:00 2001
From: chyalexcheng <chyalexcheng@gmail.com>
Date: Tue, 22 Oct 2024 00:02:28 +0200
Subject: [PATCH 21/21] solve the issue of high cost when generating samples
 with slice sampling

---
 grainlearning/sampling.py                     | 97 ++++++++++---------
 .../2D_dike_model_inference.py                |  5 +-
 2 files changed, 55 insertions(+), 47 deletions(-)

diff --git a/grainlearning/sampling.py b/grainlearning/sampling.py
index b353777c..d85ac8da 100644
--- a/grainlearning/sampling.py
+++ b/grainlearning/sampling.py
@@ -6,7 +6,7 @@
 import numpy as np
 from sklearn.mixture import BayesianGaussianMixture
 from scipy.stats.qmc import Sobol, Halton, LatinHypercube
-from scipy.optimize import minimize_scalar
+from scipy.optimize import minimize_scalar, root_scalar
 from grainlearning.dynamic_systems import DynamicSystem
 
 
@@ -90,6 +90,9 @@ class GaussianMixtureModel:
         scores.mean() - deviation_factor * scores.std(), for slice sampling, defaults to 0.0, optional
     :param gmm: The class of the Gaussian Mixture Model
     :param max_params: Current maximum values of the parameters
+    :param min_params: Current minimum values of the parameters
+    :param expanded_normalized_params: Normalized parameters expanded from the importance weights
+    :param scores: Scores of the Gaussian Mixture Model
     """
 
     def __init__(
@@ -137,10 +140,14 @@ def __init__(
 
         self.max_params = None
 
+        self.min_params = None
+
         self.expanded_normalized_params = None
 
         self.gmm = None
 
+        self.scores = None
+
     @classmethod
     def from_dict(cls: Type["GaussianMixtureModel"], obj: dict):
         """Initialize the class using a dictionary style
@@ -183,13 +190,14 @@ def expand_and_normalize_weighted_samples(self, weights: np.ndarray, system: Typ
         expanded_parameters = system.param_data[indices]
 
         max_params = np.amax(expanded_parameters, axis=0)  # find max along axis
+        min_params = np.amin(expanded_parameters, axis=0)  # find min along axis
 
-        normalized_parameters = (
-            expanded_parameters / max_params
-        )  # and do array broadcasting to divide by max
+        # Normalize the parameters from 0 to 1
+        normalized_parameters = (expanded_parameters - min_params) / (max_params - min_params)
 
         self.expanded_normalized_params = normalized_parameters
         self.max_params = max_params
+        self.min_params = min_params
 
     def train(self, weight: np.ndarray, system: Type["DynamicSystem"]):
         """Train the Gaussian mixture model.
@@ -213,6 +221,8 @@ def train(self, weight: np.ndarray, system: Type["DynamicSystem"]):
 
         self.gmm.fit(self.expanded_normalized_params)
 
+        self.scores = self.gmm.score_samples(self.expanded_normalized_params)
+
     def regenerate_params(self, weight: np.ndarray, system: Type["DynamicSystem"]):
         """Regenerating new samples by training and sampling from a Gaussian mixture model.
 
@@ -232,15 +242,15 @@ def sample_count_obj(num_samples):
             samples = self.draw_samples_within_bounds(system, num_samples)
             valid_samples_count = samples.shape[0]
             # Return the difference between required and obtained samples
-            return (minimum_num_samples - valid_samples_count) ** 2
+            return valid_samples_count - minimum_num_samples
 
         # Perform optimization to find the minimum number of samples needed
-        result = minimize_scalar(sample_count_obj, bounds=(system.num_samples_max, 100 * system.num_samples_max),
-                                 method='bounded')
-        system.num_samples_max = result.x
+        result = root_scalar(sample_count_obj, x0=system.num_samples_max, x1=100 * system.num_samples_max,
+                             method='secant')
+        system.num_samples_max = int(np.ceil(result.root))
 
         # Draw the required number of samples
-        new_params = self.draw_samples_within_bounds(system, int(np.ceil(result.x)))
+        new_params = self.draw_samples_within_bounds(system, int(np.ceil(result.root)))
 
         return new_params
 
@@ -257,15 +267,14 @@ def draw_samples_within_bounds(self, system: Type["DynamicSystem"], num: int = 1
         # use the slice sampling scheme for resampling
         else:
             # compute the minimum of score_samples as the threshold for slice sampling
-            new_params = generate_params_qmc(system, num)
-            new_params /= self.max_params
-
-            scores = self.gmm.score_samples(self.expanded_normalized_params)
+            new_params = self.scale_qmc_to_gmm_range(generate_qmc(system, num))
             new_params = new_params[np.where(
-                self.gmm.score_samples(new_params) > scores.mean() - self.deviation_factor * scores.std())]
+                self.gmm.score_samples(new_params) >= self.scores.mean() - self.deviation_factor * self.scores.std())]
 
-        new_params *= self.max_params
+        # Un-normalize the parameters
+        new_params = new_params * (self.max_params - self.min_params) + self.min_params
 
+        # Ensure the parameters are within the user-defined bounds
         if system.param_min and system.param_max:
             params_above_min = new_params > np.array(system.param_min)
             params_below_max = new_params < np.array(system.param_max)
@@ -277,33 +286,17 @@ def draw_samples_within_bounds(self, system: Type["DynamicSystem"], num: int = 1
         else:
             return new_params
 
-    # def regenerate_params_with_gmm(
-    #     self, posterior_weight: np.ndarray, system: Type["DynamicSystem"]
-    # ) -> np.ndarray:
-    #     """Regenerate the parameters by fitting the Gaussian Mixture model (for testing against the old approach)
-    #
-    #     :param posterior_weight: Posterior found by the data assimilation
-    #     :param system: Dynamic system class
-    #     :return: Expanded parameters
-    #     """
-    #
-    #     new_params, self.gmm = regenerate_params_with_gmm(
-    #         posterior_weight,
-    #         system.param_data,
-    #         system.num_samples,
-    #         self.max_num_components,
-    #         self.weight_concentration_prior,
-    #         self.covariance_type,
-    #         unweighted_resample,
-    #         system.param_min,
-    #         system.param_max,
-    #         self.n_init,
-    #         self.tol,
-    #         self.max_iter,
-    #         self.random_state,
-    #     )
-    #
-    #     return new_params
+    def scale_qmc_to_gmm_range(self, qmc_samples: np.ndarray):
+        """Scale the quasi-random number samples to the range of data generated by the sampleGaussian Mixture Model."""
+
+        # Generate trial samples to estimate the maximum and minimum samples generated by gmm
+        samples, _ = self.gmm.sample(100 * self.expanded_normalized_params.shape[0])
+        max_gmm_samples = np.amax(samples, axis=0)  # find max along axis
+        min_gmm_samples = np.amin(samples, axis=0)  # find min along axis
+
+        # Scale the quasi-random number samples to the range of data generated by the sampleGaussian Mixture Model
+        qmc_samples = qmc_samples * (max_gmm_samples - min_gmm_samples) + min_gmm_samples
+        return qmc_samples
 
     def save_gmm_to_file(self, proposal_data_file: str = "proposal_density.pkl"):
         """Save the Gaussian mixture model to a file."""
@@ -319,7 +312,7 @@ def load_gmm_from_file(self, proposal_data_file: str = "proposal_density.pkl"):
             self.max_params, self.gmm = load(f)
 
 
-def generate_params_qmc(system: Type["DynamicSystem"], num_samples: int, method: str = "halton", seed=None):
+def generate_qmc(system: Type["DynamicSystem"], num_samples: int, method: str = "halton", seed=None):
     """This is the class to uniformly draw samples in n-dimensional space from
     a low-discrepancy sequence or a Latin hypercube.
 
@@ -335,13 +328,29 @@ def generate_params_qmc(system: Type["DynamicSystem"], num_samples: int, method:
 
     if method == "sobol":
         sampler = Sobol(system.num_params, seed=seed)
-        random_base = round(np.log2(num_samples))
+        random_base = int(np.ceil(np.log2(num_samples)))
         num_samples = 2 ** random_base
 
     elif method == "LH":
         sampler = LatinHypercube(system.num_params, seed=seed)
 
     param_table = sampler.random(n=num_samples)
+    return np.array(param_table, ndmin=2)
+
+
+def generate_params_qmc(system: Type["DynamicSystem"], num_samples: int, method: str = "halton", seed=None):
+    """This is the class to uniformly draw samples in n-dimensional space from
+    a low-discrepancy sequence or a Latin hypercube.
+
+    See `Quasi-Monte Carlo <https://docs.scipy.org/doc/scipy/reference/stats.qmc.html>`_.
+
+    :param system: Dynamic system class
+    :param num_samples: Number of samples to draw
+    :param method: Method to use for Quasi-Monte Carlo sampling. Options are "halton", "sobol", and "LH"
+    :param seed: Seed for the random number generator
+    """
+
+    param_table = generate_qmc(system, num_samples, method, seed)
 
     for param_i in range(system.num_params):
         for sim_i in range(num_samples):
diff --git a/tutorials/physics_based/2D_dike_stability/2D_dike_model_inference.py b/tutorials/physics_based/2D_dike_stability/2D_dike_model_inference.py
index b5a2211d..4e41cc0b 100644
--- a/tutorials/physics_based/2D_dike_stability/2D_dike_model_inference.py
+++ b/tutorials/physics_based/2D_dike_stability/2D_dike_model_inference.py
@@ -133,15 +133,14 @@ def set_normalization_factor(self):
             "sim_name": sim_name,
             "sim_data_dir": PATH + '/sim_data/',
             "sim_data_file_ext": '.txt',
-            "sigma_tol": 0.01,
-            "sigma_min": 0.01,
+            "sigma_tol": 0.1,
+            "sigma_min": 0.1,
             "sigma_max": 10,
         },
         "calibration": {
             "inference": {"ess_target": 0.3},
             "sampling": {
                 "max_num_components": 1,
-                "n_init": 1,
                 "random_state": 0,
                 "slice_sampling": True,
             },