Skip to content

Commit

Permalink
release to adapt to pynbody v2.0 and enable parallel computation.
Browse files Browse the repository at this point in the history
  • Loading branch information
Zhiwei Shao committed Jan 21, 2025
1 parent 9476bfa commit 5a6383f
Show file tree
Hide file tree
Showing 6 changed files with 1,099 additions and 124 deletions.
57 changes: 30 additions & 27 deletions XIGrM/X_properties.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -27,33 +27,34 @@
element_masses = ppat.get_atomic_masses(element_numbers)
Z_ag89 = (element_abundances*element_masses)[2:].sum()/(element_abundances*element_masses).sum()

def cal_tweight(halo, weight_type='Lx'):
def cal_tweight(halogas, weight_type='Lx'):
"""
Calculate luminosity weighted or mass weighted temperatures.
Parameters
-----------
halo : pynbody.snapshot.SubSnap
SubSnap of the halo.
halogas : pynbody.snapshot.SubSnap
Gas subsnap to calculate.
weight_type : str
Type of weight to take. Related to the available properties
of the gas. Now available: luminousity weighted (starts
with 'l') and mass weighted (starts with 'm')
"""
try:
if weight_type[0].lower() == 'l':
weight_units = 'erg s**-1'
T, weight_sum = np.average(halo.gas['temp'].in_units('K').view(np.ndarray), \
weights=halo.gas[weight_type].in_units(weight_units).view(np.ndarray), returned=True)
elif weight_type[0].lower() == 'm':
weight_units = 'Msol'
T, weight_sum= np.average(halo.gas['temp'].in_units('K').view(np.ndarray), \
weights=halo.gas[weight_type].in_units(weight_units).view(np.ndarray), returned=True)
else:
raise Exception("weight_type Error!")
except ZeroDivisionError:
T = 0
weight_sum = 0
with halogas.immediate_mode:
try:
if weight_type[0].lower() == 'l':
weight_units = 'erg s**-1'
T, weight_sum = np.average(halogas['temp'].in_units('K').view(np.ndarray), \
weights=halogas[weight_type].in_units(weight_units).view(np.ndarray), returned=True)
elif weight_type[0].lower() == 'm':
weight_units = 'Msol'
T, weight_sum= np.average(halogas['temp'].in_units('K').view(np.ndarray), \
weights=halogas[weight_type].in_units(weight_units).view(np.ndarray), returned=True)
else:
raise Exception("weight_type Error!")
except ZeroDivisionError:
T = 0
weight_sum = 0

T = pnb.array.SimArray(T*k_B, units='erg')
T.convert_units('keV')
Expand All @@ -76,17 +77,19 @@ def cal_tspec(hdgas, cal_f, datatype):

cal_f = cal_f.encode('utf-8')
k_B_with_units = pnb.units.Unit('cm**2 g s**-2 K**-1') * k_B

with hdgas.immediate_mode:
#hdgas = gas#[pnb.filt.HighPass('temp', '5e5 K')&pnb.filt.LowPass('nh', '0.13 cm**-3')] # short for hot diffuse
tTemp_in_kev = (hdgas['temp'] * k_B_with_units).in_units('keV')
if datatype[:5] == 'gizmo':
tZ_in_solar = hdgas['metals'][:, 0]/Z_ag89
elif datatype[:5] == 'tipsy':
tZ_in_solar = hdgas['metals']/Z_ag89
else:
raise Exception("Simulation Datatype Not Accepted!")
temission_meassure = (hdgas['mass'].in_units('Msol') * \
hdgas['rho'].in_units('Msol kpc**-3')).in_units('g**2 cm**-3')
tTemp_in_kev = (hdgas['temp'] * k_B_with_units).in_units('keV')
if 'X_H' in hdgas.keys() and 'X_He' in hdgas.keys():
tZ_in_solar = (1 - hdgas['X_H']- hdgas['X_He'])/Z_ag89
elif datatype[:5] == 'gizmo':
tZ_in_solar = hdgas['metals'][:, 0]/Z_ag89
elif datatype[:5] == 'tipsy':
tZ_in_solar = hdgas['metals']/Z_ag89
else:
raise Exception("Simulation Datatype Not Accepted!")
temission_meassure = (hdgas['mass'].in_units('Msol') * \
hdgas['rho'].in_units('Msol kpc**-3')).in_units('g**2 cm**-3')

Temp_in_kev = array.array('f', tTemp_in_kev)
Z_in_solar = array.array('f', tZ_in_solar)
Expand Down
56 changes: 7 additions & 49 deletions XIGrM/calculate_R.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,7 @@
import pynbody as pnb

def get_radius(halo, overdensities = np.array([]), rho_crit=None, \
precision=1e-2, rmax=None, cen=np.array([]), prop=None):
precision=1e-2, rmax=None, cen=np.array([]), prop=None, ncpu=1):
"""
Calculate different radii of a given halo with a decreasing sphere method.
Expand Down Expand Up @@ -65,12 +65,13 @@ def get_radius(halo, overdensities = np.array([]), rho_crit=None, \
# All positions in prop are in comoving coordinates.
boxsize = prop['boxsize'].in_units('kpc')
if len(cen) == 0:
center = pnb.analysis.halo.center(halo, mode='pot', retcen=True, vel=False)
center = pnb.analysis.halo.center(halo, mode='pot', return_cen=True, with_velocity=False).in_units('kpc')
else:
center = cen.in_units('kpc')
#if prop['mass'] > 1e12: # For massive halos which spend most of time loading particle information, use numpy.
halopos = halo['pos'].in_units('kpc').view(np.ndarray) - center.view(np.ndarray)
halomass = halo['mass'].in_units('Msol').view(np.ndarray)
with halo.immediate_mode:
halopos = halo['pos'].in_units('kpc').view(np.ndarray) - center.view(np.ndarray)
halomass = halo['mass'].in_units('Msol').view(np.ndarray)

for i in range(3): # Correct the position of patricles crossing the box periodical boundary.
index1, = np.nonzero(halopos[:, i] < -boxsize/2)
Expand Down Expand Up @@ -104,64 +105,21 @@ def get_radius(halo, overdensities = np.array([]), rho_crit=None, \
while True:
temp_radius = radius
radius = ((3 / (4 * np.pi)) * mass / density)**(1/3)
# temphalo = temphalo[pnb.filt.Sphere(radius)]
# mass = temphalo['mass'].sum()
particles_within_r, = np.nonzero(halor <= radius)
halor = halor[particles_within_r]
halomass = halomass[particles_within_r]
mass = halomass.sum()
radial_difference = np.abs(temp_radius - radius) / radius
if mass == 0:
radii[overdensities[i]] = 0
masses[overdensities[i]] = 0
radii[overdensities[i]] = pnb.array.SimArray(.0 ,units='kpc')
masses[overdensities[i]] = pnb.array.SimArray(.0 ,units='Msol')
break
if radial_difference <= precision:
radii[overdensities[i]] = pnb.array.SimArray(radius)
masses[overdensities[i]] = pnb.array.SimArray(mass)
radii[overdensities[i]].units = 'kpc'
masses[overdensities[i]].units = 'Msol'
break
# else: # For small halos, pynbody built-in function is faster.
# tx = pnb.transformation.inverse_translate(halo, center)
# with tx:
# x = halo['x']
# radius = x.max() - x.min()
# mass = halo['mass'].sum()
# lunits = radius.units
# munits = mass.units
# if rho_crit == 0:
# rho_crit = pnb.analysis.cosmology.rho_crit(halo, z=prop["z"], unit=munits/lunits**3)
# if rmax != 0:
# radius = rmax.in_units(lunits)
# mass = halo[pnb.filt.Sphere(radius)]['mass'].sum()
# overdensities = np.array(overdensities)
# if len(overdensities) > 1:
# overdensities.sort() # From low density to high density
# densities = pnb.array.SimArray(overdensities * rho_crit)
# densities.units = munits/lunits**3
# temphalo = halo
# masses = {}
# radii = {}
# for i in range(len(overdensities)):
# if i > 0: # Continue calculating based on the last result.
# radius= radii[overdensities[i-1]]
# mass = masses[overdensities[i-1]]
# density = densities[i]
# while True:
# temp_radius = radius
# radius = ((3 / (4 * np.pi)) * mass / density)**(1/3)
# temphalo = temphalo[pnb.filt.Sphere(radius)]
# mass = temphalo['mass'].sum()
# radial_difference = np.abs(temp_radius - radius) / temp_radius
# if mass == 0:
# radii[overdensities[i]] = 0
# masses[overdensities[i]] = 0
# break
# if radial_difference <= precision:
# radius.units=lunits
# radii[overdensities[i]] = radius
# masses[overdensities[i]] = mass
# break
return masses, radii
#densities.units = halo.dm[0]['mass'].units/halo.dm['pos'].units**3

Expand Down
106 changes: 101 additions & 5 deletions XIGrM/gas_properties.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,8 @@
import pynbody as pnb
import astropy.constants as astroc
import numpy as np

import h5py
import os
from . import prepare_pyatomdb as ppat

# some constants
Expand Down Expand Up @@ -123,14 +124,109 @@ def calcu_luminosity(gas, filename, mode='total', elements=default_elements, ban
for i in atomicNumbers:
a_pos, = np.where(emission_atomic_numbers == i)
if len(a_pos) == 0:
raise Exception('Privided emissivity file is incomplete! Lacking Z={}'.format(i))
raise Exception('Provided emissivity file is incomplete! Lacking Z={}'.format(i))
elements_idx = np.append(elements_idx, a_pos)

specific_emissivity = emissivity[elements_idx, :]
T_index = ppat.get_index(gas['temp'].in_units('K').view(np.ndarray))
result = (specific_emissivity[:,T_index].T*gas['abundance'].view(np.ndarray)).sum(axis=1)
result *= gas['ElectronAbundance'] * (gas['nh'].in_units('cm**-3').view(np.ndarray))**2*\
gas['volume'].in_units('cm**3').view(np.ndarray)
result *= gas['ElectronAbundance'] * ((gas['nh'])**2*gas['volume']).in_units('cm**-3').view(np.ndarray)
result = pnb.array.SimArray(result)
result.units='erg s**-1'
return result
return result

def load_gas_properties(gasfile, s, hotdiffusefilter, props_to_load=['temp', 'nh', 'ne', 'volume'],
Lx_to_load=['Lx', 'Lxb', 'Lx_cont', 'Lxb_cont'],
elements_to_load=['H', 'He', 'O', 'Si', 'Fe']):
props_units = {'temp': 'K', 'nh': 'cm**-3', 'ne': 'cm**-3', 'volume': 'kpc**3'}
with h5py.File(gasfile, "r") as f:
# print(f.keys())
if 'iord' not in s.keys():
_ = s['iord']
assert np.abs(f['iord'][()] - s.gas['iord']).max() < 1, "Orders of particles doesn't agree"

for propkey in props_to_load:
s.gas[propkey] = pnb.array.SimArray(f[propkey][()], units=props_units[propkey])

igrm_filter = hotdiffusefilter
for Lxkey in Lx_to_load:
s.gas[Lxkey] = pnb.array.SimArray(f['luminosity'][Lxkey][()], units='erg s**-1')
s.gas[~igrm_filter][Lxkey] = 0.

for metal in elements_to_load:
s.gas[f'X_{metal}'] = f['mass_fraction'][metal][()]

def prepare_gas_properties(basename, s, Lxfiles,
Lx_narrowband = [.5, 2.0], Lx_broadband = [.1, 2.4],
elements=default_elements, elements_idx=None, fixElectronAbundance=False):
if elements_idx is None:
elements_idx = {'H': 0, 'He': 1, 'C': 2, 'N': 3, 'O': 4, 'Ne': 5, 'Mg': 6, 'Si': 7, 'S': 8, 'Ca': 9, 'Fe': 10}
gasfile = f"{basename}_gasproperties_Lx={Lx_narrowband[0]:.1f}-{Lx_narrowband[1]:.1f}_" + \
f"Lxb={Lx_broadband[0]:.1f}-{Lx_broadband[1]:.1f}_" + \
f"fixElectronAbundance={fixElectronAbundance}.hdf5"
if not os.path.isfile(gasfile):
s.physical_units()
s.gas['X_H'] = 1-s.gas['metals'][:,0]-s.gas['metals'][:,1] # hydrogen mass fraction
s.gas['nh'] = nh(s) # Hydrogen number density
s.gas['temp'] = temp(s) # From internal energy to temperature
if fixElectronAbundance:
s.gas['ElectronAbundance'] = 1.14
s.gas['ElectronAbundance'].units = '1'
s.gas['ne'] = s.gas['ElectronAbundance'] * s.gas['nh'].in_units('cm**-3') # eletron number density
s.gas['volume'] = s.gas['mass']/s.gas['rho'] # volume of gas particles

s.gas['mass_fraction'] = np.zeros(s.gas['metals'].shape)
s.gas['mass_fraction'][:, 0] = s.gas['X_H']
s.gas['mass_fraction'][:, 1:] = s.gas['metals'][:, 1:]
s.gas['abundance'] = abundance_to_solar(s.gas['mass_fraction'], elements=elements)
_ = s['iord']

narrow_band_file = Lxfiles['narrow']
broad_band_file = Lxfiles['broad']

Emission_type = ['Lx', 'Lxb', 'Lx_cont', 'Lxb_cont']
for i in Emission_type:
s.gas[i] = 0
s.gas[i].units = 'erg s**-1'

s.gas['Lx'] = calcu_luminosity(gas=s.gas, filename=narrow_band_file, \
mode='total', band=Lx_narrowband)
s.gas['Lxb'] = calcu_luminosity(gas=s.gas, \
filename=broad_band_file, mode='total', band=Lx_broadband)
s.gas['Lx_cont'] = calcu_luminosity(gas=s.gas, filename=narrow_band_file, \
mode='cont', band=Lx_narrowband)
s.gas['Lxb_cont'] = calcu_luminosity(gas=s.gas, \
filename=broad_band_file, mode='cont', band=Lx_broadband)

attrs = {'Lx': f'{Lx_narrowband[0]:.1f}-{Lx_narrowband[1]:.1f} keV',
'Lxb': f'{Lx_broadband[0]:.1f}-{Lx_broadband[1]:.1f} keV',
'Lx_cont': f'{Lx_narrowband[0]:.1f}-{Lx_narrowband[1]:.1f} keV, w/o line emission',
'Lxb_cont': f'{Lx_broadband[0]:.1f}-{Lx_broadband[1]:.1f} keV, w/o line emission',
'iord': 'Gas particle ids',
'nh': 'hydrogen number density in cm^-3',
'ne': 'electron number density in cm^-3',
'temp': 'temperature in K',
'volume': 'volume of the gas particle derived from mass/rho; in kpc**3'
}
with h5py.File(gasfile, "w") as f:
for key in ['iord', 'nh', 'ne', 'temp', 'volume']:
dataset = f.create_dataset(key, data=s.gas[key])
dataset.attrs['Description'] = attrs[key]

grp = f.create_group("luminosity")
for key in ['Lx', 'Lxb', 'Lx_cont', 'Lxb_cont']:
dataset = grp.create_dataset(key, data=s.gas[key])
dataset.attrs['Description'] = attrs[key]

grp = f.create_group("mass_fraction")
grp.attrs['Description'] = "Mass fractions of different elements (summed to 1)."
for i in range(len(elements)):
dataset = grp.create_dataset(elements[i], data=s.gas['mass_fraction'][:, elements_idx[elements[i]]])

grp = f.create_group("abundance_to_solar")
grp.attrs['Description'] = "Abundances of different elements relative to AG89 solar."
for i in range(len(elements)):
dataset = grp.create_dataset(elements[i], data=s.gas['abundance'][:, elements_idx[elements[i]]])
else:
print("File exists!")
return gasfile
Loading

0 comments on commit 5a6383f

Please sign in to comment.