Faster embedded bias examples

docs/examples/bias.py

@@ -12,7 +12,7 @@
 # First, let's import the required functions and classes for this example.
 
 # third party imports
-import math
+import time
 import numpy as np
 import matplotlib.pyplot as plt
 import chaospy
@@ -92,13 +92,20 @@
 
 
 class LinearModel(ForwardModelBase):
+    def __init__(self, name):
+        super().__init__(name)
+        self.pce_order = 1
+        self.bias_dist = chaospy.Normal(0.0, 1.0)
+
+        # generate the polynomial chaos expansion
+        self.expansion = chaospy.generate_expansion(self.pce_order, self.bias_dist)
+
     def interface(self):
         self.parameters = ["a", "b"]
         self.input_sensors = Sensor("x")
         self.output_sensors = Sensor("y", std_model="sigma")
 
     def response(self, inp: dict) -> dict:
-        pce_order = 1
         x = inp["x"]
         m = inp["a"]
         b = inp["b"]
@@ -110,16 +117,17 @@ def response(self, inp: dict) -> dict:
         # define the distribution for the bias term
         b_dist = chaospy.Normal(0.0, b)
         # generate quadrature nodes and weights
-        sparse_quads = chaospy.generate_quadrature(pce_order, b_dist, rule="Gaussian")
+        sparse_quads = chaospy.generate_quadrature(
+            self.pce_order, b_dist, rule="Gaussian"
+        )
         # evaluate the model at the quadrature nodes
         sparse_evals = np.array(
             [np.array((m + node) * x) for node in sparse_quads[0][0]]
         )
-        # generate the polynomial chaos expansion
-        expansion = chaospy.generate_expansion(pce_order, b_dist)
+
         # fit the polynomial chaos expansion
         fitted_sparse = chaospy.fit_quadrature(
-            expansion, sparse_quads[0], sparse_quads[1], sparse_evals
+            self.expansion, sparse_quads[0], sparse_quads[1], sparse_evals
         )
         return {"y": fitted_sparse, "dist": b_dist}
 
@@ -183,7 +191,10 @@ def response(self, inp: dict) -> dict:
 # To solve the problem, we use the specialized solver for embedded models using PCEs.
 
 solver = EmbeddedPCESolver(problem, show_progress=False)
+start_time = time.time()
 inference_data = solver.run(n_steps=200, n_initial_steps=20)
+end_time = time.time()
+print(f"Elapsed time: {end_time - start_time:.2f} seconds")
 
 # %%
 # Finally, we want to plot the results we obtained.

tests/unit_tests/inference/bias/test_likelihood_models.py

@@ -26,13 +26,19 @@ def test_embedded_likelihood_models(self):
 
         # define the forward model
         class LinRe(ForwardModelBase):
+            def __init__(self, name):
+                super().__init__(name)
+
+                self.pce_order = 1
+                self.b_dist = chaospy.Normal(0.0, 1.0)
+                self.expansion = chaospy.generate_expansion(self.pce_order, self.b_dist)
+
             def interface(self):
                 self.parameters = ["a", "b"]
                 self.input_sensors = Sensor("x")
                 self.output_sensors = Sensor("y", std_model="sigma")
 
             def response(self, inp: dict) -> dict:
-                pce_order = 1
                 x = inp["x"]
                 m = inp["a"]
                 b = inp["b"]
@@ -45,17 +51,15 @@ def response(self, inp: dict) -> dict:
                 b_dist = chaospy.Normal(0.0, b)
                 # generate quadrature nodes and weights
                 sparse_quads = chaospy.generate_quadrature(
-                    pce_order, b_dist, rule="Gaussian"
+                    self.pce_order, b_dist, rule="Gaussian"
                 )
                 # evaluate the model at the quadrature nodes
                 sparse_evals = np.array(
                     [np.array((m + node) * x) for node in sparse_quads[0][0]]
                 )
-                # generate the polynomial chaos expansion
-                expansion = chaospy.generate_expansion(pce_order, b_dist)
                 # fit the polynomial chaos expansion
                 fitted_sparse = chaospy.fit_quadrature(
-                    expansion, sparse_quads[0], sparse_quads[1], sparse_evals
+                    self.expansion, sparse_quads[0], sparse_quads[1], sparse_evals
                 )
                 return {"y": fitted_sparse, "dist": b_dist}
 

tests/unit_tests/inference/bias/test_solver.py

@@ -22,13 +22,19 @@ def test_embedded_pce_solver(self):
 
         # define the forward model
         class LinRe(ForwardModelBase):
+            def __init__(self, name):
+                super().__init__(name)
+
+                self.pce_order = 1
+                self.b_dist = chaospy.Normal(0.0, 1.0)
+                self.expansion = chaospy.generate_expansion(self.pce_order, self.b_dist)
+
             def interface(self):
                 self.parameters = ["a", "b"]
                 self.input_sensors = Sensor("x")
                 self.output_sensors = Sensor("y", std_model="sigma")
 
             def response(self, inp: dict) -> dict:
-                pce_order = 1
                 x = inp["x"]
                 m = inp["a"]
                 b = inp["b"]
@@ -41,17 +47,15 @@ def response(self, inp: dict) -> dict:
                 b_dist = chaospy.Normal(0.0, b)
                 # generate quadrature nodes and weights
                 sparse_quads = chaospy.generate_quadrature(
-                    pce_order, b_dist, rule="Gaussian"
+                    self.pce_order, b_dist, rule="Gaussian"
                 )
                 # evaluate the model at the quadrature nodes
                 sparse_evals = np.array(
                     [np.array((m + node) * x) for node in sparse_quads[0][0]]
                 )
-                # generate the polynomial chaos expansion
-                expansion = chaospy.generate_expansion(pce_order, b_dist)
                 # fit the polynomial chaos expansion
                 fitted_sparse = chaospy.fit_quadrature(
-                    expansion, sparse_quads[0], sparse_quads[1], sparse_evals
+                    self.expansion, sparse_quads[0], sparse_quads[1], sparse_evals
                 )
                 return {"y": fitted_sparse, "dist": b_dist}