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