diff --git a/moduleOutput.f90 b/moduleOutput.f90
index a54d949..896b486 100644
--- a/moduleOutput.f90
+++ b/moduleOutput.f90
@@ -112,42 +112,52 @@ module output
 
     end subroutine writeOutputPhi
 
-    subroutine writeOutputMom(t, dt, nr, r, n_i, u_i, T_i, Z)
+    subroutine writeOutputMom(t, dt, nz, nr, r, n_i, u_i, T_i, Z_list)
       use constantParameters, only: eV_to_K
       use referenceValues,    only: L_ref, t_ref, n_ref, u_ref, Temp_ref
 
       integer, intent(in):: t
       integer,  intent(in):: nr
+      integer,  intent(in):: nz
       real(dp), intent(in):: dt
       real(dp), intent(in):: r(1:nr)
-      real(dp), intent(in):: n_i(1:nr)
-      real(dp), intent(in):: u_i(1:nr)
-      real(dp), intent(in):: T_i(1:nr)
-      real(dp), intent(in):: Z(1:nr)
+      real(dp), intent(in):: n_i(1:nz,1:nr)
+      real(dp), intent(in):: u_i(1:nz,1:nr)
+      real(dp), intent(in):: T_i(1:nz,1:nr)
+      real(dp), intent(in):: Z_list(1:nz)
       character(:), allocatable:: filename
       integer:: i
+      integer:: j
       character(len=10):: timeString
+      character(len=10) :: ZString
 
-      write (timeString, formatTime) t
-      filename = 'time_' // trim(timeString)//'_mom_i.csv'
-      write (*, '(A, A)') 'Writing: ', filename
+      do j = 1, nz
+        write (timeString, formatTime) t
+        write(ZString, '(F6.3)') Z_list(j)  ! Format Z_list(j) with 3 decimal places
+        ZString = adjustl(trim(ZString))   ! Remove leading spaces
 
-      open(unit=dataPhi_id, file=pathOutput//filename)
-      write(dataPhi_id, '(A)') "t (s)"
-      write(dataPhi_id, '('//formatFloat//')') t*dt*t_ref
-      write(dataPhi_id, '(A,4('//formatSep//',A))') "r (m)","n_i (m^-3)","u_i (m s^-1)", "T_i (eV)","Zave"
-      do i = 1, nr
-        write(dataPhi_id, '('//formatFloat//',4('//formatSep //','//formatFloat//'))') &
-                          r(i)*L_ref,                                                  &
-                          n_i(i)*n_ref,                                                &
-                          u_i(i)*u_ref,                                                &
-                          T_i(i)*Temp_ref/ev_to_K,                                     &
-                          Z(i)                    
+        ZString = adjustl(ZString)
+        ZString = ZString(:index(ZString // '.', '.') - 1) // ZString(index(ZString, '.') + 1:)
+        filename = 'time_' // trim(timeString) // '_Z' // trim(ZString) // '_mom_i.csv'
+        write (*, '(A, A)') 'Writing: ', filename
 
+        open(unit=dataPhi_id, file=pathOutput//filename)
+        write(dataPhi_id, '(A)') "t (s)"
+        write(dataPhi_id, '('//formatFloat//')') t*dt*t_ref
+        write(dataPhi_id, '(A)') "Z"
+        write(dataPhi_id, '('//formatFloat//')') Z_list(j)
+        write(dataPhi_id, '(A,4('//formatSep//',A))') "r (m)","n_i (m^-3)","u_i (m s^-1)", "T_i (eV)"
+        do i = 1, nr
+          write(dataPhi_id, '('//formatFloat//',4('//formatSep //','//formatFloat//'))') &
+                            r(i)*L_ref,                                                  &
+                            n_i(j,i)*n_ref,                                                &
+                            u_i(j,i)*u_ref,                                                &
+                            T_i(j,i)*Temp_ref/ev_to_K                  
+
+        end do
+
+        close(unit=dataPhi_id)
       end do
-
-      close(unit=dataPhi_id)
-
     end subroutine writeOutputMom
 
     subroutine writeOutputBoundary(t, dt, n, u, Temp, Z)
diff --git a/scripts_python/readMom.py b/scripts_python/readMom.py
index f88fff5..6c79b64 100644
--- a/scripts_python/readMom.py
+++ b/scripts_python/readMom.py
@@ -5,12 +5,14 @@ def read(filename):
     df = pandas.read_csv(filename,skiprows=0,nrows=1)
     time = df['t (s)'].to_numpy()[0]
 
-    df = pandas.read_csv(filename,skiprows=2)
+    df = pandas.read_csv(filename,skiprows=2,nrows=1)
+    Z = df['Z'].to_numpy()[0]
+
+    df = pandas.read_csv(filename,skiprows=4)
     x   = df['r (m)'].to_numpy()
     n_i = df['n_i (m^-3)'].to_numpy()
     u_i = df['u_i (m s^-1)'].to_numpy()
     T_i = df['T_i (eV)'].to_numpy()
-    Z   = df['Zave'].to_numpy()
 
     return time, x, n_i, u_i, T_i, Z