GEAR: update zoomin

examples/GEAR/ZoomIn/add_fields.py

@@ -3,7 +3,6 @@
 import numpy as np
 import yt.utilities.physical_constants as constants
 from yt.mods import YTArray
-from yt import derived_field
 
 
 def addMassDeposit(f):

examples/GEAR/ZoomIn/cleanupSwift.py

@@ -37,6 +37,7 @@ for i in range(NPartType):
     fields = [
         ("Coordinates", "Coordinates", h),
         ("Masses", "Masses", h),
+        ("StarInitMass", "BirthMasses", h),
         ("Velocities", "Velocities", 1. / a**0.5),
         ("Density", "Densities", 1. / h**2),
         ("Entropies", "Entropies", 1.),

examples/GEAR/ZoomIn/make_image.py

@@ -12,6 +12,8 @@ import star_plot
 import profile_plot
 import matplotlib.pyplot as plt
 from mpl_toolkits.axes_grid1 import AxesGrid
+from yt.units import kpc
+import numpy as np
 
 # Parameters
 
@@ -20,19 +22,20 @@ snap = int(sys.argv[-1])
 swift = "swift/snapshot_%04i.hdf5" % snap
 gear = "gear/snapshot_%04i.hdf5" % snap
 
+width = 50 * kpc
 
 do_dmo = False
 do_hydro = False
-do_stars = True
+do_stars = False
 do_feedback = True
 do_plot = {
     # DMO
     "projection_mass": do_dmo,
     "velocity": do_dmo,
-    "halo_distribution": False,  # very slow
+    "halo_distribution": True,  # very slow
     "profile": do_dmo,
     # hydro
-    "projection_density": True, #do_hydro,
+    "projection_density": do_hydro,
     "projection_temperature": do_hydro,
     "phase_1d_density": do_hydro,
     "phase_1d_temperature": do_hydro,
@@ -44,6 +47,8 @@ do_plot = {
     # feedback
     "abundances": do_feedback,
     "projection_metals": do_feedback,
+    "iron_distribution": do_feedback,
+    "mass_distribution": do_feedback
 }
 
 # Generate the figures
@@ -103,6 +108,8 @@ data = {
     "SFR": ([], []),
     "profile": ([], []),
     "abundances": ([], []),
+    "iron_distribution": ([], []),
+    "mass_distribution": ([], []),
 }
 
 names = []
@@ -164,6 +171,14 @@ def savePlot():
         data["abundances"][1].extend(names)
         star_plot.saveAbundancesPlot(data["abundances"])
 
+    if do_plot["iron_distribution"]:
+        data["iron_distribution"][1].extend(names)
+        star_plot.saveIronDistributionPlot(data["iron_distribution"])
+
+    if do_plot["mass_distribution"]:
+        data["mass_distribution"][1].extend(names)
+        star_plot.saveMassDistributionPlot(data["mass_distribution"])
+
 
 def doPlot(filename, i, name, center):
     names.append(name)
@@ -172,7 +187,9 @@ def doPlot(filename, i, name, center):
 
     if center is None:
         center = f.find_max("Density")[1]
+
     f.center = center
+    f.width = width * f.scale_factor
 
     if (do_plot["projection_temperature"] or do_plot["phase_2d"]):
         add_fields.addTemperature(f)
@@ -258,9 +275,19 @@ def doPlot(filename, i, name, center):
         p = star_plot.doAbundancesPlot(f, name, i)
         data["abundances"][0].append(p)
 
+    if do_plot["iron_distribution"]:
+        p = star_plot.doIronDistributionPlot(f, name, i)
+        data["iron_distribution"][0].append(p)
+
+    if do_plot["mass_distribution"]:
+        p = star_plot.doMassDistributionPlot(f, name, i)
+        data["mass_distribution"][0].append(p)
+
     return center
 
 
-center = doPlot(gear, 1, "GEAR", center=None)
+center = None
+# center = np.array([1724.33547783, 1802.56263082, 1785.09893269])
+center = doPlot(gear, 1, "GEAR", center=center)
 doPlot(swift, 0, "SWIFT", center=center)
 savePlot()

examples/GEAR/ZoomIn/phase_plot.py

@@ -7,7 +7,6 @@ import matplotlib.pyplot as plt
 import matplotlib.lines as ln
 import numpy as np
 
-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)
@@ -122,7 +121,7 @@ def do2DPlot(f, name, i, fig, axes):
 
 
 def do1DPlotDensity(f, name, i):
-    sp = f.sphere(f.center, width)
+    sp = f.sphere(f.center, f.width)
     # Because ParticleProfilePlot doesn't exist, I will do the following trick.
     p = yt.create_profile(sp, ("PartType0", "density"),
                           ("PartType0", "Masses"), weight_field=None,
@@ -131,7 +130,7 @@ def do1DPlotDensity(f, name, i):
 
 
 def do1DPlotTemperature(f, name, i):
-    sp = f.sphere(f.center, width)
+    sp = f.sphere(f.center, f.width)
     # Because ParticleProfilePlot doesn't exist, I will do the following trick.
     p = yt.create_profile(sp, ("PartType0", "Temperature"),
                           ("PartType0", "Masses"), weight_field=None,

examples/GEAR/ZoomIn/profile_plot.py

@@ -4,7 +4,6 @@ import yt
 from yt.units import Msun, kpc
 import matplotlib.pyplot as plt
 
-width = 20 * kpc
 limits_mass = (1e3 * Msun, 1e7 * Msun)
 
 
@@ -34,7 +33,7 @@ def savePlot(profiles):
 
 
 def doPlot(f, name, i):
-    sp = f.sphere(f.center, width)
+    sp = f.sphere(f.center, f.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,

examples/GEAR/ZoomIn/projection_plot.py

@@ -9,8 +9,6 @@ cmap_density = "algae"
 cmap_temperature = "magma"
 cmap_mass = "inferno"
 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
@@ -58,11 +56,10 @@ 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=f.center,
-                          width=small_width * a, buff_size=(800, 800))
+                          width=f.width, buff_size=(800, 800))
 
     # Compute the limits
     p.set_unit("density", "g / cm**2")
@@ -136,7 +133,7 @@ def doTemperaturePlot(f, name, i, fig, axes):
     # compute the projection
     p = yt.ProjectionPlot(f, direction, field, center=f.center,
                           weight_field="density",
-                          width=small_width * a, buff_size=(800, 800))
+                          width=f.width, buff_size=(800, 800))
 
     # Compute the limits
     p.set_unit("Temperature", "K")
@@ -206,9 +203,7 @@ def doMassPlot(f, name, i, fig, axes, parttype):
     parttype: str
         The name of the particle type to use
     """
-    width = large_width
     if parttype == "stars":
-        width = small_width
         if name == "GEAR":
             parttype = "PartType1"
         else:
@@ -216,11 +211,10 @@ def doMassPlot(f, name, i, fig, axes, parttype):
 
     direction = "x"
     field = (parttype, "Masses")
-    a = f.scale_factor
 
     # compute the projection
     p = yt.ParticleProjectionPlot(f, direction, field, center=f.center,
-                                  width=width * a)
+                                  width=f.width)
 
     # # Compute the limits
     p.set_unit("Masses", "Msun")
@@ -290,11 +284,10 @@ def doMetalsPlot(f, name, i, fig, axes):
     """
     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),
+                          width=f.width, buff_size=(800, 800),
                           weight_field=("PartType0", "Density"))
 
     # Compute the limits
@@ -311,7 +304,8 @@ def doMetalsPlot(f, name, i, fig, axes):
     # 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])
+    if limits_metals[0] != limits_metals[1]:
+        p.set_zlim(field, limits_metals[0], limits_metals[1])
 
     # Draw it into the correct figure
     plot = p.plots[field]

examples/GEAR/ZoomIn/star_plot.py

@@ -2,13 +2,12 @@
 
 import matplotlib.pyplot as plt
 import numpy as np
-from yt.units import kpc
+from yt.units import Msun
 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
 
@@ -21,11 +20,22 @@ def saveSFRPlot(profiles):
     sfr_fig = plt.figure()
     mstar_fig = plt.figure()
 
+    # get the boundary
+    t_max = None
+    for p in profiles[0]:
+        if t_max is None:
+            t_max = p[1].max()
+        else:
+            tmp = p[1].max()
+            if t_max < tmp:
+                t_max = tmp
+    t_max = float(t_max)
+
     for i, p in enumerate(profiles[0]):
         masses = p[0]
         form_time = p[1]
-        time_range = [0, 14]  # Gyr
-        n_bins = 1000
+        time_range = [0, t_max]
+        n_bins = 100
         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
@@ -42,7 +52,7 @@ def saveSFRPlot(profiles):
 
         mstars = np.array([masses[inds == j+1].sum()
                            for j in range(len(time))])
-        mstars = np.cumsum(mstars)
+        mstars = np.cumsum(mstars) / 1e6
         plt.figure(mstar_fig.number)
         plt.plot(time, mstars, linestyle="-",
                  markeredgecolor="none", linewidth=1.2, alpha=0.8)
@@ -55,18 +65,18 @@ def saveSFRPlot(profiles):
 
     plt.figure(mstar_fig.number)
     plt.xlabel('Time  [Gyr]')
-    plt.ylabel('Stellar mass [M$_\odot$]')
+    plt.ylabel('Stellar mass [$10^6$ M$_\odot$]')
     plt.legend(profiles[1])
     plt.savefig("mstars.png")
 
 
 def doSFRPlot(f, name, i):
-    sp = f.sphere(f.center, width)
+    sp = f.sphere(f.center, f.width)
     if name == "GEAR":
-        masses = sp["PartType1", "particle_mass"].in_units("Msun")
+        masses = (sp["PartType1", "StarInitMass"] * 1e10 * Msun).in_units("Msun")
         a = sp["PartType1", "StarFormationTime"]
     else:
-        masses = sp["PartType4", "particle_mass"].in_units("Msun")
+        masses = (sp["PartType4", "StarInitMass"] * 1e10 * Msun).in_units("Msun")
         a = sp["PartType4", "BirthScaleFactors"]
     cosmo = Cosmology(hubble_constant=f.hubble_constant,
                       omega_matter=f.omega_matter,
@@ -90,18 +100,20 @@ def saveAbundancesPlot(profiles):
         fe = np.log10(p[1] / fe_sol_abund)
 
         plt.plot(fe, mg - fe, linestyle="", marker=markers[i],
-                 markeredgecolor="none", alpha=0.8)
+                 markeredgecolor="none", alpha=0.8, markersize=3)
 
-    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")
 
+    plt.xlim([-4, 0.5])
+    plt.ylim([-1, 1.5])
 
-def doAbundancesPlot(f, name, i):
-    sp = f.sphere(f.center, width)
+    plt.savefig("abundances_zoom.png")
+
+
+def getElement(f, sp, name, el):
     if name == "GEAR":
         metals = sp["PartType1", "StarMetals"]
         h = File(f.filename_template, "r")
@@ -121,7 +133,145 @@ def doAbundancesPlot(f, name, i):
             name = sub["Element %i" % i].decode()
             names.append(name)
 
-    mg = names.index("Mg")
-    fe = names.index("Fe")
+    ind = names.index(el)
+    return metals[:, ind]
+
+
+def doAbundancesPlot(f, name, i):
+    sp = f.sphere(f.center, f.width)
+
+    mg = getElement(f, sp, name, "Mg")
+    fe = getElement(f, sp, name, "Fe")
+
+    return mg, fe
+
+
+def saveIronDistributionPlot(profiles):
+    plt.figure(figsize=(8, 8))
+    plt.tight_layout(pad=0.4, w_pad=0.5, h_pad=0.5)
+
+    fe_max = max(p.max() for p in profiles[0])
+    fe_min = min(p.min() for p in profiles[0])
+    fe_min = np.log10((fe_min + 1e-10) / fe_sol_abund)
+    fe_max = np.log10(fe_max / fe_sol_abund)
+
+    for i, p in enumerate(profiles[0]):
+        fe = np.log10(p / fe_sol_abund)
+
+        plt.hist(fe, bins=20, range=[fe_min, fe_max], histtype="step",
+                 density=True, alpha=0.8)
+
+    plt.xlabel('[Fe / H]')
+    plt.ylabel("Fraction of stars")
+    plt.legend(profiles[1])
+    plt.savefig("iron_distribution.png")
+
+
+def doIronDistributionPlot(f, name, i):
+    sp = f.sphere(f.center, f.width)
+    # Because ParticleProfilePlot doesn't exist, I will do the following trick.
+    fe = getElement(f, sp, name, "Fe")
+    return fe
+
+
+def saveMassDistributionPlot(profiles):
+    # Do the time - stars
+    plt.figure(figsize=(8, 8))
+    plt.tight_layout(pad=0.4, w_pad=0.5, h_pad=0.5)
+    markers = ["s", "o"]
+    zorder = [2.5, 2.]
+
+    for i, p in enumerate(profiles[0]):
+        plt.plot(p["birth_time"], p["stars"], marker=markers[i],
+                 markeredgecolor="none", linestyle="", alpha=0.8,
+                 zorder=zorder[i])
+    plt.semilogx()
+    plt.semilogy()
+
+    plt.xlabel("Star formation time (Gyr)", fontsize='large')
+    plt.ylabel("Mass of the star $(M_{\odot})$", fontsize='large')
+    plt.legend(profiles[1], 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("mass_lost_stars.png", bbox_inches='tight',
+                pad_inches=0.03, dpi=300)
+
+    # Do the stars
+    plt.figure(figsize=(8, 8))
+    plt.tight_layout(pad=0.4, w_pad=0.5, h_pad=0.5)
+
+    mass_max = max(p["stars"].max() for p in profiles[0])
+    mass_min = min(p["stars"].min() for p in profiles[0])
+
+    for i, mass in enumerate(profiles[0]):
+        plt.hist(mass["stars"], bins=100, range=[mass_min, mass_max],
+                 histtype="step", density=True, alpha=0.8)
+    plt.semilogx()
+    plt.semilogy()
+
+    plt.xlabel("Mass of the stars $(M_{\odot})$", fontsize='large')
+    plt.ylabel("Fraction of stars", fontsize='large')
+    plt.legend(profiles[1], 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("mass_distribution_stars.png", bbox_inches='tight',
+                pad_inches=0.03, dpi=300)
+
+    # Do the gas
+    plt.figure(figsize=(8, 8))
+    plt.tight_layout(pad=0.4, w_pad=0.5, h_pad=0.5)
+
+    mass_max = max(p["gas"].max() for p in profiles[0])
+    mass_min = min(p["gas"].min() for p in profiles[0])
+
+    for i, mass in enumerate(profiles[0]):
+        plt.hist(mass["gas"], bins=100, range=[mass_min, mass_max],
+                 histtype="step", density=True, alpha=0.8)
+    plt.semilogx()
+    plt.semilogy()
+
+    plt.xlabel("Mass of the gas $(M_{\odot})$", fontsize='large')
+    plt.ylabel("Fraction of gas", fontsize='large')
+    plt.legend(profiles[1], 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("mass_distribution_gas.png", bbox_inches='tight',
+                pad_inches=0.03, dpi=300)
+
+
+def doMassDistributionPlot(f, name, i):
+    sp = f.sphere(f.center, f.width)
+
+    gas = sp[("PartType0", "Masses")].in_units("Msun").d
+
+    if name == "GEAR":
+        stars = sp[("PartType1", "Masses")].in_units("Msun").d
+        a = sp["PartType1", "StarFormationTime"]
+    else:
+        stars = sp[("PartType4", "Masses")].in_units("Msun").d
+        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 metals[:, mg], metals[:, fe]
+    d = {
+        "gas": gas,
+        "stars": stars,
+        "birth_time": formation_time
+    }
+    return d

examples/GEAR/ZoomIn/velocity_plot.py

 #!/usr/bin/env python3
 
 import yt
-from yt.units import kpc
 import matplotlib.pyplot as plt
 
-width = 1500 * kpc
 DMO = True
 
 
@@ -43,8 +41,7 @@ def do1DPlot(f, name, i):
         DMO = False
         part_type = "PartType0"
 
-    a = f.scale_factor
-    sp = f.sphere("c", width * a)
+    sp = f.sphere("c", f.width)
 
     # Because ParticleProfilePlot doesn't exist, I will do the following trick.
     p = yt.create_profile(sp, (part_type, "particle_velocity_magnitude"),