使用过滤后的 y 值创建子图

Question

I have a plot that looks like this:

And I want to create subplots with linear regressions on the y axis values from 0-5, 5-10, and 10-15 days. My code currently looks like this:

#Read in proper dataframe
storms = pd.read_csv('storms.csv', sep = ',', header = 0, ) 

#Define variables for linear regression - frequency
x_col ='season'
y_col = 'days'
x = storms[x_col]
y = storms[y_col]
x_array = np.array(x).reshape(-1,1)
y_array = np.array(y).reshape(-1,1)

linreg = LinearRegression().fit(x_array,y_array)

#Perform linear regression for frequency 
lin_reg(x_array,y_array)

#Plot
sns.set_theme(context='notebook', style='darkgrid')
sns.light_palette("#79C")

plt.scatter(x_array,y_array, alpha = 0.25,)
plt.plot(x_array,linreg.predict(x_array),  label='y=-0.0154x+38.5978')
plt.xlabel('Season')
plt.ylabel('Duration of Storms (in days)')
plt.title('Duration of Storms Over Time')
plt.legend()
plt.show

I tried defining functions to filter the y axis values which I then applied with the plt.subplot function which looked like this:

def Filter1(f):
    if f <= 5:
        return False
    else:
        return True

y_filtered = Filter1(y_array)

#Plotting subplots
plt.subplot(3,1,1)
plt.plot(x_array,y_filtered)

But no success yet. Any suggestions?

Answer 1

When you are taking your x- and y-values from the dataframe columns, you can convert them to numpy arrays by calling

x = storms[x_col].values
y = storms[y_col].values

Then, you don't have to define any function to filter the values, you can just do it with numpy features. You can obtain an array containing true and false values depending on your condition. For example, you can say x > 5 and it will give you an array with the same size as x , where for each original index in x there will be a false if the value at that index is below (or equal to) 5 and true if it is above 5.Then you can use that boolean array as a mask for other arrays of the same size, returning only the elements at the index where the mask is true . That's how you can filter with numpy:

mask = y[(y>=5) & (y<10)]
y_filtered = y[mask]
x_filtered = x[mask]

Then perform the regression with x_filtered and y_filtered .