1import os
2import shutil
3import tempfile
4import unittest
5
6import numpy as np
7
8import yt
9from yt.data_objects.particle_filters import add_particle_filter
10from yt.data_objects.profiles import Profile1D, Profile2D, Profile3D, create_profile
11from yt.testing import (
12    assert_equal,
13    assert_raises,
14    assert_rel_equal,
15    fake_random_ds,
16    fake_sph_orientation_ds,
17    requires_module,
18)
19from yt.utilities.exceptions import YTIllDefinedProfile, YTProfileDataShape
20from yt.visualization.profile_plotter import PhasePlot, ProfilePlot
21
22_fields = ("density", "temperature", "dinosaurs", "tribbles")
23_units = ("g/cm**3", "K", "dyne", "erg")
24
25
26def test_profiles():
27    ds = fake_random_ds(64, nprocs=8, fields=_fields, units=_units)
28    nv = ds.domain_dimensions.prod()
29    dd = ds.all_data()
30    rt, tt, dt = dd.quantities["TotalQuantity"](
31        [("gas", "density"), ("gas", "temperature"), ("stream", "dinosaurs")]
32    )
33
34    e1, e2 = 0.9, 1.1
35    for nb in [8, 16, 32, 64]:
36        for input_units in ["mks", "cgs"]:
37            (rmi, rma), (tmi, tma), (dmi, dma) = (
38                getattr(ex, f"in_{input_units}")()
39                for ex in dd.quantities["Extrema"](
40                    [
41                        ("gas", "density"),
42                        ("gas", "temperature"),
43                        ("stream", "dinosaurs"),
44                    ]
45                )
46            )
47            # We log all the fields or don't log 'em all.  No need to do them
48            # individually.
49            for lf in [True, False]:
50                direct_profile = Profile1D(
51                    dd,
52                    ("gas", "density"),
53                    nb,
54                    rmi * e1,
55                    rma * e2,
56                    lf,
57                    weight_field=None,
58                )
59                direct_profile.add_fields([("index", "ones"), ("gas", "temperature")])
60
61                indirect_profile_s = create_profile(
62                    dd,
63                    ("gas", "density"),
64                    [("index", "ones"), ("gas", "temperature")],
65                    n_bins=nb,
66                    extrema={("gas", "density"): (rmi * e1, rma * e2)},
67                    logs={("gas", "density"): lf},
68                    weight_field=None,
69                )
70
71                indirect_profile_t = create_profile(
72                    dd,
73                    ("gas", "density"),
74                    [("index", "ones"), ("gas", "temperature")],
75                    n_bins=nb,
76                    extrema={("gas", "density"): (rmi * e1, rma * e2)},
77                    logs={("gas", "density"): lf},
78                    weight_field=None,
79                )
80
81                for p1d in [direct_profile, indirect_profile_s, indirect_profile_t]:
82                    assert_equal(p1d["index", "ones"].sum(), nv)
83                    assert_rel_equal(tt, p1d["gas", "temperature"].sum(), 7)
84
85                p2d = Profile2D(
86                    dd,
87                    ("gas", "density"),
88                    nb,
89                    rmi * e1,
90                    rma * e2,
91                    lf,
92                    ("gas", "temperature"),
93                    nb,
94                    tmi * e1,
95                    tma * e2,
96                    lf,
97                    weight_field=None,
98                )
99                p2d.add_fields([("index", "ones"), ("gas", "temperature")])
100                assert_equal(p2d["index", "ones"].sum(), nv)
101                assert_rel_equal(tt, p2d["gas", "temperature"].sum(), 7)
102
103                p3d = Profile3D(
104                    dd,
105                    ("gas", "density"),
106                    nb,
107                    rmi * e1,
108                    rma * e2,
109                    lf,
110                    ("gas", "temperature"),
111                    nb,
112                    tmi * e1,
113                    tma * e2,
114                    lf,
115                    ("stream", "dinosaurs"),
116                    nb,
117                    dmi * e1,
118                    dma * e2,
119                    lf,
120                    weight_field=None,
121                )
122                p3d.add_fields([("index", "ones"), ("gas", "temperature")])
123                assert_equal(p3d["index", "ones"].sum(), nv)
124                assert_rel_equal(tt, p3d["gas", "temperature"].sum(), 7)
125
126        p1d = Profile1D(dd, ("index", "x"), nb, 0.0, 1.0, False, weight_field=None)
127        p1d.add_fields(("index", "ones"))
128        av = nv / nb
129        assert_equal(p1d["index", "ones"], np.ones(nb) * av)
130
131        # We re-bin ones with a weight now
132        p1d = Profile1D(
133            dd, ("index", "x"), nb, 0.0, 1.0, False, weight_field=("gas", "temperature")
134        )
135        p1d.add_fields([("index", "ones")])
136        assert_equal(p1d["index", "ones"], np.ones(nb))
137
138        # Verify we can access "ones" after adding a new field
139        # See issue 988
140        p1d.add_fields([("gas", "density")])
141        assert_equal(p1d["index", "ones"], np.ones(nb))
142
143        p2d = Profile2D(
144            dd,
145            ("index", "x"),
146            nb,
147            0.0,
148            1.0,
149            False,
150            ("index", "y"),
151            nb,
152            0.0,
153            1.0,
154            False,
155            weight_field=None,
156        )
157        p2d.add_fields(("index", "ones"))
158        av = nv / nb ** 2
159        assert_equal(p2d["index", "ones"], np.ones((nb, nb)) * av)
160
161        # We re-bin ones with a weight now
162        p2d = Profile2D(
163            dd,
164            ("index", "x"),
165            nb,
166            0.0,
167            1.0,
168            False,
169            ("index", "y"),
170            nb,
171            0.0,
172            1.0,
173            False,
174            weight_field=("gas", "temperature"),
175        )
176        p2d.add_fields([("index", "ones")])
177        assert_equal(p2d["index", "ones"], np.ones((nb, nb)))
178
179        p3d = Profile3D(
180            dd,
181            ("index", "x"),
182            nb,
183            0.0,
184            1.0,
185            False,
186            ("index", "y"),
187            nb,
188            0.0,
189            1.0,
190            False,
191            ("index", "z"),
192            nb,
193            0.0,
194            1.0,
195            False,
196            weight_field=None,
197        )
198        p3d.add_fields(("index", "ones"))
199        av = nv / nb ** 3
200        assert_equal(p3d["index", "ones"], np.ones((nb, nb, nb)) * av)
201
202        # We re-bin ones with a weight now
203        p3d = Profile3D(
204            dd,
205            ("index", "x"),
206            nb,
207            0.0,
208            1.0,
209            False,
210            ("index", "y"),
211            nb,
212            0.0,
213            1.0,
214            False,
215            ("index", "z"),
216            nb,
217            0.0,
218            1.0,
219            False,
220            weight_field=("gas", "temperature"),
221        )
222        p3d.add_fields([("index", "ones")])
223        assert_equal(p3d["index", "ones"], np.ones((nb, nb, nb)))
224
225        p2d = create_profile(
226            dd,
227            ("gas", "density"),
228            ("gas", "temperature"),
229            weight_field=("gas", "cell_mass"),
230            extrema={("gas", "density"): (None, rma * e2)},
231        )
232        assert_equal(p2d.x_bins[0], rmi - np.spacing(rmi))
233        assert_equal(p2d.x_bins[-1], rma * e2)
234        assert str(ds.field_info["gas", "cell_mass"].units) == str(p2d.weight.units)
235
236        p2d = create_profile(
237            dd,
238            ("gas", "density"),
239            ("gas", "temperature"),
240            weight_field=("gas", "cell_mass"),
241            extrema={("gas", "density"): (rmi * e2, None)},
242        )
243        assert_equal(p2d.x_bins[0], rmi * e2)
244        assert_equal(p2d.x_bins[-1], rma + np.spacing(rma))
245
246
247extrema_s = {("all", "particle_position_x"): (0, 1)}
248logs_s = {("all", "particle_position_x"): False}
249
250extrema_t = {("all", "particle_position_x"): (0, 1)}
251logs_t = {("all", "particle_position_x"): False}
252
253
254def test_particle_profiles():
255    for nproc in [1, 2, 4, 8]:
256        ds = fake_random_ds(32, nprocs=nproc, particles=32 ** 3)
257        dd = ds.all_data()
258
259        p1d = Profile1D(
260            dd, ("all", "particle_position_x"), 128, 0.0, 1.0, False, weight_field=None
261        )
262        p1d.add_fields([("all", "particle_ones")])
263        assert_equal(p1d[("all", "particle_ones")].sum(), 32 ** 3)
264
265        p1d = create_profile(
266            dd,
267            [("all", "particle_position_x")],
268            [("all", "particle_ones")],
269            weight_field=None,
270            n_bins=128,
271            extrema=extrema_s,
272            logs=logs_s,
273        )
274        assert_equal(p1d[("all", "particle_ones")].sum(), 32 ** 3)
275
276        p1d = create_profile(
277            dd,
278            [("all", "particle_position_x")],
279            [("all", "particle_ones")],
280            weight_field=None,
281            n_bins=128,
282            extrema=extrema_t,
283            logs=logs_t,
284        )
285        assert_equal(p1d[("all", "particle_ones")].sum(), 32 ** 3)
286
287        p2d = Profile2D(
288            dd,
289            ("all", "particle_position_x"),
290            128,
291            0.0,
292            1.0,
293            False,
294            ("all", "particle_position_y"),
295            128,
296            0.0,
297            1.0,
298            False,
299            weight_field=None,
300        )
301        p2d.add_fields([("all", "particle_ones")])
302        assert_equal(p2d[("all", "particle_ones")].sum(), 32 ** 3)
303
304        p3d = Profile3D(
305            dd,
306            ("all", "particle_position_x"),
307            128,
308            0.0,
309            1.0,
310            False,
311            ("all", "particle_position_y"),
312            128,
313            0.0,
314            1.0,
315            False,
316            ("all", "particle_position_z"),
317            128,
318            0.0,
319            1.0,
320            False,
321            weight_field=None,
322        )
323        p3d.add_fields([("all", "particle_ones")])
324        assert_equal(p3d[("all", "particle_ones")].sum(), 32 ** 3)
325
326
327def test_mixed_particle_mesh_profiles():
328    ds = fake_random_ds(32, particles=10)
329    ad = ds.all_data()
330    assert_raises(
331        YTIllDefinedProfile,
332        ProfilePlot,
333        ad,
334        ("index", "radius"),
335        ("all", "particle_mass"),
336    )
337    assert_raises(
338        YTIllDefinedProfile,
339        ProfilePlot,
340        ad,
341        "radius",
342        [("all", "particle_mass"), ("all", "particle_ones")],
343    )
344    assert_raises(
345        YTIllDefinedProfile,
346        ProfilePlot,
347        ad,
348        ("index", "radius"),
349        [("all", "particle_mass"), ("index", "ones")],
350    )
351    assert_raises(
352        YTIllDefinedProfile,
353        ProfilePlot,
354        ad,
355        ("all", "particle_radius"),
356        ("all", "particle_mass"),
357        ("gas", "cell_mass"),
358    )
359    assert_raises(
360        YTIllDefinedProfile,
361        ProfilePlot,
362        ad,
363        ("index", "radius"),
364        ("gas", "cell_mass"),
365        ("all", "particle_ones"),
366    )
367
368    assert_raises(
369        YTIllDefinedProfile,
370        PhasePlot,
371        ad,
372        ("index", "radius"),
373        ("all", "particle_mass"),
374        ("gas", "velocity_x"),
375    )
376    assert_raises(
377        YTIllDefinedProfile,
378        PhasePlot,
379        ad,
380        ("all", "particle_radius"),
381        ("all", "particle_mass"),
382        ("gas", "cell_mass"),
383    )
384    assert_raises(
385        YTIllDefinedProfile,
386        PhasePlot,
387        ad,
388        ("index", "radius"),
389        ("gas", "cell_mass"),
390        ("all", "particle_ones"),
391    )
392    assert_raises(
393        YTIllDefinedProfile,
394        PhasePlot,
395        ad,
396        ("all", "particle_radius"),
397        ("all", "particle_mass"),
398        ("all", "particle_ones"),
399    )
400
401
402def test_particle_profile_negative_field():
403    # see Issue #1340
404    n_particles = int(1e4)
405
406    ppx, ppy, ppz = np.random.normal(size=[3, n_particles])
407    pvx, pvy, pvz = -np.ones((3, n_particles))
408
409    data = {
410        "particle_position_x": ppx,
411        "particle_position_y": ppy,
412        "particle_position_z": ppz,
413        "particle_velocity_x": pvx,
414        "particle_velocity_y": pvy,
415        "particle_velocity_z": pvz,
416    }
417
418    bbox = 1.1 * np.array(
419        [[min(ppx), max(ppx)], [min(ppy), max(ppy)], [min(ppz), max(ppz)]]
420    )
421    ds = yt.load_particles(data, bbox=bbox)
422    ad = ds.all_data()
423
424    profile = yt.create_profile(
425        ad,
426        [("all", "particle_position_x"), ("all", "particle_position_y")],
427        ("all", "particle_velocity_x"),
428        logs={
429            ("all", "particle_position_x"): True,
430            ("all", "particle_position_y"): True,
431            ("all", "particle_position_z"): True,
432        },
433        weight_field=None,
434    )
435    assert profile[("all", "particle_velocity_x")].min() < 0
436    assert profile.x_bins.min() > 0
437    assert profile.y_bins.min() > 0
438
439    profile = yt.create_profile(
440        ad,
441        [("all", "particle_position_x"), ("all", "particle_position_y")],
442        ("all", "particle_velocity_x"),
443        weight_field=None,
444    )
445    assert profile[("all", "particle_velocity_x")].min() < 0
446    assert profile.x_bins.min() < 0
447    assert profile.y_bins.min() < 0
448
449    # can't use CIC deposition with log-scaled bin fields
450    with assert_raises(RuntimeError):
451        yt.create_profile(
452            ad,
453            [("all", "particle_position_x"), ("all", "particle_position_y")],
454            ("all", "particle_velocity_x"),
455            logs={
456                ("all", "particle_position_x"): True,
457                ("all", "particle_position_y"): False,
458                ("all", "particle_position_z"): False,
459            },
460            weight_field=None,
461            deposition="cic",
462        )
463
464    # can't use CIC deposition with accumulation or fractional
465    with assert_raises(RuntimeError):
466        yt.create_profile(
467            ad,
468            [("all", "particle_position_x"), ("all", "particle_position_y")],
469            ("all", "particle_velocity_x"),
470            logs={
471                ("all", "particle_position_x"): False,
472                ("all", "particle_position_y"): False,
473                ("all", "particle_position_z"): False,
474            },
475            weight_field=None,
476            deposition="cic",
477            accumulation=True,
478            fractional=True,
479        )
480
481
482def test_profile_zero_weight():
483    def DMparticles(pfilter, data):
484        filter = data[(pfilter.filtered_type, "particle_type")] == 1
485        return filter
486
487    def DM_in_cell_mass(field, data):
488        return data["deposit", "DM_density"] * data["index", "cell_volume"]
489
490    add_particle_filter(
491        "DM", function=DMparticles, filtered_type="io", requires=["particle_type"]
492    )
493
494    _fields = (
495        "particle_position_x",
496        "particle_position_y",
497        "particle_position_z",
498        "particle_mass",
499        "particle_velocity_x",
500        "particle_velocity_y",
501        "particle_velocity_z",
502        "particle_type",
503    )
504    _units = ("cm", "cm", "cm", "g", "cm/s", "cm/s", "cm/s", "dimensionless")
505    ds = fake_random_ds(
506        32, particle_fields=_fields, particle_field_units=_units, particles=16
507    )
508
509    ds.add_particle_filter("DM")
510    ds.add_field(
511        ("gas", "DM_cell_mass"),
512        units="g",
513        function=DM_in_cell_mass,
514        sampling_type="cell",
515    )
516
517    sp = ds.sphere(ds.domain_center, (10, "kpc"))
518
519    profile = yt.create_profile(
520        sp,
521        [("gas", "density")],
522        [("gas", "radial_velocity")],
523        weight_field=("gas", "DM_cell_mass"),
524    )
525
526    assert not np.any(np.isnan(profile["gas", "radial_velocity"]))
527
528
529def test_profile_sph_data():
530    ds = fake_sph_orientation_ds()
531    # test we create a profile without raising YTIllDefinedProfile
532    yt.create_profile(
533        ds.all_data(),
534        [("gas", "density"), ("gas", "temperature")],
535        [("gas", "kinetic_energy_density")],
536        weight_field=None,
537    )
538
539
540def test_profile_override_limits():
541    ds = fake_random_ds(64, nprocs=8, fields=_fields, units=_units)
542
543    sp = ds.sphere(ds.domain_center, (10, "kpc"))
544    obins = np.linspace(-5, 5, 10)
545    profile = yt.create_profile(
546        sp,
547        [("gas", "density")],
548        [("gas", "temperature")],
549        override_bins={("gas", "density"): (obins, "g/cm**3")},
550    )
551    assert_equal(ds.arr(obins, "g/cm**3"), profile.x_bins)
552
553    profile = yt.create_profile(
554        sp,
555        [("gas", "density"), ("stream", "dinosaurs")],
556        [("gas", "temperature")],
557        override_bins={
558            ("gas", "density"): (obins, "g/cm**3"),
559            ("stream", "dinosaurs"): obins,
560        },
561    )
562    assert_equal(ds.arr(obins, "g/cm**3"), profile.x_bins)
563    assert_equal(ds.arr(obins, "dyne"), profile.y_bins)
564
565    profile = yt.create_profile(
566        sp,
567        [("gas", "density"), (("stream", "dinosaurs")), ("stream", "tribbles")],
568        [("gas", "temperature")],
569        override_bins={
570            ("gas", "density"): (obins, "g/cm**3"),
571            (("stream", "dinosaurs")): obins,
572            ("stream", "tribbles"): (obins, "erg"),
573        },
574    )
575    assert_equal(ds.arr(obins, "g/cm**3"), profile.x_bins)
576    assert_equal(ds.arr(obins, "dyne"), profile.y_bins)
577    assert_equal(ds.arr(obins, "erg"), profile.z_bins)
578
579
580class TestBadProfiles(unittest.TestCase):
581
582    tmpdir = None
583    curdir = None
584
585    def setUp(self):
586        self.tmpdir = tempfile.mkdtemp()
587        self.curdir = os.getcwd()
588        os.chdir(self.tmpdir)
589
590    def tearDown(self):
591        os.chdir(self.curdir)
592        # clean up
593        shutil.rmtree(self.tmpdir)
594
595    @requires_module("h5py")
596    def test_unequal_data_shape_profile(self):
597        density = np.random.random(128)
598        temperature = np.random.random(128)
599        mass = np.random.random((128, 128))
600
601        my_data = {
602            ("gas", "density"): density,
603            ("gas", "temperature"): temperature,
604            ("gas", "mass"): mass,
605        }
606        fake_ds_med = {"current_time": yt.YTQuantity(10, "Myr")}
607        field_types = {field: "gas" for field in my_data.keys()}
608        yt.save_as_dataset(fake_ds_med, "mydata.h5", my_data, field_types=field_types)
609
610        ds = yt.load("mydata.h5")
611
612        with assert_raises(YTProfileDataShape):
613            yt.PhasePlot(
614                ds.data,
615                ("gas", "temperature"),
616                ("gas", "density"),
617                ("gas", "mass"),
618            )
619
620    @requires_module("h5py")
621    def test_unequal_bin_field_profile(self):
622        density = np.random.random(128)
623        temperature = np.random.random(127)
624        mass = np.random.random((128, 128))
625
626        my_data = {
627            ("gas", "density"): density,
628            ("gas", "temperature"): temperature,
629            ("gas", "mass"): mass,
630        }
631        fake_ds_med = {"current_time": yt.YTQuantity(10, "Myr")}
632        field_types = {field: "gas" for field in my_data.keys()}
633        yt.save_as_dataset(fake_ds_med, "mydata.h5", my_data, field_types=field_types)
634
635        ds = yt.load("mydata.h5")
636
637        with assert_raises(YTProfileDataShape):
638            yt.PhasePlot(
639                ds.data,
640                ("gas", "temperature"),
641                ("gas", "density"),
642                ("gas", "mass"),
643            )
644
645
646def test_index_field_units():
647    # see #1849
648    ds = fake_random_ds(16, length_unit=2)
649    ad = ds.all_data()
650    icv_units = ad["index", "cell_volume"].units
651    assert str(icv_units) == "code_length**3"
652    gcv_units = ad["gas", "cell_volume"].units
653    assert str(gcv_units) == "cm**3"
654    prof = ad.profile(
655        [("gas", "density"), ("gas", "velocity_x")],
656        [("gas", "cell_volume"), ("index", "cell_volume")],
657        weight_field=None,
658    )
659    assert str(prof["index", "cell_volume"].units) == "code_length**3"
660    assert str(prof["gas", "cell_volume"].units) == "cm**3"
661
662
663@requires_module("astropy")
664def test_export_astropy():
665    from yt.units.yt_array import YTArray
666
667    ds = fake_random_ds(64)
668    ad = ds.all_data()
669    prof = ad.profile(
670        ("index", "radius"),
671        [("gas", "density"), ("gas", "velocity_x")],
672        weight_field=("index", "ones"),
673        n_bins=32,
674    )
675    # export to AstroPy table
676    at1 = prof.to_astropy_table()
677    assert "radius" in at1.colnames
678    assert "density" in at1.colnames
679    assert "velocity_x" in at1.colnames
680    assert_equal(prof.x.d, at1["radius"].value)
681    assert_equal(prof[("gas", "density")].d, at1["density"].value)
682    assert_equal(prof[("gas", "velocity_x")].d, at1["velocity_x"].value)
683    assert prof.x.units == YTArray.from_astropy(at1["radius"]).units
684    assert prof[("gas", "density")].units == YTArray.from_astropy(at1["density"]).units
685    assert (
686        prof[("gas", "velocity_x")].units
687        == YTArray.from_astropy(at1["velocity_x"]).units
688    )
689    assert np.all(at1.mask["density"] == prof.used)
690    at2 = prof.to_astropy_table(fields=("gas", "density"), only_used=True)
691    assert "radius" in at2.colnames
692    assert "velocity_x" not in at2.colnames
693    assert_equal(prof.x.d[prof.used], at2["radius"].value)
694    assert_equal(prof[("gas", "density")].d[prof.used], at2["density"].value)
695
696
697@requires_module("pandas")
698def test_export_pandas():
699    ds = fake_random_ds(64)
700    ad = ds.all_data()
701    prof = ad.profile(
702        "radius",
703        [("gas", "density"), ("gas", "velocity_x")],
704        weight_field=("index", "ones"),
705        n_bins=32,
706    )
707    # export to pandas DataFrame
708    df1 = prof.to_dataframe()
709    assert "radius" in df1.columns
710    assert "density" in df1.columns
711    assert "velocity_x" in df1.columns
712    assert_equal(prof.x.d, df1["radius"])
713    assert_equal(prof[("gas", "density")].d, np.nan_to_num(df1["density"]))
714    assert_equal(prof["velocity_x"].d, np.nan_to_num(df1["velocity_x"]))
715    df2 = prof.to_dataframe(fields=("gas", "density"), only_used=True)
716    assert "radius" in df2.columns
717    assert "velocity_x" not in df2.columns
718    assert_equal(prof.x.d[prof.used], df2["radius"])
719    assert_equal(prof[("gas", "density")].d[prof.used], df2["density"])
720