def test_input_args(self): """ Check that scalar_field can take different input arguments """ # Check for 2D arrays as positions vectors, and a function for # the data f = lambda x, y, z: x**2 + y**2 x, y = np.mgrid[-3:3, -3:3] z = np.zeros_like(x) ss = sources.scalar_field(x, y, z, f, figure=None) self.check_positions(ss, x, y, z) s = f(x, y, z) self.check_scalars(ss, s) # Check for a 2D array as data, and no position vectors s = np.random.random((10, 10)) ss = sources.scalar_field(s, figure=None) self.check_scalars(ss, s) # Check for a 3D array as data, and no position vectors s = np.random.random((10, 10, 10)) ss = sources.scalar_field(s, figure=None) self.check_scalars(ss, s) # Check for a 3D array as data, and 3D arrays as position x, y, z = np.mgrid[-3:3, -3:3, -3:3] ss = sources.scalar_field(x, y, z, z, figure=None) self.check_positions(ss, x, y, z) self.check_scalars(ss, z)
def get_sacdata_mlab(f, cube_slice, flux=True): """ Reads in useful variables from a hdf5 file to vtk data structures Parameters ---------- f : hdf5 file handle SAC HDF5 file flux : boolean Read variables for flux calculation? cube_slice : np.slice Slice to apply to the arrays Returns ------- bfield, vfield[, density, valf, cs, beta] """ # Do this before convert_B if flux: va_f = f.get_va() cs_f = f.get_cs() thermal_p, mag_p = f.get_thermalp(beta=True) beta_f = mag_p / thermal_p # Convert B to Tesla f.convert_B() # Create TVTK datasets bfield = vector_field( f.w_sac["b3"][cube_slice] * 1e3, f.w_sac["b2"][cube_slice] * 1e3, f.w_sac["b1"][cube_slice] * 1e3, name="Magnetic Field", figure=None, ) vfield = vector_field( f.w_sac["v3"][cube_slice] / 1e3, f.w_sac["v2"][cube_slice] / 1e3, f.w_sac["v1"][cube_slice] / 1e3, name="Velocity Field", figure=None, ) if flux: density = scalar_field(f.w_sac["rho"][cube_slice], name="Density", figure=None) valf = scalar_field(va_f, name="Alven Speed", figure=None) cs = scalar_field(cs_f, name="Sound Speed", figure=None) beta = scalar_field(beta_f, name="Beta", figure=None) return bfield, vfield, density, valf, cs, beta else: return bfield, vfield
def yt_to_mlab_scalar(ds, key, cube_slice=np.s_[:, :, :], field_name=""): """ Convert obe yt keys to a mlab scalar field Parameters ---------- ds : yt dataset The dataset to get the data from. key : string yt field names for the three vector components. cube_slice : numpy slice The array slice to crop the yt fields with. field_name : string The mlab name for the field. Returns ------- field : mayavi vector field """ cg = ds.index.grids[0] return scalar_field(cg[key][cube_slice], name=field_name, figure=None)
def get_yt_mlab(ds, cube_slice, flux=True): """ Reads in useful variables from yt to vtk data structures, converts into SI units before return. Parameters ---------- ds : yt dataset with derived fields flux : boolean Read variables for flux calculation? cube_slice : np.slice Slice to apply to the arrays Returns ------- if flux: bfield, vfield, density, valf, cs, beta else: bfield, vfield """ cg = ds.index.grids[0] # Create TVTK datasets bfield = yt_to_mlab_vector( ds, "mag_field_x", "mag_field_y", "mag_field_z", cube_slice=cube_slice, field_name="Magnetic Field" ) vfield = yt_to_mlab_vector( ds, "velocity_x", "velocity_y", "velocity_z", cube_slice=cube_slice, field_name="Velocity Field" ) if flux: density = scalar_field(cg["density"][cube_slice] * 1e-3, name="Density", figure=None) valf = scalar_field(cg["alfven_speed"] * 1e-2, name="Alven Speed", figure=None) cs = scalar_field(cg["sound_speed"] * 1e-2, name="Sound Speed", figure=None) beta = scalar_field(cg["plasma_beta"], name="Beta", figure=None) return bfield, vfield, density, valf, cs, beta else: return bfield, vfield