Merge branch 'darwin/sink_example_homogeneousbox' into 'master'

Sink homogeneous box example

See merge request !1941

examples/SinkParticles/HomogeneousBox/README

+# Intro
+This example is a non cosmological homogeneous box with gas only. The cooling and small perturbations fragment the box into small compact regions that create sink particles. As the sinks accrete gas, they star spawning stars. As the first stars explode in supernovae, the first sink particles cannot accrete gas anymore. The simulations runs 282 Myr until the first stars explode and remove the gas.
+
+The default level is 5 (N_particle = 2^(3*level)). It is quick. If you want to try the example with higher levels, we recommend using HPC facilities.
+
+This example tests the impact of different parameters on the sink formation (e.g. cut_off_radius, density_threshold, ...), accretion and star spawning. It also test the effects of feedback (e.g. supernovae_efficiency) and cooling (equilibrium cooling with grackle_0 mode versu non equilibrium cooling with grackle_X, X = 1, 2 or 3) on the sink particles and star formation.
+
+This example allows you to play with the sink parameters to see how they impact the simulation. You can also change parameters in the feedback, the cooling, etc.
+
+# Configure
+To run this example with GEAR model,
+
+'''
+./configure --disable-mpi --with-chemistry=GEAR_10 --with-cooling=grackle_0 --with-stars=GEAR --with-star-formation=GEAR --with-feedback=GEAR --with-sink=GEAR --with-kernel=wendland-C2 --with-adaptive-softening --with-grackle=path/to/grackle
+
+and then
+
+make -j
+
+You can remove the adaptive softening. In this case, you may need to change the default `max_physical_baryon_softening` value.
+
+Other sink models can be probed by changing swift configuration and adding the relevant parameters in params.yml.
+
+# ICs
+The run.sh script calls `makeIC.py' script with default values. You can experiment by changing the ICs. Run `python3 makeIC.py --help` to get the list of parameters.
+
+# Run
+Type run.sh, and let's go!
+
+# Changing the cooling
+You can change grackle mode to 1,2, or 3 if you want. You may also want to change the simulation end time.
+
+For grackle_0 (equilibrium mode), the first sink form roughly at 0.260 (internal units). The end time is 0.282 Myr. The simulation runs in ~ 2 minutes.
+
+Notice that, for grackle_X with X from 1 to 3, the level 5 simulation does not have enough particles to run for long times. At 0.5 Gyr, swift ends because there are not enough particles per top-level cell. However, note how different the collapse from grackle_0 is!
+
+# For sink debugging
+This example was used to debug the sink's tasks and other sink-related bugs. The level (N_part = 2^(3*level)) was 8 during this debugging phase.

examples/SinkParticles/HomogeneousBox/getChemistryTable.sh

+#!/bin/bash
+
+wget http://virgodb.cosma.dur.ac.uk/swift-webstorage/FeedbackTables/POPIIsw.h5

examples/SinkParticles/HomogeneousBox/getGrackleCoolingTable.sh

+#!/bin/bash
+
+wget http://virgodb.cosma.dur.ac.uk/swift-webstorage/CoolingTables/CloudyData_UVB=HM2012.h5

examples/SinkParticles/HomogeneousBox/makeIC.py

+################################################################################
+# This file is part of SWIFT.
+# Copyright (c) 2022 Yves Revaz (yves.revaz@epfl.ch)
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Lesser General Public License as published
+# by the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU Lesser General Public License
+# along with this program.  If not, see <http://www.gnu.org/licenses/>.
+#
+################################################################################
+
+import h5py
+import numpy as np
+from optparse import OptionParser
+from astropy import units
+from astropy import constants
+
+
+def parse_options():
+
+    usage = "usage: %prog [options] file"
+    parser = OptionParser(usage=usage)
+
+    parser.add_option(
+        "--lJ",
+        action="store",
+        dest="lJ",
+        type="float",
+        default=0.250,
+        help="Jeans wavelength in box size unit",
+    )
+
+    parser.add_option(
+        "--rho",
+        action="store",
+        dest="rho",
+        type="float",
+        default=0.1,
+        help="Mean gas density in atom/cm3",
+    )
+
+    parser.add_option(
+        "--mass",
+        action="store",
+        dest="mass",
+        type="float",
+        default=50,
+        help="Gas particle mass in solar mass",
+    )
+
+    parser.add_option(
+        "--level",
+        action="store",
+        dest="level",
+        type="int",
+        default=5,
+        help="Resolution level: N = (2**l)**3",
+    )
+
+    parser.add_option(
+        "-o",
+        action="store",
+        dest="outputfilename",
+        type="string",
+        default="box.dat",
+        help="output filename",
+    )
+
+    (options, args) = parser.parse_args()
+
+    files = args
+
+    return files, options
+
+
+########################################
+# main
+########################################
+
+files, opt = parse_options()
+
+# define standard units
+UnitMass_in_cgs = 1.988409870698051e43  # 10^10 M_sun in grams
+UnitLength_in_cgs = 3.0856775814913673e21  # kpc in centimeters
+UnitVelocity_in_cgs = 1e5  # km/s in centimeters per second
+UnitCurrent_in_cgs = 1  # Amperes
+UnitTemp_in_cgs = 1  # Kelvin
+UnitTime_in_cgs = UnitLength_in_cgs / UnitVelocity_in_cgs
+
+UnitMass = UnitMass_in_cgs * units.g
+UnitLength = UnitLength_in_cgs * units.cm
+UnitTime = UnitTime_in_cgs * units.s
+UnitVelocity = UnitVelocity_in_cgs * units.cm / units.s
+
+np.random.seed(1)
+
+# Number of particles
+N = (2 ** opt.level) ** 3  # number of particles
+
+# Mean density
+rho = opt.rho  # atom/cc
+rho = rho * constants.m_p / units.cm ** 3
+
+# Gas particle mass
+m = opt.mass  # in solar mass
+m = m * units.Msun
+
+# Gas mass in the box
+M = N * m
+
+# Size of the box
+L = (M / rho) ** (1 / 3.0)
+
+# Jeans wavelength in box size unit
+lJ = opt.lJ
+lJ = lJ * L
+
+# Gravitational constant
+G = constants.G
+
+# Jeans wave number
+kJ = 2 * np.pi / lJ
+# Velocity dispersion
+sigma = np.sqrt(4 * np.pi * G * rho) / kJ
+
+
+print("Number of particles                   : {}".format(N))
+print("Equivalent velocity dispertion        : {}".format(sigma.to(units.m / units.s)))
+
+# Convert to code units
+m = m.to(UnitMass).value
+L = L.to(UnitLength).value
+rho = rho.to(UnitMass / UnitLength ** 3).value
+sigma = sigma.to(UnitVelocity).value
+
+# Generate the particles
+pos = np.random.random([N, 3]) * np.array([L, L, L])
+vel = np.zeros(N)
+mass = np.ones(N) * m
+u = np.ones(N) * sigma ** 2
+ids = np.arange(N)
+h = np.ones(N) * 3 * L / N ** (1 / 3.0)
+rho = np.ones(N) * rho
+
+print("Inter-particle distance (code unit)   : {}".format(L / N ** (1 / 3.0)))
+
+
+# File
+fileOutput = h5py.File(opt.outputfilename, "w")
+print("{} saved.".format(opt.outputfilename))
+
+# Header
+grp = fileOutput.create_group("/Header")
+grp.attrs["BoxSize"] = [L, L, L]
+grp.attrs["NumPart_Total"] = [N, 0, 0, 0, 0, 0]
+grp.attrs["NumPart_Total_HighWord"] = [0, 0, 0, 0, 0, 0]
+grp.attrs["NumPart_ThisFile"] = [N, 0, 0, 0, 0, 0]
+grp.attrs["Time"] = 0.0
+grp.attrs["NumFileOutputsPerSnapshot"] = 1
+grp.attrs["MassTable"] = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
+grp.attrs["Flag_Entropy_ICs"] = [0, 0, 0, 0, 0, 0]
+grp.attrs["Dimension"] = 3
+
+
+# Units
+grp = fileOutput.create_group("/Units")
+grp.attrs["Unit length in cgs (U_L)"] = UnitLength_in_cgs
+grp.attrs["Unit mass in cgs (U_M)"] = UnitMass_in_cgs
+grp.attrs["Unit time in cgs (U_t)"] = UnitTime_in_cgs
+grp.attrs["Unit current in cgs (U_I)"] = UnitCurrent_in_cgs
+grp.attrs["Unit temperature in cgs (U_T)"] = UnitTemp_in_cgs
+
+
+# Particle group
+grp = fileOutput.create_group("/PartType0")
+grp.create_dataset("Coordinates", data=pos, dtype="d")
+grp.create_dataset("Velocities", data=vel, dtype="f")
+grp.create_dataset("Masses", data=mass, dtype="f")
+grp.create_dataset("SmoothingLength", data=h, dtype="f")
+grp.create_dataset("InternalEnergy", data=u, dtype="f")
+grp.create_dataset("ParticleIDs", data=ids, dtype="L")
+grp.create_dataset("Densities", data=rho, dtype="f")
+
+fileOutput.close()

examples/SinkParticles/HomogeneousBox/params.yml

+# Define the system of units to use internally.
+InternalUnitSystem:
+  UnitMass_in_cgs:     1.988409870698051e+43   # 10^10 solar masses 
+  UnitLength_in_cgs:   3.0856775814913673e+21  # 1 kpc 
+  UnitVelocity_in_cgs: 1e5   # km/s
+  UnitCurrent_in_cgs:  1   # Amperes
+  UnitTemp_in_cgs:     1   # Kelvin
+
+# Parameters for the self-gravity scheme
+Gravity:
+  MAC:                           adaptive  # Choice of mulitpole acceptance criterion: 'adaptive' OR 'geometric'.
+  epsilon_fmm:                   0.001     # Tolerance parameter for the adaptive multipole acceptance criterion.
+  theta_cr:                      0.7       # Opening angle for the purely gemoetric criterion.
+  eta:          0.025               # Constant dimensionless multiplier for time integration.
+  theta:        0.7                 # Opening angle (Multipole acceptance criterion).
+  max_physical_baryon_softening: 0.005  # Physical softening length (in internal units)
+  mesh_side_length:        32
+
+# Parameters governing the time integration (Set dt_min and dt_max to the same value for a fixed time-step run.)
+TimeIntegration:
+  time_begin:        0.    # The starting time of the simulation (in internal units).
+  time_end:          0.282   #500 Myr # The end time of the simulation (in internal units).
+  dt_min:            1e-10  # The minimal time-step size of the simulation (in internal units).
+  dt_max:            1e-1  # The maximal time-step size of the simulation (in internal units).
+
+# Parameters governing the snapshots
+Snapshots:
+  subdir: snap
+  basename:   snapshot      # Common part of the name of output files
+  time_first: 0. #230 Myr # (Optional) Time of the first output if non-cosmological time-integration (in internal units)
+  delta_time: 1e-3        # Time difference between consecutive outputs (in internal units)
+
+Scheduler:
+  cell_extra_gparts: 10000      # (Optional) Number of spare sparts per top-level allocated at rebuild time for on-the-fly creation.
+  cell_extra_sinks: 10000       # (Optional) Number of spare sparts per top-level allocated at rebuild time for on-the-fly creation.
+  cell_extra_sparts: 10000      # (Optional) Number of spare sparts per top-level allocated at rebuild time for on-the-fly creation.
+  max_top_level_cells: 3        #
+
+# Parameters governing the conserved quantities statistics
+Statistics:
+  delta_time:           1e-3     # Time between statistics output
+  time_first:             0.     # (Optional) Time of the first stats output if non-cosmological time-integration (in internal units)
+
+# Parameters related to the initial conditions
+InitialConditions:
+  file_name:          ICs_homogeneous_box.hdf5
+  periodic:                    1    # Are we running with periodic ICs?
+  shift:              [10.0,10.0,10.0]
+
+# Parameters for the hydrodynamics scheme
+SPH:
+  resolution_eta:        1.2348   # Target smoothing length in units of the mean inter-particle separation (1.2348 == 57Ngbs with the Wendland C2 kernel).
+  CFL_condition:         0.1      # Courant-Friedrich-Levy condition for time integration.
+  h_max:                 5
+  minimal_temperature:   1
+
+
+# Cooling with Grackle 3.0
+GrackleCooling:
+  cloudy_table: CloudyData_UVB=HM2012.h5 # Name of the Cloudy Table (available on the grackle bitbucket repository)
+  with_UV_background: 0 # Enable or not the UV background
+  redshift: -1 # Redshift to use (-1 means time based redshift)
+  with_metal_cooling: 1 # Enable or not the metal cooling
+  provide_volumetric_heating_rates: 0 # User provide volumetric heating rates
+  provide_specific_heating_rates: 0 # User provide specific heating rates
+  self_shielding_method: -1 # Grackle (<= 3) or Gear self shielding method
+  self_shielding_threshold_atom_per_cm3: 0.007  # Required only with GEAR's self shielding. Density threshold of the self shielding
+  max_steps: 1000
+  convergence_limit: 1e-2
+  thermal_time_myr: 5
+  maximal_density_Hpcm3: -1    # Maximal density (in hydrogen atoms/cm^3) for cooling. Higher densities are floored to this value to ensure grackle works properly when interpolating beyond the cloudy_table maximal density. A value < 0 deactivates this parameter.
+
+GEARChemistry:
+  initial_metallicity: -5
+
+GEARFeedback:
+  supernovae_energy_erg: 1e51           # supernovae energy, used only for SNIa
+  supernovae_efficiency: 0.1           # supernovae energy efficiency, used for both SNIa and SNII
+  yields_table: POPIIsw.h5
+  yields_table_first_stars: POPIIsw.h5
+  discrete_yields: 1
+  imf_transition_metallicity: -5         # Maximal metallicity ([Fe/H]) for a first star (0 to deactivate).
+  elements: [Fe, Mg, O, C, Al, Ca, Ba, Zn, Eu]              # Elements to read in the yields table. The number of element should be one less than the number of elements (N) requested during the configuration (--with-chemistry=GEAR_N).
+
+
+GEARSink:
+  cut_off_radius: 5e-3                  # Cut off radius of all the sinks in internal units.
+  f_acc: 0.1
+  maximal_temperature:  100            # Upper limit to the temperature of a star forming particle
+  density_threshold:   1.67e-23        # Density threashold in g/cm3 (1.67e-24 =1acc)
+  size_of_calibration_sample: 100000   # Size of the calibration sample  
+  stellar_particle_mass: 50             # Mass of the stellar particle representing the low mass stars, in solar mass
+  minimal_discrete_mass: 8              # Minimal mass of stars represented by discrete particles, in solar mass
+  stellar_particle_mass_first_stars: 50  # Mass of the stellar particle representing the low mass stars, in solar mass
+  minimal_discrete_mass_first_stars: 8   # Minimal mass of stars represented by discrete particles, in solar mass
+  star_spawning_sigma_factor: 0.5        # Factor to rescale the velocity dispersion of the stars when they are spawned. (Default: 0.2)
+  sink_formation_contracting_gas_criterion: 1     # (Optional) Activate the contracting gas check for sink formation. (Default: 1)
+  sink_formation_smoothing_length_criterion: 1    # (Optional) Activate the smoothing length check for sink formation. (Default: 1)
+  sink_formation_jeans_instability_criterion: 1   # (Optional) Activate the two Jeans instability checks for sink formation. (Default: 1)
+  sink_formation_bound_state_criterion: 1         # (Optional) Activate the bound state check for sink formation. (Default: 1)
+  sink_formation_overlapping_sink_criterion: 1    # (Optional) Activate the overlapping sink check for sink formation. (Default: 1)
+  disable_sink_formation: 0                       # (Optional) Disable sink formation. (Default: 0)
+

examples/SinkParticles/HomogeneousBox/plot_gas_density.py

+"""
+Makes a gas density projection plot. Uses the swiftsimio library.
+"""
+
+import matplotlib.pyplot as plt
+import numpy as np
+
+import swiftsimio as sw
+from swiftsimio.visualisation.projection import project_gas
+
+import unyt
+from unyt import mh, cm
+from tqdm import tqdm
+from matplotlib.colors import LogNorm
+from matplotlib.animation import FuncAnimation
+
+# %%
+# Constants; these could be put in the parameter file but are rarely changed.
+image_resolution = 1024
+
+# Plotting controls
+cmap = "inferno"
+
+# %%
+
+
+def get_gas_mu(data: sw.SWIFTDataset) -> np.array:
+    """
+    Return the mean molecular weight of the gas.
+    """
+    # Get the cooling model
+    cooling = data.metadata.subgrid_scheme["Cooling Model"]
+
+    # Use a variable for better readability
+    gas = data.gas
+
+    # Get rho
+    rho = gas.densities.to(unyt.g / unyt.cm ** 3)  # self.Rho(units='g/cm3')
+
+    # hydrogen mass in gram
+    mh.convert_to_cgs()
+    mH_in_g = mh
+
+    if cooling == b"Grackle3":  # grackle3
+        nHI = gas.hi * rho / (mH_in_g)
+        nHII = gas.hii * rho / (mH_in_g)
+        nHeI = gas.he_i * rho / (4 * mH_in_g)
+        nHeII = gas.he_ii * rho / (4 * mH_in_g)
+        nHeIII = gas.he_iii * rho / (4 * mH_in_g)
+        nH2I = gas.h2_i * rho / (2 * mH_in_g)
+        nH2II = gas.h2_ii * rho / (2 * mH_in_g)
+        nHDI = gas.hdi * rho / (3 * mH_in_g)
+
+        nel = nHII + nHeII + 2 * nHeIII + nH2II
+        mu = (
+            (nHI + nHII) + (nHeI + nHeII + nHeIII) * 4 + (nH2I + nH2II) * 2 + nHDI * 3
+        ) / (nHI + nHII + nHeI + nHeII + nHeIII + nH2I + nH2II + nHDI + nel)
+        return mu
+
+    elif cooling == b"Grackle2":  # grackle2
+        nHI = gas.hi * rho / (mH_in_g)
+        nHII = gas.hii * rho / (mH_in_g)
+        nHeI = gas.he_i * rho / (4 * mH_in_g)
+        nHeII = gas.he_ii * rho / (4 * mH_in_g)
+        nHeIII = gas.he_iii * rho / (4 * mH_in_g)
+        nH2I = gas.h2_i * rho / (2 * mH_in_g)
+        nH2II = gas.h2_ii * rho / (2 * mH_in_g)
+
+        nel = nHII + nHeII + 2 * nHeIII + nH2II
+        mu = ((nHI + nHII) + (nHeI + nHeII + nHeIII) * 4 + (nH2I + nH2II) * 2) / (
+            nHI + nHII + nHeI + nHeII + nHeIII + nH2I + nH2II + nel
+        )
+        return mu
+
+    elif cooling == b"Grackle1":  # grackle1
+        nHI = gas.hi * rho / (mH_in_g)
+        nHII = gas.hii * rho / (mH_in_g)
+        nHeI = gas.he_i * rho / (4 * mH_in_g)
+        nHeII = gas.he_ii * rho / (4 * mH_in_g)
+        nHeIII = gas.he_iii * rho / (4 * mH_in_g)
+        nel = nHII + nHeII + 2 * nHeIII
+        mu = ((nHI + nHII) + (nHeI + nHeII + nHeIII) * 4) / (
+            nHI + nHII + nHeI + nHeII + nHeIII + nel
+        )
+        return mu
+
+    else:  # Grackle0
+        # print("... found info for grackle mode=0")
+
+        from unyt.physical_constants import kboltz_cgs as k_B_cgs
+
+        gamma = data.metadata.gas_gamma[0]
+
+        # Get internal energy
+        u = data.gas.internal_energies
+        u = u.to_physical()
+        u = u.to(unyt.erg / unyt.g)
+
+        # Get hydrigen fraction
+        H_frac = float(
+            data.metadata.parameters["GrackleCooling:HydrogenFractionByMass"]
+        )
+
+        # Compute T/mu
+        T_over_mu = (gamma - 1.0) * u.value * mH_in_g.value / k_B_cgs.value
+        T_trans = 1.1e4
+        mu_trans = 4.0 / (8.0 - 5.0 * (1.0 - H_frac))
+
+        # Determine if we are ionized or not
+        mu = np.ones(np.size(u))
+        mask_ionized = T_over_mu > (T_trans + 1) / mu_trans
+        mask_neutral = T_over_mu < (T_trans + 1) / mu_trans
+
+        # Give the right mu
+        mu[mask_ionized] = 4.0 / (8.0 - 5.0 * (1.0 - H_frac))
+        mu[mask_neutral] = 4.0 / (1.0 + 3.0 * H_frac)
+
+        return mu
+
+
+def get_gas_temperatures(data: sw.SWIFTDataset) -> np.array:
+    """
+        Compute the temperature of the gas.
+    """
+    from unyt.physical_constants import kboltz_cgs as k_B
+
+    # from unyt.physical_constants import mh
+
+    # Convert to cgs
+    mh.convert_to_cgs()
+
+    # Get the cooling model
+    cooling = data.metadata.subgrid_scheme["Cooling Model"]
+
+    # Get internal energy and convert to physical units in cgs
+    u = data.gas.internal_energies
+    u = u.to_physical()
+    u = u.to(unyt.erg / unyt.g)
+
+    # Get gamm and compute mu
+    gamma = data.metadata.gas_gamma[0]
+    mu = get_gas_mu(data)
+
+    # FInally compute the the temperature
+    if cooling == b"Grackle3" or cooling == b"Grackle2" or cooling == b"Grackle1":
+        T = mu * (gamma - 1.0) * u * mh / k_B
+    else:
+        a = (gamma - 1.0) * (mu * mh) / k_B * u
+        T = np.where((u.value > 0), a.value, 0) * unyt.kelvin
+    return T
+
+
+def get_sink_and_stars_positions(filename):
+    data = sw.load(filename)
+
+    sink_pos = data.sinks.coordinates
+    star_pos = data.stars.coordinates
+
+    return sink_pos, star_pos
+
+
+def make_projection(filename, image_resolution):
+    """
+    Compute a mass projection with swiftsimio.
+    """
+    data = sw.load(filename)
+
+    # Compute projected density
+    projected_mass = project_gas(
+        data,
+        resolution=image_resolution,
+        project="masses",
+        parallel=True,
+        periodic=True,
+    )
+
+    boxsize = data.metadata.boxsize
+    x_edges = np.linspace(0 * unyt.kpc, boxsize[0], image_resolution)
+    y_edges = np.linspace(0 * unyt.kpc, boxsize[1], image_resolution)
+
+    # Convert to 1/cm**2
+    projected_mass = projected_mass.to(mh / (cm ** 2))
+
+    return projected_mass.T, x_edges, y_edges
+
+
+def setup_axes():
+    """
+    Creates the figure and axis object.
+    """
+    fig, ax = plt.subplots(1, figsize=(6, 5), dpi=300)
+
+    ax.set_xlabel("x [kpc]")
+    ax.set_ylabel("y [kpc]")
+
+    return fig, ax
+
+
+def make_single_image(filename, image_resolution):
+    """
+    Makes a single image and saves it to mass_projection_{snapshot_number}.png.
+
+    Filename should be given _without_ hdf5 extension.
+    """
+    file = "{:s}.hdf5".format(filename)
+    fig, ax = setup_axes()
+    projected_mass, x_edges, y_edges = make_projection(file, image_resolution)
+
+    mappable = ax.pcolormesh(
+        x_edges, y_edges, projected_mass, cmap=cmap, norm=LogNorm()
+    )
+    fig.colorbar(mappable, label="Surface density [cm$^{-2}]$", pad=0)
+
+    sink_pos, star_pos = get_sink_and_stars_positions(file)
+
+    if star_pos.size != 0:
+        ax.scatter(star_pos[:, 0], star_pos[:, 1], c="limegreen", zorder=1, marker="*")
+
+    if sink_pos.size != 0:
+        ax.scatter(sink_pos[:, 0], sink_pos[:, 1], c="blue", zorder=2, marker=".")
+
+    ax.text(
+        0.7,
+        0.95,
+        "$N_{\mathrm{star}}" + " = {}$".format(len(star_pos)),
+        transform=ax.transAxes,
+        fontsize=7,
+        bbox=dict(facecolor="white", alpha=0.8),
+    )
+    ax.text(
+        0.1,
+        0.95,
+        "$N_{\mathrm{sink}}" + " = {}$".format(len(sink_pos)),
+        transform=ax.transAxes,
+        fontsize=7,
+        bbox=dict(facecolor="white", alpha=0.8),
+    )
+
+    fig.tight_layout()
+
+    image = "mass_projection_{:s}.png".format(filename[-4:])
+    fig.savefig(image)
+
+    return
+
+
+def make_movie(args, image_resolution):
+    """
+    Makes a movie and saves it to mass_projection_movie.mp4.
+    """
+
+    fig, ax = setup_axes()
+
+    def grab_metadata(n):
+        filename = "{:s}_{:04d}.hdf5".format(args["stub"], n)
+        data = sw.load(filename)
+
+        return data.metadata
+
+    def grab_data(n):
+        filename = "{:s}_{:04d}.hdf5".format(args["stub"], n)
+
+        H, _, _ = make_projection(filename, image_resolution)
+
+        # Need to ravel because pcolormesh's set_array takes a 1D array. Might
+        # as well do it here, because 1d arrays are easier to max() than 2d.
+        return H.ravel()
+
+    def grab_sink_star_pos(n):
+        filename = "{:s}_{:04d}.hdf5".format(args["stub"], n)
+
+        sink_pos, star_pos = get_sink_and_stars_positions(filename)
+
+        return sink_pos, star_pos
+
+    histograms = [
+        grab_data(n)
+        for n in tqdm(
+            range(args["initial"], args["final"] + 1),
+            desc="Computing gas mass projection",
+        )
+    ]
+
+    sink_star = [
+        grab_sink_star_pos(n)
+        for n in tqdm(
+            range(args["initial"], args["final"] + 1),
+            desc="Getting sink and stars positions",
+        )
+    ]
+
+    sink_pos = [s[0] for s in sink_star]
+    star_pos = [s[1] for s in sink_star]
+
+    metadata = [
+        grab_metadata(n)
+        for n in tqdm(
+            range(args["initial"], args["final"] + 1), desc="Grabbing metadata"
+        )
+    ]
+
+    # Need to get a reasonable norm so that we don't overshoot.
+    max_particles = max([x.max() for x in histograms])
+    min_particles = max([x.min() for x in histograms])
+    norm = LogNorm(vmin=min_particles, vmax=max_particles)
+
+    # First, let's make the initial frame (we need this for our d, T values that we
+    # got rid of in grab_data.
+    hist, x_edges, y_edges = make_projection(
+        "{:s}_{:04d}.hdf5".format(args["stub"], args["initial"]), image_resolution
+    )
+
+    mappable = ax.pcolormesh(x_edges, y_edges, hist, cmap=cmap, norm=norm)
+    fig.colorbar(mappable, label="Surface density [cm$^{-2}]$", pad=0)
+
+    fig.tight_layout()
+
+    # Once we've rearranged the figure with tight_layout(), we can start laying
+    # Down the metadata text.
+
+    def format_metadata(metadata: sw.SWIFTMetadata):
+        t = metadata.t
+        t.convert_to_units(unyt.Myr)
+
+        x = "$a$: {:2.2f}\n$z$: {:2.2f}\n$t$: {:2.2f}".format(metadata.a, metadata.z, t)
+
+        return x
+
+    text = ax.text(
+        0.025,
+        0.975,
+        format_metadata(metadata[0]),
+        ha="left",
+        va="top",
+        transform=ax.transAxes,
+        color="white",
+    )
+
+    ax.text(
+        0.975,
+        0.975,
+        metadata[0].code["Git Revision"].decode("utf-8"),
+        ha="right",
+        va="top",
+        transform=ax.transAxes,
+        color="white",
+    )
+
+    def animate(data):
+        mappable.set_array(histograms[data])
+        text.set_text(format_metadata(metadata[data]))
+
+        if star_pos[data].size != 0:
+            ax.scatter(
+                star_pos[data][:, 0],
+                star_pos[data][:, 1],
+                c="limegreen",
+                zorder=1,
+                marker="*",
+            )
+
+        if sink_pos[data].size != 0:
+            ax.scatter(
+                sink_pos[data][:, 0],
+                sink_pos[data][:, 1],
+                c="blue",
+                zorder=2,
+                marker=".",
+            )
+
+        return mappable
+
+    animation = FuncAnimation(
+        fig, animate, range(len(histograms)), fargs=[], interval=1000 / 10
+    )
+
+    animation.save("mass_projection_movie.mp4")
+
+    return
+
+
+# %%
+if __name__ == "__main__":
+    import argparse as ap
+
+    parser = ap.ArgumentParser(
+        description="""
+             Plotting script for making a mass projection plot.
+             Takes the filename handle, start, and (optionally) stop
+             snapshots. If stop is not given, png plot is produced for
+             that snapshot. If given, a movie is made.
+             """
+    )
+
+    parser.add_argument(
+        "-i",
+        "--initial",
+        help="""Initial snapshot number. Default: 0.""",
+        default=0,
+        required=False,
+        type=int,
+    )
+
+    parser.add_argument(
+        "-f",
+        "--final",
+        help="""Final snapshot number. Default: 0.""",
+        default=0,
+        required=False,
+        type=int,
+    )
+
+    parser.add_argument(
+        "-s",
+        "--stub",
+        help="""Root of the snapshots filenames (e.g. snapshot). This is
+                the first part of the filename for the snapshots,
+                not including the final underscore. Required.""",
+        type=str,
+        required=True,
+    )
+
+    args = vars(parser.parse_args())
+
+    if args["final"] <= args["initial"]:
+        # Run in single image mode.
+        filename = "{:s}_{:04d}".format(args["stub"], args["initial"])
+
+        make_single_image(filename, image_resolution=image_resolution)
+
+    else:
+        # Movie mode!
+        make_movie(args, image_resolution=image_resolution)

examples/SinkParticles/HomogeneousBox/rhoTPlot.py

+"""
+Makes a rho-T plot. Uses the swiftsimio library.
+"""
+
+import matplotlib.pyplot as plt
+import numpy as np
+
+import swiftsimio as sw
+
+import unyt
+from unyt import mh, cm
+from tqdm import tqdm
+from matplotlib.colors import LogNorm
+from matplotlib.animation import FuncAnimation
+
+# %%
+# Constants; these could be put in the parameter file but are rarely changed.
+density_bounds = [1e-6, 1e6]  # in nh/cm^3
+temperature_bounds = [1e0, 1e8]  # in K
+bins = 128
+
+# Plotting controls
+cmap = "viridis"
+
+# %%
+
+
+def get_gas_mu(data: sw.SWIFTDataset) -> np.array:
+    """
+    Return the mean molecular weight of the gas.
+    """
+    from unyt.physical_constants import mh
+
+    # Get the cooling model
+    cooling = data.metadata.subgrid_scheme["Cooling Model"]
+
+    # Use a variable for better readability
+    gas = data.gas
+
+    # Get rho
+    rho = gas.densities.to(unyt.g / unyt.cm ** 3)  # self.Rho(units='g/cm3')
+
+    # hydrogen mass in gram
+    mh.convert_to_cgs()
+    mH_in_g = mh
+
+    if cooling == b"Grackle3":  # grackle3
+        nHI = gas.hi * rho / (mH_in_g)
+        nHII = gas.hii * rho / (mH_in_g)
+        nHeI = gas.he_i * rho / (4 * mH_in_g)
+        nHeII = gas.he_ii * rho / (4 * mH_in_g)
+        nHeIII = gas.he_iii * rho / (4 * mH_in_g)
+        nH2I = gas.h2_i * rho / (2 * mH_in_g)
+        nH2II = gas.h2_ii * rho / (2 * mH_in_g)
+        nHDI = gas.hdi * rho / (3 * mH_in_g)
+
+        nel = nHII + nHeII + 2 * nHeIII + nH2II
+        mu = (
+            (nHI + nHII) + (nHeI + nHeII + nHeIII) * 4 + (nH2I + nH2II) * 2 + nHDI * 3
+        ) / (nHI + nHII + nHeI + nHeII + nHeIII + nH2I + nH2II + nHDI + nel)
+        return mu
+
+    elif cooling == b"Grackle2":  # grackle2
+        nHI = gas.hi * rho / (mH_in_g)
+        nHII = gas.hii * rho / (mH_in_g)
+        nHeI = gas.he_i * rho / (4 * mH_in_g)
+        nHeII = gas.he_ii * rho / (4 * mH_in_g)
+        nHeIII = gas.he_iii * rho / (4 * mH_in_g)
+        nH2I = gas.h2_i * rho / (2 * mH_in_g)
+        nH2II = gas.h2_ii * rho / (2 * mH_in_g)
+
+        nel = nHII + nHeII + 2 * nHeIII + nH2II
+        mu = ((nHI + nHII) + (nHeI + nHeII + nHeIII) * 4 + (nH2I + nH2II) * 2) / (
+            nHI + nHII + nHeI + nHeII + nHeIII + nH2I + nH2II + nel
+        )
+        return mu
+
+    elif cooling == b"Grackle1":  # grackle1
+        nHI = gas.hi * rho / (mH_in_g)
+        nHII = gas.hii * rho / (mH_in_g)
+        nHeI = gas.he_i * rho / (4 * mH_in_g)
+        nHeII = gas.he_ii * rho / (4 * mH_in_g)
+        nHeIII = gas.he_iii * rho / (4 * mH_in_g)
+        nel = nHII + nHeII + 2 * nHeIII
+        mu = ((nHI + nHII) + (nHeI + nHeII + nHeIII) * 4) / (
+            nHI + nHII + nHeI + nHeII + nHeIII + nel
+        )
+        return mu
+
+    else:  # Grackle0
+        from unyt.physical_constants import kboltz_cgs as k_B_cgs
+
+        gamma = data.metadata.gas_gamma[0]
+
+        # Get internal energy
+        u = data.gas.internal_energies
+        u = u.to_physical()
+        u = u.to(unyt.erg / unyt.g)
+
+        # Get hydrigen fraction
+        H_frac = float(
+            data.metadata.parameters["GrackleCooling:HydrogenFractionByMass"]
+        )
+
+        # Compute T/mu
+        T_over_mu = (gamma - 1.0) * u.value * mH_in_g.value / k_B_cgs.value
+        T_trans = 1.1e4
+        mu_trans = 4.0 / (8.0 - 5.0 * (1.0 - H_frac))
+
+        # Determine if we are ionized or not
+        mu = np.ones(np.size(u))
+        mask_ionized = T_over_mu > (T_trans + 1) / mu_trans
+        mask_neutral = T_over_mu < (T_trans + 1) / mu_trans
+
+        # Give the right mu
+        mu[mask_ionized] = 4.0 / (8.0 - 5.0 * (1.0 - H_frac))
+        mu[mask_neutral] = 4.0 / (1.0 + 3.0 * H_frac)
+
+        return mu
+
+
+def get_gas_temperatures(data: sw.SWIFTDataset) -> np.array:
+    """
+        Compute the temperature of the gas.
+    """
+    from unyt.physical_constants import kboltz_cgs as k_B
+    from unyt.physical_constants import mh
+
+    # Convert to cgs
+    mh.convert_to_cgs()
+
+    # Get the cooling model
+    cooling = data.metadata.subgrid_scheme["Cooling Model"]
+
+    # Get internal energy and convert to physical units in cgs
+    u = data.gas.internal_energies
+    u = u.to_physical()
+    u = u.to(unyt.erg / unyt.g)
+
+    # Get gamm and compute mu
+    gamma = data.metadata.gas_gamma[0]
+    mu = get_gas_mu(data)
+
+    # FInally compute the the temperature
+    if cooling == b"Grackle3" or cooling == b"Grackle2" or cooling == b"Grackle1":
+        T = mu * (gamma - 1.0) * u * mh / k_B
+    else:
+        a = (gamma - 1.0) * (mu * mh) / k_B * u
+        T = np.where((u.value > 0), a.value, 0) * unyt.kelvin
+    return T
+
+
+def get_data(filename):
+    """
+    Grabs the data (T in Kelvin and density in mh / cm^3).
+
+    Note: Converts data.gas.densities to mh/cm^3.
+    """
+    data = sw.SWIFTDataset(filename)
+
+    data.gas.densities = data.gas.densities.to(mh / (cm ** 3))
+    data.gas.temperatures = get_gas_temperatures(data)
+    data.gas.temperatures.convert_to_cgs()
+
+    return data.gas.densities, data.gas.temperatures
+
+
+def make_hist(filename, density_bounds, temperature_bounds, bins):
+    """
+    Makes the histogram for filename with bounds as lower, higher
+    for the bins and "bins" the number of bins along each dimension.
+
+    Also returns the edges for pcolormesh to use.
+    """
+
+    density_bins = np.logspace(
+        np.log10(density_bounds[0]), np.log10(density_bounds[1]), bins
+    )
+    temperature_bins = np.logspace(
+        np.log10(temperature_bounds[0]), np.log10(temperature_bounds[1]), bins
+    )
+
+    # print(density_bins, temperature_bins)
+
+    H, density_edges, temperature_edges = np.histogram2d(
+        *get_data(filename), bins=[density_bins, temperature_bins]
+    )
+
+    return H.T, density_edges, temperature_edges
+
+
+def setup_axes():
+    """
+    Creates the figure and axis object.
+    """
+    fig, ax = plt.subplots(1, figsize=(6, 5), dpi=300)
+
+    ax.set_xlabel("Density [$n_H$ cm$^{-3}$]")
+    ax.set_ylabel("Temperature [K]")
+
+    ax.loglog()
+
+    return fig, ax
+
+
+def make_single_image(filename, density_bounds, temperature_bounds, bins):
+    """
+    Makes a single image and saves it to rhoTPlot_{filename}.png.
+
+    Filename should be given _without_ hdf5 extension.
+    """
+
+    fig, ax = setup_axes()
+    hist, rho, T = make_hist(
+        "{:s}.hdf5".format(filename), density_bounds, temperature_bounds, bins
+    )
+
+    mappable = ax.pcolormesh(rho, T, hist, cmap=cmap, norm=LogNorm())
+    fig.colorbar(mappable, label="Number of particles", pad=0)
+
+    fig.tight_layout()
+
+    fig.savefig("rhoTPlot_{:s}.png".format(filename[-4:]))
+
+    return
+
+
+def make_movie(args, density_bounds, temperature_bounds, bins):
+    """
+    Makes a movie and saves it to rhoTPlot_{stub}.mp4.
+    """
+
+    fig, ax = setup_axes()
+
+    def grab_metadata(n):
+        filename = "{:s}_{:04d}.hdf5".format(args["stub"], n)
+        data = sw.load(filename)
+
+        return data.metadata
+
+    def grab_data(n):
+        filename = "{:s}_{:04d}.hdf5".format(args["stub"], n)
+
+        H, _, _ = make_hist(filename, density_bounds, temperature_bounds, bins)
+
+        # Need to ravel because pcolormesh's set_array takes a 1D array. Might
+        # as well do it here, beacuse 1d arrays are easier to max() than 2d.
+        return H.ravel()
+
+    histograms = [
+        grab_data(n)
+        for n in tqdm(
+            range(args["initial"], args["final"] + 1), desc="Histogramming data"
+        )
+    ]
+
+    metadata = [
+        grab_metadata(n)
+        for n in tqdm(
+            range(args["initial"], args["final"] + 1), desc="Grabbing metadata"
+        )
+    ]
+
+    # Need to get a reasonable norm so that we don't overshoot.
+    max_particles = max([x.max() for x in histograms])
+
+    norm = LogNorm(vmin=1, vmax=max_particles)
+
+    # First, let's make the initial frame (we need this for our d, T values that we
+    # got rid of in grab_data.
+    hist, d, T = make_hist(
+        "{:s}_{:04d}.hdf5".format(args["stub"], args["initial"]),
+        density_bounds,
+        temperature_bounds,
+        bins,
+    )
+
+    mappable = ax.pcolormesh(d, T, hist, cmap=cmap, norm=norm)
+    fig.colorbar(mappable, label="Number of particles", pad=0)
+
+    fig.tight_layout()
+
+    # Once we've rearranged the figure with tight_layout(), we can start laying
+    # Down the metadata text.
+
+    def format_metadata(metadata: sw.SWIFTMetadata):
+        t = metadata.t
+        t.convert_to_units(unyt.Myr)
+
+        x = "$a$: {:2.2f}\n$z$: {:2.2f}\n$t$: {:2.2f}".format(metadata.a, metadata.z, t)
+
+        return x
+
+    text = ax.text(
+        0.025,
+        0.975,
+        format_metadata(metadata[0]),
+        ha="left",
+        va="top",
+        transform=ax.transAxes,
+    )
+
+    ax.text(
+        0.975,
+        0.975,
+        metadata[0].code["Git Revision"].decode("utf-8"),
+        ha="right",
+        va="top",
+        transform=ax.transAxes,
+    )
+
+    def animate(data):
+        mappable.set_array(histograms[data])
+        text.set_text(format_metadata(metadata[data]))
+
+        return mappable
+
+    animation = FuncAnimation(
+        fig, animate, range(len(histograms)), fargs=[], interval=1000 / 25
+    )
+
+    animation.save("rhoTPlot.mp4")
+
+    return
+
+
+# %%
+if __name__ == "__main__":
+    import argparse as ap
+
+    parser = ap.ArgumentParser(
+        description="""
+             Plotting script for making a rho-T plot.
+             Takes the filename handle, start, and (optionally) stop
+             snapshots. If stop is not given, png plot is produced for
+             that snapshot. If given, a movie is made.
+             """
+    )
+
+    parser.add_argument(
+        "-i",
+        "--initial",
+        help="""Initial snapshot number. Default: 0.""",
+        default=0,
+        required=False,
+        type=int,
+    )
+
+    parser.add_argument(
+        "-f",
+        "--final",
+        help="""Final snapshot number. Default: 0.""",
+        default=0,
+        required=False,
+        type=int,
+    )
+
+    parser.add_argument(
+        "-s",
+        "--stub",
+        help="""Root of the snapshots filenames (e.g. snapshot). This is
+                the first part of the filename for the snapshots,
+                not including the final underscore. Required.""",
+        type=str,
+        required=True,
+    )
+
+    args = vars(parser.parse_args())
+
+    if args["final"] <= args["initial"]:
+        # Run in single image mode.
+        filename = "{:s}_{:04d}".format(args["stub"], args["initial"])
+
+        make_single_image(
+            filename,
+            density_bounds=density_bounds,
+            temperature_bounds=temperature_bounds,
+            bins=bins,
+        )
+
+    else:
+        # Movie mode!
+        make_movie(
+            args,
+            density_bounds=density_bounds,
+            temperature_bounds=temperature_bounds,
+            bins=bins,
+        )

examples/SinkParticles/HomogeneousBox/run.sh

+#!/bin/bash
+
+# make run.sh fail if a subcommand fails
+set -e
+
+n_threads=8 #Number of threads to use
+level=5 #Number of particles = 2^(3*level)
+jeans_lenght=0.250 #Jeans wavelenght in unit of the boxsize
+
+#Create the ICs if they do not exist
+if [ ! -e ICs_homogeneous_box.hdf5 ]
+then
+    echo "Generating initial conditions to run the example..."
+    python3 makeIC.py --level $level -o ICs_homogeneous_box.hdf5 --lJ $jeans_lenght
+fi
+
+
+# Get the Grackle cooling table
+if [ ! -e CloudyData_UVB=HM2012.h5 ]
+then
+    echo "Fetching the Cloudy tables required by Grackle..."
+    ./getGrackleCoolingTable.sh
+fi
+
+
+if [ ! -e POPIIsw.h5 ]
+then
+    echo "Fetching the chemistry tables..."
+    ./getChemistryTable.sh
+fi
+
+
+# Create output directory
+DIR=snap #First test of units conversion
+if [ -d "$DIR" ];
+then
+    echo "$DIR directory exists. Its content will be removed."
+    rm -r $DIR
+else
+    echo "$DIR directory does not exists. It will be created."
+    mkdir $DIR
+fi
+
+printf "Running simulation..."
+
+../../../swift --hydro --sinks --stars --self-gravity --feedback --cooling --sync --limiter --threads=$n_threads params.yml 2>&1 | tee output.log
+
+#Do some data analysis to show what's in this box
+python3 plot_gas_density.py -i 282 -s 'snap/snapshot'
+python3 rhoTPlot.py -i 282 -s 'snap/snapshot'
+python3 rhoTPlot.py -i 0 -f 282 -s 'snap/snapshot'
+python3 plot_gas_density.py -i 0 -f 282 -s 'snap/snapshot'