From 1f0166b33a0973efad89132a6b87d31454f2c4b8 Mon Sep 17 00:00:00 2001
From: Angel Ferran Pousa <angel.ferran.pousa@desy.de>
Date: Fri, 23 Feb 2024 16:49:45 +0100
Subject: [PATCH 01/25] Implement Weights and Biases logger

---
 optimas/loggers/__init__.py     |  4 ++
 optimas/loggers/base.py         | 23 ++++++++
 optimas/loggers/wandb_logger.py | 95 +++++++++++++++++++++++++++++++++
 3 files changed, 122 insertions(+)
 create mode 100644 optimas/loggers/__init__.py
 create mode 100644 optimas/loggers/base.py
 create mode 100644 optimas/loggers/wandb_logger.py

diff --git a/optimas/loggers/__init__.py b/optimas/loggers/__init__.py
new file mode 100644
index 00000000..7a10ea22
--- /dev/null
+++ b/optimas/loggers/__init__.py
@@ -0,0 +1,4 @@
+from .wandb_logger import WandBLogger
+
+
+__all__ = ["WandBLogger"]
diff --git a/optimas/loggers/base.py b/optimas/loggers/base.py
new file mode 100644
index 00000000..2615e7d0
--- /dev/null
+++ b/optimas/loggers/base.py
@@ -0,0 +1,23 @@
+from __future__ import annotations
+from typing import TYPE_CHECKING
+
+
+if TYPE_CHECKING:
+    from optimas.core import Trial
+    from optimas.generators.base import Generator
+    from optimas.explorations import Exploration
+
+
+class Logger:
+
+    def initialize(self, exploration: Exploration):
+        pass
+
+    def log_trial(self, trial: Trial, generator: Generator):
+        pass
+
+    def log_custom_metrics(self, last_trial: Trial, generator: Generator):
+        pass
+
+    def finish(self):
+        pass
diff --git a/optimas/loggers/wandb_logger.py b/optimas/loggers/wandb_logger.py
new file mode 100644
index 00000000..ce814467
--- /dev/null
+++ b/optimas/loggers/wandb_logger.py
@@ -0,0 +1,95 @@
+from __future__ import annotations
+import pathlib
+from typing import TYPE_CHECKING, Optional, Callable, Dict
+
+from matplotlib.figure import Figure
+import wandb
+
+from .base import Logger
+
+if TYPE_CHECKING:
+    from optimas.core import Trial
+    from optimas.generators.base import Generator
+    from optimas.explorations import Exploration
+
+
+class WandBLogger(Logger):
+    def __init__(
+        self,
+        api_key: str,
+        project: str,
+        run_name: str,
+        run_id: Optional[str] = None,
+        data_types: Optional[Dict] = None,
+        user_function: Optional[Callable] = None,
+        login_kwargs: Optional[Dict] = None,
+        init_kwargs: Optional[Dict] = None,
+    ) -> None:
+        self._api_key = api_key
+        self._project = project
+        self._run_name = run_name
+        self._run_id = run_id
+        self._data_types = {} if data_types is None else data_types
+        self._user_function = user_function
+        self._login_kwargs = {} if login_kwargs is None else login_kwargs
+        self._init_kwargs = {} if init_kwargs is None else init_kwargs
+        self._run = None
+        self._dir = None
+
+    def initialize(self, exploration: Exploration):
+        # Create dir if it doesn't exist.
+        dir = exploration.exploration_dir_path
+        pathlib.Path(dir).mkdir(parents=True, exist_ok=True)
+        self._dir = dir
+
+        # Login and initialize run.
+        wandb.login(key=self._api_key, **self._login_kwargs)
+        if self._run is None:
+            self._run = wandb.init(
+                project=self._project,
+                name=self._run_name,
+                resume=True,
+                id=self._run_id,
+                dir=self._dir,
+                **self._init_kwargs,
+            )
+            if self._run_id is None:
+                self._run_id = self._run.id
+
+    def log_trial(self, trial: Trial, generator: Generator):
+        # Get and process trial data.
+        logs = trial.data
+        for key in list(logs.keys()):
+            # Apply user-provided wandb types.
+            if key in self._data_types:
+                logs[key] = self._data_types[key]["type"](
+                    logs[key], **self._data_types[key]["type_kwargs"]
+                )
+            # By default, convert matplotlib figures to images.
+            elif isinstance(logs[key], Figure):
+                logs[key] = wandb.Image(logs[key])
+            # By default, only log scalars.
+            elif hasattr(logs[key], "__len__"):
+                del logs[key]
+
+        # Organize in sections.
+        for par in generator.varying_parameters:
+            if par.name in logs:
+                logs[f"Varying parameters/{par.name}"] = logs.pop(par.name)
+        for par in generator.objectives:
+            if par.name in logs:
+                logs[f"Objectives/{par.name}"] = logs.pop(par.name)
+        for par in generator.analyzed_parameters:
+            if par.name in logs:
+                logs[f"Analyzed parameters/{par.name}"] = logs.pop(par.name)
+
+        # Add custom user-defined logs.
+        if self._user_function is not None:
+            custom_logs = self._user_function(trial, generator)
+            logs = {**logs, **custom_logs}
+
+        # Log data.
+        self._run.log(logs)
+
+    def finish(self):
+        self._run.finish()

From e90b03dc30b6f68d7ef7337c06c50f44b9661a4e Mon Sep 17 00:00:00 2001
From: Angel Ferran Pousa <angel.ferran.pousa@desy.de>
Date: Fri, 23 Feb 2024 16:51:56 +0100
Subject: [PATCH 02/25] Add `data` property to `Trial`

---
 optimas/core/trial.py | 15 +++++++++++++++
 1 file changed, 15 insertions(+)

diff --git a/optimas/core/trial.py b/optimas/core/trial.py
index d3815d6b..3ae55f4c 100644
--- a/optimas/core/trial.py
+++ b/optimas/core/trial.py
@@ -152,6 +152,21 @@ def evaluated(self) -> bool:
         """Determine whether the trial has been evaluated."""
         return self.completed or self.failed
 
+    @property
+    def data(self) -> Dict:
+        """Get a dictionary with all the trial data."""
+        vp_dict = self.parameters_as_dict()
+        ap_dict = self.analyzed_parameters_as_dict()
+        ob_dict = self.objectives_as_dict()
+        # Do not report uncertainty. We haven't yet decided about how to
+        # report it in the history.
+        for key, val in ob_dict.items():
+            ob_dict[key] = val[0]
+        for key, val in ap_dict.items():
+            ap_dict[key] = val[0]
+        data = {**vp_dict, **ob_dict, **ap_dict}
+        return data
+
     def mark_as(self, status) -> None:
         """Set trial status.
 

From 80fd2edc7613d0da33c45a75b8e4f5888e5a8d1e Mon Sep 17 00:00:00 2001
From: Angel Ferran Pousa <angel.ferran.pousa@desy.de>
Date: Fri, 23 Feb 2024 17:05:55 +0100
Subject: [PATCH 03/25] Add `logger` argument to `Exploration`.

Other changes:
- Include `comms` in `persis_info` to detect if gen ir running o a thread.
- Update generator after calling `libE` only if not running with `threads`.
It is not needed to do it in that case, because the memory is shared.
- Fix bug in `attach_evaluations` when the fields in the data contain arrays.
- Add `completed_trials` property to `Generator`.
- Log trials in `Generator`.
- Do not run `_prepare_to_send` when the generator runs on a thread.
- Implement `_prepare_to_send_back`, which clears the logger before sending back the generator to the manager.
---
 optimas/explorations/base.py                 | 23 +++++++++---
 optimas/gen_functions.py                     |  3 +-
 optimas/generators/ax/developer/multitask.py | 10 ++++--
 optimas/generators/base.py                   | 37 ++++++++++++++++++--
 4 files changed, 63 insertions(+), 10 deletions(-)

diff --git a/optimas/explorations/base.py b/optimas/explorations/base.py
index 7c5d2f65..c937d242 100644
--- a/optimas/explorations/base.py
+++ b/optimas/explorations/base.py
@@ -21,6 +21,7 @@
 from optimas.evaluators.function_evaluator import FunctionEvaluator
 from optimas.utils.logger import get_logger
 from optimas.utils.other import convert_to_dataframe
+from optimas.loggers.base import Logger
 
 
 logger = get_logger(__name__)
@@ -78,6 +79,10 @@ class Exploration:
         manager and ``N-1`` simulation workers. In this case, the
         ``sim_workers`` parameter is ignored. By default, ``'local'`` mode
         is used.
+    logger : Logger, optional
+        A custom logger that is informed of every completed trial and can
+        report on the results. Currently, a Weights and Biases logger is
+        available.
 
     """
 
@@ -93,6 +98,7 @@ def __init__(
         exploration_dir_path: Optional[str] = "./exploration",
         resume: Optional[bool] = False,
         libe_comms: Optional[Literal["local", "threads", "mpi"]] = "local",
+        logger: Optional[Logger] = None,
     ) -> None:
         # For backward compatibility, check for old threading name.
         if libe_comms == "local_threading":
@@ -124,6 +130,10 @@ def __init__(
         self._set_default_libe_specs()
         self._libe_history = self._create_libe_history()
         self._load_history(history, resume)
+        self._logger = logger
+        if self._logger is not None:
+            self._logger.initialize(self)
+            self.generator._set_logger(self._logger)
 
     @property
     def history(self) -> pd.DataFrame:
@@ -174,6 +184,7 @@ def run(self, n_evals: Optional[int] = None) -> None:
 
         # Create persis_info.
         persis_info = add_unique_random_streams({}, self.sim_workers + 2)
+        persis_info[1]["comms"] = self.libe_comms
 
         # If specified, allocate dedicated resources for the generator.
         if self.generator.dedicated_resources and self.generator.use_cuda:
@@ -188,7 +199,7 @@ def run(self, n_evals: Optional[int] = None) -> None:
         # Get gen_specs and sim_specs.
         run_params = self.evaluator.get_run_params()
         gen_specs = self.generator.get_gen_specs(
-            self.sim_workers, run_params, sim_max
+            self.sim_workers, run_params, sim_max, self.libe_comms
         )
         sim_specs = self.evaluator.get_sim_specs(
             self.generator.varying_parameters,
@@ -214,7 +225,8 @@ def run(self, n_evals: Optional[int] = None) -> None:
         self._libe_history.H = history
 
         # Update generator with the one received from libE.
-        self.generator._update(persis_info[1]["generator"])
+        if self.libe_comms != "threads":
+            self.generator._update(persis_info[1]["generator"])
 
         # Update number of evaluation in this exploration.
         n_evals_final = self.generator.n_evaluated_trials
@@ -414,7 +426,10 @@ def attach_evaluations(
         # Fill in new rows.
         for field in fields:
             if field in history_new.dtype.names:
-                history_new[field] = evaluation_data[field]
+                # Converting to list prevent the error
+                # "ValueError: setting an array element with a sequence"
+                # when the field contains an array.
+                history_new[field] = evaluation_data[field].to_list()
 
         if not is_history:
             current_time = time.time()
@@ -502,7 +517,7 @@ def _create_libe_history(self) -> History:
         """Initialize an empty libEnsemble history."""
         run_params = self.evaluator.get_run_params()
         gen_specs = self.generator.get_gen_specs(
-            self.sim_workers, run_params, None
+            self.sim_workers, run_params, None, self.libe_comms
         )
         sim_specs = self.evaluator.get_sim_specs(
             self.generator.varying_parameters,
diff --git a/optimas/gen_functions.py b/optimas/gen_functions.py
index a90d9b41..1668be92 100644
--- a/optimas/gen_functions.py
+++ b/optimas/gen_functions.py
@@ -118,7 +118,8 @@ def persistent_generator(H, persis_info, gen_specs, libE_info):
             number_of_gen_points = 0
 
     # Add updated generator to `persis_info`.
-    generator._prepare_to_send()
+    if persis_info["comms"] != "threads":
+        generator._prepare_to_send_back()
     persis_info["generator"] = generator
 
     return H_o, persis_info, FINISHED_PERSISTENT_GEN_TAG
diff --git a/optimas/generators/ax/developer/multitask.py b/optimas/generators/ax/developer/multitask.py
index 7f21f572..2a1f346b 100644
--- a/optimas/generators/ax/developer/multitask.py
+++ b/optimas/generators/ax/developer/multitask.py
@@ -136,11 +136,17 @@ def __init__(
         self._experiment = self._create_experiment()
 
     def get_gen_specs(
-        self, sim_workers: int, run_params: Dict, sim_max: int
+        self,
+        sim_workers: int,
+        run_params: Dict,
+        max_evals: int,
+        libe_comms: str,
     ) -> Dict:
         """Get the libEnsemble gen_specs."""
         # Get base specs.
-        gen_specs = super().get_gen_specs(sim_workers, run_params, sim_max)
+        gen_specs = super().get_gen_specs(
+            sim_workers, run_params, max_evals, libe_comms
+        )
         # Add task to output parameters.
         max_length = max([len(self.lofi_task.name), len(self.hifi_task.name)])
         gen_specs["out"].append(("task", str, max_length))
diff --git a/optimas/generators/base.py b/optimas/generators/base.py
index fe62f2f6..2162e801 100644
--- a/optimas/generators/base.py
+++ b/optimas/generators/base.py
@@ -3,7 +3,7 @@
 from __future__ import annotations
 import os
 from copy import deepcopy
-from typing import List, Dict, Optional, Union
+from typing import List, Dict, Optional, Union, TYPE_CHECKING
 
 import numpy as np
 import pandas as pd
@@ -20,6 +20,8 @@
     TrialParameter,
     TrialStatus,
 )
+if TYPE_CHECKING:
+    from optimas.loggers.base import Logger
 
 logger = get_logger(__name__)
 
@@ -114,6 +116,7 @@ def __init__(
         self._queued_trials = []  # Trials queued to be given for evaluation.
         self._trial_count = 0
         self._check_parameters(self._varying_parameters)
+        self._logger = None
 
     @property
     def varying_parameters(self) -> List[VaryingParameter]:
@@ -150,6 +153,11 @@ def dedicated_resources(self) -> bool:
         """Get whether the generator has dedicated resources allocated."""
         return self._dedicated_resources
 
+    @property
+    def completed_trials(self) -> List[Trial]:
+        """Get list of completed trials."""
+        return [trial for trial in self._given_trials if trial.completed]
+
     @property
     def n_queued_trials(self) -> int:
         """Get the number of trials queued for evaluation."""
@@ -266,6 +274,8 @@ def tell(
             else:
                 log_msg = f"Failed to evaluate trial {trial.index}."
             logger.info(log_msg)
+            if self._logger is not None:
+                self._logger.log_trial(trial, self)
         if allow_saving_model and self._save_model:
             self.save_model_to_file()
 
@@ -510,7 +520,11 @@ def save_model_to_file(self) -> None:
             )
 
     def get_gen_specs(
-        self, sim_workers: int, run_params: Dict, max_evals: int
+        self,
+        sim_workers: int,
+        run_params: Dict,
+        max_evals: int,
+        libe_comms: str,
     ) -> Dict:
         """Get the libEnsemble gen_specs.
 
@@ -523,9 +537,14 @@ def get_gen_specs(
             required.
         max_evals : int
             Maximum number of evaluations to generate.
+        libe_comms : {'local', 'threads', 'mpi'}, optional.
+            The communication mode for libEnseble. Used to determine whether
+            the generator is running on a thread (and therefore in shared
+            memory).
 
         """
-        self._prepare_to_send()
+        if libe_comms != "threads":
+            self._prepare_to_send()
         gen_specs = {
             # Generator function.
             "gen_f": self._gen_function,
@@ -583,6 +602,14 @@ def _prepare_to_send(self) -> None:
         """
         pass
 
+    def _prepare_to_send_back(self) -> None:
+        """Prepare generator to send it back from the gen_f to the manager."""
+        # Removing the logger prevents the WandB error:
+        # "RuntimeError: attach in the same process is not supported currently"
+        # when using multiprocessing.
+        self._logger = None
+        self._prepare_to_send()
+
     def _update(self, new_generator: Generator) -> None:
         """Update generator with the attributes of a newer one.
 
@@ -646,3 +673,7 @@ def _check_parameters(self, parameters: List[VaryingParameter]):
                         f"{self.__class__.__name__} does not support fixing "
                         "the value of a VaryingParameter."
                     )
+
+    def _set_logger(self, logger: Logger) -> None:
+        """Set the generator logger."""
+        self._logger = logger

From ae092d29a11b5e8b9f5c76f570b5902fbb1f166b Mon Sep 17 00:00:00 2001
From: Angel Ferran Pousa <angel.ferran.pousa@desy.de>
Date: Fri, 23 Feb 2024 17:07:40 +0100
Subject: [PATCH 04/25] Add test

---
 tests/test_wandb_logger.py | 119 +++++++++++++++++++++++++++++++++++++
 1 file changed, 119 insertions(+)
 create mode 100644 tests/test_wandb_logger.py

diff --git a/tests/test_wandb_logger.py b/tests/test_wandb_logger.py
new file mode 100644
index 00000000..067fc6b3
--- /dev/null
+++ b/tests/test_wandb_logger.py
@@ -0,0 +1,119 @@
+import os
+# os.environ['WANDB_START_METHOD'] = 'thread'
+# os.environ['WANDB_DISABLE_SERVICE'] = "1"
+# API_KEY = "4b810c5f384e91c204b67371249db1a6b1e1aca4"
+
+import wandb
+import numpy as np
+import matplotlib.pyplot as plt
+from copy import deepcopy
+
+from optimas.explorations import Exploration
+from optimas.generators import RandomSamplingGenerator
+from optimas.evaluators import FunctionEvaluator
+from optimas.core import VaryingParameter, Objective, Parameter
+from optimas.loggers import WandBLogger
+
+
+def eval_func(input_params, output_params):
+    """Evaluation function used for testing"""
+    x0 = input_params["x0"]
+    x1 = input_params["x1"]
+    result = -(x0 + 10 * np.cos(x0)) * (x1 + 5 * np.cos(x1))
+    output_params["f"] = result
+    output_params["p0"] = np.array([[1, 2, 3, 4], [2, 6, 7, 4]])
+    
+    plt.figure()
+    plt.plot(output_params["p0"][0], output_params["p0"][1])
+    output_params["fig"] = deepcopy(plt.gcf())
+    plt.figure()
+    plt.imshow(output_params["p0"])
+    output_params["p1"] = deepcopy(plt.gcf())
+
+
+def user_function(last_trial, generator: RandomSamplingGenerator):    
+    all_trials = generator.completed_trials
+    n_trials = len(all_trials)
+    shape_1 = np.array(all_trials[0].data["p0"]).shape[1]
+    history = np.zeros((n_trials, shape_1))
+    for i, trial in enumerate(all_trials):
+        history[i] = np.array(trial.data["p0"]).sum(axis=0)
+    fig, ax = plt.subplots(figsize=(8, 4))
+    ax.imshow(history.T, aspect="auto")
+    return {"history": wandb.Image(fig)}
+
+
+def test_wandb_logger():
+    """Test that an exploration with a Weights and Biases logger."""
+    # Define variables and objectives.
+    var1 = VaryingParameter("x0", -50.0, 5.0)
+    var2 = VaryingParameter("x1", -5.0, 15.0)
+    obj = Objective("f", minimize=False)
+    # Test also more complex analyzed parameters.
+    p0 = Parameter("p0", dtype=(float, (2, 4)))
+    p1 = Parameter("p1", dtype="O")
+    p2 = Parameter("fig", dtype="O")
+
+    # Create generator.
+    gen = RandomSamplingGenerator(
+        varying_parameters=[var1, var2],
+        objectives=[obj],
+        analyzed_parameters=[p0, p1, p2],
+    )
+
+    # Create function evaluator.
+    ev = FunctionEvaluator(function=eval_func)
+
+    # Create exploration.
+    exploration = Exploration(
+        generator=gen,
+        evaluator=ev,
+        max_evals=10,
+        sim_workers=1,
+        exploration_dir_path="./tests_output/test_wandb_logger",
+        # libe_comms="threads",
+        logger=WandBLogger(
+            api_key="API_KEY",
+            project="project_name",
+            run_name="run_01",
+            data_types={
+                "p0": {"type": wandb.Histogram, "type_kwargs": {}},
+            },
+            user_function=user_function,
+        ),
+    )
+
+    exploration.attach_evaluations(
+        {
+            "x0": [1.0],
+            "x1": [2.0],
+            "f": [0.],
+            "p0": [np.array([[1, 2, 3, 4], [2, 6, 7, 4]])],
+            "p1": [plt.figure()],
+            "fig": [plt.figure()],
+        }
+    )
+
+    # Run exploration.
+    exploration.run(3)
+    exploration.run()
+
+    # Check that the multidimensional analyzed parameters worked as expected.
+    # for p0_data in exploration.history["p0"]:
+    #     np.testing.assert_array_equal(
+    #         np.array(p0_data), np.array([[1, 2, 3, 4], [2, 6, 7, 4]])
+    #     )
+    # for p1_data in exploration.history["p1"]:
+    #     np.testing.assert_array_equal(
+    #         np.array(p1_data), np.array([[1, 2, 3, 4], [2, 6, 7, 4]])
+    #     )
+    # for i, fig in enumerate(exploration.history["fig"]):
+    #     print(threading.current_thread().name)
+    #     fig = deepcopy(fig)
+    #     fig.savefig(
+    #         os.path.join(exploration.exploration_dir_path, f"test_fig_{i}.png")
+    #     )
+
+
+if __name__ == "__main__":
+    test_wandb_logger()

From c2b64e31d54778398957ab7ececa83a6096f6653 Mon Sep 17 00:00:00 2001
From: Angel Ferran Pousa <angel.ferran.pousa@desy.de>
Date: Fri, 23 Feb 2024 17:23:54 +0100
Subject: [PATCH 05/25] Merge branch 'main' into feature/wandb_logger

---
 .github/workflows/unix-openmpi.yml            | 37 +++++++++++++++++++
 .github/workflows/unix.yml                    | 17 ++-------
 .pre-commit-config.yaml                       |  4 +-
 README.md                                     |  9 ++++-
 doc/source/user_guide/dependencies.rst        |  2 +-
 doc/source/user_guide/installation_juwels.rst |  2 +-
 doc/source/user_guide/installation_local.rst  |  9 ++++-
 .../user_guide/installation_maxwell.rst       |  2 +-
 .../user_guide/installation_perlmutter.rst    |  2 +-
 pyproject.toml                                |  2 +-
 10 files changed, 62 insertions(+), 24 deletions(-)
 create mode 100644 .github/workflows/unix-openmpi.yml

diff --git a/.github/workflows/unix-openmpi.yml b/.github/workflows/unix-openmpi.yml
new file mode 100644
index 00000000..6ef8bf3e
--- /dev/null
+++ b/.github/workflows/unix-openmpi.yml
@@ -0,0 +1,37 @@
+name: Unix-OpenMPI
+
+on:
+  pull_request:
+  # Run daily at midnight (UTC).
+  schedule:
+    - cron: '0 0 * * *'
+
+jobs:
+  test:
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        python-version: [3.9, '3.10', 3.11]
+
+    steps:
+    - uses: actions/checkout@v4
+    - uses: conda-incubator/setup-miniconda@v3
+      name: Setup conda
+      with:
+        auto-update-conda: true
+        activate-environment: testing
+        auto-activate-base: false
+        channels: defaults
+        channel-priority: true
+        python-version: ${{ matrix.python-version }}
+
+    - shell: bash -l {0}
+      name: Install dependencies
+      run: |
+        conda install numpy pandas pytorch cpuonly -c pytorch
+        conda install -c conda-forge mpi4py openmpi
+        pip install .[test]
+    - shell: bash -l {0}
+      name: Run unit tests with openMPI
+      run: |
+        python -m pytest tests/
diff --git a/.github/workflows/unix.yml b/.github/workflows/unix.yml
index d003dddc..c728a5f3 100644
--- a/.github/workflows/unix.yml
+++ b/.github/workflows/unix.yml
@@ -1,7 +1,6 @@
-name: Unix
+name: Unix-MPICH
 
 on:
-  push:
   pull_request:
   # Run daily at midnight (UTC).
   schedule:
@@ -15,8 +14,8 @@ jobs:
         python-version: [3.9, '3.10', 3.11]
 
     steps:
-    - uses: actions/checkout@v2
-    - uses: conda-incubator/setup-miniconda@v2
+    - uses: actions/checkout@v4
+    - uses: conda-incubator/setup-miniconda@v3
       name: Setup conda
       with:
         auto-update-conda: true
@@ -36,13 +35,3 @@ jobs:
       name: Run unit tests with MPICH
       run: |
         python -m pytest tests/
-    - shell: bash -l {0}
-      name: Replace MPICH with openMPI
-      run: |
-        conda uninstall mpi4py mpich
-        conda install -c conda-forge mpi4py openmpi
-    - shell: bash -l {0}
-      name: Run unit tests with openMPI
-      run: |
-        rm -r tests_output
-        python -m pytest tests/
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index cb1844b5..46f73284 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -49,7 +49,7 @@ repos:
 
 # Changes tabs to spaces
 - repo: https://github.com/Lucas-C/pre-commit-hooks
-  rev: v1.5.4
+  rev: v1.5.5
   hooks:
   - id: remove-tabs
 
@@ -104,7 +104,7 @@ repos:
 
 # Python Formatting
 - repo: https://github.com/psf/black
-  rev: 24.1.1  # Keep in sync with blacken-docs
+  rev: 24.2.0  # Keep in sync with blacken-docs
   hooks:
   - id: black
 - repo: https://github.com/asottile/blacken-docs
diff --git a/README.md b/README.md
index 66c13ea2..d55a239f 100644
--- a/README.md
+++ b/README.md
@@ -1,4 +1,5 @@
 [![PyPI](https://img.shields.io/pypi/v/optimas)](https://pypi.org/project/optimas/)
+[![Conda Version](https://img.shields.io/conda/vn/conda-forge/optimas.svg)](https://anaconda.org/conda-forge/optimas)
 [![tests badge](https://github.com/optimas-org/optimas/actions/workflows/unix.yml/badge.svg)](https://github.com/optimas-org/optimas/actions)
 [![Documentation Status](https://readthedocs.org/projects/optimas/badge/?version=latest)](https://optimas.readthedocs.io/en/latest/?badge=latest)
 [![DOI](https://zenodo.org/badge/287560975.svg)](https://zenodo.org/badge/latestdoi/287560975)
@@ -39,11 +40,15 @@ Optimas is a Python library designed for highly scalable optimization, from lapt
 
 
 ## Installation
-You can install Optimas from PyPI:
+You can install Optimas from PyPI (recommended):
 ```sh
 pip install optimas
 ```
-Or directly from GitHub:
+from conda-forge:
+```sh
+conda install optimas --channel conda-forge
+```
+or directly from GitHub:
 ```sh
 pip install git+https://github.com/optimas-org/optimas.git
 ```
diff --git a/doc/source/user_guide/dependencies.rst b/doc/source/user_guide/dependencies.rst
index 9286a3bf..636f75be 100644
--- a/doc/source/user_guide/dependencies.rst
+++ b/doc/source/user_guide/dependencies.rst
@@ -22,7 +22,7 @@ See table below for a summary.
 
    * - Generator
      - ``pip install optimas``
-     - ``pip install optimas[all]``
+     - ``pip install 'optimas[all]'``
    * - :class:`~optimas.generators.LineSamplingGenerator`
      - ✅
      - ✅
diff --git a/doc/source/user_guide/installation_juwels.rst b/doc/source/user_guide/installation_juwels.rst
index 2cc70303..33764cab 100644
--- a/doc/source/user_guide/installation_juwels.rst
+++ b/doc/source/user_guide/installation_juwels.rst
@@ -47,7 +47,7 @@ Install ``optimas`` with all dependencies if you plan to do Bayesian optimizatio
 
 .. code::
 
-    pip install optimas[all]
+    pip install 'optimas[all]'
 
 
 Installing FBPIC and Wake-T (optional)
diff --git a/doc/source/user_guide/installation_local.rst b/doc/source/user_guide/installation_local.rst
index 9483c748..8ef70c1a 100644
--- a/doc/source/user_guide/installation_local.rst
+++ b/doc/source/user_guide/installation_local.rst
@@ -45,11 +45,18 @@ Installing with **all** dependencies:
 
 .. code::
 
-    pip install optimas[all]
+    pip install 'optimas[all]'
 
 Use this option if you plan to do Bayesian optimization
 (see :ref:`dependencies` for more details).
 
+Install from conda-forge
+~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. code::
+
+    conda install optimas --channel conda-forge
+
 Install from GitHub
 ~~~~~~~~~~~~~~~~~~~
 This will install the latest development version with all dependencies.
diff --git a/doc/source/user_guide/installation_maxwell.rst b/doc/source/user_guide/installation_maxwell.rst
index c8b32df3..7554e12e 100644
--- a/doc/source/user_guide/installation_maxwell.rst
+++ b/doc/source/user_guide/installation_maxwell.rst
@@ -56,7 +56,7 @@ Install ``optimas`` with all dependencies if you plan to do Bayesian optimizatio
 
 .. code::
 
-    pip install optimas[all]
+    pip install 'optimas[all]'
 
 
 Installing FBPIC and Wake-T (optional)
diff --git a/doc/source/user_guide/installation_perlmutter.rst b/doc/source/user_guide/installation_perlmutter.rst
index d00b944f..a9871074 100644
--- a/doc/source/user_guide/installation_perlmutter.rst
+++ b/doc/source/user_guide/installation_perlmutter.rst
@@ -17,7 +17,7 @@ environment, in which to install *optimas*.
     python3 -m venv $HOME/sw/perlmutter/gpu/venvs/optimas
     source $HOME/sw/perlmutter/gpu/venvs/optimas/bin/activate
 
-    pip install optimas[all]
+    pip install 'optimas[all]'
 
 Running an optimas job
 ~~~~~~~~~~~~~~~~~~~~~~
diff --git a/pyproject.toml b/pyproject.toml
index 3ae663d0..8a6af25c 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -22,7 +22,7 @@ classifiers = [
     'Programming Language :: Python :: 3.11',
 ]
 dependencies = [
-    'libensemble @ git+https://github.com/Libensemble/libensemble@develop',
+    'libensemble >= 1.2',
     'jinja2',
     'pandas',
     'mpi4py',

From cb92d1a593be8dc384da1c2d614a1d9806c8c63b Mon Sep 17 00:00:00 2001
From: Angel Ferran Pousa <angel.ferran.pousa@desy.de>
Date: Fri, 23 Feb 2024 17:41:34 +0100
Subject: [PATCH 06/25] Add logger back to generator after call to `libE`

---
 optimas/explorations/base.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/optimas/explorations/base.py b/optimas/explorations/base.py
index c937d242..fba08ac2 100644
--- a/optimas/explorations/base.py
+++ b/optimas/explorations/base.py
@@ -227,6 +227,8 @@ def run(self, n_evals: Optional[int] = None) -> None:
         # Update generator with the one received from libE.
         if self.libe_comms != "threads":
             self.generator._update(persis_info[1]["generator"])
+            # Restore logger (had to be removed when sending back to manager.)
+            self.generator._set_logger(self._logger)
 
         # Update number of evaluation in this exploration.
         n_evals_final = self.generator.n_evaluated_trials

From d4e550ac14e8e8ba24c8628138c58c30d5707236 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Tue, 27 Feb 2024 16:10:49 +0000
Subject: [PATCH 07/25] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 optimas/generators/base.py | 1 +
 tests/test_wandb_logger.py | 7 +++----
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/optimas/generators/base.py b/optimas/generators/base.py
index 2162e801..c2e43151 100644
--- a/optimas/generators/base.py
+++ b/optimas/generators/base.py
@@ -20,6 +20,7 @@
     TrialParameter,
     TrialStatus,
 )
+
 if TYPE_CHECKING:
     from optimas.loggers.base import Logger
 
diff --git a/tests/test_wandb_logger.py b/tests/test_wandb_logger.py
index 067fc6b3..9bfed7a7 100644
--- a/tests/test_wandb_logger.py
+++ b/tests/test_wandb_logger.py
@@ -1,4 +1,3 @@
-import os
 # os.environ['WANDB_START_METHOD'] = 'thread'
 # os.environ['WANDB_DISABLE_SERVICE'] = "1"
 # API_KEY = "4b810c5f384e91c204b67371249db1a6b1e1aca4"
@@ -22,7 +21,7 @@ def eval_func(input_params, output_params):
     result = -(x0 + 10 * np.cos(x0)) * (x1 + 5 * np.cos(x1))
     output_params["f"] = result
     output_params["p0"] = np.array([[1, 2, 3, 4], [2, 6, 7, 4]])
-    
+
     plt.figure()
     plt.plot(output_params["p0"][0], output_params["p0"][1])
     output_params["fig"] = deepcopy(plt.gcf())
@@ -31,7 +30,7 @@ def eval_func(input_params, output_params):
     output_params["p1"] = deepcopy(plt.gcf())
 
 
-def user_function(last_trial, generator: RandomSamplingGenerator):    
+def user_function(last_trial, generator: RandomSamplingGenerator):
     all_trials = generator.completed_trials
     n_trials = len(all_trials)
     shape_1 = np.array(all_trials[0].data["p0"]).shape[1]
@@ -87,7 +86,7 @@ def test_wandb_logger():
         {
             "x0": [1.0],
             "x1": [2.0],
-            "f": [0.],
+            "f": [0.0],
             "p0": [np.array([[1, 2, 3, 4], [2, 6, 7, 4]])],
             "p1": [plt.figure()],
             "fig": [plt.figure()],

From 4ecfeb68b29212bb01a4fffa3e5162d5e665a928 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=C3=81ngel=20Ferran=20Pousa?= <angel.ferran.pousa@desy.de>
Date: Fri, 8 Mar 2024 15:38:09 +0100
Subject: [PATCH 08/25] Add `wandb` to test dependencies

---
 pyproject.toml | 1 +
 1 file changed, 1 insertion(+)

diff --git a/pyproject.toml b/pyproject.toml
index 8a6af25c..a80c39b4 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -36,6 +36,7 @@ test = [
     'pytest',
     'ax-platform >= 0.3.4',
     'matplotlib',
+    'wandb',
 ]
 all = [
     'ax-platform >= 0.3.4',

From b15371a9a8c31b2f3e68e4f0af65d6fa105a2729 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=C3=81ngel=20Ferran=20Pousa?= <angel.ferran.pousa@desy.de>
Date: Fri, 8 Mar 2024 15:39:49 +0100
Subject: [PATCH 09/25] Update tests with W&B API key

---
 .github/workflows/unix-openmpi.yml |  2 ++
 .github/workflows/unix.yml         |  2 ++
 tests/test_wandb_logger.py         | 27 ++++-----------------------
 3 files changed, 8 insertions(+), 23 deletions(-)

diff --git a/.github/workflows/unix-openmpi.yml b/.github/workflows/unix-openmpi.yml
index 6ef8bf3e..1f620536 100644
--- a/.github/workflows/unix-openmpi.yml
+++ b/.github/workflows/unix-openmpi.yml
@@ -33,5 +33,7 @@ jobs:
         pip install .[test]
     - shell: bash -l {0}
       name: Run unit tests with openMPI
+      env:
+        WANDB_API_KEY: ${{ secrets.WANDB_API_KEY }}
       run: |
         python -m pytest tests/
diff --git a/.github/workflows/unix.yml b/.github/workflows/unix.yml
index c728a5f3..2a273f89 100644
--- a/.github/workflows/unix.yml
+++ b/.github/workflows/unix.yml
@@ -33,5 +33,7 @@ jobs:
         pip install .[test]
     - shell: bash -l {0}
       name: Run unit tests with MPICH
+      env:
+        WANDB_API_KEY: ${{ secrets.WANDB_API_KEY }}
       run: |
         python -m pytest tests/
diff --git a/tests/test_wandb_logger.py b/tests/test_wandb_logger.py
index 9bfed7a7..88c2d99a 100644
--- a/tests/test_wandb_logger.py
+++ b/tests/test_wandb_logger.py
@@ -1,6 +1,4 @@
-# os.environ['WANDB_START_METHOD'] = 'thread'
-# os.environ['WANDB_DISABLE_SERVICE'] = "1"
-# API_KEY = "4b810c5f384e91c204b67371249db1a6b1e1aca4"
+import os
 
 import wandb
 import numpy as np
@@ -70,11 +68,10 @@ def test_wandb_logger():
         max_evals=10,
         sim_workers=1,
         exploration_dir_path="./tests_output/test_wandb_logger",
-        # libe_comms="threads",
         logger=WandBLogger(
-            api_key="API_KEY",
-            project="project_name",
-            run_name="run_01",
+            api_key=os.getenv("WANDB_API_KEY"),
+            project="GitHub actions",
+            run_name="WandB test",
             data_types={
                 "p0": {"type": wandb.Histogram, "type_kwargs": {}},
             },
@@ -97,22 +94,6 @@ def test_wandb_logger():
     exploration.run(3)
     exploration.run()
 
-    # Check that the multidimensional analyzed parameters worked as expected.
-    # for p0_data in exploration.history["p0"]:
-    #     np.testing.assert_array_equal(
-    #         np.array(p0_data), np.array([[1, 2, 3, 4], [2, 6, 7, 4]])
-    #     )
-    # for p1_data in exploration.history["p1"]:
-    #     np.testing.assert_array_equal(
-    #         np.array(p1_data), np.array([[1, 2, 3, 4], [2, 6, 7, 4]])
-    #     )
-    # for i, fig in enumerate(exploration.history["fig"]):
-    #     print(threading.current_thread().name)
-    #     fig = deepcopy(fig)
-    #     fig.savefig(
-    #         os.path.join(exploration.exploration_dir_path, f"test_fig_{i}.png")
-    #     )
-
 
 if __name__ == "__main__":
     test_wandb_logger()

From 73732a0158523a9b462069bd9089cdae0a824f31 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=C3=81ngel=20Ferran=20Pousa?= <angel.ferran.pousa@desy.de>
Date: Fri, 8 Mar 2024 16:11:15 +0100
Subject: [PATCH 10/25] Add docstrings

---
 optimas/loggers/base.py         | 33 +++++++++++++++++++++++++++++----
 optimas/loggers/wandb_logger.py | 29 +++++++++++++++++++++++++++++
 2 files changed, 58 insertions(+), 4 deletions(-)

diff --git a/optimas/loggers/base.py b/optimas/loggers/base.py
index 2615e7d0..9d9ec737 100644
--- a/optimas/loggers/base.py
+++ b/optimas/loggers/base.py
@@ -1,5 +1,6 @@
 from __future__ import annotations
 from typing import TYPE_CHECKING
+from abc import ABC, abstractmethod
 
 
 if TYPE_CHECKING:
@@ -8,16 +9,40 @@
     from optimas.explorations import Exploration
 
 
-class Logger:
+class Logger(ABC):
+    """Base class for all loggers."""
 
     def initialize(self, exploration: Exploration):
-        pass
+        """Initialize logger.
 
-    def log_trial(self, trial: Trial, generator: Generator):
+        This method is called in `Exploration.__init__`.
+
+        Parameters
+        ----------
+        exploration : Exploration
+            The exploration instance to which the logger was attached.
+        """
         pass
 
-    def log_custom_metrics(self, last_trial: Trial, generator: Generator):
+    @abstractmethod
+    def log_trial(self, trial: Trial, generator: Generator):
+        """Log a trial.
+
+        This method is called every time an evaluated trial is given back
+        to the generator.
+
+        Parameters
+        ----------
+        trial : Trial
+            The last trial that has been evaluated.
+        generator : Generator
+            The currently active generator.
+        """
         pass
 
     def finish(self):
+        """Finish logging.
+
+        This method is meant to be called then the exploration is finished.
+        """
         pass
diff --git a/optimas/loggers/wandb_logger.py b/optimas/loggers/wandb_logger.py
index ce814467..028e217c 100644
--- a/optimas/loggers/wandb_logger.py
+++ b/optimas/loggers/wandb_logger.py
@@ -37,7 +37,20 @@ def __init__(
         self._dir = None
 
     def initialize(self, exploration: Exploration):
+        """Initialize the W&B logger.
+
+        This method logs into WandB and created a new run using the output
+        directory if the exploration.
+
+        Parameters
+        ----------
+        exploration : Exploration
+            The exploration instance to which the logger was attached.
+        """
         # Create dir if it doesn't exist.
+        # We need to do this because the logger is typically initialized
+        # before the exploration runs and, thus, before the exploration dir
+        # has been created.
         dir = exploration.exploration_dir_path
         pathlib.Path(dir).mkdir(parents=True, exist_ok=True)
         self._dir = dir
@@ -57,6 +70,18 @@ def initialize(self, exploration: Exploration):
                 self._run_id = self._run.id
 
     def log_trial(self, trial: Trial, generator: Generator):
+        """Log a trial.
+
+        This method is called every time an evaluated trial is given back
+        to the generator.
+
+        Parameters
+        ----------
+        trial : Trial
+            The last trial that has been evaluated.
+        generator : Generator
+            The currently active generator.
+        """
         # Get and process trial data.
         logs = trial.data
         for key in list(logs.keys()):
@@ -92,4 +117,8 @@ def log_trial(self, trial: Trial, generator: Generator):
         self._run.log(logs)
 
     def finish(self):
+        """Finish logging.
+
+        This method is meant to be called then the exploration is finished.
+        """
         self._run.finish()

From 4c58cc5daf93d136e1b0d2f074505004b556c60e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=C3=81ngel=20Ferran=20Pousa?= <angel.ferran.pousa@desy.de>
Date: Fri, 8 Mar 2024 16:33:49 +0100
Subject: [PATCH 11/25] Fix bug

---
 optimas/generators/ax/service/base.py | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/optimas/generators/ax/service/base.py b/optimas/generators/ax/service/base.py
index c69d90f1..d6044140 100644
--- a/optimas/generators/ax/service/base.py
+++ b/optimas/generators/ax/service/base.py
@@ -272,6 +272,10 @@ def _prepare_to_send(self) -> None:
         Delete the fitted model from the generation strategy. It can contain
         pytorch tensors that prevent serialization.
         """
+        self._delete_model()
+
+    def _delete_model(self) -> None:
+        """Delete the fitted model from the generation strategy."""
         generation_strategy = self._ax_client.generation_strategy
         if generation_strategy._model is not None:
             del generation_strategy._curr.model_spec._fitted_model
@@ -298,6 +302,7 @@ def _update(self, new_generator: Generator) -> None:
 
     def _update_parameter(self, parameter):
         """Update a parameter from the search space."""
+        self._delete_model()
         parameters = self._create_ax_parameters()
         new_search_space = InstantiationBase.make_search_space(parameters, None)
         self._ax_client.experiment.search_space.update_parameter(

From 319c2fdab95368e7fe68d952b59ea70732037fbd Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=C3=81ngel=20Ferran=20Pousa?= <angel.ferran.pousa@desy.de>
Date: Fri, 8 Mar 2024 17:12:40 +0100
Subject: [PATCH 12/25] Add `wandb` to RTD dependencies

---
 doc/environment.yaml | 1 +
 1 file changed, 1 insertion(+)

diff --git a/doc/environment.yaml b/doc/environment.yaml
index c9a5a787..9c00eb96 100644
--- a/doc/environment.yaml
+++ b/doc/environment.yaml
@@ -15,3 +15,4 @@ dependencies:
     - sphinx-copybutton
     - sphinx-design
     - sphinx-gallery
+    - wandb

From 969c3477144e1bbd1f6210b4e49757bc265796d0 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=C3=81ngel=20Ferran=20Pousa?= <angel.ferran.pousa@desy.de>
Date: Fri, 8 Mar 2024 17:13:45 +0100
Subject: [PATCH 13/25] Add docstrings

---
 optimas/loggers/base.py         |  2 ++
 optimas/loggers/wandb_logger.py | 39 ++++++++++++++++++++++++++++++++-
 2 files changed, 40 insertions(+), 1 deletion(-)

diff --git a/optimas/loggers/base.py b/optimas/loggers/base.py
index 9d9ec737..8474a851 100644
--- a/optimas/loggers/base.py
+++ b/optimas/loggers/base.py
@@ -1,3 +1,5 @@
+"""This module defines the base Logger class."""
+
 from __future__ import annotations
 from typing import TYPE_CHECKING
 from abc import ABC, abstractmethod
diff --git a/optimas/loggers/wandb_logger.py b/optimas/loggers/wandb_logger.py
index 028e217c..554f240b 100644
--- a/optimas/loggers/wandb_logger.py
+++ b/optimas/loggers/wandb_logger.py
@@ -1,3 +1,5 @@
+"""This module defines the class for logging to Weights and Biases."""
+
 from __future__ import annotations
 import pathlib
 from typing import TYPE_CHECKING, Optional, Callable, Dict
@@ -14,11 +16,46 @@
 
 
 class WandBLogger(Logger):
+    """Weights and Biases logger class.
+
+    Parameters
+    ----------
+    api_key : str
+        The API key used to log into Weight and Biases.
+    project : str
+        Project name.
+    run_name : str, optional
+        Run name. If not given, a random name will be assigned by W&B.
+    run_id : str, optional
+        A unique ID for this run, used for resuming. It must
+        be unique in the project, and if you delete a run you can't reuse
+        the ID. Use the ``run_name`` field for a short descriptive name, or
+        `config` (passed in the ``init_kwargs``)
+        for saving hyperparameters to compare across runs. The ID cannot
+        contain the following special characters: ``/\#?%:``.
+        See the [W&B guide to resuming runs](https://docs.wandb.com/guides/runs/resuming).
+    data_types : Dict, optional
+        A dictionary of the shape ``{name: DataType}``, where ``name`` is the
+        name of a varying parameter, objective or other analyzed parameter and
+        ``DataType`` is a W&B [DataType](https://docs.wandb.ai/ref/python/data-types/).
+        If provided, the given parameters will be converted to the specified
+        data types when logging.
+    user_function : Callable, optional
+        A user-defined function for creating custom logs. This function must
+        be of the shape `custom_logs(trial, generator)`, where ``trial`` is
+        the most recently evaluated trial and ``generator`` is the currently
+        active generator. The function must return a dictionary with the
+        appropriate shape to that it can be given to `wandb.log`.
+    login_kwargs : Dict, optional
+        Additional arguments to pass to ``wandb.login``.
+    init_kwargs : Dict, optional
+        Additional arguments to pass to ``wandb.init``.
+    """
     def __init__(
         self,
         api_key: str,
         project: str,
-        run_name: str,
+        run_name: Optional[str] = None,
         run_id: Optional[str] = None,
         data_types: Optional[Dict] = None,
         user_function: Optional[Callable] = None,

From 371d5784fba5444704b5e4480a4eb2392e134285 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Fri, 8 Mar 2024 16:14:32 +0000
Subject: [PATCH 14/25] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 optimas/loggers/wandb_logger.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/optimas/loggers/wandb_logger.py b/optimas/loggers/wandb_logger.py
index 554f240b..eda2d6df 100644
--- a/optimas/loggers/wandb_logger.py
+++ b/optimas/loggers/wandb_logger.py
@@ -51,6 +51,7 @@ class WandBLogger(Logger):
     init_kwargs : Dict, optional
         Additional arguments to pass to ``wandb.init``.
     """
+
     def __init__(
         self,
         api_key: str,

From d1007060213f422389a7e7f8e4f36ec8d3f928bc Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=C3=81ngel=20Ferran=20Pousa?= <angel.ferran.pousa@desy.de>
Date: Fri, 8 Mar 2024 17:29:54 +0100
Subject: [PATCH 15/25] Fix docstrings

---
 optimas/loggers/base.py         | 6 +++---
 optimas/loggers/wandb_logger.py | 4 ++--
 2 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/optimas/loggers/base.py b/optimas/loggers/base.py
index 8474a851..da25bff8 100644
--- a/optimas/loggers/base.py
+++ b/optimas/loggers/base.py
@@ -17,7 +17,7 @@ class Logger(ABC):
     def initialize(self, exploration: Exploration):
         """Initialize logger.
 
-        This method is called in `Exploration.__init__`.
+        Called in `Exploration.__init__`.
 
         Parameters
         ----------
@@ -30,7 +30,7 @@ def initialize(self, exploration: Exploration):
     def log_trial(self, trial: Trial, generator: Generator):
         """Log a trial.
 
-        This method is called every time an evaluated trial is given back
+        Called every time an evaluated trial is given back
         to the generator.
 
         Parameters
@@ -45,6 +45,6 @@ def log_trial(self, trial: Trial, generator: Generator):
     def finish(self):
         """Finish logging.
 
-        This method is meant to be called then the exploration is finished.
+        Meant to be called when the exploration is finished.
         """
         pass
diff --git a/optimas/loggers/wandb_logger.py b/optimas/loggers/wandb_logger.py
index 554f240b..5c340b72 100644
--- a/optimas/loggers/wandb_logger.py
+++ b/optimas/loggers/wandb_logger.py
@@ -16,7 +16,7 @@
 
 
 class WandBLogger(Logger):
-    """Weights and Biases logger class.
+    r"""Weights and Biases logger class.
 
     Parameters
     ----------
@@ -156,6 +156,6 @@ def log_trial(self, trial: Trial, generator: Generator):
     def finish(self):
         """Finish logging.
 
-        This method is meant to be called then the exploration is finished.
+        Call this method to finish the current run on W&B.
         """
         self._run.finish()

From c22485a5918eb4b894e401676ee100d849afa77d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=C3=81ngel=20Ferran=20Pousa?= <angel.ferran.pousa@desy.de>
Date: Fri, 8 Mar 2024 18:30:21 +0100
Subject: [PATCH 16/25] Add logging docs

---
 doc/source/api/index.rst                      |   1 +
 doc/source/api/loggers.rst                    |   9 +
 .../advanced_usage/log_to_wandb.rst           | 154 ++++++++++++++++++
 doc/source/user_guide/index.rst               |   1 +
 4 files changed, 165 insertions(+)
 create mode 100644 doc/source/api/loggers.rst
 create mode 100644 doc/source/user_guide/advanced_usage/log_to_wandb.rst

diff --git a/doc/source/api/index.rst b/doc/source/api/index.rst
index 0df28597..59d62a95 100644
--- a/doc/source/api/index.rst
+++ b/doc/source/api/index.rst
@@ -11,3 +11,4 @@ This reference manual details all classes included in optimas.
     evaluators
     exploration
     diagnostics
+    loggers
diff --git a/doc/source/api/loggers.rst b/doc/source/api/loggers.rst
new file mode 100644
index 00000000..7736c80e
--- /dev/null
+++ b/doc/source/api/loggers.rst
@@ -0,0 +1,9 @@
+Loggers
+=======
+
+.. currentmodule:: optimas.loggers
+
+.. autosummary::
+    :toctree: _autosummary
+
+    WandBLogger
diff --git a/doc/source/user_guide/advanced_usage/log_to_wandb.rst b/doc/source/user_guide/advanced_usage/log_to_wandb.rst
new file mode 100644
index 00000000..e3651867
--- /dev/null
+++ b/doc/source/user_guide/advanced_usage/log_to_wandb.rst
@@ -0,0 +1,154 @@
+Log an ``Exploration`` to Weights and Biases
+============================================
+
+[Weights and Biases](https://wandb.ai/site) (W&B) is a powerful tool for
+tracking and visualizing
+machine learning experiments. Optimas has built-in support for logging to W&B,
+allowing users to easily track and compare the performance of different
+optimization runs.
+
+This documentation provides a guide on how to use the
+:class:`~optimas.loggers.WandBLogger` class
+within Optimas to log an :class:`~optimas.explorations.Exploration`
+to Weights and Biases.
+
+
+Basic example
+-------------
+
+To log an :class:`~optimas.explorations.Exploration` to Weights and Biases,
+you first need to instantiate
+a :class:`~optimas.loggers.WandBLogger` object. This object requires several
+parameters, including
+your W&B API key, the project name, and optionally, a run name, run ID,
+data types for specific parameters, and a user-defined function for
+custom logs. For example:
+
+.. code-block:: python
+
+    from optimas.loggers import WandBLogger
+
+    logger = WandBLogger(
+        api_key="your_wandb_api_key",
+        project="your_project_name",
+        run_name="example_run",  # optional
+    )
+
+This logger can then be passed to an ``Exploration``, such as in the example
+below:
+
+.. code-block:: python
+
+    from optimas.explorations import Exploration
+    from optimas.generators import RandomSamplingGenerator
+    from optimas.evaluators import FunctionEvaluator
+    from optimas.loggers import WandBLogger
+    from optimas.core import VaryingParameter, Objective
+
+    # Define the function to be optimized
+    def objective_function(inputs, outputs):
+        x = inputs["x"]
+        y = inputs["y"]
+        outputs["result"] = x**2 + y**2
+
+    # Define the evaluator
+    evaluator = FunctionEvaluator(objective_function)
+
+    # Define the generator
+    generator = RandomSamplingGenerator(
+        parameters=[
+            VaryingParameter(name="x", lower_bound=-10, upper_bound=10),
+            VaryingParameter(name="y", lower_bound=-10, upper_bound=10)
+        ],
+        objectives=[
+            Objective(name="result", minimize=True)
+        ]
+    )
+
+    # Instantiate the WandBLogger
+    logger = WandBLogger(
+        api_key="your_wandb_api_key",
+        project="your_project_name",
+        run_name="example_run",
+    )
+
+    # Create the Exploration and pass the logger and evaluator
+    exploration = Exploration(
+        generator=generator,
+        evaluator=evaluator,
+        logger=logger
+    )
+
+    # Run the exploration
+    exploration.run(n_evals=100)
+
+
+Customizing the data type of the logger arguments
+-------------------------------------------------
+
+The `data_types` argument allows you to specify the W&B
+[DataType](https://docs.wandb.ai/ref/python/data-types/) for specific
+parameters when logging to Weights and Biases. This is useful for ensuring
+that your data is logged in the desired format. The `data_types` should be
+a dictionary where the keys are the names of the parameters you wish to
+log, and the values are dictionaries containing the `type` and
+`type_kwargs` for each parameter.
+
+For example, if you have defined two analyzed parameters called
+``"parameter_1"`` and ``"parameter_2"`` that at each evaluation store
+an image or matplotlib
+figure and a numpy array, respectively, you can tell the logger to log the
+first one as an image, and the second as a histogram:
+
+.. code-block:: python
+
+    data_types = {
+        "parameter_1": {"type": wandb.Image, "type_kwargs": {}},
+        "parameter_2": {"type": wandb.Histogram, "type_kwargs": {}}
+    }
+
+    logger = WandBLogger(
+        api_key="your_wandb_api_key",
+        project="your_project_name",
+        data_types=data_types,
+        # Other parameters...
+    )
+
+
+Defining custom logs
+--------------------
+
+By default, the ``WandBLogger`` will log the varying parameters, objectives
+and analyzed parameters of the ``Exploration``.
+If you want to include your own custom logs, you can provide a
+`user_function` that generates them.
+This function will be called every time a trial evaluation finishes.
+
+The `user_function` should take two arguments, which correspond to the most
+recently evaluated `Trial` and the currently active `Generator`.
+You do not need to use them, but they are there for convenience.
+The function must then
+return a dictionary with the appropriate shape to be given to `wandb.log`.
+
+Here's an example of how to define a `user_function` for custom logs:
+
+.. code-block:: python
+
+    def custom_logs(trial, generator):
+        # Example: Log the best score so far
+        best_score = None
+        trials = generator.completed_trials
+        for trial in trials:
+            score = trial.data["result"]
+            if best_score is None:
+                best_score = score
+            elif score < best_score:
+                best_score = score
+        return {"Best Score": best_score}
+
+    logger = WandBLogger(
+        api_key="your_wandb_api_key",
+        project="your_project_name",
+        user_function=custom_logs,
+        # Other parameters...
+    )
diff --git a/doc/source/user_guide/index.rst b/doc/source/user_guide/index.rst
index da3cc7e3..1edab3a7 100644
--- a/doc/source/user_guide/index.rst
+++ b/doc/source/user_guide/index.rst
@@ -27,6 +27,7 @@ User guide
    :caption: Advanced usage
 
    advanced_usage/build_gp_surrogates
+   advanced_usage/log_to_wandb
 
 .. toctree::
    :maxdepth: 1

From c4ffa69359721ca733ce76023dc95aa9ba51f842 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=C3=81ngel=20Ferran=20Pousa?= <angel.ferran.pousa@desy.de>
Date: Fri, 8 Mar 2024 18:32:02 +0100
Subject: [PATCH 17/25] Fix docstring

---
 optimas/loggers/wandb_logger.py | 10 +++++++---
 1 file changed, 7 insertions(+), 3 deletions(-)

diff --git a/optimas/loggers/wandb_logger.py b/optimas/loggers/wandb_logger.py
index 5c340b72..4ccbb586 100644
--- a/optimas/loggers/wandb_logger.py
+++ b/optimas/loggers/wandb_logger.py
@@ -35,9 +35,13 @@ class WandBLogger(Logger):
         contain the following special characters: ``/\#?%:``.
         See the [W&B guide to resuming runs](https://docs.wandb.com/guides/runs/resuming).
     data_types : Dict, optional
-        A dictionary of the shape ``{name: DataType}``, where ``name`` is the
-        name of a varying parameter, objective or other analyzed parameter and
-        ``DataType`` is a W&B [DataType](https://docs.wandb.ai/ref/python/data-types/).
+        A dictionary of the shape
+        ``{"name": {"type": DataType, "type_kwargs": {}}``,
+        where ``name`` is the
+        name of a varying parameter, objective or other analyzed parameter,
+        ``DataType`` is a W&B [DataType](https://docs.wandb.ai/ref/python/data-types/)
+        and ``type_kwargs`` can include additional arguments to pass to the
+        data type.
         If provided, the given parameters will be converted to the specified
         data types when logging.
     user_function : Callable, optional

From e8e06e55a9dd8bf2b58ffc8ca94b0f6a56365f65 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=C3=81ngel=20Ferran=20Pousa?= <angel.ferran.pousa@desy.de>
Date: Fri, 8 Mar 2024 18:40:22 +0100
Subject: [PATCH 18/25] Rename parameters

---
 .../user_guide/advanced_usage/log_to_wandb.rst     | 13 +++++++------
 optimas/loggers/wandb_logger.py                    | 14 +++++++-------
 tests/test_wandb_logger.py                         |  6 +++---
 3 files changed, 17 insertions(+), 16 deletions(-)

diff --git a/doc/source/user_guide/advanced_usage/log_to_wandb.rst b/doc/source/user_guide/advanced_usage/log_to_wandb.rst
index e3651867..887db533 100644
--- a/doc/source/user_guide/advanced_usage/log_to_wandb.rst
+++ b/doc/source/user_guide/advanced_usage/log_to_wandb.rst
@@ -31,7 +31,7 @@ custom logs. For example:
     logger = WandBLogger(
         api_key="your_wandb_api_key",
         project="your_project_name",
-        run_name="example_run",  # optional
+        run="example_run",  # optional
     )
 
 This logger can then be passed to an ``Exploration``, such as in the example
@@ -69,7 +69,7 @@ below:
     logger = WandBLogger(
         api_key="your_wandb_api_key",
         project="your_project_name",
-        run_name="example_run",
+        run="example_run",
     )
 
     # Create the Exploration and pass the logger and evaluator
@@ -121,16 +121,17 @@ Defining custom logs
 By default, the ``WandBLogger`` will log the varying parameters, objectives
 and analyzed parameters of the ``Exploration``.
 If you want to include your own custom logs, you can provide a
-`user_function` that generates them.
+`custom_logs` function that generates them.
 This function will be called every time a trial evaluation finishes.
 
-The `user_function` should take two arguments, which correspond to the most
+The `custom_logs` function should take two arguments, which correspond to the
+most
 recently evaluated `Trial` and the currently active `Generator`.
 You do not need to use them, but they are there for convenience.
 The function must then
 return a dictionary with the appropriate shape to be given to `wandb.log`.
 
-Here's an example of how to define a `user_function` for custom logs:
+Here's an example of how to define a `custom_logs` function:
 
 .. code-block:: python
 
@@ -149,6 +150,6 @@ Here's an example of how to define a `user_function` for custom logs:
     logger = WandBLogger(
         api_key="your_wandb_api_key",
         project="your_project_name",
-        user_function=custom_logs,
+        custom_logs=custom_logs,
         # Other parameters...
     )
diff --git a/optimas/loggers/wandb_logger.py b/optimas/loggers/wandb_logger.py
index 4ccbb586..3b85ddda 100644
--- a/optimas/loggers/wandb_logger.py
+++ b/optimas/loggers/wandb_logger.py
@@ -24,12 +24,12 @@ class WandBLogger(Logger):
         The API key used to log into Weight and Biases.
     project : str
         Project name.
-    run_name : str, optional
+    run : str, optional
         Run name. If not given, a random name will be assigned by W&B.
     run_id : str, optional
         A unique ID for this run, used for resuming. It must
         be unique in the project, and if you delete a run you can't reuse
-        the ID. Use the ``run_name`` field for a short descriptive name, or
+        the ID. Use the ``run`` field for a short descriptive name, or
         `config` (passed in the ``init_kwargs``)
         for saving hyperparameters to compare across runs. The ID cannot
         contain the following special characters: ``/\#?%:``.
@@ -44,7 +44,7 @@ class WandBLogger(Logger):
         data type.
         If provided, the given parameters will be converted to the specified
         data types when logging.
-    user_function : Callable, optional
+    custom_logs : Callable, optional
         A user-defined function for creating custom logs. This function must
         be of the shape `custom_logs(trial, generator)`, where ``trial`` is
         the most recently evaluated trial and ``generator`` is the currently
@@ -59,19 +59,19 @@ def __init__(
         self,
         api_key: str,
         project: str,
-        run_name: Optional[str] = None,
+        run: Optional[str] = None,
         run_id: Optional[str] = None,
         data_types: Optional[Dict] = None,
-        user_function: Optional[Callable] = None,
+        custom_logs: Optional[Callable] = None,
         login_kwargs: Optional[Dict] = None,
         init_kwargs: Optional[Dict] = None,
     ) -> None:
         self._api_key = api_key
         self._project = project
-        self._run_name = run_name
+        self._run_name = run
         self._run_id = run_id
         self._data_types = {} if data_types is None else data_types
-        self._user_function = user_function
+        self._user_function = custom_logs
         self._login_kwargs = {} if login_kwargs is None else login_kwargs
         self._init_kwargs = {} if init_kwargs is None else init_kwargs
         self._run = None
diff --git a/tests/test_wandb_logger.py b/tests/test_wandb_logger.py
index 88c2d99a..46c14ad1 100644
--- a/tests/test_wandb_logger.py
+++ b/tests/test_wandb_logger.py
@@ -28,7 +28,7 @@ def eval_func(input_params, output_params):
     output_params["p1"] = deepcopy(plt.gcf())
 
 
-def user_function(last_trial, generator: RandomSamplingGenerator):
+def custom_logs(last_trial, generator: RandomSamplingGenerator):
     all_trials = generator.completed_trials
     n_trials = len(all_trials)
     shape_1 = np.array(all_trials[0].data["p0"]).shape[1]
@@ -71,11 +71,11 @@ def test_wandb_logger():
         logger=WandBLogger(
             api_key=os.getenv("WANDB_API_KEY"),
             project="GitHub actions",
-            run_name="WandB test",
+            run="WandB test",
             data_types={
                 "p0": {"type": wandb.Histogram, "type_kwargs": {}},
             },
-            user_function=user_function,
+            custom_logs=custom_logs,
         ),
     )
 

From cae06897ede5462f9de1d9f00b24bb02ccfa397e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=C3=81ngel=20Ferran=20Pousa?= <angel.ferran.pousa@desy.de>
Date: Fri, 8 Mar 2024 18:48:11 +0100
Subject: [PATCH 19/25] Fix links and docs

---
 .../advanced_usage/log_to_wandb.rst           | 23 ++++++++++---------
 optimas/loggers/wandb_logger.py               |  4 ++--
 2 files changed, 14 insertions(+), 13 deletions(-)

diff --git a/doc/source/user_guide/advanced_usage/log_to_wandb.rst b/doc/source/user_guide/advanced_usage/log_to_wandb.rst
index 887db533..c076d797 100644
--- a/doc/source/user_guide/advanced_usage/log_to_wandb.rst
+++ b/doc/source/user_guide/advanced_usage/log_to_wandb.rst
@@ -1,7 +1,7 @@
 Log an ``Exploration`` to Weights and Biases
 ============================================
 
-[Weights and Biases](https://wandb.ai/site) (W&B) is a powerful tool for
+`Weights and Biases <https://wandb.ai/site>`_ (W&B) is a powerful tool for
 tracking and visualizing
 machine learning experiments. Optimas has built-in support for logging to W&B,
 allowing users to easily track and compare the performance of different
@@ -86,13 +86,13 @@ below:
 Customizing the data type of the logger arguments
 -------------------------------------------------
 
-The `data_types` argument allows you to specify the W&B
-[DataType](https://docs.wandb.ai/ref/python/data-types/) for specific
+The ``data_types`` argument allows you to specify the W&B
+`data type <https://docs.wandb.ai/ref/python/data-types/>`_ for specific
 parameters when logging to Weights and Biases. This is useful for ensuring
-that your data is logged in the desired format. The `data_types` should be
+that your data is logged in the desired format. The ``data_types`` should be
 a dictionary where the keys are the names of the parameters you wish to
-log, and the values are dictionaries containing the `type` and
-`type_kwargs` for each parameter.
+log, and the values are dictionaries containing the ``type`` and
+``type_kwargs`` for each parameter.
 
 For example, if you have defined two analyzed parameters called
 ``"parameter_1"`` and ``"parameter_2"`` that at each evaluation store
@@ -121,17 +121,18 @@ Defining custom logs
 By default, the ``WandBLogger`` will log the varying parameters, objectives
 and analyzed parameters of the ``Exploration``.
 If you want to include your own custom logs, you can provide a
-`custom_logs` function that generates them.
+``custom_logs`` function that generates them.
 This function will be called every time a trial evaluation finishes.
 
-The `custom_logs` function should take two arguments, which correspond to the
+The ``custom_logs`` function should take two arguments, which correspond to the
 most
-recently evaluated `Trial` and the currently active `Generator`.
+recently evaluated :class:`~optimas.core.Trial` and the currently active
+``Generator``.
 You do not need to use them, but they are there for convenience.
 The function must then
-return a dictionary with the appropriate shape to be given to `wandb.log`.
+return a dictionary with the appropriate shape to be given to ``wandb.log``.
 
-Here's an example of how to define a `custom_logs` function:
+Here's an example of how to define a ``custom_logs`` function:
 
 .. code-block:: python
 
diff --git a/optimas/loggers/wandb_logger.py b/optimas/loggers/wandb_logger.py
index 3b85ddda..d7a1fd39 100644
--- a/optimas/loggers/wandb_logger.py
+++ b/optimas/loggers/wandb_logger.py
@@ -33,13 +33,13 @@ class WandBLogger(Logger):
         `config` (passed in the ``init_kwargs``)
         for saving hyperparameters to compare across runs. The ID cannot
         contain the following special characters: ``/\#?%:``.
-        See the [W&B guide to resuming runs](https://docs.wandb.com/guides/runs/resuming).
+        See the `W&B guide to resuming runs <https://docs.wandb.com/guides/runs/resuming>`_.
     data_types : Dict, optional
         A dictionary of the shape
         ``{"name": {"type": DataType, "type_kwargs": {}}``,
         where ``name`` is the
         name of a varying parameter, objective or other analyzed parameter,
-        ``DataType`` is a W&B [DataType](https://docs.wandb.ai/ref/python/data-types/)
+        ``DataType`` is a W&B `DataType <https://docs.wandb.ai/ref/python/data-types/>`_
         and ``type_kwargs`` can include additional arguments to pass to the
         data type.
         If provided, the given parameters will be converted to the specified

From f29e0bdd27d754090ea41e7e3b6336adfdc1d3cc Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Fri, 8 Mar 2024 17:55:02 +0000
Subject: [PATCH 20/25] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 .../user_guide/advanced_usage/log_to_wandb.rst    | 15 +++++++--------
 1 file changed, 7 insertions(+), 8 deletions(-)

diff --git a/doc/source/user_guide/advanced_usage/log_to_wandb.rst b/doc/source/user_guide/advanced_usage/log_to_wandb.rst
index c076d797..a6582486 100644
--- a/doc/source/user_guide/advanced_usage/log_to_wandb.rst
+++ b/doc/source/user_guide/advanced_usage/log_to_wandb.rst
@@ -45,12 +45,14 @@ below:
     from optimas.loggers import WandBLogger
     from optimas.core import VaryingParameter, Objective
 
+
     # Define the function to be optimized
     def objective_function(inputs, outputs):
         x = inputs["x"]
         y = inputs["y"]
         outputs["result"] = x**2 + y**2
 
+
     # Define the evaluator
     evaluator = FunctionEvaluator(objective_function)
 
@@ -58,11 +60,9 @@ below:
     generator = RandomSamplingGenerator(
         parameters=[
             VaryingParameter(name="x", lower_bound=-10, upper_bound=10),
-            VaryingParameter(name="y", lower_bound=-10, upper_bound=10)
+            VaryingParameter(name="y", lower_bound=-10, upper_bound=10),
         ],
-        objectives=[
-            Objective(name="result", minimize=True)
-        ]
+        objectives=[Objective(name="result", minimize=True)],
     )
 
     # Instantiate the WandBLogger
@@ -74,9 +74,7 @@ below:
 
     # Create the Exploration and pass the logger and evaluator
     exploration = Exploration(
-        generator=generator,
-        evaluator=evaluator,
-        logger=logger
+        generator=generator, evaluator=evaluator, logger=logger
     )
 
     # Run the exploration
@@ -104,7 +102,7 @@ first one as an image, and the second as a histogram:
 
     data_types = {
         "parameter_1": {"type": wandb.Image, "type_kwargs": {}},
-        "parameter_2": {"type": wandb.Histogram, "type_kwargs": {}}
+        "parameter_2": {"type": wandb.Histogram, "type_kwargs": {}},
     }
 
     logger = WandBLogger(
@@ -148,6 +146,7 @@ Here's an example of how to define a ``custom_logs`` function:
                 best_score = score
         return {"Best Score": best_score}
 
+
     logger = WandBLogger(
         api_key="your_wandb_api_key",
         project="your_project_name",

From 02be9d39e1d1678174d51cf6fed104128a1a607f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=C3=81ngel=20Ferran=20Pousa?= <angel.ferran.pousa@desy.de>
Date: Fri, 8 Mar 2024 19:07:14 +0100
Subject: [PATCH 21/25] Fix docstring

---
 optimas/loggers/base.py         | 2 +-
 optimas/loggers/wandb_logger.py | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/optimas/loggers/base.py b/optimas/loggers/base.py
index da25bff8..4b4e921c 100644
--- a/optimas/loggers/base.py
+++ b/optimas/loggers/base.py
@@ -1,4 +1,4 @@
-"""This module defines the base Logger class."""
+"""Contains the definition of the base Logger class."""
 
 from __future__ import annotations
 from typing import TYPE_CHECKING
diff --git a/optimas/loggers/wandb_logger.py b/optimas/loggers/wandb_logger.py
index 8039d649..b167cbe9 100644
--- a/optimas/loggers/wandb_logger.py
+++ b/optimas/loggers/wandb_logger.py
@@ -1,4 +1,4 @@
-"""This module defines the class for logging to Weights and Biases."""
+"""Contains the definition of the class for logging to Weights and Biases."""
 
 from __future__ import annotations
 import pathlib

From 58c1df4ca36d8aaaf31d5a5f8d3a1ad804706f08 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=C3=81ngel=20Ferran=20Pousa?= <angel.ferran.pousa@desy.de>
Date: Fri, 8 Mar 2024 19:25:39 +0100
Subject: [PATCH 22/25] Warn if `wandb` is not installed

---
 optimas/loggers/wandb_logger.py | 12 +++++++++++-
 1 file changed, 11 insertions(+), 1 deletion(-)

diff --git a/optimas/loggers/wandb_logger.py b/optimas/loggers/wandb_logger.py
index b167cbe9..1fb22020 100644
--- a/optimas/loggers/wandb_logger.py
+++ b/optimas/loggers/wandb_logger.py
@@ -5,7 +5,11 @@
 from typing import TYPE_CHECKING, Optional, Callable, Dict
 
 from matplotlib.figure import Figure
-import wandb
+try:
+    import wandb
+    wandb_installed = True
+except ImportError:
+    wandb_installed = False
 
 from .base import Logger
 
@@ -67,6 +71,12 @@ def __init__(
         login_kwargs: Optional[Dict] = None,
         init_kwargs: Optional[Dict] = None,
     ) -> None:
+        if not wandb_installed:
+            raise ImportError(
+                "Logging to Weights and Biases requires `wandb` to be "
+                "installed. Please install it by running "
+                "`pip install wandb`."
+            )
         self._api_key = api_key
         self._project = project
         self._run_name = run

From 4932d09adab03576a86d6c2578d52739b026712c Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Fri, 8 Mar 2024 19:17:49 +0000
Subject: [PATCH 23/25] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 optimas/loggers/wandb_logger.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/optimas/loggers/wandb_logger.py b/optimas/loggers/wandb_logger.py
index 1fb22020..d0185583 100644
--- a/optimas/loggers/wandb_logger.py
+++ b/optimas/loggers/wandb_logger.py
@@ -5,8 +5,10 @@
 from typing import TYPE_CHECKING, Optional, Callable, Dict
 
 from matplotlib.figure import Figure
+
 try:
     import wandb
+
     wandb_installed = True
 except ImportError:
     wandb_installed = False

From 9732815871cc526f37e7a43feb9cccd9d6213819 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=C3=81ngel=20Ferran=20Pousa?= <angel.ferran.pousa@desy.de>
Date: Fri, 8 Mar 2024 20:25:34 +0100
Subject: [PATCH 24/25] Improve test description

---
 tests/test_wandb_logger.py | 12 ++++++++++--
 1 file changed, 10 insertions(+), 2 deletions(-)

diff --git a/tests/test_wandb_logger.py b/tests/test_wandb_logger.py
index 46c14ad1..a090bf23 100644
--- a/tests/test_wandb_logger.py
+++ b/tests/test_wandb_logger.py
@@ -29,6 +29,7 @@ def eval_func(input_params, output_params):
 
 
 def custom_logs(last_trial, generator: RandomSamplingGenerator):
+    """Make and log a cumulative plot of all trials."""
     all_trials = generator.completed_trials
     n_trials = len(all_trials)
     shape_1 = np.array(all_trials[0].data["p0"]).shape[1]
@@ -41,7 +42,12 @@ def custom_logs(last_trial, generator: RandomSamplingGenerator):
 
 
 def test_wandb_logger():
-    """Test that an exploration with a Weights and Biases logger."""
+    """Test an exploration with a Weights and Biases logger.
+    
+    In addition to the varying parameters and objectives, three analyzed
+    parameters of different type are added: an array and two objects. One
+    of the objects will store a matplotlib figure.
+    """
     # Define variables and objectives.
     var1 = VaryingParameter("x0", -50.0, 5.0)
     var2 = VaryingParameter("x1", -5.0, 15.0)
@@ -79,6 +85,8 @@ def test_wandb_logger():
         ),
     )
 
+    # Test also more exotic use cases where the first step is to attach
+    # manual evaluations.
     exploration.attach_evaluations(
         {
             "x0": [1.0],
@@ -90,7 +98,7 @@ def test_wandb_logger():
         }
     )
 
-    # Run exploration.
+    # Run exploration in two steps.
     exploration.run(3)
     exploration.run()
 

From 6b067866e27a73d2205f55d767cf138fae665d98 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Tue, 12 Mar 2024 09:54:21 +0000
Subject: [PATCH 25/25] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 tests/test_wandb_logger.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/test_wandb_logger.py b/tests/test_wandb_logger.py
index a090bf23..1300ca17 100644
--- a/tests/test_wandb_logger.py
+++ b/tests/test_wandb_logger.py
@@ -43,7 +43,7 @@ def custom_logs(last_trial, generator: RandomSamplingGenerator):
 
 def test_wandb_logger():
     """Test an exploration with a Weights and Biases logger.
-    
+
     In addition to the varying parameters and objectives, three analyzed
     parameters of different type are added: an array and two objects. One
     of the objects will store a matplotlib figure.