diff --git a/vlaplex.f90 b/vlaplex.f90
index 3178b5f..8510f57 100644
--- a/vlaplex.f90
+++ b/vlaplex.f90
@@ -61,6 +61,7 @@ program VlaPlEx
   real(dp), allocatable, dimension(:,:,:):: f_i, f_i_old
   real(dp), allocatable, dimension(:):: f0   ! Boundary at r = x_0
   real(dp), allocatable, dimension(:,:):: n_i
+  real(dp), allocatable, dimension(:):: b_i
   real(dp), allocatable, dimension(:,:):: u_i
   real(dp), allocatable, dimension(:):: E_i
   real(dp), allocatable, dimension(:,:):: T_i
@@ -160,6 +161,7 @@ program VlaPlEx
   ! Allocate vectors
   allocate(f_i(1:nz,1:nr,1:nv), f_i_old(1:nz,1:nr,1:nv))
   allocate(n_i(1:nz,1:nr))
+  allocate(b_i(1:nr))
   allocate(u_i(1:nz,1:nr), E_i(1:nr), T_i(1:nz,1:nr))
   allocate(Zave(1:nr))
   allocate(n_e(1:nr))
@@ -168,6 +170,7 @@ program VlaPlEx
   f_i     = 0.0_dp
   f_i_old = 0.0_dp
   n_i     = 0.0_dp
+  b_i   = 0.0_dp
   u_i     = 0.0_dp
   E_i     = 0.0_dp
   T_i     = 0.0_dp
@@ -263,6 +266,7 @@ program VlaPlEx
     ! set edge velocities to 0
     f_i_old(:,:,1)  = 0.0_dp
     f_i_old(:,:,nv) = 0.0_dp
+    b_i = 0.0_dp
     ! Advect in the r direction
     !$omp parallel do
     do iz = 1, nz
@@ -292,14 +296,12 @@ program VlaPlEx
         end if
 
       end do
+      b_i = b_i + Zave * n_i(iz,:)
     end do
     !$omp end parallel do
 
     ! Assume quasi-neutrality to start iterating
-    n_e = Zave * n_i(1,:)
-    !do iz = 1, nz
-    !  n_e = n_e + Zave * n_i(:,iz)
-    !end do
+    n_e = 1.0_dp/nz * b_i
     db_dphi = 0.0_dp
 
     ! Solve Poission (maximum number of iterations, break if convergence is reached before)
@@ -317,12 +319,7 @@ program VlaPlEx
       diag_low(nr-1) =  2.0_dp / dr**2 - db_dphi(nr) ! Neumann
 
       ! Calculate charge density
-      !b = n_e
-      !b = b - (Zave * n_i(:,1))
-      !do iz = 1, nz
-      !  b = b - (Zave * n_i(:,iz))
-      !end do
-      b = -(Zave * n_i(1,:) - n_e)
+      b = - 1.0_dp/nz * b_i + n_e
       ! Apply boundary conditions
       b(1)  = phi0 ! Dirichlet
       ! b(nr) = 0.0_dp ! Dirichlet