Skip to content

Commit

Permalink
Merge branch 'volodia99-vtk-native-coords' into develop
Browse files Browse the repository at this point in the history 
ENH: write and read native coordinates in VTK (#292)
  • Loading branch information
@glesur
glesur committed Nov 26, 2024
2 parents 9143e77 + f2d83c7 commit 6f8c961
Show file tree
Hide file tree
Showing 3 changed files with 133 additions and 10 deletions.
34 changes: 33 additions & 1 deletion pytools/vtk_io.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@
import warnings
import numpy as np
import os

import re

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

NATIVE_COORDINATE_REGEXP = re.compile(r"X(1|2|3)(L|C)_NATIVE_COORDINATES")

KNOWN_GEOMETRIES = {
0: "cartesian",
1: "polar",
Expand All @@ -33,10 +35,14 @@ class VTKDataset(object):
def __init__(self, filename, geometry=None):
self.filename = os.path.abspath(filename)
self.data = {}
self.native_coordinates = {}
with open(filename, "rb") as fh:
self._load_header(fh, geometry=geometry)
self._load(fh)

if self.native_coordinates:
self._setup_coordinates_from_native()

def _load(self, fh):
if self._dataset_type in ("RECTILINEAR_GRID", "STRUCTURED_GRID"):
self._load_hydro(fh)
Expand Down Expand Up @@ -88,6 +94,9 @@ def _load_header(self, fh, geometry=None):
self.t = np.fromfile(fh, dt, 1)
elif d.startswith("PERIODICITY"):
self.periodicity = np.fromfile(fh, dtype=dint, count=3).astype(bool)
elif NATIVE_COORDINATE_REGEXP.match(d):
native_name, _ncomp, native_dim, _dtype = d.split()
self.native_coordinates[native_name] = np.fromfile(fh, dtype=dt, count=int(native_dim))
else:
warnings.warn("Found unknown field %s" % d)
fh.readline() # skip extra linefeed (empty line)
Expand Down Expand Up @@ -341,6 +350,29 @@ def _load_hydro(self, fh):
def _load_particles(self, fh):
raise NotImplementedError("Particles vtk are not supported yet !")

def _setup_coordinates_from_native(self):
if self.geometry == "spherical":
native2attr = {
"X1L_NATIVE_COORDINATES": "rl",
"X1C_NATIVE_COORDINATES": "r",
"X2L_NATIVE_COORDINATES": "thetal",
"X2C_NATIVE_COORDINATES": "theta",
"X3L_NATIVE_COORDINATES": "phil",
"X3C_NATIVE_COORDINATES": "phi",
}
elif self.geometry in ("cartesian", "cylindrical", "polar"):
native2attr = {
"X1L_NATIVE_COORDINATES": "xl",
"X1C_NATIVE_COORDINATES": "x",
"X2L_NATIVE_COORDINATES": "yl",
"X2C_NATIVE_COORDINATES": "y",
"X3L_NATIVE_COORDINATES": "zl",
"X3C_NATIVE_COORDINATES": "z",
}

for native_field, attr in native2attr.items():
setattr(self, attr, self.native_coordinates[native_field])

def __repr__(self):
return "VTKDataset('%s')" % self.filename

Expand Down
108 changes: 99 additions & 9 deletions src/output/vtk.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -124,24 +124,51 @@ Vtk::Vtk(Input &input, DataBlock *datain, std::string filebase) {
this->xnode = new float[nx1+ioffset];
this->ynode = new float[nx2+joffset];
this->znode = new float[nx3+koffset];
this->xcenter = new float[nx1];
this->ycenter = new float[nx2];
this->zcenter = new float[nx3];

for (int32_t i = 0; i < nx1 + ioffset; i++) {
if(grid.np_tot[IDIR] == 1) // only one dimension in this direction
if(grid.np_tot[IDIR] == 1) { // only one dimension in this direction
xnode[i] = bigEndian(static_cast<float>(grid.x[IDIR](i)));
else
} else {
xnode[i] = bigEndian(static_cast<float>(grid.xl[IDIR](i + grid.nghost[IDIR])));
}
}
for (int32_t j = 0; j < nx2 + joffset; j++) {
if(grid.np_tot[JDIR] == 1) // only one dimension in this direction
if(grid.np_tot[JDIR] == 1) { // only one dimension in this direction
ynode[j] = bigEndian(static_cast<float>(grid.x[JDIR](j)));
else
} else {
ynode[j] = bigEndian(static_cast<float>(grid.xl[JDIR](j + grid.nghost[JDIR])));
}
}
for (int32_t k = 0; k < nx3 + koffset; k++) {
if(grid.np_tot[KDIR] == 1)
if(grid.np_tot[KDIR] == 1) {
znode[k] = bigEndian(static_cast<float>(grid.x[KDIR](k)));
else
} else {
znode[k] = bigEndian(static_cast<float>(grid.xl[KDIR](k + grid.nghost[KDIR])));
}
}
for (int32_t i = 0; i < nx1; i++) {
if(grid.np_tot[IDIR] == 1) { // only one dimension in this direction
xcenter[i] = xnode[i];
} else {
xcenter[i] = bigEndian(static_cast<float>(grid.x[IDIR](i + grid.nghost[IDIR])));
}
}
for (int32_t j = 0; j < nx2; j++) {
if(grid.np_tot[JDIR] == 1) { // only one dimension in this direction
ycenter[j] = ynode[j];
} else {
ycenter[j] = bigEndian(static_cast<float>(grid.x[JDIR](j + grid.nghost[JDIR])));
}
}
for (int32_t k = 0; k < nx3; k++) {
if(grid.np_tot[KDIR] == 1) {
zcenter[k] = znode[k];
} else {
zcenter[k] = bigEndian(static_cast<float>(grid.x[KDIR](k + grid.nghost[KDIR])));
}
}
#if VTK_FORMAT == VTK_STRUCTURED_GRID // VTK_FORMAT
/* -- Allocate memory for node_coord which is later used -- */
Expand Down Expand Up @@ -375,8 +402,8 @@ void Vtk::WriteHeader(IdfxFileHandler fvtk, real time) {
#elif VTK_FORMAT == VTK_STRUCTURED_GRID
ssheader << "DATASET STRUCTURED_GRID" << std::endl;
#endif
// fields: geometry, periodicity, time
int nfields = 3;
// fields: geometry, periodicity, time, 6 NativeCoordinates (x1l, x2l, x3l, x1c, x2c, x3c)
int nfields = 9;

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

// write x1l native coordinates
ssheader << "X1L_NATIVE_COORDINATES 1 " << (nx1 + ioffset) << " float" << std::endl;
// Flush the ascii header
header = ssheader.str();
WriteHeaderString(header.c_str(), fvtk);
// reset the string stream
ssheader.str(std::string());
WriteHeaderBinary(this->xnode, nx1 + ioffset, fvtk);
// Done, add cariage return for next ascii write
ssheader << std::endl;

// write x2l native coordinates
ssheader << "X2L_NATIVE_COORDINATES 1 " << (nx2 + joffset) << " float" << std::endl;
// Flush the ascii header
header = ssheader.str();
WriteHeaderString(header.c_str(), fvtk);
// reset the string stream
ssheader.str(std::string());
WriteHeaderBinary(this->ynode, nx2 + joffset, fvtk);
// Done, add cariage return for next ascii write
ssheader << std::endl;

// write x3l native coordinates
ssheader << "X3L_NATIVE_COORDINATES 1 " << (nx3 + koffset) << " float" << std::endl;
// Flush the ascii header
header = ssheader.str();
WriteHeaderString(header.c_str(), fvtk);
// reset the string stream
ssheader.str(std::string());
WriteHeaderBinary(this->znode, nx3 + koffset, fvtk);
// Done, add cariage return for next ascii write
ssheader << std::endl;

// write x1 native coordinates
ssheader << "X1C_NATIVE_COORDINATES 1 " << nx1 << " float" << std::endl;
// Flush the ascii header
header = ssheader.str();
WriteHeaderString(header.c_str(), fvtk);
// reset the string stream
ssheader.str(std::string());
WriteHeaderBinary(this->xcenter, nx1, fvtk);
// Done, add cariage return for next ascii write
ssheader << std::endl;

// write x2 native coordinates
ssheader << "X2C_NATIVE_COORDINATES 1 " << nx2 << " float" << std::endl;
// Flush the ascii header
header = ssheader.str();
WriteHeaderString(header.c_str(), fvtk);
// reset the string stream
ssheader.str(std::string());
WriteHeaderBinary(this->ycenter, nx2, fvtk);
// Done, add cariage return for next ascii write
ssheader << std::endl;

// write x3 native coordinates
ssheader << "X3C_NATIVE_COORDINATES 1 " << nx3 << " float" << std::endl;
// Flush the ascii header
header = ssheader.str();
WriteHeaderString(header.c_str(), fvtk);
// reset the string stream
ssheader.str(std::string());
WriteHeaderBinary(this->zcenter, nx3, fvtk);
// Done, add cariage return for next ascii write
ssheader << std::endl;

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


/* ********************************************************************* */
void Vtk::WriteScalar(IdfxFileHandler fvtk, float* Vin, const std::string &var_name) {
/*!
Expand Down
1 change: 1 addition & 0 deletions src/output/vtk.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -129,6 +129,7 @@ class Vtk : public BaseVtk {

// Coordinates needed by VTK outputs
float *xnode, *ynode, *znode;
float *xcenter, *ycenter, *zcenter;

IdefixHostArray4D<float> node_coord;

Expand Down

0 comments on commit 6f8c961

Please sign in to comment.