Possible PyTorch implementation of WL kernel

grakel_replace/mixed_single_task_gp_usage_example.py

@@ -15,7 +15,7 @@
 from gpytorch import ExactMarginalLogLikelihood
 from gpytorch.kernels import AdditiveKernel, MaternKernel
 from grakel_replace.optimize import optimize_acqf_graph
-from grakel_replace.torch_wl_kernel import TorchWLKernel
+from grakel_replace.torch_wl_kernel import BoTorchWLKernel
 from grakel_replace.utils import min_max_scale, seed_all
 
 if TYPE_CHECKING:
@@ -63,7 +63,7 @@
     ScaleKernel(CategoricalKernel(
         ard_num_dims=N_CATEGORICAL,
         active_dims=range(N_NUMERICAL, N_NUMERICAL + N_CATEGORICAL))),
-    ScaleKernel(TorchWLKernel(
+    ScaleKernel(BoTorchWLKernel(
         graph_lookup=train_graphs, n_iter=5, normalize=True,
         active_dims=(X.shape[1] - 1,)))
 ]
@@ -82,7 +82,7 @@
 def set_graph_lookup(_gp: SingleTaskGP, new_graphs: list[nx.Graph]) -> Iterator[None]:
     kernel_prev_graphs: list[tuple[Kernel, list[nx.Graph]]] = []
     for kern in _gp.covar_module.sub_kernels():
-        if isinstance(kern, TorchWLKernel):
+        if isinstance(kern, BoTorchWLKernel):
             kernel_prev_graphs.append((kern, kern.graph_lookup))
             kern.set_graph_lookup(new_graphs)
 

grakel_replace/optimize.py

@@ -8,7 +8,7 @@
 import networkx as nx
 import torch
 from botorch.optim import optimize_acqf_mixed
-from grakel_replace.torch_wl_kernel import TorchWLKernel
+from grakel_replace.torch_wl_kernel import BoTorchWLKernel
 
 if TYPE_CHECKING:
     from botorch.acquisition import AcquisitionFunction
@@ -33,13 +33,13 @@ def set_graph_lookup(
     kernel_prev_graphs: list[tuple[Kernel, list[nx.Graph]]] = []
 
     # Determine the modules to update based on the kernel type
-    if isinstance(kernel, TorchWLKernel):
+    if isinstance(kernel, BoTorchWLKernel):
         modules = [kernel]
     else:
         assert hasattr(
             kernel, "sub_kernels"
         ), "Kernel module must have sub_kernels method."
-        modules = [k for k in kernel.sub_kernels() if isinstance(k, TorchWLKernel)]
+        modules = [k for k in kernel.sub_kernels() if isinstance(k, BoTorchWLKernel)]
 
     # Save the current graph lookup and set the new graph lookup
     for kern in modules: