diff --git a/src/model_multi_phase_network.jl b/src/model_multi_phase_network.jl
index 9f1ab0b..d55d00a 100644
--- a/src/model_multi_phase_network.jl
+++ b/src/model_multi_phase_network.jl
@@ -228,6 +228,7 @@ from JuMP for all time steps.
 function constrain_power_balance(m, net::Network{MultiPhase})
     Sij = m[:Sij]
     Lij = m[:l]
+    w = m[:w]
     m[:loadbalcons] = Dict()
     
     for j in busses(net)
@@ -253,11 +254,18 @@ function constrain_power_balance(m, net::Network{MultiPhase})
 
             if j == net.substation_bus   # include the slack power variables
                 m[:loadbalcons][j] = @constraint(m,  [t in 1:net.Ntimesteps],
-                    m[:Sj][t][j] + Sj[t, :] - sum( diag( Sij[t][(j,k)] ) for k in j_to_k(j, net) ) .== 0
+                    m[:Sj][t][j] + Sj[t, :] 
+                    - diag(w[t][j] * conj(yj(j, net))) * net.Zbase  # put yj in per-unit
+                    - sum( diag( Sij[t][(j,k)] ) for k in j_to_k(j, net) )
+                    .== 0
                 )
+
             else  # a source node with known injection
                 m[:loadbalcons][j] = @constraint(m,  [t in 1:net.Ntimesteps],
-                    Sj[t, :] - sum( diag( Sij[t][(j,k)] ) for k in j_to_k(j, net) ) .== 0
+                    Sj[t, :] 
+                    - diag(w[t][j] * conj(yj(j, net))) * net.Zbase  # put yj in per-unit
+                    - sum( diag( Sij[t][(j,k)] ) for k in j_to_k(j, net) ) 
+                    .== 0
                 )
 
             end
@@ -273,7 +281,9 @@ function constrain_power_balance(m, net::Network{MultiPhase})
                 sum( diag( 
                     Sij[t][(i,j)] - zij_per_unit(i,j,net) * Lij[t][(i,j)]
                 ) for i in i_to_j(j, net) )
-                + Sj[t, :] .== 0
+                + Sj[t, :] 
+                - diag(w[t][j] * conj(yj(j, net))) * net.Zbase  # put yj in per-unit
+                .== 0
             )
 
         # node with lines in and out
@@ -283,11 +293,12 @@ function constrain_power_balance(m, net::Network{MultiPhase})
                     Sij[t][(i,j)] - zij_per_unit(i,j,net) * Lij[t][(i,j)]
                 ) for i in i_to_j(j, net) )
                 + Sj[t, :]
-                - sum( diag( Sij[t][(j,k)] ) for k in j_to_k(j, net) ) .== 0
+                - diag(w[t][j] * conj(yj(j, net))) * net.Zbase  # put yj in per-unit
+                - sum( diag( Sij[t][(j,k)] ) for k in j_to_k(j, net) ) 
+                .== 0
             )
         end
     end
-    # TODO add shunts, diag( openDSS-cmatrix * m[:w][t][j] ) ?
 
     nothing
 end
diff --git a/src/model_single_phase_network.jl b/src/model_single_phase_network.jl
index fd8bb47..cba3453 100644
--- a/src/model_single_phase_network.jl
+++ b/src/model_single_phase_network.jl
@@ -101,7 +101,8 @@ function constrain_power_balance(m, net::Network)
         # check for shunt admittance
         shunt_susceptance = 0.0
         if :ShuntAdmittance in keys(net[j])
-            shunt_susceptance = net[j][:ShuntAdmittance].b
+            # flip sign of susceptance b/c we want complex conjugate
+            shunt_susceptance = -net[j][:ShuntAdmittance].b
         end
 
         # check for capacitors (only fixed for now)
diff --git a/test/runtests.jl b/test/runtests.jl
index b0c8655..6b128b5 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -46,7 +46,7 @@ function build_min_loss_model(net)
 end
 
 
-function hermitian_variable_to_vector(m::JuMP.AbstractModel, var::Symbol, t::Int, bus::String)
+function hermitian_variable_to_vector(m::JuMP.AbstractModel, var::Symbol, t::Int, bus::Union{String, Tuple{String, String}}; tol=1e-5)
 
     H = value.(m[var][t][bus])
 
@@ -58,7 +58,7 @@ function hermitian_variable_to_vector(m::JuMP.AbstractModel, var::Symbol, t::Int
     eigenvectors = eig.vectors
 
     # Select a non-zero eigenvalue and corresponding eigenvector
-    non_zero_indices = findall(x -> abs(x) > 1e-2, eigenvalues)
+    non_zero_indices = findall(x -> abs(x) > tol, eigenvalues)
     lambda_i = abs(eigenvalues[non_zero_indices[1]])
     u_i = eigenvectors[:, non_zero_indices[1]]