From 1c54313a74576db0bd4d49f4fe7a6e46a2394958 Mon Sep 17 00:00:00 2001
From: Claudio Satriano <satriano@gmail.com>
Date: Wed, 26 Jun 2024 10:50:08 +0200
Subject: [PATCH] Initial support for plotting families found with template
 scan.

---
 CHANGELOG.md                      |  1 +
 requake/config/parse_arguments.py |  4 ++
 requake/families/families.py      | 63 +++++++++++++++++++++++++++++--
 3 files changed, 65 insertions(+), 3 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 5367f3c..c63e157 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -16,6 +16,7 @@ Copyright (c) 2021-2024 Claudio Satriano <satriano@ipgp.fr>
 - New option `--range` to manually specify the range of values for the color
   scale
 - Add missing `street` map style for `map_families`
+- Initial support for plotting families found with template scan
 - Improved reading of CSV catalog files:
   - avoid duplicated column guessing
   - ensure that prefectly matching column field names are correctly guessed
diff --git a/requake/config/parse_arguments.py b/requake/config/parse_arguments.py
index ab133f6..11cb1c0 100644
--- a/requake/config/parse_arguments.py
+++ b/requake/config/parse_arguments.py
@@ -307,6 +307,10 @@ def parse_arguments(progname='requake'):
         '-e', '--endtime', type=float, default=None,
         help='end time, in seconds relative to trace start, for the plot'
     )
+    plotfamilies.add_argument(
+        '-T', '--template', action='store_true',
+        help='plot family members found with template scan'
+    )
     # ---
     # --- plot_timespans
     plot_timespans = subparser.add_parser(
diff --git a/requake/families/families.py b/requake/families/families.py
index 8289ac3..e31701e 100644
--- a/requake/families/families.py
+++ b/requake/families/families.py
@@ -9,8 +9,11 @@
     GNU General Public License v3.0 or later
     (https://www.gnu.org/licenses/gpl-3.0-standalone.html)
 """
+import sys
 import logging
 import csv
+import os
+from glob import glob
 import numpy as np
 from obspy import UTCDateTime, Stream
 from obspy.geodetics import gps2dist_azimuth
@@ -149,9 +152,9 @@ def distance_from(self, lon, lat):
         return distance/1e3
 
 
-def read_families():
+def _read_families_from_catalog_scan():
     """
-    Read a list of families from file.
+    Read a list of families from the catalog scan output.
 
     :return: List of families.
     :rtype: list of Family
@@ -186,6 +189,51 @@ def read_families():
     return families
 
 
+def _read_families_from_template_scan():
+    """
+    Read a list of families from the template scan output.
+
+    :return: List of families.
+    :rtype: list of Family
+    """
+    template_catalogs = glob(
+        f'{config.args.outdir}/template_catalogs/catalog*.txt'
+    )
+    families = []
+    for template_catalog in template_catalogs:
+        fname = os.path.basename(template_catalog)
+        catalog_name = fname.split('.')[0]
+        trace_id = fname.lstrip(f'{catalog_name}.').rstrip('.txt')
+        family_number = int(catalog_name.lstrip('catalog'))
+        family = Family(family_number)
+        with open(template_catalog, 'r', encoding='utf-8') as fp:
+            for row in fp:
+                fields = row.split('|')
+                ev = RequakeEvent()
+                ev.evid = fields[0].strip()
+                ev.orig_time = UTCDateTime(fields[1].strip())
+                ev.lon = float(fields[2].strip())
+                ev.lat = float(fields[3].strip())
+                ev.depth = float(fields[4].strip())
+                ev.trace_id = trace_id
+                family.append(ev)
+        families.append(family)
+    return families
+
+
+def read_families():
+    """
+    Read families from the catalog scan output or from the template scan
+    output.
+
+    :return: List of families.
+    :rtype: list of Family
+    """
+    if getattr(config.args, 'template', False):
+        return _read_families_from_template_scan()
+    return _read_families_from_catalog_scan()
+
+
 def read_selected_families():
     """
     Read and select families based on family number, validity, length
@@ -262,11 +310,20 @@ def get_family_waveforms(family):
     :raises NoWaveformError: if no waveform is found
     """
     st = Stream()
-    for ev in family:
+    nevs = len(family)
+    clear_line = '\x1b[2K\r'  # escape sequence to clear line
+    for n, ev in enumerate(family):
+        sys.stdout.write(
+            f'{clear_line}Family {family.number}: '
+            f'reading waveform for event {ev.evid}: {n+1}/{nevs}')
         try:
             st += get_event_waveform(ev)
         except NoWaveformError as msg:
+            sys.stdout.write('\n')
             logger.error(msg)
+    sys.stdout.write(
+        f'{clear_line}Family {family.number}: reading waveforms: done.\n'
+    )
     if not st:
         raise NoWaveformError(f'No traces found for family {family.number}')
     return st