From 5735d011aa7573d5eb62e6c5ca063e92a7f4d613 Mon Sep 17 00:00:00 2001
From: rasmusvt <r.v.thogersen@kjemi.uio.no>
Date: Fri, 22 Apr 2022 15:49:02 +0200
Subject: [PATCH] Add correct formatting of x- and y-labels for EC-plots

---
 nafuma/electrochemistry/plot.py | 30 ++++++++++++++++++++++++++++++
 nafuma/plotting.py              |  5 ++++-
 2 files changed, 34 insertions(+), 1 deletion(-)

diff --git a/nafuma/electrochemistry/plot.py b/nafuma/electrochemistry/plot.py
index dad866a..a074899 100644
--- a/nafuma/electrochemistry/plot.py
+++ b/nafuma/electrochemistry/plot.py
@@ -49,12 +49,42 @@ def plot_gc(data, options=None):
 				cycle[1].plot(x=options['x_vals'], y=options['y_vals'], ax=ax, c=colours[i][1])
 
 
+	update_labels(options)
+	print(options['xunit'])
 
 	fig, ax = btp.adjust_plot(fig=fig, ax=ax, options=options)
 
 	return data['cycles'], fig, ax
 
 
+
+def update_labels(options):
+
+	if 'xlabel' not in options.keys():
+		options['xlabel'] = options['x_vals'].capitalize().replace('_', ' ')
+
+	if 'ylabel' not in options.keys():
+		options['ylabel'] = options['y_vals'].capitalize().replace('_', ' ')
+		
+
+	if 'xunit' not in options.keys():
+		if options['x_vals'] == 'capacity':
+			options['xunit'] = options['units']['capacity']
+		elif options['x_vals'] == 'specific_capacity':
+			options['xunit'] = f"{options['units']['capacity']} {options['units']['mass']}$^{{-1}}$"
+		elif options['x_vals'] == 'time':
+			options['xunit'] = options['units']['time']
+		elif options['x_vals'] == 'ions':
+			options['xunit'] = None
+		
+
+	if 'yunit' not in options.keys():
+		if options['y_vals'] == 'voltage':
+			options['yunit'] = options['units']['voltage']
+
+
+
+		
 	
 def update_cycles_list(cycles, options: dict) -> None:
 
diff --git a/nafuma/plotting.py b/nafuma/plotting.py
index 2233ab7..416c1cb 100644
--- a/nafuma/plotting.py
+++ b/nafuma/plotting.py
@@ -135,7 +135,10 @@ def adjust_plot(fig, ax, options):
         ax.set_ylabel('')
         
     if not options['hide_x_labels']:
-        ax.set_xlabel(f'{options["xlabel"]}')
+        if not options['xunit']:
+            ax.set_xlabel(f'{options["xlabel"]}')
+        else:
+            ax.set_xlabel(f'{options["xlabel"]} [{options["xunit"]}]')
     else:
         ax.set_xlabel('')