crater/serde/vtk/
serde.rs

1//! SerDe for VTK file format.
2
3// use super::evaluations::{CellOrPoint, MeshCollectionEvaluator};
4use crate::{
5    primitives::bounding::BoundingBox, primitives::mesh::MeshCollection,
6    primitives::nvector::NVector,
7};
8use burn::prelude::*;
9use itertools::Itertools;
10use vtkio::model::*;
11
12pub trait ToVtk {
13    fn to_vtk(&self) -> Vtk;
14}
15
16impl<const D: usize, B: Backend> ToVtk for Tensor<B, D> {
17    fn to_vtk(&self) -> Vtk {
18        assert_eq!(self.dims()[1], 3);
19        let points: Vec<[f32; 3]> = self
20            .clone()
21            .into_data()
22            .as_slice::<f32>()
23            .unwrap()
24            .chunks(3)
25            .map(|chunk| [chunk[0], chunk[1], chunk[2]])
26            .filter(|chunk| chunk.iter().all(|&x| x < f32::MAX))
27            .collect();
28        assert!(!points.is_empty());
29        // Create VTK data structure
30        Vtk {
31            version: Version::new_legacy(2, 0),
32            title: String::from("Points"),
33            byte_order: ByteOrder::LittleEndian,
34            file_path: None,
35            data: DataSet::inline(PolyDataPiece {
36                points: points.iter().flatten().copied().collect::<Vec<_>>().into(),
37                verts: Some(VertexNumbers::XML {
38                    connectivity: (0u64..points.len() as u64).collect(),
39                    offsets: (1..points.len() as u64 + 1).collect(),
40                }),
41                lines: None,
42                polys: None,
43                strips: None,
44                data: Attributes {
45                    point: vec![],
46                    cell: vec![],
47                },
48            }),
49        }
50    }
51}
52
53fn add_bounding_box_to_polydata(
54    bbox: &BoundingBox<3>,
55    points: &mut Vec<NVector<3>>,
56    lines: &mut Vec<Vec<u64>>,
57) {
58    let start_index = points.len() as u64;
59    // Add all the points from this bounding box to the points vector
60    points.extend(bbox.corners());
61
62    // Add all the edges
63    // Define the 12 edges of the cube. The corners are in the order:
64    // 0: (xmin, ymin, zmin)
65    // 1: (xmax, ymin, zmin)
66    // 2: (xmin, ymax, zmin)
67    // 3: (xmax, ymax, zmin)
68    // 4: (xmin, ymin, zmax)
69    // 5: (xmax, ymin, zmax)
70    // 6: (xmin, ymax, zmax)
71    // 7: (xmax, ymax, zmax)
72    let edges = [
73        (0u64, 1u64),
74        (1, 3),
75        (3, 2),
76        (2, 0), // bottom face
77        (4, 5),
78        (5, 7),
79        (7, 6),
80        (6, 4), // top face
81        (0, 4),
82        (1, 5),
83        (2, 6),
84        (3, 7), // vertical edges
85    ];
86
87    for (a, b) in &edges {
88        lines.push(vec![start_index + a, start_index + b]);
89    }
90}
91
92impl ToVtk for MeshCollection {
93    fn to_vtk(&self) -> Vtk {
94        Vtk {
95            version: Version::new(),
96            title: "Mesh Collection".into(),
97            byte_order: ByteOrder::BigEndian,
98            data: DataSet::inline(UnstructuredGridPiece {
99                points: self
100                    .vertices()
101                    .flat_map(|vertex| vertex.into_iter())
102                    .collect_vec()
103                    .into(),
104                cells: Cells {
105                    cell_verts: VertexNumbers::XML {
106                        connectivity: (0u64..self.num_vertices() as u64).collect_vec(),
107                        offsets: (0..self.num_triangles())
108                            .map(|i| (i + 1) as u64 * 3)
109                            .collect_vec(),
110                    },
111                    types: vec![CellType::Triangle; self.num_triangles()],
112                },
113                data: Attributes::new(),
114            }),
115            file_path: None,
116        }
117    }
118}
119
120impl ToVtk for BoundingBox<3> {
121    fn to_vtk(&self) -> Vtk {
122        let mut points = vec![];
123        let mut lines = vec![];
124
125        // Add just this bounding box to a running polydata buffer
126        add_bounding_box_to_polydata(self, &mut points, &mut lines);
127
128        let data = DataSet::inline(PolyDataPiece {
129            points: points.as_flattened().iter().copied().collect_vec().into(),
130            lines: Some(VertexNumbers::XML {
131                connectivity: lines.iter().flatten().copied().collect_vec(),
132                offsets: (0..lines.len()).map(|i| (i as u64 + 1) * 2).collect_vec(),
133            }),
134            ..Default::default()
135        });
136
137        Vtk {
138            version: Version::new(),
139            title: "Bounding Box".into(),
140            byte_order: ByteOrder::BigEndian,
141            data,
142            file_path: None,
143        }
144    }
145}
146
147#[cfg(test)]
148mod tests {
149    use super::*;
150
151    use crate::primitives::mesh::{MeshCollection, Triangle, TriangleMesh};
152
153    #[test]
154    fn test_triangle_mesh_to_vtk() {
155        let vertices = Triangle::new([[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [0.0, 1.0, 0.0]]);
156        let mesh: MeshCollection = TriangleMesh::new(vec![vertices]).into();
157        mesh.to_vtk();
158    }
159
160    #[test]
161    fn test_bounding_box_to_vtk() {
162        let bounds =
163            crate::primitives::bounding::BoundingBox::new([-1.0, -1.0, -1.0], [1.0, 1.0, 1.0]);
164        bounds.to_vtk();
165    }
166}
crater/serde/vtk/serde.rs

crater/serde/vtk/
serde.rs
