Merge pull request #222 from texadactyl/master

New blimpy utility, stix (Issue #221)

VERSION-HISTORY.md

@@ -3,6 +3,7 @@ This file is a version history of turbo_seti amendments, beginning with version
 <br>
 |    Date    | Version | Contents |
 | :--: | :--: | :-- |
+| 2021-08-06 | 2.0.20 | New utility, "stix" (Issue #221).  |
 | 2021-07-28 | 2.0.19 | Update fil2h5 to handle YUGE data matrixes.  |
 | 2021-07-17 | 2.0.18 | Get rid of numpy "RuntimeWarning: Mean of empty slice" messages (Issue #212).  |
 | 2021-07-13 | 2.0.17 | New utility: stax.  |

blimpy/__init__.py

@@ -11,6 +11,7 @@
     from . import calcload
     from . import rawhdr
     from . import stax
+    from . import stix
     from . import match_fils
     from blimpy.io import file_wrapper
 except:

blimpy/stix.py

@@ -0,0 +1,229 @@
+r"""
+Make waterfall plots of a large file.
+"""
+
+import os
+import sys
+from os.path import dirname, abspath, isdir
+import gc
+from argparse import ArgumentParser
+from PIL import Image
+import logging
+logger_name = "stix"
+logger = logging.getLogger(logger_name)
+logger.setLevel(logging.INFO)
+
+# Plotting packages import
+import matplotlib
+import matplotlib.pyplot as plt
+matplotlib.use("agg")
+
+# Blimpy imports
+from blimpy import Waterfall
+from blimpy.plotting import plot_waterfall
+
+
+def image_stitch(orientation, chunk_count, png_collection, path_saved_png):
+    r""" Stitch together multiple PNGs into one"""
+
+    logger.info("Stitching together the images, orientation is {}".format(orientation))
+
+    # Empty array to store loaded images.
+    img_collection = []
+
+    # Loop to load images.
+    for ii in range(chunk_count):
+        img = Image.open(png_collection[ii])
+        if ii == 0:
+            base_width, base_height = img.size
+            logger.info("Base (first) image file pixel width = {}, height = {}"
+                        .format(base_width, base_height))
+        width, height = img.size
+        if width != base_width or height != base_height:
+            logger.warning("Skipping {} because pixel width = {}, height = {}, doesn't match the base!"
+                           .format(png_collection[ii], width, height))
+            continue
+        img_collection.append(img)
+
+    # Get dimensions of the last image (same as the rest).
+
+    # Creating the template for the stitched image.
+    if orientation == "h":
+        new_img = Image.new("RGBA", (chunk_count * width, height))
+    else:
+        new_img = Image.new("RGBA", (width, chunk_count * height))
+
+    # Initalising the template before pasting the images.
+    init_width = 0
+    init_height = 0
+
+    # Pasting the images in sequence.
+    counter = 0
+    for img in img_collection:
+        box = (init_width, init_height, init_width + width, init_height + height)
+        try:
+            new_img.paste(img, box)
+        except ValueError:
+            logger.error("stix: *** Oops, Image paste error, init_width={}, init_height={}, width={}, height={}"
+                         .format(init_width, init_height, width, height))
+            sys.exit(86)
+        counter += 1
+        if orientation == "h":
+            init_width += width
+        else:
+            init_height += height
+
+    # Save the output, re-sized.
+    if orientation == "h":
+        new_img = new_img.resize((chunk_count * width, height))
+    else:
+        new_img = new_img.resize((width, chunk_count * height))
+    new_img.save(path_saved_png)
+    logger.info("Concatenated {} images to {}".format(counter, path_saved_png))
+
+
+def sort2(x, y):
+    r""" Return lowest value, highest value"""
+    if y < x:
+        return y, x
+    return x, y
+
+
+def make_waterfall_plots(input_file, chunk_count, plot_dir, width, height, dpi, source_name=None):
+    r"""
+    Make waterfall plots of a given huge-ish file.
+​
+    Parameters
+    ----------
+    input_file : str
+        Path of Filterbank or HDF5 input file to plot in a stacked mode.
+    chunk_count : int
+        The number of chunks to divide the entire bandwidth into.
+    plot_dir : str
+        Directory for storing the PNG files.
+    """
+
+    # Get directory path for storing PNG file
+    if plot_dir is None:
+        dirpath = dirname(abspath(input_file)) + "/"
+    else:
+        if not isdir(plot_dir):
+            os.mkdir(plot_dir)
+        dirpath = plot_dir
+
+    # Calculate frequency boundary variables.
+    wf1 = Waterfall(input_file, load_data=False)
+    nf = wf1.n_channels_in_file
+    if nf % chunk_count != 0:
+        msg = "Number of frequency chunks ({}) does not evenly divide the number of frequencies ({})!" \
+              .format(chunk_count, nf)
+        logger.warning(msg)
+    fchunk_f_start = wf1.file_header["fch1"]
+    fchunk_size = int(nf / chunk_count)
+    fchunk_f_offset = fchunk_size * wf1.file_header["foff"]
+    fchunk_f_stop = fchunk_f_start + (fchunk_size - 1) * wf1.file_header["foff"]
+
+    # Produce the source_name to be used in the image title.
+    if source_name is None:
+        source_name = wf1.header["source_name"]
+        if source_name.upper() == "UNKNOWN":
+            source_name = wf1.header["rawdatafile"].replace(".0000.raw", "")
+
+    # Begin PNG file collection.
+    png_collection = []
+
+    # Generate a plot/PNG for each frequency chunk.
+    logger.info("Image width = {}, height = {}, dpi = {}"
+                .format(width, height, dpi))
+    for ii in range(0, chunk_count):
+
+        ii_lowest, ii_highest = sort2(fchunk_f_start, fchunk_f_stop)
+
+        # read in data
+        wf = Waterfall(input_file, f_start=ii_lowest, f_stop=ii_highest)
+
+        # Validate frequency range.
+        logger.debug("stix: Processing chunk {} of {}, frequency lowest={}, highest={}"
+                    .format(ii, chunk_count - 1, ii_lowest, ii_highest))
+
+        # Make plot for this frequency chunk with plot_waterfall().
+        wf.header["source_name"] = source_name + " chunk " + str(ii + 1) + " of " + str(chunk_count)
+        plt.figure(wf.header["source_name"], figsize=(width, height), dpi=dpi)
+        plot_waterfall(wf, f_start=ii_lowest, f_stop=ii_highest)
+
+        # Save the figures.
+        path_png = dirpath + source_name + "_chunk_{}".format(ii + 1) + ".png"
+        png_collection.append(path_png)
+        plt.savefig(path_png, dpi=dpi)
+        logger.info("Saved plot: {}".format(path_png))
+
+        # Delete Waterfall object.
+        del wf
+        gc.collect()
+        plt.close("all")
+
+        # Prepare for next chunk.
+        fchunk_f_start += fchunk_f_offset
+        fchunk_f_stop += fchunk_f_offset
+
+    return png_collection, source_name
+
+
+def cmd_tool(args=None):
+    r"""Coomand line parser"""
+    parser = ArgumentParser(description="Make waterfall plots from a single file.")
+    parser.add_argument("input_file", type=str, default=None, help="Input file to plot")
+    parser.add_argument("chunk_count", type=int, default=0, help="Count of same-sized frequency chunks")
+    parser.add_argument("--plot_dir", "-p", type=str, default=".", help="Directory to receive the plot (.png).  Default: current directory.")
+    parser.add_argument("--source_name", "-n", type=str, default=None, help="Source name to override the header field.  Default: None.")
+    parser.add_argument("--stitch", "-s", type=str, default="n", choices=["n", "v", "h"],
+                        help="Stitch files: n(no stitching), v(vertical), or h(horizontal).  Default: n.")
+    parser.add_argument("--dpi", "-d", type=int, default=200,
+                        help="Single image file dots per inch.  Default: 200.")
+    parser.add_argument("--width", "-w", type=float, default=10,
+                        help="Single image file image width in inches.  Default: 10.0.")
+    parser.add_argument("--height", "-t", type=float, default=8,
+                        help="Single image file image height in inches.  Default: 8.0.")
+
+    if args is None:
+        args = parser.parse_args()
+    else:
+        args = parser.parse_args(args)
+    if args.input_file is None or args.chunk_count < 1:
+        os.system("stix -h")
+        sys.exit(0)
+    if args.dpi < 50:
+        logger.error("stix: --dpi must be at least 50 but I saw: {}!".format(args.dpi))
+        os.system("stix -h")
+        sys.exit(86)
+    if args.width < 6:
+        logger.error("stix: --width must be at least 6 but I saw: {}!".format(args.width))
+        os.system("stix -h")
+        sys.exit(86)
+    if args.height < 5:
+        logger.error("stix: --height must be at least 5 but I saw: {}!".format(args.height))
+        os.system("stix -h")
+        sys.exit(86)
+
+    # Make the plots.
+    logger.info("Begin")
+    args.plot_dir += "/"
+    args.input_file = os.path.abspath(args.input_file)
+    png_collection, source_name = make_waterfall_plots(args.input_file,
+                                                       args.chunk_count,
+                                                       args.plot_dir,
+                                                       args.width,
+                                                       args.height,
+                                                       args.dpi,
+                                                       source_name=args.source_name)
+    if args.stitch != "n":
+        path_saved_png = args.plot_dir + source_name + ".png"
+        image_stitch(args.stitch, args.chunk_count, png_collection, path_saved_png)
+        for file in png_collection:
+            os.remove(file)
+
+    logger.info("End")
+
+
+if __name__ == "__main__":
+    cmd_tool()

requirements.txt

@@ -8,4 +8,5 @@ h5py
 scipy
 hdf5plugin
 pandas
+Pillow
 psutil

setup.py

@@ -3,7 +3,7 @@
 """
 from setuptools import setup, find_packages
 
-__version__ = '2.0.19'
+__version__ = '2.0.20'
 
 with open("README.md", "r") as fh:
     long_description = fh.read()

tests/test_stix.py

@@ -0,0 +1,34 @@
+from os.path import dirname
+from blimpy.stix import cmd_tool
+from tests.data import voyager_fil
+
+
+PLOT_DIR = dirname(voyager_fil)
+
+
+def execute_command(args):
+    print("\ntest_stix: args:", args)
+    cmd_tool(args)
+
+
+def test_stix():
+
+    args = [voyager_fil, "16", "--plot_dir",  PLOT_DIR]
+    execute_command(args)
+
+    args = [voyager_fil, "4", "--plot_dir",  PLOT_DIR, "-s", "v"]
+    execute_command(args)
+
+    args = [voyager_fil, "4", "-p",  PLOT_DIR, "--stitch", "h"]
+    execute_command(args)
+
+    args = [voyager_fil, "4", "--plot_dir",  PLOT_DIR, "--stitch", "n",
+            "--dpi", "100", "--width", "8", "--height", "6"]
+    execute_command(args)
+
+    args = [voyager_fil, "4", "-p",  PLOT_DIR, "-s", "n", "-d", "100", "-w", "8", "-t", "6"]
+    execute_command(args)
+
+
+if __name__ == "__main__":
+    test_stix()