From 01432879d2d1a1d6060760a7378dcff3ed51b813 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Cezary=20Skrzy=C5=84ski?=
 <cezary.skrzynski@ng-analytics.com>
Date: Mon, 8 Aug 2022 18:56:55 +0200
Subject: [PATCH] #1838: TemperedWMin: compose total work model on top of
 existing one

---
 .../vrt/collection/balance/temperedwmin/temperedwmin.cc   | 8 +++++++-
 src/vt/vrt/collection/balance/temperedwmin/temperedwmin.h | 5 +++++
 2 files changed, 12 insertions(+), 1 deletion(-)

diff --git a/src/vt/vrt/collection/balance/temperedwmin/temperedwmin.cc b/src/vt/vrt/collection/balance/temperedwmin/temperedwmin.cc
index 5acf1271c3..20715c7154 100644
--- a/src/vt/vrt/collection/balance/temperedwmin/temperedwmin.cc
+++ b/src/vt/vrt/collection/balance/temperedwmin/temperedwmin.cc
@@ -43,7 +43,9 @@
 
 #include "vt/vrt/collection/balance/temperedwmin/temperedwmin.h"
 
+#include "vt/vrt/collection/balance/lb_invoke/lb_manager.h"
 #include "vt/vrt/collection/balance/model/load_model.h"
+#include "vt/vrt/collection/balance/model/weighted_communication_volume.h"
 
 namespace vt { namespace vrt { namespace collection { namespace lb {
 
@@ -93,13 +95,17 @@ void TemperedWMin::inputParams(balance::SpecEntry* spec) {
     "TemperedWMin::inputParams: alpha={}, beta={}, gamma={}\n",
     alpha_, beta_, gamma_
   );
+
+  total_work_model_ = std::make_unique<balance::WeightedCommunicationVolume>(
+    theLBManager()->getLoadModel(), alpha_, beta_, gamma_
+  );
 }
 
 TimeType TemperedWMin::getModeledValue(const elm::ElementIDStruct& obj) {
   balance::PhaseOffset when =
       {balance::PhaseOffset::NEXT_PHASE, balance::PhaseOffset::WHOLE_PHASE};
 
-  return load_model_->getModeledWork(obj, when);
+  return total_work_model_->getModeledWork(obj, when);
 }
 
 }}}} // namespace vt::vrt::collection::lb
diff --git a/src/vt/vrt/collection/balance/temperedwmin/temperedwmin.h b/src/vt/vrt/collection/balance/temperedwmin/temperedwmin.h
index 68e6bfa451..36f95b6541 100644
--- a/src/vt/vrt/collection/balance/temperedwmin/temperedwmin.h
+++ b/src/vt/vrt/collection/balance/temperedwmin/temperedwmin.h
@@ -44,8 +44,11 @@
 #if !defined INCLUDED_VT_VRT_COLLECTION_BALANCE_TEMPEREDWMIN_TEMPEREDWMIN_H
 #define INCLUDED_VT_VRT_COLLECTION_BALANCE_TEMPEREDWMIN_TEMPEREDWMIN_H
 
+#include "vt/vrt/collection/balance/model/load_model.h"
 #include "vt/vrt/collection/balance/temperedlb/temperedlb.h"
 
+#include <memory>
+
 namespace vt { namespace vrt { namespace collection { namespace lb {
 
 struct TemperedWMin : TemperedLB {
@@ -64,6 +67,8 @@ struct TemperedWMin : TemperedLB {
   TimeType getModeledValue(const elm::ElementIDStruct& obj) override;
 
 private:
+  std::unique_ptr<balance::LoadModel> total_work_model_ = nullptr;
+
   double alpha_ = 1.0;
   double beta_  = 0.0;
   double gamma_ = 0.0;