Merge pull request idaholab#28785 from GiudGiud/PR_nl_opt

framework/include/partitioner/HierarchicalGridPartitioner.h

@@ -33,12 +33,19 @@ class HierarchicalGridPartitioner : public MoosePartitioner
 protected:
   virtual void _do_partition(MeshBase & mesh, const unsigned int n) override;
 
+  /// Mesh to partition
   MooseMesh & _mesh;
 
-  const unsigned int _nx_nodes;
-  const unsigned int _ny_nodes;
-  const unsigned int _nz_nodes;
-  const unsigned int _nx_procs;
-  const unsigned int _ny_procs;
-  const unsigned int _nz_procs;
+  /// Number of nodes in the X direction
+  unsigned int _nx_nodes;
+  /// Number of nodes in the Y direction
+  unsigned int _ny_nodes;
+  /// Number of nodes in the Z direction
+  unsigned int _nz_nodes;
+  /// Number of processors on each node in the X direction
+  unsigned int _nx_procs;
+  /// Number of processors on each node in the Y direction
+  unsigned int _ny_procs;
+  /// Number of processors on each node in the Z direction
+  unsigned int _nz_procs;
 };

framework/include/positions/MultiAppPositions.h

@@ -25,6 +25,13 @@ class MultiAppPositions : public Positions
   // MultiApp positions are not known at construction
   void initialize() override;
 
+  // Since we did not initialize on construction, initialize on setup by default
+  void initialSetup() override
+  {
+    initialize();
+    finalize();
+  }
+
   /// Whether to use the subapp mesh centroids to compute the positions, further translated by the positions
   const bool _use_apps_centroid;
 };

framework/include/positions/Positions.h

@@ -11,6 +11,7 @@
 
 // Moose includes
 #include "GeneralReporter.h"
+#include "KDTree.h"
 
 /**
  * Positions objects are under the hood Reporters.
@@ -99,6 +100,10 @@ class Positions : public GeneralReporter
   /// 4D storage for all the positions : space & time
   std::vector<std::vector<std::vector<std::vector<Point>>>> _positions_4d;
 
+  /// KDTree to be able to find the nearest position to a point in a fast and scalable manner
+  std::unique_ptr<KDTree> _initial_positions_kd_tree;
+  std::unique_ptr<KDTree> _positions_kd_tree;
+
   /// Whether generation of positions is distributed or not (and therefore needs a broadcast)
   bool _need_broadcast;
 

framework/include/transfers/MultiAppGeneralFieldNearestLocationTransfer.h

@@ -60,18 +60,22 @@ class MultiAppGeneralFieldNearestLocationTransfer : public MultiAppGeneralFieldT
       const std::vector<std::pair<Point, unsigned int>> & incoming_points,
       std::vector<std::pair<Real, Real>> & outgoing_vals);
 
+  /// Pre-compute the number of sources
   /// Number of KDTrees used to hold the locations and variable value data
-  unsigned int getNumSources() const;
+  void computeNumSources();
 
-  /// Transform a point towards the local frame
-  Point getPointInLocalSourceFrame(unsigned int i_from, const Point & pt) const;
+  /// Get the index of the app in the loop over the trees and the apps contributing data to each tree
+  unsigned int getAppIndex(unsigned int kdtree_index, unsigned int app_index) const;
 
   /// Number of applications which contributed nearest-locations to each KD-tree
   unsigned int getNumAppsPerTree() const;
 
   /// Number of divisions (nearest-positions or source mesh divisions) used when building KD-Trees
   unsigned int getNumDivisions() const;
 
+  /// Transform a point towards the local frame
+  Point getPointInLocalSourceFrame(unsigned int i_from, const Point & pt) const;
+
   /**
    * @brief Examine all spatial restrictions that could preclude this source from being
    *        a valid source for this point
@@ -107,4 +111,7 @@ class MultiAppGeneralFieldNearestLocationTransfer : public MultiAppGeneralFieldT
 
   /// Whether we can just use the local zero-indexed dof to get the value from the solution
   std::vector<bool> _use_zero_dof_for_value;
+
+  /// Number of KD-Trees to create
+  unsigned int _num_sources;
 };

framework/include/utils/KDTree.h

@@ -28,7 +28,7 @@ using ResultItem = std::pair<T, U>;
 class KDTree
 {
 public:
-  KDTree(std::vector<Point> & master_points, unsigned int max_leaf_size);
+  KDTree(const std::vector<Point> & master_points, unsigned int max_leaf_size);
 
   virtual ~KDTree() = default;
 

framework/src/partitioner/GridPartitioner.C

@@ -22,14 +22,23 @@ registerMooseObject("MooseApp", GridPartitioner);
 InputParameters
 GridPartitioner::validParams()
 {
-  // These two are in this order because they are from different systems
-  // so you have to apply _this_ system's second to override the base
   InputParameters params = MoosePartitioner::validParams();
 
+  MooseEnum method("manual automatic", "manual");
+  params.addParam<MooseEnum>(
+      "grid_computation",
+      method,
+      "Whether to determine the grid manually (using nx, ny and nz) or automatically. When using "
+      "the automatic mode, the user can impose a certain value for nx, ny or nz, and the automatic "
+      "factorization will adjust the number of processors in the other directions.");
+
   // Users specify how many processors they need along each direction
-  params.addParam<unsigned int>("nx", 1, "Number of processors in the X direction");
-  params.addParam<unsigned int>("ny", 1, "Number of processors in the Y direction");
-  params.addParam<unsigned int>("nz", 1, "Number of processors in the Z direction");
+  params.addParam<unsigned int>(
+      "nx", "Number of processors in the X direction. Defaults to 1 in manual mode");
+  params.addParam<unsigned int>(
+      "ny", "Number of processors in the Y direction. Defaults to 1 in manual mode");
+  params.addParam<unsigned int>(
+      "nz", "Number of processors in the Z direction. Defaults to 1 in manual mode");
 
   params.addClassDescription("Create a uniform grid that overlays the mesh to be partitioned.  "
                              "Assign all elements within each cell of the grid to the same "
@@ -66,22 +75,65 @@ GridPartitioner::_do_partition(MeshBase & mesh, const unsigned int /*n*/)
   const auto & max = bounding_box.max();
 
   auto dim = mesh.mesh_dimension();
-  //  Need to make sure the number of cells in the grid matches the number of procs to partition for
-  nx = getParam<unsigned int>("nx");
 
+  // Gather user parameters
+  nx = isParamValid("nx") ? getParam<unsigned int>("nx") : nx;
   if (dim >= 2)
-    ny = getParam<unsigned int>("ny");
-
+    ny = isParamValid("ny") ? getParam<unsigned int>("ny") : ny;
   if (dim == 3)
-    nz = getParam<unsigned int>("nz");
+    nz = isParamValid("nz") ? getParam<unsigned int>("nz") : nz;
+
+  // simple info message unused parameters as this can be normal: we could be partitioning
+  // a 2D mesh before extruding it to 3D. The nz parameter is needed for the 3D mesh
+  if (dim < 2 && isParamValid("ny") && getParam<unsigned int>("ny") > 1)
+    paramInfo("ny", "Parameter ignored as mesh is currently of dimension less than 2.");
+  if (dim < 3 && isParamValid("nz") && getParam<unsigned int>("nz") > 1)
+    paramInfo("nz", "Parameter ignored as mesh is currently of dimension less than 3.");
+
+  // User parameters, which should match the number of partitions needed
+  if (getParam<MooseEnum>("grid_computation") == "manual")
+  {
+    //  Need to make sure the number of grid cells matches the number of procs to partition for
+    if ((nx * ny * nz) != mesh.n_partitions())
+      mooseError("Number of grid cells (" + std::to_string(nx * ny * nz) +
+                 ") does not match the number of MPI processes (" +
+                 std::to_string(mesh.n_partitions()) + ")");
+  }
 
-  // We should compute a balanced factorization so
-  // that we can assign proper processors to each direction.
-  // I just want to make grid partitioner smarter.
-  if ((nx * ny * nz) != mesh.n_partitions())
+  else if (getParam<MooseEnum>("grid_computation") == "automatic")
   {
-    // Anybody knows we are living in a 3D space.
-    int dims[] = {0, 0, 0};
+    // remove over-constraint and tell user
+    if (nx * ny * nz > mesh.n_partitions())
+    {
+      nx = 0;
+      ny = 0;
+      nz = 0;
+      paramWarning("grid_computation",
+                   "User specified (nx,ny,nz) grid exceeded number of partitions, these parameters "
+                   "will be ignored.");
+    }
+    // 0 means no restriction on which number to choose
+    int dims[] = {isParamValid("nx") ? int(nx) : 0,
+                  isParamValid("ny") ? int(ny) : 0,
+                  isParamValid("nz") ? int(nz) : 0};
+
+    if ((dims[0] > 0 && dim == 1 && dims[0] != int(mesh.n_partitions())) ||
+        (dims[0] > 0 && dims[1] > 0 && dim == 2 && dims[0] * dims[1] != int(mesh.n_partitions())) ||
+        (dims[0] > 0 && dims[1] > 0 && dims[2] > 0 && dim == 3 &&
+         dims[0] * dims[1] * dims[2] != int(mesh.n_partitions())))
+    {
+      dims[0] = 0;
+      dims[1] = 0;
+      dims[2] = 0;
+      paramWarning("grid_computation",
+                   "User specified grid for the current dimension of the mesh (" +
+                       std::to_string(dim) + ") does not fit the number of partitions (" +
+                       std::to_string(mesh.n_partitions()) +
+                       ") and constrain the grid partitioner in every direction, these parameters "
+                       "will be ignored.");
+    }
+
+    // This will error if the factorization is not possible
     MPI_Dims_create(mesh.n_partitions(), dim, dims);
 
     nx = dims[0];
@@ -147,9 +199,9 @@ GridPartitioner::_do_partition(MeshBase & mesh, const unsigned int /*n*/)
     if (dim == 3)
       k = (coordz - min(2)) / hz;
 
-    mooseAssert(i < nx, "Index caculation is wrong along x direction");
-    mooseAssert(j < ny, "Index caculation is wrong along y direction");
-    mooseAssert(k < nz, "Index caculation is wrong along z direction");
+    mooseAssert(i < nx, "Index calculation is wrong along x direction");
+    mooseAssert(j < ny || ny == 0, "Index calculation is wrong along y direction");
+    mooseAssert(k < nz || nz == 0, "Index calculation is wrong along z direction");
     // Assign processor ID to current element
     elem_ptr->processor_id() = k * nx * ny + j * nx + i;
   }

framework/src/partitioner/HierarchicalGridPartitioner.C

@@ -27,19 +27,43 @@ HierarchicalGridPartitioner::validParams()
 {
   auto params = MoosePartitioner::validParams();
 
-  params.addRequiredRangeCheckedParam<unsigned int>(
-      "nx_nodes", "nx_nodes > 0", "Number of compute nodes in the X direction");
+  // Options for automatic grid computations
+  MooseEnum method("manual automatic", "manual");
+  params.addParam<MooseEnum>(
+      "nodes_grid_computation",
+      method,
+      "Whether to determine the compute node grid manually (using nx_nodes, ny_nodes and nz_nodes) "
+      "or automatically. When using the automatic mode, the user can impose a certain value for "
+      "nx, ny or nz, and the automatic factorization will adjust the number of processors in the "
+      "other directions.");
+  params.addParam<unsigned int>(
+      "number_nodes", "Number of nodes. Used for determining the node grid automatically");
+  params.addParam<MooseEnum>(
+      "processors_grid_computation",
+      method,
+      "Whether to determine the processors grid on each node manually (using nx_procs, ny_procs "
+      "and nz_procs) or automatically. When using the automatic mode, the user can impose a "
+      "certain value for nx, ny or nz, and the automatic factorization will adjust the number of "
+      "processors in the other directions.");
+  params.addParam<unsigned int>("number_procs_per_node",
+                                "Number of processors per node. Used for determining the processor "
+                                "grid on each node automatically");
+
+  // Node grid
   params.addRangeCheckedParam<unsigned int>(
-      "ny_nodes", 0, "ny_nodes >= 0", "Number of compute nodes in the Y direction");
+      "nx_nodes", "nx_nodes >= 1", "Number of compute nodes in the X direction");
   params.addRangeCheckedParam<unsigned int>(
-      "nz_nodes", 0, "nz_nodes >= 0", "Number of compute nodes in the Z direction");
+      "ny_nodes", "ny_nodes >= 1", "Number of compute nodes in the Y direction");
+  params.addRangeCheckedParam<unsigned int>(
+      "nz_nodes", "nz_nodes >= 1", "Number of compute nodes in the Z direction");
 
-  params.addRequiredRangeCheckedParam<unsigned int>(
-      "nx_procs", "nx_procs > 0", "Number of processors in the X direction within each node");
+  // Processor grid on each node
+  params.addRangeCheckedParam<unsigned int>(
+      "nx_procs", "nx_procs >= 1", "Number of processors in the X direction within each node");
   params.addRangeCheckedParam<unsigned int>(
-      "ny_procs", 0, "ny_procs >= 0", "Number of processors in the Y direction within each node");
+      "ny_procs", "ny_procs >= 1", "Number of processors in the Y direction within each node");
   params.addRangeCheckedParam<unsigned int>(
-      "nz_procs", 0, "nz_procs >= 0", "Number of processors in the Z direction within each node");
+      "nz_procs", "nz_procs >= 1", "Number of processors in the Z direction within each node");
 
   params.addClassDescription("Partitions a mesh into sub-partitions for each computational node"
                              " then into partitions within that node.  All partitions are made"
@@ -49,14 +73,7 @@ HierarchicalGridPartitioner::validParams()
 }
 
 HierarchicalGridPartitioner::HierarchicalGridPartitioner(const InputParameters & params)
-  : MoosePartitioner(params),
-    _mesh(*getCheckedPointerParam<MooseMesh *>("mesh")),
-    _nx_nodes(getParam<unsigned int>("nx_nodes")),
-    _ny_nodes(getParam<unsigned int>("ny_nodes")),
-    _nz_nodes(getParam<unsigned int>("nz_nodes")),
-    _nx_procs(getParam<unsigned int>("nx_procs")),
-    _ny_procs(getParam<unsigned int>("ny_procs")),
-    _nz_procs(getParam<unsigned int>("nz_procs"))
+  : MoosePartitioner(params), _mesh(*getCheckedPointerParam<MooseMesh *>("mesh"))
 {
 }
 
@@ -72,36 +89,87 @@ void
 HierarchicalGridPartitioner::_do_partition(MeshBase & mesh, const unsigned int /*n*/)
 {
   const auto dim = mesh.spatial_dimension();
-  if (_ny_procs == 0 && dim > 1)
-    paramError("ny_procs", "Required for ", dim, "D meshes");
-  if (_ny_nodes == 0 && dim > 1)
-    paramError("ny_nodes", "Required for ", dim, "D meshes");
-  if (_nz_procs == 0 && dim == 3)
-    paramError("nz_procs", "Required for 3D meshes");
-  if (_nz_nodes == 0 && dim == 3)
-    paramError("nz_nodes", "Required for 3D meshes");
 
-  auto total_nodes = _nx_nodes;
+  // Process user parameters
+  _nx_nodes = isParamValid("nx_nodes") ? getParam<unsigned int>("nx_nodes") : 0;
+  _ny_nodes = isParamValid("ny_nodes") ? getParam<unsigned int>("ny_nodes") : 0;
+  _nz_nodes = isParamValid("nz_nodes") ? getParam<unsigned int>("nz_nodes") : 0;
+  _nx_procs = isParamValid("nx_procs") ? getParam<unsigned int>("nx_procs") : 0;
+  _ny_procs = isParamValid("ny_procs") ? getParam<unsigned int>("ny_procs") : 0;
+  _nz_procs = isParamValid("nz_procs") ? getParam<unsigned int>("nz_procs") : 0;
+
+  if (getParam<MooseEnum>("nodes_grid_computation") == "manual")
+  {
+    if (_nx_nodes == 0)
+      paramError("nx_nodes", "Required for manual nodes grid specification");
+    if (_ny_nodes == 0 && dim > 1)
+      paramError("ny_nodes", "Required for ", dim, "D meshes");
+    if (_nz_nodes == 0 && dim == 3)
+      paramError("nz_nodes", "Required for 3D meshes");
+  }
+  else
+  {
+    if (!isParamValid("number_nodes"))
+      paramError("number_nodes", "Required for automatic nodes grid computation");
+    // 0 means no restriction on which number to choose
+    int dims[] = {int(_nx_nodes), int(_ny_nodes), int(_nz_nodes)};
+    // This will error if the factorization is not possible
+    MPI_Dims_create(getParam<unsigned int>("number_nodes"), dim, dims);
+
+    _nx_nodes = dims[0];
+    _ny_nodes = (dim >= 2) ? dims[1] : 0;
+    _nz_nodes = (dim == 3) ? dims[2] : 0;
+  }
 
+  if (getParam<MooseEnum>("processors_grid_computation") == "manual")
+  {
+    if (_nx_procs == 0)
+      paramError("nx_procs", "Required for manual processors grid specification");
+    if (_ny_procs == 0 && dim > 1)
+      paramError("ny_procs", "Required for ", dim, "D meshes");
+    if (_nz_procs == 0 && dim == 3)
+      paramError("nz_procs", "Required for 3D meshes");
+  }
+  else
+  {
+    if (!isParamValid("number_procs_per_node"))
+      paramError("number_procs_per_node", "Required for automatic processors grid computation");
+    // 0 means no restriction on which number to choose
+    int dims[] = {int(_nx_procs), int(_ny_procs), int(_nz_procs)};
+    // This will error if the factorization is not possible
+    MPI_Dims_create(getParam<unsigned int>("number_procs_per_node"), dim, dims);
+
+    _nx_procs = dims[0];
+    _ny_procs = (dim >= 2) ? dims[1] : 0;
+    _nz_procs = (dim == 3) ? dims[2] : 0;
+  }
+
+  // Sanity checks on both grids
+  auto total_nodes = _nx_nodes;
   if (mesh.spatial_dimension() >= 2)
     total_nodes *= _ny_nodes;
-
   if (mesh.spatial_dimension() == 3)
     total_nodes *= _nz_nodes;
+  if (isParamValid("number_nodes") && total_nodes != getParam<unsigned int>("number_nodes"))
+    paramError("number_nodes",
+               "Computed number of nodes (" + std::to_string(total_nodes) + ") does not match");
 
   auto procs_per_node = _nx_procs;
-
   if (mesh.spatial_dimension() >= 2)
     procs_per_node *= _ny_procs;
-
   if (mesh.spatial_dimension() == 3)
     procs_per_node *= _nz_procs;
+  if (isParamValid("number_procs_per_node") &&
+      procs_per_node != getParam<unsigned int>("number_procs_per_node"))
+    paramError("number_procs_per_node",
+               "Computed number of processors per node (" + std::to_string(procs_per_node) +
+                   ") does not match");
 
   if (procs_per_node * total_nodes != mesh.n_partitions())
-    mooseError("Partitioning in ",
-               name(),
-               " doesn't add up to the total number of processors: ",
-               procs_per_node * total_nodes);
+    mooseError("Partitioning creates ",
+               procs_per_node * total_nodes,
+               " partitions, which does not add up to the total number of processors: ",
+               mesh.n_partitions());
 
   // Figure out the physical bounds of the given mesh
   auto nodes_bounding_box = MeshTools::create_bounding_box(mesh);

framework/src/positions/FilePositions.C

@@ -49,4 +49,6 @@ FilePositions::FilePositions(const InputParameters & parameters) : Positions(par
     }
   }
   _initialized = true;
+  // Sort (if requested) and create KDTree
+  finalize();
 }