use std::io::Write;
use energyml_geometry::{Attribute, AttributeScope, Placement, SceneObject, Topology};
use crate::error::{Error, Result};
use crate::registry::{Exporter, Options};
#[derive(Clone, Copy, PartialEq, Eq, Debug)]
pub enum Flavour {
Legacy,
PolyData,
UnstructuredGrid,
}
pub struct VtkExporter {
flavour: Flavour,
version: &'static str,
}
impl VtkExporter {
pub fn legacy() -> Self {
Self {
flavour: Flavour::Legacy,
version: "3.0",
}
}
pub fn poly_data() -> Self {
Self {
flavour: Flavour::PolyData,
version: "1.0",
}
}
pub fn unstructured_grid() -> Self {
Self {
flavour: Flavour::UnstructuredGrid,
version: "1.0",
}
}
pub fn with_version(mut self, version: &'static str) -> Self {
self.version = version;
self
}
pub fn flavour(&self) -> Flavour {
self.flavour
}
}
impl Exporter for VtkExporter {
fn name(&self) -> &'static str {
match self.flavour {
Flavour::Legacy => "vtk",
Flavour::PolyData => "vtp",
Flavour::UnstructuredGrid => "vtu",
}
}
fn extension(&self) -> &'static str {
self.name()
}
fn mime(&self) -> &'static str {
match self.flavour {
Flavour::Legacy => "application/octet-stream",
_ => "application/xml",
}
}
fn describe(&self) -> &'static str {
match self.flavour {
Flavour::Legacy => "VTK legacy ASCII, DATASET POLYDATA",
Flavour::PolyData => "VTK XML PolyData (.vtp), surfaces and lines",
Flavour::UnstructuredGrid => "VTK XML UnstructuredGrid (.vtu), the general container",
}
}
fn placement(&self) -> Placement {
Placement::Metres
}
fn handles(&self, object: &SceneObject) -> bool {
!object.geometry.is_empty()
}
fn write(&self, object: &SceneObject, out: &mut dyn Write, options: &Options<'_>) -> Result<()> {
if object.geometry.is_empty() {
return Err(Error::unsuitable(
"vtk",
format!("`{}`, which has no points", object.title),
));
}
let object = crate::registry::prepare(object, options)?;
let object = object.as_ref();
let placement = options.placement.unwrap_or_else(|| self.placement());
let points = options.points(object, placement)?;
let cells = cells_of(object);
match self.flavour {
Flavour::Legacy => self.legacy_document(object, &points, &cells, options, out),
_ => self.xml_document(object, &points, &cells, options, out),
}
}
fn write_all(&self, objects: &[SceneObject], out: &mut dyn Write, options: &Options<'_>) -> Result<()> {
let writable: Vec<&SceneObject> = objects.iter().filter(|o| self.handles(o)).collect();
match writable.len() {
0 => Ok(()),
1 => self.write(writable[0], out, options),
n => Err(Error::unsuitable(
"vtk",
format!(
"{n} objects in one file - a VTK file holds one dataset; \
write to a directory to get one file each"
),
)),
}
}
}
struct Cells {
triangles: Vec<[u32; 3]>,
polygons: Vec<Vec<u32>>,
quads: Vec<[u32; 4]>,
lines: Vec<Vec<u32>>,
vertices: Vec<u32>,
hexahedra: Vec<[u32; 8]>,
polyhedra: Vec<Vec<Vec<u32>>>,
}
impl Cells {
fn count(&self) -> usize {
self.vertices.len()
+ self.lines.len()
+ self.triangles.len()
+ self.quads.len()
+ self.polygons.len()
+ self.hexahedra.len()
+ self.polyhedra.len()
}
}
const HEXAHEDRON_FACES: [[usize; 4]; 6] = [
[0, 3, 2, 1],
[4, 5, 6, 7],
[0, 1, 5, 4],
[1, 2, 6, 5],
[2, 3, 7, 6],
[3, 0, 4, 7],
];
fn cells_of(object: &SceneObject) -> Cells {
let (_, triangles) = object.geometry.merged();
let mut lines = Vec::new();
let mut quads = Vec::new();
let mut polygons: Vec<Vec<u32>> = Vec::new();
let mut vertices = Vec::new();
let mut hexahedra = Vec::new();
let mut polyhedra: Vec<Vec<Vec<u32>>> = Vec::new();
let mut base = 0u32;
for patch in &object.geometry.patches {
match &patch.topology {
Topology::Segments(pairs) => {
lines.extend(pairs.iter().map(|pair| vec![pair[0] + base, pair[1] + base]));
}
Topology::Polylines(runs) => {
let mut start = base;
for (rank, run) in runs.iter().enumerate() {
if *run >= 2 {
let mut line: Vec<_> = (0..*run).map(|i| start + i).collect();
if patch.closed.get(rank) == Some(&true) {
line.push(start);
}
lines.push(line);
}
start += run;
}
}
Topology::Quads(quadrilaterals) => {
quads.extend(quadrilaterals.iter().map(|quad| {
let mut shifted = *quad;
for corner in &mut shifted {
*corner += base;
}
shifted
}));
}
Topology::Polygons(rings) => {
polygons.extend(
rings
.iter()
.map(|ring| ring.iter().map(|node| node + base).collect()),
);
}
Topology::Points | Topology::Grid { .. } => {
vertices.extend(base..base + patch.points.len() as u32);
}
Topology::Hexahedra { cells, .. } => {
hexahedra.extend(cells.iter().map(|cell| {
let mut shifted = *cell;
for corner in &mut shifted {
*corner += base;
}
shifted
}));
}
Topology::Polyhedra(cells) => {
polyhedra.extend((0..cells.cell_count()).map(|cell| {
cells
.faces_outward(cell)
.map(|face| face.iter().map(|node| node + base).collect())
.collect()
}));
}
Topology::Triangles(_) => {}
}
base += patch.points.len() as u32;
}
Cells {
triangles,
quads,
polygons,
lines,
vertices,
hexahedra,
polyhedra,
}
}
impl VtkExporter {
fn legacy_document(
&self,
object: &SceneObject,
points: &[[f64; 3]],
cells: &Cells,
options: &Options<'_>,
out: &mut dyn Write,
) -> Result<()> {
writeln!(out, "# vtk DataFile Version {}", self.version)?;
let title = format!("{} [{}]", one_line(&object.title), object.uuid);
writeln!(out, "{}", &title[..title.len().min(255)])?;
writeln!(out, "ASCII")?;
let volumetric = !cells.hexahedra.is_empty() || !cells.polyhedra.is_empty();
writeln!(
out,
"DATASET {}",
if volumetric {
"UNSTRUCTURED_GRID"
} else {
"POLYDATA"
}
)?;
writeln!(out, "POINTS {} double", points.len())?;
for point in points {
writeln!(out, "{} {} {}", point[0], point[1], point[2])?;
}
if volumetric {
let count = cells.hexahedra.len()
+ cells.polyhedra.len()
+ cells.triangles.len()
+ cells.quads.len()
+ cells.polygons.len()
+ cells.lines.len();
let size: usize = cells.hexahedra.len() * 9
+ cells
.polyhedra
.iter()
.map(|c| legacy_polyhedron_len(c) + 1)
.sum::<usize>()
+ cells.triangles.len() * 4
+ cells.quads.len() * 5
+ cells.polygons.iter().map(|p| p.len() + 1).sum::<usize>()
+ cells.lines.iter().map(|l| l.len() + 1).sum::<usize>();
writeln!(out, "CELLS {count} {size}")?;
for cell in &cells.hexahedra {
let corners: Vec<String> = cell.iter().map(|c| c.to_string()).collect();
writeln!(out, "8 {}", corners.join(" "))?;
}
for cell in &cells.polyhedra {
let mut stream = vec![cell.len().to_string()];
for face in cell {
stream.push(face.len().to_string());
stream.extend(face.iter().map(|node| node.to_string()));
}
writeln!(out, "{} {}", legacy_polyhedron_len(cell), stream.join(" "))?;
}
for triangle in &cells.triangles {
writeln!(out, "3 {} {} {}", triangle[0], triangle[1], triangle[2])?;
}
for quad in &cells.quads {
writeln!(out, "4 {} {} {} {}", quad[0], quad[1], quad[2], quad[3])?;
}
for ring in &cells.polygons {
let nodes: Vec<String> = ring.iter().map(|n| n.to_string()).collect();
writeln!(out, "{} {}", ring.len(), nodes.join(" "))?;
}
for line in &cells.lines {
let indices: Vec<String> = line.iter().map(|i| i.to_string()).collect();
writeln!(out, "{} {}", line.len(), indices.join(" "))?;
}
writeln!(out, "CELL_TYPES {count}")?;
for _ in &cells.hexahedra {
writeln!(out, "12")?; }
for _ in &cells.polyhedra {
writeln!(out, "42")?; }
for _ in &cells.triangles {
writeln!(out, "5")?; }
for _ in &cells.quads {
writeln!(out, "9")?; }
for _ in &cells.polygons {
writeln!(out, "7")?; }
for _ in &cells.lines {
writeln!(out, "4")?; }
}
if !volumetric
&& (!cells.triangles.is_empty() || !cells.quads.is_empty() || !cells.polygons.is_empty())
{
let size: usize = cells.triangles.len() * 4
+ cells.quads.len() * 5
+ cells.polygons.iter().map(|p| p.len() + 1).sum::<usize>();
let count = cells.triangles.len() + cells.quads.len() + cells.polygons.len();
writeln!(out, "POLYGONS {count} {size}")?;
for triangle in &cells.triangles {
writeln!(out, "3 {} {} {}", triangle[0], triangle[1], triangle[2])?;
}
for quad in &cells.quads {
writeln!(out, "4 {} {} {} {}", quad[0], quad[1], quad[2], quad[3])?;
}
for ring in &cells.polygons {
let nodes: Vec<String> = ring.iter().map(|n| n.to_string()).collect();
writeln!(out, "{} {}", ring.len(), nodes.join(" "))?;
}
}
if !volumetric && !cells.lines.is_empty() {
let size: usize = cells.lines.iter().map(|l| l.len() + 1).sum();
writeln!(out, "LINES {} {size}", cells.lines.len())?;
for line in &cells.lines {
let indices: Vec<String> = line.iter().map(|i| i.to_string()).collect();
writeln!(out, "{} {}", line.len(), indices.join(" "))?;
}
}
if !cells.vertices.is_empty() {
writeln!(
out,
"VERTICES {} {}",
cells.vertices.len(),
cells.vertices.len() * 2
)?;
for vertex in &cells.vertices {
writeln!(out, "1 {vertex}")?;
}
}
let pairs = self.pairs(object, options);
if !pairs.is_empty() {
writeln!(out, "FIELD FieldData {}", pairs.len())?;
for (key, value) in &pairs {
writeln!(out, "{} 1 1 string", sanitise(key))?;
writeln!(out, "{}", one_line(value))?;
}
}
let cell_count = cells.count();
for (scope, count) in [
(AttributeScope::Points, points.len()),
(AttributeScope::Cells, cell_count),
] {
let attributes: Vec<&Attribute> = object
.attributes
.iter()
.filter(|attribute| attribute.scope == scope)
.collect();
if attributes.is_empty() || count == 0 {
continue;
}
writeln!(
out,
"{} {count}",
match scope {
AttributeScope::Points => "POINT_DATA",
AttributeScope::Cells => "CELL_DATA",
}
)?;
for attribute in attributes {
let table = if attribute.lut.is_empty() {
"default".to_string()
} else {
format!("{}_colors", sanitise(&attribute.name))
};
writeln!(
out,
"SCALARS {} double {}",
sanitise(&attribute.name),
attribute.components.max(1)
)?;
writeln!(out, "LOOKUP_TABLE {table}")?;
for chunk in attribute.values.chunks(attribute.components.max(1) as usize) {
let numbers: Vec<String> = chunk.iter().map(number).collect();
writeln!(out, "{}", numbers.join(" "))?;
}
if !attribute.lut.is_empty() {
writeln!(out, "LOOKUP_TABLE {table} {}", attribute.lut.len())?;
for [r, g, b, a] in &attribute.lut {
writeln!(out, "{r} {g} {b} {a}")?;
}
}
}
}
Ok(())
}
fn xml_document(
&self,
object: &SceneObject,
points: &[[f64; 3]],
cells: &Cells,
options: &Options<'_>,
out: &mut dyn Write,
) -> Result<()> {
let container = match self.flavour {
Flavour::UnstructuredGrid => "UnstructuredGrid",
_ => "PolyData",
};
writeln!(out, "<?xml version=\"1.0\"?>")?;
writeln!(
out,
"<VTKFile type=\"{container}\" version=\"{}\" byte_order=\"LittleEndian\">",
self.version
)?;
let pairs = self.pairs(object, options);
if !pairs.is_empty() {
writeln!(out, " <!--")?;
for (key, value) in &pairs {
writeln!(out, " {key}: {}", one_line(value).replace("--", "- -"))?;
}
writeln!(out, " -->")?;
}
writeln!(out, " <{container}>")?;
if !pairs.is_empty() {
writeln!(out, " <FieldData>")?;
for (key, value) in &pairs {
let codes: Vec<String> = one_line(value)
.chars()
.map(|c| (c as u32).to_string())
.chain(std::iter::once("0".to_string()))
.collect();
writeln!(
out,
" <DataArray type=\"String\" Name=\"{}\" NumberOfTuples=\"1\" format=\"ascii\">",
escape(&sanitise(key))
)?;
writeln!(out, " {}", codes.join(" "))?;
writeln!(out, " </DataArray>")?;
}
writeln!(out, " </FieldData>")?;
}
match self.flavour {
Flavour::UnstructuredGrid => {
let count = cells.count();
writeln!(
out,
" <Piece NumberOfPoints=\"{}\" NumberOfCells=\"{count}\">",
points.len()
)?;
write_points(points, out)?;
write_unstructured_cells(cells, out)?;
}
_ => {
writeln!(
out,
" <Piece NumberOfPoints=\"{}\" NumberOfVerts=\"{}\" NumberOfLines=\"{}\" \
NumberOfStrips=\"0\" NumberOfPolys=\"{}\">",
points.len(),
cells.vertices.len(),
cells.lines.len(),
cells.triangles.len()
+ cells.quads.len()
+ cells.polygons.len()
+ cells.hexahedra.len() * 6
+ cells.polyhedra.iter().map(|cell| cell.len()).sum::<usize>()
)?;
write_points(points, out)?;
write_poly_cells(cells, out)?;
}
}
write_attributes(object, cells, out)?;
writeln!(out, " </Piece>")?;
writeln!(out, " </{container}>")?;
writeln!(out, "</VTKFile>")?;
Ok(())
}
fn pairs(&self, object: &SceneObject, options: &Options<'_>) -> Vec<(String, String)> {
let mut pairs: Vec<(String, String)> = vec![
("uuid".to_string(), object.uuid.clone()),
("uri".to_string(), object.uri.clone()),
("title".to_string(), object.title.clone()),
("type".to_string(), object.qualified_type.clone()),
];
if let Some(crs) = &object.crs {
let system = match crs.projected.epsg() {
Some(code) => format!("EPSG:{code}"),
None => format!("{:?}", crs.projected),
};
pairs.push(("crs".to_string(), format!("{} {system}", crs.uuid)));
}
if let Some(note) = options.time_equivalence_note(object) {
pairs.push(("vertical_axis".to_string(), note));
}
for (key, value) in object.metadata.pairs() {
if let Some(existing) = pairs.iter_mut().find(|(k, _)| k == key) {
existing.1 = value.to_string();
} else {
pairs.push((key.to_string(), value.to_string()));
}
}
pairs.retain(|(_, value)| !value.is_empty());
pairs
}
}
fn write_attributes(object: &SceneObject, cells: &Cells, out: &mut dyn Write) -> Result<()> {
for (scope, tag) in [
(AttributeScope::Points, "PointData"),
(AttributeScope::Cells, "CellData"),
] {
let attributes: Vec<&Attribute> = object
.attributes
.iter()
.filter(|attribute| attribute.scope == scope)
.collect();
if attributes.is_empty() {
continue;
}
let _ = cells;
writeln!(
out,
" <{tag} Scalars=\"{}\">",
escape(&sanitise(&attributes[0].name))
)?;
for attribute in attributes {
writeln!(
out,
" <DataArray type=\"Float64\" Name=\"{}\" NumberOfComponents=\"{}\" \
format=\"ascii\">",
escape(&sanitise(&attribute.name)),
attribute.components.max(1)
)?;
let numbers: Vec<String> = attribute.values.iter().map(number).collect();
writeln!(out, " {}", numbers.join(" "))?;
writeln!(out, " </DataArray>")?;
}
writeln!(out, " </{tag}>")?;
}
Ok(())
}
fn number(value: &f64) -> String {
if value.is_nan() {
"nan".to_string()
} else if value.is_infinite() {
if *value > 0.0 {
"inf".to_string()
} else {
"-inf".to_string()
}
} else {
value.to_string()
}
}
fn write_points(points: &[[f64; 3]], out: &mut dyn Write) -> Result<()> {
writeln!(out, " <Points>")?;
writeln!(
out,
" <DataArray type=\"Float64\" NumberOfComponents=\"3\" format=\"ascii\">"
)?;
for point in points {
writeln!(out, " {} {} {}", point[0], point[1], point[2])?;
}
writeln!(out, " </DataArray>")?;
writeln!(out, " </Points>")?;
Ok(())
}
fn legacy_polyhedron_len(cell: &[Vec<u32>]) -> usize {
1 + cell.iter().map(|face| face.len() + 1).sum::<usize>()
}
fn write_poly_cells(cells: &Cells, out: &mut dyn Write) -> Result<()> {
if !cells.vertices.is_empty() {
let connectivity: Vec<String> = cells.vertices.iter().map(|v| v.to_string()).collect();
let offsets: Vec<String> = (1..=cells.vertices.len()).map(|i| i.to_string()).collect();
section("Verts", &connectivity, &offsets, out)?;
}
if !cells.lines.is_empty() {
let mut connectivity = Vec::new();
let mut offsets = Vec::new();
let mut at = 0usize;
for line in &cells.lines {
connectivity.extend(line.iter().map(|i| i.to_string()));
at += line.len();
offsets.push(at.to_string());
}
section("Lines", &connectivity, &offsets, out)?;
}
if !cells.triangles.is_empty()
|| !cells.quads.is_empty()
|| !cells.polygons.is_empty()
|| !cells.hexahedra.is_empty()
|| !cells.polyhedra.is_empty()
{
let mut connectivity = Vec::new();
let mut offsets = Vec::new();
let mut at = 0usize;
for triangle in &cells.triangles {
connectivity.extend(triangle.iter().map(|v| v.to_string()));
at += 3;
offsets.push(at.to_string());
}
for quad in &cells.quads {
connectivity.extend(quad.iter().map(|v| v.to_string()));
at += 4;
offsets.push(at.to_string());
}
for ring in &cells.polygons {
connectivity.extend(ring.iter().map(|v| v.to_string()));
at += ring.len();
offsets.push(at.to_string());
}
for cell in &cells.hexahedra {
for face in HEXAHEDRON_FACES {
connectivity.extend(face.iter().map(|corner| cell[*corner].to_string()));
at += 4;
offsets.push(at.to_string());
}
}
for cell in &cells.polyhedra {
for face in cell {
connectivity.extend(face.iter().map(|node| node.to_string()));
at += face.len();
offsets.push(at.to_string());
}
}
section("Polys", &connectivity, &offsets, out)?;
}
Ok(())
}
fn write_unstructured_cells(cells: &Cells, out: &mut dyn Write) -> Result<()> {
let mut connectivity: Vec<String> = Vec::new();
let mut offsets: Vec<String> = Vec::new();
let mut types: Vec<String> = Vec::new();
let mut at = 0usize;
for vertex in &cells.vertices {
connectivity.push(vertex.to_string());
at += 1;
offsets.push(at.to_string());
types.push("1".to_string());
}
for line in &cells.lines {
connectivity.extend(line.iter().map(|i| i.to_string()));
at += line.len();
offsets.push(at.to_string());
types.push("4".to_string());
}
for triangle in &cells.triangles {
connectivity.extend(triangle.iter().map(|v| v.to_string()));
at += 3;
offsets.push(at.to_string());
types.push("5".to_string());
}
for quad in &cells.quads {
connectivity.extend(quad.iter().map(|v| v.to_string()));
at += 4;
offsets.push(at.to_string());
types.push("9".to_string());
}
for ring in &cells.polygons {
connectivity.extend(ring.iter().map(|v| v.to_string()));
at += ring.len();
offsets.push(at.to_string());
types.push("7".to_string());
}
for cell in &cells.hexahedra {
connectivity.extend(cell.iter().map(|corner| corner.to_string()));
at += 8;
offsets.push(at.to_string());
types.push("12".to_string());
}
let mut faces: Vec<String> = Vec::new();
let mut face_offsets: Vec<String> = Vec::new();
for cell in &cells.polyhedra {
let mut points: Vec<u32> = cell.iter().flatten().copied().collect();
points.sort_unstable();
points.dedup();
connectivity.extend(points.iter().map(|node| node.to_string()));
at += points.len();
offsets.push(at.to_string());
types.push("42".to_string());
faces.push(cell.len().to_string());
for face in cell {
faces.push(face.len().to_string());
faces.extend(face.iter().map(|node| node.to_string()));
}
face_offsets.push(faces.len().to_string());
}
writeln!(out, " <Cells>")?;
array("Int64", "connectivity", &connectivity, out)?;
array("Int64", "offsets", &offsets, out)?;
array("UInt8", "types", &types, out)?;
if !faces.is_empty() {
let others = cells.count() - cells.polyhedra.len();
let mut all: Vec<String> = vec!["-1".to_string(); others];
all.extend(face_offsets);
array("Int64", "faces", &faces, out)?;
array("Int64", "faceoffsets", &all, out)?;
}
writeln!(out, " </Cells>")?;
Ok(())
}
fn section(name: &str, connectivity: &[String], offsets: &[String], out: &mut dyn Write) -> Result<()> {
writeln!(out, " <{name}>")?;
array("Int64", "connectivity", connectivity, out)?;
array("Int64", "offsets", offsets, out)?;
writeln!(out, " </{name}>")?;
Ok(())
}
fn array(kind: &str, name: &str, values: &[String], out: &mut dyn Write) -> Result<()> {
writeln!(
out,
" <DataArray type=\"{kind}\" Name=\"{name}\" format=\"ascii\">"
)?;
writeln!(out, " {}", values.join(" "))?;
writeln!(out, " </DataArray>")?;
Ok(())
}
fn sanitise(name: &str) -> String {
let cleaned: String = name
.chars()
.map(|c| if c.is_whitespace() { '_' } else { c })
.collect();
if cleaned.is_empty() {
"field".to_string()
} else {
cleaned
}
}
fn one_line(value: &str) -> String {
value.replace(['\n', '\r'], " ")
}
fn escape(value: &str) -> String {
value
.replace('&', "&")
.replace('<', "<")
.replace('>', ">")
.replace('"', """)
}