From 317805f4c31e1e0be1153292e7403efa2e801cfb Mon Sep 17 00:00:00 2001
From: rasmusthog <code@rasmusthog.me>
Date: Sat, 15 Oct 2022 17:16:04 +0200
Subject: [PATCH] Add functions to read and plot EDS spectrum

---
 nafuma/eds/io.py   | 25 +++++++++++++++++++
 nafuma/eds/plot.py | 62 ++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 87 insertions(+)

diff --git a/nafuma/eds/io.py b/nafuma/eds/io.py
index 8d1b34f..91c8c51 100644
--- a/nafuma/eds/io.py
+++ b/nafuma/eds/io.py
@@ -1,6 +1,7 @@
 from PIL import Image
 import numpy as np
 import cv2
+import pandas as pd
 
 def read_image(path, weight=None, colour=None, crop=None, resize=None, brightness=None):
 
@@ -125,3 +126,27 @@ def change_colour(image, new_colour):
 
 
     return image
+
+
+def read_spectrum(path):
+
+    headers = find_start(path)
+
+    spectrum = pd.read_csv(path, skiprows=headers, delim_whitespace=True)
+
+
+    return spectrum
+
+
+
+def find_start(path):
+
+    with open(path, 'r') as f:
+        line = f.readline()
+        i = 0
+        while not line.startswith('Energy'):
+            line = f.readline()
+            i += 1
+
+    return i
+
diff --git a/nafuma/eds/plot.py b/nafuma/eds/plot.py
index 7726f75..be8ee39 100644
--- a/nafuma/eds/plot.py
+++ b/nafuma/eds/plot.py
@@ -71,3 +71,65 @@ def show_image(data, options={}):
     else:
         return data['image'], None, None
 
+
+
+def plot_spectrum(data: dict, options={}):
+
+    default_options = {
+        'deconvolutions': None,
+        'lines': None,
+        'colours': None,
+        'xlabel': 'Energy', 'xunit': 'keV', 'xlim': None,
+        'ylabel': 'Counts', 'yunit': 'arb. u.', 'ylim': None, 'hide_y_ticklabels': True, 'hide_y_ticks': True,
+    }
+
+    options = aux.update_options(options=options, default_options=default_options)
+
+    fig, ax = btp.prepare_plot(options=options)
+
+
+    spectrum = io.read_spectrum(data['path'])
+
+    if options['deconvolutions']:
+        
+        deconvolutions = []
+        if not isinstance(options['deconvolutions'], list):
+            options['deconvolutions'] = [options['deconvolutions']]
+
+        if options['colours'] and (len(options['colours']) != len(options['deconvolutions'])):
+            options['colours'] = None
+
+        for deconv in options['deconvolutions']:
+            df = io.read_spectrum(deconv)
+            deconvolutions.append(df)
+
+
+    
+    spectrum.plot(x='Energy', y='Counts', ax=ax, color='black')
+
+    if options['deconvolutions']:
+        if options['colours']:
+            for deconv, colour in zip(deconvolutions, options['colours']):
+                ax.fill_between(x=deconv['Energy'], y1=deconv['Counts'], y2=0, color=colour, alpha=0.4)
+        else:
+            for deconv in deconvolutions:
+                ax.fill_between(x=deconv['Energy'], y1=deconv['Counts'], y2=0, alpha=0.4)
+
+
+    if not options['xlim']:
+        options['xlim'] = [spectrum['Energy'].min(), spectrum['Energy'].max()]
+
+    if not options['ylim']:
+        options['ylim'] = [0, 1.1*spectrum['Counts'].max()]
+
+    if options['lines']:
+        for i, (line, energy) in enumerate(options['lines'].items()):
+            ax.axvline(x=energy, ls='--', lw=0.5, c='black')
+            ax.text(s=line, x=energy, y=(0.9-0.1*i)*options['ylim'][1], fontsize=8)
+
+
+    
+    fig, ax = btp.adjust_plot(fig=fig, ax=ax, options=options)
+
+
+    return spectrum, fig, ax
\ No newline at end of file