Parallel Optimisation¶

import multiprocessing
import time

import diffid
import matplotlib.pyplot as plt
import numpy as np
from scipy.integrate import solve_ivp

# Detect available cores
n_cores = multiprocessing.cpu_count()
print(f"Available CPU cores: {n_cores}")

# Sequential execution (parallel=False)
print("Running CMA-ES (sequential)...")

start = time.time()

result_seq = (
    diffid.DiffsolBuilder()
    .with_diffsl(model_str)
    .with_data(data)
    .with_parameter("alpha", 0.8)
    .with_parameter("beta", 0.3)
    .with_parameter("delta", 0.05)
    .with_parameter("gamma", 0.3)
    .with_cost(diffid.SSE())
    .with_parallel(False)  # Sequential evaluation
    .with_optimiser(diffid.CMAES().with_max_iter(100).with_step_size(0.3))
    .build()
    .optimise()
)

time_seq = time.time() - start

print("\n" + "=" * 60)
print("CMA-ES (Sequential)")
print("=" * 60)
print(f"Solution:      {result_seq.x}")
print(f"True params:   {list(true_params.values())}")
print(f"Final SSE:     {result_seq.value:.3f}")
print(f"Evaluations:   {result_seq.evaluations}")
print(f"Time:          {time_seq:.2f}s")
print(f"Time per eval: {time_seq / result_seq.evaluations * 1000:.2f}ms")

# Parallel execution (parallel=True)
print("Running CMA-ES (parallel)...")

start = time.time()

result_par = (
    diffid.DiffsolBuilder()
    .with_diffsl(model_str)
    .with_data(data)
    .with_parameter("alpha", 0.8)
    .with_parameter("beta", 0.3)
    .with_parameter("delta", 0.05)
    .with_parameter("gamma", 0.3)
    .with_cost(diffid.SSE())
    .with_parallel(True)  # Parallel evaluation!
    .with_optimiser(diffid.CMAES().with_max_iter(100).with_step_size(0.3))
    .build()
    .optimise()
)

time_par = time.time() - start

print("\n" + "=" * 60)
print("CMA-ES (Parallel)")
print("=" * 60)
print(f"Solution:      {result_par.x}")
print(f"True params:   {list(true_params.values())}")
print(f"Final SSE:     {result_par.value:.3f}")
print(f"Evaluations:   {result_par.evaluations}")
print(f"Time:          {time_par:.2f}s")
print(f"Time per eval: {time_par / result_par.evaluations * 1000:.2f}ms")

speedup = time_seq / time_par
print(f"\n🚀 Speedup: {speedup:.2f}x (with {n_cores} cores)")

# Parallel with larger population (more parallel work per generation)
print("Running CMA-ES (parallel, large population)...")

start = time.time()

result_par_large = (
    diffid.DiffsolBuilder()
    .with_diffsl(model_str)
    .with_data(data)
    .with_parameter("alpha", 0.8)
    .with_parameter("beta", 0.3)
    .with_parameter("delta", 0.05)
    .with_parameter("gamma", 0.3)
    .with_cost(diffid.SSE())
    .with_parallel(True)
    .with_optimiser(
        diffid.CMAES()
        .with_max_iter(50)
        .with_step_size(0.3)
        .with_population_size(2 * n_cores)  # Match population to available cores
    )
    .build()
    .optimise()
)

time_par_large = time.time() - start

print("\n" + "=" * 60)
print(f"CMA-ES (Parallel, Population={2 * n_cores})")
print("=" * 60)
print(f"Solution:      {result_par_large.x}")
print(f"True params:   {list(true_params.values())}")
print(f"Final SSE:     {result_par_large.value:.3f}")
print(f"Evaluations:   {result_par_large.evaluations}")
print(f"Time:          {time_par_large:.2f}s")
print(f"Time per eval: {time_par_large / result_par_large.evaluations * 1000:.2f}ms")

speedup_large = time_seq / time_par_large
print(f"\n🚀 Speedup: {speedup_large:.2f}x")

# Define an expensive objective function
def expensive_objective(params):
    """
    An expensive objective function that simulates computation.
    In practice, this could be a complex simulation, ML model, etc.
    """
    alpha, beta, delta, gamma = params

    # Simulate expensive computation (ODE integration with scipy)
    sol = solve_ivp(
        lotka_volterra,
        [0, 100],
        [10.0, 5.0],
        args=(alpha, beta, delta, gamma),
        t_eval=t_data,
        method="RK45",
    )

    if not sol.success:
        return 1e10  # Return large value for failed integrations

    # Compute SSE
    y_pred = sol.y.T
    sse = np.sum((y_pred - y_observed) ** 2)
    return sse


# Test single evaluation time
n_evals = 20
test_params = [[0.8 + i * 0.02, 0.3, 0.05, 0.3] for i in range(n_evals)]

print(f"Testing {n_evals} sequential evaluations...")

start = time.time()
seq_results = [expensive_objective(p) for p in test_params]
time_seq_python = time.time() - start
print(
    f"Sequential: {time_seq_python:.2f}s ({time_seq_python / n_evals * 1000:.1f}ms/eval)"
)

# Note: Multiprocessing in notebooks has pickling limitations
# In a regular Python script, you would use:
print("""
💡 To parallelise in a script, use multiprocessing:

from concurrent.futures import ProcessPoolExecutor

def evaluate_batch_parallel(params_list, n_workers=8):
    with ProcessPoolExecutor(max_workers=n_workers) as executor:
        return list(executor.map(expensive_objective, params_list))

# This provides near-linear speedup for expensive functions
""")

# Test DiffsolBuilder scaling with data size
data_sizes = [100, 200, 500]
scaling_results = []

print("Testing DiffsolBuilder scaling with data size...\n")

for n_points in data_sizes:
    # Generate data with different sizes
    t_test = np.linspace(0, 100, n_points)
    sol_test = solve_ivp(
        lotka_volterra,
        [t_test[0], t_test[-1]],
        [10.0, 5.0],
        args=(
            true_params["alpha"],
            true_params["beta"],
            true_params["delta"],
            true_params["gamma"],
        ),
        t_eval=t_test,
        method="RK45",
    )
    y_test = sol_test.y.T + np.random.normal(0, 0.3, (n_points, 2))
    data_test = np.column_stack((t_test, y_test))

    # Sequential
    start = time.time()
    _ = (
        diffid.DiffsolBuilder()
        .with_diffsl(model_str)
        .with_data(data_test)
        .with_parameter("alpha", 0.8)
        .with_parameter("beta", 0.3)
        .with_parameter("delta", 0.05)
        .with_parameter("gamma", 0.3)
        .with_cost(diffid.SSE())
        .with_parallel(False)
        .with_optimiser(diffid.CMAES().with_max_iter(30).with_step_size(0.3))
        .build()
        .optimise()
    )
    t_seq_scale = time.time() - start

    # Parallel
    start = time.time()
    _ = (
        diffid.DiffsolBuilder()
        .with_diffsl(model_str)
        .with_data(data_test)
        .with_parameter("alpha", 0.8)
        .with_parameter("beta", 0.3)
        .with_parameter("delta", 0.05)
        .with_parameter("gamma", 0.3)
        .with_cost(diffid.SSE())
        .with_parallel(True)
        .with_optimiser(diffid.CMAES().with_max_iter(30).with_step_size(0.3))
        .build()
        .optimise()
    )
    t_par_scale = time.time() - start

    sp = t_seq_scale / t_par_scale
    scaling_results.append(
        {
            "n_points": n_points,
            "time_seq": t_seq_scale,
            "time_par": t_par_scale,
            "speedup": sp,
        }
    )
    print(
        f"N={n_points:3d}: Sequential={t_seq_scale:.2f}s, Parallel={t_par_scale:.2f}s, Speedup={sp:.2f}x"
    )