Removing datapoints outside interval for both axes of a plot

Question

I am trying to plot some data using matplotlib.

import matplotlib.pyplot as plt

x_data = np.arange(0,100)
y_data = np.random.randint(11, size=(100,))
plt.plot(x_data, y_data)
plt.show

This, of course, works fine. However, I would like to remove the data that is outside a given interval (eg 4 < y_data < 6). For the y_data, this is done by

y_data_2 = [x for x in y_data if 4 <= x <= 6]

However, since the first dimensions are no longer equal, you are no longer able to plot y_data_2 vs. x_data. If you try to

plt.plot(x_data, y_data_2)

you will, of course, get an error stating that

ValueError: x and y must have same first dimension, but have shapes (100,) and (35,)

My question is thus twofold: is there a simple way for me to remove the equivalent datapoints in x_data? Also, is there a way I could find the indices of the points that are to be removed?

Thank you.

Answer 1

You can use masking together with indexing. Here you create a mask to capture values y values which lie between 4 and 6. You then apply this conditional mask to your x_data and y_data to get the corresponding values. This way you don't need any for loop or list comprehensions.

x_data = np.arange(0,100)
y_data = np.random.randint(11, size=(100,))
mask = (y_data>=4) & (y_data<=6)

plt.plot(x_data[mask], y_data[mask], 'bo')

Answer 2

First, you can get the index of y_data_2 in y_data, and then get the subarray x_data_2 of x_data. Then, plot the x_data_2, y_data_2.

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd

x_data = np.arange(0,100)
y_data = np.random.randint(11, size=(100,))
y = pd.Series(y_data)

y_data_2 = [x for x in y_data if 4 <= x <= 6]
index = y[y.isin(y_data_2)].index
print(index)
x_data_2 = x_data[index]
plt.plot(x_data, y_data)
plt.scatter(x_data_2, y_data_2)
plt.show()