Rounding option

giacomomarchioro · Jan 7, 2022 · ccac18f · ccac18f
1 parent 11d05bf
commit ccac18f
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/energydiagram/energydiagram.py b/energydiagram/energydiagram.py
@@ -28,6 +28,7 @@ def __init__(self, aspect='equal'):
         self.offset_ratio = 0.02
         self.color_bottom_text = 'blue'
         self.aspect = aspect
+        self.round_energies_at_digit = "keep all digits"
         # data
         self.pos_number = 0
         self.energies = []
@@ -98,7 +99,10 @@ def add_level(self, energy, bottom_text='', position=None, color='k',
             raise ValueError(
                 "Position must be None or 'last' (abrv. 'l') or in case an integer or float specifing the position. It was: %s" % position)
         if top_text == 'Energy':
-            top_text = energy
+            if self.round_energies_at_digit == "keep all digits":
+                top_text = energy
+            else:
+                top_text = round(energy,self.round_energies_at_digit)
 
         link = []
         self.colors.append(color)