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.

  • 1
    $\begingroup$ What prevents you from writing down the actual minimization problem in $a$ and $b$? The integral can be rewritten as a finite sum over the steps with analytically known integrals. Input to the minimization problem are only the steps and step heights as constants. $\endgroup$
    – Bort
    Commented Oct 11, 2021 at 12:19
  • $\begingroup$ If you mean evaluating $\int_{E_n}^{E_{n+1}} [n(E) - a - bE]^2$ analytically and then taking partial derivatives wrt $a$ and $b$, then yes I have done that. But sometimes that analytic expression becomes a complicated mess. I was hoping that there might be a nice way to handle step functions numerically. $\endgroup$
    – adch99
    Commented Oct 14, 2021 at 12:14

3 Answers 3


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

  • $\begingroup$ Maybe I misunderstood your question, is the height or the width of each element a constant 1? My answer is for a constant width. In any case, the answer here can be modified but you need to change the boundaries of the integral accordingly. It would match the "weighted least square" case. $\endgroup$
    – PC1
    Commented Oct 12, 2021 at 19:28
  • $\begingroup$ The step function $n(E)$ has constant height, not constant width in my case i.e. $n(E_{i+1}) - n(E_{i})$ is constant but $E_{i+1} - E_i$ is not constant. $\endgroup$
    – adch99
    Commented Oct 14, 2021 at 12:18

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)
plt.plot(xdata, ydata, lw=0, marker=".")
plt.plot(xdata, fun(xdata, *popt))

With answer

[4.50074129 2.99666613]

And the following plot.

enter image description here


I think that you can also solve the problem analytically.

We can write the problem as

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


$$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}


\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}

  • $\begingroup$ This is indeed what I ended up doing and was hoping to avoid as it leads to a decent amount of complexity while writing the code and doesn't scale very well if in the future I have to move from fitting a line to fitting a polynomial. I was hoping for a numerical method like PC1's answer. The analytical approach is probably the most accurate answer for this problem, but unfortunately, not the one I was looking for. $\endgroup$
    – adch99
    Commented Oct 17, 2021 at 5:59
  • $\begingroup$ @ach99, my other answer was purely numerical. $\endgroup$
    – nicoguaro
    Commented Oct 17, 2021 at 12:03

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.