diff --git a/framework/mesh/mesh_generator/orthogonal_mesh_generator.cc b/framework/mesh/mesh_generator/orthogonal_mesh_generator.cc index f00360ae6d..1da289af53 100644 --- a/framework/mesh/mesh_generator/orthogonal_mesh_generator.cc +++ b/framework/mesh/mesh_generator/orthogonal_mesh_generator.cc @@ -2,16 +2,467 @@ // SPDX-License-Identifier: MIT #include "framework/mesh/mesh_generator/orthogonal_mesh_generator.h" +#include "framework/graphs/graph_partitioner.h" +#include "framework/mesh/mesh_continuum/mesh_continuum.h" +#include "framework/mpi/mpi_utils.h" #include "framework/object_factory.h" +#include "framework/runtime.h" #include "framework/logging/log.h" +#include +#include +#include +#include namespace opensn { +namespace +{ + +struct OrthoCellInfo +{ + unsigned int dimension = 0; + std::array cell_counts = {1, 1, 1}; + std::array node_counts = {2, 2, 2}; +}; + +OrthoCellInfo +MakeOrthoCellInfo(const std::vector>& node_sets) +{ + OrthoCellInfo info; + info.dimension = static_cast(node_sets.size()); + + if (info.dimension == 1) + { + info.cell_counts = {1, 1, node_sets[0].size() - 1}; + info.node_counts = {1, 1, node_sets[0].size()}; + } + else if (info.dimension == 2) + { + info.cell_counts = {node_sets[0].size() - 1, node_sets[1].size() - 1, 1}; + info.node_counts = {node_sets[0].size(), node_sets[1].size(), 1}; + } + else if (info.dimension == 3) + { + info.cell_counts = {node_sets[0].size() - 1, node_sets[1].size() - 1, node_sets[2].size() - 1}; + info.node_counts = {node_sets[0].size(), node_sets[1].size(), node_sets[2].size()}; + } + else + throw std::logic_error("Unsupported orthogonal mesh dimension."); + + return info; +} + +uint64_t +CellGlobalID(const OrthoCellInfo& info, const size_t i, const size_t j, const size_t k) +{ + const auto [nx, ny, nz] = info.cell_counts; + if (info.dimension == 1) + return k; + if (info.dimension == 2) + return j * nx + i; + return (j * nx + i) * nz + k; +} + +std::array +CellIJK(const OrthoCellInfo& info, const uint64_t cell_global_id) +{ + const auto [nx, ny, nz] = info.cell_counts; + if (info.dimension == 1) + return {0, 0, static_cast(cell_global_id)}; + if (info.dimension == 2) + return {static_cast(cell_global_id % nx), static_cast(cell_global_id / nx), 0}; + + const auto xy = static_cast(cell_global_id / nz); + return {xy % nx, xy / nx, static_cast(cell_global_id % nz)}; +} + +uint64_t +VertexGlobalID(const OrthoCellInfo& info, const size_t i, const size_t j, const size_t k) +{ + const auto [nx, ny, nz] = info.node_counts; + if (info.dimension == 1) + return k; + if (info.dimension == 2) + return j * nx + i; + return (j * nx + i) * nz + k; +} + +Vector3 +Vertex(const std::vector>& node_sets, + const unsigned int dimension, + const size_t i, + const size_t j, + const size_t k) +{ + if (dimension == 1) + return {0.0, 0.0, node_sets[0][k]}; + if (dimension == 2) + return {node_sets[0][i], node_sets[1][j], 0.0}; + return {node_sets[0][i], node_sets[1][j], node_sets[2][k]}; +} + +Vector3 +CellCentroid(const std::vector>& node_sets, + const unsigned int dimension, + const size_t i, + const size_t j, + const size_t k) +{ + if (dimension == 1) + return {0.0, 0.0, 0.5 * (node_sets[0][k] + node_sets[0][k + 1])}; + if (dimension == 2) + return {0.5 * (node_sets[0][i] + node_sets[0][i + 1]), + 0.5 * (node_sets[1][j] + node_sets[1][j + 1]), + 0.0}; + return {0.5 * (node_sets[0][i] + node_sets[0][i + 1]), + 0.5 * (node_sets[1][j] + node_sets[1][j + 1]), + 0.5 * (node_sets[2][k] + node_sets[2][k + 1])}; +} + +size_t +TotalCellCount(const OrthoCellInfo& info) +{ + return info.cell_counts[0] * info.cell_counts[1] * info.cell_counts[2]; +} + +size_t +TotalVertexCount(const OrthoCellInfo& info) +{ + return info.node_counts[0] * info.node_counts[1] * info.node_counts[2]; +} + +void +RebalanceOrthogonalPartitions(std::vector& cell_pids, const int num_partitions) +{ + std::vector cells_per_partition(num_partitions, 0); + for (const auto partition : cell_pids) + ++cells_per_partition[partition]; + + if (std::none_of(cells_per_partition.begin(), + cells_per_partition.end(), + [](const int count) { return count == 0; })) + return; + + const auto total_cells = static_cast(cell_pids.size()); + const int target = total_cells / num_partitions; + + for (int partition = 0; partition < num_partitions; ++partition) + { + while (cells_per_partition[partition] > target) + { + const auto it = std::min_element(cells_per_partition.begin(), cells_per_partition.end()); + const auto min_partition = static_cast(std::distance(cells_per_partition.begin(), it)); + if (min_partition == partition or cells_per_partition[min_partition] >= target) + break; + + for (auto& cell_pid : cell_pids) + { + if (cell_pid == partition) + { + cell_pid = min_partition; + --cells_per_partition[partition]; + ++cells_per_partition[min_partition]; + break; + } + } + } + } +} + +void +SetOrthogonalBoundaryMaps(const std::shared_ptr& grid, const unsigned int dimension) +{ + if (dimension >= 2) + { + grid->GetBoundaryIDMap()[XMIN] = "xmin"; + grid->GetBoundaryNameMap()["xmin"] = XMIN; + grid->GetBoundaryIDMap()[XMAX] = "xmax"; + grid->GetBoundaryNameMap()["xmax"] = XMAX; + grid->GetBoundaryIDMap()[YMIN] = "ymin"; + grid->GetBoundaryNameMap()["ymin"] = YMIN; + grid->GetBoundaryIDMap()[YMAX] = "ymax"; + grid->GetBoundaryNameMap()["ymax"] = YMAX; + } + + if (dimension == 1 or dimension == 3) + { + grid->GetBoundaryIDMap()[ZMIN] = "zmin"; + grid->GetBoundaryNameMap()["zmin"] = ZMIN; + grid->GetBoundaryIDMap()[ZMAX] = "zmax"; + grid->GetBoundaryNameMap()["zmax"] = ZMAX; + } +} + +std::vector +CellGraphNode(const OrthoCellInfo& info, const size_t i, const size_t j, const size_t k) +{ + std::vector node; + const auto [nx, ny, nz] = info.cell_counts; + + if (info.dimension == 1) + { + if (k > 0) + node.push_back(CellGlobalID(info, i, j, k - 1)); + if (k + 1 < nz) + node.push_back(CellGlobalID(info, i, j, k + 1)); + } + else if (info.dimension == 2) + { + if (i > 0) + node.push_back(CellGlobalID(info, i - 1, j, k)); + if (i + 1 < nx) + node.push_back(CellGlobalID(info, i + 1, j, k)); + if (j > 0) + node.push_back(CellGlobalID(info, i, j - 1, k)); + if (j + 1 < ny) + node.push_back(CellGlobalID(info, i, j + 1, k)); + } + else + { + if (i > 0) + node.push_back(CellGlobalID(info, i - 1, j, k)); + if (i + 1 < nx) + node.push_back(CellGlobalID(info, i + 1, j, k)); + if (j > 0) + node.push_back(CellGlobalID(info, i, j - 1, k)); + if (j + 1 < ny) + node.push_back(CellGlobalID(info, i, j + 1, k)); + if (k > 0) + node.push_back(CellGlobalID(info, i, j, k - 1)); + if (k + 1 < nz) + node.push_back(CellGlobalID(info, i, j, k + 1)); + } + + return node; +} + +std::shared_ptr +MakeLightWeightCell1D(const OrthoCellInfo& info, + const std::vector>& node_sets, + const size_t k) +{ + auto cell = std::make_shared(CellType::SLAB, CellType::SLAB); + + cell->centroid = CellCentroid(node_sets, 1, 0, 0, k); + cell->vertex_ids = {VertexGlobalID(info, 0, 0, k), VertexGlobalID(info, 0, 0, k + 1)}; + + UnpartitionedMesh::LightWeightFace left_face; + left_face.vertex_ids = {cell->vertex_ids[0]}; + left_face.has_neighbor = k != 0; + left_face.neighbor = k == 0 ? ZMIN : CellGlobalID(info, 0, 0, k - 1); + + UnpartitionedMesh::LightWeightFace right_face; + right_face.vertex_ids = {cell->vertex_ids[1]}; + right_face.has_neighbor = k + 1 != info.cell_counts[2]; + right_face.neighbor = right_face.has_neighbor ? CellGlobalID(info, 0, 0, k + 1) : ZMAX; + + cell->faces.push_back(left_face); + cell->faces.push_back(right_face); + + return cell; +} + +std::shared_ptr +MakeLightWeightCell2D(const OrthoCellInfo& info, + const std::vector>& node_sets, + const size_t i, + const size_t j) +{ + auto cell = std::make_shared(CellType::POLYGON, + CellType::QUADRILATERAL); + + const auto v00 = VertexGlobalID(info, i, j, 0); + const auto v10 = VertexGlobalID(info, i + 1, j, 0); + const auto v11 = VertexGlobalID(info, i + 1, j + 1, 0); + const auto v01 = VertexGlobalID(info, i, j + 1, 0); + cell->centroid = CellCentroid(node_sets, 2, i, j, 0); + cell->vertex_ids = {v00, v10, v11, v01}; + + const auto max_i = info.cell_counts[0] - 1; + const auto max_j = info.cell_counts[1] - 1; + + for (int f = 0; f < 4; ++f) + { + UnpartitionedMesh::LightWeightFace face; + if (f < 3) + face.vertex_ids = {cell->vertex_ids[f], cell->vertex_ids[f + 1]}; + else + face.vertex_ids = {cell->vertex_ids[f], cell->vertex_ids[0]}; + + if (f == 1) + { + face.has_neighbor = i != max_i; + face.neighbor = face.has_neighbor ? CellGlobalID(info, i + 1, j, 0) : XMAX; + } + else if (f == 3) + { + face.has_neighbor = i != 0; + face.neighbor = face.has_neighbor ? CellGlobalID(info, i - 1, j, 0) : XMIN; + } + else if (f == 2) + { + face.has_neighbor = j != max_j; + face.neighbor = face.has_neighbor ? CellGlobalID(info, i, j + 1, 0) : YMAX; + } + else + { + face.has_neighbor = j != 0; + face.neighbor = face.has_neighbor ? CellGlobalID(info, i, j - 1, 0) : YMIN; + } + + cell->faces.push_back(face); + } + + return cell; +} + +std::shared_ptr +MakeLightWeightCell3D(const OrthoCellInfo& info, + const std::vector>& node_sets, + const size_t i, + const size_t j, + const size_t k) +{ + auto cell = std::make_shared(CellType::POLYHEDRON, + CellType::HEXAHEDRON); + + const auto v000 = VertexGlobalID(info, i, j, k); + const auto v100 = VertexGlobalID(info, i + 1, j, k); + const auto v110 = VertexGlobalID(info, i + 1, j + 1, k); + const auto v010 = VertexGlobalID(info, i, j + 1, k); + const auto v001 = VertexGlobalID(info, i, j, k + 1); + const auto v101 = VertexGlobalID(info, i + 1, j, k + 1); + const auto v111 = VertexGlobalID(info, i + 1, j + 1, k + 1); + const auto v011 = VertexGlobalID(info, i, j + 1, k + 1); + + cell->centroid = CellCentroid(node_sets, 3, i, j, k); + cell->vertex_ids = {v000, v100, v110, v010, v001, v101, v111, v011}; + + const auto max_i = info.cell_counts[0] - 1; + const auto max_j = info.cell_counts[1] - 1; + const auto max_k = info.cell_counts[2] - 1; + + auto add_face = + [&cell](std::vector vertex_ids, const bool has_neighbor, const uint64_t neighbor) + { + UnpartitionedMesh::LightWeightFace face; + face.vertex_ids = std::move(vertex_ids); + face.has_neighbor = has_neighbor; + face.neighbor = neighbor; + cell->faces.push_back(std::move(face)); + }; + + add_face( + {v100, v110, v111, v101}, i != max_i, i == max_i ? XMAX : CellGlobalID(info, i + 1, j, k)); + add_face({v000, v001, v011, v010}, i != 0, i == 0 ? XMIN : CellGlobalID(info, i - 1, j, k)); + add_face( + {v010, v011, v111, v110}, j != max_j, j == max_j ? YMAX : CellGlobalID(info, i, j + 1, k)); + add_face({v000, v100, v101, v001}, j != 0, j == 0 ? YMIN : CellGlobalID(info, i, j - 1, k)); + add_face( + {v001, v101, v111, v011}, k != max_k, k == max_k ? ZMAX : CellGlobalID(info, i, j, k + 1)); + add_face({v000, v010, v110, v100}, k != 0, k == 0 ? ZMIN : CellGlobalID(info, i, j, k - 1)); + + return cell; +} + +std::shared_ptr +MakeLightWeightCell(const OrthoCellInfo& info, + const std::vector>& node_sets, + const uint64_t cell_global_id) +{ + const auto [i, j, k] = CellIJK(info, cell_global_id); + if (info.dimension == 1) + return MakeLightWeightCell1D(info, node_sets, k); + if (info.dimension == 2) + return MakeLightWeightCell2D(info, node_sets, i, j); + return MakeLightWeightCell3D(info, node_sets, i, j, k); +} + +template +void +ForEachCellClosureCell(const OrthoCellInfo& info, const uint64_t cell_gid, Callback callback) +{ + const auto [i, j, k] = CellIJK(info, cell_gid); + const auto [nx, ny, nz] = info.cell_counts; + + const int i_min = info.dimension >= 2 and i > 0 ? -1 : 0; + const int i_max = info.dimension >= 2 and i + 1 < nx ? 1 : 0; + const int j_min = info.dimension >= 2 and j > 0 ? -1 : 0; + const int j_max = info.dimension >= 2 and j + 1 < ny ? 1 : 0; + const int k_min = (info.dimension == 1 or info.dimension == 3) and k > 0 ? -1 : 0; + const int k_max = (info.dimension == 1 or info.dimension == 3) and k + 1 < nz ? 1 : 0; + + for (int dj = j_min; dj <= j_max; ++dj) + for (int di = i_min; di <= i_max; ++di) + for (int dk = k_min; dk <= k_max; ++dk) + callback(CellGlobalID(info, + static_cast(static_cast(i) + di), + static_cast(static_cast(j) + dj), + static_cast(static_cast(k) + dk))); +} + +void +AddUniquePartition(std::array& partitions, size_t& num_partitions, const int partition_id) +{ + for (size_t i = 0; i < num_partitions; ++i) + if (partitions[i] == partition_id) + return; + + partitions[num_partitions++] = partition_id; +} + +void +SortUnique(std::vector& values) +{ + std::sort(values.begin(), values.end()); + values.erase(std::unique(values.begin(), values.end()), values.end()); +} + +std::vector +BuildLocalCellPayload(const OrthoCellInfo& info, + const std::vector& cell_pids, + const int num_partitions, + std::vector& partI_num_cells) +{ + partI_num_cells.assign(num_partitions, 0); + std::map> send_payloads; + + for (uint64_t cell_gid = 0; cell_gid < cell_pids.size(); ++cell_gid) + { + const int cell_pid = cell_pids[cell_gid]; + ++partI_num_cells[cell_pid]; + + std::array recipient_pids = {}; + size_t num_recipient_pids = 0; + ForEachCellClosureCell( + info, + cell_gid, + [&cell_pids, &recipient_pids, &num_recipient_pids](const uint64_t gid) + { AddUniquePartition(recipient_pids, num_recipient_pids, cell_pids[gid]); }); + + for (size_t i = 0; i < num_recipient_pids; ++i) + { + auto& payload = send_payloads[recipient_pids[i]]; + payload.push_back(cell_gid); + payload.push_back(static_cast(cell_pid)); + } + } + + const auto received_payloads = MapAllToAll(send_payloads, mpi_comm); + const auto it = received_payloads.find(0); + if (it == received_payloads.end()) + return {}; + return it->second; +} + +} // namespace + OrthogonalMeshGenerator::OrthogonalMeshGenerator(const InputParameters& params) : MeshGenerator(params), coord_sys_(params.GetParamValue("coord_sys") == "cartesian" ? CARTESIAN : params.GetParamValue("coord_sys") == "cylindrical" ? CYLINDRICAL - : SPHERICAL) + : SPHERICAL), + distributed_generation_(params.GetParamValue("distributed_generation")) { // Parse the node_sets param if (params.IsParameterValid("node_sets")) @@ -96,6 +547,131 @@ OrthogonalMeshGenerator::GenerateUnpartitionedMesh( throw std::logic_error(""); } +std::shared_ptr +OrthogonalMeshGenerator::Execute() +{ + if (not distributed_generation_) + return MeshGenerator::Execute(); + + if (not inputs_.empty()) + throw std::invalid_argument("OrthogonalMeshGenerator with distributed_generation=true cannot " + "be preceded by another mesh generator."); + + if (replicated_) + throw std::invalid_argument("OrthogonalMeshGenerator with distributed_generation=true does not " + "support replicated_mesh=true."); + + const auto info = MakeOrthoCellInfo(node_sets_); + const auto total_num_cells = TotalCellCount(info); + const auto num_partitions = mpi_comm.size(); + + std::vector cell_pids; + std::vector partI_num_cells; + std::vector local_cell_payload; + + if (mpi_comm.rank() == 0) + { + std::vector> cell_graph; + std::vector cell_centroids; + cell_graph.reserve(total_num_cells); + cell_centroids.reserve(total_num_cells); + + for (uint64_t cell_gid = 0; cell_gid < total_num_cells; ++cell_gid) + { + const auto [i, j, k] = CellIJK(info, cell_gid); + cell_graph.push_back(CellGraphNode(info, i, j, k)); + cell_centroids.push_back(CellCentroid(node_sets_, info.dimension, i, j, k)); + } + + cell_pids = partitioner_->Partition(cell_graph, cell_centroids, num_partitions); + RebalanceOrthogonalPartitions(cell_pids, num_partitions); + } + + if (mpi_comm.rank() == 0 and cell_pids.size() != total_num_cells) + throw std::logic_error("OrthogonalMeshGenerator distributed generation received an invalid " + "partition vector."); + + local_cell_payload = BuildLocalCellPayload(info, cell_pids, num_partitions, partI_num_cells); + mpi_comm.broadcast(partI_num_cells, 0); + + size_t avg_num_cells = 0; + auto max_num_cells = partI_num_cells.front(); + auto min_num_cells = partI_num_cells.front(); + for (size_t count : partI_num_cells) + { + max_num_cells = std::max(max_num_cells, count); + min_num_cells = std::min(min_num_cells, count); + avg_num_cells += count; + } + avg_num_cells /= num_partitions; + + if (mpi_comm.rank() == 0) + log.Log() << "Number of cells per partition (max,min,avg) = " << max_num_cells << "," + << min_num_cells << "," << avg_num_cells; + if (min_num_cells == 0) + throw std::runtime_error("Partitioning failed. At least one partition contains no cells."); + + if (local_cell_payload.size() % 2 != 0) + throw std::logic_error("OrthogonalMeshGenerator received an invalid local cell payload."); + + std::vector cells_needed; + std::unordered_map local_cell_pids; + cells_needed.reserve(local_cell_payload.size() / 2); + local_cell_pids.reserve(local_cell_payload.size() / 2); + for (size_t i = 0; i < local_cell_payload.size(); i += 2) + { + const uint64_t cell_gid = local_cell_payload[i]; + const int cell_pid = static_cast(local_cell_payload[i + 1]); + cells_needed.push_back(cell_gid); + local_cell_pids.emplace(cell_gid, cell_pid); + } + + auto grid_ptr = MeshContinuum::New(); + SetOrthogonalBoundaryMaps(grid_ptr, info.dimension); + + std::vector vertices_needed; + for (const auto cell_gid : cells_needed) + { + const auto raw_cell = MakeLightWeightCell(info, node_sets_, cell_gid); + for (const auto vid : raw_cell->vertex_ids) + vertices_needed.push_back(vid); + + grid_ptr->cells.PushBack(SetupCell(*raw_cell, cell_gid, local_cell_pids.at(cell_gid))); + } + + SortUnique(vertices_needed); + for (const auto vid : vertices_needed) + { + const auto [i, j, k] = [&info, vid]() + { + const auto [nx, ny, nz] = info.node_counts; + if (info.dimension == 1) + return std::array{0, 0, static_cast(vid)}; + if (info.dimension == 2) + return std::array{ + static_cast(vid % nx), static_cast(vid / nx), 0}; + + const auto xy = static_cast(vid / nz); + return std::array{xy % nx, xy / nx, static_cast(vid % nz)}; + }(); + + grid_ptr->vertices.Insert(vid, Vertex(node_sets_, info.dimension, i, j, k)); + } + + grid_ptr->SetDimension(info.dimension); + grid_ptr->SetCoordinateSystem(coord_sys_); + grid_ptr->SetType(ORTHOGONAL); + grid_ptr->SetExtruded(false); + grid_ptr->SetOrthoAttributes({info.cell_counts[0], info.cell_counts[1], info.cell_counts[2]}); + grid_ptr->SetGlobalVertexCount(TotalVertexCount(info)); + grid_ptr->ComputeGeometricInfo(); + + ComputeAndPrintStats(grid_ptr); + mpi_comm.barrier(); + + return grid_ptr; +} + OpenSnRegisterObjectInNamespace(mesh, OrthogonalMeshGenerator); InputParameters @@ -111,6 +687,10 @@ OrthogonalMeshGenerator::GetInputParameters() params.AddOptionalParameter("coord_sys", "cartesian", "The coordinate system of the mesh."); params.ConstrainParameterRange( "coord_sys", AllowableRangeList::New({"cartesian", "cylindrical", "spherical"})); + params.AddOptionalParameter("distributed_generation", + false, + "Generate cells directly on each MPI rank after partitioning " + "the orthogonal cell graph on rank 0."); return params; } diff --git a/framework/mesh/mesh_generator/orthogonal_mesh_generator.h b/framework/mesh/mesh_generator/orthogonal_mesh_generator.h index 799bbbacce..0fdbb86ca9 100644 --- a/framework/mesh/mesh_generator/orthogonal_mesh_generator.h +++ b/framework/mesh/mesh_generator/orthogonal_mesh_generator.h @@ -13,12 +13,15 @@ class OrthogonalMeshGenerator : public MeshGenerator public: explicit OrthogonalMeshGenerator(const InputParameters& params); + std::shared_ptr Execute() override; + protected: std::shared_ptr GenerateUnpartitionedMesh(std::shared_ptr input_umesh) override; const CoordinateSystemType coord_sys_; std::vector> node_sets_; + const bool distributed_generation_; public: static InputParameters GetInputParameters(); diff --git a/python/lib/mesh.cc b/python/lib/mesh.cc index df594e2a1e..95d16c30bf 100644 --- a/python/lib/mesh.cc +++ b/python/lib/mesh.cc @@ -402,6 +402,9 @@ WrapMeshGenerator(py::module& mesh) Sets of nodes per dimension. Node values must be monotonically increasing. coord_sys: {'cartesian', 'cylindrical', 'spherical'} The coordinate system of the mesh. + distributed_generation: bool, default=False + When true, partition the orthogonal cell graph on rank 0 and build each + rank's local and ghost cells directly on that rank. )" );