Make Ax optional

doc/source/user_guide/dependencies.rst

@@ -1,9 +1,46 @@
+.. _dependencies:
+
 Dependencies
 ============
 
 Optimas relies on the following packages:
 
-* `mpi4py <https://pypi.org/project/mpi4py/>`_ - Python bindings for MPI. Required for launching parallel simulations.
-* `libEnsemble <https://pypi.org/project/libensemble/>`_ - The backbone of optimas, orchestrates the concurrent evaluation of simulations, the resource detection and allocation, and the communication between simulations and manager.
 * `jinja2 <https://pypi.org/project/jinja2/>`_ - Needed to generate simulation scripts from templates.
-* `Ax <https://pypi.org/project/ax-platform/>`_ - Algorithms for Bayesian optimization.
+* `libEnsemble <https://pypi.org/project/libensemble/>`_ - The backbone of optimas, orchestrates the concurrent evaluation of simulations, the resource detection and allocation, and the communication between simulations and manager.
+* `mpi4py <https://pypi.org/project/mpi4py/>`_ - Python bindings for MPI. Required for launching parallel simulations.
+* `pandas <https://pypi.org/project/pandas/>`_ - Enable output as pandas DataFrames.
+* `pydantic <https://pypi.org/project/pydantic/>`_ - Input validation and object serialization.
+* (optional) `Ax <https://pypi.org/project/ax-platform/>`_ - Algorithms for Bayesian optimization.
+
+
+The installed dependencies will determine which generators are available for use.
+See table below for a summary.
+
+.. list-table:: Available generators and their dependencies
+   :widths: 35 25 25
+   :header-rows: 1
+
+   * - Generator
+     - ``pip install optimas``
+     - ``pip install optimas[all]``
+   * - :class:`~optimas.generators.LineSamplingGenerator`
+     - ✅
+     - ✅
+   * - :class:`~optimas.generators.GridSamplingGenerator`
+     - ✅
+     - ✅
+   * - :class:`~optimas.generators.RandomSamplingGenerator`
+     - ✅
+     - ✅
+   * - :class:`~optimas.generators.AxSingleFidelityGenerator`
+     - ❌
+     - ✅
+   * - :class:`~optimas.generators.AxMultiFidelityGenerator`
+     - ❌
+     - ✅
+   * - :class:`~optimas.generators.AxMultitaskGenerator`
+     - ❌
+     - ✅
+   * - :class:`~optimas.generators.AxClientGenerator`
+     - ❌
+     - ✅

doc/source/user_guide/installation_juwels.rst

@@ -42,11 +42,12 @@ Activate the environment
     source $PROJECT/<username>/pyenvs/optimas_env/bin/activate
 
 
-Install ``optimas``
+Install ``optimas`` with all dependencies if you plan to do Bayesian optimization
+(see :ref:`dependencies` for more details).
 
 .. code::
 
-    pip install optimas
+    pip install optimas[all]
 
 
 Installing FBPIC and Wake-T (optional)

doc/source/user_guide/installation_local.rst

@@ -31,10 +31,29 @@ On Windows:
 
     conda install -c conda-forge mpi4py msmpi
 
-Install optimas
-~~~~~~~~~~~~~~~
-Install the latest release from PyPI
+Install optimas from PyPI (recommended)
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+This will install the latest stable release.
+
+Installing with only the **basic** dependencies:
 
 .. code::
 
     pip install optimas
+
+Installing with **all** dependencies:
+
+.. code::
+
+    pip install optimas[all]
+
+Use this option if you plan to do Bayesian optimization
+(see :ref:`dependencies` for more details).
+
+Install from GitHub
+~~~~~~~~~~~~~~~~~~~
+This will install the latest development version with all dependencies.
+
+.. code::
+
+    pip install "optimas[all] @ git+https://github.com/optimas-org/optimas.git"

doc/source/user_guide/installation_maxwell.rst

@@ -51,11 +51,12 @@ Install ``mpi4py``
     pip install mpi4py --no-cache-dir
 
 
-Install ``optimas``
+Install ``optimas`` with all dependencies if you plan to do Bayesian optimization
+(see :ref:`dependencies` for more details).
 
 .. code::
 
-    pip install optimas
+    pip install optimas[all]
 
 
 Installing FBPIC and Wake-T (optional)

optimas/generators/__init__.py

@@ -1,7 +1,24 @@
-from .ax.service.single_fidelity import AxSingleFidelityGenerator
-from .ax.service.multi_fidelity import AxMultiFidelityGenerator
-from .ax.service.ax_client import AxClientGenerator
-from .ax.developer.multitask import AxMultitaskGenerator
+# Import Ax generators
+try:
+    from .ax.service.single_fidelity import AxSingleFidelityGenerator
+    from .ax.service.multi_fidelity import AxMultiFidelityGenerator
+    from .ax.service.ax_client import AxClientGenerator
+    from .ax.developer.multitask import AxMultitaskGenerator
+except ImportError as e:
+    if e.__str__() == "No module named 'ax'":
+        # Replace generators by dummy generators that will
+        # raise an error only if the user tries to instantiate them
+        # and tell them to install ax-platform
+        from .ax.import_error_dummy_generator import AxImportErrorDummyGenerator
+
+        AxSingleFidelityGenerator = AxImportErrorDummyGenerator
+        AxMultiFidelityGenerator = AxImportErrorDummyGenerator
+        AxClientGenerator = AxImportErrorDummyGenerator
+        AxMultitaskGenerator = AxImportErrorDummyGenerator
+    else:
+        raise (e)
+
+# Import optimas native generators
 from .grid_sampling import GridSamplingGenerator
 from .line_sampling import LineSamplingGenerator
 from .random_sampling import RandomSamplingGenerator

optimas/generators/ax/import_error_dummy_generator.py

@@ -0,0 +1,16 @@
+"""Contains the definition of dummy generator that raises an import error."""
+
+
+class AxImportErrorDummyGenerator(object):
+    """Class that raises an error when instantiated, telling the user to install ax-platform.
+
+    This class replaces all other Ax-based classes,
+    when Ax is not installed
+    """
+
+    def __init__(self, *args, **kwargs) -> None:
+        raise RuntimeError(
+            "You need to install ax-platform, in order "
+            "to use Ax-based generators in optimas.\n"
+            "e.g. with `pip install ax-platform >= 0.3.4`"
+        )

pyproject.toml

@@ -24,7 +24,7 @@ classifiers = [
 dependencies = [
     'libensemble @ git+https://github.com/Libensemble/libensemble@develop',
     'jinja2',
-    'ax-platform >= 0.3.4',
+    'pandas',
     'mpi4py',
     'pydantic >= 2.0',
 ]
@@ -34,8 +34,12 @@ dynamic = ['version']
 test = [
     'flake8',
     'pytest',
+    'ax-platform >= 0.3.4',
     'matplotlib',
 ]
+all = [
+    'ax-platform >= 0.3.4',
+]
 
 [project.urls]
 Documentation = 'https://optimas.readthedocs.io/'