Xcorr updates (#267)

* Added more robust subset-stacking
* Miscellaneous Changes
* Improved data-coverage function in FederatedASDFDataSet
* Improved handling of location codes in cross-correlation routines

Author: geojunky

seismic/ASDFdatabase/FederatedASDFDataSet.py

@@ -138,13 +138,15 @@ def get_global_time_range(self, network, station=None, location=None, channel=No
 
     # end func
 
-    def get_nslc_list(self):
+    def get_nslc_coverage(self):
         """
-        Get a list of all net, sta, loc, cha combinations featured in the database
+        Get a structured numpy array with named columns
+        'net', 'sta', 'loc', 'cha', 'min_st', 'max_et'
+        representing contents of the database
         @return:
         """
 
-        results = self.fds.get_nslc_list()
+        results = self.fds.get_nslc_coverage()
         return results
     # end if
 

seismic/ASDFdatabase/_FederatedASDFDataSetImpl.py

@@ -490,11 +490,17 @@ def get_global_time_range(self, network, station=None, location=None, channel=No
         return min, max
     # end func
 
-    def get_nslc_list(self):
-        query = "select net, sta, loc, cha from nslc"
+    def get_nslc_coverage(self):
+        query = "select net, sta, loc, cha, st, et from nslc"
         rows = self.conn.execute(query).fetchall()
 
-        return rows
+        fields = {'names': ['net', 'sta', 'loc', 'cha', 'min_st', 'max_et'],
+                  'formats': ['U10', 'U10', 'U10', 'U10', 'f8', 'f8']}
+        result = np.zeros(len(rows), dtype=fields)
+
+        for i, row in enumerate(rows): result[i] = row
+
+        return result
     # end if
 
     def get_stations(self, starttime, endtime, network=None, station=None, location=None, channel=None):

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/ASDFdatabase/viewer/FederatedASDFViewer.py

@@ -482,9 +482,9 @@ def addWidget(emitter):
         # populate net, sta, loc, cha dict
         self.nslc_dict = defaultdict(lambda: defaultdict(lambda: defaultdict(list)))
 
-        nslc_list = self.fds.get_nslc_list()
-        for row in nslc_list:
-            net, sta, loc, cha = row
+        nslc_coverage = self.fds.get_nslc_coverage()
+        for row in nslc_coverage:
+            net, sta, loc, cha, _, _ = row
             self.nslc_dict[net][sta][loc].append(cha)
         # end for
 

seismic/misc_p.py

@@ -31,7 +31,11 @@ def __init__(self, output_folder, restart_mode=False):
 
         if(self.restart_mode):
             if(os.path.exists(self.proc_fn)):
-                self.prev_progress = int(open(self.proc_fn).read())
+                try:
+                    self.prev_progress = int(open(self.proc_fn).read())
+                except:
+                    pass
+                # end try
             # end if
         # end if
     # end func
@@ -41,12 +45,13 @@ def increment(self):
         if(self.restart_mode and (self.prev_progress > 0) and (self.progress < self.prev_progress)):
             return False
         else:
-            tmpfn = self.proc_fn + '.tmp'
-            f = open(tmpfn, 'w+')
-            f.write(str(self.progress))
-            f.close()
-            os.rename(tmpfn, self.proc_fn)
-
+            try:
+                f = open(self.proc_fn, 'w+')
+                f.write(str(self.progress))
+                f.close()
+            except:
+                pass
+            # end try
             return True
         # end if
     # end func

seismic/xcorqc/correlator.py

@@ -193,7 +193,6 @@ def cull_pairs(pairs, keep_list_fn):
 
     startTime = UTCDateTime(start_time)
     endTime = UTCDateTime(end_time)
-    stationsCache = defaultdict(list)
     for pair in proc_stations[rank]:
         netsta1, netsta2 = pair
 
@@ -212,7 +211,9 @@ def get_loccha(cha1, cha2):
                      cha1 and cha2 -- e.g. ['.SHZ', '00.BHZ'], ['01.HHZ']
             """
             result = [[], []]
-            for chidx, (netsta, cha, ds) in enumerate(zip((netsta1, netsta2), (cha1, cha2), (ds1, ds2))):
+            for chidx, (netsta, cha, ds) in enumerate(zip((netsta1, netsta2),
+                                                          (cha1, cha2),
+                                                          (ds1, ds2))):
                 if('*' in cha1):
                     cha = cha.replace('*', '.*')  # hack to capture simple regex comparisons
                 # end if
@@ -225,12 +226,14 @@ def get_loccha(cha1, cha2):
 
                 net, sta = netsta.split('.')
 
-                if((start_time, end_time, net, sta) in stationsCache):
-                    stations = stationsCache[(start_time, end_time, net, sta)]
-                else:
-                    stations = ds.fds.get_stations(start_time, end_time, net, sta)
-                    stationsCache[(start_time, end_time, net, sta)] = stations
-                # end if
+                # find a list of entries where network and station names match and
+                # start- and end-times overlap with data coverage. Note that this is
+                # an approximate estimate and an actual cross-correlation may not be
+                # computed due to gaps in data
+                stations = ds.nslc_coverage[(ds.nslc_coverage['net'] == net) & \
+                                            (ds.nslc_coverage['sta'] == sta) & \
+                                            (ds.nslc_coverage['max_et'] >= startTime.timestamp) & \
+                                            (ds.nslc_coverage['min_st'] <= endTime.timestamp)]
 
                 loc_pref = location_preferences_dict[netsta]
                 ulocs = set()

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,20 +192,15 @@ 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
         eids_outside_az = ~(eids_inside_az)
 
         if(True):
             # sanity check
-            test = (eids1 & ((eaz1 < (baz - self.AZ_TOL)) | (eaz1 > (baz + self.AZ_TOL)))) & \
-                   (eids2 & ((eaz2 < (az - self.AZ_TOL)) | (eaz2 > (az + self.AZ_TOL))))
-
-            assert np.alltrue(eids_outside_az == test)
             assert len(set(np.where(eids1_inside_az | eids2_inside_az)[0]).intersection( \
                                     set(np.where(eids_outside_az)[0]))) == 0
         # end if
@@ -236,18 +250,18 @@ def get_affected_indices(source_eids, pat, swat, swet):
         wc_XeiUXec = np.sum(idsXeiUXec)
         wc_Xeo = np.sum(idsXeo)
 
-        mean /= float(wc)
-        mean_Xei /= float(wc_Xei)
-        mean_Xec /= float(wc_Xec)
-        mean_XeiUXec /= float(wc_XeiUXec)
-        mean_Xeo /= float(wc_Xeo)
+        if(wc > 0): mean /= float(wc)
+        if(wc_Xei > 0): mean_Xei /= float(wc_Xei)
+        if(wc_Xec > 0): mean_Xec /= float(wc_Xec)
+        if(wc_XeiUXec > 0): mean_XeiUXec /= float(wc_XeiUXec)
+        if(wc_Xeo > 0): 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/utils.py

@@ -20,6 +20,8 @@ def __init__(self, asdf_file_name, netsta_list='*'):
         self._earth_radius = 6371  # km
 
         self.fds = FederatedASDFDataSet(asdf_file_name)
+        self.nslc_coverage = self.fds.get_nslc_coverage()
+
         # Gather station metadata
         netsta_list_subset = set(netsta_list.split(' ')) if netsta_list != '*' else netsta_list
         self.netsta_list = []
@@ -179,6 +181,7 @@ def read_subset_stacker_config()->dict:
             "EMAG_MAX",
             "AZ_TOL"]
     fn = os.path.join(os.getcwd(), 'subset_stack.conf')
+    d = {}
     try:
         d = read_key_value_pairs(fn, keys, strict=True)
         for k in keys[1:]:
@@ -400,7 +403,7 @@ def from_nc(cls, nc_file):
         """
         try:
             ds = ncDataset(nc_file)
-            xcorr = np.array(ds.variables['xcorr'])
+            xcorr = np.array(ds.variables['X'])
             shp = xcorr.shape
             ncols = 0