Add plotting, saving options

diffpy/pdfmorph/pdfmorphapp.py

@@ -71,13 +71,13 @@ def custom_error(self, msg):
         action="store_true",
         help=f"""Changes usage to \'{prog_short} [options] FILE DIRECTORY\'. FILE
 will be morphed with each file in DIRECTORY as target.
-Files in directory are sorted by alphabetical order unless
---temperature is enabled. Plotting and saving are disabled
-when this option is enabled.""",
+Files in DIRECTORY are sorted by alphabetical order unless
+--temperature is enabled. Plotting and saving options are for
+a Rw plot and table respectively.""",
     )
     parser.add_option(
         '--temperature',
-        dest="temperature_sort",
+        dest="temp",
         action="store_true",
         help="""Used with --sequence to sort files in DIRECTORY by temperature.
 File names in DIRECTORY should end in _#K.gr or _#K.cgr
@@ -240,8 +240,6 @@ def main():
     parser = create_option_parser()
     (opts, pargs) = parser.parse_args()
     if opts.sequence:
-        opts.plot = False  # Disable plotting
-        opts.savefile = None  # Disable saving
         multiple_morphs(parser, opts, pargs, stdout_flag=True)
     else:
         single_morph(parser, opts, pargs, stdout_flag=True)
@@ -264,32 +262,48 @@ def multiple_morphs(parser, opts, pargs, stdout_flag):
 
     # Sort files in directory
     target_list = list(target_directory.iterdir())
-    if opts.temperature_sort:
+    if opts.temp:
         # Sort by temperature
-        target_list = tools.temperature_sort(target_list)
+        try:
+            target_list = tools.temperature_sort(target_list)
+        except ValueError:
+            parser.custom_error("All file names in directory must end in _###K.gr or _###K.cgr to use "
+                                "the --temperature option.")
     else:
         # Default is alphabetical sort
         target_list.sort()
 
+    # Disable single morph plotting and saving
+    plot_opt = opts.plot
+    opts.plot = False
+    save_opt = opts.savefile
+    opts.savefile = None
+
+    # Do not morph morph_file against itself if it is in the same directory
+    if morph_file in target_list:
+        target_list.remove(morph_file)
+
     # Morph morph_file against all other files in target_directory
     results = []
     for target_file in target_list:
-        # Only morph morph_file against different files in target_directory
-        if target_file.is_file and morph_file != target_file:
+        if target_file.is_file:
             results.append([
                 target_file.name,
                 single_morph(parser, opts, [morph_file, target_file], stdout_flag=False),
             ])
 
+    # Input parameters used for every morph
+    inputs = [None, None]
+    inputs[0] = f"# Morphed file: {morph_file.name}"
+    inputs[1] = "\n# Input morphing parameters:"
+    inputs[1] += f"\n# scale = {opts.scale}"
+    inputs[1] += f"\n# stretch = {opts.stretch}"
+    inputs[1] += f"\n# smear = {opts.smear}"
+
     # If print enabled
     if stdout_flag:
-        # Input parameters used for every morph
-        inputs = "\n# Input morphing parameters:"
-        inputs += f"\n# scale = {opts.scale}"
-        inputs += f"\n# stretch = {opts.stretch}"
-        inputs += f"\n# smear = {opts.smear}"
-
-        print(inputs)
+        # Separated for saving
+        print(f"\n{inputs[0]}{inputs[1]}")
 
         # Results from each morph
         for entry in results:
@@ -298,6 +312,49 @@ def multiple_morphs(parser, opts, pargs, stdout_flag):
             outputs += "\n".join(f"# {i[0]} = {i[1]:.6f}" for i in entry[1])
             print(outputs)
 
+    rws = []
+    target_labels = []
+    results_length = len(results)
+    for entry in results:
+        if opts.temp:
+            name = entry[0]
+            target_labels.append(tools.extract_temperatures([name])[0])
+        else:
+            target_labels.append(entry[0])
+        for item in entry[1]:
+            if item[0] == "Rw":
+                rws.append(item[1])
+
+    if save_opt is not None:
+        # Save table of Rw values
+        try:
+            with open(save_opt, 'w') as outfile:
+                # Header
+                print(f"{inputs[0]}\n{inputs[1]}", file=outfile)
+                if opts.temp:
+                    print(f"\n# L T(K) Rw", file=outfile)
+                else:
+                    print(f"\n# L Target Rw", file=outfile)
+
+                # Table
+                for idx in range(results_length):
+                    print(f"{target_labels[idx]} {rws[idx]}", file=outfile)
+                outfile.close()
+
+                # Output to stdout
+                path_name = Path(outfile.name).absolute()
+                save_message = f"\n# Rw table saved to {path_name}"
+                print(save_message)
+
+        # Save failed
+        except FileNotFoundError as e:
+            save_fail_message = "\nUnable to save to designated location"
+            print(save_fail_message)
+            parser.custom_error(str(e))
+
+    if plot_opt:
+        pdfplot.plot_rws(target_labels, rws, opts.temp)
+
     return results
 
 

diffpy/pdfmorph/pdfplot.py

@@ -15,6 +15,7 @@
 """Collection of plotting functions for PDFs."""
 
 import matplotlib.pyplot as plt
+import diffpy.pdfmorph.tools as tools
 from bg_mpl_stylesheet.bg_mpl_stylesheet import bg_mpl_style
 import numpy
 
@@ -207,6 +208,52 @@ def comparePDFs(
     return
 
 
+def plot_rws(target_labels, rws, temp_flag=False):
+    """
+    Plot Rw values for multiple morphs.
+    :param target_labels: Names (or temperatures if --temperature is enabled) of each file acting as target for the morph.
+    :type target_labels: list
+    :param rws: Contains the Rw values corresponding to each file.
+    :type rws: list
+    :param temp_flag: When True, temperature is extracted from file names in target_files. Then, a line chart of
+    Rw versus temperature is made. Otherwise (default), it plots a bar chart of Rw values per file.
+    :type temp_flag: bool
+    """
+
+    # If we can extract temperature data
+    if temp_flag:
+        temps = target_labels
+
+        # Plot Rw vs Temperature
+        plt.plot(temps, rws, linestyle='-', marker='o')
+        plt.ylabel(r"$R_w$")
+        plt.xlabel(r"Temperature ($K$)")
+        plt.minorticks_on()
+
+    # Create bar chart for each file
+    else:
+        file_names = target_labels
+        # Ensure file names do not crowd
+        bar_size = 5
+        max_len = bar_size
+        for name in file_names:
+            max_len = max(max_len, len(name))
+        angle = numpy.arccos(bar_size / max_len)
+        angle *= 180 / numpy.pi  # Convert to degrees
+        plt.xticks(rotation=angle)
+
+        # Plot Rw for each file
+        plt.bar(file_names, rws)
+        plt.ylabel(r"$R_w$")
+        plt.xlabel(r"Target File")
+
+    # Show plot
+    plt.tight_layout()
+    plt.show()
+
+    return
+
+
 def truncatePDFs(r, gr, rmin=None, rmax=None):
     """Truncate a PDF to specified bounds.
 

diffpy/pdfmorph/tools.py

@@ -113,21 +113,36 @@ def readPDF(fname):
 
 
 def nn_value(val, name):
-    # Convenience function for ensuring certain non-negative inputs
+    """Convenience function for ensuring certain non-negative inputs."""
     if val < 0:
         negative_value_warning = f"\n# Negative value for {name} given. Using absolute value instead."
         print(negative_value_warning)
         return -val
     return val
 
 
-def temperature_sort(filenames):
+def temperature_sort(filepaths):
+    """Sort a list of files by temperatures encoded in their name. Files should all end in _###K.gr or _###K.cgr."""
+
+    # Get temperature from file names
+    filenames = []
+    for path in filepaths:
+        filenames.append(path.name)
+    temps = extract_temperatures(filenames)
+
     # Sort files (whose paths are contained in filenames) ending in _###K.gr/_###K.cgr by ###
-    for idx in range(len(filenames)):
-        filename = filenames[idx].name
-        s_index = filename.rfind("_")  # Start of temperature value
-        e_index = filename.rfind("K")  # End of temperature value
-        temp = float(filename[s_index + 1: e_index])
-        filenames[idx] = [filenames[idx], temp]
-    filenames.sort(key=lambda entry: entry[1])
-    return [entry[0] for entry in filenames]
+    for idx in range(len(filepaths)):
+        filepaths[idx] = [filepaths[idx], temps[idx]]
+    filepaths.sort(key=lambda entry: entry[1])
+    return [entry[0] for entry in filepaths]
+
+
+def extract_temperatures(filenames):
+    """Convenience function to extract temperatures from file names."""
+
+    temps = []
+    for name in filenames:
+        s_index = name.rfind("_")  # Start of temperature value
+        e_index = name.rfind("K")  # End of temperature value
+        temps.append(float(name[s_index + 1: e_index]))
+    return temps