We Keep Coding

Problem with electric potential - Matplotlib

What I want to do is to create a simple contour plot from point charges at different coordinates and strengths with matplotlib.
I copied the code from this blog post: https://pythonmatplotlibtips.blogspot.com/2017/12/draw-beautiful-electric-field-lines.html and it works beautifully.
However, the color at the location of the point charges do not seems to scale properly. There is no different in color for point at coordinate (8,5) with a charge of +0.5 and point at coordinate (16,4.5) with a charge of +8.0 for instance. I am not sure what I am doing wrong, as I am a novice in both physic and coding. Is there a way to force the colors to match the scaling of the color bar?the radius scale with the magnitude of the charge but not the color...
import matplotlib.pyplot as plt
import numpy as np
from scipy.integrate import ode as ode
from matplotlib import cm
from itertools import product
class charge:
def __init__(self, q, pos):
self.q=q
self.pos=pos
def E_point_charge(q, a, x, y):
return q*(x-a[0])/((x-a[0])**2+(y-a[1])**2)**(1.5), \
q*(y-a[1])/((x-a[0])**2+(y-a[1])**2)**(1.5)
def E_total(x, y, charges):
Ex, Ey=0, 0
for C in charges:
E=E_point_charge(C.q, C.pos, x, y)
Ex=Ex+E[0]
Ey=Ey+E[1]
return [ Ex, Ey ]
def E_dir(t, y, charges):
Ex, Ey=E_total(y[0], y[1], charges)
n=np.sqrt(Ex**2+Ey*Ey)
return [Ex/n, Ey/n]
def V_point_charge(q, a, x, y):
return q/((x-a[0])**2+(y-a[1])**2)**(0.5)
def V_total(x, y, charges):
V=0
for C in charges:
Vp=V_point_charge(C.q, C.pos, x, y)
V = V+Vp
return V
#define graph boundary
x0 = 0
x1 = 20
y0 = 0
y1 = 8
# charges and positions
charges=[ charge(-0.5, [10, 4]),
charge(-1, [3, 3]),
charge(0.5, [8, 5]),
charge(-2, [1.5, 1.5]),
charge(0.5,[14,5]),
charge(8,[16,4.5])]
qq = []
for c in charges:
qq.append(c.q)
qq.sort()
q0 = qq[0]
q1 = qq[-1]
# calculate electric potential
vvs = []
xxs = []
yys = []
numcalcv = 300
for xx,yy in product(np.linspace(x0,x1,numcalcv),np.linspace(y0,y1,numcalcv)):
xxs.append(xx)
yys.append(yy)
vvs.append(V_total(xx,yy,charges))
xxs = np.array(xxs)
yys = np.array(yys)
vvs = np.array(vvs)
#output chart
# plot electric potential
clim0,clim1 = q0,q1
vvs[np.where(vvs<clim0)] = clim0*0.999999 # to avoid error
vvs[np.where(vvs>clim1)] = clim1*0.999999 # to avoid error
#field color
plt.tricontourf(xxs,yys,vvs,100,cmap=cm.jet)
#color scale
cbar = plt.colorbar()
#chart bound
plt.axis('tight')
plt.axis('equal')
plt.xlim(x0,x1)
plt.ylim(y0,y1)
#plt.savefig('out22.jpg',dpi = 300)
plt.show()

Well, after a while I figured out what the problem was, so might as well post an update on what I found. Point charge for each coordinate created by np.linspace is calculated by charge/distance^2 function. The problem arises when the distance to the charge is too small, so the calculated point charge approaches infinity. The solution I came up with is to replace the distance value with a minimum value (line spacing) when the distance is too small.
My code now look like this:
import matplotlib.pyplot as plt
import numpy as np
from scipy.integrate import ode as ode
from matplotlib import cm
from itertools import product
class charge:
def __init__(self, q, pos):
self.q=q
self.pos=pos
def V_point_charge(q, a, x, y, xinc, yinc):
xx = (x-a[0])**2
yy = (y-a[1])**2
if xx < xinc:
xx = xinc
if yy < yinc:
yy = yinc
return q/((xx)+(yy))**(0.5)
def V_total(x, y, charges,xinc,yinc):
V=0
for C in charges:
Vp=V_point_charge(C.q, C.pos, x, y, xinc,yinc)
V = V+Vp
return V
#define graph boundary
x0 = 0
x1 = 20
y0 = 0
y1 = 8
# charges and positions (charge, [x,y])
charges=[ charge(-0.5, [10, 4]),
charge(-1, [3, 3]),
charge(0.5, [8, 5]),
charge(-10, [1.5, 1.5]),
charge(2,[14,5]),
charge(8,[16,4.5])]
qq = []
for c in charges:
qq.append(c.q)
qq.sort()
q0 = qq[0]
q1 = qq[-1]
# calculate electric potential
vvs = []
xxs = []
yys = []
numcalcv = 300
xinc = np.sqrt((x1-x0)**2)/(numcalcv-1)
yinc = np.sqrt((y1-y0)**2)/(numcalcv-1)
for xx,yy in product(np.linspace(x0,x1,numcalcv),np.linspace(y0,y1,numcalcv)):
xxs.append(xx)
yys.append(yy)
vvs.append(V_total(xx,yy,charges,xinc,yinc))
xxs = np.array(xxs)
yys = np.array(yys)
vvs = np.array(vvs)
#output chart
# plot electric potential
clim0,clim1 = q0,q1
vvs[np.where(vvs<clim0)] = clim0*0.999999 # to avoid error
vvs[np.where(vvs>clim1)] = clim1*0.999999 # to avoid error
#field contour
#plt.tricontour(xxs,yys,vvs,10,colors="0.3")
#field color
plt.tricontourf(xxs,yys,vvs,100,cmap=cm.jet)
#chart bound
plt.axis('tight')
plt.xlim(x0,x1)
plt.ylim(y0,y1)
#color scale
cbar = plt.colorbar()
#save image
plt.axis('equal')
#plt.axis('off')
#plt.savefig('out6.jpg',dpi = 300, bbox_inches = 'tight', pad_inches = 0)
plt.show()

I tried implementing RK4 to N-body simulation. But the look from the graph the code is not working properly

as I said in the title, even though the code is running without any error messages, the code is not working properly. I'm applying RK4 method to the N-body simulation code, particularly solar system. The graph looks strange like there is no gravity between planets. Is my RK4 code running properly? I think the error is there because EULER method is just working fine.
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
class Body():
"""
This class contains adjustable parameters as attributes. These
parameters include current and previous positions, velocities, and
accelerations.
"""
def __init__(self,
id, facecolor, pos,
mass=1, vel=None, acc=None):
self.id = id
self.facecolor = facecolor
self.pos = np.array(pos, dtype=float)
self.mass = mass
self.vel = self.autocorrect_parameter(vel)
self.acc = self.autocorrect_parameter(acc)
self.vector_pos = [self.pos]
self.vector_vel = [self.vel]
self.vector_acc = [self.acc]
def autocorrect_parameter(self, args):
if args is None:
return np.zeros(self.pos.shape, dtype=float)
return np.array(args, dtype=float)
def scalar_pos(self):
return np.sqrt(np.sum(np.square(self.vector_pos), axis=1))
def partial_step(point1, point2, dt):
ret=0
ret=point1+point2*dt
return ret
class computinggravity():
"""
This class contains methods to run a simulation of N bodies that interact
gravitationally over some time. Each body in the simulation keeps a record
of parameters (pos, vel, and acc) as time progresses.
"""
def __init__(self, bodies, t=0, gravitational_constant=6.67e-11):
self.bodies = bodies
self.nbodies = len(bodies)
self.ndim = len(bodies[0].pos)
self.t = t
self.moments = [t]
self.gravitational_constant = gravitational_constant
def get_acc(self, ibody, jbody):
dpos = ibody.pos - jbody.pos
r = np.sum(dpos**2)
acc = -1 * self.gravitational_constant * jbody.mass*dpos / np.sqrt(r**3)
return acc
def update_accelerations(self,dt):
for ith_body, ibody in enumerate(self.bodies):
ibody.acc *= 0
k1 =0
k2 =0
k3 =0
k4 =0
acc_pos =0
acc_vel =0
for jth_body, jbody in enumerate(self.bodies):
if ith_body != jth_body:
acc = self.get_acc(ibody, jbody)
k1=acc*(jbody.pos-ibody.pos)
#acc_vel=partial_step(ibody.vel,k1,0.5)
acc_vel=ibody.vel+k1*0.5
#acc_pos=partial_step(ibody.pos,acc_vel,0.5)
acc_pos=ibody.pos+acc_vel*0.5
k2=(acc_pos-(acc_pos+acc_vel*0.5*dt))*acc
#acc_vel=partial_step(ibody.vel, k2, 0.5)
acc_vel=ibody.vel+k2*0.5
k3=(acc_pos-(acc_pos+acc_vel*0.5*dt))*acc
#acc_vel=partial_step(ibody.vel, k3, 1)
acc_vel=ibody.vel+k3*1
#acc_pos=partial_step(ibody.pos, acc_vel, 0.5)
acc_pos=ibody.pos+acc_vel*0.5
k4=(jbody.pos-(acc_pos+acc_vel*dt))*acc
#ibody.acc=ibody.acc+acc
ibody.acc=ibody.acc+(k1+k2*2+k3*2+k4)/6
ibody.vector_acc.append(np.copy(ibody.acc))
def updatingvelocity(self,dt):
for ibody in self.bodies:
#ibody.acc=self.update_accelerations(self)
ibody.vel=ibody.vel+ibody.acc*dt
ibody.vector_vel.append(np.copy(ibody.vel))
def updatingposition(self,dt):
for ibody in self.bodies:
ibody.pos=ibody.pos+ibody.vel * dt
ibody.vector_pos.append(np.copy(ibody.pos))
def simulation(self, dt, duration):
nsteps = int(duration / dt)
for ith_step in range(nsteps):
self.update_accelerations(dt)
self.updatingvelocity(dt)
self.updatingposition(dt)
self.t= self.t+ dt
self.moments.append(self.t)
# Masses
SOLAR_MASS = 1.988e30
EARTH_MASS = 5.9722e24
# Distances
ASTRO_UNIT = 1.496e11
MILE = 1609
# Durations
HOUR = 3600
DAY = 24 * HOUR
YEAR = 365 * DAY
def main():
m_sun = 1 * SOLAR_MASS
sun = Body('Sun', 'yellow', [0, 0, 0], m_sun)
m_mercury = 0.05227 * EARTH_MASS
d_mercury = 0.4392 * ASTRO_UNIT
v_mercury = (106_000 * MILE) / (1 * HOUR)
mercury = Body('Mercury', 'gray',
[d_mercury, 0, 0], m_mercury,
[0, v_mercury, 0])
m_earth = 1 * EARTH_MASS
d_earth = 1 * ASTRO_UNIT
v_earth = (66_600 * MILE) / (1 * HOUR)
earth = Body('Earth', 'blue', [d_earth, 0, 0], m_earth, [0, v_earth, 0])
m_mars = 0.1704 * EARTH_MASS
d_mars = 1.653 * ASTRO_UNIT
v_mars = (53_900 * MILE) / (1 * HOUR)
mars = Body('Mars', 'darkred', [0, d_mars, 0], m_mars, [v_mars, 0, 0])
m_jupiter = 318* EARTH_MASS
d_jupister= 5 * ASTRO_UNIT
v_jupiter = (28_000 * MILE)/ (1* HOUR)
jupiter = Body('Jupiter', 'red',[d_jupister,0,0], m_jupiter, [0,v_jupiter,0])
m_saturn = 95* EARTH_MASS
d_saturn= 9.5 * ASTRO_UNIT
v_saturn = (21_675 * MILE)/ (1* HOUR)
saturn = Body('Saturn', 'brown',[0,d_saturn,0], m_saturn, [v_saturn,0,0])
m_uranus = 14.5 * EARTH_MASS
d_uranus = 19.8 * ASTRO_UNIT
v_uranus = (15_233* MILE)/(1 * HOUR)
uranus = Body('Uranus', 'cyan', [d_uranus,0,0], m_uranus, [0,v_uranus,0])
m_neptune = 17 * EARTH_MASS
d_neptune = 30 * ASTRO_UNIT
v_neptune = (12_146* MILE)/(1*HOUR)
neptune = Body('Neptune', 'blue', [0, d_neptune,0], m_neptune, [v_neptune,0,0])
bodies = [sun, mercury, earth, mars, jupiter,saturn, uranus, neptune]
dt = 1 * DAY
duration = 40 * YEAR
gd = computinggravity(bodies)
gd.simulation(dt, duration)
fig = plt.figure(figsize=(11, 7))
ax_left = fig.add_subplot(1, 2, 1, projection='3d')
ax_left.set_title('3-D Position')
ax_right = fig.add_subplot(1, 2, 2)
ax_right.set_title('Displacement vs Time')
ts = np.array(gd.moments) / YEAR
xticks = np.arange(max(ts)+1).astype(int)
for body in gd.bodies:
vector = np.array(body.vector_pos)
vector_coordinates = vector / ASTRO_UNIT
scalar = body.scalar_pos()
scalar_coordinates = scalar / ASTRO_UNIT
ax_left.scatter(*vector_coordinates.T, marker='.',
c=body.facecolor, label=body.id)
ax_right.scatter(ts, scalar_coordinates, marker='.',
c=body.facecolor, label=body.id)
ax_right.set_xticks(xticks)
ax_right.grid(color='k', linestyle=':', alpha=0.3)
fig.subplots_adjust(bottom=0.3)
fig.legend(loc='lower center', mode='expand', ncol=len(gd.bodies))
plt.show()
# plt.close(fig)
if __name__ == '__main__':
main()

Your problem is a typical one and occurs when the structured and appropriate approach to implement Euler/symplectic Euler/Verlet on a particle-centered philosophy is extended to higher order methods. In short, it does not work.
The big difference is that for higher-order methods one uses temporary states (while keeping the original state from the start of the time step) that are not part of the numerical result and even further slightly aside the trajectory, see Visualization of Third Order Runge-Kutta for how one could imagine that. Or in other words, you need to complete stage 1 of the Runge-Kutta method for all particles, then compute the new (virtual) positions for the computations of stage 2, then complete stage 2 for all particles before starting anything of stage 3 etc.
The overall aim should be to keep the physics of the model and the numerical integration method apart. The model should only implement the general support functionality, for instance cloning itself at a state that is shifted by the scaled derivatives of another state, or by the linear combinations of a set of derivatives. So that RK4 is implemented as something like
state.compute_derivatives()
temp1 = state.clone_shift_single(0.5*dt,state)
temp1.compute_derivatives()
temp2 = state.clone_shift_single(0.5*dt,temp1)
temp2.compute_derivatives()
temp3 = state.clone_shift_single(1.0*dt,temp2)
temp3.compute_derivatives()
newstate = state.clone_shift_multiple([dt/6,dt/3,dt/3,dt/6],[state,temp1,temp2,temp3])
Another variant is for the model to implement functions that copy the state and derivative components of the particles to and from a flat vector in a previously fixed order, so that the vector operations inside the RK method can be implemented by numpy arrays or the likes. This means that the derivatives function that is passed to the solver method, be it RK4 in a "cook book" implementation or some other solver as scipy.integrate.solve_ivp, is a wrapper for the model functions that can look like
def derivs(t,u):
model.copy_state_from(u)
model.compute_derivatives(t) # if there is some time dependent component
du_dt = np.zeros_like(u)
model.copy_derivatives_to(du_dt)
return du_dt

Calculating random sample points using polar coordinates on cartesian map

I'm trying to generate random sample points on a cartesian plane using polar coordinates. I have a cartesian map with polar sectors, I'd like to put a random sample point within each of the sectors.
Problem Visual Description
I've added a sample point in the first sector. The problem is I don't know how to set the min and max limits for each sector as it's a cartesian plane (using cartesian min and max of the sector corners will give you boxes instead of the entire polar sector).
Code is commented for clarity. Final output posted below.
import numpy as np
import matplotlib.pyplot as plt
plt.rcParams['figure.figsize'] = [10, 10]
import math
import pylab as pl
from matplotlib import collections as mc
import pprint
from IPython.utils import io
from random import randrange, uniform
#convertes cartesian x,y coordinates to polar r, theta coordinates
def cart2pol(x, y):
rho = np.sqrt(x**2 + y**2)
phi = np.arctan2(y, x)
return np.array([rho, phi])
#convertes polar r,theta coordinates to cartesian x,y coordinates
def pol2cart(r, theta): #r is distance
x = r * np.cos(theta)
y = r * np.sin(theta)
return np.array([x, y])
#cooks delicious pie
pi = np.pi
#no idea what this does
theta = np.linspace(0,2*pi,100)
#x theta
def x_size(r):
return r*np.cos(theta)
#y theta
def y_size(r):
return r*np.sin(theta)
#calculates distribution of sectors on a circle in radians
#eg. sub_liner(3) = array([0. , 2.0943951, 4.1887902])
def sub_liner(k):
sub_lines = []
for c,b in enumerate(range(0,k)):
sub_lines = np.append(sub_lines,((12*pi/6)/k)*c)
return sub_lines
#calculates all distribution sectors for every ring and puts them in a list
def mlp(i):
master_lines = []
k = 3
for a in range(0,i):
master_lines.append(sub_liner(k))
k += 3
return master_lines
#calculates all four corners of each sector for a ring
#(ring,ring points,number of rings)
def cg(r,rp,n):
return [[[pol2cart(r-1,mlp(n)[r-1][i])[0],pol2cart(r-1,mlp(n)[r-1][i])[1]]\
,[pol2cart(r,mlp(n)[r-1][i])[0],pol2cart(r,mlp(n)[r-1][i])[1]]] for i in range(0,rp)]
#generates all corners for the ring sectors
def rg(n):
cgl = []
k = 3
for r in range(1,11):
cgl.append(cg(r,k,n))
k += 3
output = cgl[0]
for q in range(1,10):
output = np.concatenate((output,cgl[q]))
return output
#print(cg(1,3,10)[0][0][0])
#print(cg(1,3,10))
# randrange gives you an integral value
irand = randrange(0, 10)
# uniform gives you a floating-point value
frand = uniform(0, 10)
#define ring sectors
ring_sectors = rg(10)
#define node points
nx = 0.5
ny = 0.5
#define ring distance
ymin = [0]
ymax = [1]
#generate rings
ring_r = np.sqrt(1.0)
master_array = np.array([[x_size(i),y_size(i)] for i in range(0,11)])
#plot rings
fig, ax = plt.subplots(1)
[ax.plot(master_array[i][0],master_array[i][1]) for i in range(0,11)]
ax.set_aspect(1)
size = 10
plt.xlim(-size,size)
plt.ylim(-size,size)
#generate nodes
ax.plot(nx, ny, 'o', color='black');
#ring lines
lc = mc.LineCollection(ring_sectors, color='black', linewidths=2)
ax.add_collection(lc)
plt.grid(linestyle='--')
plt.title('System Generator', fontsize=8)
plt.show()
Sample output can be viewed at.
Edit:
What I've tried:
Based on feedback, I implemented a system which gets random uniform values between the polar coordinates, and it works, but the points aren't neatly distributed within their sectors as they should be, and I'm not sure why. Maybe my math is off or I made a mistake in the generator functions. If anyone has any insight, I'm all ears.
Output with points
def ngx(n):
rmin = 0
rmax = 1
nxl = []
s1 = 0
s2 = 1
k = 0
for i in range(0,n):
for a in range(0,rmax*3):
nxl.append(pol2cart(np.random.uniform(rmin,rmax),\
np.random.uniform(sub_liner(rmax*3)[(s1+k)%(rmax*3)],sub_liner(rmax*3)[(s2+k)%(rmax*3)]))[0])
k += 1
rmin += 1
rmax += 1
return nxl
def ngy(n):
rmin = 0
rmax = 1
nyl = []
s1 = 0
s2 = 1
k = 0
for i in range(0,n):
for a in range(0,rmax*3):
nyl.append(pol2cart(np.random.uniform(rmin,rmax),\
np.random.uniform(sub_liner(rmax*3)[(s1+k)%(rmax*3)],sub_liner(rmax*3)[(s2+k)%(rmax*3)]))[1])
k += 1
rmin += 1
rmax += 1
return nyl
#define node points
nx = ngx(10)
ny = ngy(10)

Mapping surface curvature to face?

I am trying to map surface curvature (mean, gaussian and principle curvature) values to surface faces. I have computed the curvature values for an artificially generated 3D surface (eg. cylinder). The resulting 3D surface that I am trying to get is something like this mean curvature mapped to surface. Can somebody guide me in how to get this?
The code for the surface I am creating is:
import math
import matplotlib.pyplot as plt
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
xindex = []
yindex = []
zindex = []
x = []
y = []
z = []
count = 1
surfaceSt = []
import numpy
numpy.set_printoptions(threshold=numpy.nan)
#surfaceStX = numpy.empty((10,36))
#surfaceStY = numpy.empty((10,36))
#surfaceStZ = numpy.empty((10,36))
surfaceStZ = []
surfaceStX = []
surfaceStY = []
for i in range(1,21):
if i < 11:
x = []
y = []
z = []
pt = []
ptX = []
ptY = []
ptZ = []
for t in range(0,360,10):
x = i*math.sin(math.radians(t))
y = i*math.cos(math.radians(t))
z = i-1
ptX.append(x)
ptY.append(y)
ptZ.append(z)
pt.append([x,y,z])
ptX.append(ptX[0])
ptY.append(ptY[0])
ptZ.append(ptZ[0])
surfaceStX.append(ptX)
surfaceStY.append(ptY)
surfaceStZ.append(ptZ)
# numpy.append(surfaceStX,ptX)
# numpy.append(surfaceStY,ptY)
# numpy.append(surfaceStZ,ptZ)
#ax.scatter(x,y,z)
elif i >= 11:
x = []
y = []
z = []
pt = []
ptX = []
ptY = []
ptZ = []
for t in range(0,360,10):
x = (i-count)*math.sin(math.radians(t))
y = (i-count)*math.cos(math.radians(t))
z = i-1
ptX.append(x)
ptY.append(y)
ptZ.append(z)
pt.append([x,y,z])
ptX.append(ptX[0])
ptY.append(ptY[0])
ptZ.append(ptZ[0])
surfaceStX.append(ptX)
surfaceStY.append(ptY)
surfaceStZ.append(ptZ)
count +=2
X = numpy.array(surfaceStX)
Y = numpy.array(surfaceStY)
Z = numpy.array(surfaceStZ)
ax = fig.add_subplot(111, projection='3d')
ax.plot_surface(X, Y, Z, rstride=1, cstride=1,shade = 'True' )
from surfaceCurvature import surface_curvature
Pcurvature,Gcurvature,Mcurvature = surface_curvature(X,Y,Z)
plt.show()
My surface curvature computation is given below (courtesy: https://github.com/sujithTSR/surface-curvature):
def surface_curvature(X,Y,Z):
(lr,lb)=X.shape
#print lr
#print "awfshss-------------"
#print lb
#First Derivatives
Xv,Xu=np.gradient(X)
Yv,Yu=np.gradient(Y)
Zv,Zu=np.gradient(Z)
#Second Derivatives
Xuv,Xuu=np.gradient(Xu)
Yuv,Yuu=np.gradient(Yu)
Zuv,Zuu=np.gradient(Zu)
Xvv,Xuv=np.gradient(Xv)
Yvv,Yuv=np.gradient(Yv)
Zvv,Zuv=np.gradient(Zv)
#2D to 1D conversion
#Reshape to 1D vectors
Xu=np.reshape(Xu,lr*lb)
Yu=np.reshape(Yu,lr*lb)
Zu=np.reshape(Zu,lr*lb)
Xv=np.reshape(Xv,lr*lb)
Yv=np.reshape(Yv,lr*lb)
Zv=np.reshape(Zv,lr*lb)
Xuu=np.reshape(Xuu,lr*lb)
Yuu=np.reshape(Yuu,lr*lb)
Zuu=np.reshape(Zuu,lr*lb)
Xuv=np.reshape(Xuv,lr*lb)
Yuv=np.reshape(Yuv,lr*lb)
Zuv=np.reshape(Zuv,lr*lb)
Xvv=np.reshape(Xvv,lr*lb)
Yvv=np.reshape(Yvv,lr*lb)
Zvv=np.reshape(Zvv,lr*lb)
Xu=np.c_[Xu, Yu, Zu]
Xv=np.c_[Xv, Yv, Zv]
Xuu=np.c_[Xuu, Yuu, Zuu]
Xuv=np.c_[Xuv, Yuv, Zuv]
Xvv=np.c_[Xvv, Yvv, Zvv]
# First fundamental Coeffecients of the surface (E,F,G)
E=np.einsum('ij,ij->i', Xu, Xu)
F=np.einsum('ij,ij->i', Xu, Xv)
G=np.einsum('ij,ij->i', Xv, Xv)
m=np.cross(Xu,Xv,axisa=1, axisb=1)
p=np.sqrt(np.einsum('ij,ij->i', m, m))
n=m/np.c_[p,p,p]
# Second fundamental Coeffecients of the surface (L,M,N), (e,f,g)
L= np.einsum('ij,ij->i', Xuu, n) #e
M= np.einsum('ij,ij->i', Xuv, n) #f
N= np.einsum('ij,ij->i', Xvv, n) #g
# Gaussian Curvature
K=(L*N-M**2)/(E*G-F**2)
K=np.reshape(K,lr*lb)
# Mean Curvature
H = (E*N + G*L - 2*F*M)/((E*G - F**2))
H = np.reshape(H,lr*lb)
# Principle Curvatures
Pmax = H + np.sqrt(H**2 - K)
Pmin = H - np.sqrt(H**2 - K)
#[Pmax, Pmin]
Principle = [Pmax,Pmin]
return Principle,K,H
EDIT 1:
I tried a few things based on the link provided by armatita. Following is my code:
'''
Creat half cylinder
'''
import numpy
import matplotlib.pyplot as plt
import math
ptX= []
ptY = []
ptZ = []
ptX1 = []
ptY1 = []
ptZ1 = []
for i in range(0,10):
x = []
y = []
z = []
for t in range(0,200,20):
x.append(10*math.cos(math.radians(t)))
y.append(10*math.sin(math.radians(t)))
z.append(i)
x1= 5*math.cos(math.radians(t))
y1 = 5*math.sin(math.radians(t))
z1 = i
ptX1.append(x1)
ptY1.append(y1)
ptZ1.append(z1)
ptX.append(x)
ptY.append(y)
ptZ.append(z)
X = numpy.array(ptX)
Y = numpy.array(ptY)
Z = numpy.array(ptZ)
fig = plt.figure()
ax = fig.add_subplot(111,projection = '3d')
from surfaceCurvature import surface_curvature
p,g,m= surface_curvature(X,Y,Z)
n = numpy.reshape(m,numpy.shape(X))
ax.plot_surface(X,Y,Z, rstride=1, cstride=1)
plt.show()
'''
Map mean curvature to color
'''
import numpy as np
X1 = X.ravel()
Y1 = Y.ravel()
Z1 = Z.ravel()
from scipy.interpolate import RectBivariateSpline
# Define the points at the centers of the faces:
y_coords, x_coords = np.unique(Y1), np.unique(X1)
y_centers, x_centers = [ arr[:-1] + np.diff(arr)/2 for arr in (y_coords, x_coords)]
# Convert back to a 2D grid, required for plot_surface:
#Y1 = Y.reshape(y_coords.size, -1)
#X1 = X.reshape(-1, x_coords.size)
#Z1 = Z.reshape(X.shape)
C = m.reshape(X.shape)
C -= C.min()
C /= C.max()
interp_func = RectBivariateSpline(x_coords, y_coords, C.T, kx=1, ky=1)
I get the following error:
raise TypeError('y dimension of z must have same number of y')
TypeError: y dimension of z must have same number of elements as y
All the dimensions are same. Can anybody tell what's going wrong with my implementation?

I think you need to figure out exactly what you need. Looking at your code I notice you are producing variables that have no use. Also you seem to have a function to calculate the surface curvature but than you try to make some calculations using the np.unique function for which I cannot see the purpose here (and that is why that error appears).
So let's assume this:
You have a function that returns the curvature value for each cell.
You have the X,Y and Z meshes to plot that surface.
Using your code, and assuming you m variable is the curvature (again this is in your code), if I do this:
import numpy
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from matplotlib import cm
import math
# Here would be the surface_curvature function
X = numpy.array(ptX)
Y = numpy.array(ptY)
Z = numpy.array(ptZ)
p,g,m= surface_curvature(X,Y,Z)
C = m.reshape(X.shape)
C -= C.min()
C /= C.max()
fig = plt.figure()
ax = fig.add_subplot(111,projection = '3d')
n = numpy.reshape(m,numpy.shape(X))
ax.plot_surface(X,Y,Z,facecolors = cm.jet(C), rstride=1, cstride=1)
plt.show()
, I obtain this:
Which is a value mapped to color in a matplotlib surface. If that C you've built is not the actual curvature you need to replace it by the one that is.

Python Solving 1D wave equation and animating

I adapted a code for solving 1d wave equation. My problem is animating the time dependent results. I tried it with the following code, any idea what I made wrong? I used the animation package, but somehiw it always says that more than one element is ambigous...
from __future__ import division
import numpy as np
import matplotlib.pyplot as plt
from matplotlib import rc
import matplotlib.animation as animation
#Set Helvetica Font
rc('font',**{'family':'sans-serif','sans-serif':['Helvetica']})
## for Palatino and other serif fonts use:
#rc('font',**{'family':'serif','serif':['Palatino']})
rc('text', usetex=True)
class wave1d(object):
def __init__(self,width,T,nx,nt,c):
self.x = np.linspace(-0.5*width,0.5*width,nx)
self.t = np.linspace(0,T,nt+1)
self.dx = self.x[1]-self.x[0]
self.dt = self.t[1]-self.t[0]
self.xx = self.x
# Gamma_x squared
self.gx2 = c*self.dt/self.dx
# 2*(1-gamma_x^2-gamma_y^2)
self.gamma = 2*(1 - self.gx2)
def solve(self,ffun,gfun):
f = ffun(self.xx)
g = gfun(self.xx)
u = np.zeros((nx,nt+1))
# Set initial condition
u[:,0] = f
""" Compute first time step """
u[:,1] = 0.5*self.gamma*f+self.dt*g
u[1:-1,1] += 0.5*self.gx2*(f[2:]+f[:-2])
for k in range(1,nt):
# Every point contains these terms
u[:,k+1] = self.gamma*u[:,k] - u[:,k-1]
# Interior points
u[1:-1,k+1] += self.gx2*(u[2:,k]+u[:-2,k])
# Right boundary
u[-1,k+1] += 2*self.gx2*u[-2,k]
# Left boundary
u[0,k+1] += 2*self.gx2*u[1,k]
return u
if __name__ == '__main__':
# Final time
T = 2
# Domain dimensions
width = 8
# Wave speed
c = 1
# Number of time steps
nt = 400
# Grid points in x direction
nx = 500
wave_eq = wave1d(width,T,nx,nt,c)
# Initial value functions
f = lambda x: np.sin(x-width)
g = lambda x: 0
u = wave_eq.solve(f,g)
x = wave_eq.x
frames = []
fig = plt.figure(1,(16,8))
for k in range(nt+1):
frame = plt.show(u[:,k])
frames.append([frame])
ani = animation.ArtistAnimation(fig,frames,interval=50, blit=True,repeat_delay=1000)
# ani.save('wave1d.mp4',dpi=300)
plt.show()