Merge branch 'ncr-rad-tmp'

pyathena/load_sim.py

@@ -131,9 +131,12 @@ def __init__(self, basedir, savdir=None, load_method='pyathena',
 
         # Get config time
         try:
-            self.config_time = pd.to_datetime(self.par['configure']['config_date'])
-            if 'PDT' in self.par['configure']['config_date']:
-                self.config_time = self.config_time.tz_localize('US/Pacific')
+            config_time = self.par['configure']['config_date']
+            # Avoid un-recognized timezone FutureWarning
+            config_time = config_time.replace('PDT ', '')
+            config_time = config_time.replace('EDT ', '')
+            self.config_time = pd.to_datetime(config_time).tz_localize('US/Pacific')
+            #self.config_time = self.config_time
         except:
             try:
                 # set it using hst file creation time
@@ -645,13 +648,13 @@ def _find_files(self):
 
         athinput_patterns = [
             ('athinput.runtime',),
-            ('athinput.*',),
             ('out.txt',),
             ('out*.txt',),
             ('stdout.txt',),
             ('log.txt',),
             ('*.par',),
             ('*.out',),
+            ('athinput.*',),
             ('slurm-*',)]
 
         hst_patterns = [('id0', '*.hst'),

pyathena/tigress_dig/zprof.py

@@ -201,7 +201,7 @@ def read_zprof_partially_ionized(self, num, prefix='pionized',
 
         return r
 
-    def get_zprof_partially_ionzied_all(self, nums,
+    def get_zprof_partially_ionized_all(self, nums,
                                         savdir, force_override=False):
 
         rr = dict()

pyathena/tigress_ncr/do_tasks_job_array.py

@@ -0,0 +1,58 @@
+#!/usr/bin/env python
+
+import os
+import os.path as osp
+import gc
+import time
+from mpi4py import MPI
+import matplotlib.pyplot as plt
+import pprint
+import sys
+import pickle
+
+import pyathena as pa
+from pyathena.util.split_container import split_container
+
+if __name__ == '__main__':
+    COMM = MPI.COMM_WORLD
+
+    basedir = '/projects/EOSTRIKE/TIGRESS-NCR'
+    # All subdirectories
+    pp = [osp.join(basedir, p) for p in next(os.walk(basedir))[1]]
+
+    idx = int(os.environ["SLURM_ARRAY_TASK_ID"])
+    basedir = pp[idx]
+    model = osp.basename(basedir)
+    savdir = osp.join('/tigress/jk11/NCR-RAD-LOWZ', model, 'pionized')
+    if not osp.exists(savdir):
+        os.makedirs(savdir)
+
+    s = pa.LoadSimTIGRESSNCR(basedir, savdir=savdir, verbose=False)
+    nums = s.nums
+
+    if COMM.rank == 0:
+        print('basedir, nums', s.basedir, nums)
+        nums = split_container(nums, COMM.size)
+    else:
+        nums = None
+
+    mynums = COMM.scatter(nums, root=0)
+    print('[rank, mynums]:', COMM.rank, mynums)
+
+    time0 = time.time()
+    for num in mynums:
+        print(num, end=' ')
+        rr = s.read_zprof_partially_ionized(num, savdir=savdir)
+        n = gc.collect()
+        print('Unreachable objects:', n, end=' ')
+        print('Remaining Garbage:', end=' ')
+        pprint.pprint(gc.garbage)
+
+    COMM.barrier()
+    if COMM.rank == 0:
+        print('')
+        print('################################################')
+        print('# Do tasks')
+        print('# Execution time [sec]: {:.1f}'.format(time.time()-time0))
+        print('################################################')
+        print('')

pyathena/tigress_ncr/load_sim_tigress_ncr.py

@@ -14,9 +14,11 @@
 from .starpar import StarPar
 from .snapshot_HIH2EM import Snapshot_HIH2EM
 from .profile_1d import Profile1D
+from .rad_and_pionized import RadiationAndPartiallyIonized
 
 class LoadSimTIGRESSNCR(LoadSim, Hst, Zprof, SliceProj,
-                        StarPar, PDF, Hist2d, H2, Profile1D, Snapshot_HIH2EM):
+                        StarPar, PDF, Hist2d, H2, Profile1D, Snapshot_HIH2EM,
+                        RadiationAndPartiallyIonized):
     """LoadSim class for analyzing TIGRESS-RT simulations.
     """
 

pyathena/tigress_ncr/ncr_rad.py

@@ -0,0 +1,79 @@
+import numpy as np
+
+def get_nums_starpar(s, tMyr_range=np.array([250, 451])):
+    """Function to get starpar vtk snapshot numbers
+
+    Time range hardcoded
+    """
+
+    if s.par['output3']['out_fmt'] == 'starpar_vtk':
+        dt_starpar = s.par['output3']['dt']
+    else:
+        raise ValueError('Cannot find starpar_vtk output dt')
+
+    tMyr_range_def = {'R8-4pc': np.array([250, 451])/s.u.Myr/dt_starpar,
+                      'LGR4-2pc': np.array([250, 351])/s.u.Myr/dt_starpar,
+                      }
+
+    if s.basename.startswith('R8_4pc_NCR.full.xy2048.eps0.np768.has'):
+        tMyr_range = tMyr_range_def['R8-4pc']
+    elif s.basename.startswith('LGR4_2pc_NCR.full'):
+        tMyr_range = tMyr_range_def['LGR4-2pc']
+    else:
+        raise ValueError('Cannot find matching model name')
+
+    nums_sp = [num for num in range(*tuple([int(t) for t in tMyr_range]))]
+
+    return nums_sp
+
+def get_Qi_L_FUV_distribution_all(s, force_override=False):
+    """Function to calculate radiation source statistics
+    """
+    nums = get_nums_starpar(s)
+    spa = s.read_starpar_all(nums=nums, savdir=s.savdir,
+                             force_override=force_override)
+    # Instantaneous Qi,sp and L_FUV,sp in all snapshots
+    Qi = []
+    L_FUV = []
+    # Total number of radiation sources
+    ntot_Qi = []
+    ntot_FUV = []
+    # Qi,sp and L_FUV,sp that account for >90% of the total
+    Qi_90 = []
+    L_FUV_90 = []
+    # Number of such sources
+    n90_Qi = []
+    n90_FUV = []
+    for i, sp in spa['sp'].items():
+        # print(i, end=' ')
+        Qi.append(list(sp['Qi'].values))
+        L_FUV.append(sp['L_FUV'].values)
+
+        Qi_srt = sp['Qi'].sort_values(ascending=False)
+        L_FUV_srt = sp['L_FUV'].sort_values(ascending=False)
+
+        idx_Qi = Qi_srt.cumsum() < 0.9*Qi_srt.sum()
+        idx_FUV = L_FUV_srt.cumsum() < 0.9*L_FUV_srt.sum()
+
+        n90_Qi.append(idx_Qi.sum())
+        n90_FUV.append(idx_FUV.sum())
+        Qi_90.append(Qi_srt[idx_Qi])
+        L_FUV_90.append(L_FUV_srt[idx_FUV])
+
+        ntot_Qi.append(len(sp['Qi'].values))
+        ntot_FUV.append(len(sp['L_FUV'].values))
+
+    # Convert list of list to 1d array
+    import itertools
+    Qi = np.array(list(itertools.chain.from_iterable(Qi)))
+    L_FUV = np.array(list(itertools.chain.from_iterable(L_FUV)))
+    Qi_90 = np.array(list(itertools.chain.from_iterable(Qi_90)))
+    L_FUV_90 = np.array(list(itertools.chain.from_iterable(L_FUV_90)))
+
+    time_code = spa['time'].values
+    time_Myr = spa['time'].values*s.u.Myr
+    r = dict(spa=spa, time_code=time_code, time_Myr=time_Myr,
+             Qi=Qi, L_FUV=L_FUV, ntot_Qi=ntot_Qi, ntot_FUV=ntot_FUV,
+             n90_Qi=n90_Qi, n90_FUV=n90_FUV, Qi_90=Qi_90, L_FUV_90=L_FUV_90)
+
+    return r