Self-shielded Grackle table

extern/cooling/resample_grackle_cooling_tables.py

@@ -4,12 +4,12 @@
 Outputs include cooling rates, temperatures, sound speeds, pressures, and entropies.
 
 Example usage:
-Resample using the default CloudyData_UVB=HM2012.h5 data (auto-downloaded via wget) at solar metallicity
-    ./resample_grackle_cooling_tables.py
-The CloudyData_UVB=HM2012_resampled.h5 file in this folder was generated via this command:
+Resample using the default CloudyData_UVB=HM2012.h5 data (auto-downloaded via wget) at solar metallicity. This command was used to generate the CloudyData_UVB=HM2012_resampled.h5 file in this folder.
     ./resample_grackle_cooling_tables.py --output CloudyData_UVB=HM2012_resampled.h5
-Resample using the default CloudyData_UVB=HM2012.h5 data (auto-downloaded via wget) at 0.5 solar metallicity
-    ./resample_grackle_cooling_tables.py --zmet 0.5
+Resample using the default CloudyData_UVB=HM2012.h5 data at 0.5 solar metallicity
+    ./resample_grackle_cooling_tables.py --zmet 0.5 --output CloudyData_UVB=HM2012_resampled_zmet0.5.h5
+Resample using CloudyData_UVB=HM2012_shielded.h5 data at solar metallicity, excluding photoelectric heating. This command was used to generate the CloudyData_UVB=HM2012_shielded_resampled_noPE.h5 file in this folder.
+    ./resample_grackle_cooling_tables.py --shield --exclude_pe --output CloudyData_UVB=HM2012_shielded_resampled_noPE.h5
 Resample with user-provided tables and non-default spacing parameters
     ./resample_grackle_cooling_tables.py /path/to/CloudyData.h5 --n_rho 50 --n_eint 400 --output my_tables.h5
 Validate fast_log2 <-> inverse_fast_log2
@@ -23,17 +23,24 @@
 
 import numpy as np
 import h5py
+import matplotlib.pyplot as plt
+import matplotlib.colors as colors
 
 from grackle_tables import (
     read_tables, cooling_rate, interpolate_mu, compute_temperature_from_nH_e,
     m_H, boltzmann_constant_cgs_, cloudy_H_mass_fraction, specific_energy_from_temperature
 )
 
+KPC = 3.086e21 # cm
 
 DEFAULT_GRACKLE_DATA_URL = (
     "https://github.com/grackle-project/grackle_data_files/raw/"
     "928696482fbe15d9bac4382de6134d95568f099c/input/CloudyData_UVB=HM2012.h5"
 )
+DEFAULT_GRACKLE_DATA_URL_SHIELDED = (
+    "https://github.com/grackle-project/grackle_data_files/raw/"
+    "928696482fbe15d9bac4382de6134d95568f099c/input/CloudyData_UVB=HM2012_shielded.h5"
+)
 
 
 def fast_log2(x):
@@ -178,6 +185,20 @@ def compute_entropy(rho, T, mu):
     return K
 
 
+def compute_cooling_time(rho, e_int, Edot):
+    """Compute cooling time from density, specific internal energy, and net *heating* rate.
+
+    Args:
+        rho: density (g/cm^3)
+        e_int: specific internal energy (erg/g)
+        Edot: net *heating* rate (erg/cm^3/s)
+
+    Returns:
+        t_cool: cooling time (s)
+    """
+    return -(rho * e_int) / Edot
+
+
 def find_eint_range(tables):
     """Find the range of specific internal energies for a given table.
 
@@ -246,6 +267,12 @@ def resample_cooling_tables(grackle_file, n_rho=100, n_eint=100, zmet=1.0,
     rho_grid = inverse_fast_log2(fast_log_rho_scaled)
     eint_grid = inverse_fast_log2(fast_log_eint_scaled)
 
+    # Broadcast nH to 2D
+    # x: n_H, y: T/mu = 2/3 * eint * m_H / k_B
+    x = rho_grid * cloudy_H_mass_fraction / m_H
+    y = 2./3. * eint_grid * m_H / boltzmann_constant_cgs_
+    X, Y = np.meshgrid(x, y)
+
     # Check that the grid boundaries are correct
     print("Verifying grid boundaries...")
     print(f"  rho_grid[0] = {rho_grid[0]:.15e}, rho_min = {rho_min:.15e}")
@@ -265,6 +292,7 @@ def resample_cooling_tables(grackle_file, n_rho=100, n_eint=100, zmet=1.0,
 
     # Initialize output arrays
     cooling_rates = np.zeros((n_rho, n_eint))
+    cooling_lengths = np.zeros((n_rho, n_eint))
     temperatures = np.zeros((n_rho, n_eint))
     sound_speeds = np.zeros((n_rho, n_eint))
     pressures = np.zeros((n_rho, n_eint))
@@ -285,6 +313,7 @@ def resample_cooling_tables(grackle_file, n_rho=100, n_eint=100, zmet=1.0,
         nH = rho * cloudy_H_mass_fraction / m_H
 
         for j, e_int in enumerate(eint_grid):
+            # e_int: erg/g
             try:
                 T = compute_temperature_from_nH_e(nH, e_int, tables=tables)
                 Edot = cooling_rate(nH, T, zmet, redshift=0., tables=tables, include_pe=include_pe)
@@ -293,17 +322,25 @@ def resample_cooling_tables(grackle_file, n_rho=100, n_eint=100, zmet=1.0,
                 P = compute_pressure(rho, T, mu)
                 K = compute_entropy(rho, T, mu)
                 temperatures[i, j] = T
-                cooling_rates[i, j] = Edot / rho**2
-                sound_speeds[i, j] = cs
+                cooling_rates[i, j] = Edot / rho**2 # 'erg/cm^3/s/(g/cm^3)^2'
+                sound_speeds[i, j] = cs # 'cm/s'
                 pressures[i, j] = P
                 entropies[i, j] = K
+                ct = compute_cooling_time(rho, e_int, Edot) # s
+                cooling_lengths[i, j] = cs * ct # cm
+
             except ValueError:
                 # Handle extrapolation errors by setting to NaN
                 temperatures[i, j] = np.nan
                 cooling_rates[i, j] = np.nan
                 sound_speeds[i, j] = np.nan
                 pressures[i, j] = np.nan
                 entropies[i, j] = np.nan
+                cooling_lengths[i, j] = np.nan
+
+    print(f"cooling_lengths.shape = {cooling_lengths.shape}")
+    # midrow values in kpc
+    print(f"cooling_lengths[n_rho/2, :] = {cooling_lengths[n_rho//2, :] / KPC} kpc")
 
     # Save resampled tables to HDF5 file
     print(f"\nSaving resampled tables to {output_file}")
@@ -332,6 +369,7 @@ def resample_cooling_tables(grackle_file, n_rho=100, n_eint=100, zmet=1.0,
         data_group.create_dataset('sound_speeds', data=sound_speeds)
         data_group.create_dataset('pressures', data=pressures)
         data_group.create_dataset('entropies', data=entropies)
+        data_group.create_dataset('cooling_lengths', data=cooling_lengths)
 
         # Store metadata as attributes
         metadata_group.attrs['n_rho'] = n_rho
@@ -343,6 +381,7 @@ def resample_cooling_tables(grackle_file, n_rho=100, n_eint=100, zmet=1.0,
         metadata_group.attrs['cloudy_H_mass_fraction'] = cloudy_H_mass_fraction
         metadata_group.attrs['description'] = 'Cooling rates resampled on (rho, e_int) grid using not-quite-logarithmic spacing'
         metadata_group.attrs['spacing_method'] = 'not-quite-logarithmic (fast log2 approximation)'
+        metadata_group.attrs['include_pe'] = 1 if include_pe else 0
 
         # Store units as attributes
         units_group.attrs['rho'] = 'g/cm^3'
@@ -352,6 +391,7 @@ def resample_cooling_tables(grackle_file, n_rho=100, n_eint=100, zmet=1.0,
         units_group.attrs['sound_speed'] = 'cm/s'
         units_group.attrs['pressure'] = 'dyne/cm^2'
         units_group.attrs['entropy'] = 'erg*cm^2'
+        units_group.attrs['cooling_length'] = 'cm'
 
     print("Done!")
 
@@ -365,6 +405,40 @@ def resample_cooling_tables(grackle_file, n_rho=100, n_eint=100, zmet=1.0,
         print(f"Sound speed range: {np.min(sound_speeds[valid_mask]):.2e} to {np.max(sound_speeds[valid_mask]):.2e} cm/s")
         print(f"Pressure range: {np.min(pressures[valid_mask]):.2e} to {np.max(pressures[valid_mask]):.2e} dyne/cm^2")
         print(f"Entropy range: {np.min(entropies[valid_mask]):.2e} to {np.max(entropies[valid_mask]):.2e} erg*cm^2")
+        print(f"Cooling length range: {np.min(cooling_lengths[valid_mask]):.2e} to {np.max(cooling_lengths[valid_mask]):.2e} cm")
+
+    # Generate phase plot
+    print("\nGenerating phase plot...")
+
+    # Create plot
+    plt.figure(figsize=(5, 4))
+
+    # Convert to kpc for plotting
+    L_cool_kpc = cooling_lengths / KPC
+    Z = L_cool_kpc.T # transpose to get (n_H, T/mu)
+
+    # mask out negative Z
+    # Z_masked = np.ma.masked_where(Z_abs > 1e6, Z_abs)
+    Z_masked = np.ma.masked_where(Z < 0, Z)
+
+    assert X.shape == Y.shape == Z.shape, f"X.shape={X.shape} != Y.shape={Y.shape} != Z.shape={Z.shape}"
+
+    # Use pcolormesh with linear scale
+    pcm = plt.pcolormesh(X, Y, Z_masked, shading='auto', cmap='viridis', norm=colors.LogNorm())
+
+    plt.xscale('log')
+    plt.yscale('log')
+    plt.xlim(1e-6, x.max())
+    plt.xlabel(r'Number Density $n_H$ (cm$^{-3}$)')
+    plt.ylabel(r'$T/\mu$ (K)')
+    plt.title('Cooling Length $L_{cool} = c_s t_{cool}$ (kpc)')
+    cbar = plt.colorbar(pcm, label='Cooling Length (kpc)')
+    cbar.ax.set_facecolor("0.8")
+
+    plot_file = os.path.splitext(output_file)[0] + '_cooling_length_phase_plot.png'
+    plt.savefig(plot_file, dpi=300)
+    plt.close()
+    print(f"Saved phase plot to {plot_file}")
 
 
 def test_inverse_fast_log2():
@@ -415,11 +489,19 @@ def test_inverse_fast_log2():
     return max_error < eps * tolerance_factor
 
 
-def download_default_grackle_tables(url: str = DEFAULT_GRACKLE_DATA_URL) -> str:
+def download_default_grackle_tables(url: str = DEFAULT_GRACKLE_DATA_URL, local: str = None) -> str:
     """Download the default Grackle tables via wget and return the temporary file path."""
+    if local is not None:
+        if os.path.exists(local):
+            return local
+
     if shutil.which('wget') is None:
         raise RuntimeError("wget is required to download default Grackle tables automatically.")
 
+    if local is not None:
+        subprocess.run(['wget', '-O', local, url], check=True)
+        return local
+
     temp_file = tempfile.NamedTemporaryFile(delete=False, suffix='.h5')
     temp_path = temp_file.name
     temp_file.close()
@@ -461,6 +543,12 @@ def main():
     parser.add_argument('--zmet', type=float, default=1.0,
                         help='Gas Metallicity scaled to Solar (default: 1.0)')
 
+    parser.add_argument('--shield', action='store_true',
+                        help='Account for self-shielding (default: False)')
+
+    parser.add_argument('--local', action='store_true',
+                        help='Store the Grackle table in the local directory (default: False)')
+
     args = parser.parse_args()
 
     if args.test:
@@ -471,8 +559,18 @@ def main():
     cleanup_path = None
     grackle_file = args.grackle_file
     if not grackle_file:
+        if args.shield:
+            url = DEFAULT_GRACKLE_DATA_URL_SHIELDED
+        else:
+            url = DEFAULT_GRACKLE_DATA_URL
+
+        if args.local:
+            grackle_file = 'grackle.h5'
+        else:
+            grackle_file = None
+
         try:
-            cleanup_path = download_default_grackle_tables()
+            cleanup_path = download_default_grackle_tables(url, local=grackle_file)
             grackle_file = cleanup_path
         except RuntimeError as err:
             parser.error(str(err))
@@ -487,7 +585,7 @@ def main():
             include_pe=not args.exclude_pe
         )
     finally:
-        if cleanup_path:
+        if cleanup_path and not args.local:
             try:
                 os.unlink(cleanup_path)
             except FileNotFoundError:

src/QuokkaSimulation.hpp

@@ -528,6 +528,7 @@ template <typename problem_t> void QuokkaSimulation<problem_t>::readParmParse()
 	}
 
 	// set cooling runtime parameters
+	bool cooling_table_include_pe = false;
 	{
 		amrex::ParmParse const hpp("cooling");
 		int alwaysReadTables = 0;
@@ -542,7 +543,7 @@ template <typename problem_t> void QuokkaSimulation<problem_t>::readParmParse()
 			if (coolingTableType_ == "resampled") {
 				// read resampled cooling tables
 				amrex::Print() << "Reading resampled cooling tables...\n";
-				quokka::ResampledCooling::readResampledData(coolingTableFilename_, resampledTables_);
+				cooling_table_include_pe = quokka::ResampledCooling::readResampledData(coolingTableFilename_, resampledTables_);
 			} else {
 				amrex::Abort("Invalid cooling table type! Only 'resampled' is supported.");
 			}
@@ -558,6 +559,9 @@ template <typename problem_t> void QuokkaSimulation<problem_t>::readParmParse()
 		pp.query("const_sfr_Msun_per_year_per_kpc2", const_sfr_Msun_per_year_per_kpc2_);
 		// It's allowed to turn on sfh and not turn on use_sfh_based_pe_heating, but the opposite is not allowed.
 		if (use_sfh_based_pe_heating_) {
+			AMREX_ALWAYS_ASSERT_WITH_MESSAGE(
+			    !cooling_table_include_pe,
+			    "When use_sfh_based_pe_heating is set to true, please use a Grackle cooling table that does NOT include photoelectric heating.");
 			AMREX_ALWAYS_ASSERT_WITH_MESSAGE(
 			    !sfh_to_pe_heating_table_filename_.empty(),
 			    "When use_sfh_based_pe_heating is set to true, a PE heating table must be specified via sfh_to_pe_heating_table");

src/cooling/ResampledCooling.cpp

@@ -19,20 +19,28 @@ namespace quokka::ResampledCooling
 
 constexpr double cloudy_H_mass_fraction = 1. / (1. + 0.1 * 3.971);
 
-void readResampledData(std::string const &hdf5_file, resampled_tables &resampledTables)
+// return: is_include_pe
+auto readResampledData(std::string const &hdf5_file, resampled_tables &resampledTables) -> bool
 {
 	amrex::Print() << "Initializing resampled cooling.\n";
 	amrex::Print() << fmt::format("resampled_table_file: {}.\n", hdf5_file);
 
+	// Check if file exists
+	if (!amrex::FileSystem::Exists(hdf5_file)) {
+		amrex::Abort("Resampled cooling table file does not exist!");
+	}
+
 	// Define coordinate names and fast_log setting
 	const std::vector<std::string> coord_names = {"rho", "eint"};
 	const int is_fast_log = 1;
 
 	// Coordinate bounds will be read by H5Reader
 	std::array<std::pair<amrex::Real, amrex::Real>, 2> coord_bounds;
+	bool is_pe_enabled = false;
 
 	// Read all 2D datasets using generic DataTable H5Reader (file path-based interface)
-	resampledTables.cooling_rates = quokka::DataTable<2, 1>::H5Reader(hdf5_file, "/data/cooling_rates", coord_names, is_fast_log, &coord_bounds);
+	resampledTables.cooling_rates =
+	    quokka::DataTable<2, 1>::H5Reader(hdf5_file, "/data/cooling_rates", coord_names, is_fast_log, &coord_bounds, &is_pe_enabled);
 	resampledTables.temperatures = quokka::DataTable<2, 1>::H5Reader(hdf5_file, "/data/temperatures", coord_names, is_fast_log);
 	resampledTables.sound_speeds = quokka::DataTable<2, 1>::H5Reader(hdf5_file, "/data/sound_speeds", coord_names, is_fast_log);
 	resampledTables.pressures = quokka::DataTable<2, 1>::H5Reader(hdf5_file, "/data/pressures", coord_names, is_fast_log);
@@ -52,6 +60,9 @@ void readResampledData(std::string const &hdf5_file, resampled_tables &resampled
 
 	amrex::Print() << fmt::format("\tDensity range: {} to {} g/cm^3 ({} steps).\n", resampledTables.rho_min, resampledTables.rho_max, n_rho);
 	amrex::Print() << fmt::format("\tSpecific energy range: {} to {} erg/g ({} steps).\n", resampledTables.eint_min, resampledTables.eint_max, n_eint);
+	amrex::Print() << fmt::format("\tPhotoelectric heating: {}.\n", is_pe_enabled ? "enabled" : "disabled");
+
+	return is_pe_enabled;
 }
 
 auto resampled_tables::const_tables() const -> resampledGpuConstTables

src/cooling/ResampledCooling.hpp

@@ -236,7 +236,7 @@ auto computeCooling(amrex::MultiFab &mf, const Real dt_in, resampled_tables &res
 	return true; // success
 }
 
-void readResampledData(std::string const &hdf5_file, resampled_tables &resampledTables);
+auto readResampledData(std::string const &hdf5_file, resampled_tables &resampledTables) -> bool;
 
 } // namespace quokka::ResampledCooling
 

src/util/DataTable.hpp

@@ -1287,8 +1287,9 @@ template <int Ndim, int Nout = 1, OutOfBounds oob_policy = OutOfBounds::clamp> c
 	// H5Reader: Generic static method to read n-dimensional data from HDF5 file and create DataTable
 	// Reads metadata, coordinates, and data all from the HDF5 file
 	// Optionally returns coordinate bounds via coord_bounds parameter
+	// Optionally returns whether photoelectric heating is enabled via include_pe parameter
 	static auto H5Reader(const std::string &file_path, const std::string &dataset_path, const std::vector<std::string> &coord_names, int is_fast_log = 0,
-			     std::array<std::pair<amrex::Real, amrex::Real>, Ndim> *coord_bounds = nullptr) -> DataTable
+			     std::array<std::pair<amrex::Real, amrex::Real>, Ndim> *coord_bounds = nullptr, bool *include_pe = nullptr) -> DataTable
 	{
 		static_assert(Ndim >= 1 && Ndim <= 4, "H5Reader supports 1D-4D tables");
 		AMREX_ALWAYS_ASSERT_WITH_MESSAGE(
@@ -1339,6 +1340,35 @@ template <int Ndim, int Nout = 1, OutOfBounds oob_policy = OutOfBounds::clamp> c
 			}
 		}
 
+		// Read include_pe if requested
+		if (include_pe != nullptr) {
+			if (H5Aexists(metadata_group, "include_pe") > 0) {
+				attr_id = H5Aopen(metadata_group, "include_pe", H5P_DEFAULT);
+				const hid_t attr_type = H5Aget_type(attr_id);
+				if (H5Tget_class(attr_type) == H5T_INTEGER) {
+					int cooling_include_pe = 0;
+					status = H5Aread(attr_id, H5T_NATIVE_INT, &cooling_include_pe);
+					AMREX_ALWAYS_ASSERT_WITH_MESSAGE(status != h5_error, "Failed to read include_pe (integer)!");
+					*include_pe = (cooling_include_pe != 0);
+				} else if (H5Tget_class(attr_type) == H5T_STRING) {
+					// Handle string written by older Python script versions ('0' or '1')
+					std::array<char, 4> buf{};
+					const hid_t mem_type = H5Tcopy(H5T_C_S1);
+					H5Tset_size(mem_type, 4);
+					status = H5Aread(attr_id, mem_type, buf.data());
+					AMREX_ALWAYS_ASSERT_WITH_MESSAGE(status != h5_error, "Failed to read include_pe (string)!");
+					*include_pe = (buf[0] == '1');
+					H5Tclose(mem_type);
+				} else {
+					AMREX_ALWAYS_ASSERT_WITH_MESSAGE(false, "include_pe attribute is neither integer nor string!");
+				}
+				H5Tclose(attr_type);
+				H5Aclose(attr_id);
+			} else {
+				*include_pe = false;
+			}
+		}
+
 		H5Gclose(metadata_group);
 
 		// Read coordinate grids