How to Fit a Step Function in Python

Step functions are methods with graphs that look like a series of steps. They consist of a series of horizontal line segments with intervals in between and can also be referred to as staircase functions.

At any given interval, step functions have a constant value, creating a horizontal line on the graph. The intervals make the jumps in between each line segment.

Step functions are helpful when generating discrete plots and are widely used in vectorized plotting in Python. They can be implemented in Python using numpy.

A simple demonstration of step functions is given in this article.

Fit a Step Function in Python

For this example will be using Python version 3.10.6. Furthermore, we need to install the required libraries, which, in our case, is numpy.

Numpy can be installed by running the following command.

pip install numpy

Now we need scipy to optimize and fit the data to the graphs. It can be installed using the command below.

pip install scipy

We will generate a simple step function using a dataset for this example. Starting, import numpy and scipy to the environment using the following statements:

import numpy as np
import scipy

We will use the numpy linspace method for this example to generate a small dataset. The following code snippet can generate this dataset:

x = np.linspace(0, 10, 101)

We need to generate a heaviside function to showcase the staircase plotting. Depending on the use case requirement, it can be either generated using numpy or a custom method.

To create the function, we will use the numpy method heaviside for this example.

y = np.heaviside((x - 5), 0.0)

This method generates the heaviside value according to the following plot.

                      0   if x1 < 0
heaviside(x1, x2) =  x2   if x1 == 0
                      1   if x1 > 0

We will use curve_fit from the scipy library to generate the args with optimal data points.

curve_fit(sigmoid, x, y)

Now that we have cleared up and understood the flow, the final script will look something like this:

import numpy as np
from scipy.special import expit
from scipy.optimize import curve_fit

x = np.linspace(0, 10, 101)
y = np.heaviside((x - 5), 0.0)


def sigmoid(x, x0, b):
    return expit((x - x0) * b)


args, cov = curve_fit(sigmoid, x, y)
print(args)

We can monitor the results of this plotting using matplotlib. Adding the plotting snippet and the final code will be as below.

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

x = np.linspace(0, 10, 101)
y = np.heaviside((x - 5), 0.0)


def sigmoid(x, x0, b):
    return expit((x - x0) * b)


args, cov = curve_fit(sigmoid, x, y)
plt.scatter(x, y)
plt.plot(x, sigmoid(x, *args))
plt.show()
print(args)

Output: