def test_load_phantom_snap():
"""Testing reading Phantom HDF5 snapshots."""
# Read from pathlib.Path
snap = plonk.load_snap(TEST_FILE)
snap.close_file()
# Read from str
snap = plonk.load_snap(str(TEST_FILE))
snap.close_file()
# Not exists
with pytest.raises(FileNotFoundError):
plonk.load_snap('does_not_exist.h5')
def test_load_phantom_snap(snaptype):
"""Testing reading Phantom HDF5 snapshots."""
# Read from Path
filename = DIR / snaptype.filename
snap = plonk.load_snap(filename)
snap.close_file()
# Read from str
snap = plonk.load_snap(str(filename))
snap.close_file()
# Not exists
with pytest.raises(FileNotFoundError):
plonk.load_snap('does_not_exist.h5')
def test_plot(snaptype):
"""Test particle plot."""
filename = DIR / snaptype.filename
snap = plonk.load_snap(filename)
plonk.plot(snap=snap)
snap.close_file()
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()
def test_image_projection(snaptype):
"""Test image projection."""
filename = DIR / snaptype.filename
snap = plonk.load_snap(filename)
plonk.image(snap=snap, quantity='density', num_pixels=(32, 32))
snap.close_file()
def test_bulk_load(snaptype):
"""Testing bulk loading arrays."""
filename = DIR / snaptype.filename
snap = plonk.load_snap(filename)
snap.bulk_load()
snap.close_file()
def test_image_on_snap(snaptype):
"""Test image projection as method on Snap."""
filename = DIR / snaptype.filename
snap = plonk.load_snap(filename)
snap.image(quantity='density', num_pixels=(32, 32))
snap.close_file()
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_read_properties_from_phantom():
"""Testing reading Phantom HDF5 snapshot properties."""
snap = plonk.load_snap(TEST_FILE)
for key, value in properties.items():
np.testing.assert_allclose(snap.properties[key], value)
snap.close_file()
def test_get_extent(snaptype):
"""Test getting extent from percentile."""
filename = DIR / snaptype.filename
snap = plonk.load_snap(filename)
visualize.get_extent_from_percentile(snap=snap, x='x', y='y')
snap.close_file()
def test_plot_smoothing_length(snaptype):
"""Test plot smoothing length as circle."""
filename = DIR / snaptype.filename
snap = plonk.load_snap(filename)
visualize.plot_smoothing_length(snap=snap, indices=[0, 1])
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_read_particle_arrays_from_phantom():
"""Testing reading Phantom HDF5 snapshot particle arrays."""
snap = plonk.load_snap(TEST_FILE)
for array in mean_array_values.keys():
np.testing.assert_allclose(snap[array_name_map[array]].mean(),
mean_array_values[array])
snap.close_file()
def test_read_write_extra(snaptype):
"""Testing read write extra arrays."""
filename = DIR / snaptype.filename
snap = plonk.load_snap(filename)
_filename = Path('tmp.h5')
arr = np.arange(len(snap)) * plonk.units('dimensionless')
snap['my_array'] = arr
snap.write_extra_arrays(arrays=['my_array'], filename=_filename)
snap = None
snap = plonk.load_snap(filename)
snap.read_extra_arrays(filename=_filename)
np.testing.assert_allclose(snap['my_array'], arr, rtol=RTOL)
_filename.unlink()
snap.close_file()
def test_image_slice(snaptype):
"""Test image slice."""
filename = DIR / snaptype.filename
snap = plonk.load_snap(filename)
plonk.image(snap=snap,
quantity='density',
interp='slice',
num_pixels=(32, 32))
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_particles(snaptype):
"""Test particles analysis functions."""
filename = DIR / snaptype.filename
snap = plonk.load_snap(filename)
snap.set_molecular_weight(2.381)
_test_particles(snap=snap, ignore=False)
_test_particles(snap=snap, ignore=True)
snap.close_file()
def test_read_particle_arrays_from_phantom(snaptype):
"""Testing reading Phantom HDF5 snapshot particle arrays."""
filename = DIR / snaptype.filename
snap = plonk.load_snap(filename)
_check_arrays(
snap,
snaptype.array_name_map,
snaptype.mean_array_values,
snaptype.std_array_values,
)
snap.close_file()
def test_write_to_dataframe(snaptype):
"""Testing writing Snap to DataFrame."""
filename = DIR / snaptype.filename
snap = plonk.load_snap(filename)
columns = ['position', 'density', 'smoothing_length']
snap.to_dataframe(columns=columns)
columns = ['position', 'density', 'smoothing_length']
snap.to_dataframe(columns=columns, units=['au', 'g/cm^3', 'au'])
snap.close_file()
def test_sinks(snaptype):
"""Testing getting sink particles."""
filename = DIR / snaptype.filename
snap = plonk.load_snap(filename)
if snap.num_sinks > 0:
sinks = snap.sinks
assert snap.num_sinks == snaptype.num_sinks
assert len(sinks) == snaptype.num_sinks
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 test_sph_derivative():
"""Test SPH derivative."""
filename = DIR / dustseparate.filename
snap = plonk.load_snap(filename)
analysis.sph.derivative(
snap=snap,
derivative='div',
quantity='velocity',
kernel='cubic',
)
snap.close_file()
def test_set_array(snaptype):
"""Testing setting array on particles."""
filename = DIR / snaptype.filename
snap = plonk.load_snap(filename)
particle_array = np.arange(len(snap)) * plonk.units('dimensionless')
snap['array'] = particle_array
np.testing.assert_allclose(snap['array'].m, particle_array.m, rtol=RTOL)
sink_array = np.arange(len(snap.sinks)) * plonk.units('dimensionless')
snap.sinks['array'] = sink_array
np.testing.assert_allclose(snap.sinks['array'].m, sink_array.m, rtol=RTOL)
snap.close_file()
def test_array_code_unit(snaptype):
"""Testing getting array code unit."""
filename = DIR / snaptype.filename
snap = plonk.load_snap(filename)
position_unit = snaptype.length_unit * plonk.units('meter')
assert snap.array_code_unit('position') == position_unit
for arr in ['position', 'position_x', 'x']:
snap.array_code_unit(arr)
with pytest.raises(ValueError):
snap.array_code_unit('does_not_exist')
snap.close_file()
def test_total(snaptype):
"""Test total analysis functions."""
filename = DIR / snaptype.filename
snap = plonk.load_snap(filename)
total.accreted_mass(snap=snap)
total.angular_momentum(snap=snap)
total.center_of_mass(snap=snap)
total.kinetic_energy(snap=snap)
total.mass(snap=snap)
total.momentum(snap=snap)
total.specific_angular_momentum(snap=snap)
total.specific_kinetic_energy(snap=snap)
snap.close_file()
def test_plot_as_methods(snaptype):
"""Testing plot methods."""
filename = DIR / snaptype.filename
snap = plonk.load_snap(filename)
snap.image('density', num_pixels=(16, 16))
snap.plot()
if snap.num_sinks > 0:
sinks = snap.sinks
sinks.plot()
subsnap = snap['gas']
subsnap.image('density', num_pixels=(16, 16))
subsnap.plot()
snap.close_file()
def test_plot_with_kwargs(snaptype):
"""Test particle plot with kwargs."""
filename = DIR / snaptype.filename
snap = plonk.load_snap(filename)
plonk.plot(
snap=snap,
x='x',
y='y',
c='density',
units={
'position': 'au',
'density': 'g/cm^3',
'projection': 'cm'
},
)
snap.close_file()
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_translate_snap(snaptype):
"""Testing translating Snap."""
filename = DIR / snaptype.filename
snap = plonk.load_snap(filename)
unit = f"{snap.code_units['length'].m} {snap.code_units['length'].u}"
snap['position']
if snap.num_sinks > 0:
snap.sinks['position']
snap.translate(translation=(100, 200, 300), unit=unit)
snap.translate(translation=(-100, -200, -300), unit=unit)
_check_arrays(
snap,
snaptype.array_name_map,
snaptype.mean_array_values,
snaptype.std_array_values,
)
snap.translate(translation=(100, 200, 300), unit=unit)
snap.reset()
_check_arrays(
snap,
snaptype.array_name_map,
snaptype.mean_array_values,
snaptype.std_array_values,
)
with pytest.raises(ValueError):
snap.translate(translation=(100, 200, 300))
with pytest.raises(ValueError):
snap.translate(translation=(100, 200))
snap.translate(translation=(100, 200, 300) * plonk.units(unit), unit=unit)
snap.translate(translation=(-100, -200, -300) * plonk.units(unit),
unit=unit)
_check_arrays(
snap,
snaptype.array_name_map,
snaptype.mean_array_values,
snaptype.std_array_values,
)
snap.close_file()
def main(parameters):
dirname = parameters['dirname']
filename = parameters['filename']
debug = parameters['debug']
dust_species_to_plot = parameters['dust_species_to_plot']
radius_min = parameters['radius_min']
radius_max = parameters['radius_max']
scale_height_fac = parameters['scale_height_fac']
n_bins = parameters['n_bins']
path = Path(dirname).expanduser().resolve()
if not path.exists():
raise FileNotFoundError(f'{dirname} does not exist')
in_file = path / 'radialdrift.in'
if not in_file.exists():
raise FileNotFoundError('{dirname}/radialdrift.in does not exist')
snap_file = path / filename
if not snap_file.exists():
raise FileNotFoundError('{dirname}/{filename} does not exist')
print(f' file name = {dirname}/{filename}')
print(f' radius min = {radius_min:~}')
print(f' radius max = {radius_max:~}')
print(f'scale height factor = {scale_height_fac}')
print(f' n_bins = {n_bins}')
snap = plonk.load_snap(filename=snap_file)
profs = calculate_profiles(
snap=snap,
radius_min=radius_min,
radius_max=radius_max,
scale_height_fac=scale_height_fac,
n_bins=n_bins,
)
ax = plot_profiles(
snap=snap, profs=profs, dust_species_to_plot=dust_species_to_plot, debug=debug
)
plt.show()
return snap, profs, ax
