From 4d0b46677571d0f1856cde402a1aa0e9da5403ce Mon Sep 17 00:00:00 2001
From: Henry Weller <http://cfd.direct>
Date: Wed, 29 May 2024 12:34:15 +0100
Subject: [PATCH] decompositionMethods: Added checks for multiple constraints

Multi-constraint load-balancing is only supported by the Zoltan rcb method and
ParMETIS.
---
 .../decompositionMethod/decompositionMethod.C | 54 +++++++++++++++++--
 .../decompositionMethod/decompositionMethod.H | 18 +++++++
 .../hierarchical/hierarchical.C               |  6 ++-
 .../decompositionMethods/random/random.C      |  2 +
 .../decompositionMethods/simple/simple.C      |  2 +
 src/parallel/decompose/metis/metis.C          |  7 ++-
 src/parallel/decompose/parMetis/parMetis.C    | 25 +++++----
 src/parallel/decompose/parMetis/parMetis.H    |  1 +
 src/parallel/decompose/ptscotch/ptscotch.C    |  3 ++
 src/parallel/decompose/scotch/scotch.C        |  5 ++
 src/parallel/decompose/zoltan/zoltan.C        | 16 +++++-
 11 files changed, 121 insertions(+), 18 deletions(-)

diff --git a/src/parallel/decompose/decompositionMethods/decompositionMethod/decompositionMethod.C b/src/parallel/decompose/decompositionMethods/decompositionMethod/decompositionMethod.C
index 207b94defe..bb2593fb8f 100644
--- a/src/parallel/decompose/decompositionMethods/decompositionMethod/decompositionMethod.C
+++ b/src/parallel/decompose/decompositionMethods/decompositionMethod/decompositionMethod.C
@@ -49,6 +49,51 @@ namespace Foam
 }
 
 
+// * * * * * * * * * * * * * Protected Member Functions  * * * * * * * * * * //
+
+Foam::label Foam::decompositionMethod::nWeights
+(
+    const pointField& points,
+    const scalarField& pointWeights
+) const
+{
+    const label localnWeights =
+        points.size() ? pointWeights.size()/points.size() : 0;
+
+    const label nWeights = returnReduce(localnWeights, maxOp<label>());
+
+    if (localnWeights && localnWeights != nWeights)
+    {
+        FatalErrorInFunction
+            << "Number of weights on this processor " << localnWeights
+            << " does not equal the maximum number of weights " << nWeights
+            << exit(FatalError);
+    }
+
+    return nWeights;
+}
+
+
+Foam::label Foam::decompositionMethod::checkWeights
+(
+    const pointField& points,
+    const scalarField& pointWeights
+) const
+{
+    const label nWeights = this->nWeights(points, pointWeights);
+
+    if (nWeights > 1)
+    {
+        FatalErrorInFunction
+            << "decompositionMethod " << type()
+            << " does not support multiple constraints"
+            << exit(FatalError);
+    }
+
+    return nWeights;
+}
+
+
 // * * * * * * * * * * * * * * * * Constructors  * * * * * * * * * * * * * * //
 
 Foam::decompositionMethod::decompositionMethod
@@ -758,7 +803,8 @@ Foam::labelList Foam::decompositionMethod::decompose
 )
 {
     // Any weights specified?
-    label nWeights = returnReduce(cellWeights.size(), sumOp<label>());
+    const bool hasWeights =
+        returnReduce(cellWeights.size(), sumOp<label>()) > 0;
 
     // Any processor sets?
     label nProcSets = 0;
@@ -796,7 +842,7 @@ Foam::labelList Foam::decompositionMethod::decompose
     {
         // No constraints, possibly weights
 
-        if (nWeights > 0)
+        if (hasWeights)
         {
             finalDecomp = decompose
             (
@@ -860,11 +906,11 @@ Foam::labelList Foam::decompositionMethod::decompose
 
         scalarField regionWeights(localRegion.nLocalRegions(), 0);
 
-        if (nWeights > 0)
+        if (hasWeights)
         {
             forAll(localRegion, celli)
             {
-                label regionI = localRegion[celli];
+                const label regionI = localRegion[celli];
 
                 regionWeights[regionI] += cellWeights[celli];
             }
diff --git a/src/parallel/decompose/decompositionMethods/decompositionMethod/decompositionMethod.H b/src/parallel/decompose/decompositionMethods/decompositionMethod/decompositionMethod.H
index a27300a6bc..ab30813f7f 100644
--- a/src/parallel/decompose/decompositionMethods/decompositionMethod/decompositionMethod.H
+++ b/src/parallel/decompose/decompositionMethods/decompositionMethod/decompositionMethod.H
@@ -61,6 +61,24 @@ protected:
         PtrList<decompositionConstraint> constraints_;
 
 
+    // Protected member functions
+
+        //- Return the number of weights per point
+        label nWeights
+        (
+            const pointField& points,
+            const scalarField& pointWeights
+        ) const;
+
+        //- Check the weights against the points
+        //  and return the number of weights per point
+        label checkWeights
+        (
+            const pointField& points,
+            const scalarField& pointWeights
+        ) const;
+
+
 public:
 
     //- Runtime type information
diff --git a/src/parallel/decompose/decompositionMethods/hierarchical/hierarchical.C b/src/parallel/decompose/decompositionMethods/hierarchical/hierarchical.C
index 4bf71c3026..4060f9f76a 100644
--- a/src/parallel/decompose/decompositionMethods/hierarchical/hierarchical.C
+++ b/src/parallel/decompose/decompositionMethods/hierarchical/hierarchical.C
@@ -731,8 +731,7 @@ Foam::labelList Foam::decompositionMethods::hierarchical::decompose
     const scalarField& weights
 )
 {
-    // construct a list for the final result
-    labelList finalDecomp(points.size(), 0);
+    checkWeights(points, weights);
 
     // Start off with every point sorted onto itself.
     labelList slice(points.size());
@@ -751,6 +750,9 @@ Foam::labelList Foam::decompositionMethods::hierarchical::decompose
 
     const label sizeTol = max(1, label(1e-3*allSize/nProcessors_));
 
+    // Construct a list for the final result
+    labelList finalDecomp(points.size(), 0);
+
     // Sort recursive
     sortComponent
     (
diff --git a/src/parallel/decompose/decompositionMethods/random/random.C b/src/parallel/decompose/decompositionMethods/random/random.C
index a2d508604d..407aa9b718 100644
--- a/src/parallel/decompose/decompositionMethods/random/random.C
+++ b/src/parallel/decompose/decompositionMethods/random/random.C
@@ -77,6 +77,8 @@ Foam::labelList Foam::decompositionMethods::random::decompose
     const scalarField& pointWeights
 )
 {
+    checkWeights(points, pointWeights);
+
     randomGenerator rndGen(seed_);
 
     // Simple random integer. No guarantees about balance.
diff --git a/src/parallel/decompose/decompositionMethods/simple/simple.C b/src/parallel/decompose/decompositionMethods/simple/simple.C
index 1f58740bf1..30b97e5b2d 100644
--- a/src/parallel/decompose/decompositionMethods/simple/simple.C
+++ b/src/parallel/decompose/decompositionMethods/simple/simple.C
@@ -404,6 +404,8 @@ Foam::labelList Foam::decompositionMethods::simple::decompose
     const scalarField& weights
 )
 {
+    checkWeights(points, weights);
+
     if (!Pstream::parRun())
     {
         return decomposeOneProc(points, weights);
diff --git a/src/parallel/decompose/metis/metis.C b/src/parallel/decompose/metis/metis.C
index 15745d86a9..fe0468f374 100644
--- a/src/parallel/decompose/metis/metis.C
+++ b/src/parallel/decompose/metis/metis.C
@@ -53,7 +53,6 @@ Foam::label Foam::decompositionMethods::metis::decompose
     const labelList& adjncy,
     const labelList& xadj,
     const scalarField& cellWeights,
-
     labelList& decomp
 )
 {
@@ -217,6 +216,8 @@ Foam::labelList Foam::decompositionMethods::metis::decompose
             << exit(FatalError);
     }
 
+    checkWeights(points, pointWeights);
+
     // Make Metis CSR (Compressed Storage Format) storage
     //   adjncy      : contains neighbours (= edges in graph)
     //   xadj(celli) : start of information in adjncy for celli
@@ -260,6 +261,8 @@ Foam::labelList Foam::decompositionMethods::metis::decompose
             << exit(FatalError);
     }
 
+    checkWeights(regionPoints, regionWeights);
+
     // Make Metis CSR (Compressed Storage Format) storage
     //   adjncy      : contains neighbours (= edges in graph)
     //   xadj(celli) : start of information in adjncy for celli
@@ -299,6 +302,8 @@ Foam::labelList Foam::decompositionMethods::metis::decompose
             << ")." << exit(FatalError);
     }
 
+    checkWeights(cellCentres, cellWeights);
+
     // Make Metis CSR (Compressed Storage Format) storage
     //   adjncy      : contains neighbours (= edges in graph)
     //   xadj(celli) : start of information in adjncy for celli
diff --git a/src/parallel/decompose/parMetis/parMetis.C b/src/parallel/decompose/parMetis/parMetis.C
index 1fb051441d..37c1b045ff 100644
--- a/src/parallel/decompose/parMetis/parMetis.C
+++ b/src/parallel/decompose/parMetis/parMetis.C
@@ -55,6 +55,7 @@ Foam::label Foam::decompositionMethods::parMetis::decompose
     const labelList& xadj,
     const labelList& adjncy,
     const pointField& cellCentres,
+    const label nWeights,
     const labelList& cellWeights,
     const labelList& faceWeights,
     labelList& decomp
@@ -63,9 +64,6 @@ Foam::label Foam::decompositionMethods::parMetis::decompose
     // C style numbering
     label numFlag = 0;
 
-    // Number of weights or balance constraints
-    label nWeights = cellWeights.size()/cellCentres.size();
-
     List<real_t> processorWeights;
 
     if (processorWeights_.size())
@@ -140,7 +138,7 @@ Foam::label Foam::decompositionMethods::parMetis::decompose
             const_cast<label*>(adjwgtPtr),
             &wgtFlag,
             &numFlag,
-            &nWeights,
+            const_cast<label*>(&nWeights),
             &nProcessors_,
             processorWeights.begin(),
             ubvec.begin(),
@@ -177,7 +175,7 @@ Foam::label Foam::decompositionMethods::parMetis::decompose
             &numFlag,
             &nDims,
             xyz.begin(),
-            &nWeights,
+            const_cast<label*>(&nWeights),
             &nProcessors_,
             processorWeights.begin(),
             ubvec.begin(),
@@ -202,7 +200,7 @@ Foam::label Foam::decompositionMethods::parMetis::decompose
             const_cast<label*>(adjwgtPtr),
             &wgtFlag,
             &numFlag,
-            &nWeights,
+            const_cast<label*>(&nWeights),
             &nProcessors_,
             processorWeights.begin(),
             ubvec.begin(),
@@ -329,13 +327,16 @@ Foam::labelList Foam::decompositionMethods::parMetis::decompose
         cellCells
     );
 
+    const label nWeights = this->nWeights(points, pointWeights);
+
     labelList decomp;
     decompose
     (
         cellCells.offsets(),
         cellCells.m(),
         points,
-        scaleWeights(pointWeights, pointWeights.size()/points.size()),
+        nWeights,
+        scaleWeights(pointWeights, nWeights),
         labelList(),
         decomp
     );
@@ -381,6 +382,8 @@ Foam::labelList Foam::decompositionMethods::parMetis::decompose
         cellCells
     );
 
+    const label nWeights = this->nWeights(regionPoints, pointWeights);
+
     // Decompose using weights
     labelList decomp;
     decompose
@@ -388,7 +391,8 @@ Foam::labelList Foam::decompositionMethods::parMetis::decompose
         cellCells.m(),
         cellCells.offsets(),
         regionPoints,
-        scaleWeights(pointWeights, pointWeights.size()/regionPoints.size()),
+        nWeights,
+        scaleWeights(pointWeights, nWeights),
         labelList(),
         decomp
     );
@@ -433,13 +437,16 @@ Foam::labelList Foam::decompositionMethods::parMetis::decompose
     //   xadj(celli) : start of information in adjncy for celli
     CompactListList<label> cellCells(globalCellCells);
 
+    const label nWeights = this->nWeights(cellCentres, cellWeights);
+
     labelList decomp;
     decompose
     (
         cellCells.offsets(),
         cellCells.m(),
         cellCentres,
-        scaleWeights(cellWeights, cellWeights.size()/cellCentres.size()),
+        nWeights,
+        scaleWeights(cellWeights, nWeights),
         labelList(),
         decomp
     );
diff --git a/src/parallel/decompose/parMetis/parMetis.H b/src/parallel/decompose/parMetis/parMetis.H
index 30b8b126a8..a85e3e7c59 100644
--- a/src/parallel/decompose/parMetis/parMetis.H
+++ b/src/parallel/decompose/parMetis/parMetis.H
@@ -115,6 +115,7 @@ class parMetis
             const labelList& xadj,
             const labelList& adjncy,
             const pointField& cellCentres,
+            const label nWeights,
             const labelList& cellWeights,
             const labelList& faceWeights,
             labelList& finalDecomp
diff --git a/src/parallel/decompose/ptscotch/ptscotch.C b/src/parallel/decompose/ptscotch/ptscotch.C
index 5756384d18..e43a8da4e9 100644
--- a/src/parallel/decompose/ptscotch/ptscotch.C
+++ b/src/parallel/decompose/ptscotch/ptscotch.C
@@ -420,6 +420,7 @@ Foam::labelList Foam::decompositionMethods::ptscotch::decompose
             << exit(FatalError);
     }
 
+    checkWeights(points, pointWeights);
 
     // Make Metis CSR (Compressed Storage Format) storage
     //   adjncy      : contains neighbours (= edges in graph)
@@ -465,6 +466,7 @@ Foam::labelList Foam::decompositionMethods::ptscotch::decompose
             << exit(FatalError);
     }
 
+    checkWeights(agglomPoints, pointWeights);
 
     // Make Metis CSR (Compressed Storage Format) storage
     //   adjncy      : contains neighbours (= edges in graph)
@@ -517,6 +519,7 @@ Foam::labelList Foam::decompositionMethods::ptscotch::decompose
             << ")." << exit(FatalError);
     }
 
+    checkWeights(cellCentres, cellWeights);
 
     // Make Metis CSR (Compressed Storage Format) storage
     //   adjncy      : contains neighbours (= edges in graph)
diff --git a/src/parallel/decompose/scotch/scotch.C b/src/parallel/decompose/scotch/scotch.C
index 1c3da2cbc4..6a53254356 100644
--- a/src/parallel/decompose/scotch/scotch.C
+++ b/src/parallel/decompose/scotch/scotch.C
@@ -445,6 +445,8 @@ Foam::labelList Foam::decompositionMethods::scotch::decompose
             << exit(FatalError);
     }
 
+    checkWeights(points, pointWeights);
+
     // Calculate local or global (if Pstream::parRun()) connectivity
     CompactListList<label> cellCells;
     calcCellCells
@@ -487,6 +489,8 @@ Foam::labelList Foam::decompositionMethods::scotch::decompose
             << exit(FatalError);
     }
 
+    checkWeights(agglomPoints, pointWeights);
+
     // Calculate local or global (if Pstream::parRun()) connectivity
     CompactListList<label> cellCells;
     calcCellCells
@@ -536,6 +540,7 @@ Foam::labelList Foam::decompositionMethods::scotch::decompose
             << ")." << exit(FatalError);
     }
 
+    checkWeights(cellCentres, cellWeights);
 
     // Make Metis CSR (Compressed Storage Format) storage
     //   adjncy      : contains neighbours (= edges in graph)
diff --git a/src/parallel/decompose/zoltan/zoltan.C b/src/parallel/decompose/zoltan/zoltan.C
index 6fe081dcb6..79d33219bd 100644
--- a/src/parallel/decompose/zoltan/zoltan.C
+++ b/src/parallel/decompose/zoltan/zoltan.C
@@ -280,7 +280,7 @@ Foam::label Foam::decompositionMethods::zoltan::decompose
 
     struct Zoltan_Struct *zz = Zoltan_Create(PstreamGlobals::MPI_COMM_FOAM);
 
-    const int nWeights = pWeights.size()/points.size();
+    const int nWeights = this->nWeights(points, pWeights);
 
     // Set internal parameters
     Zoltan_Set_Param(zz, "return_lists", "export");
@@ -291,8 +291,10 @@ Foam::label Foam::decompositionMethods::zoltan::decompose
     Zoltan_Set_Param(zz, "debug_level", "0");
     Zoltan_Set_Param(zz, "imbalance_tol", "1.05");
 
+    word lb_method("graph");
+
     // Set default method parameters
-    Zoltan_Set_Param(zz, "lb_method", "graph");
+    Zoltan_Set_Param(zz, "lb_method", lb_method.c_str());
     Zoltan_Set_Param(zz, "lb_approach", "repartition");
 
     if (decompositionDict_.found("zoltanCoeffs"))
@@ -300,6 +302,8 @@ Foam::label Foam::decompositionMethods::zoltan::decompose
         const dictionary& coeffsDict_ =
             decompositionDict_.subDict("zoltanCoeffs");
 
+        coeffsDict_.readIfPresent("lb_method", lb_method);
+
         forAllConstIter(IDLList<entry>, coeffsDict_, iter)
         {
             if (!iter().isDict())
@@ -318,6 +322,14 @@ Foam::label Foam::decompositionMethods::zoltan::decompose
         }
     }
 
+    if (nWeights > 1 && lb_method != "rcb")
+    {
+        FatalIOErrorInFunction(decompositionDict_)
+            << "Multiple constraints specified for lb_method " << lb_method
+            << " but is only supported by the rcb method"
+            << exit(FatalIOError);
+    }
+
     void* pointsPtr = &const_cast<pointField&>(points);
     Zoltan_Set_Num_Obj_Fn(zz, get_number_of_vertices, pointsPtr);