diff --git a/src/vlaplex.f90 b/src/vlaplex.f90
index 21b3b6c..4c5d4fa 100644
--- a/src/vlaplex.f90
+++ b/src/vlaplex.f90
@@ -14,7 +14,6 @@ program VlaPlEx
   real(dp), parameter:: gamma_e_exp  =            1.0_dp /(gamma_e - 1.0_dp) ! Exponent for polytropic electrons
   real(dp), parameter:: gamma_e_dexp = (2.0_dp - gamma_e)/(gamma_e - 1.0_dp) ! Exponent for polytropic db_dphi
   real(dp), parameter:: n_epsilon = 1.0e-16_dp
-  real(dp), parameter:: cosTheta = 0.995_dp
 
   real(dp):: r0, rf
   real(dp), allocatable, dimension(:):: r
@@ -164,8 +163,8 @@ program VlaPlEx
   b         = 0.0_dp
   db_dphi   = 0.0_dp
   diag           = -2.0_dp / dr**2
-  diag_low       =  1.0_dp / dr**2! - 1.0_dp / (r(2:nr)   * dr)
-  diag_high      =  1.0_dp / dr**2! + 1.0_dp / (r(1:nr-1) * dr)
+  diag_low       =  1.0_dp / dr**2 - 1.0_dp / (r(2:nr)   * dr)
+  diag_high      =  1.0_dp / dr**2 + 1.0_dp / (r(1:nr-1) * dr)
   diag(1)        =  1.0_dp ! Dirichlet
   diag_high(1)   =  0.0_dp ! Dirichlet
   ! diag_high(1)   =  2.0_dp / dr**2  ! Neumann
@@ -248,13 +247,13 @@ program VlaPlEx
       do i = 1, nr
         ! Advect negative velocity
         if (i < nr) then
-          f_i(iz,i,1:j0-1) = f_i_old(iz,i,1:j0-1) - v(1:j0-1)*cosTheta*dt/dr*(f_i_old(iz,i+1,1:j0-1) - &
-                                                                              f_i_old(iz,i  ,1:j0-1))
+          f_i(iz,i,1:j0-1) = f_i_old(iz,i,1:j0-1) - v(1:j0-1)*dt/dr/r(i)**2*(r(i+1)**2*f_i_old(iz,i+1,1:j0-1) - &
+                                                                    r(i  )**2*f_i_old(iz,i  ,1:j0-1))
         end if
         ! Advect positive velocity
         if (i >  1) then
-          f_i(iz,i,j0:nv)  = f_i_old(iz,i, j0:nv) - v( j0:nv)*cosTheta*dt/dr*(f_i_old(iz,i  , j0:nv) - &
-                                                                              f_i_old(iz,i-1, j0:nv))
+          f_i(iz,i,j0:nv)  = f_i_old(iz,i, j0:nv) - v( j0:nv)*dt/dr/r(i)**2*(r(i  )**2*f_i_old(iz,i  , j0:nv) - &
+                                                                    r(i-1)**2*f_i_old(iz,i-1, j0:nv))
         end if
 
         n_i(iz,i) = sum(f_i(iz,i,:))*dv
@@ -284,8 +283,8 @@ program VlaPlEx
       phi_old = phi
 
       diag           = -2.0_dp / dr**2 - db_dphi
-      diag_low       =  1.0_dp / dr**2! - 1.0_dp / (r(2:nr)   * dr)
-      diag_high      =  1.0_dp / dr**2! + 1.0_dp / (r(1:nr-1) * dr)
+      diag_low       =  1.0_dp / dr**2 - 1.0_dp / (r(2:nr)   * dr)
+      diag_high      =  1.0_dp / dr**2 + 1.0_dp / (r(1:nr-1) * dr)
       diag(1)        =  1.0_dp ! Dirichlet
       diag_high(1)   =  0.0_dp ! Dirichlet
       ! diag(nr)       =  1.0_dp ! Dirichlet
@@ -323,6 +322,15 @@ program VlaPlEx
 
       end if
 
+      ! ! Calculate new potential to ensure 0 current at the edge
+      ! if (n_i(nr) > n_epsilon) then
+      !   phiF = phi0 + T_e0 * log((2.0_dp*sqrt(pi)*Zave(nr)*n_i(nr)*u_i(nr)) / (Zave(1)*n_i(1)*sqrt(m_i*T_e0/m_e)))
+      !
+      ! else
+      !   phiF = phi(nr-5)
+      !
+      ! end if
+
     end do
 
     ! Calculate electric field