diff --git a/diffractsim/visualization/plot_intensity.py b/diffractsim/visualization/plot_intensity.py
index 60b4c59..92dc4ae 100644
--- a/diffractsim/visualization/plot_intensity.py
+++ b/diffractsim/visualization/plot_intensity.py
@@ -116,8 +116,11 @@ def plot_intensity(self, I, square_root = False, figsize=(7, 6),
             x = self.x
             y = self.y
 
-        ax_slice.plot(x/units, I[np.argmin(abs(y-slice_y_pos)),:]**2)
-        ax_slice.set_ylabel(r'Intensity $\left[W / m^2 \right]$')
+        ax_slice.plot(x/units, I[np.argmin(abs(y-slice_y_pos)),:])
+        if square_root == False:
+            ax_slice.set_ylabel(r'Intensity $\left[W / m^2 \right]$')
+        else:
+            ax_slice.set_ylabel(r'Square Root Intensity $\left[ \sqrt{W / m^2 } \right]$')
 
         if grid == True:
             ax_slice.grid(alpha =0.2)
@@ -149,8 +152,11 @@ def plot_intensity(self, I, square_root = False, figsize=(7, 6),
             x = self.x
             y = self.y
 
-        ax_slice.plot(y/units, I[:, np.argmin(abs(x-slice_x_pos))]**2)
-        ax_slice.set_ylabel(r'Intensity $\left[W / m^2 \right]$')
+        ax_slice.plot(y/units, I[:, np.argmin(abs(x-slice_x_pos))])
+        if square_root == False:
+            ax_slice.set_ylabel(r'Intensity $\left[W / m^2 \right]$')
+        else:
+            ax_slice.set_ylabel(r'Square Root Intensity $\left[ \sqrt{W / m^2 } \right]$')
 
         if grid == True:
             ax_slice.grid(alpha =0.2)