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 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 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