add a flag to allow initializing the spectral kernel from data

- also change the default number of mixtures to 3

xopt/generators/bayesian/models/time_dependent.py

@@ -13,12 +13,14 @@
 from pydantic import Field
 
 from xopt.generators.bayesian.models.standard import StandardModelConstructor
+from xopt.generators.bayesian.utils import get_training_data
 from xopt.vocs import VOCS
 
 
 class TimeDependentModelConstructor(StandardModelConstructor):
     name: str = Field("time_dependent", frozen=True)
     use_spectral_mixture_kernel: bool = True
+    initialize_spectral_kernel_from_data: bool = False
 
     def build_model(
         self,
@@ -39,9 +41,9 @@ def build_model(
         new_input_bounds["time"] = [min_t, max_t]
 
         # set covar modules if not specified -- use SpectralMixtureKernel for time axis
-        # see Kuklev, N., et al. "Online accelerator tuning with adaptive bayesian optimization."
-        # Proc. NAPAC 22 (2022): 842.
-        if len(self.covar_modules) == 0 and self.use_spectral_mixture_kernel:
+        # see Kuklev, N., et al. "Online accelerator tuning with adaptive
+        # bayesian optimization." Proc. NAPAC 22 (2022): 842.
+        if self.use_spectral_mixture_kernel:
             covar_modules = {}
             for name in outcome_names:
                 if len(input_names) == 1:
@@ -55,12 +57,29 @@ def build_model(
                     active_dims=matern_dims,
                     lengthscale_prior=GammaPrior(3.0, 6.0),
                 )
-
-                covar_modules[name] = ProductKernel(
-                    SpectralMixtureKernel(num_mixtures=5, active_dims=time_dim),
-                    matern_kernel,
+                spectral_kernel = SpectralMixtureKernel(
+                    num_mixtures=3, active_dims=time_dim
                 )
-                self.covar_modules = covar_modules
+
+                if self.initialize_spectral_kernel_from_data:
+                    train_X, train_Y, train_Yvar = get_training_data(
+                        new_input_names, name, data
+                    )
+
+                    # can only initialize spectral kernel from data if there are
+                    # more than one training data point
+                    if len(train_X) > 1:
+                        spectral_kernel.initialize_from_data(train_X, train_Y)
+                    else:
+                        raise RuntimeWarning(
+                            "cannot initialize spectral kernel from a "
+                            "single data sample, may negatively impact"
+                            " performance"
+                        )
+
+                covar_modules[name] = ProductKernel(spectral_kernel, matern_kernel)
+
+            self.covar_modules = covar_modules
 
         return super().build_model(
             new_input_names, outcome_names, data, new_input_bounds, dtype, device

xopt/tests/generators/bayesian/test_time_dependent_bo.py

@@ -19,6 +19,31 @@ def test_init(self):
 
     @patch.multiple(TimeDependentBayesianGenerator, __abstractmethods__=set())
     def test_model_generation(self):
+        # test single dim variable space
+        vocs = deepcopy(TEST_VOCS_BASE)
+        vocs.variables = {"x1": [0, 1.0]}
+        gen = TimeDependentBayesianGenerator(vocs=vocs)
+        test_data = deepcopy(TEST_VOCS_DATA)
+        test_data.drop("x2", axis=1, inplace=True)
+
+        time_array = []
+        for i in range(len(test_data)):
+            time_array.append(time.time())
+            time.sleep(0.01)
+
+        test_data["time"] = np.array(time_array)
+
+        model = gen.train_model(test_data)
+
+        # make sure time data is in the last model
+        assert np.all(
+            model.models[-1]
+            .input_transform.untransform(model.models[-1].train_inputs[0])[:, -1]
+            .numpy()
+            == test_data["time"].to_numpy().flatten()
+        )
+
+        # test multi-dim variable space
         gen = TimeDependentBayesianGenerator(vocs=TEST_VOCS_BASE)
         test_data = deepcopy(TEST_VOCS_DATA)