Remove extra normalize in pairwise probit (facebookresearch#823)

Summary:
Pull Request resolved: facebookresearch#823

This diff modifies the `PairwiseProbitModel` by removing an extra normalization step in the `pairwise_probit.py` file. The normalization should now be handled by transforms. Additionally, the `test_pairwise_probit.py` file is updated to reflect this change.

Reviewed By: phigua

Differential Revision: D84107038

fbshipit-source-id: 6b253c150126174a2ee1eaa7a7000c68e7348b12

Author: JasonKChow

aepsych/models/pairwise_probit.py

@@ -17,7 +17,6 @@
 from aepsych.utils_logging import getLogger
 from botorch.fit import fit_gpytorch_mll
 from botorch.models import PairwiseGP, PairwiseLaplaceMarginalLogLikelihood
-from botorch.models.transforms.input import Normalize
 from torch.distributions import Normal
 
 logger = getLogger()
@@ -93,8 +92,6 @@ def __init__(
 
         self.max_fit_time = max_fit_time
 
-        bounds = torch.stack((self.lb, self.ub))
-        input_transform = Normalize(d=dim, bounds=bounds)
         if covar_module is None:
             factory = DefaultMeanCovarFactory(
                 dim=dim, stimuli_per_trial=self.stimuli_per_trial
@@ -106,7 +103,6 @@ def __init__(
             comparisons=None,
             covar_module=covar_module,
             jitter=1e-3,
-            input_transform=input_transform,
         )
 
         self.dim = dim  # The Pairwise constructor sets self.dim = None.

tests/models/test_pairwise_probit.py

@@ -183,7 +183,14 @@ def test_pairwise_memorize_rescaled(self):
         xrescaled[:, 0, :] = xrescaled[:, 0, :] / 500 + 1
         xrescaled[:, 1, :] = xrescaled[:, 1, :] / 5e-6 - 1
         y = torch.Tensor(f_pairwise(new_novel_det, xrescaled) > 0.5).int()
-        model = PairwiseProbitModel(lb=lb, ub=ub)
+        lb_tensor = torch.tensor(lb)
+        ub_tensor = torch.tensor(ub)
+        transforms = ParameterTransforms(
+            normalize=NormalizeScale(d=2, bounds=torch.stack([lb_tensor, ub_tensor]))
+        )
+        model = ParameterTransformedModel(
+            PairwiseProbitModel, lb=lb_tensor, ub=ub_tensor, transforms=transforms
+        )
         model.fit(x[:18], y[:18])
         with torch.no_grad():
             f0, _ = model.predict(x[18:, ..., 0])
@@ -341,29 +348,48 @@ def test_2d_pairwise_probit(self):
         ub = torch.tensor([1.0, 1.0])
         extra_acqf_args = {"beta": 3.84}
 
+        transforms = ParameterTransforms(
+            normalize=NormalizeScale(d=2, bounds=torch.stack([lb, ub]))
+        )
+        sobol_gen = ParameterTransformedGenerator(
+            generator=SobolGenerator,
+            lb=lb,
+            ub=ub,
+            seed=seed,
+            stimuli_per_trial=2,
+            transforms=transforms,
+        )
+        acqf_gen = ParameterTransformedGenerator(
+            generator=OptimizeAcqfGenerator,
+            acqf=qUpperConfidenceBound,
+            acqf_kwargs=extra_acqf_args,
+            stimuli_per_trial=2,
+            transforms=transforms,
+            lb=lb,
+            ub=ub,
+        )
+        probit_model = ParameterTransformedModel(
+            model=PairwiseProbitModel, lb=lb, ub=ub, transforms=transforms
+        )
         model_list = [
             Strategy(
                 lb=lb,
                 ub=ub,
-                generator=SobolGenerator(lb=lb, ub=ub, seed=seed, stimuli_per_trial=2),
+                generator=sobol_gen,
                 min_asks=n_init,
                 stimuli_per_trial=2,
                 outcome_types=["binary"],
+                transforms=transforms,
             ),
             Strategy(
                 lb=lb,
                 ub=ub,
-                model=PairwiseProbitModel(lb=lb, ub=ub),
-                generator=OptimizeAcqfGenerator(
-                    acqf=qUpperConfidenceBound,
-                    acqf_kwargs=extra_acqf_args,
-                    stimuli_per_trial=2,
-                    lb=lb,
-                    ub=ub,
-                ),
+                model=probit_model,
+                generator=acqf_gen,
                 min_asks=n_opt,
                 stimuli_per_trial=2,
                 outcome_types=["binary"],
+                transforms=transforms,
             ),
         ]