在python中为三个向量创建一个总向量

时间:2016-10-28 14:04:40

标签: python numpy

我们有一个有三个函数的类(Bdisk,Bhalo和BX)。 所有这些函数都接受数组(例如形状(1000))而不是矩阵(例如形状(2,1000))。 我想得到所有这些函数的总和(total = Bdisk + Bhalo + BX),总共这些函数给出了千位坐标点(r,phi,z)的所有三个分量(B_r,B_phi,B_z)中的磁场

代码在这里:

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))	

然后,我创建了三个数组,一个用于r,一个用于phi,一个用于z。这些阵列中的每一个都具有(例如:千元素)。像这样:

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 

但是当我在2d矩阵中添加3行100列的三个函数Btotal = Bdisk + Bhalo + BX时,我得到了。

我的问题是如何将这三个函数组合在一起以获得Btotal的形状(n,),例如(shape(100,) 因为正如我在开头所说的那样,三个函数接受接受数组(例如形状(1000))然后当我们将三个函数加在一起时,我们必须得到总数也是相同的形状(形状(n,)? 我不知道怎么办,你能告诉我怎样才能做到。 谢谢您的合作。

1 个答案:

答案 0 :(得分:0)

您需要更正缩进,例如在def Bdisk方法中。

更重要的是

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 

你是为每次迭代做这个添加,还是在循环结束时做一次?您不会在迭代中累积任何值。你只是扔掉旧的,留下最后的迭代。

至于添加数组 - 似乎所有数组的初始化都是:

Bdisk = np.zeros((3,r.shape[0]))

如果这是方法返回的内容,那么

Bdisk + Bhalo + BX

将只对每个数组的相应元素求和,从而得到具有相同形状的Btotal。如果您不喜欢Btotal的形状,请更改Bdisk的计算方式,因为它具有相同的形状。