Search code examples
pythonnumpymatplotlibpiecewise

Condition checking when using numpy.piecewise to generate a piecewise periodic plot


I'm trying to generate a piecewise periodic plot using Numpy and matplotlib, like this:

import numpy as np
import matplotlib.pyplot as plt

Q1 = lambda t, f, Q_max: Q_max * np.sin(2 * np.pi *f * t)
Q2 = lambda t, f, Q_max: 0

def Q_true(t, f, stat):
    while(t >= 1/f):
        t -= 1/f 
    while(t < 0): 
        t += 1/f 
    return ((t <= 1/(2*f)) == stat)

Q = lambda t, f, Q_max: np.piecewise(t, [Q_true(t, f, True) , Q_true(t,f, False)], [Q1, Q2], f, Q_max)

Q_max = 225 # mL/sec
f = 1.25 # Hz
t = np.linspace(0,4,101) # secs
plt.plot(t, Q(t, f, Q_max))

The problem is that Q_true is receiving the entire t array instead of individual points. This isn't an issue if I just use less than/greater than statements in the numpy.piecewise condlist, but it's much easier to determine whether it's true or false using Q_true.

The plot should look something like this:

Example Plot

Any ideas?

Thanks!


Solution

  • The following version of Q_true works:

    def Q_true(t, f, stat):
        period = 1/f
        return (t % period < period/2) == stat
    

    Note that you are naming your anonymous functions (Q1 = lambda ...). In this case you should just define the function using def.