Related
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()
Error Message
File "/Applications/Spyder.app/Contents/Resources/lib/python3.9/matplotlib/axes/_axes.py", line 5575, in _pcolorargs
raise TypeError('Dimensions of C %s are incompatible with'
TypeError: Dimensions of C (101, 101) are incompatible with X (101) and/or Y (101); see help(pcolormesh)
Problem
https://www.chegg.com/homework-help/questions-and-answers/problem-1-compute-steady-state-concentration-umol-um3-distribution-channel-shown-figure--p-q96105354
Code
# Import the required modules
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits import mplot3d
# Constants
D = 0.5 # (μm^2)/s
k = 0.1 # 1/s
Lx = 30 # μm
Ly = 10 # μm
Ca = 40 # Initial concentration (μmol/μm^3)
Cb = 100 # Final concentration (μmol/μm^3)
Cguess = 70 # Guess concentration (μmol/μm^3)
# Iteration parameters
maxit = 2000
tol = 0.0001 # Relative tolerance
merr = 1e5
lam = 1.4 # Parameter for convergence rate
# Setup grid
npts = 100
dx = Lx/npts;
xx = np.linspace(0,Lx,num=npts+1,endpoint=True);
nx = npts+1;
dy = Ly/npts;
yy = np.linspace(0,Ly,num=npts+1,endpoint=True);
ny = npts+1;
# Step 1 - Initial Guesses
M = np.ones((ny,nx)); # set matrix size
M = Cguess*M; # initial guess for Ca in all nodes
# Step 2 - Apply Boundary Conditions
M[0,:] = Ca
M[:,0] = Cb
# Step 3 and 4 - Apply LDE, walking over nodes
cc = 0 # Counter
a = 1+k*dx*dx/(4*D)
while merr>tol:
Mold = np.copy(M) # Save current values to old
M[-1,-1] = (2*M[-2,-1]+2*M[-1,-2])/(4*a)
for j in range(1,ny-1):
for i in range(1,nx-1):
M[j,i] = (M[j,i-1]+M[j,i+1]+M[j-1,i]+M[j+1,i])/(4*a);
for j in range(1,ny-1):
M[j,-1] = (2*M[j,-2]+M[j-1,-1]+M[j+1,-1])/(4*a);
for i in range(1,nx-1):
M[-1,i] = (M[-1,i-1]+M[-1,i+1]+2*M[-2,i])/(4*a);
for i in range(1,68):
M[0,i] = (M[0,i-1]+M[0,i+1]+2*M[1,i])/(4*a);
M = lam*M+(1-lam)*Mold # Adjust for convergence rate
cc = cc+1
ea = np.abs((M-Mold)/M)
merr = np.max(ea)
if cc>maxit: break
avg = np.mean(M[:,0])
# Plot 2D
X,Y = np.meshgrid(xx,yy)
ax = plt.axes(projection="3d")
ax.plot_surface(X,Y,M,edgecolor="none")
ax.view_init(elev=40,azim=-35)
plt.show()
# Plot color mesh
p = plt.pcolormesh(xx, yy, M, cmap="RdBu", shading="flat", vmin=0, vmax=100)
ct = plt.contour(X, Y, M, cmap="gray", levels=10, vmin=0, vmax=100)
c = plt.colorbar(p)
plt.xlabel("Lx (μm)")
plt.ylabel("Ly (μm)")
c.set_label("Concentration (μmol/μm3)")
plt.show()
# Print Results
print("Converged in %d iterations" % cc)
print("Max error is %f" % merr)
print("Mean concentration along central axis = %f μmol/μm^3" % avg)
I have the following code. I know it's long and complex, however it takes 1.5 mins on my laptop to run. I would greatly appreciate any help towards finding the problem causing the error at the end - the plotting part.I didn't find anything on Google related to this error message:
TypeError: unsupported operand type(s) for : 'QuadMesh' and 'float'
from scipy import interpolate
from scipy.fft import fft, ifft
from scipy.constants import c, epsilon_0
import numpy as np, math
from matplotlib import pyplot as plt
lambda_0 = 800 * 10**(-9)
omega_0 = 2*np.pi*c / lambda_0
delta_lambda = 50.0 * 10**(-9)
delta_Tau = (1.47 * 10**(-3) * (lambda_0*10**9)**2 / (delta_lambda*10**9)) * 10**(-15) #
delta_omega_PFT = (4*np.log(2) / delta_Tau) # equivalent (equal) to: ((4*ln(2)) / (2*pi)) * (2*pi/delta_Tau) = 0.441 * (2*pi/delta_Tau)
F = 2.0 # focal length in meters, as from Liu 2017
def G(omegas):
# return ( np.sqrt(np.pi/(2*np.log(2))) * tau * np.exp( -(tau**2/(8*np.log(2))) * (omegas-omega_0)**2 ) ) # why is this here?
return np.exp( -(delta_Tau**2/(8*np.log(2))) * (omegas-omega_0)**2 )
xsi = 0.1 * (1.0 * 10**(-15)) / (1.0 * 10**(-3))
def phase_c(xs, omegas):
return ( xsi * np.reshape(xs, (xs.shape[0], 1)) * np.reshape((omegas-omega_0), (1, omegas.shape[0])) )
E0 = np.sqrt( (0.2*10**4 * 8 * np.sqrt(np.log(2))) / (delta_Tau*np.sqrt(np.pi)*c*epsilon_0/2.0) ) * np.sqrt(2*np.pi*np.log(2)) / delta_omega_PFT
def f(xi, omega): # the prefactors from Eq. (5) of Li et al. (2017) (the ones pre-multiplying the Fraunhoffer integral)
first = omega * np.exp(1j * (omega/c) * F) / (1j * 2*np.pi*c*F) # only function of omega. first is shape (omega.shape[0], )
omega = np.reshape(omega, (1, omega.shape[0]))
xi = np.reshape(xi, (xi.shape[0], 1))
second = np.exp(1j * (omega/c) * xi**2 / (2*F)) # second is shape (xi.shape[0], omega.shape[0]).
return (first * second) # returned is shape (xi.shape[0], omega.shape[0])
x0 = 0.0
delta_x = 196.0 # obtained from N=10, N_x=8*10^3, xi_max=10um, F=2m
xmin_PFT = x0 - delta_x #
xmax_PFT = x0 + delta_x #
num_xs_PFT = 8 * 10**3
xs_PFT = np.linspace(xmin_PFT, xmax_PFT, num_xs_PFT)
sampling_spacing_xs_PFT = np.true_divide( (xmax_PFT-xmin_PFT), num_xs_PFT)
num_omegas_focus = 5 * 10**2
maximum_time = 100.0 * 10**(-15)
N = math.ceil( (np.pi*num_omegas_focus)/(2*delta_omega_PFT*maximum_time) ) - 1
omega_max_focus = omega_0 + N*delta_omega_PFT
omega_min_focus = omega_0 - N*delta_omega_PFT
omegas_focus = np.linspace(omega_min_focus, omega_max_focus, num_omegas_focus) # shape (num_omegas_focus, )
sampling_spacing_omegas_focus = np.true_divide((omega_max_focus-omega_min_focus) , num_omegas_focus)
Es_x_omega = np.multiply( (E0 * G(omegas_focus)) ,
(np.exp(1j*phase_c(xs_PFT, omegas_focus))) # phase_c uses xsi, the PFT coefficient
)
# Es_x_omega holds across columns (vertically downwards) the x-dependence and across rows (horizontally) the omega-dependence
# Es_x_omega is shape (num_xs_PFT, num_omegas_focus)
Bprime_data_real = np.empty((Es_x_omega.shape[0], Es_x_omega.shape[1])) # this can be rewritten in a more Pythonic way
Bprime_data_imag = np.empty((Es_x_omega.shape[0], Es_x_omega.shape[1]))
for i in range(Es_x_omega.shape[1]): # for all the columns (all omegas)
# Perform FFT wrt x (so go from x to Kappa (a scaled spatial frequency))
intermediate = fft(Es_x_omega[:, i])
Bprime_data_real[:, i] = np.real(intermediate) * sampling_spacing_xs_PFT # multiplication by \Delta, see my docu above
Bprime_data_imag[:, i] = np.imag(intermediate) * sampling_spacing_xs_PFT # multiplication by \Delta, see my docu above
if i % 10000 == 0:
print("We have done fft number {}".format(i) + " out of {}".format(Es_x_omega.shape[1]) + "ffts")
# Bprime is function of (Kappa, omega): across rows the omega dependence, across columns the Kappa dependence.
# Get the Kappas:
returned_freqs = np.fft.fftfreq(num_xs_PFT, sampling_spacing_xs_PFT) # shape (num_xs_PFT, )
Kappas_ugly = 2*np.pi * returned_freqs # shape (num_xs_PFT, ), but unordered in terms of the magnitude of the values! see https://numpy.org/doc/stable/reference/generated/numpy.fft.fftfreq.html
Kappas_pretty = 2*np.pi * np.fft.fftshift(returned_freqs)
indices = (Kappas_ugly == Kappas_pretty[:, None]).argmax(1) # shape (num_xs_PFT, )
indices = indices.reshape((indices.shape[0], 1)) # needed for adapting Dani Mesejo's answer: he reordered based on 1D slices laid horizontally, here I reorder based on 1D slices laid vertically.
# see my notebook for visuals, 22-23 Nov 2021
hold_real = Bprime_data_real.shape[1]
hold_imag = Bprime_data_imag.shape[1]
Bprime_data_real_pretty = np.take_along_axis(Bprime_data_real, np.tile(indices, (1, hold_real)), axis=0) # adapted from Dani Mesejo's answer
Bprime_data_imag_pretty = np.take_along_axis(Bprime_data_imag, np.tile(indices, (1, hold_imag)), axis=0) # adapted from Dani Mesejo's answer
print(Bprime_data_real_pretty.shape) # shape (num_xs_PFT, num_omegas_focus), similarly for Bprime_data_imag_pretty
Bprime_real = interpolate.RectBivariateSpline(Kappas_pretty, omegas_focus, Bprime_data_real_pretty) # this CTOR creates an object faster (which can also be queried faster)
Bprime_imag = interpolate.RectBivariateSpline(Kappas_pretty, omegas_focus, Bprime_data_imag_pretty) # than interpolate.interp2d() does.
print("We have the interpolators!")
# Prepare for the aim: plot E versus time (horizontal axis) and xi (vertical axis).
xi_min = -5.0 * 10**(-6) # um
xi_max = 5.0 * 10**(-6) # um
num_xis = 5000
xis = np.linspace(xi_min, xi_max, num_xis)
print("We are preparing now!")
Es_Kappa_omega_without_prefactor = np.empty((xis.shape[0], omegas_focus.shape[0]), dtype=complex)
for j in range(Es_Kappa_omega_without_prefactor.shape[0]): # for each row
for i in range(Es_Kappa_omega_without_prefactor.shape[1]): # for each column
Es_Kappa_omega_without_prefactor[j, i] = Bprime_real(omegas_focus[i]*xis[j] /(c*F), omegas_focus[i]) + 1j*Bprime_imag(omegas_focus[i]*xis[j] /(c*F), omegas_focus[i])
if ((i + j*Es_Kappa_omega_without_prefactor.shape[1]) % 30000 == 0):
print("We have done iter number {}".format(i + j*Es_Kappa_omega_without_prefactor.shape[1])
+ " out of {}".format(Es_Kappa_omega_without_prefactor.shape[0] * Es_Kappa_omega_without_prefactor.shape[1]) + " iterations in querying the interpolators")
Es_Kappa_omega = np.multiply( f(xis, omegas_focus), # f(xis, omegas_focus) is shape (xis.shape[0], omegas_focus.shape[0])
Es_Kappa_omega_without_prefactor # Es_Kappa_omega_without_prefactor is shape (xis.shape[0], omegas_focus.shape[0])
) # the obtained variable is shape (xis.shape[0], omegas_focus.shape[0])
# Do IFT of Es_Kappa_omega w.r.t. omega to go from FD (omega) to TD (time t).
Es_Kappa_time = np.empty_like(Es_Kappa_omega, dtype=complex) # shape (xis.shape[0], omegas_focus.shape[0])
# Along columns (so vertically) the xi dependence, along rows (horizontally), the omega dependence
for i in range(Es_Kappa_omega.shape[0]): # for each row (for each xi)
Es_Kappa_time[i, :] = ifft(Es_Kappa_omega[i, :]) * (sampling_spacing_omegas_focus/(2*np.pi)) * num_omegas_focus # 1st multiplication is by Delta, 2nd multiplication is by N
if i % 10000 == 0:
print("We have done ifft number {}".format(i) + " out of a total of {}".format(Es_Kappa_omega.shape[0]) + " iffts")
returned_times_ugly = np.fft.fftfreq(num_omegas_focus, d=(sampling_spacing_omegas_focus/(2*np.pi))) # shape (num_omegas_focus, )
returned_times_pretty = np.fft.fftshift(returned_times_ugly) # order the returned "frequencies" (here "frequencies" = times because it's IFT (so from FD to TD))
indices = (returned_times_ugly == returned_times_pretty[:, None]).argmax(1)
Es_Kappa_time = np.take_along_axis(Es_Kappa_time, np.tile(indices, (Es_Kappa_time.shape[0], 1)), axis=1) # this is purely Dani Mesejo's answer
returned_times_pretty_mesh, xis_mesh = np.meshgrid(returned_times_pretty, xis)
fig, ax = plt.subplots()
c = ax.pcolormesh(returned_times_pretty_mesh, xis_mesh, np.real(Es_Kappa_time), cmap='viridis')
fig.colorbar(c, ax=ax, label=r'$[V/m]$')
ax.set_xlabel("t [s]")
ax.set_ylabel("xi [m]")
plt.show()
fig, ax = plt.subplots()
ax.imshow(np.multiply(np.square(np.real(Es_Kappa_time)), (c*epsilon_0)), cmap='viridis')
ax.set_xlabel("t [s]")
ax.set_ylabel("xi [m]")
plt.show()
I have tried many forms to be introduced in the plotting part. It fails with:
c = ax.pcolormesh(returned_times_pretty_mesh, xis_mesh, (c*epsilon_0)*np.real(Es_Kappa_time)**2, cmap='viridis')
fig.colorbar(c, ax=ax, label=r'$[V/m]$')
Fails with:
160 fig, ax = plt.subplots()
--> 161 ax.imshow(np.multiply(np.real(Es_Kappa_time)**2, (c*epsilon_0)), cmap='viridis')
I ran out of ideas of what I might introduce there.
Thank you!
It looks like you're reassigning c to the return value from ax.pcolormesh(...) after you import c from scipy.constants.
I have implemented a simple randomized, population-based optimization method - Grey Wolf optimizer. I am having some trouble with properly capturing the Matplotlib plots at each iteration using the camera package.
I am running GWO for the objective function f(x,y) = x^2 + y^2. I can only see the candidate solutions converging to the minima, but the contour plot doesn't show up.
Do you have any suggestions, how can I display the contour plot in the background?
GWO Algorithm implementation
%matplotlib notebook
import matplotlib.pyplot as plt
import numpy as np
from celluloid import Camera
import ffmpeg
import pillow
# X : Position vector of the initial population
# n : Initial population size
def gwo(f,max_iterations,LB,UB):
fig = plt.figure()
camera = Camera(fig)
def random_population_uniform(m,a,b):
dims = len(a)
x = [list(a + np.multiply(np.random.rand(dims),b - a)) for i in range(m)]
return np.array(x)
def search_agent_fitness(fitness):
alpha = 0
if fitness[1] < fitness[alpha]:
alpha, beta = 1, alpha
else:
beta = 1
if fitness[2] > fitness[alpha] and fitness[2] < fitness[beta]:
beta, delta = 2, beta
elif fitness[2] < fitness[alpha]:
alpha,beta,delta = 2,alpha,beta
else:
delta = 2
for i in range(3,len(fitness)):
if fitness[i] <= fitness[alpha]:
alpha, beta,delta = i, alpha, beta
elif fitness[i] > fitness[alpha] and fitness[i]<= fitness[beta]:
beta,delta = i,beta
elif fitness[i] > fitness[beta] and fitness[i]<= fitness[delta]:
delta = i
return alpha, beta, delta
def plot_search_agent_positions(f,X,alpha,beta,delta,a,b):
# Plot the positions of search agents
x = X[:,0]
y = X[:,1]
s = plt.scatter(x,y,c='gray',zorder=1)
s = plt.scatter(x[alpha],y[alpha],c='red',zorder=1)
s = plt.scatter(x[beta],y[beta],c='blue',zorder=1)
s = plt.scatter(x[delta],y[delta],c='green',zorder=1)
camera.snap()
# Initialize the position of the search agents
X = random_population_uniform(50,np.array(LB),np.array(UB))
n = len(X)
l = 1
# Plot the first image on screen
x = np.linspace(LB[0],LB[1],1000)
y = np.linspace(LB[0],UB[1],1000)
X1,X2 = np.meshgrid(x,y)
Z = f(X1,X2)
cont = plt.contour(X1,X2,Z,20,linewidths=0.75)
while (l < max_iterations):
# Take the x,y coordinates of the initial population
x = X[:,0]
y = X[:,1]
# Calculate the objective function for each search agent
fitness = list(map(f,x,y))
# Update alpha, beta and delta
alpha,beta,delta = search_agent_fitness(fitness)
# Plot search agent positions
plot_search_agent_positions(f,X,alpha,beta,delta,LB,UB)
# a decreases linearly from 2 to 0
a = 2 - l *(2 / max_iterations)
# Update the position of search agents including the Omegas
for i in range(n):
x_prey = X[alpha]
r1 = np.random.rand(2) #r1 is a random vector in [0,1] x [0,1]
r2 = np.random.rand(2) #r2 is a random vector in [0,1] x [0,1]
A1 = 2*a*r1 - a
C1 = 2*r2
D_alpha = np.abs(C1 * x_prey - X[i])
X_1 = x_prey - A1*D_alpha
x_prey = X[beta]
r1 = np.random.rand(2)
r2 = np.random.rand(2)
A2 = 2*a*r1 - a
C2 = 2*r2
D_beta = np.abs(C2 * x_prey - X[i])
X_2 = x_prey - A2*D_beta
x_prey = X[delta]
r1 = np.random.rand(2)
r2 = np.random.rand(2)
A3 = 2*a*r1 - a
C3 = 2*r2
D_delta = np.abs(C3 * x_prey - X[i])
X_3 = x_prey - A3*D_delta
X[i] = (X_1 + X_2 + X_3)/3
l = l + 1
return X[alpha],camera
Function call
# define the objective function
def f(x,y):
return x**2 + y**2
minimizer,camera = gwo(f,7,[-10,-10],[10,10])
animation = camera.animate(interval = 1000, repeat = True,
repeat_delay = 500)
Is it possible that the line x = np.linspace(LB[0],LB[1],1000) should be x = np.linspace(LB[0],UB[1],1000) instead? With your current definition of x, x is an array only filled with the value -10 which means that you are unlikely to find a contour.
Another thing that you might want to do is to move the cont = plt.contour(X1,X2,Z,20,linewidths=0.75) line inside of your plot_search_agent_positions function to ensure that the contour is plotted at each iteration of the animation.
Once you make those changes, the code looks like that:
import matplotlib.pyplot as plt
import numpy as np
from celluloid import Camera
import ffmpeg
import PIL
from matplotlib import animation, rc
from IPython.display import HTML, Image # For GIF
from scipy.interpolate import griddata
rc('animation', html='html5')
# X : Position vector of the initial population
# n : Initial population size
def gwo(f,max_iterations,LB,UB):
fig = plt.figure()
fig.gca(aspect='equal')
camera = Camera(fig)
def random_population_uniform(m,a,b):
dims = len(a)
x = [list(a + np.multiply(np.random.rand(dims),b - a)) for i in range(m)]
return np.array(x)
def search_agent_fitness(fitness):
alpha = 0
if fitness[1] < fitness[alpha]:
alpha, beta = 1, alpha
else:
beta = 1
if fitness[2] > fitness[alpha] and fitness[2] < fitness[beta]:
beta, delta = 2, beta
elif fitness[2] < fitness[alpha]:
alpha,beta,delta = 2,alpha,beta
else:
delta = 2
for i in range(3,len(fitness)):
if fitness[i] <= fitness[alpha]:
alpha, beta,delta = i, alpha, beta
elif fitness[i] > fitness[alpha] and fitness[i]<= fitness[beta]:
beta,delta = i,beta
elif fitness[i] > fitness[beta] and fitness[i]<= fitness[delta]:
delta = i
return alpha, beta, delta
def plot_search_agent_positions(f,X,alpha,beta,delta,a,b,X1,X2,Z):
# Plot the positions of search agents
x = X[:,0]
y = X[:,1]
s = plt.scatter(x,y,c='gray',zorder=1)
s = plt.scatter(x[alpha],y[alpha],c='red',zorder=1)
s = plt.scatter(x[beta],y[beta],c='blue',zorder=1)
s = plt.scatter(x[delta],y[delta],c='green',zorder=1)
Z=f(X1,X2)
cont=plt.contour(X1,X2,Z,levels=20,colors='k',norm=True)
plt.clabel(cont, cont.levels, inline=True, fontsize=10)
camera.snap()
# Initialize the position of the search agents
X = random_population_uniform(50,np.array(LB),np.array(UB))
n = len(X)
l = 1
# Plot the first image on screen
x = np.linspace(LB[0],UB[1],1000)
y = np.linspace(LB[0],UB[1],1000)
X1,X2 = np.meshgrid(x,y)
Z=f(X1,X2)
while (l < max_iterations):
# Take the x,y coordinates of the initial population
x = X[:,0]
y = X[:,1]
# Calculate the objective function for each search agent
fitness = list(map(f,x,y))
# Update alpha, beta and delta
alpha,beta,delta = search_agent_fitness(fitness)
# Plot search agent positions
plot_search_agent_positions(f,X,alpha,beta,delta,LB,UB,X1,X2,Z)
# a decreases linearly from 2 to 0
a = 2 - l *(2 / max_iterations)
# Update the position of search agents including the Omegas
for i in range(n):
x_prey = X[alpha]
r1 = np.random.rand(2) #r1 is a random vector in [0,1] x [0,1]
r2 = np.random.rand(2) #r2 is a random vector in [0,1] x [0,1]
A1 = 2*a*r1 - a
C1 = 2*r2
D_alpha = np.abs(C1 * x_prey - X[i])
X_1 = x_prey - A1*D_alpha
x_prey = X[beta]
r1 = np.random.rand(2)
r2 = np.random.rand(2)
A2 = 2*a*r1 - a
C2 = 2*r2
D_beta = np.abs(C2 * x_prey - X[i])
X_2 = x_prey - A2*D_beta
x_prey = X[delta]
r1 = np.random.rand(2)
r2 = np.random.rand(2)
A3 = 2*a*r1 - a
C3 = 2*r2
D_delta = np.abs(C3 * x_prey - X[i])
X_3 = x_prey - A3*D_delta
X[i] = (X_1 + X_2 + X_3)/3
l = l + 1
return X[alpha],camera
# define the objective function
def f(x,y):
return x**2 + y**2
minimizer,camera = gwo(f,7,[-10,-10],[10,10])
animation = camera.animate(interval = 1000, repeat = True,repeat_delay = 500)
And the output gives:
I have following python code, and would like to:
Plot the same function in 1 (only one) figure with many different (lets say 4) 'v0' and 'theta' values, each trajectory in a different color.
Make 4 plots in 4 different figures, so that it looks like a square with 4 plots of 4 different 'v0' and 'theta' values
Make a widget to vary the v0 and theta values as the user wants with the mouse.
import numpy as np
import scipy.integrate as integrate
import matplotlib.pyplot as plt
%matplotlib inline
theta = 45.
theta = theta * np.pi/180.
v0 = 20.0
g = 9.81
R = 0.035
m = 0.057
rho = 1.2041
C = 0.5
k = (0.5*np.pi*R**2*C*rho)/m
x0=0
y0=10
vx0 = v0*np.sin(theta)
vy0 =
v0*np.cos(theta)
print(vx0)
print(vy0)
def f_func(X_vek,time):
f = np.zeros(4)
f[0] = X_vek[2]
f[1] = X_vek[3]
f[2] = - k*(f[0]**2 + f[1]**2)**(0.5)*f[0]
f[3] = -g - k*(f[0]**2 + f[1]**2)**(0.5)*f[1]
return f
X0 = [ x0, y0, vx0, vy0]
t0 = 0. tf = 10
tau = 0.05
t = np.arange(t0,tf,tau)
X = integrate.odeint(f_func,X0,t)
x = X[:,0]
y = X[:,1]
vx = X[:,2]
vy = X[:,3]
mask = y >= 0
plt.scatter(x[mask],y[mask])
plt.scatter(x[mask],y[mask])
plt.xlabel('x') plt.ylabel('y') plt.show()
I could do point 1 and 2 of my question with changing the values after plotting, then calculate vx0 and vy0 again and then call the integrate function and finally plot again, but that's kinda weird and not clean. Is there any better way to do that? like an array of different v0 and theta values or something?
Thanks!
Make your code as a function:
def func(theta=45, v0=20):
theta = theta * np.pi/180.
g = 9.81
R = 0.035
m = 0.057
rho = 1.2041
C = 0.5
k = (0.5*np.pi*R**2*C*rho)/m
x0=0
y0=10
vx0 = v0*np.sin(theta)
vy0 = v0*np.cos(theta)
def f_func(X_vek,time):
f0, f1 = X_vek[2:4].tolist()
f2 = - k*(f0**2 + f1**2)**(0.5)*f0
f3 = -g - k*(f0**2 + f1**2)**(0.5)*f1
return [f0, f1, f2, f3]
X0 = [ x0, y0, vx0, vy0]
t0 = 0.
tf = 10
tau = 0.05
t = np.arange(t0,tf,tau)
X = integrate.odeint(f_func,X0,t)
x = X[:,0]
y = X[:,1]
vx = X[:,2]
vy = X[:,3]
mask = y >= 0
return x[mask], y[mask]
then you can plot it with different parameters:
plt.plot(*func())
plt.plot(*func(theta=30))
plt.xlabel('x')
plt.ylabel('y')
plt.show()
I suggest you use Holoviews to make dynamic graph:
import holoviews as hv
hv.extension("bokeh")
hv.DynamicMap(
lambda theta, v0:hv.Curve(func(theta, v0)).redim.range(x=(0, 50), y=(0, 50)),
kdims=[hv.Dimension("theta", range=(0, 80), default=40),
hv.Dimension("v0", range=(1, 40), default=20)])
Here is the result: