diff --git a/.github/workflows/xray_build_and_test.yml b/.github/workflows/xray_build_and_test.yml new file mode 100644 index 0000000..8e33185 --- /dev/null +++ b/.github/workflows/xray_build_and_test.yml @@ -0,0 +1,61 @@ +# +# Build the module and run tests. +# +# Repeats tests on local and remote files by running an instance of the +# hdfstream server in a container. +# +name: Xray branch only build and test + +on: + push: + branches: [ "lc_xrays" ] + pull_request: + branches: [ "lc_xrays" ] + +permissions: + contents: read + +jobs: + build: + + runs-on: ubuntu-latest + + steps: + - uses: actions/checkout@v4 + + - name: Set up Python 3.10 + uses: actions/setup-python@v3 + with: + python-version: "3.10" + + - name: Start hdfstream service with ./tests/data mounted + uses: ./.github/actions/start-hdfstream + with: + data-dir: ./tests/data + virtual-prefix: tests/data + image: ghcr.io/jchelly/hdfstream-api:0.0.6 + + - name: Install dependencies + run: | + python -m pip install --upgrade pip + pip install flake8 pytest + if [ -f requirements.txt ]; then pip install -r requirements.txt; fi + + - name: Lint with flake8 + run: | + # stop the build if there are Python syntax errors or undefined names + flake8 . --count --select=E9,F63,F7,F82 --show-source --statistics + # exit-zero treats all errors as warnings. The GitHub editor is 127 chars wide + # flake8 . --count --exit-zero --max-complexity=10 --max-line-length=127 --statistics + + - name: Install module + run: | + pip install -e . + + - name: Test with pytest + run: | + pytest --hdfstream-server=http://localhost:8080/hdfstream + + - name: Stop hdfstream service + uses: ./.github/actions/stop-hdfstream + if: always() diff --git a/examples/lightcone_shell_volume.py b/examples/lightcone_shell_volume.py new file mode 100644 index 0000000..83fabc7 --- /dev/null +++ b/examples/lightcone_shell_volume.py @@ -0,0 +1,107 @@ +#!/bin/env python +import sys +import glob +import numpy as np +import lightcone_io.particle_reader as pr +import lightcone_io.lc_xray_calculator as Xcalc +from lightcone_io.xray_utils import Snapshot_Cosmology_For_Lightcone, Xray_Filter +import unyt +import h5py + +import matplotlib.pyplot as plt +import matplotlib.gridspec as gridspec + +########## plot settings ########## +import matplotlib as mpl +# plot settings can be removed if needed +# set plotting params +# Ticks inside plots; more space devoted to plot. +plt.rcParams["xtick.direction"] ='in' +plt.rcParams["ytick.direction"] ='in' +plt.rcParams["xtick.top"]=True +plt.rcParams["ytick.right"]= True + +# axes font settings +mpl.rcParams['axes.labelsize']=10 +mpl.rcParams['xtick.labelsize']=9 +mpl.rcParams['ytick.labelsize']=9 + +plt.rcParams["legend.fontsize"]=9 +################################### + + +# for a given sim, find the comoving volume of a lightcone shell using +# the redshift values of the outer and inner edges of each shell, +# then make a simple diagram. + + +# simulation +sim_name="HYDRO_FIDUCIAL" +sim="L1000N1800/"+sim_name +# sim base dir +base_dir="/cosma8/data/dp004/flamingo/Runs/{sim}".format(sim="L1000N1800/"+sim_name) + + +# create astropy cosmology object for simulation using any of the snapshot files +snapshot_dir = "{base_name}/snapshots".format(base_name=base_dir) +cosmo=Snapshot_Cosmology_For_Lightcone(snapshot_dir) + +# get comoving radii for each shell +redshift_filename = "/cosma8/data/dp004/flamingo/Runs/{sim}/shell_redshifts_z3.txt".format(sim="L1000N1800/"+sim_name) +r_shells = cosmo.shell_comoving_raddii(redshift_filename) +z_shells = cosmo.shell_redshifts + +# show a diagram for the first 10 shells +from matplotlib.collections import PatchCollection + +fig = plt.figure(figsize=(3.21, 2.5)) +plot_grid = fig.add_gridspec(nrows=1, ncols=1, width_ratios=[1], height_ratios=[1]) +axs=plot_grid.subplots(sharex='row', sharey='row') + +n_shells=10 +colours = mpl.colormaps['terrain_r'](np.linspace(0, 1, n_shells+1)) +shell_volumes = np.zeros(n_shells) +for shell_nr in range(n_shells): + print( + "Shell: {shell_nr:d}\n\tredshift: {zmin:.3f} -> {zmax:.3f}\n\tradii: {rmin:.3e} -> {rmax:.3e}".format( + shell_nr=shell_nr, + zmin=z_shells[shell_nr, 0], zmax=z_shells[shell_nr, -1], + rmin=r_shells[shell_nr, 0], rmax=r_shells[shell_nr, -1] + ) + ) + v_shell=cosmo.lightcone_volume(r_shells[shell_nr, 0], r_shells[shell_nr, -1], use_redshift=False) + print("\tcomoving volume: {v_shell:.3e}".format(v_shell=v_shell)) + shell_volumes[shell_nr]=v_shell + + axs.add_artist(plt.Circle((0, 0), r_shells[shell_nr, -1].to_value("Mpc"), fc=colours[int(shell_nr)+1], ec='black', zorder=-100-shell_nr)) + +axs.scatter(0,0, fc='crimson', ec='black', marker='X', zorder=100, label='observer') + +cmap=mpl.colormaps['terrain_r'] +bounds = shell_volumes +print(bounds) +norm = mpl.colors.LogNorm(bounds[0],bounds[-1]) +fig.colorbar(mpl.cm.ScalarMappable(norm=norm, cmap=cmap), + orientation='vertical', shrink=1., ax=axs, extend="both", + label=r"shell volume [cMpc$^3$]") + +axs.set_xlim(-100, r_shells[n_shells+1, -1]) +axs.set_ylim(-0.5*r_shells[n_shells+1, -1].to_value("Mpc") -50, 0.5*r_shells[n_shells+1, -1].to_value("Mpc")+50) + +axs.set_xlabel("radius [cMpc]") +axs.set_ylabel("radius [cMpc]") + +axs.legend( + loc='best', + frameon=True, fancybox=False, + edgecolor='black', + borderpad=0.3,borderaxespad=0.7, + columnspacing=0.75, + labelspacing=0.25, handletextpad=0.25, + handleheight=1, handlelength=1.8, +) + +fig_name="shells_{fig_numb}".format(fig_numb=n_shells) +plt.savefig("./"+fig_name+".png",dpi=300, bbox_inches='tight') + +plt.close() diff --git a/examples/xray_emission_in_a_pencil_beam.py b/examples/xray_emission_in_a_pencil_beam.py new file mode 100644 index 0000000..f61375f --- /dev/null +++ b/examples/xray_emission_in_a_pencil_beam.py @@ -0,0 +1,161 @@ + +#!/bin/env python +import sys +import glob +import numpy as np +import lightcone_io.particle_reader as pr +import lightcone_io.lc_xray_calculator as Xcalc +from lightcone_io.xray_utils import Snapshot_Cosmology_For_Lightcone, Xray_Filter +import unyt +import h5py + +import matplotlib.pyplot as plt +import matplotlib.gridspec as gridspec + +########## plot settings ########## +import matplotlib as mpl +# plot settings can be removed if needed +# set plotting params +# Ticks inside plots; more space devoted to plot. +plt.rcParams["xtick.direction"] ='in' +plt.rcParams["ytick.direction"] ='in' +plt.rcParams["xtick.top"]=True +plt.rcParams["ytick.right"]= True + +# axes font settings +mpl.rcParams['axes.labelsize']=10 +mpl.rcParams['xtick.labelsize']=9 +mpl.rcParams['ytick.labelsize']=9 + +plt.rcParams["legend.fontsize"]=9 +################################### + +# simulation +sim_name="HYDRO_FIDUCIAL" +sim="L1000N1800/"+sim_name +# sim base dir +base_dir="/cosma8/data/dp004/flamingo/Runs/{sim}".format(sim="L1000N1800/"+sim_name) + +# Specify one file from the spatially indexed lightcone particle data +lightcone_nr=0 +input_filename = "{base_name}/particle_lightcones/lightcone{lc_nr}_particles/lightcone{lc_nr}_0000.0.hdf5".format(base_name=base_dir, lc_nr=lightcone_nr) + + + +vector = (1.0, 0.0, 0.0) # Vector pointing at a spot on the sky +radius = np.deg2rad(1.0) # Angular radius around that spot +redshift_range = (0., 0.02) #redshift range to load + +# Open the lightcone +lightcone = pr.IndexedLightcone(input_filename) + +# define required part type and properties +ptype = "Gas" +property_names = ( + "Coordinates", + "Masses", + "ExpansionFactors", + "Densities", + "SmoothedElementMassFractions", + "Temperatures", + "LastAGNFeedbackScaleFactors", + "StarFormationRates" + ) + +# read particle data from lightcone +particle_data = lightcone[ptype].read( + property_names=property_names, + redshift_range=redshift_range, + vector=vector, radius=radius +) + +nr_parts = len(particle_data["ExpansionFactors"]) +print("\nread {n_all} gas particles".format(n_all=nr_parts)) +# filter out: +# 1) starforming particles +# 2) too cold +# 3) particles that have recently been heated by AGN feedback + +# get cosmology information +# specify the snapshot to use the cosmology information from +snapshot_dir = "{base_name}/snapshots".format(base_name=base_dir) +cosmo=Snapshot_Cosmology_For_Lightcone(snapshot_dir) +particle_filter = Xray_Filter(particle_data, nr_parts, cosmo=cosmo) + +keep_particles = particle_filter.KEEP_FOR_XRAY +print("number of particles with x-ray values / total: {n_xray}/{n_all} ({n_frac:.5f})\n".format(n_xray=np.sum(keep_particles), n_all=nr_parts, n_frac=(nr_parts/np.sum(keep_particles)))) + + +# X-ray observation bands and types to compute values for +observation_bands=['erosita-high','erosita-low','ROSAT'] +observation_types=['photons_intrinsic', 'photons_convolved', 'energies_intrinsic', 'energies_convolved'] + +xray_emissivity_table_filename = "/cosma8/data/dp004/flamingo/Tables/Xray/X_Ray_table_combined.hdf5" + +Xray_flux, Xray_names = Xcalc.particle_xray_values_for_map( + observation_bands, + observation_types, + particle_data, + xray_emissivity_table_filename, + part_mask=keep_particles) + +# Sanity check for the output values +for i in range(len(Xray_names)): + for j in range(len(Xray_names[i])): + + map_str="{xray_value_name}\n\tparticle total: {value_total:.4e}".format(xray_value_name=Xray_names[i][j], value_total=np.sum(Xray_flux[i][:, j])) + value_range_str="\n\trange (min>0, max): ({min_val:.4e}, {max_val:.4e}) [{val_units}]".format(min_val=np.min(Xray_flux[i][:, j][Xray_flux[i][:, j]>0].value), max_val=np.max(Xray_flux[i][:, j].value), val_units=Xray_flux[i][:, j].units) + summary_str=map_str+value_range_str + print(summary_str) + +# show distribution of particle values + + +# create large range of bins +log_bin_min=-50 +log_bin_max=1 +log_bin_width=0.5 +bin_edges=10**np.arange(log_bin_min-log_bin_width/2, log_bin_max+(log_bin_width*1.5), log_bin_width) +midpoints=0.5*(bin_edges[:-1]+bin_edges[1:]) + +figure_xray_values = [['photons_intrinsic', 'photons_convolved'], ['energies_intrinsic', 'energies_convolved']] +for fig_idx in range(2): + numb_cols = 2 + fig = plt.figure(figsize=(5, 2.5)) + plot_grid = fig.add_gridspec(nrows=1, ncols=numb_cols, width_ratios=np.ones(numb_cols, dtype=int), height_ratios=[1]) + axs=plot_grid.subplots(sharex='row', sharey='row') + + for col_nr, xray_type in enumerate(figure_xray_values[fig_idx]): + ax=axs[col_nr] + for j, xray_band in enumerate(observation_bands): + n_pdf, __ = np.histogram(Xray_flux[col_nr+(2*fig_idx)][:, j].value, bins=bin_edges, density=True) + + m = n_pdf>0 + ax.plot(midpoints[m], n_pdf[m], label=xray_band, color=f'C{9-j}', zorder=-200) + + + ax.set_xscale("log") + ax.set_yscale("log") + + xray_units=unyt.unit_object.Unit(Xray_flux[col_nr+(2*fig_idx)][:, j].units) + xray_label_str=rf'$[{xray_units.latex_repr}]$' + ax.set_xlabel(xray_label_str) + ax.set_title(xray_type, fontsize=9) + + ax.legend( + loc='best', + frameon=False, fancybox=False, + borderpad=0.3,borderaxespad=0.7, + columnspacing=0.75, + labelspacing=0.25, handletextpad=0.25, + handleheight=1, handlelength=1.8, + ) + + axs[0].set_ylabel(r"PDF") + + fig_name="Xray_value_example_{fig_numb}".format(fig_numb=fig_idx) + plt.savefig("./"+fig_name+".png",dpi=300, bbox_inches='tight') + + plt.close() + + diff --git a/lightcone_io/__init__.py b/lightcone_io/__init__.py index de44b94..d1920bd 100644 --- a/lightcone_io/__init__.py +++ b/lightcone_io/__init__.py @@ -1,4 +1,4 @@ -__all__ = ["ShellArray", "Shell", "HealpixMap", "ParticleLightcone", "IndexedLightconeParticleType", "HaloLightconeFile"] +__all__ = ["ShellArray", "Shell", "HealpixMap", "ParticleLightcone", "IndexedLightconeParticleType", "HaloLightconeFile", "XrayCalculator_LC", "Snapshot_Cosmology_For_Lightcone", "Xray_Filter"] # Use setuptools_scm to get version from git tags from importlib.metadata import version, PackageNotFoundError @@ -16,3 +16,9 @@ # Classes for reading halo lightcones from .halo_reader import HaloLightconeFile + +# Classes for computing X-ray values +from .lc_xray_calculator import XrayCalculator_LC + +# Classes for using cosmology objects and filtering particles for X-rays +from .xray_utils import Snapshot_Cosmology_For_Lightcone, Xray_Filter diff --git a/lightcone_io/lc_xray_calculator.py b/lightcone_io/lc_xray_calculator.py new file mode 100644 index 0000000..db13bbe --- /dev/null +++ b/lightcone_io/lc_xray_calculator.py @@ -0,0 +1,634 @@ +#!/bin/env python +import sys +import h5py +import numpy as np +from numba import jit +import unyt +from unyt import g, cm, mp, erg, s, photons + +# +COMBINED_XRAY_EMISSIVITY_TABLE_FILENAME = "/cosma8/data/dp004/flamingo/Tables/Xray/X_Ray_table_combined.hdf5" + +xray_map_observation_type_units_cgs={ + 'photons_intrinsic':unyt.photons * unyt.cm**-2 * unyt.s**-1, + 'energies_intrinsic':unyt.erg * unyt.cm**-2 * unyt.s**-1, + 'photons_convolved':unyt.photons * unyt.s**-1, + 'energies_convolved':unyt.erg * unyt.s**-1, + } + + +class XrayCalculator_LC: + def __init__(self, redshifts, table, bands, observing_types): + if bands == None: + print('Please specify the band you would like to generate emissivities for\n \ + Using the "band = " keyword\n\n \ + Available options are:\n \ + "erosita-low" (0.2-2.3 keV)\n \ + "erosita-high" (2.3-8.0 keV)\n \ + "ROSAT" (0.5-2.0 keV)') + raise KeyError + + if observing_types == None: + print('Please specify whether you would like to generate photon or energie emissivities\n \ + Using the "observing_type = " keyword\n\n \ + Available options are:\n \ + "energies_intrinsic"\n \ + "photons_intrinsic"') + raise KeyError + + # we do not allow restframe X-ray observation types + + observing_types, bands = self.check_for_restframe(observing_types, bands) + + + if (bands != None) & (observing_types != None): + assert len(bands) == len(observing_types) + + self.tables = self.load_all_tables(redshifts, table, bands, observing_types) + + self.observation_type_luminosities_cgs_units={ + 'photons_intrinsic':photons * s**-1, + 'energies_intrinsic':erg * s**-1, + 'photons_convolved':photons * cm**2 * s**-1, + 'energies_convolved':erg *cm**2 * s**-1 + } + + + def check_for_restframe(self, observing_types, bands): + """ + remove the restframe X-ray observation types and corresponding bands + """ + for i, xray_type in enumerate(observing_types): + if xray_type in ["energies_intrinsic_restframe", "photons_intrinsic_restframe"]: + del observing_types[i] + del bands[i] + return observing_types, bands + + + def load_all_tables(self, redshifts, table, bands, observing_types): + ''' + Load the x-ray tables for the specified bands and observing-types + + Params: + redshifts: of the particles + table: either a path to the table to be read or a dictionary containing the tables themselves + bands, observing_types: the bands and observation types, within the band to add to tables. + ''' + + # given a table path, attempt to read the table: + if isinstance(table, str)==False: + try: + assert isinstance(table, dict) # confirm that the table is a dictionary + assert "Bins" in table # confirm the table has redshift bins + except: + raise ValueError("You must pass a working dictionary with the x-ray table values") + + # table has been passed in a dictionary and doesn't need to be read as h5file. + self.redshift_bins = table['Bins']['Redshift_bins'].astype(np.float32) + idx_z, _= self.get_index_1d(self.redshift_bins, redshifts) + ############ make it always load min redshift index value. + #min_idx_z = np.min(idx_z) + min_idx_z = 0 + ############ + max_idx_z = np.max(idx_z) + 2 + + self.He_bins = table['Bins']['He_bins'].astype(np.float32) + self.missing_elements = table['Bins']['Missing_element'] + self.element_masses = table['Bins']['Element_masses'].astype(np.float32) + + self.density_bins = table['Bins']['Density_bins'].astype(np.float32) + self.temperature_bins = table['Bins']['Temperature_bins'].astype(np.float32) + self.redshift_bins = table['Bins']['Redshift_bins'].astype(np.float32) + + self.log10_solar_metallicity = table['Bins']['Solar_metallicities'].astype(np.float32) + self.solar_metallicity = np.power(10, self.log10_solar_metallicity) + + elif isinstance(table, str): + # table is not a dictionary, therefore try read it as hdf5 file. + try: + table = h5py.File(table, "r") + except ValueError as e: + raise Exception("You must pass a working x-ray table path") from e + + self.redshift_bins = table['/Bins/Redshift_bins'][()].astype(np.float32) + idx_z, _= self.get_index_1d(self.redshift_bins, redshifts) + ############ make it always load min redshift index value. + #min_idx_z = np.min(idx_z) + min_idx_z = 0 + ############ + max_idx_z = np.max(idx_z) + 2 + + self.He_bins = table['/Bins/He_bins'][()].astype(np.float32) + self.missing_elements = table['/Bins/Missing_element'][()] + self.element_masses = table['Bins/Element_masses'][()].astype(np.float32) + + self.density_bins = table['/Bins/Density_bins/'][()].astype(np.float32) + self.temperature_bins = table['/Bins/Temperature_bins/'][()].astype(np.float32) + self.redshift_bins = table['/Bins/Redshift_bins'][()].astype(np.float32) + + self.log10_solar_metallicity = table['/Bins/Solar_metallicities/'][()].astype(np.float32) + self.solar_metallicity = np.power(10, self.log10_solar_metallicity) + + tables = {} + for band in bands: + tables[band] = {} + for observing_type in observing_types: + temp = table[band][observing_type][int(min_idx_z):int(max_idx_z), :, :, :, :].astype(np.float32) + # temp = np.swapaxes(temp, 2, 4) + # self.table_shape = temp.shape[:-1] + # temp = temp.reshape((temp.shape[0] * temp.shape[1] * temp.shape[2] * temp.shape[3], temp.shape[4])) + tables[band][observing_type] = temp + + return tables + + + + @staticmethod + @jit(nopython = True) + def get_index_1d(bins, subdata): + ''' + Find the closest bin index below the specified value, and the relative offset compared to that bin. + ''' + eps = 1e-4 + delta = (len(bins) - 1) / (bins[-1] - bins[0]) + + idx = np.zeros_like(subdata) + dx = np.zeros_like(subdata, dtype = np.float32) + for i, x in enumerate(subdata): + if x < bins[0] + eps: + # We are below the first element + idx[i] = 0 + dx[i] = 0 + elif x < bins[-1] - eps: + # Normal case + idx[i] = int((x - bins[0]) * delta) + dx[i] = (x - bins[int(idx[i])]) * delta + else: + # We are after the last element + idx[i] = len(bins) - 2 + dx[i] = 1 + + return idx, dx + + @staticmethod + @jit(nopython = True) + def get_index_1d_irregular(bins, subdata): + ''' + Find the closest bin index below the specified value, and the relative offset compared to that bin. + Unlike get_index_1d, this allows for irregular bin spacing + ''' + eps = 1e-6 + idx = np.zeros_like(subdata) + dx = np.zeros_like(subdata, dtype = np.float32) + + for i, x in enumerate(subdata): + if x < bins[0] + eps: + idx[i] = 0 + dx[i] = 0 + elif x < bins[-1] - eps: + min_idx = -1 + + ''' + Do this the hard way: Search the table + for the smallest index i in bins[i] such + that table[i] < x + ''' + for j in range(len(bins)): + if x - bins[j] <= 0: + # Found the first entry that is larger than x, go back by 1 + min_idx = j - 1 + break + + idx[i] = min_idx + dx[i] = (x - bins[min_idx]) / (bins[min_idx + 1] - bins[min_idx]) + else: + idx[i] = len(bins) - 2 + dx[i] = 1 + + return idx, dx + + @staticmethod + # @jit(nopython = True) + def get_table_interp(idx_z, idx_he, idx_T, idx_n, t_z, d_z, t_T, d_T, t_nH, d_nH, t_He, d_He, X_Ray, abundance_to_solar): + ''' + 4D interpolate the x-ray table for each traced metal + Scale the metals with their respective relative solar abundances + Add the metals to the background case without metals + ''' + + f_n_T = np.zeros((t_nH.shape[0], X_Ray.shape[1]), dtype = np.float32) + + f_n_T += (t_nH * t_He * t_T * t_z)[:, None] * X_Ray[idx_z, idx_he, :, idx_T, idx_n] + f_n_T += (t_nH * t_He * d_T * t_z)[:, None] * X_Ray[idx_z, idx_he, :, idx_T + 1, idx_n] + f_n_T += (t_nH * d_He * t_T * t_z)[:, None] * X_Ray[idx_z, idx_he + 1, :, idx_T, idx_n] + f_n_T += (d_nH * t_He * t_T * t_z)[:, None] * X_Ray[idx_z, idx_he, :, idx_T, idx_n + 1] + + + f_n_T += (t_nH * d_He * d_T * t_z)[:, None] * X_Ray[idx_z, idx_he + 1, :, idx_T + 1, idx_n] + f_n_T += (d_nH * t_He * d_T * t_z)[:, None] * X_Ray[idx_z, idx_he, :, idx_T + 1, idx_n + 1] + f_n_T += (d_nH * d_He * t_T * t_z)[:, None] * X_Ray[idx_z, idx_he + 1, :, idx_T, idx_n + 1] + f_n_T += (d_nH * d_He * d_T * t_z)[:, None] * X_Ray[idx_z, idx_he + 1, :, idx_T + 1, idx_n + 1] + + + f_n_T += (t_nH * t_He * t_T * d_z)[:, None] * X_Ray[idx_z + 1, idx_he, :, idx_T, idx_n] + f_n_T += (t_nH * t_He * d_T * d_z)[:, None] * X_Ray[idx_z + 1, idx_he, :, idx_T + 1, idx_n] + f_n_T += (t_nH * d_He * t_T * d_z)[:, None] * X_Ray[idx_z + 1, idx_he + 1, :, idx_T, idx_n] + f_n_T += (d_nH * t_He * t_T * d_z)[:, None] * X_Ray[idx_z + 1, idx_he, :, idx_T, idx_n + 1] + + f_n_T += (t_nH * d_He * d_T * d_z)[:, None] * X_Ray[idx_z + 1, idx_he + 1, :, idx_T + 1, idx_n] + f_n_T += (d_nH * t_He * d_T * d_z)[:, None] * X_Ray[idx_z + 1, idx_he, :, idx_T + 1, idx_n + 1] + f_n_T += (d_nH * d_He * t_T * d_z)[:, None] * X_Ray[idx_z + 1, idx_he + 1, :, idx_T, idx_n + 1] + f_n_T += (d_nH * d_He * d_T * d_z)[:, None] * X_Ray[idx_z + 1, idx_he + 1, :, idx_T + 1, idx_n + 1] + + + + # Add each metal contribution individually + f_n_T_Z_temp = np.power(10, f_n_T[:, -1], dtype=np.float64) + for j in range(f_n_T.shape[1] - 1): + f_n_T_Z_temp += np.power(10, f_n_T[:, j]) * abundance_to_solar[:, j] + + f_n_T_Z = np.log10(f_n_T_Z_temp) + + + + + return f_n_T_Z + + def find_indices(self, densities, temperatures, element_mass_fractions, masses, redshifts, fill_value = 0): + ''' + Check interpolation table bounds + Compute all interpolation bin indices, and the offsets from those bins + Compute all the indices for the flattened x-ray table + ''' + scale_factors = 1 / (1 + redshifts) + data_n = np.log10(element_mass_fractions[:, 0] * (1 / scale_factors**3) * densities.to(g * cm**-3) / mp) + data_T = np.log10(temperatures) + volumes = (masses.astype(np.float64) / ((1 / scale_factors**3) * densities.astype(np.float64))).to(cm**3) + + # Create density mask, round to avoid numerical errors + density_mask = (data_n >= np.round(self.density_bins.min(), 1)) & (data_n <= np.round(self.density_bins.max(), 1)) + # Create temperature mask, round to avoid numerical errors + temperature_mask = (data_T >= np.round(self.temperature_bins.min(), 1)) & (data_T <= np.round(self.temperature_bins.max(), 1)) + + # Combine masks + joint_mask = density_mask & temperature_mask + + + # Check if within density and temperature bounds + density_bounds = np.sum(density_mask) == density_mask.shape[0] + temperature_bounds = np.sum(temperature_mask) == temperature_mask.shape[0] + if ~(density_bounds & temperature_bounds): + #If no fill_value is set, return an error with some explanation + if fill_value == None: + raise ValueError( + "Temperature or density are outside of the interpolation range and no fill_value is supplied\n \ + Temperature ranges between log(T) = 5 and log(T) = 9.5\n \ + Density ranges between log(nH) = -8 and log(nH) = 6\n \ + Set the kwarg 'fill_value = some value' to set all particles outside of the interpolation range to 'some value'\n \ + Or limit your particle data set to be within the interpolation range" + ) + else: + pass + + #get individual mass fraction + mass_fraction = element_mass_fractions[joint_mask] + + # Find redshift offsets + idx_z, dx_z = self.get_index_1d(self.redshift_bins, redshifts[joint_mask]) + idx_z = idx_z.astype(int) + + # Find density offsets + idx_n, dx_n = self.get_index_1d(self.density_bins, data_n[joint_mask]) + idx_n = idx_n.astype(int) + + # Find temperature offsets + idx_T, dx_T = self.get_index_1d(self.temperature_bins, data_T[joint_mask]) + idx_T = idx_T.astype(int) + + # Calculate the abundance wrt to solar + abundances = (mass_fraction / np.expand_dims(mass_fraction[:, 0], axis = 1)) * (self.element_masses[0] / np.array(self.element_masses)) + + # Calculate abundance offsets using solar abundances + abundance_to_solar = abundances / self.solar_metallicity + + # Add columns for Calcium and Sulphur and move Iron to the end + abundance_to_solar = np.c_[abundance_to_solar[:, :-1], abundance_to_solar[:, -2], abundance_to_solar[:, -2], abundance_to_solar[:, -1]] + + #Find helium offsets + idx_he, dx_he = self.get_index_1d_irregular(self.He_bins, np.log10(abundances[:, 1])) + idx_he = idx_he.astype(int) + + + t_z = 1 - dx_z + d_z = dx_z + + # Compute temperature offset relative to bin + t_T = 1 - dx_T + d_T = dx_T + + # Compute density offset relative to bin + t_nH = 1 - dx_n + d_nH = dx_n + + # Compute Helium offset relative to bin + t_He = 1 - dx_he + d_He = dx_he + + return idx_z, idx_he, idx_T, idx_n, t_z, d_z, t_T, d_T, t_nH, d_nH, t_He, d_He, abundance_to_solar, joint_mask, volumes, data_n + + + def interpolate_X_Ray(self, idx_z, idx_he, idx_T, idx_n, t_z, d_z, t_T, d_T, t_nH, d_nH, t_He, d_He, abundance_to_solar, joint_mask, volumes, data_n, bands = None, observing_types = None, fill_value = None): + ''' + Main function + Calculate the particle emissivities through interpolation + Convert to luminosity using the particle volume + ''' + # Initialise the emissivity array which will be returned + emissivities = np.zeros((joint_mask.shape[0], len(bands)), dtype = float) + luminosities = np.zeros_like(emissivities) + emissivities[~joint_mask] = fill_value + + # Interpolate the table for each specified band + for i_interp, band, observing_type in zip(range(len(bands)), bands, observing_types): + + emissivities[joint_mask, i_interp] = self.get_table_interp(idx_z, idx_he, idx_T, idx_n, t_z, d_z, t_T, d_T, t_nH, d_nH, t_He, d_He, self.tables[band][observing_type], abundance_to_solar[:, 2:]) + + # Convert from erg cm^3 s^-1 to erg cm^-3 s^-1 + # To do so we multiply by nH^2, this is the actual nH not the nearest bin + # It allows to extrapolate in density space without too much worry + # log(emissivity * nH^2) = log(emissivity) + 2*log(nH) + emissivities[joint_mask, i_interp] += 2*data_n[joint_mask] + + luminosities[joint_mask, i_interp] = np.power(10, emissivities[joint_mask, i_interp]) * volumes[joint_mask] + + # get X-ray observation type units + luminosities_cgs_unyts = self.observation_type_luminosities_cgs_units[observing_types[0]] + + if 'energies' in observing_types[0]: + return luminosities * luminosities_cgs_unyts + elif 'photon' in observing_types[0]: + return luminosities * luminosities_cgs_unyts + + +def xray_map_names(observation_band, observation_type): + dataset_names={ + "ROSAT_photons_intrinsic":"XrayROSATIntrinsicPhotons", + "ROSAT_photons_convolved":"XrayROSATConvolvedPhotons", + "ROSAT_energies_intrinsic":"XrayROSATIntrinsicEnergies", + "ROSAT_energies_convolved":"XrayROSATConvolvedEnergies", + "erosita-high_photons_intrinsic":"XrayErositaHighIntrinsicPhotons", + "erosita-high_photons_convolved":"XrayErositaHighConvolvedPhotons", + "erosita-high_energies_intrinsic":"XrayErositaHighIntrinsicEnergies", + "erosita-high_energies_convolved":"XrayErositaHighConvolvedEnergies", + "erosita-low_photons_intrinsic":"XrayErositaLowIntrinsicPhotons", + "erosita-low_photons_convolved":"XrayErositaLowConvolvedPhotons", + "erosita-low_energies_intrinsic":"XrayErositaLowIntrinsicEnergies", + "erosita-low_energies_convolved":"XrayErositaLowConvolvedEnergies" + } + return dataset_names[observation_band+'_'+observation_type] + + + +def get_observation_type_per_band( + unique_observation_types=['photons_intrinsic', 'photons_convolved', 'energies_intrinsic', 'energies_convolved'], + unique_observation_bands=['erosita-high','erosita-low','ROSAT']): + + observation_bands=unique_observation_bands*len(unique_observation_types) + all_observation_type_per_band=[] + for obvs_type in unique_observation_types: + for i in range(len(unique_observation_bands)): + all_observation_type_per_band.append(obvs_type) + return observation_bands, all_observation_type_per_band + + +def compute_luminosities( + observation_bands, observation_types, particle_data, + emissivity_table_filename, part_mask=None): + """ + Compute X-ray luminosities in the given observation bands for each of the given observation types. + + Params + observation_bands [list] : ['erosita-high','erosita-low','ROSAT'] + observation_types [list] : ['photons_intrinsic', 'photons_convolved', 'energies_intrinsic', 'energies_convolved'] + particle_data : lightcone particle data, as read with lightcone_io.particle_reader + emissivity_table_filename : path to emissivity table + part_mask [boolean array] : apply this mask to the particle data, compute luminosities for particlse where True. + + Returns + Nested List of X-ray luminosities, Nested List of xray names "band_type" + """ + + # check all required properties are in particle data. + required_particle_properties = ["ExpansionFactors","Masses", "Densities", "SmoothedElementMassFractions","Temperatures", "StarFormationRates"] + missing_properties = [] + for prop_name in required_particle_properties: + if prop_name not in particle_data: + missing_properties.append(prop_name) + raise ValueError("particle {required_property} not found ....".format(required_property=prop_name)) + + # get all combinations of observation bands and types + all_observation_bands_per_type, all_observation_types_per_band = get_observation_type_per_band(observation_types, observation_bands) + + # determine if mask is required + nparts=len(particle_data["ExpansionFactors"]) + if part_mask is None: + part_mask = np.ones(nparts, dtype=bool) + + #create X-ray indices + xray_calc = XrayCalculator_LC( + (1./particle_data["ExpansionFactors"][part_mask])-1., # need to interpolate over the redshift range of particles + emissivity_table_filename, + bands=all_observation_bands_per_type, + observing_types=all_observation_types_per_band) + + # X-ray indicies + idx_z, idx_he, idx_T, idx_n, t_z, d_z, t_T, d_T, t_nH, d_nH, t_He, d_He, abundance_to_solar, joint_mask, volumes, data_n = xray_calc.find_indices( + particle_data["Densities"][part_mask], + particle_data["Temperatures"][part_mask], + particle_data["SmoothedElementMassFractions"][part_mask], + particle_data["Masses"][part_mask], + (1./particle_data["ExpansionFactors"][part_mask])-1., + fill_value = 0 + ) + + # make empty nested list for each observation type + xray_luminosities_units_cgs=xray_calc.observation_type_luminosities_cgs_units + + XRAY_VALUES = [ [] for i in range(len(observation_types))] + XRAY_NAMES = [ [] for i in range(len(observation_types))] + + #iterate through each observation type and compute X-ray values in required bands + for type_idx, observation_type in enumerate(observation_types): + if observation_type not in xray_luminosities_units_cgs: + raise ValueError("Cannot find observation type and its matching units") + + # compute luminosities in all bands for given X-ray type + if len(observation_bands) == 1: + input_types=[observation_type] + elif len(observation_bands) > 1: + input_types = [observation_type]*len(observation_bands) + + luminosities = xray_calc.interpolate_X_Ray( + idx_z, idx_he, idx_T, idx_n, t_z, d_z, t_T, d_T, t_nH, d_nH, t_He, d_He, abundance_to_solar, joint_mask, volumes, data_n, + bands = observation_bands, observing_types = input_types, fill_value = 0) + + # unit check + assert luminosities.units == xray_calc.observation_type_luminosities_cgs_units[observation_type] + + + # check correct number of x-ray values are returned per band + assert np.shape(luminosities)[0] == np.sum(part_mask) + + if len(observation_bands) == 1: + luminosities.flatten() + + XRAY_VALUES[type_idx] = unyt.unyt_array( + np.zeros(nparts), + units=xray_calc.observation_type_luminosities_cgs_units[observation_type], dtype=float) + + XRAY_VALUES[type_idx][part_mask]+=luminosities + XRAY_NAMES[type_idx]=[observation_bands[0]+"_"+observation_type] + + elif len(observation_bands) > 1: + XRAY_VALUES[type_idx] = unyt.unyt_array( + np.zeros((nparts, len(observation_bands))), + units=xray_calc.observation_type_luminosities_cgs_units[observation_type], dtype=float) + for axis_idx in range(len(observation_bands)): + XRAY_VALUES[type_idx][part_mask, axis_idx] += luminosities[:, axis_idx] + + XRAY_NAMES[type_idx]=[band+"_"+observation_type for band in observation_bands] + + #clean up per loop + del luminosities + + return XRAY_VALUES, XRAY_NAMES + + +def particle_xray_values_for_map( + observation_bands, observation_types, particle_data, + emissivity_table_filename, part_mask=None): + """ + Compute X-ray values for all-sky in the given observation bands for each of the given observation types. + + Params + observation_bands [list] : ['erosita-high','erosita-low','ROSAT'] + observation_types [list] : ['photons_intrinsic', 'photons_convolved', 'energies_intrinsic', 'energies_convolved'] + particle_data : lightcone particle data, as read with lightcone_io.particle_reader + emissivity_table_filename : path to emissivity table + part_mask [boolean array] : apply this mask to the particle data, compute luminosities for particlse where True. + + Returns + Nested List of X-ray values, Nested List of xray map names + """ + + required_particle_properties = ["ExpansionFactors","Coordinates","Masses", "Densities", "SmoothedElementMassFractions","Temperatures", "StarFormationRates"] + missing_properties = [] + for prop_name in required_particle_properties: + if prop_name not in particle_data: + missing_properties.append(prop_name) + raise ValueError("particle {required_property} not found ....".format(required_property=prop_name)) + + # compute the luminosity distance + cdist_cross = np.sqrt( + particle_data["Coordinates"][part_mask, 0]**2 + + particle_data["Coordinates"][part_mask, 1]**2 + + particle_data["Coordinates"][part_mask, 2]**2 + ) + luminosity_distance_no_z = 4 * np.pi * (cdist_cross**2) + + # map units in cgs + xray_map_observation_type_units_cgs={ + 'photons_intrinsic':unyt.photons * unyt.cm**-2 * unyt.s**-1, + 'energies_intrinsic':unyt.erg * unyt.cm**-2 * unyt.s**-1, + 'photons_convolved':unyt.photons * unyt.s**-1, + 'energies_convolved':unyt.erg * unyt.s**-1, + } + + # indices of redshift component in luminosity distance + luminosity_distance_redshift_power_dict={ + 'photons_intrinsic':1, + 'photons_convolved':1, + 'energies_intrinsic':2, + 'energies_convolved':2 + } + + # get all combinations of observation bands and types + all_observation_bands_per_type, all_observation_types_per_band = get_observation_type_per_band(observation_types, observation_bands) + + # determine if mask is required + nparts=len(particle_data["ExpansionFactors"]) + if part_mask is None: + part_mask = np.ones(nparts, dtype=bool) + + #create X-ray indices + xray_calc = XrayCalculator_LC( + (1./particle_data["ExpansionFactors"][part_mask])-1., # need to interpolate over the redshift range of particles + emissivity_table_filename, + bands=all_observation_bands_per_type, + observing_types=all_observation_types_per_band) + + # X-ray indicies + idx_z, idx_he, idx_T, idx_n, t_z, d_z, t_T, d_T, t_nH, d_nH, t_He, d_He, abundance_to_solar, joint_mask, volumes, data_n = xray_calc.find_indices( + particle_data["Densities"][part_mask], + particle_data["Temperatures"][part_mask], + particle_data["SmoothedElementMassFractions"][part_mask], + particle_data["Masses"][part_mask], + (1./particle_data["ExpansionFactors"][part_mask])-1., + fill_value = 0 + ) + + + XRAY_VALUES = [ [] for i in range(len(observation_types))] + XRAY_NAMES = [ [] for i in range(len(observation_types))] + + #iterate through each observation type and compute X-ray values in required bands + for type_idx, observation_type in enumerate(observation_types): + + #check observation type has units + if observation_type not in xray_map_observation_type_units_cgs: + raise ValueError("Cannot find observation type and its matching units") + + # check observation type has known luminosity distances + if observation_type in luminosity_distance_redshift_power_dict: + luminosity_distance_redshift_power=luminosity_distance_redshift_power_dict[observation_type] + else: + raise ValueError("Cannot find observation type and its needed luminosity distance redshift power!!!") + + # compute luminosities in all bands for given X-ray type + if len(observation_bands) == 1: + input_types=[observation_type] + elif len(observation_bands) > 1: + input_types = [observation_type]*len(observation_bands) + + luminosities = xray_calc.interpolate_X_Ray( + idx_z, idx_he, idx_T, idx_n, t_z, d_z, t_T, d_T, t_nH, d_nH, t_He, d_He, abundance_to_solar, joint_mask, volumes, data_n, + bands = observation_bands, observing_types = input_types, fill_value = 0) + + # unit check + assert (luminosities * unyt.cm**-2).units == xray_map_observation_type_units_cgs[observation_type] + + # check correct number of x-ray values are returned per band + assert np.shape(luminosities)[0] == np.sum(part_mask) + + if len(observation_bands) == 1: + luminosities.flatten() + + XRAY_VALUES[type_idx] = unyt.unyt_array( + np.zeros(nparts), + units=xray_map_observation_type_units_cgs[observation_type], dtype=float) + + XRAY_VALUES[type_idx][part_mask]+=luminosities/ (luminosity_distance_no_z * ((1/particle_data["ExpansionFactors"][part_mask].value)**luminosity_distance_redshift_power)) + XRAY_NAMES[type_idx]=[observation_bands[0]+"_"+observation_type] + + elif len(observation_bands) > 1: + XRAY_VALUES[type_idx] = unyt.unyt_array( + np.zeros((nparts, len(observation_bands))), + units=xray_map_observation_type_units_cgs[observation_type], dtype=float) + for axis_idx in range(len(observation_bands)): + XRAY_VALUES[type_idx][part_mask, axis_idx] += luminosities[:, axis_idx] / (luminosity_distance_no_z * ((1/particle_data["ExpansionFactors"][part_mask].value)**luminosity_distance_redshift_power)) + + XRAY_NAMES[type_idx]=[xray_map_names(band, observation_type) for band in observation_bands] + + + return XRAY_VALUES, XRAY_NAMES + diff --git a/lightcone_io/smoothed_map.py b/lightcone_io/smoothed_map.py index 96b72a6..1841565 100644 --- a/lightcone_io/smoothed_map.py +++ b/lightcone_io/smoothed_map.py @@ -7,7 +7,7 @@ import numpy as np import h5py import unyt - +import datetime as dt import lightcone_io.particle_reader as pr import lightcone_io.kernel as kernel @@ -154,12 +154,30 @@ def create_empty_array(arr, comm): return arr -def message(m): +def message(m, show_time=False): comm.barrier() if comm_rank == 0: + + if show_time == True: # add current time to message + current_time=dt.datetime.now() + time_str=current_time.strftime("%H:%M:%S") + m="[{print_time}]\t".format(print_time=time_str)+m print(m) +def rank_message(m, rank): + """ + Print a new message on each rank. + No rank barrier as we want independant messages + from each rank with time of arrival shown. + """ + #comm.barrier() + current_time=dt.datetime.now() + time_str=current_time.strftime("%H:%M:%S") + print('\t[Rank {rank_nr:03d}] [@{print_time}]'.format(rank_nr=rank,print_time=time_str) + m) + + + def distribute_pixels(comm, nside): """ Decide how to assign HEALPix pixels to MPI ranks. @@ -210,8 +228,60 @@ def distribute_pixels(comm, nside): return nr_total_pixels, nr_local_pixels, local_offset, theta_boundary +def rotate_particle_coordinates(input_coords, theta, phi): + """ + Params + input_coords (unyt array): the coordinates of particles within the + lightcone shell [co-moving Mpc], as read + from particle data of particle lightcone + theta, phi (unyt quantity): angles to rotate particle coordinates + about [units=degree] + + Returns + output_coords (unyt array): the particle coordiantes (input_coords) + rotated on the sky via healpy rotator + functions + """ + + chunk_size=65536 # from SWIFTSIMIO softenning length code, this seems to be a good value ?? + number_of_parts=np.shape(input_coords)[0] + number_of_chunks=1+number_of_parts // chunk_size + + rotated_coordinates=unyt.unyt_array( + np.zeros((number_of_parts, 3), dtype=np.float64), + units=input_coords.units) + + def rotate_vec(coords, theta, phi): + + if (theta.value != 0.) or (phi.value != 0.): + rot_custom = hp.Rotator( + rot=[theta.to_value(unyt.deg), phi.to_value(unyt.deg)], + inv=True, deg=True, eulertype='ZYX') + rot_matrix = rot_custom._matrix + output_arr=hp.rotator.rotateVector( + rot_matrix, + vec=coords[:,0].value, vy=coords[:,1].value, vz=coords[:,-1].value).T + return output_arr + else: + return input_coords + + rank_message("rotating particles on the sky ({ang0:.3f}, {ang1:.3f})".format(ang0=theta, ang1=phi), comm_rank) + + for chunk in range(number_of_chunks): + start_id = chunk * chunk_size + end_id = (chunk+1) * chunk_size + + if end_id > number_of_parts: + end_id = number_of_parts+1 + if start_id >= end_id: + break + + rotated_coordinates[start_id:end_id, :] = rotate_vec(input_coords[start_id:end_id, :], theta, phi) + + return rotated_coordinates + def make_sky_map(input_filename, ptype, property_names, particle_value_function, - zmin, zmax, nside, vector = None, radius = None, smooth=True, progress=False): + zmin, zmax, nside, rot_theta=None, rot_phi=None, vector = None, radius = None, smooth=True, progress=False): """ Make a new HEALPix map from lightcone particle data @@ -224,6 +294,7 @@ def make_sky_map(input_filename, ptype, property_names, particle_value_function, zmax : maximum redshift to use nside : HEALPix resolution parameter smooth : whether to smooth the map + rot_theta,rot_phi: angles to rotate particle coorinates on the sky [units=Degree] """ if progress and comm_rank == 0: @@ -231,14 +302,14 @@ def make_sky_map(input_filename, ptype, property_names, particle_value_function, progress_bar = tqdm else: progress_bar = lambda x: x - + # Ensure property_names list contains coordinates and smoothing lengths property_names = list(property_names) if "Coordinates" not in property_names: property_names.append("Coordinates") if smooth and ("SmoothingLengths" not in property_names): property_names.append("SmoothingLengths") - + # Open the lightcone lightcone = pr.IndexedLightcone(input_filename, comm=comm) @@ -262,8 +333,19 @@ def make_sky_map(input_filename, ptype, property_names, particle_value_function, nr_parts_tot = comm.allreduce(particle_data["Coordinates"].shape[0]) message(f"Read in {nr_parts_tot} particles") - # Find the particle positions and smoothing lengths - part_pos_send = particle_data["Coordinates"] + # Find the particle positions and rotate if needed. + if (rot_theta is None) and (rot_phi is None): + part_pos_send = particle_data["Coordinates"] # no rotation + else: + # give rotation value of 0 degree if rotation angle is not given + if rot_theta is None: + rot_theta=0.*unyt.degree + if rot_phi is None: + rot_phi=0.*unyt.degree + # apply rotation to particle coordinates + part_pos_send = rotate_particle_coordinates(particle_data, rot_theta, rot_phi) + + # Find Find the particle positions and smoothing lengths if smooth: part_hsml_send = find_angular_smoothing_length(part_pos_send, particle_data["SmoothingLengths"]) else: diff --git a/lightcone_io/units.py b/lightcone_io/units.py index 773dc65..1fe749c 100644 --- a/lightcone_io/units.py +++ b/lightcone_io/units.py @@ -7,7 +7,9 @@ # UNIT_CONV_NO_UNITS = {"U_"+base_unit : 0 for base_unit in "MLtIT"} UNIT_CONV_PHOTON_FLUX_PER_UNIT_SURFACE = dict(UNIT_CONV_NO_UNITS, U_L=-2.0, U_t=-1.0) -UNIT_CONV_ENERGY_FLUX_PER_UNIT_SURFACE = dict(UNIT_CONV_NO_UNITS, U_M=1.0, U_t=-3.0) +UNIT_CONV_ENERGY_FLUX_PER_UNIT_SURFACE = dict(UNIT_CONV_NO_UNITS, U_M= 1.0, U_t=-3.0) +UNIT_CONV_PHOTON_CONVOLVED_FLUX_PER_UNIT_SURFACE = dict(UNIT_CONV_NO_UNITS, U_t=-1.0) +UNIT_CONV_ENERGY_CONVOLVED_FLUX_PER_UNIT_SURFACE = dict(UNIT_CONV_NO_UNITS, U_M=1.0, U_L=2.0, U_t=-3.0) missing_units = { "XrayErositaLowIntrinsicPhotons" : UNIT_CONV_PHOTON_FLUX_PER_UNIT_SURFACE, "XrayErositaHighIntrinsicPhotons" : UNIT_CONV_PHOTON_FLUX_PER_UNIT_SURFACE, @@ -15,6 +17,12 @@ "XrayErositaLowIntrinsicEnergies" : UNIT_CONV_ENERGY_FLUX_PER_UNIT_SURFACE, "XrayErositaHighIntrinsicEnergies" : UNIT_CONV_ENERGY_FLUX_PER_UNIT_SURFACE, "XrayROSATIntrinsicEnergies" : UNIT_CONV_ENERGY_FLUX_PER_UNIT_SURFACE, + "XrayErositaLowConvolvedPhotons" : UNIT_CONV_PHOTON_CONVOLVED_FLUX_PER_UNIT_SURFACE, + "XrayErositaHighConvolvedPhotons" : UNIT_CONV_PHOTON_CONVOLVED_FLUX_PER_UNIT_SURFACE, + "XrayROSATConvolvedPhotons" : UNIT_CONV_PHOTON_CONVOLVED_FLUX_PER_UNIT_SURFACE, + "XrayErositaLowConvolvedEnergies" : UNIT_CONV_ENERGY_CONVOLVED_FLUX_PER_UNIT_SURFACE, + "XrayErositaHighConvolvedEnergies" : UNIT_CONV_ENERGY_CONVOLVED_FLUX_PER_UNIT_SURFACE, + "XrayROSATConvolvedEnergies" : UNIT_CONV_ENERGY_CONVOLVED_FLUX_PER_UNIT_SURFACE, } # Name of attributes in InternalCodeUnits and Units groups diff --git a/lightcone_io/xray_map_all_bands.py b/lightcone_io/xray_map_all_bands.py new file mode 100644 index 0000000..00552e3 --- /dev/null +++ b/lightcone_io/xray_map_all_bands.py @@ -0,0 +1,685 @@ +#!/bin/env python +# +# This script is adapted from lightcone_io.smoothed_map +# +# This script uses the lightcone particle data to create a new HEALPix map. +# Parallelised with MPI. +# +# general +import sys +import os +import healpy as hp +import numpy as np +import h5py +import unyt +from unyt import mp +import datetime +import psutil +import argparse +# lightcone io +import lightcone_io.particle_reader as pr +import lightcone_io.kernel as kernel +import lightcone_io.lc_xray_calculator as Xcalc +from lightcone_io.xray_utils import Snapshot_Cosmology_For_Lightcone, Xray_Filter +from lightcone_io.smoothed_map import find_angular_smoothing_length, exchange_particles, distribute_pixels, rotate_particle_coordinates, message, rank_message +# virgo dc +import virgo.mpi.parallel_sort as psort +import virgo.mpi.parallel_hdf5 as phdf5 +# mpi +from mpi4py import MPI + + +comm = MPI.COMM_WORLD +comm_size = comm.Get_size() +comm_rank = comm.Get_rank() +projected_kernel = kernel.ProjectedKernel() + +def explode_particle_in_all_bands(nside, + part_pos, + angular_smoothing_length, + part_val_0, part_val_1, part_val_2): + """ + Same as for smoothed map, but now returns values in 3 seperate bands. + + Given a particle's position vector and angular radius, + return indexes of the pixels it will update and the values + to add to the pixels. + + Params: + part_pos: particle's position vector + part_val: particle's contribution to the map, [0,1,2]=x-ray bands + angular_smoothing_length: smoothing length in radians + + Returns: + pix_index: array of indexes of the pixels to update + pix_val: array of values to add to the pixels + """ + + # Normalize position vector + part_pos = part_pos / np.sqrt(np.sum(part_pos**2, dtype=float)) + + # Find radius containing the pixels to update + angular_search_radius = angular_smoothing_length*kernel.kernel_gamma + + # Get pixel indexes to update + pix_index = hp.query_disc(nside, part_pos, angular_search_radius) + assert len(pix_index) >= 1 + + # For each pixel, find angle between pixel centre and the particle + pix_vec_x, pix_vec_y, pix_vec_z = hp.pixelfunc.pix2vec(nside, pix_index) + dp = part_pos[0]*pix_vec_x + part_pos[1]*pix_vec_y + part_pos[2]*pix_vec_z + dp = np.clip(dp, a_min=None, a_max=1.0) + pix_angle = np.arccos(dp) + + del pix_vec_x + del pix_vec_y + del pix_vec_z + + # Evaluate the projected kernel for each pixel + pix_weight = projected_kernel(pix_angle/angular_smoothing_length) + del pix_angle + + # find number of bands being used: + assert part_val_0 is not None # must have atleast 1 band given as an input + + # remove pixels with a weight of 0 (e.g. pix_angle>=search_radius) + # can occur due to roudning or problem of single particles not being removed + in_range = pix_weight>0 + n_in_range=np.sum(in_range) + + if n_in_range==len(in_range): + # all pixels to update have a weight > 0 + del in_range # mask not needed + pix_weight=pix_weight + pix_index=pix_index + pix_weight_total = np.sum(pix_weight, dtype=float) + + pix_val_0 = part_val_0 * pix_weight / pix_weight_total + pix_val_1 = part_val_1 * pix_weight / pix_weight_total + pix_val_2 = part_val_2 * pix_weight / pix_weight_total + + elif n_in_range0: + # skip pixels if weight == 0 + pix_weight=pix_weight[in_range] #update pixel weights + pix_index=pix_index[in_range] #update pixel index + del in_range # mask no longer needed + + pix_weight_total = np.sum(pix_weight, dtype=float) #weights total + + pix_val_0 = part_val_0 * pix_weight / pix_weight_total + pix_val_1 = part_val_1 * pix_weight / pix_weight_total + pix_val_2 = part_val_2 * pix_weight / pix_weight_total + + else: #there are no pixels to update + nr_out_of_range=np.sum(in_range==False) + del in_range + pix_index=None + pix_val_0=None + pix_val_1=None + pix_val_2=None + + return pix_index, pix_val_0, pix_val_1, pix_val_2 + + +def get_map_attributes(nside, units): + + """ + Given a unyt.Unit object, generate SWIFT dataset attributes + + units: the Unit object + + Returns a dict with the maps attributes + """ + attrs = {} + # Get CGS conversion factor. Note that this is the conversion to physical units, + # because unyt multiplies out the dimensionless a factor. + cgs_factor, offset = units.get_conversion_factor(units.get_cgs_equivalent()) + + # values are physical and therefore we do not care about the scale factor + physical=True + a_exponent_in_units=0 + a_val=1 + a_exponent=None + + # Ignore h exponenent in map output + h_exponent=0 + h_val=1. + + # Check a_exponent is consistent + if a_exponent is None: + assert physical + else: + if physical: + assert a_exponent_in_units == 0 + else: + assert float(a_exponent_in_units) == a_exponent + # Something has gone wrong if we have h factors + assert h_exponent == 0 + + # Find the power associated with each dimension + powers = units.get_mks_equivalent().dimensions.as_powers_dict() + # Set the attributes + attrs["Conversion factor to CGS (not including cosmological corrections)"] = [ + float(cgs_factor / (a_val**a_exponent_in_units) / (h_val**h_exponent)) + ] + attrs["Conversion factor to physical CGS (including cosmological corrections)"] = [ + float(cgs_factor) + ] + attrs["U_I exponent"] = [float(powers[unyt.dimensions.current_mks])] + attrs["U_L exponent"] = [float(powers[unyt.dimensions.length])] + attrs["U_M exponent"] = [float(powers[unyt.dimensions.mass])] + attrs["U_T exponent"] = [float(powers[unyt.dimensions.temperature])] + attrs["U_t exponent"] = [float(powers[unyt.dimensions.time])] + attrs["h-scale exponent"] = [float(h_exponent)] + attrs["a-scale exponent"] = [0.0 if a_exponent is None else a_exponent] + attrs["Value stored as physical"] = [1 if physical else 0] + attrs["Property can be converted to comoving"] = [0 if a_exponent is None else 1] + attrs["nside"]=[int(nside)] + attrs["number_of_pixels"]=[int(hp.nside2npix(nside))] + attrs["pixel_ordering_scheme"]=bytes('ring',"utf-8") + attrs["Expression for physical CGS units"]=[str(units.get_cgs_equivalent())] + attrs["comoving_inner_radius"]=[-1] + attrs["comoving_outer_radius"]=[-1] + + return attrs + + +def get_xray_table(observation_bands, table_filename): + table_dict={} + with h5py.File(table_filename, 'r') as f: + table_dict['Bins']={} + for k, v in f['Bins'].items(): + if k =='Missing_element': + table_dict['Bins'][k]=v[()] + else: + table_dict['Bins'][k]=v[()].astype(np.float32) + for observation_band_name in observation_bands: + table_dict[observation_band_name]={} + for k, v in f[observation_band_name].items(): + table_dict[observation_band_name][k]=v[()].astype(np.float64) + return table_dict + + +def get_cosmology_object(snapshot_root_dir): + """ + Return the cosmology object + """ + return Snapshot_Cosmology_For_Lightcone(snapshot_root_dir) + + + +def write_and_restart_output_shell(output_filename, new_map_names, xray_table_filename, input_filename): + """ + Make sure there is a .hdf5 file to write new X-ray maps into. + + - If the outputfile .hdf5 file does not exist then create new output file. + - If outputfile does exist, remove X-ray maps with the same name as those + being made and record the restart date. + + + !! Should only be written with a single rank !! + + """ + start_date = datetime.datetime.now().strftime("%d-%B-%Y") # track date + + try: + # check if the shell exists, then update with restart info for new maps. + if os.path.isfile(output_filename): + print(f"\nWriting to existing shell: {output_filename}\n", flush=True) + # we will add to the output file + with h5py.File(output_filename, "a") as outfile: + # delete dataset with same name as new maps being made + for map_name in new_map_names: + if map_name in outfile: + del outfile[map_name] + + # write Restart information statement + if "__xrayInfo" not in outfile: + NewMapInfo = outfile.create_group("__xrayInfo") #create group that new map info can be written too + elif "__xrayInfo" in outfile: + NewMapInfo = outfile["__xrayInfo"] + # clear old attrs if they exist for new map name + for map_name in new_map_names: + if map_name in NewMapInfo.attrs: + del NewMapInfo.attrs[map_name] + for map_name in new_map_names: + # write restart date and emissivity table + NewMapInfo.attrs[map_name] = (start_date, xray_table_filename) + else: + # write a new shell + with h5py.File(output_filename, "w") as outfile: + print(f"\nWriting to empty shell: {output_filename}\n", flush=True) + NewMapInfo = outfile.create_group("__xrayInfo") # create group that new map info can be written too + NewMapInfo.attrs["ParticleLightcone"]=input_filename #path to particle lightcones used + NewMapInfo.attrs["Created"]=start_date #date shell was written + # write date and emissivity table + for map_name in new_map_names: + NewMapInfo.attrs[map_name] = (start_date, xray_table_filename) + ShellInfo=outfile.create_group("Shell") + ShellInfo.attrs["nr_files_per_shell"]=1 # should only have 1 file per shell, with multiple maps. + ShellInfo.attrs["comoving_inner_radius"]=-1 # leave as non-physical for now + ShellInfo.attrs["comoving_outer_radius"]=-1 # leave as non-physical for now + except: + raise FileNotFoundError("Cannot find existing or create new output .hdf5 file") + + return None + + +def write_smoothed_xray_map_in_all_bands( + input_filename, property_names, + zmin, zmax, nside, cosmo, xray_tables, output_filename, + xray_type, + vector = None, radius = None, theta=0., phi=0., + ): + """ + Compute the X-ray emission from hot gas in the particle lightcone and + write a HEALPix map for each X-ray observation band for a given X-ray + observation type. + + input_filename : name of one file from the particle lightcone output + property_names : gas particle properties to read in, e.g., + ("Coordinates","ExpansionFactors","Temperatures") + zmin : minimum redshift of gas particles used + zmax : maximum redshift of gas particles used + nside : HEALPix resolution parameter + cosmo : cosmology object + xray_tables : X-ray interpolation table, either as path to + file or dictionary + output_filename : path to h5file that the maps will be written to + xray_type : the X-ray observation type, determines units + of output X-ray flux. + theta, phi : angles to rotate particle positions on the + sky [degree] (must be a unyt quantity) + """ + + # Ensure property_names list contains required properties for X-ray maps + property_names = list(property_names) + required_particle_properties = [ + "SmoothingLengths", "Coordinates","ExpansionFactors", "Masses", + "Densities", "SmoothedElementMassFractions","Temperatures", + "StarFormationRates"] + for prop_name in required_particle_properties: + if prop_name not in property_names: + property_names.append(prop_name) + + # Open the lightcone + lightcone = pr.IndexedLightcone(input_filename, comm=comm) + + # Create an empty HEALPix map, distributed over MPI ranks. + # Here we assume we can put equal sized chunks of the map on each rank. + nr_total_pixels, nr_local_pixels, local_offset, theta_boundary = distribute_pixels(comm, nside) + max_pixrad = hp.pixelfunc.max_pixrad(nside) + message(f"Total number of pixels = {nr_total_pixels}") + + # Will read the full sky within the redshift range + redshift_range = (zmin, zmax) + + # Determine number of particles to read on this MPI rank + nr_particles_local = lightcone["Gas"].count_particles(redshift_range=redshift_range, + vector=vector, radius=radius,) + nr_particles_total = comm.allreduce(nr_particles_local) + message(f"Total number of particles in selected cells = {nr_particles_total:.4e}") + + # Read in the particle data + particle_data = lightcone["Gas"].read_exact(property_names, vector, radius, redshift_range) + nr_parts_tot = comm.allreduce(particle_data["Coordinates"].shape[0]) + message(f"Read in {nr_parts_tot:.4e} particles") + + # remove particles that are not able to give an X-ray value + # X-ray filter particles: exclude based on SFR, temp, RHP, density + particle_filter = Xray_Filter(particle_data, particle_data["Coordinates"].shape[0], cosmo=cosmo) + keep_particles = particle_filter.KEEP_FOR_XRAY + for prop_name in property_names: + particle_data[prop_name] = particle_data[prop_name][keep_particles] + + nr_parts_tot = comm.allreduce(particle_data["Coordinates"].shape[0]) + message(f"Read in {nr_parts_tot:.4e} particles [Post Filtering]") + del keep_particles # clean up + + # Find the particle positions and smoothing lengths + if theta !=0. or phi !=0.: + message(f"Rotating by (theta, phi): ({theta}, {phi})") + part_pos_send = rotate_particle_coordinates(particle_data, theta, phi) + else: + part_pos_send = particle_data["Coordinates"] + + part_hsml_send = find_angular_smoothing_length(part_pos_send, particle_data["SmoothingLengths"]) + + + # Determine range of colatitudes each particle will update. + # Smoothing kernel drops to zero at kernel_gamma * smoothing length. + # Note that particles with radius < max_pixrad can still update + # pixels up to max_pixrad away because they update whatever pixel + # they are in. + + radius = np.maximum(kernel.kernel_gamma*part_hsml_send, max_pixrad) # Might update pixels with centres within this radius + theta, phi = hp.pixelfunc.vec2ang(part_pos_send) + part_min_theta = np.clip(theta-radius, 0.0, np.pi) # Minimum central theta of pixels each particle might update + part_max_theta = np.clip(theta+radius, 0.0, np.pi) # Maximum central theta of pixels each particle might update + + # Determine what range of MPI ranks each particle needs to be sent to + part_first_rank = np.searchsorted(theta_boundary, part_min_theta, side="left") - 1 + part_first_rank = np.clip(part_first_rank, 0, comm_size-1) + part_last_rank = np.searchsorted(theta_boundary, part_max_theta, side="right") - 1 + part_last_rank = np.clip(part_last_rank, 0, comm_size-1) + assert np.all(theta_boundary[part_first_rank] <= part_min_theta) + assert np.all(theta_boundary[part_last_rank+1] >= part_max_theta) + + # Determine how many ranks each particle needs to be sent to + nr_copies = part_last_rank - part_first_rank + 1 + assert np.all(nr_copies>=1) and np.all(nr_copies<=comm_size) + + # Duplicate the particles + nr_parts = part_pos_send.shape[0] + index = np.repeat(np.arange(nr_parts, dtype=int), nr_copies) + part_pos_send = part_pos_send[index,...] + #part_val_send = part_val_send[index] # not needed, 3 different values inside loop + part_hsml_send = part_hsml_send[index] + #del index # do not remove index as we use it inside the loop + + # Determine destination rank for each particle copy + nr_parts = np.sum(nr_copies) # Total number of copied particles + offset = np.cumsum(nr_copies) - nr_copies # Offset to first copy of each particle in array of copies + part_dest = -np.ones(nr_parts, dtype=int) # Destination rank for each copied particle + for pfr, nrc, off in zip(part_first_rank, nr_copies, offset): + part_dest[off:off+nrc] = np.arange(pfr, pfr+nrc, dtype=int) + assert np.all(part_dest >=0) & np.all(part_dest=0) & (local_pix_index < nr_local_pixels) + + # add single pixel updates for each map + np.add.at(map_view_eROSITA_high, local_pix_index[local], part_val_recv_view_eROSITA_high[single_pixel][local]) + np.add.at(map_view_eROSITA_low, local_pix_index[local], part_val_recv_view_eROSITA_low[single_pixel][local]) + np.add.at(map_view_ROSAT, local_pix_index[local], part_val_recv_view_ROSAT[single_pixel][local]) + + del part_pos_recv_view + del part_val_recv_view_eROSITA_high + del part_val_recv_view_eROSITA_low + del part_val_recv_view_ROSAT + + message(f"\nApplied {nr_single_pixel} single pixel updates\nStarting Smoothing and multi-pixel updates") + + ## Discard single pixel particles + nr_parts = np.sum(np.invert(single_pixel)) + nr_parts_tot = comm.allreduce(nr_parts) + + # view particles that update multiple pixels + part_pos_view = part_pos_recv[single_pixel==False,:].ndarray_view() # view array where multipixel updates are needed + + # remove single pixel particles from x-ray values and view only the multi-pixel particles + part_val_view_eROSITA_high = part_val_recv_eROSITA_high[single_pixel==False].ndarray_view() + part_val_view_eROSITA_low = part_val_recv_eROSITA_low[single_pixel==False].ndarray_view() + part_val_view_ROSAT = part_val_recv_ROSAT[single_pixel==False].ndarray_view() + + # sanity check + assert len(part_hsml_recv)==len(part_pos_view[:, 0]) + + # track how long it takes to apply multipixel updates + if comm_rank==0: + smoothing_time_0=datetime.datetime.now() + else: + smoothing_time_0=None + + #local_nr_pixels_updated=0 + local_nr_pixels_out_of_range=0 + for part_nr in range(nr_parts): + + # add smoothing filter here; If shell_nr==0 & density or value in pixel is negligable remove this pixel + + # skip particle if there is no value >0 in all X-ray bands + if (part_val_view_eROSITA_high[part_nr]+part_val_view_eROSITA_low[part_nr]+part_val_view_ROSAT[part_nr])>0 : + + pix_index, pix_val_eROSITA_high, pix_val_eROSITA_low, pix_val_ROSAT = explode_particle_in_all_bands( + nside, part_pos_view[part_nr,:], part_hsml_recv[part_nr], + part_val_view_eROSITA_high[part_nr], + part_val_view_eROSITA_low[part_nr], + part_val_view_ROSAT[part_nr] + ) + + # if for some there are no pixles to update continue to next particle. + if pix_index is None: + local_nr_pixels_out_of_range+=1 + rank_message(f" All searched pixels out of range, n(pixels searched)={-1*pix_val_ROSAT[0]}", comm_rank) + continue + + + local_pix_index = pix_index - local_offset + local = (local_pix_index >=0) & (local_pix_index < nr_local_pixels) + + # Don't need to use np.add.at here because pixel indexes are unique + map_view_eROSITA_high[local_pix_index[local]] += pix_val_eROSITA_high[local] + map_view_eROSITA_low[local_pix_index[local]] += pix_val_eROSITA_low[local] + map_view_ROSAT[local_pix_index[local]] += pix_val_ROSAT[local] + + del part_pos_view + del part_val_view_eROSITA_high + del part_val_view_eROSITA_low + del part_val_view_ROSAT + + # give total number of multi-pixel updates done, i.e. the number of particles that have been smoothed + message(f"\nAdded {nr_parts_tot} multi-pixel particles to the map") + + if comm_rank==0: + smoothing_time_1=datetime.datetime.now() + dt_str = str(smoothing_time_1-smoothing_time_0) + rank_message("time spent applying multipixel updates = " + dt_str, comm_rank) + else: + smoothing_time_1=None + + del smoothing_time_0 + del smoothing_time_1 + + comm.barrier() # gather at barrier to then give sanity check update + + #################################################################################################### + # Sanity chcek & Write maps to output .hdf5 file + + # Sanity check: + # Sum over the map should equal sum of values to be accumulated to the map. + particle_total_str="\nparticle totals:\n\teROSITA-high: {total0:.5e}\n\teROSITA-low: {total1:.5e}\n\tROSAT: {total2:.5e}\n" + message(particle_total_str.format(total0=val_total_global_eROSITA_high, total1=val_total_global_eROSITA_low, total2=val_total_global_ROSAT)) + + # sum over each map and give post map making total + map_sum_eROSITA_high = comm.allreduce(np.sum(map_data_eROSITA_high)) + map_sum_eROSITA_low = comm.allreduce(np.sum(map_data_eROSITA_low)) + map_sum_ROSAT = comm.allreduce(np.sum(map_data_ROSAT)) + + map_total_str="\nmap totals:\n\teROSITA-high: {total0:.5e}\n\teROSITA-low: {total1:.5e}\n\tROSAT: {total2:.5e}\n" + message(map_total_str.format(total0=map_sum_eROSITA_high, total1=map_sum_eROSITA_low, total2=map_sum_ROSAT)) + + ratio_total_str="\ntotals ratio (map/particles):\n\teROSITA-high: {total0:.5f}\n\teROSITA-low: {total1:.5f}\n\tROSAT: {total2:.5f}\n" + message(ratio_total_str.format(total0=map_sum_eROSITA_high/val_total_global_eROSITA_high, total1=map_sum_eROSITA_low/val_total_global_eROSITA_low, total2=map_sum_ROSAT/val_total_global_ROSAT)) + + # message ensures all ranks are gathered after making sanity checks so map writing can start. + + # can now delete all additional total values + del map_sum_eROSITA_high + del map_sum_eROSITA_low + del map_sum_ROSAT + + del val_total_global_eROSITA_high + del val_total_global_eROSITA_low + del val_total_global_ROSAT + + + if comm_rank==0: + # start tracking the time spent on writing the maps to .hdf5 file + writing_time_0=datetime.datetime.now() + + # write maps to .hdf5 file + # messages enforce comm.barrier() between writing outputs. + + # eROSITA-high + dataset_name=Xcalc.xray_map_names('erosita-high', observation_type) + message(f"\nwriting dataset ({dataset_name}) to file:\n {output_filename}") + with h5py.File(output_filename, "a", driver="mpio", comm=comm) as outfile: + phdf5.collective_write(outfile, dataset_name, map_data_eROSITA_high, comm) + message(f"\tcompleted writing {dataset_name}") + del map_data_eROSITA_high + del map_view_eROSITA_high + + # eROSITA-low + dataset_name=Xcalc.xray_map_names('erosita-low', observation_type) + message(f"writing dataset ({dataset_name}) to file:\n {output_filename}") + with h5py.File(output_filename, "a", driver="mpio", comm=comm) as outfile: + phdf5.collective_write(outfile, dataset_name, map_data_eROSITA_low, comm) + message(f"\tcompleted writing {dataset_name}") + del map_data_eROSITA_low + del map_view_eROSITA_low + + # ROSAT + dataset_name=Xcalc.xray_map_names('ROSAT', observation_type) + message(f"writing dataset ({dataset_name}) to file:\n {output_filename}") + with h5py.File(output_filename, "a", driver="mpio", comm=comm) as outfile: + phdf5.collective_write(outfile, dataset_name, map_data_ROSAT, comm) + message(f"\tcompleted writing datasets {dataset_name}") + del map_data_ROSAT + del map_view_ROSAT + + # define the map dataset attributes and add to each bands output maps. + if comm_rank==0: + map_attrs=get_map_attributes(nside, map_units) # compute atttrs + with h5py.File(output_filename, "a") as outfile: + for band in ['erosita-high', 'erosita-low','ROSAT']: + # load map dataset object + dataset_name = Xcalc.xray_map_names(band, observation_type) + Xray_map=outfile[dataset_name] + # write attrs for a given map + for k, v in map_attrs.items(): + Xray_map.attrs[k]=v + + # give update on time spent writing outputs + writing_time_1=datetime.datetime.now() + dt_writing_str=str(writing_time_1-writing_time_0) + rank_message("Writing maps end time: "+str(writing_time_1)+"\nTime Spent Writing: "+dt_writing_str, comm_rank) + del writing_time_0 + del writing_time_1 + + + return None + + diff --git a/lightcone_io/xray_utils.py b/lightcone_io/xray_utils.py new file mode 100644 index 0000000..12d955e --- /dev/null +++ b/lightcone_io/xray_utils.py @@ -0,0 +1,428 @@ +#!/bin/env python +import os +import h5py +import numpy as np +import unyt +import re +from astropy.cosmology import w0waCDM, z_at_value + +# create cosmology object with snapshot information that is needed to filter the X-ray particles +class Snapshot_Cosmology_For_Lightcone: + def __init__(self, snapshot_root_dir, search_snapshot_pattern=None): + """ + Make a cosmology object from the simuations snapshot cosmology metadata, create + functions used for common tasks with the lightcone and store information needed + for generating all-sky maps from the particle lightcones. + + Requires FLAMINGO snapshots to be accessible and contain cosmology metadata. + + Params: + snapshot_root_dir : string with path to FLAMINGO snapshot directory or + to snapshot .hdf5 file. + i.e. "./path/to/flamingo/runs/LBoxsizeNresolution/SimName/snapshots" + search_snapshot_pattern: only required if FLAMINGO snapshot file naming conventions + has been changed from + "/snapshots/flamingo_4DigitInt/flamingo_4DigitInt.hdf5" + + + """ + + # try access snapshot file + self.all_snapshot_filenames=self.get_snapshot_filename(snapshot_root_dir, search_snapshot_pattern) + + # to be consistent use the lowest redshift (highest number) snapshot. + snapshot_path = self.all_snapshot_filenames[-1] + + # collect cosmology info from hdf5 + self.raw_cosmo = self.cosmology_from_hdf5(snapshot_path) + + # collect internal units + self.internal_units = self.get_internal_units(snapshot_path) + + # make cosmology object + self.COSMO = w0waCDM( + H0=unyt.unyt_quantity(self.raw_cosmo['H0 [internal units]'][0], units='km / Mpc / s').to("1/s").to_astropy(), + Om0=self.raw_cosmo['Omega_m'][0], + Ode0=self.raw_cosmo['Omega_lambda'][0], + Ob0=self.raw_cosmo['Omega_b'][0], + w0=self.raw_cosmo['w_0'][0], + wa=self.raw_cosmo['w_a'][0], + Tcmb0=unyt.unyt_quantity(self.raw_cosmo['T_CMB_0 [K]'][0], units=unyt.K), + ) + + # needed for filtering of recently heated gas particles + self.AGN_delta_T_K = self.get_AGN_delta_T_K(snapshot_path) + + def get_snapshot_filename(self, rootdir, re_search_snapshot_pattern=None): + """ + Relies on FLAMINGO snapshot virtual files to be named as: + "/snapshots/flamingo_4DigitInt/flamingo_4DigitInt.hdf5" + + Params: + rootdir: path to directory with flamingo snapshots + "./path/to/flamingo/runs/LBoxsizeNresolution/SimName/snapshots" + re_search_snapshot_pattern: file naming pattern in regex terms. + + Returns: + path to lowest redshift snapshot file. + """ + try: + if re_search_snapshot_pattern is None: + regex = re.compile('(flamingo_00[0-9][0-9][.]hdf5$)|(flamingo_[0-9][0-9][0-9][0-9].*hdf5$)') + else: + regex = re.compile(re_search_snapshot_pattern) + filenames=[] + for root, dirs, files in os.walk(rootdir): + for file in files: + if regex.match(file): + filenames.append( os.path.join(root, file)) + filenames.sort() + if len(filenames)==0: + raise FileNotFoundError("cannot find snapshot files from root directory") + return filenames + except FileNotFoundError as fnfe: + raise Exception("cannot find snapshot files from root directory") from fnfe + except ValueError as ve: + raise Exception("cannot locate the ") from ve + + + def cosmology_from_hdf5(self, snapshot_path): + try: + with h5py.File(snapshot_path, "r") as sn: + raw_cosmo = {} + for k,v in sn['Cosmology'].attrs.items(): + raw_cosmo[f'{k}']=v + return raw_cosmo + except FileNotFoundError as fnfe: + raise Exception("cannot access snapshot filename") from fnfe + except KeyError as ke: + raise Exception("Cosmology not found in snapshot .hdf5 file") from ke + except AttributeError as ae: + raise Exception("Cosmology attributes not found in snapshot .hdf5 file") from ae + + + def get_internal_units(self, snapshot_path): + try: + internal_units={} + with h5py.File(snapshot_path, "r") as sn: + for k, v in sn['InternalCodeUnits'].attrs.items(): + internal_units[f'{k}'] = v + return internal_units + except FileNotFoundError as fnfe: + raise Exception("cannot access snapshot filename") from fnfe + except KeyError as ke: + raise Exception("InternalCodeUnits not found in snapshot .hdf5 file") from ke + except AttributeError as ae: + raise Exception("InternalCodeUnits attributes not found in snapshot .hdf5 file") from ae + + + def z2Myr(self, z): + """ + Return a unyt_array [Myrs] with age of the universe at the input redshift(s). + """ + t_age=self.COSMO.age(z).to("Myr") + return unyt.unyt_array.from_astropy(t_age) + + def redshift_with_time_offset(self, z, dt): + """ + Params: + z: redshift of observer + dt: length of time [Myr] + + Returns: + The redshift equal to the input redshift - length of time (dt) + """ + time_from_bb = self.COSMO.age(z).to('Myr') + time_with_dt_bb = (time_from_bb.to_value("Myr") - dt)*unyt.Myr + z_min = z_at_value(self.COSMO.age, time_with_dt_bb.to_astropy()) + + return z_min.value + + def z2r(self, z): + """ + Returns the comoving radius for a given redshift in cMpc + """ + return unyt.unyt_array.from_astropy(self.COSMO.comoving_distance(z)).to("Mpc") + + def shell_comoving_raddii(self, redshift_filename): + """ + Params: + redshift_filename: path to .txt file with shell redshift ranges + Returns: + N x 2 array of comoving radii, per shell returns [inner, outer] + """ + try: + self.shell_redshifts = np.loadtxt(redshift_filename, delimiter=",") + zmin = self.shell_redshifts[:,0] + zmax = self.shell_redshifts[:,1] + comoving_radii = unyt.unyt_array(np.zeros((len(zmin), 2)), units="Mpc") + comoving_radii[:, 0]+=self.z2r(zmin) + comoving_radii[:, 1]+=self.z2r(zmax) + except: + raise ValueError("cannot access shell redshifts .txt file") + + return comoving_radii + + + def get_AGN_delta_T_K(self, snapshot_path): + """ + Return the Delta_T_AGN param required to identify recently heated particles + """ + # 'AGN_delta_T_K' Change in temperature to apply to the gas particle in an AGN feedback event [K] + with h5py.File(snapshot_path, "r") as sn: + if 'EAGLEAGN:AGN_delta_T_K' in sn['Parameters'].attrs: # check if thermal mode AGN_delta_T_K exists + AGN_delta_T_K = np.float64(sn['Parameters'].attrs['EAGLEAGN:AGN_delta_T_K'])*unyt.K + elif 'SPINJETAGN:AGN_delta_T_K' in sn['Parameters'].attrs: # check if Jet mode AGN_delta_T_K exists + AGN_delta_T_K = np.float64(sn['Parameters'].attrs['SPINJETAGN:AGN_delta_T_K'])*unyt.K + else: # cannot find thermal or Jet mode AGN_delta_T_k param + raise NameError("\n!!!\nCannot find AGN_delta_T_K value:\n\t'EAGLEAGN:AGN_delta_T_K' or 'SPINJETAGN:AGN_delta_T_K' not in snapshot 'Parameters'\n!!!") + AGN_delta_T_K=-1*unyt.K # set as an error value + + return AGN_delta_T_K + + def lightcone_volume(self, comoving_inner_radius, comoving_outer_radius, use_redshift=False, phi=None): + """ + Compute the comoving volume of the lightcone given some inner and outer radius + + Params: + comoving_inner_radius, comoving_outer_radius: the radius of the lightcone in Mpc + use_redshift [boolean]: if False, then radii are in comoving distances. if True, then + input radii are redshift values. + phi: 1/2 angle of the cones apature [radian] + + """ + if use_redshift: + # redshift -> comoving distance + comoving_inner_radius = unyt.unyt_array.from_astropy(self.COSMO.comoving_distance(comoving_inner_radius)) + comoving_outer_radius = unyt.unyt_array.from_astropy(self.COSMO.comoving_distance(comoving_outer_radius)) + + if phi is None: + # spherical shells + shell_volume = (4/3)*np.pi * (comoving_outer_radius**3 - comoving_inner_radius**3) + else: + # cone shells + r_outer=comoving_outer_radius*(1-np.cos(phi)) + r_inner=comoving_inner_radius*(1-np.cos(phi)) + shell_volume = (2/3)*np.pi * ((comoving_outer_radius**2 * r_outer) - (comoving_outer_radius**2 * r_outer)) + + return shell_volume + + +class Xray_Filter: + def __init__(self, particle_data, numb_particles, remove_SF=True, log_T_min_K=5, remove_rhp=True, cosmo=None): + + + # create a mask that tracks all particles which can be used for X-ray values + self.KEEP_FOR_XRAY=np.ones(numb_particles, dtype=bool) + + if remove_SF: + if "StarFormationRates" in particle_data: + self.filter_by_StarFormationRates(particle_data["StarFormationRates"].value) + else: + raise ValueError("particle StarFormationRates not found ....") + + if "Temperatures" in particle_data: + self.filter_by_Temperatures(particle_data["Temperatures"].to_value("K"), log_T_min_K) + else: + raise ValueError("particle Temperatures not found ....") + + if remove_rhp and cosmo is not None: + # check all the needed properties exist in particle data + for prop_name in ["ExpansionFactors", "LastAGNFeedbackScaleFactors", "Temperatures", "Densities"]: + if prop_name not in particle_data: + ValueError("particle {required_property} not found ....".format(required_property=prop_name)) + + if cosmo is not None: + + + # identify all recently heated particles and keep mask for future use + is_recently_heated = self.identify_recently_heated( + cosmo, + scalefactors=particle_data["ExpansionFactors"][self.KEEP_FOR_XRAY].value, + last_agn_feedback_scalefactors=particle_data["LastAGNFeedbackScaleFactors"][self.KEEP_FOR_XRAY].value, + temperatures=particle_data["Temperatures"][self.KEEP_FOR_XRAY], # keep units + densities=particle_data["Densities"][self.KEEP_FOR_XRAY].to("g * cm**-3"), # convert the units + RHP_filter_max_time_Myr=70, + RHP_filter_log_density_cm3=-2.25) + + self.numb_recently_heated = np.sum(is_recently_heated) + + # update mask + self.KEEP_FOR_XRAY[self.KEEP_FOR_XRAY][is_recently_heated]=0 + + else: + print("Need to provide cosmology information to filter recently heated particles") + + + + + def filter_by_StarFormationRates(self, sfr): + """ + remove all star forming gas particles & update mask + + Params + sfr: gas particle star formation rates + + """ + self.KEEP_FOR_XRAY[sfr>0.]=0 + + def filter_by_Temperatures(self, temperatures, log_T_min_K): + """ + keep only `hot' gas particles & update mask + + Params + temperatures: temperatures of gas particles + log_T_min_K: minimum temperature to be considered a hot gas particle + + """ + self.KEEP_FOR_XRAY[temperatures<10**log_T_min_K]=0 + + + def time_from_last_AGN_feedback(self, + scalefactors, last_agn_feedback_scalefactors, + f_z2t, handle_negatives=0.): + """ + Compute the difference in time from current scalefactor and the scalefactor + when last heated by AGN feedback. + Due to compression on LastAGNFeedbackScaleFactors values, sometimes + LastAGNFeedbackScaleFactors > scalefactor, therefore we must allow negative values. + + Params + scalefactors: scalefactor of gas particles + last_agn_feedback_scalefactor: scalefactor that particle was last + directly heated by AGN feedback and -1 + if it has never been heated by AGN + feedback + f_z2t: function to give age of universe for a given redshift + handle_negatives: time difference given to where + last_agn_feedback_scalefactor > scalefactors + """ + + # make output array of time differences + dt = np.zeros(len(scalefactors), dtype=float) -1. + dt[last_agn_feedback_scalefactors<0]=999 # cannot ever be recently heated + dt[last_agn_feedback_scalefactors >= scalefactors] = handle_negatives # assume its due to compression + + mask = (last_agn_feedback_scalefactors < scalefactors) & (last_agn_feedback_scalefactors>0) + + if np.sum(mask)>0: + + z_part = 1./scalefactors[mask] -1. + z_last_agn_feedback = 1./last_agn_feedback_scalefactors[mask] -1. + + particle_time_from_big_bang_Myr = f_z2t(z_part) + last_heating_time_from_big_bang_Myr = f_z2t(z_last_agn_feedback) + + dt[mask]=particle_time_from_big_bang_Myr.to_value("Myr")-last_heating_time_from_big_bang_Myr.to_value("Myr") + + return dt + + + def identify_recently_heated( + self, + cosmo, + scalefactors, + last_agn_feedback_scalefactors, + temperatures, + densities, + RHP_filter_max_time_Myr, + RHP_filter_log_density_cm3=-2.25): + """ + Identify recently heated particles. + Due to the BFloat16 lossy compression filter applied to + LastAGNFeedbackScaleFactors we cannot reliably determine if a particle + was last heated by AGN feedback within the last 15 Myr. + Note the minimum density criterion is included to compensate for the uncertanty + about LastAGNFeedbackScaleFactors values due to the lossy compression filter used for the + particle lightcones. + To reproduce snapshot selection of recently heated particles: + RHP_filter_max_time_Myr=15 + RHP_filter_log_density=None + + + Params + cosmo: cosmology object, i.e. Snapshot_Cosmology_For_Lightcone() + scalefactors: scalefactor of gas particles + last_agn_feedback_scalefactor: scalefactor that particle was last + directly heated by AGN feedback and -1 + if it has never been heated by AGN + feedback + temperatures (unyt_array): temperatures of gas particles [K] + densities (unyt_array): densities of gas particles [g/cm^3] + RHP_filter_max_time_Myr: maximum time in Myrs that has elapsed from the + LastAGNFeedbackScaleFactors to be considered a recently heated + particle. If == None then adopt 70 Myr as compromise for uncertanty. + RHP_filter_log_density: log10(minimum gas density [1/cm^3]), the minimum density + to be considered a recently heated particle. If == None, then + do not use density to select recently heated particles. + -2.25 selected as default for McDonald+26 based on + HYDRO_FIDUCIAL L1_m9 snapshots & ROSAT photon intrinsic maps + Returns + Boolean mask, True if RHP, otherwise False + """ + + + # find particles that have never been heated by AGN feedback, LastAGNFeedbackScaleFactors==-1 + never_heated = last_agn_feedback_scalefactors<0 + + # if no particles have been heated by AGN then no need to continue filtering + npart = len(scalefactors) + if np.sum(never_heated) == npart: + return np.zeros(npart, dtype=bool) + + # get temperature range for AGN heating + AGN_delta_T_K = cosmo.AGN_delta_T_K + xray_maps_recent_AGN_injection_delta_logT_min=-1 #from IC file + xray_maps_recent_AGN_injection_delta_logT_max=0.3 #from IC file + lower_T_K = AGN_delta_T_K*10**xray_maps_recent_AGN_injection_delta_logT_min + upper_T_K = AGN_delta_T_K*10**xray_maps_recent_AGN_injection_delta_logT_max + + # check if temperatures are within the range expected from recent AGN heating + is_in_temp_bounds = (temperatures>lower_T_K) & (temperatures= min_gas_density_cm3 + + # combine masks and clean up + combined_mask = is_in_temp_bounds & is_above_min_density & np.invert(never_heated) + + # if there are no more particles remaining then no need to check time since heating + if np.sum(combined_mask) == 0: + return np.zeros(npart, dtype=bool) + + del never_heated + del is_in_temp_bounds + del is_above_min_density + del densities + del temperatures + + # for remaining particles check the time since they where last heated by AGN feedback + time_from_last_heating=np.zeros(npart, dtype=float) + time_from_last_heating[combined_mask==False] = 999. # cannot possibly be recently heated. + time_from_last_heating[combined_mask] += self.time_from_last_AGN_feedback( + scalefactors[combined_mask], + last_agn_feedback_scalefactors[combined_mask], + f_z2t=cosmo.z2Myr, + handle_negatives=0.) + + # combine masks again to flag recently heated particles + is_recently_heated = combined_mask & (time_from_last_heating<=RHP_filter_max_time_Myr) + + return is_recently_heated + + + + + + diff --git a/pyproject.toml b/pyproject.toml index 0d65b14..638535c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -19,6 +19,7 @@ dependencies = [ "tqdm", "pytest", "hdfstream>=0.0.22", + "numba" ] readme = "README.md" classifiers = [ @@ -40,4 +41,4 @@ packages = ["lightcone_io",] [tool.setuptools_scm] version_scheme = "guess-next-dev" -local_scheme = "node-and-date" \ No newline at end of file +local_scheme = "node-and-date" diff --git a/requirements-mpi.txt b/requirements-mpi.txt index 047997e..e7a892a 100644 --- a/requirements-mpi.txt +++ b/requirements-mpi.txt @@ -6,3 +6,4 @@ healpy virgodc scipy tqdm +numba diff --git a/requirements.txt b/requirements.txt index fd570db..8a8d17b 100644 --- a/requirements.txt +++ b/requirements.txt @@ -4,3 +4,5 @@ unyt healpy scipy tqdm +numba +astropy