OCFE Parameter Estimator

A Python class for Ordinary Differential Equation (ODE) parameter estimation using Orthogonal Collocation on Finite Elements (OCFE).

This tool transforms dynamic optimization problems into nonlinear programming problems by discretizing state trajectories using Lagrange polynomials. Using scipy.optimize one can simultaneously estimate model parameters while minimizing the difference between model predictions and observed data, which makes an iterative ODE integration obsolete.

Features ✨

• 😶‍🌫️ Missing Data Tolerant: Safely handles missing observations (None or np.nan) in both the objective functions and visualizations without crashing.
• ↔️ Flexible Collocation Schemes: Supports both Gauss-Legendre (gauss) and Radau Right-endpoint (radau) collocation schemes with configurable polynomial degrees and finite element grids.
• 🛡️ Robust Scaling: Built-in standard and min-max scaling for both states and parameters to improve the numerical stability of the solver.
• 🧮 Multiple Solvers: Seamlessly switch between local least-squares (trf, lm) and global stochastic optimizers (Differential Evolution, Dual Annealing, Basinhopping).
• 📊 Built-in Diagnostics: Includes automated methods for plotting parity charts, raw observations, and final model predictions against data.
• 📈 Forward Prediction: Easily integrate the estimated ODE forward in time using SciPy's solve_ivp.

Requirements

• numpy
• scipy
• matplotlib

Quick Start

A minimal working example demonstrating how to estimate the decay constant of a simple first-order reaction.

import numpy as np
import matplotlib.pyplot as plt
from estimator import OCFEParameterEstimator

# Define the ODE: dz/dt = -k * z
def decay_ode(t, z, k):
    return [-k * z[0]]

# Generate some synthetic noisy data
t_data = np.linspace(0, 5, 10)
y_true = np.exp(-1.5 * t_data).reshape(-1, 1) # True k = 1.5
y_obs = y_true + np.random.normal(0, 0.05, y_true.shape)

# Introduce a missing value to demonstrate robustness
y_obs[4] = np.nan 

# Initialize the Estimator
estimator = OCFEParameterEstimator(degree=3, nfe=10, scheme='gauss')

# Register the ODE and Data
estimator.add_ode(func=decay_ode, n_states=1, param_names=['k'])
estimator.add_data(t=t_data, y=y_obs, y_true=y_true)

# Plot observations
estimator.plot_observations(save=True)

# Solve for parameters (Initial guess k=1.0)
result = estimator.solve(
    theta0=[1.0],
    z0=[1.0], # Initial state
    bounds={'theta': ([0.1], [5.0])}
)

Console output:

   Iteration     Total nfev        Cost      Cost reduction    Step norm     Optimality
       0              1         8.1015e-01                                    3.99e+12
       1              2         2.6480e-02      7.84e-01       1.45e+05       3.33e+10
       2              3         2.2463e-02      4.02e-03       2.89e+05       1.86e+10
       3              4         2.2279e-02      1.85e-04       5.78e+05       1.30e+08
       4              5         2.2278e-02      8.62e-07       1.15e+06       2.00e+07
       5              6         2.2278e-02      1.24e-08       2.31e+06       1.36e+06
       6              7         2.2278e-02      3.11e-10       4.63e+06       3.94e+05
       7              8         2.2278e-02      6.09e-12       9.25e+06       1.97e+06
       8              9         2.2278e-02      3.86e-12       9.25e+06       2.55e+05
       9             12         2.2278e-02      7.63e-14       1.16e+06       4.91e+04
`ftol` termination condition is satisfied.
Function evaluations 12, initial cost 8.1015e-01, final cost 2.2278e-02, first-order optimality 4.91e+04.

# Predict and Plot
estimator.predict()
estimator.plot_prediction(show_legend=True)
plt.show()

Medium article

A more in-depth article was published on Medium, which goes into a bit more detail of the mathematics in the background.

LLM support

This framework was just a proof of concept, and it is not perfect. LLM support was used for the coding, so take it with a grain of salt. ✨🤓