From 1f81500580178f2bd03c6777eeb0debaff33317e Mon Sep 17 00:00:00 2001
From: Frederik Hennig <frederik.hennig@fau.de>
Date: Thu, 3 Sep 2020 19:20:07 +0200
Subject: [PATCH] advanced LBM codegen WIP

---
 .../codegen/03_AdvancedLBMCodegen.cpp         | 203 ++++++++++++++++++
 .../codegen/03_AdvancedLBMCodegen.py          |  78 +++++++
 apps/tutorials/codegen/CMakeLists.txt         |  14 ++
 3 files changed, 295 insertions(+)
 create mode 100644 apps/tutorials/codegen/03_AdvancedLBMCodegen.cpp
 create mode 100644 apps/tutorials/codegen/03_AdvancedLBMCodegen.py

diff --git a/apps/tutorials/codegen/03_AdvancedLBMCodegen.cpp b/apps/tutorials/codegen/03_AdvancedLBMCodegen.cpp
new file mode 100644
index 000000000..4dedd5a61
--- /dev/null
+++ b/apps/tutorials/codegen/03_AdvancedLBMCodegen.cpp
@@ -0,0 +1,203 @@
+//======================================================================================================================
+//
+//  This file is part of waLBerla. waLBerla is free software: you can
+//  redistribute it and/or modify it under the terms of the GNU General Public
+//  License as published by the Free Software Foundation, either version 3 of
+//  the License, or (at your option) any later version.
+//
+//  waLBerla is distributed in the hope that it will be useful, but WITHOUT
+//  ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
+//  FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License
+//  for more details.
+//
+//  You should have received a copy of the GNU General Public License along
+//  with waLBerla (see COPYING.txt). If not, see <http://www.gnu.org/licenses/>.
+//
+//! \file 03_AdvancedLBMCodegen.cpp
+//! \author Frederik Hennig <frederik.hennig@fau.de>
+//
+//======================================================================================================================
+
+#include "blockforest/all.h"
+
+#include "core/all.h"
+
+#include "domain_decomposition/all.h"
+
+#include "field/all.h"
+
+#include "geometry/all.h"
+
+#include "lbm/boundary/factories/DefaultBoundaryHandling.h"
+#include "lbm/communication/PdfFieldPackInfo.h"
+#include "lbm/field/AddToStorage.h"
+#include "lbm/field/PdfField.h"
+#include "lbm/field/initializer/all.h"
+#include "lbm/vtk/VTKOutput.h"
+
+#include "stencil/D2Q9.h"
+
+#include "timeloop/all.h"
+
+//    Codegen Includes
+#include "CumulantMRTNoSlip.h"
+#include "CumulantMRTPackInfo.h"
+#include "CumulantMRTSweep.h"
+#include "DensityAndVelocityFieldSetter.h"
+
+namespace walberla
+{
+///////////////////////
+/// Typedef Aliases ///
+///////////////////////
+
+// Communication Pack Info
+typedef pystencils::CumulantMRTPackInfo PackInfo_T;
+
+// LB Method Stencil
+typedef stencil::D2Q9 Stencil_T;
+
+// PDF field type
+typedef field::GhostLayerField< real_t, Stencil_T::Size > PdfField_T;
+
+// Velocity Field Type
+typedef field::GhostLayerField< real_t, Stencil_T::D > VectorField_T;
+
+// Boundary Handling
+typedef walberla::uint8_t flag_t;
+typedef FlagField< flag_t > FlagField_T;
+typedef lbm::CumulantMRTNoSlip NoSlip_T;
+
+//////////////////////////////////////////
+/// Shear Flow Velocity Initialization ///
+//////////////////////////////////////////
+
+bool initShearFlowVelocityField(const shared_ptr< StructuredBlockForest >& blocks, const BlockDataID& velocityFieldId,
+                                const Config::BlockHandle& config)
+{
+   math::RealRandom< real_t > rng(config.getParameter< std::mt19937::result_type >("noise_seed", 42));
+
+   real_t velocityMagnitude = config.getParameter< real_t >("VelocityMagnitude", real_c(0.08));
+   real_t noiseMagnitude    = config.getParameter< real_t >("NoiseMagnitude", real_c(0.1) * velocityMagnitude);
+
+   real_t n_y = real_c(blocks->getNumberOfYCells());
+
+   for (auto blockIt = blocks->begin(); blockIt != blocks->end(); ++blockIt)
+   {
+      auto u = (*blockIt).getData< VectorField_T >(velocityFieldId);
+
+      for (auto cellIt = u->beginWithGhostLayerXYZ(); cellIt != u->end(); ++cellIt)
+      {
+         Cell globalCell(cellIt.cell());
+         blocks->transformBlockLocalToGlobalCell(globalCell, *blockIt);
+
+         real_t relative_y = real_c(globalCell.y()) / n_y;
+
+         u->get(cellIt.cell(), 0) = relative_y < 0.3 || relative_y > 0.7 ? velocityMagnitude : -velocityMagnitude;
+
+         u->get(cellIt.cell(), 1) = noiseMagnitude * rng();
+      }
+   }
+}
+
+/////////////////////
+/// Main Function ///
+/////////////////////
+
+int main(int argc, char** argv)
+{
+   walberla::Environment walberlaEnv(argc, argv);
+
+   if (!walberlaEnv.config()) { WALBERLA_ABORT("No configuration file specified!"); }
+
+   ///////////////////////////////////////////////////////
+   /// Block Storage Creation and Simulation Parameter ///
+   ///////////////////////////////////////////////////////
+
+   auto blocks = blockforest::createUniformBlockGridFromConfig(walberlaEnv.config());
+
+   // read parameters
+   auto parameters = walberlaEnv.config()->getOneBlock("Parameters");
+
+   const uint_t timesteps = parameters.getParameter< uint_t >("timesteps", uint_c(10));
+   const real_t omega     = parameters.getParameter< real_t >("omega", real_c(1.8));
+   const double remainingTimeLoggerFrequency =
+      parameters.getParameter< double >("remainingTimeLoggerFrequency", 3.0); // in seconds
+
+   ////////////////////////////
+   /// Velocity Field Setup ///
+   ////////////////////////////
+
+   BlockDataID velocityFieldId = field::addToStorage< VectorField_T >(blocks, "velocity", real_c(0.0), field::fzyx);
+
+   ///////////////////////
+   /// PDF Field Setup ///
+   ///////////////////////
+
+   BlockDataID pdfFieldId  = field::addToStorage< PdfField_T >(blocks, "pdf field", real_c(0.0), field::fzyx);
+   BlockDataID flagFieldId = field::addFlagFieldToStorage< FlagField_T >(blocks, "flag field");
+
+   // First, set up the initial shear flow in the velocity field...
+
+   auto shearFlowSetup = walberlaEnv.config()->getOneBlock("ShearFlowSetup");
+   initShearFlowVelocityField(blocks, velocityFieldId, shearFlowSetup);
+
+   real_t rho = shearFlowSetup.getParameter("rho", real_c(1.8));
+
+   // ... and then use the generated PDF setter to initialize the PDF field.
+   pystencils::DensityAndVelocityFieldSetter pdfSetter(pdfFieldId, velocityFieldId, rho);
+   
+   for(auto blockIt = blocks->begin(); blockIt != blocks->end(); ++blockIt){
+      pdfSetter(&(*blockIt));
+   }
+
+   /////////////////////////
+   /// Boundary Handling ///
+   /////////////////////////
+
+   const FlagUID fluidFlagUID("Fluid");
+
+   auto boundariesConfig = walberlaEnv.config()->getOneBlock("Boundaries");
+
+   // BlockDataID boundaryHandlingId =
+   //    BHFactory::addBoundaryHandlingToStorage(blocks, "boundary handling", flagFieldId, pdfFieldId, fluidFlagUID,
+   //                                            Vector3< real_t >(), Vector3< real_t >(), real_c(0.0), real_c(0.0));
+
+   // geometry::initBoundaryHandling< BHFactory::BoundaryHandling >(*blocks, boundaryHandlingId, boundariesConfig);
+   // geometry::setNonBoundaryCellsToDomain< BHFactory::BoundaryHandling >(*blocks, boundaryHandlingId);
+
+   /////////////////
+   /// Time Loop ///
+   /////////////////
+
+   SweepTimeloop timeloop(blocks->getBlockStorage(), timesteps);
+
+   // Communication
+   blockforest::communication::UniformBufferedScheme< Stencil_T > communication(blocks);
+   communication.addPackInfo(make_shared< PackInfo_T >(pdfFieldId));
+
+   // Timeloop
+   timeloop.add() << BeforeFunction(communication, "communication");
+
+   // Stability Checker
+   timeloop.addFuncAfterTimeStep(makeSharedFunctor(field::makeStabilityChecker< PdfField_T, FlagField_T >(
+                                    walberlaEnv.config(), blocks, pdfFieldId, flagFieldId, fluidFlagUID)),
+                                 "LBM stability check");
+
+   // Time logger
+   timeloop.addFuncAfterTimeStep(timing::RemainingTimeLogger(timeloop.getNrOfTimeSteps(), remainingTimeLoggerFrequency),
+                                 "remaining time logger");
+
+   // TODO: Replace
+   // LBM VTK Output
+   // lbm::VTKOutput< LatticeModel_T, FlagField_T >::addToTimeloop(timeloop, blocks, walberlaEnv.config(), pdfFieldId,
+   //                                                              flagFieldId, fluidFlagUID);
+
+   timeloop.run();
+
+   return EXIT_SUCCESS;
+}
+
+} // namespace walberla
+
+int main(int argc, char** argv) { return walberla::main(argc, argv); }
\ No newline at end of file
diff --git a/apps/tutorials/codegen/03_AdvancedLBMCodegen.py b/apps/tutorials/codegen/03_AdvancedLBMCodegen.py
new file mode 100644
index 000000000..e3113c92c
--- /dev/null
+++ b/apps/tutorials/codegen/03_AdvancedLBMCodegen.py
@@ -0,0 +1,78 @@
+import sympy as sp
+import pystencils as ps
+
+from lbmpy.creationfunctions import create_lb_update_rule
+from lbmpy.macroscopic_value_kernels import macroscopic_values_setter
+from lbmpy.boundaries import NoSlip
+
+from pystencils_walberla import CodeGeneration, generate_sweep, generate_pack_info_from_kernel
+from lbmpy_walberla import generate_boundary
+
+
+#   ========================
+#      General Parameters
+#   ========================
+
+STENCIL = 'D2Q9'
+OMEGA = sp.Symbol('omega')
+LAYOUT = 'fzyx'
+
+#   PDF Fields
+pdfs, pdfs_tmp = ps.fields('pdfs(9), pdfs_tmp(9): [2D]', layout=LAYOUT)
+
+#   Velocity Output Field
+velocity = ps.fields("velocity(2): [2D]", layout=LAYOUT)
+OUTPUT = {'velocity': velocity}
+
+#   Optimization
+OPT = {'target': 'cpu',
+       'cse_global': True,
+       'symbolic_field': pdfs,
+       'symbolic_temporary_field': pdfs_tmp,
+       'field_layout': LAYOUT}
+
+
+#   ==================
+#      Method Setup
+#   ==================
+
+lbm__params = {'stencil': STENCIL,
+               'method': 'mrt_raw',
+               'relaxation_rates': [0, 0, 0, OMEGA, OMEGA, OMEGA, 1, 1, 1],
+               'cumulant': True,
+               'compressible': True}
+
+lbm_update_rule = create_lb_update_rule(optimization=OPT,
+                                        output=OUTPUT,
+                                        **lbm__params)
+
+lbm_method = lbm_update_rule.method
+
+#   ========================
+#      PDF Initialization
+#   ========================
+
+initial_rho = sp.Symbol('rho_0')
+
+pdfs_setter = macroscopic_values_setter(lbm_method, 
+                                        initial_rho, 
+                                        velocity.center_vector, 
+                                        pdfs.center_vector)
+
+#   =====================
+#      Code Generation
+#   =====================
+
+with CodeGeneration() as ctx:
+
+    #   LBM Sweep
+    generate_sweep(ctx, "CumulantMRTSweep", lbm_update_rule, field_swaps=[(pdfs, pdfs_tmp)])
+
+    #   Pack Info
+    generate_pack_info_from_kernel(ctx, "CumulantMRTPackInfo", lbm_update_rule)
+
+    #   Macroscopic Values Setter
+    generate_sweep(ctx, "DensityAndVelocityFieldSetter", pdfs_setter)
+
+    #   NoSlip Boundary
+    generate_boundary(ctx, "CumulantMRTNoSlip", NoSlip(), lbm_method)
diff --git a/apps/tutorials/codegen/CMakeLists.txt b/apps/tutorials/codegen/CMakeLists.txt
index 84d49bb28..cc3a2939e 100644
--- a/apps/tutorials/codegen/CMakeLists.txt
+++ b/apps/tutorials/codegen/CMakeLists.txt
@@ -23,4 +23,18 @@ if( WALBERLA_BUILD_WITH_CODEGEN )
                     FILES 02_LBMLatticeModelGeneration.cpp
                     DEPENDS blockforest core domain_decomposition field geometry timeloop lbm stencil vtk 02_LBMLatticeModelGenerationPython )
 
+    #   Tutorial 3: Advanced lbmpy Code Generation
+
+    walberla_generate_target_from_python( NAME 03_AdvancedLBMCodegenPython
+        FILE 03_AdvancedLBMCodegen.py
+        OUT_FILES   CumulantMRTSweep.cpp CumulantMRTSweep.h
+                    CumulantMRTPackInfo.cpp CumulantMRTPackInfo.h
+                    DensityAndVelocityFieldSetter.cpp DensityAndVelocityFieldSetter.h
+                    CumulantMRTNoSlip.cpp CumulantMRTNoSlip.h)
+
+    walberla_add_executable ( NAME 03_AdvancedLBMCodegenApp 
+                    FILES 03_AdvancedLBMCodegen.cpp
+                    DEPENDS blockforest core domain_decomposition field geometry timeloop lbm stencil vtk 03_AdvancedLBMCodegenPython )
+
+
 endif()
\ No newline at end of file
-- 
GitLab