def _add_mesh_sampling_particle_field(self, deposit_field, ftype, ptype): units = self.ds.field_info[ftype, deposit_field].units take_log = self.ds.field_info[ftype, deposit_field].take_log field_name = f"cell_{ftype}_{deposit_field}" def _mesh_sampling_particle_field(field, data): pos = data[ptype, "particle_position"] field_values = data[ftype, deposit_field] if isinstance(data, FieldDetector): return np.zeros(pos.shape[0]) i, j, k = np.floor( (pos - data.LeftEdge) / data.dds).astype("int64").T # Make sure all particles are within the current grid, otherwise return nan maxi, maxj, maxk = field_values.shape mask = (i < maxi) & (j < maxj) & (k < maxk) mask &= (i >= 0) & (j >= 0) & (k >= 0) result = np.full(len(pos), np.nan, dtype="float64") if result.shape[0] > 0: result[mask] = field_values[i[mask], j[mask], k[mask]] return data.ds.arr(result, field_values.units) self.ds.add_field( (ptype, field_name), function=_mesh_sampling_particle_field, sampling_type="particle", units=units, take_log=take_log, validators=[ValidateSpatial()], )
def add_particle_average(registry, ptype, field_name, weight="particle_mass", density=True): field_units = registry[ptype, field_name].units def _pfunc_avg(field, data): pos = data[ptype, "particle_position"] f = data[ptype, field_name] wf = data[ptype, weight] f *= wf v = data.deposit(pos, [f], method="sum") w = data.deposit(pos, [wf], method="sum") v /= w if density: v /= data["index", "cell_volume"] v[np.isnan(v)] = 0.0 return v fn = ("deposit", "%s_avg_%s" % (ptype, field_name)) registry.add_field(fn, function=_pfunc_avg, validators=[ValidateSpatial(0)], particle_type=False, units=field_units) return fn
def setup_counts_fields(ds, ebounds, ftype="gas"): r""" Create deposited image fields from X-ray count data in energy bands. Parameters ---------- ds : Dataset The FITS events file dataset to add the counts fields to. ebounds : list of tuples A list of tuples, one for each field, with (emin, emax) as the energy bounds for the image. ftype : string, optional The field type of the resulting field. Defaults to "gas". Examples -------- >>> ds = yt.load("evt.fits") >>> ebounds = [(0.1,2.0),(2.0,3.0)] >>> setup_counts_fields(ds, ebounds) """ for (emin, emax) in ebounds: cfunc = _make_counts(emin, emax) fname = "counts_%s-%s" % (emin, emax) mylog.info("Creating counts field %s." % fname) ds.add_field((ftype, fname), function=cfunc, units="counts/pixel", validators=[ValidateSpatial()], display_name="Counts (%s-%s keV)" % (emin, emax))
def add_particle_average(registry, ptype, field_name, weight=None, density=True): if weight is None: weight = (ptype, "particle_mass") field_units = registry[ptype, field_name].units def _pfunc_avg(field, data): pos = data[ptype, "particle_position"] f = data[ptype, field_name] wf = data[ptype, weight] f *= wf v = data.deposit(pos, [f], method="sum") w = data.deposit(pos, [wf], method="sum") v /= w if density: v /= data["index", "cell_volume"] v[np.isnan(v)] = 0.0 return v fn = ("deposit", f"{ptype}_avg_{field_name}") registry.add_field( fn, sampling_type="cell", function=_pfunc_avg, validators=[ValidateSpatial(0)], units=field_units, ) return fn
def add_nearest_neighbor_value_field(ptype, coord_name, sampled_field, registry): """ This adds a nearest-neighbor field, where values on the mesh are assigned based on the nearest particle value found. This is useful, for instance, with voronoi-tesselations. """ field_name = ("deposit", f"{ptype}_nearest_{sampled_field}") field_units = registry[ptype, sampled_field].units unit_system = registry.ds.unit_system def _nearest_value(field, data): pos = data[ptype, coord_name] pos = pos.convert_to_units("code_length") value = data[ptype, sampled_field].in_base(unit_system.name) rv = data.smooth( pos, [value], method="nearest", create_octree=True, nneighbors=1 ) rv = data.apply_units(rv, field_units) return rv registry.add_field( field_name, sampling_type="cell", function=_nearest_value, validators=[ValidateSpatial(0)], units=field_units, ) return [field_name]
def add_volume_weighted_smoothed_field(ptype, coord_name, mass_name, smoothing_length_name, density_name, smoothed_field, registry, nneighbors = None): field_name = ("deposit", "%s_smoothed_%s" % (ptype, smoothed_field)) field_units = registry[ptype, smoothed_field].units def _vol_weight(field, data): pos = data[ptype, coord_name].in_units("code_length") mass = data[ptype, mass_name].in_cgs() dens = data[ptype, density_name].in_cgs() quan = data[ptype, smoothed_field].in_units(field_units) if smoothing_length_name is None: hsml = np.zeros(quan.shape, dtype='float64') - 1 hsml = data.apply_units(hsml, "code_length") else: hsml = data[ptype, smoothing_length_name].in_units("code_length") kwargs = {} if nneighbors: kwargs['nneighbors'] = nneighbors rv = data.smooth(pos, [mass, hsml, dens, quan], method="volume_weighted", create_octree = True)[0] rv[np.isnan(rv)] = 0.0 # Now some quick unit conversions. rv = data.apply_units(rv, field_units) return rv registry.add_field(field_name, function = _vol_weight, validators = [ValidateSpatial(0)], units = field_units) return [field_name]
def _add_mesh_sampling_particle_field(self, deposit_field, ftype, ptype): units = self.ds.field_info[ftype, deposit_field].units take_log = self.ds.field_info[ftype, deposit_field].take_log field_name = f"cell_{ftype}_{deposit_field}" def _mesh_sampling_particle_field(field, data): pos = data[ptype, "particle_position"] field_values = data[ftype, deposit_field] if isinstance(data, FieldDetector): return np.zeros(pos.shape[0]) i, j, k = np.floor( (pos - data.LeftEdge) / data.dds).astype("int64").T return field_values[i, j, k] self.ds.add_field( (ptype, field_name), function=_mesh_sampling_particle_field, sampling_type="particle", units=units, take_log=take_log, validators=[ValidateSpatial()], )
def add_density_kernel(ptype, coord_name, mass_name, registry, nneighbors=64): field_name = (ptype, "smoothed_density") field_units = registry[ptype, mass_name].units def _nth_neighbor(field, data): pos = data[ptype, coord_name] pos.convert_to_units("code_length") mass = data[ptype, mass_name] mass.convert_to_units("g") densities = mass * 0.0 data.particle_operation(pos, [mass, densities], method="density", nneighbors=nneighbors) ones = pos.prod(axis=1) # Get us in code_length**3 ones[:] = 1.0 densities /= ones # Now some quick unit conversions. return densities registry.add_field(field_name, function=_nth_neighbor, validators=[ValidateSpatial(0)], particle_type=True, units="g/cm**3") return [field_name]
def add_deposited_particle_field(self, deposit_field, method): """Add a new deposited particle field Creates a new deposited field based on the particle *deposit_field*. Parameters ---------- deposit_field : tuple The field name tuple of the particle field the deposited field will be created from. This must be a field name tuple so yt can appropriately infer the correct particle type. method : one of 'count', 'sum', or 'cic' The particle deposition method to use. Returns ------- The field name tuple for the newly created field. """ self.index if isinstance(deposit_field, tuple): ptype, deposit_field = deposit_field[0], deposit_field[1] else: raise RuntimeError units = self.field_info[ptype, deposit_field].units def _deposit_field(field, data): """ Create a grid field for particle wuantities weighted by particle mass, using cloud-in-cell deposition. """ pos = data[ptype, "particle_position"] # get back into density if method != 'count': pden = data[ptype, "particle_mass"] top = data.deposit(pos, [data[(ptype, deposit_field)] * pden], method=method) bottom = data.deposit(pos, [pden], method=method) top[bottom == 0] = 0.0 bnz = bottom.nonzero() top[bnz] /= bottom[bnz] d = data.ds.arr(top, input_units=units) else: d = data.ds.arr( data.deposit(pos, [data[ptype, deposit_field]], method=method)) return d name_map = {"cic": "cic", "sum": "nn", "count": "count"} field_name = "%s_" + name_map[method] + "_%s" field_name = field_name % (ptype, deposit_field.replace( 'particle_', '')) self.add_field(("deposit", field_name), function=_deposit_field, units=units, take_log=False, validators=[ValidateSpatial()]) return ("deposit", field_name)
def add_volume_weighted_smoothed_field(ptype, coord_name, mass_name, smoothing_length_name, density_name, smoothed_field, registry, nneighbors=64, kernel_name='cubic'): unit_system = registry.ds.unit_system if kernel_name == 'cubic': field_name = ("deposit", "%s_smoothed_%s" % (ptype, smoothed_field)) else: field_name = ("deposit", "%s_%s_smoothed_%s" % (ptype, kernel_name, smoothed_field)) field_units = registry[ptype, smoothed_field].units def _vol_weight(field, data): pos = data[ptype, coord_name] pos = pos.convert_to_units("code_length") mass = data[ptype, mass_name].in_base(unit_system.name) dens = data[ptype, density_name].in_base(unit_system.name) quan = data[ptype, smoothed_field] if hasattr(quan, "units"): quan = quan.convert_to_units(field_units) if smoothing_length_name is None: hsml = np.zeros(quan.shape, dtype='float64') - 1 hsml = data.apply_units(hsml, "code_length") else: hsml = data[ptype, smoothing_length_name] hsml.convert_to_units("code_length") # This is for applying cutoffs, similar to in the SPLASH paper. smooth_cutoff = data["index", "cell_volume"]**(1. / 3) smooth_cutoff.convert_to_units("code_length") # volume_weighted smooth operations return lists of length 1. rv = data.smooth(pos, [mass, hsml, dens, quan], index_fields=[smooth_cutoff], method="volume_weighted", create_octree=True, nneighbors=nneighbors, kernel_name=kernel_name)[0] rv[np.isnan(rv)] = 0.0 # Now some quick unit conversions. # This should be used when seeking a non-normalized value: rv /= hsml.uq**3 / hsml.uq.in_base(unit_system.name).uq**3 rv = data.apply_units(rv, field_units) return rv registry.add_field(field_name, sampling_type="cell", function=_vol_weight, validators=[ValidateSpatial(0)], units=field_units) registry.find_dependencies((field_name, )) return [field_name]
def add_nearest_neighbor_field(ptype, coord_name, registry, nneighbors = 64): field_name = (ptype, "nearest_neighbor_distance_%s" % (nneighbors)) def _nth_neighbor(field, data): pos = data[ptype, coord_name] pos.convert_to_units("code_length") distances = 0.0 * pos[:,0] data.particle_operation(pos, [distances], method="nth_neighbor", nneighbors = nneighbors) # Now some quick unit conversions. return distances registry.add_field(field_name, sampling_type="particle", function = _nth_neighbor, validators = [ValidateSpatial(0)], units = "code_length") return [field_name]
def add_contour_field(ds, contour_key): def _contours(field, data): fd = data.get_field_parameter("contour_slices_%s" % contour_key) vals = data["index", "ones"] * -1 if fd is None or fd == 0.0: return vals for sl, v in fd.get(data.id, []): vals[sl] = v return vals ds.add_field(("index", "contours_%s" % contour_key), function=_contours, validators=[ValidateSpatial(0)], take_log=False, display_field=False)
ValidateSpatial, \ NeedsParameter import h5py from yt.funcs import \ just_one from common_functions import * sl_left = slice(None, -2, None) sl_right = slice(2, None, None) div_fac = 2.0 sl_center = slice(1, -1, None) ftype='gas' vort_validators = [ValidateSpatial(1, [(ftype, "velocity_x"), (ftype, "velocity_y"), (ftype, "velocity_z")])] def _Disk_H(field, data): center = data.get_field_parameter('center') z = data["z"] - center[2] return np.abs(z) def _vturb(field, data): fx = data[ftype, "velocity_x"][sl_right,sl_center,sl_center]/6.0 fx += data[ftype, "velocity_x"][sl_left,sl_center,sl_center]/6.0 fx += data[ftype, "velocity_x"][sl_center,sl_right,sl_center]/6.0 fx += data[ftype, "velocity_x"][sl_center,sl_left,sl_center]/6.0 fx += data[ftype, "velocity_x"][sl_center,sl_center,sl_right]/6.0
def func(field, data): return -data['spitzer_conduction_coefficient']*yt.physical_constants.kb*data['H_nuclei_density']*data['temperature_gradient_%s' % ax] return func for ax in 'xyz': f = _spitzer_heat_flux(ax) yt.add_field('%s_%s' % (basename, ax), function=f, sampling_type='cell', units='erg/s/cm**2') sl_left = slice(None, -2, None) sl_right = slice(2, None, None) div_fac = 2 @yt.derived_field(('gas', "%s_divergence" % basename), sampling_type="cell", units='erg/s/cm**3', validators=[ValidateSpatial(2)]) def _divergence(field, data): ds = div_fac * just_one(data["index", "dx"]) f = data[xn][sl_right,1:-1,1:-1]/ds f -= data[xn][sl_left ,1:-1,1:-1]/ds ds = div_fac * just_one(data["index", "dy"]) f += data[yn][1:-1,sl_right,1:-1]/ds f -= data[yn][1:-1,sl_left ,1:-1]/ds ds = div_fac * just_one(data["index", "dz"]) f += data[zn][1:-1,1:-1,sl_right]/ds f -= data[zn][1:-1,1:-1,sl_left ]/ds new_field = data.ds.arr(np.zeros(data[xn].shape, dtype=np.float64), f.units) new_field[1:-1,1:-1,1:-1] = f return new_field
def add_synchrotron_dtau_emissivity(ds, ptype='lobe', nu=(1.4, 'GHz'), method='nearest_weighted', proj_axis='x', extend_cells=32): me = yt.utilities.physical_constants.mass_electron #9.109E-28 c = yt.utilities.physical_constants.speed_of_light #2.998E10 e = yt.utilities.physical_constants.elementary_charge #4.803E-10 esu gamma_min = yt.YTQuantity(10, 'dimensionless') # Index for electron power law distribution p = 2.0 pol_ratio = (p + 1.) / (p + 7. / 3.) # Fitted constants for the approximated power-law + exponential spectra # Integral of 2*F(x) -> tot_const*(nu**-2)*exp(-nu/nuc) # 2*F(x) for the total intensity (parallel + perpendicular) tot_const = 4.1648 stokes = StokesFieldName(ptype, nu, proj_axis) nu = yt.YTQuantity(*nu) if proj_axis == 'x': los = [1., 0., 0.] xvec = [0., 1., 0.] yvec = [0., 0., 1.] elif proj_axis == 'y': los = [0., 1., 0.] xvec = [0., 0., 1.] yvec = [1., 0., 0.] elif proj_axis == 'z': los = [0., 0., 1.] xvec = [1., 0., 0.] yvec = [0., 1., 0.] elif type(proj_axis) is list: los = proj_axis if los[0] != 0.: # not perpendicular to z-axis xvec = [0., 1., 0.] yvec = [0., 0., 1.] # TODO: xvec and yvec for arbitrary proj_axis else: raise NotImplementedError else: raise NotImplementedError los = np.array(los) xvec = np.array(xvec) yvec = np.array(yvec) los = los / np.sqrt(np.sum(los * los)) # Determine the version of the simulation if ('io', 'particle_tau1') in ds.field_list: version = 2018 mylog.info( 'Field particle_tau1 detected. This is a version 2018 simulation.') elif ('io', 'particle_type') in ds.field_list: version = 2016 mylog.info( 'Field particle_type detected. This is a version 2016 simulation.') else: version = 2015 # Update the particle file handler in yt; raise exception if not successful success = setup_part_file(ds) mylog.info('Assuming this is a version 2015 simulation.') if not success: raise IOError if ('gas', 'jet_volume_fraction') not in ds.derived_field_list: ds.add_field(('gas', 'jet_volume_fraction'), function=_jet_volume_fraction, display_name="Jet Volume Fraction", sampling_type='cell') #def _gamc(field, data): # # The new cutoff gamma # # Note that particle_dens could be negative due to quadratic interpolation! # gamc = (np.abs(data['particle_dens'] / data['particle_den0']))**(1./3.) \ # / (data['particle_dtau'] + np.finfo(np.float64).tiny) # ind = np.where(gamc < 0.0)[0] # if ind.shape[0] > 0: # print(ind) # print(gamc) # return gamc #pfname = 'particle_gamc_dtau' #ds.add_field(pfname, function=_gamc, sampling_type='particle', # units='', force_override=True) #deposit_field = 'particle_gamc_dtau' def _synchrotron_spec(field, data): # To convert from FLASH "none" unit to cgs unit, times the B field from FLASH by sqrt(4*pi) Bvec = np.array([data['particle_magx'],\ data['particle_magy'],\ data['particle_magz']])*np.sqrt(4.0*np.pi) Bvec = data.apply_units(Bvec, 'gauss') # Calculate sin(a), in which a is the pitch angle of the electrons relative to B field. # See _nn_emissivity_i for more comments cross = np.cross(los, Bvec, axisb=0) Bsina = np.sqrt(np.sum(cross * cross, axis=-1)) Bsina = data.apply_units(Bsina, 'gauss') #B = np.sqrt(np.sum(Bvec*Bvec, axis=0)) if version == 2018: # Return for the FieldDetector; do nothing if isinstance(data, FieldDetector): return data['particle_dens']/data['particle_den0']**(1./3.)/ \ (data['particle_tau1']+data['particle_cmb1']) gamc = (np.abs(data['particle_dens'] / data['particle_den0']))**(1./3.) \ / (data['particle_tau1'] + data['particle_cmb1'] + np.finfo(np.float64).tiny) else: # Return for the FieldDetector; do nothing if isinstance(data, FieldDetector): return data['particle_dens']/data['particle_den0']**(1./3.)/ \ (data['particle_dtau']) if np.any(data['particle_dtau'] < 0.0): print('negative tau!') print(data) print(data['particle_tau']) print(data['particle_dtau']) # The new cutoff gamma # Note that particle_dens could be negative due to quadratic interpolation! gamc = (np.abs(data['particle_dens'] / data['particle_den0']))**(1./3.) \ / (data['particle_dtau'] + np.finfo(np.float64).tiny) ind = np.where(gamc <= 0.0)[0] if ind.shape[0] > 0: print(ind) print(gamc[ind]) #gamc = data[(ptype, 'particle_gamc')] # Cutoff frequency nuc = 3.0 * gamc**2 * e * Bsina / (4.0 * np.pi * me * c) #nu = data.get_field_parameter("frequency", default=yt.YTQuantity(1.4, 'GHz')) # B**1.5 is taken from the grid data norm = 3.0 / 8.0 * e**3.5 / (c**2.5 * me**1.5 * (np.pi)**0.5) # P is taken from the grid data N0 = 3.0 / me / c / c / (np.log(gamc / gamma_min)) / yt.YTQuantity( 4. * np.pi, 'sr') # Fix where the cutoff gamma < 0 N0[ind] = 0.0 return np.clip(N0 * norm * nu**(-0.5) * np.exp(-nu / nuc), np.finfo(np.float64).tiny, None) # particle field name pfname = 'particle_sync_spec_%s' % stokes.nu_str ds.add_field(pfname, function=_synchrotron_spec, sampling_type='particle', units='cm**(3/4)*s**(3/2)/g**(3/4)/sr', force_override=True) deposit_field = pfname #try: ds.add_particle_filter(ptype) #except: # raise NotImplementedError ########################################################################### ## Nearest Neighbor method ########################################################################### #fname_nn = ds.add_deposited_particle_field( # (ptype, 'particle_sync_spec_%s' % stokes.nu_str), 'nearest', extend_cells=extend_cells) sync_unit = ds.field_info[deposit_field].units if method == "nearest": field_name = "%s_nn_%s" elif method == "nearest_weighted": field_name = "%s_nnw_%s" else: raise NotImplementedError field_name = field_name % (ptype, deposit_field.replace('particle_', '')) ad = ds.all_data() #print(ad[ptype, "particle_position"]) def _nnw_deposit_field(field, data): if isinstance(data, FieldDetector): jetfluid = data['velocity_magnitude'] > lobe_v d = data.deposit(data[ptype, "particle_position"], (ptype, deposit_field), method=method) d[jetfluid] = 0.0 return d #pos = ad[ptype, "particle_position"] if ptype == 'lobe': # Calling other fields must preceed the more intensive # deposit function to prevent repeated iterations jetfluid = data['velocity_magnitude'] > lobe_v alldata = True if alldata: pos = ad[ptype, "particle_position"] # Deposit using the distance weighted log field #fields = [np.log(ad[ptype, deposit_field])] fields = [ad[ptype, deposit_field]] fields = [np.ascontiguousarray(f) for f in fields] else: left_edge = np.maximum(data.LeftEdge-data.dds*extend_cells,\ data.ds.domain_left_edge) right_edge = np.minimum(data.RightEdge+data.dds*extend_cells,\ data.ds.domain_right_edge) box = data.ds.box(left_edge, right_edge) pos = box[ptype, "particle_position"] fields = [box[ptype, deposit_field]] d = data.deposit(pos, fields, method=method, extend_cells=extend_cells) if np.all(d == 0.0): d = data.ds.arr(d, input_units=sync_unit) else: # Conver the log back to real value #d = data.ds.arr(np.exp(d), input_units=sync_unit) d = data.ds.arr(d, input_units=sync_unit) if ptype == 'lobe': d[jetfluid] = 0.0 return d fname_nn = ("deposit", field_name) ds.add_field(fname_nn, function=_nnw_deposit_field, sampling_type="cell", units=sync_unit, take_log=True, force_override=True, validators=[ValidateSpatial()]) def _nn_emissivity_i(field, data): ''' Emissivity using nearest neighbor. Integrate over line of sight to get intensity. ''' Bvec = np.array([data[('gas', 'magnetic_field_x')],\ data[('gas', 'magnetic_field_y')],\ data[('gas', 'magnetic_field_z')]]) # Calculate sin(a), in which a is the pitch angle of the electrons relative to B field. # We only see the radiation from electrons with pitch angles pointing to line of sight. cross = np.cross(los, Bvec, axisb=0) # B * sin(alpha) = (B * |(los x Bvec)|/|los|/|Bvec|) # = |(los x Bvec)| Bsina = np.sqrt(np.sum(cross * cross, axis=-1)) Bsina = data.apply_units(Bsina, 'gauss') # P * (B*sina)^1.5 PBsina = data['gas', 'pressure'] * Bsina**1.5 frac = data['gas', 'jet_volume_fraction'] # Use gamc from deposit field ############################# # fname_nn = 'ptype_nnw_gamc_dtau' #if ('flash', fname_nn[1]) in data.ds.field_list: # gamc = data[('flash', fname_nn[1])] #else: # gamc = data[fname_nn] #bad_mask = gamc <= gamma_min ## Cutoff frequency #nuc = 3.0*gamc**2*e*Bsina/(4.0*np.pi*me*c) ##nu = data.get_field_parameter("frequency", default=yt.YTQuantity(1.4, 'GHz')) #norm = 3.0*Bsina**1.5/8.0*e**3.5/(c**2.5*me**1.5*(np.pi)**0.5) #N0 = 3.0*data['gas', 'pressure']/me/c/c/(np.log(gamc/gamma_min))/yt.YTQuantity(4.*np.pi, 'sr') #N0[bad_mask] = 0.0 return PBsina * frac * tot_const * data[fname_nn] ds.add_field(stokes.I, function=_nn_emissivity_i, sampling_type='cell', display_name=stokes.display_name('I'), units='Jy/cm/arcsec**2', take_log=True, force_override=True, validators=[ValidateSpatial()]) def _cos_theta(field, data): Bvec = np.stack([data[('gas', 'magnetic_field_x')],\ data[('gas', 'magnetic_field_y')],\ data[('gas', 'magnetic_field_z')]], axis=-1) Bproj = Bvec - np.expand_dims(np.inner(Bvec, los), -1) * los # cos = cos(theta), theta is the angle between projected B and xvec # Ignore invalid 0/0 warnings with np.errstate(invalid='ignore'): cos = np.inner(Bproj, xvec) / np.sqrt( np.sum(Bproj * Bproj, axis=-1)) return cos ds.add_field('cos', function=_cos_theta, sampling_type='cell', display_name='cos theta', units='', take_log=False, force_override=True) def _nn_emissivity_q(field, data): # pol_ratio = (perp - para) / (perp + para) # The minus accounts for the perpendicular polarization rtn = -data[stokes.I] * pol_ratio * (2 * data['cos']**2 - 1.0) rtn[~np.isfinite(data['cos'])] = 0.0 return rtn ds.add_field(stokes.Q, function=_nn_emissivity_q, sampling_type='cell', display_name=stokes.display_name('Q'), units='Jy/cm/arcsec**2', take_log=False, force_override=True) def _nn_emissivity_u(field, data): sin = np.sqrt(1.0 - data['cos']**2) rtn = -data[stokes.I] * pol_ratio * 2 * sin * data['cos'] rtn[~np.isfinite(data['cos'])] = 0.0 return rtn ds.add_field(stokes.U, function=_nn_emissivity_u, sampling_type='cell', display_name=stokes.display_name('U'), units='Jy/cm/arcsec**2', take_log=False, force_override=True) return pfname, fname_nn, stokes.I, stokes.nu_str
def create_fields(ds): units_override = { "length_unit": (1, "Rsun"), "time_unit": (6.955e+05, "s"), "mass_unit": (3.36427433875e+17, "g"), "magnetic_unit": (1.121e-02, "G") } def _radialvelocity(field, data): return data['velocity_x']*data['x']/data['radius'] + \ data['velocity_y']*data['y']/data['radius'] + \ data['velocity_z']*data['z']/data['radius'] ds.add_field(('gas', "radialvelocity"), function=_radialvelocity, units="cm/s", take_log=False) def _sound_speed(field, data): gamma = 1.05 ftype = field.name[0] tr = gamma * data[ftype, "pressure"] / data[ftype, "density"] return np.sqrt(tr) ds.add_field(('gas', "sound_speed"), function=_sound_speed, units="cm/s", take_log=False) def _mach_number(field, data): """ M{|v|/c_sound} """ ftype = field.name[0] return data[ftype, "velocity_magnitude"] / data[ftype, "sound_speed"] ds.add_field(('gas', "mach_number"), function=_mach_number, units="", take_log=False) def _temperature(field, data): return (data["gas", "pressure"] * mp) / (2.0 * data["gas", "density"] * kb) ds.add_field(('gas', "temperature"), function=_temperature, units="K", take_log=True) def _radius_planet(field, data): a = 0.047 * 1.496e+13 / 6.955e+10 shift = data.ds.arr(np.ones_like(data['x'])) * a x_planet = data['x'] - shift return np.sqrt(x_planet*x_planet \ + data['y']*data['y'] \ + data['z']*data['z']) ds.add_field(('index', "radius_planet"), function=_radius_planet, units="cm", take_log=False) def _ni(field, data): return data["density"] / (1.09 * mp) ds.add_field(("gas", "ni"), function=_ni, units="cm**-3") def _BGx1(field, data): B0s = YTQuantity(2.0, "G") B0p = YTQuantity(1.0, "G") Rs = YTQuantity(6.955e+10, "cm") Rp = YTQuantity(1.5 * 0.10045 * Rs, "cm") a = YTQuantity(0.047, "au").in_units("cm") center = data.get_field_parameter('center') x1 = data["x"].in_units('cm') x2 = data["y"].in_units('cm') x3 = data["z"].in_units('cm') rs = np.sqrt(x1 * x1 + x2 * x2 + x3 * x3) rp = np.sqrt((x1 - a) * (x1 - a) + x2 * x2 + x3 * x3) BGx1 = data.ds.arr(np.zeros_like(data["magnetic_field_x"]), "G") BGx1 = 3.0 * x1 * x3 * B0s * Rs**3 * rs**(-5) + 3.0 * ( x1 - a) * x3 * B0p * Rp**3 * rp**(-5) BGx1[rs <= Rs] = 3.0 * x1[rs <= Rs] * x3[rs <= Rs] * B0s * Rs**3 * rs[ rs <= Rs]**(-5) BGx1[rs <= 0.5 * Rs] = 0.0 BGx1[rp <= Rp] = 3.0*(x1[rp <= Rp] - a)*x3[rp <= Rp]\ *B0p*Rp**3*rp[rp <= Rp]**(-5) BGx1[rp <= 0.5 * Rp] = 0.0 return BGx1 ds.add_field(("gas", "BGx1"), function=_BGx1, units="G", take_log=False) def _BGx2(field, data): B0s = YTQuantity(2.0, "G") B0p = YTQuantity(1.0, "G") Rs = YTQuantity(6.955e+10, "cm") Rp = YTQuantity(1.5 * 0.10045 * Rs, "cm") a = YTQuantity(0.047, "au").in_units("cm") center = data.get_field_parameter('center') x1 = data["x"].in_units('cm') x2 = data["y"].in_units('cm') x3 = data["z"].in_units('cm') rs = np.sqrt(x1 * x1 + x2 * x2 + x3 * x3) rp = np.sqrt((x1 - a) * (x1 - a) + x2 * x2 + x3 * x3) BGx2 = data.ds.arr(np.zeros_like(data["magnetic_field_y"]), "G") BGx2 = 3.0 * x3 * x2 * B0s * Rs**3 * rs**( -5) + 3.0 * x3 * x2 * B0p * Rp**3 * rp**(-5) BGx2[rs <= Rs] = 3.0*x3[rs <= Rs]*x2[rs <= Rs]\ *B0s*Rs**3*rs[rs <= Rs]**(-5) BGx2[rs <= 0.5 * Rs] = 0.0 BGx2[rp <= Rp] = 3.0*x3[rp <= Rp]*x2[rp <= Rp]\ *B0p*Rp**3*rp[rp <= Rp]**(-5) BGx2[rp <= 0.5 * Rp] = 0.0 return BGx2 ds.add_field(("gas", "BGx2"), function=_BGx2, units="G", take_log=False) def _BGx3(field, data): B0s = YTQuantity(2.0, "G") B0p = YTQuantity(1.0, "G") Rs = YTQuantity(6.955e+10, "cm") Rp = YTQuantity(1.5 * 0.10045 * Rs, "cm") a = YTQuantity(0.047, "au").in_units("cm") x1 = data["x"].in_units('cm') x2 = data["y"].in_units('cm') x3 = data["z"].in_units('cm') rs = np.sqrt(x1 * x1 + x2 * x2 + x3 * x3) rp = np.sqrt((x1 - a) * (x1 - a) + x2 * x2 + x3 * x3) BG_z = data.ds.arr(np.zeros_like(data["magnetic_field_z"]), "G") BGx3 = (3.0*x3*x3 - rs*rs)*B0s*Rs**3*rs**(-5) \ + (3.0*x3*x3 - rp*rp)*B0p*Rp**3*rp**(-5) BGx3[rs <= Rs] = (3.0*x3[rs <= Rs]*x3[rs <= Rs] - \ rs[rs <= Rs]*rs[rs <= Rs])*B0s*Rs**3*rs[rs <= Rs]**(-5) BGx3[rs <= 0.5 * Rs] = 16.0 * B0s BGx3[rp <= Rp] = (3.0*x3[rp <= Rp]*x3[rp <= Rp] - \ rp[rp <= Rp]*rp[rp <= Rp])*B0p*Rp**3*rp[rp <= Rp]**(-5) BGx3[rp <= 0.5 * Rp] = 16.0 * B0p return BGx3 ds.add_field(("gas", "BGx3"), function=_BGx3, units="G", take_log=False) def _Bx1(field, data): return data["gas", "magnetic_field_x"] + data["gas", "BGx1"] ds.add_field(("gas", "Bx1"), function=_Bx1, units="G", take_log=False) def _Bx2(field, data): return data["gas", "magnetic_field_y"] + data["gas", "BGx2"] ds.add_field(("gas", "Bx2"), function=_Bx2, units="G", take_log=False) def _Bx3(field, data): return data["gas", "magnetic_field_z"] + data["gas", "BGx3"] ds.add_field(("gas", "Bx3"), function=_Bx3, units="G", take_log=False) def _mag_energy(field, data): return (data["Bx1"]**2 + data["Bx2"]**2 + data["Bx3"]**2) / (8 * np.pi) ds.add_field(("gas", "mag_energy"), function=_mag_energy, units="g*cm**-1*s**-2", take_log=True) def _mag_field_strength(field, data): return np.sqrt(8. * np.pi * data["mag_energy"]) ds.add_field(("gas", "mag_field_strength"), function=_mag_field_strength, units="G", take_log=True) def _mag_field_magnitude(field, data): return np.sqrt(data["Bx1"]**2 + data["Bx2"]**2 + data["Bx3"]**2) ds.add_field(("gas", "mag_field_magnitude"), function=_mag_field_magnitude, units="G", take_log=True) def _plasma_b(field, data): return data['pressure'] / data['mag_energy'] ds.add_field(("gas", "plasma_b"), function=_plasma_b, units="") def _B_divergence(field, data): sl_right = slice(None, -2, None) sl_left = slice(2, None, None) div_fac = 2.0 ds = div_fac * just_one(data["index", "dx"]) f = data["Bx1"][sl_right, 1:-1, 1:-1] / ds f -= data["Bx1"][sl_left, 1:-1, 1:-1] / ds ds = div_fac * just_one(data["index", "dy"]) f += data["Bx2"][1:-1, sl_right, 1:-1] / ds f -= data["Bx2"][1:-1, sl_left, 1:-1] / ds ds = div_fac * just_one(data["index", "dz"]) f += data["Bx3"][1:-1, 1:-1, sl_right] / ds f -= data["Bx3"][1:-1, 1:-1, sl_left] / ds new_field = data.ds.arr(np.zeros(data["Bx1"].shape, dtype=np.float64), f.units) new_field[1:-1, 1:-1, 1:-1] = f return np.abs(new_field) ds.add_field(("gas", "B_divergence"), function=_B_divergence, units="G/code_length", validators=[ValidateSpatial(1)], take_log=True) def _divB_measure(field, data): return data["index", "dx"]*np.abs(data['B_divergence'])\ /data['mag_field_magnitude'] ds.add_field(("gas", "divB_measure"), function=_divB_measure, units="dimensionless", take_log=True) def _mag_alfven_speed(field, data): ftype = field.name[0] B = data[ftype, 'mag_field_strength'] return B / np.sqrt(4.0 * np.pi * data[ftype, 'density']) ds.add_field(("gas", "mag_alfven_speed"), function=_mag_alfven_speed, units="cm/s", take_log=True) def _mag_mach_alfven(field, data): ftype = field.name[0] return data[ftype, 'velocity_magnitude'] / data[ftype, 'mag_alfven_speed'] ds.add_field(("gas", "mag_mach_alfven"), function=_mag_mach_alfven, units="", take_log=True) def _fc(field, data): return YTQuantity(2.8, 'G**-1*MHz') * data["gas", "mag_field_strength"] ds.add_field(('gas', "fc"), function=_fc, units="MHz", take_log=True) def _fp(field, data): return YTQuantity(8.98e-3, 'cm**(3/2)*MHz')*\ np.sqrt(1.01*data["gas", "ni"]) ds.add_field(('gas', "fp"), function=_fp, units="MHz", take_log=True) def _f_ratio(field, data): return data['gas', 'fc'] / data['gas', 'fp'] ds.add_field(('gas', "fc/fp"), function=_f_ratio, units="", take_log=True) def _specific_angular_momentum_density_x(field, data): return data["gas", "specific_angular_momentum_x"] * data["gas", "density"] ds.add_field(('index', "specific_angular_momentum_density_x"), function=_specific_angular_momentum_density_x, units="cm**-1*g*s**-1", take_log=False, force_override=True) def _specific_angular_momentum_density_y(field, data): return data["gas", "specific_angular_momentum_y"] * data["gas", "density"] ds.add_field(('index', "specific_angular_momentum_density_y"), function=_specific_angular_momentum_density_y, units="cm**-1*g*s**-1", take_log=False, force_override=True) def _specific_angular_momentum_density_z(field, data): return data["gas", "specific_angular_momentum_z"] * data["gas", "density"] ds.add_field(('index', "specific_angular_momentum_density_z"), function=_specific_angular_momentum_density_z, units="cm**-1*g*s**-1", take_log=False, force_override=True) ds.periodicity = (True, True, True)
def setup_fluid_vector_fields(registry, ftype="gas", slice_info=None): # Current implementation for gradient is not valid for curvilinear geometries if is_curvilinear(registry.ds.geometry): return unit_system = registry.ds.unit_system # slice_info would be the left, the right, and the factor. # For example, with the old Enzo-ZEUS fields, this would be: # slice(None, -2, None) # slice(1, -1, None) # 1.0 # Otherwise, we default to a centered difference. if slice_info is None: sl_left = slice(None, -2, None) sl_right = slice(2, None, None) div_fac = 2.0 else: sl_left, sl_right, div_fac = slice_info sl_center = slice(1, -1, None) def _baroclinic_vorticity_x(field, data): rho2 = data[ftype, "density"].astype(np.float64)**2 return (data[ftype, "pressure_gradient_y"] * data[ftype, "density_gradient_z"] - data[ftype, "pressure_gradient_z"] * data[ftype, "density_gradient_z"]) / rho2 def _baroclinic_vorticity_y(field, data): rho2 = data[ftype, "density"].astype(np.float64)**2 return (data[ftype, "pressure_gradient_z"] * data[ftype, "density_gradient_x"] - data[ftype, "pressure_gradient_x"] * data[ftype, "density_gradient_z"]) / rho2 def _baroclinic_vorticity_z(field, data): rho2 = data[ftype, "density"].astype(np.float64)**2 return (data[ftype, "pressure_gradient_x"] * data[ftype, "density_gradient_y"] - data[ftype, "pressure_gradient_y"] * data[ftype, "density_gradient_x"]) / rho2 bv_validators = [ ValidateSpatial(1, [(ftype, "density"), (ftype, "pressure")]) ] for ax in 'xyz': n = "baroclinic_vorticity_%s" % ax registry.add_field((ftype, n), sampling_type="cell", function=eval("_%s" % n), validators=bv_validators, units=unit_system["frequency"]**2) create_magnitude_field(registry, "baroclinic_vorticity", unit_system["frequency"]**2, ftype=ftype, slice_info=slice_info, validators=bv_validators) def _vorticity_x(field, data): vz = data[ftype, "relative_velocity_z"] vy = data[ftype, "relative_velocity_y"] f = ((vz[sl_center, sl_right, sl_center] - vz[sl_center, sl_left, sl_center]) / (div_fac * just_one(data["index", "dy"]))) f -= ((vy[sl_center, sl_center, sl_right] - vy[sl_center, sl_center, sl_left]) / (div_fac * just_one(data["index", "dz"]))) new_field = data.ds.arr(np.zeros_like(vz, dtype=np.float64), f.units) new_field[sl_center, sl_center, sl_center] = f return new_field def _vorticity_y(field, data): vx = data[ftype, "relative_velocity_x"] vz = data[ftype, "relative_velocity_z"] f = ((vx[sl_center, sl_center, sl_right] - vx[sl_center, sl_center, sl_left]) / (div_fac * just_one(data["index", "dz"]))) f -= ((vz[sl_right, sl_center, sl_center] - vz[sl_left, sl_center, sl_center]) / (div_fac * just_one(data["index", "dx"]))) new_field = data.ds.arr(np.zeros_like(vz, dtype=np.float64), f.units) new_field[sl_center, sl_center, sl_center] = f return new_field def _vorticity_z(field, data): vx = data[ftype, "relative_velocity_x"] vy = data[ftype, "relative_velocity_y"] f = ((vy[sl_right, sl_center, sl_center] - vx[sl_left, sl_center, sl_center]) / (div_fac * just_one(data["index", "dx"]))) f -= ((vx[sl_center, sl_right, sl_center] - vx[sl_center, sl_left, sl_center]) / (div_fac * just_one(data["index", "dy"]))) new_field = data.ds.arr(np.zeros_like(vy, dtype=np.float64), f.units) new_field[sl_center, sl_center, sl_center] = f return new_field vort_validators = [ ValidateSpatial(1, [(ftype, "velocity_%s" % d) for d in 'xyz']), ValidateParameter('bulk_velocity') ] for ax in 'xyz': n = "vorticity_%s" % ax registry.add_field((ftype, n), sampling_type="cell", function=eval("_%s" % n), units=unit_system["frequency"], validators=vort_validators) create_magnitude_field(registry, "vorticity", unit_system["frequency"], ftype=ftype, slice_info=slice_info, validators=vort_validators) create_squared_field(registry, "vorticity", unit_system["frequency"]**2, ftype=ftype, slice_info=slice_info, validators=vort_validators) def _vorticity_stretching_x(field, data): return data[ftype, "velocity_divergence"] * data[ftype, "vorticity_x"] def _vorticity_stretching_y(field, data): return data[ftype, "velocity_divergence"] * data[ftype, "vorticity_y"] def _vorticity_stretching_z(field, data): return data[ftype, "velocity_divergence"] * data[ftype, "vorticity_z"] for ax in 'xyz': n = "vorticity_stretching_%s" % ax registry.add_field((ftype, n), sampling_type="cell", function=eval("_%s" % n), units=unit_system["frequency"]**2, validators=vort_validators) create_magnitude_field(registry, "vorticity_stretching", unit_system["frequency"]**2, ftype=ftype, slice_info=slice_info, validators=vort_validators) def _vorticity_growth_x(field, data): return -data[ftype, "vorticity_stretching_x"] - \ data[ftype, "baroclinic_vorticity_x"] def _vorticity_growth_y(field, data): return -data[ftype, "vorticity_stretching_y"] - \ data[ftype, "baroclinic_vorticity_y"] def _vorticity_growth_z(field, data): return -data[ftype, "vorticity_stretching_z"] - \ data[ftype, "baroclinic_vorticity_z"] for ax in 'xyz': n = "vorticity_growth_%s" % ax registry.add_field((ftype, n), sampling_type="cell", function=eval("_%s" % n), units=unit_system["frequency"]**2, validators=vort_validators) def _vorticity_growth_magnitude(field, data): result = np.sqrt(data[ftype, "vorticity_growth_x"]**2 + data[ftype, "vorticity_growth_y"]**2 + data[ftype, "vorticity_growth_z"]**2) dot = data.ds.arr(np.zeros(result.shape), "") for ax in "xyz": dot += (data[ftype, "vorticity_%s" % ax] * data[ftype, "vorticity_growth_%s" % ax]).to_ndarray() result = np.sign(dot) * result return result registry.add_field((ftype, "vorticity_growth_magnitude"), sampling_type="cell", function=_vorticity_growth_magnitude, units=unit_system["frequency"]**2, validators=vort_validators, take_log=False) def _vorticity_growth_magnitude_absolute(field, data): return np.sqrt(data[ftype, "vorticity_growth_x"]**2 + data[ftype, "vorticity_growth_y"]**2 + data[ftype, "vorticity_growth_z"]**2) registry.add_field((ftype, "vorticity_growth_magnitude_absolute"), sampling_type="cell", function=_vorticity_growth_magnitude_absolute, units=unit_system["frequency"]**2, validators=vort_validators) def _vorticity_growth_timescale(field, data): domegax_dt = data[ftype, "vorticity_x"] / data[ftype, "vorticity_growth_x"] domegay_dt = data[ftype, "vorticity_y"] / data[ftype, "vorticity_growth_y"] domegaz_dt = data[ftype, "vorticity_z"] / data[ftype, "vorticity_growth_z"] return np.sqrt(domegax_dt**2 + domegay_dt**2 + domegaz_dt**2) registry.add_field((ftype, "vorticity_growth_timescale"), sampling_type="cell", function=_vorticity_growth_timescale, units=unit_system["time"], validators=vort_validators) ######################################################################## # With radiation pressure ######################################################################## def _vorticity_radiation_pressure_x(field, data): rho = data[ftype, "density"].astype(np.float64) return (data[ftype, "radiation_acceleration_y"] * data[ftype, "density_gradient_z"] - data[ftype, "radiation_acceleration_z"] * data[ftype, "density_gradient_y"]) / rho def _vorticity_radiation_pressure_y(field, data): rho = data[ftype, "density"].astype(np.float64) return (data[ftype, "radiation_acceleration_z"] * data[ftype, "density_gradient_x"] - data[ftype, "radiation_acceleration_x"] * data[ftype, "density_gradient_z"]) / rho def _vorticity_radiation_pressure_z(field, data): rho = data[ftype, "density"].astype(np.float64) return (data[ftype, "radiation_acceleration_x"] * data[ftype, "density_gradient_y"] - data[ftype, "radiation_acceleration_y"] * data[ftype, "density_gradient_x"]) / rho vrp_validators = [ ValidateSpatial(1, [(ftype, "density"), (ftype, "radiation_acceleration_x"), (ftype, "radiation_acceleration_y"), (ftype, "radiation_acceleration_z")]) ] for ax in 'xyz': n = "vorticity_radiation_pressure_%s" % ax registry.add_field((ftype, n), sampling_type="cell", function=eval("_%s" % n), units=unit_system["frequency"]**2, validators=vrp_validators) create_magnitude_field(registry, "vorticity_radiation_pressure", unit_system["frequency"]**2, ftype=ftype, slice_info=slice_info, validators=vrp_validators) def _vorticity_radiation_pressure_growth_x(field, data): return -data[ftype, "vorticity_stretching_x"] - \ data[ftype, "baroclinic_vorticity_x"] \ -data[ftype, "vorticity_radiation_pressure_x"] def _vorticity_radiation_pressure_growth_y(field, data): return -data[ftype, "vorticity_stretching_y"] - \ data[ftype, "baroclinic_vorticity_y"] \ -data[ftype, "vorticity_radiation_pressure_y"] def _vorticity_radiation_pressure_growth_z(field, data): return -data[ftype, "vorticity_stretching_z"] - \ data[ftype, "baroclinic_vorticity_z"] \ -data[ftype, "vorticity_radiation_pressure_z"] for ax in 'xyz': n = "vorticity_radiation_pressure_growth_%s" % ax registry.add_field((ftype, n), sampling_type="cell", function=eval("_%s" % n), units=unit_system["frequency"]**2, validators=vrp_validators) def _vorticity_radiation_pressure_growth_magnitude(field, data): result = np.sqrt( data[ftype, "vorticity_radiation_pressure_growth_x"]**2 + data[ftype, "vorticity_radiation_pressure_growth_y"]**2 + data[ftype, "vorticity_radiation_pressure_growth_z"]**2) dot = data.ds.arr(np.zeros(result.shape), "") for ax in "xyz": dot += (data[ftype, "vorticity_%s" % ax] * data[ftype, "vorticity_growth_%s" % ax]).to_ndarray() result = np.sign(dot) * result return result registry.add_field( (ftype, "vorticity_radiation_pressure_growth_magnitude"), sampling_type="cell", function=_vorticity_radiation_pressure_growth_magnitude, units=unit_system["frequency"]**2, validators=vrp_validators, take_log=False) def _vorticity_radiation_pressure_growth_magnitude_absolute(field, data): return np.sqrt( data[ftype, "vorticity_radiation_pressure_growth_x"]**2 + data[ftype, "vorticity_radiation_pressure_growth_y"]**2 + data[ftype, "vorticity_radiation_pressure_growth_z"]**2) registry.add_field( (ftype, "vorticity_radiation_pressure_growth_magnitude_absolute"), sampling_type="cell", function=_vorticity_radiation_pressure_growth_magnitude_absolute, units="s**(-2)", validators=vrp_validators) def _vorticity_radiation_pressure_growth_timescale(field, data): domegax_dt = data[ftype, "vorticity_x"] / \ data[ftype, "vorticity_radiation_pressure_growth_x"] domegay_dt = data[ftype, "vorticity_y"] / \ data[ftype, "vorticity_radiation_pressure_growth_y"] domegaz_dt = data[ftype, "vorticity_z"] / \ data[ftype, "vorticity_radiation_pressure_growth_z"] return np.sqrt(domegax_dt**2 + domegay_dt**2 + domegaz_dt**2) registry.add_field( (ftype, "vorticity_radiation_pressure_growth_timescale"), sampling_type="cell", function=_vorticity_radiation_pressure_growth_timescale, units=unit_system["time"], validators=vrp_validators) def _shear(field, data): """ Shear is defined as [(dvx/dy + dvy/dx)^2 + (dvz/dy + dvy/dz)^2 + (dvx/dz + dvz/dx)^2 ]^(0.5) where dvx/dy = [vx(j-1) - vx(j+1)]/[2dy] and is in units of s^(-1) (it's just like vorticity except add the derivative pairs instead of subtracting them) """ if data.ds.dimensionality > 1: vx = data[ftype, "relative_velocity_x"] vy = data[ftype, "relative_velocity_y"] dvydx = ((vy[sl_right, sl_center, sl_center] - vy[sl_left, sl_center, sl_center]) / (div_fac * just_one(data["index", "dx"]))) dvxdy = ((vx[sl_center, sl_right, sl_center] - vx[sl_center, sl_left, sl_center]) / (div_fac * just_one(data["index", "dy"]))) f = (dvydx + dvxdy)**2.0 del dvydx, dvxdy if data.ds.dimensionality > 2: vz = data[ftype, "relative_velocity_z"] dvzdy = ((vz[sl_center, sl_right, sl_center] - vz[sl_center, sl_left, sl_center]) / (div_fac * just_one(data["index", "dy"]))) dvydz = ((vy[sl_center, sl_center, sl_right] - vy[sl_center, sl_center, sl_left]) / (div_fac * just_one(data["index", "dz"]))) f += (dvzdy + dvydz)**2.0 del dvzdy, dvydz dvxdz = ((vx[sl_center, sl_center, sl_right] - vx[sl_center, sl_center, sl_left]) / (div_fac * just_one(data["index", "dz"]))) dvzdx = ((vz[sl_right, sl_center, sl_center] - vz[sl_left, sl_center, sl_center]) / (div_fac * just_one(data["index", "dx"]))) f += (dvxdz + dvzdx)**2.0 del dvxdz, dvzdx np.sqrt(f, out=f) new_field = data.ds.arr(np.zeros_like(data[ftype, "velocity_x"]), f.units) new_field[sl_center, sl_center, sl_center] = f return new_field registry.add_field((ftype, "shear"), sampling_type="cell", function=_shear, validators=[ ValidateSpatial(1, [(ftype, "velocity_x"), (ftype, "velocity_y"), (ftype, "velocity_z")]), ValidateParameter('bulk_velocity') ], units=unit_system["frequency"]) def _shear_criterion(field, data): """ Divide by c_s to leave shear in units of length**-1, which can be compared against the inverse of the local cell size (1/dx) to determine if refinement should occur. """ return data[ftype, "shear"] / data[ftype, "sound_speed"] registry.add_field((ftype, "shear_criterion"), sampling_type="cell", function=_shear_criterion, units=unit_system["length"]**-1, validators=[ ValidateSpatial(1, [(ftype, "sound_speed"), (ftype, "velocity_x"), (ftype, "velocity_y"), (ftype, "velocity_z")]) ]) def _shear_mach(field, data): """ Dimensionless shear (shear_mach) is defined nearly the same as shear, except that it is scaled by the local dx/dy/dz and the local sound speed. So it results in a unitless quantity that is effectively measuring shear in mach number. In order to avoid discontinuities created by multiplying by dx/dy/dz at grid refinement boundaries, we also multiply by 2**GridLevel. Shear (Mach) = [(dvx + dvy)^2 + (dvz + dvy)^2 + (dvx + dvz)^2 ]^(0.5) / c_sound """ if data.ds.dimensionality > 1: vx = data[ftype, "relative_velocity_x"] vy = data[ftype, "relative_velocity_y"] dvydx = (vy[sl_right, sl_center, sl_center] - vy[sl_left, sl_center, sl_center]) / div_fac dvxdy = (vx[sl_center, sl_right, sl_center] - vx[sl_center, sl_left, sl_center]) / div_fac f = (dvydx + dvxdy)**2.0 del dvydx, dvxdy if data.ds.dimensionality > 2: vz = data[ftype, "relative_velocity_z"] dvzdy = (vz[sl_center, sl_right, sl_center] - vz[sl_center, sl_left, sl_center]) / div_fac dvydz = (vy[sl_center, sl_center, sl_right] - vy[sl_center, sl_center, sl_left]) / div_fac f += (dvzdy + dvydz)**2.0 del dvzdy, dvydz dvxdz = (vx[sl_center, sl_center, sl_right] - vx[sl_center, sl_center, sl_left]) / div_fac dvzdx = (vz[sl_right, sl_center, sl_center] - vz[sl_left, sl_center, sl_center]) / div_fac f += (dvxdz + dvzdx)**2.0 del dvxdz, dvzdx f *= ( 2.0**data["index", "grid_level"][sl_center, sl_center, sl_center] / data[ftype, "sound_speed"][sl_center, sl_center, sl_center])**2.0 np.sqrt(f, out=f) new_field = data.ds.arr(np.zeros_like(vx), f.units) new_field[sl_center, sl_center, sl_center] = f return new_field vs_fields = [(ftype, "sound_speed"), (ftype, "velocity_x"), (ftype, "velocity_y"), (ftype, "velocity_z")] registry.add_field((ftype, "shear_mach"), sampling_type="cell", function=_shear_mach, units="", validators=[ ValidateSpatial(1, vs_fields), ValidateParameter('bulk_velocity') ])
def particle_deposition_functions(ptype, coord_name, mass_name, registry): unit_system = registry.ds.unit_system orig = set(registry.keys()) ptype_dn = ptype.replace("_", " ").title() def particle_count(field, data): pos = data[ptype, coord_name] d = data.deposit(pos, method="count") d = data.ds.arr(d, input_units="cm**-3") return data.apply_units(d, field.units) registry.add_field(("deposit", "%s_count" % ptype), function=particle_count, validators=[ValidateSpatial()], units='', display_name=r"\mathrm{%s Count}" % ptype_dn) def particle_mass(field, data): pos = data[ptype, coord_name] pmass = data[ptype, mass_name] pmass.convert_to_units(field.units) d = data.deposit(pos, [pmass], method="sum") return data.apply_units(d, field.units) registry.add_field(("deposit", "%s_mass" % ptype), function=particle_mass, validators=[ValidateSpatial()], display_name=r"\mathrm{%s Mass}" % ptype_dn, units=unit_system["mass"]) def particle_density(field, data): pos = data[ptype, coord_name].convert_to_units("code_length") mass = data[ptype, mass_name].convert_to_units("code_mass") d = data.deposit(pos, [mass], method="sum") d = data.ds.arr(d, "code_mass") d /= data["index", "cell_volume"] return d registry.add_field(("deposit", "%s_density" % ptype), function=particle_density, validators=[ValidateSpatial()], display_name=r"\mathrm{%s Density}" % ptype_dn, units=unit_system["density"]) def particle_cic(field, data): pos = data[ptype, coord_name] d = data.deposit(pos, [data[ptype, mass_name]], method="cic") d = data.apply_units(d, data[ptype, mass_name].units) d /= data["index", "cell_volume"] return d registry.add_field(("deposit", "%s_cic" % ptype), function=particle_cic, validators=[ValidateSpatial()], display_name=r"\mathrm{%s CIC Density}" % ptype_dn, units=unit_system["density"]) def _get_density_weighted_deposit_field(fname, units, method): def _deposit_field(field, data): """ Create a grid field for particle quantities weighted by particle mass, using cloud-in-cell deposit. """ pos = data[ptype, "particle_position"] # Get back into density pden = data[ptype, 'particle_mass'] top = data.deposit(pos, [data[(ptype, fname)] * pden], method=method) bottom = data.deposit(pos, [pden], method=method) top[bottom == 0] = 0.0 bnz = bottom.nonzero() top[bnz] /= bottom[bnz] d = data.ds.arr(top, input_units=units) return d return _deposit_field for ax in 'xyz': for method, name in zip(("cic", "sum"), ("cic", "nn")): function = _get_density_weighted_deposit_field( "particle_velocity_%s" % ax, "cm/s", method) registry.add_field( ("deposit", ("%s_" + name + "_velocity_%s") % (ptype, ax)), function=function, units=unit_system["velocity"], take_log=False, validators=[ValidateSpatial(0)]) for method, name in zip(("cic", "sum"), ("cic", "nn")): function = _get_density_weighted_deposit_field("age", "s", method) registry.add_field(("deposit", ("%s_" + name + "_age") % (ptype)), function=function, units=unit_system["time"], take_log=False, validators=[ValidateSpatial(0)]) # Now some translation functions. def particle_ones(field, data): v = np.ones(data[ptype, mass_name].shape, dtype="float64") return data.apply_units(v, field.units) registry.add_field((ptype, "particle_ones"), function=particle_ones, particle_type=True, units="", display_name=r"Particle Count") def particle_mesh_ids(field, data): pos = data[ptype, coord_name] ids = np.zeros(pos.shape[0], dtype="float64") - 1 # This is float64 in name only. It will be properly cast inside the # deposit operation. #_ids = ids.view("float64") data.deposit(pos, [ids], method="mesh_id") return data.apply_units(ids, "") registry.add_field((ptype, "mesh_id"), function=particle_mesh_ids, validators=[ValidateSpatial()], units='', particle_type=True) return list(set(registry.keys()).difference(orig))
import fastcluster import scipy.cluster.hierarchy as sch from mpi4py import MPI import os, sys from astropy.table import Table, Column, vstack, hstack from common_functions import * sl_left = slice(None, -2, None) sl_right = slice(2, None, None) div_fac = 2.0 sl_center = slice(1, -1, None) ftype = 'gas' vort_validators = [ ValidateSpatial(1, [(ftype, "velocity_x"), (ftype, "velocity_y"), (ftype, "velocity_z")]) ] def mad_based_outlier(points, thresh=5): if len(points.shape) == 1: points = points[:, None] median = np.median(points, axis=0) diff = np.sum((points - median)**2, axis=-1) diff = np.sqrt(diff) med_abs_deviation = np.median(diff) modified_z_score = 0.6745 * diff / med_abs_deviation return modified_z_score > thresh