examples.py

# Import requried packages
import numpy as np
from scipy.integrate import solve_ivp
from typing import Callable, List
import matplotlib.pyplot as plt

from ode_collection import (
    monod_growth,
    oregonator
)
from estimator import OCFEParameterEstimator

# Create a helper to derive noisy data from underlying ODE system
def get_noisy_obs_data(
    ode_function: Callable,
    true_theta: np.ndarray, 
    t_obs: np.ndarray,
    y0: np.ndarray,
    noise_percentage: float = 10
) -> np.ndarray:
    # 1. Generate Ground Truth
    sol = solve_ivp(ode_function, [t_obs[0], t_obs[-1]], y0, 
                     t_eval=t_obs, args=tuple(true_theta))
    y_true = sol.y.T 
    
    # 2. Calculate Standard Deviation for each state variable
    # This ensures noise is proportional to the "action" of that specific state
    stds = np.std(y_true, axis=0)
    
    # 3. Add Gaussian Noise
    # noise_percentage/100 * stds acts as the 'sigma' for the normal distribution
    noise = np.random.normal(loc=0, 
                             scale=(noise_percentage / 100) * stds, 
                             size=y_true.shape)
    
    y_obs_noisy = y_true + noise
    return y_true, y_obs_noisy

def run_full_example(
    ode_function: Callable,
    param_names: List,
    t_obs: np.ndarray,
    y0: list,
    true_theta: np.ndarray,
    bounds: dict = None,
    theta0: List = None,
    noise_percentage: float = 10,
    affected_species_none: dict = None,
    # OCFE settings
    degree: int = 3,
    n_finite_elements: int = 5,
    scheme: str = 'gauss',
    verbose: bool = True,
    solveroptions: dict = None,
    # Figure saving
    save: bool = False
):
        
    # Demo data
    np.random.seed(42)
    Y_true, Y = get_noisy_obs_data(
        ode_function, true_theta, t_obs,
        y0=y0, noise_percentage=noise_percentage
    )

    # If apply none is given, delete some values to make it more realistic
    if affected_species_none is not None:
        for i__ in affected_species_none:
            if (
                affected_species_none[i__] is None 
                or affected_species_none[i__] is False
                or affected_species_none[i__] == []
            ):
                continue
            if affected_species_none[i__] == 'secondhalf':
                Y[int(len(t_obs)/2):, i__] = np.nan
            else:
                Y[affected_species_none[i__], i__] = np.nan

    # Get number of states and parameters
    n_params = len(param_names)
    n_time_points = len(t_obs)
    n_states = len(y0)

    # Print summary
    print(f'[+] Number of observed state variables: {n_states}')
    print(f'[+] Number of observed time points: {n_time_points}')
    print(f'[+] Number of estimatable parameters: {n_params}')

    # Build estimator
    est = OCFEParameterEstimator(
        degree=degree, nfe=n_finite_elements, scheme=scheme, verbose=verbose
    )

    # Add the ODE callable
    est.add_ode(
        ode_function, n_states=n_states, param_names=param_names
    )

    # Add the observed data
    est.add_data(t_obs, y=Y, y_true=Y_true, sigma=0.5)

    # Explicitly scale the observed data
    est.apply_scaling(method='standard')  # or 'minmax'

    # Give an initial guess of the estimatable parameters 
    if theta0 is None:
        theta0 = list(np.ones(n_params))

    # Give the initial condition for the ODEs
    # (basically our first observed time point)
    z0 = Y[0]  

    # Plot the observations to see the data
    ax1 = est.plot_observations(save=save)

    # Solve the parameter estimation with the OCFE estimator
    res = est.solve(
        theta0=theta0,
        z0=z0,
        bounds=bounds,
        state_scaling={'mode': 'standard'},
        param_scaling={'mode': 'auto'},
        method='trf',
        options={'max_nfev': 600, 'ftol': 1e-9, 'xtol': 1e-9, 'gtol': 1e-9} or solveroptions
    )

    print("\nSuccess:", res['success'])
    print("Message:", res['message'])
    print("Estimated parameters:")
    for name, val in zip(res['theta_names'], res['theta']):
        print(f"  {name} = {val:.6f}")

    # Plot the parity plot (observed-versus-predicted values)
    ax2 = est.plot_parity(save=save)

    # Plot the time series in physical units
    est.predict()
    ax3 = est.plot_prediction(save=save)  # markers = observations, lines = predictions

    return ax1, ax2, ax3, est

if __name__ == '__main__':

    CASE_STUDY = 'monod'

    # Get CS data
    if CASE_STUDY == 'monod':
        case = monod_growth
        true_theta = [0.5, 0.1, 0.6] # [mu_max, Ks, Yxs]
        z0 = [0.1, 10.0] # Initial [Biomass, Substrate]
        t_obs = np.linspace(0, 10,20)
        bounds={
            'theta': ([0,0,0], [10,10,10]),
            'z': ([0, 0], [50, 50])
        }
        param_names = ['mu_max', 'Ks', 'Yxs']
        affected_species_none = {
            0: [1,2,3,4],
            1: None,
        }

    elif CASE_STUDY == 'oregonator':
        case = oregonator
        true_theta = [0.01, 0.0001, 1.0, 1.0] # [epsilon, q, f, p]
        z0 = [0.5, 0.001, 0.5]
        t_obs = np.linspace(0, 10, 25)
        bounds={'theta': ([0,0,0,0], [1,1,1,1]),
                'z': ([0,0,0], [5,5,5])},
        param_names = ['epsilon', 'q', 'f', 'p']
        affected_species_none = {
            0: [5,6,7,10,11,12],
            1: 'secondhalf',
            2: [2,4,6,8]
        }

    # Run example
    ax1, ax2, ax3, est = run_full_example(
        ode_function = case,
        t_obs = t_obs,
        affected_species_none = affected_species_none,
        param_names = param_names,
        true_theta = true_theta,
        bounds=bounds,
        y0 = z0,
        noise_percentage=10,
        # OCFE settings
        verbose = True,
        degree = 2,
        n_finite_elements = 5,
    )

    plt.show()