Merge pull request #609 from eqcorrscan/cc-squared

[WIP] Add option to output cross-correlation * abs(cross-correlation)

Author: calum-chamberlain

CHANGES.md

@@ -10,6 +10,8 @@
   - More stable extraction of data for correlation
 - utils.findpeaks
   - Added pick-time declustering.
+- utils.correlate
+  - Added option to produce "cc-squared" (cc * abs(cc)) detection statistic
 
 ## 0.5.0
 * core.match_filter.tribe

eqcorrscan/core/match_filter/helpers/tribe.py

@@ -311,6 +311,7 @@ def _corr_and_peaks(
     concurrency: str,
     cores: int,
     i: int,
+    cc_squared: bool,
     export_cccsums: bool,
     parallel: bool,
     peak_cores: int,
@@ -336,6 +337,9 @@ def _corr_and_peaks(
     :param concurrency: Concurrency of cross-correlation function
     :param cores: Cores (threads) to use for cross-correlation
     :param i: Group-id (internal book-keeping)
+    :param cc_squared:
+        Whether to detect using "cc_squared" (actually cc * abs(cc)) or
+        just using cc.
     :param export_cccsums: Whether to export the raw cross-correlation sums
     :param parallel: Whether to compute peaks in parallel
     :param peak_cores: Number of cores (threads) to use for peak finding
@@ -358,6 +362,7 @@ def _corr_and_peaks(
         f"Starting correlation run for template group {i}")
     tic = default_timer()
     if prepped and xcorr_func == "fmf":
+        assert cc_squared == False, "FMF does not support squared correlation"
         assert isinstance(templates, np.ndarray)
         assert isinstance(stream, np.ndarray)
         # These need to be passed from queues.
@@ -387,7 +392,8 @@ def _corr_and_peaks(
         cccsums, tr_chans = fftw_multi_normxcorr(
             template_array=templates, stream_array=stream,
             pad_array=pads, seed_ids=seed_ids, cores_inner=num_cores_inner,
-            cores_outer=num_cores_outer, stack=True, **kwargs)
+            cores_outer=num_cores_outer, stack=True, cc_squared=cc_squared,
+            **kwargs)
         n_templates = len(cccsums)
         # Post processing
         no_chans = np.sum(np.array(tr_chans).astype(int), axis=0)
@@ -408,7 +414,8 @@ def _corr_and_peaks(
         multichannel_normxcorr = get_stream_xcorr(xcorr_func, concurrency)
         Logger.debug(f"Calling {multichannel_normxcorr}")
         cccsums, no_chans, chans = multichannel_normxcorr(
-            templates=templates, stream=stream, cores=cores, **kwargs
+            templates=templates, stream=stream, cores=cores,
+            cc_squared=cc_squared, **kwargs
         )
     if len(cccsums[0]) == 0:
         raise MatchFilterError(

eqcorrscan/core/match_filter/matched_filter.py

@@ -61,8 +61,8 @@ def __str__(self):
 # Note: maintained for backwards compatability. All efforts now in tribes
 def match_filter(template_names, template_list, st, threshold,
                  threshold_type, trig_int, plot=False, plotdir=None,
-                 xcorr_func=None, concurrency=None, cores=None,
-                 plot_format='png', output_cat=False,
+                 xcorr_func=None, concurrency=None, cc_squared=False,
+                 cores=None, plot_format='png', output_cat=False,
                  extract_detections=False, arg_check=True, full_peaks=False,
                  peak_cores=None, export_cccsums=False, **kwargs):
     """
@@ -115,6 +115,10 @@ def match_filter(template_names, template_list, st, threshold,
         The type of concurrency to apply to the xcorr function. Options are
         'multithread', 'multiprocess', 'concurrent'. For more details see
         :func:`eqcorrscan.utils.correlate.get_stream_xcorr`
+    :type cc_squared: bool
+    :param cc_squared:
+        Whether to detect using "cc_squared" (actually cc * abs(cc)) or
+        just using cc.
     :type cores: int
     :param cores: Number of cores to use
     :type plot_format: str
@@ -319,7 +323,9 @@ def match_filter(template_names, template_list, st, threshold,
         overlap="calculate", full_peaks=full_peaks, save_progress=False,
         process_cores=None, pre_processed=True, check_processing=False,
         return_stream=False, plot_format=plot_format,
-        peak_cores=peak_cores, export_cccsums=export_cccsums, **kwargs
+        peak_cores=peak_cores, export_cccsums=export_cccsums,
+        cc_squared=cc_squared,
+        **kwargs
     )
     detections = [d for f in party for d in f]
 

eqcorrscan/core/match_filter/tribe.py

@@ -675,11 +675,11 @@ def cluster(self, method, **kwargs):
     def detect(self, stream, threshold, threshold_type, trig_int, plot=False,
                plotdir=None, parallel_process=True,
                xcorr_func=None, concurrency=None, cores=None,
-               concurrent_processing=False, ignore_length=False,
-               ignore_bad_data=False, group_size=None, overlap="calculate",
-               full_peaks=False, save_progress=False, process_cores=None,
-               pre_processed=False, check_processing=True, make_events=True,
-               min_stations=0, **kwargs):
+               cc_squared=False, concurrent_processing=False,
+               ignore_length=False, ignore_bad_data=False, group_size=None,
+               overlap="calculate", full_peaks=False, save_progress=False,
+               process_cores=None, pre_processed=False, check_processing=True,
+               make_events=True, min_stations=0, **kwargs):
         """
         Detect using a Tribe of templates within a continuous stream.
 
@@ -720,6 +720,10 @@ def detect(self, stream, threshold, threshold_type, trig_int, plot=False,
             :func:`eqcorrscan.utils.correlate.get_stream_xcorr`
         :type cores: int
         :param cores: Number of workers for processing and detection.
+        :type cc_squared: bool
+        :param cc_squared:
+            Whether to detect using "cc_squared" (actually cc * abs(cc)) or
+            just using cc.
         :type concurrent_processing: bool
         :param concurrent_processing:
             Whether to process steps in detection workflow concurrently or not.
@@ -895,8 +899,9 @@ def detect(self, stream, threshold, threshold_type, trig_int, plot=False,
                 process_cores, ignore_length, overlap,
                 ignore_bad_data, group_size, groups, sampling_rate, threshold,
                 threshold_type, save_progress, xcorr_func, concurrency, cores,
-                export_cccsums, parallel, peak_cores, trig_int, full_peaks,
-                plot, plotdir, plot_format, make_events, min_stations,)
+                cc_squared, export_cccsums, parallel, peak_cores, trig_int,
+                full_peaks, plot, plotdir, plot_format, make_events,
+                min_stations,)
 
         if concurrent_processing:
             party = self._detect_concurrent(*args, **inner_kwargs)
@@ -922,9 +927,9 @@ def _detect_serial(
         self, stream, template_ids, pre_processed, parallel_process,
         process_cores, ignore_length, overlap, ignore_bad_data,
         group_size, groups, sampling_rate, threshold, threshold_type,
-        save_progress, xcorr_func, concurrency, cores, export_cccsums,
-        parallel, peak_cores, trig_int, full_peaks, plot, plotdir, plot_format,
-        make_events, min_stations, **kwargs
+        save_progress, xcorr_func, concurrency, cores, cc_squared,
+        export_cccsums, parallel, peak_cores, trig_int, full_peaks, plot,
+        plotdir, plot_format, make_events, min_stations, **kwargs
     ):
         """ Internal serial detect workflow. """
         from eqcorrscan.core.match_filter.helpers.tribe import (
@@ -977,7 +982,8 @@ def _detect_serial(
                     threshold=threshold, threshold_type=threshold_type,
                     trig_int=trig_int, sampling_rate=sampling_rate,
                     full_peaks=full_peaks, plot=plot, plotdir=plotdir,
-                    plot_format=plot_format, prepped=False, **kwargs)
+                    plot_format=plot_format, prepped=False,
+                    cc_squared=cc_squared, **kwargs)
                 Logger.debug(chans)
 
                 detections = _detect(
@@ -1014,9 +1020,9 @@ def _detect_concurrent(
         self, stream, template_ids, pre_processed, parallel_process,
         process_cores, ignore_length, overlap, ignore_bad_data,
         group_size, groups, sampling_rate, threshold, threshold_type,
-        save_progress, xcorr_func, concurrency, cores, export_cccsums,
-        parallel, peak_cores, trig_int, full_peaks, plot, plotdir, plot_format,
-        make_events, min_stations, **kwargs
+        save_progress, xcorr_func, concurrency, cores, cc_squared,
+        export_cccsums, parallel, peak_cores, trig_int, full_peaks, plot,
+        plotdir, plot_format, make_events, min_stations, **kwargs
     ):
         """ Internal concurrent detect workflow. """
         from eqcorrscan.core.match_filter.helpers.processes import (
@@ -1180,7 +1186,8 @@ def _detect_concurrent(
                     threshold=threshold, threshold_type=threshold_type,
                     trig_int=trig_int, sampling_rate=sampling_rate,
                     full_peaks=full_peaks, plot=plot, plotdir=plotdir,
-                    plot_format=plot_format, **inner_kwargs
+                    plot_format=plot_format, cc_squared=cc_squared,
+                    **inner_kwargs
                 )
                 peaks_queue.put(
                     (starttime, all_peaks, thresholds, no_chans, chans,
@@ -1251,9 +1258,9 @@ def client_detect(self, client, starttime, endtime, threshold,
                       threshold_type, trig_int, plot=False, plotdir=None,
                       min_gap=None, parallel_process=True,
                       xcorr_func=None, concurrency=None, cores=None,
-                      concurrent_processing=False, ignore_length=False,
-                      ignore_bad_data=False, group_size=None,
-                      return_stream=False, full_peaks=False,
+                      cc_squared=False, concurrent_processing=False,
+                      ignore_length=False, ignore_bad_data=False,
+                      group_size=None, return_stream=False, full_peaks=False,
                       save_progress=False, process_cores=None, retries=3,
                       check_processing=True, make_events=True,
                       min_stations=0, **kwargs):
@@ -1304,6 +1311,10 @@ def client_detect(self, client, starttime, endtime, threshold,
             :func:`eqcorrscan.utils.correlate.get_stream_xcorr`
         :type cores: int
         :param cores: Number of workers for processing and detection.
+        :type cc_squared: bool
+        :param cc_squared:
+            Whether to detect using "cc_squared" (actually cc * abs(cc)) or
+            just using cc.
         :type concurrent_processing: bool
         :param concurrent_processing:
             Whether to process steps in detection workflow concurrently or not.
@@ -1462,7 +1473,8 @@ def client_detect(self, client, starttime, endtime, threshold,
             return_stream=return_stream, check_processing=False,
             poison_queue=poison_queue, shutdown=False,
             concurrent_processing=concurrent_processing, groups=groups,
-            make_events=make_events, min_stations=min_stations)
+            make_events=make_events, min_stations=min_stations,
+            cc_squared=cc_squared)
         detector_kwargs.update(kwargs)
 
         if not concurrent_processing:

eqcorrscan/tests/correlate_test.py

@@ -490,6 +490,45 @@ def gappy_real_cc_dict(gappy_real_cc_output_dict):
     return {name: (result[0][0], result[1])
             for name, result in gappy_real_cc_output_dict.items()}
 
+# ------------------------------------ correlation ** abs(correlation)
+
+@pytest.fixture(scope='module')
+def stream_cc_output_dict_corrsq(
+        multichannel_templates, multichannel_stream):
+    """ return a dict of outputs from all stream_xcorr functions """
+    # corr._get_array_dicts(multichannel_templates, multichannel_stream)
+    for tr in multichannel_stream:
+        tr.data = tr.data[0:unstacked_stream_len]
+    multichannel_templates = multichannel_templates[0:5]
+    out = {}
+    for name, func in stream_funcs.items():
+        if name.startswith("fmf"):
+            print("Skipping fmf - unstacked not implemented")
+            continue
+        for cores in [1, cpu_count()]:
+            print("Running {0} with {1} cores".format(name, cores))
+
+            cc_out = time_func(func, name, multichannel_templates,
+                               multichannel_stream, cores=cores, stack=True,
+                               cc_squared=True)
+            out["{0}.{1}".format(name, cores)] = cc_out
+            if "fftw" in name and cores > 1:
+                print("Running outer core parallel")
+                # Make sure that using both parallel methods gives the same
+                # result
+                cc_out = time_func(
+                    func, name, multichannel_templates, multichannel_stream,
+                    cores=1, cores_outer=cores, stack=True, cc_squared=True)
+                out["{0}.{1}_outer".format(name, cores)] = cc_out
+    return out
+
+
+@pytest.fixture(scope='module')
+def stream_cc_dict_corrsq(stream_cc_output_dict_corrsq):
+    """ return just the cc arrays from the stream_cc functions """
+    return {name: result[0]
+            for name, result in stream_cc_output_dict_corrsq.items()}
+
 
 # ------------------------------------ unstacked setup
 
@@ -759,6 +798,24 @@ def test_gappy_real_multi_channel_xcorr(self, gappy_real_cc_dict):
                 assert np.allclose(cc_1, cc, atol=self.atol * 100)
 
 
+@pytest.mark.serial
+class TestStreamCorrelateFunctionsCorrSq:
+    """ same thing as TestArrayCorrelateFunction but for stream interface """
+    atol = TestArrayCorrelateFunctions.atol
+
+    def test_multi_channel_xcorr(self, stream_cc_dict_corrsq):
+        """ test various correlation methods with multiple channels """
+        # get correlation results into a list
+        cc_names = list(stream_cc_dict_corrsq.keys())
+        cc_list = [stream_cc_dict_corrsq[cc_name] for cc_name in cc_names]
+        cc_1 = cc_list[0]
+        # loop over correlations and compare each with the first in the list
+        # this will ensure all cc are "close enough"
+        for cc_name, cc in zip(cc_names[2:], cc_list[2:]):
+            assert np.allclose(cc_1, cc, atol=self.atol)
+
+
+
 @pytest.mark.serial
 class TestStreamCorrelateFunctionsUnstacked:
     """ same thing as TestArrayCorrelateFunction but for stream interface """
@@ -900,10 +957,10 @@ def test_callable_registered(self, multichannel_templates,
         """ ensure a callable can be registered """
         small_count = {}
 
-        def some_callable(template_array, stream_array, pad_array):
+        def some_callable(template_array, stream_array, pad_array, cc_squared):
             small_count['name'] = 1
             return corr.numpy_normxcorr(template_array, stream_array,
-                                        pad_array)
+                                        pad_array, cc_squared)
 
         func = corr.get_stream_xcorr(some_callable)
         func(multichannel_templates, multichannel_stream)
@@ -929,7 +986,7 @@ def func(templates, streams, pads):
     def test_using_custom_function_doesnt_change_default(self):
         """ ensure a custom function will not change the default """
 
-        def func(templates, streams, pads):
+        def func(templates, streams, pads, cc_squared):
             pass
 
         default = corr.get_array_xcorr(None)