import numpy as np
x = [0,1,2]
y = [0,1,2]
xx,yy = np.meshgrid(x,y)
                   

def f(x,y):
    return x+y
z = f(xx,yy)
                   

def HrHz(r,z,z1,rr,tt,deltaS,M):
    b2 = (z-z1)**2
    cos = np.cos(tt);
    D = (r**2+rr**2-2*r*rr*cos+b2)**1.5
    Hr = (r-rr*cos)/D
    Hz = (z-z1)/D
    Hr = Hr.sum()*deltaS/(2*np.pi)*M
    Hz = Hz.sum()*deltaS/(2*np.pi)*M
    return (Hr,Hz)
    
def grille(Nr,Ntheta,a):
    delta_theta = np.pi/Ntheta
    delta_r = a/np.sqrt(Nr)
    theta1 = (0.5+np.arange(Ntheta))*delta_theta
    r1 = np.sqrt(0.5+np.arange(Nr))*delta_r
    deltaS = delta_theta*delta_r**2*0.5
    rr,tt = np.meshgrid(r1,theta1)
    return (rr,tt,deltaS)
               
def champH(Nr,Ntheta,r,z,z1,a,M):
    (rr,tt,deltaS) = grille(Nr,Ntheta,a)
    return HrHz(r,z,z1,rr,tt,deltaS,M)
                 

Nr = 100
Ntheta = 50
a = 1
z = 1
z1 = 0
r = 0
M = 1
(Hr,Hz) = champH(Nr,Ntheta,r,z,z1,a,M)

                 

def Hz_axe(z,z1,a,M):
    return M/2*((z-z1)/abs(z-z1)-(z-z1)/(a**2+(z-z1)**2)**0.5)

Hz = Hz_axe(z,z1,a,M)
                 

from matplotlib.pyplot import *
z = np.linspace(0.01,5,100)
Nr = 100
Ntheta = 50
a = 1
z1 = 0
r = 0
M = 1
(rr,tt,deltaS) = grille(Nr,Ntheta,a)
Hz1 = np.zeros(len(z))
for k in range(len(z)):
    Hr,Hz1[k] = HrHz(r,z[k],z1,rr,tt,deltaS,M)
Hz2 = Hz_axe(z,z1,a,M)
figure()
plot(z,Hz1,"r-",label='approché')
plot(z,Hz2,"k",label='exact')
grid()
xlabel('z')
ylabel('Hz')
legend(loc='upper right')
ylim(0,1)
                 


z = np.logspace(-4,-2,10)
Nr = 1000
Ntheta = 50
a = 1
z1 = 0
r = 0
M = 1
(rr,tt,deltaS) = grille(Nr,Ntheta,a)
Hz1 = np.zeros(len(z))
for k in range(len(z)):
    Hr,Hz1[k] = HrHz(r,z[k],z1,rr,tt,deltaS,M)
Hz2 = Hz_axe(z,z1,a,M)
figure()
plot(z,Hz1,"r-",label='approché')
plot(z,Hz2,"k",label='exact')
grid()
xlabel('z')
ylabel('Hz')
legend(loc='upper right')
xscale('log')
                 

z = np.linspace(0.01,5,100)
Ntheta = 50
a = 1
z1 = 0
r = 0
M = 1
Hz1 = np.zeros(len(z))
p = 3
Nr_min = 10
for k in range(len(z)):
    Nr = max(int((p*a/abs(z[k]))**2),Nr_min)
    (rr,tt,deltaS) = grille(Nr,Ntheta,a)
    Hr,Hz1[k] = HrHz(r,z[k],z1,rr,tt,deltaS,M)
Hz2 = Hz_axe(z,z1,a,M)
figure()
plot(z,Hz1,"r-",label='approché')
plot(z,Hz2,"k",label='exact')
grid()
xlabel('z')
ylabel('Hz')
legend(loc='upper right')
ylim(0,1)
                 

r = np.linspace(0,2,50)
z = 1e-1
a = 1
z1 = 0
M = 1
Hz1 = np.zeros(len(r))
Hr1 = np.zeros(len(r))
Nr_min = 10
p = 2
Nr = max(int((p*a/abs(z))**2),Nr_min)
(rr,tt,deltaS) = grille(Nr,Ntheta,a)
for k in range(len(r)):    
    Hr1[k],Hz1[k] = HrHz(r[k],z,z1,rr,tt,deltaS,M)
H1 = np.sqrt(Hr1**2+Hz1**2)
Hz2 = np.zeros(len(r))
Hr2 = np.zeros(len(r))
p = 3
Nr = max(int((p*a/abs(z))**2),Nr_min)
(rr,tt,deltaS) = grille(Nr,Ntheta,a)
for k in range(len(r)):    
    Hr2[k],Hz2[k] = HrHz(r[k],z,z1,rr,tt,deltaS,M)
H2 = np.sqrt(Hr2**2+Hz2**2)
figure()
subplot(311)
plot(r,Hr1,'r-')
plot(r,Hr2,'b-')
ylabel('Hr')
grid()
subplot(312)
plot(r,Hz1,'r-')
plot(r,Hz2,'b-')
ylabel('Hz')
grid()
subplot(313)
plot(r,H1,'r')
plot(r,H2,'b-')
ylabel('H')
xlabel('r')
grid()
                  

def champAimant(r,z,a,zs,zn,M,mu0=1):
    p = 3
    Nr_min = 10
    Ntheta = 50
    Nr = Nr_min
    if abs(r)<=a:
        Nr = max(int((p*a/abs(z-zs))**2),Nr_min)
    (rr,tt,deltaS) = grille(Nr,Ntheta,a)
    (Hr1,Hz1) = HrHz(r,z,zs,rr,tt,deltaS,-M)
    Nr = Nr_min
    if abs(r)<=a:
        Nr = max(int((p*a/abs(z-zn))**2),Nr_min)
    (rr,tt,deltaS) = grille(Nr,Ntheta,a)
    (Hr2,Hz2) = HrHz(r,z,zn,rr,tt,deltaS,M)
    Br1 = mu0*Hr1
    Br2 = mu0*Hr2
    if abs(r)<a and z>min(zs,zn) and z<max(zs,zn):
        Bz1 = mu0*(Hz1+M)
        Bz2 = mu0*(Hz2+M)
    else:
        Bz1 = mu0*Hz1
        Bz2 = mu0*Hz2
    return (Hr1+Hr2,Hz1+Hz2,Br1+Br2,Bz1+Bz2)
                    

def fleche(x,y,sens,long,style='k-'):
    n = len(x)//2
    xa = x[n]
    xb = x[n+sens]
    ya = y[n]
    yb = y[n+sens]
    z = (xb-xa)+1j*(yb-ya)
    phi = np.angle(z)
    a = np.pi/2+np.pi/3
    xc = xb+long*np.cos(phi-a)
    yc = yb+long*np.sin(phi-a)
    xd = xb+long*np.cos(phi+a)
    yd = yb+long*np.sin(phi+a)
    plot([xb,xc],[yb,yc],style)
    plot([xb,xd],[yb,yd],style)
                    

def ligneH(a,zs,zn,M,ri,zi,sens,rmax,zmax,dmin):
    ds = 0.01*sens
    r = ri
    z = zi
    ligne_z = []
    ligne_r = []
    continuer = True
    while continuer:
        ligne_z.append(z)
        ligne_r.append(r)
        (Hr,Hz,Br,Bz) = champAimant(r,z,a,zs,zn,M)
        H = np.sqrt(Hr**2+Hz**2)
        r += Hr/H*ds                                    
        z += Hz/H*ds       
        if (abs(z-zs) < dmin and abs(r)<a) or (abs(z-zn)< dmin and abs(r)<a) or (abs(r)>rmax) or (abs(z)>zmax) : continuer=False
    return (np.array(ligne_r),np.array(ligne_z))
a=1
M = 1 
zs=-2
zn=2
ds = 0.05
figure(figsize=(8,8))
zmax = 10
rmax = 10
longfleche = 0.2
dmin = 0.1 # distance minimale d'approche des faces
for ri in [0,0.2,0.4,0.6,0.8,1.0]:
    sens = 1
    ligne_r,ligne_z = ligneH(a,zs,zn,M,ri,zn+dmin,sens,rmax,zmax,dmin)
    plot(ligne_z,ligne_r,'b-')
    fleche(ligne_z,ligne_r,sens,longfleche,style='b-')
    plot(ligne_z,-ligne_r,'b-')
    fleche(ligne_z,-ligne_r,sens,longfleche,style='b-')
    sens = -1
    ligne_r,ligne_z = ligneH(a,zs,zn,M,ri,zs-dmin,sens,rmax,zmax,dmin)
    plot(ligne_z,ligne_r,'b-')
    fleche(ligne_z,ligne_r,sens,longfleche,style='b-')
    plot(ligne_z,-ligne_r,'b-')
    fleche(ligne_z,-ligne_r,sens,longfleche,style='b-')
    sens = 1
    ligne_r,ligne_z = ligneH(a,zs,zn,M,ri,zn-dmin,sens,rmax,zmax,dmin)
    plot(ligne_z,ligne_r,'b-')
    fleche(ligne_z,ligne_r,sens,longfleche,style='b-')
    plot(ligne_z,-ligne_r,'b-')
    fleche(ligne_z,-ligne_r,sens,longfleche,style='b-')
axis('equal')
xlabel('z')
ylabel('r')
xlim(-zmax,zmax)
ylim(-rmax,rmax)
plot([-zs,zs,zs,-zs,-zs],[a,a,-a,-a,a],'r-')
grid()
title('Lignes de H')
                     

def ligneB(a,zs,zn,M,ri,zi,sens,rmax,zmax,dmin):
    ds = 0.01*sens
    r = ri
    z = zi
    ligne_z = []
    ligne_r = []
    continuer = True
    while continuer:
        ligne_z.append(z)
        ligne_r.append(r)
        (Hr,Hz,Br,Bz) = champAimant(r,z,a,zs,zn,M)
        B = np.sqrt(Br**2+Bz**2)
        r += Br/B*ds
        z += Bz/B*ds
        if (abs(z-zs) < dmin and abs(r)<a) or (abs(z-zn)< dmin and abs(r)<a) or (abs(r)>rmax) or (abs(z)>zmax) : continuer=False
    return (np.array(ligne_r),np.array(ligne_z))
    

figure(figsize=(8,8))
for ri in [0,0.2,0.4,0.6,0.8,1.0]: 
    sens = 1
    ligne_r,ligne_z = ligneB(a,zs,zn,M,ri,zn+dmin,sens,rmax,zmax,dmin)
    plot(ligne_z,ligne_r,'b-')
    fleche(ligne_z,ligne_r,sens,longfleche,style='b-')
    plot(ligne_z,-ligne_r,'b-')
    fleche(ligne_z,-ligne_r,sens,longfleche,style='b-')
    sens = -1 
    ligne_r,ligne_z = ligneB(a,zs,zn,M,ri,zs-dmin,sens,rmax,zmax,dmin)
    plot(ligne_z,ligne_r,'b-')
    fleche(ligne_z,ligne_r,sens,longfleche,style='b-')
    plot(ligne_z,-ligne_r,'b-')
    fleche(ligne_z,-ligne_r,sens,longfleche,style='b-')
    sens = -1
    ligne_r,ligne_z = ligneB(a,zs,zn,M,ri,zn-dmin,sens,rmax,zmax,dmin)
    plot(ligne_z,ligne_r,'b-')
    fleche(ligne_z,ligne_r,sens,longfleche,style='b-')
    plot(ligne_z,-ligne_r,'b-')
    fleche(ligne_z,-ligne_r,sens,longfleche,style='b-')
axis('equal')
xlabel('z')
ylabel('r')
xlim(-zmax,zmax)
ylim(-rmax,rmax)
plot([-zs,zs,zs,-zs,-zs],[a,a,-a,-a,a],'r-')
grid()
title('Lignes de B')
                     

def energieAimant(a,zs,zn,M,Zmin,Zmax,R,Nz,Nr,n,mu0=1):
    Delta_z = (Zmax-Zmin)/Nz
    Delta_r = R/Nr**n
    W = 0
    for i in range(Nz):
        for j in range(Nr):
            rm = (j**n+(j+1)**n)*Delta_r/2
            zm = Zmin + (i+0.5)*Delta_z
            DeltaS = ((j+1)**n-j**n)*Delta_r*Delta_z
            (Hr,Hz,Br,Bz) = champAimant(rm,zm,a,zs,zn,M,mu0) 
            W += rm*(Hz**2+Hr**2)*DeltaS
    W *= mu0*np.pi
    return W
    
                          

a=1
M = 1 
zs=-2
zn=2
Nz = 300
Nr = 100
W = energieAimant(a,zs,zn,M,-10,10,10,Nz,Nr,3)
                          

Nz = 400
Nr = 100
W = energieAimant(a,zs,zn,M,-10,10,10,Nz,Nr,3)                          
                         

Nz = 500
Nr = 150
W = energieAimant(a,zs,zn,M,-10,10,10,Nz,Nr,3)                          
                         

Nz = 500
Nr = 200
W = energieAimant(a,zs,zn,M,-10,10,10,Nz,Nr,3)