Brings tuning changes on LaMBO2 (#55)

* add lambo ehrlich example, tune params

* lambo tweaks

* add geometric recency

* Adds beignet as a dependency

* Removes the hardcoded GPU acceleration

* Runs the lambo test for less time

* Makes it even smaller for testing

* Updates the example's readme

---------

Co-authored-by: Samuel Stanton <[email protected]>

Author: miguelgondu

.gitignore

@@ -11,8 +11,10 @@ __pycache__/
 
 .tox/
 build/
+oracle/
 
 examples/05_optimizing_qed_using_latent_space_bo/VAESelfies_TINY-CID-SELFIES-20_latent_dim_2.pt
+examples/06_running_lambo2_on_rasp/directed_evolution_rasp*.npz
 examples/07_optimizing_logp_on_chembl/VAE_CHEMBL.pt
 examples/08_pymoo_nsga_ii_on_foldx/history.json
 examples/09_replicating_nsga_ii_of_lambo_by_hand/pdbs

examples/06_running_lambo2_on_rasp/readme.md

@@ -1,4 +1,6 @@
-This folder includes an example in which we optimize the thermal stability of red fluorescent proteins (RFPs) using `LaMBO2`.
+This folder includes an example in which we optimize the thermal stability of red fluorescent proteins (RFPs), measured using an additive version of RaSP, using `LaMBO2`.
+
+As a pre-requisite, [we encourage you to set-up `poli` for RaSP](https://machinelearninglifescience.github.io/poli-docs/using_poli/objective_repository/RaSP.html).
 
 It includes the following assets:
 - Several PDB files for these RFPs, based on the Pareto front found by [LaMBO](https://arxiv.org/abs/2203.12742).
@@ -10,7 +12,6 @@ We recommend running it inside the environment of `LaMBO2`, which you can find i
 
 ```bash
 # From the root of the poli-baselines directory
-conda env create --file src/poli_baselines/solvers/bayesian_optimization/lambo2/environment.lambo2.yml
-conda activate poli__lambo2
+pip install -e .[lambo2]
 python examples/06_running_lambo2_on_rasp/run.py
 ```

examples/06_running_lambo2_on_rasp/run.py

@@ -3,6 +3,7 @@
 
 import matplotlib.pyplot as plt
 import numpy as np
+import torch
 from poli.objective_repository import RaspProblemFactory
 
 from poli_baselines.solvers.bayesian_optimization.lambo2 import LaMBO2
@@ -33,20 +34,25 @@ def run_with_default_hyperparameters():
     observer.x_s.append(x0.reshape(-1, 1))
     observer.y_s.append(y0)
 
+    torch.set_float32_matmul_precision("medium")
     lambo2 = LaMBO2(
         black_box=black_box,
         x0=x0,
         y0=y0,
+        overrides=["max_epochs=2"],
+        max_epochs_for_retraining=8,
     )
-
-    lambo2.solve(max_iter=10)
+    lambo2.solve(max_iter=32)
 
     fig, (ax1, ax2) = plt.subplots(1, 2)
     plot_best_y(observer, ax1)
     plot_best_y(observer, ax2, start_from=x0.shape[0])
     ax1.axvline(x0.shape[0], color="red")
     plt.show()
 
+    print("Best starting obj value: ", np.max(y0))
+    print("Best final obj value: ", np.max(lambo2.history_for_training["y"]))
+
     black_box.terminate()
 
 
@@ -107,5 +113,5 @@ def run_with_modified_hyperparameters():
 
 
 if __name__ == "__main__":
-    # run_with_default_hyperparameters()
-    run_with_modified_hyperparameters()
+    run_with_default_hyperparameters()
+    # run_with_modified_hyperparameters()

examples/07_running_lambo2_on_ehrlich/run.py

@@ -0,0 +1,72 @@
+import sys
+from pathlib import Path
+
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+from poli.objective_repository import EhrlichProblemFactory
+
+from poli_baselines.solvers.bayesian_optimization.lambo2 import LaMBO2
+from poli_baselines.solvers.simple.genetic_algorithm import FixedLengthGeneticAlgorithm
+
+THIS_DIR = Path(__file__).resolve().parent
+sys.path.append(str(THIS_DIR))
+
+from simple_observer import SimpleObserver, plot_best_y  # noqa: E402
+
+
+def run_with_default_hyperparameters():
+    problem = EhrlichProblemFactory().create(
+        sequence_length=32,
+        motif_length=4,
+        n_motifs=2,
+        quantization=4,
+        return_value_on_unfeasible=-1.0,
+    )
+    black_box = problem.black_box
+    x0 = problem.x0
+    random_seqs = np.array(
+        [list(black_box._sample_random_sequence()) for _ in range(127)]
+    )
+    x0 = np.concatenate([problem.x0, random_seqs], axis=0)
+    y0 = black_box(x0)
+
+    observer = SimpleObserver()
+    black_box.set_observer(observer)
+
+    # arr = np.load(THIS_DIR / "rasp_seed_data.npz")
+    # x0 = arr["x0"]
+    # y0 = arr["y0"]
+
+    observer.x_s.append(x0.reshape(-1, 1))
+    observer.y_s.append(y0)
+
+    presolver = FixedLengthGeneticAlgorithm(
+        black_box=black_box, x0=x0, y0=y0, population_size=128, prob_of_mutation=0.4
+    )
+    presolver.solve(max_iter=1)
+    presolver_x = np.array(presolver.history["x"])
+    presolver_x = presolver_x.reshape(presolver_x.shape[0], -1)
+
+    # import pdb; pdb.set_trace()
+    torch.set_float32_matmul_precision("medium")
+    lambo2 = LaMBO2(
+        black_box=black_box,
+        x0=presolver_x,  # inconsistent API
+        overrides=["max_epochs=2"],
+        max_epochs_for_retraining=8,
+    )
+
+    lambo2.solve(max_iter=32)
+
+    fig, (ax1, ax2) = plt.subplots(1, 2)
+    plot_best_y(observer, ax1)
+    plot_best_y(observer, ax2, start_from=x0.shape[0])
+    ax1.axvline(x0.shape[0], color="red")
+    plt.show()
+
+    black_box.terminate()
+
+
+if __name__ == "__main__":
+    run_with_default_hyperparameters()

examples/07_running_lambo2_on_ehrlich/simple_observer.py

@@ -0,0 +1,27 @@
+import numpy as np
+import matplotlib.pyplot as plt
+
+from poli.core.black_box_information import BlackBoxInformation
+from poli.core.util.abstract_observer import AbstractObserver
+
+
+class SimpleObserver(AbstractObserver):
+    def __init__(self) -> None:
+        self.x_s = []
+        self.y_s = []
+        super().__init__()
+
+    def initialize_observer(
+        self, problem_setup_info: BlackBoxInformation, caller_info: object, seed: int
+    ) -> object: ...
+
+    def observe(self, x: np.ndarray, y: np.ndarray, context=None) -> None:
+        self.x_s.append(x)
+        self.y_s.append(y)
+
+
+def plot_best_y(obs: SimpleObserver, ax: plt.Axes, start_from: int = 0):
+    best_y = np.maximum.accumulate(np.vstack(obs.y_s).flatten())
+    ax.plot(best_y.flatten()[start_from:])
+    ax.set_xlabel("Number of evaluations")
+    ax.set_ylabel("Best value found")

pyproject.toml

@@ -77,7 +77,7 @@ bounce = [
     "pytest>=7.3.1,<7.4.0",
     "bounce @ git+https://github.com/miguelgondu/bounce.git@main"
 ]
-lambo2 = ["pytorch-cortex"]
+lambo2 = ["pytorch-cortex", "beignet"]
 
 [project.urls]
 Homepage = "https://github.com/MachineLearningLifeScience/poli-baselines"

src/poli_baselines/solvers/bayesian_optimization/lambo2/environment.lambo2.small.yml

@@ -7,6 +7,7 @@ dependencies:
   - pip:
       - --extra-index-url https://download.pytorch.org/whl/cpu
       - numpy<2
+      - beignet
       - "git+https://github.com/prescient-design/cortex.git"
       - "git+https://github.com/MachineLearningLifeScience/poli.git"
       - "git+https://github.com/MachineLearningLifeScience/poli-baselines.git@main"

src/poli_baselines/solvers/bayesian_optimization/lambo2/environment.lambo2.yml

@@ -6,6 +6,7 @@ dependencies:
   - pip
   - pip:
       - numpy<2
+      - beignet
       - "git+https://github.com/prescient-design/cortex.git"
       - "git+https://github.com/MachineLearningLifeScience/poli.git"
       - "git+https://github.com/MachineLearningLifeScience/poli-baselines.git@main"

src/poli_baselines/solvers/bayesian_optimization/lambo2/hydra_configs/branches/protein_constraint.yaml

@@ -0,0 +1,6 @@
+protein_constraint:
+  _target_: cortex.model.branch.Conv1dBranch
+  out_dim: 8
+  channel_dim: ${feature_dim}
+  num_blocks: 1
+  kernel_size: ${kernel_size}

src/poli_baselines/solvers/bayesian_optimization/lambo2/hydra_configs/branches/protein_property.yaml

@@ -2,5 +2,5 @@ protein_property:
   _target_: cortex.model.branch.Conv1dBranch
   out_dim: 8
   channel_dim: ${feature_dim}
-  num_blocks: 0
+  num_blocks: 1
   kernel_size: ${kernel_size}