Skip to content

Commit 6f8c961

Browse files
glesurglesur
committed
Merge branch 'volodia99-vtk-native-coords' into develop
ENH: write and read native coordinates in VTK (#292)
2 parents 9143e77 + f2d83c7 commit 6f8c961

3 files changed

Lines changed: 133 additions & 10 deletions

File tree

‎pytools/vtk_io.py‎

Lines changed: 33 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@
77
import warnings
88
import numpy as np
99
import os
10-
10+
import re
1111

1212
# restrict what's included with `import *` to public API
1313
__all__ = [
@@ -21,6 +21,8 @@
2121
dt = np.dtype(">f") # Big endian single precision floats
2222
dint = np.dtype(">i4") # Big endian integer
2323

24+
NATIVE_COORDINATE_REGEXP = re.compile(r"X(1|2|3)(L|C)_NATIVE_COORDINATES")
25+
2426
KNOWN_GEOMETRIES = {
2527
0: "cartesian",
2628
1: "polar",
@@ -33,10 +35,14 @@ class VTKDataset(object):
3335
def __init__(self, filename, geometry=None):
3436
self.filename = os.path.abspath(filename)
3537
self.data = {}
38+
self.native_coordinates = {}
3639
with open(filename, "rb") as fh:
3740
self._load_header(fh, geometry=geometry)
3841
self._load(fh)
3942

43+
if self.native_coordinates:
44+
self._setup_coordinates_from_native()
45+
4046
def _load(self, fh):
4147
if self._dataset_type in ("RECTILINEAR_GRID", "STRUCTURED_GRID"):
4248
self._load_hydro(fh)
@@ -88,6 +94,9 @@ def _load_header(self, fh, geometry=None):
8894
self.t = np.fromfile(fh, dt, 1)
8995
elif d.startswith("PERIODICITY"):
9096
self.periodicity = np.fromfile(fh, dtype=dint, count=3).astype(bool)
97+
elif NATIVE_COORDINATE_REGEXP.match(d):
98+
native_name, _ncomp, native_dim, _dtype = d.split()
99+
self.native_coordinates[native_name] = np.fromfile(fh, dtype=dt, count=int(native_dim))
91100
else:
92101
warnings.warn("Found unknown field %s" % d)
93102
fh.readline() # skip extra linefeed (empty line)
@@ -341,6 +350,29 @@ def _load_hydro(self, fh):
341350
def _load_particles(self, fh):
342351
raise NotImplementedError("Particles vtk are not supported yet !")
343352

353+
def _setup_coordinates_from_native(self):
354+
if self.geometry == "spherical":
355+
native2attr = {
356+
"X1L_NATIVE_COORDINATES": "rl",
357+
"X1C_NATIVE_COORDINATES": "r",
358+
"X2L_NATIVE_COORDINATES": "thetal",
359+
"X2C_NATIVE_COORDINATES": "theta",
360+
"X3L_NATIVE_COORDINATES": "phil",
361+
"X3C_NATIVE_COORDINATES": "phi",
362+
}
363+
elif self.geometry in ("cartesian", "cylindrical", "polar"):
364+
native2attr = {
365+
"X1L_NATIVE_COORDINATES": "xl",
366+
"X1C_NATIVE_COORDINATES": "x",
367+
"X2L_NATIVE_COORDINATES": "yl",
368+
"X2C_NATIVE_COORDINATES": "y",
369+
"X3L_NATIVE_COORDINATES": "zl",
370+
"X3C_NATIVE_COORDINATES": "z",
371+
}
372+
373+
for native_field, attr in native2attr.items():
374+
setattr(self, attr, self.native_coordinates[native_field])
375+
344376
def __repr__(self):
345377
return "VTKDataset('%s')" % self.filename
346378

‎src/output/vtk.cpp‎

Lines changed: 99 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -124,24 +124,51 @@ Vtk::Vtk(Input &input, DataBlock *datain, std::string filebase) {
124124
this->xnode = new float[nx1+ioffset];
125125
this->ynode = new float[nx2+joffset];
126126
this->znode = new float[nx3+koffset];
127+
this->xcenter = new float[nx1];
128+
this->ycenter = new float[nx2];
129+
this->zcenter = new float[nx3];
127130

128131
for (int32_t i = 0; i < nx1 + ioffset; i++) {
129-
if(grid.np_tot[IDIR] == 1) // only one dimension in this direction
132+
if(grid.np_tot[IDIR] == 1) { // only one dimension in this direction
130133
xnode[i] = bigEndian(static_cast<float>(grid.x[IDIR](i)));
131-
else
134+
} else {
132135
xnode[i] = bigEndian(static_cast<float>(grid.xl[IDIR](i + grid.nghost[IDIR])));
136+
}
133137
}
134138
for (int32_t j = 0; j < nx2 + joffset; j++) {
135-
if(grid.np_tot[JDIR] == 1) // only one dimension in this direction
139+
if(grid.np_tot[JDIR] == 1) { // only one dimension in this direction
136140
ynode[j] = bigEndian(static_cast<float>(grid.x[JDIR](j)));
137-
else
141+
} else {
138142
ynode[j] = bigEndian(static_cast<float>(grid.xl[JDIR](j + grid.nghost[JDIR])));
143+
}
139144
}
140145
for (int32_t k = 0; k < nx3 + koffset; k++) {
141-
if(grid.np_tot[KDIR] == 1)
146+
if(grid.np_tot[KDIR] == 1) {
142147
znode[k] = bigEndian(static_cast<float>(grid.x[KDIR](k)));
143-
else
148+
} else {
144149
znode[k] = bigEndian(static_cast<float>(grid.xl[KDIR](k + grid.nghost[KDIR])));
150+
}
151+
}
152+
for (int32_t i = 0; i < nx1; i++) {
153+
if(grid.np_tot[IDIR] == 1) { // only one dimension in this direction
154+
xcenter[i] = xnode[i];
155+
} else {
156+
xcenter[i] = bigEndian(static_cast<float>(grid.x[IDIR](i + grid.nghost[IDIR])));
157+
}
158+
}
159+
for (int32_t j = 0; j < nx2; j++) {
160+
if(grid.np_tot[JDIR] == 1) { // only one dimension in this direction
161+
ycenter[j] = ynode[j];
162+
} else {
163+
ycenter[j] = bigEndian(static_cast<float>(grid.x[JDIR](j + grid.nghost[JDIR])));
164+
}
165+
}
166+
for (int32_t k = 0; k < nx3; k++) {
167+
if(grid.np_tot[KDIR] == 1) {
168+
zcenter[k] = znode[k];
169+
} else {
170+
zcenter[k] = bigEndian(static_cast<float>(grid.x[KDIR](k + grid.nghost[KDIR])));
171+
}
145172
}
146173
#if VTK_FORMAT == VTK_STRUCTURED_GRID // VTK_FORMAT
147174
/* -- Allocate memory for node_coord which is later used -- */
@@ -375,8 +402,8 @@ void Vtk::WriteHeader(IdfxFileHandler fvtk, real time) {
375402
#elif VTK_FORMAT == VTK_STRUCTURED_GRID
376403
ssheader << "DATASET STRUCTURED_GRID" << std::endl;
377404
#endif
378-
// fields: geometry, periodicity, time
379-
int nfields = 3;
405+
// fields: geometry, periodicity, time, 6 NativeCoordinates (x1l, x2l, x3l, x1c, x2c, x3c)
406+
int nfields = 9;
380407

381408
// Write grid geometry in the VTK file
382409
ssheader << "FIELD FieldData " << nfields << std::endl;
@@ -426,7 +453,71 @@ void Vtk::WriteHeader(IdfxFileHandler fvtk, real time) {
426453
// Done, add cariage return for next ascii write
427454
ssheader << std::endl;
428455

456+
// write x1l native coordinates
457+
ssheader << "X1L_NATIVE_COORDINATES 1 " << (nx1 + ioffset) << " float" << std::endl;
458+
// Flush the ascii header
459+
header = ssheader.str();
460+
WriteHeaderString(header.c_str(), fvtk);
461+
// reset the string stream
462+
ssheader.str(std::string());
463+
WriteHeaderBinary(this->xnode, nx1 + ioffset, fvtk);
464+
// Done, add cariage return for next ascii write
465+
ssheader << std::endl;
466+
467+
// write x2l native coordinates
468+
ssheader << "X2L_NATIVE_COORDINATES 1 " << (nx2 + joffset) << " float" << std::endl;
469+
// Flush the ascii header
470+
header = ssheader.str();
471+
WriteHeaderString(header.c_str(), fvtk);
472+
// reset the string stream
473+
ssheader.str(std::string());
474+
WriteHeaderBinary(this->ynode, nx2 + joffset, fvtk);
475+
// Done, add cariage return for next ascii write
476+
ssheader << std::endl;
477+
478+
// write x3l native coordinates
479+
ssheader << "X3L_NATIVE_COORDINATES 1 " << (nx3 + koffset) << " float" << std::endl;
480+
// Flush the ascii header
481+
header = ssheader.str();
482+
WriteHeaderString(header.c_str(), fvtk);
483+
// reset the string stream
484+
ssheader.str(std::string());
485+
WriteHeaderBinary(this->znode, nx3 + koffset, fvtk);
486+
// Done, add cariage return for next ascii write
487+
ssheader << std::endl;
488+
489+
// write x1 native coordinates
490+
ssheader << "X1C_NATIVE_COORDINATES 1 " << nx1 << " float" << std::endl;
491+
// Flush the ascii header
492+
header = ssheader.str();
493+
WriteHeaderString(header.c_str(), fvtk);
494+
// reset the string stream
495+
ssheader.str(std::string());
496+
WriteHeaderBinary(this->xcenter, nx1, fvtk);
497+
// Done, add cariage return for next ascii write
498+
ssheader << std::endl;
499+
500+
// write x2 native coordinates
501+
ssheader << "X2C_NATIVE_COORDINATES 1 " << nx2 << " float" << std::endl;
502+
// Flush the ascii header
503+
header = ssheader.str();
504+
WriteHeaderString(header.c_str(), fvtk);
505+
// reset the string stream
506+
ssheader.str(std::string());
507+
WriteHeaderBinary(this->ycenter, nx2, fvtk);
508+
// Done, add cariage return for next ascii write
509+
ssheader << std::endl;
429510

511+
// write x3 native coordinates
512+
ssheader << "X3C_NATIVE_COORDINATES 1 " << nx3 << " float" << std::endl;
513+
// Flush the ascii header
514+
header = ssheader.str();
515+
WriteHeaderString(header.c_str(), fvtk);
516+
// reset the string stream
517+
ssheader.str(std::string());
518+
WriteHeaderBinary(this->zcenter, nx3, fvtk);
519+
// Done, add cariage return for next ascii write
520+
ssheader << std::endl;
430521

431522
ssheader << "DIMENSIONS " << nx1 + ioffset << " " << nx2 + joffset << " " << nx3 + koffset
432523
<< std::endl;
@@ -486,7 +577,6 @@ void Vtk::WriteHeader(IdfxFileHandler fvtk, real time) {
486577
#undef VTK_STRUCTURED_GRID
487578
#undef VTK_RECTILINEAR_GRID
488579

489-
490580
/* ********************************************************************* */
491581
void Vtk::WriteScalar(IdfxFileHandler fvtk, float* Vin, const std::string &var_name) {
492582
/*!

‎src/output/vtk.hpp‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -129,6 +129,7 @@ class Vtk : public BaseVtk {
129129

130130
// Coordinates needed by VTK outputs
131131
float *xnode, *ynode, *znode;
132+
float *xcenter, *ycenter, *zcenter;
132133

133134
IdefixHostArray4D<float> node_coord;
134135

0 commit comments

Comments
 (0)