diff --git a/src/vlaplex.f90 b/src/vlaplex.f90
index 21b3b6c..7fc9ebe 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
@@ -54,7 +53,7 @@ program VlaPlEx
   real(dp), allocatable, dimension(:):: phi, phi_old, E, db_dphi
   real(dp):: phiConv
   real(dp):: phi0
-  real(dp):: T_e0
+  real(dp):: T_e
   ! real(dp):: phiF
   integer:: k
 
@@ -146,7 +145,7 @@ program VlaPlEx
   E_i         = 0.0_dp
   T_i         = 0.0_dp
   n_e         = 0.0_dp
-  T_e0        = 0.0_dp
+  T_e         = 0.0_dp
   Zave        = 0.0_dp
   Zave_bc_old = 0.0_dp
   phi         = 0.0_dp
@@ -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
@@ -212,7 +211,7 @@ program VlaPlEx
 
     ! Get boundary conditions for specific time
     call boundaryConditions%get(time, n_bc, u_bc, Temp_bc)
-    ! Find new \bar{Z}_i based on T_e0 = Temp_bc and n_e = n_bc
+    ! Find new \bar{Z}_i based on T_e = Temp_bc and n_e = n_bc
     call Tene_to_Z%get(Temp_bc, n_bc, Zave_bc)
     ! Assign Z(T,n) to bin
     z_inj = minloc(abs(Zlist - Zave_bc),1)
@@ -230,7 +229,7 @@ program VlaPlEx
     f_i_old(z_inj,1,j0:nv) = f0
     f_i(:,1,j0:nv) = f_i_old(:,1,j0:nv)
 
-    T_e0 = Temp_bc
+    T_e = Temp_bc
 
     ! r = rf, v<0
     f_i_old(:,nr,1:j0-1) = 0.0_dp
@@ -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
@@ -309,12 +308,12 @@ program VlaPlEx
       ! phi0=phi(1) ! Neumann
 
       ! Calculate distribution of electrons
-      ! n_e = sum_ni(1) * exp((phi- phi0) / T_e0) ! Isothermal (Boltzmann)
-      n_e = sum_ni(1) * (1.0_dp + (gamma_e - 1.0_dp)/gamma_e*(phi-phi0)/T_e0)**gamma_e_exp !Polytropic
+      ! n_e = sum_ni(1) * exp((phi- phi0) / T_e) ! Isothermal (Boltzmann)
+      n_e = sum_ni(1) * (1.0_dp + (gamma_e - 1.0_dp)/gamma_e*(phi-phi0)/T_e)**gamma_e_exp !Polytropic
       ! Diagonal matrix for Newton integration scheme
-      ! db_dphi        = n_e / T_e0 ! Isothermal (Boltzmann)
-      db_dphi        =  sum_ni(1) / (gamma_e * T_e0) * &
-                        (1.0_dp + (gamma_e - 1.0_dp)/gamma_e*(phi-phi0)/T_e0)**gamma_e_dexp !Polytropic
+      ! db_dphi        = n_e / T_e ! Isothermal (Boltzmann)
+      db_dphi        =  sum_ni(1) / (gamma_e * T_e) * &
+                        (1.0_dp + (gamma_e - 1.0_dp)/gamma_e*(phi-phi0)/T_e)**gamma_e_dexp !Polytropic
 
       ! Check if the solution has converged
       phiConv = maxval(abs(Res),1)
@@ -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_e * log((2.0_dp*sqrt(pi)*Zave(nr)*n_i(nr)*u_i(nr)) / (Zave(1)*n_i(1)*sqrt(m_i*T_e/m_e)))
+      !
+      ! else
+      !   phiF = phi(nr-5)
+      !
+      ! end if
+
     end do
 
     ! Calculate electric field