N-D interpolation for equally-spaced data - numpy

I'm trying to copy the Scipy Cookbook function:
from scipy import ogrid, sin, mgrid, ndimage, array
x,y = ogrid[-1:1:5j,-1:1:5j]
fvals = sin(x)*sin(y)
newx,newy = mgrid[-1:1:100j,-1:1:100j]
x0 = x[0,0]
y0 = y[0,0]
dx = x[1,0] - x0
dy = y[0,1] - y0
ivals = (newx - x0)/dx
jvals = (newy - y0)/dy
coords = array([ivals, jvals])
newf = ndimage.map_coordinates(fvals, coords)
by using my own function that has to work for many scenarios
import scipy
import numpy as np
"""N-D interpolation for equally-spaced data"""
x = np.c_[plist['modx']]
y = np.transpose(np.c_[plist['mody']])
pdb.set_trace()
#newx,newy = np.meshgrid(plist['newx'],plist['newy'])
newx,newy = scipy.mgrid[plist['modx'][0]:plist['modx'][-1]:-plist['remapto'],
plist['mody'][0]:plist['mody'][-1]:-plist['remapto']]
x0 = x[0,0]
y0 = y[0,0]
dx = x[1,0] - x0
dy = y[0,1] - y0
ivals = (newx - x0)/dx
jvals = (newy - y0)/dy
coords = scipy.array([ivals, jvals])
for i in np.arange(ivals.shape[0]):
nvals[i] = scipy.ndimage.map_coordinates(ivals[i], coords)
I'm having difficulty getting this code to work properly. The problem areas are:
1.) Recreating this line: newx,newy = mgrid[-1:1:100j,-1:1:100j]. In my case I have a dictionary with the grid in vector form. I've tried to recreate this line using np.meshgrid but then I get an error on line coords = scipy.array([ivals, jvals]). I'm looking for some help in recreating this Cookbook function and making it more dynamic
any help is greatly appreciated.
/M

You should have a look at the documentation for map_coordinates. I don't see where the actual data you are trying to interpolate is in your code. What I mean is, presumably you have some data input which is a function of x and y; i.e. input = f(x,y) that you want to interpolate. In the first example you show, this is the array fvals. This should be your first argument to map_coordinates.
For example, if the data you are trying to inperpolate is input, which should be a 2-dimensional array of shape (len(x),len(y)), then the interpolated data would be:
interpolated_data = map_coordinates(input, coords)

Related

Individual axes titles in a facet plot?

I have a Pandas dataframe that I would like to plot with different labels on each yaxis using plotly.
For example
fig = dataframe.plot(
x="t", y="value",
facet_col="variable_label",
facet_col_wrap=2,
yaxes_titles="multiplier_label", # This is unfortunately not a valid keyword argument
).update_yaxes(matches=None, showticklabels=True)
Is there a way to update each yaxis with a different title?
I found a solution - it is quite complicated due to the ordering of yaxes being from bottom to top, left to right! I needed the following utility function:
def get_yaxis(fig, i, facet_col_wrap):
"""Get the yaxis for the ith facet in a facet plot"""
yaxes = list(fig.select_yaxes())
n = len(yaxes)
n_cols = min(facet_col_wrap, n)
n_rows = n // n_cols
idx = [[n-((1+r)*n_cols-c) for c in range(n_cols)] for r in range(n_rows)]
idx = [item for sublist in idx for item in sublist]
return yaxes[idx[i]]
With this function and a list of my desired yaxes_titles at hand, I can finally do:
for i, yaxis_title in enumerate(yaxis_titles):
get_yaxis(fig, i, facet_col_wrap).update(title_text=yaxis_title)

How do I stop python from drawing a new colorbar every iteration

I made this code as a CFD of sorts for fun, and I want to add a color bar to show the velocity of the fluid in different places. Unfortunately, every time it plots a new frame it also plots a new colorbar rather than refreshing the old one. I'd like to get it to refresh rather than draw a new one entirely. Any help would be appreciated. Plotting Begins on line 70
import numpy as np
from matplotlib import pyplot
plot_every = 100
def distance(x1,y1,x2,y2):
return np.sqrt((x2-x1)**2 + (y2-y1)**2)
def main():
Nx = 400 #Cells Across x direction
Ny = 100 #Cells Across y direction
#CELL <> NODE
tau = .53 #kinimatic viscosity
tymestep = tau
Nt = 30000 #total iterations
#Lattice Speeds and Velcoties
NL = 9 #There are 9 differnct velocites, (up, down, left, right, up-left diag, up-right diag, down-left diag, down-right diag, and zero)
#NL would be 27 in 3D flow
cxs = np.array([0,0,1,1,1,0,-1,-1,-1]) #I don't know what this is
cys = np.array([0,1,1,0,-1,-1,-1,0,1]) #I don't know what this is
weights = np.array([4/9,1/9,1/36,1/9,1/36,1/9,1/36,1/9,1/36])
#COMPLETELY DIFFERNT WEIGTS FOR 2D AND 3D FLOW
#Initial Conditions
F = np.ones((Ny,Nx,NL)) + 0.01*np.random.randn(Ny,Nx,NL)
F[:,:,3] = 2.3 #Assigning an inital speed in x direction with right as posative
#Drawing Our cylinder
cylinder = np.full((Ny,Nx), False)
radius = 13
for y in range(0,Ny):
for x in range(0,Nx):
if (distance(Nx//4,Ny//2,x,y) < radius):
cylinder[y][x] = True
#main loop
for it in range(Nt):
#print(it)
F[:,-1, [6,7,8]] = F[:,-2, [6,7,8]] #without this, fluid will bounce off of outside walls (you may want this to happen)
F[:,0, [2,3,4]] = F[:,1, [2,3,4]] #without this, fluid will bounce off of outside walls (you may want this to happen)
for i, cx, cy in zip(range(NL),cxs, cys): #this line is sligtly differnt than his because I think he made a typo
F[:,:,i] = np.roll(F[:,:,i], cx, axis = 1)
F[:,:,i] = np.roll(F[:,:,i], cy, axis = 0)
bndryF = F[cylinder,:]
bndryF = bndryF[:, [0,5,6,7,8,1,2,3,4]] #defines what happens in a colsion (reverse the velocity). This works by setting the up vel to down vel etc
#Fluid Variables
rho = np.sum(F,2) #density
ux = np.sum(F * cxs, 2)/rho #x velocity (momentum/mass)
uy = np.sum(F * cys, 2)/rho #y velocity
F[cylinder,: ] = bndryF
ux[cylinder] = 0 #set all velocities in cylinder = 0
uy[cylinder] = 0 #set all velocities in cylinder = 0
#collisions
Feq = np.zeros(F.shape)
for i, cx, cy, w in zip(range(NL), cxs, cys, weights):
Feq[:, :, i] = rho * w * (
1 + 3*(cx*ux + cy*uy) + 9*(cx*ux + cy*uy)**2/2 - 3*(ux**2 + uy**2)/2
)
F += -1/tau * (F-Feq)
if(it%plot_every == 0):
dfydx = ux[2:, 1:-1] - ux[0:-2, 1: -1]
dfxdy = uy[1: -1, 2:] - uy[1: -1, 0: -2]
curl = dfydx - dfxdy
pyplot.imshow(np.sqrt(ux**2+uy**2),cmap = "bwr")
#pyplot.imshow(curl, cmap = "bwr")
pyplot.colorbar(label="Velocity", orientation="horizontal")
pyplot.pause(0.01)
pyplot.cla()
if __name__ == "__main__":
main()
In your code you are adding a new colorbar at every iteration.
As far as I know, it is impossible to update a colorbar. The workaround is to delete the colorbar of the previous time step, and replace it with a new one.
This is achieved by the update_colorbar function in the code below.
import numpy as np
from matplotlib import pyplot
from matplotlib.cm import ScalarMappable
from matplotlib.colors import Normalize
plot_every = 100
def distance(x1,y1,x2,y2):
return np.sqrt((x2-x1)**2 + (y2-y1)**2)
def update_colorbar(fig, cmap, param, norm=None):
"""The name is misleading: here we create a new colorbar which will be
placed on the same colorbar axis as the original.
"""
# colorbar axes
cax = None
if len(fig.axes) > 1:
cax = fig.axes[-1]
# remove the previous colorbar, if present
if cax is not None:
cax.clear()
if norm is None:
norm = Normalize(vmin=np.amin(param), vmax=np.amax(param))
mappable = ScalarMappable(cmap=cmap, norm=norm)
fig.colorbar(mappable, orientation="horizontal", label="Velocity", cax=cax)
def main():
Nx = 400 #Cells Across x direction
Ny = 100 #Cells Across y direction
#CELL <> NODE
tau = .53 #kinimatic viscosity
tymestep = tau
Nt = 30000 #total iterations
#Lattice Speeds and Velcoties
NL = 9 #There are 9 differnct velocites, (up, down, left, right, up-left diag, up-right diag, down-left diag, down-right diag, and zero)
#NL would be 27 in 3D flow
cxs = np.array([0,0,1,1,1,0,-1,-1,-1]) #I don't know what this is
cys = np.array([0,1,1,0,-1,-1,-1,0,1]) #I don't know what this is
weights = np.array([4/9,1/9,1/36,1/9,1/36,1/9,1/36,1/9,1/36])
#COMPLETELY DIFFERNT WEIGTS FOR 2D AND 3D FLOW
#Initial Conditions
F = np.ones((Ny,Nx,NL)) + 0.01*np.random.randn(Ny,Nx,NL)
F[:,:,3] = 2.3 #Assigning an inital speed in x direction with right as posative
#Drawing Our cylinder
cylinder = np.full((Ny,Nx), False)
radius = 13
for y in range(0,Ny):
for x in range(0,Nx):
if (distance(Nx//4,Ny//2,x,y) < radius):
cylinder[y][x] = True
fig, ax = pyplot.subplots()
cmap = "bwr"
#main loop
for it in range(Nt):
# clear previous images
ax.images.clear()
#print(it)
F[:,-1, [6,7,8]] = F[:,-2, [6,7,8]] #without this, fluid will bounce off of outside walls (you may want this to happen)
F[:,0, [2,3,4]] = F[:,1, [2,3,4]] #without this, fluid will bounce off of outside walls (you may want this to happen)
for i, cx, cy in zip(range(NL),cxs, cys): #this line is sligtly differnt than his because I think he made a typo
F[:,:,i] = np.roll(F[:,:,i], cx, axis = 1)
F[:,:,i] = np.roll(F[:,:,i], cy, axis = 0)
bndryF = F[cylinder,:]
bndryF = bndryF[:, [0,5,6,7,8,1,2,3,4]] #defines what happens in a colsion (reverse the velocity). This works by setting the up vel to down vel etc
#Fluid Variables
rho = np.sum(F,2) #density
ux = np.sum(F * cxs, 2)/rho #x velocity (momentum/mass)
uy = np.sum(F * cys, 2)/rho #y velocity
F[cylinder,: ] = bndryF
ux[cylinder] = 0 #set all velocities in cylinder = 0
uy[cylinder] = 0 #set all velocities in cylinder = 0
#collisions
Feq = np.zeros(F.shape)
for i, cx, cy, w in zip(range(NL), cxs, cys, weights):
Feq[:, :, i] = rho * w * (
1 + 3*(cx*ux + cy*uy) + 9*(cx*ux + cy*uy)**2/2 - 3*(ux**2 + uy**2)/2
)
F += -1/tau * (F-Feq)
if(it%plot_every == 0):
dfydx = ux[2:, 1:-1] - ux[0:-2, 1: -1]
dfxdy = uy[1: -1, 2:] - uy[1: -1, 0: -2]
curl = dfydx - dfxdy
img = np.sqrt(ux**2+uy**2)
ax.imshow(img ,cmap = cmap)
#pyplot.imshow(curl, cmap = "bwr")
update_colorbar(fig, cmap, param=img)
pyplot.pause(0.01)
if __name__ == "__main__":
main()
One thing you can definitely improve is the following line of code, which defines the values visible in the colorbar:
norm = Normalize(vmin=np.amin(param), vmax=np.amax(param))
Specifically, you'd have to choose a wise (conservative) value for vmax=. Currently, vmax=np.amax(param), but the maximum is going to change at every iteration. If I were you, I would chose a value big enough such that np.amax(param) < your_value, in order to ensure consistent colors for each time step.

Shortest rotation between two vectors not working like expected

def signed_angle_between_vecs(target_vec, start_vec, plane_normal=None):
start_vec = np.array(start_vec)
target_vec = np.array(target_vec)
start_vec = start_vec/np.linalg.norm(start_vec)
target_vec = target_vec/np.linalg.norm(target_vec)
if plane_normal is None:
arg1 = np.dot(np.cross(start_vec, target_vec), np.cross(start_vec, target_vec))
else:
arg1 = np.dot(np.cross(start_vec, target_vec), plane_normal)
arg2 = np.dot(start_vec, target_vec)
return np.arctan2(arg1, arg2)
from scipy.spatial.transform import Rotation as R
world_frame_axis = input_rotation_object.apply(canonical_axis)
angle = signed_angle_between_vecs(canonical_axis, world_frame_axis)
axis_angle = np.cross(world_frame_axis, canonical_axis) * angle
C = R.from_rotvec(axis_angle)
transformed_world_frame_axis_to_canonical = C.apply(world_frame_axis)
I am trying to align world_frame_axis to canonical_axis by performing a rotation around the normal vector generated by the cross product between the two vectors, using the signed angle between the two axes.
However, this code does not work. If you start with some arbitrary rotation as input_rotation_object you will see that transformed_world_frame_axis_to_canonical does not match canonical_axis.
What am I doing wrong?
not a python coder so I might be wrong but this looks suspicious:
start_vec = start_vec/np.linalg.norm(start_vec)
from the names I would expect that np.linalg.norm normalizes the vector already so the line should be:
start_vec = np.linalg.norm(start_vec)
and all the similar lines too ...
Also the atan2 operands are not looking right to me. I would (using math):
a = start_vec / |start_vec | // normalized start
b = target_vec / |target_vec| // normalized end
u = a // normalized one axis of plane
v = cross(u ,b)
v = cross(v ,u)
v = v / |v| // normalized second axis of plane perpendicular to u
dx = dot(u,b) // target vector in 2D aligned to start
dy = dot(v,b)
ang = atan2(dy,dx)
beware the ang might negated (depending on your notations) if the case either add minus sign or reverse the order in cross(u,v) to cross(v,u) Also you can do sanity check with comparing result to unsigned:
ang' = acos(dot(a,b))
in absolute values they should be the same (+/- rounding error).

Scipy Optimize minimize returns the initial value

I am building machine learning models for a certain data set. Then, based on the constraints and bounds for the outputs and inputs, I am trying to find the input parameters for the most minimized answer.
The problem which I am facing is that, when the model is a linear regression model or something like lasso, the minimization works perfectly fine.
However, when the model is "Decision Tree", it constantly returns the very initial value that is given to it. So basically, it does not enforce the constraints.
import numpy as np
import pandas as pd
from scipy.optimize import minimize
I am using the very first sample from the input data set for the optimization. As it is only one sample, I need to reshape it to (1,-1) as well.
x = df_in.iloc[0,:]
x = np.array(x)
x = x.reshape(1,-1)
This is my Objective function:
def objective(x):
x = np.array(x)
x = x.reshape(1,-1)
y = 0
for n in range(df_out.shape[1]):
y = Model[n].predict(x)
Y = y[0]
return Y
Here I am defining the bounds of inputs:
range_max = pd.DataFrame(range_max)
range_min = pd.DataFrame(range_min)
B_max=[]
B_min =[]
for i in range(range_max.shape[0]):
b_max = range_max.iloc[i]
b_min = range_min.iloc[i]
B_max.append(b_max)
B_min.append(b_min)
B_max = pd.DataFrame(B_max)
B_min = pd.DataFrame(B_min)
bnds = pd.concat([B_min, B_max], axis=1)
These are my constraints:
con_min = pd.DataFrame(c_min)
con_max = pd.DataFrame(c_max)
Here I am defining the constraint function:
def const(x):
x = np.array(x)
x = x.reshape(1,-1)
Y = []
for n in range(df_out.shape[1]):
y = Model[n].predict(x)[0]
Y.append(y)
Y = pd.DataFrame(Y)
a4 =[]
for k in range(Y.shape[0]):
a1 = Y.iloc[k,0] - con_min.iloc[k,0]
a2 = con_max.iloc[k, 0] - Y.iloc[k,0]
a3 = [a2,a1]
a4 = np.concatenate([a4, a3])
return a4
c = const(x)
con = {'type': 'ineq', 'fun': const}
This is where I try to minimize. I do not pick a method as the automatically picked model has worked so far.
sol = minimize(fun = objective, x0=x,constraints=con, bounds=bnds)
So the actual constraints are:
c_min = [0.20,1000]
c_max = [0.3,1600]
and the max and min range for the boundaries are:
range_max = [285,200,8,85,0.04,1.6,10,3.5,20,-5]
range_min = [215,170,-1,60,0,1,6,2.5,16,-18]
I think you should check the output of 'sol'. At times, the algorithm is not able to perform line search completely. To check for this, you should check message associated with 'sol'. In such a case, the optimizer returns initial parameters itself. There may be various reasons of this behavior. In a nutshell, please check the output of sol and act accordingly.
Arad,
If you have not yet resolved your issue, try using scipy.optimize.differential_evolution instead of scipy.optimize.minimize. I ran into similar issues, particularly with decision trees because of their step-like behavior resulting in infinite gradients.

Is there any way to use bivariate colormaps in matplotlib?

In other words, I want to make a heatmap (or surface plot) where the color varies as a function of 2 variables. (Specifically, luminance = magnitude and hue = phase.) Is there any native way to do this?
Some examples of similar plots:
Several good examples of exactly(?) what I want to do.
More examples from astronomy, but with non-perceptual hue
Edit: This is what I did with it: https://github.com/endolith/complex_colormap
imshow can take an array of [r, g, b] entries. So you can convert the absolute values to intensities and phases - to hues.
I will use as an example complex numbers, because for it it makes the most sense. If needed, you can always add numpy arrays Z = X + 1j * Y.
So for your data Z you can use e.g.
imshow(complex_array_to_rgb(Z))
where (EDIT: made it quicker and nicer thanks to this suggestion)
def complex_array_to_rgb(X, theme='dark', rmax=None):
'''Takes an array of complex number and converts it to an array of [r, g, b],
where phase gives hue and saturaton/value are given by the absolute value.
Especially for use with imshow for complex plots.'''
absmax = rmax or np.abs(X).max()
Y = np.zeros(X.shape + (3,), dtype='float')
Y[..., 0] = np.angle(X) / (2 * pi) % 1
if theme == 'light':
Y[..., 1] = np.clip(np.abs(X) / absmax, 0, 1)
Y[..., 2] = 1
elif theme == 'dark':
Y[..., 1] = 1
Y[..., 2] = np.clip(np.abs(X) / absmax, 0, 1)
Y = matplotlib.colors.hsv_to_rgb(Y)
return Y
So, for example:
Z = np.array([[3*(x + 1j*y)**3 + 1/(x + 1j*y)**2
for x in arange(-1,1,0.05)] for y in arange(-1,1,0.05)])
imshow(complex_array_to_rgb(Z, rmax=5), extent=(-1,1,-1,1))
imshow(complex_array_to_rgb(Z, rmax=5, theme='light'), extent=(-1,1,-1,1))
imshow will take an NxMx3 (rbg) or NxMx4 (grba) array so you can do your color mapping 'by hand'.
You might be able to get a bit of traction by sub-classing Normalize to map your vector to a scaler and laying out a custom color map very cleverly (but I think this will end up having to bin one of your dimensions).
I have done something like this (pdf link, see figure on page 24), but the code is in MATLAB (and buried someplace in my archives).
I agree a bi-variate color map would be useful (primarily for representing very dense vector fields where your kinda up the creek no matter what you do).
I think the obvious extension is to let color maps take complex arguments. It would require specialized sub-classes of Normalize and Colormap and I am going back and forth on if I think it would be a lot of work to implement. I suspect if you get it working by hand it will just be a matter of api wrangling.
I created an easy to use 2D colormap class, that takes 2 NumPy arrays and maps them to an RGB image, based on a reference image.
I used #GjjvdBurg's answer as a starting point. With a bit of work, this could still be improved, and possibly turned into a proper Python module - if you want, feel free to do so, I grant you all credits.
TL;DR:
# read reference image
cmap_2d = ColorMap2D('const_chroma.jpeg', reverse_x=True) # , xclip=(0,0.9))
# map the data x and y to the RGB space, defined by the image
rgb = cmap_2d(data_x, data_y)
# generate a colorbar image
cbar_rgb = cmap_2d.generate_cbar()
The ColorMap2D class:
class ColorMap2D:
def __init__(self, filename: str, transpose=False, reverse_x=False, reverse_y=False, xclip=None, yclip=None):
"""
Maps two 2D array to an RGB color space based on a given reference image.
Args:
filename (str): reference image to read the x-y colors from
rotate (bool): if True, transpose the reference image (swap x and y axes)
reverse_x (bool): if True, reverse the x scale on the reference
reverse_y (bool): if True, reverse the y scale on the reference
xclip (tuple): clip the image to this portion on the x scale; (0,1) is the whole image
yclip (tuple): clip the image to this portion on the y scale; (0,1) is the whole image
"""
self._colormap_file = filename or COLORMAP_FILE
self._img = plt.imread(self._colormap_file)
if transpose:
self._img = self._img.transpose()
if reverse_x:
self._img = self._img[::-1,:,:]
if reverse_y:
self._img = self._img[:,::-1,:]
if xclip is not None:
imin, imax = map(lambda x: int(self._img.shape[0] * x), xclip)
self._img = self._img[imin:imax,:,:]
if yclip is not None:
imin, imax = map(lambda x: int(self._img.shape[1] * x), yclip)
self._img = self._img[:,imin:imax,:]
if issubclass(self._img.dtype.type, np.integer):
self._img = self._img / 255.0
self._width = len(self._img)
self._height = len(self._img[0])
self._range_x = (0, 1)
self._range_y = (0, 1)
#staticmethod
def _scale_to_range(u: np.ndarray, u_min: float, u_max: float) -> np.ndarray:
return (u - u_min) / (u_max - u_min)
def _map_to_x(self, val: np.ndarray) -> np.ndarray:
xmin, xmax = self._range_x
val = self._scale_to_range(val, xmin, xmax)
rescaled = (val * (self._width - 1))
return rescaled.astype(int)
def _map_to_y(self, val: np.ndarray) -> np.ndarray:
ymin, ymax = self._range_y
val = self._scale_to_range(val, ymin, ymax)
rescaled = (val * (self._height - 1))
return rescaled.astype(int)
def __call__(self, val_x, val_y):
"""
Take val_x and val_y, and associate the RGB values
from the reference picture to each item. val_x and val_y
must have the same shape.
"""
if val_x.shape != val_y.shape:
raise ValueError(f'x and y array must have the same shape, but have {val_x.shape} and {val_y.shape}.')
self._range_x = (np.amin(val_x), np.amax(val_x))
self._range_y = (np.amin(val_y), np.amax(val_y))
x_indices = self._map_to_x(val_x)
y_indices = self._map_to_y(val_y)
i_xy = np.stack((x_indices, y_indices), axis=-1)
rgb = np.zeros((*val_x.shape, 3))
for indices in np.ndindex(val_x.shape):
img_indices = tuple(i_xy[indices])
rgb[indices] = self._img[img_indices]
return rgb
def generate_cbar(self, nx=100, ny=100):
"generate an image that can be used as a 2D colorbar"
x = np.linspace(0, 1, nx)
y = np.linspace(0, 1, ny)
return self.__call__(*np.meshgrid(x, y))
Usage:
Full example, using the constant chroma reference taken from here as a screenshot:
# generate data
x = y = np.linspace(-2, 2, 300)
xx, yy = np.meshgrid(x, y)
ampl = np.exp(-(xx ** 2 + yy ** 2))
phase = (xx ** 2 - yy ** 2) * 6 * np.pi
data = ampl * np.exp(1j * phase)
data_x, data_y = np.abs(data), np.angle(data)
# Here is the 2D colormap part
cmap_2d = ColorMap2D('const_chroma.jpeg', reverse_x=True) # , xclip=(0,0.9))
rgb = cmap_2d(data_x, data_y)
cbar_rgb = cmap_2d.generate_cbar()
# plot the data
fig, plot_ax = plt.subplots(figsize=(8, 6))
plot_extent = (x.min(), x.max(), y.min(), y.max())
plot_ax.imshow(rgb, aspect='auto', extent=plot_extent, origin='lower')
plot_ax.set_xlabel('x')
plot_ax.set_ylabel('y')
plot_ax.set_title('data')
# create a 2D colorbar and make it fancy
plt.subplots_adjust(left=0.1, right=0.65)
bar_ax = fig.add_axes([0.68, 0.15, 0.15, 0.3])
cmap_extent = (data_x.min(), data_x.max(), data_y.min(), data_y.max())
bar_ax.imshow(cbar_rgb, extent=cmap_extent, aspect='auto', origin='lower',)
bar_ax.set_xlabel('amplitude')
bar_ax.set_ylabel('phase')
bar_ax.yaxis.tick_right()
bar_ax.yaxis.set_label_position('right')
for item in ([bar_ax.title, bar_ax.xaxis.label, bar_ax.yaxis.label] +
bar_ax.get_xticklabels() + bar_ax.get_yticklabels()):
item.set_fontsize(7)
plt.show()
I know this is an old post, but want to help out others that may arrive late. Below is a python function to implement complex_to_rgb from sage. Note: This implementation isn't optimal, but it is readable. See links: (examples)(source code)
Code:
import numpy as np
def complex_to_rgb(z_values):
width = z_values.shape[0]
height = z_values.shape[1]
rgb = np.zeros(shape=(width, height, 3))
for i in range(width):
row = z_values[i]
for j in range(height):
# define value, real(value), imag(value)
zz = row[j]
x = np.real(zz)
y = np.imag(zz)
# define magnitued and argument
magnitude = np.hypot(x, y)
arg = np.arctan2(y, x)
# define lighness
lightness = np.arctan(np.log(np.sqrt(magnitude) + 1)) * (4 / np.pi) - 1
if lightness < 0:
bot = 0
top = 1 + lightness
else:
bot = lightness
top = 1
# define hue
hue = 3 * arg / np.pi
if hue < 0:
hue += 6
# set ihue and use it to define rgb values based on cases
ihue = int(hue)
# case 1
if ihue == 0:
r = top
g = bot + hue * (top - bot)
b = bot
# case 2
elif ihue == 1:
r = bot + (2 - hue) * (top - bot)
g = top
b = bot
# case 3
elif ihue == 2:
r = bot
g = top
b = bot + (hue - 2) * (top - bot)
# case 4
elif ihue == 3:
r = bot
g = bot + (4 - hue) * (top - bot)
b = top
# case 5
elif ihue == 4:
r = bot + (hue - 4) * (top - bot)
g = bot
b = top
# case 6
else:
r = top
g = bot
b = bot + (6 - hue) * (top - bot)
# set rgb array values
rgb[i, j, 0] = r
rgb[i, j, 1] = g
rgb[i, j, 2] = b
return rgb