Merge pull request #263 from optimas-org/feature/optional_ax

Use postponed annotations (PEP 563)

Author: RemiLehe

.github/workflows/unix-noax.yml

@@ -0,0 +1,39 @@
+name: Unix-Without-Ax
+
+on:
+  pull_request:
+  # Run daily at midnight (UTC).
+  schedule:
+    - cron: '0 0 * * *'
+
+jobs:
+  test:
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        python-version: ['3.10', 3.11, 3.12]
+
+    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 mpich
+        pip install .[test]
+        pip uninstall --yes ax-platform # Run without Ax
+    - shell: bash -l {0}
+      name: Run unit tests without Ax
+      run: |
+        python -m pytest tests/ --ignore=tests/test_ax_generators.py --ignore=tests/test_ax_model_manager.py  --ignore=tests/test_gpu_resources.py
+        mpirun -np 3 python -m pytest --with-mpi tests/test_grid_sampling_mpi.py

optimas/utils/ax/ax_model_manager.py

@@ -1,6 +1,17 @@
 """Contains the definition of the AxModelManager class."""
 
-from typing import Optional, Union, List, Tuple, Dict, Any, Literal
+from __future__ import annotations
+
+from typing import (
+    Optional,
+    Union,
+    List,
+    Tuple,
+    Dict,
+    Any,
+    Literal,
+    TYPE_CHECKING,
+)
 
 import numpy as np
 from numpy.typing import NDArray
@@ -34,6 +45,10 @@
 from optimas.core import VaryingParameter, Objective
 from optimas.utils.other import convert_to_dataframe
 
+if TYPE_CHECKING:
+    from ax.service.ax_client import AxClient
+    from ax.modelbridge.torch import TorchModelBridge
+
 
 class AxModelManager:
     """Class for building and exploring GP models using an ``AxClient``.