Add functions to plot and format STM from .dat file

utils/stm_tools.py

@@ -1,9 +1,36 @@
 import numpy as np 
 from matplotlib import pyplot as plt
+import matplotlib.font_manager as fm
 
-def plot_stm_from_dat(file: str, gridsize: int = 300):
+def get_plot_stm_from_dat(file: str, gridsize: int = 300) -> plt:
     '''Plots a STM from a .dat file generated by Critic2'''
     data = np.loadtxt(file)
     plt.hexbin(data[:,2], data[:,3], C=data[:,4], gridsize=gridsize, cmap='Greys_r')
-    plt.colorbar(label='Relative Height (arb. units)')
-    plt.show()
+
+    return plt
+
+def format_stm_plot(plt: plt, title: str = 'Simulated STM', mode: str = 'current') -> plt:
+    '''Formats the plot of a STM'''
+    # Set font to serif
+    plt.rcParams['font.family'] = 'serif'
+    # Set title
+    plt.title(title)
+    # Replace \AA with angstrom symbol
+    plt.xlabel(r'x ($\mathrm{\AA}$)')
+    plt.ylabel(r'y ($\mathrm{\AA}$)')
+
+    if mode == 'current':
+        plt.colorbar(label=r'Tip Height ($\mathrm{\AA}$)')
+    elif mode == 'height':
+        plt.colorbar(label='Current (a.u.)')
+
+    return plt
+
+def plot_stm_from_dat(file: str, title: str = 'Simulated STM', mode: str = 'current', gridsize: int = 300) -> plt:
+    '''Plots a STM from a .dat file generated by Critic2'''
+    plt = get_plot_stm_from_dat(file, gridsize)
+    plt = format_stm_plot(plt, title, mode)
+
+    return plt.show()
+
+