From 61126718a43d743cc2974d30083127138c544451 Mon Sep 17 00:00:00 2001 From: Geethapranay1 Date: Tue, 31 Mar 2026 22:26:47 +0530 Subject: [PATCH 1/2] Implement local VTU export for load balancing --- CHANGELOG.md | 1 + micro_manager/micro_manager.py | 115 +++++++++++++++++++++++++- micro_manager/tools/vtu_export.py | 132 ++++++++++++++++++++++++++++++ 3 files changed, 244 insertions(+), 4 deletions(-) create mode 100644 micro_manager/tools/vtu_export.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 503f8b50..1d3255c3 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,7 @@ ## latest +- Added native VTU export functionality to support dynamic load balancing [#189](https://github.com/precice/micro-manager/issues/189) - Fixed lazy initialization for ranks without (active) micro simulations [#238](https://github.com/precice/micro-manager/pull/238) - Added coverage testing and simulation interface tests [#225](https://github.com/precice/micro-manager/pull/225) - Added `--test-dependencies` CLI flag to check if all required dependencies are correctly installed, with clear error messages listing missing packages and how to fix them [#221](https://github.com/precice/micro-manager/pull/221) diff --git a/micro_manager/micro_manager.py b/micro_manager/micro_manager.py index 8f1515c6..70650836 100644 --- a/micro_manager/micro_manager.py +++ b/micro_manager/micro_manager.py @@ -33,6 +33,7 @@ from .tasking.connection import spawn_local_workers from .micro_simulation import create_simulation_class, load_backend_class from .tools.logging_wrapper import Logger +from .tools.vtu_export import write_vtu, write_pvtu from .load_balancing import create_load_balancer try: @@ -73,6 +74,48 @@ def __init__(self, config_file: str, log_file: str = "") -> None: else: self._output_dir = os.path.abspath(os.getcwd()) + "/" + self._export_vtu = False + self._export_vtu_dir = "output_vtu" + self._export_vtu_n = 1 + + self._is_load_balancing = ( + self._config.turn_on_load_balancing() and self._is_parallel + ) + + precice_config_file = self._config.get_precice_config_file_name() + + import xml.etree.ElementTree as ET + try: + tree = ET.parse(precice_config_file) + root = tree.getroot() + tag_found = False + for participant in root.findall('participant'): + if participant.get('name') == 'Micro-Manager': + for child in list(participant): + if child.tag.endswith('export:vtu'): + self._export_vtu = True + if child.get('directory'): + self._export_vtu_dir = child.get('directory') + if child.get('every-n-time-windows'): + self._export_vtu_n = int(child.get('every-n-time-windows')) + participant.remove(child) + tag_found = True + + if tag_found and self._is_load_balancing: + modified_xml = precice_config_file + f".lb-modified-rank{self._rank}.xml" + tree.write(modified_xml, xml_declaration=True, encoding="UTF-8") + precice_config_file = modified_xml + self._logger.log_info_rank_zero("Intercepted preCICE export:vtu for custom load-balancing aware VTU export.") + elif tag_found and not self._is_load_balancing: + self._export_vtu = False + except Exception as e: + self._logger.log_warning_rank_zero(f"Failed to parse precice config for VTU interception: {e}") + + if self._export_vtu: + if not os.path.isabs(self._export_vtu_dir): + self._export_vtu_dir = os.path.join(self._output_dir, self._export_vtu_dir) + os.makedirs(self._export_vtu_dir, exist_ok=True) + # Data names of data to output to the snapshot database self._write_data_names = self._config.get_write_data_names() @@ -138,7 +181,7 @@ def __init__(self, config_file: str, log_file: str = "") -> None: # Define the preCICE Participant self._participant = precice.Participant( "Micro-Manager", - self._config.get_precice_config_file_name(), + precice_config_file, self._rank, self._size, ) @@ -153,9 +196,6 @@ def __init__(self, config_file: str, log_file: str = "") -> None: if self._is_model_adaptivity_on: self.state_loader = lambda sim: sim.attachments self.state_setter = lambda sim, state: sim.attachments.update(state) - self._is_load_balancing = ( - self._config.turn_on_load_balancing() and self._is_parallel - ) self._load_balancing_n = self._config.get_load_balancing_n() self.load_balancing = None @@ -335,6 +375,9 @@ def solve(self) -> None: if sim: sim.output() + if self._export_vtu and (self._n % self._export_vtu_n == 0): + self._export_vtu_data(micro_sims_input, micro_sims_output) + if ( self._is_adaptivity_on and self._adaptivity_output_type @@ -420,6 +463,70 @@ def solve(self) -> None: self._conn.close() self._participant.finalize() + def _export_vtu_data(self, micro_sims_input: list, micro_sims_output: list) -> None: + """ + Exports the current micro simulation data locally owned by this rank to VTU files. + """ + data_local = {} + coords_local = np.empty((0, len(self._macro_bounds) // 2)) + + if not self._is_rank_empty: + coords_local = self._mesh_vertex_coords[self._global_ids_of_local_sims] + + all_fields = list(self._read_data_names) + list(self._write_data_names) + all_fields = list(dict.fromkeys(all_fields)) + + for field in all_fields: + field_data = [] + for i in range(self._local_number_of_sims): + if micro_sims_input and i < len(micro_sims_input) and field in micro_sims_input[i]: + field_data.append(micro_sims_input[i][field]) + elif micro_sims_output and i < len(micro_sims_output) and field in micro_sims_output[i]: + field_data.append(micro_sims_output[i][field]) + else: + field_data.append(0.0) + + if len(field_data) > 0: + val = np.asarray(field_data[0]) + shape = val.shape + arr = np.zeros((self._local_number_of_sims,) + shape, dtype=np.float64) + for i, v in enumerate(field_data): + arr[i] = v + data_local[field] = arr + + filename = os.path.join( + self._export_vtu_dir, + f"Micro-Manager_Macro-Mesh_{self._n}_rank{self._rank}.vtu" + ) + write_vtu(filename, coords_local, data_local) + + if self._is_parallel: + pvtu_filename = os.path.join( + self._export_vtu_dir, + f"Micro-Manager_Macro-Mesh_{self._n}.pvtu" + ) + local_keys = {} + for k, v in data_local.items(): + val_np = np.asarray(v) + comp = 1 if val_np.ndim == 1 else val_np.shape[1] + if comp == 2: + comp = 3 + local_keys[k] = comp + + all_keys = self._comm.gather(local_keys, root=0) + + if self._rank == 0: + global_keys = {} + for r_keys in all_keys: + if r_keys: + global_keys.update(r_keys) + + source_files = [ + f"Micro-Manager_Macro-Mesh_{self._n}_rank{r}.vtu" + for r in range(self._size) + ] + write_pvtu(pvtu_filename, source_files, global_keys) + def initialize(self) -> None: """ Initialize the Micro Manager by performing the following tasks: diff --git a/micro_manager/tools/vtu_export.py b/micro_manager/tools/vtu_export.py new file mode 100644 index 00000000..33ba174b --- /dev/null +++ b/micro_manager/tools/vtu_export.py @@ -0,0 +1,132 @@ +import os +import xml.etree.ElementTree as ET +import numpy as np + + +def write_vtu(filename: str, coords: np.ndarray, data: dict) -> None: + """ + Writes a VTU file (UnstructuredGrid with VTK_VERTEX cells) containing + points (coords) and associated scalar/vector point data. + + Parameters + ---------- + filename : str + Output file path (e.g., "output.vtu"). + coords : numpy array + 2D or 3D numpy array of shape (N, 2) or (N, 3). + data : dict + Dictionary of point data fields. Keys are names, values are scalar (N,) or vector (N, d) arrays. + """ + n_points = coords.shape[0] + + if n_points == 0: + return + + dim = coords.shape[1] + + if dim == 2: + coords_3d = np.zeros((n_points, 3), dtype=np.float64) + coords_3d[:, :2] = coords + else: + coords_3d = np.asarray(coords, dtype=np.float64) + + vtk_file = ET.Element( + "VTKFile", type="UnstructuredGrid", version="0.1", byte_order="LittleEndian" + ) + unstr_grid = ET.SubElement(vtk_file, "UnstructuredGrid") + piece = ET.SubElement( + unstr_grid, "Piece", NumberOfPoints=str(n_points), NumberOfCells=str(n_points) + ) + + points = ET.SubElement(piece, "Points") + pts_arr = ET.SubElement( + points, + "DataArray", + type="Float64", + NumberOfComponents="3", + format="ascii", + ) + pts_arr.text = " ".join(map(str, coords_3d.ravel())) + + cells = ET.SubElement(piece, "Cells") + + conn_arr = ET.SubElement( + cells, "DataArray", type="Int32", Name="connectivity", format="ascii" + ) + conn_arr.text = " ".join(map(str, range(n_points))) + + off_arr = ET.SubElement( + cells, "DataArray", type="Int32", Name="offsets", format="ascii" + ) + off_arr.text = " ".join(map(str, range(1, n_points + 1))) + + type_arr = ET.SubElement( + cells, "DataArray", type="UInt8", Name="types", format="ascii" + ) + type_arr.text = " ".join(["1"] * n_points) + + point_data = ET.SubElement(piece, "PointData") + for key, val in data.items(): + val_np = np.asarray(val, dtype=np.float64) + if val_np.ndim == 1: + n_comp = 1 + else: + n_comp = val_np.shape[1] + if n_comp == 2: + val_3d = np.zeros((n_points, 3), dtype=np.float64) + val_3d[:, :2] = val_np + val_np = val_3d + n_comp = 3 + + data_arr = ET.SubElement( + point_data, + "DataArray", + type="Float64", + Name=key, + NumberOfComponents=str(n_comp), + format="ascii", + ) + data_arr.text = " ".join(map(str, val_np.ravel())) + + tree = ET.ElementTree(vtk_file) + os.makedirs(os.path.dirname(os.path.abspath(filename)), exist_ok=True) + tree.write(filename, xml_declaration=True, encoding="utf-8") + + +def write_pvtu(filename: str, source_files: list, data_keys: dict) -> None: + """ + Writes a Parallel VTU (.pvtu) file referencing the multiple subset .vtu files. + + Parameters + ---------- + filename : str + Output file path for the PVTU. + source_files : list + List of VTU file names that this PVTU references. + data_keys : dict + Dictionary mapping data array names to their number of components. + """ + vtk_file = ET.Element( + "VTKFile", type="PUnstructuredGrid", version="0.1", byte_order="LittleEndian" + ) + p_grid = ET.SubElement(vtk_file, "PUnstructuredGrid", GhostLevel="0") + + p_points = ET.SubElement(p_grid, "PPoints") + ET.SubElement(p_points, "PDataArray", type="Float64", NumberOfComponents="3") + + p_point_data = ET.SubElement(p_grid, "PPointData") + for key, num_comp in data_keys.items(): + ET.SubElement( + p_point_data, + "PDataArray", + type="Float64", + Name=key, + NumberOfComponents=str(num_comp), + ) + + for sf in source_files: + ET.SubElement(p_grid, "Piece", Source=os.path.basename(sf)) + + tree = ET.ElementTree(vtk_file) + os.makedirs(os.path.dirname(os.path.abspath(filename)), exist_ok=True) + tree.write(filename, xml_declaration=True, encoding="utf-8") From f55d4738b99f0db642bcccf084af895d126e1dea Mon Sep 17 00:00:00 2001 From: Geethapranay1 Date: Tue, 5 May 2026 08:13:02 +0530 Subject: [PATCH 2/2] Fix load balancing VTU export --- micro_manager/micro_manager.py | 294 ++++++++++++++++++++---------- micro_manager/tools/vtu_export.py | 123 ++++++++----- tests/unit/test_vtu_export.py | 81 ++++++++ 3 files changed, 356 insertions(+), 142 deletions(-) create mode 100644 tests/unit/test_vtu_export.py diff --git a/micro_manager/micro_manager.py b/micro_manager/micro_manager.py index 70650836..ce3d8960 100644 --- a/micro_manager/micro_manager.py +++ b/micro_manager/micro_manager.py @@ -12,29 +12,28 @@ Detailed documentation: https://precice.org/tooling-micro-manager-overview.html """ -import os -import sys -from typing import Callable -import numpy as np -from psutil import Process import csv +import os +import re import subprocess +import sys from functools import partial +from typing import Callable +import numpy as np import precice +from psutil import Process -from .model_manager import ModelManager -from .micro_manager_base import MicroManager - -from .adaptivity.model_adaptivity import ModelAdaptivity from .adaptivity.adaptivity_selection import create_adaptivity_calculator - +from .adaptivity.model_adaptivity import ModelAdaptivity from .domain_decomposition import DomainDecomposer -from .tasking.connection import spawn_local_workers +from .load_balancing import create_load_balancer +from .micro_manager_base import MicroManager from .micro_simulation import create_simulation_class, load_backend_class +from .model_manager import ModelManager +from .tasking.connection import spawn_local_workers from .tools.logging_wrapper import Logger -from .tools.vtu_export import write_vtu, write_pvtu -from .load_balancing import create_load_balancer +from .tools.vtu_export import write_pvtu, write_vtu try: from .interpolation import Interpolation @@ -75,7 +74,7 @@ def __init__(self, config_file: str, log_file: str = "") -> None: self._output_dir = os.path.abspath(os.getcwd()) + "/" self._export_vtu = False - self._export_vtu_dir = "output_vtu" + self._export_vtu_dir = "." self._export_vtu_n = 1 self._is_load_balancing = ( @@ -83,37 +82,16 @@ def __init__(self, config_file: str, log_file: str = "") -> None: ) precice_config_file = self._config.get_precice_config_file_name() - - import xml.etree.ElementTree as ET - try: - tree = ET.parse(precice_config_file) - root = tree.getroot() - tag_found = False - for participant in root.findall('participant'): - if participant.get('name') == 'Micro-Manager': - for child in list(participant): - if child.tag.endswith('export:vtu'): - self._export_vtu = True - if child.get('directory'): - self._export_vtu_dir = child.get('directory') - if child.get('every-n-time-windows'): - self._export_vtu_n = int(child.get('every-n-time-windows')) - participant.remove(child) - tag_found = True - - if tag_found and self._is_load_balancing: - modified_xml = precice_config_file + f".lb-modified-rank{self._rank}.xml" - tree.write(modified_xml, xml_declaration=True, encoding="UTF-8") - precice_config_file = modified_xml - self._logger.log_info_rank_zero("Intercepted preCICE export:vtu for custom load-balancing aware VTU export.") - elif tag_found and not self._is_load_balancing: - self._export_vtu = False - except Exception as e: - self._logger.log_warning_rank_zero(f"Failed to parse precice config for VTU interception: {e}") + if precice_config_file is None: + raise ValueError("No preCICE configuration file has been configured.") + if self._is_load_balancing: + precice_config_file = self._configure_load_balancing_vtu_export( + precice_config_file + ) if self._export_vtu: if not os.path.isabs(self._export_vtu_dir): - self._export_vtu_dir = os.path.join(self._output_dir, self._export_vtu_dir) + self._export_vtu_dir = os.path.abspath(self._export_vtu_dir) os.makedirs(self._export_vtu_dir, exist_ok=True) # Data names of data to output to the snapshot database @@ -232,7 +210,6 @@ def solve(self) -> None: self._adaptivity_controller.log_metrics(self._n) while self._participant.is_coupling_ongoing(): - dt = min(self._participant.get_max_time_step_size(), self._micro_dt) if self._is_adaptivity_on: @@ -463,69 +440,188 @@ def solve(self) -> None: self._conn.close() self._participant.finalize() + def _configure_load_balancing_vtu_export(self, precice_config_file: str) -> str: + """ + Configure Micro Manager's native VTU export for load balancing + + preCICE XML files commonly use prefixed tags such as ``export:vtu`` + without declaring XML namespaces. ``xml.etree.ElementTree`` therefore + cannot parse these files reliably. For this specific interception, a + text-based replacement is sufficient and keeps all other preCICE tags + untouched. + """ + try: + with open(precice_config_file, encoding="utf-8") as file: + config_text = file.read() + except OSError as error: + self._logger.log_warning_rank_zero( + f"Failed to read preCICE config for VTU interception: {error}" + ) + return precice_config_file + + participant_pattern = re.compile( + r"(]*\bname\s*=\s*(['\"])Micro-Manager\2[^>]*>)" + r"(?P.*?)" + r"()", + re.DOTALL, + ) + participant_match = participant_pattern.search(config_text) + if participant_match is None: + return precice_config_file + + export_pattern = re.compile( + r"\s*[^>]*)/>\s*" + r"|\s*[^>]*)>.*?\s*", + re.DOTALL, + ) + + body = participant_match.group("body") + export_matches = list(export_pattern.finditer(body)) + if not export_matches: + return precice_config_file + + attrs = self._parse_xml_attributes( + export_matches[0].group("self_closing_attrs") + or export_matches[0].group("paired_attrs") + or "" + ) + + self._export_vtu = True + if attrs.get("directory"): + self._export_vtu_dir = attrs["directory"] + if attrs.get("every-n-time-windows"): + self._export_vtu_n = int(attrs["every-n-time-windows"]) + if self._export_vtu_n == -1: + self._export_vtu = False + elif self._export_vtu_n <= 0: + raise ValueError( + "export:vtu every-n-time-windows must be positive or -1." + ) + if attrs.get("every-iteration", "false").lower() == "true": + self._logger.log_warning_rank_zero( + "Micro Manager's load-balancing aware VTU export supports exports " + "at completed time windows. The preCICE export:vtu option " + 'every-iteration="true" is ignored.' + ) + + stripped_body = export_pattern.sub("\n", body) + stripped_config = ( + config_text[: participant_match.start("body")] + + stripped_body + + config_text[participant_match.end("body") :] + ) + + modified_xml = precice_config_file + f".lb-modified-rank{self._rank}.xml" + with open(modified_xml, "w", encoding="utf-8") as file: + file.write(stripped_config) + + self._logger.log_info_rank_zero( + "Intercepted preCICE export:vtu for custom load-balancing aware VTU export." + ) + return modified_xml + + @staticmethod + def _parse_xml_attributes(attribute_text: str) -> dict[str, str]: + """Parse XML-style attributes from a single tag.""" + attributes = {} + attr_pattern = re.compile(r"([\w:.-]+)\s*=\s*(['\"])(.*?)\2", re.DOTALL) + for match in attr_pattern.finditer(attribute_text): + attributes[match.group(1)] = match.group(3) + return attributes + + @staticmethod + def _get_export_field_value( + field: str, index: int, micro_sims_input: list, micro_sims_output: list + ): + if micro_sims_input and index < len(micro_sims_input): + if field in micro_sims_input[index]: + return micro_sims_input[index][field] + if micro_sims_output and index < len(micro_sims_output): + if field in micro_sims_output[index]: + return micro_sims_output[index][field] + return None + + @staticmethod + def _get_vtu_components(values: np.ndarray) -> int: + if values.ndim == 1: + return 1 + if values.ndim == 2: + return 3 if values.shape[1] == 2 else values.shape[1] + raise ValueError("VTU export supports only scalar or vector point data.") + def _export_vtu_data(self, micro_sims_input: list, micro_sims_output: list) -> None: """ Exports the current micro simulation data locally owned by this rank to VTU files. """ data_local = {} - coords_local = np.empty((0, len(self._macro_bounds) // 2)) + dimension = len(self._macro_bounds) // 2 + coords_local = np.empty((0, dimension)) if not self._is_rank_empty: coords_local = self._mesh_vertex_coords[self._global_ids_of_local_sims] - + all_fields = list(self._read_data_names) + list(self._write_data_names) all_fields = list(dict.fromkeys(all_fields)) for field in all_fields: field_data = [] + template = None for i in range(self._local_number_of_sims): - if micro_sims_input and i < len(micro_sims_input) and field in micro_sims_input[i]: - field_data.append(micro_sims_input[i][field]) - elif micro_sims_output and i < len(micro_sims_output) and field in micro_sims_output[i]: - field_data.append(micro_sims_output[i][field]) - else: - field_data.append(0.0) + value = self._get_export_field_value( + field, i, micro_sims_input, micro_sims_output + ) + if value is not None and template is None: + template = np.asarray(value, dtype=np.float64) + field_data.append(value) + + if template is None: + continue + + arr = np.zeros( + (self._local_number_of_sims,) + template.shape, dtype=np.float64 + ) + for i, value in enumerate(field_data): + if value is not None: + arr[i] = value + data_local[field] = arr - if len(field_data) > 0: - val = np.asarray(field_data[0]) - shape = val.shape - arr = np.zeros((self._local_number_of_sims,) + shape, dtype=np.float64) - for i, v in enumerate(field_data): - arr[i] = v - data_local[field] = arr + local_schema = { + key: self._get_vtu_components(np.asarray(values)) + for key, values in data_local.items() + } + + if self._is_parallel: + all_schemas = self._comm.allgather(local_schema) + data_schema = {} + for schema in all_schemas: + for key, components in schema.items(): + if key in data_schema and data_schema[key] != components: + raise ValueError( + f"Inconsistent component count for VTU data '{key}'." + ) + data_schema[key] = components + else: + data_schema = local_schema + mesh_file_name = self._macro_mesh_name.replace(os.sep, "_") filename = os.path.join( self._export_vtu_dir, - f"Micro-Manager_Macro-Mesh_{self._n}_rank{self._rank}.vtu" + f"Micro-Manager_{mesh_file_name}_{self._n}_rank{self._rank}.vtu", ) - write_vtu(filename, coords_local, data_local) + write_vtu(filename, coords_local, data_local, data_schema) if self._is_parallel: - pvtu_filename = os.path.join( - self._export_vtu_dir, - f"Micro-Manager_Macro-Mesh_{self._n}.pvtu" - ) - local_keys = {} - for k, v in data_local.items(): - val_np = np.asarray(v) - comp = 1 if val_np.ndim == 1 else val_np.shape[1] - if comp == 2: - comp = 3 - local_keys[k] = comp - - all_keys = self._comm.gather(local_keys, root=0) - + self._comm.Barrier() if self._rank == 0: - global_keys = {} - for r_keys in all_keys: - if r_keys: - global_keys.update(r_keys) - + pvtu_filename = os.path.join( + self._export_vtu_dir, + f"Micro-Manager_{mesh_file_name}_{self._n}.pvtu", + ) source_files = [ - f"Micro-Manager_Macro-Mesh_{self._n}_rank{r}.vtu" - for r in range(self._size) + f"Micro-Manager_{mesh_file_name}_{self._n}_rank{rank}.vtu" + for rank in range(self._size) ] - write_pvtu(pvtu_filename, source_files, global_keys) + write_pvtu(pvtu_filename, source_files, data_schema) def initialize(self) -> None: """ @@ -860,9 +956,9 @@ def initialize(self) -> None: # Save initial data from first micro simulation as we anyway have it for name in initial_micro_output.keys(): if name in self._data_for_adaptivity: - self._data_for_adaptivity[name][ - first_id - ] = initial_micro_output[name] + self._data_for_adaptivity[name][first_id] = ( + initial_micro_output[name] + ) else: raise Exception( "The initialize() method needs to return data which is required for the adaptivity calculation." @@ -875,17 +971,17 @@ def initialize(self) -> None: initial_data[i] ) for name in self._adaptivity_micro_data_names: - self._data_for_adaptivity[name][ - i - ] = initial_micro_output[name] + self._data_for_adaptivity[name][i] = ( + initial_micro_output[name] + ) initial_micro_data[name][i] = initial_micro_output[name] else: for i in micro_sims_to_init: initial_micro_output = self._micro_sims[i].initialize() for name in self._adaptivity_micro_data_names: - self._data_for_adaptivity[name][ - i - ] = initial_micro_output[name] + self._data_for_adaptivity[name][i] = ( + initial_micro_output[name] + ) initial_micro_data[name][i] = initial_micro_output[name] else: self._logger.log_warning_rank_zero( @@ -964,7 +1060,7 @@ def _read_data_from_precice(self, dt) -> list: read_data: dict[str, list] = dict() if self._is_load_balancing: - read_vertex_ids = self._global_ids_of_local_sims + read_vertex_ids = self._mesh_vertex_ids[self._global_ids_of_local_sims] else: read_vertex_ids = self._mesh_vertex_ids @@ -996,7 +1092,7 @@ def _write_data_to_precice(self, data: list) -> None: List of dicts in which keys are names of data and the values are the data to be written to preCICE. """ if self._is_load_balancing: - write_vertex_ids = self._global_ids_of_local_sims + write_vertex_ids = self._mesh_vertex_ids[self._global_ids_of_local_sims] else: write_vertex_ids = self._mesh_vertex_ids @@ -1155,9 +1251,9 @@ def _solve_micro_simulations_with_adaptivity( # Mark the micro sim as active for export micro_sims_output[lid]["Active-State"] = 1 gid = self._global_ids_of_local_sims[lid] - micro_sims_output[lid][ - "Active-Steps" - ] = self._micro_sims_active_steps[gid] + micro_sims_output[lid]["Active-Steps"] = ( + self._micro_sims_active_steps[gid] + ) # If simulation crashes, log the error and keep the output constant at the previous iteration's output except Exception as error_message: @@ -1211,9 +1307,9 @@ def _solve_micro_simulations_with_adaptivity( for inactive_lid in inactive_sim_lids: micro_sims_output[inactive_lid]["Active-State"] = 0 gid = self._global_ids_of_local_sims[inactive_lid] - micro_sims_output[inactive_lid][ - "Active-Steps" - ] = self._micro_sims_active_steps[gid] + micro_sims_output[inactive_lid]["Active-Steps"] = ( + self._micro_sims_active_steps[gid] + ) # Collect micro sim output for adaptivity calculation for i in range(self._local_number_of_sims): diff --git a/micro_manager/tools/vtu_export.py b/micro_manager/tools/vtu_export.py index 33ba174b..037a9ab0 100644 --- a/micro_manager/tools/vtu_export.py +++ b/micro_manager/tools/vtu_export.py @@ -1,34 +1,76 @@ import os import xml.etree.ElementTree as ET +from typing import Dict, Optional + import numpy as np -def write_vtu(filename: str, coords: np.ndarray, data: dict) -> None: - """ - Writes a VTU file (UnstructuredGrid with VTK_VERTEX cells) containing - points (coords) and associated scalar/vector point data. - - Parameters - ---------- - filename : str - Output file path (e.g., "output.vtu"). - coords : numpy array - 2D or 3D numpy array of shape (N, 2) or (N, 3). - data : dict - Dictionary of point data fields. Keys are names, values are scalar (N,) or vector (N, d) arrays. - """ - n_points = coords.shape[0] +def _as_3d_points(coords: np.ndarray) -> np.ndarray: + coords_np = np.asarray(coords, dtype=np.float64) + + if coords_np.ndim != 2 or coords_np.shape[1] not in (2, 3): + raise ValueError("VTU coordinates must have shape (N, 2) or (N, 3).") + + if coords_np.shape[1] == 3: + return coords_np - if n_points == 0: - return + coords_3d = np.zeros((coords_np.shape[0], 3), dtype=np.float64) + coords_3d[:, :2] = coords_np + return coords_3d - dim = coords.shape[1] - if dim == 2: - coords_3d = np.zeros((n_points, 3), dtype=np.float64) - coords_3d[:, :2] = coords +def _data_array_and_components( + values: np.ndarray, n_points: int, n_components: Optional[int] = None +) -> tuple[np.ndarray, int]: + values_np = np.asarray(values, dtype=np.float64) + + if values_np.ndim == 0: + values_np = np.full(n_points, values_np, dtype=np.float64) + + if values_np.shape[0] != n_points: + raise ValueError( + "VTU point data arrays must have the same first dimension as coordinates." + ) + + if values_np.ndim == 1: + inferred_components = 1 + elif values_np.ndim == 2: + inferred_components = values_np.shape[-1] else: - coords_3d = np.asarray(coords, dtype=np.float64) + raise ValueError("VTU point data arrays must be scalar or vector arrays.") + + component_count = n_components + if component_count is None: + component_count = 3 if inferred_components == 2 else inferred_components + + if inferred_components == 2 and component_count == 3: + values_3d = np.zeros((n_points, 3), dtype=np.float64) + values_3d[:, :2] = values_np + values_np = values_3d + elif inferred_components != component_count: + raise ValueError( + "VTU point data component count does not match the provided schema." + ) + + return values_np, component_count + + +def write_vtu( + filename: str, + coords: np.ndarray, + data: dict, + data_schema: Optional[Dict[str, int]] = None, +) -> None: + """ + Writes a VTU file containing locally owned points and point data + + The file is an ``UnstructuredGrid`` with one ``VTK_VERTEX`` cell per point. + ``data_schema`` can be used to force all ranks, including empty ranks, to + write the same point-data arrays. This keeps the corresponding ``.pvtu`` + file valid even if a rank currently owns no points. + """ + coords_3d = _as_3d_points(coords) + n_points = coords_3d.shape[0] vtk_file = ET.Element( "VTKFile", type="UnstructuredGrid", version="0.1", byte_order="LittleEndian" @@ -65,25 +107,29 @@ def write_vtu(filename: str, coords: np.ndarray, data: dict) -> None: ) type_arr.text = " ".join(["1"] * n_points) + schema = dict(data_schema or {}) + if not schema: + for key, val in data.items(): + _, n_components = _data_array_and_components(np.asarray(val), n_points) + schema[key] = n_components + point_data = ET.SubElement(piece, "PointData") - for key, val in data.items(): - val_np = np.asarray(val, dtype=np.float64) - if val_np.ndim == 1: - n_comp = 1 + for key, n_components in schema.items(): + if key in data: + val_np, n_components = _data_array_and_components( + np.asarray(data[key]), n_points, n_components + ) + elif n_components == 1: + val_np = np.zeros(n_points, dtype=np.float64) else: - n_comp = val_np.shape[1] - if n_comp == 2: - val_3d = np.zeros((n_points, 3), dtype=np.float64) - val_3d[:, :2] = val_np - val_np = val_3d - n_comp = 3 + val_np = np.zeros((n_points, n_components), dtype=np.float64) data_arr = ET.SubElement( point_data, "DataArray", type="Float64", Name=key, - NumberOfComponents=str(n_comp), + NumberOfComponents=str(n_components), format="ascii", ) data_arr.text = " ".join(map(str, val_np.ravel())) @@ -95,16 +141,7 @@ def write_vtu(filename: str, coords: np.ndarray, data: dict) -> None: def write_pvtu(filename: str, source_files: list, data_keys: dict) -> None: """ - Writes a Parallel VTU (.pvtu) file referencing the multiple subset .vtu files. - - Parameters - ---------- - filename : str - Output file path for the PVTU. - source_files : list - List of VTU file names that this PVTU references. - data_keys : dict - Dictionary mapping data array names to their number of components. + Writes a Parallel VTU (.pvtu) file referencing rank-local .vtu files. """ vtk_file = ET.Element( "VTKFile", type="PUnstructuredGrid", version="0.1", byte_order="LittleEndian" diff --git a/tests/unit/test_vtu_export.py b/tests/unit/test_vtu_export.py new file mode 100644 index 00000000..0961d21e --- /dev/null +++ b/tests/unit/test_vtu_export.py @@ -0,0 +1,81 @@ +import importlib.util +import os +import xml.etree.ElementTree as ET +from pathlib import Path +from tempfile import TemporaryDirectory +from unittest import TestCase + +import numpy as np + +vtu_export_path = ( + Path(__file__).parents[2] / "micro_manager" / "tools" / "vtu_export.py" +) +spec = importlib.util.spec_from_file_location("vtu_export", vtu_export_path) +assert spec is not None +assert spec.loader is not None +vtu_export = importlib.util.module_from_spec(spec) +spec.loader.exec_module(vtu_export) +write_pvtu = vtu_export.write_pvtu +write_vtu = vtu_export.write_vtu + + +class TestVTUExport(TestCase): + def test_write_vtu_with_empty_rank_and_schema(self): + with TemporaryDirectory() as tmpdir: + filename = os.path.join(tmpdir, "empty.vtu") + + write_vtu( + filename, + np.empty((0, 2)), + {}, + {"Scalar": 1, "Vector": 3}, + ) + + root = ET.parse(filename).getroot() + piece = root.find("./UnstructuredGrid/Piece") + assert piece is not None + self.assertEqual(piece.get("NumberOfPoints"), "0") + self.assertEqual(piece.get("NumberOfCells"), "0") + + data_arrays = piece.findall("./PointData/DataArray") + self.assertEqual( + [array.get("Name") for array in data_arrays], ["Scalar", "Vector"] + ) + self.assertEqual( + [array.get("NumberOfComponents") for array in data_arrays], ["1", "3"] + ) + + def test_write_vtu_pads_two_component_vectors(self): + with TemporaryDirectory() as tmpdir: + filename = os.path.join(tmpdir, "rank.vtu") + + write_vtu( + filename, + np.array([[1.0, 2.0]]), + {"Vector": np.array([[3.0, 4.0]])}, + {"Vector": 3}, + ) + + root = ET.parse(filename).getroot() + vector_data = root.find("./UnstructuredGrid/Piece/PointData/DataArray") + assert vector_data is not None + self.assertEqual(vector_data.get("NumberOfComponents"), "3") + self.assertEqual(vector_data.text, "3.0 4.0 0.0") + + def test_write_pvtu_references_rank_files(self): + with TemporaryDirectory() as tmpdir: + filename = os.path.join(tmpdir, "all.pvtu") + + write_pvtu(filename, ["rank0.vtu", "rank1.vtu"], {"Scalar": 1}) + + root = ET.parse(filename).getroot() + pieces = root.findall("./PUnstructuredGrid/Piece") + self.assertEqual( + [piece.get("Source") for piece in pieces], ["rank0.vtu", "rank1.vtu"] + ) + + +if __name__ == "__main__": + import unittest + + unittest.main()