Merge pull request #3 from rasmusvt/rasmus_heatmap

Rasmus heatmap

beamtime/auxillary.py

@@ -1,4 +1,5 @@
 import json
+import numpy as np
 
 def update_options(options, required_options, default_options):
     ''' Takes a dictionary of options along with a list of required options and dictionary of default options, and sets all keyval-pairs of options that is not already defined to the default values'''
@@ -37,6 +38,18 @@ def swap_values(dict, key1, key2):
 
 
 
-def hello_world2(a=1, b=2):
+def ceil(a, roundto=1):
 
-    print(f'Halla, MAFAKKAS! a = {a} og b = {b}')
+    fac = 1/roundto
+
+    a = np.ceil(a*fac) / fac
+
+    return a
+
+def floor(a, roundto=1):
+
+    fac = 1/roundto
+
+    a = np.floor(a*fac) / fac
+
+    return a

beamtime/plotting.py

@@ -182,7 +182,7 @@ def adjust_plot(fig, ax, options):
         ax.xaxis.set_minor_locator(MultipleLocator(options['x_tick_locators'][1]))
 
     
-    # THIS NEEDS REWORK FOR IT TO FUNCTION PROPERLY!
+    # FIXME THIS NEEDS REWORK FOR IT TO FUNCTION PROPERLY!
     if options['xticks']:
         ax.set_xticks(np.arange(plot_data['start'], plot_data['end']+1))
         ax.set_xticklabels(options['xticks'])

beamtime/xrd/io.py

@@ -40,7 +40,7 @@ def integrate_1d(data, options={}, index=0):
     df: DataFrame contianing 1D diffractogram if option 'return' is True
     ''' 
 
-    required_options = ['unit', 'nbins', 'save', 'save_filename', 'save_extension', 'save_folder', 'overwrite']
+    required_options = ['unit', 'nbins', 'save', 'save_filename', 'save_extension', 'save_folder', 'overwrite', 'extract_folder']
 
     default_options = {
         'unit': '2th_deg', 
@@ -358,9 +358,18 @@ def read_data(data, options={}, index=0):
 
 def apply_offset(diffractogram, wavelength, index, options):
 
-    options['current_offset_y'] = options['offset_y']
+    if 'current_offset_y' not in options.keys():
+        options['current_offset_y'] = options['offset_y']
+    else:
+        if options['current_offset_y'] != options['offset_y']:
+            options['offset_change'] = True
+        
+        options['current_offset_y'] = options['offset_y']
+
     options['current_offset_x'] = options['offset_x']
 
+    
+
     #Apply offset along y-axis
     diffractogram['I'] = diffractogram['I_org'] # Reset intensities
 
@@ -391,7 +400,7 @@ def revert_offset(diffractogram,which=None):
 
     return diffractogram
 
-def load_reflection_table(data, options={}):
+def load_reflection_table(data: dict, reflections_params: dict, options={}):
 
     required_options = ['ref_wavelength', 'to_wavelength']
 
@@ -404,12 +413,12 @@ def load_reflection_table(data, options={}):
 
     # 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(data['path'], delim_whitespace=True)
+    reflections = pd.read_csv(reflections_params['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(data['path'], 'r') as f:
+    with open(reflections_params['path'], 'r') as f:
         line = f.readline()
 
         headers = line.split()
@@ -425,13 +434,28 @@ def load_reflection_table(data, options={}):
 
     reflections = translate_wavelengths(data=reflections, wavelength=options['ref_wavelength'], to_wavelength=options['to_wavelength'])
 
-    #print(reflections)
+    if 'heatmap' in data.keys():
+        
+        start_2th, stop_2th = data['diffractogram'][0]['2th'].min(), data['diffractogram'][0]['2th'].max()
+        len_2th = stop_2th - start_2th 
+        #print(start_2th, stop_2th, len_2th)
+        
+        start_heatmap, stop_heatmap = 0, data['heatmap'].shape[1]
+        len_heatmap = stop_heatmap - start_heatmap
+        #print(start_heatmap, stop_heatmap, len_heatmap)
+
+        scale = len_heatmap/len_2th
+
+        #print(scale)
+        #print(stop_2th * scale)
+
+        reflections['heatmap'] = (reflections['2th']-start_2th) * scale
 
     return reflections
 
 
 
-def translate_wavelengths(data, wavelength, to_wavelength=None):
+def translate_wavelengths(data: pd.DataFrame, wavelength: float, to_wavelength=None) -> pd.DataFrame:
     # FIXME Somewhere here there is an invalid arcsin-argument. Not sure where.
 
     pd.options.mode.chained_assignment = None

beamtime/xrd/plot.py

@@ -1,3 +1,4 @@
+import seaborn as sns
 import matplotlib.pyplot as plt
 from matplotlib.ticker import (MultipleLocator, FormatStrFormatter,AutoMinorLocator)
 
@@ -19,8 +20,8 @@ def plot_diffractogram(data, options={}):
     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', 'line', 'scatter', 'xlim', 'ylim', 'normalise', 'offset', 'offset_x', 'offset_y',
+    required_options = ['x_vals', 'y_vals', 'ylabel', 'xlabel', 'xunit', 'yunit', 'line', 'scatter', 'xlim', 'ylim', 'normalise', 'offset', 'offset_x', 'offset_y', 'offset_change',
-    'reflections_plot', 'reflections_indices', 'reflections_data', 'plot_kind', 'palettes', 'interactive', 'rc_params', 'format_params', 'interactive_session_active']
+    'reflections_plot', 'reflections_indices', 'reflections_data', 'heatmap', 'cmap', 'plot_kind', 'palettes', 'interactive', 'rc_params', 'format_params', 'interactive_session_active']
 
     default_options = {
         'x_vals': '2th', 
@@ -32,11 +33,14 @@ def plot_diffractogram(data, options={}):
         'offset': True,
         'offset_x': 0,
         'offset_y': 1,
+        'offset_change': False,
         'line': True, # whether or not to plot diffractogram as a line plot
         'scatter': False, # whether or not to plot individual data points
         'reflections_plot': False, # whether to plot reflections as a plot
         'reflections_indices': False, # whether to plot the reflection indices
         'reflections_data': None, # Should be passed as a list of dictionaries on the form {path: rel_path, reflection_indices: number of indices, colour: [r,g,b], min_alpha: 0-1]
+        'heatmap': False,
+        'cmap': 'viridis',
         'plot_kind': None,
         'palettes': [('qualitative', 'Dark2_8')],
         'interactive': False,
@@ -50,6 +54,7 @@ def plot_diffractogram(data, options={}):
             default_options['offset_y'] = 0.05
 
     options = aux.update_options(options=options, required_options=required_options, default_options=default_options)
+    #options['current_offset_y'] = options['offset_y']
 
     # Convert data['path'] to list to allow iteration over this to accommodate both single and multiple diffractograms
     if not isinstance(data['path'], list):
@@ -69,16 +74,22 @@ def plot_diffractogram(data, options={}):
             data['diffractogram'][index] = diffractogram
             data['wavelength'][index] = wavelength
 
-            
+        # Sets the xlim if this has not bee specified
+        if not options['xlim']:
+            options['xlim'] = [data['diffractogram'][0][options['x_vals']].min(), data['diffractogram'][0][options['x_vals']].max()]
+
+        # Generate heatmap data
+        data['heatmap'], data['heatmap_xticks'], data['heatmap_xticklabels'] = generate_heatmap(data=data, options=options)
+        options['heatmap_loaded'] = True
+
+        if options['heatmap']:
+            options['xlim'] = options['heatmap_xlim']
 
     else:
         if not isinstance(data['diffractogram'], list):
             data['diffractogram'] = [data['diffractogram']]
             data['wavelength'] = [data['wavelength']]
 
-    # Sets the xlim if this has not bee specified
-    if not options['xlim']:
-        options['xlim'] = [diffractogram[options['x_vals']].min(), diffractogram[options['x_vals']].max()]
 
 
     if options['interactive_session_active']:
@@ -129,13 +140,19 @@ def plot_diffractogram(data, options={}):
     else:
         colours = btp.generate_colours(['black'], kind='single')
 
+    if options['heatmap']:
+        sns.heatmap(data['heatmap'], cmap=options['cmap'], cbar=False, ax=ax)
+        ax.set_xticks(data['heatmap_xticks'][options['x_vals']])
+        ax.set_xticklabels(data['heatmap_xticklabels'][options['x_vals']])
+        ax.tick_params(axis='x', which='minor', bottom=False, top=False)
 
-    for diffractogram in data['diffractogram']:
+    else:
-        if options['line']:
+        for diffractogram in data['diffractogram']:
-            diffractogram.plot(x=options['x_vals'], y=options['y_vals'], ax=ax, c=next(colours), zorder=1)
+            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)]))
+            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)]))
 
 
 
@@ -148,27 +165,102 @@ def plot_diffractogram(data, options={}):
         options['xlim'] = ax.get_xlim()
         options['to_wavelength'] = data['wavelength'][0]
         
-        for reference, axis in zip(options['reflections_data'], ref_axes):
+        for reflections_params, axis in zip(options['reflections_data'], ref_axes):
-            plot_reflection_table(data=reference, ax=axis, options=options)
+            plot_reflection_table(data=data, reflections_params=reflections_params, ax=axis, options=options)
 
     # Print the reflection indices. By default, the wavelength of the first diffractogram will be used for this.
     if options['reflections_indices'] and options['reflections_data']:
         options['xlim'] = ax.get_xlim()
         options['to_wavelength'] = data['wavelength'][0]
 
-        for reference in options['reflections_data']:
+        for reflections_params in options['reflections_data']:
-            plot_reflection_indices(data=reference, ax=indices_ax, options=options)
+            plot_reflection_indices(data=data, reflections_params=reflections_params, ax=indices_ax, options=options)
 
 
     if options['interactive_session_active']:
-        btp.update_widgets(options=options)
+        options['current_y_offset'] = options['widget'].kwargs['offset_y']
+        update_widgets(data=data, options=options)
 
+
+
+    return data['diffractogram'], fig, ax
+
+
+
+def generate_heatmap(data, options={}):
+
+    required_options = ['x_tick_locators']
+
+    default_options = {
+        'x_tick_locators': [0.5, 0.1]
+    }
+
+    options = aux.update_options(options=options, required_options=required_options, default_options=default_options)
+
+    twotheta = []
+    intensities = []
+    scans = []
+
+    for i, d in enumerate(data['diffractogram']):
+        twotheta = np.append(twotheta, d['2th'].to_numpy())
+        intensities = np.append(intensities, d['I'].to_numpy())
+        scans = np.append(scans, np.full(len(d['2th'].to_numpy()), int(i)))
+
+
+    heatmap = pd.DataFrame({'2th': twotheta, 'scan': scans, 'I': intensities})
+    xrd.io.translate_wavelengths(data=heatmap, wavelength=data['wavelength'][0])
+    min_dict = {'2th': heatmap['2th'].min(), '2th_cuka': heatmap['2th_cuka'].min(), '2th_moka': heatmap['2th_moka'].min(),
+                        'q': heatmap['q'].min(), 'q2': heatmap['q2'].min(), 'q4': heatmap['q4'].min(), '1/d': heatmap['1/d'].min()}
+
+    max_dict = {'2th': heatmap['2th'].max(), '2th_cuka': heatmap['2th_cuka'].max(), '2th_moka': heatmap['2th_moka'].max(),
+                        'q': heatmap['q'].max(), 'q2': heatmap['q2'].max(), 'q4': heatmap['q4'].max(), '1/d': heatmap['1/d'].max()}
+
+
+    ndatapoints = len(data['diffractogram'][0]['2th'])
+
+    xlims = [0, ndatapoints, 0, ndatapoints] # 0: xmin, 1: xmax, 2: xmin_start, 3: xmax_start
+    xticks = {}
+    xticklabels = {}
+
+    for xval in min_dict.keys():
+       
+       # Add xticks labels
         
+        label_max = aux.floor(max_dict[xval], roundto=options['x_tick_locators'][0])
+        label_min = aux.ceil(min_dict[xval], roundto=options['x_tick_locators'][0])
+        label_steps = (label_max - label_min)/options['x_tick_locators'][0]
+
+        xticklabels[xval] = np.linspace(label_min, label_max, num=int(label_steps)+1)
+
+        # Add xticks
+        xval_span = max_dict[xval] - min_dict[xval]
+        steps = xval_span / ndatapoints
+        
+        
+        xticks_xval = []
+
+        for tick in xticklabels[xval]:
+            xticks_xval.append((tick-min_dict[xval])/steps)
+
+        xticks[xval] = xticks_xval
+
+
+    options['x_tick_locators'] = None
+
+    heatmap = heatmap.reset_index().pivot_table(index='scan', columns='2th', values='I')
+
+    options['heatmap_xlim'] = xlims
+
+
+    return heatmap, xticks, xticklabels
 
 
 
 
-    return diffractogram, fig, ax
+
+
+
+# #results = np.transpose(np.vstack([twotheta, scans, intensities]))
 
 
 def determine_grid_layout(options):
@@ -187,57 +279,76 @@ def determine_grid_layout(options):
 
 
 
+
+
 def plot_diffractogram_interactive(data, options):
 
 
-    minmax = {'2th': [None, None], '2th_cuka': [None, None], '2th_moka': [None, None], 'd': [None, None], '1/d': [None, None], 'q': [None, None], 'q2': [None, None], 'q4': [None, None]}
+    # Format here is xminmax[0]: xmin, xminmax[1]: xmax, xminmax[2]: xmin_start, xminmax[3]: xmax_start, where "_start" denotes starting value of the slider
+    xminmax = { '2th':      [None, None, None, None], '2th_cuka':   [None, None, None, None],   '2th_moka':     [None, None, None, None], 
+                'd':        [None, None, None, None], '1/d':        [None, None, None, None], 
+                'q':        [None, None, None, None], 'q2':         [None, None, None, None],   'q4':           [None, None, None, None], 
+                'heatmap':  [None, None, None, None], 'start':      [None, None, None, None]}
+
+    yminmax = { 'diff':     [None, None, None, None], 'heatmap':    [None, None, None, None],   'start':        [None, None, None, None]}
     
-    update_minmax(minmax, data)
+    update_xminmax(xminmax=xminmax, data=data, options=options)
+    update_yminmax(yminmax=yminmax, data=data, options=options)
 
-    ymin, ymax = None, None
+    options['xminmax'], options['yminmax'] = xminmax, yminmax
-    for index, diffractogram in enumerate(data['diffractogram']):
-        if not ymin or (ymin > (diffractogram['I'].min())): #+index*options['offset_y'])):
-            ymin = diffractogram['I'].min()#+index*options['offset_y']
 
-        if not ymax or (ymax < (diffractogram['I'].max())):#+index*options['offset_y'])):
+    # Get start values for ylim slider based on choice (FIXME This can be impleneted into update_yminmax). Can also make a 'start' item that stores the start values, instead of having 4 items in 'diff' as it is now.
-            ymax = diffractogram['I'].max()#+index*options['offset_y']
+    if options['heatmap']:
+        ymin = yminmax['heatmap'][0]
+        ymax = yminmax['heatmap'][1]
+        ymin_start = yminmax['heatmap'][0]
+        ymax_start = yminmax['heatmap'][1]
 
+    elif not options['heatmap']:
+        ymin = yminmax['diff'][0]
+        ymax = yminmax['diff'][1]
+        ymin_start = yminmax['diff'][2]
+        ymax_start = yminmax['diff'][3]
+
+    
+    # FIXME The start values for xlim should probably also be decided by initial value of x_vals, and can likewise be implemented in update_xminmax() 
 
-    ymin_start = ymin - 0.1*ymax
-    ymax_start = ymax+0.2*ymax
-    ymin = ymin - 5*ymax
-    ymax = ymax*5
 
         
-
-
     options['widgets'] = {
         'xlim': {
-            'w': widgets.FloatRangeSlider(value=[minmax['2th'][0], minmax['2th'][1]], min=minmax['2th'][0], max=minmax['2th'][1], step=0.5, layout=widgets.Layout(width='95%')),
+            'w': widgets.FloatRangeSlider(value=[xminmax['start'][2], xminmax['start'][3]], min=xminmax['start'][0], max=xminmax['start'][1], step=0.5, layout=widgets.Layout(width='95%')),
-            'state': '2th',
+            'state': options['x_vals'],
-            '2th_default':      {'min': minmax['2th'][0],       'max': minmax['2th'][1],        'value': [minmax['2th'][0],         minmax['2th'][1]],      'step': 0.5},
+            '2th_default':      {'min': xminmax['2th'][0],          'max': xminmax['2th'][1],           'value': [xminmax['2th'][0],            xminmax['2th'][1]],             'step': 0.5},
-            '2th_cuka_default': {'min': minmax['2th_cuka'][0],  'max': minmax['2th_cuka'][1],   'value': [minmax['2th_cuka'][0],    minmax['2th_cuka'][1]], 'step': 0.5},
+            '2th_cuka_default': {'min': xminmax['2th_cuka'][0],     'max': xminmax['2th_cuka'][1],      'value': [xminmax['2th_cuka'][0],       xminmax['2th_cuka'][1]],        'step': 0.5},
-            '2th_moka_default': {'min': minmax['2th_moka'][0],  'max': minmax['2th_moka'][1],   'value': [minmax['2th_moka'][0],    minmax['2th_moka'][1]], 'step': 0.5},
+            '2th_moka_default': {'min': xminmax['2th_moka'][0],     'max': xminmax['2th_moka'][1],      'value': [xminmax['2th_moka'][0],       xminmax['2th_moka'][1]],        'step': 0.5},
-            'd_default':        {'min': minmax['d'][0],         'max': minmax['d'][1],          'value': [minmax['d'][0],           minmax['d'][1]],        'step': 0.5},
+            'd_default':        {'min': xminmax['d'][0],            'max': xminmax['d'][1],             'value': [xminmax['d'][0],              xminmax['d'][1]],               'step': 0.5},
-            '1/d_default':      {'min': minmax['1/d'][0],       'max': minmax['1/d'][1],        'value': [minmax['1/d'][0],         minmax['1/d'][1]],      'step': 0.5},
+            '1/d_default':      {'min': xminmax['1/d'][0],          'max': xminmax['1/d'][1],           'value': [xminmax['1/d'][0],            xminmax['1/d'][1]],             'step': 0.5},
-            'q_default':        {'min': minmax['q'][0],         'max': minmax['q'][1],          'value': [minmax['q'][0],           minmax['q'][1]],        'step': 0.5},
+            'q_default':        {'min': xminmax['q'][0],            'max': xminmax['q'][1],             'value': [xminmax['q'][0],              xminmax['q'][1]],               'step': 0.5},
-            'q2_default':       {'min': minmax['q2'][0],        'max': minmax['q2'][1],         'value': [minmax['q2'][0],          minmax['q2'][1]],       'step': 0.5},
+            'q2_default':       {'min': xminmax['q2'][0],           'max': xminmax['q2'][1],            'value': [xminmax['q2'][0],             xminmax['q2'][1]],              'step': 0.5},
-            'q4_default':       {'min': minmax['q4'][0],        'max': minmax['q4'][1],         'value': [minmax['q4'][0],          minmax['q4'][1]],       'step': 0.5}
+            'q4_default':       {'min': xminmax['q4'][0],           'max': xminmax['q4'][1],            'value': [xminmax['q4'][0],             xminmax['q4'][1]],              'step': 0.5},
+            'heatmap_default':  {'min': xminmax['heatmap'][0],      'max': xminmax['heatmap'][1],       'value': [xminmax['heatmap'][0],        xminmax['heatmap'][1]],         'step': 10}
+        },
+        'ylim': { 
+            'w': widgets.FloatRangeSlider(value=[yminmax['start'][2], yminmax['start'][3]], min=yminmax['start'][0], max=yminmax['start'][1], step=0.5, layout=widgets.Layout(width='95%')),
+            'state': 'heatmap' if options['heatmap'] else 'diff',
+            'diff_default':     {'min': yminmax['diff'][0],         'max': yminmax['diff'][1],          'value': [yminmax['diff'][2],           yminmax['diff'][3]],            'step': 0.1},
+            'heatmap_default':  {'min': yminmax['heatmap'][0],      'max': yminmax['heatmap'][1],       'value': [yminmax['heatmap'][0],        yminmax['heatmap'][1]],         'step': 0.1}
         }
     }
 
-
     if options['reflections_data']:
         w = widgets.interactive(btp.ipywidgets_update, func=widgets.fixed(plot_diffractogram), data=widgets.fixed(data), options=widgets.fixed(options), 
         scatter=widgets.ToggleButton(value=False), 
         line=widgets.ToggleButton(value=True), 
         reflections_plot=widgets.ToggleButton(value=True),
         reflections_indices=widgets.ToggleButton(value=False),
+        heatmap=widgets.ToggleButton(value=options['heatmap']),
         x_vals=widgets.Dropdown(options=['2th', 'd', '1/d', 'q', 'q2', 'q4', '2th_cuka', '2th_moka'], value='2th', description='X-values'),
         xlim=options['widgets']['xlim']['w'],
-        ylim=widgets.FloatRangeSlider(value=[ymin_start, ymax_start], min=ymin, max=ymax, step=0.5, layout=widgets.Layout(width='95%')),
+        ylim=options['widgets']['ylim']['w'],
-        offset_y=widgets.BoundedFloatText(value=options['offset_y'], min=-5, max=5, step=0.01),
+        offset_y=widgets.BoundedFloatText(value=options['offset_y'], min=-5, max=5, step=0.01, description='offset_y'),
-        offset_x=widgets.BoundedFloatText(value=options['offset_x'], min=-1, max=1, step=0.01)
+        offset_x=widgets.BoundedFloatText(value=options['offset_x'], min=-1, max=1, step=0.01, description='offset_x')
         )
     
     else:
@@ -247,45 +358,184 @@ def plot_diffractogram_interactive(data, options):
         xlim=options['widgets']['xlim']['w'])
     
     
+    options['widget'] = w
+
     display(w)
 
 
-def update_minmax(minmax, data):
+def update_xminmax(xminmax, data, options={}):
     ''' Finds minimum and maximum values of each column and updates the minmax dictionary to contain the correct values.
     
     Input:
     minmax (dict): contains '''
 
+    xminmax['2th'] = [None, None, None, None]
     for index, diffractogram in enumerate(data['diffractogram']):
-        if not minmax['2th'][0] or diffractogram['2th'].min() < minmax['2th'][0]:
+        
-            minmax['2th'][0] = diffractogram['2th'].min()
+        if not xminmax['2th'][0] or diffractogram['2th'].min() < xminmax['2th'][0]:
+            xminmax['2th'][0] = diffractogram['2th'].min()
             min_index = index
 
-        if not minmax['2th'][1] or diffractogram['2th'].max() > minmax['2th'][1]:
+        if not xminmax['2th'][1] or diffractogram['2th'].max() > xminmax['2th'][1]:
-            minmax['2th'][1] = diffractogram['2th'].max()
+            xminmax['2th'][1] = diffractogram['2th'].max()
             max_index = index
 
-    minmax['2th_cuka'][0], minmax['2th_cuka'][1] = data['diffractogram'][min_index]['2th_cuka'].min(),  data['diffractogram'][max_index]['2th_cuka'].max()
-    minmax['2th_moka'][0], minmax['2th_moka'][1] = data['diffractogram'][min_index]['2th_moka'].min(),  data['diffractogram'][max_index]['2th_moka'].max() 
-    minmax['d'][0], minmax['d'][1]               = data['diffractogram'][max_index]['d'].min(),         data['diffractogram'][min_index]['d'].max() # swapped, intended
-    minmax['1/d'][0], minmax['1/d'][1]           = data['diffractogram'][min_index]['1/d'].min(),       data['diffractogram'][max_index]['1/d'].max() 
-    minmax['q'][0], minmax['q'][1]               = data['diffractogram'][min_index]['q'].min(),         data['diffractogram'][max_index]['q'].max()
-    minmax['q2'][0], minmax['q2'][1]             = data['diffractogram'][min_index]['q2'].min(),        data['diffractogram'][max_index]['q2'].max() 
-    minmax['q4'][0], minmax['q4'][1]             = data['diffractogram'][min_index]['q4'].min(),        data['diffractogram'][max_index]['q4'].max() 
 
-def update_widgets(options):
+    xminmax['2th'][2], xminmax['2th'][3] = xminmax['2th'][0], xminmax['2th'][1]
 
-    for widget in options['widgets'].values():
+    xminmax['2th_cuka'][0], xminmax['2th_cuka'][1]  = data['diffractogram'][min_index]['2th_cuka'].min(),  data['diffractogram'][max_index]['2th_cuka'].max()
+    xminmax['2th_cuka'][2], xminmax['2th_cuka'][3]  = xminmax['2th_cuka'][0], xminmax['2th_cuka'][1]
 
-        if widget['state'] != options['x_vals']:
+    xminmax['2th_moka'][0], xminmax['2th_moka'][1]  = data['diffractogram'][min_index]['2th_moka'].min(),  data['diffractogram'][max_index]['2th_moka'].max()
-            for arg in widget[f'{options["x_vals"]}_default']:
+    xminmax['2th_moka'][2], xminmax['2th_moka'][3]  = xminmax['2th_moka'][0], xminmax['2th_moka'][1]
-                setattr(widget['w'], arg, widget[f'{options["x_vals"]}_default'][arg])
+
+    xminmax['d'][0], xminmax['d'][1]                = data['diffractogram'][max_index]['d'].min(),         data['diffractogram'][min_index]['d'].max() # swapped, intended
+    xminmax['d'][2], xminmax['d'][3]                = xminmax['d'][0], xminmax['d'][1]
+
+    xminmax['1/d'][0], xminmax['1/d'][1]            = data['diffractogram'][min_index]['1/d'].min(),       data['diffractogram'][max_index]['1/d'].max()
+    xminmax['1/d'][2], xminmax['1/d'][3]            = xminmax['1/d'][0], xminmax['1/d'][1]
+
+    xminmax['q'][0], xminmax['q'][1]                = data['diffractogram'][min_index]['q'].min(),         data['diffractogram'][max_index]['q'].max()
+    xminmax['q'][2], xminmax['q'][3]                = xminmax['q'][0], xminmax['q'][1]
+
+    xminmax['q2'][0], xminmax['q2'][1]              = data['diffractogram'][min_index]['q2'].min(),        data['diffractogram'][max_index]['q2'].max()
+    xminmax['q2'][2], xminmax['q2'][3]              = xminmax['q2'][0], xminmax['q2'][1]
+
+    xminmax['q4'][0], xminmax['q4'][1]              = data['diffractogram'][min_index]['q4'].min(),        data['diffractogram'][max_index]['q4'].max()
+    xminmax['q4'][2], xminmax['q4'][3]              = xminmax['q4'][0], xminmax['q4'][1]
+
+
+    xminmax['heatmap'] = options['heatmap_xlim'] # This value is set in the generate_heatmap()-function
+
+
+    xminmax['start'][0], xminmax['start'][1] = xminmax[options['x_vals']][0], xminmax[options['x_vals']][1]
+    xminmax['start'][2], xminmax['start'][3] = xminmax[options['x_vals']][2], xminmax[options['x_vals']][3]
+
+
+def update_yminmax(yminmax: dict, data: dict, options={}) -> None:
+    
+    yminmax['diff'] = [None, None, None, None]
+    # Go through diffractograms and find the minimum and maximum intensity values
+    for diffractogram in data['diffractogram']:
+        if not yminmax['diff'][0] or (yminmax['diff'][0] > (diffractogram['I'].min())): 
+            yminmax['diff'][0] = diffractogram['I'].min()
+
+        if not yminmax['diff'][1] or (yminmax['diff'][1] < (diffractogram['I'].max())):
+            yminmax['diff'][1] = diffractogram['I'].max()
+
+
+    # Set start values of ymin and ymax to be slightly below lowest data points and slightly above highest data points to give some whitespace around the plot
+    yminmax['diff'][2] = yminmax['diff'][0] - 0.1*yminmax['diff'][1]
+    yminmax['diff'][3] = yminmax['diff'][1] + 0.2*yminmax['diff'][1]
+    
+    # Allow for adjustment up to five times ymax above and below data
+    yminmax['diff'][0] = yminmax['diff'][0] - 5*yminmax['diff'][1]
+    yminmax['diff'][1] = yminmax['diff'][1]*5
+
+
+    # Set start values to the edges of the dataset
+    yminmax['heatmap'][0], yminmax['heatmap'][1] = 0, data['heatmap'].shape[0]
+    yminmax['heatmap'][2], yminmax['heatmap'][3] = yminmax['heatmap'][0], yminmax['heatmap'][1]
+
+
+    if options['heatmap']:
+        yminmax['start'][0], yminmax['start'][1] = yminmax['heatmap'][0], yminmax['heatmap'][1]
+        yminmax['start'][2], yminmax['start'][3] = yminmax['heatmap'][0], yminmax['heatmap'][1]
+
+    else:
+        # The third and fourth index are different here to not be zoomed completely out to begin with.
+        yminmax['start'][0], yminmax['start'][1] = yminmax['diff'][0], yminmax['diff'][1]
+        yminmax['start'][2], yminmax['start'][3] = yminmax['diff'][2], yminmax['diff'][3]
+    
+
+def update_defaults(widget: dict, minmax: dict) -> None:
+    ''' Updates the default x- or y-limits of a given widget. Refer to plot_diffractogram_interactive() to see the form of the widget that is passed in. An update of the min/max-values is done just prior to calling this function.
+    Changes dictionaries in place.
+    
+    Input:
+    widget (dict): A dictionary containing the widget itself (widget['w']) and all its default-values (e.g. widget['2th_default'])
+    minmax (dict): A dictionary containing min and max values, as well as min_start and max_start values. (e.g. minmax['2th'] is a list with four elements: [xmin, xmax, xmin_start, xmax_start])
+    
+    Output:
+    None.'''
+
+    for name, attr in widget.items():
+        if name.endswith('default'):
+            attr['min'] = minmax[name.replace('_default', '')][0]
+            attr['max'] = minmax[name.replace('_default', '')][1]           
+            attr['value'] = [minmax[name.replace('_default', '')][2], minmax[name.replace('_default', '')][3]]
+
+
+def update_widgets(data, options):
+
+
+    for widget_name, widget in options['widgets'].items():
+
+        # Make changes to xlim-widget
+        if widget_name == 'xlim':
+            # First update the min and max values
+            update_xminmax(xminmax=options['xminmax'], data=data, options=options)
+            update_defaults(widget=widget, minmax=options['xminmax'])
+           
+
+            if options['heatmap'] and (widget['state'] != 'heatmap'):
+
+
+                setattr(widget['w'], 'min', widget['heatmap_default']['min'])
+                setattr(widget['w'], 'max', widget['heatmap_default']['max'])
+                setattr(widget['w'], 'value', widget['heatmap_default']['value'])
+                setattr(widget['w'], 'step', widget['heatmap_default']['step'])
+
+                widget['state'] = 'heatmap'
             
-            widget['state'] = options['x_vals']
+            elif not options['heatmap'] and (widget['state'] != options['x_vals']):              
+                # Then loop through all attributes in the widget and change to current mode.
+                for arg in widget[f'{options["x_vals"]}_default']:
+                    
+                    # If new min value is larger than previous max, or new max value is smaller than previous min, set the opposite first
+                    if arg == 'min':
+                        if widget[f'{options["x_vals"]}_default']['min'] > getattr(widget['w'], 'max'):
+                            setattr(widget['w'], 'max', widget[f'{options["x_vals"]}_default']['max'])
+                    
+                    elif arg == 'max':
+                        if widget[f'{options["x_vals"]}_default']['max'] < getattr(widget['w'], 'min'):
+                            setattr(widget['w'], 'min', widget[f'{options["x_vals"]}_default']['min'])
+
+                
+                    setattr(widget['w'], arg, widget[f'{options["x_vals"]}_default'][arg])
+                
+                
+                widget['state'] = options['x_vals']
+
+        # Make changes to ylim-widget
+        elif widget_name == 'ylim':
+            update_yminmax(yminmax=options['yminmax'], data=data, options=options)
+            update_defaults(widget=widget, minmax=options['yminmax'])
+            
+            state = 'heatmap' if options['heatmap'] else 'diff' 
+    
+            if widget['state'] != state or options['offset_change']:
+
+                for arg in widget[f'{state}_default']:
+                    # If new min value is larger than previous max, or new max value is smaller than previous min, set the opposite first
+                    if arg == 'min':                     
+                        if widget[f'{state}_default']['min'] > getattr(widget['w'], 'max'):
+                            setattr(widget['w'], 'max', widget[f'{state}_default']['max'])
+                    
+                    elif arg == 'max':
+                        if widget[f'{state}_default']['max'] < getattr(widget['w'], 'min'):
+                            setattr(widget['w'], 'min', widget[f'{state}_default']['min'])
+
+                
+                    setattr(widget['w'], arg, widget[f'{state}_default'][arg])
+                
+                options['offset_change'] = False
+                widget['state'] = state
 
 
 
-def plot_reflection_indices(data, ax, options={}):
+
+def plot_reflection_indices(data, reflections_params, ax, options={}):
     ''' Print reflection indices from output generated by VESTA.
     
     Required contents of data:
@@ -299,20 +549,21 @@ def plot_reflection_indices(data, ax, options={}):
         'hide_indices': False
     }
 
-    data = aux.update_options(options=data, required_options=required_options, default_options=default_options)
+    reflections_params = aux.update_options(options=reflections_params, required_options=required_options, default_options=default_options)
 
-    if not data['hide_indices']:
+    if not reflections_params['hide_indices']:
-        reflection_table = xrd.io.load_reflection_table(data=data, options=options)
+        reflection_table = xrd.io.load_reflection_table(data=data, reflections_params=reflections_params, options=options)
         
-        if data['reflection_indices'] > 0:
+        if reflections_params['reflection_indices'] > 0:
 
             # Get the data['reflection_indices'] number of highest reflections within the subrange options['xlim']
-            reflection_indices = reflection_table.loc[(reflection_table[options['x_vals']] > options['xlim'][0]) & (reflection_table[options['x_vals']] < options['xlim'][1])].nlargest(options['reflection_indices'], 'I')
+            x_vals = 'heatmap' if options['heatmap'] else options['x_vals']
+            reflection_indices = reflection_table.loc[(reflection_table[x_vals] > options['xlim'][0]) & (reflection_table[x_vals] < options['xlim'][1])].nlargest(options['reflection_indices'], 'I')
 
             # Plot the indices
-            for i in range(data['reflection_indices']):
+            for i in range(reflections_params['reflection_indices']):
                 if reflection_indices.shape[0] > i:
-                    ax.text(s=f'({reflection_indices["h"].iloc[i]} {reflection_indices["k"].iloc[i]} {reflection_indices["l"].iloc[i]})', x=reflection_indices[options['x_vals']].iloc[i], y=0, fontsize=2.5, rotation=90, va='bottom', ha='center', c=data['text_colour'])    
+                    ax.text(s=f'({reflection_indices["h"].iloc[i]} {reflection_indices["k"].iloc[i]} {reflection_indices["l"].iloc[i]})', x=reflection_indices[x_vals].iloc[i], y=0, fontsize=2.5, rotation=90, va='bottom', ha='center', c=reflections_params['text_colour'])    
 
         
         if options['xlim']:
@@ -323,7 +574,7 @@ def plot_reflection_indices(data, ax, options={}):
 
     return
     
-def plot_reflection_table(data, ax=None, options={}):
+def plot_reflection_table(data, reflections_params, ax=None, options={}):
     ''' Plots a reflection table from output generated by VESTA.
     
     Required contents of data:
@@ -342,15 +593,15 @@ def plot_reflection_table(data, ax=None, options={}):
     }
 
     if 'colour' in data.keys():
-        options['reflections_colour'] = data['colour']
+        options['reflections_colour'] = reflections_params['colour']
-    if 'min_alpha' in data.keys():
+    if 'min_alpha' in reflections_params.keys():
-        options['min_alpha'] = data['min_alpha']
+        options['min_alpha'] = reflections_params['min_alpha']
-    if 'reflection_indices' in data.keys():
+    if 'reflection_indices' in reflections_params.keys():
-        options['reflection_indices'] = data['reflection_indices']
+        options['reflection_indices'] = reflections_params['reflection_indices']
-    if 'label' in data.keys():
+    if 'label' in reflections_params.keys():
-        options['label'] = data['label']
+        options['label'] = reflections_params['label']
-    if 'wavelength' in data.keys():
+    if 'wavelength' in reflections_params.keys():
-        options['wavelength'] = data['wavelength']
+        options['wavelength'] = reflections_params['wavelength']
 
     options = aux.update_options(options=options, required_options=required_options, default_options=default_options)
 
@@ -358,9 +609,10 @@ def plot_reflection_table(data, ax=None, options={}):
     if not ax:
         _, ax = btp.prepare_plot(options)
 
-    reflection_table = xrd.io.load_reflection_table(data=data, options=options)
+    x_vals = 'heatmap' if options['heatmap'] else options['x_vals']
 
-    reflections, intensities  = reflection_table[options['x_vals']], reflection_table['I']
+    reflection_table = xrd.io.load_reflection_table(data=data, reflections_params=reflections_params, options=options)
+    reflections, intensities  = reflection_table[x_vals], reflection_table['I']
 
 
     
@@ -390,7 +642,7 @@ def plot_reflection_table(data, ax=None, options={}):
         xlim_range = ax.get_xlim()[1] - ax.get_xlim()[0]
         ylim_avg = (ax.get_ylim()[0]+ax.get_ylim()[1])/2
 
-        ax.text(s=data['label'], x=(ax.get_xlim()[0]-0.01*xlim_range), y=ylim_avg, ha = 'right', va = 'center')
+        ax.text(s=reflections_params['label'], x=(ax.get_xlim()[0]-0.01*xlim_range), y=ylim_avg, ha = 'right', va = 'center')