def makeHeatContent(salt, temp, destMask, thetao, pressure):
"""
The makeHeatContent() function takes 3D (not temporal) arguments and creates
heat content which is then mapped to a destination grid and written to a
specified variable
Author: Paul J. Durack : [email protected] : @durack1.
Created on Tue Nov 24 15:34:30 2015.
Inputs:
------
- salt(lev,lat,lon) - 3D array.
- temp(lev,lat,lon) - 3D array either in-situ or potential temperature.
- destGridArray(str) - 2D array with valid grid and mask.
- thetao(bool) - boolean value specifying either in-situ or potential temperature arrays provided.
- pressure(bool) - boolean value specifying whether lev-coordinate is pressure (dbar) or depth (m).
Usage:
------
>>> from oceanLib import makeHeatContent
>>> makeHeatContent(salt,temp,destGridArray,thetao=True,pressure=False)
Notes:
-----
- PJD 24 Nov 2015 - Migrated into new oceanLib from heatContentLib
- TODO: Better deal with insitu vs thetao variables
- TODO:
"""
# Remap variables to short names
#print salt.getAxisIds()
s = salt(squeeze=1)
# Trim off singleton time dimension
#print s.getAxisIds()
t = temp(squeeze=1)
mask = destMask
#print mask.getAxisIds()
del (salt, temp, destMask)
gc.collect()
depthInd = 0
# Set depth coordinate index
#print 's: ',s.min(),s.max()
#print 't: ',t.min(),t.max()
# Fix out of bounds values
t = mv.where(t < -2.6, -2.6, t)
# Fix for NaN values
# Calculate pressure - inputs depth & lat
# Create z-coordinate from salinity input
if not pressure:
zCoord = s.getAxis(depthInd)
# Assume time,depth,latitude,longitude grid
yCoord = s.getAxis(depthInd + 1)
yCoord = tile(yCoord, (s.shape[depthInd + 2], 1)).transpose()
depthLevels = tile(
zCoord.getValue(),
(s.shape[depthInd + 2], s.shape[depthInd + 1], 1)).transpose()
pressureLevels = sw.pres(np.array(depthLevels), np.array(yCoord))
del (zCoord, yCoord, depthLevels)
gc.collect()
else:
pressureLevels = s.getAxis(depthInd)
#print pressureLevels.getValue()
pressureLevels = transpose(
tile(pressureLevels,
(s.shape[depthInd + 2], s.shape[depthInd + 1], 1)))
pressureLevels = cdm.createVariable(pressureLevels, id='pressureLevels')
pressureLevels.setAxis(0, s.getAxis(depthInd))
pressureLevels.setAxis(1, s.getAxis(depthInd + 1))
pressureLevels.setAxis(2, s.getAxis(depthInd + 2))
pressureLevels.units_long = 'decibar (pressure)'
pressureLevels.positive = 'down'
pressureLevels.long_name = 'sea_water_pressure'
pressureLevels.standard_name = 'sea_water_pressure'
pressureLevels.units = 'decibar'
pressureLevels.axis = 'Z'
#print 'pres: ',pressureLevels.min(),pressureLevels.max()
#print pressureLevels.shape
#print s.shape
#print t.shape
#print mask.shape
# Calculate temp,rho,cp - inputs temp,salt,pressure
if thetao:
# Process potential temperature to in-situ
temp = sw.temp(np.array(s), np.array(t), np.array(pressureLevels))
# units degrees C
rho = sw.dens(np.array(s), np.array(temp), np.array(pressureLevels))
# units kg m-3
cp = sw.cp(np.array(s), np.array(temp), np.array(pressureLevels))
# units J kg-1 C-1
# Correct instances of NaN values and fix masks - applied before cdms variables are created otherwise names/ids/attributes are reset
temp = scrubNaNAndMask(temp, s)
rho = scrubNaNAndMask(rho, s)
cp = scrubNaNAndMask(cp, s)
#print 'temp: ',temp.min(),temp.max()
#print 'rho: ',rho.min(),rho.max()
#print 'cp: ',cp.min(),cp.max()
# Calculate heatContent - inputs temp,rho,cp
heatContent = np.array(temp) * np.array(rho) * np.array(cp)
# units J
# Correct instances of NaN values and fix masks - applied before cdms variables are created otherwise names/ids/attributes are reset
heatContent = scrubNaNAndMask(heatContent, s)
#print 'hc: ',heatContent.min(),heatContent.max()
# Interpolate to standard levels - inputs heatContent,levels
newDepth = np.array([
5, 10, 20, 30, 40, 50, 75, 100, 125, 150, 200, 300, 500, 700, 1000,
1500, 1800, 2000
]).astype('f')
newDepth_bounds = np.array([[0, 5], [5, 10], [10, 20], [20, 30], [30, 40],
[40, 50], [50, 75], [75, 100], [100, 125],
[125, 150], [150, 200], [200, 300], [300, 500],
[500, 700], [700, 1000], [1000, 1500],
[1500, 1800], [1800, 2000]]).astype('f')
# Interpolate to standard levels
#print heatContent.shape
#print heatContent.getAxisIds()
#print pressureLevels.shape
#print pressureLevels.getAxisIds()
# Reset variable axes
heatContent.setAxis(0, s.getAxis(0))
#heatContent.setAxis(1,s.getAxis(1))
#heatContent.setAxis(2,s.getAxis(2))
pdb.set_trace()
heatContent.setGrid(s.getGrid())
#print heatContent.shape
#print heatContent.getAxisIds()
#print pressureLevels.shape
#print pressureLevels.getAxisIds()
heatContent_depthInterp = cdu.linearInterpolation(heatContent,
pressureLevels,
levels=newDepth)
# Fix bounds
newDepth = heatContent_depthInterp.getAxis(0)
newDepth.setBounds(newDepth_bounds)
del (newDepth_bounds)
newDepth.id = 'depth2'
newDepth.units_long = 'decibar (pressure)'
newDepth.positive = 'down'
newDepth.long_name = 'sea_water_pressure'
newDepth.standard_name = 'sea_water_pressure'
newDepth.units = 'decibar'
newDepth.axis = 'Z'
#print 'hc_interp:',heatContent_depthInterp.min(),heatContent_depthInterp.max()
# Integrate to 700 dbar - inputs heatContent
heatContent_depthInteg = cdu.averager(
heatContent_depthInterp[0:14, ...],
axis=0,
weights='weighted',
action='sum')(squeeze=1) # Calculate depth-weighted-integrated thetao
# Assign all axis info
#print heatContent_depthInteg.shape
pdb.set_trace()
# Interpolate in x,y - inputs heatContent
#tmp1 = heatContent_depthInteg.regrid(mask.getGrid(),regridTool='esmf',regridMethod='linear') ; # Use defaults - ,coordSys='deg',diag = {},periodicity=1)
tmp1 = heatContent_depthInteg.regrid(mask,
regridTool='esmf',
regridMethod='linear')
# Use defaults - ,coordSys='deg',diag = {},periodicity=1)
#print tmp1.shape
tmp1 = mv.where(tmp1 < 0, 0, tmp1)
# Fix for negative values
# Infill - inputs heatContent
# Create inputs for interpolation
points = np.zeros([(mask.shape[0] * mask.shape[1]), 2])
# Create 25380 vectors of lon/lat
latcounter = 0
loncounter = 0
for count, data in enumerate(points):
if not np.mod(count, 180) and not count == 0:
latcounter = latcounter + 1
loncounter = 0
points[count, 0] = mask.getLatitude().getValue()[latcounter]
points[count, 1] = mask.getLongitude().getValue()[loncounter]
loncounter = loncounter + 1
del (count, data, latcounter, loncounter)
gc.collect()
valid = np.logical_not(tmp1.mask)
# Get inverted-logic boolean mask from variable
if valid.size == 1:
print '** No valid mask found, skipping **'
return
valid = valid.flatten()
# Flatten 2D to 1D
#maskFilled = mask(tmp,points,valid)
interpolant = interpolate.LinearNDInterpolator(
points[valid, :],
np.array(tmp1.flatten())[valid])
# Create interpolant
maskFill = interpolant(points[:, 0].squeeze(), points[:, 1].squeeze())
# Use interpolant to create filled matrix
maskFill = np.reshape(maskFill, mask.shape)
# Resize to original dimensions
# Fix issues with interpolant
tmp2 = mv.where(np.isnan(maskFill), 1e+20, maskFill)
# Fix for NaN values
tmp2 = mv.where(tmp2 > tmp1.max(), 0, tmp2)
# Fix for max values
tmp2 = mv.where(tmp2 < tmp1.min(), 0, tmp2)
# Fix for min values
tmp = mv.masked_where(mask.mask, tmp2)
#print tmp.shape
# Redress variable
heatContent = cdm.createVariable([tmp], id='heatContent')
depthInt = cdm.createAxis([350], id='depth')
depthInt.setBounds(np.array([0, 700]))
depthInt.units_long = 'decibar (pressure)'
depthInt.positive = 'down'
depthInt.long_name = 'sea_water_pressure'
depthInt.standard_name = 'sea_water_pressure'
depthInt.units = 'decibar'
depthInt.axis = 'Z'
heatContent.setAxis(0, depthInt)
heatContent.setAxis(1, mask.getAxis(0))
heatContent.setAxis(2, mask.getAxis(1))
heatContent.units_long = 'Joules'
heatContent.long_name = 'sea_water_heat_content'
heatContent.standard_name = 'sea_water_heat_content'
heatContent.units = 'J'
return heatContent
def __init__(self, x_vector, y):
# x_vector is of dimension (n_points, n_dims)
super().__init__(x_vector, y)
self.NDfunction = interpolate.LinearNDInterpolator(x_vector, y)
nodeDisplacementX = np.empty(nodes.shape[0])
nodeDisplacementY = np.empty(nodes.shape[0])
nodeDisplacementZ = np.empty(nodes.shape[0])
for inode in range(nodes.shape[0]):
cord = np.array([nodes[inode][0],nodes[inode][1],nodes[inode][2],1])
new_cord = affineMatrix.dot(cord)
nodeDisplacementX[inode] = new_cord[0]-nodes[inode][0]
nodeDisplacementY[inode] = new_cord[1]-nodes[inode][1]
nodeDisplacementZ[inode] = new_cord[2]-nodes[inode][2]
newim3_map = ndimage.map_coordinates(im2_matched, [mapX, mapY, mapZ], order=1)
surf = plot.surf_points
inter_z = interpolate.LinearNDInterpolator(nodes, nodeDisplacementZ, fill_value=0.0, rescale=True)
extrapo_z = inter_z.__call__(surf)
plot.plot3dpoint(extrapo_z)
#####
plot3 = Visualization.DataVisulization(ndimage.gaussian_filter(newim3_map,5), 110)
plot3.contour3d()
plot4 = Visualization.DataVisulization(ndimage.gaussian_filter(newim3,5), 110)
plot4.contour3d()
#src3 = mlab.pipeline.scalar_field(ndimage.gaussian_filter(newim3,3))
#mlab.pipeline.iso_surface(src3, contours=[70], colormap='Oranges')
save = SaveLoad.saveDataBase("/Users/junchaowei/Desktop/SpaceRegistration_000_125/")
datadict={}
def active_from_xyz(mesh, xyz, grid_reference='CC', method='linear'):
"""Returns an active cell index array below a surface
Get active cells in the `mesh` that are below the surface create by
interpolating over the last dimension of the input points. This method will
uses scipy's interpolation routine to interpolate between input points. This
will use a nearest neighbour interpolation for cell values outside the convex
hull of points.
For `grid_reference='CC'`, this will check if the center of a given cell in
the mesh is below the interpolation surface to determine if that cell is
active.
For `grid_reference='N'`, this will check if **every** node of a given cell
in the mesh is below the interpolation surface.
For the
Parameters
----------
mesh : TensorMesh or TreeMesh or CylMesh
Mesh object, (if CylMesh: mesh must be symmetric).
xyz : numpy.ndarray
Points coordinates shape (*, mesh.dim).
grid_reference : {'CC', 'N'}
Use cell coordinates from cells-center 'CC' or nodes 'N'.
method : {'linear', 'nearest'}
Interpolation method for the xyz points.
Returns
-------
active : numpy.ndarray
1D mask array of `bool` for the active cells below xyz.
Examples
--------
.. plot::
:include-source:
import matplotlib.pyplot as plt
import numpy as np
from discretize import TensorMesh
from discretize.utils import active_from_xyz
mesh = TensorMesh([5, 5])
topo_func = lambda x: -3*(x-0.2)*(x-0.8)+.5
topo_points = np.linspace(0, 1)
topo_vals = topo_func(topo_points)
active_cc = active_from_xyz(mesh, np.c_[topo_points, topo_vals], grid_reference='CC')
ax = plt.subplot(121)
mesh.plotImage(active_cc, ax=ax)
mesh.plotGrid(centers=True, ax=ax)
ax.plot(np.linspace(0,1), topo_func(np.linspace(0,1)), color='C3')
ax.set_title("CC")
active_n = active_from_xyz(mesh, np.c_[topo_points, topo_vals], grid_reference='N')
ax = plt.subplot(122)
mesh.plotImage(active_n, ax=ax)
mesh.plotGrid(nodes=True, ax=ax)
ax.plot(np.linspace(0,1), topo_func(np.linspace(0,1)), color='C3')
ax.set_title("N")
plt.show()
"""
if isinstance(mesh, discretize.CylMesh) and not mesh.isSymmetric:
raise NotImplementedError('Unsymmetric CylMesh is not yet supported')
if not isinstance(mesh, (discretize.TensorMesh, discretize.TreeMesh, discretize.CylMesh)):
raise TypeError("Mesh must be either TensorMesh, TreeMesh, or Symmetric CylMesh")
if grid_reference not in ["N", "CC"]:
raise ValueError("Value of grid_reference must be 'N' (nodal) or 'CC' (cell center)")
dim = mesh.dim - 1
if mesh.dim == 3:
if xyz.shape[1] != 3:
raise ValueError("xyz locations of shape (*, 3) required for 3D mesh")
if method == 'linear':
tri2D = Delaunay(xyz[:, :2])
z_interpolate = interpolate.LinearNDInterpolator(tri2D, xyz[:, 2])
else:
z_interpolate = interpolate.NearestNDInterpolator(xyz[:, :2], xyz[:, 2])
elif mesh.dim == 2:
if xyz.shape[1] != 2:
raise ValueError("xyz locations of shape (*, 2) required for 2D mesh")
z_interpolate = interpolate.interp1d(
xyz[:, 0], xyz[:, 1], bounds_error=False, fill_value=np.nan, kind=method
)
else:
if xyz.ndim != 1:
raise ValueError("xyz locations of shape (*, ) required for 1D mesh")
if grid_reference == 'CC':
locations = mesh.gridCC
if mesh.dim == 1:
active = np.zeros(mesh.nC, dtype='bool')
active[np.searchsorted(mesh.vectorCCx, xyz).max():] = True
return active
elif grid_reference == 'N':
if mesh.dim == 3:
locations = np.vstack([
mesh.gridCC + (np.c_[-1, 1, 1][:, None] * mesh.h_gridded / 2.).squeeze(),
mesh.gridCC + (np.c_[-1, -1, 1][:, None] * mesh.h_gridded / 2.).squeeze(),
mesh.gridCC + (np.c_[1, 1, 1][:, None] * mesh.h_gridded / 2.).squeeze(),
mesh.gridCC + (np.c_[1, -1, 1][:, None] * mesh.h_gridded / 2.).squeeze()
])
elif mesh.dim == 2:
locations = np.vstack([
mesh.gridCC + (np.c_[-1, 1][:, None] * mesh.h_gridded / 2.).squeeze(),
mesh.gridCC + (np.c_[1, 1][:, None] * mesh.h_gridded / 2.).squeeze(),
])
else:
active = np.zeros(mesh.nC, dtype='bool')
active[np.searchsorted(mesh.vectorNx, xyz).max():] = True
return active
# Interpolate z values on CC or N
z_xyz = z_interpolate(locations[:, :-1]).squeeze()
# Apply nearest neighbour if in extrapolation
ind_nan = np.isnan(z_xyz)
if any(ind_nan):
tree = cKDTree(xyz)
_, ind = tree.query(locations[ind_nan, :])
z_xyz[ind_nan] = xyz[ind, dim]
# Create an active bool of all True
active = np.all(
(locations[:, dim] < z_xyz).reshape((mesh.nC, -1), order='F'), axis=1
)
return active.ravel()
#dispz = np.zeros(mesh2.points.T.shape[1])
colors = (dispz) * 0.0495
#colors = np.zeros(mesh2.points.T.shape[1])
#sio.savemat('/Users/junchaowei/Desktop/0_90_z.mat', {'p':colors})
#sio.savemat('/Users/junchaowei/Desktop/0_90_element.mat',{'p':mesh2.ntri})
@mlab.show
def main():
mesh2.view(mesh2.polydata(mesh2.points.T, colors))
main()
mayavi.mlab.colorbar(object=None,
title="U(Z)(mm)",
orientation="vertical",
nb_labels=20,
nb_colors=None,
label_fmt=None)
"colorbar reference: http://docs.enthought.com/mayavi/mayavi/auto/mlab_decorations.html"
plot = Visualization.DataVisulization(image1, threshold1 + 5000)
plot.contour3d()
surf = plot.surf_points
inter_z = interpolate.LinearNDInterpolator(mesh2.points,
colors,
fill_value=0.0,
rescale=True)
extrapo_z = inter_z.__call__(surf)
plot.plot3dpoint(extrapo_z)
def _ip(self):
"""A `scipy.interpolate.LinearNDInterpolator` instance
containing the learner's data."""
# XXX: take our own triangulation into account when generating the _ip
return interpolate.LinearNDInterpolator(self.points, self.values)
def ax_scalp(
v,
channels,
ax=None,
annotate=False,
vmin=None,
vmax=None,
cmap=cm.coolwarm,
scalp_line_width=1,
scalp_line_style="solid",
chan_pos_list=CHANNEL_10_20_APPROX,
interpolation="bilinear",
fontsize=8,
):
"""Draw a scalp plot.
Draws a scalp plot on an existing axes. The method takes an array of
values and an array of the corresponding channel names. It matches
the channel names with an channel position list
to project them correctly on the scalp.
Parameters
----------
v : 1d-array of floats
The values for the channels
channels : 1d array of strings
The corresponding channel names for the values in ``v``
ax : Axes, optional
The axes to draw the scalp plot on. If not provided, the
currently activated axes (i.e. ``gca()``) will be taken
annotate : Boolean, optional
Draw the channel names next to the channel markers.
vmin, vmax : float, optional
The display limits for the values in ``v``. If the data in ``v``
contains values between -3..3 and ``vmin`` and ``vmax`` are set
to -1 and 1, all values smaller than -1 and bigger than 1 will
appear the same as -1 and 1. If not set, the maximum absolute
value in ``v`` is taken to calculate both values.
cmap : matplotlib.colors.colormap, optional
A colormap to define the color transitions.
scalp_line_width: float
Line width for outline of scalp
scalp_line_style: float
Line style for outline of scalp
chan_pos_list: iterable of tuples
First entry should be 'angle' or 'cartesian',
remaining entries 2-tuples of x and y.
interpolation: str
Returns
-------
ax : Axes
the axes on which the plot was drawn
Notes
-----
Code adapted from Wyrm [1]_ toolbox https://github.com/bbci/wyrm.
References
----------
.. [1] Schirrmeister, R. T., Springenberg, J. T., Fiederer, L. D. J.,
Glasstetter, M., Eggensperger, K., Tangermann, M., Hutter, F. & Ball, T. (2017).
Deep learning with convolutional neural networks for EEG decoding and
visualization.
Human Brain Mapping , Aug. 2017. Online: http://dx.doi.org/10.1002/hbm.23730
"""
if ax is None:
ax = plt.gca()
assert len(v) == len(channels), "Should be as many values as channels"
assert interpolation == "bilinear" or interpolation == "nearest"
if vmin is None:
# added by me ([email protected])
assert vmax is None
vmin, vmax = -np.max(np.abs(v)), np.max(np.abs(v))
# what if we have an unknown channel?
points = [get_channelpos(c, chan_pos_list) for c in channels]
for c in channels:
assert get_channelpos(
c, chan_pos_list) is not None, ("Expect " + c +
" to exist in positions")
z = [v[i] for i in range(len(points))]
# calculate the interpolation
x = [i[0] for i in points]
y = [i[1] for i in points]
# interpolate the in-between values
xx = np.linspace(min(x), max(x), 500)
yy = np.linspace(min(y), max(y), 500)
if interpolation == "bilinear":
xx_grid, yy_grid = np.meshgrid(xx, yy)
f = interpolate.LinearNDInterpolator(list(zip(x, y)), z)
zz = f(xx_grid, yy_grid)
else:
assert interpolation == "nearest"
f = interpolate.NearestNDInterpolator(list(zip(x, y)), z)
assert len(xx) == len(yy)
zz = np.ones((len(xx), len(yy)))
for i_x in range(len(xx)):
for i_y in range(len(yy)):
# somehow this is correct. don't know why :(
zz[i_y, i_x] = f(xx[i_x], yy[i_y])
# zz[i_x,i_y] = f(xx[i_x], yy[i_y])
assert not np.any(np.isnan(zz))
# plot map
image = ax.imshow(
zz,
vmin=vmin,
vmax=vmax,
cmap=cmap,
extent=[min(x), max(x), min(y), max(y)],
origin="lower",
interpolation=interpolation,
)
if scalp_line_width > 0:
# paint the head
ax.add_artist(
plt.Circle(
(0, 0),
1,
linestyle=scalp_line_style,
linewidth=scalp_line_width,
fill=False,
))
# add a nose
ax.plot(
[-0.1, 0, 0.1],
[1, 1.1, 1],
color="black",
linewidth=scalp_line_width,
linestyle=scalp_line_style,
)
# add ears
_add_ears(ax, scalp_line_width, scalp_line_style)
# add markers at channels positions
# set the axes limits, so the figure is centered on the scalp
ax.set_ylim([-1.05, 1.15])
ax.set_xlim([-1.15, 1.15])
# hide the frame and ticks
ax.set_frame_on(False)
ax.set_xticks([])
ax.set_yticks([])
# draw the channel names
if annotate:
for i in zip(channels, list(zip(x, y))):
ax.annotate(
" " + i[0],
i[1],
horizontalalignment="center",
verticalalignment="center",
fontsize=fontsize,
)
ax.set_aspect(1)
return image
def create_model_maps(
altitude,
variable=None,
model=None,
file=None,
numContours=25,
fill=False,
ct='viridis', # https://matplotlib.org/examples/color/colormaps_reference.html
transparency=1,
nearest=False,
linear=True,
saveFig=True):
import matplotlib
#matplotlib.use('tkagg')
import matplotlib.pyplot as plt
if model == None and file == None:
print("Please input either a model or the file path/name to a model.")
return
if file != None:
model = read_model_results(file)
if (variable == None) or (variable not in model.keys()):
print("Variable not entered or not found, Please select one from: ")
index = 0
name_index_dict = {}
for name in model:
if name.lower() == 'dim':
continue
if name.lower() == 'meta':
continue
print(index, ": ", name)
name_index_dict[index] = name
index += 1
i_choice = int(input("Enter Selection: "))
dataname = name_index_dict[i_choice].lower()
print(dataname)
else:
dataname = variable.lower()
mars_radius = model['meta']['mars_radius']
lats = np.arange(181) - 90
lons = np.arange(361) - 180
sc_lat_mso = np.repeat(lats, len(lons))
sc_lon_mso = np.tile(lons, len(lats))
r = np.full(len(sc_lon_mso), altitude + mars_radius)
sc_alt_array = np.full(len(sc_lon_mso), altitude)
sc_mso_x = r * np.sin(
(90 - sc_lat_mso) * (np.pi / 180)) * np.cos(sc_lon_mso * (np.pi / 180))
sc_mso_y = r * np.sin(
(90 - sc_lat_mso) * (np.pi / 180)) * np.sin(sc_lon_mso * (np.pi / 180))
sc_mso_z = r * np.cos((90 - sc_lat_mso) * (np.pi / 180))
sc_path = np.array([sc_mso_x, sc_mso_y, sc_mso_z]).T
if nearest:
interp_method = 'nearest'
else:
interp_method = 'linear'
if model == None and file == None:
print(
"Please input either a model dictionary from read_model_results, or a model file."
)
return
if 'lon' in model['dim']:
if 'mso' == model['meta']['coord_sys'].lower():
#Build a big matrix with dimensions 3 columns by (num lat * num lon * num alt) rows
lat_mso_model = model['dim']['lat']
lon_mso_model = model['dim']['lon']
alt_mso_model = model['dim']['alt']
lat_array = np.repeat(lat_mso_model, len(lon_mso_model))
lon_array = np.tile(lon_mso_model, len(lat_mso_model))
data_points = np.transpose(np.array([lon_array, lat_array]))
index = 0
for point in sc_path:
r_mso = np.sqrt(point[0]**2 + point[1]**2 + point[2]**2)
alt_mso = altitude
lat_mso = 90 - (np.arccos(point[2] / r_mso) / (np.pi / 180))
lon_mso = np.arctan2(point[1], point[0]) / (np.pi / 180)
sc_path[index] = np.array([lon_mso, lat_mso, alt_mso])
index += 1
latlon_triangulation = spatial.Delaunay(data_points)
for var in model:
if var.lower() != dataname:
continue
print("Interpolating variable " + var)
if var == 'dim' or var == 'meta':
continue
#Rearrange the data to lon/lat/alt
data = model[var]['data']
dim_order_array = [0, 1, 2]
for j in [0, 1, 2]:
if model[var]['dim_order'][j] == 'longitude':
dim_order_array[0] = j
elif model[var]['dim_order'][j] == 'latitude':
dim_order_array[1] = j
elif model[var]['dim_order'][j] == 'altitude':
dim_order_array[2] = j
data_new = np.transpose(data, dim_order_array)
#Build an array of values that correspond to the points in data point
values = np.empty([
len(lat_mso_model) * len(lon_mso_model),
len(alt_mso_model)
])
index = 0
for alt in range(0, len(alt_mso_model)):
for lat in range(0, len(lat_mso_model)):
for lon in range(0, len(lon_mso_model)):
values[index, alt] = data_new[lon][lat][alt]
index += 1
index = 0
x = np.empty(len(sc_path))
index = 0
for sc_pos in sc_path:
if sc_pos[2] > np.max(alt_mso_model):
x[index] = np.NaN
index += 1
continue
if sc_pos[2] < np.min(alt_mso_model):
x[index] = np.NaN
index += 1
continue
sorted_x_distance = np.argsort(
np.abs(alt_mso_model - sc_pos[2]))
alti1 = sorted_x_distance[0]
if nearest:
x[index] = interpolate.griddata(data_points,
values[:, alti1],
[sc_pos[0], sc_pos[1]],
method='nearest')
else:
if alt_mso_model[alti1] < sc_pos[2]:
alti2 = alti1 + 1
else:
temp = alti1 - 1
alti2 = alti1
alti1 = temp
#Interpolate through space
first_val_calc = interpolate.LinearNDInterpolator(
latlon_triangulation, values[:, alti1])
second_val_calc = interpolate.LinearNDInterpolator(
latlon_triangulation, values[:, alti2])
first_val = first_val_calc([sc_pos[0], sc_pos[1]])
second_val = second_val_calc([sc_pos[0], sc_pos[1]])
delta_1 = sc_pos[2] - alt_mso_model[alti1]
delta_2 = alt_mso_model[alti2] - sc_pos[2]
delta_tot = alt_mso_model[alti2] - alt_mso_model[alti1]
x[index] = ((first_val * delta_2) +
(second_val * delta_1)) / (delta_tot)
index += 1
tracer = np.array(x)
if 'geo' == model['meta']['coord_sys'].lower():
#Build the Matrix that transforms GEO to MSO coordinates
ls_rad = model['meta']['ls'] * np.pi / 180
rads_tilted_y = 25.19 * np.sin(ls_rad) * np.pi / 180
rads_tilted_x = -25.19 * np.cos(ls_rad) * np.pi / 180
lonsubsol_rad = -model['meta']['longsubsol'] * np.pi / 180
z_rotation = np.matrix(
[[np.cos(lonsubsol_rad), -np.sin(lonsubsol_rad), 0],
[np.sin(lonsubsol_rad),
np.cos(lonsubsol_rad), 0], [0, 0, 1]])
y_rotation = np.matrix(
[[np.cos(rads_tilted_y), 0,
np.sin(rads_tilted_y)], [0, 1, 0],
[-np.sin(rads_tilted_y), 0,
np.cos(rads_tilted_y)]])
x_rotation = np.matrix(
[[1, 0, 0], [0,
np.cos(rads_tilted_x), -np.sin(rads_tilted_x)],
[0, np.sin(rads_tilted_x),
np.cos(rads_tilted_x)]])
geo_to_mso_matrix = np.dot(x_rotation,
np.dot(y_rotation, z_rotation))
#Build a big matrix with dimensions 3 columns by (num lat * num lon * num alt) rows
lat_geo_model = model['dim']['lat']
lon_geo_model = model['dim']['lon']
alt_geo_model = model['dim']['alt']
alt_array = np.repeat(alt_geo_model,
len(lon_geo_model) * len(lat_geo_model))
lat_array = np.tile(np.repeat(lat_geo_model, len(lon_geo_model)),
len(alt_geo_model))
lon_array = np.tile(lon_geo_model,
len(lat_geo_model) * len(alt_geo_model))
data_points = np.transpose(
np.array([lon_array, lat_array, alt_array]))
#Convert to GEO coordinates, then to MSO
index = 0
for point in data_points:
r = point[2] + mars_radius
x = r * np.sin((90 - point[1]) * np.pi / 180) * np.cos(
point[0] * np.pi / 180)
y = r * np.sin((90 - point[1]) * np.pi / 180) * np.sin(
point[0] * np.pi / 180)
z = r * np.cos((90 - point[1]) * np.pi / 180)
data_points[index] = np.dot(geo_to_mso_matrix,
np.array([x, y, z]))
index += 1
#Convert to MSO Lon/Lat/Alt in order to weight the interpolation better
lat_mso = np.empty(len(lon_geo_model) * len(lat_geo_model))
lon_mso = np.empty(len(lon_geo_model) * len(lat_geo_model))
index = 0
for point in data_points:
r_mso = np.sqrt(point[0]**2 + point[1]**2 + point[2]**2)
lat_mso[index] = 90 - (np.arccos(point[2] / r_mso) /
(np.pi / 180))
lon_mso[index] = np.arctan2(point[1], point[0]) / (np.pi / 180)
index += 1
if index >= len(lon_geo_model) * len(lat_geo_model):
break
latlon_points = np.transpose(np.array([lon_mso, lat_mso]))
index = 0
for point in sc_path:
r_mso = np.sqrt(point[0]**2 + point[1]**2 + point[2]**2)
alt_mso = altitude
lat_mso = 90 - (np.arccos(point[2] / r_mso) / (np.pi / 180))
lon_mso = np.arctan2(point[1], point[0]) / (np.pi / 180)
sc_path[index] = np.array([lon_mso, lat_mso, alt_mso])
index += 1
latlon_triangulation = spatial.Delaunay(latlon_points)
#Loop through the variables in the model
for var in model:
if var.lower() != dataname:
continue
print("Interpolating variable " + var)
#Rearrange the data to lon/lat/alt
data = model[var]['data']
dim_order_array = [0, 1, 2]
for j in [0, 1, 2]:
if model[var]['dim_order'][j] == 'longitude':
dim_order_array[0] = j
elif model[var]['dim_order'][j] == 'latitude':
dim_order_array[1] = j
elif model[var]['dim_order'][j] == 'altitude':
dim_order_array[2] = j
data_new = np.transpose(data, dim_order_array)
#Build an array of values that correspond to the points in data point
values = np.empty([
len(lat_geo_model) * len(lon_geo_model),
len(alt_geo_model)
])
index = 0
for alt in range(0, len(alt_geo_model)):
for lat in range(0, len(lat_geo_model)):
for lon in range(0, len(lon_geo_model)):
values[index, alt] = data_new[lon][lat][alt]
index += 1
index = 0
x = np.empty(len(sc_path))
index = 0
for sc_pos in sc_path:
if sc_pos[2] > np.max(alt_geo_model):
x[index] = np.NaN
index += 1
continue
if sc_pos[2] < np.min(alt_geo_model):
x[index] = np.NaN
index += 1
continue
sorted_x_distance = np.argsort(
np.abs(alt_geo_model - sc_pos[2]))
alti1 = sorted_x_distance[0]
if nearest:
x[index] = interpolate.griddata(latlon_points,
values[:, alti1],
[sc_pos[0], sc_pos[1]],
method='nearest')
else:
if alt_geo_model[alti1] < sc_pos[2]:
alti2 = alti1 + 1
else:
temp = alti1 - 1
alti2 = alti1
alti1 = temp
#Interpolate through space
first_val_calc = interpolate.LinearNDInterpolator(
latlon_triangulation, values[:, alti1])
second_val_calc = interpolate.LinearNDInterpolator(
latlon_triangulation, values[:, alti2])
first_val = first_val_calc([sc_pos[0], sc_pos[1]])
second_val = second_val_calc([sc_pos[0], sc_pos[1]])
delta_1 = sc_pos[2] - alt_geo_model[alti1]
delta_2 = alt_geo_model[alti2] - sc_pos[2]
delta_tot = alt_geo_model[alti2] - alt_geo_model[alti1]
x[index] = ((first_val * delta_2) +
(second_val * delta_1)) / (delta_tot)
index += 1
tracer = np.array(x)
else:
if 'mso' == model['meta']['coord_sys'].lower():
#Build a big matrix with dimensions 3 columns by (num lat * num lon * num alt) rows
x_mso_model = model['dim']['x']
y_mso_model = model['dim']['y']
z_mso_model = model['dim']['z']
#Loop through the variables in the model
for var in model:
if var.lower() != dataname:
continue
#Rearrange the data to lon/lat/alt
data = model[var]['data']
dim_order_array = [0, 1, 2]
for j in [0, 1, 2]:
if model[var]['dim_order'][j] == 'x':
dim_order_array[0] = j
elif model[var]['dim_order'][j] == 'y':
dim_order_array[1] = j
elif model[var]['dim_order'][j] == 'z':
dim_order_array[2] = j
data_new = np.transpose(data, dim_order_array)
#Interpolate through space
tracer = mvn_kp_sc_traj_xyz(x_mso_model,
y_mso_model,
z_mso_model,
data_new,
sc_mso_x,
sc_mso_y,
sc_mso_z,
nn=interp_method)
xi, yi = np.linspace(sc_lon_mso.min(), sc_lon_mso.max(),
300), np.linspace(sc_lat_mso.min(), sc_lat_mso.max(),
300)
xi, yi = np.meshgrid(xi, yi)
zi = scipy.interpolate.griddata((sc_lon_mso, sc_lat_mso),
tracer, (xi, yi),
method=interp_method)
if saveFig:
fig = plt.figure()
ax = fig.add_subplot(1, 1, 1)
if fill:
plt.contourf(xi,
yi,
zi,
numContours,
alpha=transparency,
cmap=ct,
extent=(-180, -90, 180, 90))
else:
CS = plt.contour(xi,
yi,
zi,
numContours,
alpha=transparency,
cmap=ct)
plt.clabel(CS, inline=1, fontsize=7, fmt='%1.0f')
extent = ax.get_window_extent().transformed(
fig.dpi_scale_trans.inverted())
plt.axis('off')
save_name = "ModelData_" + dataname + "_" + str(altitude) + "km"
if fill:
save_name = save_name + "_filled"
print("Saving file as:" +
os.path.join(os.path.dirname(file), save_name + ".png"))
plt.savefig(os.path.join(os.path.dirname(file), save_name + ".png"),
transparent=False,
bbox_inches=extent,
pad_inches=0,
dpi=150)
plt.show(block=True)
else:
# convert back to east longitude (as in model file)
return dict(lon=xi[0, :],
elon=((xi[0, :] + 360) % 360),
longsubsol=model['meta']['longsubsol'],
dec=model['meta']['dec'],
Ls=model['meta']['ls'],
lat=yi[:, 0],
param=zi)
from scipy import interpolate
import matplotlib.pyplot as plt
from scipy.integrate import solve_ivp
# http://www.ece.uci.edu/docs/hspice/hspice_2001_2-87.html
# i guess we can do optimization ?
try:
import cPickle as pickle
except ImportError:
import pickle
print("loading data")
vd = np.genfromtxt('backward_vd.csv', delimiter=',').reshape(-1, 1)
vg = np.genfromtxt('backward_vg.csv', delimiter=',').reshape(-1, 1)
vc = np.genfromtxt('backward_vc.csv', delimiter=',').reshape(-1, 1)
i = np.genfromtxt('backward_i.csv', delimiter=',').reshape(-1, 1)
print("building interpolator")
# points = np.concatenate((vd, vg), axis=1)
points = np.concatenate((vc, vg), axis=1)
x = points
y = i
fit = interpolate.LinearNDInterpolator(x, y, fill_value=0.0, rescale=True)
with open('backward.pkl', 'wb') as f:
pickle.dump(fit, f)
"""
this module creates image transformation or project for two images' registration
"""
import SaveLoad
import Visulization
from scipy import ndimage
from scipy import interpolate
import numpy as np
import scipy.ndimage as ndimage
inputDVC = SaveLoad.DVCdata("/Users/junchaowei/Desktop/MultiStrainGage11182016/") # import the DVC database
################### save the displacement, and create interpolation ##############
inter_x = interpolate.LinearNDInterpolator(inputDVC.getPoints1(), inputDVC.getDisplacementX()[0], fill_value=0.0, rescale=True)
inter_y = interpolate.LinearNDInterpolator(inputDVC.getPoints1(), inputDVC.getDisplacementY()[0], fill_value=0.0, rescale=True)
inter_z = interpolate.LinearNDInterpolator(inputDVC.getPoints1(), inputDVC.getDisplacementZ()[0], fill_value=0.0, rescale=True)
image1 = inputDVC.getImage1()
image2 = inputDVC.getImage2()
sx, sy = image1[:,:,0].shape # get dimensionality of image
mx, my = np.meshgrid(np.linspace(1,sx,sx),np.linspace(1,sy,sy))
# meshgrid change the orientation of x,y: x becomes the horizontal axis, y becomes the vertical
# this change would affect the return statement of reshape.
Cord_xy = np.vstack([mx.flatten(), my.flatten()]) # column major vector
displacementDataX = np.zeros([sy,sx, image1.shape[2]])
displacementDataY = np.zeros([sy,sx, image1.shape[2]])
displacementDataZ = np.zeros([sy,sx, image1.shape[2]])
def trace_mline(self, init_state, time, direction='FWD', \
length_line=None, stp=None, ripple=True):
'''
Traces the field line given a starting point.
Integration step defined by stp and maximum length of the field line
defined by s.
Collision with the wall stops integration.
input:
- init_state [[r_1, r_2, ...], [phi_1, phi_2, ...], [z_1, z_2, ...]] :
the coordinates of the starting points (list or numpy array)
OR list of lists OR numpy 2D array
- time : array of times requested (list or numpy array)
- direction : direction of integration 'FWD' forward or 'REV' reverse
(string)
- ripple : take into account magnetic ripple or not (boolean)
output:
returns a dictionary containing:
- r : radial coordinate (np.array)
- p : toridal coordinate (np.array)
- z : vertical coordinate (np.array)
- x : x coordinate (np.array)
- y : y coordinate (np.array)
- cp : collision point with the wall (list)
'''
# Step for the integration
if (stp is None):
stp = 0.001
# Length of the field line
if (length_line is None):
s = 100
else:
s = length_line
ds = np.linspace(0, s, int(s / stp))
init_state = np.squeeze(np.asarray(init_state))
if (init_state.ndim < 2):
init_state = init_state[:, np.newaxis]
self.ar_time = np.atleast_1d(np.squeeze(np.asarray([time])))
br_intp_t = np.atleast_2d(np.squeeze(self.f_intp_br(self.ar_time)))
bt_intp_t = np.atleast_2d(np.squeeze(self.f_intp_bt(self.ar_time)))
bz_intp_t = np.atleast_2d(np.squeeze(self.f_intp_bz(self.ar_time)))
# !!!!!!!!!!!!!!!!!!!!!
# HARD CODED CORRECTION
if (np.nanmean(bt_intp_t) > 0):
bt_intp_t *= -1
# !!!!!!!!!!!!!!!!!!!!!
# Interpolate current
itor_intp_t_vect = np.interp(self.ar_time, self.t_itor[:, 0], \
self.itor[:, 0])
b0_intp_t_vect = np.interp(self.ar_time, self.equi.time[self.mask], \
self.equi.vacuum_toroidal_field.b0[self.mask])
outMagLine = []
for ii in range(self.ar_time.size):
self.br_lin_intp = \
interpolate.LinearNDInterpolator(self.delaunay, br_intp_t[ii])
self.bt_lin_intp = \
interpolate.LinearNDInterpolator(self.delaunay, bt_intp_t[ii])
self.bz_lin_intp = \
interpolate.LinearNDInterpolator(self.delaunay, bz_intp_t[ii])
# Interpolated current and b0
self.itor_intp_t = itor_intp_t_vect[ii]
self.b0_intp_t = b0_intp_t_vect[ii]
out = []
for jj in range(init_state.shape[1]):
out_prime = self.integrate_solve_ivp(init_state[:, jj], \
direction, s, ds, ripple)
out_prime['init_point'] = init_state[:, jj]
out_prime['time'] = self.ar_time[ii]
# Return list of dict
out.append(out_prime)
outMagLine.append(out)
return outMagLine
def interp_quantity(quantity, interp_points, points, time_points, equi, \
ar_time, ar_R, ar_Phi, ar_Z, mask, mask_time, \
firstSpaceInterp, itor, t_itor, t_ignitron, \
no_ripple, ind_mid, equiDict):
value_interpolated = np.full((ar_time.size, ar_R.size), np.nan)
if (firstSpaceInterp):
value_interpSpace = np.full((time_points.size, ar_R.size), np.nan)
# Computation of requested quantities
if (quantity == 'b_field_norm'):
if (no_ripple):
b_field_norm = np.sqrt(equiDict['b_field_r']**2. \
+ equiDict['b_field_z']**2. \
+ equiDict['b_field_tor']**2.)
if (firstSpaceInterp):
# Space interpolation
for ii in range(time_points.size):
lin_intp = interpolate.LinearNDInterpolator(points, \
b_field_norm[mask, :][mask_time][ii, :])
value_interpSpace[ii, :] = lin_intp.__call__(interp_points)
else:
# Time interpolation
f_intp = interpolate.interp1d(equi.time[mask], \
b_field_norm[mask, :], axis=0, \
bounds_error=False)
else: # B_norm with ripple calculation
# Declaration arrays
br_intp = np.full((ar_time.size, ar_R.size), np.nan)
bt_intp = np.full((ar_time.size, ar_R.size), np.nan)
bz_intp = np.full((ar_time.size, ar_R.size), np.nan)
if (firstSpaceInterp):
br_Sintp = np.full((time_points.size, ar_R.size), np.nan)
bt_Sintp = np.full((time_points.size, ar_R.size), np.nan)
bz_Sintp = np.full((time_points.size, ar_R.size), np.nan)
# Space interpolation
for ii in range(time_points.size):
lin_intp = interpolate.LinearNDInterpolator(points, \
equiDict['b_field_r'][mask, :][mask_time][ii, :])
br_Sintp[ii, :] = lin_intp.__call__(interp_points)
lin_intp = interpolate.LinearNDInterpolator(points, \
equiDict['b_field_tor'][mask, :][mask_time][ii, :])
bt_Sintp[ii, :] = lin_intp.__call__(interp_points)
lin_intp = interpolate.LinearNDInterpolator(points, \
equiDict['b_field_z'][mask, :][mask_time][ii, :])
bz_Sintp[ii, :] = lin_intp.__call__(interp_points)
# Time interpolation
f_intp = interpolate.interp1d(time_points, \
br_Sintp, axis=0, \
bounds_error=False)
br_intp = np.atleast_2d(np.squeeze(f_intp(ar_time)))
f_intp = interpolate.interp1d(time_points, \
bt_Sintp, axis=0, \
bounds_error=False)
bt_intp = np.atleast_2d(np.squeeze(f_intp(ar_time)))
f_intp = interpolate.interp1d(time_points, \
bz_Sintp, axis=0, \
bounds_error=False)
bz_intp = np.atleast_2d(np.squeeze(f_intp(ar_time)))
else:
# Time interpolation
f_intp_br = interpolate.interp1d(equi.time[mask], \
equiDict['b_field_r'][mask, :], axis=0, \
bounds_error=False)
f_intp_bt = interpolate.interp1d(equi.time[mask], \
equiDict['b_field_tor'][mask, :], axis=0, \
bounds_error=False)
f_intp_bz = interpolate.interp1d(equi.time[mask], \
equiDict['b_field_z'][mask, :], axis=0, \
bounds_error=False)
br_intp_t = np.atleast_2d(np.squeeze(f_intp_br(ar_time)))
bt_intp_t = np.atleast_2d(np.squeeze(f_intp_bt(ar_time)))
bz_intp_t = np.atleast_2d(np.squeeze(f_intp_bz(ar_time)))
# Space interpolation
for ii in range(ar_time.size):
lin_intp = interpolate.LinearNDInterpolator(
points, br_intp_t[ii])
br_intp[ii, :] = lin_intp.__call__(interp_points)
lin_intp = interpolate.LinearNDInterpolator(
points, bt_intp_t[ii])
bt_intp[ii, :] = lin_intp.__call__(interp_points)
lin_intp = interpolate.LinearNDInterpolator(
points, bz_intp_t[ii])
bz_intp[ii, :] = lin_intp.__call__(interp_points)
# Interpolate current
itor_intp_t = np.interp(ar_time, t_itor[:, 0], itor[:, 0])
b0_intp_t = np.interp(ar_time, equi.time[mask], \
equi.vacuum_toroidal_field.b0[mask])
# Compute reference vaccuum magnetic field
bt_vac = equi.vacuum_toroidal_field.r0*b0_intp_t[:, np.newaxis] \
/ ar_R[np.newaxis, :]
# Compute magnetic field for given Phi
br_ripple, bt_ripple, bz_ripple = mr.mag_ripple(ar_R, ar_Phi, \
ar_Z, itor_intp_t)
# Check and correct Br if needed
r_ax = equi.time_slice[ind_mid].global_quantities.magnetic_axis.r
z_mid_ax = \
equi.time_slice[ind_mid].global_quantities.magnetic_axis.z \
+ 0.5*equi.time_slice[ind_mid].boundary.minor_radius
ind_p = np.abs((equiDict['r'] - r_ax)**2. \
+ (equiDict['z'] - z_mid_ax)**2).argmin()
if (equiDict['b_field_r'][ind_mid, ind_p] > 0):
br_intp *= -1
warnings.warn(
'Correcting b_field_r in b_field_norm, for negative toroidal current COCOS 11'
)
# Check and correct Bz if needed
r_mid_ax = \
equi.time_slice[ind_mid].global_quantities.magnetic_axis.r \
+ 0.5*equi.time_slice[ind_mid].boundary.minor_radius
z_ax = equi.time_slice[ind_mid].global_quantities.magnetic_axis.z
ind_p = np.abs((equiDict['r'] - r_mid_ax)**2. \
+ (equiDict['z'] - z_ax)**2).argmin()
if (equiDict['b_field_z'][ind_mid, ind_p] < 0):
bz_intp *= -1
warnings.warn(
'Correcting b_field_z in b_field_norm, for negative toroidal current COCOS 11'
)
# Check and correct Btor if needed
ind_p = np.abs((equiDict['r'] - r_ax)**2. \
+ (equiDict['z'] - z_ax)**2).argmin()
if (equiDict['b_field_tor'][ind_mid, ind_p] > 0):
bt_intp *= -1
warnings.warn(
'Correcting b_field_tor in b_field_norm, for negative toroidal field COCOS 11'
)
# Value interpolated
value_interpolated = \
np.sqrt((br_intp - br_ripple)**2. \
+ (np.abs(bt_intp - bt_ripple) - np.abs(bt_vac))**2. \
+ (bz_intp - bz_ripple)**2.)
return np.squeeze(value_interpolated)
elif (not no_ripple and re.match('b_field_*', quantity)):
# Declaration arrays
br_ripple = np.full((ar_time.size, ar_R.size), np.nan)
bt_ripple = np.full((ar_time.size, ar_R.size), np.nan)
bz_ripple = np.full((ar_time.size, ar_R.size), np.nan)
if (firstSpaceInterp):
# Space interpolation
quant_mask = eval('equiDict["' + quantity +
'"][mask, :][mask_time]')
for ii in range(time_points.size):
lin_intp = interpolate.LinearNDInterpolator(points, \
quant_mask[ii, :])
value_interpSpace[ii, :] = lin_intp.__call__(interp_points)
# Time interpolation
f_intp = interpolate.interp1d(time_points, \
value_interpSpace, axis=0, \
bounds_error=False)
value_interpolated = np.atleast_2d(np.squeeze(f_intp(ar_time)))
else:
# Time interpolation
f_intp = interpolate.interp1d(equi.time[mask], \
eval('equiDict["'+quantity+'"][mask, :]'), axis=0, \
bounds_error=False)
b_intp_t = np.atleast_2d(np.squeeze(f_intp(ar_time)))
# Space interpolation
for ii in range(ar_time.size):
lin_intp = interpolate.LinearNDInterpolator(
points, b_intp_t[ii])
value_interpolated[ii, :] = lin_intp.__call__(interp_points)
# Interpolate current
itor_intp_t = np.interp(ar_time, t_itor[:, 0], itor[:, 0])
# Compute magnetic field for given Phi
br_ripple, bt_ripple, bz_ripple = mr.mag_ripple(ar_R, ar_Phi, \
ar_Z, itor_intp_t)
if (quantity == 'b_field_r'):
r_ax = equi.time_slice[ind_mid].global_quantities.magnetic_axis.r
z_mid_ax = \
equi.time_slice[ind_mid].global_quantities.magnetic_axis.z \
+ 0.5*equi.time_slice[ind_mid].boundary.minor_radius
ind_p = np.abs((equiDict['r'] - r_ax)**2. \
+ (equiDict['z'] - z_mid_ax)**2).argmin()
if (equiDict['b_field_r'][ind_mid, ind_p] > 0):
value_interpolated *= -1
value_interpolated -= br_ripple
warnings.warn(
'Correcting b_field_r, for negative toroidal current COCOS 11'
)
else:
value_interpolated -= br_ripple
elif (quantity == 'b_field_z'):
r_mid_ax = \
equi.time_slice[ind_mid].global_quantities.magnetic_axis.r \
+ 0.5*equi.time_slice[ind_mid].boundary.minor_radius
z_ax = equi.time_slice[ind_mid].global_quantities.magnetic_axis.z
ind_p = np.abs((equiDict['r'] - r_mid_ax)**2. \
+ (equiDict['z'] - z_ax)**2).argmin()
if (equiDict['b_field_z'][ind_mid, ind_p] < 0):
value_interpolated *= -1
value_interpolated -= bz_ripple
warnings.warn(
'Correcting b_field_z, for negative toroidal current COCOS 11'
)
else:
value_interpolated -= bz_ripple
elif (quantity == 'b_field_tor'):
b0_intp_t = np.interp(ar_time, equi.time[mask], \
equi.vacuum_toroidal_field.b0[mask])
# Compute reference vaccuum magnetic field
bt_vac = equi.vacuum_toroidal_field.r0*b0_intp_t[:, np.newaxis] \
/ ar_R[np.newaxis, :]
r_ax = equi.time_slice[ind_mid].global_quantities.magnetic_axis.r
z_ax = equi.time_slice[ind_mid].global_quantities.magnetic_axis.z
ind_p = np.abs((equiDict['r'] - r_ax)**2. \
+ (equiDict['z'] - z_ax)**2).argmin()
if (equiDict['b_field_tor'][ind_mid, ind_p] > 0):
value_interpolated *= -1
value_interpolated -= (bt_ripple - bt_vac)
warnings.warn(
'Correcting b_field_tor, for negative toroidal field COCOS 11'
)
else:
value_interpolated -= (bt_ripple - bt_vac)
else:
print()
print('ERROR: not valid quantity input:', quantity)
print()
raise SyntaxError
return np.squeeze(value_interpolated)
elif (quantity == 'rho_pol_norm'):
rho_pol_norm = np.sqrt((equiDict['psi'] \
- equiDict['prof_1d_psi'][:, 0, np.newaxis]) \
/ (equiDict['prof_1d_psi'][:, -1, np.newaxis] \
- equiDict['prof_1d_psi'][:, 0, np.newaxis]))
if (firstSpaceInterp):
# Space interpolation
for ii in range(time_points.size):
lin_intp = interpolate.LinearNDInterpolator(points, \
rho_pol_norm[mask, :][mask_time][ii, :])
value_interpSpace[ii, :] = lin_intp.__call__(interp_points)
else:
# Time interpolation
f_intp = interpolate.interp1d(equi.time[mask], rho_pol_norm[mask, :], \
axis=0, bounds_error=False)
elif (quantity == 'rho_tor_norm' or quantity == 'rho_tor'):
if (firstSpaceInterp):
# Space interpolation
for ii in range(time_points.size):
lin_intp = interpolate.LinearNDInterpolator(points, \
equiDict['psi'][mask, :][mask_time][ii, :])
value_interpSpace[ii, :] = lin_intp.__call__(interp_points)
else:
# Time interpolation
f_intp = interpolate.interp1d(equi.time[mask], equiDict['psi'][mask, :], \
axis=0, bounds_error=False)
elif (quantity == 'theta'):
mag_ax_r = np.atleast_1d(np.squeeze(np.interp(ar_time, \
equi.time[mask], equiDict['mag_axis_r'][mask])))
mag_ax_z = np.atleast_1d(np.squeeze(np.interp(ar_time, \
equi.time[mask], equiDict['mag_axis_z'][mask])))
for ii in range(ar_time.size):
delta_R = ar_R - mag_ax_r[ii, np.newaxis]
delta_Z = ar_Z - mag_ax_z[ii, np.newaxis]
val_arctan = np.arctan(delta_Z / delta_R)
mask_theta = (delta_R >= 0) & (delta_Z >= 0)
value_interpolated[ii, mask_theta] = val_arctan[mask_theta]
mask_theta = (delta_R >= 0) & (delta_Z < 0)
value_interpolated[ii,
mask_theta] = 2 * np.pi + val_arctan[mask_theta]
mask_theta = (delta_R < 0)
value_interpolated[ii, mask_theta] = np.pi + val_arctan[mask_theta]
return np.squeeze(2. * np.pi - value_interpolated)
else:
if (firstSpaceInterp):
# Space interpolation
for ii in range(time_points.size):
lin_intp = interpolate.LinearNDInterpolator(points, \
eval('equiDict["'+quantity+'"][mask, :][mask_time][ii, :]'))
value_interpSpace[ii, :] = lin_intp.__call__(interp_points)
else:
# Time interpolation
f_intp = interpolate.interp1d(equi.time[mask], \
eval('equiDict["'+quantity+'"][mask, :]'), \
axis=0, bounds_error=False)
if (firstSpaceInterp):
# Time interpolation
f_intp = interpolate.interp1d(time_points, \
value_interpSpace, axis=0, \
bounds_error=False)
value_interpolated = np.atleast_2d(np.squeeze(f_intp(ar_time)))
else:
# Time interpolation
out_time_interp = np.atleast_2d(np.squeeze(f_intp(ar_time)))
# Space interpolation
for ii in range(ar_time.size):
lin_intp = interpolate.LinearNDInterpolator(
points, out_time_interp[ii])
value_interpolated[ii, :] = lin_intp.__call__(interp_points)
# Extra calculations for rho_tor, rho_tor_norm, b_field_r and b_field_z
if (quantity == 'rho_tor_norm' or quantity == 'rho_tor'):
value_interp_rho_tor = np.full(value_interpolated.shape, np.nan)
if (quantity == 'rho_tor_norm'):
try:
rho_tor_norm = equiDict['prof_1d_rho_tor'] \
/ equiDict['prof_1d_rho_tor'][:, -1, np.newaxis]
except ZeroDivisionError as err:
print('Division by zero for rho_tor_norm interpolation:', err)
raise
f_intp_rho_tor = interpolate.interp1d(equi.time[mask], \
rho_tor_norm[mask, :], axis=0, bounds_error=False)
elif (quantity == 'rho_tor'):
f_intp_rho_tor = interpolate.interp1d(equi.time[mask], \
equiDict['prof_1d_rho_tor'][mask, :], \
axis=0, bounds_error=False)
rho_tor1D = np.atleast_2d(np.squeeze(f_intp_rho_tor(ar_time)))
f_intp_psi = interpolate.interp1d(equi.time[mask], \
equiDict['prof_1d_psi'][mask, :], axis=0, bounds_error=False)
psi1D = np.atleast_2d(np.squeeze(f_intp_psi(ar_time)))
for ii in range(ar_time.size):
f_intp_psi_rho_tor = interpolate.interp1d(psi1D[ii], rho_tor1D[ii], \
bounds_error=False)
value_interp_rho_tor[ii] = f_intp_psi_rho_tor(
value_interpolated[ii])
return np.squeeze(value_interp_rho_tor)
elif (quantity == 'b_field_r'):
r_ax = equi.time_slice[ind_mid].global_quantities.magnetic_axis.r
z_mid_ax = equi.time_slice[ind_mid].global_quantities.magnetic_axis.z \
+ 0.5*equi.time_slice[ind_mid].boundary.minor_radius
ind_p = np.abs((equiDict['r'] - r_ax)**2. \
+ (equiDict['z'] - z_mid_ax)**2).argmin()
if (equiDict['b_field_r'][ind_mid, ind_p] > 0):
value_interpolated *= -1
warnings.warn('Correcting b_field_r, if Ip negative COCOS 11')
return np.squeeze(value_interpolated)
elif (quantity == 'b_field_z'):
r_mid_ax = equi.time_slice[ind_mid].global_quantities.magnetic_axis.r \
+ 0.5*equi.time_slice[ind_mid].boundary.minor_radius
z_ax = equi.time_slice[ind_mid].global_quantities.magnetic_axis.z
ind_p = np.abs((equiDict['r'] - r_mid_ax)**2. \
+ (equiDict['z'] - z_ax)**2).argmin()
if (equiDict['b_field_z'][ind_mid, ind_p] < 0):
value_interpolated *= -1
warnings.warn('Correcting b_field_z, if Ip negative COCOS 11')
return np.squeeze(value_interpolated)
elif (quantity == 'b_field_tor'):
r_ax = equi.time_slice[ind_mid].global_quantities.magnetic_axis.r
z_ax = equi.time_slice[ind_mid].global_quantities.magnetic_axis.z
ind_p = np.abs((equiDict['r'] - r_ax)**2. \
+ (equiDict['z'] - z_ax)**2).argmin()
if (equiDict['b_field_tor'][ind_mid, ind_p] > 0):
value_interpolated *= -1
warnings.warn(
'Correcting b_field_tor, if b_field_tor negative COCOS 11')
return np.squeeze(value_interpolated)
else:
return np.squeeze(value_interpolated)
nfiles = len(filenames)
beam_E = np.zeros((hp.nside2npix(nside), len(freqs)))
thetai, phii = hp.pix2ang(nside, np.arange(hp.nside2npix(nside)))
for fi, f in enumerate(filenames):
data = np.loadtxt(f, skiprows=2)
if opts.linear:
lat = data[:, 0] * np.pi / 180.0
nlat = len(lat)
lon = data[:, 1] * np.pi / 180.0
nlon = len(lon)
coords = np.stack((lat, lon)).T
gain = data[:, 2]
lut = interpolate.LinearNDInterpolator(coords, gain)
else:
lat = np.unique(data[:, 0]) * np.pi / 180.0
nlat = len(lat)
lon = np.unique(data[:, 1]) * np.pi / 180.0
nlon = len(lon)
gain = data[:, 2].reshape(nlon, nlat).transpose()
lut = interpolate.RectBivariateSpline(lat, lon, gain)
for i in np.arange(hp.nside2npix(nside)):
beam_E[i, fi] = lut(thetai[i], phii[i])
# Write to .fits file
if opts.linear:
filename = 'healpix_beam_%.0f-%.0fMHz_nside-%d_linear-interp.fits' %(freqs[0], freqs[-1], nside)
else:
filename = 'healpix_beam_%.0f-%.0fMHz_nside-%d.fits' %(freqs[0], freqs[-1], nside)
def compute_disps_on_grid(E_grid_data,
disp_sensor_coords,
L=1.,
w=0.2,
Nx=200,
Ny=40,
nu=0.3,
traction=0.1,
output_paraview=False):
'''
E_grid_data: #grid pts x 3 array with {x_pts, y_pts, E_vals}
disp_sensor_coords: # sensors x 2 array with x/y coordinates of u sensor loc
L: length
w: width
Nx: number of elements in x direction
Ny: ...
nu: poisson ratio
traction: applied traction value on right side
returns u_x, u_y - arrays with x/y components of displacement evaluated
at the 2d grid defined by disp_x/y_sensors
'''
mesh = RectangleMesh(Point(0., 0.), Point(L, w), Nx, Ny, "crossed")
interp = scipy_interp.LinearNDInterpolator(E_grid_data[:, 0:2],
E_grid_data[:, 2])
E = StiffnessInterp(interp)
D = FunctionSpace(mesh, "DG", 0)
E_field = interpolate(E, D)
nu = Constant(nu)
def eps(v):
return sym(grad(v))
# https://en.wikipedia.org/wiki/Hooke%27s_law#Plane_stress:
def sigma(v):
return E_field / (1. - nu**2) * (
(1. - nu) * eps(v) + nu * tr(eps(v)) * Identity(2))
V = VectorFunctionSpace(mesh, 'Lagrange', degree=2)
du = TrialFunction(V)
u_ = TestFunction(V)
a = inner(sigma(du), eps(u_)) * dx
# Handle boundary conditions
# Dirichlet displacement BC on the left side of beam:
def left(x, on_boundary):
return near(x[0], 0.)
def bottomleftcorner(x, on_boundary):
return (near(x[0], 0.) and near(x[1], 0.))
bc1 = DirichletBC(V.sub(0), Constant(0.), left)
bc2 = DirichletBC(V,
Constant((0., 0.)),
bottomleftcorner,
method='pointwise')
bcs = [bc1, bc2]
# Traction boundary condition on the right side of beam:
class Right(SubDomain):
def inside(self, x, on_boundary):
return near(x[0], L) and on_boundary
# Define traction:
T = Constant((traction, 0.0))
facets = MeshFunction("size_t", mesh, 1)
facets.set_all(0)
Right().mark(facets, 1)
ds = Measure('ds')(subdomain_data=facets)
# Define right hand side of variational form
l = dot(T, u_) * ds(1)
# Solve problem
u = Function(V, name="Displacement")
set_log_level(50)
solve(a == l, u, bcs)
if output_paraview:
file_results = XDMFFile("elasticity_results.xdmf")
file_results.parameters["flush_output"] = True
file_results.parameters["functions_share_mesh"] = True
file_results.write(u, 0.)
u_x = np.zeros(len(disp_sensor_coords))
u_y = np.zeros(len(disp_sensor_coords))
for i, coord in enumerate(disp_sensor_coords):
p = Point(coord[0], coord[1])
disp = u(p)
u_x[i] = disp[0]
u_y[i] = disp[1]
return u_x, u_y
def get_interpolation(cls, geo_data, **kwargs):
"""读等值线数据文件，并返回插值对象"""
x = geo_data.x
y = geo_data.y
z = geo_data.z
return interpolate.LinearNDInterpolator(np.array([x, y]).T, z)
def interpolate_heatmap(x, y, z, n: int=None, interp_method:str='linear'):
"""
Args:
x (array): x data points
y (array): y data points
z (array): z data points
n (int): number of points for each dimension on the interpolated
grid
interp_method {"linear", "nearest", "deg"} determines what interpolation
method is used.
Returns:
x_grid : N*1 array of x-values of the interpolated grid
y_grid : N*1 array of x-values of the interpolated grid
z_grid : N*N array of z-values that form a grid.
The output of this method can directly be used for
plt.imshow(z_grid, extent=extent, aspect='auto')
where the extent is determined by the min and max of the x_grid and
y_grid.
The output can also be used as input for
ax.pcolormesh(x, y, Z,**kw)
"""
points = list(zip(x, y))
lbrt = np.min(points, axis=0), np.max(points, axis=0)
lbrt = lbrt[0][0], lbrt[0][1], lbrt[1][0], lbrt[1][1]
xy_mean = np.mean([lbrt[0], lbrt[2]]), np.mean([lbrt[1], lbrt[3]])
xy_scale = np.ptp([lbrt[0], lbrt[2]]), np.ptp([lbrt[1], lbrt[3]])
# interpolation needs to happen on a rescaled grid, this is somewhat akin to an
# assumption in the interpolation that the scale of the experiment is chosen sensibly.
# N.B. even if interp_method == "nearest" the linear interpolation is used
# to determine the amount of grid points. Could be improved.
ip = interpolate.LinearNDInterpolator(
scale(points, xy_mean=xy_mean, xy_scale=xy_scale), z)
if n is None:
# Calculate how many grid points are needed.
# factor from A=√3/4 * a² (equilateral triangle)
# N.B. a factor 4 was added as there were to few points for uniform
# grid otherwise.
n = int(0.658 / np.sqrt(areas(ip).min()))*4
n = max(n, 10)
if n > 500:
logging.warning('n: {} larger than 500'.format(n))
n=500
x_lin = y_lin = np.linspace(-0.5, 0.5, n)
if interp_method == 'linear':
z_grid = ip(x_lin[:, None], y_lin[None, :]).squeeze()
elif interp_method == "nearest":
ip = interpolate.NearestNDInterpolator(
scale(points, xy_mean=xy_mean, xy_scale=xy_scale), z)
z_grid = ip(x_lin[:, None], y_lin[None, :]).squeeze()
elif interp_method == "deg":
# Circular interpolation in deg units
phases=np.deg2rad(z)
newdata_cos=np.cos(phases)
newdata_sin=np.sin(phases)
ip_cos = interpolate.LinearNDInterpolator(
scale(points, xy_mean=xy_mean, xy_scale=xy_scale), newdata_cos)
newdata_cos = ip_cos(x_lin[:, None], y_lin[None, :]).squeeze()
ip_sin = interpolate.LinearNDInterpolator(
scale(points, xy_mean=xy_mean, xy_scale=xy_scale), newdata_sin)
newdata_sin = ip_sin(x_lin[:, None], y_lin[None, :]).squeeze()
z_grid = (np.rad2deg(np.arctan2(newdata_sin, newdata_cos)) % 360).squeeze()
# x and y grid points need to be rescaled from the linearly chosen points
points_grid = unscale(list(zip(x_lin, y_lin)),
xy_mean=xy_mean, xy_scale=xy_scale)
x_grid = points_grid[:, 0]
y_grid = points_grid[:, 1]
return x_grid, y_grid, (z_grid).T
S = []
sig = []
#This here makes a grid for a given Re (alfa,S) and for each point sigma
for case in caseList:
# print case.alfa, case.Re
for p in zip(case.S, case.sigma):
alfa.append(float(case.alfa))
S.append(p[0])
sig.append(p[1])
p = [] #this grid (alfa,S) is stored in p, and used to find an interpolant
for pair in zip(alfa, S):
p.append([pair[0], pair[1]])
if InterType == 1:
interp = interpolate.LinearNDInterpolator(p, sig)
elif InterType == 2:
interp = interpolate.CloughTocher2DInterpolator(p, sig)
else:
raise "Unknown Interpolation type"
if with3D:
Anew = np.linspace(min(alfa), max(alfa), num=200, endpoint=True)
Snew = np.linspace(min(S), max(S), num=200, endpoint=True)
Anew, Snew = np.meshgrid(Anew, Snew)
zz = interp(Anew, Snew)
fig = plt.figure()
ax = fig.gca(projection='3d')
ax.set_xlabel('alpha')
ax.set_ylabel('S')
ax.set_zlabel('Sigma')
# lat_values_1d = lat_values.flatten()
# data_values_1d = data_values.flatten()
# data_interp_values = interpolate.griddata((lon_values_1d, lat_values_1d), data_values_1d, (lon0_values, lat0_values), method='nearest')
########
### For consistency I have to bring Vector of lats and cords to
### scipy.interpolate.LinearNDInterpolator
if lon_values.ndim == 1:
lon_values, lat_values = np.meshgrid(lon_values, lat_values)
###
lon_values_1d = lon_values.flatten()
lon_values_1d[lon_values_1d < 0] = lon_values_1d[lon_values_1d < 0] + 360
lat_values_1d = lat_values.flatten()
data_values_1d = data_values.flatten()
points = list(zip(lon_values_1d, lat_values_1d))
interp_surface = interpolate.LinearNDInterpolator(points, data_values_1d)
data_interp_values = interp_surface(lon0_grid, lat0_grid)
# data_interp_values = scipy.interpolate.interpn((lat, lon), data_values, (lat0, lon0), method='nearest')
# xr.concat([data0, data_interp], 'ensembles') # stopped using xarray
data_ensembles = np.dstack((data_ensembles, data_interp_values))
print('Finished #' + str(iexperiment) + ' From ' + str(len(cmip_list)))
fig, ax = plt.subplots(3)
# data.plot(ax=ax[0], vmin=5, vmax=30)
ax[0].pcolormesh(data, cmap='RdBu_r', vmin=5, vmax=30)
ax[0].set_title('Data : ' + str(ds.source_id))
ax[1].pcolormesh(lon0_values,
lat0_values,
data0_values,
vmin=5,
dis_map_h = np.load('../data/distortion/194-437-xa_high-centroids.npy')
#dis_map_l = np.load('../data/distortion/5-185-xa_low-centroids.npy')
#dis_map_h = np.load('../data/distortion/5-185-xa_high-centroids.npy')
dis_map_l = dis_map_l.reshape(
50 * 50, 2) / detector_size[0] * 36000 * 800 * 0.001666
dis_map_h = dis_map_h.reshape(
50 * 50, 2) / detector_size[0] * 36000 * 800 * 0.001666
cen_x = np.arange(detector_size[0]) + 0.5
cen_y = np.arange(detector_size[1]) + 0.5
xi = np.repeat(detector2gon(cen_x, detector_size[0]), detector_size[1])
eta = np.tile(detector2gon(cen_y, detector_size[1]), detector_size[0])
points = np.array([xi, eta]).T
fxi_l = interpolate.LinearNDInterpolator(points, dis_map_l[:, 0])
feta_l = interpolate.LinearNDInterpolator(points, dis_map_l[:, 1])
fxi_h = interpolate.LinearNDInterpolator(points, dis_map_h[:, 0])
feta_h = interpolate.LinearNDInterpolator(points, dis_map_h[:, 1])
dis_map = {
'xi_l': fxi_l,
'xi_h': fxi_h,
'eta_l': feta_l,
'eta_h': feta_h
}
manager = Manager()
return_dict = manager.dict()
p_list = []
pid = 0
def _get_t_isochrone(self,
age,
metal=None,
FeH=None,
masses=None,
*args,
**kwargs):
""" Retrieve isochrone from the original source
internal use to adapt any library
"""
# make sure unit is in years and then only give the value (no units)
_age = int(units.Quantity(age, units.year).value)
# if hasUnit(age):
# _age = int(age.to('yr').magnitude)
# else:
# _age = int(age * inputUnit.to('yr').magnitude)
_logA = np.log10(_age)
assert ((metal is not None) | (FeH is not None)),\
"Need a chemical par. value."
if (metal is not None) & (FeH is not None):
print("Warning: both Z & [Fe/H] provided, ignoring [Fe/H].")
if metal is None:
metal = self.FeHtometal(FeH)
if self.interpolation():
# Do the actual nd interpolation
# Maybe already exists?
if (metal in self.Z) & (_age in self.ages):
t = self.selectWhere(
'*',
'(round(Z, 6) == {0}) & (round(logA, 6) == {1})'.format(
metal, _logA))
if t.nrows > 0:
return t
# apparently not
# find 2 closest metal values
ca1 = (self.ages <= _age)
ca2 = (self.ages > _age)
cz1 = (self.Z <= metal)
cz2 = (self.Z > metal)
if (metal in self.Z):
# perfect match in metal, need to find ages
if (_age in self.ages):
return self.selectWhere(
'*', '(round(Z, 6) == {0}) & (round(logA, 6) == {1})'.
format(metal, _logA))
elif (True in ca1) & (True in ca2):
# bracket on _age: closest values
a1, a2 = (np.log10(max(self.ages[ca1])),
np.log10(min(self.ages[ca2])))
iso = self.selectWhere(
'*',
'(Z == 0.02) & ( (abs(logA - {0}) < 1e-4) | (abs(logA - {1}) < 1e-4 ) )'
.format(a1, a2))
if masses is None:
_logM = np.unique(iso['logM'])
else:
_logM = masses
# define interpolator
points = np.array([self[k]
for k in 'logA logM Z'.split()]).T
values = np.array(
[self[k] for k in list(self.data.keys())]).T
_ifunc = interpolate.LinearNDInterpolator(points, values)
pts = np.array([(_logA, logMk, metal) for logMk in _logM])
r = _ifunc(pts)
return Table(r)
else:
raise Exception('Age not covered by the isochrones')
elif (True in cz1) & (True in cz2):
# need to find closest Z
pass
return
else:
# find the closest match
_Z = self.Z[((metal - self.Z)**2).argmin()]
# _logA = np.log10(self.ages[((_age - self.ages) ** 2).argmin()])
_logA = self.logages[((np.log10(_age) - self.logages)**2).argmin()]
tab = self.data.selectWhere(
'*', "(round(Z, 6) == {0}) & (round(logA,6) == {1})".format(
_Z, _logA))
# mass selection
if masses is not None:
# masses are expected in logM for interpolation
# if masses.max() > 2.3:
# _m = np.log10(masses)
# else:
_m = masses
data_logM = tab['logM'][:]
# refuse extrapolation!
# ind = np.where(_m <= max(data_logM))
data = {}
for kn in list(tab.keys()):
data[kn] = interp(_m,
data_logM,
tab[kn],
left=np.nan,
right=np.nan)
return Table(data)
def refine_tree_xyz(
mesh, xyz,
method="radial",
octree_levels=[1, 1, 1],
octree_levels_padding=None,
finalize=False,
min_level=0,
max_distance=np.inf
):
"""
Refine a TreeMesh based on xyz point locations
Parameters
----------
mesh: TreeMesh
The TreeMesh object to be refined
xyz: numpy.ndarray
2D array of points
method: str
Method used to refine the mesh based on xyz locations
- "radial": Based on radial distance xyz and cell centers
- "surface": Along triangulated surface repeated vertically
- "box": Inside limits defined by outer xyz locations
octree_levels: list
Minimum number of cells around points in each k octree level
[N(k), N(k-1), ...]
octree_levels_padding: list
Padding cells added to the outer limits of the data each octree levels
used for method= "surface" and "box" [N(k), N(k-1), ...].
finalize: bool
True | [False] Finalize the TreeMesh
max_distance: float
Maximum refinement distance from xyz locations.
Used for method="surface" to reduce interpolation distance.
Returns
--------
discretize.TreeMesh
mesh
"""
if octree_levels_padding is not None:
if len(octree_levels_padding) != len(octree_levels):
raise ValueError(
"'octree_levels_padding' must be the length %i" % len(octree_levels)
)
else:
octree_levels_padding = np.zeros_like(octree_levels)
# Prime the refinement against large cells
mesh.insert_cells(
xyz,
[mesh.max_level - np.nonzero(octree_levels)[0][0]]*xyz.shape[0],
finalize=False
)
# Trigger different refine methods
if method.lower() == "radial":
# Build a cKDTree for fast nearest lookup
tree = cKDTree(xyz)
# Compute the outer limits of each octree level
rMax = np.cumsum(
mesh.hx.min() *
np.asarray(octree_levels) *
2**np.arange(len(octree_levels))
)
# Radial function
def inBall(cell):
xyz = cell.center
r, ind = tree.query(xyz)
for ii, nC in enumerate(octree_levels):
if r < rMax[ii]:
return mesh.max_level-ii
return min_level
mesh.refine(inBall, finalize=finalize)
elif method.lower() == 'surface':
# Compute centroid
centroid = np.mean(xyz, axis=0)
if mesh.dim == 2:
rOut = np.abs(centroid[0]-xyz).max()
hz = mesh.hy.min()
else:
# Largest outer point distance
rOut = np.linalg.norm(np.r_[
np.abs(centroid[0]-xyz[:, 0]).max(),
np.abs(centroid[1]-xyz[:, 1]).max()
]
)
hz = mesh.hz.min()
# Compute maximum depth of refinement
zmax = np.cumsum(
hz *
np.asarray(octree_levels) *
2**np.arange(len(octree_levels))
)
# Compute maximum horizontal padding offset
padWidth = np.cumsum(
mesh.hx.min() *
np.asarray(octree_levels_padding) *
2**np.arange(len(octree_levels_padding))
)
# Increment the vertical offset
zOffset = 0
xyPad = -1
depth = zmax[-1]
# Cycle through the Tree levels backward
for ii in range(len(octree_levels)-1, -1, -1):
dx = mesh.hx.min() * 2**ii
if mesh.dim == 3:
dy = mesh.hy.min() * 2**ii
dz = mesh.hz.min() * 2**ii
else:
dz = mesh.hy.min() * 2**ii
# Increase the horizontal extent of the surface
if xyPad != padWidth[ii]:
xyPad = padWidth[ii]
# Calculate expansion for padding XY cells
expansion_factor = (rOut + xyPad) / rOut
xLoc = (xyz - centroid)*expansion_factor + centroid
if mesh.dim == 3:
# Create a new triangulated surface
tri2D = Delaunay(xLoc[:, :2])
F = interpolate.LinearNDInterpolator(tri2D, xLoc[:, 2])
else:
F = interpolate.interp1d(xLoc[:, 0], xLoc[:, 1], fill_value='extrapolate')
limx = np.r_[xLoc[:, 0].max(), xLoc[:, 0].min()]
nCx = int(np.ceil((limx[0]-limx[1]) / dx))
if mesh.dim == 3:
limy = np.r_[xLoc[:, 1].max(), xLoc[:, 1].min()]
nCy = int(np.ceil((limy[0]-limy[1]) / dy))
# Create a grid at the octree level in xy
CCx, CCy = np.meshgrid(
np.linspace(
limx[1], limx[0], nCx
),
np.linspace(
limy[1], limy[0], nCy
)
)
xy = np.c_[CCx.reshape(-1), CCy.reshape(-1)]
# Only keep points within triangulation
indexTri = tri2D.find_simplex(xy)
else:
xy = np.linspace(
limx[1], limx[0], nCx
)
indexTri = np.ones_like(xy, dtype='bool')
# Interpolate the elevation linearly
z = F(xy[indexTri != -1])
newLoc = np.c_[xy[indexTri != -1], z]
# Only keep points within max_distance
tree = cKDTree(xyz)
r, ind = tree.query(newLoc)
# Apply vertical padding for current octree level
dim = mesh.dim - 1
zOffset = 0
while zOffset < depth:
indIn = r < (max_distance + padWidth[ii])
nnz = int(np.sum(indIn))
if nnz > 0:
mesh.insert_cells(
np.c_[
newLoc[indIn, :dim],
newLoc[indIn, -1]-zOffset],
np.ones(nnz)*mesh.max_level-ii,
finalize=False
)
zOffset += dz
depth -= dz * octree_levels[ii]
if finalize:
mesh.finalize()
elif method.lower() == 'box':
# Define the data extend [bottom SW, top NE]
bsw = np.min(xyz, axis=0)
tne = np.max(xyz, axis=0)
hx = mesh.hx.min()
if mesh.dim == 2:
hz = mesh.hy.min()
else:
hz = mesh.hz.min()
# Pre-calculate max depth of each level
zmax = np.cumsum(
hz * np.asarray(octree_levels) *
2**np.arange(len(octree_levels))
)
if mesh.dim == 2:
# Pre-calculate outer extent of each level
padWidth = np.cumsum(
mesh.hx.min() *
np.asarray(octree_levels_padding) *
2**np.arange(len(octree_levels_padding))
)
# Make a list of outer limits
BSW = [
bsw - np.r_[padWidth[ii], zmax[ii]]
for ii, (octZ, octXY) in enumerate(
zip(octree_levels, octree_levels_padding)
)
]
TNE = [
tne + np.r_[padWidth[ii], zmax[ii]]
for ii, (octZ, octXY) in enumerate(
zip(octree_levels, octree_levels_padding)
)
]
else:
hy = mesh.hy.min()
# Pre-calculate outer X extent of each level
padWidth_x = np.cumsum(
hx * np.asarray(octree_levels_padding) *
2**np.arange(len(octree_levels_padding))
)
# Pre-calculate outer Y extent of each level
padWidth_y = np.cumsum(
hy * np.asarray(octree_levels_padding) *
2**np.arange(len(octree_levels_padding))
)
# Make a list of outer limits
BSW = [
bsw - np.r_[padWidth_x[ii], padWidth_y[ii], zmax[ii]]
for ii, (octZ, octXY) in enumerate(
zip(octree_levels, octree_levels_padding)
)
]
TNE = [
tne + np.r_[padWidth_x[ii], padWidth_y[ii], zmax[ii]]
for ii, (octZ, octXY) in enumerate(
zip(octree_levels, octree_levels_padding)
)
]
def inBox(cell):
xyz = cell.center
for ii, (nC, bsw, tne) in enumerate(zip(octree_levels, BSW, TNE)):
if np.all([xyz > bsw, xyz < tne]):
return mesh.max_level-ii
return cell._level
mesh.refine(inBox, finalize=finalize)
else:
raise NotImplementedError(
"Only method= 'radial', 'surface'"
" or 'box' have been implemented"
)
return mesh
def write2netCDF(filename,
lon,
lat,
z,
increments=None,
nSamples=None,
title='CSI product',
name='z',
scale=1.0,
offset=0.0,
mask=None,
xyunits=['Lon', 'Lat'],
units='None',
interpolation=True,
verbose=True,
noValues=np.nan):
'''
Creates a netCDF file with the arrays in Z.
Z can be list of array or an array, the size of lon.
.. Args:
* filename -> Output file name
* lon -> 1D Array of lon values
* lat -> 1D Array of lat values
* z -> 2D slice to be saved
* mask -> if not None, must be a 2d-array of a polynome to mask
what is outside of it. This option is really long, so I don't
use it...
.. Kwargs:
* title -> Title for the grd file
* name -> Name of the field in the grd file
* scale -> Scale value in the grd file
* offset -> Offset value in the grd file
.. Returns:
* None
'''
if interpolation:
# Check
if nSamples is not None:
if type(nSamples) is int:
nSamples = [nSamples, nSamples]
dlon = (lon.max() - lon.min()) / nSamples[0]
dlat = (lat.max() - lat.min()) / nSamples[1]
if increments is not None:
dlon, dlat = increments
# Resample on a regular grid
olon, olat = np.meshgrid(np.arange(lon.min(), lon.max(), dlon),
np.arange(lat.min(), lat.max(), dlat))
else:
# Get lon lat
olon = lon
olat = lat
if increments is not None:
dlon, dlat = increments
else:
dlon = olon[0, 1] - olon[0, 0]
dlat = olat[1, 0] - olat[0, 0]
# Create a file
fid = netcdf(filename, 'w')
# Create a dimension variable
fid.createDimension('side', 2)
if verbose:
print('Create dimension xysize with size {}'.format(np.prod(
olon.shape)))
fid.createDimension('xysize', np.prod(olon.shape))
# Range variables
fid.createVariable('x_range', 'd', ('side', ))
fid.variables['x_range'].units = xyunits[0]
fid.createVariable('y_range', 'd', ('side', ))
fid.variables['y_range'].units = xyunits[1]
# Spacing
fid.createVariable('spacing', 'd', ('side', ))
fid.createVariable('dimension', 'i4', ('side', ))
# Informations
if title is not None:
fid.title = title
fid.source = 'CSI.utils.write2netCDF'
# Filing rnage and spacing
if verbose:
print('x_range from {} to {} with spacing {}'.format(
olon[0, 0], olon[0, -1], dlon))
fid.variables['x_range'][0] = olon[0, 0]
fid.variables['x_range'][1] = olon[0, -1]
fid.variables['spacing'][0] = dlon
if verbose:
print('y_range from {} to {} with spacing {}'.format(
olat[0, 0], olat[-1, 0], dlat))
fid.variables['y_range'][0] = olat[0, 0]
fid.variables['y_range'][1] = olat[-1, 0]
fid.variables['spacing'][1] = dlat
if interpolation:
# Interpolate
interpZ = sciint.LinearNDInterpolator(np.vstack((lon, lat)).T,
z,
fill_value=noValues)
oZ = interpZ(olon, olat)
else:
# Get values
oZ = z
# Masking?
if mask is not None:
# Import matplotlib.path
import matplotlib.path as path
# Create the path
poly = path.Path([[lo, la] for lo, la in zip(mask[:, 0], mask[:, 1])],
closed=False)
# Create the list of points
xy = np.vstack((olon.flatten(), olat.flatten())).T
# Findthose outside
bol = poly.contains_points(xy)
# Mask those out
oZ = oZ.flatten()
oZ[bol] = np.nan
oZ = oZ.reshape(olon.shape)
# Range
zmin = np.nanmin(oZ)
zmax = np.nanmax(oZ)
fid.createVariable('{}_range'.format(name), 'd', ('side', ))
fid.variables['{}_range'.format(name)].units = units
fid.variables['{}_range'.format(name)][0] = zmin
fid.variables['{}_range'.format(name)][1] = zmax
# Create Variable
fid.createVariable(name, 'd', ('xysize', ))
fid.variables[name].long_name = name
fid.variables[name].scale_factor = scale
fid.variables[name].add_offset = offset
fid.variables[name].node_offset = 0
# Fill it
fid.variables[name][:] = np.flipud(oZ).flatten()
# Set dimension
fid.variables['dimension'][:] = oZ.shape[::-1]
# Synchronize and close
fid.sync()
fid.close()
# All done
return
def ax_scalp(v,
channels,
ax=None,
annotate=False,
vmin=None,
vmax=None,
colormap=None):
"""Draw a scalp plot.
Draws a scalp plot on an existing axes. The method takes an array of
values and an array of the corresponding channel names. It matches
the channel names with an internal list of known channels and their
positions to project them correctly on the scalp.
.. warning:: The behaviour for unkown channels is undefined.
Parameters
----------
v : 1d-array of floats
The values for the channels
channels : 1d array of strings
The corresponding channel names for the values in ``v``
ax : Axes, optional
The axes to draw the scalp plot on. If not provided, the
currently activated axes (i.e. ``gca()``) will be taken
annotate : Boolean, optional
Draw the channel names next to the channel markers.
vmin, vmax : float, optional
The display limits for the values in ``v``. If the data in ``v``
contains values between -3..3 and ``vmin`` and ``vmax`` are set
to -1 and 1, all values smaller than -1 and bigger than 1 will
appear the same as -1 and 1. If not set, the maximum absolute
value in ``v`` is taken to calculate both values.
colormap : matplotlib.colors.colormap, optional
A colormap to define the color transitions.
Returns
-------
ax : Axes
the axes on which the plot was drawn
See Also
--------
ax_colorbar
"""
if ax is None:
ax = plt.gca()
# what if we have an unknown channel?
points = [get_channelpos(c) for c in channels]
# calculate the interpolation
x = [i[0] for i in points]
y = [i[1] for i in points]
z = v
# interplolate the in-between values
xx = np.linspace(min(x), max(x), 500)
yy = np.linspace(min(y), max(y), 500)
xx, yy = np.meshgrid(xx, yy)
f = interpolate.LinearNDInterpolator(list(zip(x, y)), z)
zz = f(xx, yy)
# draw the contour map
ctr = ax.contourf(xx, yy, zz, 20, vmin=vmin, vmax=vmax, cmap=colormap)
ax.contour(xx, yy, zz, 5, colors="k", vmin=vmin, vmax=vmax, linewidths=.1)
# paint the head
ax.add_artist(
plt.Circle((0, 0), 1, linestyle='solid', linewidth=2, fill=False))
# add a nose
ax.plot([-0.1, 0, 0.1], [1, 1.1, 1], 'k-')
# add markers at channels positions
ax.plot(x, y, 'k+')
# set the axes limits, so the figure is centered on the scalp
ax.set_ylim([-1.05, 1.15])
ax.set_xlim([-1.15, 1.15])
# hide the frame and axes
# hiding the axes might be too much, as this will also hide the x/y
# labels :/
ax.set_frame_on(False)
ax.get_xaxis().set_visible(False)
ax.get_yaxis().set_visible(False)
# draw the channel names
if annotate:
for i in zip(channels, list(zip(x, y))):
ax.annotate(" " + i[0], i[1])
ax.set_aspect(1)
plt.sci(ctr)
return ax
def refine_tree_xyz(
mesh,
xyz,
method="radial",
octree_levels=[1, 1, 1],
octree_levels_padding=None,
finalize=False,
min_level=0,
max_distance=np.inf,
):
"""Refine region within a :class:`~discretize.TreeMesh`
This function refines the specified region of a tree mesh using
one of several methods. These are summarized below:
**radial:** refines based on radial distances from a set of xy[z] locations.
Consider a tree mesh whose smallest cell size has a width of *h* . And
*octree_levels = [nc1, nc2, nc3, ...]* . Within a distance of *nc1 x h*
from any of the points supplied, the smallest cell size is used. Within a distance of
*nc2 x (2h)* , the cells will have a width of *2h* . Within a distance of *nc3 x (4h)* ,
the cells will have a width of *4h* . Etc...
**surface:** refines downward from a triangulated surface.
Consider a tree mesh whose smallest cell size has a width of *h*. And
*octree_levels = [nc1, nc2, nc3, ...]* . Within a downward distance of *nc1 x h*
from the topography (*xy[z]* ) supplied, the smallest cell size is used. The
topography is triangulated if the points supplied are coarser than the cell
size. No refinement is done above the topography. Within a vertical distance of
*nc2 x (2h)* , the cells will have a width of *2h* . Within a vertical distance
of *nc3 x (4h)* , the cells will have a width of *4h* . Etc...
**box:** refines inside the convex hull defined by the xy[z] locations.
Consider a tree mesh whose smallest cell size has a width of *h*. And
*octree_levels = [nc1, nc2, nc3, ...]* . Within the convex hull defined by *xyz* ,
the smallest cell size is used. Within a distance of *nc2 x (2h)* from that convex
hull, the cells will have a width of *2h* . Within a distance of *nc3 x (4h)* ,
the cells will have a width of *4h* . Etc...
Parameters
----------
mesh : discretize.TreeMesh
The tree mesh object to be refined
xyz : numpy.ndarray
2D array of points (n, dim)
method : {'radial', 'surface', 'box'}
Method used to refine the mesh based on xyz locations.
- *radial:* Based on radial distance xy[z] and cell centers
- *surface:* Refines downward from a triangulated surface
- *box:* Inside limits defined by outer xy[z] locations
octree_levels : list of int, optional
Minimum number of cells around points in each *k* octree level
starting from the smallest cells size; i.e. *[nc(k), nc(k-1), ...]* .
Note that you *can* set entries to 0; e.g. you don't want to discretize
using the smallest cell size.
octree_levels_padding : list of int, optional
Padding cells added to extend the region of refinement at each level.
Used for *method = surface* and *box*. Has the form *[nc(k), nc(k-1), ...]*
finalize : bool, optional
Finalize the tree mesh.
min_level : int, optional
Sets the largest cell size allowed in the mesh. The default (*0*),
allows the largest cell size to be used.
max_distance : float
Maximum refinement distance from xy[z] locations.
Used if *method* = "surface" to reduce interpolation distance
Returns
-------
discretize.TreeMesh
The refined tree mesh
Examples
--------
Here we use the **refine_tree_xyz** function refine a tree mesh
based on topography as well as a cluster of points.
>>> from discretize import TreeMesh
>>> from discretize.utils import mkvc, refine_tree_xyz
>>> import matplotlib.pyplot as plt
>>> import numpy as np
>>> dx = 5 # minimum cell width (base mesh cell width) in x
>>> dy = 5 # minimum cell width (base mesh cell width) in y
>>> x_length = 300.0 # domain width in x
>>> y_length = 300.0 # domain width in y
Compute number of base mesh cells required in x and y
>>> nbcx = 2 ** int(np.round(np.log(x_length / dx) / np.log(2.0)))
>>> nbcy = 2 ** int(np.round(np.log(y_length / dy) / np.log(2.0)))
Define the base mesh
>>> hx = [(dx, nbcx)]
>>> hy = [(dy, nbcy)]
>>> mesh = TreeMesh([hx, hy], x0="CC")
Refine surface topography
>>> xx = mesh.vectorNx
>>> yy = -3 * np.exp((xx ** 2) / 100 ** 2) + 50.0
>>> pts = np.c_[mkvc(xx), mkvc(yy)]
>>> mesh = refine_tree_xyz(
... mesh, pts, octree_levels=[2, 4], method="surface", finalize=False
... )
Refine mesh near points
>>> xx = np.array([-10.0, 10.0, 10.0, -10.0])
>>> yy = np.array([-40.0, -40.0, -60.0, -60.0])
>>> pts = np.c_[mkvc(xx), mkvc(yy)]
>>> mesh = refine_tree_xyz(
... mesh, pts, octree_levels=[4, 2], method="radial", finalize=True
... )
Plot the mesh
>>> fig = plt.figure(figsize=(6, 6))
>>> ax = fig.add_subplot(111)
>>> mesh.plotGrid(ax=ax)
>>> ax.set_xbound(mesh.x0[0], mesh.x0[0] + np.sum(mesh.hx))
>>> ax.set_ybound(mesh.x0[1], mesh.x0[1] + np.sum(mesh.hy))
>>> ax.set_title("QuadTree Mesh")
>>> plt.show()
"""
if octree_levels_padding is not None:
if len(octree_levels_padding) != len(octree_levels):
raise ValueError("'octree_levels_padding' must be the length %i" %
len(octree_levels))
else:
octree_levels_padding = np.zeros_like(octree_levels)
octree_levels = np.asarray(octree_levels)
octree_levels_padding = np.asarray(octree_levels_padding)
# Trigger different refine methods
if method.lower() == "radial":
# Compute the outer limits of each octree level
rMax = np.cumsum(mesh.h[0].min() * octree_levels *
2**np.arange(len(octree_levels)))
rs = np.ones(xyz.shape[0])
level = np.ones(xyz.shape[0], dtype=np.int32)
for ii, nC in enumerate(octree_levels):
# skip "zero" sized balls
if rMax[ii] > 0:
mesh.refine_ball(xyz,
rs * rMax[ii],
level * (mesh.max_level - ii),
finalize=False)
if finalize:
mesh.finalize()
elif method.lower() == "surface":
# Compute centroid
centroid = np.mean(xyz, axis=0)
if mesh.dim == 2:
rOut = np.abs(centroid[0] - xyz).max()
hz = mesh.h[1].min()
else:
# Largest outer point distance
rOut = np.linalg.norm(np.r_[np.abs(centroid[0] - xyz[:, 0]).max(),
np.abs(centroid[1] -
xyz[:, 1]).max(), ])
hz = mesh.h[2].min()
# Compute maximum depth of refinement
zmax = np.cumsum(hz * octree_levels * 2**np.arange(len(octree_levels)))
# Compute maximum horizontal padding offset
padWidth = np.cumsum(mesh.h[0].min() * octree_levels_padding *
2**np.arange(len(octree_levels_padding)))
# Increment the vertical offset
zOffset = 0
xyPad = -1
depth = zmax[-1]
# Cycle through the Tree levels backward
for ii in range(len(octree_levels) - 1, -1, -1):
dx = mesh.h[0].min() * 2**ii
if mesh.dim == 3:
dy = mesh.h[1].min() * 2**ii
dz = mesh.h[2].min() * 2**ii
else:
dz = mesh.h[1].min() * 2**ii
# Increase the horizontal extent of the surface
if xyPad != padWidth[ii]:
xyPad = padWidth[ii]
# Calculate expansion for padding XY cells
expansion_factor = (rOut + xyPad) / rOut
xLoc = (xyz - centroid) * expansion_factor + centroid
if mesh.dim == 3:
# Create a new triangulated surface
tri2D = Delaunay(xLoc[:, :2])
F = interpolate.LinearNDInterpolator(tri2D, xLoc[:, 2])
else:
F = interpolate.interp1d(xLoc[:, 0],
xLoc[:, 1],
fill_value="extrapolate")
limx = np.r_[xLoc[:, 0].max(), xLoc[:, 0].min()]
nCx = int(np.ceil((limx[0] - limx[1]) / dx))
if mesh.dim == 3:
limy = np.r_[xLoc[:, 1].max(), xLoc[:, 1].min()]
nCy = int(np.ceil((limy[0] - limy[1]) / dy))
# Create a grid at the octree level in xy
CCx, CCy = np.meshgrid(
np.linspace(limx[1], limx[0], nCx),
np.linspace(limy[1], limy[0], nCy),
)
xy = np.c_[CCx.reshape(-1), CCy.reshape(-1)]
# Only keep points within triangulation
indexTri = tri2D.find_simplex(xy)
else:
xy = np.linspace(limx[1], limx[0], nCx)
indexTri = np.ones_like(xy, dtype="bool")
# Interpolate the elevation linearly
z = F(xy[indexTri != -1])
newLoc = np.c_[xy[indexTri != -1], z]
# Only keep points within max_distance
tree = cKDTree(xyz)
r, ind = tree.query(newLoc)
# Apply vertical padding for current octree level
dim = mesh.dim - 1
zOffset = 0
while zOffset < depth:
indIn = r < (max_distance + padWidth[ii])
nnz = int(np.sum(indIn))
if nnz > 0:
mesh.insert_cells(
np.c_[newLoc[indIn, :dim],
newLoc[indIn, -1] - zOffset],
np.ones(nnz) * mesh.max_level - ii,
finalize=False,
)
zOffset += dz
depth -= dz * octree_levels[ii]
if finalize:
mesh.finalize()
elif method.lower() == "box":
# Define the data extent [bottom SW, top NE]
bsw = np.min(xyz, axis=0)
tne = np.max(xyz, axis=0)
hs = np.asarray([h.min() for h in mesh.h])
hx = hs[0]
hz = hs[-1]
# Pre-calculate outer extent of each level
# x_pad
padWidth = np.cumsum(hx * octree_levels_padding *
2**np.arange(len(octree_levels)))
if mesh.dim == 3:
# y_pad
hy = hs[1]
padWidth = np.c_[padWidth,
np.cumsum(hy * octree_levels_padding *
2**np.arange(len(octree_levels))), ]
# Pre-calculate max depth of each level
padWidth = np.c_[padWidth,
np.cumsum(hz * np.maximum(octree_levels - 1, 0) *
2**np.arange(len(octree_levels))), ]
levels = []
BSW = []
TNE = []
for ii, octZ in enumerate(octree_levels):
if octZ > 0:
levels.append(mesh.max_level - ii)
BSW.append(bsw - padWidth[ii])
TNE.append(tne + padWidth[ii])
mesh.refine_box(BSW, TNE, levels, finalize=finalize)
else:
raise NotImplementedError("Only method= 'radial', 'surface'"
" or 'box' have been implemented")
return mesh
matData = sio.loadmat(sys.argv[1])
cellWidth = matData['cellWidth']
LonPos = matData['lon'].T * dtor
LatPos = matData['lat'].T * dtor
minCellWidth = cellWidth.min()
meshDensityVsLatLon = (minCellWidth / cellWidth)**4
print 'minimum cell width in grid definition:', minCellWidth
print 'maximum cell width in grid definition:', cellWidth.max()
print 'cellWidth, south to north:'
print cellWidth
print 'meshDensityVsLatLon, south to north, smallest cell gets a 1:'
print meshDensityVsLatLon
LON, LAT = np.meshgrid(LonPos, LatPos)
ds = nc4.Dataset(sys.argv[2], 'r+')
meshDensity = ds.variables["meshDensity"][:]
print "Preparing interpolation of meshDensity from lat/lon to mesh..."
meshDensityInterp = interpolate.LinearNDInterpolator(
np.vstack((LAT.ravel(), LON.ravel())).T, meshDensityVsLatLon.ravel())
print "Interpolating and writing meshDensity..."
ds.variables['meshDensity'][:] = meshDensityInterp(
np.vstack(
(ds.variables['latCell'][:],
np.mod(ds.variables['lonCell'][:] + np.pi, 2 * np.pi) - np.pi)).T)
ds.close()
def active_from_xyz(mesh, xyz, grid_reference="CC", method="linear"):
"""Return boolean array indicating which cells are below surface
For a set of locations defining a surface, **active_from_xyz** outputs a
boolean array indicating which mesh cells like below the surface points.
This method uses SciPy's interpolation routine to interpolate between
location points defining the surface. Nearest neighbour interpolation
is used for cells outside the convex hull of the surface points.
Parameters
----------
mesh : discretize.TensorMesh or discretize.TreeMesh or discretize.CylindricalMesh
Mesh object. If *mesh* is a cylindrical mesh, it must be symmetric
xyz : (N, dim) numpy.ndarray
Points defining the surface topography.
grid_reference : {'CC', 'N'}
Define where the cell is defined relative to surface. Choose between {'CC','N'}
- If 'CC' is used, cells are active if their centers are below the surface.
- If 'N' is used, cells are active if they lie entirely below the surface.
method : {'linear', 'nearest'}
Interpolation method for locations between the xyz points.
Returns
-------
(n_cells) numpy.ndarray of bool
1D mask array of *bool* for the active cells below xyz.
Examples
--------
Here we define the active cells below a parabola. We demonstrate the differences
that appear when using the 'CC' and 'N' options for *reference_grid*.
>>> import matplotlib.pyplot as plt
>>> import numpy as np
>>> from discretize import TensorMesh
>>> from discretize.utils import active_from_xyz
Determine active cells for a given mesh and topography
>>> mesh = TensorMesh([5, 5])
>>> topo_func = lambda x: -3*(x-0.2)*(x-0.8)+.5
>>> topo_points = np.linspace(0, 1)
>>> topo_vals = topo_func(topo_points)
>>> active_cc = active_from_xyz(mesh, np.c_[topo_points, topo_vals], grid_reference='CC')
>>> active_n = active_from_xyz(mesh, np.c_[topo_points, topo_vals], grid_reference='N')
Plot visual representation
.. collapse:: Expand to show scripting for plot
>>> ax = plt.subplot(121)
>>> mesh.plot_image(active_cc, ax=ax)
>>> mesh.plot_grid(centers=True, ax=ax)
>>> ax.plot(np.linspace(0,1), topo_func(np.linspace(0,1)), color='C3')
>>> ax.set_title("CC")
>>> ax = plt.subplot(122)
>>> mesh.plot_image(active_n, ax=ax)
>>> mesh.plot_grid(nodes=True, ax=ax)
>>> ax.plot(np.linspace(0,1), topo_func(np.linspace(0,1)), color='C3')
>>> ax.set_title("N")
>>> plt.show()
"""
try:
if not mesh.is_symmetric:
raise NotImplementedError(
"Unsymmetric CylindricalMesh is not yet supported")
except AttributeError:
pass
if grid_reference not in ["N", "CC"]:
raise ValueError(
"Value of grid_reference must be 'N' (nodal) or 'CC' (cell center)"
)
dim = mesh.dim - 1
if mesh.dim == 3:
if xyz.shape[1] != 3:
raise ValueError(
"xyz locations of shape (*, 3) required for 3D mesh")
if method == "linear":
tri2D = Delaunay(xyz[:, :2])
z_interpolate = interpolate.LinearNDInterpolator(tri2D, xyz[:, 2])
else:
z_interpolate = interpolate.NearestNDInterpolator(
xyz[:, :2], xyz[:, 2])
elif mesh.dim == 2:
if xyz.shape[1] != 2:
raise ValueError(
"xyz locations of shape (*, 2) required for 2D mesh")
z_interpolate = interpolate.interp1d(xyz[:, 0],
xyz[:, 1],
bounds_error=False,
fill_value=np.nan,
kind=method)
else:
if xyz.ndim != 1:
raise ValueError(
"xyz locations of shape (*, ) required for 1D mesh")
if grid_reference == "CC":
# this should work for all 4 mesh types...
locations = mesh.cell_centers
if mesh.dim == 1:
active = np.zeros(mesh.nC, dtype="bool")
active[np.searchsorted(mesh.cell_centers_x, xyz).max():] = True
return active
elif grid_reference == "N":
try:
# try for Cyl, Tensor, and Tree operations
if mesh.dim == 3:
locations = np.vstack([
mesh.cell_centers + (np.c_[-1, 1, 1][:, None] *
mesh.h_gridded / 2.0).squeeze(),
mesh.cell_centers + (np.c_[-1, -1, 1][:, None] *
mesh.h_gridded / 2.0).squeeze(),
mesh.cell_centers +
(np.c_[1, 1, 1][:, None] * mesh.h_gridded / 2.0).squeeze(),
mesh.cell_centers + (np.c_[1, -1, 1][:, None] *
mesh.h_gridded / 2.0).squeeze(),
])
elif mesh.dim == 2:
locations = np.vstack([
mesh.cell_centers +
(np.c_[-1, 1][:, None] * mesh.h_gridded / 2.0).squeeze(),
mesh.cell_centers +
(np.c_[1, 1][:, None] * mesh.h_gridded / 2.0).squeeze(),
])
else:
active = np.zeros(mesh.nC, dtype="bool")
active[np.searchsorted(mesh.nodes_x, xyz).max():] = True
return active
except AttributeError:
# Try for Curvilinear Mesh
gridN = mesh.gridN.reshape((*mesh.vnN, mesh.dim), order="F")
if mesh.dim == 3:
locations = np.vstack([
gridN[:-1, 1:, 1:].reshape((-1, mesh.dim), order="F"),
gridN[:-1, :-1, 1:].reshape((-1, mesh.dim), order="F"),
gridN[1:, 1:, 1:].reshape((-1, mesh.dim), order="F"),
gridN[1:, :-1, 1:].reshape((-1, mesh.dim), order="F"),
])
elif mesh.dim == 2:
locations = np.vstack([
gridN[:-1, 1:].reshape((-1, mesh.dim), order="F"),
gridN[1:, 1:].reshape((-1, mesh.dim), order="F"),
])
# Interpolate z values on CC or N
z_xyz = z_interpolate(locations[:, :-1]).squeeze()
# Apply nearest neighbour if in extrapolation
ind_nan = np.isnan(z_xyz)
if any(ind_nan):
tree = cKDTree(xyz)
_, ind = tree.query(locations[ind_nan, :])
z_xyz[ind_nan] = xyz[ind, dim]
# Create an active bool of all True
active = np.all((locations[:, dim] < z_xyz).reshape((mesh.nC, -1),
order="F"),
axis=1)
return active.ravel()
def ax_scalp(v,
channels,
ax=None,
annotate=True,
vmin=None,
vmax=None,
cmap=cm.coolwarm,
scalp_line_width=1,
scalp_line_style='solid',
chan_pos_list=CHANNEL_9_APPROX,
interpolation='bilinear',
fontsize=8):
"""
:param v: 输入通道对应的相关性值
:param channels: 通道的坐标位置
:param ax: 最后的存到ax这个框架里
:param annotate:
:param vmin:
:param vmax:
:param cmap:
:param scalp_line_width:头皮轮廓的线宽
:param scalp_line_style:头皮轮廓的线的形状
:param chan_pos_list:整个电极通道的位置
:param interpolation:双线性插值
:param fontsize:字体大小
:return: 一张头皮图
"""
if ax is None:
ax = plt.gca()
assert len(v) == len(channels), "Should be as many values as channels"
assert interpolation == 'bilinear' or interpolation == 'nearest'
if vmin is None:
# added by me ([email protected])
assert vmax is None
vmin, vmax = -np.max(np.abs(v)), np.max(np.abs(v))
# what if we have an unknown channel?得到对应电极通道在列表中正确的位置
points = [get_channelpos(c, chan_pos_list) for c in channels]
for c in channels:
assert get_channelpos(
c, chan_pos_list) is not None, ("Expect " + c +
" to exist in positions")
z = [v[i] for i in range(len(points))]
# calculate the interpolation
x = [i[0] for i in points]
y = [i[1] for i in points]
# interpolate the in-between values插值填充
xx = np.linspace(min(x), max(x), 1000)
yy = np.linspace(min(y), max(y), 1000)
if interpolation == 'bilinear':
xx_grid, yy_grid = np.meshgrid(xx, yy) #生成网格点坐标矩阵
#对输入数据进行三角测量，并在每个执行线性重心插值的三角形上构造插值
f = interpolate.LinearNDInterpolator(list(zip(x, y)), z) #
zz = f(xx_grid, yy_grid)
else:
assert interpolation == 'nearest'
f = interpolate.NearestNDInterpolator(list(zip(x, y)), z)
assert len(xx) == len(yy)
zz = np.ones((len(xx), len(yy)))
for i_x in xrange(len(xx)):
for i_y in xrange(len(yy)):
# somehow this is correct. don't know why :(
zz[i_y, i_x] = f(xx[i_x], yy[i_y])
# zz[i_x,i_y] = f(xx[i_x], yy[i_y])
assert not np.any(np.isnan(zz))
# plot map开始画头皮图
image = ax.imshow(zz,
vmin=vmin,
vmax=vmax,
cmap=cmap,
extent=[min(x), max(x), min(y),
max(y)],
origin='lower',
interpolation=interpolation)
if scalp_line_width > 0:
# paint the head
ax.add_artist(
plt.Circle((0, 0),
1,
linestyle=scalp_line_style,
linewidth=scalp_line_width,
fill=False))
# add a nose
ax.plot([-0.1, 0, 0.1], [1, 1.1, 1],
color='black',
linewidth=scalp_line_width,
linestyle=scalp_line_style)
# add ears
_add_ears(ax, scalp_line_width, scalp_line_style)
# add markers at channels positions
# set the axes limits, so the figure is centered on the scalp
ax.set_ylim([-1.05, 1.15])
ax.set_xlim([-1.15, 1.15])
# hide the frame and ticks
ax.set_frame_on(False)
ax.set_xticks([])
ax.set_yticks([])
# draw the channel names
if annotate:
for i in zip(channels, list(zip(x, y))):
ax.annotate(" " + i[0],
i[1],
horizontalalignment="center",
verticalalignment='center',
fontsize=fontsize)
ax.set_aspect(1)
plt.show()
return image
for m in trange(m_+1,desc='Day'):
tau = np.busday_count(t_m[m], t_end)/252
if tau < tau_implvol[0]:
tau = tau_implvol[0]
for j in range(j_):
# compute shadow yield
x_y = ytm_shadowrates(np.array([y]))
x_y = np.atleast_1d(x_y)
# compute call option value
v_call_thor[j, m] = \
call_option_value(x_proj[j, m, 0], x_y, tau,
x_proj[j, m, 1:], m_moneyness, tau_implvol,
k_strk, t_end, t_m[m])
# compute log-implied volatility at the moneyness
log_sigma_atm[j, m] = \
interpolate.LinearNDInterpolator(points,
x_proj[j, m, 1:])(*np.r_[tau, 0])
# -
# ## [Step 2](https://www.arpm.co/lab/redirect.php?permalink=s_pricing_calloption-implementation-step02): Scenario-probability expectations and standard deviations
# +
mu_v = np.zeros(m_+1)
sig_v = np.zeros(m_+1)
for m in range(len(t_m)):
mu_v[m], sig1 = meancov_sp(v_call_thor[:, m].reshape(-1, 1))
sig_v[m] = np.sqrt(sig1)
# -
# ## [Step 3](https://www.arpm.co/lab/redirect.php?permalink=s_pricing_calloption-implementation-step03): Save databases
def interpf(x, z, method=None, extrap=False, arrays=False):
'''interpolate to find f, where z = f(*x)
Input:
x: an array of shape (npts, dim) or (npts,)
z: an array of shape (npts,)
method: string for kind of interpolator
extrap: if True, extrapolate a bounding box (can reduce # of nans)
arrays: if True, z = f(*x) is a numpy array; otherwise don't use arrays
Output:
interpolated function f, where z = f(*x)
NOTE:
if scipy is not installed, will use np.interp for 1D (non-rbf),
or mystic's rbf otherwise. default method is 'nearest' for
1D and 'linear' otherwise. method can be one of ('rbf','linear','cubic',
'nearest','inverse','gaussian','multiquadric','quintic','thin_plate').
NOTE:
if extrap is True, extrapolate using interpf with method='thin_plate'
(or 'rbf' if scipy is not found). Alternately, any one of ('rbf',
'linear','cubic','nearest','inverse','gaussian','multiquadric',
'quintic','thin_plate') can be used. If extrap is a cost function
z = f(x), then directly use it in the extrapolation.
''' #XXX: return f(*x) or f(x)?
if arrays:
_f, _fx = _to_array, _array
else:
_f, _fx = _to_nonarray, _nonarray
if len(x) > len(z): # if len(x) < len(z), points are dropped
x = x[:len(z)] # drop points
x, z = extrapolate(x, z, method=extrap)
x, z = _unique(x, z, sort=True)
# avoid nan as first value #XXX: better choice than 'nearest'?
if method is None: method = 'nearest' if x.ndim == 1 else 'linear'
methods = dict(rbf=0, linear=1, nearest=2, cubic=3)
functions = {0: 'multiquadric', 1: 'linear', 2: 'nearest', 3: 'cubic'}
# also: ['thin_plate','inverse','gaussian','quintic']
kind = methods[method] if method in methods else None
function = functions[kind] if (not kind is None) else method
if kind is None: kind = 0 #XXX: or raise KeyError ?
try:
import scipy.interpolate as si
except ImportError:
if not kind == 0: # non-rbf
if x.ndim == 1: # is 1D, so use np.interp
import numpy as np
return lambda xn: _fx(np.interp(xn, xp=x, fp=z))
kind = 0 # otherwise, utilize mystic's rbf
si = _rbf
if kind == 0: # 'rbf'
import numpy as np
return _f(si.Rbf(*np.vstack((x.T, z)), function=function, smooth=0))
elif x.ndim == 1:
return _f(
si.interp1d(x,
z,
fill_value='extrapolate',
bounds_error=False,
kind=method))
elif kind == 1: # 'linear'
return _f(si.LinearNDInterpolator(x, z, rescale=False))
elif kind == 2: # 'nearest'
return _f(si.NearestNDInterpolator(x, z, rescale=False))
#elif x.ndim == 1: # 'cubic'
# return lambda xn: _fx(si.spline(x, z, xn))
return _f(si.CloughTocher2DInterpolator(x, z, rescale=False))
def __init__(self,
x_values,
y_values,
z_values,
classification,
scalar_ans=None,
x_vector_ans=None,
y_vector_ans=None,
z_vector_ans=None):
"""A class representing either a scalar or vector table.
If a lookup table is to be used instead of a function for the model
observing method, it is required to standardize the information given
by the user's lookup table, as is the purpose of this class.
Arguments
---------
x_values : array_like
The values of the x points for the table. Array must be parallel
with y_values and z_values along with the scalar/vector answer.
y_values : array_like
The values of the x points for the table. Array must be parallel
with x_values and z_values along with the scalar/vector answer.
z_values : array_like
The values of the x points for the table. Array must be parallel
with x_values and y_values along with the scalar/vector answer.
classification : string
The classification of this table, either as a scalar lookup table
or a vector lookup table. Should be one of:
- 'scalar' A scalar based lookup table.
- 'vector' A vector based lookup table.
scalar_ans : array_like, {for | classification == 'scalar'}
The scalar answers to the (x,y,z) point given by the input values.
Must be parallel with x_values, y_values, and z_values. Ignored if
classification == 'vector'.
x_vector_ans : array_like, {for | classification == 'vector'}
The x component of the answer vector that exists at the point
(x,y,z) given by the input values. Must be parallel with x_values,
y_values, and z_values along with other components. Ignored if
classification == 'scalar'
y_vector_ans : array_like, {for | classification == 'vector'}
The y component of the answer vector that exists at the point
(x,y,z) given by the input values. Must be parallel with x_values,
y_values, and z_values along with other components. Ignored if
classification == 'scalar'
z_vector_ans : array_like, {for | classification == 'vector'}
The z component of the answer vector that exists at the point
(x,y,z) given by the input values. Must be parallel with x_values,
y_values, and z_values along with other components. Ignored if
classification == 'scalar'
"""
# Type check.
x_values = Robust.valid.validate_float_array(x_values)
y_values = Robust.valid.validate_float_array(y_values)
z_values = Robust.valid.validate_float_array(z_values)
classification = Robust.valid.validate_string(classification).lower()
# Decide on the type before type checking.
if (classification == 'scalar'):
if (scalar_ans is not None):
scalar_ans = Robust.valid.validate_float_array(scalar_ans)
else:
raise TypeError('Scalar answer array must be provided if '
'table classification is set to scalar.'
' --Kyubey')
elif (classification == 'vector'):
if (x_vector_ans is not None):
x_vector_ans = Robust.valid.validate_float_array(x_vector_ans)
else:
raise TypeError('The x component of the vector answer array '
'must be provided if table classification is '
'set to vector.'
' --Kyubey')
if (y_vector_ans is not None):
y_vector_ans = Robust.valid.validate_float_array(y_vector_ans)
else:
raise TypeError('The y component of the vector answer array '
'must be provided if table classification is '
'set to vector.'
' --Kyubey')
if (z_vector_ans is not None):
z_vector_ans = Robust.valid.validate_float_array(z_vector_ans)
else:
raise TypeError('The z component of the vector answer array '
'must be provided if table classification is '
'set to vector.'
' --Kyubey')
else:
raise Robust.InputError(
'Table classification must be one of the '
'following: \n'
'scalar, vector \n'
'It is currently: < {table_cls} >'
' --Kyubey'.format(table_cls=classification))
# Precompute the Delaunay triangulation, this is done under the
# assumption that the table should not be changed after data is
# put into it.
# pylint: disable=maybe-no-member
try:
Delanuay_tri = sp_spt.Delaunay(
np.array([x_values, y_values, z_values]).T)
except (TypeError, ValueError):
raise
except Exception:
# If there is a Qhull error, we don't want to deal with it. Does
# not currently know how to specify the QHull error on its own.
Delanuay_tri = None
# pylint: enable=maybe-no-member
# Attempt to make the linear interpolators.
if (Delanuay_tri is not None):
if (classification == 'scalar'):
linear_interp_scalar = \
sp_inter.LinearNDInterpolator(Delanuay_tri,
scalar_ans,fill_value=0)
# Just for safety reasons.
linear_interp_x_axis = None
linear_interp_y_axis = None
linear_interp_z_axis = None
elif (classification == 'vector'):
linear_interp_x_axis = \
sp_inter.LinearNDInterpolator(Delanuay_tri,
x_vector_ans,
fill_value=0)
linear_interp_y_axis = \
sp_inter.LinearNDInterpolator(Delanuay_tri,
y_vector_ans,
fill_value=0)
linear_interp_z_axis = \
sp_inter.LinearNDInterpolator(Delanuay_tri,
z_vector_ans,
fill_value=0)
# For safety reasons.
linear_interp_scalar = None
else:
linear_interp_scalar = None
linear_interp_x_axis = None
linear_interp_y_axis = None
linear_interp_z_axis = None
# Assign variables. Depending on the actual
self.x_values = x_values
self.y_values = y_values
self.z_values = z_values
self.classification = classification
self.scalar_ans = scalar_ans
self.x_vector_ans = x_vector_ans
self.y_vector_ans = y_vector_ans
self.z_vector_ans = z_vector_ans
self._Delanuay_triangulation = Delanuay_tri
self._linear_interp_scalar = linear_interp_scalar
self._linear_interp_x_axis = linear_interp_x_axis
self._linear_interp_y_axis = linear_interp_y_axis
self._linear_interp_z_axis = linear_interp_z_axis