Chi-squared fitting: from least squares to Bayesian inference

import numpy as np
np.random.seed(1234)  # For reproducibility

class Measurement:
    """
    Represents a measurement with local position, global reference, and uncertainty.

    This class encapsulates data points used in fitting algorithms, including:
    - The observed measurement value (local_pos)
    - The independent variable or reference position (global_pos) 
    - The uncertainty/covariance of the measurement (local_cov)

    This structure is particularly useful for weighted fitting methods where
    measurement uncertainties need to be properly propagated.
    """

    def __init__(self, local_pos: np.ndarray, global_pos: np.ndarray, local_cov: np.ndarray):
        """ 
        Initialize a measurement object.

        Parameters:
        -----------
        local_pos : np.ndarray
            Local position of the measurement (observed y-value)
        global_pos : np.ndarray
            Global reference position (independent variable, e.g., x-coordinate)
        local_cov : np.ndarray
            Local covariance matrix representing measurement uncertainty
            For 1D measurements, this is typically a 1x1 matrix containing σ²

        Raises:
        -------
        TypeError
            If local_pos, global_pos, or local_cov are not numpy arrays.

        Example:
        --------
        >>> # Create a measurement: y = 5.2 ± 0.1 at x = 3.0
        >>> measurement = Measurement(np.array([5.2]), np.array([3.0]), np.array([[0.01]]))
        """
        if not isinstance(local_pos, np.ndarray):
            raise TypeError("local_pos must be a numpy array.")
        if not isinstance(global_pos, np.ndarray):
            raise TypeError("global_pos must be a numpy array.")
        if not isinstance(local_cov, np.ndarray):
            raise TypeError("local_cov must be a numpy array.")
        self.local_pos = local_pos    # Measured value (dependent variable)
        self.global_pos = global_pos  # Reference position (independent variable)
        self.local_cov = local_cov    # Measurement uncertainty (covariance)

    def __repr__(self):
        """String representation with limited precision for readability."""
        formatter = {'float_kind': lambda x: "%.2f" % x}
        return (f"Measurement("
                f"local_pos={np.array2string(self.local_pos, formatter=formatter)}, "
                f"global_pos={np.array2string(self.global_pos, formatter=formatter)}, "
                f"local_cov={np.array2string(self.local_cov, formatter=formatter)})")


def generate_data(true_coeffs: np.ndarray, num_points: int, noise_level: float) -> tuple:
    """
    Generate synthetic polynomial data with Gaussian noise for testing fitting algorithms.

    This function creates realistic test data by:
    1. Generating evenly spaced x-coordinates
    2. Evaluating a polynomial at these points
    3. Adding Gaussian noise to simulate measurement errors
    4. Packaging the data into Measurement objects with proper uncertainties

    Parameters:
    -----------
    true_coeffs : np.ndarray
        True polynomial coefficients (highest degree first)
        For a linear fit: [slope, intercept]
        For a quadratic: [a, b, c] representing ax² + bx + c
    num_points : int
        Number of data points to generate
    noise_level : float
        Standard deviation of Gaussian noise to add to measurements

    Returns:
    --------
    tuple
        (measurements, y_true) where:
        - measurements: list of Measurement objects with noise
        - y_true: numpy array of noiseless polynomial values

    Example:
    --------
    >>> # Generate noisy linear data: y = 2x + 3
    >>> measurements, y_true = generate_data([2, 3], 10, 0.5)
    """
    # Generate evenly spaced x values over a reasonable range
    x_data = np.linspace(1, 10, num_points)

    # Calculate true polynomial values (no noise)
    # np.polyval expects coefficients in descending order of powers
    y_data = np.polyval(true_coeffs, x_data)

    # Add Gaussian noise to create realistic measurements
    # This simulates real-world measurement errors
    y_data_noisy = y_data + np.random.normal(0, noise_level, size=y_data.shape)

    # Create Measurement objects with uncertainty information
    # Each measurement knows its uncertainty (noise_level²)
    measurements = [
        Measurement(np.array([y_noisy]), np.array([x]), np.eye(1) * (noise_level**2)) 
        for x, y_noisy in zip(x_data, y_data_noisy)
    ]

    return measurements, y_data

# Known, externally specified uncertainties: they do not depend on observed y.
true_coeffs = np.array([3., 2.])  # np.polyval uses [slope, intercept].
x_data = np.linspace(1, 10, 30)
sigma_data = np.linspace(0.5, 5.0, len(x_data))
y_true = np.polyval(true_coeffs, x_data)
y_observed = y_true + np.random.normal(size=len(x_data)) * sigma_data
measurements = [Measurement(np.array([yy]), np.array([xx]), np.array([[ss**2]]))
                for xx, yy, ss in zip(x_data, y_observed, sigma_data)]
print(f"Generated {len(measurements)} measurements; sigma ranges from "
      f"{sigma_data.min():.1f} to {sigma_data.max():.1f}.")

import plotly.graph_objects as go
import plotly.io as pio
pio.renderers.default = "notebook_connected"  # HTML-compatible interactive output

# ===== LEAST SQUARES POLYNOMIAL FITTING =====
# This section demonstrates classical least squares fitting and compares
# results with the true underlying relationship

# Extract x and y coordinates from measurements for fitting
x = np.array([m.global_pos[0] for m in measurements])
y = np.array([m.local_pos[0] for m in measurements])

print("Data extracted from measurements:")
print(f"x-coordinates: {x}")
# print y positions with limited precision for readability
formatter = {'float_kind': lambda x: "%.2f" % x}
print(f"y-coordinates: {np.array2string(y, formatter=formatter)}")
print()

# Perform least squares polynomial fit (degree 1: linear)
# Known heteroscedastic errors require the sandwich covariance below.
design_ols = np.column_stack([x, np.ones_like(x)])
coefficients = np.linalg.lstsq(design_ols, y, rcond=None)[0]
ols_map = np.linalg.solve(design_ols.T @ design_ols, design_ols.T)
cov = ols_map @ np.diag(sigma_data**2) @ ols_map.T

# Display fitting results with uncertainties
print("=== LEAST SQUARES FITTING RESULTS ===")
print(f"True coefficients:     slope = {true_coeffs[0]:.2f}, intercept = {true_coeffs[1]:.2f}")
print(f"Fitted coefficients:   slope = {coefficients[0]:.2f} ± {np.sqrt(cov[0,0]):.2f}")
print(f"                       intercept = {coefficients[1]:.2f} ± {np.sqrt(cov[1,1]):.2f}")

# Calculate fitted y-values using estimated coefficients
y_fit = np.polyval(coefficients, x)

# Calculate residuals and a simple fit-quality metric
residuals = y - y_fit
rms_residual = np.sqrt(np.mean(residuals**2))

print("\nFit quality:")
print(f"RMS residual: {rms_residual:.2f}")
print(f"Max residual: {np.max(np.abs(residuals)):.2f}")
print(f"Parameter correlation: {cov[0,1]/np.sqrt(cov[0,0]*cov[1,1]):.3f}")

# Create interactive plot using Plotly
fig = go.Figure()

# Add scatter plot of noisy measurement data
fig.add_trace(go.Scatter(
    x=x, y=y, 
    mode='markers', 
    name='Data Points',
    marker=dict(size=6, color='blue'),
    error_y=dict(type='data', array=sigma_data, visible=True)
))

# Add the least squares fit line (sorted for clean plotting)
sort_idx = np.argsort(x)
x_sorted = x[sort_idx]
y_fit_sorted = y_fit[sort_idx]
fig.add_trace(go.Scatter(
    x=x_sorted, y=y_fit_sorted, 
    mode='lines', 
    name='Least Squares Fit', 
    line=dict(color='red', width=2)
))

# Add the true underlying line for comparison (on sorted x)
y_true_sorted = y_true[sort_idx]
fig.add_trace(go.Scatter(
    x=x_sorted, y=y_true_sorted, 
    mode='lines', 
    name='True Line', 
    line=dict(color='blue', dash='dash', width=2)
))

def sigma_x(x_val, cov_matrix):
    """
    Calculate uncertainty in fitted polynomial at given x-value.
    Design vector for linear model is [x, 1].
    """
    Xv = np.array([x_val, 1.0])
    return float(np.sqrt(Xv @ cov_matrix @ Xv.T))

# Dense, sorted grid for smooth uncertainty bands
x_dense = np.linspace(float(np.min(x_sorted)), float(np.max(x_sorted)), 200)
y_fit_dense = np.polyval(coefficients, x_dense)
std_dense = np.array([sigma_x(xi, cov) for xi in x_dense])

# 1σ filled band
fig.add_trace(go.Scatter(
    x=np.concatenate([x_dense, x_dense[::-1]]),
    y=np.concatenate([y_fit_dense - std_dense, (y_fit_dense + std_dense)[::-1]]),
    fill='toself',
    fillcolor='rgba(255, 0, 0, 0.25)',
    line=dict(color='rgba(255, 0, 0, 0)'),
    name='Uncertainty - 1σ'
))

# 2σ filled band
fig.add_trace(go.Scatter(
    x=np.concatenate([x_dense, x_dense[::-1]]),
    y=np.concatenate([y_fit_dense - 2*std_dense, (y_fit_dense + 2*std_dense)[::-1]]),
    fill='toself',
    fillcolor='rgba(255, 0, 0, 0.12)',
    line=dict(color='rgba(255, 0, 0, 0)'),
    name='Uncertainty - 2σ'
))

# Add residual lines showing difference between data and fit
for xi, yi, yfi in zip(x, y, y_fit):
    fig.add_trace(go.Scatter(
        x=[xi, xi],
        y=[yi, yfi],
        mode='lines',
        line=dict(color='gray', dash='dot', width=1),
        showlegend=False
    ))

# Update plot layout
fig.update_layout(
    title='Least squares: fit and uncertainty',
    xaxis_title='x',
    yaxis_title='y',
    legend=dict(x=0.01, y=0.99),
    height=500,
    showlegend=True
)

fig.show()

def weighted_least_squares(measurements: list, degree: int) -> tuple:
    """
    Perform weighted least squares fitting for polynomial regression.

    Minimizes χ² = (y - Xθ)ᵀ W (y - Xθ), with W = Σ⁻¹.
    Returns parameter estimates, their covariance, and fit-quality metrics.
    """
    n = len(measurements)
    print("\n=== WEIGHTED LEAST SQUARES FITTING ===")
    print(f"Processing {n} measurements with individual uncertainties")

    # Data vectors and covariance
    global_pos = np.zeros((n, 1))
    local_pos = np.zeros((n, 1))
    V = np.zeros((n, n))

    for i, measurement in enumerate(measurements):
        local_pos[i, 0] = measurement.local_pos[0]
        global_pos[i, 0] = measurement.global_pos[0]
        V[i, i] = measurement.local_cov[0, 0]

    sig_min, sig_max = np.sqrt(np.min(np.diag(V))), np.sqrt(np.max(np.diag(V)))
    print(f"Measurement uncertainties range: σ = {sig_min:.2f} to {sig_max:.2f}")

    # Design matrix for polynomial y = a + bx + cx² + ...
    X = np.zeros((n, degree + 1))
    X[:, 0] = 1
    for i in range(1, degree + 1):
        X[:, i] = global_pos.reshape(-1) ** i

    # Simple shape sanity checks
    assert X.shape == (n, degree + 1)
    assert local_pos.shape == (n, 1)
    assert V.shape == (n, n)

    print(f"Design matrix shape: {X.shape}")
    print(f"Fitting polynomial degree {degree}: y = a + b x + ...")

    # Weight matrix and normal equations (robust via pseudo-inverse)
    W = np.linalg.solve(V, np.eye(n))
    A = X.T @ W @ X
    CovB = np.linalg.solve(A, np.eye(degree + 1))
    B = CovB @ X.T @ W @ local_pos

    # Fit quality
    y_pred = X @ B
    residuals = local_pos - y_pred
    chi2 = (residuals.T @ W @ residuals).item()
    dof = n - (degree + 1)
    red_chi2 = chi2 / max(dof, 1)

    print("\nFit Results:")
    print(f"χ² = {chi2:.2f}, DOF = {dof}, χ²/DOF = {red_chi2:.2f}")

    return B, CovB, chi2, dof, red_chi2

# Apply weighted least squares to our data
B, CovB, chi2_wls, dof_wls, rchi2_wls = weighted_least_squares(measurements, 1)

# Display results
print("\nParameter Estimates:")
print(f"True coefficients:     slope = {true_coeffs[0]:.2f}, intercept = {true_coeffs[1]:.2f}")
print("WLS fitted coefficients:")
print(f"  Constant (a):  {B[0, 0]:.2f} ± {np.sqrt(CovB[0, 0]):.2f}")
print(f"  Linear (b):    {B[1, 0]:.2f} ± {np.sqrt(CovB[1, 1]):.2f}")

# Compare with simple least squares
simple_coeffs = np.polyfit([m.global_pos[0] for m in measurements], 
                           [m.local_pos[0] for m in measurements], 1)
print("\nComparison with Simple Least Squares:")
print(f"Simple LS: slope = {simple_coeffs[0]:.2f}, intercept = {simple_coeffs[1]:.2f}")
print(f"Weighted LS: slope = {B[1,0]:.2f}, intercept = {B[0,0]:.2f}")
print("Note: Differences arise because WLS accounts for measurement uncertainties\n")
# Independently check the fit with a whitened least-squares solve.
X_check = np.column_stack([np.ones_like(x), x])
reference_wls = np.linalg.lstsq(X_check / sigma_data[:, None], y / sigma_data, rcond=None)[0]
np.testing.assert_allclose(B.ravel(), reference_wls, rtol=1e-10)

# Chi2 minimisation fit for the measurement 

# ===== BAYESIAN LINEAR REGRESSION =====
# This section demonstrates Bayesian parameter estimation, which combines
# prior knowledge with measurement data to estimate posterior distributions

def bayesian_linear_regression(measurements: list, prior_coeffs: np.ndarray, prior_cov: np.ndarray) -> tuple:
    """
    Bayesian linear regression with Gaussian prior and likelihood.

    Posterior: θ|y ~ N(θ*, Σ*), where
      Σ* = (XᵀΣ⁻¹X + Σ₀⁻¹)⁻¹
      θ* = Σ* (XᵀΣ⁻¹y + Σ₀⁻¹θ₀)
    """
    n = len(measurements)

    print("\n=== BAYESIAN LINEAR REGRESSION ===")
    print(f"Processing {n} measurements with Bayesian inference")

    # Extract measurement data
    global_pos = np.zeros((n, 1))
    local_pos = np.zeros((n, 1))
    Sigma = np.zeros((n, n))

    for i, measurement in enumerate(measurements):
        local_pos[i, 0] = measurement.local_pos[0]
        global_pos[i, 0] = measurement.global_pos[0]
        Sigma[i, i] = measurement.local_cov[0, 0]

    # Design matrix from prior size
    degree = prior_coeffs.shape[0] - 1
    X = np.zeros((n, degree + 1))
    X[:, 0] = 1
    for i in range(1, degree + 1):
        X[:, i] = global_pos.reshape(-1) ** i

    print(f"Design matrix shape: {X.shape}")

    # Precisions (robust pseudo-inverse)
    Sigma_inv = np.linalg.pinv(Sigma)
    Sigma0_inv = np.linalg.pinv(prior_cov)

    # Posterior
    Sigma_post_inv = X.T @ Sigma_inv @ X + Sigma0_inv
    Sigma_post = np.linalg.pinv(Sigma_post_inv)
    theta_post = Sigma_post @ (X.T @ Sigma_inv @ local_pos + Sigma0_inv @ prior_coeffs)

    return theta_post, Sigma_post

# Apply Bayesian linear regression with uninformative prior
prior_mean = np.zeros((2, 1))      # Neutral prior mean
prior_cov = np.eye(2) * 1e6        # Large prior variance (uninformative)

print("Using uninformative prior (large variance = low confidence)")
coeffs_post, cov_post = bayesian_linear_regression(measurements, prior_mean, prior_cov)

print("\n=== BAYESIAN RESULTS ===")
print("Posterior coefficients:")
print(f"  Constant (a):  {coeffs_post[0, 0]:.2f} ± {np.sqrt(cov_post[0, 0]):.2f}")
print(f"  Linear (b):    {coeffs_post[1, 0]:.2f} ± {np.sqrt(cov_post[1, 1]):.2f}")

# Compare with classical methods
print("\nComparison with other methods:")
print(f"True:        slope = {true_coeffs[0]:.2f}, intercept = {true_coeffs[1]:.2f}")
print(f"Least Squares: slope = {coefficients[0]:.2f}, intercept = {coefficients[1]:.2f}")
print(f"Weighted LS:   slope = {B[1,0]:.2f}, intercept = {B[0,0]:.2f}")
print(f"Bayesian:      slope = {coeffs_post[1,0]:.2f}, intercept = {coeffs_post[0,0]:.2f}")

print("\nNote: With uninformative prior, Bayesian ≈ Maximum Likelihood ≈ Weighted Least Squares")
print("The power of Bayesian methods becomes apparent with informative priors.")
np.testing.assert_allclose(coeffs_post, B, atol=1e-4)
# An informative prior can shift the estimate and reduce posterior uncertainty.
informative_mean = np.array([[2.], [3.]])
informative_cov = np.diag([0.25, 0.04])
mean_informed, cov_informed = bayesian_linear_regression(measurements, informative_mean, informative_cov)
print("Informative-prior posterior [intercept, slope]:", mean_informed.ravel())
print("Posterior standard deviations:", np.sqrt(np.diag(cov_informed)))

# Residual diagnostics: histogram and Q–Q plot
import scipy.stats as stats

resid_std = (y - X_check @ B.ravel()) / sigma_data

fig_resid = go.Figure()

# Histogram of standardized residuals
hist_y = resid_std
fig_resid.add_trace(go.Histogram(x=hist_y, histnorm='probability density', name='Standardized residuals', nbinsx=12, opacity=0.7))

# Overlay standard normal PDF if scipy is available
xs = np.linspace(-3.5, 3.5, 300)
fig_resid.add_trace(go.Scatter(x=xs, y=stats.norm.pdf(xs), mode='lines', name='N(0,1) PDF'))

fig_resid.update_layout(title='Residual histogram (standardized)', xaxis_title='r/σ', yaxis_title='Density')
fig_resid.show()

# Q–Q plot
theo_q, samp_q = stats.probplot(resid_std, dist='norm', fit=False)
# probplot returns (osm, osr) arrays when fit=False
fig_qq = go.Figure()
fig_qq.add_trace(go.Scatter(x=theo_q, y=samp_q, mode='markers', name='Residual quantiles'))
line_min = min(np.min(theo_q), np.min(samp_q))
line_max = max(np.max(theo_q), np.max(samp_q))
fig_qq.add_trace(go.Scatter(x=[line_min, line_max], y=[line_min, line_max], mode='lines', name='y=x'))
fig_qq.update_layout(title='Q–Q plot vs N(0,1)', xaxis_title='Theoretical quantiles', yaxis_title='Sample quantiles')
fig_qq.show()

def chi2_minimize(model_func, theta_init, measurements, max_iter=100, tol=1e-6, lambda0=1e-3):
    """
    Perform chi-square minimization using a Levenberg–Marquardt-style damping.

    Parameters:
    - model_func: returns (f(x; θ), J) for given x and θ
    - theta_init: initial parameters
    - measurements: list of Measurement
    - max_iter: maximum iterations
    - tol: convergence tolerance on parameter update norm
    - lambda0: initial damping factor

    Returns:
    - theta: best-fit parameters
    - cov_matrix: covariance matrix (approximate)
    - chi2: final chi-square value
    """
    theta = np.array(theta_init, dtype=float)
    x = np.array([m.global_pos[0] for m in measurements])
    y = np.array([m.local_pos[0] for m in measurements])
    sigma2 = np.array([m.local_cov[0, 0] for m in measurements])
    W = np.diag(1.0 / np.clip(sigma2, 1e-12, None))

    lam = lambda0
    prev_chi2 = float(np.sum((y - model_func(x, theta)[0]) ** 2 / sigma2))
    converged = False

    for it in range(max_iter):
        f_pred, J = model_func(x, theta)
        r = y - f_pred
        JT_W = J.T @ W
        H = JT_W @ J
        g = JT_W @ r

        # LM damping
        H_damped = H + lam * np.eye(H.shape[0])
        try:
            delta = np.linalg.solve(H_damped, g)
        except np.linalg.LinAlgError:
            # Fall back to pseudo-inverse
            delta = np.linalg.pinv(H_damped) @ g

        theta_candidate = theta + delta
        f_cand, _ = model_func(x, theta_candidate)
        chi2_cand = float(np.sum((y - f_cand) ** 2 / sigma2))

        if chi2_cand <= prev_chi2:
            theta = theta_candidate
            prev_chi2 = chi2_cand
            lam = max(lam / 3, 1e-12)
            if np.linalg.norm(delta) < tol:
                converged = True
                print(f"Converged in {it+1} iterations.")
                break
        else:
            lam = min(lam * 10, 1e12)

    if not converged:
        raise RuntimeError("Fit did not converge; inspect the starting point and model.")

    # Final evaluation
    f_final, J_final = model_func(x, theta)
    chi2 = float(np.sum((y - f_final) ** 2 / sigma2))

    # Covariance estimate (use pseudo-inverse for robustness)
    cov_matrix = np.linalg.pinv(J_final.T @ W @ J_final)

    return theta, cov_matrix, chi2 

def linear_model(x, theta):
    """
    Linear model function: f(x; theta) = theta[0] + theta[1] * x
    Returns predictions and Jacobian.
    """
    f = theta[0] + theta[1] * x
    J = np.column_stack((np.ones_like(x), x))
    return f, J

# Initial guess for parameters: [intercept, slope]
theta_init = [0.0, 1.0]
# Perform chi-square minimization
theta_best, cov_best, chi2_final = chi2_minimize(linear_model, theta_init, measurements)
print("\n=== CHI-SQUARE MINIMIZATION RESULTS ===")
print(f"Best-fit parameters: intercept = {theta_best[0]:.2f}, slope = {theta_best[1]:.2f}")
print(f"Final chi-square value: {chi2_final:.2f}")

# Exponential model example: y = a * exp(b x)

def exp_model(x, theta):
    a, b = theta
    f = a * np.exp(b * x)
    J = np.column_stack((np.exp(b * x), a * x * np.exp(b * x)))
    return f, J

def exp_model_error(x, theta, Sigma):
    a, b = theta
    g = np.array([np.exp(b*x), a*x*np.exp(b*x)])  # gradient
    return np.sqrt(g @ Sigma @ g)

# Generate synthetic data for the exponential model
np.random.seed(42)
x_exp = np.linspace(0, 5, 100)
true_a, true_b = 1.0, 0.4
sigma_exp = 0.1
y_exp_true = true_a * np.exp(true_b * x_exp)
sigma_y_exp = sigma_exp * y_exp_true
y_exp = y_exp_true + np.random.normal(size=x_exp.shape) * sigma_y_exp
meas_exp = [Measurement(np.array([yy]), np.array([xx]), np.array([[ss**2]])) for xx, yy, ss in zip(x_exp, y_exp, sigma_y_exp)]


# Fit with chi2_minimize
theta0 = [1.0, 1.0]
th_exp, cov_exp, chi2_exp = chi2_minimize(exp_model, theta0, meas_exp, max_iter=100, tol=1e-8)
res_exp = np.array([m.local_pos[0] - exp_model(m.global_pos[0], th_exp)[0] for m in meas_exp])
# In-sample residuals: subtract the fitted-mean variance (linear approximation).
_, J_exp = exp_model(x_exp, th_exp)
mean_var_exp = np.einsum('ij,jk,ik->i', J_exp, cov_exp, J_exp)
residual_var_exp = sigma_y_exp**2 - mean_var_exp
assert np.all(residual_var_exp > 0)
pull_exp = res_exp / np.sqrt(residual_var_exp)
from scipy.optimize import curve_fit
reference_exp, reference_cov = curve_fit(lambda xx, a, b: a*np.exp(b*xx),
    x_exp, y_exp, p0=[1., 0.4], sigma=sigma_y_exp, absolute_sigma=True)
np.testing.assert_allclose(th_exp, reference_exp, rtol=1e-5)
np.testing.assert_allclose(cov_exp, reference_cov, rtol=1e-4)

print("\n=== NON-LINEAR EXPONENTIAL FIT ===")
print(f"True:      a={true_a:.2f}, b={true_b:.2f}")
print(f"Estimated: a={th_exp[0]:.2f} ± {np.sqrt(cov_exp[0,0]):.2f}, b={th_exp[1]:.2f} ± {np.sqrt(cov_exp[1,1]):.2f}")
print(f"χ²: {chi2_exp:.2f}")

# Plot
x_grid = np.linspace(0, 5, 200)
y_fit_exp, _ = exp_model(x_grid, th_exp)
fig_exp = go.Figure()
# Add scatter plot of noisy measurement data with error bars
fig_exp.add_trace(go.Scatter(
    x=x_exp, y=y_exp,
    mode='markers',
    name='Data Points',
    marker=dict(size=8, color='blue'),
    error_y=dict(type='data', array= [np.sqrt(m.local_cov[0, 0]) for m in meas_exp], visible=True)
))

# Add uncertainty bands for the fit
std_exp = np.array([ exp_model_error(xi, th_exp, cov_exp) for xi in x_grid ])
fig_exp.add_trace(go.Scatter(
    x=np.concatenate([x_grid, x_grid[::-1]]),
    y=np.concatenate([y_fit_exp - std_exp*2, (y_fit_exp + std_exp*2)[::-1]]),
    fill='toself',
    fillcolor='rgba(255, 0, 0, 0.25)',
    line=dict(color='rgba(255, 0, 0, 0)'),
    name='Uncertainty - 2σ'
))

fig_exp.add_trace(go.Scatter(x=x_grid, y=y_fit_exp, mode='lines', name='Fit'))
fig_exp.add_trace(go.Scatter(x=x_grid, y=true_a*np.exp(true_b*x_grid), mode='lines', name='True', line=dict(dash='dash')))
fig_exp.update_layout(title='Non-linear fit: y = a exp(bx)', xaxis_title='x', yaxis_title='y')
fig_exp.show()

# plot residuals
fig_resid_exp = go.Figure()
fig_resid_exp.add_trace(go.Histogram(x=pull_exp, histnorm='probability density', name='Standardized residuals', nbinsx=14, opacity=0.7))

# Overlay standard normal PDF if scipy is available
xs = np.linspace(-3.5, 3.5, 300)
fig_resid_exp.add_trace(go.Scatter(x=xs, y=stats.norm.pdf(xs), mode='lines', name='N(0,1) PDF'))

fig_resid_exp.update_layout(title='Residual histogram (standardized)', xaxis_title='r/σ', yaxis_title='Density')
fig_resid_exp.show()

# Answer: constant weights cancel out of the normal equations.
equal_weight_fit = np.linalg.lstsq(X_check / 2., y / 2., rcond=None)[0]
ordinary_fit = np.linalg.lstsq(X_check, y, rcond=None)[0]
np.testing.assert_allclose(equal_weight_fit, ordinary_fit)
print("Constant-weight WLS equals OLS:", equal_weight_fit)

from pathlib import Path
import sys
interactive_project = next(p / "projects" for p in [Path.cwd(), *Path.cwd().parents]
                           if (p / "projects" / "src").is_dir())
sys.path.insert(0, str(interactive_project / "src"))
from notebook_interactive import display_interactive
display_interactive("fit")