fix: ruff format

neps/optimizers/models/graphs/context_managers.py

@@ -33,15 +33,20 @@ def set_graph_lookup(
 
     # Determine the modules to update based on the input type
     if isinstance(kernel_or_gp, SingleTaskGP):
-        modules = [k for k in kernel_or_gp.covar_module.sub_kernels() if
-                   isinstance(k, BoTorchWLKernel)]
+        modules = [
+            k
+            for k in kernel_or_gp.covar_module.sub_kernels()
+            if isinstance(k, BoTorchWLKernel)
+        ]
     elif isinstance(kernel_or_gp, BoTorchWLKernel):
         modules = [kernel_or_gp]
     else:
-        assert hasattr(kernel_or_gp,
-                       "sub_kernels"), "Kernel module must have sub_kernels method."
-        modules = [k for k in kernel_or_gp.sub_kernels() if
-                   isinstance(k, BoTorchWLKernel)]
+        assert hasattr(kernel_or_gp, "sub_kernels"), (
+            "Kernel module must have sub_kernels method."
+        )
+        modules = [
+            k for k in kernel_or_gp.sub_kernels() if isinstance(k, BoTorchWLKernel)
+        ]
 
     # Save the current graph lookup and set the new graph lookup
     for kern in modules:

neps/optimizers/models/graphs/kernels.py

@@ -88,6 +88,7 @@ class BoTorchWLKernel(Kernel):
         adjacency_cache (list[Tensor]): Cached adjacency matrices of the graphs.
         label_cache (list[Tensor]): Cached initial node labels of the graphs.
     """
+
     has_lengthscale = False
 
     def __init__(

neps/optimizers/models/graphs/optimization.py

@@ -96,5 +96,5 @@ def optimize_acqf_graph(
     return (
         best_candidates[best_idx][:, :-1],
         best_graphs[best_idx],
-        best_scores[best_idx].item()
+        best_scores[best_idx].item(),
     )

neps/optimizers/models/graphs/utils.py

@@ -26,8 +26,7 @@ def min_max_scale(tensor: torch.Tensor) -> torch.Tensor:
 
 
 def graphs_to_tensors(
-    graphs: list[nx.Graph],
-    device: torch.device | None = None
+    graphs: list[nx.Graph], device: torch.device | None = None
 ) -> tuple[list[torch.sparse.Tensor], list[torch.Tensor]]:
     """Convert a list of NetworkX graphs into sparse adjacency matrices and label tensors.
 

tests/test_graphs/test_optimization_over_graphs.py

@@ -32,12 +32,19 @@ def setup_data(self) -> dict:
         N_GRAPH = 1
 
         # Generate random data
-        X = torch.cat([
-            torch.rand((TOTAL_CONFIGS, N_NUMERICAL), dtype=torch.float64),
-            torch.randint(0, N_CATEGORICAL_VALUES_PER_CATEGORY,
-                          (TOTAL_CONFIGS, N_CATEGORICAL), dtype=torch.float64),
-            torch.arange(TOTAL_CONFIGS, dtype=torch.float64).unsqueeze(1)
-        ], dim=1)
+        X = torch.cat(
+            [
+                torch.rand((TOTAL_CONFIGS, N_NUMERICAL), dtype=torch.float64),
+                torch.randint(
+                    0,
+                    N_CATEGORICAL_VALUES_PER_CATEGORY,
+                    (TOTAL_CONFIGS, N_CATEGORICAL),
+                    dtype=torch.float64,
+                ),
+                torch.arange(TOTAL_CONFIGS, dtype=torch.float64).unsqueeze(1),
+            ],
+            dim=1,
+        )
 
         # Generate random graphs
         graphs = [nx.erdos_renyi_graph(5, 0.5) for _ in range(TOTAL_CONFIGS)]
@@ -76,23 +83,36 @@ def test_gp_fit_and_predict(self, setup_data: dict) -> None:
 
         # Define the kernels
         kernels = [
-            ScaleKernel(MaternKernel(nu=2.5, ard_num_dims=setup_data["N_NUMERICAL"],
-                                     active_dims=range(setup_data["N_NUMERICAL"]))),
             ScaleKernel(
-                CategoricalKernel(ard_num_dims=setup_data["N_CATEGORICAL"],
-                                  active_dims=range(setup_data["N_NUMERICAL"],
-                                                    setup_data["N_NUMERICAL"] +
-                                                    setup_data["N_CATEGORICAL"])
-                                  )
+                MaternKernel(
+                    nu=2.5,
+                    ard_num_dims=setup_data["N_NUMERICAL"],
+                    active_dims=range(setup_data["N_NUMERICAL"]),
+                )
+            ),
+            ScaleKernel(
+                CategoricalKernel(
+                    ard_num_dims=setup_data["N_CATEGORICAL"],
+                    active_dims=range(
+                        setup_data["N_NUMERICAL"],
+                        setup_data["N_NUMERICAL"] + setup_data["N_CATEGORICAL"],
+                    ),
+                )
             ),
             ScaleKernel(
-                BoTorchWLKernel(graph_lookup=train_graphs, n_iter=5, normalize=True,
-                                active_dims=(train_x.shape[1] - 1,)))
+                BoTorchWLKernel(
+                    graph_lookup=train_graphs,
+                    n_iter=5,
+                    normalize=True,
+                    active_dims=(train_x.shape[1] - 1,),
+                )
+            ),
         ]
 
         # Create the GP model
-        gp = SingleTaskGP(train_X=train_x, train_Y=train_y,
-                          covar_module=AdditiveKernel(*kernels))
+        gp = SingleTaskGP(
+            train_X=train_x, train_Y=train_y, covar_module=AdditiveKernel(*kernels)
+        )
 
         # Fit the GP
         mll = ExactMarginalLogLikelihood(gp.likelihood, gp)
@@ -119,24 +139,36 @@ def test_acquisition_function_optimization(self, setup_data: dict) -> None:
 
         # Define the kernels
         kernels = [
-            ScaleKernel(MaternKernel(nu=2.5, ard_num_dims=setup_data["N_NUMERICAL"],
-                                     active_dims=range(setup_data["N_NUMERICAL"]))),
+            ScaleKernel(
+                MaternKernel(
+                    nu=2.5,
+                    ard_num_dims=setup_data["N_NUMERICAL"],
+                    active_dims=range(setup_data["N_NUMERICAL"]),
+                )
+            ),
             ScaleKernel(
                 CategoricalKernel(
                     ard_num_dims=setup_data["N_CATEGORICAL"],
-                    active_dims=range(setup_data["N_NUMERICAL"],
-                                      setup_data["N_NUMERICAL"] +
-                                      setup_data["N_CATEGORICAL"])
+                    active_dims=range(
+                        setup_data["N_NUMERICAL"],
+                        setup_data["N_NUMERICAL"] + setup_data["N_CATEGORICAL"],
+                    ),
                 )
             ),
             ScaleKernel(
-                BoTorchWLKernel(graph_lookup=train_graphs, n_iter=5, normalize=True,
-                                active_dims=(train_x.shape[1] - 1,)))
+                BoTorchWLKernel(
+                    graph_lookup=train_graphs,
+                    n_iter=5,
+                    normalize=True,
+                    active_dims=(train_x.shape[1] - 1,),
+                )
+            ),
         ]
 
         # Create the GP model
-        gp = SingleTaskGP(train_X=train_x, train_Y=train_y,
-                          covar_module=AdditiveKernel(*kernels))
+        gp = SingleTaskGP(
+            train_X=train_x, train_Y=train_y, covar_module=AdditiveKernel(*kernels)
+        )
 
         # Fit the GP
         mll = ExactMarginalLogLikelihood(gp.likelihood, gp)
@@ -151,21 +183,30 @@ def test_acquisition_function_optimization(self, setup_data: dict) -> None:
         )
 
         # Define bounds for optimization
-        bounds = torch.tensor([
-            [0.0] * setup_data["N_NUMERICAL"] + [0.0] * setup_data["N_CATEGORICAL"] + [
-                -1.0] * setup_data["N_GRAPH"],
-            [1.0] * setup_data["N_NUMERICAL"] + [
-                float(setup_data["N_CATEGORICAL_VALUES_PER_CATEGORY"] - 1)] * setup_data[
-                "N_CATEGORICAL"] + [len(train_x) - 1] * setup_data["N_GRAPH"],
-        ])
+        bounds = torch.tensor(
+            [
+                [0.0] * setup_data["N_NUMERICAL"]
+                + [0.0] * setup_data["N_CATEGORICAL"]
+                + [-1.0] * setup_data["N_GRAPH"],
+                [1.0] * setup_data["N_NUMERICAL"]
+                + [float(setup_data["N_CATEGORICAL_VALUES_PER_CATEGORY"] - 1)]
+                * setup_data["N_CATEGORICAL"]
+                + [len(train_x) - 1] * setup_data["N_GRAPH"],
+            ]
+        )
 
         # Define fixed categorical features
-        cats_per_column = {i: list(range(setup_data["N_CATEGORICAL_VALUES_PER_CATEGORY"]))
-                           for i in range(setup_data["N_NUMERICAL"],
-                                          setup_data["N_NUMERICAL"] + setup_data[
-                                              "N_CATEGORICAL"])}
-        fixed_cats = [dict(zip(cats_per_column.keys(), combo, strict=False)) for combo in
-                      product(*cats_per_column.values())]
+        cats_per_column = {
+            i: list(range(setup_data["N_CATEGORICAL_VALUES_PER_CATEGORY"]))
+            for i in range(
+                setup_data["N_NUMERICAL"],
+                setup_data["N_NUMERICAL"] + setup_data["N_CATEGORICAL"],
+            )
+        }
+        fixed_cats = [
+            dict(zip(cats_per_column.keys(), combo, strict=False))
+            for combo in product(*cats_per_column.values())
+        ]
 
         # Optimize the acquisition function
         best_candidate, best_graph, best_score = optimize_acqf_graph(
@@ -180,16 +221,19 @@ def test_acquisition_function_optimization(self, setup_data: dict) -> None:
         )
 
         # Assertions for the acquisition function optimization
-        assert isinstance(best_candidate,
-                          torch.Tensor), "Best candidate should be a tensor"
-        assert best_candidate.shape == (1, train_x.shape[1] - 1), \
+        assert isinstance(best_candidate, torch.Tensor), (
+            "Best candidate should be a tensor"
+        )
+        assert best_candidate.shape == (1, train_x.shape[1] - 1), (
             "Best candidate should have the correct shape (excluding the graph index)"
+        )
         assert isinstance(best_graph, nx.Graph), "Best graph should be a NetworkX graph"
         assert isinstance(best_score, float), "Best score should be a float"
 
         # Ensure the best candidate does not contain the graph index column
-        assert best_candidate.shape[1] == train_x.shape[1] - 1, \
+        assert best_candidate.shape[1] == train_x.shape[1] - 1, (
             "Best candidate should not include the graph index column"
+        )
 
     def test_graph_sampling(self, setup_data: dict) -> None:
         """Test the graph sampling functionality."""
@@ -200,12 +244,15 @@ def test_graph_sampling(self, setup_data: dict) -> None:
         sampled_graphs = sample_graphs(train_graphs, num_samples=num_samples)
 
         # Basic checks
-        assert len(sampled_graphs) == num_samples, \
+        assert len(sampled_graphs) == num_samples, (
             f"Expected {num_samples} sampled graphs, got {len(sampled_graphs)}"
-        assert all(isinstance(graph, nx.Graph) for graph in sampled_graphs), \
+        )
+        assert all(isinstance(graph, nx.Graph) for graph in sampled_graphs), (
             "All sampled graphs should be NetworkX graphs"
-        assert all(nx.is_connected(graph) for graph in sampled_graphs), \
+        )
+        assert all(nx.is_connected(graph) for graph in sampled_graphs), (
             "All sampled graphs should be connected"
+        )
 
     def test_min_max_scaling(self, setup_data: dict) -> None:
         """Test the min-max scaling utility."""
@@ -217,26 +264,29 @@ def test_min_max_scaling(self, setup_data: dict) -> None:
         # Assertions for min-max scaling
         assert torch.all(scaled_train_x >= 0), "Scaled values should be >= 0"
         assert torch.all(scaled_train_x <= 1), "Scaled values should be <= 1"
-        assert scaled_train_x.shape == train_x.shape, \
+        assert scaled_train_x.shape == train_x.shape, (
             "Scaled data should have the same shape as the input data"
+        )
 
         # Check that the scaling is correct
         for i in range(train_x.shape[1]):
             col_min = torch.min(train_x[:, i])
             col_max = torch.max(train_x[:, i])
             if col_min != col_max:  # Avoid division by zero
                 expected_scaled_col = (train_x[:, i] - col_min) / (col_max - col_min)
-                assert torch.allclose(scaled_train_x[:, i], expected_scaled_col), \
+                assert torch.allclose(scaled_train_x[:, i], expected_scaled_col), (
                     f"Scaling is incorrect for column {i}"
+                )
 
     def test_set_graph_lookup(self, setup_data: dict) -> None:
         """Test the set_graph_lookup context manager."""
         train_graphs = setup_data["train_graphs"]
         test_graphs = setup_data["test_graphs"]
 
         # Define the kernel
-        kernel = BoTorchWLKernel(graph_lookup=train_graphs, n_iter=5, normalize=True,
-                                 active_dims=(0,))
+        kernel = BoTorchWLKernel(
+            graph_lookup=train_graphs, n_iter=5, normalize=True, active_dims=(0,)
+        )
 
         # Use the context manager to temporarily set the graph lookup
         with set_graph_lookup(kernel, test_graphs, append=True):