Add just-in-time mapping

CMakeLists.txt

@@ -121,7 +121,7 @@ add_test(NAME read_write_test  COMMAND test-precice-aste WORKING_DIRECTORY "${CM
 
 # Detect and register examples as tests
 
-set(_examples lci_2d lci_3d nn nng_scalar nng_vector mapping_tester mapping_tester_serial mapping_tester_runtime mapping_tester_conv replay_mode)
+set(_examples lci_2d lci_3d nn nn_just_in_time nng_scalar nng_vector mapping_tester mapping_tester_serial mapping_tester_runtime mapping_tester_conv replay_mode)
 
 foreach(example IN LISTS _examples)
   add_test(NAME aste.example.${example}.setup

examples/mapping_tester/setup-test.json

@@ -33,8 +33,12 @@
             "executor": "cpu",
             "executor-options": ""
           },
+         "np": {
+            "kind": "nearest-projection"
+          },
          "nn": {
-            "kind": "nearest-neighbor"
+            "kind": "nearest-neighbor",
+            "batch-size": "5"
           }
         }
       },

examples/nn/precice-config.xml

@@ -26,6 +26,9 @@
     <provide-mesh name="B-Mesh" />
     <read-data name="Data" mesh="B-Mesh" />
     <mapping:nearest-neighbor constraint="consistent" direction="read" from="A-Mesh" to="B-Mesh" />
+    <!--mapping:rbf constraint="consistent" direction="read" from="A-Mesh" to="B-Mesh" >
+     <basis-function:compact-polynomial-c4 support-radius="0.1" />
+    </mapping:rbf -->
     <!-- <mapping:nearest-projection constraint="consistent" direction="read" from="A-Mesh" to="B-Mesh" /> -->
     <!-- <mapping:rbf-gaussian shape-parameter="4" constraint="consistent" direction="read" from="A-Mesh" to="B-Mesh" x-dead="false" y-dead="false" z-dead="false" /> -->
     <!-- <mapping:rbf-gaussian shape-parameter="5600" solver-rtol="1e-9" constraint="consistent" direction="read" from="A-Mesh" to="B-Mesh" x-dead="true" y-dead="false" z-dead="false" polynomial="separate" preallocation="tree"/> -->

examples/nn_just_in_time/clean.sh

@@ -0,0 +1,10 @@
+#!/usr/bin/env bash
+set -e -x
+
+rm -f -r precice-run
+rm -f -r precice-profiling
+rm -f profiling.json
+rm -f trace.json
+rm -f *.log
+rm -f map*.vtk
+rm -f fine_mesh_*.vtk

examples/nn_just_in_time/precice-config.xml

@@ -0,0 +1,47 @@
+<?xml version="1.0" encoding="UTF-8" ?>
+<precice-configuration experimental="true">
+  <log>
+    <sink
+      filter="%Severity% > info and %Participant% contains B"
+      format="---[%Participant%] %Function% %Message%"
+      type="file"
+      output="debug.log"
+      enabled="true" />
+  </log>
+
+  <profiling mode="all" />
+
+  <data:scalar name="Data" />
+
+  <mesh name="A-Mesh" dimensions="3">
+    <use-data name="Data" />
+  </mesh>
+
+  <m2n:sockets acceptor="A" connector="B" exchange-directory="." />
+
+  <participant name="A">
+    <provide-mesh name="A-Mesh" />
+    <write-data name="Data" mesh="A-Mesh" />
+  </participant>
+
+  <participant name="B">
+    <receive-mesh name="A-Mesh" from="A" direct-access="true" />
+    <read-data name="Data" mesh="A-Mesh" />
+    <mapping:nearest-neighbor constraint="consistent" direction="read" from="A-Mesh" />
+    <!--mapping:rbf constraint="consistent" direction="read" from="A-Mesh" >
+     <basis-function:compact-polynomial-c4 support-radius="0.1" />
+    </mapping:rbf -->
+    <!-- <mapping:nearest-projection constraint="consistent" direction="read" from="A-Mesh" to="B-Mesh" /> -->
+    <!-- <mapping:rbf-gaussian shape-parameter="4" constraint="consistent" direction="read" from="A-Mesh" to="B-Mesh" x-dead="false" y-dead="false" z-dead="false" /> -->
+    <!-- <mapping:rbf-gaussian shape-parameter="5600" solver-rtol="1e-9" constraint="consistent" direction="read" from="A-Mesh" to="B-Mesh" x-dead="true" y-dead="false" z-dead="false" polynomial="separate" preallocation="tree"/> -->
+    <!-- <mapping:rbf-volume-splines solver-rtol="1e-9" constraint="consistent" direction="read" from="A-Mesh" to="B-Mesh" x-dead="false" y-dead="false" z-dead="false" /> -->
+    <!-- <export:vtk every-n-time-windows="1" directory="vtkB/" normals="0"/> -->
+  </participant>
+
+  <coupling-scheme:parallel-explicit>
+    <participants first="A" second="B" />
+    <max-time value="1.0" />
+    <time-window-size value="1" />
+    <exchange data="Data" mesh="A-Mesh" from="A" to="B" />
+  </coupling-scheme:parallel-explicit>
+</precice-configuration>

examples/nn_just_in_time/run.sh

@@ -0,0 +1,12 @@
+#!/usr/bin/env bash
+set -e -x
+
+# Calculate franke function on fine mesh
+precice-aste-evaluate -m ../fine_mesh.vtk -f "franke3d" -d "Franke Function" -o "fine_mesh_nn.vtk"
+
+# Map from the finer mesh to coarser mesh
+precice-aste-run -v -p A --mesh fine_mesh_nn --data "Franke Function" &
+precice-aste-run -v -p B --mesh ../coarse_mesh --output map_nn --data "InterpolatedData" --read-just-in-time 0
+
+# Calculate statistics
+precice-aste-evaluate -m map_nn.vtk -f "franke3d" -d difference --diffdata "InterpolatedData" --diff

src/common.cpp

@@ -22,7 +22,7 @@ OptionMap getOptions(int argc, char *argv[])
       "precice-aste-run will look for timeseries as well as distributed meshes (e.g. from preCICE exports) "
       "automatically and load them if required.")(
       "output", po::value<std::string>(),
-      "Output file name.")("vector", po::bool_switch(), "Distinguish between vector valued data and scalar data")("verbose,v", po::bool_switch(), "Enable verbose output")("all,a", po::bool_switch(), "Enable output from secondary ranks");
+      "Output file name.")("vector", po::bool_switch(), "Distinguish between vector valued data and scalar data")("verbose,v", po::bool_switch(), "Enable verbose output")("all,a", po::bool_switch(), "Enable output from secondary ranks")("read-just-in-time,jit", po::value<int>(), "Instead of the conventional mapping in preCICE, make use of a just-in-time mapping, where values are read in batches of <int> values.");
 
   po::variables_map vm;
 

src/modes.cpp

@@ -146,6 +146,15 @@ void aste::runMapperMode(const aste::ExecutionContext &context, const OptionMap
   const bool        isVector        = options["vector"].as<bool>();
   const std::string preciceConfig   = options["precice-config"].as<std::string>();
 
+  constexpr int conventionalReadData = -1;
+  // potentially make use of a just-in-time mapping using a prescribed batch size
+  const int batchSize = options.count("read-just-in-time")
+                            ? options["read-just-in-time"].as<int>()
+                            : conventionalReadData;
+  if (batchSize < -1) {
+    ASTE_ERROR << "Batch size needs to be greater or equal to zero.";
+  }
+
   addLogIdentity(participantName, context.rank);
   ASTE_INFO << "ASTE Running in mapping test mode";
 
@@ -186,7 +195,11 @@ void aste::runMapperMode(const aste::ExecutionContext &context, const OptionMap
     asteConfiguration.participantName       = "B";
     asteConfiguration.preciceConfigFilename = preciceConfig;
     aste::asteInterface asteInterface;
-    asteInterface.meshName       = "B-Mesh";
+    if (batchSize == conventionalReadData) {
+      asteInterface.meshName = "B-Mesh";
+    } else {
+      asteInterface.meshName = "A-Mesh";
+    }
     asteInterface.meshFilePrefix = meshname;
     asteInterface.meshes         = aste::BaseName(meshname).findAll(context);
 
@@ -212,7 +225,14 @@ void aste::runMapperMode(const aste::ExecutionContext &context, const OptionMap
   asteInterface.meshes.front().loadMesh(asteInterface.mesh, dim, requireConnectivity);
   asteInterface.meshes.front().loadData(asteInterface.mesh);
   ASTE_INFO << "The loaded mesh " << asteInterface.meshes.front().filename() << " contains: " << asteInterface.mesh.summary();
-  auto vertexIDs = aste::setupMesh(preciceInterface, asteInterface.mesh, asteInterface.meshName);
+
+  std::vector<int>    vertexIDs;
+  std::vector<double> coordinates;
+  if (batchSize == conventionalReadData) {
+    vertexIDs = aste::setupMesh(preciceInterface, asteInterface.mesh, asteInterface.meshName);
+  } else {
+    coordinates = aste::setupDirectMeshAccess(preciceInterface, asteInterface.mesh, asteInterface.meshName, context);
+  }
   ASTE_DEBUG << "Mesh setup completed on Rank " << context.rank;
 
   if (preciceInterface.requiresInitialData()) {
@@ -263,8 +283,33 @@ void aste::runMapperMode(const aste::ExecutionContext &context, const OptionMap
     for (auto &asteInterface : asteConfiguration.asteInterfaces) {
       for (auto &meshdata : asteInterface.mesh.meshdata) {
         if (meshdata.type == aste::datatype::READ) {
-          meshdata.dataVector.resize(vertexIDs.size() * meshdata.numcomp);
-          preciceInterface.readData(asteInterface.meshName, "Data", vertexIDs, dt, meshdata.dataVector);
+          if (batchSize == conventionalReadData) {
+            meshdata.dataVector.resize(vertexIDs.size() * meshdata.numcomp);
+            preciceInterface.readData(asteInterface.meshName, "Data", vertexIDs, dt, meshdata.dataVector);
+          } else {
+            int dim       = preciceInterface.getMeshDimensions(asteInterface.meshName);
+            int nVertices = coordinates.size() / dim;
+            meshdata.dataVector.resize(nVertices * meshdata.numcomp);
+
+            // pass everything in one go
+            if (batchSize == 0) {
+              preciceInterface.mapAndReadData(asteInterface.meshName, "Data", coordinates, dt, meshdata.dataVector);
+            } else { // pass in batches
+              for (int i = 0; i < nVertices; i += batchSize) {
+                int availableVertices = std::min(i + batchSize, nVertices);
+                int coordStart        = i * dim;
+                int coordEnd          = availableVertices * dim;
+                int dataStart         = i * meshdata.numcomp;
+                int dataEnd           = availableVertices * meshdata.numcomp;
+
+                ::precice::span<const double> coordinateBatch(coordinates.data() + coordStart, coordEnd - coordStart);
+                ::precice::span<double>       dataBatch(meshdata.dataVector.data() + dataStart, dataEnd - dataStart);
+
+                // Call the API function with the current batch
+                preciceInterface.mapAndReadData(asteInterface.meshName, "Data", coordinateBatch, dt, dataBatch);
+              }
+            }
+          }
           ASTE_DEBUG << "Data read: " << asteInterface.mesh.previewData(meshdata);
         }
       }

src/utilities.cpp

@@ -72,6 +72,51 @@ void aste::setupEdgeIDs(precice::Participant &interface, const aste::Mesh &mesh,
   }
 }
 
+std::vector<double> aste::setupDirectMeshAccess(precice::Participant &interface, const aste::Mesh &mesh, const std::string &meshName, const aste::ExecutionContext &exec)
+{
+  ASTE_DEBUG << "Setting up direct access for: " << meshName;
+  ASTE_DEBUG << "Setting up coordinates vector for mesh...";
+
+  const auto          dim       = interface.getMeshDimensions(meshName);
+  const auto          nvertices = mesh.positions.size();
+  std::vector<double> coordinates(dim * nvertices);
+  for (unsigned long i = 0; i < nvertices; ++i) {
+    const auto &pos = mesh.positions[i];
+    assert(pos.size() == static_cast<unsigned int>(dim));
+    std::copy(pos.begin(), pos.end(), &coordinates[i * dim]);
+  }
+
+  ASTE_DEBUG << "Computing the (rank-local) bounding box:";
+
+  std::vector<double> bb(2 * dim);
+
+  // for parallel runs, we define a proper bounding box
+  if (exec.isParallel()) {
+    // Initialize min and max values
+    for (int i = 0; i < dim; ++i) {
+      bb[2 * i]     = std::numeric_limits<double>::max();    // min values
+      bb[2 * i + 1] = std::numeric_limits<double>::lowest(); // max values
+    }
+
+    // Iterate through the coordinates
+    for (size_t i = 0; i < coordinates.size(); i += dim) {
+      for (int j = 0; j < dim; ++j) {
+        double coord  = coordinates[i + j];
+        bb[2 * j]     = std::min(coord - 1e-12, bb[2 * j]);     // Update min
+        bb[2 * j + 1] = std::max(coord + 1e-12, bb[2 * j + 1]); // Update max
+      }
+    }
+  } else {
+    // for serial runs, we just take everything (to account for potential discretization imbalances)
+    for (int i = 0; i < dim; ++i) {
+      bb[2 * i]     = std::numeric_limits<double>::lowest(); // min values
+      bb[2 * i + 1] = std::numeric_limits<double>::max();    // max values
+    }
+  }
+  interface.setMeshAccessRegion(meshName, bb);
+  return coordinates;
+}
+
 std::vector<int> aste::setupMesh(precice::Participant &interface, const aste::Mesh &mesh, const std::string &meshName)
 {
   auto tstart = std::chrono::steady_clock::now();

src/utilities.hpp

@@ -60,4 +60,6 @@ void setupEdgeIDs(precice::Participant &interface, const aste::Mesh &mesh, const
  * @return std::vector<int> a vector of vertexIDs in preCICE
  */
 std::vector<int> setupMesh(precice::Participant &interface, const aste::Mesh &mesh, const std::string &meshName);
+
+std::vector<double> setupDirectMeshAccess(precice::Participant &interface, const aste::Mesh &mesh, const std::string &meshName, const aste::ExecutionContext &exec);
 } // namespace aste

tools/mapping-tester/config-template.xml

@@ -1,6 +1,6 @@
 <?xml version="1.0"?>
 
-<precice-configuration>
+<precice-configuration experimental="true">
 
   <log>
     <sink type="file" output="debug.log" filter="" format='%TimeStamp(format="%H:%M:%S.%f")%|%Participant%|%Rank%|%Module%|l%Line%|%Function%|%Severity%%Message%' />
@@ -15,10 +15,12 @@
     <mesh name="A-Mesh" dimensions="3">
       <use-data name="Data" />
     </mesh>
-
+    <!-- batch size option works only for consistent mappings -->
+    {% if mapping.get('batch-size') == "-1" %}
     <mesh name="B-Mesh" dimensions="3">
       <use-data name="Data" />
     </mesh>
+    {% endif %}
 
     <m2n:sockets acceptor="A" connector="B"{% if network %} network="{{ network }}"{% endif %} exchange-directory="." />
 
@@ -41,13 +43,16 @@
     </participant>
 
     <participant name="B">
+    {% if mapping.get('batch-size') == "-1" %}
       <provide-mesh name="B-Mesh" />
-      <read-data name="Data" mesh="B-Mesh" />
+    {% endif %}
+
+      <read-data name="Data" {% if mapping.get('batch-size') == "-1" %} mesh="B-Mesh" {% else %} mesh="A-Mesh" {% endif %} />
 
       {% if mapping.constraint == "consistent" %}
-      <receive-mesh name="A-Mesh" from="A" />
+      <receive-mesh name="A-Mesh" from="A" {% if not mapping.get('batch-size') == "-1" %}direct-access="true"{% endif %} />
 
-      <mapping:{{ mapping.kind }} {{ mapping.options }} constraint="consistent" direction="read" from="A-Mesh" to="B-Mesh" >
+      <mapping:{{ mapping.kind }} {{ mapping.options }} constraint="consistent" direction="read" from="A-Mesh" {% if mapping.get('batch-size') == "-1" %} to="B-Mesh" {% endif %} >
         {% if mapping.kind.startswith("rbf") %}
           {% if mapping.executor %}
         <executor:{{ mapping.executor }} {{ mapping.executoroptions }} />

tools/mapping-tester/generate.py

@@ -55,6 +55,7 @@ def generateCases(setup):
                                         "executoroptions": mapping.get(
                                             "executor-options", ""
                                         ),
+                                        "batch-size": mapping.get("batch-size", "-1"),
                                     },
                                     "A": {
                                         "ranks": ranksA,
@@ -158,9 +159,15 @@ def createRunScript(outdir: pathlib.Path, path: pathlib.Path, case):
         path, walk_up=True
     )
     mapped_data_name = case["function"] + "(mapped)"
+
+    # Handle batch_size condition
+    batch_size_flag = ""
     output = "--output mapped" if case["computeAccuracy"] else ""
-    bcmd = f'env {time_command} -f %M -a -o memory-B.log precice-aste-run -v -a -p B --data "{mapped_data_name}" --mesh {bmeshLocation} {output} || kill 0 &'
+    if case["mapping"]["batch-size"] != "-1":
+        batch_size_flag = "--read-just-in-time {}".format(case["mapping"]["batch-size"])
 
+    # Generate runner script for participant B
+    bcmd = f'env {time_command} -f %M -a -o memory-B.log precice-aste-run -v -a -p B --data "{mapped_data_name}" --mesh {bmeshLocation} {output} {batch_size_flag} || kill 0 &'
     if branks > 1:
         bcmd = "mpirun -n {} $ASTE_B_MPIARGS {}".format(branks, bcmd)