Fix import of AxModelManager

Author: RemiLehe

optimas/generators/__init__.py

@@ -4,8 +4,8 @@
     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'":
+except (ImportError, ModuleNotFoundError) as e:
+    if e.__str__().startswith("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

optimas/utils/ax/ax_model_manager.py

@@ -29,8 +29,12 @@
 
     ax_installed = True
 except ImportError:
+    from .import_error_dummy_class import AxImportErrorDummyClass
+    AxClient = AxImportErrorDummyClass
+    TorchModelBridge = AxImportErrorDummyClass
     ax_installed = False
 
+
 from optimas.core import VaryingParameter, Objective
 from optimas.utils.other import convert_to_dataframe
 

optimas/utils/ax/import_error_dummy_class.py

@@ -0,0 +1,15 @@
+"""Contains the definition of dummy class that raises an import error."""
+
+
+class AxImportErrorDummyClass(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.5.0`"
+        )