Coverage for pyrc \ visualization \ vtk_parser.py: 11%
104 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-04-13 16:59 +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(t_idx,
21 name: str,
22 time_steps: list,
23 matrix,
24 output_folder: str,
25 format_layout: str,
26 points: np.ndarray,
27 cells: np.array,
28 cell_type: str = "hexahedron"):
29 """
30 Writes a single .vtu file for a specific time step. Is used in multiprocessing pool.starmap
32 Parameters
33 ----------
34 t_idx : int
35 The index of the time step.
36 name : str
37 The name of the parameter.
38 time_steps : list
39 A list with all time steps.
40 matrix : np.ndarray
41 A numpy matrix with the values.
42 Shape: time_steps, nodes
43 For one time step one row is used (not one column)!
44 output_folder : str
45 The output folder.
46 format_layout : str
47 A format to style the index of the temperature. Should be: "0{len(time_steps)}d"
50 Returns
51 -------
52 str :
53 The line to add to the pvd file that links to the created temperature file.
54 """
55 vtu_filename = os.path.join(output_folder, f"{name}_{t_idx:{format_layout}}.vtu")
57 mesh = meshio.Mesh(
58 points=points,
59 cells=[(cell_type, cells)],
60 cell_data={name: [matrix[t_idx, :]]}, # Fast NumPy slicing
61 )
63 mesh.write(vtu_filename)
64 return f'<DataSet timestep="{time_steps[t_idx]}" file="{name}_{t_idx:{format_layout}}.vtu" />'
67# Function to write VTU files
68def write_static_cells_vtu(nodes: list[Node],
69 output_folder="output") -> tuple[np.ndarray, np.ndarray]:
70 os.makedirs(os.path.normpath(output_folder), exist_ok=True)
72 # Define hexahedral cells (each node is a box-like cell)
73 cells = []
74 all_points = []
75 for i, node in enumerate(nodes):
76 # Compute the 8 corner points of the hexahedral cell
77 min_corner = node.position - node.delta / 2
78 max_corner = node.position + node.delta / 2
79 cell_points = np.array([
80 min_corner,
81 [max_corner[0], min_corner[1], min_corner[2]],
82 [max_corner[0], max_corner[1], min_corner[2]],
83 [min_corner[0], max_corner[1], min_corner[2]],
84 [min_corner[0], min_corner[1], max_corner[2]],
85 [max_corner[0], min_corner[1], max_corner[2]],
86 [max_corner[0], max_corner[1], max_corner[2]],
87 [min_corner[0], max_corner[1], max_corner[2]],
88 ])
90 # Append points and define cell connectivity
91 start_idx = len(all_points)
92 all_points.extend(cell_points)
93 cells.append(list(range(start_idx, start_idx + 8)))
95 all_points = np.array(all_points) # Convert to NumPy array
97 # --- 2️⃣ Save the Base Geometry as `nodes.vtu` ---
98 cells = np.array(cells)
99 base_mesh = meshio.Mesh(
100 all_points,
101 [("hexahedron", cells)],
102 )
103 base_vtu_file = os.path.join(output_folder, "nodes.vtu")
104 base_mesh.write(base_vtu_file)
105 print(f"Base geometry saved: {base_vtu_file}")
107 pvd_static = ['<DataSet timestep="0.0" file="nodes.vtu" />']
108 write_pvd(pvd_static, output_folder=output_folder, name="nodes")
110 return all_points, cells
113# Function to write VTU files
114def write_temperature_vtu(result_temperatures: np.ndarray,
115 time_steps: list,
116 points: np.ndarray,
117 cells: np.ndarray,
118 output_folder="output",
119 step_interval: int = 1, ):
120 os.makedirs(os.path.normpath(output_folder), exist_ok=True)
122 # --- 1️⃣ Generate Time-Step VTU Files ---
123 temperatures = result_temperatures # Shape: (num_time_steps, num_nodes)
125 num_time_steps = temperatures.shape[0]
127 format_number = len(str(num_time_steps))
128 format_layout = f"0{format_number}d"
129 if step_interval is None:
130 step_interval = max(1, len(time_steps) // 2000)
131 selected_indices = range(0, num_time_steps, step_interval)
133 # --- 2️⃣ Parallel File Writing ---
134 pool = multiprocessing.Pool(processes=multiprocessing.cpu_count()) # Use all CPU cores
135 tasks = [(t_idx, "temperature", time_steps, temperatures, output_folder, format_layout, points, cells) for t_idx in
136 selected_indices]
138 pvd_entries = pool.starmap(write_vtu_parallel, tasks)
139 pool.close()
140 pool.join()
142 # for t_idx in selected_indices:
143 # temperatures = np.array([node.temperature_vector.iloc[t_idx] for node in nodes])
144 #
145 # timestep_mesh = meshio.Mesh(
146 # all_points,
147 # [("hexahedron", np.array(cells))],
148 # cell_data={"Temperature": [temperatures]},
149 # )
150 #
151 # vtu_filename = os.path.join(output_folder, f"temperature_{t_idx:{format_layout}}.vtu")
152 # timestep_mesh.write(vtu_filename)
153 #
154 # pvd_entries.append(f'<DataSet timestep="{time_steps[t_idx]}" file="temperature_{t_idx:{format_layout}}.vtu" />')
156 write_pvd(pvd_entries, output_folder)
159def write_pvd(pvd_entries: list, output_folder="output", name="00_simulation"):
160 os.makedirs(output_folder, exist_ok=True)
162 pvd_content = f"""<?xml version="1.0"?>
163 <VTKFile type="Collection" version="0.1">
164 <Collection>
165 {"".join(pvd_entries)}
166 </Collection>
167 </VTKFile>"""
169 pvd_file_path = os.path.join(output_folder, f'{name}.pvd')
170 with open(pvd_file_path, "w") as f:
171 f.write(pvd_content)
173 print(f"PVD file saved: {pvd_file_path}")
176def create_heat_flux_lines(resistors,
177 time_steps: list,
178 output_folder="output",
179 step_interval=None):
180 os.makedirs(output_folder, exist_ok=True)
181 n_time_steps = len(time_steps)
183 if step_interval is None:
184 step_interval = max(1, n_time_steps // 2000)
185 format_number = len(str(n_time_steps))
186 format_layout = f"0{format_number}d"
188 selected_indices = range(0, n_time_steps, step_interval)
190 line_points = []
191 line_cells = []
192 flux_vectors = []
193 sources = []
194 sinks = []
196 point_idx = 0
198 from pyrc.core.resistors import MassTransport
199 for resistor in resistors:
200 if isinstance(resistor, MassTransport):
201 source, sink = resistor.source, resistor.sink
202 else:
203 source, sink = resistor.neighbours[0], resistor.neighbours[1]
204 sources.append(source)
205 sinks.append(sink)
207 direction = sink.position - source.position
208 direction /= np.linalg.norm(direction) # Normalize vector
209 flux_vectors.append(direction) # Store unit vector for later scaling
211 line_points.extend([source.position, sink.position])
212 line_cells.append([point_idx, point_idx + 1]) # Positive direction
213 line_cells.append([point_idx + 1, point_idx]) # Negative direction
215 point_idx += 2
217 line_points = np.array(line_points)
218 line_cells = np.array(line_cells)
219 flux_vectors = np.array(flux_vectors)
221 line_mesh = meshio.Mesh(
222 line_points,
223 [("line", np.array(line_cells))],
224 )
225 line_vtu_file = os.path.join(output_folder, "line_cells.vtu")
226 line_mesh.write(line_vtu_file)
227 print(f"Line cell geometry saved: {line_vtu_file}")
229 temperatures_source = np.array(
230 [source.temperature_vector for source in sources]) # shape[resistors,timesteps]
231 temperatures_sink = np.array([sink.temperature_vector for sink in sinks]) # shape[resistors, timesteps]
232 resistances = np.array([resistor.resistance for resistor in resistors])[:, np.newaxis] # shape[resistors, 1]
233 temp_delta = temperatures_source - temperatures_sink # shape[resistors, timesteps]
234 fluxes = temp_delta[:, selected_indices] / resistances # shape[resistors, timesteps]
236 flux_positive = np.maximum(fluxes, 0)[..., np.newaxis] * flux_vectors[:, np.newaxis, :]
237 flux_negative = -np.maximum(-fluxes, 0)[..., np.newaxis] * flux_vectors[:, np.newaxis, :]
239 flux_vector_data = np.empty((len(selected_indices), len(resistors) * 2, 3))
240 flux_vector_data[:, ::2, :] = np.transpose(flux_positive, (1, 0, 2))
241 flux_vector_data[:, 1::2, :] = np.transpose(flux_negative, (1, 0, 2))
243 pool = multiprocessing.Pool(processes=multiprocessing.cpu_count()) # Use all CPU cores
244 tasks = [(t_idx,
245 "HeatFluxVector",
246 time_steps[selected_indices],
247 flux_vector_data,
248 output_folder,
249 format_layout,
250 line_points,
251 line_cells,
252 "line") for t_idx in range(len(selected_indices))]
254 pvd_entries = pool.starmap(write_vtu_parallel, tasks)
255 pool.close()
256 pool.join()
258 # for t_idx in selected_indices:
259 # fluxes = (temperatures_source[:, t_idx] - temperatures_sink[:, t_idx]) / resistances.flatten()
260 #
261 # flux_vector_data = np.zeros((len(resistors) * 2, 3)) # 3D vector field
262 # flux_vector_data[::2] = flux_vectors * np.maximum(fluxes[:, np.newaxis], 0) # Positive flux
263 # flux_vector_data[1::2] = -flux_vectors * np.maximum(-fluxes[:, np.newaxis], 0) # Negative flux
264 #
265 # mesh = meshio.Mesh(
266 # points=line_points,
267 # cells=[("line", np.array(line_cells))],
268 # cell_data={"HeatFluxVector": [flux_vector_data]}
269 # )
270 # vtu_filename = os.path.join(output_folder, f"heat_flux_{t_idx:{format_layout}}.vtu")
271 # mesh.write(vtu_filename)
272 #
273 # pvd_entries.append(f'<DataSet timestep="{time_steps[t_idx]}" file="heat_flux_{t_idx:{format_layout}}.vtu" />')
275 write_pvd(pvd_entries, output_folder, name="00_HeatFluxVector")