implement AffinePosteriorTransform

aepsych/acquisition/objective/__init__.py

@@ -16,6 +16,7 @@
     ProbitObjective,
 )
 from .semi_p import SemiPProbabilityObjective, SemiPThresholdObjective
+from .multi_outcome import AffinePosteriorTransform
 
 __all__ = [
     "AEPsychObjective",
@@ -25,6 +26,7 @@
     "ProbitObjective",
     "SemiPProbabilityObjective",
     "SemiPThresholdObjective",
+    "AffinePosteriorTransform"
 ]
 
 Config.register_module(sys.modules[__name__])

aepsych/acquisition/objective/multi_outcome.py

@@ -0,0 +1,11 @@
+#!/usr/bin/env python3
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+from botorch.acquisition.objective import ScalarizedPosteriorTransform
+from aepsych.config import ConfigurableMixin
+
+class AffinePosteriorTransform(ScalarizedPosteriorTransform, ConfigurableMixin):
+    pass

aepsych/models/model_protocol.py

@@ -50,8 +50,7 @@ def predict_probability(self, x: torch.Tensor, **kwargs) -> torch.Tensor:
     def predict_transform(
         self,
         x: torch.Tensor,
-        transformed_posterior_cls: Optional[type[TransformedPosterior]] = None,
-        **transform_kwargs,
+        **kwargs,
     ):
         pass
 

tests/acquisition/test_objective.py

@@ -14,10 +14,14 @@
     FloorGumbelObjective,
     FloorLogitObjective,
     FloorProbitObjective,
+    AffinePosteriorTransform
 )
 from parameterized import parameterized
 from scipy.stats import gumbel_l, logistic, norm
 
+from aepsych.config import Config
+import numpy.testing as npt
+
 objective_pairs = [
     (FloorLogitObjective, logistic),
     (FloorProbitObjective, norm),
@@ -42,3 +46,15 @@ def test_floor_links(self, objectives, floor):
 
         our_inverse = our_link.inverse(our_answer)
         self.assertTrue(np.allclose(x, our_inverse.numpy()))
+
+class AffinePosteriorTransformTests(unittest.TestCase):
+    def test_from_config(self):
+        config_str = """
+        [AffinePosteriorTransform]
+        weights = [1, -1]
+        offset = 0
+        """
+        config = Config(config_str=config_str)
+        apt = AffinePosteriorTransform.from_config(config=config)
+        npt.assert_array_equal(apt.weights, torch.tensor([1., -1.]))
+        self.assertEqual(apt.offset, 0.0)