From 04b8d269d1013878403bd191e549f842bf7cae5b Mon Sep 17 00:00:00 2001
From: "Fabian P. Schmidt" <kerel-fs@gmx.de>
Date: Fri, 9 Nov 2018 03:13:44 +0100
Subject: [PATCH] contrib: Add python script 'plot_spectrum'

---
 contrib/plot_spectrum.py   | 16 ++++++++++++
 contrib/spectrum_parser.py | 52 ++++++++++++++++++++++++++++++++++++++
 contrib/spectrum_plot.py   | 44 ++++++++++++++++++++++++++++++++
 3 files changed, 112 insertions(+)
 create mode 100755 contrib/plot_spectrum.py
 create mode 100644 contrib/spectrum_parser.py
 create mode 100644 contrib/spectrum_plot.py

diff --git a/contrib/plot_spectrum.py b/contrib/plot_spectrum.py
new file mode 100755
index 0000000..cb20a8e
--- /dev/null
+++ b/contrib/plot_spectrum.py
@@ -0,0 +1,16 @@
+#!/usr/bin/env python3
+import argparse
+
+from spectrum_parser import read_spectrum
+from spectrum_plot import plot_spectrum
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(description='Plot RF observations')
+    parser.add_argument('path',
+                        type=str,
+                        help='Path to the directory with the input files $PATH/x_???.bin')
+    args = parser.parse_args()
+
+    z, headers = read_spectrum(args.path)
+    plot_spectrum(z, headers)
diff --git a/contrib/spectrum_parser.py b/contrib/spectrum_parser.py
new file mode 100644
index 0000000..bab6852
--- /dev/null
+++ b/contrib/spectrum_parser.py
@@ -0,0 +1,52 @@
+from datetime import datetime
+import numpy as np
+import glob
+import os
+import re
+from pathlib import Path
+
+
+def read_spectrum(path):
+    # Get the number of files
+    filenames = glob.glob(os.path.join(path, '*.bin'))
+    datestr = Path(filenames[0]).stem.split('_')[0]
+
+    # Read first file to get the number of channels and number of "samples"
+    filename = os.path.join(path, '{:s}_{:06}.bin'.format(datestr, 0))
+    with open(filename, 'rb') as f:
+        header = parse_header(f.read(256))
+
+    zs = []
+    headers = []
+    for i_file in range(len(filenames)):
+        filename = os.path.join(path, '{:s}_{:06}.bin'.format(datestr, i_file))
+        # i_sub = 0
+        with open(filename, 'rb') as f:
+            next_header = f.read(256)
+            while(next_header):
+                headers.append(parse_header(next_header))
+                zs.append(np.fromfile(f, dtype=np.float32, count=400))
+                next_header = f.read(256)
+    return np.transpose(np.vstack(zs)), headers
+
+
+def parse_header(header_b):
+    # TODO. Support files with the additional fields
+    #       - NBITS
+    #       - MEAN
+    #       - RMS
+    # "HEADER\nUTC_START    %s\nFREQ         %lf Hz\nBW           %lf Hz\nLENGTH       %f s\nNCHAN        %d\nNSUB         %d\n"
+
+    header_s = header_b.decode('ASCII').strip('\x00')
+
+    regex = r"^HEADER\nUTC_START    (.*)\nFREQ         (.*) Hz\nBW           (.*) Hz\nLENGTH       (.*) s\nNCHAN        (.*)\nNSUB         (.*)\nEND\n$"
+    match = re.fullmatch(regex, header_s, re.MULTILINE)
+
+    utc_start = datetime.strptime(match.group(1), '%Y-%m-%dT%H:%M:%S.%f')
+
+    return {'utc_start': utc_start,
+            'freq': float(match.group(2)),
+            'bw': float(match.group(3)),
+            'length': float(match.group(4)),
+            'nchan': int(match.group(5)),
+            'nsub': int(match.group(6))}
diff --git a/contrib/spectrum_plot.py b/contrib/spectrum_plot.py
new file mode 100644
index 0000000..bddefd0
--- /dev/null
+++ b/contrib/spectrum_plot.py
@@ -0,0 +1,44 @@
+from datetime import datetime, timedelta
+
+import numpy as np
+import matplotlib.pyplot as plt
+import matplotlib.dates as mdates
+
+
+def plot_spectrum(z, headers):
+    z_mean = np.mean(z)
+    z_std = np.std(z)
+    zmin = z_mean - 2 * z_std
+    zmax = z_mean + 6 * z_std
+    
+    f_min = -headers[0]['bw']*1e-3
+    f_max = +headers[0]['bw']*1e-3
+    
+    utc_start = headers[0]['utc_start']
+    t_min = mdates.date2num(utc_start + timedelta(seconds=0))
+    t_max = mdates.date2num(utc_start + timedelta(seconds=z.shape[1]))
+    
+    fig, ax = plt.subplots(figsize=(12, 6))
+    
+    ax.imshow(z,
+              origin='lower',
+              aspect='auto',
+              extent=[t_min, t_max, f_min, f_max],
+              vmin=zmin,
+              vmax=zmax,
+              interpolation='None',
+              cmap='viridis')
+    
+    ax.xaxis_date()
+    date_format = mdates.DateFormatter('%H:%M:%S')
+    ax.xaxis.set_major_formatter(date_format)
+    
+    # diagonal xaxis labels
+    fig.autofmt_xdate()
+    
+    plt.xlabel('Time;  start: {:%Y-%m-%d %H:%M:%S}Z'.format(headers[0]['utc_start']))
+    plt.ylabel('Freqeuncy / kHz; center: {} MHz'.format(headers[0]['freq'] * 1e-6))
+
+    plt.tight_layout()
+
+    plt.show()