DNN surrogate improvements

CHANGELOG.md

@@ -7,6 +7,7 @@
 - Docstrings and typing to Explainer methods
 - Campaign.out property to dynamically capture measured responses "y" or objectives as appropriate
 - Campaign.evaluate method to map optimizer.evaluate method
+- DNN to PyTests
 
 ### Modified
 - Fixed SHAP explainer analysis and visualization functions
@@ -15,6 +16,7 @@
 - Moved Explainer testing from optimizer pytests to campaign pytests
 - Generalized plotting function MOO_results and renamed optim_progress
 - Campaign analysis and plotting methods fixed for 
+- Greatly increased the number of samples used for DNNPosterior, increasing the stability of posterior predictions
 
 ### Removed
 - Removed code chunks regarding unused optional inputs to PDP ICE function imported from SHAP GitHub

obsidian/surrogates/custom_torch.py

@@ -7,6 +7,16 @@
 from torch import Tensor
 
 
+class DNNPosterior(EnsemblePosterior):
+
+    def __init__(self, values: Tensor):
+        super().__init__(values)
+
+    def quantile(self, value: Tensor) -> Tensor:
+        """Quantile of the ensemble posterior"""
+        return self.values.quantile(q=value.to(self.values), dim=-3, interpolation='linear')
+
+
 class DNN(Model):
     def __init__(self,
                  train_X: Tensor,
@@ -50,7 +60,7 @@ def forward(self,
 
     def posterior(self,
                   X: Tensor,
-                  n_sample: int = 1024,
+                  n_sample: int = 16384,
                   output_indices: list[int] = None,
                   observation_noise: bool | Tensor = False) -> Posterior:
         """Calculates the posterior distribution of the model at X"""
@@ -68,7 +78,7 @@ def posterior(self,
         y_out = self.forward(X_sample)
         self.eval()
 
-        post = EnsemblePosterior(y_out[..., output_indices])
+        post = DNNPosterior(y_out[..., output_indices])
 
         return post
 

obsidian/tests/test_optimizer_MOO.py

@@ -35,18 +35,24 @@ def Z0(X_space):
 ]
 
 
-@pytest.mark.parametrize('surrogate', [pytest.param('GP', marks=pytest.mark.fast), 'GPflat', pytest.param('DKL', marks=pytest.mark.slow)])
+@pytest.mark.parametrize('surrogate', [pytest.param('GP', marks=pytest.mark.fast),
+                                       'GPflat',
+                                       pytest.param('DKL', marks=pytest.mark.slow),
+                                       'DNN'])
 def test_optimizer_fit(X_space, surrogate, Z0, serial_test=True):
     optimizer = BayesianOptimizer(X_space, surrogate=surrogate, seed=0, verbose=0)
+
+    tol = 1e-2 if surrogate == 'DNN' else 1e-5
+
     optimizer.fit(Z0, target=target)
     if serial_test:
         save = optimizer.save_state()
         optimizer_2 = BayesianOptimizer.load_state(save)
         optimizer_2.__repr__()
         y_pred = optimizer.predict(optimizer.X_train)
         y_pred_2 = optimizer_2.predict(optimizer.X_train)
-        y_error = ((y_pred_2-y_pred)/y_pred).values
-        assert abs(y_error).max() < 1e-5, 'Prediction error in loading parameters of saved optimizer'
+        y_error = ((y_pred_2-y_pred)/y_pred.max(axis=0)).values
+        assert abs(y_error).max() < tol, 'Prediction error in loading parameters of saved optimizer'
 
 
 # Generate a baseline optimizer to use for future tests

obsidian/tests/test_optimizer_SOO.py

@@ -46,8 +46,11 @@ def test_f_transform(X_space, Z0, f_transform):
 target = Target(name='Response', f_transform='Standard', aim='max')
 
 
-@pytest.mark.parametrize('surrogate', [pytest.param('GP', marks=pytest.mark.fast), 'GPflat',
-                                       'GPprior', pytest.param('DKL', marks=pytest.mark.slow)])
+@pytest.mark.parametrize('surrogate', [pytest.param('GP', marks=pytest.mark.fast),
+                                       'GPflat',
+                                       'GPprior',
+                                       pytest.param('DKL', marks=pytest.mark.slow),
+                                       'DNN'])
 def test_optimizer_fit(X_space, surrogate, Z0, serial_test=True):
 
     optimizer = BayesianOptimizer(X_space, surrogate=surrogate, seed=0, verbose=0)
@@ -56,15 +59,17 @@ def test_optimizer_fit(X_space, surrogate, Z0, serial_test=True):
         # GPflat will fail will a purely categorical space because the design matrix is not p.d.
         return
 
+    tol = 1e-2 if surrogate == 'DNN' else 1e-5
+
     optimizer.fit(Z0, target=target)
     if serial_test:
         save = optimizer.save_state()
         optimizer_2 = BayesianOptimizer.load_state(save)
         optimizer_2.__repr__()
         y_pred = optimizer.predict(optimizer.X_train)
         y_pred_2 = optimizer_2.predict(optimizer.X_train)
-        y_error = ((y_pred_2-y_pred)/y_pred).values
-        assert abs(y_error).max() < 1e-5, 'Prediction error in loading parameters of saved optimizer'
+        y_error = ((y_pred_2-y_pred)/y_pred.max(axis=0)).values
+        assert abs(y_error).max() < tol, 'Prediction error in loading parameters of saved optimizer'
 
 
 # Generate a baseline optimizer to use for future tests