#!/usr/bin/env python
from __future__ import print_function
import numpy as np
import numexpr as ne
import viscid
from viscid import field
from viscid import coordinate
# ne.set_num_threads(1) # for performance checking
def _normalize_scalar_dtype(s, arrs):
# cast python scalars to an appropriate numpy dtype
if isinstance(s, (int, float, complex)):
ndarrs = [_a for _a in arrs if hasattr(_a, 'dtype')]
flt_arrs = [_a for _a in ndarrs if _a.dtype.kind in 'fc']
int_arrs = [_a for _a in ndarrs if _a.dtype.kind in 'i']
if flt_arrs and isinstance(s, (int, float, complex)):
s = np.asarray(s).astype(np.common_type(*flt_arrs))
elif int_arrs and isinstance(s, (int, )):
s = np.asarray(s).astype(max([_a.dtype for _a in int_arrs]))
return s
[docs]class VneCalc(object):
expr = None
local_dict = None
slice_dict = None
ret_type = None
ret_context = None
def __init__(self, expr, local_dict,
ret_type=None, ret_context=None):
self.expr = expr
self.local_dict = local_dict
# self.slice_dict = slice_dict
self.ret_type = ret_type
self.ret_context = ret_context
def __call__(self, *args):
ldict = {}
for name, val in self.local_dict.items():
if isinstance(val, int):
ldict[name] = _normalize_scalar_dtype(args[val], args)
else:
ldict[name] = val
# apply slices if given
# if self.slice_dict is not None:
# for name, slc in self.slice_dict.items():
# ldict[name] = ldict[name][slc].data
# print(ldict)
result = ne.evaluate(self.expr, local_dict=ldict)
return args[0].wrap(result, context=self.ret_context,
fldtype=self.ret_type)
neg = VneCalc("-a", {'a': 0})
add = VneCalc("a + b", {'a': 0, 'b': 1})
diff = VneCalc("a - b", {'a': 0, 'b': 1})
mul = VneCalc("a * b", {'a': 0, 'b': 1})
relative_diff = VneCalc("(a - b) / a", {'a': 0, 'b': 1})
abs_diff = VneCalc("abs(a - b)", {'a': 0, 'b': 1})
abs_val = VneCalc("abs(a)", {'a': 0})
[docs]def scale(a, fld):
a = np.asarray(a, dtype=fld.dtype)
b = fld.data # pylint: disable=unused-variable
return fld.wrap(ne.evaluate("a * b"))
[docs]def axpby(a, x, b, y):
a = _normalize_scalar_dtype(a, [x, y])
b = _normalize_scalar_dtype(b, [x, y])
return x.wrap(ne.evaluate("a * x + b * y"))
[docs]def abs_max(fld):
a = fld.data # pylint: disable=W0612
return np.max(ne.evaluate("abs(a)"))
[docs]def abs_min(fld):
a = fld.data # pylint: disable=W0612
return np.min(ne.evaluate("abs(a)"))
# def scalar_mul(s, fld):
# sarr = np.array([s], dtype=fld.dtype) # pylint: disable=W0612
# fldarr = fld.data # pylint: disable=W0612
# result = ne.evaluate("sarr * fldarr")
# return fld.wrap(result)
[docs]def magnitude(fld):
vx, vy, vz = fld.component_views() # pylint: disable=W0612
mag = ne.evaluate("sqrt((vx**2) + (vy**2) + (vz**2))")
return fld.wrap(mag, fldtype="Scalar")
[docs]def dot(fld_a, fld_b):
""" 3d dot product of two vector fields """
ax, ay, az = fld_a.component_views() # pylint: disable=W0612
bx, by, bz = fld_b.component_views() # pylint: disable=W0612
prod = ne.evaluate("(ax * bx) + (ay * by) + (az * bz)")
return fld_a.wrap(prod, fldtype="Scalar")
[docs]def cross(fld_a, fld_b):
""" cross product of two vector fields """
ax, ay, az = fld_a.component_views() # pylint: disable=W0612
bx, by, bz = fld_b.component_views() # pylint: disable=W0612
prodx = ne.evaluate("ay * bz - az * by")
prody = ne.evaluate("-ax * bz + az * bx")
prodz = ne.evaluate("ax * by - ay * bx")
return fld_a.wrap([prodx, prody, prodz])
[docs]def project(fld_a, fld_b):
""" project a along b (a dot b / magnitude(b)) """
# ax, ay, az = fld_a.component_views() # pylint: disable=W0612
# bx, by, bz = fld_b.component_views() # pylint: disable=W0612
# prod = ne.evaluate("(ax * bx) + (ay * by) + (az * bz)")
# mag = ne.evaluate("sqrt(bx**2 + by**2 + bz**2)")
# prod = prod / mag
# return fld_a.wrap(prod, fldtype="Scalar")
ax, ay, az = fld_a.component_views() # pylint: disable=W0612
bx, by, bz = fld_b.component_views() # pylint: disable=W0612
projection = ne.evaluate("((ax * bx) + (ay * by) + (az * bz)) / "
"sqrt((bx**2) + (by**2) + (bz**2))")
return fld_a.wrap(projection, fldtype="Scalar")
[docs]def normalize(fld_a):
""" normalize a vector field """
a = fld_a.as_flat().data
ax, ay, az = fld_a.component_views() # pylint: disable=W0612
normed = ne.evaluate("a / sqrt((ax**2) + (ay**2) + (az**2))")
return fld_a.wrap(normed, fldtype="Vector")
[docs]def grad(fld, bnd=True):
"""2nd order centeral diff, 1st order @ boundaries if bnd"""
# vx, vy, vz = fld.component_views()
if bnd:
fld = viscid.extend_boundaries(fld, order=0, crd_order=0)
if fld.iscentered("Cell"):
crdx, crdy, crdz = fld.get_crds_cc(shaped=True)
# divcenter = "Cell"
# divcrds = coordinate.NonuniformCartesianCrds(fld.crds.get_clist(np.s_[1:-1]))
# divcrds = fld.crds.slice_keep(np.s_[1:-1, 1:-1, 1:-1])
elif fld.iscentered("Node"):
crdx, crdy, crdz = fld.get_crds_nc(shaped=True)
# divcenter = "Node"
# divcrds = coordinate.NonuniformCartesianCrds(fld.crds.get_clist(np.s_[1:-1]))
# divcrds = fld.crds.slice_keep(np.s_[1:-1, 1:-1, 1:-1])
else:
raise NotImplementedError("Can only do cell and node centered gradients")
v = fld.data
g = viscid.zeros(fld['x=1:-1, y=1:-1, z=1:-1'].crds, nr_comps=3)
xp, xm = crdx[2:, :, :], crdx[:-2, : , : ] # pylint: disable=bad-whitespace
yp, ym = crdy[ :, 2:, :], crdy[: , :-2, : ] # pylint: disable=bad-whitespace
zp, zm = crdz[ :, :, 2:], crdz[: , : , :-2] # pylint: disable=bad-whitespace
vxp, vxm = v[2: , 1:-1, 1:-1], v[ :-2, 1:-1, 1:-1] # pylint: disable=bad-whitespace
vyp, vym = v[1:-1, 2: , 1:-1], v[1:-1, :-2, 1:-1] # pylint: disable=bad-whitespace
vzp, vzm = v[1:-1, 1:-1, 2: ], v[1:-1, 1:-1, :-2] # pylint: disable=bad-whitespace
g['x'].data[...] = ne.evaluate("(vxp-vxm)/(xp-xm)")
g['y'].data[...] = ne.evaluate("(vyp-vym)/(yp-ym)")
g['z'].data[...] = ne.evaluate("(vzp-vzm)/(zp-zm)")
return g
[docs]def div(fld, bnd=True):
"""2nd order centeral diff, 1st order @ boundaries if bnd"""
if fld.iscentered("Face"):
# dispatch fc div immediately since that does its own pre-processing
return viscid.div_fc(fld, bnd=bnd)
if bnd:
fld = viscid.extend_boundaries(fld, order=0, crd_order=0)
vx, vy, vz = fld.component_views()
if fld.iscentered("Cell"):
crdx, crdy, crdz = fld.get_crds_cc(shaped=True)
divcenter = "Cell"
# divcrds = coordinate.NonuniformCartesianCrds(fld.crds.get_clist(np.s_[1:-1]))
divcrds = fld.crds.slice_keep(np.s_[1:-1, 1:-1, 1:-1])
elif fld.iscentered("Node"):
crdx, crdy, crdz = fld.get_crds_nc(shaped=True)
divcenter = "Node"
# divcrds = coordinate.NonuniformCartesianCrds(fld.crds.get_clist(np.s_[1:-1]))
divcrds = fld.crds.slice_keep(np.s_[1:-1, 1:-1, 1:-1])
else:
raise NotImplementedError("Can only do cell and node centered divs")
xp, xm = crdx[2:, :, :], crdx[:-2, : , : ] # pylint: disable=bad-whitespace
yp, ym = crdy[ :, 2:, :], crdy[: , :-2, : ] # pylint: disable=bad-whitespace
zp, zm = crdz[ :, :, 2:], crdz[: , : , :-2] # pylint: disable=bad-whitespace
vxp, vxm = vx[2: , 1:-1, 1:-1], vx[ :-2, 1:-1, 1:-1] # pylint: disable=bad-whitespace
vyp, vym = vy[1:-1, 2: , 1:-1], vy[1:-1, :-2, 1:-1] # pylint: disable=bad-whitespace
vzp, vzm = vz[1:-1, 1:-1, 2: ], vz[1:-1, 1:-1, :-2] # pylint: disable=bad-whitespace
div_arr = ne.evaluate("(vxp-vxm)/(xp-xm) + (vyp-vym)/(yp-ym) + "
"(vzp-vzm)/(zp-zm)")
return field.wrap_field(div_arr, divcrds, name="div " + fld.name,
center=divcenter, time=fld.time, parents=[fld])
[docs]def curl(fld, bnd=True):
"""2nd order centeral diff, 1st order @ boundaries if bnd"""
if bnd:
fld = viscid.extend_boundaries(fld, order=0, crd_order=0)
vx, vy, vz = fld.component_views()
if fld.iscentered("Cell"):
crdx, crdy, crdz = fld.get_crds_cc(shaped=True)
curlcenter = "cell"
# curlcrds = coordinate.NonuniformCartesianCrds(fld.crds.get_clist(np.s_[1:-1]))
curlcrds = fld.crds.slice_keep(np.s_[1:-1, 1:-1, 1:-1])
elif fld.iscentered("Node"):
crdx, crdy, crdz = fld.get_crds_nc(shaped=True)
curlcenter = "node"
# curlcrds = coordinate.NonuniformCartesianCrds(fld.crds.get_clist(np.s_[1:-1]))
curlcrds = fld.crds.slice_keep(np.s_[1:-1, 1:-1, 1:-1])
else:
raise NotImplementedError("Can only do cell and node centered divs")
xp, xm = crdx[2:, :, :], crdx[:-2, : , : ] # pylint: disable=C0326
yp, ym = crdy[ :, 2:, :], crdy[: , :-2, : ] # pylint: disable=C0326
zp, zm = crdz[ :, :, 2:], crdz[: , : , :-2] # pylint: disable=C0326
vxpy, vxmy = vx[1:-1, 2: , 1:-1], vx[1:-1, :-2, 1:-1] # pylint: disable=C0326
vxpz, vxmz = vx[1:-1, 1:-1, 2: ], vx[1:-1, 1:-1, :-2] # pylint: disable=C0326
vypx, vymx = vy[2: , 1:-1, 1:-1], vy[ :-2, 1:-1, 1:-1] # pylint: disable=C0326
vypz, vymz = vy[1:-1, 1:-1, 2: ], vy[1:-1, 1:-1, :-2] # pylint: disable=C0326
vzpx, vzmx = vz[2: , 1:-1, 1:-1], vz[ :-2, 1:-1, 1:-1] # pylint: disable=C0326
vzpy, vzmy = vz[1:-1, 2: , 1:-1], vz[1:-1, :-2, 1:-1] # pylint: disable=C0326
curl_x = ne.evaluate("(vzpy-vzmy)/(yp-ym) - (vypz-vymz)/(zp-zm)")
curl_y = ne.evaluate("-(vzpx-vzmx)/(xp-xm) + (vxpz-vxmz)/(zp-zm)")
curl_z = ne.evaluate("(vypx-vymx)/(xp-xm) - (vxpy-vxmy)/(yp-ym)")
return field.wrap_field([curl_x, curl_y, curl_z], curlcrds,
name="curl " + fld.name, fldtype="Vector",
center=curlcenter, time=fld.time,
parents=[fld])
##
## EOF
##