LavaPlanet/run_ktable_vis.py at master · gitlinffff/LavaPlanet · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
#! /usr/bin/env python3
import sys, os
sys.path.append('/home/linfel/LavaPlanet/build/python')

from pyathena import FileMode, io_wrapper, parameter_input
from pyharp import radiation, init_index_map
from multiprocessing import Pool, cpu_count
import os, re, subprocess, argparse

cmake_source_dir = '/home/linfel/canoe'
cmake_binary_dir = '/home/linfel/LavaPlanet/build'

# prepare parse and add argument
parser = argparse.ArgumentParser()
parser.add_argument('-i', '--input',
                    help = 'athena input file'
                    )
args = vars(parser.parse_args())

# file architecture
hitbin  = f"{cmake_binary_dir}/bin/hitbin.release"
rfm     = f"{cmake_binary_dir}/bin/rfm.release"
run_rfm = f"{cmake_binary_dir}/bin/run_rfm.py"
kcoeff  = f"{cmake_binary_dir}/bin/kcoeff.release"
cktable = f"{cmake_binary_dir}/bin/cktable.py"

parfile = ""
hitfile = "/home/linfel/LavaPlanet/co_HITRAN2020/lb_28Si-16O__SiOUVenIR__14285-25000__296K.par"
#hitfile = "/home/linfel/LavaPlanet/co_HITRAN2020/lb_28Si-16O__SiOUVenIR__100-14285__296K.par"

base_name = os.path.basename(args['input'])
path_name = os.path.dirname(args['input'])
inpfile = base_name.split('.')[0]
# number of threads
max_threads = cpu_count() // 2

# ktable specifics
atm     = f"{cmake_source_dir}/data/lava_SiO_atm_isothermal.txt"

# temperature grid
#temp    = "-50 50 3"
temp    = "-400 400 7"

# flags
generate_tab = True
generate_nc  = True
generate_cktable = True

# spectral bands
file = io_wrapper()
pin = parameter_input()

if path_name != '':
    file.open(f"{path_name}/{inpfile}.inp", FileMode.read)
else:
    file.open(f"{inpfile}.inp", FileMode.read)
pin.load_from_file(file)
file.close()

init_index_map(pin)
rad = radiation()
rad.load_all_radiation_bands(pin)

bands_info = []
for i in range(rad.get_num_bands()):
    info = {}
    band = rad.get_band(i)
    info['name'] = band.get_name()
    info['wmin'] = band.get_wavenumber_min()
    info['wmax'] = band.get_wavenumber_max()
    info['wres'] = band.get_wavenumber_res()
    cia_list, hitran_list = [], []
    for j in range(band.get_num_absorbers()):
        if band.get_absorber(j).get_category() == "cia":
            cia_list.append(band.get_absorber(j).get_name())
        if band.get_absorber(j).get_category() == "hitran":
            hitran_list.append(band.get_absorber(j).get_name())
    info['cia'] = '\"' + ' '.join(cia_list) + '\"'
    info['hitran'] = ' '.join(hitran_list)
    bands_info.append(info)

# number of parallel threads
nthreads = min(max_threads, len(bands_info))

# run ktable in single thread
def RunSingleKtable(info):
  bname = info['name']
  print(f"working on band {bname} ...")
  wmin, wmax, wres = info['wmin'], info['wmax'], info['wres']
  wave = f'{wmin} {wmax} {wres}'
  tab_folder = f'{wmin}-{wmax}'
  kinp = f'kcoeff.inp-{bname}'
  kncfile = f'kcoeff.{inpfile}-{bname}.nc'
  # create tab files and kcoeff.inp
  if generate_tab:
    script = ['python', run_rfm,'--hitbin', hitbin,'--rfm', rfm,
              '--par', parfile, '--hit', hitfile,
              '--atm',atm,'--wave',wave,'--temp',temp,
              '--molecule',info['hitran'],'--output',kinp,'--rundir',tab_folder]
    out, err = subprocess.Popen(script,
        stdout = subprocess.PIPE,
        stderr = subprocess.PIPE).communicate()
    print(out.decode(), err.decode())
    #if not ('Successful' in out.decode()):
    #  raise RuntimeError("Error in generating tab files.")

  # run kcoeff
  if generate_nc:
    script = [kcoeff,'-i',kinp,'-o',kncfile]
    out, err = subprocess.Popen(script,
        stdout = subprocess.PIPE,
        stderr = subprocess.PIPE).communicate()
    print(out.decode(), err.decode())
    #if err != b'':
    #    raise RuntimeError("Error in generating kcoeff.**.nc.")
  print("band %s finishes." % wave)

  # create correlated-K table
  if generate_cktable:
    script = [cktable,'--kcoeff',kncfile,'--atm',atm,'--cia',info['cia'],
              '--input', kinp, '--output', f'{inpfile}-{bname}',
              f' > log.cktable-{inpfile}-{bname} &']
    with open('run_cktable.sh', 'a') as f:
      f.write(' '.join(script))
      f.write('\n\n')
  print("band %s finishes." % wave)

if __name__ == '__main__':
    if generate_cktable:
        with open('run_cktable.sh', 'w') as f:
            f.write('#!/bin/bash\n')

    print(bands_info)
# parallel on spectral bands
    pool = Pool(nthreads)
    pool.map(RunSingleKtable, bands_info)