mass format codebase (#436)

Summary:
Pull Request resolved: #436

Format entire codebase using ufmt (black + usort).

Reviewed By: crasanders

Differential Revision: D65457850

.github/workflows/build-lint-test.yml

@@ -20,6 +20,7 @@ jobs:
         run: |
           python -m pip install --upgrade pip
           pip install ".[dev]"
+          pip install -r requirements-fmt.txt
       - name: Lint with flake8
         run: |
           # stop the build if there are Python syntax errors or undefined names
@@ -29,6 +30,11 @@ jobs:
       - name: Type check with mypy
         run: |
           mypy --config-file mypy.ini
+      - uses: omnilib/ufmt@action-v1
+        with:
+          path: aespych tests tests_gpu
+          requirements: requirements-fmt.txt
+          python-version: "3.10"
 
   build-test:
 

aepsych/__init__.py

@@ -8,7 +8,6 @@
 import sys
 
 import torch
-
 from gpytorch.likelihoods import BernoulliLikelihood, GaussianLikelihood
 
 from . import (

aepsych/acquisition/bvn.py

@@ -7,8 +7,8 @@
 
 from math import pi as _pi
 from typing import Tuple
-import torch
 
+import torch
 
 inv_2pi = 1 / (2 * _pi)
 _neg_inv_sqrt2 = -1 / (2**0.5)

aepsych/acquisition/lookahead.py

@@ -5,7 +5,7 @@
 # This source code is licensed under the license found in the
 # LICENSE file in the root directory of this source tree.
 
-from typing import Any, Callable, Dict, Optional, Tuple, cast
+from typing import Any, Callable, cast, Dict, Optional, Tuple
 
 import numpy as np
 import torch
@@ -185,7 +185,7 @@ def _compute_acqf(self, Px: Tensor, P1: Tensor, P0: Tensor, py1: Tensor) -> Tens
 def construct_inputs_local_lookahead(
     model: GPyTorchModel,
     training_data: None,
-    lookahead_type: str ="levelset",
+    lookahead_type: str = "levelset",
     target: Optional[float] = None,
     posterior_transform: Optional[PosteriorTransform] = None,
     **kwargs,
@@ -350,7 +350,6 @@ def __init__(
         Xq: Optional[Tensor] = None,
         k: Optional[float] = 20.0,
     ) -> None:
-
         super().__init__(
             model=model,
             target=target,

aepsych/acquisition/lse.py

@@ -97,7 +97,6 @@ def construct_inputs_lse(
     sampler: Optional[MCSampler] = None,
     **kwargs,
 ) -> Dict[str, Any]:
-
     return {
         "model": model,
         "objective": objective,

aepsych/acquisition/mutual_information.py

@@ -5,8 +5,7 @@
 # This source code is licensed under the license found in the
 # LICENSE file in the root directory of this source tree.
 
-r"""
-"""
+r""" """
 
 from __future__ import annotations
 

aepsych/acquisition/objective/__init__.py

@@ -17,7 +17,6 @@
 )
 from .semi_p import SemiPProbabilityObjective, SemiPThresholdObjective
 
-
 __all__ = [
     "AEPsychObjective",
     "FloorGumbelObjective",

aepsych/acquisition/objective/objective.py

@@ -60,7 +60,7 @@ class FloorLinkObjective(AEPsychObjective):
     the probability is known not to go below it.
     """
 
-    def __init__(self, floor: float=0.5) -> None:
+    def __init__(self, floor: float = 0.5) -> None:
         self.floor = floor
         super().__init__()
 

aepsych/acquisition/objective/semi_p.py

@@ -66,7 +66,6 @@ def forward(self, samples: Tensor, X: Tensor) -> Tensor:
 
     @classmethod
     def from_config(cls, config: Config) -> SemiPProbabilityObjective:
-
         classname = cls.__name__
 
         likelihood_cls = config.getobj(classname, "likelihood", fallback=None)
@@ -87,7 +86,13 @@ class SemiPThresholdObjective(SemiPObjectiveBase):
     that gives the threshold distribution.
     """
 
-    def __init__(self, target: float, likelihood: Optional[LinearBernoulliLikelihood] =None, *args, **kwargs):
+    def __init__(
+        self,
+        target: float,
+        likelihood: Optional[LinearBernoulliLikelihood] = None,
+        *args,
+        **kwargs,
+    ):
         """Evaluates the probability objective.
 
         Args:
@@ -116,7 +121,6 @@ def forward(self, samples: Tensor, X: Optional[Tensor] = None) -> Tensor:
 
     @classmethod
     def from_config(cls, config: Config) -> SemiPThresholdObjective:
-
         classname = cls.__name__
 
         likelihood_cls = config.getobj(classname, "likelihood", fallback=None)

aepsych/benchmark/__init__.py

@@ -6,14 +6,14 @@
 # LICENSE file in the root directory of this source tree.
 
 from .benchmark import Benchmark, DerivedValue
-from .pathos_benchmark import PathosBenchmark, run_benchmarks_with_checkpoints
-from .problem import LSEProblem, LSEProblemWithEdgeLogging, Problem
 from .example_problems import (
-    DiscrimLowDim,
+    ContrastSensitivity6d,
     DiscrimHighDim,
+    DiscrimLowDim,
     Hartmann6Binary,
-    ContrastSensitivity6d,
 )
+from .pathos_benchmark import PathosBenchmark, run_benchmarks_with_checkpoints
+from .problem import LSEProblem, LSEProblemWithEdgeLogging, Problem
 from .test_functions import (
     discrim_highdim,
     make_songetal_testfun,

aepsych/benchmark/example_problems.py

@@ -1,16 +1,16 @@
 # Copyright (c) Meta Platforms, Inc. and affiliates. All rights reserved.
 import os
 from typing import List, Optional, Union
+
 import numpy as np
 import torch
-from aepsych.models import GPClassificationModel
+from aepsych.benchmark.problem import LSEProblemWithEdgeLogging
 from aepsych.benchmark.test_functions import (
-    modified_hartmann6,
     discrim_highdim,
+    modified_hartmann6,
     novel_discrimination_testfun,
 )
-from aepsych.benchmark.problem import LSEProblemWithEdgeLogging
-
+from aepsych.models import GPClassificationModel
 
 """The DiscrimLowDim, DiscrimHighDim, ContrastSensitivity6d, and Hartmann6Binary classes
 are copied from bernoulli_lse github repository (https://github.com/facebookresearch/bernoulli_lse)
@@ -21,7 +21,9 @@ class DiscrimLowDim(LSEProblemWithEdgeLogging):
     name = "discrim_lowdim"
     bounds = torch.tensor([[-1, 1], [-1, 1]], dtype=torch.double).T
 
-    def __init__(self, thresholds: Union[float, List, torch.Tensor, None] = None) -> None:
+    def __init__(
+        self, thresholds: Union[float, List, torch.Tensor, None] = None
+    ) -> None:
         thresholds = 0.75 if thresholds is None else thresholds
         super().__init__(thresholds=thresholds)
 
@@ -45,7 +47,9 @@ class DiscrimHighDim(LSEProblemWithEdgeLogging):
         dtype=torch.double,
     ).T
 
-    def __init__(self, thresholds: Union[float, List, torch.Tensor, None] = None) -> None:
+    def __init__(
+        self, thresholds: Union[float, List, torch.Tensor, None] = None
+    ) -> None:
         thresholds = 0.75 if thresholds is None else thresholds
         super().__init__(thresholds=thresholds)
 
@@ -62,7 +66,9 @@ class Hartmann6Binary(LSEProblemWithEdgeLogging):
         )
     )
 
-    def __init__(self, thresholds: Union[float, List, torch.Tensor, None] = None) -> None:
+    def __init__(
+        self, thresholds: Union[float, List, torch.Tensor, None] = None
+    ) -> None:
         thresholds = 0.5 if thresholds is None else thresholds
         super().__init__(thresholds=thresholds)
 
@@ -83,7 +89,9 @@ class ContrastSensitivity6d(LSEProblemWithEdgeLogging):
         dtype=torch.double,
     ).T
 
-    def __init__(self, thresholds: Union[float, List, torch.Tensor, None] = None) -> None:
+    def __init__(
+        self, thresholds: Union[float, List, torch.Tensor, None] = None
+    ) -> None:
         thresholds = 0.75 if thresholds is None else thresholds
         super().__init__(thresholds=thresholds)
 

aepsych/benchmark/pathos_benchmark.py

@@ -96,7 +96,6 @@ def run_experiment(
         try:
             return super().run_experiment(problem, local_config, seed, rep)
         except Exception as e:
-
             logging.error(
                 f"Error on config {config_dict}: {e}!"
                 + f"Traceback follows:\n{traceback.format_exc()}"
@@ -229,7 +228,7 @@ def run_benchmarks_with_checkpoints(
         final_results = bench.pandas()
         final_results.to_csv(out_fname)
     else:
-        for chunk in range(start_idx, n_chunks+start_idx):
+        for chunk in range(start_idx, n_chunks + start_idx):
             out_fname = Path(f"{out_path}/{benchmark_name}_chunk{chunk}_out.csv")
 
             intermediate_fname = Path(

aepsych/benchmark/problem.py

@@ -5,14 +5,14 @@
 # This source code is licensed under the license found in the
 # LICENSE file in the root directory of this source tree.
 from functools import cached_property
-from typing import Any, Dict, Union, List
+from typing import Any, Dict, List, Union
 
 import aepsych
 import numpy as np
 import torch
-from scipy.stats import bernoulli
 from aepsych.strategy import SequentialStrategy, Strategy
 from aepsych.utils import make_scaled_sobol
+from scipy.stats import bernoulli
 
 
 class Problem:
@@ -64,8 +64,9 @@ def p(self, x: torch.Tensor) -> torch.Tensor:
         normal_dist = torch.distributions.Normal(0, 1)  # Standard normal distribution
         return normal_dist.cdf(self.f(x))  # Use PyTorch's CDF equivalent
 
-
-    def sample_y(self, x: torch.Tensor) -> np.ndarray: # TODO: This can be done with  torch.bernoulli(self.p(x)), but Strategy.add_data() expects a numpy array for now
+    def sample_y(
+        self, x: torch.Tensor
+    ) -> np.ndarray:  # TODO: This can be done with  torch.bernoulli(self.p(x)), but Strategy.add_data() expects a numpy array for now
         """Sample a response from test function.
 
         Args:
@@ -104,8 +105,8 @@ def p_true(self) -> torch.Tensor:
         Returns:
             torch.Tensor: Values of true response probability over evaluation grid.
         """
-        normal_dist = torch.distributions.Normal(0, 1) 
-        return normal_dist.cdf(self.f_true)  
+        normal_dist = torch.distributions.Normal(0, 1)
+        return normal_dist.cdf(self.f_true)
 
     def p_hat(self, model: aepsych.models.base.ModelProtocol) -> torch.Tensor:
         """Generate mean predictions from the model over the evaluation grid.
@@ -155,24 +156,29 @@ def evaluate(
         mae_f = torch.mean(torch.abs(self.f_true - f_hat))
         mse_f = torch.mean((self.f_true - f_hat) ** 2)
         max_abs_err_f = torch.max(torch.abs(self.f_true - f_hat))
-        corr_f = torch.corrcoef(torch.stack((self.f_true.flatten(), f_hat.flatten())))[0, 1]
+        corr_f = torch.corrcoef(torch.stack((self.f_true.flatten(), f_hat.flatten())))[
+            0, 1
+        ]
         mae_p = torch.mean(torch.abs(self.p_true - p_hat))
         mse_p = torch.mean((self.p_true - p_hat) ** 2)
         max_abs_err_p = torch.max(torch.abs(self.p_true - p_hat))
-        corr_p = torch.corrcoef(torch.stack((self.p_true.flatten(), p_hat.flatten())))[0, 1]
+        corr_p = torch.corrcoef(torch.stack((self.p_true.flatten(), p_hat.flatten())))[
+            0, 1
+        ]
         brier = torch.mean(2 * torch.square(self.p_true - p_hat))
 
         # eval in samp-based expectation over posterior instead of just mean
         fsamps = model.sample(self.eval_grid, num_samples=1000)
         try:
             psamps = (
                 model.sample(self.eval_grid, num_samples=1000, probability_space=True)  # type: ignore
-
             )
         except (
             TypeError
         ):  # vanilla models don't have proba_space samps, TODO maybe we should add them
-            normal_dist = torch.distributions.Normal(0, 1)  # Standard normal distribution
+            normal_dist = torch.distributions.Normal(
+                0, 1
+            )  # Standard normal distribution
             psamps = normal_dist.cdf(fsamps)
 
         ferrs = fsamps - self.f_true[None, :]
@@ -186,21 +192,21 @@ def evaluate(
         expected_brier = torch.mean((2 * torch.square(self.p_true[None, :] - psamps)))
 
         metrics = {
-        "mean_abs_err_f": mae_f.item(),
-        "mean_integrated_abs_err_f": miae_f.item(),
-        "mean_square_err_f": mse_f.item(),
-        "mean_integrated_square_err_f": mise_f.item(),
-        "max_abs_err_f": max_abs_err_f.item(),
-        "pearson_corr_f": corr_f.item(),
-        "mean_abs_err_p": mae_p.item(),
-        "mean_integrated_abs_err_p": miae_p.item(),
-        "mean_square_err_p": mse_p.item(),
-        "mean_integrated_square_err_p": mise_p.item(),
-        "max_abs_err_p": max_abs_err_p.item(),
-        "pearson_corr_p": corr_p.item(),
-        "brier": brier.item(),
-        "expected_brier": expected_brier.item(),
-    }
+            "mean_abs_err_f": mae_f.item(),
+            "mean_integrated_abs_err_f": miae_f.item(),
+            "mean_square_err_f": mse_f.item(),
+            "mean_integrated_square_err_f": mise_f.item(),
+            "max_abs_err_f": max_abs_err_f.item(),
+            "pearson_corr_f": corr_f.item(),
+            "mean_abs_err_p": mae_p.item(),
+            "mean_integrated_abs_err_p": miae_p.item(),
+            "mean_square_err_p": mse_p.item(),
+            "mean_integrated_square_err_p": mise_p.item(),
+            "max_abs_err_p": max_abs_err_p.item(),
+            "pearson_corr_p": corr_p.item(),
+            "brier": brier.item(),
+            "expected_brier": expected_brier.item(),
+        }
 
         return metrics
 
@@ -237,16 +243,13 @@ def inverse_link(x):
                 return inverse_torch(x)
 
         except AttributeError:
+
             def inverse_link(x):
-                normal_dist = torch.distributions.Normal(0, 1)  
+                normal_dist = torch.distributions.Normal(0, 1)
                 return normal_dist.icdf(x)
 
-
         return inverse_link(self.thresholds).float()  # Return as float32 tensor
 
-
-
-
     @cached_property
     def true_below_threshold(self) -> torch.Tensor:
         """
@@ -284,9 +287,7 @@ def evaluate(self, strat: SequentialStrategy) -> Dict[str, float]:
         # define what "threshold" means in high-dim.
 
         # Brier score on level-set probabilities
-        p_l = model.p_below_threshold(
-            self.eval_grid, self.f_threshold(model)
-        )
+        p_l = model.p_below_threshold(self.eval_grid, self.f_threshold(model))
         true_p_l = self.true_below_threshold
         assert (
             p_l.ndim == 2
@@ -298,12 +299,16 @@ def evaluate(self, strat: SequentialStrategy) -> Dict[str, float]:
         brier_p_below_thresh = torch.mean(2 * torch.square(true_p_l - p_l), dim=1)
         # Classification error
         misclass_on_thresh = torch.mean(
-        p_l * (1 - true_p_l) + (1 - p_l) * true_p_l, dim=1
+            p_l * (1 - true_p_l) + (1 - p_l) * true_p_l, dim=1
         )
 
         for i_threshold, threshold in enumerate(self.thresholds):
-            metrics[f"brier_p_below_{threshold}"] = brier_p_below_thresh.detach().cpu().numpy()[i_threshold]
-            metrics[f"misclass_on_thresh_{threshold}"] = misclass_on_thresh.detach().cpu().numpy()[i_threshold]
+            metrics[f"brier_p_below_{threshold}"] = (
+                brier_p_below_thresh.detach().cpu().numpy()[i_threshold]
+            )
+            metrics[f"misclass_on_thresh_{threshold}"] = (
+                misclass_on_thresh.detach().cpu().numpy()[i_threshold]
+            )
         return metrics