We Keep Coding

Convergence tests of Leapfrog method for vectorial wave equation in Python

Considering the following Leapfrog scheme used to discretize a vectorial wave equation with given initial conditions and periodic boundary conditions. I have implemented the scheme and now I want to make numerical convergence tests to show that the scheme is of second order in space and time.
I'm mainly struggling with two points here:
I'm not 100% sure if I implemented the scheme correctly. I really wanted to use slicing because it is so much faster than using loops.
I don't really know how to get the right error plot, because I'm not sure which norm to use. In the examples I have found (they were in 1D) we've always used the L2-Norm.
import numpy as np
import matplotlib.pyplot as plt
# Initial conditions
def p0(x):
return np.cos(2 * np.pi * x)
def u0(x):
return -np.cos(2 * np.pi * x)
# exact solution
def p_exact(x, t):
# return np.cos(2 * np.pi * (x + t))
return p0(x + t)
def u_exact(x, t):
# return -np.cos(2 * np.pi * (x + t))
return u0(x + t)
# function for doing one time step, considering the periodic boundary conditions
def leapfrog_step(p, u):
p[1:] += CFL * (u[:-1] - u[1:])
p[0] = p[-1]
u[:-1] += CFL * (p[:-1] - p[1:])
u[-1] = u[0]
return p, u
# Parameters
CFL = 0.3
LX = 1 # space length
NX = 100 # number of space steps
T = 2 # end time
NN = np.array(range(50, 1000, 50)) # list of discretizations
Ep = []
Eu = []
for NX in NN:
print(NX)
errorsp = []
errorsu = []
x = np.linspace(0, LX, NX) # space grid
dx = x[1] - x[0] # spatial step
dt = CFL * dx # time step
t = np.arange(0, T, dt) # time grid
# TEST
# time loop
for time in t:
if time == 0:
p = p0(x)
u = u0(x)
else:
p, u = leapfrog_step(p, u)
errorsp.append(np.linalg.norm((p - p_exact(x, time)), 2))
errorsu.append(np.linalg.norm((u - u_exact(x, time)), 2))
errorsp = np.array(errorsp) * dx ** (1 / 2)
errorsu = np.array(errorsu) * dx ** (1 / 2)
Ep.append(errorsp[-1])
Eu.append(errorsu[-1])
# plot the error
plt.figure(figsize=(8, 5))
plt.xlabel("$Nx$")
plt.ylabel(r'$\Vert p-\bar{p}\Vert_{L_2}$')
plt.loglog(NN, 15 / NN ** 2, "green", label=r'$O(\Delta x^{2})$')
plt.loglog(NN, Ep, "o", label=r'$E_p$')
plt.loglog(NN, Eu, "o", label=r'$E_u$')
plt.legend()
plt.show()
I would really appreciate it if someone could quickly check the implementation of the scheme and an indication on how to get the error plot.

Apart from the initialization, I see no errors in your code.
As to the initialization, consider the first step. There you should compute, per the method description, approximations for p(dt,j*dx) from the values of p(0,j*dx) and u(0.5*dt, (j+0.5)*dx). This means that you need to initialize at time==0
u = u_exact(x+0.5*dx, 0.5*dt).
and also need to compare the then obtained solution against u_exact(x+0.5*dx, time+0.5*dt).
That you obtained the correct order is IMO more an artefact of the test problem than an accidentially still correct algorithm.
If no exact solution is known, or if you want to use a more realistic algorithm in the test, you would need to compute the initial u values from p(0,x) and u(0,x) via Taylor expansions
u(t,x) = u(0,x) + t*u_t(0,x) + 0.5*t^2*u_tt(0,x) + ...
u(0.5*dt,x) = u(0,x) - 0.5*dt*p_x(0,x) + 0.125*dt^2*u_xx(0,x) + ...
= u(0,x) - 0.5*CFL*(p(0,x+0.5*dx)-p(0,x-0.5*dx))
+ 0.125*CFL^2*(u(0,x+dx)-2*u(0,x)+u(0,x-dx)) + ...
It might be sufficient to take just the linear expansion,
u[j] = u0(x[j]+0.5*dx) - 0.5*CFL*(p0(x[j]+dx)-p0(x[j])
or with array operations
p = p0(x)
u = u0(x+0.5*dx)
u[:-1] -= 0.5*CFL*(p[1:]-p[:-1])
u[-1]=u[0]
as then the second order error in the initial data just adds to the general second order error.
You might want to change the space grid to x = np.linspace(0, LX, NX+1) to have dx = LX/NX.
I would define the exact solution and the initial condition the other way around, as that allows more flexibility in the test problems.
# components of the solution
def f(x): return np.cos(2 * np.pi * x)
def g(x): return 2*np.sin(6 * np.pi * x)
# exact solution
def u_exact(x,t): return f(x+t)+g(x-t)
def p_exact(x,t): return -f(x+t)+g(x-t)
# Initial conditions
def u0(x): return u_exact(x,0)
def p0(x): return p_exact(x,0)

is there a way to examine if an angle is a right angle programmatically?

i plot an angle in python
here is the code
x = [0,0.5,1]
y = [0,0.5,0]
plt.scatter(x,y)
plt.plot(x,y)
plt.show()
is there a way to examine if the angle is a right angle programmatically?

The easiest way is to test if the dot product of the vectors is 0.
In your case, you simply compute:
v1 = ( (x[1]-x[0]), (y[1]-y[0]) ) <- (0.5, 0.5)
v2 = ( (x[2]-x[1]), (y[2]-y[1]) ) <- (0.5, -0.5)
dot_product = v1[0]*v2[0] + v1[1]*v2[1] <- 0.5² - 0.5² = 0

The other answers do not really care about possible inaccuracies and truncation errors, nor efficiency.
Rather than an exact comparison to 90° (or 0° in case of dot product), it is wiser to check for a small angle difference to 90° (resp. 0°).
Also wise to avoid divisions, square roots and trigonometric functions. The cross-product method is among the most attractive.
Compute the cross-product of the sides of the angle and their squared lengths, and compare
with a precomputed tolerance:
(ABx . BCy - ABy . BCx)² ≥ α.(ABx² + ABy²).(BCx² + BCy²)
with α = cos²δ where δ is the angle tolerance.

You can try to calculate the angle, but an easier way could be to check whether the Pythagorean Theorem applies. For that you'll need to calculate the size of the three edges and then check whether A^2 + B^2 ~= C^2

Yes, there is.
x = [0,0.5,1]
y = [0,0.5,0]
points = [np.array(point) for point in zip(x,y)]
a, b, c = points
ba = a - b
bc = c - b
cosine_angle = np.dot(ba, bc) / (np.linalg.norm(ba) * np.linalg.norm(bc))
angle_rad = np.arccos(cosine_angle)
angle_deg = np.rad2deg(angle_rad)
print(angle_deg) # 90.0

You can compute the angle between the two vectors as following: first, get the two vectors v1 and v2 and then use np.arccos() which returns the angle in radians. Convert it to degrees to check if it is 90 degrees. The formulae for computing angle between two vectors can be found on this Wiki link
import numpy as np
x = np.array([0,0.5,1])
y = np.array([0,0.5,0])
vecs = np.vstack((x, y))
v1 = vecs[:, 1] - vecs[:, 0]
v2 = vecs[:, 2] - vecs[:, 1]
angle_rad = np.arccos(np.dot(v1, v2) / (np.linalg.norm(v1) * np.linalg.norm(v2)))
angle_deg = np.rad2deg(angle_rad)
# 90.0

How can I vectorize this numpy rotation code?

The Goal:
I would like to vectorize (or otherwise speed up) this code. It rotates a 3d numpy model around its center point (let x,y,z denote the dimensions; then we want to rotate around the z-axis). The np model is binary voxels that are either "on" or "off"
I bet some basic matrix operation could do it, like take a layer and apply the rotation matrix to each element. The only issue with that is decimals; where should I have the new value land since cos(pi / 6) == sqrt(3) / 2?
The Code:
def rotate_model(m, theta):
'''
theta in degrees
'''
n =np.zeros(m.shape)
for i,layer in enumerate(m):
rotated = rotate(layer,theta)
n[i] = rotated
return n
where rotate() is:
def rotate(arr, theta):
'''
Rotates theta clockwise
rotated.shape == arr.shape, unlike scipy.ndimage.rotate(), which inflates size and also does some strange mixing
'''
if theta == int(theta):
theta *= pi / 180
theta = -theta
# theta=-theta b/c clockwise. Otherwise would default to counterclockwise
rotated =np.zeros(arr.shape)
#print rotated.shape[0], rotated.shape[1]
y_mid = arr.shape[0]//2
x_mid = arr.shape[1]//2
val = 0
for x_new in range(rotated.shape[1]):
for y_new in range(rotated.shape[0]):
x_centered = x_new - x_mid
y_centered = y_new - y_mid
x = x_centered*cos(theta) - y_centered*sin(theta)
y = x_centered*sin(theta) + y_centered*cos(theta)
x += x_mid
y += y_mid
x = int(round(x)); y = int(round(y)) # cast so range() picks it up
# lossy rotation
if x in range(arr.shape[1]) and y in range(arr.shape[0]):
val = arr[y,x]
rotated[y_new,x_new] = val
#print val
#print x,y
return rotated

You have a couple of problems in your code. First, if you want to fit the original image onto a rotated grid then you need a larger grid (usually). Alternatively, imagine a regular grid but the shape of your object - a rectangle - is rotated, thus becoming a "rhomb". It is obvious if you want to fit the entire rhomb - you need a larger output grid (array). On the other hand, you say in the code "rotated.shape == arr.shape, unlike scipy.ndimage.rotate(), which inflates size". If that is the case, maybe you do not want to fit the entire object? So, maybe it is OK to do this: rotated=np.zeros(arr.shape). But in general, yeah, one has to have a larger grid in order to fit the entire input image after it is rotated.
Another issue is angle conversion that you are doing:
if theta == int(theta):
theta *= pi / 180
theta = -theta
Why??? What will happen when I want to rotate the image by 1 radian? Or 2 radians? Am I forbidden to use integer number of radians? I think you are trying to do too much in this function and therefore it will be very confusing to do use it. Just require the caller to convert angles to radians. Or, you can do it inside this function if input theta is always in degrees. Or, you can add another parameter called, e.g., units and caller could set it to radians or degrees. Don't try to guess it based on "integer-ness" of input!
Now, let's rewrite your code a little bit:
rotated = np.zeros_like(arr) # instead of np.zero(arr.shape)
y_mid = arr.shape[0] // 2
x_mid = arr.shape[1] // 2
# val = 0 <- this is unnecessary
# pre-compute cos(theta) and sin(theta):
cs = cos(theta)
sn = sin(theta)
for x_new in range(rotated.shape[1]):
for y_new in range(rotated.shape[0]):
x = int(round((x_new - x_mid) * cs - (y_new - y_mid) * sn + x_mid)
y = int(round((x_new - x_mid) * sn - (y_new - y_mid) * cs + y_mid)
# just use comparisons, don't search through many values!
if 0 <= x < arr.shape[1] and 0 <= y < arr.shape[0]:
rotated[y_new, x_new] = arr[y, x]
So, now I can see (more easily) that for each pixel from the output array is mapped to a location in the input array. Yes, you can vectorize this.
import numpy as np
def rotate(arr, theta, unit='rad'):
# deal with theta units:
if unit.startswith('deg'):
theta = np.deg2rad(theta)
# for convenience, store array size:
ny, nx = arr.shape
# generate arrays of indices and flatten them:
y_new, x_new = np.indices(arr.shape)
x_new = x_new.ravel()
y_new = y_new.ravel()
# compute center of the array:
x0 = nx // 2
y0 = ny // 2
# compute old coordinates
xc = x_new - x0
yc = y_new - y0
x = np.round(np.cos(theta) * xc - np.sin(theta) * yc + x0).astype(np.int)
y = np.round(np.sin(theta) * xc - np.cos(theta) * yc + y0).astype(np.int)
# main idea to deal with indices is to create a mask:
mask = (x >= 0) & (x < nx) & (y >= 0) & (y < ny)
# ... and then select only those coordinates (both in
# input and "new" coordinates) that satisfy the above condition:
x = x[mask]
y = y[mask]
x_new = x_new[mask]
y_new = y_new[mask]
# map input values to output pixels *only* for selected "good" pixels:
rotated = np.zeros_like(arr)
rotated[y_new, x_new] = arr[y, x]
return rotated

Here is some code for anyone also doing 3d modeling. It solved my specific use-case pretty well. Still figuring out how to rotate in the proper plane. Hope it's helpful to you as well:
def rotate_model(m, theta):
'''
Redefines the prev 'rotate_model()' method
theta has to be in degrees
'''
rotated = scipy.ndimage.rotate(m, theta, axes=(1,2))
# have tried (1,0), (2,0), and now (1,2)
# ^ z is "up" and "2"
# scipy.ndimage.rotate() shrinks the model
# TODO: regrow it back
x_r = rotated.shape[1]
y_r = rotated.shape[0]
x_m = m.shape[1]
y_m = m.shape[0]
x_diff = abs(x_r - x_m)
y_diff = abs(y_r - y_m)
if x_diff%2==0 and y_diff%2==0:
return rotated[
x_diff//2 : x_r-x_diff//2,
y_diff//2 : y_r-y_diff//2,
:
]
elif x_diff%2==0 and y_diff%2==1:
# if this shift ends up turning the model to shit in a few iterations,
# change the following lines to include a flag that alternates cutting off the top and bottom bits of the array
return rotated[
x_diff//2 : x_r-x_diff//2,
y_diff//2+1 : y_r-y_diff//2,
:
]
elif x_diff%2==1 and y_diff%2==0:
return rotated[
x_diff//2+1 : x_r-x_diff//2,
y_diff//2 : y_r-y_diff//2,
:
]
else:
# x_diff%2==1 and y_diff%2==1:
return rotated[
x_diff//2+1 : x_r-x_diff//2,
y_diff//2+1 : y_r-y_diff//2,
:
]

3D Rotation Maths in Python

I'm writing code to make a machine orient a cutting tool in angle of a given vector. I need to rotate the 3D vector coordinates about the Y axis, I've got the equation figured out, but the code just breaks when I use it. Let me show you.
class Vector:
def __init__(self, i, j, k):
self.i = i
self.j = j
self.k = k
v0 = Vector(0.7844645405527362, 0.19611613513818404, 0.5883484054145521)
def findB():
if v0.i != 0:
angle = atan(v0.i/v0.k)
else:
angle = 0
print(angle, degrees(angle)) --> 0.92729... 53.13010...
v1 = Vector(v0.i, v0.j, v0.k)
# this formula rotates the vector about the j axis (comparable to y axis rotation)
v1.i = v0.i * cos(angle) - v0.k * sin(angle)
v1.k = v0.k * cos(angle) + v0.i * sin(angle)
return angle, v1
machine.b, v1 = findB()
print(v1.i, v1.j, v1.k) --> 1.1102230246251565e-16 0.19611613513818404 0.9805806756909203
I've tested that formula on a calculator and it rotates the coordinates correctly, I even hard coded the values to be:
v1.i = 0.784 * cos(0.927) - 0.588 * sin(0.927)
v1.k = 0.588 * cos(0.927) + 0.784 * sin(0.927)
Hard coding the values produces the correct result, but as soon as I use 'angle', v0.i and v0.k as the variables it gets v1.i incorrect. v1.i should be zero (or nearly zero from rounding). I've tested the same formula on a calculator and it works fine. Does anyone know what's going on?

Make a total vector for three vectors in python

We have a class that have three functions called(Bdisk, Bhalo,and BX).
all of these functions accept arrays (e.g. shape (1000))not matrices (e.g. shape (2,1000)).
I want to get the total of all these functions( total= Bdisk + Bhalo+BX), total these all functions give the magnetic field in all three components (B_r, B_phi, B_z) for thousand coordinate points (r, phi, z).
the code is here:
import numpy as np
import logging
import warnings
import gmf
signum = lambda x: (x < 0.) * -1. + (x >= 0) * 1.
pi = np.pi
#Class with analytical functions that describe the GMF according to the model of JF12
class GMF(object):
def __init__(self): # self:is automatically set to reference the newly created object that needs to be initialized
self.Rsun = -8.5 # position of the sun along the x axis in kpc
############################################################################
# Disk Parameters
############################################################################
self.bring, self.bring_unc = 0.1,0.1 # floats, field strength in ring at 3 kpc < r < 5 kpc
self.hdisk, self.hdisk_unc = 0.4, 0.03 # float, disk/halo transition height
self.wdisk, self.wdisk_unc = 0.27,0.08 # floats, transition width
self.b = np.array([0.1,3.,-0.9,-0.8,-2.0,-4.2,0.,2.7]) # (8,1)-dim np.arrays, field strength of spiral arms at 5 kpc
self.b_unc = np.array([1.8,0.6,0.8,0.3,0.1,0.5,1.8,1.8]) # uncertainty
self.rx = np.array([5.1,6.3,7.1,8.3,9.8,11.4,12.7,15.5])# (8,1)-dim np.array,dividing lines of spiral lines coordinates of neg. x-axes that intersect with arm
self.idisk = 11.5 * pi/180. # float, spiral arms pitch angle
#############################################################################
# Halo Parameters
#############################################################################
self.Bn, self.Bn_unc = 1.4,0.1 # floats, field strength northern halo
self.Bs, self.Bs_unc = -1.1,0.1 # floats, field strength southern halo
self.rn, self.rn_unc = 9.22,0.08 # floats, transition radius south, lower limit
self.rs, self.rs_unc = 16.7,0. # transition radius south, lower limit
self.whalo, self.whalo_unc = 0.2,0.12 # floats, transition width
self.z0, self.z0_unc = 5.3, 1.6 # floats, vertical scale height
##############################################################################
# Out of plaxe or "X" component Parameters
##############################################################################
self.BX0, self.BX_unc = 4.6,0.3 # floats, field strength at origin
self.ThetaX0, self.ThetaX0_unc = 49. * pi/180., pi/180. # elev. angle at z = 0, r > rXc
self.rXc, self.rXc_unc = 4.8, 0.2 # floats, radius where thetaX = thetaX0
self.rX, self.rX_unc = 2.9, 0.1 # floats, exponential scale length
# striated field
self.gamma, self.gamma_unc = 2.92,0.14 # striation and/or rel. elec. number dens. rescaling
return
##################################################################################
##################################################################################
# Transition function given by logistic function eq.5
##################################################################################
def L(self,z,h,w):
if np.isscalar(z):
z = np.array([z]) # scalar or numpy array with positions (height above disk, z; distance from center, r)
ones = np.ones(z.shape[0])
return 1./(ones + np.exp(-2. *(np.abs(z)- h)/w))
####################################################################################
# return distance from center for angle phi of logarithmic spiral
# r(phi) = rx * exp(b * phi) as np.array
####################################################################################
def r_log_spiral(self,phi):
if np.isscalar(phi): #Returns True if the type of num is a scalar type.
phi = np.array([phi])
ones = np.ones(phi.shape[0])
# self.rx.shape = 8
# phi.shape = p
# then result is given as (8,p)-dim array, each row stands for one rx
# vstack : Take a sequence of arrays and stack them vertically to make a single array
# tensordot(a, b, axes=2):Compute tensor dot product along specified axes for arrays >=1D.
result = np.tensordot(self.rx , np.exp((phi - 3.*pi*ones) / np.tan(pi/2. - self.idisk)),axes = 0)
result = np.vstack((result, np.tensordot(self.rx , np.exp((phi - pi*ones) / np.tan(pi/2. - self.idisk)),axes = 0) ))
result = np.vstack((result, np.tensordot(self.rx , np.exp((phi + pi*ones) / np.tan(pi/2. - self.idisk)),axes = 0) ))
return np.vstack((result, np.tensordot(self.rx , np.exp((phi + 3.*pi*ones) / np.tan(pi/2. - self.idisk)),axes = 0) ))
#############################################################################################
# Disk component in galactocentric cylindrical coordinates (r,phi,z)
#############################################################################################
def Bdisk(self,r,phi,z):
# Bdisk is purely azimuthal (toroidal) with the field strength b_ring
"""
r: N-dim np.array, distance from origin in GC cylindrical coordinates, is in kpc
z: N-dim np.array, height in kpc in GC cylindrical coordinates
phi:N-dim np.array, polar angle in GC cylindircal coordinates, in radian
Bdisk: (3,N)-dim np.array with (r,phi,z) components of disk field for each coordinate tuple
|Bdisk|: N-dim np.array, absolute value of Bdisk for each coordinate tuple
"""
if (not r.shape[0] == phi.shape[0]) and (not z.shape[0] == phi.shape[0]):
warnings.warn("List do not have equal shape! returning -1", RuntimeWarning)
return -1
# Return a new array of given shape and type, filled with zeros.
Bdisk = np.zeros((3,r.shape[0])) # Bdisk vector in r, phi, z
ones = np.ones(r.shape[0])
r_center = (r >= 3.) & (r < 5.1)
r_disk = (r >= 5.1) & (r <= 20.)
Bdisk[1,r_center] = self.bring
# Determine in which arm we are
# this is done for each coordinate individually
if np.sum(r_disk):
rls = self.r_log_spiral(phi[r_disk])
rls = np.abs(rls - r[r_disk])
arms = np.argmin(rls, axis = 0) % 8
# The magnetic spiral defined at r=5 kpc and fulls off as 1/r ,the field direction is given by:
Bdisk[0,r_disk] = np.sin(self.idisk)* self.b[arms] * (5. / r[r_disk])
Bdisk[1,r_disk] = np.cos(self.idisk)* self.b[arms] * (5. / r[r_disk])
Bdisk *= (ones - self.L(z,self.hdisk,self.wdisk)) # multiplied by L
return Bdisk, np.sqrt(np.sum(Bdisk**2.,axis = 0)) # the Bdisk, the normalization
# axis=0 : sum over index 0(row)
# axis=1 : sum over index 1(columns)
##############################################################################################
# Halo component
###############################################################################################
def Bhalo(self,r,z):
# Bhalo is purely azimuthal (toroidal), i.e. has only a phi component
if (not r.shape[0] == z.shape[0]):
warnings.warn("List do not have equal shape! returning -1", RuntimeWarning)
return -1
Bhalo = np.zeros((3,r.shape[0])) # Bhalo vector in r, phi, z rows: r, phi and z component
ones = np.ones(r.shape[0])
m = ( z != 0. )
# SEE equation 6.
Bhalo[1,m] = np.exp(-np.abs(z[m])/self.z0) * self.L(z[m], self.hdisk, self.wdisk) * \
( self.Bn * (ones[m] - self.L(r[m], self.rn, self.whalo)) * (z[m] > 0.) \
+ self.Bs * (ones[m] - self.L(r[m], self.rs, self.whalo)) * (z[m] < 0.) )
return Bhalo , np.sqrt(np.sum(Bhalo**2.,axis = 0))
##############################################################################################
# BX component (OUT OF THE PLANE)
###############################################################################################
def BX(self,r,z):
#BX is purely ASS and poloidal, i.e. phi component = 0
if (not r.shape[0] == z.shape[0]):
warnings.warn("List do not have equal shape! returning -1", RuntimeWarning)
return -1
BX= np.zeros((3,r.shape[0])) # BX vector in r, phi, z rows: r, phi and z component
m = np.sqrt(r**2. + z**2.) >= 1.
bx = lambda r_p: self.BX0 * np.exp(-r_p / self.rX) # eq.7
thetaX = lambda r,z,r_p: np.arctan(np.abs(z)/(r - r_p)) # eq.10
r_p = r[m] *self.rXc/(self.rXc + np.abs(z[m] ) / np.tan(self.ThetaX0)) # eq 9
m_r_b = r_p > self.rXc # region with constant elevation angle
m_r_l = r_p <= self.rXc # region with varying elevation angle
theta = np.zeros(z[m].shape[0])
b = np.zeros(z[m].shape[0])
r_p0 = (r[m])[m_r_b] - np.abs( (z[m])[m_r_b] ) / np.tan(self.ThetaX0) # eq.8
b[m_r_b] = bx(r_p0) * r_p0/ (r[m])[m_r_b] # the field strength in the constant elevation angle (b_x(r_p)r_p/r)
theta[m_r_b] = self.ThetaX0 * np.ones(theta.shape[0])[m_r_b]
b[m_r_l] = bx(r_p[m_r_l]) * (r_p[m_r_l]/(r[m])[m_r_l] )**2. # the field strength with varying elevation angle (b_x(r_p)(r_p/r)**2)
theta[m_r_l] = thetaX((r[m])[m_r_l] ,(z[m])[m_r_l] ,r_p[m_r_l])
mz = (z[m] == 0.)
theta[mz] = np.pi/2.
BX[0,m] = b * (np.cos(theta) * (z[m] >= 0) + np.cos(pi*np.ones(theta.shape[0]) - theta) * (z[m] < 0))
BX[2,m] = b * (np.sin(theta) * (z[m] >= 0) + np.sin(pi*np.ones(theta.shape[0]) - theta) * (z[m] < 0))
return BX, np.sqrt(np.sum(BX**2.,axis=0))
then, I create three arrays, one for r, one for phi, one for z. Each of these arrays has (e.g: thousand elements). like this:
import gmf
gmfm = gmf.GMF()
x = np.linspace(-20.,20.,100)
y = np.linspace(-20.,20.,100)
z = np.linspace(-1.,1.,x.shape[0])
xx,yy = np.meshgrid(x,y)
rr = np.sqrt(xx**2. + yy**2.)
theta = np.arctan2(yy,xx)
for i,r in enumerate(rr[:]):
Bdisk, Babs_d = gmfm.Bdisk(r,theta[i],z)
Bhalo, Babs_h = gmfm.Bhalo(r,z)
BX, Babs_x = gmfm.BX(r,z)
Btotal = Bdisk + Bhalo + BX
but I am getting when I make the addition of the three functions Btotal= Bdisk + Bhalo+BX) in 2d matrix with 3 rows and 100 columns.
My question is how can I add these three functions together to get Btotal in shape (n,) e.g( shape(100,)
because as I said in the beginning the three functions accept accept arrays (e.g. shape (1000) )then when we adding the three functions together we have to get the total also in the same shape (shape (n,)?
I do not know how can I do it, could you please tell me how can I make it.
thank you for your cooperation.

You need to correct the indention, for example in the def Bdisk method.
More importantly in
for i,r in enumerate(rr[:]):
Bdisk, Babs_d = gmfm.Bdisk(r,theta[i],z)
Bhalo, Babs_h = gmfm.Bhalo(r,z)
BX, Babs_x = gmfm.BX(r,z)
Btotal = Bdisk + Bhalo + BX
are you doing this addition for each iteration, or once at the end of the loop? You aren't accumulating any values over iterations. You are just throwing away the old ones, leaving you with the final iteration.
As for adding the array - it appears that all your arrays are initialed like:
Bdisk = np.zeros((3,r.shape[0]))
If that's what the method returns, then
Bdisk + Bhalo + BX
will just sum the corresponding elements of each array, resulting in a Btotal with the same shape. If you don't not like the shape of Btotal then change how Bdisk is calculated, because it has the same shape.