from .error import *
import numpy as np
import logging
logger = logging.getLogger(__name__)
__all__ = ['VectorField','RGrid3DVectorField']
[docs]class VectorField():
"""Vector field model class
Returns zero field at all coordinates."""
def __init__(self):
"""Initialize vector field"""
[docs] def field(self,x):
"""Return field at coordinates x,y,z"""
return(np.array([0,0,0]))
[docs]class RGrid3DVectorField(VectorField):
"""Vector field with data on a regular 3D grid"""
ERROR = 0
ZERO = 1
NORMAL = 2
TANGENTIAL = 3
def __init__(self,data):
"""
Initialize vector field from numpy array.
Expecting columns x,y,z,Fx,Fy,Fz in SI units.
"""
# Read data and extract unique coordinates
#data = np.loadtxt(filename)
x = [np.unique(data[:,0]),np.unique(data[:,1]),np.unique(data[:,2])]
self.dx = np.zeros(3)
self.Nx = np.zeros(3,dtype=int)
self.x0 = np.zeros(3)
self.x1 = np.zeros(3)
# Define grid
for a in range(3):
self.dx[a] = x[a][1]-x[a][0]
self.Nx[a] = int(round((x[a][-1]-x[a][0])/self.dx[a]) + 1)
self.dx[a] = (x[a][-1]-x[a][0])/(self.Nx[a]-1)
self.x0[a] = x[a][0]
self.x1[a] = x[a][-1]
assert(self.Nx[a] == len(x[a]))
# Check grid definition
assert(self.Nx[0]*self.Nx[1]*self.Nx[2] == data.shape[0])
# Store field data in correct order
# Sort primarily according to x, then y, then z
self.F = data[np.lexsort(data.T[[0,1,2]]),3:]
self.scale = 1.0
self.extrpl = np.array(6*[self.ERROR])
[docs] def set_scale(self,factor):
"""Set field scaling factor."""
self.scale = factor
[docs] def field(self,x):
"""Return field at coordinates x
Uses linear interpolation at x. If point x is outside the
defined domain the behaviour is defined by the extrapolation
setting.
"""
# Make a local copy
x = x.copy()
sign = np.ones(3)
i = np.zeros(3,dtype=int)
u = np.zeros(3)
# Boundary checks for 3 coordinate directions
for a in range(3):
if x[a] < self.x0[a]:
if self.extrpl[2*a] == self.ZERO:
return( np.array([0,0,0]) )
elif self.extrpl[2*a] == self.ERROR:
raise FieldError('Coordinate out of field bounds at x[{:d}]={:e}<{:e}'.format(a,x[a],self.x0[a]))
elif x[a] < self.x0[a]-self.Nx[a]*self.dx[a]:
#return( np.array([0,0,0]) )
raise FieldError('Coordinate out of field bounds at x[{:d}]={:e}<{:e}'.format(a,x[a],self.x0[a]))
elif self.extrpl[2*a] == self.TANGENTIAL:
sign[a] *= -1
x[a] = 2*self.x0[a] - x[a]
elif self.extrpl[2*a] == self.NORMAL:
sign *= -1
sign[a] *= -1
x[a] = 2*self.x0[a] - x[a]
elif x[a] > self.x1[a]:
if self.extrpl[2*a+1] == self.ZERO:
return( np.array([0,0,0]) )
elif self.extrpl[2*a+1] == self.ERROR:
raise FieldError('Coordinate out of field bounds at x[{:d}]={:e}>{:e}'.format(a,x[a],self.x1[a]))
elif x[a] > self.x0[a]+2*self.Nx[a]*self.dx[a]:
#return( np.array([0,0,0]) )
raise FieldError('Coordinate out of field bounds at x[{:d}]={:e}<{:e}'.format(a,x[a],self.x0[a]))
elif self.extrpl[2*a+1] == self.TANGENTIAL:
sign[a] *= -1
x[a] = 2*self.x1[a] - x[a]
elif self.extrpl[2*a+1] == self.NORMAL:
sign *= -1
sign[a] *= -1
x[a] = 2*self.x1[a] - x[a]
# Find the cell index where (x,y) is located
i[a] = int((x[a]-self.x0[a])/self.dx[a])
# Check bounds
if i[a] < 0:
i[a] = 0
elif i[a] >= self.Nx[a]-1:
i[a] = self.Nx[a]-2
u[a] = (x[a]-self.x0[a])/self.dx[a]-i[a]
base = (i[2]*self.Nx[1] + i[1])*self.Nx[0] + i[0]
dj = self.Nx[0]
R = (1-u[2])*( (1.0-u[0])*(1.0-u[1])*self.F[base,:] +
( u[0])*(1.0-u[1])*self.F[base+1,:] +
(1.0-u[0])*( u[1])*self.F[base+dj,:] +
( u[0])*( u[1])*self.F[base+1+dj,:] )
base += self.Nx[0]*self.Nx[1]
R += u[2]*( (1.0-u[0])*(1.0-u[1])*self.F[base,:] +
( u[0])*(1.0-u[1])*self.F[base+1,:] +
(1.0-u[0])*( u[1])*self.F[base+dj,:] +
( u[0])*( u[1])*self.F[base+1+dj,:] )
return(self.scale*sign*R)
[docs] def plot_field(self,c,v1,v2,N1,N2):
"""
Plot the field magnitude on a plane with colormap
The plane is defined by point c in 3D space and two vectors v1 and v2.
The points used to produce the colormap are c+i*v1+j*v2, where i=[0,N1]
and j=[0,N2]
"""
data = np.zeros((N1,N2))
for i in range(N1):
for j in range(N2):
x = c + i*v1 + j*v2
data[i,j] = np.linalg.norm(self.field(x))
plt.imshow(data.T,origin='lower')
plt.colorbar()
[docs] def debug_print(self):
"""
Print information for debugging purposes
"""
logger.debug('Rgrid3DVectorField')
logger.debug('x0 = {:s}'.format(np.array2string(self.x0)))
logger.debug('x1 = {:s}'.format(np.array2string(self.x1)))
logger.debug('dx = {:s}'.format(np.array2string(self.dx)))
logger.debug('Nx = {:s}'.format(np.array2string(self.Nx)))
logger.debug('scale = {:f}'.format(self.scale))
