import configparser
import subprocess
import numpy as np
import h5py as h5
from .column import Column
[docs]
def get_touts(tf: float) -> list[float]:
"""Create the default list of output times
Args:
tf (float): final time to create touts until
Returns:
list[float]: list of output times
"""
touts = []
base = 1.0
exp = 3
time = 0
while time < tf:
time = base*10**exp
touts.append(time)
if exp < 5:
base += 1.0
else:
base += 0.5
if base >= 10:
base = base/10
exp += 1
return touts
[docs]
def parse_list(s: str) -> list:
"""Read in a string and convert to a python list. Used to read in
touts from input file.
Args:
s (str): list string
Returns:
list: parsed list
"""
assert s.startswith('[') and s.endswith(']'), "touts format must be as a list with square brackets"
ts = list(map(float,s[1:-1].split(',')))
return ts
[docs]
def get_defaults() -> tuple[dict, dict, dict]:
"""Returns the default parameters for CANDY run as a tuple of
dictionaries.
Returns:
tuple[dict, dict, dict]: tuple of dictionaries containing default model params, physical params, and abundances
"""
# default model parameters
model_inputs = {
'chmfile' : 'astrochem/networks/umist12_x.chm',
'pebfile' : 'pebcomp.out',
'ti' : 0,
'tf' : 1e6,
'chemtime' : 500.0,
'difftime' : 100.0,
'outputdir' : 'r00',
'outfile' : 'output_abuns.npz',
'ncells' : 50
}
model_inputs['touts'] = get_touts(model_inputs['tf'])
# default physical paramters
phys_inputs = {
'r' : 30,
'r_units' : 'au',
'alpha' : 1.e-3,
'chi' : 50.,
'cosmic' : 1.3e-17,
'grain_size' : 0.1,
'dg0' : 0.01,
'opacity' : 1.0e5,
'growth_timescale_factor' : 1.,
'growth_height' : 1.,
'growth_delay_time' : 0.,
}
# default initial chemical abundances
abundances = {
'H2' : 0.5,
'He' : 9.75e-2,
'NH3': 1.45e-6,
'H2O': 1.18e-4,
'CO' : 6.00e-5,
'N2' : 2.00e-5,
'CH4': 2.00e-6,
'CH3OH' : 1.00e-6,
'H2S' : 1.91e-8,
'CO2' : 5.00e-5,
'HCN' : 3.50e-7,
'grain' : 2.2e-12
}
return model_inputs,phys_inputs,abundances
[docs]
def read_infile(fin: str) -> tuple[dict, dict, dict]:
"""Reads the input file to create input parameters dictionaries for CANDY
Args:
fin (str): name of input file
Returns:
tuple[dict, dict, dict]: dictionaries of model params, physical params, and abundances
"""
config = configparser.ConfigParser()
# the next line retains case sesnitivity for input parameters
config.optionxform = str
config.read(fin)
model,phys,abuns = get_defaults()
for key in config['model']:
if key in ['chmfile','pebfile','outputdir','outfile']:
model[key] = config['model'][key]
elif key in ['touts']:
model[key] = parse_list(config['model'][key])
elif key in ['ncells']:
model[key] = int(config['model'][key])
else:
model[key] = float(config['model'][key])
if 'touts' not in config['model'].keys():
model['touts'] = get_touts(model['tf'])
if model['tf'] not in model['touts']:
print('Warning: tf is not in touts, including tf in list of output times')
model['touts'].append(model['tf'])
dterror = 'Diffusion time must be less than or equal to chemtime.\n' \
+f'chemtime = {model["chemtime"]}\n' \
+f'difftime = {model["difftime"]}'
assert model['chemtime'] >= model['difftime'], dterror
model['touts'].sort()
print('touts = ',model['touts'])
for key in config['phys']:
if key != 'r_units':
phys[key] = float(config['phys'][key])
else:
phys[key] = config['phys'][key]
if 'abundances' in config:
abuns = {key:float(val) for (key,val) in config['abundances'].items()}
return model,phys,abuns
[docs]
def save_outputs(col: Column, nout: int, outputdirr: str) -> None:
"""Copy and paste astrochem outputs to store as backup or to continue a run
Args:
col (Column): CANDY column of cells
nout (int): the output index
outputdirr (str): directory where savefiles should be stored
"""
for j in range(col.ncells):
dirr = f'{outputdirr}/z{j:0>2}'
subprocess.run(['cp','astrochem_output.h5',
f'astrochem_output_t{nout}.h5'], cwd=dirr)
subprocess.run(['cp','source.mdl',f'source_t{nout}.mdl'],cwd=dirr)
subprocess.run(['cp','input.ini',f'input_t{nout}.ini'],cwd=dirr)
[docs]
def new_save_outputs(col: Column, nout: int, time: float,
outputdirr: str, outputfile: str,
delete_subdir: bool = False) -> None:
"""Save the current abundances of ice and gas into npz file. This
saves from the memory and accounts for ices lost due to grain growth
Args:
col (Column): CANDY column of cells
nout (int): output index
time (float): the current time
outputdirr (str): output directory
outputfile (str): output file name
delete_subdir (bool, optional): Delete subdirectories after values are saved. Defaults to False.
"""
save_outputs(col,nout,outputdirr)
nzs = col.ncells
if nout==0:
species = get_specs_array(outputdirr)
nspec = len(species)
times = read_touts_from_file(f'{outputdirr}/times.out')
nts = len(times)
gas_dens = np.zeros(nzs)
extinctions = np.zeros((nts,nzs))
abundances = np.zeros((nts,nzs,nspec))
heights = np.zeros(nzs)
for j in range(nzs):
cell = col.cells[j]
gas_dens[j] = cell.nh
extinctions[0,j] = cell.av
heights[j] = cell.z
for k in range(nspec):
abundances[0,j,k] = cell.abundances[species[k]]
else:
oldoutput = np.load(f'{outputdirr}/{outputfile}')
# get things that don't change
times = oldoutput['times']
species = oldoutput['species']
heights = oldoutput['heights']
gas_dens = oldoutput['density']
# update the things that do
extinctions = oldoutput['extinctions']
abundances = oldoutput['abundances']
nspec = len(species)
for j in range(nzs):
cell = col.cells[j]
extinctions[nout,j] = cell.av
for k in range(nspec):
abundances[nout,j,k] = cell.abundances[species[k]]
np.savez(f'{outputdirr}/{outputfile}',times=times,species=species,
abundances=abundances,heights=heights,density=gas_dens,
extinctions=extinctions)
if delete_subdir:
print('removing sub directories')
print(f'rm -rf {outputdirr}/z*')
subprocess.run(f'rm -rf {outputdirr}/z*', shell=True)
[docs]
def get_nh_and_av(source_file: str) -> tuple[float, float]:
"""Read the number denisty of hydrogen and visual extinction from
the astrochem source file
Args:
source_file (str): name of astrochem source file
Returns:
tuple[float, float]: number density of hydrogen and visual extinction
"""
with open(source_file,'r') as f:
f.readline() # skip header
params = f.readline().split()
nh = float(params[2])
av = float(params[1])
return nh,av
[docs]
def get_specs_array(outputdirr: str) -> np.ndarray:
"""Returns an array of species used from astrochem output. Reads in
from outputdirr/z00/astrochem_output.h5
Args:
outputdirr (str): name of output directory
Returns:
np.ndarray: array of species names
"""
_,abundict = get_abundict(outputdirr+f'/z00/astrochem_output.h5')
specs = np.empty(len(list(abundict.keys())),dtype='U25')
for k,spec in enumerate(list(abundict.keys())):
specs[k] = spec
return specs
[docs]
def read_touts_from_file(fname: str) -> np.ndarray:
"""Helper to read in output times saved in times.out
Args:
fname (str): name of file
Returns:
np.ndarray: array of output times
"""
touts = []
with open (fname,'r') as f:
f.readline() # header
for line in f:
t,time = line.split('\t')
touts.append(float(time))
return np.array(touts)
[docs]
def gather_all_abundances(
col: Column, outputdirr: str, outputfile: str,
tf: float, delete_subdir: bool) -> None:
"""(DEPRECATED) reads in all of the astrochem outputs from a column
and puts all the output params into a single .npz file
Args:
col (Column): CANDY column
outputdirr (str): output directory
outputfile (str): output file
tf (float): final time of simulation
delete_subdir (bool): option to delete subdirectories after
gathering abundances
"""
specs = get_specs_array(outputdirr)
nspec = len(specs)
touts = read_touts_from_file(f'{outputdirr}/times.out')
touts = touts[touts<=tf]
nts = len(touts)
nzs = col.ncells
gas_dens = np.zeros(nzs)
extinctions = np.zeros((nts,nzs))
abundances = np.zeros((nts,nzs,nspec))
heights = np.zeros(nzs)
for j in range(nzs):
dirr = f'{outputdirr}/z{j:0>2}'
for t,time in enumerate(touts):
nh,av = get_nh_and_av(dirr+f'/source_t{t}.mdl')
extinctions[t,j] = av
if t==0:
gas_dens[j] = nh
heights[j] = col.cells[j].z
final_dict = get_final_abuns(dirr+f'/astrochem_output_t{t}.h5')
for k,spec in enumerate(list(final_dict.keys())):
abundances[t,j,k] = final_dict[spec]
np.savez(f'{outputdirr}/{outputfile}',times=touts,species=specs,
abundances=abundances,heights=heights,density=gas_dens,
extinctions=extinctions)
if delete_subdir:
print('removing sub directories')
print(f'rm -rf {outputdirr}/z*')
subprocess.run(f'rm -rf {outputdirr}/z*', shell=True)
[docs]
def gather_static_abuns(
col: Column, outputdirr: str, outputfile: str,
tf: float, delete_subdir: bool):
"""Gather outputs from column into single .npz file for a static
run of CANDY
Args:
col (Column): CANDY column
outputdirr (str): output directory
outputfile (str): output filename
tf (float): final time from the simulation
delete_subdir (bool): option to delete subdirectories after
creating .npz file
"""
specs = get_specs_array(outputdirr)
nspec = len(specs)
times,_ = get_abundict(f'{outputdirr}/z00/astrochem_output.h5')
nts = len(times)
nzs = col.ncells
gas_dens = np.zeros(nzs)
heights = np.zeros(nzs)
extinctions = np.zeros((nts,nzs))
abundances = np.zeros((nts,nzs,nspec))
for j in range(nzs):
dirr = f'{outputdirr}/z{j:0>2}'
nh,av = get_nh_and_av(dirr+'/source.mdl')
extinctions[:,j] = np.ones_like(times)*av
gas_dens[j] = nh
heights[j] = col.cells[j].z
ts,abundict = get_abundict(dirr+'/astrochem_output.h5')
for k,spec in enumerate(list(abundict.keys())):
abundances[:,j,k] = abundict[spec]
np.savez(f'{outputdirr}/{outputfile}',times=times,species=specs,
abundances=abundances,heights=heights,density=gas_dens,
extinctions=extinctions)
if delete_subdir:
print('removing sub directories')
print(f'rm -rf {outputdirr}/z*')
subprocess.run(f'rm -rf {outputdirr}/z*', shell=True)
[docs]
def readh5file(filename: str='astrochem_output.h5') -> tuple[h5.File, h5.Dataset, h5.Dataset, h5.Dataset]:
"""Helper function to read astrochem output and return relevent information
Args:
filename (str, optional): name of .h5 file to read. Defaults to 'astrochem_output.h5'.
Returns:
tuple[h5.File, dict, dict, dict]: tuple of outputs
opened h5.File object (be sure to close this later!)
list of species from astrochem
array of abundances for each species over time
list of times
"""
f = h5.File(filename,'r')
spec = f['Species']
abun = f['Abundances']
time = f['TimeSteps']
return f,spec,abun,time
[docs]
def get_abundict(filename='astrochem_output.h5',specs = 'all') -> tuple[np.ndarray, dict]:
"""Create a dictionary of species abundances from an astrochem output
Args:
filename (str, optional): astrochem output h5 file. Defaults to 'astrochem_output.h5'.
specs (list[str] | str, optional): list of species to include in the output dictionary. Defaults to 'all'.
Returns:
tuple[np.ndarray, dict]: array of times and the dictionary of species and abundances
"""
f,spec,abun,time = readh5file(filename)
d = {}
time = np.array(time)
for i in range(spec.size):
specie = spec[i].decode('utf-8')
if specie in specs or specs == 'all':
abund = abun[0,:,i]
d[specie] = abund
f.close()
return time,d
[docs]
def get_final_abuns(f_name: str,specs: list[str]|str='all') -> dict:
"""Create a dictionary of the species abundances at the last time of the astrochem output
Args:
f_name (str): name of astrochem .h5 file
specs (list[str] | str, optional): list of species to include. Defaults to 'all'.
Returns:
dict: dictionary of species and abundances
"""
final_abuns = {}
times,d = get_abundict(f_name,specs=specs)
for spec in d:
final_abuns[spec]=d[spec][-1]
return final_abuns
[docs]
def save_pebbles(
col: Column, pebcomp: dict, f_pebout: str, time: float) -> None:
"""Save the pebble compistion to text file
Args:
col (Column): CANDY column
pebcomp (dict): dictionary of pebble compositions
f_pebout (str): name of pebble text file
time (float): the current time
"""
with open(f_pebout, 'a+') as f:
for spec in list(pebcomp.keys()):
peb_comp_norm = pebcomp[spec]/col.cells[0].NH
f.write(f'{time}\t{spec}\t{peb_comp_norm:.10e}\n')
if __name__ == '__main__':
import argparse
from pprint import pprint
parser = argparse.ArgumentParser()
parser.add_argument('infile',default='example.ini')
args = parser.parse_args()
model,phys,abuns = read_infile(args.infile)
pprint(model)
pprint(phys)
pprint(abuns)