GEAR: small improvement to the example

examples/GEAR/ZoomIn/add_fields.py

@@ -3,6 +3,7 @@
 import numpy as np
 import yt.utilities.physical_constants as constants
 from yt.mods import YTArray
+from yt import derived_field
 
 
 def addMassDeposit(f):
@@ -65,3 +66,25 @@ def addTemperature(f):
         return YTArray(convert_T_over_mu_to_T(T_over_mu), 'K')  # now T
     f.add_field(("PartType0", "Temperature"), function=_Temperature_3,
                 particle_type=True, force_override=True, units="K")
+
+
+def addMetals(f):
+    def _metallicity(field, data):
+        if len(data["PartType0", "Metals"].shape) == 1:
+            return data["PartType0", "Metals"]
+        else:
+            # in_units("") turned out to be crucial!;
+            # otherwise code_metallicity will be used
+            return data["PartType0", "Metals"][:, -1].in_units("")
+
+    # We are creating ("Gas", "Metallicity") here, different from
+    # ("Gas", "metallicity") which is auto-generated by yt
+    # but doesn't work properly
+    f.add_field(("PartType0", "Metallicity"), function=_metallicity,
+                display_name="Metal mass fraction", particle_type=True,
+                take_log=True, units="")
+
+    def _density_squared(field, data):
+        return data[("PartType0", "density")]**2
+    f.add_field(("PartType0", "density_squared"), function=_density_squared,
+                particle_type=True, units="g**2/cm**6")

examples/GEAR/ZoomIn/cleanupSwift.py

@@ -11,7 +11,7 @@ out = sys.argv[-1]
 
 copyfile(filename, out)
 
-f = File(out)
+f = File(out, "a")
 
 # read cosmological parameters
 a = f["Cosmology"].attrs["Scale-factor"][0]
@@ -23,12 +23,6 @@ f["Header"].attrs["BoxSize"] = f["Header"].attrs["BoxSize"][0] * h
 # avoid the snapshot to be taken for SWIFT
 del f["Header"].attrs["Code"]
 
-# Delete problematic fields
-for i in range(NPartType):
-    name = "PartType{}/ElementAbundances".format(i)
-    if name in f:
-        del f[name]
-
 # Update the fields (name + cosmo)
 for i in range(NPartType):
     name = "PartType%i" % i
@@ -43,11 +37,12 @@ for i in range(NPartType):
     fields = [
         ("Coordinates", "Coordinates", h),
         ("Masses", "Masses", h),
-        ("Velocities", "Velocities", a**0.5),
+        ("Velocities", "Velocities", 1. / a**0.5),
         ("Density", "Densities", 1. / h**2),
         ("Entropies", "Entropies", 1.),
         ("InternalEnergy", "InternalEnergies", 1. / a**2),
-        ("SmoothingLength", "SmoothingLengths", h)
+        ("SmoothingLength", "SmoothingLengths", h),
+        ("Metals", "SmoothedElementAbundances", 1.)
     ]
 
     # create links
@@ -67,7 +62,7 @@ cosmo = f["Cosmology"].attrs
 head = f["Header"].attrs
 head["Redshift"] = float(cosmo["Redshift"])
 head["OmegaLambda"] = cosmo["Omega_lambda"]
-head["Omega0"] = cosmo["Omega_b"]
+head["Omega0"] = cosmo["Omega_m"]
 head["HubbleParam"] = cosmo["h"][0]
 head["Time"] = float(a)
 

examples/GEAR/ZoomIn/make_image.py

 #!/usr/bin/env python3
 
 import yt
-from yt.units import Msun, kpc, s, km
-import unyt
 import sys
-import matplotlib
 import projection_plot
 import velocity_plot
 import phase_plot
 import halo_distribution
+import metal_plot
 import add_fields
+import star_plot
+import profile_plot
 import matplotlib.pyplot as plt
 from mpl_toolkits.axes_grid1 import AxesGrid
-from matplotlib.offsetbox import AnchoredText
-matplotlib.use('Agg')
 
 # Parameters
 
@@ -23,14 +21,29 @@ swift = "swift/snapshot_%04i.hdf5" % snap
 gear = "gear/snapshot_%04i.hdf5" % snap
 
 
+do_dmo = False
+do_hydro = False
+do_stars = True
+do_feedback = True
 do_plot = {
-    "projection_density": False,
-    "projection_temperature": False,
-    "projection_mass": False,
-    "halo_distribution": False,
-    "phase_1d": False,
-    "phase_2d": False,
-    "velocity": True
+    # DMO
+    "projection_mass": do_dmo,
+    "velocity": do_dmo,
+    "halo_distribution": False,  # very slow
+    "profile": do_dmo,
+    # hydro
+    "projection_density": True, #do_hydro,
+    "projection_temperature": do_hydro,
+    "phase_1d_density": do_hydro,
+    "phase_1d_temperature": do_hydro,
+    "phase_2d": do_hydro,
+    "metals": do_hydro,
+    # stars
+    "SFR": do_stars,
+    "projection_mass_stars": do_stars,
+    # feedback
+    "abundances": do_feedback,
+    "projection_metals": do_feedback,
 }
 
 # Generate the figures
@@ -39,7 +52,9 @@ figsize = (100, 100)
 figures = {
     "projection_density": plt.figure(figsize=figsize),
     "projection_temperature": plt.figure(figsize=figsize),
+    "projection_metals": plt.figure(figsize=figsize),
     "projection_mass": plt.figure(figsize=figsize),
+    "projection_mass_stars": plt.figure(figsize=figsize),
     "phase_2d": plt.figure(figsize=figsize)
 }
 
@@ -63,18 +78,31 @@ axes = {
         figures["projection_mass"], fig_size, subgrid,
         add_all=True, share_all=True, cbar_mode="single",
         cbar_size="2%", cbar_pad=0.02),
+    "projection_mass_stars": AxesGrid(
+        figures["projection_mass_stars"], fig_size, subgrid,
+        add_all=True, share_all=True, cbar_mode="single",
+        cbar_size="2%", cbar_pad=0.02),
     "phase_2d": AxesGrid(
         figures["phase_2d"], fig_size, subgrid, axes_pad=0.05,
         add_all=True, share_all=True, cbar_mode="single",
-        cbar_size="2%", cbar_pad=0.05, aspect=False)
+        cbar_size="2%", cbar_pad=0.05, aspect=False),
+    "projection_metals": AxesGrid(
+        figures["projection_metals"], fig_size, subgrid,
+        add_all=True, share_all=True, cbar_mode="single",
+        cbar_size="2%", cbar_pad=0.02),
 }
 
 
 # Data
 data = {
-    "phase_1d": ([], []),
+    "phase_1d_density": ([], []),
+    "phase_1d_temperature": ([], []),
     "halo_distribution": ([], []),
     "velocity": ([], []),
+    "metals": ([], []),
+    "SFR": ([], []),
+    "profile": ([], []),
+    "abundances": ([], []),
 }
 
 names = []
@@ -88,13 +116,25 @@ def savePlot():
         projection_plot.savePlot(figures["projection_temperature"],
                                  "temperature")
 
+    if do_plot["projection_metals"]:
+        projection_plot.savePlot(figures["projection_metals"],
+                                 "metals")
+
     if do_plot["projection_mass"]:
         projection_plot.savePlot(figures["projection_mass"],
                                  "mass")
 
-    if do_plot["phase_1d"]:
-        data["phase_1d"][1].extend(names)
-        phase_plot.save1DPlot(data["phase_1d"])
+    if do_plot["projection_mass_stars"]:
+        projection_plot.savePlot(figures["projection_mass_stars"],
+                                 "mass_stars")
+
+    if do_plot["phase_1d_density"]:
+        data["phase_1d_density"][1].extend(names)
+        phase_plot.save1DPlotDensity(data["phase_1d_density"])
+
+    if do_plot["phase_1d_temperature"]:
+        data["phase_1d_temperature"][1].extend(names)
+        phase_plot.save1DPlotTemperature(data["phase_1d_temperature"])
 
     if do_plot["phase_2d"]:
         phase_plot.save2DPlot(figures["phase_2d"])
@@ -108,13 +148,38 @@ def savePlot():
         data["velocity"][1].extend(names)
         velocity_plot.save1DPlot(data["velocity"])
 
+    if do_plot["metals"]:
+        data["metals"][1].extend(names)
+        metal_plot.save1DPlot(data["metals"])
+
+    if do_plot["SFR"]:
+        data["SFR"][1].extend(names)
+        star_plot.saveSFRPlot(data["SFR"])
+
+    if do_plot["profile"]:
+        data["profile"][1].extend(names)
+        profile_plot.savePlot(data["profile"])
 
-def doPlot(filename, i, name):
+    if do_plot["abundances"]:
+        data["abundances"][1].extend(names)
+        star_plot.saveAbundancesPlot(data["abundances"])
+
+
+def doPlot(filename, i, name, center):
     names.append(name)
 
     f = yt.load(filename)
+
+    if center is None:
+        center = f.find_max("Density")[1]
+    f.center = center
+
     if (do_plot["projection_temperature"] or do_plot["phase_2d"]):
         add_fields.addTemperature(f)
+    add_fields.addMassDeposit(f)
+
+    if (do_plot["projection_metals"] or do_plot["metals"]):
+        add_fields.addMetals(f)
 
     global scale_factor
     if scale_factor is None:
@@ -135,17 +200,31 @@ def doPlot(filename, i, name):
             f, name, i, figures["projection_temperature"],
             axes["projection_temperature"])
 
+    if do_plot["projection_metals"]:
+        projection_plot.doMetalsPlot(
+            f, name, i, figures["projection_metals"],
+            axes["projection_metals"])
+
     # Do mass projection plot
     if do_plot["projection_mass"]:
-        add_fields.addMassDeposit(f)
         projection_plot.doMassPlot(
             f, name, i, figures["projection_mass"],
-            axes["projection_mass"])
+            axes["projection_mass"], "all")
+
+    # Do stellar mass projection plot
+    if do_plot["projection_mass_stars"]:
+        projection_plot.doMassPlot(
+            f, name, i, figures["projection_mass_stars"],
+            axes["projection_mass_stars"], "stars")
 
     # 1D Phase plot
-    if do_plot["phase_1d"]:
-        p = phase_plot.do1DPlot(f, name, i)
-        data["phase_1d"][0].append(p)
+    if do_plot["phase_1d_density"]:
+        p = phase_plot.do1DPlotDensity(f, name, i)
+        data["phase_1d_density"][0].append(p)
+
+    if do_plot["phase_1d_temperature"]:
+        p = phase_plot.do1DPlotTemperature(f, name, i)
+        data["phase_1d_temperature"][0].append(p)
 
     # 2D Phase plot
     if do_plot["phase_2d"]:
@@ -162,7 +241,26 @@ def doPlot(filename, i, name):
         p = velocity_plot.do1DPlot(f, name, i)
         data["velocity"][0].append(p)
 
+    # metal plot
+    if do_plot["metals"]:
+        p = metal_plot.do1DPlot(f, name, i)
+        data["metals"][0].append(p)
+
+    if do_plot["SFR"]:
+        p = star_plot.doSFRPlot(f, name, i)
+        data["SFR"][0].append(p)
+
+    if do_plot["profile"]:
+        p = profile_plot.doPlot(f, name, i)
+        data["profile"][0].append(p)
+
+    if do_plot["abundances"]:
+        p = star_plot.doAbundancesPlot(f, name, i)
+        data["abundances"][0].append(p)
+
+    return center
+
 
-doPlot(swift, 0, "SWIFT")
-doPlot(gear, 1, "GEAR")
+center = doPlot(gear, 1, "GEAR", center=None)
+doPlot(swift, 0, "SWIFT", center=center)
 savePlot()

examples/GEAR/ZoomIn/phase_plot.py

@@ -4,10 +4,12 @@ import yt
 from yt.units import Msun, amu, cm, K, kpc
 from matplotlib.offsetbox import AnchoredText
 import matplotlib.pyplot as plt
+import matplotlib.lines as ln
+import numpy as np
 
-width = 5 * kpc
-limits_temperature = (1e1 * K, 1e7 * K)
-limits_density = (1e-10 * amu / cm**3, 1e3 * amu / cm**3)
+width = 10 * kpc
+limits_temperature = (1e1 * K, 1e5 * K)
+limits_density = (1e-7 * amu / cm**3, 1e7 * amu / cm**3)
 limits_mass = (1e3 * Msun, 1e7 * Msun)
 
 
@@ -16,27 +18,54 @@ def save2DPlot(fig):
                 pad_inches=0.03, dpi=300)
 
 
-def save1DPlot(profiles):
+def save1DPlotDensity(profiles):
     plt.figure(figsize=(8, 8))
     plt.tight_layout(pad=0.4, w_pad=0.5, h_pad=0.5)
     markers = ["s", "o"]
     for i, p in enumerate(profiles[0]):
         rho = p.x.in_units("g/cm**3").d
         mass = p["Masses"].in_units("Msun").d
+        mass[mass == 0] = np.nan
         plt.plot(rho, mass, linestyle="-", marker=markers[i],
                  markeredgecolor='none', linewidth=1.2, alpha=0.8)
     plt.semilogx()
     plt.semilogy()
 
     plt.xlabel("$\mathrm{Density\ (g/cm^3)}$", fontsize='large')
-    plt.ylabel(r"$\mathrm{Mass,}\/\mathrm{d}M\mathrm{/dlog}\/\mathrm{\rho}\/\mathrm{(M_{\odot})}$", fontsize='large')
+    plt.ylabel(r"$\mathrm{Mass}\/\mathrm{(M_{\odot})}$", fontsize='large')
     plt.legend(profiles[1], loc=4, frameon=True, ncol=2, fancybox=True)
     leg = plt.gca().get_legend()
     ltext = leg.get_texts()
     plt.setp(ltext, fontsize='small')
     plt.grid(True)
 
-    plt.savefig("phase_1d.png", bbox_inches='tight', pad_inches=0.03, dpi=300)
+    plt.savefig("phase_1d_density.png", bbox_inches='tight',
+                pad_inches=0.03, dpi=300)
+
+
+def save1DPlotTemperature(profiles):
+    plt.figure(figsize=(8, 8))
+    plt.tight_layout(pad=0.4, w_pad=0.5, h_pad=0.5)
+    markers = ["s", "o"]
+    for i, p in enumerate(profiles[0]):
+        rho = p.x.in_units("K").d
+        mass = p["Masses"].in_units("Msun").d
+        mass[mass == 0] = np.nan
+        plt.plot(rho, mass, linestyle="-", marker=markers[i],
+                 markeredgecolor='none', linewidth=1.2, alpha=0.8)
+    plt.semilogx()
+    plt.semilogy()
+
+    plt.xlabel("$\mathrm{Temperature\ K}$", fontsize='large')
+    plt.ylabel(r"$\mathrm{Mass}\/\mathrm{(M_{\odot})}$", fontsize='large')
+    plt.legend(profiles[1], loc=4, frameon=True, ncol=2, fancybox=True)
+    leg = plt.gca().get_legend()
+    ltext = leg.get_texts()
+    plt.setp(ltext, fontsize='small')
+    plt.grid(True)
+
+    plt.savefig("phase_1d_temperature.png", bbox_inches='tight',
+                pad_inches=0.03, dpi=300)
 
 
 def do2DPlot(f, name, i, fig, axes):
@@ -77,12 +106,35 @@ def do2DPlot(f, name, i, fig, axes):
     at = AnchoredText(name, loc=2, prop=dict(size=6), frameon=True)
     plot.axes.add_artist(at)
 
+    # Make a grid
+    x_ticks = plot.axes.get_xticks()
+    y_ticks = plot.axes.get_yticks()
+    for x in x_ticks:
+        N = len(y_ticks)
+        line = ln.Line2D([x]*N, y_ticks, linestyle=":", linewidth=0.6,
+                         color="k", alpha=0.7)
+        plot.axes.add_line(line)
+    for y in y_ticks:
+        N = len(x_ticks)
+        line = ln.Line2D(x_ticks, [y]*N, linestyle=":", linewidth=0.6,
+                         color="k", alpha=0.7)
+        plot.axes.add_line(line)
+
+
+def do1DPlotDensity(f, name, i):
+    sp = f.sphere(f.center, width)
+    # Because ParticleProfilePlot doesn't exist, I will do the following trick.
+    p = yt.create_profile(sp, ("PartType0", "density"),
+                          ("PartType0", "Masses"), weight_field=None,
+                          n_bins=40, accumulation=False)
+    return p
+
 
-def do1DPlot(f, name, i):
-    sp = f.sphere("max", width)
+def do1DPlotTemperature(f, name, i):
+    sp = f.sphere(f.center, width)
     # Because ParticleProfilePlot doesn't exist, I will do the following trick.
-    p = yt.create_profile(sp, ("PartType0", "density"),  ("PartType0", "Masses"),
-                          weight_field=None, n_bins=50,
-                          accumulation=False)
+    p = yt.create_profile(sp, ("PartType0", "Temperature"),
+                          ("PartType0", "Masses"), weight_field=None,
+                          n_bins=40, accumulation=False)
 
     return p

examples/GEAR/ZoomIn/profile_plot.py

+#!/usr/bin/env python3
+
+import yt
+from yt.units import Msun, kpc
+import matplotlib.pyplot as plt
+
+width = 20 * kpc
+limits_mass = (1e3 * Msun, 1e7 * Msun)
+
+
+def savePlot(profiles):
+    plt.figure(figsize=(8, 8))
+    plt.tight_layout(pad=0.4, w_pad=0.5, h_pad=0.5)
+    markers = ["s", "o"]
+    for i, p in enumerate(profiles[0]):
+        r = p.x.in_units("pc").d
+        mass = p["Masses"].in_units("Msun").d
+        plt.plot(r, mass, linestyle="-", marker=markers[i],
+                 markeredgecolor='none', linewidth=1.2, alpha=0.8)
+    plt.semilogx()
+    plt.semilogy()
+
+    plt.xlabel(r"$\mathrm{Radius}\/\mathrm{(pc)}$", fontsize='large')
+    plt.ylabel("$\mathrm{Integrated}\ \mathrm{mass}\ (Msun)}$",
+               fontsize='large')
+    plt.legend(profiles[1], loc=4, frameon=True, ncol=2, fancybox=True)
+    leg = plt.gca().get_legend()
+    ltext = leg.get_texts()
+    plt.setp(ltext, fontsize='small')
+    plt.grid(True)
+
+    plt.savefig("integrated_mass.png", bbox_inches='tight',
+                pad_inches=0.03, dpi=300)
+
+
+def doPlot(f, name, i):
+    sp = f.sphere(f.center, width)
+    # Because ParticleProfilePlot doesn't exist, I will do the following trick.
+    p = yt.create_profile(sp, "particle_radius", "Masses",
+                          weight_field=None, n_bins=50,
+                          accumulation=True)
+
+    return p

examples/GEAR/ZoomIn/projection_plot.py

@@ -3,17 +3,17 @@
 import yt
 from yt.units import kpc, g, cm, K
 from matplotlib.offsetbox import AnchoredText
+import numpy as np
 
 cmap_density = "algae"
 cmap_temperature = "magma"
 cmap_mass = "inferno"
-center = (1441.9954833984375000,
-          1666.1169433593750000,
-          1891.6248779296875000)
-small_width = 250 * kpc
-large_width = 3000 * kpc
-limits_density = (1e-10 * g / cm**2, 1e-2 * g / cm**2)
+cmap_metals = "coolwarm"
+small_width = 400 * kpc
+large_width = 400 * kpc
+limits_density = (1e-10 * g / cm**2, 1e2 * g / cm**2)
 limits_temperature = (10 * K, 1e5 * K)
+limits_metals = None
 limits_mass = None
 
 
@@ -58,10 +58,11 @@ def doDensityPlot(f, name, i, fig, axes):
     """
     direction = "x"
     field = ("PartType0", "density")
+    a = f.scale_factor
 
     # compute the projection
-    p = yt.ProjectionPlot(f, direction, field, center=center,
-                          width=small_width, buff_size=(800, 800))
+    p = yt.ProjectionPlot(f, direction, field, center=f.center,
+                          width=small_width * a, buff_size=(800, 800))
 
     # Compute the limits
     p.set_unit("density", "g / cm**2")
@@ -130,11 +131,12 @@ def doTemperaturePlot(f, name, i, fig, axes):
     """
     direction = "x"
     field = ("PartType0", "Temperature")
+    a = f.scale_factor
 
     # compute the projection
-    p = yt.ProjectionPlot(f, direction, field, center=center,
+    p = yt.ProjectionPlot(f, direction, field, center=f.center,
                           weight_field="density",
-                          width=small_width, buff_size=(800, 800))
+                          width=small_width * a, buff_size=(800, 800))
 
     # Compute the limits
     p.set_unit("Temperature", "K")
@@ -179,7 +181,7 @@ def doTemperaturePlot(f, name, i, fig, axes):
     plot.axes.add_artist(at)
 
 
-def doMassPlot(f, name, i, fig, axes):
+def doMassPlot(f, name, i, fig, axes, parttype):
     """
     Generate a mass projection (including dark matter) plot.
 
@@ -200,36 +202,116 @@ def doMassPlot(f, name, i, fig, axes):
 
     axes: list
         The list of axes to use for the plot
+
+    parttype: str
+        The name of the particle type to use
     """
-    direction = "y"
-    field = ("all", "Masses")
+    width = large_width
+    if parttype == "stars":
+        width = small_width
+        if name == "GEAR":
+            parttype = "PartType1"
+        else:
+            parttype = "PartType4"
+
+    direction = "x"
+    field = (parttype, "Masses")
+    a = f.scale_factor
 
     # compute the projection
-    p = yt.ParticleProjectionPlot(f, direction, field, center=center,
-                                  width=large_width)
+    p = yt.ParticleProjectionPlot(f, direction, field, center=f.center,
+                                  width=width * a)
 
     # # Compute the limits
-    # p.set_unit("masses", "Msun * kpc")
-    # data = p.to_fits_data()["masses"].data
-    # global limits_mass
-    # if limits_mass is None:
-    #     dmin = data.min()
-    #     if data.min() == 0:
-    #         dmin = 1e-6 * data.max()
-    #     else:
-    #         dmin *= 0.5
-    #     dmax = data.max() * 2
-    #     limits_mass = (dmin, dmax)
-    # else:
-    #     if limits_mass[0] > data.min():
-    #         print("WARNING: data below min", data.min())
-    #     if limits_mass[1] < data.max():
-    #         print("WARNING: data above max", data.max())
+    p.set_unit("Masses", "Msun")
+    data = p.to_fits_data()["Masses"].data
+    global limits_mass
+    if limits_mass is None:
+        dmin = np.nanmin(data)
+        if np.nanmin(data) == 0:
+            dmin = 1e-6 * np.nanmax(data)
+        else:
+            dmin *= 0.5
+        dmax = np.nanmax(data) * 2
+        limits_mass = (dmin, dmax)
+    else:
+        if limits_mass[0] > np.nanmin(data):
+            print("WARNING: data below min", np.nanmin(data))
+        if limits_mass[1] < np.nanmax(data):
+            print("WARNING: data above max", np.nanmax(data))
 
     # Adjust the plot
     p.set_cmap(field=field, cmap=cmap_mass)
     # p.set_log(field, True)
-    # p.set_zlim(field, limits_mass[0], limits_mass[1])
+    p.set_zlim(field, limits_mass[0], limits_mass[1])
+
+    # Draw it into the correct figure
+    plot = p.plots[field]
+
+    plot.figure = fig
+    plot.axes = axes[i].axes
+
+    # Add code name
+    at = AnchoredText(name, loc=2, prop=dict(size=6), frameon=True)
+    plot.axes.add_artist(at)
+
+    # plot color bar
+    if i != 0:
+        plot.cax = axes.cbar_axes[0]
+        p.hide_axes()
+    p._setup_plots()
+
+    # Add code name
+    at = AnchoredText(name, loc=2, prop=dict(size=6), frameon=True)
+    plot.axes.add_artist(at)
+
+
+def doMetalsPlot(f, name, i, fig, axes):
+    """
+    Generate a metal projection plot.
+
+    Parameters
+    ----------
+
+    f: yt dataset
+        The data set to use in the projection
+
+    name: str
+        The name to print on the plot
+
+    i: int
+        The index of the plot
+
+    fig: matplotlib figure
+        The figure to use
+
+    axes: list
+        The list of axes to use for the plot
+    """
+    direction = "x"
+    field = ("PartType0", "Metallicity")
+    a = f.scale_factor
+
+    # compute the projection
+    p = yt.ProjectionPlot(f, direction, field, center=f.center,
+                          width=small_width * a, buff_size=(800, 800),
+                          weight_field=("PartType0", "Density"))
+
+    # Compute the limits
+    # p.set_unit("metallicity", "g / cm**2")
+    data = p.to_fits_data()["Metallicity"].data
+    global limits_metals
+    if limits_metals is None:
+        limits_metals = (data.min(), data.max())
+    if limits_metals[0] > data.min():
+        print("WARNING: data below min", data.min())
+    if limits_metals[1] < data.max():
+        print("WARNING: data above max", data.max())
+
+    # Adjust the plot
+    p.set_cmap(field=field, cmap=cmap_metals)
+    p.set_log(field, True)
+    #p.set_zlim(field, limits_metals[0], limits_metals[1])
 
     # Draw it into the correct figure
     plot = p.plots[field]
@@ -247,6 +329,15 @@ def doMassPlot(f, name, i, fig, axes):
         p.hide_axes()
     p._setup_plots()
 
+    if i == 0:
+        z = 1. / f.scale_factor - 1.
+        text = "Redshift = %.2g" % z
+        prop = {
+            "color": "w"
+        }
+        at = AnchoredText(text, loc=3, prop=prop, frameon=False)
+        plot.axes.add_artist(at)
+
     # Add code name
     at = AnchoredText(name, loc=2, prop=dict(size=6), frameon=True)
     plot.axes.add_artist(at)

examples/GEAR/ZoomIn/star_plot.py

+#!/usr/bin/env python3
+
+import matplotlib.pyplot as plt
+import numpy as np
+from yt.units import kpc
+from yt.utilities.cosmology import Cosmology
+from multiprocessing import Pool
+from functools import partial
+from h5py import File
+
+width = 10 * kpc
+fe_sol_abund = 1.76603e-3
+mg_sol_abund = 9.24316e-4
+
+
+def saveSFRPlot(profiles):
+    plt.figure(figsize=(8, 8))
+    plt.tight_layout(pad=0.4, w_pad=0.5, h_pad=0.5)
+    markers = ["s", "o"]
+
+    sfr_fig = plt.figure()
+    mstar_fig = plt.figure()
+
+    for i, p in enumerate(profiles[0]):
+        masses = p[0]
+        form_time = p[1]
+        time_range = [0, 14]  # Gyr
+        n_bins = 1000
+        hist, bins = np.histogram(form_time, bins=n_bins, range=time_range)
+        inds = np.digitize(form_time, bins=bins)
+        time = (bins[:-1] + bins[1:])/2
+
+        sfr = np.array([masses[inds == j+1].sum()/(bins[j+1]-bins[j])
+                        for j in range(len(time))])
+        sfr[sfr == 0] = np.nan
+
+        sfr /= 1e9  # convert M / Gyr -> M / yr
+
+        plt.figure(sfr_fig.number)
+        plt.plot(time, sfr, linestyle="-", marker=markers[i],
+                 markeredgecolor="none", linewidth=1.2, alpha=0.8)
+
+        mstars = np.array([masses[inds == j+1].sum()
+                           for j in range(len(time))])
+        mstars = np.cumsum(mstars)
+        plt.figure(mstar_fig.number)
+        plt.plot(time, mstars, linestyle="-",
+                 markeredgecolor="none", linewidth=1.2, alpha=0.8)
+
+    plt.figure(sfr_fig.number)
+    plt.xlabel('Time  [Gyr]')
+    plt.ylabel('SFR  [M$_\odot$ yr$^{-1}$]')
+    plt.legend(profiles[1])
+    plt.savefig("sfr.png")
+
+    plt.figure(mstar_fig.number)
+    plt.xlabel('Time  [Gyr]')
+    plt.ylabel('Stellar mass [M$_\odot$]')
+    plt.legend(profiles[1])
+    plt.savefig("mstars.png")
+
+
+def doSFRPlot(f, name, i):
+    sp = f.sphere(f.center, width)
+    if name == "GEAR":
+        masses = sp["PartType1", "particle_mass"].in_units("Msun")
+        a = sp["PartType1", "StarFormationTime"]
+    else:
+        masses = sp["PartType4", "particle_mass"].in_units("Msun")
+        a = sp["PartType4", "BirthScaleFactors"]
+    cosmo = Cosmology(hubble_constant=f.hubble_constant,
+                      omega_matter=f.omega_matter,
+                      omega_lambda=f.omega_lambda)
+    z = np.array(1. / a - 1.)
+    with Pool() as p:
+        formation_time = p.map(
+            partial(cosmo.lookback_time, z_f=1e6), z)
+    formation_time = f.arr(list(formation_time), "s").in_units("Gyr")
+
+    return masses, formation_time
+
+
+def saveAbundancesPlot(profiles):
+    plt.figure(figsize=(8, 8))
+    plt.tight_layout(pad=0.4, w_pad=0.5, h_pad=0.5)
+    markers = ["s", "o"]
+
+    for i, p in enumerate(profiles[0]):
+        mg = np.log10(p[0] / mg_sol_abund)
+        fe = np.log10(p[1] / fe_sol_abund)
+
+        plt.plot(fe, mg - fe, linestyle="", marker=markers[i],
+                 markeredgecolor="none", alpha=0.8)
+
+    plt.xlim([-4, 0.5])
+    plt.ylim([-1, 1.5])
+    plt.xlabel('[Fe / H]')
+    plt.ylabel('[Mg / Fe]')
+    plt.legend(profiles[1])
+    plt.savefig("abundances.png")
+
+
+def doAbundancesPlot(f, name, i):
+    sp = f.sphere(f.center, width)
+    if name == "GEAR":
+        metals = sp["PartType1", "StarMetals"]
+        h = File(f.filename_template, "r")
+        names = h["Header"].attrs["ChimieLabels"]
+        names = names.split(b",")
+        # remove garbage
+        names = names[:-2]
+        # bytes to unicode
+        names = [name.decode() for name in names]
+        h.close()
+    else:
+        metals = sp["PartType4", "ElementAbundances"]
+        h = File(f.filename_template, "r")
+        sub = h["SubgridScheme"].attrs
+        names = []
+        for i in range(metals.shape[1]):
+            name = sub["Element %i" % i].decode()
+            names.append(name)
+
+    mg = names.index("Mg")
+    fe = names.index("Fe")
+
+    return metals[:, mg], metals[:, fe]

examples/GEAR/ZoomIn/velocity_plot.py

@@ -5,6 +5,7 @@ from yt.units import kpc
 import matplotlib.pyplot as plt
 
 width = 1500 * kpc
+DMO = True
 
 
 def save1DPlot(profiles):
@@ -20,7 +21,12 @@ def save1DPlot(profiles):
     plt.semilogy()
 
     plt.xlabel("$\mathrm{Velocity\ (km/s)}$", fontsize='large')
-    plt.ylabel(r"$\mathrm{Mass,}\/\mathrm{d}M\mathrm{/dlog}\/\mathrm{\rho}\/\mathrm{(M_{\odot})}$", fontsize='large')
+    if DMO:
+        plt.ylabel(r"$\mathrm{Mass}\/\mathrm{all}\/\mathrm{(M_{\odot})}$",
+                   fontsize='large')
+    else:
+        plt.ylabel(r"$\mathrm{Mass}\/\mathrm{Gas}\/\mathrm{(M_{\odot})}$",
+                   fontsize='large')
     plt.legend(profiles[1], loc=4, frameon=True, ncol=2, fancybox=True)
     leg = plt.gca().get_legend()
     ltext = leg.get_texts()
@@ -31,10 +37,18 @@ def save1DPlot(profiles):
 
 
 def do1DPlot(f, name, i):
-    sp = f.sphere("max", width)
+    global DMO
+    part_type = "all"
+    if f.particle_type_counts["PartType0"] != 0:
+        DMO = False
+        part_type = "PartType0"
+
+    a = f.scale_factor
+    sp = f.sphere("c", width * a)
 
     # Because ParticleProfilePlot doesn't exist, I will do the following trick.
-    p = yt.create_profile(sp, ("PartType0", "velocity_magnitude"),  ("PartType0", "Masses"),
+    p = yt.create_profile(sp, (part_type, "particle_velocity_magnitude"),
+                          (part_type, "Masses"),
                           weight_field=None, n_bins=50,
                           accumulation=False)
 

tools/task_plots/iplot_tasks.py

@@ -40,7 +40,7 @@ matplotlib.use('TkAgg')
 import numpy as np
 import matplotlib.backends.backend_tkagg as tkagg
 from matplotlib.figure import Figure
-import Tkinter as tk
+import tkinter as tk
 import matplotlib.collections as collections
 import matplotlib.ticker as plticker
 import pylab as pl