Add reflection table and interactive mode to xrd

beamtime/xrd/plot.py

@@ -5,8 +5,9 @@ import pandas as pd
 import numpy as np
 import math
 
-import beamtime.xrd as xrd
+import ipywidgets as widgets
 
+import beamtime.xrd as xrd
 import beamtime.auxillary as aux
 import beamtime.plotting as btp
 
@@ -18,42 +19,168 @@ def plot_diffractogram(plot_data, options={}):
     plot_data (dict): Must include path = string to diffractogram data, and plot_kind = (recx, beamline, image)'''
 
     # Update options
-    required_options = ['x_vals', 'y_vals', 'ylabel', 'xlabel', 'xunit', 'yunit', 'scatter', 'plot_kind', 'rc_params', 'format_params']
+    required_options = ['x_vals', 'y_vals', 'ylabel', 'xlabel', 'xunit', 'yunit', 'line', 'scatter', 
+    'reflections', 'plot_kind', 'palettes', 'interactive', 'rc_params', 'format_params']
 
     default_options = {
         'x_vals': '2th', 
         'y_vals': 'I',
         'ylabel': 'Intensity', 'xlabel': '2theta', 
         'xunit': 'deg', 'yunit': 'a.u.',
+        'line': True,
         'scatter': False,
+        'reflections': False,
+        'reflection_data': None, # Should be passed as a dictionary on the form {path: rel_path, colour: [r,g,b], min_alpha: 0-1]
         'plot_kind': None,
+        'palettes': [('qualitative', 'Dark2_8')],
+        'interactive': False,
         'rc_params': {},
-        'format_params': {}
+        'format_params': {},
         }
 
     options = aux.update_options(options=options, required_options=required_options, default_options=default_options)
 
 
+    if options['interactive']:
+        options['interactive'] = False
+        plot_diffractogram_interactive(plot_data=plot_data, options=options)
+        return
+    
+    # Make adjustments to parameters if reflections data is passed
+    if options['reflections']:
+        options['format_params']['nrows'] = 2
+
+        if not 'grid_ratio_height' in options['format_params'].keys():
+            options['format_params']['grid_ratio_height'] = [1,10]
+
+    else:
+        options['format_params']['nrows'] = 1
+
+
     # Prepare plot, and read and process data
+
     fig, ax = btp.prepare_plot(options=options)
+    
+    if options['reflections']:
+        reflection_ax = ax[0]
+        ax = ax[1]
+
+    colours = btp.generate_colours(options['palettes'])
+
+
     diffractogram = xrd.io.read_data(path=plot_data['path'], kind=plot_data['plot_kind'], options=options)
 
-    if options['scatter']:
-        ax.scatter(x= diffractogram[options['x_vals']], y = diffractogram[options['y_vals']])
-        
-        #diffractogram.plot(x=options['x_vals'], y=options['y_vals'], ax=ax, kind='scatter')
-    
-    else:
-        diffractogram.plot(x=options['x_vals'], y=options['y_vals'], ax=ax)
 
+    if options['line']:
+        diffractogram.plot(x=options['x_vals'], y=options['y_vals'], ax=ax, c=next(colours), zorder=1)
+    
+    if options['scatter']:
+        ax.scatter(x= diffractogram[options['x_vals']], y = diffractogram[options['y_vals']], c=[(1,1,1,0)], edgecolors=[next(colours)], linewidths=plt.rcParams['lines.markeredgewidth'], zorder=2) #, edgecolors=np.array([next(colours)]))
 
 
     fig, ax = btp.adjust_plot(fig=fig, ax=ax, options=options)
 
+    if options['reflections'] and options['reflection_data']:
+        plot_reflection_table(plot_data=options['reflection_data'], ax=reflection_ax, options=options)
+
 
     return diffractogram, fig, ax
 
 
+def plot_diffractogram_interactive(plot_data, options):
+
+    if options['reflections']:
+        w = widgets.interactive(btp.ipywidgets_update, func=widgets.fixed(plot_diffractogram), plot_data=widgets.fixed(plot_data), options=widgets.fixed(options), 
+        scatter=widgets.ToggleButton(value=False), 
+        line=widgets.ToggleButton(value=True), 
+        reflections=widgets.ToggleButton(value=True))
+    
+    else:
+        w = widgets.interactive(btp.ipywidgets_update, func=widgets.fixed(plot_diffractogram), plot_data=widgets.fixed(plot_data), options=widgets.fixed(options), 
+        scatter=widgets.ToggleButton(value=False), 
+        line=widgets.ToggleButton(value=True))
+    
+    
+    display(w)
+
+
+    
+def plot_reflection_table(plot_data, ax=None, options={}):
+    ''' Plots a reflection table from output generated by VESTA.
+    
+    Required contents of plot_data:
+    path (str): relative path to reflection table file'''
+
+    required_options = ['reflections_colour', 'min_alpha', 'format_params', 'rc_params']
+
+    default_options = {
+        'reflections_colour': [0,0,0],
+        'min_alpha': 0,
+        'format_params': {},
+        'rc_params': {}
+    }
+
+    if 'colour' in plot_data.keys():
+        options['reflections_colour'] = plot_data['colour']
+    if 'min_alpha' in plot_data.keys():
+        options['min_alpha'] = plot_data['min_alpha']
+
+
+    options = aux.update_options(options=options, required_options=required_options, default_options=default_options)
+
+    
+    
+    if not ax:
+        _, ax = btp.prepare_plot(options)
+
+    reflection_table = load_reflection_table(plot_data['path'])
+
+    reflections, intensities  = reflection_table['2th'], reflection_table['I']
+
+    for ref, intensity in zip(reflections, intensities):
+
+        rel_intensity = (intensity / intensities.max())*(1-options['min_alpha']) + options['min_alpha']
+
+
+        ax.axvline(x=ref, c=options['reflections_colour'], alpha=rel_intensity)
+        
+
+
+    ax.tick_params(which='both', bottom=False, labelbottom=False, right=False, labelright=False, left=False, labelleft=False, top=False, labeltop=False)
+    
+    if options['xlim']:
+        ax.set_xlim(options['xlim'])
+
+
+
+
+def load_reflection_table(path):
+
+    # VESTA outputs the file with a header that has a space between the parameter and units - so there is some extra code to rectify the issue
+    # that ensues from this formatting
+    reflections = pd.read_csv(path, delim_whitespace=True)
+
+    # Remove the extra column that appears from the headers issue
+    reflections.drop(reflections.columns[-1], axis=1, inplace=True)
+
+    with open(path, 'r') as f:
+        line = f.readline()
+
+        headers = line.split()
+
+        # Delete the fourth element which is '(Å)'
+        del headers[4]
+
+        # Change name of column to avoid using greek letters
+        headers[7] = '2th'
+
+    # Set the new modified headers as the headers of 
+    reflections.columns = headers
+
+    return reflections
+
+    
+
 
 def prepare_diffractogram_plot(options=None):
     # First take care of the options for plotting - set any values not specified to the default values