Coverage for pyrc \ visualization \ vtk_parser.py: 11%
104 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-05-07 08:20 +0200
1# -------------------------------------------------------------------------------
2# Copyright (C) 2026 Joel Kimmich, Tim Jourdan
3# ------------------------------------------------------------------------------
4# License
5# This file is part of PyRC, distributed under GPL-3.0-or-later.
6# ------------------------------------------------------------------------------
8from __future__ import annotations
9from typing import TYPE_CHECKING
11import numpy as np
12import meshio
13import os
14import multiprocessing
16if TYPE_CHECKING:
17 from pyrc.core.nodes import Node
20def write_vtu_parallel(
21 t_idx,
22 name: str,
23 time_steps: list,
24 matrix,
25 output_folder: str,
26 format_layout: str,
27 points: np.ndarray,
28 cells: np.array,
29 cell_type: str = "hexahedron",
30):
31 """
32 Writes a single .vtu file for a specific time step. Is used in multiprocessing pool.starmap
34 Parameters
35 ----------
36 t_idx : int
37 The index of the time step.
38 name : str
39 The name of the parameter.
40 time_steps : list
41 A list with all time steps.
42 matrix : np.ndarray
43 A numpy matrix with the values.
44 Shape: time_steps, nodes
45 For one time step one row is used (not one column)!
46 output_folder : str
47 The output folder.
48 format_layout : str
49 A format to style the index of the temperature. Should be: "0{len(time_steps)}d"
52 Returns
53 -------
54 str :
55 The line to add to the pvd file that links to the created temperature file.
56 """
57 vtu_filename = os.path.join(output_folder, f"{name}_{t_idx:{format_layout}}.vtu")
59 mesh = meshio.Mesh(
60 points=points,
61 cells=[(cell_type, cells)],
62 cell_data={name: [matrix[t_idx, :]]}, # Fast NumPy slicing
63 )
65 mesh.write(vtu_filename)
66 return f'<DataSet timestep="{time_steps[t_idx]}" file="{name}_{t_idx:{format_layout}}.vtu" />'
69# Function to write VTU files
70def write_static_cells_vtu(nodes: list[Node], output_folder="output") -> tuple[np.ndarray, np.ndarray]:
71 os.makedirs(os.path.normpath(output_folder), exist_ok=True)
73 # Define hexahedral cells (each node is a box-like cell)
74 cells = []
75 all_points = []
76 for i, node in enumerate(nodes):
77 # Compute the 8 corner points of the hexahedral cell
78 min_corner = node.position - node.delta / 2
79 max_corner = node.position + node.delta / 2
80 cell_points = np.array(
81 [
82 min_corner,
83 [max_corner[0], min_corner[1], min_corner[2]],
84 [max_corner[0], max_corner[1], min_corner[2]],
85 [min_corner[0], max_corner[1], min_corner[2]],
86 [min_corner[0], min_corner[1], max_corner[2]],
87 [max_corner[0], min_corner[1], max_corner[2]],
88 [max_corner[0], max_corner[1], max_corner[2]],
89 [min_corner[0], max_corner[1], max_corner[2]],
90 ]
91 )
93 # Append points and define cell connectivity
94 start_idx = len(all_points)
95 all_points.extend(cell_points)
96 cells.append(list(range(start_idx, start_idx + 8)))
98 all_points = np.array(all_points) # Convert to NumPy array
100 # --- 2️⃣ Save the Base Geometry as ``nodes.vtu`` ---
101 cells = np.array(cells)
102 base_mesh = meshio.Mesh(
103 all_points,
104 [("hexahedron", cells)],
105 )
106 base_vtu_file = os.path.join(output_folder, "nodes.vtu")
107 base_mesh.write(base_vtu_file)
108 print(f"Base geometry saved: {base_vtu_file}")
110 pvd_static = ['<DataSet timestep="0.0" file="nodes.vtu" />']
111 write_pvd(pvd_static, output_folder=output_folder, name="nodes")
113 return all_points, cells
116# Function to write VTU files
117def write_temperature_vtu(
118 result_temperatures: np.ndarray,
119 time_steps: list,
120 points: np.ndarray,
121 cells: np.ndarray,
122 output_folder="output",
123 step_interval: int = 1,
124 run_parallel=False,
125):
126 """
127 Creates multiple vtu files containing cell temperatures.
129 Parameters
130 ----------
131 result_temperatures : np.narray
132 The temperatures with shape (num_time_steps, num_nodes)
133 time_steps : list
134 The time steps.
135 points : np.ndarray
136 The points to build the mesh.
137 First return of ``write_static_cells_vtu``
138 cells : np.ndarray
139 A list defining the cells out of the points using their number/increment.
140 Second return of ``write_static_cells_vtu``
141 output_folder : str, optional
142 The location to save the files in (folder with path).
143 step_interval : int, optional
144 An increment which time steps should be converted to vtu files. If 1, every time step is used, if 2,
145 each second, and so on.
146 run_parallel : bool
147 If True, the creation of the vtu files is done using all cores. But it only works on Linux because of the
148 different mechanisms how processes are created (spawn/fork).
149 """
150 os.makedirs(os.path.normpath(output_folder), exist_ok=True)
152 # --- Generate Time-Step VTU Files ---
153 temperatures = result_temperatures # Shape: (num_time_steps, num_nodes)
155 num_time_steps = temperatures.shape[0]
157 format_number = len(str(num_time_steps))
158 format_layout = f"0{format_number}d"
159 if step_interval is None:
160 step_interval = max(1, len(time_steps) // 2000)
161 selected_indices = range(0, num_time_steps, step_interval)
163 # --- File Writing ---
164 tasks = [
165 (t_idx, "temperature", time_steps, temperatures, output_folder, format_layout, points, cells)
166 for t_idx in selected_indices
167 ]
169 if run_parallel:
170 with multiprocessing.Pool(processes=multiprocessing.cpu_count()) as pool:
171 pvd_entries = pool.starmap(write_vtu_parallel, tasks)
172 else:
173 pvd_entries = [write_vtu_parallel(*task) for task in tasks]
175 # for t_idx in selected_indices:
176 # temperatures = np.array([node.temperature_vector.iloc[t_idx] for node in nodes])
177 #
178 # timestep_mesh = meshio.Mesh(
179 # all_points,
180 # [("hexahedron", np.array(cells))],
181 # cell_data={"Temperature": [temperatures]},
182 # )
183 #
184 # vtu_filename = os.path.join(output_folder, f"temperature_{t_idx:{format_layout}}.vtu")
185 # timestep_mesh.write(vtu_filename)
186 #
187 # pvd_entries.append(f'<DataSet timestep="{time_steps[t_idx]}" file="temperature_{t_idx:{format_layout}}.vtu" />')
189 write_pvd(pvd_entries, output_folder)
192def write_pvd(pvd_entries: list, output_folder="output", name="00_simulation"):
193 os.makedirs(output_folder, exist_ok=True)
195 pvd_content = f"""<?xml version="1.0"?>
196 <VTKFile type="Collection" version="0.1">
197 <Collection>
198 {"".join(pvd_entries)}
199 </Collection>
200 </VTKFile>"""
202 pvd_file_path = os.path.join(output_folder, f"{name}.pvd")
203 with open(pvd_file_path, "w") as f:
204 f.write(pvd_content)
206 print(f"PVD file saved: {pvd_file_path}")
209def create_heat_flux_lines(resistors, time_steps: list, output_folder="output", step_interval=None):
210 os.makedirs(output_folder, exist_ok=True)
211 n_time_steps = len(time_steps)
213 if step_interval is None:
214 step_interval = max(1, n_time_steps // 2000)
215 format_number = len(str(n_time_steps))
216 format_layout = f"0{format_number}d"
218 selected_indices = range(0, n_time_steps, step_interval)
220 line_points = []
221 line_cells = []
222 flux_vectors = []
223 sources = []
224 sinks = []
226 point_idx = 0
228 from pyrc.core.resistors import MassTransport
230 for resistor in resistors:
231 if isinstance(resistor, MassTransport):
232 source, sink = resistor.source, resistor.sink
233 else:
234 source, sink = resistor.neighbours[0], resistor.neighbours[1]
235 sources.append(source)
236 sinks.append(sink)
238 direction = sink.position - source.position
239 direction /= np.linalg.norm(direction) # Normalize vector
240 flux_vectors.append(direction) # Store unit vector for later scaling
242 line_points.extend([source.position, sink.position])
243 line_cells.append([point_idx, point_idx + 1]) # Positive direction
244 line_cells.append([point_idx + 1, point_idx]) # Negative direction
246 point_idx += 2
248 line_points = np.array(line_points)
249 line_cells = np.array(line_cells)
250 flux_vectors = np.array(flux_vectors)
252 line_mesh = meshio.Mesh(
253 line_points,
254 [("line", np.array(line_cells))],
255 )
256 line_vtu_file = os.path.join(output_folder, "line_cells.vtu")
257 line_mesh.write(line_vtu_file)
258 print(f"Line cell geometry saved: {line_vtu_file}")
260 temperatures_source = np.array([source.temperature_vector for source in sources]) # shape[resistors,timesteps]
261 temperatures_sink = np.array([sink.temperature_vector for sink in sinks]) # shape[resistors, timesteps]
262 resistances = np.array([resistor.resistance for resistor in resistors])[:, np.newaxis] # shape[resistors, 1]
263 temp_delta = temperatures_source - temperatures_sink # shape[resistors, timesteps]
264 fluxes = temp_delta[:, selected_indices] / resistances # shape[resistors, timesteps]
266 flux_positive = np.maximum(fluxes, 0)[..., np.newaxis] * flux_vectors[:, np.newaxis, :]
267 flux_negative = -np.maximum(-fluxes, 0)[..., np.newaxis] * flux_vectors[:, np.newaxis, :]
269 flux_vector_data = np.empty((len(selected_indices), len(resistors) * 2, 3))
270 flux_vector_data[:, ::2, :] = np.transpose(flux_positive, (1, 0, 2))
271 flux_vector_data[:, 1::2, :] = np.transpose(flux_negative, (1, 0, 2))
273 pool = multiprocessing.Pool(processes=multiprocessing.cpu_count()) # Use all CPU cores
274 tasks = [
275 (
276 t_idx,
277 "HeatFluxVector",
278 time_steps[selected_indices],
279 flux_vector_data,
280 output_folder,
281 format_layout,
282 line_points,
283 line_cells,
284 "line",
285 )
286 for t_idx in range(len(selected_indices))
287 ]
289 pvd_entries = pool.starmap(write_vtu_parallel, tasks)
290 pool.close()
291 pool.join()
293 # for t_idx in selected_indices:
294 # fluxes = (temperatures_source[:, t_idx] - temperatures_sink[:, t_idx]) / resistances.flatten()
295 #
296 # flux_vector_data = np.zeros((len(resistors) * 2, 3)) # 3D vector field
297 # flux_vector_data[::2] = flux_vectors * np.maximum(fluxes[:, np.newaxis], 0) # Positive flux
298 # flux_vector_data[1::2] = -flux_vectors * np.maximum(-fluxes[:, np.newaxis], 0) # Negative flux
299 #
300 # mesh = meshio.Mesh(
301 # points=line_points,
302 # cells=[("line", np.array(line_cells))],
303 # cell_data={"HeatFluxVector": [flux_vector_data]}
304 # )
305 # vtu_filename = os.path.join(output_folder, f"heat_flux_{t_idx:{format_layout}}.vtu")
306 # mesh.write(vtu_filename)
307 #
308 # pvd_entries.append(f'<DataSet timestep="{time_steps[t_idx]}" file="heat_flux_{t_idx:{format_layout}}.vtu" />')
310 write_pvd(pvd_entries, output_folder, name="00_HeatFluxVector")