diff --git a/ext/JutulDarcyGLMakieExt/well_plots.jl b/ext/JutulDarcyGLMakieExt/well_plots.jl
index 89c722ec..650ad22a 100644
--- a/ext/JutulDarcyGLMakieExt/well_plots.jl
+++ b/ext/JutulDarcyGLMakieExt/well_plots.jl
@@ -100,7 +100,7 @@ function JutulDarcy.plot_well_results(well_data::Vector, time = nothing; start_d
         else
             c_key = :Paired_10
         end
-        cmap = cgrad(c_key, nw, categorical=true)
+        cmap = cgrad(c_key, max(nw, 2), categorical=true)
     end
     wellstr = [String(x) for x in wells]
 
diff --git a/ext/JutulDarcyPartitionedArraysExt/JutulDarcyPartitionedArraysExt.jl b/ext/JutulDarcyPartitionedArraysExt/JutulDarcyPartitionedArraysExt.jl
index 566a538c..0d2832eb 100644
--- a/ext/JutulDarcyPartitionedArraysExt/JutulDarcyPartitionedArraysExt.jl
+++ b/ext/JutulDarcyPartitionedArraysExt/JutulDarcyPartitionedArraysExt.jl
@@ -40,7 +40,7 @@ module JutulDarcyPartitionedArraysExt
     function Jutul.parray_preconditioner_apply!(global_out, main_prec::CPRPreconditioner{<:BoomerAMGPreconditioner, <:Any}, X, preconditioners, simulator, arg...)
         global_cell_vector = simulator.storage.distributed_cell_buffer
         global_buf = simulator.storage.distributed_residual_buffer
-        map(local_values(X), preconditioners, ghost_values(X)) do x, prec, x_g
+        @tic "cpr first stage" map(local_values(X), preconditioners, ghost_values(X)) do x, prec, x_g
             @. x_g = 0.0
             JutulDarcy.apply_cpr_first_stage!(prec, x, arg...)
             nothing
@@ -49,7 +49,7 @@ module JutulDarcyPartitionedArraysExt
         # copy!(global_cell_vector, main_prec.p)
         p_h = main_prec.p
         @assert !isnothing(p_h) "CPR is not properly initialized."
-        map(own_values(global_cell_vector), preconditioners) do ov, prec
+        @tic "hypre GetValues" map(own_values(global_cell_vector), preconditioners) do ov, prec
             helper = prec.pressure_precond.data[:assembly_helper]
             indices = helper.indices
             indices::Vector{HYPRE.HYPRE_BigInt}
@@ -59,29 +59,27 @@ module JutulDarcyPartitionedArraysExt
         # End unsafe shenanigans
 
         # consistent!(global_cell_vector) |> wait
-        map(own_values(global_buf), own_values(global_cell_vector), preconditioners) do dx, dp, prec
+        @tic "set dp" map(own_values(global_buf), own_values(global_cell_vector), preconditioners) do dx, dp, prec
             bz = prec.block_size
             for i in eachindex(dp)
                 JutulDarcy.set_dp!(dx, bz, dp, i)
             end
             nothing
         end
-        if main_prec.full_system_correction
-            mul_ix = nothing
-        else
-            mul_ix = 1
+
+        @tic "correct residual" begin
+            mul!(X, main_prec.A_ps, global_buf, -1.0, true)
+            nothing
         end
-        full_op = Jutul.parray_linear_system_operator(simulator.storage.simulators, global_buf)
-        mul!(X, full_op, global_buf, -1.0, true)
 
-        map(local_values(global_out), local_values(X), preconditioners, local_values(global_cell_vector), ghost_values(X)) do y, x, prec, dp, x_g
+        @tic "increment dp" map(local_values(global_out), local_values(X), preconditioners, local_values(global_cell_vector), ghost_values(X)) do y, x, prec, dp, x_g
             @. x_g = 0.0
             apply!(y, prec.system_precond, x, arg...)
             bz = prec.block_size
             JutulDarcy.increment_pressure!(y, dp, bz)
             nothing
         end
-        consistent!(global_out) |> wait
+        @tic "communication" consistent!(global_out) |> wait
         global_out
     end
 
@@ -104,8 +102,16 @@ module JutulDarcyPartitionedArraysExt
             cpr.r_p = create_hypre_vector()
             cpr.p = create_hypre_vector()
             cpr.np = n
+            if cpr.full_system_correction
+                mul_ix = nothing
+            else
+                mul_ix = 1
+            end
+            global_buf = sim.storage.distributed_residual_buffer
+            cpr.A_ps = Jutul.parray_linear_system_operator(sim.storage.simulators, global_buf)
         end
         A_p = cpr.A_p
+        A_ps = cpr.A_ps
         r_p = cpr.r_p
         x_p = cpr.p
 
@@ -114,6 +120,7 @@ module JutulDarcyPartitionedArraysExt
             model = sim.model
             storage = sim.storage
             prec.A_p = A_p
+            prec.A_ps = A_ps
             prec.p = x_p
             prec.r_p = r_p
             prec.np = n
diff --git a/src/facility/wells/wells.jl b/src/facility/wells/wells.jl
index f50b0bcf..d982d749 100644
--- a/src/facility/wells/wells.jl
+++ b/src/facility/wells/wells.jl
@@ -46,7 +46,7 @@ function common_well_setup(nr; dz = nothing, WI = nothing, gravity = gravity_con
         @warn "No well indices provided. Using 1e-12."
         WI = repeat(1e-12, nr)
     else
-        @assert length(WI) == nr  "Must have one well index per perforated cell"
+        @assert length(WI) == nr  "Must have one well index per perforated cell ($(length(WI)) well indices, $nr reservoir cells))"
     end
     return (WI, gdz)
 end
diff --git a/src/flux.jl b/src/flux.jl
index 9f92ab05..fedb8e8a 100644
--- a/src/flux.jl
+++ b/src/flux.jl
@@ -43,24 +43,34 @@ end
 end
 
 @inline function darcy_phase_kgrad_potential(face, phase, state, model, flux_type, tpfa::TPFA, upw, common = flux_primitives(face, state, model, flux_type, upw, tpfa))
-    ρ = state.PhaseMassDensities
     pc, ref_index = capillary_pressure(model, state)
     ∇p, T_f, gΔz = common
     l = tpfa.left
     r = tpfa.right
 
     Δpc = capillary_gradient(pc, l, r, phase, ref_index)
-    @inbounds ρ_c = ρ[phase, l]
-    @inbounds ρ_i = ρ[phase, r]
-    ρ_avg = 0.5*(ρ_i + ρ_c)
+    ρ_avg = face_average_density(model, state, tpfa, phase)
     q = -T_f*(∇p + Δpc + gΔz*ρ_avg)
     return q
 end
 
+function face_average_density(model, state, tpfa, phase)
+    ρ = state.PhaseMassDensities
+    l = tpfa.left
+    r = tpfa.right
+    @inbounds ρ_c = ρ[phase, l]
+    @inbounds ρ_i = ρ[phase, r]
+    return 0.5*(ρ_i + ρ_c)
+end
+
 @inline function gradient(X, tpfa::TPFA)
     return @inbounds X[tpfa.right] - X[tpfa.left]
 end
 
+@inline function gradient(X::AbstractMatrix, i, tpfa::TPFA)
+    return @inbounds X[i, tpfa.right] - X[i, tpfa.left]
+end
+
 pressure_gradient(state, disc) = gradient(state.Pressure, disc)
 
 @inline function upwind(upw::SPU, F, q)
diff --git a/src/forces/bc.jl b/src/forces/bc.jl
index 7e3935d7..1b4bef46 100644
--- a/src/forces/bc.jl
+++ b/src/forces/bc.jl
@@ -128,3 +128,90 @@ function apply_flow_bc!(acc, q, bc, model::SimulationModel{<:Any, T}, state, tim
         end
     end
 end
+
+function Jutul.vectorization_length(bc::FlowBoundaryCondition, variant)
+    if variant == :all
+        n = 4 # pressure, temperature, T_flow, T_thermal
+        f = bc.fractional_flow
+        if !isnothing(f)
+            n += length(f)
+        end
+        if !isnothing(bc.density)
+            n += 1
+        end
+        return n
+    elseif variant == :control
+        return 1
+    else
+        error("Variant $variant not supported")
+    end
+end
+
+function Jutul.vectorize_force!(v, bc::FlowBoundaryCondition, variant)
+    names = []
+    if variant == :all
+        v[1] = bc.pressure
+        push!(names, :pressure)
+        v[2] = bc.temperature
+        push!(names, :temperature)
+        v[3] = bc.trans_flow
+        push!(names, :trans_flow)
+        v[4] = bc.trans_thermal
+        push!(names, :trans_thermal)
+        offset = 4
+        f = bc.fractional_flow
+        if !isnothing(f)
+            for (i, f_i) in enumerate(f)
+                offset += 1
+                v[offset] = f_i
+                push!(names, Symbol("fractional_flow$i"))
+            end
+        end
+        if !isnothing(bc.density)
+            offset += 1
+            v[offset] = bc.density
+            push!(names, :density)
+        end
+    elseif variant == :control
+        v[1] = bc.pressure
+        push!(names, :pressure)
+    else
+        error("Variant $variant not supported")
+    end
+    return (names = names, )
+end
+
+function Jutul.devectorize_force(bc::FlowBoundaryCondition, X, meta, variant)
+    p = X[1]
+    T = bc.temperature
+    trans_flow = bc.trans_flow
+    trans_thermal = bc.trans_thermal
+    f = bc.fractional_flow
+    ρ = bc.density
+    if variant == :all
+        T = X[2]
+        trans_flow = X[3]
+        trans_thermal = X[4]
+
+        offset = 4
+        if !isnothing(f)
+            f = X[(offset+1):(offset+length(f))]
+        end
+        if !isnothing(ρ)
+            ρ = X[end]
+        end
+    elseif variant == :control
+        # DO nothing
+    else
+        error("Variant $variant not supported")
+    end
+    return FlowBoundaryCondition(
+        bc.cell,
+        p,
+        T,
+        trans_flow,
+        trans_thermal,
+        f,
+        ρ
+    )
+end
diff --git a/src/forces/sources.jl b/src/forces/sources.jl
index 721c84d1..69e88bdd 100644
--- a/src/forces/sources.jl
+++ b/src/forces/sources.jl
@@ -39,3 +39,53 @@ function Jutul.subforce(s::AbstractVector{S}, model) where S<:SourceTerm
     return s[keep]
 end
 
+function Jutul.vectorization_length(src::SourceTerm, variant)
+    n = 1
+    if variant == :all
+        f = src.fractional_flow
+        if !isnothing(f)
+            n += length(f)
+        end
+    elseif variant == :control
+        # Do nothing
+    else
+        error("Variant $variant not supported")
+    end
+    return n
+end
+
+function Jutul.vectorize_force!(v, src::SourceTerm, variant)
+    v[1] = src.value
+    names = [:value]
+    if variant == :all
+        offset = 1
+        f = src.fractional_flow
+        if !isnothing(f)
+            for (i, f_i) in enumerate(f)
+                offset += 1
+                v[offset] = f_i
+                push!(names, Symbol("fractional_flow$i"))
+            end
+        end
+    elseif variant == :control
+        # Do nothing
+    else
+        error("Variant $variant not supported")
+    end
+    return (names = names, )
+end
+
+function Jutul.devectorize_force(src::SourceTerm, X, meta, variant)
+    val = X[1]
+    f = src.fractional_flow
+    if variant == :all
+        if !isnothing(f)
+            f = tuple(X[2:end]...)
+        end
+    elseif variant == :control
+        # Do nothing
+    else
+        error("Variant $variant not supported")
+    end
+    return SourceTerm(src.cell, val, f, src.type)
+end
diff --git a/src/mrst_input.jl b/src/mrst_input.jl
index 71e08af2..eb4a4e9e 100644
--- a/src/mrst_input.jl
+++ b/src/mrst_input.jl
@@ -37,11 +37,6 @@ function reservoir_domain_from_mrst(name::String; extraout = false)
     @debug "File read complete. Unpacking data..."
     g = MRSTWrapMesh(exported["G"])
     has_trans = haskey(exported, "T") && length(exported["T"]) > 0
-    if haskey(exported, "N")
-        N = Int64.(exported["N"]')
-    else
-        N = nothing
-    end
 
     function get_vec(d)
         if isa(d, AbstractArray)
@@ -54,6 +49,19 @@ function reservoir_domain_from_mrst(name::String; extraout = false)
     perm = copy((exported["rock"]["perm"])')
     domain = reservoir_domain(g, porosity = poro, permeability = perm)
 
+    nf = number_of_faces(domain)
+    if haskey(exported, "N")
+        N = Int64.(exported["N"]')
+        nf = size(N, 2)
+        if nf != number_of_faces(domain)
+            d = Jutul.dim(g)
+            domain.entities[Faces()] = nf
+            domain[:areas, Faces()] = fill!(zeros(nf), NaN)
+            domain[:normals, Faces()] = fill!(zeros(d, nf), NaN)
+            domain[:face_centroids, Faces()] = fill!(zeros(d, nf), NaN)
+        end
+        domain[:neighbors, Faces()] = N
+    end
     # Deal with face data
     if has_trans
         @debug "Found precomputed transmissibilities, reusing"
@@ -156,6 +164,13 @@ function get_well_from_mrst_data(
             L = vec(segs["length"])
             D = vec(segs["diameter"])
             rough = vec(segs["roughness"])
+            n_segs = size(well_topo, 2)
+            if length(L) != n_segs
+                @warn "Inconsistent segments. Adding averages."
+                L = repeat([mean(L)], n_segs)
+                D = repeat([mean(D)], n_segs)
+                rough = repeat([mean(rough)], n_segs)
+            end
             @assert size(well_topo, 2) == length(L) == length(D) == length(rough)
             segment_models = map(SegmentWellBoreFrictionHB, L, rough, D)
         else
diff --git a/src/multicomponent/flux.jl b/src/multicomponent/flux.jl
index b6b28ae2..e8cb21aa 100644
--- a/src/multicomponent/flux.jl
+++ b/src/multicomponent/flux.jl
@@ -42,3 +42,15 @@ end
     end
     return q
 end
+
+function face_average_density(model::CompositionalModel, state, tpfa, phase)
+    ρ = state.PhaseMassDensities
+    s = state.Saturations
+    l = tpfa.left
+    r = tpfa.right
+    @inbounds s_l = s[phase, l]
+    @inbounds s_r = s[phase, r]
+    @inbounds ρ_l = ρ[phase, l]
+    @inbounds ρ_r = ρ[phase, r]
+    return (s_l*ρ_r + s_r*ρ_l)/max(s_l + s_r, 1e-8)
+end
diff --git a/src/variables/pvt.jl b/src/variables/pvt.jl
index 8cd7f5ae..81d263a3 100644
--- a/src/variables/pvt.jl
+++ b/src/variables/pvt.jl
@@ -123,7 +123,15 @@ function Jutul.default_parameter_values(data_domain, model, param::FluidVolume,
     return copy(vol)
 end
 
-struct Temperature <: ScalarVariable end
+Base.@kwdef struct Temperature{T} <: ScalarVariable
+    min::T = 0.0
+    max::T = Inf
+    max_rel::Union{T, Nothing} = nothing
+    max_abs::Union{T, Nothing} = nothing
+end
 
-Jutul.default_value(model, ::Temperature) = 303.15 # 30.15 C°
-Jutul.minimum_value(::Temperature) = 0.0
+Jutul.default_value(model, T::Temperature) = 303.15 # 30.15 C°
+Jutul.minimum_value(T::Temperature) = T.min
+Jutul.maximum_value(T::Temperature) = T.max
+Jutul.absolute_increment_limit(T::Temperature) = T.max_abs
+Jutul.relative_increment_limit(T::Temperature) = T.max_rel