From eb937e8ba259ebce93bd66cbc5ad98bc990ff9ed Mon Sep 17 00:00:00 2001
From: Joakim Hove <joakim.hove@gmail.com>
Date: Tue, 1 Jun 2021 08:10:58 +0200
Subject: [PATCH] Use SegmentState as backing store for segment pressure

---
 opm/simulators/wells/WellState.cpp | 15 +++++++--------
 opm/simulators/wells/WellState.hpp |  9 ++++-----
 2 files changed, 11 insertions(+), 13 deletions(-)

diff --git a/opm/simulators/wells/WellState.cpp b/opm/simulators/wells/WellState.cpp
index d6f5ec48c..c44504172 100644
--- a/opm/simulators/wells/WellState.cpp
+++ b/opm/simulators/wells/WellState.cpp
@@ -527,7 +527,6 @@ void WellState::init(const std::vector<double>& cellPressures,
         nseg_ = nw;
         top_segment_index_.resize(nw);
         seg_number_.resize(nw);
-        seg_press_.resize(nw);
         for (int w = 0; w < nw; ++w) {
             top_segment_index_[w] = w;
             seg_number_[w] = 1; // Top segment is segment #1
@@ -821,7 +820,6 @@ void WellState::initWellStateMSWell(const std::vector<Well>& wells_ecl,
     const int np = pu.num_phases;
 
     top_segment_index_.clear();
-    seg_press_.clear();
     seg_rates_.clear();
     seg_number_.clear();
     nseg_ = 0;
@@ -840,7 +838,6 @@ void WellState::initWellStateMSWell(const std::vector<Well>& wells_ecl,
         if ( !well_ecl.isMultiSegment() ) { // not multi-segment well
             nseg_ += 1;
             seg_number_.push_back(1); // Assign single segment (top) as number 1.
-            seg_press_.push_back(0);
             for (int p = 0; p < np; ++p) {
                 seg_rates_.push_back(rates[p]);
             }
@@ -917,19 +914,18 @@ void WellState::initWellStateMSWell(const std::vector<Well>& wells_ecl,
             // improved during the solveWellEq process
             {
                 // top segment is always the first one, and its pressure is the well bhp
-                seg_press_.push_back(bhp(w));
-                const int top_segment = top_segment_index_[w];
+                auto& segment_pressure = this->segments(w).pressure;
+                segment_pressure[0] = bhp(w);
                 const auto& perf_press = this->perfPress(w);
                 for (int seg = 1; seg < well_nseg; ++seg) {
                     if ( !segment_perforations[seg].empty() ) {
                         const int first_perf = segment_perforations[seg][0];
-                        seg_press_.push_back(perf_press[first_perf]);
+                        segment_pressure[seg] = perf_press[first_perf];
                     } else {
                         // seg_press_.push_back(bhp); // may not be a good decision
                         // using the outlet segment pressure // it needs the ordering is correct
                         const int outlet_seg = segment_set[seg].outletSegment();
-                        seg_press_.push_back(
-                            seg_press_[top_segment + segment_set.segmentNumberToIndex(outlet_seg)]);
+                        segment_pressure[seg] = segment_pressure[segment_set.segmentNumberToIndex(outlet_seg)];
                     }
                 }
             }
@@ -946,6 +942,9 @@ void WellState::initWellStateMSWell(const std::vector<Well>& wells_ecl,
             if (well.getStatus() == Well::Status::SHUT)
                 continue;
 
+            if ( !well.isMultiSegment() )
+                continue;
+
             const auto& it = prev_well_state->wellMap().find( wells_ecl[w].name() );
             if (it != end) { // the well is found in the prev_well_state
                 // TODO: the well with same name can change a lot, like they might not have same number of segments
diff --git a/opm/simulators/wells/WellState.hpp b/opm/simulators/wells/WellState.hpp
index 7d30eb447..3103699d9 100644
--- a/opm/simulators/wells/WellState.hpp
+++ b/opm/simulators/wells/WellState.hpp
@@ -210,14 +210,14 @@ public:
 
     double * segPress(std::size_t well_index)
     {
-        const int top_segment_index = this->top_segment_index_[well_index];
-        return &seg_press_[top_segment_index];
+        auto& segment = this->segments(well_index);
+        return segment.pressure.data();
     }
 
     const double * segPress(std::size_t well_index) const
     {
-        const int top_segment_index = this->top_segment_index_[well_index];
-        return &seg_press_[top_segment_index];
+        const auto& segment = this->segments(well_index);
+        return segment.pressure.data();
     }
 
 
@@ -503,7 +503,6 @@ private:
     // MS well related
     // for StandardWell, the number of segments will be one
     std::vector<double> seg_rates_;
-    std::vector<double> seg_press_;
     // the index of the top segments, which is used to locate the
     // multisegment well related information in WellState
     std::vector<int> top_segment_index_;