diff --git a/micro_manager/micro_manager.py b/micro_manager/micro_manager.py index 64239a35..4e5e6b27 100644 --- a/micro_manager/micro_manager.py +++ b/micro_manager/micro_manager.py @@ -12,28 +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 .load_balancing import create_load_balancer +from .tools.vtu_export import write_pvtu, write_vtu try: from .interpolation import Interpolation @@ -73,6 +73,27 @@ 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 = "." + 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() + 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.abspath(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() @@ -133,7 +154,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, ) @@ -148,9 +169,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 @@ -187,7 +205,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: @@ -330,6 +347,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 @@ -415,6 +435,189 @@ 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 = {} + 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): + 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 + + 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_{mesh_file_name}_{self._n}_rank{self._rank}.vtu", + ) + write_vtu(filename, coords_local, data_local, data_schema) + + if self._is_parallel: + self._comm.Barrier() + if self._rank == 0: + pvtu_filename = os.path.join( + self._export_vtu_dir, + f"Micro-Manager_{mesh_file_name}_{self._n}.pvtu", + ) + source_files = [ + f"Micro-Manager_{mesh_file_name}_{self._n}_rank{rank}.vtu" + for rank in range(self._size) + ] + write_pvtu(pvtu_filename, source_files, data_schema) + def initialize(self) -> None: """ Initialize the Micro Manager by performing the following tasks: @@ -750,9 +953,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." @@ -765,17 +968,17 @@ def initialize(self) -> None: initial_macro_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: # If adaptivity is off, the returned initial data from the initialize() method will be ignored self._logger.log_warning_rank_zero( @@ -854,7 +1057,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 @@ -886,7 +1089,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 @@ -1045,9 +1248,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: @@ -1101,9 +1304,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 new file mode 100644 index 00000000..037a9ab0 --- /dev/null +++ b/micro_manager/tools/vtu_export.py @@ -0,0 +1,169 @@ +import os +import xml.etree.ElementTree as ET +from typing import Dict, Optional + +import numpy as np + + +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 + + coords_3d = np.zeros((coords_np.shape[0], 3), dtype=np.float64) + coords_3d[:, :2] = coords_np + return coords_3d + + +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: + 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" + ) + 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) + + 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, 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: + 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_components), + 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 rank-local .vtu files. + """ + 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") 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()