def density_error_norm(snap, drag_coefficients, x_shock, xrange, n_bins=50):
n_dust = len(drag_coefficients)
d_gas, d_dusts = density_exact(
x_shock=x_shock,
x_width=X_WIDTH_EXACT,
n_dust=n_dust,
dust_to_gas=DUST_TO_GAS,
drag_coefficient=drag_coefficients,
density_left=DENSITY_LEFT,
velocity_left=VELOCITY_LEFT,
mach_number=MACH_NUMBER,
)
d_exact = [d_gas] + d_dusts
subsnaps = [snap['gas']] + snap['dust']
error_squared = 0.0
for idx, subsnap in enumerate(subsnaps):
prof = plonk.load_profile(
snap=subsnap,
radius_min=xrange[0],
radius_max=xrange[1],
ndim=1,
n_bins=n_bins,
)
x, d_numerical = prof['radius'], prof['velocity_x']
error_squared += np.sum((d_exact[idx](x) - d_numerical)**2)
return np.sqrt(error_squared)
def test_check_data(snaptype):
"""Test Profile data accuracy."""
filename = DIR / snaptype.filename
snap = plonk.load_snap(filename)
prof = plonk.load_profile(snap)
columns = [
'density',
'mass',
'number',
'radius',
'scale_height',
'size',
'smoothing_length',
'sound_speed',
]
df = prof.to_dataframe(columns=columns)
units = ['g/cm^3', 'g', '', 'au', 'au', 'au^2', 'au', 'km/s']
df = prof.to_dataframe(columns=columns, units=units)
profile_file = DIR / snaptype.profile_file
pd.testing.assert_frame_equal(df, pd.read_csv(profile_file, index_col=0))
snap.close_file()
def test_to_function(snaptype):
"""Test to_function."""
filename = DIR / snaptype.filename
snap = plonk.load_snap(filename)
prof = plonk.load_profile(snap=snap)
fn = prof.to_function(profile='scale_height')
assert np.allclose(fn(prof['radius']), prof['scale_height'])
snap.close_file()
def test_alias(snaptype):
"""Test using profile aliases."""
filename = DIR / snaptype.filename
snap = plonk.load_snap(filename)
prof = plonk.load_profile(snap=snap)
prof.add_alias('scale_height', 'H')
prof['H']
snap.close_file()
def test_set_data(snaptype):
"""Test setting array on Profile."""
filename = DIR / snaptype.filename
snap = plonk.load_snap(filename)
prof = plonk.load_profile(snap=snap)
prof['array'] = np.arange(len(prof)) * plonk.units.au
with pytest.raises(ValueError):
prof['array'] = np.arange(len(prof) - 1)
with pytest.raises(ValueError):
prof['array'] = 1.0
snap.close_file()
def plot_quantity_profile_subsnaps(snap, quantity, ax, xrange, n_bins):
subsnaps = [snap['gas']] + snap['dust']
for idx, subsnap in enumerate(subsnaps):
label = 'Gas' if idx == 0 else f'Dust {idx}'
prof = plonk.load_profile(
snap=subsnap,
radius_min=xrange[0],
radius_max=xrange[1],
ndim=1,
n_bins=n_bins,
)
x, y = prof['radius'], prof[quantity]
ax.plot(x, y, 'o', ms=4, label=label, fillstyle='none')
ax.grid(b=True)
ax.set(xlim=xrange)
return ax
def test_profile_plot(snaptype):
"""Test loading Profile."""
filename = DIR / snaptype.filename
snap = plonk.load_snap(filename)
prof = plonk.load_profile(snap=snap)
prof.plot(x='radius', y='surface_density')
prof.plot(
x='radius',
y='density',
units={
'position': 'au',
'density': 'g/cm^3'
},
std='shading',
)
snap.close_file()
def test_animation_profiles():
"""Test animation of profiles."""
sim = plonk.load_simulation(prefix=PREFIX, directory=DIR_PATH)
snaps = [sim.snaps[0], sim.snaps[0], sim.snaps[0]]
profiles = [plonk.load_profile(snap) for snap in snaps]
filename = Path('animation.mp4')
visualize.animation_profiles(
filename=filename,
profiles=profiles,
x='radius',
y='surface_density',
units={
'position': 'au',
'surface_density': 'g/cm^2'
},
)
filename.unlink()
def density_error(snap,
drag_coefficients,
x_shock,
xrange,
error_type='absolute',
n_bins=50):
if error_type not in ['absolute', 'relative']:
raise ValueError('Wrong error type: must be "absolute" or "relative"')
n_dust = len(drag_coefficients)
d_gas, d_dusts = density_exact(
x_shock=x_shock,
x_width=X_WIDTH_EXACT,
n_dust=n_dust,
dust_to_gas=DUST_TO_GAS,
drag_coefficient=drag_coefficients,
density_left=DENSITY_LEFT,
velocity_left=VELOCITY_LEFT,
mach_number=MACH_NUMBER,
)
d_exact = [d_gas] + d_dusts
subsnaps = [snap['gas']] + snap['dust']
error = list()
for idx, subsnap in enumerate(subsnaps):
prof = plonk.load_profile(
snap=subsnap,
radius_min=xrange[0],
radius_max=xrange[1],
ndim=1,
n_bins=n_bins,
)
x, d_numerical = prof['radius'], prof['density']
if error_type == 'absolute':
error.append(np.abs(d_exact[idx](x) - d_numerical))
elif error_type == 'relative':
error.append(
np.abs((d_exact[idx](x) - d_numerical) / d_exact[idx](x)))
return x, error[0], error[1:]
def calculate_profiles(
snap: Snap,
radius_min: Quantity,
radius_max: Quantity,
scale_height_fac: float,
n_bins: int = 50,
) -> Dict[str, List[Profile]]:
"""Calculate radial drift velocity profiles.
Parameters
----------
snap
The Snap object.
radius_min
The minimum radius for the profiles.
radius_max
The maximum radius for the profiles.
scale_height_fac
A factor of the scale height within which to average over.
n_bins
The number of bins in the profile. Default is 50.
Returns
-------
Dict
A dictionary of list of profiles. The keys are 'gas' and 'dust'
and the values are lists of profiles, one per sub-type.
"""
print('Calculating profiles...')
snap.add_quantities('disc')
snap.set_gravitational_parameter(0)
gamma = snap.properties['adiabatic_index']
num_dust = snap.num_dust_species
# Use particles in the midplane only
# Choose particles such that they are within a factor of the gas scale height
gas = snap.family('gas')
prof = plonk.load_profile(snap=gas, cmin=radius_min, cmax=radius_max)
scale_height = prof.to_function('scale_height')
snap_midplane = snap[np.abs(snap['z']) < scale_height_fac * scale_height(snap['R'])]
subsnaps = snap_midplane.subsnaps_as_dict()
# Create radial profiles for the gas and each dust species
cmin, cmax = radius_min, radius_max
profs: Dict[str, List[Profile]] = {'gas': list(), 'dust': list()}
profs['gas'] = [
plonk.load_profile(subsnaps['gas'], cmin=cmin, cmax=cmax, n_bins=n_bins)
]
for subsnap in subsnaps['dust']:
profs['dust'].append(
plonk.load_profile(subsnap, cmin=cmin, cmax=cmax, n_bins=n_bins)
)
p = profs['gas'][0]
# velocity_pressure is (15) in Dipierro+2018
p['velocity_pressure'] = np.gradient(p['pressure'], p['radius']) / (
p['density'] * p['keplerian_frequency']
)
# shear_viscosity is between (16) an (17) in Dipierro+2018
p['shear_viscosity'] = p['disc_viscosity'] * p['density']
# velocity_visc is (16) in Dipierro+2018
p['velocity_visc'] = np.gradient(
p['shear_viscosity']
* p['radius'] ** 3
* np.gradient(p['keplerian_frequency'], p['radius']),
p['radius'],
) / (
p['radius']
* p['density']
* np.gradient(p['radius'] ** 2 * p['keplerian_frequency'], p['radius'])
)
for idx, prof_dust in enumerate(profs['dust']):
p[f'midplane_dust_to_gas_{idx+1:03}'] = prof_dust['density'] / p['density']
p[f'_midplane_stokes_number_{idx+1:03}'] = (
np.sqrt(np.pi * gamma / 8)
* snap.properties['grain_density'][idx]
* snap.properties['grain_size'][idx]
* p['keplerian_frequency']
/ (p['density'] * p['sound_speed'])
)
# lambda_0 and lambda_1 are (17) in Dipierro+2018
l0 = np.zeros(len(p)) * plonk.units['dimensionless']
l1 = np.zeros(len(p)) * plonk.units['dimensionless']
for idx in range(num_dust):
St = p[f'_midplane_stokes_number_{idx+1:03}']
eps = p[f'midplane_dust_to_gas_{idx+1:03}']
l0 = l0 + 1 / (1 + St ** 2) * eps
l1 = l1 + St / (1 + St ** 2) * eps
p['lambda_0'] = l0
p['lambda_1'] = l1
v_P = p['velocity_pressure']
v_visc = p['velocity_visc']
l0 = p['lambda_0']
l1 = p['lambda_1']
# velocity_radial_gas is (11) in Dipierro+2018
p['gas_velocity_radial'] = (-l1 * v_P + (1 + l0) * v_visc) / (
(1 + l0) ** 2 + l1 ** 2
)
# velocity_azimuthal_gas is (12) in Dipierro+2018
p['gas_velocity_azimuthal'] = (
1 / 2 * (v_P * (1 + l0) + v_visc * l1) / ((1 + l0) ** 2 + l1 ** 2)
)
# velocity_radial_dust_i is (13) in Dipierro+2018
# velocity_azimuthal_dust_i is (14) in Dipierro+2018
for idx in range(num_dust):
St = p[f'_midplane_stokes_number_{idx+1:03}']
eps = p[f'midplane_dust_to_gas_{idx+1:03}']
numerator_R = v_P * ((1 + l0) * St - l1) + v_visc * (1 + l0 + St * l1)
numerator_phi = 0.5 * v_P * (1 + l0 + St * l1) - v_visc * ((1 + l0) * St - l1)
denominator = ((1 + l0) ** 2 + l1 ** 2) * (1 + St ** 2)
p[f'dust_velocity_radial_{idx+1:03}'] = numerator_R / denominator
p[f'dust_velocity_azimuthal_{idx+1:03}'] = numerator_phi / denominator
# Divide by |v_P| for comparison with Figure B1 in Dipierro+2018
# "Analytical" solution
v_R = p['gas_velocity_radial']
p['gas_velocity_radial_analytical'] = v_R / np.abs(v_P)
for idx in range(num_dust):
v_R = p[f'dust_velocity_radial_{idx+1:03}']
p[f'dust_velocity_radial_analytical_{idx+1:03}'] = v_R / np.abs(v_P)
# "Numerical" solution
v_R = p['velocity_radial_cylindrical']
v_R_std = p['velocity_radial_cylindrical_std']
p['velocity_radial_numerical'] = v_R / np.abs(v_P)
p['velocity_radial_numerical_std'] = v_R_std / np.abs(v_P)
for prof in profs['dust']:
v_R = prof['velocity_radial_cylindrical']
v_R_std = prof['velocity_radial_cylindrical_std']
prof['velocity_radial_numerical'] = v_R / np.abs(v_P)
prof['velocity_radial_numerical_std'] = v_R_std / np.abs(v_P)
return profs
def test_load_profile(snaptype):
"""Test loading Profile."""
filename = DIR / snaptype.filename
snap = plonk.load_snap(filename)
prof = plonk.load_profile(snap=snap)
for p in [
'aspect_ratio', 'angular_momentum_phi', 'angular_momentum_theta'
]:
prof[p]
with pytest.raises(ValueError):
prof['does_not_exist']
plonk.load_profile(snap=snap,
ndim=3,
cmin='10 au',
cmax='100 au',
n_bins=30)
plonk.load_profile(snap=snap, spacing='log', ignore_accreted=False)
plonk.load_profile(snap=snap, ndim=1, coordinate='x')
plonk.load_profile(snap=snap, ndim=1, coordinate='y')
plonk.load_profile(snap=snap, ndim=1, coordinate='z')
with pytest.raises(ValueError):
plonk.load_profile(snap=snap, ndim=1, coordinate='does_not_exist')
with pytest.raises(ValueError):
plonk.load_profile(snap=snap, cmin=10, cmax=100)
snap.close_file()