1 // Copyright (c) 2017-2021, Lawrence Livermore National Security, LLC and 2 // other Axom Project Developers. See the top-level LICENSE file for details. 3 // 4 // SPDX-License-Identifier: (BSD-3-Clause) 5 6 /** 7 * \file SamplingShaper.hpp 8 * 9 * \brief Helper class for sampling-based shaping queries 10 */ 11 12 #ifndef AXOM_QUEST_SAMPLING_SHAPER__HPP_ 13 #define AXOM_QUEST_SAMPLING_SHAPER__HPP_ 14 15 #include "axom/config.hpp" 16 #include "axom/core.hpp" 17 #include "axom/slic.hpp" 18 #include "axom/slam.hpp" 19 #include "axom/primal.hpp" 20 #include "axom/mint.hpp" 21 #include "axom/spin.hpp" 22 #include "axom/klee.hpp" 23 24 #ifndef AXOM_USE_MFEM 25 #error Shaping functionality requires Axom to be configured with MFEM and the AXOM_ENABLE_MFEM_SIDRE_DATACOLLECTION option 26 #endif 27 28 #include "axom/quest/Shaper.hpp" 29 #include "axom/quest/InOutOctree.hpp" 30 #include "axom/quest/interface/internal/mpicomm_wrapper.hpp" 31 #include "axom/quest/interface/internal/QuestHelpers.hpp" 32 #include "axom/quest/detail/shaping/shaping_helpers.hpp" 33 34 #include "mfem.hpp" 35 36 #include "axom/fmt.hpp" 37 #include "axom/fmt/locale.h" 38 39 namespace axom 40 { 41 namespace quest 42 { 43 namespace shaping 44 { 45 using QFunctionCollection = mfem::NamedFieldsMap<mfem::QuadratureFunction>; 46 using DenseTensorCollection = mfem::NamedFieldsMap<mfem::DenseTensor>; 47 48 enum class VolFracSampling : int 49 { 50 SAMPLE_AT_DOFS, 51 SAMPLE_AT_QPTS 52 }; 53 54 template <int NDIMS> 55 class InOutSampler 56 { 57 public: 58 static constexpr int DIM = NDIMS; 59 using InOutOctreeType = quest::InOutOctree<DIM>; 60 61 using GeometricBoundingBox = typename InOutOctreeType::GeometricBoundingBox; 62 using SpacePt = typename InOutOctreeType::SpacePt; 63 using SpaceVector = typename InOutOctreeType::SpaceVector; 64 using GridPt = typename InOutOctreeType::GridPt; 65 using BlockIndex = typename InOutOctreeType::BlockIndex; 66 67 public: 68 /** 69 * \brief Constructor for a SamplingShaper 70 * 71 * \param shapeName The name of the shape; will be used for the field for the associated samples 72 * \param surfaceMesh Pointer to the surface mesh 73 * 74 * \note Does not take ownership of the surface mesh 75 */ InOutSampler(const std::string & shapeName,mint::Mesh * surfaceMesh)76 InOutSampler(const std::string& shapeName, mint::Mesh* surfaceMesh) 77 : m_shapeName(shapeName) 78 , m_surfaceMesh(surfaceMesh) 79 { } 80 ~InOutSampler()81 ~InOutSampler() { delete m_octree; } 82 getSurfaceMesh() const83 mint::Mesh* getSurfaceMesh() const { return m_surfaceMesh; } 84 85 /// Computes the bounding box of the surface mesh computeBounds()86 void computeBounds() 87 { 88 SLIC_ASSERT(m_surfaceMesh != nullptr); 89 90 m_bbox.clear(); 91 SpacePt pt; 92 93 for(int i = 0; i < m_surfaceMesh->getNumberOfNodes(); ++i) 94 { 95 m_surfaceMesh->getNode(i, pt.data()); 96 m_bbox.addPoint(pt); 97 } 98 99 SLIC_ASSERT(m_bbox.isValid()); 100 101 SLIC_INFO("Mesh bounding box: " << m_bbox); 102 } 103 initSpatialIndex(double vertexWeldThreshold)104 void initSpatialIndex(double vertexWeldThreshold) 105 { 106 // Create octree over mesh's bounding box 107 m_octree = new InOutOctreeType(m_bbox, m_surfaceMesh); 108 m_octree->setVertexWeldThreshold(vertexWeldThreshold); 109 m_octree->generateIndex(); 110 } 111 sampleInOutField(mfem::DataCollection * dc,shaping::QFunctionCollection & inoutQFuncs,int sampleRes)112 void sampleInOutField(mfem::DataCollection* dc, 113 shaping::QFunctionCollection& inoutQFuncs, 114 int sampleRes) 115 { 116 auto* mesh = dc->GetMesh(); 117 SLIC_ASSERT(mesh != nullptr); 118 const int NE = mesh->GetNE(); 119 const int dim = mesh->Dimension(); 120 121 // Generate a Quadrature Function with the geometric positions, if not already available 122 if(!inoutQFuncs.Has("positions")) 123 { 124 shaping::generatePositionsQFunction(mesh, inoutQFuncs, sampleRes); 125 } 126 127 // Access the positions QFunc and associated QuadratureSpace 128 mfem::QuadratureFunction* pos_coef = inoutQFuncs.Get("positions"); 129 mfem::QuadratureSpace* sp = pos_coef->GetSpace(); 130 const int nq = sp->GetElementIntRule(0).GetNPoints(); 131 132 // Sample the in/out field at each point 133 // store in QField which we register with the QFunc collection 134 const std::string inoutName = axom::fmt::format("inout_{}", m_shapeName); 135 const int vdim = 1; 136 auto* inout = new mfem::QuadratureFunction(sp, vdim); 137 inoutQFuncs.Register(inoutName, inout, true); 138 139 mfem::DenseMatrix m; 140 mfem::Vector res; 141 142 axom::utilities::Timer timer(true); 143 for(int i = 0; i < NE; ++i) 144 { 145 pos_coef->GetElementValues(i, m); 146 inout->GetElementValues(i, res); 147 148 for(int p = 0; p < nq; ++p) 149 { 150 const SpacePt pt(m.GetColumn(p), dim); 151 const bool in = m_octree->within(pt); 152 res(p) = in ? 1. : 0.; 153 154 // SLIC_INFO(axom::fmt::format("[{},{}] Pt: {}, In: {}", i,p,pt, (in? "yes" : "no") )); 155 } 156 } 157 timer.stop(); 158 159 SLIC_INFO( 160 axom::fmt::format(std::locale("en_US.UTF-8"), 161 "\t Sampling inout field '{}' took {} seconds (@ " 162 "{:L} queries per second)", 163 inoutName, 164 timer.elapsed(), 165 static_cast<int>((NE * nq) / timer.elapsed()))); 166 } 167 168 /** 169 * Compute volume fractions function for shape on a grid of resolution \a gridRes 170 * in region defined by bounding box \a queryBounds 171 */ computeVolumeFractionsBaseline(mfem::DataCollection * dc,int AXOM_UNUSED_PARAM (sampleRes),int outputOrder)172 void computeVolumeFractionsBaseline(mfem::DataCollection* dc, 173 int AXOM_UNUSED_PARAM(sampleRes), 174 int outputOrder) 175 { 176 // Step 1 -- generate a QField w/ the spatial coordinates 177 mfem::Mesh* mesh = dc->GetMesh(); 178 const int NE = mesh->GetNE(); 179 const int dim = mesh->Dimension(); 180 181 if(NE < 1) 182 { 183 SLIC_WARNING("Mesh has no elements!"); 184 return; 185 } 186 187 mfem::L2_FECollection* coll = 188 new mfem::L2_FECollection(outputOrder, dim, mfem::BasisType::Positive); 189 mfem::FiniteElementSpace* fes = new mfem::FiniteElementSpace(mesh, coll); 190 mfem::GridFunction* volFrac = new mfem::GridFunction(fes); 191 volFrac->MakeOwner(coll); 192 auto volFracName = axom::fmt::format("vol_frac_{}", m_shapeName); 193 dc->RegisterField(volFracName, volFrac); 194 195 auto* fe = fes->GetFE(0); 196 auto& ir = fe->GetNodes(); 197 198 // Assume all elements have the same integration rule 199 const int nq = ir.GetNPoints(); 200 const auto* geomFactors = 201 mesh->GetGeometricFactors(ir, mfem::GeometricFactors::COORDINATES); 202 203 mfem::DenseTensor pos_coef(dim, nq, NE); 204 205 // Rearrange positions into quadrature function 206 { 207 for(int i = 0; i < NE; ++i) 208 { 209 for(int j = 0; j < dim; ++j) 210 { 211 for(int k = 0; k < nq; ++k) 212 { 213 pos_coef(j, k, i) = geomFactors->X((i * nq * dim) + (j * nq) + k); 214 } 215 } 216 } 217 } 218 219 // Step 2 -- sample the in/out field at each point -- store directly in volFrac grid function 220 mfem::Vector res(nq); 221 mfem::Array<int> dofs; 222 for(int i = 0; i < NE; ++i) 223 { 224 mfem::DenseMatrix& m = pos_coef(i); 225 for(int p = 0; p < nq; ++p) 226 { 227 const SpacePt pt(m.GetColumn(p), dim); 228 const bool in = m_octree->within(pt); 229 res(p) = in ? 1. : 0.; 230 } 231 232 fes->GetElementDofs(i, dofs); 233 volFrac->SetSubVector(dofs, res); 234 } 235 } 236 237 private: 238 DISABLE_COPY_AND_ASSIGNMENT(InOutSampler); 239 DISABLE_MOVE_AND_ASSIGNMENT(InOutSampler); 240 241 std::string m_shapeName; 242 243 GeometricBoundingBox m_bbox; 244 mint::Mesh* m_surfaceMesh {nullptr}; 245 InOutOctreeType* m_octree {nullptr}; 246 }; 247 248 } // end namespace shaping 249 250 /*! 251 * \brief Concrete class for sample based shaping 252 */ 253 class SamplingShaper : public Shaper 254 { 255 public: SamplingShaper(const klee::ShapeSet & shapeSet,sidre::MFEMSidreDataCollection * dc)256 SamplingShaper(const klee::ShapeSet& shapeSet, 257 sidre::MFEMSidreDataCollection* dc) 258 : Shaper(shapeSet, dc) 259 { } 260 261 //@{ 262 //! @name Functions to get and set shaping parameters related to sampling; supplements parameters in base class 263 setSamplingType(shaping::VolFracSampling vfSampling)264 void setSamplingType(shaping::VolFracSampling vfSampling) 265 { 266 m_vfSampling = vfSampling; 267 } 268 setQuadratureOrder(int quadratureOrder)269 void setQuadratureOrder(int quadratureOrder) 270 { 271 m_quadratureOrder = quadratureOrder; 272 } 273 setVolumeFractionOrder(int volfracOrder)274 void setVolumeFractionOrder(int volfracOrder) 275 { 276 m_volfracOrder = volfracOrder; 277 } 278 279 //@} 280 281 private: getShapeDimension() const282 klee::Dimensions getShapeDimension() const 283 { 284 const bool has2D = (m_inoutSampler2D != nullptr); 285 const bool has3D = (m_inoutSampler3D != nullptr); 286 SLIC_ERROR_IF(!(has2D || has3D), "Shape not initialized"); 287 SLIC_ERROR_IF(has2D && has3D, "Cannot have concurrent 2D and 3D shapes"); 288 289 return has2D ? klee::Dimensions::Two : klee::Dimensions::Three; 290 } 291 292 public: 293 //@{ 294 //! @name Functions related to the stages for a given shape 295 296 /// Initializes the spatial index for shaping prepareShapeQuery(klee::Dimensions shapeDimension,const klee::Shape & shape)297 void prepareShapeQuery(klee::Dimensions shapeDimension, 298 const klee::Shape& shape) override 299 { 300 const auto& shapeName = shape.getName(); 301 302 SLIC_INFO(axom::fmt::format("{:-^80}", " Generating the octree ")); 303 304 internal::ScopedLogLevelChanger logLevelChanger( 305 this->isVerbose() ? slic::message::Debug : slic::message::Warning); 306 307 switch(shapeDimension) 308 { 309 case klee::Dimensions::Two: 310 m_inoutSampler2D = new shaping::InOutSampler<2>(shapeName, m_surfaceMesh); 311 m_inoutSampler2D->computeBounds(); 312 m_inoutSampler2D->initSpatialIndex(this->m_vertexWeldThreshold); 313 m_surfaceMesh = m_inoutSampler2D->getSurfaceMesh(); 314 break; 315 316 case klee::Dimensions::Three: 317 m_inoutSampler3D = new shaping::InOutSampler<3>(shapeName, m_surfaceMesh); 318 m_inoutSampler3D->computeBounds(); 319 m_inoutSampler3D->initSpatialIndex(this->m_vertexWeldThreshold); 320 m_surfaceMesh = m_inoutSampler3D->getSurfaceMesh(); 321 break; 322 323 default: 324 SLIC_ERROR( 325 "Shaping dimension must be 2 or 3, but requested dimension was " 326 << static_cast<int>(shapeDimension)); 327 break; 328 } 329 330 // Check that one of sampling shapers (2D or 3D) is null and the other is not 331 SLIC_ASSERT((m_inoutSampler2D == nullptr && m_inoutSampler3D != nullptr) || 332 (m_inoutSampler3D == nullptr && m_inoutSampler2D != nullptr)); 333 334 // Output some logging info and dump the mesh 335 if(this->isVerbose() && this->getRank() == 0) 336 { 337 const int nVerts = m_surfaceMesh->getNumberOfNodes(); 338 const int nCells = m_surfaceMesh->getNumberOfCells(); 339 340 SLIC_INFO(axom::fmt::format( 341 "After welding, surface mesh has {} vertices and {} triangles.", 342 nVerts, 343 nCells)); 344 mint::write_vtk(m_surfaceMesh, 345 axom::fmt::format("meldedTriMesh_{}.vtk", shapeName)); 346 } 347 } 348 runShapeQuery(const klee::Shape & shape)349 void runShapeQuery(const klee::Shape& shape) override 350 { 351 SLIC_INFO(axom::fmt::format( 352 "{:-^80}", 353 axom::fmt::format(" Querying the octree for shape '{}'", shape.getName()))); 354 355 internal::ScopedLogLevelChanger logLevelChanger( 356 this->isVerbose() ? slic::message::Debug : slic::message::Warning); 357 358 switch(getShapeDimension()) 359 { 360 case klee::Dimensions::Two: 361 runShapeQueryImpl(m_inoutSampler2D); 362 break; 363 case klee::Dimensions::Three: 364 runShapeQueryImpl(m_inoutSampler3D); 365 break; 366 } 367 } 368 applyReplacementRules(const klee::Shape & shape)369 void applyReplacementRules(const klee::Shape& shape) override 370 { 371 using axom::utilities::string::rsplitN; 372 373 const auto& shapeName = shape.getName(); 374 SLIC_INFO(axom::fmt::format( 375 "{:-^80}", 376 axom::fmt::format("Applying replacement rules over for shape '{}'", 377 shapeName))); 378 379 internal::ScopedLogLevelChanger logLevelChanger( 380 this->isVerbose() ? slic::message::Debug : slic::message::Warning); 381 382 // Get inout qfunc for this shape 383 auto* shapeQFunc = 384 m_inoutShapeQFuncs.Get(axom::fmt::format("inout_{}", shapeName)); 385 SLIC_ASSERT_MSG( 386 shapeQFunc != nullptr, 387 axom::fmt::format("Missing inout samples for shape '{}'", shapeName)); 388 389 // Create a copy of the inout samples for this shape 390 // Replacements will be applied to this and then copied into our shape's material 391 auto* shapeQFuncCopy = new mfem::QuadratureFunction(*shapeQFunc); 392 393 // apply replacement rules to all other materials 394 const auto& thisMatName = shape.getMaterial(); 395 for(auto& mat : m_inoutMaterialQFuncs) 396 { 397 const std::string otherMatName = rsplitN(mat.first, 2, '_')[1]; 398 399 // We'll handle the current shape's material at the end 400 if(otherMatName == thisMatName) 401 { 402 continue; 403 } 404 405 const bool shouldReplace = shape.replaces(otherMatName); 406 SLIC_DEBUG(axom::fmt::format( 407 "Should we replace material '{}' with shape '{}' of material '{}'? {}", 408 otherMatName, 409 shapeName, 410 thisMatName, 411 shouldReplace ? "yes" : "no")); 412 413 auto* otherMatQFunc = mat.second; 414 SLIC_ASSERT_MSG( 415 otherMatQFunc != nullptr, 416 axom::fmt::format("Missing inout samples for material '{}'", 417 otherMatName)); 418 419 quest::shaping::replaceMaterial(shapeQFuncCopy, otherMatQFunc, shouldReplace); 420 } 421 422 // Get inout qfunc for the current material 423 const std::string materialQFuncName = 424 axom::fmt::format("mat_inout_{}", thisMatName); 425 if(!m_inoutMaterialQFuncs.Has(materialQFuncName)) 426 { 427 // initialize material from shape inout, the QFunc registry takes ownership 428 m_inoutMaterialQFuncs.Register(materialQFuncName, shapeQFuncCopy, true); 429 } 430 else 431 { 432 // copy shape data into current material and delete the copy 433 auto* matQFunc = m_inoutMaterialQFuncs.Get(materialQFuncName); 434 SLIC_ASSERT_MSG( 435 matQFunc != nullptr, 436 axom::fmt::format("Missing inout samples for material '{}'", thisMatName)); 437 438 quest::shaping::copyShapeIntoMaterial(shapeQFuncCopy, matQFunc); 439 440 delete shapeQFuncCopy; 441 shapeQFuncCopy = nullptr; 442 } 443 } 444 finalizeShapeQuery()445 void finalizeShapeQuery() override 446 { 447 delete m_inoutSampler2D; 448 m_inoutSampler2D = nullptr; 449 450 delete m_inoutSampler3D; 451 m_inoutSampler3D = nullptr; 452 453 delete m_surfaceMesh; 454 m_surfaceMesh = nullptr; 455 } 456 457 //@} 458 459 public: adjustVolumeFractions()460 void adjustVolumeFractions() override 461 { 462 internal::ScopedLogLevelChanger logLevelChanger( 463 this->isVerbose() ? slic::message::Debug : slic::message::Warning); 464 465 for(auto& mat : m_inoutMaterialQFuncs) 466 { 467 const std::string matName = mat.first; 468 SLIC_INFO( 469 axom::fmt::format("Generating volume fraction fields for '{}' material", 470 matName)); 471 472 // Sample the InOut field at the mesh quadrature points 473 switch(m_vfSampling) 474 { 475 case shaping::VolFracSampling::SAMPLE_AT_QPTS: 476 quest::shaping::computeVolumeFractions(matName, 477 m_dc, 478 m_inoutMaterialQFuncs, 479 m_volfracOrder); 480 break; 481 case shaping::VolFracSampling::SAMPLE_AT_DOFS: 482 /* no-op for now */ 483 break; 484 } 485 } 486 } 487 488 private: 489 // Handles 2D or 3D shaping, based on the template and associated parameter 490 template <typename InOutSamplerType> runShapeQueryImpl(InOutSamplerType * shaper)491 void runShapeQueryImpl(InOutSamplerType* shaper) 492 { 493 // Sample the InOut field at the mesh quadrature points 494 switch(m_vfSampling) 495 { 496 case shaping::VolFracSampling::SAMPLE_AT_QPTS: 497 shaper->sampleInOutField(m_dc, m_inoutShapeQFuncs, m_quadratureOrder); 498 break; 499 case shaping::VolFracSampling::SAMPLE_AT_DOFS: 500 shaper->computeVolumeFractionsBaseline(m_dc, 501 m_quadratureOrder, 502 m_volfracOrder); 503 break; 504 } 505 } 506 507 private: 508 // TODO: Use MfemSidreDataCollection QFuncs for this when we upgrade to post mfem@4.3 509 shaping::QFunctionCollection m_inoutShapeQFuncs; 510 shaping::QFunctionCollection m_inoutMaterialQFuncs; 511 shaping::DenseTensorCollection m_inoutDofs; 512 513 shaping::InOutSampler<2>* m_inoutSampler2D {nullptr}; 514 shaping::InOutSampler<3>* m_inoutSampler3D {nullptr}; 515 516 shaping::VolFracSampling m_vfSampling {shaping::VolFracSampling::SAMPLE_AT_QPTS}; 517 int m_quadratureOrder {5}; 518 int m_volfracOrder {2}; 519 }; 520 521 } // end namespace quest 522 } // end namespace axom 523 524 #endif // AXOM_QUEST_SAMPLING_SHAPER__HPP_ 525