diff --git a/pytools/vtk_io.py b/pytools/vtk_io.py index 1f693318..000a9b63 100644 --- a/pytools/vtk_io.py +++ b/pytools/vtk_io.py @@ -59,6 +59,7 @@ def _load_header(self, fh, geometry=None): self._dataset_type = line.split()[1] self.geometry = geometry + self._has_native_coordinates = False ref_position = fh.tell() line = fh.readline().decode("utf-8") # DIMENSIONS NX NY NZ or FIELD @@ -88,6 +89,8 @@ 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 d.startswith("HAS_NATIVE_COORDINATES"): + self._has_native_coordinates = np.fromfile(fh, dtype=dint, count=1).astype(bool) else: warnings.warn("Found unknown field %s" % d) fh.readline() # skip extra linefeed (empty line) @@ -338,6 +341,29 @@ def _load_hydro(self, fh): fh.readline() # extra line feed + if self._has_native_coordinates: + if self.geometry == "spherical": + native2attr = { + "X1L_NATIVE_COORDS": "rl", + "X1C_NATIVE_COORDS": "r", + "X2L_NATIVE_COORDS": "thetal", + "X2C_NATIVE_COORDS": "theta", + "X3L_NATIVE_COORDS": "phil", + "X3C_NATIVE_COORDS": "phi", + } + elif self.geometry in ("cartesian", "cylindrical", "polar"): + native2attr = { + "X1L_NATIVE_COORDS": "xl", + "X1C_NATIVE_COORDS": "x", + "X2L_NATIVE_COORDS": "yl", + "X2C_NATIVE_COORDS": "y", + "X3L_NATIVE_COORDS": "zl", + "X3C_NATIVE_COORDS": "z", + } + + for native_field, attr in native2attr.items(): + setattr(self, attr, self.data[native_field]) + def _load_particles(self, fh): raise NotImplementedError("Particles vtk are not supported yet !")