3.2. Custom Models

Learn how to write your own model functions for curve fitting.

3.2.1. Basic Structure

A model function must:

  1. Take x as the first argument

  2. Take fit parameters as subsequent arguments

  3. Return the predicted y values

  4. Use jax.numpy for mathematical operations

import jax.numpy as jnp


def my_model(x, param1, param2, param3):
    """Model description."""
    return param1 * jnp.exp(-param2 * x) + param3

3.2.2. Example: Damped Oscillation

import jax.numpy as jnp
from nlsq import fit


def damped_oscillation(x, amplitude, decay, frequency, phase, offset):
    """Damped sinusoidal oscillation.

    y = A * exp(-gamma * x) * cos(omega * x + phi) + c
    """
    return amplitude * jnp.exp(-decay * x) * jnp.cos(frequency * x + phase) + offset


# Fit
p0 = [1.0, 0.1, 2.0, 0.0, 0.0]
popt, pcov = fit(damped_oscillation, x, y, p0=p0)

3.2.3. Example: Sum of Gaussians

import jax.numpy as jnp


def double_gaussian(x, a1, c1, w1, a2, c2, w2, offset):
    """Sum of two Gaussian peaks."""
    g1 = a1 * jnp.exp(-0.5 * ((x - c1) / w1) ** 2)
    g2 = a2 * jnp.exp(-0.5 * ((x - c2) / w2) ** 2)
    return g1 + g2 + offset


# 7 parameters: 2 peaks (3 each) + offset
p0 = [1, 2, 0.5, 1, 5, 0.5, 0]
bounds = ([0, 0, 0.1, 0, 0, 0.1, -1], [10, 10, 3, 10, 10, 3, 1])

popt, pcov = fit(double_gaussian, x, y, p0=p0, bounds=bounds)

3.2.4. Example: Michaelis-Menten Kinetics

import jax.numpy as jnp


def michaelis_menten(S, Vmax, Km):
    """Enzyme kinetics: v = Vmax * S / (Km + S)"""
    return Vmax * S / (Km + S)


# Fit enzyme kinetics data
popt, pcov = fit(michaelis_menten, substrate_conc, reaction_rate, p0=[100.0, 10.0])
Vmax, Km = popt

3.2.5. Using Constants

If your model needs fixed constants, use closures:

import jax.numpy as jnp


def create_model(wavelength):
    """Create a model with fixed wavelength."""

    def model(x, amplitude, phase):
        k = 2 * jnp.pi / wavelength
        return amplitude * jnp.sin(k * x + phase)

    return model


# Create model with wavelength = 5.0
wave_model = create_model(wavelength=5.0)
popt, pcov = fit(wave_model, x, y, p0=[1.0, 0.0])

3.2.6. Array Parameters

For models with array-like parameters:

import jax.numpy as jnp


def sum_of_exponentials(x, *params):
    """Sum of N exponential terms.

    params = [A1, k1, A2, k2, ..., offset]
    """
    n_terms = (len(params) - 1) // 2
    result = jnp.zeros_like(x)

    for i in range(n_terms):
        A = params[2 * i]
        k = params[2 * i + 1]
        result = result + A * jnp.exp(-k * x)

    return result + params[-1]  # Add offset


# Two exponential terms + offset = 5 parameters
popt, pcov = fit(sum_of_exponentials, x, y, p0=[1, 0.1, 0.5, 0.5, 0])

3.2.7. Common Patterns

Avoid divisions by zero:

def safe_model(x, a, b):
    # Bad: may divide by zero
    # return a / (x - b)

    # Good: add small epsilon
    return a / (x - b + 1e-10)

Use stable functions:

# Numerically stable log-sum-exp
def log_sum_exp(x, a, b):
    max_val = jnp.maximum(a, b)
    return max_val + jnp.log(jnp.exp(a - max_val) + jnp.exp(b - max_val))

Vectorized operations:

# Good: vectorized
def good_model(x, a, b):
    return a * jnp.exp(-b * x)


# Bad: Python loops (slow)
def bad_model(x, a, b):
    result = []
    for xi in x:
        result.append(a * jnp.exp(-b * xi))
    return jnp.array(result)

3.2.8. Testing Your Model

Before fitting, verify your model works:

import numpy as np
import jax.numpy as jnp


def my_model(x, a, b):
    return a * jnp.exp(-b * x)


# Test with sample data
x_test = np.linspace(0, 10, 50)
y_test = my_model(x_test, 2.0, 0.5)

print(f"Output shape: {y_test.shape}")
print(f"Output type: {type(y_test)}")
print(f"Sample values: {y_test[:5]}")

3.2.9. Next Steps