How do I plot this function in Python?

Question

How can I plot this with numpy and matplotlib?

f(t) = t+2, if t is an integer
f(t) = 2**t, otherwise

Say my function is defined just for t > 0.

Is this possible to plot?

The approach I tried does not work... Not sure how to fix it.

    import numpy as np
    import matplotlib.pyplot as plt
    import matplotlib
    from math import floor
    
    # csfont = {'fontname':'Comic Sans MS'}
    # hfont = {'fontname':'Helvetica'}
    
    # plt.rcParams["font.family"] = "Comic Sans MS" # "Comic Sans MS"
    
    font = {'family' : 'Comic Sans MS',
            'weight' : 'bold',
            'style'  : 'italic',
            'size'   : 12}
    
    matplotlib.rc('font', **font)
    
    # a,b - the two ends of our interval
    # [-16, 57]
    a = 0.000001
    b = 0.001
    
    def f(t):
        return ( t + 2 if (t==int(t)) else 2**t )
    
    def g(x):
        return np.floor(1/x)
    
    t = np.linspace(a, b, 400)
    f1 = f(t)
    g1 = g(t)
    
    plt.plot(t, f1, 'red', label="f (x)", linewidth=2) # plotting t, f1 separately 
    plt.plot(t, g1, 'green', label="g (x)", linewidth=2) # plotting t, g1 separately
    
    plt.axhline(0, color='k')
    plt.axvline(0, color='k')
    
    # plt.xlabel('X axis', fontsize=14)
    # plt.ylabel('Y axis', fontsize=14)
    
    # plt.title("Illustration of the mean value theorem proof")
    
    # plt.savefig('FFF.png')
    
    plt.legend()
    plt.show()

I also tried t==np.floor(t) but this says: The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()
I know why but that's not what I want. I just want to compute this function on that ndarray.

I also tried `t==np.floor(t)` but this says: `The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()` — peter.petrov, Jun 26 '20 at 22:42
Does this solve your problem? https://stackoverflow.com/questions/31638508/plot-piecewise-function-in-python — trish_s, Jun 26 '20 at 22:53
Does https://stackoverflow.com/questions/26254946/piecewise-functions-on-numpy-arrays help? — Karl Knechtel, Jun 26 '20 at 22:53

score 1 · Answer 1 · answered Jun 26 '20 at 22:59

You need to vectorize the function to perform element by element inspection on a numpy array.

import numpy as np

def main():
    a = 0.000001
    b = 0.001
    
    def f(t):
        if int(t) == t:
            return t + 2
        else:
            return 2 ** t

    t = np.linspace(a, b, 400)
    new_f = np.vectorize(f)
    f1 = new_f(t)
    print(f1)

if __name__ == '__main__':
    main()

score 1 · Answer 2 · answered Jun 26 '20 at 22:59

We can start by figuring out how to create a mask of where the integers are as outlined at this post. Next we can review how np.linspace() works and realize you were not generating a list with any integers at all (the largest number being 0.001). We can arrive at the code below

import numpy as np
import matplotlib.pyplot as plt

t = np.linspace(0, 10, 201) # 201 numbers from 0 to 10 (the extra one is needed for the endpoint)
data = np.where(np.mod(t, 1), 2**t, t + 2)  # np.mod will be 0 where there is an int, based on this information we use one of the two formulas
plt.plot(t, data) # plot results
plt.show()

Sanity checking our plot looks ok

How do I plot this function in Python?

2 Answers2