Skip to main content

del_msh_cpu/
io_vtk.rs

1//! method for VTK files
2
3/// "#[repr(u32)]" is for "as" operator: t = VtkElementType::TETRA as u32 // t == 5
4#[repr(u32)]
5pub enum VtkElementType {
6    TRIANGLE = 5,
7    QUAD = 9,
8    TETRA = 10,
9    HEXAHEDRON = 12,
10    WEDGE = 13,
11    PYRAMID = 14,
12}
13
14pub fn write_vtk_points<T>(
15    file: &mut std::fs::File,
16    name: &str,
17    vtx2xyz: &[T],
18    ndim: usize,
19) -> std::io::Result<()>
20where
21    T: std::fmt::Display,
22{
23    use std::io::Write;
24    let np = vtx2xyz.len() / ndim;
25    writeln!(file, "# vtk DataFile Version 2.0")?;
26    writeln!(file, "{name}")?;
27    writeln!(file, "ASCII")?;
28    writeln!(file, "DATASET UNSTRUCTURED_GRID")?;
29    writeln!(file, "POINTS {np} float")?;
30    if ndim == 3 {
31        for xyz in vtx2xyz.chunks(3) {
32            writeln!(file, "{} {} {}", xyz[0], xyz[1], xyz[2])?;
33        }
34    } else if ndim == 2 {
35        for xy in vtx2xyz.chunks(2) {
36            writeln!(file, "{} {}", xy[0], xy[1])?;
37        }
38    } else {
39        panic!();
40    }
41    Ok(())
42}
43
44pub fn write_vtk_cells(
45    file: &mut std::fs::File,
46    vtk_elem_type: VtkElementType,
47    elem2vtx: &[usize],
48) -> std::io::Result<()> {
49    let num_node = match vtk_elem_type {
50        VtkElementType::TRIANGLE => 3,
51        VtkElementType::QUAD => 4,
52        VtkElementType::TETRA => 4,
53        VtkElementType::HEXAHEDRON => 8,
54        VtkElementType::PYRAMID => 5,
55        VtkElementType::WEDGE => 6,
56    };
57    let num_elem = elem2vtx.len() / num_node;
58    assert_eq!(elem2vtx.len(), num_elem * num_node);
59    use std::io::Write;
60    let mut writer = std::io::BufWriter::new(file);
61    writeln!(writer, "CELLS {} {}", num_elem, num_elem * (num_node + 1))?;
62    for av in elem2vtx.chunks(num_node) {
63        write!(writer, "{num_node}")?;
64        for v in av {
65            write!(writer, " {v}")?;
66        }
67        writeln!(writer)?;
68    }
69    writeln!(writer, "CELL_TYPES {num_elem}")?;
70    {
71        let id_elem = vtk_elem_type as usize;
72        for _ in 0..num_elem {
73            writeln!(writer, "{id_elem}")?;
74        }
75    }
76    writer.flush()?;
77    Ok(())
78}
79
80pub fn write_vtk_cells_mix<IDX>(
81    file: &mut std::fs::File,
82    tet2vtx: &[IDX],
83    pyramid2vtx: &[IDX],
84    prism2vtx: &[IDX],
85    hex2vtx: &[IDX],
86) -> std::io::Result<()>
87where
88    IDX: num_traits::PrimInt + std::fmt::Display,
89{
90    let num_tet = tet2vtx.len() / 4;
91    let num_pyramid = pyramid2vtx.len() / 5;
92    let num_prism = prism2vtx.len() / 6;
93    let num_hex = hex2vtx.len() / 8;
94    let num_elem = num_tet + num_pyramid + num_prism + num_hex;
95    let num_idx = 5 * num_tet + 6 * num_pyramid + 7 * num_prism + 9 * num_hex;
96    use std::io::Write;
97    let mut writer = std::io::BufWriter::new(file);
98    writeln!(writer, "CELLS {num_elem} {num_idx}")?;
99    for av in tet2vtx.chunks(4) {
100        write!(writer, "4")?;
101        av.iter().for_each(|v| write!(writer, " {v}").unwrap());
102        writeln!(writer)?;
103    }
104    for av in pyramid2vtx.chunks(5) {
105        write!(writer, "5")?;
106        av.iter().for_each(|v| write!(writer, " {v}").unwrap());
107        writeln!(writer)?;
108    }
109    for av in prism2vtx.chunks(6) {
110        write!(writer, "6")?;
111        //av.iter().for_each(|v| write!(writer, " {v}").unwrap());
112        writeln!(
113            writer,
114            " {} {} {} {} {} {}",
115            av[0], av[2], av[1], av[3], av[5], av[4]
116        )
117        .unwrap();
118    }
119    for av in hex2vtx.chunks(8) {
120        write!(writer, "8")?;
121        //av.iter().for_each(|v| write!(writer, " {v}").unwrap());
122        writeln!(
123            writer,
124            " {} {} {} {} {} {} {} {}",
125            av[0], av[1], av[2], av[3], av[4], av[5], av[6], av[7]
126        )
127        .unwrap();
128    }
129    writeln!(writer, "CELL_TYPES {num_elem}")?;
130    for _ in 0..num_tet {
131        writeln!(writer, "{}", VtkElementType::TETRA as u32).unwrap();
132    }
133    for _ in 0..num_pyramid {
134        writeln!(writer, "{}", VtkElementType::PYRAMID as u32).unwrap();
135    }
136    for _ in 0..num_prism {
137        writeln!(writer, "{}", VtkElementType::WEDGE as u32).unwrap();
138    }
139    for _ in 0..num_hex {
140        writeln!(writer, "{}", VtkElementType::HEXAHEDRON as u32).unwrap();
141    }
142    writer.flush()?;
143    Ok(())
144}
145
146/// Write mixed polyhedron mesh cells from a CSR (jagged) array.
147/// Element type is inferred from node count: 4→tet, 5→pyramid, 6→prism, 8→hex.
148pub fn write_vtk_cells_polyhedron<IDX>(
149    file: &mut std::fs::File,
150    elem2idx_offset: &[IDX],
151    idx2vtx: &[IDX],
152) -> std::io::Result<()>
153where
154    IDX: num_traits::PrimInt + std::fmt::Display,
155{
156    use std::io::Write;
157    let num_elem = elem2idx_offset.len() - 1;
158    let num_idx: usize = elem2idx_offset
159        .iter()
160        .enumerate()
161        .skip(1)
162        .map(|(i, &e)| {
163            let num_node = (e - elem2idx_offset[i - 1]).to_usize().unwrap();
164            num_node + 1
165        })
166        .sum();
167    let mut writer = std::io::BufWriter::new(file);
168    writeln!(writer, "CELLS {num_elem} {num_idx}")?;
169    for i_elem in 0..num_elem {
170        let i0 = elem2idx_offset[i_elem].to_usize().unwrap();
171        let i1 = elem2idx_offset[i_elem + 1].to_usize().unwrap();
172        let av = &idx2vtx[i0..i1];
173        match av.len() {
174            4 => writeln!(writer, "4 {} {} {} {}", av[0], av[1], av[2], av[3])?,
175            5 => writeln!(
176                writer,
177                "5 {} {} {} {} {}",
178                av[0], av[1], av[2], av[3], av[4]
179            )?,
180            6 => writeln!(
181                writer,
182                "6 {} {} {} {} {} {}",
183                av[0], av[2], av[1], av[3], av[5], av[4]
184            )?,
185            8 => writeln!(
186                writer,
187                "8 {} {} {} {} {} {} {} {}",
188                av[0], av[1], av[2], av[3], av[4], av[5], av[6], av[7]
189            )?,
190            n => panic!("unsupported element with {n} nodes"),
191        }
192    }
193    writeln!(writer, "CELL_TYPES {num_elem}")?;
194    for i_elem in 0..num_elem {
195        let num_node = (elem2idx_offset[i_elem + 1] - elem2idx_offset[i_elem])
196            .to_usize()
197            .unwrap();
198        let cell_type = match num_node {
199            4 => VtkElementType::TETRA as u32,
200            5 => VtkElementType::PYRAMID as u32,
201            6 => VtkElementType::WEDGE as u32,
202            8 => VtkElementType::HEXAHEDRON as u32,
203            n => panic!("unsupported element with {n} nodes"),
204        };
205        writeln!(writer, "{cell_type}")?;
206    }
207    writer.flush()?;
208    Ok(())
209}
210
211pub fn write_vtk_data_point_scalar<T>(
212    file: &mut std::fs::File,
213    vtx2data: &[T],
214    num_vtx: usize,
215    num_stride: usize,
216) -> std::io::Result<()>
217where
218    T: std::fmt::Display,
219{
220    use std::io::Write;
221    writeln!(file, "SCALARS pointvalue float 1")?;
222    writeln!(file, "LOOKUP_TABLE default")?;
223    for ip in 0..num_vtx {
224        writeln!(file, "{}", vtx2data[ip * num_stride])?;
225    }
226    Ok(())
227}
228
229pub fn write_vtk_points_with_velocity(
230    file: &mut std::fs::File,
231    points: &[f32],
232    velocities: &[f32],
233) -> std::io::Result<()> {
234    assert_eq!(points.len() % 3, 0);
235    assert_eq!(velocities.len() % 3, 0);
236    assert_eq!(points.len(), velocities.len());
237
238    let n = points.len() / 3;
239
240    use std::io::Write;
241    let mut w = std::io::BufWriter::new(file);
242
243    writeln!(w, "# vtk DataFile Version 3.0")?;
244    writeln!(w, "points with velocity")?;
245    writeln!(w, "ASCII")?;
246    writeln!(w, "DATASET POLYDATA")?;
247
248    writeln!(w, "POINTS {} double", n)?;
249    for i in 0..n {
250        let x = points[3 * i];
251        let y = points[3 * i + 1];
252        let z = points[3 * i + 2];
253        writeln!(w, "{} {} {}", x, y, z)?;
254    }
255
256    writeln!(w, "VERTICES {} {}", n, 2 * n)?;
257    for i in 0..n {
258        writeln!(w, "1 {}", i)?;
259    }
260
261    writeln!(w, "POINT_DATA {}", n)?;
262    writeln!(w, "VECTORS velocity double")?;
263    for i in 0..n {
264        let vx = velocities[3 * i];
265        let vy = velocities[3 * i + 1];
266        let vz = velocities[3 * i + 2];
267        writeln!(w, "{} {} {}", vx, vy, vz)?;
268    }
269
270    Ok(())
271}
272
273#[cfg(test)]
274mod test {
275    use crate::io_vtk::VtkElementType;
276
277    #[test]
278    fn trimesh3_scalardata() {
279        let (tri2vtx, vtx2xyz) = crate::trimesh3_primitive::hemisphere_zup::<f64>(1., 16, 32);
280        let mut file = std::fs::File::create("../target/trimesh3.vtk").expect("file not found.");
281        crate::io_vtk::write_vtk_points(&mut file, "hoge", &vtx2xyz, 3).unwrap();
282        crate::io_vtk::write_vtk_cells(&mut file, VtkElementType::TRIANGLE, &tri2vtx).unwrap();
283        let vtx2data = {
284            let mut vtx2data = Vec::<f64>::with_capacity(vtx2xyz.len() / 3);
285            for i_vtx in 0..vtx2xyz.len() / 3 {
286                let z = vtx2xyz[i_vtx * 3 + 2];
287                vtx2data.push(z.powi(3));
288            }
289            vtx2data
290        };
291        use std::io::Write;
292        let _ = writeln!(file, "POINT_DATA {}", vtx2xyz.len() / 3);
293        let _ =
294            crate::io_vtk::write_vtk_data_point_scalar(&mut file, &vtx2data, vtx2xyz.len() / 3, 1);
295    }
296}