1#![allow(clippy::manual_div_ceil, clippy::manual_strip)]
2use std::fmt;
14use std::io::{self, BufRead, Write};
15
16#[derive(Debug, Clone, Copy, PartialEq, Eq)]
24pub enum Su2ElementType {
25 Line = 3,
27 Triangle = 5,
29 Quadrilateral = 9,
31 Tetrahedron = 10,
33 Hexahedron = 12,
35 Prism = 13,
37 Pyramid = 14,
39}
40
41impl Su2ElementType {
42 pub fn n_nodes(self) -> usize {
44 match self {
45 Su2ElementType::Line => 2,
46 Su2ElementType::Triangle => 3,
47 Su2ElementType::Quadrilateral => 4,
48 Su2ElementType::Tetrahedron => 4,
49 Su2ElementType::Hexahedron => 8,
50 Su2ElementType::Prism => 6,
51 Su2ElementType::Pyramid => 5,
52 }
53 }
54
55 pub fn from_vtk(code: u32) -> Option<Self> {
57 match code {
58 3 => Some(Su2ElementType::Line),
59 5 => Some(Su2ElementType::Triangle),
60 9 => Some(Su2ElementType::Quadrilateral),
61 10 => Some(Su2ElementType::Tetrahedron),
62 12 => Some(Su2ElementType::Hexahedron),
63 13 => Some(Su2ElementType::Prism),
64 14 => Some(Su2ElementType::Pyramid),
65 _ => None,
66 }
67 }
68}
69
70impl fmt::Display for Su2ElementType {
71 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
72 write!(f, "{}", *self as u32)
73 }
74}
75
76#[derive(Debug, Clone)]
82pub struct Su2Element {
83 pub element_type: Su2ElementType,
85 pub nodes: Vec<usize>,
87 pub global_id: Option<usize>,
89}
90
91impl Su2Element {
92 pub fn new(element_type: Su2ElementType, nodes: Vec<usize>) -> Self {
94 Self {
95 element_type,
96 nodes,
97 global_id: None,
98 }
99 }
100}
101
102#[derive(Debug, Clone)]
111pub struct Su2BoundaryMarker {
112 pub tag: String,
114 pub elements: Vec<Su2Element>,
116}
117
118impl Su2BoundaryMarker {
119 pub fn new(tag: impl Into<String>) -> Self {
121 Self {
122 tag: tag.into(),
123 elements: Vec::new(),
124 }
125 }
126
127 pub fn n_elements(&self) -> usize {
129 self.elements.len()
130 }
131
132 pub fn add_element(&mut self, elem: Su2Element) {
134 self.elements.push(elem);
135 }
136}
137
138#[derive(Debug, Clone)]
146pub struct Su2Mesh {
147 pub ndim: usize,
149 pub coords: Vec<f64>,
153 pub elements: Vec<Su2Element>,
155 pub markers: Vec<Su2BoundaryMarker>,
157}
158
159impl Su2Mesh {
160 pub fn new(ndim: usize) -> Self {
162 Self {
163 ndim,
164 coords: Vec::new(),
165 elements: Vec::new(),
166 markers: Vec::new(),
167 }
168 }
169
170 pub fn n_nodes(&self) -> usize {
172 self.coords.len().checked_div(self.ndim).unwrap_or(0)
173 }
174
175 pub fn n_elements(&self) -> usize {
177 self.elements.len()
178 }
179
180 pub fn add_node(&mut self, coords: &[f64]) -> usize {
184 debug_assert_eq!(coords.len(), self.ndim);
185 let idx = self.n_nodes();
186 self.coords.extend_from_slice(coords);
187 idx
188 }
189
190 pub fn add_element(&mut self, elem: Su2Element) {
192 self.elements.push(elem);
193 }
194
195 pub fn add_marker(&mut self, marker: Su2BoundaryMarker) {
197 self.markers.push(marker);
198 }
199
200 pub fn boundary_marker(&self, tag: &str) -> Option<&Su2BoundaryMarker> {
202 self.markers.iter().find(|m| m.tag == tag)
203 }
204
205 pub fn node_coords(&self, i: usize) -> &[f64] {
207 let start = i * self.ndim;
208 &self.coords[start..start + self.ndim]
209 }
210}
211
212pub fn element_connectivity(elem_type: Su2ElementType) -> u32 {
221 elem_type as u32
222}
223
224pub fn read_su2<R: BufRead>(reader: R) -> io::Result<Su2Mesh> {
236 let lines: Vec<String> = reader.lines().map(|l| l.unwrap_or_default()).collect();
237
238 let mut ndim = 2usize;
239 let mut mesh = Su2Mesh::new(2);
240 let mut pos = 0;
241
242 while pos < lines.len() {
243 let line = lines[pos].trim().to_string();
244 pos += 1;
245
246 if line.starts_with('%') || line.is_empty() {
247 continue;
248 }
249
250 if let Some(val) = parse_kv(&line, "NDIME=") {
251 ndim = val.parse().unwrap_or(2);
252 mesh = Su2Mesh::new(ndim);
253 continue;
254 }
255
256 if let Some(val) = parse_kv(&line, "NPOIN=") {
257 let n_nodes: usize = val
258 .split_whitespace()
259 .next()
260 .unwrap_or("0")
261 .parse()
262 .unwrap_or(0);
263 for _ in 0..n_nodes {
264 if pos >= lines.len() {
265 break;
266 }
267 let node_line = lines[pos].trim().to_string();
268 pos += 1;
269 let tokens: Vec<f64> = node_line
270 .split_whitespace()
271 .take(ndim)
272 .filter_map(|t| t.parse().ok())
273 .collect();
274 if tokens.len() == ndim {
275 mesh.coords.extend_from_slice(&tokens);
276 }
277 }
278 continue;
279 }
280
281 if let Some(val) = parse_kv(&line, "NELEM=") {
282 let n_elem: usize = val.parse().unwrap_or(0);
283 for _ in 0..n_elem {
284 if pos >= lines.len() {
285 break;
286 }
287 let el_line = lines[pos].trim().to_string();
288 pos += 1;
289 if let Some(elem) = parse_element_line(&el_line) {
290 mesh.elements.push(elem);
291 }
292 }
293 continue;
294 }
295
296 if let Some(val) = parse_kv(&line, "NMARK=") {
297 let _n_mark: usize = val.parse().unwrap_or(0);
298 let _ = _n_mark;
300 continue;
301 }
302
303 if let Some(tag) = parse_kv(&line, "MARKER_TAG=") {
304 let mut marker = Su2BoundaryMarker::new(tag.trim());
305 while pos < lines.len() {
307 let next = lines[pos].trim().to_string();
308 pos += 1;
309 if next.is_empty() || next.starts_with('%') {
310 continue;
311 }
312 if let Some(val) = parse_kv(&next, "MARKER_ELEMS=") {
313 let n_mel: usize = val.parse().unwrap_or(0);
314 for _ in 0..n_mel {
315 if pos >= lines.len() {
316 break;
317 }
318 let mel_line = lines[pos].trim().to_string();
319 pos += 1;
320 if let Some(elem) = parse_element_line(&mel_line) {
321 marker.elements.push(elem);
322 }
323 }
324 break;
325 }
326 }
327 mesh.markers.push(marker);
328 continue;
329 }
330 }
331
332 Ok(mesh)
333}
334
335fn parse_kv<'a>(line: &'a str, key: &str) -> Option<&'a str> {
337 let trimmed = line.trim();
338 if trimmed.starts_with(key) {
339 Some(trimmed[key.len()..].trim())
340 } else {
341 None
342 }
343}
344
345fn parse_element_line(line: &str) -> Option<Su2Element> {
347 let mut tokens = line.split_whitespace();
348 let vtk_type: u32 = tokens.next()?.parse().ok()?;
349 let elem_type = Su2ElementType::from_vtk(vtk_type)?;
350 let n = elem_type.n_nodes();
351 let nodes: Vec<usize> = tokens.take(n).filter_map(|t| t.parse().ok()).collect();
352 if nodes.len() == n {
353 Some(Su2Element::new(elem_type, nodes))
354 } else {
355 None
356 }
357}
358
359pub fn write_su2<W: Write>(writer: &mut W, mesh: &Su2Mesh) -> io::Result<()> {
368 writeln!(writer, "%")?;
369 writeln!(writer, "% SU2 mesh written by OxiPhysics")?;
370 writeln!(writer, "%")?;
371 writeln!(writer, "NDIME= {}", mesh.ndim)?;
372 writeln!(writer)?;
373
374 writeln!(writer, "NELEM= {}", mesh.n_elements())?;
376 for (i, elem) in mesh.elements.iter().enumerate() {
377 let nodes_str: Vec<String> = elem.nodes.iter().map(|n| n.to_string()).collect();
378 writeln!(
379 writer,
380 "{} {} {}",
381 elem.element_type,
382 nodes_str.join(" "),
383 i
384 )?;
385 }
386 writeln!(writer)?;
387
388 writeln!(writer, "NPOIN= {}", mesh.n_nodes())?;
390 for i in 0..mesh.n_nodes() {
391 let coords = mesh.node_coords(i);
392 let coord_str: Vec<String> = coords.iter().map(|&c| format!("{c:.10e}")).collect();
393 writeln!(writer, "{}", coord_str.join("\t"))?;
394 }
395 writeln!(writer)?;
396
397 writeln!(writer, "NMARK= {}", mesh.markers.len())?;
399 for marker in &mesh.markers {
400 writeln!(writer, "MARKER_TAG= {}", marker.tag)?;
401 writeln!(writer, "MARKER_ELEMS= {}", marker.n_elements())?;
402 for elem in &marker.elements {
403 let nodes_str: Vec<String> = elem.nodes.iter().map(|n| n.to_string()).collect();
404 writeln!(writer, "{} {}", elem.element_type, nodes_str.join(" "))?;
405 }
406 }
407
408 Ok(())
409}
410
411pub fn solution_file<W: Write>(
425 writer: &mut W,
426 mesh: &Su2Mesh,
427 field_names: &[&str],
428 fields: &[&[f64]],
429) -> io::Result<()> {
430 let mut header = vec!["PointID".to_string()];
432 match mesh.ndim {
433 2 => {
434 header.push("x".to_string());
435 header.push("y".to_string());
436 }
437 _ => {
438 header.push("x".to_string());
439 header.push("y".to_string());
440 header.push("z".to_string());
441 }
442 }
443 for name in field_names {
444 header.push(name.to_string());
445 }
446 writeln!(writer, "{}", header.join(","))?;
447
448 for i in 0..mesh.n_nodes() {
449 let mut row = vec![i.to_string()];
450 let coords = mesh.node_coords(i);
451 for &c in coords {
452 row.push(format!("{c:.10e}"));
453 }
454 for field in fields {
455 if i < field.len() {
456 row.push(format!("{:.10e}", field[i]));
457 } else {
458 row.push("0.0".to_string());
459 }
460 }
461 writeln!(writer, "{}", row.join(","))?;
462 }
463 Ok(())
464}
465
466#[cfg(test)]
471mod tests {
472 use super::*;
473 use std::io::BufReader;
474
475 #[test]
478 fn test_element_type_n_nodes_triangle() {
479 assert_eq!(Su2ElementType::Triangle.n_nodes(), 3);
480 }
481
482 #[test]
483 fn test_element_type_n_nodes_hexa() {
484 assert_eq!(Su2ElementType::Hexahedron.n_nodes(), 8);
485 }
486
487 #[test]
488 fn test_element_type_n_nodes_tetra() {
489 assert_eq!(Su2ElementType::Tetrahedron.n_nodes(), 4);
490 }
491
492 #[test]
493 fn test_element_type_n_nodes_line() {
494 assert_eq!(Su2ElementType::Line.n_nodes(), 2);
495 }
496
497 #[test]
498 fn test_element_type_from_vtk_triangle() {
499 assert_eq!(Su2ElementType::from_vtk(5), Some(Su2ElementType::Triangle));
500 }
501
502 #[test]
503 fn test_element_type_from_vtk_unknown() {
504 assert_eq!(Su2ElementType::from_vtk(99), None);
505 }
506
507 #[test]
508 fn test_element_type_from_vtk_hexa() {
509 assert_eq!(
510 Su2ElementType::from_vtk(12),
511 Some(Su2ElementType::Hexahedron)
512 );
513 }
514
515 #[test]
516 fn test_element_connectivity_triangle() {
517 assert_eq!(element_connectivity(Su2ElementType::Triangle), 5);
518 }
519
520 #[test]
521 fn test_element_type_display() {
522 assert_eq!(Su2ElementType::Triangle.to_string(), "5");
523 }
524
525 #[test]
528 fn test_mesh_new_2d() {
529 let m = Su2Mesh::new(2);
530 assert_eq!(m.ndim, 2);
531 assert_eq!(m.n_nodes(), 0);
532 assert_eq!(m.n_elements(), 0);
533 }
534
535 #[test]
536 fn test_mesh_add_node_2d() {
537 let mut m = Su2Mesh::new(2);
538 let idx = m.add_node(&[1.0, 2.0]);
539 assert_eq!(idx, 0);
540 assert_eq!(m.n_nodes(), 1);
541 }
542
543 #[test]
544 fn test_mesh_node_coords() {
545 let mut m = Su2Mesh::new(2);
546 m.add_node(&[3.0, 4.0]);
547 let c = m.node_coords(0);
548 assert!((c[0] - 3.0).abs() < 1e-12);
549 assert!((c[1] - 4.0).abs() < 1e-12);
550 }
551
552 #[test]
553 fn test_mesh_add_element() {
554 let mut m = Su2Mesh::new(2);
555 m.add_node(&[0.0, 0.0]);
556 m.add_node(&[1.0, 0.0]);
557 m.add_node(&[0.0, 1.0]);
558 m.add_element(Su2Element::new(Su2ElementType::Triangle, vec![0, 1, 2]));
559 assert_eq!(m.n_elements(), 1);
560 }
561
562 #[test]
563 fn test_mesh_boundary_marker_lookup() {
564 let mut m = Su2Mesh::new(2);
565 m.add_marker(Su2BoundaryMarker::new("wall"));
566 assert!(m.boundary_marker("wall").is_some());
567 assert!(m.boundary_marker("farfield").is_none());
568 }
569
570 #[test]
573 fn test_marker_new_empty() {
574 let mk = Su2BoundaryMarker::new("inlet");
575 assert_eq!(mk.tag, "inlet");
576 assert_eq!(mk.n_elements(), 0);
577 }
578
579 #[test]
580 fn test_marker_add_element() {
581 let mut mk = Su2BoundaryMarker::new("wall");
582 mk.add_element(Su2Element::new(Su2ElementType::Line, vec![0, 1]));
583 assert_eq!(mk.n_elements(), 1);
584 }
585
586 fn make_simple_2d_mesh() -> Su2Mesh {
589 let mut m = Su2Mesh::new(2);
590 m.add_node(&[0.0, 0.0]);
591 m.add_node(&[1.0, 0.0]);
592 m.add_node(&[0.0, 1.0]);
593 m.add_element(Su2Element::new(Su2ElementType::Triangle, vec![0, 1, 2]));
594 let mut mk = Su2BoundaryMarker::new("wall");
595 mk.add_element(Su2Element::new(Su2ElementType::Line, vec![0, 1]));
596 m.add_marker(mk);
597 m
598 }
599
600 #[test]
601 fn test_write_su2_contains_ndime() {
602 let m = make_simple_2d_mesh();
603 let mut buf = Vec::new();
604 write_su2(&mut buf, &m).unwrap();
605 let s = String::from_utf8(buf).unwrap();
606 assert!(s.contains("NDIME= 2"), "missing NDIME");
607 }
608
609 #[test]
610 fn test_write_su2_contains_nelem() {
611 let m = make_simple_2d_mesh();
612 let mut buf = Vec::new();
613 write_su2(&mut buf, &m).unwrap();
614 let s = String::from_utf8(buf).unwrap();
615 assert!(s.contains("NELEM= 1"));
616 }
617
618 #[test]
619 fn test_write_su2_contains_marker_tag() {
620 let m = make_simple_2d_mesh();
621 let mut buf = Vec::new();
622 write_su2(&mut buf, &m).unwrap();
623 let s = String::from_utf8(buf).unwrap();
624 assert!(s.contains("MARKER_TAG= wall"));
625 }
626
627 #[test]
628 fn test_roundtrip_n_nodes() {
629 let m = make_simple_2d_mesh();
630 let mut buf = Vec::new();
631 write_su2(&mut buf, &m).unwrap();
632 let reader = BufReader::new(buf.as_slice());
633 let m2 = read_su2(reader).unwrap();
634 assert_eq!(m2.n_nodes(), 3);
635 }
636
637 #[test]
638 fn test_roundtrip_n_elements() {
639 let m = make_simple_2d_mesh();
640 let mut buf = Vec::new();
641 write_su2(&mut buf, &m).unwrap();
642 let reader = BufReader::new(buf.as_slice());
643 let m2 = read_su2(reader).unwrap();
644 assert_eq!(m2.n_elements(), 1);
645 }
646
647 #[test]
648 fn test_roundtrip_node_coords() {
649 let m = make_simple_2d_mesh();
650 let mut buf = Vec::new();
651 write_su2(&mut buf, &m).unwrap();
652 let reader = BufReader::new(buf.as_slice());
653 let m2 = read_su2(reader).unwrap();
654 let c = m2.node_coords(0);
655 assert!(c[0].abs() < 1e-6, "x0={}", c[0]);
656 assert!(c[1].abs() < 1e-6, "y0={}", c[1]);
657 }
658
659 #[test]
660 fn test_roundtrip_element_type() {
661 let m = make_simple_2d_mesh();
662 let mut buf = Vec::new();
663 write_su2(&mut buf, &m).unwrap();
664 let reader = BufReader::new(buf.as_slice());
665 let m2 = read_su2(reader).unwrap();
666 assert_eq!(m2.elements[0].element_type, Su2ElementType::Triangle);
667 }
668
669 #[test]
670 fn test_roundtrip_element_nodes() {
671 let m = make_simple_2d_mesh();
672 let mut buf = Vec::new();
673 write_su2(&mut buf, &m).unwrap();
674 let reader = BufReader::new(buf.as_slice());
675 let m2 = read_su2(reader).unwrap();
676 assert_eq!(m2.elements[0].nodes, vec![0, 1, 2]);
677 }
678
679 #[test]
680 fn test_roundtrip_boundary_marker_tag() {
681 let m = make_simple_2d_mesh();
682 let mut buf = Vec::new();
683 write_su2(&mut buf, &m).unwrap();
684 let reader = BufReader::new(buf.as_slice());
685 let m2 = read_su2(reader).unwrap();
686 assert!(m2.boundary_marker("wall").is_some());
687 }
688
689 #[test]
690 fn test_roundtrip_boundary_marker_n_elements() {
691 let m = make_simple_2d_mesh();
692 let mut buf = Vec::new();
693 write_su2(&mut buf, &m).unwrap();
694 let reader = BufReader::new(buf.as_slice());
695 let m2 = read_su2(reader).unwrap();
696 let mk = m2.boundary_marker("wall").unwrap();
697 assert_eq!(mk.n_elements(), 1);
698 }
699
700 #[test]
703 fn test_solution_file_header() {
704 let m = make_simple_2d_mesh();
705 let pressure = vec![1.0, 2.0, 3.0];
706 let mut buf = Vec::new();
707 solution_file(&mut buf, &m, &["Pressure"], &[&pressure]).unwrap();
708 let s = String::from_utf8(buf).unwrap();
709 assert!(s.contains("PointID"), "missing PointID");
710 assert!(s.contains("Pressure"), "missing Pressure");
711 assert!(s.contains("x"), "missing x");
712 assert!(s.contains("y"), "missing y");
713 }
714
715 #[test]
716 fn test_solution_file_row_count() {
717 let m = make_simple_2d_mesh();
718 let pressure = vec![1.0, 2.0, 3.0];
719 let mut buf = Vec::new();
720 solution_file(&mut buf, &m, &["p"], &[&pressure]).unwrap();
721 let s = String::from_utf8(buf).unwrap();
722 assert_eq!(s.lines().count(), 4);
724 }
725
726 #[test]
727 fn test_solution_file_first_data_row() {
728 let m = make_simple_2d_mesh();
729 let pressure = vec![101325.0, 101000.0, 100500.0];
730 let mut buf = Vec::new();
731 solution_file(&mut buf, &m, &["p"], &[&pressure]).unwrap();
732 let s = String::from_utf8(buf).unwrap();
733 let first_data = s.lines().nth(1).unwrap();
735 assert!(first_data.starts_with("0,"), "row: {first_data}");
736 }
737
738 #[test]
739 fn test_solution_file_multiple_fields() {
740 let m = make_simple_2d_mesh();
741 let pressure = vec![1.0, 2.0, 3.0];
742 let temperature = vec![300.0, 310.0, 320.0];
743 let mut buf = Vec::new();
744 solution_file(&mut buf, &m, &["p", "T"], &[&pressure, &temperature]).unwrap();
745 let s = String::from_utf8(buf).unwrap();
746 assert!(s.contains("p"), "missing p");
747 assert!(s.contains("T"), "missing T");
748 }
749
750 #[test]
753 fn test_read_su2_empty_input() {
754 let reader = BufReader::new("".as_bytes());
755 let m = read_su2(reader).unwrap();
756 assert_eq!(m.n_nodes(), 0);
757 assert_eq!(m.n_elements(), 0);
758 }
759
760 #[test]
761 fn test_read_su2_comments_ignored() {
762 let input = "% This is a comment\n% Another comment\n";
763 let reader = BufReader::new(input.as_bytes());
764 let m = read_su2(reader).unwrap();
765 assert_eq!(m.n_nodes(), 0);
766 }
767}