Graph velocity and distance from calculus functions in Python

Question

Is there an more code-light way I could have implemented graphing the following piece-wise functions in the calculus problem more easily? In my approach, I used matplotlib and just combined the graphs into two main graphs to show the discontinuity.

import matplotlib.pyplot as plt

def v(time_range):
    velocity_val = []
    for i in time_range:
        if i < .2:
            velocity_val.append(20)
        elif i > .2:
            velocity_val.append(0)
    return velocity_val

def f(time_range):
    distance_val = []
    for i in time_range:
        if i <= .2:
            distance_val.append(20*i)
        if i >= .2:
            distance_val.append(4)
    return distance_val

def time_vals(time_range):
    decimal = 100
    time_val = []
    for i in time_range:
        num = i / decimal
        time_val.append(num)
    return time_val

#convert time into decimal
time_range_1 = range(1,20,1)
time_range_2 = range(21,40,1)
t_1 = time_vals(time_range_1)
t_2 = time_vals(time_range_2)

#get x, y for plot
v_1 = v(t_1)
v_2 = v(t_2)
f_1 = f(t_1)
f_2 = f(t_2)

#plot values into two graphs.
plt.subplot(2, 1, 1)
plt.plot(t_1, v_1)
plt.plot(t_2, v_2)
plt.title(' Problem 9')
plt.ylabel('Velocity')

plt.subplot(2, 1, 2)
plt.plot(t_1, f_1)
plt.plot(t_2, f_2)
plt.xlabel('time (t)')
plt.ylabel('Velocity');

Answer 1

You can make use of np.where to assign your v(t) and f(t) depending on your the conditions. You don't need any for loops. Vectorized approaches make your code much neat and short. In np.where, you first check the condition and then the first value after the condition is assigned to the indices where the condition holds True and the second value is assigned to the indices where the condition holds False.

Below is an example:

import numpy as np
import matplotlib.pyplot as plt

# Initialise time
t_1 = np.arange(1,20,1)/100
t_2 = np.arange(21,40,1)/100

# Compute v(t)
v_1 = np.where(t_1<0.2, 20, 0)
v_2 = np.where(t_2<0.2, 20, 0)

# Compute f(t)
f_1 = np.where(t_1<=0.2, 20*t_1, 4)
f_2 = np.where(t_2<=0.2, 20*t_2, 4)

# Plotting as you are doing