1import numpy as np
2
3from yt.frontends.athena_pp.api import AthenaPPDataset
4from yt.loaders import load
5from yt.testing import (
6    assert_allclose,
7    assert_equal,
8    requires_file,
9    units_override_check,
10)
11from yt.utilities.answer_testing.framework import (
12    GenericArrayTest,
13    data_dir_load,
14    requires_ds,
15    small_patch_amr,
16)
17
18_fields_disk = ("density", "velocity_r")
19
20disk = "KeplerianDisk/disk.out1.00000.athdf"
21
22
23@requires_ds(disk)
24def test_disk():
25    ds = data_dir_load(disk)
26    assert_equal(str(ds), "disk.out1.00000")
27    dd = ds.all_data()
28    vol = (ds.domain_right_edge[0] ** 3 - ds.domain_left_edge[0] ** 3) / 3.0
29    vol *= np.cos(ds.domain_left_edge[1]) - np.cos(ds.domain_right_edge[1])
30    vol *= ds.domain_right_edge[2].v - ds.domain_left_edge[2].v
31    assert_allclose(dd.quantities.total_quantity(("gas", "cell_volume")), vol)
32    for field in _fields_disk:
33
34        def field_func(name):
35            return dd[field]
36
37        yield GenericArrayTest(ds, field_func, args=[field])
38
39
40_fields_AM06 = ("temperature", "density", "velocity_magnitude", "magnetic_field_x")
41
42AM06 = "AM06/AM06.out1.00400.athdf"
43
44
45@requires_ds(AM06)
46def test_AM06():
47    ds = data_dir_load(AM06)
48    assert_equal(str(ds), "AM06.out1.00400")
49    for test in small_patch_amr(ds, _fields_AM06):
50        test_AM06.__name__ = test.description
51        yield test
52
53
54uo_AM06 = {
55    "length_unit": (1.0, "kpc"),
56    "mass_unit": (1.0, "Msun"),
57    "time_unit": (1.0, "Myr"),
58}
59
60
61@requires_file(AM06)
62def test_AM06_override():
63    # verify that overriding units causes derived unit values to be updated.
64    # see issue #1259
65    ds = load(AM06, units_override=uo_AM06)
66    assert_equal(float(ds.magnetic_unit.in_units("gauss")), 9.01735778342523e-08)
67
68
69@requires_file(AM06)
70def test_units_override():
71    units_override_check(AM06)
72
73
74@requires_file(AM06)
75def test_AthenaPPDataset():
76    assert isinstance(data_dir_load(AM06), AthenaPPDataset)
77