Fitting line to a staircase function

Question

I have a staircase/step function $n(E)$. I know the points $\{E_i\}$ at which each "step" occurs and all steps are of constant height 1. I need to fit a line $a + bE$ to this function and find the least-squares deviation. In particular, I have to calculate the quantity

$$ \Delta = \min_{a,b} \int_{E_i}^{E_f} [n(E) - a - bE]^2 dE $$

In Python/NumPy, I could try and recreate the function $n(E)$ with np.heaviside and then try to fit a line to it, but that feels inefficient. Is there a better way to fit a line to a staircase function?

One approach might be to break up the integral into parts between each "step" and then optimize the resultant expression wrt $a,b$. But I'd like to know if there is a cleaner, more efficient way to fit a line to a step function numerically.

Answer 1

Edited to correct assumption on step width

If your fitting function is simply a linear relation like $a + bE$ then it's very similar to a simple linear regression (least squares) with weights. For each element (of width $w$ and centred at 0), it is simple to show that the terms in $E$ vanish so: $$\int_{-w/2}^{w/2} \left[n(E) - a - bE\right]^2 dE = w\left(n(0) - a\right)^2 + b^2\int_{-w/2}^{w/2}E^2dE$$

The first term is the least squares contribution and the second term (the one proportional to $b^2$) is a constant for all the terms. It will act as a penalty on $b^2$. As you add more and more step functions, it will become largely insignificant, especially if the standard deviation of your data is large.

If you can ignore this (small) contribution then maybe just using the least squares can prove simpler to solve your problem as it's readily implemented in many languages.

For more information on weighted least squares: https://en.wikipedia.org/wiki/Weighted_least_squares

Answer 2

You could use the floor function

$$n(E) = \lfloor a + b E\rfloor\, .$$

Following is an example with $a=5$ and $b=3$.

import numpy as np
import matplotlib.pyplot as plt
from scipy.optimize import curve_fit


def fun(x, a, b):
    return a + b*x
    
xdata = np.linspace(-5, 5, 2001, endpoint=False)
ydata = np.floor(5 + 3*xdata)

popt, pcov = curve_fit(fun, xdata, ydata)
print(popt)
plt.plot(xdata, ydata, lw=0, marker=".")
plt.plot(xdata, fun(xdata, *popt))
plt.show()

With answer

[4.50074129 2.99666613]

And the following plot.

Answer 3

I think that you can also solve the problem analytically.

We can write the problem as

$$\min_{a, b} f$$

with

$$f = \sum_{i = 1}^N \int\limits_{E_i}^{E_{i + 1}}(n_i - a - bE)^2 dE \, .$$

And the minimum should happen when $\nabla f = 0$ or (I would double check)

\begin{align} &\sum_{i = 1}^N \int\limits_{E_i}^{E_{i + 1}}(n_i - a - bE) dE = 0\, ,\\ &\sum_{i = 1}^N \int\limits_{E_i}^{E_{i + 1}}(n_i - a - bE) E dE = 0\, . \end{align}

When you solve the system of equations you end up with (I would also double check this)

\begin{align} a = \frac{\gamma_{1} \lambda_{3} - \gamma_{2} \lambda_{2}}{\lambda_{1} \lambda_{3} - \lambda_{2}^{2}}\, ,\\ b = \frac{- \gamma_{1} \lambda_{2} + \gamma_{2} \lambda_{1}}{\lambda_{1} \lambda_{3} - \lambda_{2}^{2}}\, , \end{align}

where

\begin{align} &\lambda_1 = \sum_{i=1}^N (E_{i + 1} - E_i)\, , &\lambda_2 = \sum_{i=1}^N \frac{(E_{i + 1}^2 - E_i^2)}{2}\, ,\\ &\lambda_3 = \sum_{i=1}^N \frac{(E_{i + 1}^3 - E_i^3)}{3}\, , &\gamma_1 = \sum_{i=1}^N n_i (E_{i + 1} - E_i)\, ,\\ &\gamma_2 = \sum_{i=1}^N \frac{n_i (E_{i + 1}^2 - E_i^2)}{2}\, . \end{align}