Line data    Source code

       1             : /**********************************************************************/
       2             : /*                    DO NOT MODIFY THIS HEADER                       */
       3             : /*             Swift, a Fourier spectral solver for MOOSE             */
       4             : /*                                                                    */
       5             : /*            Copyright 2024 Battelle Energy Alliance, LLC            */
       6             : /*                        ALL RIGHTS RESERVED                         */
       7             : /**********************************************************************/
       8             : 
       9             : #include "DomainPartitioner.h"
      10             : 
      11             : #include "GeneratedMesh.h"
      12             : #include "MooseApp.h"
      13             : 
      14             : #include "libmesh/mesh_tools.h"
      15             : #include "libmesh/elem.h"
      16             : 
      17             : registerMooseObject("SwiftApp", DomainPartitioner);
      18             : 
      19             : #include <memory>
      20             : 
      21             : InputParameters
      22           0 : DomainPartitioner::validParams()
      23             : {
      24             :   // These two are in this order because they are from different systems
      25             :   // so you have to apply _this_ system's second to override the base
      26           0 :   InputParameters params = MoosePartitioner::validParams();
      27             : 
      28             :   // Users specify how many processors they need along each direction
      29           0 :   params.addParam<unsigned int>("nx", 1, "Number of processors in the X direction");
      30           0 :   params.addParam<unsigned int>("ny", 1, "Number of processors in the Y direction");
      31           0 :   params.addParam<unsigned int>("nz", 1, "Number of processors in the Z direction");
      32             : 
      33           0 :   params.addClassDescription("Create a uniform grid that overlays the mesh to be partitioned.  "
      34             :                              "Assign all elements within each cell of the grid to the same "
      35             :                              "processor.");
      36             : 
      37           0 :   return params;
      38           0 : }
      39             : 
      40           0 : DomainPartitioner::DomainPartitioner(const InputParameters & params)
      41           0 :   : MoosePartitioner(params), _mesh(*getCheckedPointerParam<MooseMesh *>("mesh"))
      42             : {
      43           0 : }
      44             : 
      45           0 : DomainPartitioner::~DomainPartitioner() {}
      46             : 
      47             : std::unique_ptr<Partitioner>
      48           0 : DomainPartitioner::clone() const
      49             : {
      50           0 :   return _app.getFactory().clone(*this);
      51             : }
      52             : 
      53             : void
      54           0 : DomainPartitioner::_do_partition(MeshBase & mesh, const unsigned int /*n*/)
      55             : {
      56             :   // By default, there are one processor along each direction
      57             :   // nx: the number of processors along x direction
      58             :   // ny: the number of processors along y direction
      59             :   // nz: the number of processors along z direction
      60             :   unsigned int nx = 1, ny = 1, nz = 1;
      61             : 
      62             :   // Figure out the physical bounds of the given mesh
      63           0 :   auto bounding_box = MeshTools::create_bounding_box(mesh);
      64             :   const auto & min = bounding_box.min();
      65             :   const auto & max = bounding_box.max();
      66             : 
      67           0 :   auto dim = mesh.mesh_dimension();
      68             :   //  Need to make sure the number of cells in the grid matches the number of procs to partition for
      69           0 :   nx = getParam<unsigned int>("nx");
      70             : 
      71           0 :   if (dim >= 2)
      72           0 :     ny = getParam<unsigned int>("ny");
      73             : 
      74           0 :   if (dim == 3)
      75           0 :     nz = getParam<unsigned int>("nz");
      76             : 
      77             :   // We should compute a balanced factorization so
      78             :   // that we can assign proper processors to each direction.
      79             :   // I just want to make grid partitioner smarter.
      80           0 :   if ((nx * ny * nz) != mesh.n_partitions())
      81             :   {
      82             :     // Anybody knows we are living in a 3D space.
      83           0 :     int dims[] = {0, 0, 0};
      84           0 :     MPI_Dims_create(mesh.n_partitions(), dim, dims);
      85             : 
      86           0 :     nx = dims[0];
      87           0 :     if (dim >= 2)
      88           0 :       ny = dims[1];
      89           0 :     if (dim == 3)
      90           0 :       nz = dims[2];
      91             :   }
      92             : 
      93             :   // hx: grid interval along x direction
      94             :   // hy: grid interval along y direction
      95             :   // hz: grid interval along z direction
      96             :   // Lx: domain length along x direction
      97             :   // Ly: domain length along y direction
      98             :   // Lz: domain length along z direction
      99             :   Real hx = 1., hy = 1., hz = 1., Lx = 1., Ly = 1., Lz = 1.;
     100           0 :   Lx = max(0) - min(0);
     101           0 :   hx = Lx / nx;
     102           0 :   if (dim >= 2)
     103             :   {
     104           0 :     Ly = max(1) - min(1);
     105           0 :     hy = Ly / ny;
     106             :   }
     107             : 
     108           0 :   if (dim == 3)
     109             :   {
     110           0 :     Lz = max(2) - min(2);
     111           0 :     hz = Lz / nz;
     112             :   }
     113             : 
     114             :   // Processor coordinates along x, y, z directions
     115             :   unsigned int k = 0, j = 0, i = 0;
     116             :   // Coordinates for current element centroid
     117             :   Real coordx = 0, coordy = 0, coordz = 0;
     118             : 
     119             :   // Loop over all of the elements in the given mesh
     120           0 :   for (auto & elem_ptr : mesh.active_element_ptr_range())
     121             :   {
     122             :     // Find the element it lands in in the GeneratedMesh
     123           0 :     auto centroid = elem_ptr->vertex_average();
     124             : 
     125           0 :     coordx = centroid(0);
     126             :     mooseAssert(coordx >= min(0) && coordy <= max(0),
     127             :                 "element is outside of bounding box along x direction");
     128           0 :     if (dim >= 2)
     129             :     {
     130           0 :       coordy = centroid(1);
     131             :       mooseAssert(coordy >= min(1) && coordy <= max(1),
     132             :                   "element is outside of bounding box along y direction");
     133             :     }
     134           0 :     if (dim == 3)
     135             :     {
     136           0 :       coordz = centroid(2);
     137             :       mooseAssert(coordz >= min(2) && coordz <= max(2),
     138             :                   "element is outside of bounding box along z direction");
     139             :     }
     140             : 
     141             :     // Compute processor coordinates
     142             :     j = k = 0;
     143           0 :     i = (coordx - min(0)) / hx;
     144           0 :     if (dim >= 2)
     145           0 :       j = (coordy - min(1)) / hy;
     146           0 :     if (dim == 3)
     147           0 :       k = (coordz - min(2)) / hz;
     148             : 
     149             :     mooseAssert(i < nx, "Index caculation is wrong along x direction");
     150             :     mooseAssert(j < ny, "Index caculation is wrong along y direction");
     151             :     mooseAssert(k < nz, "Index caculation is wrong along z direction");
     152             :     // Assign processor ID to current element
     153           0 :     elem_ptr->processor_id() = k * nx * ny + j * nx + i;
     154           0 :   }
     155           0 : }