Coverage for src/beamme/core/vtk

1# The MIT License (MIT)

5# Permission is hereby granted, free of charge, to any person obtaining a copy

6# of this software and associated documentation files (the "Software"), to deal

7# in the Software without restriction, including without limitation the rights

8# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell

9# copies of the Software, and to permit persons to whom the Software is

10# furnished to do so, subject to the following conditions:

11#

12# The above copyright notice and this permission notice shall be included in

13# all copies or substantial portions of the Software.

14#

15# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR

16# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,

17# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE

18# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER

19# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,

20# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN

21# THE SOFTWARE.

22"""This module provides a class that is used to write VTK files."""

24import numbers as _numbers

25import os as _os

26import warnings as _warnings

27from enum import Enum as _Enum

28from enum import auto as _auto

30import numpy as _np

31import vtk as _vtk

33from beamme.core.conf import bme as _bme

35# Number of digits for node set output (this will be set in the

36# Mesh.get_unique_geometry_sets() method.

37VTK_NODE_SET_FORMAT = "{:05}"

39# Nan values for vtk data, since we currently can't set nan explicitly.

40VTK_NAN_INT = -1

41VTK_NAN_FLOAT = 0.0

44class VTKGeometry(_Enum):

45 """Enum for VTK geometry types (for now cells and points)."""

47 point = _auto()

48 cell = _auto()

51class VTKType(_Enum):

52 """Enum for VTK value types."""

54 int = _auto()

55 float = _auto()

58class VTKTensor(_Enum):

59 """Enum for VTK tensor types."""

61 scalar = _auto()

62 vector = _auto()

65def add_point_data_node_sets(point_data, nodes, *, extra_points=0):

66 """Add the information if a node is part of a set to the point_data vector

67 for all nodes in the list 'nodes'.

69 The extra_points argument specifies how many additional

70 visualization points there are, i.e., points that are not based on

71 nodes, but are only used for visualization purposes.

72 """

74 # Get list with node set indices of the given nodes

75 geometry_set_list = []

76 for node in nodes:

77 geometry_set_list.extend(node.node_sets_link)

79 # Remove double entries of list.

80 geometry_set_list = list(set(geometry_set_list))

82 # Loop over the geometry sets.

83 n_nodes = len(nodes)

84 for geometry_set in geometry_set_list:

85 # Check which nodes are connected to a geometry set.

86 data_vector = _np.zeros(n_nodes + extra_points)

87 for i, node in enumerate(nodes):

88 if geometry_set in node.node_sets_link:

89 data_vector[i] = 1

90 else:

91 data_vector[i] = VTK_NAN_INT

92 for i in range(extra_points):

93 data_vector[n_nodes + i] = (

94 1 if geometry_set.geometry_type is _bme.geo.line else VTK_NAN_INT

95 )

97 # Get the name of the geometry type.

98 if geometry_set.geometry_type is _bme.geo.point:

99 geometry_name = "geometry_point"

100 elif geometry_set.geometry_type is _bme.geo.line:

101 geometry_name = "geometry_line"

102 elif geometry_set.geometry_type is _bme.geo.surface:

103 geometry_name = "geometry_surface"

104 elif geometry_set.geometry_type is _bme.geo.volume:

105 geometry_name = "geometry_volume"

106 else:

107 raise TypeError("The geometry type is wrong!")

108

109 # Add the data vector.

110 set_name = f"{geometry_name}_set_{_bme.vtk_node_set_format.format(geometry_set.i_global + 1)}"

111 point_data[set_name] = (data_vector, VTKType.int)

112

113

114def _get_data_value_and_type(data):

115 """Return the data and its type if one was given.

116

117 The default type, if none was given is float.

118 """

119 if isinstance(data, tuple):

120 return data[0], data[1]

121 else:

122 return data, VTKType.float

123

124

125def _get_vtk_array_type(data):

126 """Return the corresponding beamme type."""

127 data_type = data.GetDataTypeAsString()

128 if data_type == "int":

129 return VTKType.int

130 elif data_type == "double":

131 return VTKType.float

132 raise ValueError(f'Got unexpected type "{data_type}"!')

133

134

135class VTKWriter:

136 """A class that manages VTK cells and data and can also create them."""

137

138 def __init__(self):

139 # Initialize VTK objects.

140 self.points = _vtk.vtkPoints()

141 self.points.SetDataTypeToDouble()

142 self.grid = _vtk.vtkUnstructuredGrid()

143

144 # Link points to grid.

145 self.grid.SetPoints(self.points)

146

147 # Container for output data.

148 self.data = {}

149 for key1 in VTKGeometry:

150 for key2 in VTKTensor:

151 self.data[key1, key2] = {}

152

153 def add_points(self, points, *, point_data=None):

154 """Add points to the data stored in this object.

155

156 Args

157 ----

158 points: [3d vector]

159 Coordinates of points for this cell.

160 point_data: dic

161 A dictionary containing data that will be added for the newly added points.

162 If a field exists in the global data but not in the one added here, that field

163 will be set to bme.vtk_nan for the newly added points.

164

165 Return:

166 ----

167 indices: [int]

168 A list with the global indices of the added points.

169 """

170

171 n_points = len(points)

172

173 # Check if point data containers are of the correct size

174 if point_data is not None:

175 for key, item_value in point_data.items():

176 value, _data_type = _get_data_value_and_type(item_value)

177 if not len(value) == n_points:

178 raise IndexError(

179 f"The length of coordinates is {n_points},"

180 f"the length of {key} is {len(value)}, does not match!"

181 )

182

183 # Add point data

184 self._add_data(point_data, VTKGeometry.point, n_new_items=n_points)

185

186 # Add point coordinates

187 n_grid_points = self.points.GetNumberOfPoints()

188 for point in points:

189 # Add the coordinate to the global list of coordinates.

190 self.points.InsertNextPoint(*point)

191

192 return _np.array(

193 [n_grid_points + i_point for i_point in range(len(points))], dtype=int

194 )

195

196 def add_cell(self, cell_type, topology, *, cell_data=None):

197 """Create a cell and add it to the global array.

198

199 Args

200 ----

201 cell_type: VTK_type

202 Type of cell that will be created.

203 topology: [int]

204 The connectivity between the cell and the global points.

205 cell_data: dic

206 A dictionary containing data that will be added for the newly added cell.

207 If a field exists in the global data but not in the one added here, that field

208 will be set to bme.vtk_nan for the newly added cell.

209 """

210

211 # Add the data entries.

212 self._add_data(cell_data, VTKGeometry.cell)

213

214 # Create the cell.

215 geometry_item = cell_type()

216 geometry_item.GetPointIds().SetNumberOfIds(len(topology))

217

218 # Set the connectivity

219 for i_local, i_global in enumerate(topology):

220 geometry_item.GetPointIds().SetId(i_local, i_global)

221

222 # Add to global cells

223 self.grid.InsertNextCell(

224 geometry_item.GetCellType(), geometry_item.GetPointIds()

225 )

226

227 def _add_data(self, data_container, vtk_geom_type, *, n_new_items=1):

228 """Add a data container to the existing global data container of this

229 object.

230

231 Args

232 ----

233 data_container: see self.add_cell

234 vtk_geom_type: bme.vtk_geo

235 Type of data container that is added

236 n_new_items: int

237 Number of new items added. This is needed to fill up data fields that are in the

238 global data but not in the one that is added.

239 """

240

241 # Check if data container already exists. If not, add it and also add

242 # previous entries.

243 if data_container is not None:

244 if vtk_geom_type == VTKGeometry.cell:

245 n_items = self.grid.GetNumberOfCells()

246 else:

247 n_items = self.grid.GetNumberOfPoints()

248

249 for key, item_value in data_container.items():

250 # Get the data and the value type (int or float).

251 value, data_type = _get_data_value_and_type(item_value)

252

253 # Data type.

254 if vtk_geom_type == VTKGeometry.cell:

255 vtk_tensor_type = self._get_vtk_data_type(value)

256 else:

257 for item in value:

258 vtk_tensor_type = self._get_vtk_data_type(item)

259

260 # Check if key already exists.

261 if key not in self.data[vtk_geom_type, vtk_tensor_type].keys():

262 # Set up the VTK data array.

263 if data_type is VTKType.float:

264 data = _vtk.vtkDoubleArray()

265 else:

266 data = _vtk.vtkIntArray()

267 data.SetName(key)

268 if vtk_tensor_type == VTKTensor.scalar:

269 data.SetNumberOfComponents(1)

270 else:

271 data.SetNumberOfComponents(3)

272

273 # Add the empty values for all previous cells / points.

274

275 for i in range(n_items):

276 self._add_single_data_item(data, vtk_tensor_type)

277 self.data[vtk_geom_type, vtk_tensor_type][key] = data

278

279 else:

280 # In this case we just check that the already existing

281 # data has the same type.

282 data_array = self.data[vtk_geom_type, vtk_tensor_type][key]

283 if not _get_vtk_array_type(data_array) == data_type:

284 raise ValueError(

285 (

286 'The existing data with the key "{}"'

287 + ' is of type "{}", but the type you tried to add'

288 + ' is "{}"!'

289 ).format(key, data_array.GetDataTypeAsString(), data_type)

290 )

291

292 # Add to global data. Check if there is something to be added. If not an empty value

293 # is added.

294 for key_tensor in VTKTensor:

295 global_data = self.data[vtk_geom_type, key_tensor]

296 if data_container is None:

297 data_container = {}

298

299 for key, value in global_data.items():

300 # Check if an existing field is also given for this function.

301 if key in data_container.keys():

302 # Get the data and the value type (int or float).

303 data_values, _ = _get_data_value_and_type(data_container[key])

304

305 # Add the given data.

306 if vtk_geom_type == VTKGeometry.cell:

307 self._add_single_data_item(

308 value, key_tensor, non_zero_data=data_values

309 )

310 else:

311 for item in data_values:

312 self._add_single_data_item(

313 value, key_tensor, non_zero_data=item

314 )

315 else:

316 # Add empty data.

317 if vtk_geom_type == VTKGeometry.cell:

318 self._add_single_data_item(value, key_tensor)

319 else:

320 for item in range(n_new_items):

321 self._add_single_data_item(value, key_tensor)

322

323 @staticmethod

324 def _get_vtk_data_type(data):

325 """Return the type of data.

326

327 Check if data matches an expected case.

328 """

329

330 if isinstance(data, (list, _np.ndarray)):

331 if len(data) == 3:

332 return VTKTensor.vector

333 raise IndexError(

334 f"Only 3d vectors are implemented yet! Got len(data) = {len(data)}"

335 )

336 elif isinstance(data, _numbers.Number):

337 return VTKTensor.scalar

338

339 raise ValueError(f"Data {data} did not match any expected case!")

340

341 @staticmethod

342 def _add_single_data_item(data, vtk_tensor_type, non_zero_data=None):

343 """Add data to a VTK data array."""

344

345 if _get_vtk_array_type(data) == VTKType.int:

346 nan_value = VTK_NAN_INT

347 elif _get_vtk_array_type(data) == VTKType.float:

348 nan_value = VTK_NAN_FLOAT

349

350 if vtk_tensor_type == VTKTensor.scalar:

351 if non_zero_data is None:

352 data.InsertNextTuple1(nan_value)

353 else:

354 data.InsertNextTuple1(non_zero_data)

355 else:

356 if non_zero_data is None:

357 data.InsertNextTuple3(nan_value, nan_value, nan_value)

358 else:

359 data.InsertNextTuple3(

360 non_zero_data[0], non_zero_data[1], non_zero_data[2]

361 )

362

363 def complete_data(self):

364 """Add the stored data to the vtk grid."""

365 for (key_geom, _key_data), value in self.data.items():

366 for vtk_data in value.values():

367 if key_geom == VTKGeometry.cell:

368 self.grid.GetCellData().AddArray(vtk_data)

369 else:

370 self.grid.GetPointData().AddArray(vtk_data)

371

372 def write_vtk(self, filepath, *, binary=True):

373 """Write the VTK geometry and data to a file.

374

375 Args

376 ----

377 filepath: str

378 Path to output file. The file extension should be vtu.

379 binary: bool

380 If the data should be written encoded in binary or in human readable text.

381 """

382

383 # Check if directory for file exits.

384 file_directory = _os.path.dirname(filepath)

385 if not _os.path.isdir(file_directory):

386 raise ValueError(f"Directory {file_directory} does not exist!".format())

387

388 # Initialize VTK writer.

389 writer = _vtk.vtkXMLUnstructuredGridWriter()

390

391 # Set the ascii flag.

392 if not binary:

393 writer.SetDataModeToAscii()

394

395 # Check the file extension.

396 _filename, file_extension = _os.path.splitext(filepath)

397 if not file_extension.lower() == ".vtu":

398 _warnings.warn(f'The extension should be "vtu", got {file_extension}!')

399

400 # Write geometry and data to file.

401 writer.SetFileName(filepath)

402 writer.SetInputData(self.grid)

403 writer.Write()

Coverage for src/beamme/core/vtk_writer.py: 95%

166 statements