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