def _evaluate(self,*args,**kwargs):
"""
NAME:
__call__ (_evaluate)
PURPOSE:
evaluate the actions (jr,lz,jz)
INPUT:
Either:
a) R,vR,vT,z,vz[,phi]:
1) floats: phase-space value for single object (phi is optional) (each can be a Quantity)
2) numpy.ndarray: [N] phase-space values for N objects (each can be a Quantity)
b) Orbit instance: initial condition used if that's it, orbit(t) if there is a time given as well as the second argument
cumul= if True, return the cumulative average actions (to look
at convergence)
OUTPUT:
(jr,lz,jz)
HISTORY:
2013-09-10 - Written - Bovy (IAS)
"""
R,vR,vT,z,vz,phi= self._parse_args(False,False,*args)
if self._c: #pragma: no cover
pass
else:
#Use self._aAI to calculate the actions and angles in the isochrone potential
acfs= self._aAI._actionsFreqsAngles(R.flatten(),
vR.flatten(),
vT.flatten(),
z.flatten(),
vz.flatten(),
phi.flatten())
jrI= nu.reshape(acfs[0],R.shape)[:,:-1]
jzI= nu.reshape(acfs[2],R.shape)[:,:-1]
anglerI= nu.reshape(acfs[6],R.shape)
anglezI= nu.reshape(acfs[8],R.shape)
if nu.any((nu.fabs(nu.amax(anglerI,axis=1)-_TWOPI) > _ANGLETOL)\
*(nu.fabs(nu.amin(anglerI,axis=1)) > _ANGLETOL)): #pragma: no cover
warnings.warn("Full radial angle range not covered for at least one object; actions are likely not reliable",galpyWarning)
if nu.any((nu.fabs(nu.amax(anglezI,axis=1)-_TWOPI) > _ANGLETOL)\
*(nu.fabs(nu.amin(anglezI,axis=1)) > _ANGLETOL)): #pragma: no cover
warnings.warn("Full vertical angle range not covered for at least one object; actions are likely not reliable",galpyWarning)
danglerI= ((nu.roll(anglerI,-1,axis=1)-anglerI) % _TWOPI)[:,:-1]
danglezI= ((nu.roll(anglezI,-1,axis=1)-anglezI) % _TWOPI)[:,:-1]
if kwargs.get('cumul',False):
sumFunc= nu.cumsum
else:
sumFunc= nu.sum
jr= sumFunc(jrI*danglerI,axis=1)/sumFunc(danglerI,axis=1)
jz= sumFunc(jzI*danglezI,axis=1)/sumFunc(danglezI,axis=1)
if _isNonAxi(self._pot):
lzI= nu.reshape(acfs[1],R.shape)[:,:-1]
anglephiI= nu.reshape(acfs[7],R.shape)
danglephiI= ((nu.roll(anglephiI,-1,axis=1)-anglephiI) % _TWOPI)[:,:-1]
if nu.any((nu.fabs(nu.amax(anglephiI,axis=1)-_TWOPI) > _ANGLETOL)\
*(nu.fabs(nu.amin(anglephiI,axis=1)) > _ANGLETOL)): #pragma: no cover
warnings.warn("Full azimuthal angle range not covered for at least one object; actions are likely not reliable",galpyWarning)
lz= sumFunc(lzI*danglephiI,axis=1)/sumFunc(danglephiI,axis=1)
else:
lz= R[:,0]*vT[:,0]
return (jr,lz,jz)
def _evaluate(self,*args,**kwargs):
"""
NAME:
__call__ (_evaluate)
PURPOSE:
evaluate the actions (jr,lz,jz)
INPUT:
Either:
a) R,vR,vT,z,vz[,phi]:
1) floats: phase-space value for single object (phi is optional) (each can be a Quantity)
2) numpy.ndarray: [N] phase-space values for N objects (each can be a Quantity)
b) Orbit instance: initial condition used if that's it, orbit(t) if there is a time given as well as the second argument
cumul= if True, return the cumulative average actions (to look
at convergence)
OUTPUT:
(jr,lz,jz)
HISTORY:
2013-09-10 - Written - Bovy (IAS)
"""
R,vR,vT,z,vz,phi= self._parse_args(False,False,*args)
if self._c: #pragma: no cover
pass
else:
#Use self._aAI to calculate the actions and angles in the isochrone potential
acfs= self._aAI._actionsFreqsAngles(R.flatten(),
vR.flatten(),
vT.flatten(),
z.flatten(),
vz.flatten(),
phi.flatten())
jrI= nu.reshape(acfs[0],R.shape)[:,:-1]
jzI= nu.reshape(acfs[2],R.shape)[:,:-1]
anglerI= nu.reshape(acfs[6],R.shape)
anglezI= nu.reshape(acfs[8],R.shape)
if nu.any((nu.fabs(nu.amax(anglerI,axis=1)-_TWOPI) > _ANGLETOL)\
*(nu.fabs(nu.amin(anglerI,axis=1)) > _ANGLETOL)): #pragma: no cover
warnings.warn("Full radial angle range not covered for at least one object; actions are likely not reliable",galpyWarning)
if nu.any((nu.fabs(nu.amax(anglezI,axis=1)-_TWOPI) > _ANGLETOL)\
*(nu.fabs(nu.amin(anglezI,axis=1)) > _ANGLETOL)): #pragma: no cover
warnings.warn("Full vertical angle range not covered for at least one object; actions are likely not reliable",galpyWarning)
danglerI= ((nu.roll(anglerI,-1,axis=1)-anglerI) % _TWOPI)[:,:-1]
danglezI= ((nu.roll(anglezI,-1,axis=1)-anglezI) % _TWOPI)[:,:-1]
if kwargs.get('cumul',False):
sumFunc= nu.cumsum
else:
sumFunc= nu.sum
jr= sumFunc(jrI*danglerI,axis=1)/sumFunc(danglerI,axis=1)
jz= sumFunc(jzI*danglezI,axis=1)/sumFunc(danglezI,axis=1)
if _isNonAxi(self._pot):
lzI= nu.reshape(acfs[1],R.shape)[:,:-1]
anglephiI= nu.reshape(acfs[7],R.shape)
danglephiI= ((nu.roll(anglephiI,-1,axis=1)-anglephiI) % _TWOPI)[:,:-1]
if nu.any((nu.fabs(nu.amax(anglephiI,axis=1)-_TWOPI) > _ANGLETOL)\
*(nu.fabs(nu.amin(anglephiI,axis=1)) > _ANGLETOL)): #pragma: no cover
warnings.warn("Full azimuthal angle range not covered for at least one object; actions are likely not reliable",galpyWarning)
lz= sumFunc(lzI*danglephiI,axis=1)/sumFunc(danglephiI,axis=1)
else:
lz= R[:,0]*vT[:,0]
return (jr,lz,jz)
def __init__(self, *args, **kwargs):
"""
NAME:
__init__
PURPOSE:
initialize an actionAngleTorus object
INPUT:
pot= potential or list of potentials (3D)
tol= default tolerance to use when fitting tori (|dJ|/J)
dJ= default action difference when computing derivatives (Hessian or Jacobian)
OUTPUT:
instance
HISTORY:
2015-08-07 - Written - Bovy (UofT)
"""
if not 'pot' in kwargs: #pragma: no cover
raise IOError("Must specify pot= for actionAngleTorus")
self._pot = kwargs['pot']
if _isNonAxi(self._pot):
raise RuntimeError(
"actionAngleTorus for non-axisymmetric potentials is not supported"
)
if self._pot == MWPotential:
warnings.warn(
"Use of MWPotential as a Milky-Way-like potential is deprecated; galpy.potential.MWPotential2014, a potential fit to a large variety of dynamical constraints (see Bovy 2015), is the preferred Milky-Way-like potential in galpy",
galpyWarning)
if ext_loaded:
self._c = _check_c(self._pot)
if not self._c:
raise RuntimeError(
'The given potential is not fully implemented in C; using the actionAngleTorus code is not supported in pure Python'
)
else: # pragma: no cover
raise RuntimeError(
'actionAngleTorus instances cannot be used, because the actionAngleTorus_c extension failed to load'
)
self._tol = kwargs.get('tol', 0.001)
self._dJ = kwargs.get('dJ', 0.001)
return None
def __init__(self,*args,**kwargs):
"""
NAME:
__init__
PURPOSE:
initialize an actionAngleTorus object
INPUT:
pot= potential or list of potentials (3D)
tol= default tolerance to use when fitting tori (|dJ|/J)
dJ= default action difference when computing derivatives (Hessian or Jacobian)
OUTPUT:
instance
HISTORY:
2015-08-07 - Written - Bovy (UofT)
"""
if not 'pot' in kwargs: #pragma: no cover
raise IOError("Must specify pot= for actionAngleTorus")
self._pot= kwargs['pot']
if _isNonAxi(self._pot):
raise RuntimeError("actionAngleTorus for non-axisymmetric potentials is not supported")
if self._pot == MWPotential:
warnings.warn("Use of MWPotential as a Milky-Way-like potential is deprecated; galpy.potential.MWPotential2014, a potential fit to a large variety of dynamical constraints (see Bovy 2015), is the preferred Milky-Way-like potential in galpy",
galpyWarning)
if ext_loaded:
self._c= _check_c(self._pot)
if not self._c:
raise RuntimeError('The given potential is not fully implemented in C; using the actionAngleTorus code is not supported in pure Python')
else:# pragma: no cover
raise RuntimeError('actionAngleTorus instances cannot be used, because the actionAngleTorus_c extension failed to load')
self._tol= kwargs.get('tol',0.001)
self._dJ= kwargs.get('dJ',0.001)
return None
def _actionsFreqsAngles(self,*args,**kwargs):
"""
NAME:
_actionsFreqsAngles
PURPOSE:
evaluate the actions, frequencies, and angles
(jr,lz,jz,Omegar,Omegaphi,Omegaz,angler,anglephi,anglez)
INPUT:
Either:
a) R,vR,vT,z,vz:
1) floats: phase-space value for single object
2) numpy.ndarray: [N] phase-space values for N objects
3) numpy.ndarray: [N,M] phase-space values for N objects at M
times
b) Orbit instance or list thereof; can be integrated already
maxn= (default: object-wide default) Use a grid in vec(n) up to this n (zero-based)
ts= if set, the phase-space points correspond to these times (IF NOT SET, WE ASSUME THAT ts IS THAT THAT IS ASSOCIATED WITH THIS OBJECT)
_firstFlip= (False) if True and Orbits are given, the backward part of the orbit is integrated first and stored in the Orbit object
OUTPUT:
(jr,lz,jz,Omegar,Omegaphi,Omegaz,angler,anglephi,anglez)
HISTORY:
2013-09-10 - Written - Bovy (IAS)
"""
from galpy.orbit import Orbit
_firstFlip= kwargs.get('_firstFlip',False)
#If the orbit was already integrated, set ts to the integration times
if isinstance(args[0],Orbit) and hasattr(args[0]._orb,'orbit') \
and not 'ts' in kwargs:
kwargs['ts']= args[0]._orb.t
elif (isinstance(args[0],list) and isinstance(args[0][0],Orbit)) \
and hasattr(args[0][0]._orb,'orbit') \
and not 'ts' in kwargs:
kwargs['ts']= args[0][0]._orb.t
R,vR,vT,z,vz,phi= self._parse_args(True,_firstFlip,*args)
if 'ts' in kwargs and not kwargs['ts'] is None:
ts= kwargs['ts']
if _APY_LOADED and isinstance(ts,units.Quantity):
ts= ts.to(units.Gyr).value\
/time_in_Gyr(self._vo,self._ro)
else:
ts= nu.empty(R.shape[1])
ts[self._ntintJ-1:]= self._tsJ
ts[:self._ntintJ-1]= -self._tsJ[1:][::-1]
maxn= kwargs.get('maxn',self._maxn)
if self._c: #pragma: no cover
pass
else:
#Use self._aAI to calculate the actions and angles in the isochrone potential
if '_acfs' in kwargs: acfs= kwargs['_acfs']
else:
acfs= self._aAI._actionsFreqsAngles(R.flatten(),
vR.flatten(),
vT.flatten(),
z.flatten(),
vz.flatten(),
phi.flatten())
jrI= nu.reshape(acfs[0],R.shape)[:,:-1]
jzI= nu.reshape(acfs[2],R.shape)[:,:-1]
anglerI= nu.reshape(acfs[6],R.shape)
anglezI= nu.reshape(acfs[8],R.shape)
if nu.any((nu.fabs(nu.amax(anglerI,axis=1)-_TWOPI) > _ANGLETOL)\
*(nu.fabs(nu.amin(anglerI,axis=1)) > _ANGLETOL)): #pragma: no cover
warnings.warn("Full radial angle range not covered for at least one object; actions are likely not reliable",galpyWarning)
if nu.any((nu.fabs(nu.amax(anglezI,axis=1)-_TWOPI) > _ANGLETOL)\
*(nu.fabs(nu.amin(anglezI,axis=1)) > _ANGLETOL)): #pragma: no cover
warnings.warn("Full vertical angle range not covered for at least one object; actions are likely not reliable",galpyWarning)
danglerI= ((nu.roll(anglerI,-1,axis=1)-anglerI) % _TWOPI)[:,:-1]
danglezI= ((nu.roll(anglezI,-1,axis=1)-anglezI) % _TWOPI)[:,:-1]
jr= nu.sum(jrI*danglerI,axis=1)/nu.sum(danglerI,axis=1)
jz= nu.sum(jzI*danglezI,axis=1)/nu.sum(danglezI,axis=1)
if _isNonAxi(self._pot): #pragma: no cover
lzI= nu.reshape(acfs[1],R.shape)[:,:-1]
anglephiI= nu.reshape(acfs[7],R.shape)
if nu.any((nu.fabs(nu.amax(anglephiI,axis=1)-_TWOPI) > _ANGLETOL)\
*(nu.fabs(nu.amin(anglephiI,axis=1)) > _ANGLETOL)): #pragma: no cover
warnings.warn("Full azimuthal angle range not covered for at least one object; actions are likely not reliable",galpyWarning)
danglephiI= ((nu.roll(anglephiI,-1,axis=1)-anglephiI) % _TWOPI)[:,:-1]
lz= nu.sum(lzI*danglephiI,axis=1)/nu.sum(danglephiI,axis=1)
else:
lz= R[:,len(ts)//2]*vT[:,len(ts)//2]
#Now do an 'angle-fit'
angleRT= dePeriod(nu.reshape(acfs[6],R.shape))
acfs7= nu.reshape(acfs[7],R.shape)
negFreqIndx= nu.median(acfs7-nu.roll(acfs7,1,axis=1),axis=1) < 0. #anglephi is decreasing
anglephiT= nu.empty(acfs7.shape)
anglephiT[negFreqIndx,:]= dePeriod(_TWOPI-acfs7[negFreqIndx,:])
negFreqPhi= nu.zeros(R.shape[0],dtype='bool')
negFreqPhi[negFreqIndx]= True
anglephiT[True-negFreqIndx,:]= dePeriod(acfs7[True-negFreqIndx,:])
angleZT= dePeriod(nu.reshape(acfs[8],R.shape))
#Write the angle-fit as Y=AX, build A and Y
nt= len(ts)
no= R.shape[0]
#remove 0,0,0 and half-plane
if _isNonAxi(self._pot):
nn= (2*maxn-1)**2*maxn-(maxn-1)*(2*maxn-1)-maxn
else:
nn= maxn*(2*maxn-1)-maxn
A= nu.zeros((no,nt,2+nn))
A[:,:,0]= 1.
A[:,:,1]= ts
#sorting the phi and Z grids this way makes it easy to exclude the origin
phig= list(nu.arange(-maxn+1,maxn,1))
phig.sort(key = lambda x: abs(x))
phig= nu.array(phig,dtype='int')
if _isNonAxi(self._pot):
grid= nu.meshgrid(nu.arange(maxn),phig,phig)
else:
grid= nu.meshgrid(nu.arange(maxn),phig)
gridR= grid[0].T.flatten()[1:] #remove 0,0,0
gridZ= grid[1].T.flatten()[1:]
mask = nu.ones(len(gridR),dtype=bool)
# excludes axis that is not in half-space
if _isNonAxi(self._pot):
gridphi= grid[2].T.flatten()[1:]
mask= True\
-(gridR == 0)*((gridphi < 0)+((gridphi==0)*(gridZ < 0)))
else:
mask[:2*maxn-3:2]= False
gridR= gridR[mask]
gridZ= gridZ[mask]
tangleR= nu.tile(angleRT.T,(nn,1,1)).T
tgridR= nu.tile(gridR,(no,nt,1))
tangleZ= nu.tile(angleZT.T,(nn,1,1)).T
tgridZ= nu.tile(gridZ,(no,nt,1))
if _isNonAxi(self._pot):
gridphi= gridphi[mask]
tgridphi= nu.tile(gridphi,(no,nt,1))
tanglephi= nu.tile(anglephiT.T,(nn,1,1)).T
sinnR= nu.sin(tgridR*tangleR+tgridphi*tanglephi+tgridZ*tangleZ)
else:
sinnR= nu.sin(tgridR*tangleR+tgridZ*tangleZ)
A[:,:,2:]= sinnR
#Matrix magic
atainv= nu.empty((no,2+nn,2+nn))
AT= nu.transpose(A,axes=(0,2,1))
for ii in range(no):
atainv[ii,:,:,]= linalg.inv(nu.dot(AT[ii,:,:],A[ii,:,:]))
ATAR= nu.sum(AT*nu.transpose(nu.tile(angleRT,(2+nn,1,1)),axes=(1,0,2)),axis=2)
ATAT= nu.sum(AT*nu.transpose(nu.tile(anglephiT,(2+nn,1,1)),axes=(1,0,2)),axis=2)
ATAZ= nu.sum(AT*nu.transpose(nu.tile(angleZT,(2+nn,1,1)),axes=(1,0,2)),axis=2)
angleR= nu.sum(atainv[:,0,:]*ATAR,axis=1)
OmegaR= nu.sum(atainv[:,1,:]*ATAR,axis=1)
anglephi= nu.sum(atainv[:,0,:]*ATAT,axis=1)
Omegaphi= nu.sum(atainv[:,1,:]*ATAT,axis=1)
angleZ= nu.sum(atainv[:,0,:]*ATAZ,axis=1)
OmegaZ= nu.sum(atainv[:,1,:]*ATAZ,axis=1)
Omegaphi[negFreqIndx]= -Omegaphi[negFreqIndx]
anglephi[negFreqIndx]= _TWOPI-anglephi[negFreqIndx]
if kwargs.get('_retacfs',False):
return (jr,lz,jz,OmegaR,Omegaphi,OmegaZ, #pragma: no cover
angleR % _TWOPI,
anglephi % _TWOPI,
angleZ % _TWOPI,acfs)
else:
return (jr,lz,jz,OmegaR,Omegaphi,OmegaZ,
angleR % _TWOPI,
anglephi % _TWOPI,
angleZ % _TWOPI)
def _actionsFreqsAngles(self,*args,**kwargs):
"""
NAME:
actionsFreqsAngles (_actionsFreqsAngles)
PURPOSE:
evaluate the actions, frequencies, and angles (jr,lz,jz,Omegar,Omegaphi,Omegaz,angler,anglephi,anglez)
INPUT:
Either:
a) R,vR,vT,z,vz[,phi]:
1) floats: phase-space value for single object (phi is optional) (each can be a Quantity)
2) numpy.ndarray: [N] phase-space values for N objects (each can be a Quantity)
b) Orbit instance: initial condition used if that's it, orbit(t) if there is a time given as well as the second argument
maxn= (default: object-wide default) Use a grid in vec(n) up to this n (zero-based)
ts= if set, the phase-space points correspond to these times (IF NOT SET, WE ASSUME THAT ts IS THAT THAT IS ASSOCIATED WITH THIS OBJECT)
_firstFlip= (False) if True and Orbits are given, the backward part of the orbit is integrated first and stored in the Orbit object
OUTPUT:
(jr,lz,jz,Omegar,Omegaphi,Omegaz,angler,anglephi,anglez)
HISTORY:
2013-09-10 - Written - Bovy (IAS)
"""
from galpy.orbit import Orbit
_firstFlip= kwargs.get('_firstFlip',False)
#If the orbit was already integrated, set ts to the integration times
if isinstance(args[0],Orbit) and hasattr(args[0]._orb,'orbit') \
and not 'ts' in kwargs:
kwargs['ts']= args[0]._orb.t
elif (isinstance(args[0],list) and isinstance(args[0][0],Orbit)) \
and hasattr(args[0][0]._orb,'orbit') \
and not 'ts' in kwargs:
kwargs['ts']= args[0][0]._orb.t
R,vR,vT,z,vz,phi= self._parse_args(True,_firstFlip,*args)
if 'ts' in kwargs and not kwargs['ts'] is None:
ts= kwargs['ts']
if _APY_LOADED and isinstance(ts,units.Quantity):
ts= ts.to(units.Gyr).value\
/time_in_Gyr(self._vo,self._ro)
else:
ts= nu.empty(R.shape[1])
ts[self._ntintJ-1:]= self._tsJ
ts[:self._ntintJ-1]= -self._tsJ[1:][::-1]
maxn= kwargs.get('maxn',self._maxn)
if self._c: #pragma: no cover
pass
else:
#Use self._aAI to calculate the actions and angles in the isochrone potential
if '_acfs' in kwargs: acfs= kwargs['_acfs']
else:
acfs= self._aAI._actionsFreqsAngles(R.flatten(),
vR.flatten(),
vT.flatten(),
z.flatten(),
vz.flatten(),
phi.flatten())
jrI= nu.reshape(acfs[0],R.shape)[:,:-1]
jzI= nu.reshape(acfs[2],R.shape)[:,:-1]
anglerI= nu.reshape(acfs[6],R.shape)
anglezI= nu.reshape(acfs[8],R.shape)
if nu.any((nu.fabs(nu.amax(anglerI,axis=1)-_TWOPI) > _ANGLETOL)\
*(nu.fabs(nu.amin(anglerI,axis=1)) > _ANGLETOL)): #pragma: no cover
warnings.warn("Full radial angle range not covered for at least one object; actions are likely not reliable",galpyWarning)
if nu.any((nu.fabs(nu.amax(anglezI,axis=1)-_TWOPI) > _ANGLETOL)\
*(nu.fabs(nu.amin(anglezI,axis=1)) > _ANGLETOL)): #pragma: no cover
warnings.warn("Full vertical angle range not covered for at least one object; actions are likely not reliable",galpyWarning)
danglerI= ((nu.roll(anglerI,-1,axis=1)-anglerI) % _TWOPI)[:,:-1]
danglezI= ((nu.roll(anglezI,-1,axis=1)-anglezI) % _TWOPI)[:,:-1]
jr= nu.sum(jrI*danglerI,axis=1)/nu.sum(danglerI,axis=1)
jz= nu.sum(jzI*danglezI,axis=1)/nu.sum(danglezI,axis=1)
if _isNonAxi(self._pot): #pragma: no cover
lzI= nu.reshape(acfs[1],R.shape)[:,:-1]
anglephiI= nu.reshape(acfs[7],R.shape)
if nu.any((nu.fabs(nu.amax(anglephiI,axis=1)-_TWOPI) > _ANGLETOL)\
*(nu.fabs(nu.amin(anglephiI,axis=1)) > _ANGLETOL)): #pragma: no cover
warnings.warn("Full azimuthal angle range not covered for at least one object; actions are likely not reliable",galpyWarning)
danglephiI= ((nu.roll(anglephiI,-1,axis=1)-anglephiI) % _TWOPI)[:,:-1]
lz= nu.sum(lzI*danglephiI,axis=1)/nu.sum(danglephiI,axis=1)
else:
lz= R[:,len(ts)//2]*vT[:,len(ts)//2]
#Now do an 'angle-fit'
angleRT= dePeriod(nu.reshape(acfs[6],R.shape))
acfs7= nu.reshape(acfs[7],R.shape)
negFreqIndx= nu.median(acfs7-nu.roll(acfs7,1,axis=1),axis=1) < 0. #anglephi is decreasing
anglephiT= nu.empty(acfs7.shape)
anglephiT[negFreqIndx,:]= dePeriod(_TWOPI-acfs7[negFreqIndx,:])
negFreqPhi= nu.zeros(R.shape[0],dtype='bool')
negFreqPhi[negFreqIndx]= True
anglephiT[True^negFreqIndx,:]= dePeriod(acfs7[True^negFreqIndx,:])
angleZT= dePeriod(nu.reshape(acfs[8],R.shape))
#Write the angle-fit as Y=AX, build A and Y
nt= len(ts)
no= R.shape[0]
#remove 0,0,0 and half-plane
if _isNonAxi(self._pot):
nn= (2*maxn-1)**2*maxn-(maxn-1)*(2*maxn-1)-maxn
else:
nn= maxn*(2*maxn-1)-maxn
A= nu.zeros((no,nt,2+nn))
A[:,:,0]= 1.
A[:,:,1]= ts
#sorting the phi and Z grids this way makes it easy to exclude the origin
phig= list(nu.arange(-maxn+1,maxn,1))
phig.sort(key = lambda x: abs(x))
phig= nu.array(phig,dtype='int')
if _isNonAxi(self._pot):
grid= nu.meshgrid(nu.arange(maxn),phig,phig)
else:
grid= nu.meshgrid(nu.arange(maxn),phig)
gridR= grid[0].T.flatten()[1:] #remove 0,0,0
gridZ= grid[1].T.flatten()[1:]
mask = nu.ones(len(gridR),dtype=bool)
# excludes axis that is not in half-space
if _isNonAxi(self._pot):
gridphi= grid[2].T.flatten()[1:]
mask= True\
^(gridR == 0)*((gridphi < 0)+((gridphi==0)*(gridZ < 0)))
else:
mask[:2*maxn-3:2]= False
gridR= gridR[mask]
gridZ= gridZ[mask]
tangleR= nu.tile(angleRT.T,(nn,1,1)).T
tgridR= nu.tile(gridR,(no,nt,1))
tangleZ= nu.tile(angleZT.T,(nn,1,1)).T
tgridZ= nu.tile(gridZ,(no,nt,1))
if _isNonAxi(self._pot):
gridphi= gridphi[mask]
tgridphi= nu.tile(gridphi,(no,nt,1))
tanglephi= nu.tile(anglephiT.T,(nn,1,1)).T
sinnR= nu.sin(tgridR*tangleR+tgridphi*tanglephi+tgridZ*tangleZ)
else:
sinnR= nu.sin(tgridR*tangleR+tgridZ*tangleZ)
A[:,:,2:]= sinnR
#Matrix magic
atainv= nu.empty((no,2+nn,2+nn))
AT= nu.transpose(A,axes=(0,2,1))
for ii in range(no):
atainv[ii,:,:,]= linalg.inv(nu.dot(AT[ii,:,:],A[ii,:,:]))
ATAR= nu.sum(AT*nu.transpose(nu.tile(angleRT,(2+nn,1,1)),axes=(1,0,2)),axis=2)
ATAT= nu.sum(AT*nu.transpose(nu.tile(anglephiT,(2+nn,1,1)),axes=(1,0,2)),axis=2)
ATAZ= nu.sum(AT*nu.transpose(nu.tile(angleZT,(2+nn,1,1)),axes=(1,0,2)),axis=2)
angleR= nu.sum(atainv[:,0,:]*ATAR,axis=1)
OmegaR= nu.sum(atainv[:,1,:]*ATAR,axis=1)
anglephi= nu.sum(atainv[:,0,:]*ATAT,axis=1)
Omegaphi= nu.sum(atainv[:,1,:]*ATAT,axis=1)
angleZ= nu.sum(atainv[:,0,:]*ATAZ,axis=1)
OmegaZ= nu.sum(atainv[:,1,:]*ATAZ,axis=1)
Omegaphi[negFreqIndx]= -Omegaphi[negFreqIndx]
anglephi[negFreqIndx]= _TWOPI-anglephi[negFreqIndx]
if kwargs.get('_retacfs',False):
return (jr,lz,jz,OmegaR,Omegaphi,OmegaZ, #pragma: no cover
angleR % _TWOPI,
anglephi % _TWOPI,
angleZ % _TWOPI,acfs)
else:
return (jr,lz,jz,OmegaR,Omegaphi,OmegaZ,
angleR % _TWOPI,
anglephi % _TWOPI,
angleZ % _TWOPI)
