Skip to content

Instantly share code, notes, and snippets.

@chrishavlin
Created November 19, 2021 23:44
Show Gist options
  • Save chrishavlin/2fd39c054b481b47a3a62e216a02eb70 to your computer and use it in GitHub Desktop.
Save chrishavlin/2fd39c054b481b47a3a62e216a02eb70 to your computer and use it in GitHub Desktop.
# based on callbacks from https://yt-project.org/doc/visualizing/callbacks.html#overplot-halo-annotations
import yt
# annotation markers and text
ds = yt.load("IsolatedGalaxy/galaxy0030/galaxy0030")
s = yt.SlicePlot(ds, "x", ("gas", "density"))
s.set_axes_unit("kpc")
# Plot marker and text in data coords
s.annotate_marker((0.2, 0.5, 0.9), coord_system="data")
s.annotate_text((0.2, 0.5, 0.9), "data: (0.2, 0.5, 0.9)", coord_system="data")
# Plot marker and text in plot coords
s.annotate_marker((200, -300), coord_system="plot")
s.annotate_text((200, -300), "plot: (200, -300)", coord_system="plot")
# Plot marker and text in axis coords
s.annotate_marker((0.1, 0.2), coord_system="axis")
s.annotate_text((0.1, 0.2), "axis: (0.1, 0.2)", coord_system="axis")
# Plot marker and text in figure coords
# N.B. marker will not render outside of axis bounds
s.annotate_marker((0.1, 0.2), coord_system="figure", plot_args={"color": "black"})
s.annotate_text(
(0.1, 0.2),
"figure: (0.1, 0.2)",
coord_system="figure",
text_args={"color": "black"},
)
s.save("annotation_coords.png")
s.swap_axes()
s.save("annotation_coords_swapped.png")
# arrows
slc = yt.SlicePlot(ds, "z", ("gas", "density"), center="c", width=(30,'kpc'))
slc.annotate_arrow((0.4975, 0.502, 0.5), length=0.06, plot_args={"color": "blue"})
slc.save("arrow.png")
slc.swap_axes()
slc.save("arrow_swapped.png")
# clump finder
import numpy as np
from yt.data_objects.level_sets.api import Clump, find_clumps
data_source = ds.disk([0.5, 0.5, 0.5], [0.0, 0.0, 1.0], (8.0, "kpc"), (1.0, "kpc"))
c_min = 10 ** np.floor(np.log10(data_source[("gas", "density")]).min())
c_max = 10 ** np.floor(np.log10(data_source[("gas", "density")]).max() + 1)
master_clump = Clump(data_source, ("gas", "density"))
master_clump.add_validator("min_cells", 20)
find_clumps(master_clump, c_min, c_max, 2.0)
leaf_clumps = master_clump.leaves
prj = yt.ProjectionPlot(ds, "z", ("gas", "density"), center="c", width=(20, "kpc"))
prj.annotate_clumps(leaf_clumps)
prj.save("clumps.png")
prj.swap_axes()
prj.save("clumps_swapped.png")
# quivers
p = yt.ProjectionPlot(
ds,
"z",
("gas", "density"),
center=[0.5, 0.5, 0.5],
weight_field="density",
width=(20, "kpc"),
)
p.annotate_quiver(("gas", "velocity_x"), ("gas", "velocity_y"), factor=16,
plot_args={"color": "purple"})
p.save("quiver.png")
p.swap_axes()
p.save("quiver_swapped.png")
p.flip_horizontal()
p.save("quiver_swapped_flip_h.png")
# grid
slc = yt.SlicePlot(ds, "z", ("gas", "density"), width=(10, "kpc"), center="max")
slc.annotate_grids()
slc.save("grid.png")
slc.swap_axes()
slc.save("grid_swapped.png")
# cell edges
slc = yt.SlicePlot(ds, "z", ("gas", "density"), width=(10, "kpc"), center="max")
slc.annotate_cell_edges()
slc.save("cells.png")
slc.swap_axes()
slc.save("cells_swapped.png")
slc.flip_horizontal()
slc.save("cells_swapped_flip_h.png")
# streamlines
s = yt.SlicePlot(ds, "z", ("gas", "density"), center="c", width=(20, "kpc"))
s.annotate_streamlines(("gas", "velocity_x"), ("gas", "velocity_y"))
s.save("streamlines.png")
s.swap_axes()
s.save("streamlines_swapped.png")
s.flip_horizontal()
s.save("streamlines_swapped_flip_h.png")
# line integral convolution
s = yt.SlicePlot(ds, "z", ("gas", "density"), center="c", width=(20, "kpc"))
s.annotate_line_integral_convolution(("gas", "velocity_x"), ("gas", "velocity_y"), lim=(0.5, 0.65))
s.save("lineint.png")
s.swap_axes()
s.save("lineint_swapped.png")
# line annotation (not working??)
p = yt.ProjectionPlot(ds, "z", ("gas", "density"), center="m", width=(10, "kpc"))
p.annotate_line((0.3, 0.4), (0.8, 0.9), coord_system="axis")
p.save("line.png")
p.swap_axes()
p.save("line_swapped.png")
# ray annotation (also not working cause it is a line)
oray = ds.ortho_ray(0, (0.3, 0.4))
ray = ds.ray((0.1, 0.2, 0.3), (0.6, 0.7, 0.8))
p = yt.ProjectionPlot(ds, "z", ("gas", "density"))
p.annotate_ray(oray)
p.annotate_ray(ray)
p.save("ray.png")
p.swap_axes()
p.save("ray_swapped.png")
# cquiver
ds = yt.load("Enzo_64/DD0043/data0043")
s = yt.OffAxisSlicePlot(ds, [1, 1, 0], [("gas", "density")], center="c")
s.annotate_cquiver(
("gas", "cutting_plane_velocity_x"),
("gas", "cutting_plane_velocity_y"),
factor=10,
plot_args={"color": "orange"},
)
s.zoom(1.5)
s.save("cquiver.png")
s.swap_axes()
s.save("cquiver_swapped.png")
s.flip_vertical()
s.save("cquiver_swapped_flip_v.png")
# contours
s = yt.SlicePlot(ds, "x", ("gas", "density"), center="max")
s.annotate_contour(("gas", "temperature"))
s.save("contour.png")
s.swap_axes()
s.save("contour_swapped.png")
# particles
p = yt.ProjectionPlot(ds, "x", ("gas", "density"), center="m", width=(10, "Mpc"))
p.annotate_particles((10, "Mpc"))
p.save("particles.png")
p.swap_axes()
p.save("particles_swapped.png")
# https://yt-project.org/doc/visualizing/callbacks.html#overplot-halo-annotations
# need a new dataset
data_ds = yt.load("Enzo_64/RD0006/RedshiftOutput0006")
halos_ds = yt.load("rockstar_halos/halos_0.0.bin")
prj = yt.ProjectionPlot(data_ds, "z", ("gas", "density"))
prj.annotate_halos(halos_ds, annotate_field="particle_identifier")
prj.save("halo.png")
prj.swap_axes()
prj.save("halo_swapped.png")
# annotate_triangle_facets not implemented, so this fails for now:
import yt
ds = yt.load("MOOSE_sample_data/out.e")
sl = yt.SlicePlot(ds, "z", ("connect1", "nodal_aux"))
sl.annotate_mesh_lines(plot_args={"color": "black"})
sl.save("mesh_lines.png")
sl.swap_axes()
sl.save("mesh_lines_swapped.png")
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment