def calc_dip(pts, m=(0, 0, -DEFAULT_STRENGTH), strength=None, crd_system='gse',
dtype=None):
"""Calculate a dipole field at various points
Args:
pts (ndarray): Nx3 array of points at which to calculate the
dipole. Should use the same crd system as `m`
m (sequence, datetime): dipole moment
strength (None, float): If given, rescale m to this magnitude
crd_system (str): Something from which cotr can divine the
coordinate system for both `pts` and `m`. This is only used
if m is given as a datetime and we need to figure out the
dipole moment at a given time in a given crd system
dtype (str, np.dtype): dtype of the result, defaults to
the same datatype as `pts`
Returns:
ndarray: Nx3 dipole field vectors for N points
"""
pts = np.asarray(pts, dtype=dtype)
if len(pts.shape) == 1:
pts = pts.reshape(1, 3)
single_pt = True
else:
single_pt = False
if dtype is None:
dtype = pts.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=crd_system)
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
m = m.reshape(1, 3)
# geneate a dipole field for the entire grid
# Note: this is almost the same as fill_dipole, but all components
# are calculated simultaneously, and so this uses more memory
if _HAS_NUMEXPR:
_X, _Y, _Z = pts.T
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 = ne.evaluate("((three * pts * mdotr / rsq) - m) / rsq**1.5")
else:
_X, _Y, _Z = pts.T
rsq = _X**2 + _Y**2 + _Z**2
mdotr = mx * _X + my * _Y + mz * _Z
Bdip = ((three * pts * mdotr / rsq) - m) / rsq**1.5
if single_pt:
Bdip = Bdip[0, :]
return Bdip
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 time(self, val):
if viscid.is_datetime_like(val):
val = viscid.as_timedelta(self.basetime - viscid.as_datetime64(val))
val = val.total_seconds()
elif viscid.is_timedelta_like(val, conservative=True):
val = viscid.as_timedelta(val).total_seconds()
elif val is not None:
self.set_info('time', float(val))
def time(self, val):
if viscid.is_datetime_like(val):
val = viscid.as_timedelta(self.basetime -
viscid.as_datetime64(val))
val = val.total_seconds()
elif viscid.is_timedelta_like(val, conservative=True):
val = viscid.as_timedelta(val).total_seconds()
elif val is not None:
self.set_info('time', float(val))
def convert_timelike(value):
if value in (None, ''):
return None
elif viscid.is_datetime_like(value, conservative=True):
return viscid.as_datetime64(value)
elif viscid.is_timedelta_like(value, conservative=True):
return viscid.as_timedelta64(value)
else:
return value
def _apply_rotation(obj, from_system, rotate=None, crd_system='gse',
notilt1967=True):
if hasattr(rotate, "get_rotation_wxyz"):
rotate = rotate.get_rotation_wxyz(from_system, crd_system)
elif viscid.is_datetime_like(rotate):
cotr = viscid.Cotr(rotate, notilt1967=notilt1967) # pylint: disable=not-callable
rotate = cotr.get_rotation_wxyz(from_system, crd_system)
else:
cotr = viscid.Cotr() # pylint: disable=not-callable
rotate = cotr.get_rotation_wxyz(from_system, crd_system)
if len(rotate) != 4:
raise ValueError("Rotate should be [angle, ux, uy, uz], got {0}"
"".format(rotate))
obj.actor.actor.rotate_wxyz(*rotate)
def as_floating_t(self, t, none_passthrough=False):
t_as_s = None
try:
t = vutil.str_to_value(t)
if viscid.is_timedelta_like(t, conservative=True):
t_as_s = viscid.as_timedelta(t).total_seconds()
elif viscid.is_datetime_like(t, conservative=True):
delta_t = viscid.as_datetime64(t) - self.basetime
t_as_s = viscid.as_timedelta(delta_t).total_seconds()
elif not isinstance(t, (int, np.integer, type(None))):
t_as_s = float(t)
except AttributeError:
if t is None:
if none_passthrough:
pass
else:
t = 0.0
else:
t_as_s = float(t)
return t_as_s
def as_floating_t(self, t, none_passthrough=False):
t_as_s = None
try:
t = vutil.str_to_value(t)
if viscid.is_timedelta_like(t, conservative=True):
t_as_s = viscid.as_timedelta(t).total_seconds()
elif viscid.is_datetime_like(t, conservative=True):
delta_t = viscid.as_datetime64(t) - self.basetime
t_as_s = viscid.as_timedelta(delta_t).total_seconds()
elif not isinstance(t, (int, np.integer, type(None))):
t_as_s = float(t)
except AttributeError:
if t is None:
if none_passthrough:
pass
else:
t = 0.0
else:
t_as_s = float(t)
return t_as_s
def extract_index(arr, start=None, stop=None, step=None, endpoint=True,
tol=100):
"""Get integer indices for slice parts
If start, stop, or step are strings, they are either cast to
integers or used for a float lookup if they have a trailing 'f'.
An example float lookup is::
>>> [0.1, 0.2, 0.3, 0.4, 0.5, 0.6]['0.2f:0.6f:2']
[0.2, 0.4, 0.6]
The rules for float lookup endpoints are:
- The slice will never include an element whose value in arr
is < start (or > if the slice is backward)
- The slice will never include an element whose value in arr
is > stop (or < if the slice is backward)
- !! The slice WILL INCLUDE stop if you don't change endpoint.
This is different from normal slicing, but
it's more natural when specifying a slice as a float.
To this end, an epsilon tolerance can be given to
determine what's close enough.
- TODO: implement floating point steps, this is tricky
since arr need not be uniformly spaced, so step is
ambiguous in this case
Args:
arr (ndarray): filled with floats to do the lookup
start (None, int, str): like slice().start
stop (None, int, str): like slice().stop
step (None, int): like slice().step
endpoint (bool): iff True then include stop in the slice.
Set to False to get python slicing symantics when it
comes to excluding stop, but fair warning, python
symantics feel awkward here. Consider the case
[0.1, 0.2, 0.3][:0.25]. If you think this should include
0.2, then leave keep endpoint=True.
tol (int): number of machine epsilons to consider
"close enough"
Returns:
start, stop, step after floating point vals have been
converted to integers
"""
float_err_msg = ("Slicing by floats is no longer supported. If you "
"want to slice by location, suffix the value with "
"'f', as in 'x = 0f'.")
arr = np.asarray(arr)
try:
epsilon = tol * np.finfo(arr.dtype).eps
except ValueError:
# array is probably of type numpy.int*
epsilon = 0.01
_step = 1 if step is None else int(step)
epsilon_step = epsilon if _step > 0 else -epsilon
start = convert_timelike(start)
stop = convert_timelike(stop)
start = convert_deprecated_floats(start, "start")
stop = convert_deprecated_floats(stop, "stop")
# print("?!? |{0}| |{1}|".format(start, stop))
startstop = [start, stop]
eps_sign = [1, -1]
# if start or stop is not an int, try to make it one
for i in range(2):
byval = None
s = startstop[i]
_epsilon_step = epsilon_step
if viscid.is_datetime_like(s, conservative=True):
byval = s.astype(arr.dtype)
_epsilon_step = 0
elif viscid.is_timedelta_like(s, conservative=True):
byval = s.astype(arr.dtype)
_epsilon_step = 0
else:
try:
s = s.strip()
if len(s) == 0:
startstop[i] = None
elif s[-1] == 'f':
byval = float(s[:-1])
except AttributeError:
pass
if byval is not None:
if _epsilon_step:
diff = arr - byval + (eps_sign[i] * _epsilon_step)
else:
diff = arr - byval
zero = np.array([0]).astype(diff.dtype)[0]
# FIXME: there is far too much decision making here
if i == 0:
# start
if _step > 0:
diff = np.ma.masked_less(diff, zero)
else:
diff = np.ma.masked_greater(diff, zero)
if np.ma.count(diff) == 0:
# start value is past the wrong end of the array
if _step > 0:
startstop[i] = len(arr)
else:
# start = -len(arr) - 1
# having a value < -len(arr) won't play
# nice with make_fwd_slice, but in this
# case, the slice will have no data, so...
return 0, 0, step
else:
startstop[i] = int(np.argmin(np.abs(diff)))
else:
# stop
if _step > 0:
diff = np.ma.masked_greater(diff, zero)
if np.ma.count(diff) == 0:
# stop value is past the wong end of the array
startstop[i] = 0
else:
startstop[i] = int(np.argmin(np.abs(diff)))
if endpoint:
startstop[i] += 1
else:
diff = np.ma.masked_less(diff, zero)
if np.ma.count(diff) == 0:
# stop value is past the wrong end of the array
startstop[i] = len(arr)
else:
startstop[i] = int(np.argmin(np.abs(diff)))
if endpoint:
if startstop[i] > 0:
startstop[i] -= 1
else:
# 0 - 1 == -1 which would wrap to the end of
# of the array... instead, just make it None
startstop[i] = None
start, stop = startstop
# turn start, stop, step into indices
sss = [start, stop, step]
for i, s in enumerate(sss):
if s is None:
pass
elif isinstance(s, string_types):
sss[i] = int(s)
else:
sss[i] = s.__index__()
return sss
