We Keep Coding

Plotting a 3-dimensional superball shape in matplotlib

I'm trying to plot a 3D superball in python matplotlib, where a superball is defined as a general mathematical shape that can be used to describe rounded cubes using a shape parameter p, where for p = 1 the shape is equal to that of a sphere.
This paper claims that the superball is defined by using modified spherical coordinates with:
x = r*cos(u)**1/p * sin(v)**1/p
y = r*cos(u)**1/p * sin(v)**1/p
z = r*cos(v)**1/p
with u = phi and v = theta.
I managed to get the code running, at least for p = 1 which generates a sphere - exactly as it should do:
import matplotlib.pyplot as plt
import numpy as np
fig = plt.figure()
ax = fig.add_subplot(projection='3d')
r, p = 1, 1
# Make data
u = np.linspace(0, 2 * np.pi, 100)
v = np.linspace(0, np.pi, 100)
u, v = np.meshgrid(u, v)
x = r * np.cos(u)**(1/p) * np.sin(v)**(1/p)
y = r * np.sin(u)**(1/p) * np.sin(v)**(1/p)
z = r * np.cos(v)**(1/p)
# Plot the surface
ax.plot_surface(x, y, z)
plt.show()
This is a 3D plot of the code above for p = 1.
However, as I put in any other value for p, e.g. 2, it's giving me only a partial shape, while it should actually give me a full superball.
This is a 3D plot of the code above for p = 2.
I believe the fix is more of mathematical nature, but how can this be fixed?

When plotting a regular sphere, we transform positive and negative coordinates differently:
Positives: x**0.5
Negatives: -1 * abs(x)**0.5
For the superball variants, apply the same logic using np.sign and np.abs:
power = lambda base, exp: np.sign(base) * np.abs(base)**exp
x = r * power(np.cos(u), 1/p) * power(np.sin(v), 1/p)
y = r * power(np.sin(u), 1/p) * power(np.sin(v), 1/p)
z = r * power(np.cos(v), 1/p)
Full example for p = 4:
import matplotlib.pyplot as plt
import numpy as np
fig, ax = plt.subplots(subplot_kw={'projection': '3d'})
r, p = 1, 4
# Make the data
u = np.linspace(0, 2 * np.pi)
v = np.linspace(0, np.pi)
u, v = np.meshgrid(u, v)
# Transform the coordinates
# Positives: base**exp
# Negatives: -abs(base)**exp
power = lambda base, exp: np.sign(base) * np.abs(base)**exp
x = r * power(np.cos(u), 1/p) * power(np.sin(v), 1/p)
y = r * power(np.sin(u), 1/p) * power(np.sin(v), 1/p)
z = r * power(np.cos(v), 1/p)
# Plot the surface
ax.plot_surface(x, y, z)
plt.show()

How to stop the roots of a cubic from becoming mixed up when plotting the 3 roots as contour plots?

I have a function which determines the roots of a complex cubic. I am solving the cubic for a variety of k0 and k1 values and showing the solutions as contour plots. Since the cubic has three roots, I produce 3 contour plots for the real parts and 3 for the imaginary parts. However, sometimes you can clearly see that sections of the contour plots for one root really should be swapped with a different contour plot - all the contours should be continuous. I have tried various "sorting methods" which you can see, but none of them fully fix it. What can I do so that the roots don't get mixed up resulting in non-continuous contours.
import numpy as np
import matplotlib.pyplot as plt
# Constants
Ra = 2e4
Pr = 0.1
Omega = 1e5
zeta = 1e-4
deltaN = 0.05
L = 55
def polynomial(k):
m = 1
delta_k = m**2 * np.pi**2 + k[0]**2
a_3 = delta_k
a_2 = 1j*(Ra * Pr * delta_k * k[0])/Omega + (Pr + zeta + 1)*delta_k**2
a_1 = 1j*(Ra * Pr * delta_k**2 * k[0] * (Pr + zeta)/Omega) + k[1] * Pr * zeta * (delta_k**2/L**2 + delta_k) - deltaN * Ra * Pr * k[0]**2 + (Pr * zeta + Pr + zeta) * delta_k**3
a_0 = 1j*(Pr * zeta * k[0] * (Ra * Pr * delta_k**3/Omega + k[1] * Omega * deltaN * delta_k / L**2)) + Pr * zeta * (k[1] * (Pr * delta_k**3 / L**2 + delta_k**2) - deltaN * Ra * delta_k * k[0]**2 + delta_k**4)
x_K = np.roots([a_3, a_2, a_1, a_0])
# x_K = np.sort_complex(x_K)
x_K = sorted(x_K, key=lambda x: x.imag)
# x_K = sorted(x_K, key=lambda x: x.real)
# if x_K[2].imag >= 0:
# x_K[-1], x_K[-2] = x_K[-2], x_K[-1]
# if x_K[0].imag >= x_K[2].imag:
# x_K[0], x_K[-1] = x_K[-1], x_K[0]
if x_K[0].real >= x_K[1].real:
x_K[0], x_K[1] = x_K[1], x_K[0]
# if x_K[1].real >= x_K[2].real:
# x_K[1], x_K[2] = x_K[2], x_K[1]
return x_K
# Create arrays of k[0] and k[1] values for contour plot
k0, k1 = np.linspace(0, 5, 100), np.linspace(0, 5e2, 100)
K0, K1 = np.meshgrid(k0, k1)
# Get roots for each pair of k[0], k[1] value
roots = np.array([polynomial([K0[i, j], K1[i, j]]) for i in range(100) for j in range(100)], dtype=complex)
ky_max = []
Qz_max = []
# Plot real and imaginary parts of roots separately in one figure
fig, axs = plt.subplots(2, 3, figsize=(13.6, 7.6), constrained_layout=True)
axs = axs.ravel()
for i in range(3):
cnt = axs[i].contourf(K0, K1, roots[:, i].real.reshape(K0.shape), levels=20, cmap='coolwarm')
axs[i].set_title(f'Real part of root {i+1}')
axs[i].set_xlabel('$k_y$')
axs[i].set_ylabel('$Q_z$')
# axs[i].set_yscale('log')
fig.colorbar(cnt, ax=axs[i])
cnt = axs[i+3].contourf(K0, K1, roots[:, i].imag.reshape(K0.shape), levels=20, cmap='coolwarm')
axs[i+3].set_title(f'Imaginary part of root {i+1}')
axs[i+3].set_xlabel('$k_y$')
axs[i+3].set_ylabel('$Q_z$')
# axs[i+3].set_yscale('log')
cbar1 = fig.colorbar(cnt, ax=axs[i+3])
cbar1.formatter.set_powerlimits((0, 0))
max_val = np.max(roots[:, i].real)
print(f'Maximum value for real part of root {i+1} is: {max_val}')
max_val = np.max(roots[:, i].real)
max_index = np.argmax(roots[:, i].real)
k0_max, k1_max = K0.flatten()[max_index], K1.flatten()[max_index]
axs[i].scatter(k0_max, k1_max, s=150, color='yellow', marker='x', label=f'Max value {max_val:.4f}')
axs[i].legend(loc=0)
ky_max.append(K0.flatten()[max_index])
Qz_max.append(K1.flatten()[max_index])
print(f'k_y for root {i+1} is: {k0_max}')
print(f'Q_z for root {i+1} is: {k1_max}')
for axis in ['top','bottom','left','right']:
axs[2].spines[axis].set_linewidth(3)
axs[2].spines[axis].set_color("green")
axs[5].spines[axis].set_linewidth(3)
axs[5].spines[axis].set_color("green")
# Create a caption
caption = f'Contour plot showing the real and imaginary components of the roots of the cubic for a range of $k_y$ and $Q_z$ values. Where the other variables are given by: Ra$^* = $ {Ra:.1e}, $\Delta N =$ {deltaN}, Pr = {Pr:.1e}, $\zeta =$ {zeta:.1e}, $\Omega =$ {Omega:.1e}, $L$ = {L}.'
# Create a file name
figure_name = f'decay_contour_Ra={Ra:.1e}_Pr={Pr:.1e}_dN={deltaN}'
pdf_file = f'{figure_name}.pdf'
tex_file = f'{figure_name}.tex'
# save the plot as a PDF
plt.savefig(pdf_file)
# create a text file containing the LaTeX code to include the figure
with open(tex_file, 'w') as f:
f.write("\\begin{figure}[h]\n")
f.write("\\centering\n")
f.write("\\includegraphics[width=0.85\linewidth]{"+ pdf_file+"}\n")
f.write("\\caption{"+ caption +"}\n")
f.write("\\end{figure}\n")
fig2, axs2 = plt.subplots(2, 3, figsize=(11, 8), constrained_layout=True)
for idx_1 in range(3):
k1_slice = 0
indices = np.where(K1.flatten() == k1_slice)
root_slice = roots[indices][:,idx_1].real
k1_slice = K0.flatten()[indices]
root_slice = roots[indices][:,idx_1].real
axs2[0][idx_1].plot(k1_slice, root_slice, color = 'red')
k1_slice_imag = K0.flatten()[indices]
root_slice_imag = roots[indices][:,idx_1].imag
axs2[1][idx_1].plot(k1_slice, root_slice_imag, color = 'red')
axs2[1][idx_1].set_xlabel('$k_y$')
axs2[0][0].set_ylabel('Re$(s)$')
axs2[1][0].set_ylabel('Im$(s)$')
for idx_1 in range(3):
axs2[0][idx_1].plot(k0, -zeta*(np.pi**2 + k0**2), 'x', markevery=10, color = 'black')
# Create a caption
caption = f'Profiles at the $k_y$ at $Q_z = 0$ showing the real and imaginary components of the roots of the cubic for a range of $k_y$ and $Q_z$ values. Where the other variables are given by: Ra$^* = $ {Ra:.1e}, $\Delta N =$ {deltaN}, Pr = {Pr:.1e}, $\zeta =$ {zeta:.1e}, $\Omega =$ {Omega:.1e}, $L$ = {L}.'
# Create a file name
figure_name = f'decay_profiles_Ra={Ra:.1e}_Pr={Pr:.1e}_dN={deltaN}'
pdf_file = f'{figure_name}.pdf'
tex_file = f'{figure_name}.tex'
# create a text file containing the LaTeX code to include the figure
with open(tex_file, 'w') as f:
f.write("\\begin{figure}[h]\n")
f.write("\\centering\n")
f.write("\\includegraphics[width=0.99\linewidth]{"+ pdf_file+"}\n")
f.write("\\caption{"+ caption +"}\n")
f.write("\\end{figure}\n")
for axis in ['top','bottom','left','right']:
axs2[0][2].spines[axis].set_linewidth(3)
axs2[0][2].spines[axis].set_color("green")
axs2[1][2].spines[axis].set_linewidth(3)
axs2[1][2].spines[axis].set_color("green")
# save the plot as a PDF
plt.savefig(pdf_file)
plt.show()
I've tried np.sort, np.sorted, flapping the roots using if statements etc, nothing works 100%

For two successive polynomials P and Q, I suggest simply solving the assignment problem to pair each root of P to the closest root of Q.
You can use scipy's linear_sum_assignment along with distance_matrix to find the best assignment of P's roots with Q's roots.
import numpy as np
from scipy.optimize import linear_sum_assignment
from scipy.spatial import distance_matrix
import matplotlib.pyplot as plt
def get_root_sequence(sequence_of_polynomials):
r0 = np.roots(sequence_of_polynomials[0])
roots = [r0]
for P in sequence_of_polynomials[1:]:
r1 = np.roots(P)
_, idx = linear_sum_assignment(distance_matrix(r0.reshape(3, 1), r1.reshape(3,1)))
r1 = r1[idx]
roots.append(r1)
r0 = r1
return np.array(roots)
sequence_of_polynomials = np.linspace((1,0,0,-1), (1,-7-2j,15+9j,-10-10j), 100)
roots = get_root_sequence(sequence_of_polynomials)
plt.axes().set_aspect('equal')
for i in range(3):
r = roots[:, i]
ordinal = ('first', 'second', 'third')[i]
plt.plot(r.real, r.imag, label=f'{ordinal} root')
for triangle, label in zip((roots[0], roots[-1]), ('x³-1', '(x-2)(x-2-i)(x-3-i)')):
triangle = triangle[[0,1,2,0]]
plt.plot(triangle.real, triangle.imag, label=label)
plt.legend(loc='best')
plt.xlabel('Real part')
plt.ylabel('Imaginary part')
plt.show()

Mayavi - Add clipping planes

Is it possible to add clipping planes to a Mayavi scene?
Here is an example:
import numpy as np
import mayavi.mlab as mlab
mlab.figure()
u, v = np.mgrid[0:2*np.pi:150j, 0:np.pi:150j]
r = 2 + np.sin(7 * u + 5 * v)
x = r * np.cos(u) * np.sin(v)
y = r * np.sin(u) * np.sin(v)
z = r * np.cos(v)
mlab.mesh(x, y, z)
mlab.show()
which produces:
But I would like to add a couple of clipping planes to obtain something similar to this:

Using matplotlib to generate 3D surfaces from 2D matrices

Using a 2d matrix in python, how can I create a 3d surface plot, where columns=x, rows=y and the values are the heights in z?
I can't understand how to creat 3D surface plot using matplotlib.
Maybe it's different from MatLab.
example:
from pylab import *
from mpl_toolkits.mplot3d import Axes3D
def p(eps=0.9, lmd=1, err=10e-3, m=60, n=40):
delta_phi = 2 * np.pi / m
delta_lmd = 2 / n
k = 1
P0 = np.zeros([m + 1, n + 1])
P = np.zeros([m + 1, n + 1])
GAP = 1
while GAP >= err:
k = k + 1
for i in range(0, m):
for j in range(0, n):
if (i == 1) or (j == 1) or (i == m + 1) or (i == n + 1):
P[i,j] = 0
else:
A = (1+eps*np.cos((i+1/2)*delta_phi))**3
B = (1+eps*np.cos((i-1/2)*delta_phi))**3
C = (lmd*delta_phi/delta_lmd)**2 * (1+eps*np.cos((i)*delta_phi))**3
D = C
E = A + B + C + D
F = 3*delta_phi*((1+eps*np.cos((i+1/2)*delta_phi))-(1+eps*np.cos((i-1/2)*delta_phi)))
P[i,j] = (A*P[i+1,j] + B*P[i-1,j] + C*P[i,j+1] + D*P[i,j-1] - F)/E
if P[i,j] < 0:
P[i,j] = 0
S = P.sum() - P0.sum()
T = P.sum()
GAP = S / T
P0 = P.copy()
return P, k
def main():
start = time.time()
eps = 0.9
lmd = 1
err = 10e-8
m = 60
n = 40
P, k = p()
fig = figure()
ax = Axes3D(fig)
X = np.linspace(0, 2*np.pi, m+1)
Y = np.linspace(-1, 1, n+1)
X, Y = np.meshgrid(X, Y)
#Z = P[0:m, 0:n]
#Z = Z.reshape(X.shape)
ax.set_xticks([0, np.pi/2, np.pi, np.pi*1.5, 2*np.pi])
ax.set_yticks([-1, -0.5, 0, 0.5, 1])
ax.plot_surface(X, Y, P)
show()
if __name__ == '__main__':
main()
ValueError: shape mismatch: objects cannot be broadcast to a single
shape
And the pic
pic by matplotlic
And I also use MatLab to generate,the pic:
pic by MatLab

I should think this is a problem of getting the notaton straight. A m*n matrix is a matrix with m rows and n columns. Hence Y should be of length m and X of length n, such that after meshgridding X,Y and P all have shape (m,n).
At this point there would be no need to reshape of reindex and just plotting
ax.plot_surface(X, Y, P)
would give your the desired result.

Let's assume if you have a matrix mat.
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits import mplot3d
h, w = mat.shape
plt.figure(figsize=(16, 8))
ax = plt.axes(projection='3d')
X, Y = np.meshgrid(np.arange(w), np.arange(h))
ax.plot_surface(X, Y, mat, rstride=1, cstride=1, cmap='viridis', edgecolor='none', antialiased=False)

Still having trouble with curve fitting

I already opened a question on this topic, but I wasn't sure, if I should post it there, so I opened a new question here.
I have trouble again when fitting two or more peaks. First problem occurs with a calculated example function.
xg = np.random.uniform(0,1000,500)
mu1 = 200
sigma1 = 20
I1 = -2
mu2 = 800
sigma2 = 20
I2 = -1
yg3 = 0.0001*xg
yg1 = (I1 / (sigma1 * np.sqrt(2 * np.pi))) * np.exp( - (xg - mu1)**2 / (2 * sigma1**2) )
yg2 = (I2 / (sigma2 * np.sqrt(2 * np.pi))) * np.exp( - (xg - mu2)**2 / (2 * sigma2**2) )
yg=yg1+yg2+yg3
plt.figure(0, figsize=(8,8))
plt.plot(xg, yg, 'r.')
I tried two different approaches, I found in the documentation, which are shown below (modified for my data), but both give me wrong fitting data and a messy chaos of graphs (I guess one line per fitting step).
1st attempt:
import numpy as np
from lmfit.models import PseudoVoigtModel, LinearModel, GaussianModel, LorentzianModel
import sys
import matplotlib.pyplot as plt
gauss1 = PseudoVoigtModel(prefix='g1_')
pars.update(gauss1.make_params())
pars['g1_center'].set(200)
pars['g1_sigma'].set(15, min=3)
pars['g1_amplitude'].set(-0.5)
pars['g1_fwhm'].set(20, vary=True)
#pars['g1_fraction'].set(0, vary=True)
gauss2 = PseudoVoigtModel(prefix='g2_')
pars.update(gauss2.make_params())
pars['g2_center'].set(800)
pars['g2_sigma'].set(15)
pars['g2_amplitude'].set(-0.4)
pars['g2_fwhm'].set(20, vary=True)
#pars['g2_fraction'].set(0, vary=True)
mod = gauss1 + gauss2 + LinearModel()
pars.add('intercept', value=0, vary=True)
pars.add('slope', value=0.0001, vary=True)
init = mod.eval(pars, x=xg)
out = mod.fit(yg, pars, x=xg)
print(out.fit_report(min_correl=0.5))
plt.figure(5, figsize=(8,8))
out.plot_fit()
When I include the 'fraction'-parameter, I often get
'NameError: name 'pv1_fraction' is not defined in expr='<_ast.Module object at 0x00000000165E03C8>'.
although it should be defined. I get this Error for real data with this approach, too.
2nd attempt:
import pandas as pd
import matplotlib.pyplot as plt
import numpy as np
import lmfit
def gauss(x, sigma, mu, A):
return A*np.exp(-(x-mu)**2/(2*sigma**2))
def linear(x, m, n):
return m*x + n
peak1 = lmfit.model.Model(gauss, prefix='p1_')
peak2 = lmfit.model.Model(gauss, prefix='p2_')
lin = lmfit.model.Model(linear, prefix='l_')
model = peak1 + lin + peak2
params = model.make_params()
params['p1_mu'] = lmfit.Parameter(value=200, min=100, max=250)
params['p2_mu'] = lmfit.Parameter(value=800, min=100, max=1000)
params['p1_sigma'] = lmfit.Parameter(value=15, min=0.01)
params['p2_sigma'] = lmfit.Parameter(value=20, min=0.01)
params['p1_A'] = lmfit.Parameter(value=-2, min=-3)
params['p2_A'] = lmfit.Parameter(value=-2, min=-3)
params['l_m'] = lmfit.Parameter(value=0)
params['l_n'] = lmfit.Parameter(value=0)
out = model.fit(yg, params, x=xg)
print out.fit_report()
plt.figure(8, figsize=(8,8))
out.plot_fit()
So the result looks like this, in both cases. It seems to plot all fitting attempts, but never solves it correctly. The best fitted parameters are in the range that I gave it.
Anyone knows this type of error? Or has any solutions for this? And does anyone know how to avoid the NameError when calling a model function from lmfit with those approaches?

I have a somewhat tolerable solution for you. Since I don't know how variable your data is, I cannot say that it will work in a general sense but should get you started. If your data is along 0-1000 and has two peaks or dips along a line as you showed, then it should work.
I used the scipy curve_fit and put all of the components of the function together into one function. One can pass starting locations into the curve_fit function. (you can probably do this with the lib you're using but I'm not familiar with it) There is a loop in loop where I vary the mu parameters to find the ones with the lowest squared error. If you are needing to fit your data many times or in some real-time scenario then this is not for you but if you just need to fit some data, launch this code and grab a coffee.
from scipy.optimize import curve_fit
import numpy as np
import matplotlib.pyplot as plt
import pylab
from matplotlib import cm as cm
import time
def my_function_big(x, m, n, #lin vars
sigma1, mu1, I1, #gaussian 1
sigma2, mu2, I2): #gaussian 2
y = m * x + n + (I1 / (sigma1 * np.sqrt(2 * np.pi))) * np.exp( - (x - mu1)**2 / (2 * sigma1**2) ) + (I2 / (sigma2 * np.sqrt(2 * np.pi))) * np.exp( - (x - mu2)**2 / (2 * sigma2**2) )
return y
#make some data
xs = np.random.uniform(0,1000,500)
mu1 = 200
sigma1 = 20
I1 = -2
mu2 = 800
sigma2 = 20
I2 = -1
yg3 = 0.0001 * xs
yg1 = (I1 / (sigma1 * np.sqrt(2 * np.pi))) * np.exp( - (xs - mu1)**2 / (2 * sigma1**2) )
yg2 = (I2 / (sigma2 * np.sqrt(2 * np.pi))) * np.exp( - (xs - mu2)**2 / (2 * sigma2**2) )
ys = yg1 + yg2 + yg3
xs = np.array(xs)
ys = np.array(ys)
#done making data
#start a double loop...very expensive but this is quick and dirty
#it would seem that the regular optimizer has trouble finding the minima so i
#found that having the near proper mu values helped it zero in much better
start = time.time()
serr = []
_x = []
_y = []
for x in np.linspace(0, 1000, 61):
for y in np.linspace(0, 1000, 61):
cfiti = curve_fit(my_function_big, xs, ys, p0=[0, 0, 1, x, 1, 1, y, 1], maxfev=20000000)
serr.append(np.sum((my_function_big(xs, *cfiti[0]) - ys) ** 2))
_x.append(x)
_y.append(y)
serr = np.array(serr)
_x = np.array(_x)
_y = np.array(_y)
print 'done loop in loop fitting'
print 'time: %0.1f' % (time.time() - start)
gridsize=20
plt.subplot(111)
plt.hexbin(_x, _y, C=serr, gridsize=gridsize, cmap=cm.jet, bins=None)
plt.axis([_x.min(), _x.max(), _y.min(), _y.max()])
cb = plt.colorbar()
cb.set_label('SE')
plt.show()
ix = np.argmin(serr.ravel())
mustart1 = _x.ravel()[ix]
mustart2 = _y.ravel()[ix]
print mustart1
print mustart2
cfit = curve_fit(my_function_big, xs, ys, p0=[0, 0, 1, mustart1, 1, 1, mustart2, 1], maxfev=2000000000)
xp = np.linspace(0, 1000, 1001)
plt.figure()
plt.scatter(xs, ys) #plot synthetic dat
plt.plot(xp, my_function_big(xp, *cfit[0]), '-', label='fit function') #plot data evaluated along 0-1000
plt.legend(loc=3, numpoints=1, prop={'size':12})
plt.show()
pylab.close()
Good luck!

In your first attempt:
pars['g1_fraction'].set(0, vary=True)
The fraction must be a value between 0 and 1, but I believe that cannot be zero. Try to put something like 0.000001, and it will work.