TDE initializers (#534)

* Read / Write starData, tdeOrbitInit, polytrope

* add std-relax propagator

* Move polytrope creation functions to physics/. Code-reuse in Evrard and Polytrope initializer. TDE initializers compile switch.

---------

Co-authored-by: Sebastian Keller <sebastian.f.keller@gmail.com>

Author: nknk567

CMakeLists.txt

@@ -109,6 +109,7 @@ if (SPH_EXA_WITH_GRACKLE)
 endif()
 
 option(SPH_EXA_WITH_DISKS "Enable disk physics and propagator" OFF)
+option(SPH_EXA_WITH_TDE_INIT "Enable TDE initializers" OFF)
 
 add_subdirectory(domain)
 add_subdirectory(ryoanji)

main/CMakeLists.txt

@@ -4,6 +4,7 @@ set(COOLING_DIR ${PROJECT_SOURCE_DIR}/physics/cooling/include)
 set(RYOANJI_DIR ${PROJECT_SOURCE_DIR}/ryoanji/src)
 set(SPH_DIR ${PROJECT_SOURCE_DIR}/sph/include)
 set(DISK_DIR ${PROJECT_SOURCE_DIR}/physics/Disk/include)
+set(POLYTROPE_DIR ${PROJECT_SOURCE_DIR}/physics/polytrope/include)
 
 add_subdirectory(src)
 if (BUILD_TESTING)

main/src/init/CMakeLists.txt

@@ -1,16 +1,27 @@
 add_library(sim_init sim_init.cpp)
 target_include_directories(sim_init PRIVATE ${PROJECT_SOURCE_DIR}/main/src ${COOLING_DIR} ${CSTONE_DIR} ${SPH_DIR}
-        ${MPI_CXX_INCLUDE_PATH})
+        ${MPI_CXX_INCLUDE_PATH} ${DISK_DIR})
 target_link_libraries(sim_init PRIVATE ${MPI_CXX_LIBRARIES} OpenMP::OpenMP_CXX)
+
+function(enableTDEInit exename)
+    if (SPH_EXA_WITH_TDE_INIT)
+        target_include_directories(${exename} PRIVATE ${POLYTROPE_DIR})
+        target_compile_definitions(${exename} PRIVATE SPH_EXA_HAVE_TDE_INIT)
+        target_link_libraries(${exename} PRIVATE polytrope)
+    endif ()
+endfunction()
+
+enableTDEInit(sim_init)
 enableGrackle(sim_init)
 
 if (CMAKE_CUDA_COMPILER OR CMAKE_HIP_COMPILER)
     add_library(sim_init_gpu sim_init.cpp)
     target_compile_definitions(sim_init_gpu PRIVATE USE_CUDA)
     target_include_directories(sim_init_gpu PRIVATE ${PROJECT_SOURCE_DIR}/main/src ${COOLING_DIR} ${CSTONE_DIR}
-            ${SPH_DIR} ${MPI_CXX_INCLUDE_PATH})
+            ${SPH_DIR} ${MPI_CXX_INCLUDE_PATH} ${DISK_DIR})
     target_link_libraries(sim_init_gpu PRIVATE ${MPI_CXX_LIBRARIES} OpenMP::OpenMP_CXX)
     enableGrackle(sim_init_gpu)
+    enableTDEInit(sim_init_gpu)
 endif ()
 
 if (CMAKE_CUDA_COMPILER)

main/src/init/evrard_cooling_init.hpp

@@ -30,11 +30,8 @@
 
 #pragma once
 
-#include "cstone/primitives/primitives_acc.hpp"
-
 #include "evrard_init.hpp"
 #include "cooling/cooler.hpp"
-
 #include "cooling/init_chemistry.h"
 
 namespace sphexa
@@ -63,31 +60,26 @@ InitSettings evrardCoolingConstants()
 }
 
 template<class Dataset>
-class EvrardGlassSphereCooling : public EvrardGlassSphere<Dataset>
+class EvrardGlassSphereCooling : public RadialProfile<Dataset>
 {
-    using Base = EvrardGlassSphere<Dataset>;
-    mutable InitSettings settings_;
+    using Base = RadialProfile<Dataset>;
+    using Base::settings_;
 
 public:
     EvrardGlassSphereCooling(std::string initBlock, std::string settingsFile, IFileReader* reader)
-        : EvrardGlassSphere<Dataset>(initBlock, settingsFile, reader)
+        : Base(std::move(initBlock), evrardCoolingConstants(), std::move(settingsFile), reader)
     {
-        Dataset d;
-        settings_ = buildSettings(d, evrardCoolingConstants(), settingsFile, reader);
-        Base::resetConstants(settings_);
     }
 
     cstone::Box<typename Dataset::RealType> init(int rank, int numRanks, size_t cbrtNumPart, Dataset& simData,
                                                  IFileReader* reader) const override
     {
-        constexpr bool gpu = cstone::HaveGpu<typename Dataset::AcceleratorType>{};
-        auto           box = Base::init(rank, numRanks, cbrtNumPart, simData, reader);
-        cstone::fill<gpu>(simData.hydro.u.begin(), simData.hydro.u.end(), settings_.at("u0"));
+        auto radialTransform = [](auto r) { return std::sqrt(r); };
+        auto globalBox = Base::init(rank, numRanks, cbrtNumPart, simData, reader, settings_.at("r"), radialTransform);
+        initEvrardFields(simData.hydro, settings_);
         cooling::initChemistryData(simData.chem, simData.hydro.x.size());
-        return box;
+        return globalBox;
     }
-
-    [[nodiscard]] const InitSettings& constants() const override { return settings_; }
 };
 
 } // namespace sphexa

main/src/init/evrard_init.hpp

@@ -1,26 +1,10 @@
 /*
- * MIT License
+ * SPH-EXA
  *
- * Copyright (c) 2021 CSCS, ETH Zurich
- *               2021 University of Basel
+ * Copyright (c) 2026 CSCS, ETH Zurich, University of Zurich, University of Basel
  *
- * Permission is hereby granted, free of charge, to any person obtaining a copy
- * of this software and associated documentation files (the "Software"), to deal
- * in the Software without restriction, including without limitation the rights
- * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
- * copies of the Software, and to permit persons to whom the Software is
- * furnished to do so, subject to the following conditions:
- *
- * The above copyright notice and this permission notice shall be included in all
- * copies or substantial portions of the Software.
- *
- * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
- * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
- * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
- * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
- * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
- * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
- * SOFTWARE.
+ * Please, refer to the LICENSE file in the root directory.
+ * SPDX-License-Identifier: MIT License
  */
 
 /*! @file
@@ -34,19 +18,14 @@
 #include <map>
 
 #include "cstone/primitives/primitives_acc.hpp"
-#include "cstone/sfc/box.hpp"
 #include "cstone/tree/continuum.hpp"
-#include "sph/eos.hpp"
 
-#include "isim_init.hpp"
-#include "early_sync.hpp"
-#include "grid.hpp"
-#include "utils.hpp"
+#include "radial_profile.hpp"
 
 namespace sphexa
 {
 
-std::map<std::string, double> evrardConstants()
+InitSettings evrardConstants()
 {
     return {{"gravConstant", 1.}, {"r", 1.},          {"mTotal", 1.}, {"gamma", 5. / 3.}, {"u0", 0.05},
             {"minDt", 1e-4},      {"minDt_m1", 1e-4}, {"mui", 10},    {"ng0", 100},       {"ngmax", 150}};
@@ -101,22 +80,6 @@ void initEvrardFields(Dataset& d, const std::map<std::string, double>& constants
     d.h = std::move(h);
 }
 
-template<class Vector>
-void contractRhoProfile(Vector& x, Vector& y, Vector& z)
-{
-#pragma omp parallel for schedule(static)
-    for (size_t i = 0; i < x.size(); i++)
-    {
-        auto radius0 = std::sqrt(x[i] * x[i] + y[i] * y[i] + z[i] * z[i]);
-
-        // multiply coordinates by sqrt(r) to generate a density profile ~ 1/r
-        auto contraction = std::sqrt(radius0);
-        x[i] *= contraction;
-        y[i] *= contraction;
-        z[i] *= contraction;
-    }
-}
-
 //! @brief Estimate SFC partition of the Evrard sphere based on approximate continuum particle counts
 template<class KeyType, class T>
 std::tuple<KeyType, KeyType> estimateEvrardSfcPartition(size_t cbrtNumPart, const cstone::Box<T>& box, int rank,
@@ -142,74 +105,25 @@ std::tuple<KeyType, KeyType> estimateEvrardSfcPartition(size_t cbrtNumPart, cons
 }
 
 template<class Dataset>
-class EvrardGlassSphere : public ISimInitializer<Dataset>
+class EvrardGlassSphere : public RadialProfile<Dataset>
 {
-    std::string          glassBlock;
-    mutable InitSettings settings_;
+    using Base = RadialProfile<Dataset>;
+    using Base::settings_;
 
 public:
     explicit EvrardGlassSphere(std::string initBlock, std::string settingsFile, IFileReader* reader)
-        : glassBlock(std::move(initBlock))
+        : Base(std::move(initBlock), evrardConstants(), std::move(settingsFile), reader)
     {
-        Dataset d;
-        settings_ = buildSettings(d, evrardConstants(), settingsFile, reader);
     }
 
     cstone::Box<typename Dataset::RealType> init(int rank, int numRanks, size_t cbrtNumPart, Dataset& simData,
                                                  IFileReader* reader) const override
     {
-        auto& d       = simData.hydro;
-        using KeyType = typename Dataset::KeyType;
-        using T       = typename Dataset::RealType;
-
-        std::vector<T> xBlock, yBlock, zBlock;
-        readTemplateBlock(glassBlock, reader, xBlock, yBlock, zBlock);
-        size_t blockSize = xBlock.size();
-
-        int               multi1D      = std::rint(cbrtNumPart / std::cbrt(blockSize));
-        cstone::Vec3<int> multiplicity = {multi1D, multi1D, multi1D};
-
-        T              r = settings_.at("r");
-        cstone::Box<T> globalBox(-r, r, cstone::BoundaryType::open);
-
-        auto [keyStart, keyEnd] = equiDistantSfcSegments<KeyType>(rank, numRanks, 100);
-
-        std::vector<T> x, y, z;
-        auto           t0 = std::chrono::high_resolution_clock::now();
-        assembleCuboid<T>(keyStart, keyEnd, globalBox, multiplicity, xBlock, yBlock, zBlock, x, y, z);
-        cutSphere(r, x, y, z);
-        auto t1 = std::chrono::high_resolution_clock::now();
-        if (rank == 0) std::cout << "assembly " << std::chrono::duration<float>(t1 - t0).count() << std::endl;
-
-        size_t numParticlesGlobal = x.size();
-        MPI_Allreduce(MPI_IN_PLACE, &numParticlesGlobal, 1, MpiType<size_t>{}, MPI_SUM, simData.comm);
-
-        contractRhoProfile(x, y, z);
-
-        t0  = std::chrono::high_resolution_clock::now();
-        d.x = x; // uploads to GPU if active
-        d.y = y;
-        d.z = z;
-        syncCoords<KeyType>(rank, numRanks, numParticlesGlobal, d.x, d.y, d.z, globalBox);
-        // 2nd call needed to reduce imbalance, 1st call is not able to fully balance number of particles per rank
-        syncCoords<KeyType>(rank, numRanks, numParticlesGlobal, d.x, d.y, d.z, globalBox);
-        t1 = std::chrono::high_resolution_clock::now();
-        if (rank == 0) std::cout << "earlySync " << std::chrono::duration<float>(t1 - t0).count() << std::endl;
-
-        d.resize(d.x.size());
-
-        settings_["numParticlesGlobal"] = double(numParticlesGlobal);
-        BuiltinWriter attributeSetter(settings_);
-        d.loadOrStoreAttributes(&attributeSetter);
-
-        initEvrardFields(d, settings_);
-
+        auto radialTransform = [](auto r) { return std::sqrt(r); };
+        auto globalBox = Base::init(rank, numRanks, cbrtNumPart, simData, reader, settings_.at("r"), radialTransform);
+        initEvrardFields(simData.hydro, settings_);
         return globalBox;
     }
-
-    void resetConstants(InitSettings newSettings) { settings_ = std::move(newSettings); }
-
-    [[nodiscard]] const InitSettings& constants() const override { return settings_; }
 };
 
 } // namespace sphexa

main/src/init/factory.hpp

@@ -92,6 +92,26 @@ std::unique_ptr<ISimInitializer<Dataset>> initializerFactory(std::string testCas
         if (glassBlock.empty()) { throw std::runtime_error("need a valid glass block for evrard-cooling\n"); }
         return SimInitializers<Dataset>::makeEvrardCooling(glassBlock, settingsFile, reader);
     }
+    if (testNamedBase == "polytrope")
+    {
+        if (glassBlock.empty()) { throw std::runtime_error("need a valid glass block to create polytrope\n"); }
+        else { return SimInitializers<Dataset>::makePolytrope(glassBlock, settingsFile, reader); }
+    }
+    if (testCase.starts_with("tde-orbit"))
+    {
+        // use as: --init tde-orbit:file.h5:2 to read from step 2 of file.h5
+        const std::string filePathAndStep = strAfterSign(testCase, ":");
+        const std::string filePath        = strBeforeSign(filePathAndStep, ":");
+        const int         step            = numberAfterSign(filePathAndStep, ":");
+        std::printf("%s\n", testCase.c_str());
+        std::printf("%s\n", filePath.c_str());
+        std::printf("step: %d\n", step);
+        if (std::filesystem::exists(filePath))
+        {
+            return SimInitializers<Dataset>::makeTDEOrbitInit(filePath, step, reader);
+        }
+    }
+
     if (std::filesystem::exists(strBeforeSign(testCase, ":")))
     {
         return SimInitializers<Dataset>::makeFile(strBeforeSign(testCase, ":"), numberAfterSign(testCase, ":"), reader);

main/src/init/file_init.hpp

@@ -244,13 +244,16 @@ class FileSplitInit : public ISimInitializer<Dataset>
         replicateField(reader, "vx", d.vx, T(1));
         replicateField(reader, "vy", d.vy, T(1));
         replicateField(reader, "vz", d.vz, T(1));
-        replicateField(reader, "temp", d.temp, T(1));
+        if (d.isAllocated("temp")) { replicateField(reader, "temp", d.temp, T(1)); }
+        else if (d.isAllocated("u")) { replicateField(reader, "u", d.u, T(1)); }
         cstone::fill<gpu>(d.du_m1.begin(), d.du_m1.end(), 0);
         cstone::fill<gpu>(d.rung.begin(), d.rung.end(), 0);
         cstone::scaleGpuAcc<gpu>(d.vx.data(), d.vx.data() + d.vx.size(), d.x_m1.data(), d.minDt);
         cstone::scaleGpuAcc<gpu>(d.vy.data(), d.vy.data() + d.vy.size(), d.y_m1.data(), d.minDt);
         cstone::scaleGpuAcc<gpu>(d.vz.data(), d.vz.data() + d.vz.size(), d.z_m1.data(), d.minDt);
 
+        generateParticleIDs<gpu>(d.id);
+
         if (d.isAllocated("alpha"))
         {
             try

main/src/init/isim_init.hpp

@@ -72,6 +72,8 @@ struct SimInitializers
     static InitPtr makeSedovGrid();
     static InitPtr makeTurbulence(std::string glassBlock, std::string settingsFile, IFileReader* reader);
     static InitPtr makeWindShock(std::string glassBlock, std::string settingsFile, IFileReader* reader);
+    static InitPtr makePolytrope(std::string glassBlock, std::string settingsFile, IFileReader* reader);
+    static InitPtr makeTDEOrbitInit(const std::string& filePath, int initStep, IFileReader* reader);
 };
 
 extern template struct SimInitializers<SimulationData<cstone::CpuTag>>;

main/src/init/kelvin_helmholtz_init.hpp

@@ -33,12 +33,11 @@
 
 #include "cstone/primitives/primitives_acc.hpp"
 #include "cstone/sfc/box.hpp"
-#include "cstone/sfc/sfc.hpp"
-#include "cstone/primitives/gather.hpp"
 
-#include "isim_init.hpp"
 #include "grid.hpp"
 #include "utils.hpp"
+#include "early_sync.hpp"
+#include "isim_init.hpp"
 
 namespace sphexa
 {