Colorize based on condition

I am having the following code to display a plot. I would like to conditionally format the values based on a threshold of 0.05.

from matplotlib.colors import to_rgba

# Generate x-axis
x = np.linspace(0, len(data_formula), len(data_formula))
colors = np.where(data_formula <= 0.05, "blue", "green")
plt.plot(x, data_formula, c=colors)

# Add labels and title
plt.ylabel('Volume')
plt.xlabel('time')
plt.title('Energy')

# Display the plot
plt.show()

Unfortunately I do receive the error: array(['blue', 'blue', 'blue', ..., 'blue', 'blue', 'blue'], dtype='<U5') is not a valid value for color, which indicates that I have passed the wrong value to the color parameter. I have tried it with lists, etc. But it doesn't seem to work. What went wrong? Does the argument color just not accept any datastructure or is the format wrong?

For reference: data_formula is defined as follow:

def energy(x):
    
    return (x**2)

data_formula = np.apply_along_axis(energy, axis=0, arr=data_normalized)

It is of datatype: numpy.ndarray.

For using different colors you should split data segments each one with its color:

#%matplotlib inline
import numpy as np
import matplotlib.pyplot as plt

x = np.linspace(0, len(data_formula), len(data_formula))
crossings = np.where(np.diff(data_formula > 0.05))[0]
start = 0

# Plot segments
for ind in crossings:
    plt.plot(x[start:ind+1], data_formula[start:ind+1], 
             color='blue' if data_formula[start] <= 0.05 else 'green')
    start = ind+1

plt.plot(x[start:], data_formula[start:], 
         color='blue' if data_formula[start] <= 0.05 else 'green')

plt.ylabel('Volume')
plt.xlabel('time')
plt.title('Energy')
plt.show()

I assumed your data is like:

data_formula = np.array([0, 0.01, 0.02, 0.04, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0])

If instead of segments you just want the points do the following:

import numpy as np
import matplotlib.pyplot as plt


for x, y in enumerate(data_formula):
    plt.scatter(x, y, color='blue' if y <= 0.05 else 'green')

    
plt.ylabel('Volume')
plt.xlabel('time')
plt.title('Energy')
plt.show()

this last one inspired from stackoverflow: different colors series scatter plot on matplotlib.

If the previous answer takes to much time try:

import numpy as np
import matplotlib.pyplot as plt

# data_formula = np.array([0, 0.01, 0.02, 0.04, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0])

xy_blue = [(x, y) for x, y in enumerate(data_formula) if y <= 0.05]
xy_gree = [(x, y) for x, y in enumerate(data_formula) if y > 0.05]

plt.scatter([t[0] for t in xy_blue], [t[1] for t in xy_blue], color='blue')
plt.scatter([t[0] for t in xy_gree], [t[1] for t in xy_gree], color='green')

    
plt.ylabel('Volume')
plt.xlabel('time')
plt.title('Energy')
plt.show()

Anyway, I've never see a plot to take more than 30sec so maybe you have to reload your environment, you have way too much data to plot and should make different plots, or you have a low RAM issue.