Update progress bar (#19)

* refactor: remove continuous optimization features

* ci: update python versions to 3.10-3.13

* ci: add notebook execution workflow

* docs: add notebooks workflow badge

* docs: update README to reflect removal of continuous optimization

* Add notebook dependencies and update CI workflows to include main branch

* test: enforce strict MPI requirement in distributed notebook

* ci: install libopenmpi-dev for notebook workflow

* Update notebooks workflow and dependencies

Author: ASKabalan

.github/workflows/notebooks.yml

@@ -0,0 +1,46 @@
+name: Notebooks
+
+on:
+  push:
+    branches: [ "main" ]
+  pull_request:
+    branches: [ "main" ]
+
+jobs:
+  test-notebooks:
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        python-version: ["3.10", "3.11", "3.12", "3.13"]
+    steps:
+    - uses: actions/checkout@v4
+
+    - name: Set up Python ${{ matrix.python-version }}
+      uses: actions/setup-python@v5
+      with:
+        python-version: ${{ matrix.python-version }}
+
+    - name: Install system dependencies
+      run: |
+        sudo apt-get update
+        sudo apt-get install -y libopenmpi-dev
+        python -m pip install --upgrade pip
+        pip install setuptools cython mpi4py
+
+    - name: Install dependencies
+      run: |
+        pip install .[notebooks]
+        pip install nbconvert ipykernel
+
+    - name: Execute notebooks
+      run: |
+        mkdir -p executed_notebooks
+        find examples -name "*.ipynb" -print0 | xargs -0 -I {} jupyter nbconvert --to notebook --execute {} --output-dir executed_notebooks
+
+
+    - name: Upload executed notebooks
+      if: always()
+      uses: actions/upload-artifact@v4
+      with:
+        name: executed-notebooks-${{ matrix.python-version }}
+        path: executed_notebooks/*.ipynb

.github/workflows/tests.yml

@@ -12,7 +12,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: ['3.10', '3.11', '3.12']
+        python-version: ['3.10', '3.11', '3.12', '3.13']
 
     steps:
       - name: Checkout code

README.md

@@ -1,34 +1,32 @@
-# Distributed Grid Search & Continuous Optimization using JAX
+# Distributed Grid Search using JAX
 
-[![Testing](https://github.com/CMBSciPol/jax-grid-search/actions/workflows/tests.yml/badge.svg)](https://github.com/CMBSciPol/jax-grid-search/actions/workflows/tests.yml)
+[![Tests](https://github.com/CMBSciPol/jax-grid-search/actions/workflows/tests.yml/badge.svg)](https://github.com/CMBSciPol/jax-grid-search/actions/workflows/tests.yml)
+[![Notebooks](https://img.shields.io/github/actions/workflow/status/CMBSciPol/jax-grid-search/notebooks.yml?logo=jupyter&label=notebooks)](https://github.com/CMBSciPol/jax-grid-search/actions/workflows/notebooks.yml)
 [![Code Formatting](https://github.com/CMBSciPol/jax-grid-search/actions/workflows/formatting.yml/badge.svg)](https://github.com/CMBSciPol/jax-grid-search/actions/workflows/formatting.yml)
 [![Upload Python Package](https://github.com/CMBSciPol/jax-grid-search/actions/workflows/python-publish.yml/badge.svg)](https://github.com/CMBSciPol/jax-grid-search/actions/workflows/python-publish.yml)
 [![PyPI version](https://badge.fury.io/py/jax-grid-search.svg)](https://badge.fury.io/py/jax-grid-search)
 [![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)
 <a href="https://doi.org/10.5281/zenodo.17674777"><img src="https://zenodo.org/badge/917061582.svg" alt="DOI"></a>
 
-
-
 ## About
 
 This package is designed to minimize likelihoods computed by [FURAX](https://github.com/CMBSciPol/furax), a JAX-based CMB analysis framework. It provides distributed grid search capabilities specifically optimized for:
 
 - **Spatial spectral index variability:** Efficiently explore parameter spaces for spatially-varying spectral indices in foreground models
 - **Foreground component optimization:** Test and compare different foreground component configurations to find the optimal model choice
-- **Likelihood model optimization:** Systematically search through discrete model configurations and continuously optimize their parameters
+- **Likelihood model optimization:** Systematically search through discrete model configurations
 
-The distributed grid search is built to handle the computational demands of CMB likelihood analysis, leveraging JAX's performance and enabling efficient parallel exploration of both discrete and continuous parameter spaces.
+The distributed grid search is built to handle the computational demands of CMB likelihood analysis, leveraging JAX's performance and enabling efficient parallel exploration of discrete parameter spaces.
+
+> **Note:** Continuous optimization features (formerly `optimize`) have been moved to [furax-cs](https://github.com/CMBSciPol/furax-cs). Please use `furax_cs.minimize` for gradient-based optimization.
 
 ---
 
-This repository provides two complementary optimization tools:
+This repository provides:
 
 1. **Distributed Grid Search for Discrete Optimization:**
    Explore a parameter space by evaluating a user-defined objective function on a grid of discrete values. The search runs in parallel across available processes, automatically handling batching, progress tracking, and result aggregation.
 
-2. **Continuous Optimization with Optax:**
-   Minimize continuous functions using gradient-based methods (such as LBFGS). This routine leverages Optax for iterative parameter updates and includes built-in progress monitoring.
-
 ---
 
 ## Getting Started
@@ -47,7 +45,7 @@ pip install jax_grid_search
 
 For comprehensive tutorials and hands-on examples, see the **[examples directory](./examples/)** which contains:
 
-- **5 interactive Jupyter notebooks** covering basic to advanced concepts
+- **Interactive Jupyter notebooks** covering basic to advanced concepts
 - **Distributed computing examples** with MPI setup
 - **Complete API demonstrations** with visualization
 
@@ -57,7 +55,7 @@ For comprehensive tutorials and hands-on examples, see the **[examples directory
 
 ## Usage Examples
 
-### 1. Distributed Grid Search (Discrete Optimization)
+### Distributed Grid Search (Discrete Optimization)
 
 Define your objective function and parameter grid, then run a distributed grid search. The objective function must return a dictionary with a `"value"` key.
 
@@ -192,217 +190,6 @@ def image_filter_objective(kernel, bias_matrix):
 
 See [02-advanced-grid-search.ipynb](./examples/02-advanced-grid-search.ipynb) for complete examples with visualization.
 
-### 2. Continuous Optimization using Optax
-
-Use the continuous optimization routine to minimize a function with gradient-based methods (e.g., LBFGS). The example below minimizes a simple quadratic function.
-
-```python
-import jax.numpy as jnp
-import optax
-from jax_grid_search import optimize
-from jax_progress import TqdmProgressMeter
-
-# Define a continuous objective function (e.g., quadratic)
-def quadratic(x):
-    return jnp.sum((x - 3.0) ** 2)
-
-# Initial parameters and an optimizer (e.g., LBFGS)
-init_params = jnp.array([0.0])
-optimizer = optax.lbfgs()
-
-# Run continuous optimization with progress monitoring (optional)
-best_params, opt_state = optimize(
-    init_params,
-    quadratic,
-    opt=optimizer,
-    max_iter=50,
-    tol=1e-10,
-    progress=TqdmProgressMeter(total=50)
-)
-
-print("Optimized Parameters:", best_params)
-```
-
-#### Using Different Optimizers
-
-The library supports various Optax optimizers beyond LBFGS:
-
-```python
-import optax
-from jax_grid_search import optimize
-from jax_progress import TqdmProgressMeter
-
-def rosenbrock(x):
-    # Classic optimization test function
-    return 100 * (x[1] - x[0]**2)**2 + (1 - x[0])**2
-
-init_params = jnp.array([-1.0, 1.0])
-
-# Try different optimizers
-optimizers = {
-    "LBFGS": optax.lbfgs(),
-    "Adam": optax.adam(learning_rate=0.01),
-    "SGD": optax.sgd(learning_rate=0.1),
-    "RMSprop": optax.rmsprop(learning_rate=0.01)
-}
-
-for name, optimizer in optimizers.items():
-    result, state = optimize(
-        init_params, rosenbrock, optimizer,
-        max_iter=1000, tol=1e-8,
-        progress=TqdmProgressMeter(total=1000)
-    )
-    print(f"{name}: {result}, final value: {rosenbrock(result)}")
-```
-
-#### Parameter Bounds and Constraints
-
-Use box constraints to limit parameter values during optimization:
-
-```python
-from jax_progress import TqdmProgressMeter
-
-# Constrain parameters to [0, 10] range
-lower_bounds = jnp.array([0.0, 0.0])
-upper_bounds = jnp.array([10.0, 10.0])
-
-result, state = optimize(
-    init_params,
-    objective_function,
-    optax.adam(0.1),
-    max_iter=100,
-    tol=1e-6,
-    progress=TqdmProgressMeter(total=100),
-    lower_bound=lower_bounds,
-    upper_bound=upper_bounds
-)
-```
-
-#### Update History and Debugging
-
-Track optimization progress for analysis and debugging:
-
-```python
-from jax_progress import TqdmProgressMeter
-
-result, state = optimize(
-    init_params,
-    objective_function,
-    optax.lbfgs(),
-    max_iter=100,
-    tol=1e-8,
-    progress=TqdmProgressMeter(total=100),
-    log_updates=True  # Enable update history logging
-)
-
-# Plot optimization history
-import matplotlib.pyplot as plt
-if state.update_history is not None:
-    history = state.update_history
-    plt.figure(figsize=(12, 4))
-
-    plt.subplot(1, 2, 1)
-    plt.plot(history[:, 0])
-    plt.ylabel('Update Norm')
-    plt.xlabel('Iteration')
-    plt.yscale('log')
-
-    plt.subplot(1, 2, 2)
-    plt.plot(history[:, 1])
-    plt.ylabel('Objective Value')
-    plt.xlabel('Iteration')
-    plt.show()
-```
-
-#### Running multiple optimization tasks with vmap
-
-You can run multiple optimization tasks in parallel using `jax.vmap`. This is useful when optimizing multiple functions or parameters simultaneously.
-
-(This is very useful for simulating multiple noise realizations for example)
-
-The jax-progress library automatically handles vmap tracking internally.
-
-```python
-import jax
-import jax.numpy as jnp
-import optax
-from jax_grid_search import optimize
-from jax_progress import TqdmProgressMeter
-
-# Define multiple objective functions
-def objective_fn(x, normal):
-    return jnp.sum(((x - 3.0) ** 2) + normal)
-
-progress = TqdmProgressMeter(total=50)
-
-def solve_one(seed):
-    init_params = jnp.array([0.0])
-    normal = jax.random.normal(jax.random.PRNGKey(seed), init_params.shape)
-    optimizer = optax.lbfgs()
-    # Run continuous optimization with progress monitoring (optional)
-    best_params, opt_state = optimize(
-        init_params,
-        objective_fn,
-        opt=optimizer,
-        max_iter=50,
-        tol=1e-4,
-        progress=progress,
-        normal=normal
-    )
-
-    return best_params
-
-jax.vmap(solve_one)(jnp.arange(10))
-```
-
-### 3. Function Conditioning
-
-Improve optimization performance by transforming parameters to similar scales and normalizing outputs. This is essential for problems with parameters in different ranges (e.g., temperature vs spectral index) or large objective values (e.g., chi-square with many pixels).
-
-```python
-import jax.numpy as jnp
-import optax
-from jax_grid_search import condition, optimize
-from jax_progress import TqdmProgressMeter
-
-# Function with parameters in different scales and large output
-npix = 12 * 64**2  # HEALPix pixels
-
-def objective(params):
-    # Simulate chi-square scaled by number of pixels
-    return npix * ((params['temp'] - 20)**2 + (params['beta'] - 1.5)**2)
-
-# Apply conditioning: parameter scaling + output normalization
-lower = {'temp': 10.0, 'beta': 0.5}
-upper = {'temp': 40.0, 'beta': 3.0}
-
-conditioned_fn, to_opt, from_opt = condition(
-    objective,
-    lower=lower,
-    upper=upper,
-    factor=npix  # Normalize by problem size
-)
-
-# Transform parameters to [0,1] space, optimize, then transform back
-init_opt = to_opt({'temp': 15.0, 'beta': 1.0})
-
-result_opt, _ = optimize(
-    init_opt,
-    conditioned_fn,
-    optax.lbfgs(),
-    max_iter=100,
-    tol=1e-6,
-    progress=TqdmProgressMeter(total=100)
-)
-
-result = from_opt(result_opt)  # Back to physical space
-print(f"Optimized: temp={result['temp']:.2f}, beta={result['beta']:.3f}")
-```
-### 4. Optimizing Likelihood parameters and models
-
-You can use the continuous optimization to optimize the parameters of a model that is defined in a function.
-For performance purposes, you need to make sure that the discrete parameters that can control the likelihood model can be jitted (using `lax.cond` for example or other lax control flow functions).
-
 ## Citation
 
 ```