Color volume plot with data from a different function

📊 Plotly Python
1

Hello, I have built a code that develops a volume plot based on a given function, this all works well and good. The problem I have is that it is all one color and I want to give it a color based on values from a second function.
Function 1
Wave = Psi_2 → computes values from calculation
Function 2
colours = np.sign(psi)–> computes values from a different calculation (these are either +1 or -1)
Function 3 → gives 3d volume plot (x,y,z data with the ‘values’ coming from function 1)
I want function 3 to also give colour to the plot based on the output of function 2.
I know you cant do this for the isosurface function but I was hoping volume would work.
In matlab it is relatively straight forward by just adding the ‘colors’ term.

I have added some of the code for context,
def psi(n,l,m,r,theta,phi):
psi = R(n,l,r)*Y(l,m,theta,phi)
return psi

def psi_2(n,l,m,r,theta,phi):
psi_2 = (r*R(n,l,r)*Y(l,m,theta,phi))**2
return psi_2

def Plot_psi(n, l, m):
probabilitydensity = 1e-6
size = np.round((2.5 * n ** 2 + 1.5 * n + 1) / 5) * 5
border = 80
accuracy = 150
raster = np.linspace(-border,border,accuracy)
[x,y,z] = np.meshgrid(raster,raster,raster)
r = np.sqrt(x**2+y ** 2 + z ** 2)
theta = np.arccos(z/r)
phi = np.arctan2(y,x)
colours = np.sign(psi(n, l, m, r, theta, phi))
Wave=psi_2(n,l,m,r,theta,phi)

fig = go.Figure(data=go.Volume(
    x=x.flatten(),
    y=y.flatten(),
    z=z.flatten(),
    value=Wave.flatten(),
    opacity=0.5,
    isomin=probabilitydensity,
    isomax=probabilitydensity,
    surface_count=1,
    caps=dict(x_show=True, y_show=True)
))
fig.show()

Ideally, I would be able to say something like colormap=colours like in Matlab.

Any help would be awesome.
Thanks.