def fill_dipole(B, m=(0, 0, -DEFAULT_STRENGTH), strength=None, mask=None):
"""set B to a dipole with magnetic moment m
Args:
B (Field): Field to fill with a dipole
m (ndarray, or datetime64-like): Description
strength (float): if given, rescale the dipole moment
even if it was given explicitly
mask (Field): boolean field as mask, B will be filled where
the mask is True
Returns:
Field: B
"""
# FIXME: should really be taking the curl of a vector field
if mask:
Bdip = field.empty_like(B)
else:
Bdip = B
# Xcc, Ycc, Zcc = B.get_crds_cc(shaped=True) # pylint: disable=W0612
Xv, Yv, Zv = B.get_crds_vector(shaped=True) # pylint: disable=W0612
_crd_lst = [[_x, _y, _z] for _x, _y, _z in zip(Xv, Yv, Zv)]
dtype = B.dtype
one = np.array([1.0], dtype=dtype) # pylint: disable=W0612
three = np.array([3.0], dtype=dtype) # pylint: disable=W0612
if viscid.is_datetime_like(m):
m = viscid.get_dipole_moment(m, crd_system=B)
else:
m = np.asarray(m, dtype=dtype)
if strength is not None:
m = (strength / np.linalg.norm(m)) * m
mx, my, mz = m # pylint: disable=W0612
# geneate a dipole field for the entire grid
# Note: this is almost the exact same as calc_dip, but since components
# are done one-at-a-time, it requires less memory since it copies the
# result of each component into Bdip separately
if _HAS_NUMEXPR:
for i, cn in enumerate("xyz"):
_X, _Y, _Z = _crd_lst[i]
_XI = _crd_lst[i][i]
_mi = m[i]
rsq = ne.evaluate("_X**2 + _Y**2 + _Z**2") # pylint: disable=W0612
mdotr = ne.evaluate("mx * _X + my * _Y + mz * _Z") # pylint: disable=W0612
Bdip[cn] = ne.evaluate("((three * _XI * mdotr / rsq) - _mi) / rsq**1.5")
else:
for i, cn in enumerate("xyz"):
_X, _Y, _Z = _crd_lst[i]
_XI = _crd_lst[i][i]
_mi = m[i]
rsq = _X**2 + _Y**2 + _Z**2
mdotr = mx * _X + my * _Y + mz * _Z
Bdip[cn] = ((three * _XI * mdotr / rsq) - _mi) / rsq**1.5
if mask:
B.data[...] = np.choose(mask.astype('i'), [B, Bdip])
return B
def set_in_region(a, b, alpha=1.0, beta=1.0, mask=None, out=None):
"""set `ret = alpha * a + beta * b` where mask is True"""
alpha = np.asarray(alpha, dtype=a.dtype)
beta = np.asarray(beta, dtype=a.dtype)
a_dat = a.data if isinstance(a, viscid.field.Field) else a
b_dat = b.data if isinstance(b, viscid.field.Field) else b
b = None
if _HAS_NUMEXPR:
vals = ne.evaluate("alpha * a_dat + beta * b_dat")
else:
vals = alpha * a_dat + beta * b_dat
a_dat = b_dat = None
if out is None:
out = field.empty_like(a)
if mask is None:
out.data[...] = vals
else:
if hasattr(mask, "nr_comps") and mask.nr_comps:
mask = mask.as_centered(a.center).as_layout(a.layout)
try:
out.data[...] = np.choose(mask, [out.data, vals])
except ValueError:
out.data[...] = np.choose(mask.data.reshape(list(mask.sshape) + [1]),
[out.data, vals])
return out
