From 381fed0733fbb6322f6a4164ec628700ad5dc478 Mon Sep 17 00:00:00 2001 From: Ishaan Desai Date: Thu, 12 Dec 2024 17:34:06 +0100 Subject: [PATCH] Refactoring adaptivity data handling to make it more cleaner and help with logging --- micro_manager/adaptivity/adaptivity.py | 3 + micro_manager/adaptivity/global_adaptivity.py | 71 ++++--- micro_manager/adaptivity/local_adaptivity.py | 53 +++-- micro_manager/micro_manager.py | 186 +++++++----------- 4 files changed, 162 insertions(+), 151 deletions(-) diff --git a/micro_manager/adaptivity/adaptivity.py b/micro_manager/adaptivity/adaptivity.py index b1240505..68320a9f 100644 --- a/micro_manager/adaptivity/adaptivity.py +++ b/micro_manager/adaptivity/adaptivity.py @@ -35,6 +35,9 @@ def __init__(self, configurator, logger) -> None: configurator.get_adaptivity_similarity_measure() ) + self._active_sim_ids = np.array([], dtype=int) + self._inactive_sim_ids = np.array([], dtype=int) + def _get_similarity_dists( self, dt: float, similarity_dists: np.ndarray, data: dict ) -> np.ndarray: diff --git a/micro_manager/adaptivity/global_adaptivity.py b/micro_manager/adaptivity/global_adaptivity.py index e30b29b4..d12b2ba4 100644 --- a/micro_manager/adaptivity/global_adaptivity.py +++ b/micro_manager/adaptivity/global_adaptivity.py @@ -70,9 +70,7 @@ def compute_adaptivity( self, dt: float, micro_sims: list, - similarity_dists_nm1: np.ndarray, - is_sim_active_nm1: np.ndarray, - sim_is_associated_to_nm1: np.ndarray, + adaptivity_data_nm1: list, data_for_adaptivity: dict, ) -> tuple: """ @@ -84,21 +82,15 @@ def compute_adaptivity( Current time step of the macro-micro coupled problem micro_sims : list List of objects of class MicroProblem, which are the micro simulations - similarity_dists_nm1 : numpy array - 2D array having similarity distances between each micro simulation pair - is_sim_active_nm1 : numpy array - 1D array having state (active or inactive) of each micro simulation on this rank - sim_is_associated_to_nm1 : numpy array - 1D array with values of associated simulations of inactive simulations. Active simulations have None + adaptivity_data_nm1 : list + List of numpy arrays: similarity_dists (2D array having similarity distances between each micro simulation pair), is_sim_active (1D array having state (active or inactive) of each micro simulation), sim_is_associated_to (1D array with values of associated simulations of inactive simulations. Active simulations have None) data_for_adaptivity : dict Dictionary with keys as names of data to be used in the similarity calculation, and values as the respective data for the micro simulations Results ------- - similarity_dists : numpy array - 2D array having similarity distances between each micro simulation pair - is_sim_active : numpy array - 1D array having state (active or inactive) of each micro simulation + list + List of numpy arrays: similarity_dists (2D array having similarity distances between each micro simulation pair), is_sim_active (1D array having state (active or inactive) of each micro simulation), sim_is_associated_to (1D array with values of associated simulations of inactive simulations. Active simulations have None) """ for name in data_for_adaptivity.keys(): if name not in self._adaptivity_data_names: @@ -115,18 +107,27 @@ def compute_adaptivity( global_data_for_adaptivity[name] = np.concatenate((data_as_list[:]), axis=0) similarity_dists = self._get_similarity_dists( - dt, similarity_dists_nm1, global_data_for_adaptivity + dt, adaptivity_data_nm1[0], global_data_for_adaptivity ) - is_sim_active = self._update_active_sims(similarity_dists, is_sim_active_nm1) + is_sim_active = self._update_active_sims( + similarity_dists, adaptivity_data_nm1[1] + ) is_sim_active, sim_is_associated_to = self._update_inactive_sims( - similarity_dists, is_sim_active, sim_is_associated_to_nm1, micro_sims + similarity_dists, is_sim_active, adaptivity_data_nm1[2], micro_sims ) sim_is_associated_to = self._associate_inactive_to_active( similarity_dists, is_sim_active, sim_is_associated_to ) + self._active_sim_ids = np.where( + is_sim_active[self._global_ids[0] : self._global_ids[-1] + 1] + )[0] + self._inactive_sim_ids = np.where( + is_sim_active[self._global_ids[0] : self._global_ids[-1] + 1] == False + )[0] + self._logger.info( "{} active simulations, {} inactive simulations".format( np.count_nonzero( @@ -139,12 +140,11 @@ def compute_adaptivity( ) ) - return similarity_dists, is_sim_active, sim_is_associated_to + return [similarity_dists, is_sim_active, sim_is_associated_to] def communicate_micro_output( self, - is_sim_active: np.ndarray, - sim_is_associated_to: np.ndarray, + adaptivity_data: list, micro_output: list, ) -> None: """ @@ -155,13 +155,16 @@ def communicate_micro_output( ---------- micro_sims : list List of objects of class MicroProblem, which are the micro simulations - is_sim_active : numpy array - 1D array having state (active or inactive) of each micro simulation on this rank - sim_is_associated_to : numpy array - 1D array with values of associated simulations of inactive simulations. Active simulations have -2 + adaptivity_data : list + List of numpy arrays: + is_sim_active (1D array having state (active or inactive) of each micro simulation) + sim_is_associated_to (1D array with values of associated simulations of inactive simulations. Active simulations have None) micro_output : list List of dicts having individual output of each simulation. Only the active simulation outputs are entered. """ + is_sim_active = adaptivity_data[0] + sim_is_associated_to = adaptivity_data[1] + inactive_local_ids = np.where( is_sim_active[self._global_ids[0] : self._global_ids[-1] + 1] == False )[0] @@ -198,6 +201,28 @@ def communicate_micro_output( for local_id in active_to_inactive_map[assoc_active_ids[count]]: micro_output[local_id] = deepcopy(output) + def get_active_sim_ids(self) -> np.ndarray: + """ + Get the ids of active simulations. + + Returns + ------- + numpy array + 1D array of active simulation ids + """ + return self._active_sim_ids + + def get_inactive_sim_ids(self) -> np.ndarray: + """ + Get the ids of inactive simulations. + + Returns + ------- + numpy array + 1D array of inactive simulation ids + """ + return self._inactive_sim_ids + def _update_inactive_sims( self, similarity_dists: np.ndarray, diff --git a/micro_manager/adaptivity/local_adaptivity.py b/micro_manager/adaptivity/local_adaptivity.py index a2defb8c..2aa7f24c 100644 --- a/micro_manager/adaptivity/local_adaptivity.py +++ b/micro_manager/adaptivity/local_adaptivity.py @@ -26,9 +26,7 @@ def compute_adaptivity( self, dt, micro_sims, - similarity_dists_nm1: np.ndarray, - is_sim_active_nm1: np.ndarray, - sim_is_associated_to_nm1: np.ndarray, + adaptivity_data_nm1: list, data_for_adaptivity: dict, ) -> tuple: """ @@ -40,22 +38,16 @@ def compute_adaptivity( Current time step micro_sims : list List containing simulation objects - similarity_dists_nm1 : numpy array - 2D array having similarity distances between each micro simulation pair. - is_sim_active_nm1 : numpy array - 1D array having True if sim is active, False if sim is inactive. - sim_is_associated_to_nm1 : numpy array - 1D array with values of associated simulations of inactive simulations. Active simulations have None. + adaptivity_data_nm1 : list + List of numpy arrays: similarity_dists (2D array having similarity distances between each micro simulation pair), is_sim_active (1D array having state (active or inactive) of each micro simulation), sim_is_associated_to (1D array with values of associated simulations of inactive simulations. Active simulations have None) data_for_adaptivity : dict A dictionary containing the names of the data to be used in adaptivity as keys and information on whether the data are scalar or vector as values. Returns ------- - similarity_dists : numpy array - 2D array having similarity distances between each micro simulation pair. - is_sim_active : numpy array - 1D array, True is sim is active, False if sim is inactive. + list + List of numpy arrays: similarity_dists (2D array having similarity distances between each micro simulation pair), is_sim_active (1D array having state (active or inactive) of each micro simulation), sim_is_associated_to (1D array with values of associated simulations of inactive simulations. Active simulations have None) """ for name in data_for_adaptivity.keys(): if name not in self._adaptivity_data_names: @@ -66,14 +58,16 @@ def compute_adaptivity( ) similarity_dists = self._get_similarity_dists( - dt, similarity_dists_nm1, data_for_adaptivity + dt, adaptivity_data_nm1[0], data_for_adaptivity ) # Operation done globally if global adaptivity is chosen - is_sim_active = self._update_active_sims(similarity_dists, is_sim_active_nm1) + is_sim_active = self._update_active_sims( + similarity_dists, adaptivity_data_nm1[1] + ) is_sim_active, sim_is_associated_to = self._update_inactive_sims( - similarity_dists, is_sim_active, sim_is_associated_to_nm1, micro_sims + similarity_dists, is_sim_active, adaptivity_data_nm1[2], micro_sims ) sim_is_associated_to = self._associate_inactive_to_active( @@ -87,7 +81,32 @@ def compute_adaptivity( ) ) - return similarity_dists, is_sim_active, sim_is_associated_to + self._active_sim_ids = np.where(is_sim_active)[0] + self._inactive_sim_ids = np.where(is_sim_active == False)[0] + + return [similarity_dists, is_sim_active, sim_is_associated_to] + + def get_active_sim_ids(self) -> np.ndarray: + """ + Get the ids of active simulations. + + Returns + ------- + numpy array + 1D array of active simulation ids + """ + return self._active_sim_ids + + def get_inactive_sim_ids(self) -> np.ndarray: + """ + Get the ids of inactive simulations. + + Returns + ------- + numpy array + 1D array of inactive simulation ids + """ + return self._inactive_sim_ids def _update_inactive_sims( self, diff --git a/micro_manager/micro_manager.py b/micro_manager/micro_manager.py index 62ad3d6f..34b34fda 100644 --- a/micro_manager/micro_manager.py +++ b/micro_manager/micro_manager.py @@ -20,6 +20,7 @@ from copy import deepcopy from typing import Dict from warnings import warn +from typing import Callable import numpy as np import precice @@ -118,41 +119,42 @@ def solve(self) -> None: """ t, n = 0, 0 t_checkpoint, n_checkpoint = 0, 0 - similarity_dists_cp = None - is_sim_active_cp = None - sim_is_associated_to_cp = None + adaptivity_data_cp: list = [] sim_states_cp = [None] * self._local_number_of_sims + micro_sim_solve = self._get_solve_variant() + dt = min(self._participant.get_max_time_step_size(), self._micro_dt) + # adaptivity_data is a of numpy arrays: similarity_dists (2D array having similarity distances between each micro simulation pair), is_sim_active (1D array having state (active or inactive) of each micro simulation), sim_is_associated_to (1D array with values of associated simulations of inactive simulations. Active simulations have None) + adaptivity_data: list = [] + if self._is_adaptivity_on: - similarity_dists = np.zeros( - ( - self._number_of_sims_for_adaptivity, - self._number_of_sims_for_adaptivity, + adaptivity_data.append( + np.zeros( + ( + self._number_of_sims_for_adaptivity, + self._number_of_sims_for_adaptivity, + ) ) ) # Start adaptivity calculation with all sims active - is_sim_active = np.array([True] * self._number_of_sims_for_adaptivity) + adaptivity_data.append( + np.array([True] * self._number_of_sims_for_adaptivity) + ) # Active sims do not have an associated sim - sim_is_associated_to = np.full( - (self._number_of_sims_for_adaptivity), -2, dtype=np.intc + adaptivity_data.append( + np.full((self._number_of_sims_for_adaptivity), -2, dtype=np.intc) ) # If micro simulations have been initialized, compute adaptivity before starting the coupling if self._micro_sims_init: - ( - similarity_dists, - is_sim_active, - sim_is_associated_to, - ) = self._adaptivity_controller.compute_adaptivity( + adaptivity_data = self._adaptivity_controller.compute_adaptivity( dt, self._micro_sims, - similarity_dists, - is_sim_active, - sim_is_associated_to, + adaptivity_data, self._data_for_adaptivity, ) @@ -169,77 +171,29 @@ def solve(self) -> None: if self._is_adaptivity_on: if not self._adaptivity_in_every_implicit_step: - ( - similarity_dists, - is_sim_active, - sim_is_associated_to, - ) = self._adaptivity_controller.compute_adaptivity( - dt, - self._micro_sims, - similarity_dists, - is_sim_active, - sim_is_associated_to, - self._data_for_adaptivity, + adaptivity_data = ( + self._adaptivity_controller.compute_adaptivity( + dt, + self._micro_sims, + adaptivity_data, + self._data_for_adaptivity, + ) ) # Only checkpoint the adaptivity configuration if adaptivity is computed # once in every time window - similarity_dists_cp = np.copy(similarity_dists) - is_sim_active_cp = np.copy(is_sim_active) - sim_is_associated_to_cp = np.copy(sim_is_associated_to) - - if self._adaptivity_type == "local": - active_sim_ids = np.where(is_sim_active)[0] - elif self._adaptivity_type == "global": - active_sim_ids = np.where( - is_sim_active[ - self._global_ids_of_local_sims[ - 0 - ] : self._global_ids_of_local_sims[-1] - + 1 - ] - )[0] + adaptivity_data_cp = deepcopy(adaptivity_data) + + active_sim_ids = ( + self._adaptivity_controller.get_active_sim_ids() + ) for active_id in active_sim_ids: self._micro_sims_active_steps[active_id] += 1 micro_sims_input = self._read_data_from_precice(dt) - if self._is_adaptivity_on: - if self._adaptivity_in_every_implicit_step: - ( - similarity_dists, - is_sim_active, - sim_is_associated_to, - ) = self._adaptivity_controller.compute_adaptivity( - dt, - self._micro_sims, - similarity_dists, - is_sim_active, - sim_is_associated_to, - self._data_for_adaptivity, - ) - - if self._adaptivity_type == "local": - active_sim_ids = np.where(is_sim_active)[0] - elif self._adaptivity_type == "global": - active_sim_ids = np.where( - is_sim_active[ - self._global_ids_of_local_sims[ - 0 - ] : self._global_ids_of_local_sims[-1] - + 1 - ] - )[0] - - for active_id in active_sim_ids: - self._micro_sims_active_steps[active_id] += 1 - - micro_sims_output = self._solve_micro_simulations_with_adaptivity( - micro_sims_input, is_sim_active, sim_is_associated_to, dt - ) - else: - micro_sims_output = self._solve_micro_simulations(micro_sims_input, dt) + micro_sims_output = micro_sim_solve(micro_sims_input, dt, adaptivity_data) # Check if more than a certain percentage of the micro simulations have crashed and terminate if threshold is exceeded if self._interpolate_crashed_sims: @@ -282,9 +236,7 @@ def solve(self) -> None: # If adaptivity is computed only once per time window, the states of sims need to be reset too if self._is_adaptivity_on: if not self._adaptivity_in_every_implicit_step: - similarity_dists = np.copy(similarity_dists_cp) - is_sim_active = np.copy(is_sim_active_cp) - sim_is_associated_to = np.copy(sim_is_associated_to_cp) + adaptivity_data = deepcopy(adaptivity_data_cp) if ( self._participant.is_time_window_complete() @@ -643,7 +595,9 @@ def _write_data_to_precice(self, data: list) -> None: self._macro_mesh_name, dname, [], np.array([]) ) - def _solve_micro_simulations(self, micro_sims_input: list, dt: float) -> list: + def _solve_micro_simulations( + self, micro_sims_input: list, dt: float, adaptivity_data=None + ) -> list: """ Solve all micro simulations and assemble the micro simulations outputs in a list of dicts format. @@ -654,6 +608,8 @@ def _solve_micro_simulations(self, micro_sims_input: list, dt: float) -> list: solve a micro simulation. dt : float Time step size. + adaptivity_data : list + Dummy parameter to match the signature of the function with the function _solve_micro_simulations_with_adaptivity. Returns ------- @@ -720,9 +676,8 @@ def _solve_micro_simulations(self, micro_sims_input: list, dt: float) -> list: def _solve_micro_simulations_with_adaptivity( self, micro_sims_input: list, - is_sim_active: np.ndarray, - sim_is_associated_to: np.ndarray, dt: float, + adaptivity_data: list, ) -> list: """ Solve all micro simulations and assemble the micro simulations outputs in a list of dicts format. @@ -732,12 +687,10 @@ def _solve_micro_simulations_with_adaptivity( micro_sims_input : list List of dicts in which keys are names of data and the values are the data which are required inputs to solve a micro simulation. - is_sim_active : numpy array - 1D array having state (active or inactive) of each micro simulation - sim_is_associated_to : numpy array - 1D array with values of associated simulations of inactive simulations. Active simulations have None dt : float Time step size. + adaptivity_data : list + List of numpy arrays: similarity_dists (2D array having similarity distances between each micro simulation pair), is_sim_active (1D array having state (active or inactive) of each micro simulation), sim_is_associated_to (1D array with values of associated simulations of inactive simulations. Active simulations have None) Returns ------- @@ -745,27 +698,22 @@ def _solve_micro_simulations_with_adaptivity( List of dicts in which keys are names of data and the values are the data of the output of the micro simulations. """ - if self._adaptivity_type == "global": - active_sim_ids = np.where( - is_sim_active[ - self._global_ids_of_local_sims[0] : self._global_ids_of_local_sims[ - -1 - ] - + 1 - ] - )[0] - inactive_sim_ids = np.where( - is_sim_active[ - self._global_ids_of_local_sims[0] : self._global_ids_of_local_sims[ - -1 - ] - + 1 - ] - == False - )[0] - elif self._adaptivity_type == "local": - active_sim_ids = np.where(is_sim_active)[0] - inactive_sim_ids = np.where(is_sim_active == False)[0] + if self._adaptivity_in_every_implicit_step: + adaptivity_data = self._adaptivity_controller.compute_adaptivity( + dt, + self._micro_sims, + adaptivity_data, + self._data_for_adaptivity, + ) + + sim_is_associated_to = adaptivity_data[2] + + active_sim_ids = self._adaptivity_controller.get_active_sim_ids() + + for active_id in active_sim_ids: + self._micro_sims_active_steps[active_id] += 1 + + inactive_sim_ids = self._adaptivity_controller.get_inactive_sim_ids() micro_sims_output = [None] * self._local_number_of_sims @@ -834,7 +782,7 @@ def _solve_micro_simulations_with_adaptivity( # For each inactive simulation, copy data from most similar active simulation if self._adaptivity_type == "global": self._adaptivity_controller.communicate_micro_output( - is_sim_active, sim_is_associated_to, micro_sims_output + adaptivity_data[1:3], micro_sims_output ) elif self._adaptivity_type == "local": for inactive_id in inactive_sim_ids: @@ -859,6 +807,22 @@ def _solve_micro_simulations_with_adaptivity( return micro_sims_output + def _get_solve_variant(self) -> Callable[[list, float, list], list]: + """ + Get the solve variant function based on the adaptivity type. + + Returns + ------- + solve_variant : Callable + Solve variant function based on the adaptivity type. + """ + if self._is_adaptivity_on: + solve_variant = self._solve_micro_simulations_with_adaptivity + else: + solve_variant = self._solve_micro_simulations + + return solve_variant + def _interpolate_output_for_crashed_sim( self, micro_sims_input: list,