Merge pull request #45 from robin-janssen/improve-unit-tests

Improve unit tests

Author: robin-janssen

.github/workflows/ci.yaml

@@ -21,32 +21,31 @@ jobs:
 
     strategy:
       matrix:
-        python-version: ['3.10']
+        python-version: ["3.10"]
 
     steps:
       - name: Check out the repository
         uses: actions/checkout@v4
         with:
-          fetch-depth: 0  # Fetch all history to support dependency checks
+          fetch-depth: 0 # Fetch all history to support dependency checks
           ref: ${{ github.head_ref || github.ref }}
 
       - name: Set up Python
         uses: actions/setup-python@v5
         with:
           python-version: ${{ matrix.python-version }}
 
-      - name: Install Poetry
+      - name: Install uv
         run: |
-          curl -sSL https://install.python-poetry.org | python3 -
-          echo "$HOME/.local/bin" >> $GITHUB_PATH
+          curl -LsSf https://astral.sh/uv/install.sh | sh
+          echo "$HOME/.cargo/bin" >> $GITHUB_PATH
 
-      - name: Install Poetry Plugin for Export
+      - name: Setup Python environment with uv
         run: |
-          poetry self add poetry-plugin-export
-
-      - name: Install project with dev dependencies
-        run: |
-          poetry install --with dev
+          uv venv
+          source .venv/bin/activate
+          uv pip install -e .
+          uv pip install --group dev
 
       - name: Check for dependency changes
         id: deps
@@ -66,8 +65,8 @@ jobs:
 
           echo "Changed files: $CHANGED_FILES"
 
-          # Check if pyproject.toml or poetry.lock has changed
-          if echo "$CHANGED_FILES" | grep -qE '(^|/)pyproject\.toml$|(^|/)poetry\.lock$'; then
+          # Check if pyproject.toml has changed
+          if echo "$CHANGED_FILES" | grep -qE '(^|/)pyproject\.toml$'; then
             echo "dependencies-changed=true" >> $GITHUB_ENV
           else
             echo "dependencies-changed=false" >> $GITHUB_ENV
@@ -76,7 +75,8 @@ jobs:
       - name: Generate requirements.txt
         if: env.dependencies-changed == 'true'
         run: |
-          poetry export -f requirements.txt --output requirements.txt --without-hashes
+          source .venv/bin/activate
+          uv pip freeze > requirements.txt
 
       - name: Commit and push updated requirements.txt
         if: |
@@ -92,58 +92,66 @@ jobs:
 
       - name: Run Black (auto-reformat)
         run: |
-          poetry run black .
+          source .venv/bin/activate
+          black .
 
       - name: Run isort (auto-reformat)
         run: |
-          poetry run isort .
+          source .venv/bin/activate
+          isort .
 
       - name: Run pytest and generate coverage report
         run: |
-          poetry run pytest --cov-report=term-missing:skip-covered --cov=codes test/ --cov-report=xml:coverage.xml
+          source .venv/bin/activate
+          pytest --cov-report=term-missing:skip-covered --cov=codes test/ --cov-report=xml:coverage.xml
 
       - name: Upload results to Codecov
         uses: codecov/codecov-action@v4
         with:
-          token: ${{ secrets.CODECOV_TOKEN }}  
-          files: ./coverage.xml  # Path to the coverage XML file, adjust if necessary
-          fail_ci_if_error: true  # Optional, ensures the CI fails if Codecov upload fails
+          token: ${{ secrets.CODECOV_TOKEN }}
+          files: ./coverage.xml # Path to the coverage XML file, adjust if necessary
+          fail_ci_if_error: true # Optional, ensures the CI fails if Codecov upload fails
 
   docs:
     if: github.ref == 'refs/heads/main' || (github.event_name == 'pull_request' && github.event.pull_request.base.ref == 'main')
     runs-on: ubuntu-latest
 
     strategy:
       matrix:
-        python-version: ['3.10']
+        python-version: ["3.10"]
 
     steps:
       - name: Check out the repository
         uses: actions/checkout@v4
         with:
-          fetch-depth: 0  # Fetch all history if necessary
+          fetch-depth: 0 # Fetch all history if necessary
 
       - name: Set up Python
         uses: actions/setup-python@v5
         with:
           python-version: ${{ matrix.python-version }}
 
-      - name: Install Poetry
+      - name: Install uv
         run: |
-          curl -sSL https://install.python-poetry.org | python3 -
-          echo "$HOME/.local/bin" >> $GITHUB_PATH
+          curl -LsSf https://astral.sh/uv/install.sh | sh
+          echo "$HOME/.cargo/bin" >> $GITHUB_PATH
 
-      - name: Install project with dev dependencies
+      - name: Setup Python environment with uv
         run: |
-          poetry install --with dev
+          uv venv
+          source .venv/bin/activate
+          uv pip install -e .
+          uv pip install --group dev
 
       - name: Generate API Documentation with Sphinx
         run: |
-          poetry run sphinx-apidoc -o docs/ codes
+          source .venv/bin/activate
+          sphinx-apidoc -o docs/ codes
 
       - name: Build HTML with Sphinx
         run: |
-          poetry run sphinx-build -b html docs/ docs/_build
+          source .venv/bin/activate
+          sphinx-build -b html docs/ docs/_build
 
       - name: Deploy Sphinx API docs to gh-pages
         uses: peaceiris/actions-gh-pages@v4

.python-version

@@ -0,0 +1 @@
+3.10

AGENT.md

@@ -0,0 +1,26 @@
+# CODES Benchmark - Agent Guide
+
+## Build/Test Commands
+- Install dependencies: `poetry install` or `pip install -r requirements.txt`
+- Run tests: `pytest test/` or `python -m pytest test/`
+- Run single test: `pytest test/test_data.py::test_function_name`
+- Run training: `python run_training.py --config config.yaml`
+- Run evaluation: `python run_eval.py --config config.yaml`
+- Run hyperparameter tuning: `python run_tuning.py --config config.yaml`
+
+## Architecture & Structure
+- `codes/` - Main package with 4 modules: benchmark, surrogates, train, tune, utils
+- `codes/surrogates/` - Neural network surrogate models (NN, DeepONet, NeuralODE, Polynomial)
+- `codes/benchmark/` - Core benchmarking logic and evaluation metrics
+- `codes/train/` - Parallel/sequential training infrastructure with task queues
+- `codes/utils/` - Data handling, config management, progress bars, seeding
+- `datasets/` - Training data in HDF5 format, `trained/` - Model checkpoints, `results/` - Benchmark outputs
+- Uses PyTorch, torchode for ODE solving, Optuna for hyperparameter tuning, PostgreSQL for study storage
+
+## Code Style & Conventions
+- Python 3.10+, type hints required (e.g., `str | None`, `int | None`)
+- Use `__all__` in `__init__.py` files for explicit exports
+- Function docstrings with Args/Returns sections
+- Classes use PascalCase, functions/variables use snake_case
+- Threading with explicit locks, use DummyLock() for single-threaded contexts
+- Config-driven architecture with YAML files, use `read_yaml_config()` utility

codes/benchmark/bench_fcts.py

@@ -1019,37 +1019,6 @@ def load_losses(model_identifier: str):
     )
 
 
-# def compare_MAE(metrics: dict, config: dict) -> None:
-#     """
-#     Compare the MAE of different surrogate models over the course of training.
-
-#     Args:
-#         metrics (dict): dictionary containing the benchmark metrics for each surrogate model.
-#         config (dict): Configuration dictionary.
-
-#     Returns:
-#         None
-#     """
-#     MAE = []
-#     labels = []
-#     train_durations = []
-#     device = config["devices"]
-#     device = device[0] if isinstance(device, list) else device
-
-#     for surr_name, _ in metrics.items():
-#         training_id = config["training_id"]
-#         surrogate_class = get_surrogate(surr_name)
-#         n_timesteps = metrics[surr_name]["timesteps"].shape[0]
-#         n_quantities = metrics[surr_name]["accuracy"]["absolute_errors"].shape[2]
-#         model_config = get_model_config(surr_name, config)
-#         model = surrogate_class(device, n_quantities, n_timesteps, model_config)
-#         model_identifier = f"{surr_name.lower()}_main"
-#         model.load(training_id, surr_name, model_identifier=model_identifier)
-#         MAE.append(model.MAE)
-#         labels.append(surr_name)
-#         train_durations.append(model.train_duration)
-
-
 def compare_relative_errors(metrics: dict[str, dict], config: dict) -> None:
     """
     Compare the relative errors over time for different surrogate models.

codes/surrogates/DeepONet/deeponet.py

@@ -359,11 +359,15 @@ def create_dataloader(
         shuffle: bool,
         dataset_params: np.ndarray | None,
         params_in_branch: bool,
-        num_workers: int = 0,  # will usually stay 0 here
+        num_workers: int = 0,
         pin_memory: bool = True,
     ):
         n_samples, n_timesteps, n_quantities = data.shape
 
+        if pin_memory:
+            if "cuda" not in self.device:
+                pin_memory = False
+
         branch = np.repeat(data[:, 0, :], n_timesteps, axis=0)  # (total, n_q)
         trunk = np.tile(timesteps.reshape(1, -1), (n_samples, 1)).reshape(
             -1, 1
@@ -419,21 +423,23 @@ def prepare_data(
             dataset_train,
             timesteps,
             batch_size,
-            True,
+            shuffle=shuffle,
             dataset_params=dataset_train_params,
             params_in_branch=self.config.params_branch,
             num_workers=nw,
         )
 
-        test_loader = self.create_dataloader(
-            dataset_test,
-            timesteps,
-            batch_size,
-            False,
-            dataset_params=dataset_test_params,
-            params_in_branch=self.config.params_branch,
-            num_workers=nw,
-        )
+        test_loader = None
+        if dataset_test is not None:
+            test_loader = self.create_dataloader(
+                dataset_test,
+                timesteps,
+                batch_size,
+                False,
+                dataset_params=dataset_test_params,
+                params_in_branch=self.config.params_branch,
+                num_workers=nw,
+            )
 
         val_loader = None
         if dataset_val is not None:

codes/surrogates/FCNN/fcnn.py

@@ -310,6 +310,10 @@ def create_dataloader(
         n_samples, n_timesteps, n_quantities = data.shape
         total = n_samples * n_timesteps
 
+        if pin_memory:
+            if "cuda" not in self.device:
+                pin_memory = False
+
         init_states = data[:, 0, :]  # (n_samples, n_quantities)
         rep_init = np.repeat(
             init_states[:, None, :], n_timesteps, 1

codes/surrogates/LatentNeuralODE/latent_neural_ode.py

@@ -132,6 +132,10 @@ def create_dataloader(
         num_workers: int = 0,
         pin_memory: bool = True,
     ):
+        if pin_memory:
+            if "cuda" not in self.device:
+                pin_memory = False
+
         data_t = torch.from_numpy(data).float()
         t_t = torch.from_numpy(timesteps).float()
         if dataset_params is not None:

codes/surrogates/LatentPolynomial/latent_poly.py

@@ -92,6 +92,10 @@ def create_dataloader(
         num_workers: int = 0,
         pin_memory: bool = True,
     ):
+        if pin_memory:
+            if "cuda" not in self.device:
+                pin_memory = False
+
         data_t = torch.from_numpy(data).float()
         t_t = torch.from_numpy(timesteps).float()
         params_t = (

codes/train/__init__.py

@@ -4,6 +4,7 @@
     train_and_save_model,
     create_task_list_for_surrogate,
     worker,
+    DummyLock,
 )
 
 __all__ = [
@@ -12,4 +13,5 @@
     "train_and_save_model",
     "create_task_list_for_surrogate",
     "worker",
+    "DummyLock",
 ]

codes/tune/__init__.py

@@ -11,7 +11,15 @@
     maybe_set_runtime_threshold,
     training_run,
 )
-from .postgres_fcts import _make_db_url, initialize_optuna_database
+from .postgres_fcts import (
+    _make_db_url,
+    initialize_optuna_database,
+    _check_postgres_running_local,
+    _start_postgres_server_local,
+    _check_remote_reachable,
+    _initialize_postgres_local,
+    _initialize_postgres_remote,
+)
 from .tune_utils import (
     build_study_names,
     copy_config,
@@ -37,4 +45,9 @@
     "yes_no",
     "_make_db_url",
     "initialize_optuna_database",
+    "_check_postgres_running_local",
+    "_start_postgres_server_local",
+    "_check_remote_reachable",
+    "_initialize_postgres_local",
+    "_initialize_postgres_remote",
 ]