Coverage for /usr/local/lib/python3.8/site-packages/spam/helpers/vtkio.py: 100%

1# Library of SPAM functions for reading and writing VTK files with meshio sometimes

4# This program is free software: you can redistribute it and/or modify it

5# under the terms of the GNU General Public License as published by the Free

6# Software Foundation, either version 3 of the License, or (at your option)

7# any later version.

9# This program is distributed in the hope that it will be useful, but WITHOUT

10# ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or

11# FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License for

12# more details.

13#

14# You should have received a copy of the GNU General Public License along with

15# this program. If not, see <http://www.gnu.org/licenses/>.

17"""

18This module offers a set functions in order to read and writes VTK files.

20This module is mainly made of ``meshio`` wrappers adapted to ``spam``.

21The types of data currently handled are:

23 * Structured meshes of cuboids

24 * Unstructured meshes of tetrahedrs

25 * List of objects (TODO)

27VTK files are assumed to be in X,Y,Z format, and numpy arrays are in Z,Y,X, so the following

28policy is followed:

30 - position arrays are simply reversed so Z comes out first

31 - scalar arrays are left as they are

32 - vectors and matrices are assumed to be positional (a surface vector for example) and are reversed

33"""

34import meshio

35import numpy

37# Structured meshes

40def readStructuredVTK(f, fieldAsked=None):

41 """

42 Read a plain text VTK.

44 By default read the full text and put all fields in a dictionnay.

45 Specific fields can be selected with `fieldAsked`.

47 Parameters

48 ----------

49 f : string

50 Name of the VTK file.

52 fieldAsked : array of string, default=None

53 List of the field names.

55 Returns

56 -------

57 fields : dict

58 Dictionnary of all the fields without nodal or cell distinctions.

60 .. code-block:: python

62 fields = {

63 "field1name": [field1value1, field1value2, ...],

64 "field2name": [field2value1, field2value2, ...],

65 # ...

66 }

68 Note

69 ----

70 The outputs are kept flattened.

71 """

73 # read the full file

74 with open(f) as content_file:

75 content = [line.strip() for line in content_file if line.strip()]

77 # find line numbers of principal key words

78 pl = {line.split()[0]: i for i, line in enumerate(content) if line.split()[0] in ["CELLS", "POINTS", "CELL_DATA", "POINT_DATA"]}

80 # find keys in the file

81 nl = {}

82 keys_found = {}

83 if "POINT_DATA" in pl:

84 nl.update({"POINT_DATA": int(content[pl["POINT_DATA"]].split()[1])})

85 keys_found.update(

86 {

87 line.split()[1]: [

88 pl["POINT_DATA"] + i,

89 content[pl["POINT_DATA"] + i].split()[0],

90 "POINT_DATA",

91 ]

92 for i, line in enumerate(content[pl["POINT_DATA"] :])

93 if line.split()[0] in ["SCALARS", "VECTORS", "TENSORS"]

94 }

95 )

97 # not used in regular mesh (implicit node position)

98 # if('POINTS' in pl):

99 # nl.update({'POINTS': int(content[pl['POINTS']].split()[1])})

100 # keys_found.update({'POINTS': [pl['POINTS'], 'VECTORS', 'POINTS']})

101

102 if "CELL_DATA" in pl:

103 nl.update({"CELL_DATA": int(content[pl["CELL_DATA"]].split()[1])})

104 keys_found.update(

105 {

106 l.split()[1]: [

107 pl["CELL_DATA"] + i,

108 content[pl["CELL_DATA"] + i].split()[0],

109 "CELL_DATA",

110 ]

111 for i, l in enumerate(content[pl["CELL_DATA"] :])

112 if l.split()[0] in ["SCALARS", "VECTORS", "TENSORS"]

113 }

114 )

115

116 # not used in regular mesh (implicit connectivity)

117 # if('CELLS' in pl):

118 # nl.update({'CELLS': int(content[pl['CELLS']].split()[1])})

119 # keys_found.update({'CELLS': [pl['CELLS'], 'TETRA', 'CELLS']})

120

121 # get relevant keys

122 if fieldAsked is None:

123 keys_valid = keys_found

124 elif isinstance(fieldAsked, list):

125 keys_valid = {k: keys_found[k] for k in fieldAsked if k in keys_found}

126 elif isinstance(fieldAsked, str) and fieldAsked in keys_found:

127 keys_valid = {fieldAsked: keys_found[fieldAsked]}

128 else:

129 keys_valid = {}

130

131 # fill output

132 fields = {}

133 for key in keys_valid.items():

134 dict_skip = {

135 "CELLS": {"TETRA": 1},

136 "POINTS": {"VECTORS": 1},

137 "CELL_DATA": {"SCALARS": 2, "VECTORS": 1, "TENSORS": 1},

138 "POINT_DATA": {"SCALARS": 2, "VECTORS": 1, "TENSORS": 1},

139 }

140 l_skip = dict_skip[key[1][2]][key[1][1]]

141 l_start = l_skip + key[1][0]

142 l_end = l_start + nl[key[1][2]]

143 # hack for the scalar case to avoir list of single element lists

144 if key[1][1] == "SCALARS":

145 fields.update({key[0]: numpy.array([float(r) for r in content[l_start:l_end]])})

146 else:

147 fields.update({key[0]: numpy.array([[float(c) for c in reversed(r.split())] for r in content[l_start:l_end]])})

148

149 return fields

150

151

152def writeStructuredVTK(

153 aspectRatio=[1.0, 1.0, 1.0],

154 origin=[0.0, 0.0, 0.0],

155 cellData={},

156 pointData={},

157 fileName="spam.vtk",

158):

159 """Write a plain text regular grid vtk from

160

161 - 3D arrays for 3D scalar fields

162 - 4D arrays for 3D vector fields

163

164 Parameters

165 ----------

166 aspectRatio : size 3 array, float

167 Length between two nodes in every direction `e.i.` size of a cell

168 Default = [1, 1, 1]

169

170 origin : size 3 array float

171 Origin of the grid

172 Default = [0, 0, 0]

173

174 cellData : dict ``{"field1name": field1, "field2name": field2, ...}``

175 Cell fields, not interpolated by paraview.

176 The field values are reshaped into a flat array in the lexicographic order.

177 ``field1`` and ``field2`` are ndimensional array

178 (3D arrays are scalar fields and 4D array are vector valued fields).

179

180 pointData : dict ``{"field1name": field1, "field2name": field2, ...}``

181 Nodal fields, interpolated by paraview. ``pointData`` has the same shape as ``cellData``.

182

183 fileName : string

184 Name of the output file.

185 Default = 'spam.vtk'

186

187 WARNING

188 -------

189 This function deals with structured mesh thus ``x`` and ``z`` axis are swapped **in python**.

190 """

191

192 dimensions = []

193

194 # Check dimensions

195 if len(cellData) + len(pointData) == 0:

196 print("spam.helpers.writeStructuredVTK() Empty files. Not writing {}".format(fileName))

197 return 0

198

199 if len(cellData):

200 dimensionsCell = list(cellData.values())[0].shape[:3]

201 for k, v in cellData.items():

202 if set(dimensionsCell) != set(v.shape[:3]):

203 print("spam.helpers.writeStructuredVTK() Inconsistent cell field sizes {} != {}".format(dimensionsCell, v.shape[:3]))

204 return 0

205 dimensions = [n + 1 for n in dimensionsCell]

206

207 if len(pointData):

208 dimensionsPoint = list(pointData.values())[0].shape[:3]

209 for k, v in pointData.items():

210 if set(dimensionsPoint) != set(v.shape[:3]):

211 print("spam.helpers.writeStructuredVTK() Inconsistent point field sizes {} != {}".format(dimensionsPoint, v.shape[:3]))

212 return 0

213 dimensions = dimensionsPoint

214

215 if len(cellData) and len(pointData):

216 if {n + 1 for n in dimensionsCell} != set(dimensionsPoint):

217 print(

218 "spam.helpers.writeStructuredVTK() Inconsistent point VS cell field sizes.\

219 Point size should be +1 for each axis."

220 )

221

222 with open(fileName, "w") as f:

223 # header

224 f.write("# vtk DataFile Version 2.0\n")

225 f.write("VTK file from spam: {}\n".format(fileName))

226 f.write("ASCII\n\n")

227 f.write("DATASET STRUCTURED_POINTS\n")

228

229 f.write("DIMENSIONS {} {} {}\n".format(*reversed(dimensions)))

230 f.write("ASPECT_RATIO {} {} {}\n".format(*reversed(aspectRatio)))

231 f.write("ORIGIN {} {} {}\n\n".format(*reversed(origin)))

232

233 # pointData

234 if len(pointData) == 1:

235 f.write("POINT_DATA {}\n\n".format(dimensions[0] * dimensions[1] * dimensions[2]))

236 _writeFieldInVtk(pointData, f)

237 elif len(pointData) > 1:

238 f.write("POINT_DATA {}\n\n".format(dimensions[0] * dimensions[1] * dimensions[2]))

239 for k in pointData:

240 _writeFieldInVtk({k: pointData[k]}, f)

241

242 # cellData

243 if len(cellData) == 1:

244 f.write("CELL_DATA {}\n\n".format((dimensions[0] - 1) * (dimensions[1] - 1) * (dimensions[2] - 1)))

245 _writeFieldInVtk(cellData, f)

246 elif len(cellData) > 1:

247 f.write("CELL_DATA {}\n\n".format((dimensions[0] - 1) * (dimensions[1] - 1) * (dimensions[2] - 1)))

248 for k in cellData:

249 _writeFieldInVtk({k: cellData[k]}, f)

250

251 f.write("\n")

252

253

254def writeGlyphsVTK(coordinates, pointData, fileName="spam.vtk"):

255 """

256 Write a plain text glyphs vtk.

257

258 Parameters

259 ----------

260 coordinates : n by 3 array of float

261 Coordinates of the centre of all n glyphs

262

263 pointData : dict ``{"field1name": field1, "field2name": field2, ...}``

264 Value attached to each glyph.

265 ``field1`` and ``field2`` are n by 1 arrays for scalar values and n by 3 for vector values.

266

267 fileName : string, default='spam.vtk'

268 Name of the output file.

269

270 """

271 # WARNING

272 # -------

273 # This function deals with structured mesh thus ``x`` and ``z`` axis are swapped **in python**.

274

275 # check dimensions

276

277 dimension = coordinates.shape[0]

278

279 if len(pointData):

280 for k, v in pointData.items():

281 if dimension != v.shape[0]:

282 print("spam.helpers.writeGlyphsVTK() Inconsistent point field sizes {} != {}".format(dimension, v.shape[0]))

283 return 0

284 else:

285 print("spam.helpers.writeGlyphsVTK() Empty files. Not writing {}".format(fileName))

286 return

287

288 with open(fileName, "w") as f:

289 # header

290 f.write("# vtk DataFile Version 2.0\n")

291 # f.write('VTK file from spam: {}\n'.format(fileName))

292 f.write("Unstructured grid legacy vtk file with point scalar data\n")

293 f.write("ASCII\n\n")

294

295 # coordinates

296 f.write("DATASET UNSTRUCTURED_GRID\n")

297 f.write("POINTS {:.0f} float\n".format(dimension))

298 for coord in coordinates:

299 f.write(" {} {} {}\n".format(*reversed(coord)))

300

301 f.write("\n")

302

303 # pointData

304 if len(pointData) == 1:

305 f.write("POINT_DATA {}\n\n".format(dimension))

306 _writeFieldInVtk(pointData, f, flat=True)

307 elif len(pointData) > 1:

308 f.write("POINT_DATA {}\n\n".format(dimension))

309 for k in pointData:

310 _writeFieldInVtk({k: pointData[k]}, f, flat=True)

311

312 f.write("\n")

313

314

315def _writeFieldInVtk(data, f, flat=False):

316 """

317 Private helper function for writing vtk fields

318 """

319

320 for key in data:

321 field = data[key]

322

323 if flat:

324 # SCALAR flatten (n by 1)

325 if len(field.shape) == 1:

326 f.write("SCALARS {} float\n".format(key.replace(" ", "_")))

327 f.write("LOOKUP_TABLE default\n")

328 for item in field:

329 f.write(" {}\n".format(item))

330 f.write("\n")

331

332 # VECTORS flatten (n by 3)

333 elif len(field.shape) == 2 and field.shape[1] == 3:

334 f.write("VECTORS {} float\n".format(key.replace(" ", "_")))

335 for item in field:

336 f.write(" {} {} {}\n".format(*reversed(item)))

337 f.write("\n")

338

339 else:

340 # SCALAR not flatten (n1 by n2 by n3)

341 if len(field.shape) == 3:

342 f.write("SCALARS {} float\n".format(key.replace(" ", "_")))

343 f.write("LOOKUP_TABLE default\n")

344 for item in field.reshape(-1):

345 f.write(" {}\n".format(item))

346 f.write("\n")

347

348 # VECTORS (n1 by n2 by n3 by 3)

349 elif len(field.shape) == 4 and field.shape[3] == 3:

350 f.write("VECTORS {} float\n".format(key.replace(" ", "_")))

351 for item in field.reshape((field.shape[0] * field.shape[1] * field.shape[2], field.shape[3])):

352 f.write(" {} {} {}\n".format(*reversed(item)))

353 f.write("\n")

354

355 # TENSORS (n1 by n2 by n3 by 3 by 3)

356 elif len(field.shape) == 5 and field.shape[3] * field.shape[4] == 9:

357 f.write("TENSORS {} float\n".format(key.replace(" ", "_")))

358 for item in field.reshape(

359 (

360 field.shape[0] * field.shape[1] * field.shape[2],

361 field.shape[3] * field.shape[4],

362 )

363 ):

364 f.write(" {} {} {}\n {} {} {}\n {} {} {}\n\n".format(*reversed(item)))

365 f.write("\n")

366 else:

367 print("spam.helpers.vtkio._writeFieldInVtk(): I'm in an unknown condition!")

368

369

370# Unstructured meshes

371def writeUnstructuredVTK(

372 points,

373 connectivity,

374 elementType="tetra",

375 pointData={},

376 cellData={},

377 fileName="spam.vtk",

378):

379 """Writes a binary VTK using ``meshio`` selecting only the tetrahedra.

380

381 Parameters

382 ----------

383 points : 2D numpy array

384 The coordinates of the mesh nodes (zyx)

385 Each line is [zPos, yPos, xPos]

386

387 connectivity : 2D numpy array

388 The connectivity matrix of the tetrahedra elements

389 Each line is [node1, node2, node3, node4]

390

391 elementType : string, optional, default="tetra"

392 The type of element used for the mesh

393

394 cellData : dict, optional, default={}

395 Cell fields, not interpolated by paraview.

396 With ``field1`` and ``field2`` as 1D or 2D arrays. 1D arrays are scalar fields

397 and 2D array are vector valued fields.

398

399 .. code-block:: python

400

401 cellData = {

402 "field1name": field1,

403 "field2name": field2,

404 # ...

405 }

406

407 pointData : dict, optional, default={}

408 Nodal fields, interpolated by paraview. ``pointData`` has the same shape as ``cellData``.

409

410 .. code-block:: python

411

412 pointData = {

413 "field1name": field1,

414 "field2name": field2,

415 # ...

416 }

417

418 fileName : string, optional, default="spam.vtk"

419 Name of the output file.

420

421 Note

422 -----

423 VTK files are assumed to be in X,Y,Z format, and numpy arrays are in Z,Y,X, so the following

424 policy is follwed:

425 - position arrays are simply reversed so Z comes out first

426 - scalar arrays are left as they are

427 - vectors and matrices are assumed to be positional (a surface vector for example) and are reversed

428 - node numbering in connectivity matrix is changed to ensure a positive Jacobian for each element

429

430 """

431

432 # Anything bigger than a scalar field in pointData should be flipped

433 for key, value in pointData.items():

434 # print("writeUnstructuredVTK(): pointData",key,pointData[key].shape)

435 if len(pointData[key].shape) > 1:

436 if pointData[key].shape[1] > 1:

437 pointData[key] = pointData[key][:, ::-1]

438

439 # Anything bigger than a scalar field in cellData should be flipped

440 for key, value in cellData.items():

441 # print("writeUnstructuredVTK(): cellData",key,cellData[key].shape)

442 if len(cellData[key].shape) > 1:

443 if cellData[key].shape[1] > 1:

444 cellData[key] = cellData[key][:, ::-1]

445

446 # change node numbering in connectivity matrix for tetrahedra elements

447 # if elementType == "tetra":

448 # tmp = connectivity.copy()

449 # connectivity[:, 1] = tmp[:, 3]

450 # connectivity[:, 3] = tmp[:, 1]

451

452 meshio.write_points_cells(

453 fileName,

454 points[:, ::-1],

455 {elementType: connectivity},

456 point_data=pointData if len(pointData) else None,

457 cell_data={k: [v] for k, v in cellData.items()} if len(cellData) else None,

458 )

459

460

461def readUnstructuredVTK(fileName):

462 """

463 Read a binary VTK using ``meshio`` selecting only the tetrahedra.

464

465 Parameters

466 ----------

467 fileName : string

468 Name of the input file.

469

470 Returns

471 -------

472 numpy array

473 The list of node coordinates (zyx).

474

475 numpy array

476 The connectivity matrix (cell id starts at 0).

477

478 dict

479 Nodal fields, interpolated by paraview.

480

481 dict

482 Cell fields, not interpolated by paraview.

483

484 """

485 # VTK files are assumed to be in X,Y,Z format, and numpy arrays are in Z,Y,X, so the following

486 # policy is follwed:

487 # - position arrays are simply reversed so Z comes out first

488 # - scalar arrays are left as they are

489 # - vectors and matrices are assumed to be positional (a surface vector for example) and are reversed

490 # - node numbering in connectivity matrix is changed to ensure a positive Jacobian for each element

491

492 mesh = meshio.read(fileName)

493

494 # Reverse points order into z-first

495 points = mesh.points[:, ::-1]

496 cellData = dict({})

497 pointData = dict({})

498

499 # Make sure that any non-scalar fields are flipped so that z is in the right place in numpy

500 for key, values in mesh.cell_data.items():

501 for value in values:

502 # Inspect first item and make sure it's nut just a number

503 if isinstance(value[0], (list, tuple, numpy.ndarray)):

504 cellData[key] = value[:, ::-1]

505 else:

506 cellData[key] = value

507

508 # Make sure that any non-scalar fields are flipped so that z is in the right place in numpy

509 for key, value in mesh.point_data.items():

510 # Inspect first item and make sure it's nut just a number

511 if isinstance(value[0], (list, tuple, numpy.ndarray)):

512 pointData[key] = value[:, ::-1]

513 else:

514 pointData[key] = value

515

516 connectivity = mesh.cells_dict["tetra"]

517 # change node numbering in connectivity matrix

518 tmp = connectivity.copy()

519 connectivity[:, 1] = tmp[:, 3]

520 connectivity[:, 3] = tmp[:, 1]

521

522 return points, connectivity, pointData, cellData

523

524

525def TIFFtoVTK(fileName, voxelSize=1.0):

526 """

527 Convert a tifffile to a vtk file for paraview.

528 It saves it with the same base name.

529

530 Parameters

531 ----------

532 fileName : string

533 The name of the tif file to convert

534 """

535 import os

536

537 import spam.mesh

538 import tifffile

539

540 # convert to array [fileName] if single name

541 fileName = fileName if isinstance(fileName, (list, tuple)) else [fileName]

542

543 # read the tifffile

544 im = tifffile.imread(fileName[0])

545 print("spam.helpers.TIFFtoVTK: image shape: ", im.shape)

546 nCells = numpy.array(im.shape)

547

548 # create coordinates

549 points = spam.mesh.createLexicoCoordinates([voxelSize * n for n in nCells], [n + 1 for n in nCells])

550 print("spam.helpers.TIFFtoVTK: points shape: ", points.shape)

551 print("spam.helpers.TIFFtoVTK: 0 min max: ", points[:, 0].min(), points[:, 0].max())

552 print("spam.helpers.TIFFtoVTK: 1 min max: ", points[:, 1].min(), points[:, 1].max())

553 print("spam.helpers.TIFFtoVTK: 2 min max: ", points[:, 2].min(), points[:, 2].max())

554

555 # create connectivity

556 cells = _loop_for_tifToVTK(nCells)

557 print("spam.helpers.TIFFtoVTK: cells shape: ", cells.shape)

558

559 # create VTK

560 f = os.path.splitext(fileName[0])[0]

561 if len(fileName) == 1:

562 meshio.write_points_cells(

563 f"{f}.vtk",

564 points,

565 {"hexahedron": cells},

566 cell_data={"grey": [im.T.ravel()]},

567 )

568 else:

569 with meshio.xdmf.TimeSeriesWriter(f"{f}.xmf") as writer:

570 writer.write_points_cells(points, [("hexahedron", cells)])

571 for i, name in enumerate(fileName):

572 im = tifffile.imread(name)

573 writer.write_data(i, cell_data={"grey": [im.T.ravel()]})

574

575

576# @jit(nopython=True)

577def _loop_for_tifToVTK(nCells):

578 cells = numpy.zeros((numpy.prod(nCells), 8), dtype=numpy.uint)

579 nx, ny, nz = nCells

580 i = 0

581 for z in range(nz):

582 for y in range(ny):

583 for x in range(nx):

584 n = x + ((nx + 1) * y) + ((nx + 1) * (ny + 1)) * z

585 cells[i][0] = n

586 cells[i][1] = n + 1

587 cells[i][2] = n + nx + 2

588 cells[i][3] = n + nx + 1

589 for j in range(4):

590 cells[i][4 + j] = cells[i][j] + ((nx + 1) * (ny + 1))

591 i += 1

592 return cells

593

594

595# if __name__ == "__main__":

596# tifToVTK(["/home/roubin/nx/xr_bin16.tif", "/home/roubin/nx/ne_bin16.tif"])