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