Python vcmq.N类代码示例

def gene_bathy(xc, yc, xr, yr, n=500, amp=30.0):
    noise = N.random.random(n)
    a = N.random.random(n) * N.pi * 2
    r = N.random.random(n)
    x = xc + xr * r * N.cos(a)
    y = yc + yr * r * N.sin(a)
    return N.asarray([x, y, N.exp(-(x - xc) ** 2 / xr ** 2 - (y - yc) ** 2 / yr ** 2) * amp + noise])

def plot(xx, yy, target, label, figfiles, figfile, lon=None, lat=None, show=False):
    xs, ys, mask = coord2slice(target, lon=lon, lat=lat)
    P.figure(figsize=(6, 3.5))
    P.title('Target=%(label)s / select: lon=%(lon)s, lat=%(lat)s'%locals())
    add_grid((xx, yy))
    xx = xx.asma()
    yy = yy.asma()
    if isinstance(lon, tuple): 
        P.axvline(lon[0], color='m', ls='--', lw=2)
        P.axvline(lon[1], color='m', ls='--', lw=2)
    elif isinstance(lon, slice):
        i, j, k = lon.indices(xx.shape[1])
        P.plot(xx[:, i], yy[:, i], 'c--', lw=2)
        P.plot(xx[:, j-1], yy[:, j-1], 'c--', lw=2)
    if isinstance(lat, tuple): 
        P.axhline(lat[0], color='m', ls='--', lw=2)
        P.axhline(lat[1], color='m', ls='--', lw=2)
    elif isinstance(lat, slice):
        i, j, k = lat.indices(yy.shape[0])
        P.plot(xx[i], yy[i], 'c--', lw=2)
        P.plot(xx[j-1], yy[j-1], 'c--', lw=2)
    P.xticks(N.arange(xx.min()-1, xx.max()+1))
    P.yticks(N.arange(yy.min()-1, yy.max()+1))
    xxi, yyi = xx, yy
    xx = xx[ys, xs]
    yy = yy[ys, xs]
#    mask = mask[ys, xs]
    xxb, yyb = meshbounds(xx, yy)
    P.pcolor(xxb, yyb, mask, shading='faceted')
    P.scatter(xx.ravel(), yy.ravel(), c=(0, 1, 0))
    P.grid('on')
    P.axis('image')
    P.tight_layout()
    i = len(figfiles)
    savefig = figfile%i
    if os.path.exists(savefig): os.remove(savefig)
    P.savefig(savefig)
    figfiles.append(savefig)
    if show: P.show()
    else: P.close()

"""Compare CDAT regridding speed with rectangular and rectangular grids"""

config = {
   'esmf':['linear', 'patch', 'conserv'], 
   'libcf':['linear'], 
}

# Imports
from vcmq import MV2, create_grid2d, code_file_name, os, CDATRegridder, N, set_grid, psinfo
from vacumm.misc.grid import rotate_grid
from time import time

# Input
nx = ny = 300
vari = MV2.array(N.arange(nx*ny*1.).reshape(ny, nx))
gridi = create_grid2d(vari.getAxis(1)[:]*50/nx,  vari.getAxis(0)[:]*50/nx)
set_grid(vari, gridi)

# Output grid
gridor = create_grid2d(vari.getAxis(1)[:]*0.09*50/nx,  
    vari.getAxis(0)[:]*0.09*50/nx)
gridoc = rotate_grid(gridi, 30)

# Log
logfile = code_file_name(ext='log')
if os.path.exists(logfile): os.remove(logfile)
f = open(logfile, 'w')
print >>f, 'NY=%(ny)i, NX=%(nx)i'%locals()

# Loop on methods
for tool, methods in config.items():

# Lecture de la température
f =cdms2.open(data_sample('mars3d.tsigyx.nc'))
data_in = f('TEMP', time=slice(0,2)) # T-YX (- = level)


# Détermination des profondeurs d'après les sigma
sigma_converter = NcSigma.factory(f)                # détection auto et initialisation du convertisseur
# -> VERIFIER QUE sigma_class EST BIEN SigmaGeneralized
depths_in = sigma_converter.sigma_to_depths(selector=dict(time=slice(0,2))).filled()          # lecture eta, etc + conversion
# (Equivalent à depths_in = sigma_converter(selector=dict(time=slice(0,2))).filled())
f.close()


# Creation de l'axe des profondeurs cibles
depths = N.array([0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,
    22,24,26,28,30,32,34,36,38,40,45,50,55,60,65,70,75,80,85,90,95,100,120,140,160])
depths = -depths[::-1] # croissantes et négatives
depth_out = create_depth(depths)


# Interpolation
data_out = regrid1d(data_in, depth_out, axi=depths_in, axis=1, method='linear', extrap=1)


# Plot
kw = dict(vmin=10, vmax=14, xhide='auto', add_grid=True, ymax=0, fill='contourf')   # FILL=PCOLOR ?
section2(data_in[0, :, 10], yaxis=depths_in[0, :, 10], subplot=211, title='Sigma', show=False, **kw)
s = section2(data_out[0, :, 10], subplot=212, title='Z', show=False, savefig=__file__, **kw)

"""Test the :func:`~vacumm.misc.grid.regridding.regrid2d` function"""
from vcmq import P, N, MV2, code_file_name, os, add_grid, rotate_grid, set_grid, \
    create_grid, rc, rcdefaults, plot2d
from vacumm.misc.grid.regridding import regrid2d


# Input grid and data
nxi = 20
nyi = 15
# - rect
xi = N.arange(nxi*1.)
yi = N.arange(nyi*1.)
gridri = create_grid(xi, yi)
xxri, yyri = N.meshgrid(xi, yi)
zzri = N.ma.array(yyri)
zzri[int(nyi*0.3):int(nyi*0.6), int(nxi*0.3):int(nxi*0.6)] = N.ma.masked
varri = MV2.asarray(zzri)
set_grid(varri, gridri)
# - curv
gridci = rotate_grid(gridri, 30)
xxci = gridci.getLongitude().getValue()
yyci = gridci.getLatitude().getValue()
zzci = N.ma.array(yyci)
zzci[int(nyi*0.3):int(nyi*0.6), int(nxi*0.3):int(nxi*0.6)] = N.ma.masked
varci = MV2.asarray(zzci)
set_grid(varci, gridci)

# Output positions
nxo = 25
nyo = 18
# - rect

"""Test the fortran function :f:func:`extrap1d`"""
from vcmq import N, P,meshcells, minmax
from vacumm.misc.grid._interp_ import extrap1d


mv = 999.
vari = N.zeros((5,5))+mv
vari[1,2:4] = [2,3]
vari[2,1] = 1
vari[2,3] = 3
vari[3,3:] = [3,4]
vari[4,:2] = [0,1]
vari = N.asfortranarray(vari)

varo0 = extrap1d(vari, mv, 0)
varop1 = extrap1d(vari, mv, 1)
varom1 = extrap1d(vari, mv, -1)
varop2 = extrap1d(vari, mv, 2)

result = [
    ('AssertTrue', N.allclose(vari, varo0)), 
    ('AssertTrue', N.allclose(varop1[:, -1], [999., 3., 3., 4., 1.])), 
    ('AssertTrue', N.allclose(varom1[:, 0], [999., 2., 1., 3., 0.])), 
    ('AssertTrue', N.allclose(varop1[:, -1], varop2[:, -1])), 
    ('AssertTrue', N.allclose(varom1[:, 0], varop2[:, 0])), 
    ]
    

"""Test the fortran function :f:func:`bilin`"""
from vcmq import N, P, meshcells, minmax, code_file_name, os
from vacumm.misc.grid._interp_ import dstwgt


nxi = 15
nyi = 10
mv = 1.e20
u, v = N.mgrid[-3:3:nyi*1j, -3:3:nxi*1j]-2
vari = N.ma.asarray(u**2+v**2)
vari.set_fill_value(mv)
xi = N.arange(nxi)
yi = N.arange(nyi)
vari[int(nyi*0.4):int(nyi*0.4)+3, int(nxi*0.4):int(nxi*0.4)+2] = N.ma.masked
xxib, yyib  = meshcells(xi, yi)

nxo = 40
nyo = 25
xo = N.linspace(int(nxi*0.2),int(nxi*1.2),nxo)
yo = N.linspace(int(-nyi*0.2),int(nyi*0.8),nyo)
xxob, yyob  = meshcells(xo, yo)

vari.shape = (1, )+vari.shape
varo = N.ma.masked_values(dstwgt(vari.filled(), xi, yi, xo, yo, mv, 0), mv)

kw = dict(vmin=vari.min(), vmax=vari.max())
axlims = [min(xi.min(), xo.min()), max(xi.max(), xo.max()), 
    min(yi.min(), yo.min()), max(yi.max(), yo.max())]
P.figure(figsize=(8, 4))
P.subplot(211)
P.pcolor(xxib, yyib, vari[0], **kw)

sp = f("speed")
spe = f("speed_error")
f.close()

# Create hourly time axis
taxi = sp.getTime()
taxi.toRelativeTime("hours since 2000")
ctimesi = taxi.asComponentTime()
ct0 = round_date(ctimesi[0], "hour")
ct1 = round_date(ctimesi[-1], "hour")
taxo = create_time(lindates(ct0, ct1, 1, "hour"), taxi.units)

# Lag error
# - estimation
els = []
lags = N.arange(1, 6)
for lag in lags:
    els.append(N.sqrt(((sp[lag:] - sp[:-lag]) ** 2).mean()))
els = N.array(els)
a, b, _, _, _ = linregress(lags, els)
# - plot
P.figure(figsize=(6, 6))
P.subplot(211)
P.plot(lags, els, "o")
P.plot([0, lags[-1]], [b, a * lags[-1] + b], "g")
P.axhline(b, color="0.8", ls="--")
P.ylim(ymin=0)
P.xlabel("Lag [hour]")
P.ylabel("Error [m s-1]")
add_key(1)
P.title("Linear lag error model")

lon1d = create_lon((0, 10))
result.append(('AssertEqual', (coord2slice(lon1d, lon=(2.5, 4., 'cc')), slice(3, 5, 1))))
result.append(('AssertEqual', (coord2slice(lon1d, lon=(2.5, 4., 'ccb')), slice(2, 5, 1))))
result.append(('AssertEqual', (coord2slice(lon1d, lon=slice(3, 6)), slice(3, 6, None))))
result.append(('AssertEqual', (coord2slice(lon1d, lat=(6, 8)), slice(0, 10, 1))))
result.append(('AssertEqual', (coord2slice(lon1d, lon=(60, 70)), None)))

# Rect grid
grid = create_grid((0, 10.), (20, 30.))
result.append(('AssertEqual', (coord2slice(grid, lon=(0., 3.5), lat=slice(3, 5)), 
    (slice(0, 4, 1), slice(3, 5, None), None))))
result.append(('AssertEqual', (coord2slice(grid, lat=(21,21, 'ccb')),
    (slice(0, 10, 1), slice(1, 2, 1), None))))

# 2D axis
lon2d = N.empty((10, 10.))
for i in xrange(10): 
    lon2d[i] = lon1d[:]+i
lat2d = N.resize((N.arange(10)+20), (10, 10)).T
lon2d, lat2d = create_axes2d(lon2d, lat2d)
kw = dict(show=False)
plot(lon2d, lat2d, lon2d, 'lon2d', figfiles, figfile, lon=(2, 4), **kw)
plot(lon2d, lat2d, lon2d, 'lon2d', figfiles, figfile, lon=(2, 4), lat=slice(0, 2), **kw)
plot(lon2d, lat2d, lat2d,  'lat2d', figfiles, figfile, lat=(22, 26.6,'ccb'), **kw)

# Curv grid
grid = create_grid(lon2d, lat2d)
plot(lon2d, lat2d, grid, 'grid', figfiles, figfile, lon=(8, 11, 'cc'), lat=(21.9, 26., 'cc'), **kw)
plot(lon2d, lat2d, grid, 'grid', figfiles, figfile, lon=slice(2, 5), lat=(23.4, 23.6, 'ccb'), **kw)
res = coord2slice(grid,lon=(8,8,'ccb'),lat=(24,24,'ccb'))
result.append(('AssertEqual', (res[:2], (slice(3, 6, 1), slice(4, 5, 1)))))

"""Test fortran function :f:func:`dstwgt2dto1dc`"""

from vcmq import P, N, code_file_name, os, add_grid, rotate_grid
from vacumm.misc.grid._interp_ import dstwgt2dto1dc


# Input grid and data
nxy = 15
xi = N.arange(nxy*1.)
yi = N.arange(nxy*1.)
gridi = rotate_grid((xi, yi), 30)
xxi = gridi.getLongitude().getValue()
yyi = gridi.getLatitude().getValue()
zzi = N.ma.array(yyi)
zzi[int(nxy*0.3):int(nxy*0.8), int(nxy*0.3):int(nxy*0.8)] = N.ma.masked
zzi.shape = 1, nxy, nxy

# Output positions
no = 1000
xo = N.random.uniform(-nxy/4., nxy+nxy/4., no)
yo = N.random.uniform(-nxy/4., nxy+nxy/4., no)

# Interpolate
mv = zzi.get_fill_value()
zo = dstwgt2dto1dc(xxi,yyi,zzi.filled(mv),xo,yo,mv)
zo = N.ma.masked_values(zo, mv) 

# Plot
kw = dict(vmin=zzi.min(), vmax=zzi.max())
P.figure(figsize=(6, 6))
P.subplot(111, aspect=1)

"""Test the fortran function :f:func:`cellerr1d`, :f:func:`cellerr1dx` et :f:func:`cellerr1dxx`"""
from vcmq import cdms2, data_sample, N, P, meshcells, minmax, code_file_name, os
from vacumm.misc.grid._interp_ import cellerr1d, cellerr1dx, cellerr1dxx

# Academic example pure 1d
# - data
yi = N.array([-2., 2.,  6. , 12., 15.])
vari = N.ma.array(yi)
vari[-1] = N.ma.masked
errm = N.ma.array([1, 2, 1., 2., 3.])
yo = N.array([5.,  15])
mv = 1e20
vari = vari.reshape(1, -1).filled(mv)
errm = errm.reshape(1, -1).filled(mv)
errl = N.ones(1)
# - interp
varo, erro = cellerr1d(vari, yi, yo, mv, errm, errl)
# - truth
errot =  N.array([
    1/N.sqrt(
        1/(errm[0, 1]**2+N.abs(yo[0]-yi[1])*errl[0]) +
        1/(errm[0, 2]**2+N.abs(yo[0]-yi[2])*errl[0])),
    N.sqrt(errm[0, 3]**2+N.abs(yo[1]-yi[3])*errl[0])
    ])
varot = N.array([
    vari[0, 1]/(errm[0, 1]**2+N.abs(yo[0]-yi[1])*errl[0]) +
    vari[0, 2]/(errm[0, 2]**2+N.abs(yo[0]-yi[2])*errl[0]),
    vari[0, 3]/(errm[0, 3]**2+N.abs(yo[1]-yi[3])*errl[0])
    ])*errot**2
# - check
result = [(N.testing.assert_allclose, [erro.ravel(), errot]),

"""Test :meth:`vacumm.data.misc.arakawa.CGrid.interp`"""

from vcmq import MV2, N, create_grid, create_dep, set_grid, map2, \
    code_file_name, CGrid, minmax, curve2, add_grid

# Initial variable
grid = create_grid(N.arange(-7, 0.), N.arange(43, 50.))
dep = create_dep([-5000, -3000, -2000, -1000, -500, -300, -200, -100.])
var = {}
var['t'] = MV2.reshape(N.arange(grid.size()*len(dep))*1., (len(dep), )+grid.shape)
set_grid(var['t'], grid)
var['t'].setAxis(0, dep)


# Arakawa manager
ag = CGrid()

# Interpolations
for p in 'u', 'v', 'f', 'w':
    var[p] = ag.interp(var['t'], 't', p, mode='extrap')

# Surface plots
vmin, vmax = minmax(*[var[p][-1] for p in ['u', 'v', 'f']])
kw = dict(show=False, res=None, vmin=vmin, vmax=vmax, colorbar=False, grid=False, cmap='jet')
m = map2(var['t'][-1], fill='pcolor', 
    title='Interpolations on an Arakawa C grid: T->U/V/F', **kw)
add_grid(var['t'], linestyle='-')
kw.update(fill='scatter', contour=False, fill_s=60)
markers = dict(u='>', v='^', f='D', t='o')
for p in 't', 'u', 'v', 'f':
    m = map2(var[p][-1], fill_marker=markers[p], shadow=True, zorder=100, **kw)

"""Test the fortran function :f:func:`curv2rel_single`"""
from vcmq import N, P, code_file_name, P, os, meshbounds
from vacumm.misc.grid._interp_ import curv2rel_single


# Input grid
x0, y0 = 0., 2.
nxi = 3
nyi = 3
dxi = (2., 2.)
dyi = (-2., 2.)
xxi = N.zeros((nyi, nxi))
xxi[0] = x0 + N.arange(nxi) * dxi[0]
for j in range(1, nyi):
    xxi[j] = xxi[j-1] + dyi[0]
yyi = N.zeros((nyi, nxi))
yyi[:, 0] = y0 + N.arange(nyi) * dyi[1]
for j in range(1, nyi):
    yyi[:, j] = yyi[:, j-1] + dxi[1]
xxbi, yybi = meshbounds(xxi, yyi)
relpos2index = lambda fi, fj, nyi: fj * nyi + fi
ii, jj = N.meshgrid(N.arange(nxi)+.5, N.arange(nyi)+.5)
iib, jjb = meshbounds(ii, jj)
zzi = relpos2index(ii, jj, nyi)
zzbi = relpos2index(iib, jjb, nyi)

# Input random points
N.random.seed(0)
np = 100
xxo = N.random.random(np)*(xxbi.max()-xxbi.min()) + xxbi.min()
yyo = N.random.random(np)*(yybi.max()-yybi.min()) + yybi.min()

"""Test :func:`~vacumm.misc.plot.section2` with a Z- variable"""

# Imports
from vcmq import N, MV2, cdms2, create_dep, rc, section2, code_file_name, os

# Init data with z 1D
nz = 8
nd = 10
var = N.dot(N.hanning(nz).reshape(nz, 1), N.hanning(nd).reshape(1, nd))
var = MV2.array(var)
d = cdms2.createAxis(N.arange(nd))
d.units='km'
d.long_name='Distance'
z1d = create_dep((-nz+1, 1.))
var.setAxis(0, z1d)
var.setAxis(1, d)
z2d = N.resize(z1d[:].reshape(1, nz), (nd, nz)).T
z2d *= N.arange(1., nd+1)/nd

# Plot with z 1D
rc('font', size=8)
kw = dict(show=False, bgcolor='0.5')
section2(var, subplot=211, **kw)

# Plot with z 2D
figfile = code_file_name(ext='png')
if os.path.exists(figfile): os.remove(figfile)
section2(var, yaxis=z2d, subplot=212, savefig=figfile, close=True, **kw)

# Result
result = dict(files=figfile)

"""Test the fortran function :f:func:`curv2rect`"""
from vcmq import N, P, code_file_name, P, os
from vacumm.misc.grid._interp_ import curv2rect

# Input
x1, y1 = 0., 0.
x2, y2 = 3., 1.
x3, y3 = 2., 4.
x4, y4 = -1., 2.

# Format and convert
xx, yy = N.meshgrid(N.arange(-2, 4, 0.25), N.arange(-1, 5, 0.25))
nxy = xx.shape
xx.shape = -1
yy.shape = -1
pp, qq = [], []
for x, y in zip(xx, yy):
    p, q = curv2rect(x1,x2,x3,x4,y1,y2,y3,y4,x,y)
    pp.append(p)
    qq.append(q)
pp = N.array(pp)
qq = N.array(qq)

# Plot
xp = [x1, x2, x3, x4, x1]
yp = [y1, y2, y3, y4, y1]
P.subplot(211)
levels = N.array([-10, 0, 1, 10.])
o = P.contourf(xx.reshape(nxy), yy.reshape(nxy), pp.reshape(nxy), levels=levels)
P.colorbar(o)
P.plot(xp, yp, 'k')