Implement Adaptivity interpolation using RBF (#242)

* Default behaviour of Adaptivity still copies results from nearest representatives.
* Functional dependencies can be specified. Those will be interpolated using RBF.

---------

Co-authored-by: Ishaan Desai <ishaandesai@gmail.com>

Author: Snapex2409

.github/workflows/check-coverage.yml

@@ -83,6 +83,7 @@ jobs:
           . ../../.venv/bin/activate
           mpirun -n 2 --allow-run-as-root -x PYTHONPATH=. python3 -m coverage run --parallel-mode --source=micro_manager -m unittest test_adaptivity_parallel
           mpirun -n 2 --allow-run-as-root -x PYTHONPATH=. python3 -m coverage run --parallel-mode --source=micro_manager -m unittest test_load_balancing
+          mpirun -n 2 --allow-run-as-root -x PYTHONPATH=. python3 -m coverage run --parallel-mode --source=micro_manager -m unittest test_interpolation
 
       - name: Combine coverage data
         working-directory: micro-manager/tests/unit

.github/workflows/run-adaptivity-tests-parallel.yml

@@ -113,3 +113,13 @@ jobs:
           . .venv/bin/activate
           cd tests/unit
           mpiexec -n 4 --oversubscribe --allow-run-as-root python3 -m unittest test_load_balancing.py
+
+      - name: Run interpolation unit tests with 2 ranks
+        timeout-minutes: 3
+        working-directory: micro-manager
+        run: |
+          . .venv/bin/activate
+          pip install .[sklearn]
+          pip uninstall -y pyprecice
+          cd tests/unit
+          mpiexec -n 2 --allow-run-as-root python3 -m unittest test_interpolation.py

CHANGELOG.md

@@ -1,5 +1,9 @@
 # Micro Manager changelog
 
+## latest
+
+- Added RBF interpolation, currently used within adaptivity for output interpolation [#242](https://github.com/precice/micro-manager/pull/242)
+
 ## v0.10.1
 
 - Fixed parameter inspection of the `initialize()` routine when micro simulation is pybind11 wrapped [#281](https://github.com/precice/micro-manager/pull/281)

docs/configuration.md

@@ -113,6 +113,57 @@ To turn on adaptivity, set `"adaptivity": true` in `simulation_params`. Then und
 | `similarity_measure`              | Similarity measure to be used for adaptivity. Can be either `L1`, `L2`, `L1rel` or `L2rel`. By default, `L1` is used. The `rel` variants calculate the respective relative norms. This parameter is *optional*.                                                                                     | `L2rel`       |
 | `lazy_initialization`             | Set to `true` to lazily create and initialize micro simulations. If selected, micro simulation objects are created only when the micro simulation is activated for the first time.                                                                                                                  | `false`       |
 | `load_balancing`                  | Set to `true` to dynamically balance simulations for parallel runs. See [load balancing settings](#load-balancing) below.                                                                                                                                                                           | `false`       |
+| `mappings`                        | Optional interpolation of results. Set to list of mapping configurations. See below for further details.                                                                                                                                                                                            | `[]`          |
+
+Results of inactive simulations can be interpolated from active simulations using radial basis function interpolation. For data in `write_data_names`, a function
+can be defined from `read_data_names` to `write_data_names`. When using multiple functions, their interpolation target, i.e., fields
+of `write_data_names` must be mutually disjunct. Mappings can be defined as:
+
+```json
+"mappings": [
+    {
+        "src_fields": ["input1", "input2"],
+        "dst_fields": ["output1", "output2"],
+        "n_neighbors": 50,
+        "rbf_config": {
+            "basis": {
+                "type": "c6"
+            }
+        },
+        "domain_config": {
+            "max_filling": 8,
+            "coarsening_factor": 2,
+            "projection": {
+                "type": "std",
+                "target_dims": 3
+            }
+        }
+    },
+]
+```
+
+| Parameter       | Description                             | Default |
+|-----------------|-----------------------------------------|---------|
+| `src_fields`    | List of entries from `read_data_names`  | `None`  |
+| `dst_fields`    | List of entries from `write_data_names` | `None`  |
+| `n_neighbours`  | The minimum amount of support points.   | `50`    |
+| `rbf_config`    | RBF interpolation configuration.        | `None`  |
+| `domain_config` | Function source domain description.     |         |
+
+A selection of basis functions is available: `c0`, `c2`, `c4`, `c6`.
+The domain must be described/further configured as input data is shared across ranks and must be redistributed for interpolation.
+Towards this, spatial discretization techniques are used. For better performance, data can be projected to a lower dimensional space
+using the fields with the highest standard deviation.
+
+| Parameter           | Description                                                               | Default    |
+|---------------------|---------------------------------------------------------------------------|------------|
+| `use_pu`            | Enables PU-RBF. (currently not supported)                                 | `False`    |
+| `pu_overlap`        | Controlls overlap radius for PU decomposition.                            | `0.1`      |
+| `basis`             | RBF basis function: `c0`, `c2`, `c4`, `c6`                                | `None`     |
+| `max_filling`       | Tunes maximum filling of tree nodes used during decomposition.            | `8`        |
+| `coarsening_factor` | Adjusts the fidelity of the discretized domain. Only integer values >= 1. | `2`        |
+| `projection`        | Either `std` or `identity`.                                               | `identity` |
+| `target_dims`       | Only if `std` is used. Denotes the target dimension after projection.     | `None`     |
 
 Example of adaptivity configuration is
 

micro_manager/adaptivity/adaptivity.py

@@ -58,6 +58,13 @@ def __init__(
 
         self._max_similarity_dist = 0.0
 
+        self._interpolation = None
+        self._interp_min = -1
+        self._mappings = []
+        self._mapping_configs = []
+        mappings = configurator.get_adaptivity_mapping_configs()
+        self._load_mappings(mappings)
+
         # is_sim_active: 1D array having state (active or inactive) of each micro simulation
         # Start adaptivity calculation with all sims active
         # This array is modified in place via the function update_active_sims and update_inactive_sims
@@ -109,6 +116,65 @@ def __init__(
 
             self._metrics_logger.log_info("n|n active|n inactive|assoc ranks")
 
+    def _load_mappings(self, mappings: list) -> None:
+        """
+        Translates the mapping information provided from the configuration file into a
+        interpolation method parseable structure.
+
+        This will populate the self._mappings and self._mapping_configs buffers.
+        Called once during __init__.
+
+        Parameters
+        ----------
+        mappings : list
+            List of mappings as provided by the configuration file.
+        """
+        for mapping in mappings:
+            src_fields = mapping["src_fields"]
+            dst_fields = mapping["dst_fields"]
+            n_neighbors = mapping["n_neighbors"]
+            if self._interp_min == -1:
+                self._interp_min = n_neighbors
+            else:
+                self._interp_min = min(n_neighbors, self._interp_min)
+
+            self._mappings.append((src_fields, dst_fields))
+            config = {}
+            if "use_pu" in mapping["rbf_config"]:
+                config["use_pu"] = mapping["rbf_config"]["use_pu"]
+            if "pu_overlap" in mapping["rbf_config"]:
+                config["pu_overlap"] = mapping["rbf_config"]["pu_overlap"]
+            config["pu_cluster_size"] = n_neighbors
+            if "basis" in mapping["rbf_config"]:
+                if "type" in mapping["rbf_config"]["basis"]:
+                    config["basis"] = mapping["rbf_config"]["basis"]["type"]
+                if (
+                    config["basis"] == "gauss"
+                    and "eps" in mapping["rbf_config"]["basis"]
+                ):
+                    config["gauss_eps"] = mapping["rbf_config"]["basis"]["eps"]
+
+            dom_config = {}
+            dom_config["n_neighbors"] = n_neighbors
+            if "max_filling" in mapping["domain_config"]:
+                dom_config["max_filling"] = mapping["domain_config"]["max_filling"]
+            if "coarsening_factor" in mapping["domain_config"]:
+                dom_config["coarsening_factor"] = mapping["domain_config"][
+                    "coarsening_factor"
+                ]
+            if "projection" in mapping["domain_config"]:
+                if "type" in mapping["domain_config"]["projection"]:
+                    dom_config["projection_type"] = mapping["domain_config"][
+                        "projection"
+                    ]["type"]
+                if "target_dims" in mapping["domain_config"]["projection"]:
+                    dom_config["projection_std_dims"] = mapping["domain_config"][
+                        "projection"
+                    ]["target_dims"]
+
+            config["domain_config"] = dom_config
+            self._mapping_configs.append(config)
+
     def _update_similarity_dists(self, dt: float, data: dict) -> None:
         """
         Calculate metric which determines if two micro simulations are similar enough to have one of them deactivated.
@@ -199,6 +265,89 @@ def _check_for_deactivation(self, active_id: int, active_ids: list) -> bool:
                     return True
         return False
 
+    def _interpolate_output(self, micro_input, micro_sims_output) -> None:
+        """
+        Interpolates the micro output based on the available inputs and outputs using the selected
+        interpolation method and desired mappings.
+        Will compute functions f1 ... fN described in the config.
+        fi: X -> Y, X and Y must be subsets of the coupled fields.
+        Every output field may only be used once as interpolation target, meaning there may not be
+        a function fi and fj with shared Yi and Yj.
+
+        This method will edit the output buffer, instead of returning a new buffer.
+
+        Parameters
+        ----------
+        micro_input : list
+            List of all local micro simulation inputs.
+
+        micro_sims_output : list
+            List of all local micro simulation outputs. (current state)
+        """
+        targets = []
+        for _, target_args in self._mappings:
+            targets.extend(target_args)
+        assert len(targets) == len(set(targets))
+
+        # precompute arg sizes
+        active_lids = self.get_active_sim_local_ids()
+        inactive_lids = self.get_inactive_sim_local_ids()
+        arg_sizes = {}
+        for name, value in micro_input[-1].items():
+            arg_sizes[name] = (
+                1 if type(value) != np.ndarray and type(value) != list else len(value)
+            )
+        for name, value in micro_sims_output[-1].items():
+            arg_sizes[name] = (
+                1 if type(value) != np.ndarray and type(value) != list else len(value)
+            )
+
+        # create interpolation data structures
+        n_points = len(active_lids)
+        n_points_inactive = len(inactive_lids)
+        for m_idx, fun in enumerate(self._mappings):
+            src_args, dst_args = fun
+            src_size = np.array([arg_sizes[name] for name in src_args]).sum()
+            dst_size = np.array([arg_sizes[name] for name in dst_args]).sum()
+            input_data = np.zeros((n_points, src_size))
+            output_data = np.zeros((n_points, dst_size))
+            for idx, lid in enumerate(active_lids):
+                offset = 0
+                for src_arg in src_args:
+                    input_data[idx, offset : offset + arg_sizes[src_arg]] = micro_input[
+                        lid
+                    ][src_arg]
+                    offset += arg_sizes[src_arg]
+                offset = 0
+                for dst_arg in dst_args:
+                    output_data[
+                        idx, offset : offset + arg_sizes[dst_arg]
+                    ] = micro_sims_output[lid][dst_arg]
+                    offset += arg_sizes[dst_arg]
+            input_data_inactive = np.zeros((n_points_inactive, src_size))
+            for idx, lid in enumerate(inactive_lids):
+                offset = 0
+                for src_arg in src_args:
+                    input_data_inactive[
+                        idx, offset : offset + arg_sizes[src_arg]
+                    ] = micro_input[lid][src_arg]
+                    offset += arg_sizes[src_arg]
+
+            # use interpolant
+            self._interpolation.configure(self._mappings[m_idx])
+            self._interpolation.set_local_data(
+                input_data, input_data_inactive, output_data
+            )
+            output_data_inactive = self._interpolation.interpolate()
+
+            for idx, lid in enumerate(inactive_lids):
+                offset = 0
+                for dst_arg in dst_args:
+                    micro_sims_output[lid][dst_arg] = output_data_inactive[
+                        idx, offset : offset + arg_sizes[dst_arg]
+                    ]
+                    offset += arg_sizes[dst_arg]
+
     def _get_similarity_measure(
         self, similarity_measure: str
     ) -> Callable[[np.ndarray], np.ndarray]:

micro_manager/adaptivity/global_adaptivity.py

@@ -15,6 +15,7 @@
 from micro_manager.tools.logging_wrapper import Logger
 from micro_manager.micro_simulation import MicroSimulationClass
 from micro_manager.model_manager import ModelManager
+from micro_manager.interpolation import RBF_PU
 
 from micro_manager.tools.p2p import p2p_comm, get_ranks_of_sims
 
@@ -68,6 +69,9 @@ def __init__(
         self._global_ids = global_ids
         self._comm = comm
 
+        self._interpolation = RBF_PU(
+            base_logger, comm, self._rank, self._comm.Get_size()
+        )
         rank_of_sim = get_ranks_of_sims(global_ids, rank, comm, global_number_of_sims)
 
         self._is_sim_on_this_rank = [False] * global_number_of_sims  # DECLARATION
@@ -240,12 +244,16 @@ def get_inactive_sim_global_ids(self) -> np.ndarray:
 
         return np.array(inactive_sim_ids)
 
-    def get_full_field_micro_output(self, micro_output: list) -> list:
+    def get_full_field_micro_output(
+        self, micro_input: list, micro_output: list
+    ) -> list:
         """
         Get the full field micro output from active simulations to inactive simulations.
 
         Parameters
         ----------
+        micro_input : list
+            List of dicts containing the input data for each simulation.
         micro_output : list
             List of dicts having individual output of each simulation. Only the active simulation outputs are entered.
 
@@ -259,7 +267,17 @@ def get_full_field_micro_output(self, micro_output: list) -> list:
         )
 
         micro_sims_output = deepcopy(micro_output)
+        num_active = np.sum(self._is_sim_active)
+        if num_active == self._is_sim_active.shape[0]:
+            self._precice_participant.stop_last_profiling_section()
+            return micro_sims_output
+
         self._communicate_micro_output(micro_sims_output)
+        if num_active <= self._interp_min:
+            self._precice_participant.stop_last_profiling_section()
+            return micro_sims_output
+
+        self._interpolate_output(micro_input, micro_sims_output)
 
         self._precice_participant.stop_last_profiling_section()
 

micro_manager/adaptivity/local_adaptivity.py

@@ -12,6 +12,7 @@
 from micro_manager.micro_simulation import MicroSimulationClass
 from micro_manager.tools.logging_wrapper import Logger
 from micro_manager.model_manager import ModelManager
+from micro_manager.interpolation import RBF_PU
 
 
 class LocalAdaptivityCalculator(AdaptivityCalculator):
@@ -49,6 +50,12 @@ def __init__(
             configurator, num_sims, micro_problem_cls, model_manager, base_logger, rank
         )
         self._comm = comm
+        self._interpolation = RBF_PU(
+            base_logger,
+            MPI.COMM_SELF,
+            MPI.COMM_SELF.Get_rank(),
+            MPI.COMM_SELF.Get_size(),
+        )
 
         # similarity_dists: 2D array having similarity distances between each micro simulation pair
         # This matrix is modified in place via the function update_similarity_dists
@@ -146,12 +153,16 @@ def get_inactive_sim_global_ids(self) -> np.ndarray:
         inactive_sim_ids = self.get_inactive_sim_local_ids()
         return inactive_sim_ids
 
-    def get_full_field_micro_output(self, micro_output: list) -> list:
+    def get_full_field_micro_output(
+        self, micro_input: list, micro_output: list
+    ) -> list:
         """
         Get the full field micro output from active simulations to inactive simulations.
 
         Parameters
         ----------
+        micro_input : list
+            List of dicts containing the input data for each simulation.
         micro_output : list
             List of dicts having individual output of each simulation. Only the active simulation outputs are entered.
 
@@ -168,6 +179,7 @@ def get_full_field_micro_output(self, micro_output: list) -> list:
             micro_sims_output[inactive_id] = deepcopy(
                 micro_sims_output[self._sim_is_associated_to[inactive_id]]
             )
+        self._interpolate_output(micro_input, micro_sims_output)
 
         return micro_sims_output