Added more robust subset-stacking

Author: geojunky

seismic/ASDFdatabase/cwb2asdf/cwb2asdf.py

@@ -28,8 +28,8 @@
 from ordered_set import OrderedSet as set
 from tqdm import tqdm
 from seismic.misc import split_list
-from seismic.misc import recursive_glob
 from seismic.ASDFdatabase.utils import cleanse_inventory
+from seismic.misc import recursive_glob
 
 def make_ASDF_tag(tr, tag):
     # def make_ASDF_tag(ri, tag):
@@ -45,7 +45,7 @@ def make_ASDF_tag(tr, tag):
 # end func
 
 CONTEXT_SETTINGS = dict(help_option_names=['-h', '--help'])
-BUFFER_LENGTH = 1000
+BUFFER_LENGTH = 2000
 
 @click.command(context_settings=CONTEXT_SETTINGS)
 @click.argument('input-folder', required=True,
@@ -60,7 +60,7 @@ def make_ASDF_tag(tr, tag):
 @click.option('--min-length-sec', type=int, default=None, help="Minimum length in seconds")
 @click.option('--merge-threshold', type=int, default=None, help="Merge traces if the number of traces fetched for an "
                                                                 "interval exceeds this threshold")
-@click.option('--ntraces-per-file', type=int, default=3600, help="Maximum number of traces per file; if exceeded, the "
+@click.option('--ntraces-per-file', type=int, default=600, help="Maximum number of traces per file; if exceeded, the "
                                                                  "file is ignored.")
 @click.option('--dry-run', default=False, is_flag=True, show_default=True,
               help="Dry run only reports stations that were not found in the stationXML files, for "
@@ -75,6 +75,7 @@ def process(input_folder, inventory_folder, output_file_name, file_pattern,
     OUTPUT_FILE_NAME: Name of output ASDF file \n
     """
 
+    inventory_added = defaultdict(bool)
     def _read_inventories(inventory_folder):
         inv_files = recursive_glob(inventory_folder, '*.xml')
 
@@ -103,8 +104,11 @@ def _write(ds, ostream, inventory_dict, netsta_set):
         # end try
 
         for item in netsta_set:
+            if(inventory_added[item]): continue
+
             try:
                 ds.add_stationxml(inventory_dict[item])
+                inventory_added[item] = True
             except Exception as e:
                 print(e)
                 print('Failed to append inventory:')
@@ -129,10 +133,7 @@ def _write(ds, ostream, inventory_dict, netsta_set):
         inv = _read_inventories(inventory_folder)
 
         # generate a list of files
-        paths = [i for i in os.listdir(input_folder) if os.path.isfile(os.path.join(input_folder, i))]
-        expr = re.compile(fnmatch.translate(file_pattern), re.IGNORECASE)
-        files = [os.path.join(input_folder, j) for j in paths if re.match(expr, j)]
-
+        files = recursive_glob(input_folder, file_pattern)
         files = np.array(files)
         random.Random(nproc).shuffle(files)
         #print(files); exit(0)
@@ -143,21 +144,30 @@ def _write(ds, ostream, inventory_dict, netsta_set):
         stationlist = []
         filtered_files = []
         for file in tqdm(files, desc='Reading trace headers: '):
-            #_, _, net, sta, _ = file.split('.')
-            #tokens = os.path.basename(file).split('.')
-            #net, sta = tokens[0], tokens[1]
+            net = sta = None
 
-            st = []
-            try:
-                st = read(file, headonly=True)
-            except Exception as e:
-                print(e)
-                continue
-            # end try
-            if(len(st) == 0): continue
+            if(True):
+                try:
+                    fn = os.path.basename(file)
+                    net, sta = fn.split('.')[:2]
+                except:
+                    continue
+                # end try
+                #tokens = os.path.basename(file).split('.')
+                #net, sta = tokens[0], tokens[1]
+            else:
+                st = []
+                try:
+                    st = read(file, headonly=True)
+                except Exception as e:
+                    print(e)
+                    continue
+                # end try
+                if(len(st) == 0): continue
 
-            net = st[0].meta.network
-            sta = st[0].meta.station
+                net = st[0].meta.network
+                sta = st[0].meta.station
+            # end if
 
             ustations.add('%s.%s' % (net, sta))
             networklist.append(net)

seismic/xcorqc/subset_stacker.py

@@ -61,6 +61,7 @@ def __init__(self):
         self.DIST_MINMAX = [d['DIST_MIN'], d['DIST_MAX']]
         self.EMAG_MINMAX = [d['EMAG_MIN'], d['EMAG_MAX']]
         self.AZ_TOL = d['AZ_TOL']
+        self.param_dict = d
 
         self.gc = None
         if(self.rank == 0):
@@ -91,6 +92,24 @@ def stack(self, spooled_matrix:SpooledMatrix,
                  as defined in the manuscript
         """
 
+        def circular_select(angles, min_angle, max_angle):
+            # Normalize angles to [0, 360)
+            angles = np.mod(angles, 360);
+            min_angle = np.mod(min_angle, 360);
+            max_angle = np.mod(max_angle, 360);
+
+            result = None
+            if min_angle <= max_angle:
+                # Linear range (e.g., 10-20)
+                result = (angles >= min_angle) & (angles <= max_angle)
+            else:
+                # Circular wraparound range (e.g., 350-10)
+                result = (angles >= min_angle) | (angles <= max_angle)
+            # end if
+
+            return result
+        # end func
+
         def get_affected_indices(source_eids, pat, swat, swet):
             """
             Finds indices of CC windows affected by P and SW energy
@@ -173,9 +192,8 @@ def get_affected_indices(source_eids, pat, swat, swet):
         eids2 = (edistdeg2 >= self.DIST_MINMAX[0]) & (edistdeg2 <= self.DIST_MINMAX[1])
 
         # find event indices within azimuth of stations 1 and 2
-        eids1_inside_az = eids1 & ((eaz1 >= (baz - self.AZ_TOL)) & (eaz1 <= (baz + self.AZ_TOL)))
-        eids2_inside_az = eids2 & ((eaz2 >= (az - self.AZ_TOL)) & (eaz2 <= (az + self.AZ_TOL)))
-
+        eids1_inside_az = eids1 & circular_select(eaz1, (baz - self.AZ_TOL), (baz + self.AZ_TOL))
+        eids2_inside_az = eids2 & circular_select(eaz2, (az - self.AZ_TOL), (az + self.AZ_TOL))
         eids_inside_az = eids1_inside_az | eids2_inside_az
 
         # find event indices outside azimuth of both stations
@@ -242,12 +260,12 @@ def get_affected_indices(source_eids, pat, swat, swet):
         mean_XeiUXec /= float(wc_XeiUXec)
         mean_Xeo /= float(wc_Xeo)
 
-        #"""
+        """
         np.savez('stack3outputs.npz', xcf=mean,
                  xcf1=mean_Xei, xcf2=mean_Xec, xcf3=mean_XeiUXec,
                  xcf4=mean_Xeo, idsXei=idsXei, idsXec=idsXec,
                  idsXeiUXec=idsXeiUXec, idsXeo=idsXeo)
-        #"""
+        """
 
         return mean, mean_Xei, mean_Xec, mean_XeiUXec, mean_Xeo, \
                wc, wc_Xei, wc_Xec, wc_XeiUXec, wc_Xeo

seismic/xcorqc/xcorqc.py

@@ -809,15 +809,15 @@ def IntervalStackXCorr(refds, tempds,
         if(memoryTracker): memoryTracker.update()
     # wend (loop over time range)
 
-    x = None
+    times = None
     flattenedWindowCounts = None
     flattenedIntervalStartTimes = None
     flattenedIntervalEndTimes = None
     flattenedWindowStartTimes = None
     flattenedWindowEndTimes = None
     if(spooledXcorr and spooledXcorr.nrows):
         dt = 1./sr
-        x = np.linspace(-window_seconds + dt, window_seconds - dt, spooledXcorr.ncols)
+        times = np.linspace(-window_seconds + dt, window_seconds - dt, spooledXcorr.ncols)
 
         flattenedWindowCounts = np.concatenate(windowCounts).ravel()
         flattenedIntervalStartTimes = np.concatenate(intervalStartTimes).ravel()
@@ -856,11 +856,11 @@ def IntervalStackXCorr(refds, tempds,
 
         # Add data
         if ensemble_stack:
-            nsw = root_grp.createVariable('NumStackedWindows', 'i8')
+            nsw = root_grp.createVariable('nX', 'i8')
             avgnsw = root_grp.createVariable('AvgNumStackedWindowsPerInterval', 'f4')
             ist = root_grp.createVariable('IntervalStartTime', 'i8')
             iet = root_grp.createVariable('IntervalEndTime', 'i8')
-            xc = root_grp.createVariable('xcorr', XCORR_FORMAT, ('lag',))
+            xc = root_grp.createVariable('X', XCORR_FORMAT, ('lag',))
 
             totalIntervalCount = int(np.sum(flattenedWindowCounts > 0))
             totalWindowCount = int(np.sum(flattenedWindowCounts))
@@ -897,26 +897,26 @@ def IntervalStackXCorr(refds, tempds,
             wet[:] = flattenedWindowEndTimes
 
             if(apply_simple_stacking):
-                nsw = root_grp.createVariable('NumStackedWindows', 'f4', ('interval',))
-                xc = root_grp.createVariable('xcorr', XCORR_FORMAT, ('interval', 'lag',),
+                nsw = root_grp.createVariable('nX', 'f4', ('interval',))
+                xc = root_grp.createVariable('X', XCORR_FORMAT, ('interval', 'lag',),
                                          chunksizes=(1, spooledXcorr.ncols),
                                          zlib=True)
                 nsw[:] = flattenedWindowCounts
                 for irow in np.arange(spooledXcorr.nrows):
                     xc[irow, :] = spooledXcorr.read_row(irow)
                 # end for
             elif subset_stacker is not None:
-                xc = root_grp.createVariable('xcorr', XCORR_FORMAT, ('lag',))
-                xc_Xei = root_grp.createVariable('xcorr_Xei', XCORR_FORMAT, ('lag',))
-                xc_Xec = root_grp.createVariable('xcorr_Xec', XCORR_FORMAT, ('lag',))
-                xc_XeiUXec = root_grp.createVariable('xcorr_XeiUXec', XCORR_FORMAT, ('lag',))
-                xc_Xeo = root_grp.createVariable('xcorr_Xeo', XCORR_FORMAT, ('lag',))
-
-                xc_wc = root_grp.createVariable('xcorr_NumStackedWindows', 'i8')
-                xc_Xei_wc = root_grp.createVariable('xcorr_Xei_NumStackedWindows', 'i8')
-                xc_Xec_wc = root_grp.createVariable('xcorr_Xec_NumStackedWindows', 'i8')
-                xc_XeiUXec_wc = root_grp.createVariable('xcorr_XeiUXec_NumStackedWindows', 'i8')
-                xc_Xeo_wc = root_grp.createVariable('xcorr_Xeo_NumStackedWindows', 'i8')
+                xc = root_grp.createVariable('X', XCORR_FORMAT, ('lag',))
+                xc_Xei = root_grp.createVariable('Xei', XCORR_FORMAT, ('lag',))
+                xc_Xec = root_grp.createVariable('Xec', XCORR_FORMAT, ('lag',))
+                xc_XeiUXec = root_grp.createVariable('XeiUXec', XCORR_FORMAT, ('lag',))
+                xc_Xeo = root_grp.createVariable('Xeo', XCORR_FORMAT, ('lag',))
+
+                xc_wc = root_grp.createVariable('nX', 'i8')
+                xc_Xei_wc = root_grp.createVariable('nXei', 'i8')
+                xc_Xec_wc = root_grp.createVariable('nXec', 'i8')
+                xc_XeiUXec_wc = root_grp.createVariable('nXeiUXec', 'i8')
+                xc_Xeo_wc = root_grp.createVariable('nXeo', 'i8')
 
                 slon1, slat1 = refds.unique_coordinates[ref_net_sta]
                 slon2, slat2 = tempds.unique_coordinates[temp_net_sta]
@@ -938,7 +938,7 @@ def IntervalStackXCorr(refds, tempds,
                 xc_XeiUXec_wc[:] = wc_XeiUXec
                 xc_Xeo_wc[:] = wc_Xeo
             else:
-                xc = root_grp.createVariable('xcorr', XCORR_FORMAT, ('window', 'lag',),
+                xc = root_grp.createVariable('X', XCORR_FORMAT, ('window', 'lag',),
                                              chunksizes=(1, spooledXcorr.ncols),
                                              zlib=True)
                 for irow in np.arange(spooledXcorr.nrows):
@@ -947,7 +947,7 @@ def IntervalStackXCorr(refds, tempds,
             # end if
         # end if
 
-        lag[:] = x
+        lag[:] = times
 
         # Add and populate a new group for parameters used
         pg = root_grp.createGroup('Parameters')
@@ -972,6 +972,12 @@ def IntervalStackXCorr(refds, tempds,
                   'ensemble_stack': int(ensemble_stack),
                   'simple_stack': int(apply_simple_stacking),
                   'subset_stack': int(subset_stacker is not None)}
+        if(subset_stacker is not None):
+            for k, v in subset_stacker.param_dict.items():
+                k = 'subset_stack.{}'.format(k.lower())
+                params.update({k: v})
+            # end for
+        # end if
 
         if whitening:
             params['whitening_window_frequency'] = whitening_window_frequency

tests/test_seismic/xcorqc/data/expected/expected.tar.gz

@@ -83,13 +83,13 @@ def test_correlator():
         fn = os.path.join(curr_output_folder, '%s.%s.%s.%s.%s.%s.nc'%(netsta1, loc_code, cha,
                                                                       netsta2, '', cha))
         dc = Dataset(fn)
-        xcorr_c = dc.variables['xcorr'][:]
+        xcorr_c = dc.variables['X'][:]
 
         # Read expected
         fn = '%s/%s.%s.%s.%s.%s.%s.nc'%(curr_expected_folder, netsta1, loc_code, cha,
                                         netsta2, '', cha)
         de = Dataset(fn)
-        xcorr_e = de.variables['xcorr'][:]
+        xcorr_e = de.variables['X'][:]
 
         rtol = 1e-5
         atol = 1e-5
@@ -116,13 +116,13 @@ def test_correlator():
         fn = os.path.join(curr_output_folder, '%s.%s.%s.%s.%s.%s.nc'%(netsta1, loc_code, cha,
                                                                       netsta2, '', cha))
         dc = Dataset(fn)
-        xcorr_c = dc.variables['xcorr'][:]
+        xcorr_c = dc.variables['X'][:]
 
         # Read expected
         fn = '%s/%s.%s.%s.%s.%s.%s.nc'%(curr_expected_folder, netsta1, loc_code, cha,
                                         netsta2, '', cha)
         de = Dataset(fn)
-        xcorr_e = de.variables['xcorr'][:]
+        xcorr_e = de.variables['X'][:]
 
         rtol = 1e-5
         atol = 1e-5

tests/test_seismic/xcorqc/data/expected/stacked/II.WRAB.00.BHZ.AU.MAW..BHZ.nc

@@ -129,13 +129,13 @@ def test_interval_stack_xcorr(loccha, inv1, inv2, interval_seconds, window_secon
     fn = os.path.join(output_folder, '%s.%s.%s.%s.%s.%s.%s.nc'%(netsta1, loc, cha,
                                                                 netsta2, loc, cha, tag))
     dc = Dataset(fn)
-    xcorr_c = dc.variables['xcorr'][:]
+    xcorr_c = dc.variables['X'][:]
 
     # Read expected
     fn = '%s/%s.%s.%s.%s.%s.%s.%s.nc'%(expected_folder, netsta1, loc, cha,
                                        netsta2, loc, cha, tag)
     de = Dataset(fn)
-    xcorr_e = de.variables['xcorr'][:]
+    xcorr_e = de.variables['X'][:]
 
     rtol = 1e-3
     atol = 1e-3