From fce7247da5609509198b2dc16f2ae9f27664fe36 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Atgeirr=20Fl=C3=B8=20Rasmussen?= <atgeirr@sintef.no>
Date: Mon, 12 Oct 2015 12:49:30 +0200
Subject: [PATCH] Update the SolutionState's qs and bhp members.

This way they can be used instead of repeatedly extracting the top segment data,
which also could make it possible to use more of the existing implementation.
---
 opm/autodiff/BlackoilMultiSegmentModel.hpp    |  2 ++
 .../BlackoilMultiSegmentModel_impl.hpp        | 33 ++++++++++++++-----
 2 files changed, 27 insertions(+), 8 deletions(-)

diff --git a/opm/autodiff/BlackoilMultiSegmentModel.hpp b/opm/autodiff/BlackoilMultiSegmentModel.hpp
index 212a8251d..f8292c364 100644
--- a/opm/autodiff/BlackoilMultiSegmentModel.hpp
+++ b/opm/autodiff/BlackoilMultiSegmentModel.hpp
@@ -202,6 +202,8 @@ namespace Opm {
 
         const std::vector<WellMultiSegmentConstPtr> wells_multisegment_;
 
+        std::vector<int> top_well_segments_;
+
         // return wells object
         // TODO: remove this wells structure
         using Base::wells;
diff --git a/opm/autodiff/BlackoilMultiSegmentModel_impl.hpp b/opm/autodiff/BlackoilMultiSegmentModel_impl.hpp
index f5acd1141..04fad2597 100644
--- a/opm/autodiff/BlackoilMultiSegmentModel_impl.hpp
+++ b/opm/autodiff/BlackoilMultiSegmentModel_impl.hpp
@@ -96,13 +96,15 @@ namespace Opm {
     BlackoilMultiSegmentModel<Grid>::
     prepareStep(const double dt,
                 ReservoirState& reservoir_state,
-                WellState& /* well_state */)
+                WellState& well_state)
     {
         pvdt_ = geo_.poreVolume() / dt;
         if (active_[Gas]) {
             updatePrimalVariableFromState(reservoir_state);
         }
 
+        top_well_segments_ = well_state.topSegmentLoc();
+
         //TODO: handle the volume related.
     }
 
@@ -154,25 +156,40 @@ namespace Opm {
         }
     }
 
-    // TODO: using the original Qs for the segment rates for now.
-    // TODO: using the original Bhp for the segment pressures for now.
-    //
+
 
 
     template <class Grid>
     void
     BlackoilMultiSegmentModel<Grid>::variableStateExtractWellsVars(const std::vector<int>& indices,
-                                                                          std::vector<ADB>& vars,
-                                                                          SolutionState& state) const
+                                                                   std::vector<ADB>& vars,
+                                                                   SolutionState& state) const
     {
+        // TODO: using the original Qs for the segment rates for now.
+        // TODO: using the original Bhp for the segment pressures for now.
+
         // segment phase rates in surface volume
         state.segqs = std::move(vars[indices[Qs]]);
 
         // segment pressures
         state.segp = std::move(vars[indices[Bhp]]);
 
-        // TODO: should the bhp and qs also be updated?
-
+        // The qs and bhp are no longer primary variables, but could
+        // still be used in computations. They are identical to the
+        // pressures and flows of the top segments.
+        const int np = numPhases();
+        const int ns = state.segp.size();
+        const int nw = top_well_segments_.size();
+        state.qs = ADB::constant(ADB::V::Zero(np*nw));
+        for (int phase = 0; phase < np; ++phase) {
+            // Extract segment fluxes for this phase (ns consecutive elements).
+            ADB segqs_phase = subset(state.segqs, Span(ns, 1, ns*phase));
+            // Extract top segment fluxes (= well fluxes)
+            ADB wellqs_phase = subset(segqs_phase, top_well_segments_);
+            // Expand to full size of qs (which contains all phases) and add.
+            state.qs += superset(wellqs_phase, Span(nw, 1, nw*phase), nw*np);
+        }
+        state.bhp = subset(state.segp, top_well_segments_);
     }