diff --git a/src/parser.rs b/src/parser.rs index 891ddb8..52f4b09 100644 --- a/src/parser.rs +++ b/src/parser.rs @@ -290,8 +290,14 @@ impl VtkParser { fn attribute_lookup_table(input: &[u8], ft: FileType) -> IResult<&[u8], Attribute> { tagged_block(tag_no_case("LOOKUP_TABLE"), |input| { let (input, (name, num_elements)) = line(tuple((sp(name), sp(parse_u32))))(input)?; - let (input, data) = - Self::attribute_data(input, 4 * num_elements as usize, ScalarType::F32, ft)?; + let (input, data) = match ft { + FileType::Ascii => { + Self::attribute_data(input, 4 * num_elements as usize, ScalarType::F32, ft)? + } + FileType::Binary => { + Self::attribute_data(input, 4 * num_elements as usize, ScalarType::U8, ft)? + } + }; let (input, _) = opt(meta)(input)?; Ok(( diff --git a/src/writer.rs b/src/writer.rs index 0a52b02..72e8e41 100644 --- a/src/writer.rs +++ b/src/writer.rs @@ -247,6 +247,9 @@ mod write_vtk_impl { data: Vec, ) -> Result; fn write_buf(&mut self, data: IOBuffer) -> Result; + fn write_color_scalars_buf(&mut self, buf: IOBuffer) -> Result { + self.write_buf::(buf) + } fn write_attributes( &mut self, @@ -296,7 +299,7 @@ mod write_vtk_impl { )) }, )?; - self.write_buf::(data).map_err(|e| { + self.write_color_scalars_buf::(data).map_err(|e| { Error::Attribute(AttributeError::ColorScalars(EntryPart::Data( e.into(), ))) @@ -308,7 +311,7 @@ mod write_vtk_impl { Error::Attribute(AttributeError::LookupTable(EntryPart::Header)) }, )?; - self.write_buf::(data).map_err(|e| { + self.write_color_scalars_buf::(data).map_err(|e| { Error::Attribute(AttributeError::LookupTable(EntryPart::Data( e.into(), ))) @@ -844,6 +847,54 @@ mod write_vtk_impl { let buf = IOBuffer::from(data); self.write_buf::(buf) } + fn write_color_scalars_buf(&mut self, buf: IOBuffer) -> Result { + fn write_buf_impl(vec: Vec, writer: &mut W, convert_to_u8: F) -> Result + where + W: WriteBytesExt, + F: Fn(T) -> u8, + { + for elem in vec { + W::write_u8(writer, convert_to_u8(elem))?; + } + Ok(()) + } + fn convert_float>(input: T) -> u8 { + (input.into() * 255.0).round().clamp(0.0, 255.0) as u8 + } + + match buf { + IOBuffer::Bit(v) => write_buf_impl(v, &mut self.0, |x| x)?, + IOBuffer::U8(v) => write_buf_impl(v, &mut self.0, |x| x)?, + IOBuffer::I8(v) => write_buf_impl(v, &mut self.0, |x| x as u8)?, + IOBuffer::U16(v) => { + write_buf_impl(v, &mut self.0, |x| x as u8)?; + } + IOBuffer::I16(v) => { + write_buf_impl(v, &mut self.0, |x| x as u8)?; + } + IOBuffer::U32(v) => { + write_buf_impl(v, &mut self.0, |x| x as u8)?; + } + IOBuffer::I32(v) => { + write_buf_impl(v, &mut self.0, |x| x as u8)?; + } + IOBuffer::U64(v) => { + write_buf_impl(v, &mut self.0, |x| x as u8)?; + } + IOBuffer::I64(v) => { + write_buf_impl(v, &mut self.0, |x| x as u8)?; + } + IOBuffer::F32(v) => { + write_buf_impl(v, &mut self.0, convert_float)?; + } + IOBuffer::F64(v) => { + write_buf_impl(v, &mut self.0, convert_float)?; + } + } + + writeln!(&mut self.0)?; + Ok(()) + } fn write_buf(&mut self, buf: IOBuffer) -> Result { fn write_buf_impl(vec: Vec, writer: &mut W, elem_writer: E) -> Result where @@ -910,6 +961,9 @@ mod write_vtk_impl { fn write_buf(&mut self, buf: IOBuffer) -> Result { BinaryWriter(self).write_buf::(buf) } + fn write_color_scalars_buf(&mut self, buf: IOBuffer) -> Result { + BinaryWriter(self).write_color_scalars_buf::(buf) + } } impl WriteVtkImpl for AsciiWriter { @@ -950,6 +1004,56 @@ mod write_vtk_impl { writeln!(&mut self.0, "{}", data)?; Ok(()) } + fn write_color_scalars_buf(&mut self, data: IOBuffer) -> Result { + fn write_buf_impl(vec: Vec, writer: &mut W, convert_to_float: F) -> Result + where + W: std::fmt::Write, + F: Fn(T) -> f32, + { + let mut iter = vec.into_iter().peekable(); + while let Some(elem) = iter.next() { + write!(writer, "{}", convert_to_float(elem))?; + if iter.peek().is_some() { + write!(writer, " ")?; + } + } + Ok(()) + } + fn convert_int>(input: T) -> f32 { + 0.max(input.into()) as f32 / 255.0 + } + + match data { + IOBuffer::Bit(v) => write_buf_impl(v, &mut self.0, convert_int)?, + IOBuffer::U8(v) => write_buf_impl(v, &mut self.0, convert_int)?, + IOBuffer::I8(v) => write_buf_impl(v, &mut self.0, convert_int)?, + IOBuffer::U16(v) => { + write_buf_impl(v, &mut self.0, convert_int)?; + } + IOBuffer::I16(v) => { + write_buf_impl(v, &mut self.0, convert_int)?; + } + IOBuffer::U32(v) => { + write_buf_impl(v, &mut self.0, convert_int)?; + } + IOBuffer::I32(v) => { + write_buf_impl(v, &mut self.0, convert_int)?; + } + IOBuffer::U64(v) => { + write_buf_impl(v, &mut self.0, |x| 0.max(x) as f32 / 255.0)?; + } + IOBuffer::I64(v) => { + write_buf_impl(v, &mut self.0, convert_int)?; + } + IOBuffer::F32(v) => write_buf_impl(v, &mut self.0, |x| x)?, + IOBuffer::F64(v) => { + write_buf_impl(v, &mut self.0, |x| x as f32)?; + } + } + + writeln!(&mut self.0)?; + Ok(()) + } } impl WriteVtkImpl for String { @@ -971,6 +1075,10 @@ mod write_vtk_impl { fn write_buf(&mut self, buf: IOBuffer) -> Result { AsciiWriter(self).write_buf::(buf) } + + fn write_color_scalars_buf(&mut self, buf: IOBuffer) -> Result { + AsciiWriter(self).write_color_scalars_buf::(buf) + } } } diff --git a/tests/legacy.rs b/tests/legacy.rs index 6f93012..dfa47d1 100644 --- a/tests/legacy.rs +++ b/tests/legacy.rs @@ -530,7 +530,14 @@ fn cube_complex_test() -> Result { ], }); - let mut attributes = Attributes { + let to_bin = |v: Vec| -> Vec { v.into_iter().map(|x| (x * 255.0) as u8).collect() }; + + let my_table = vec![ + 0.0f32, 0.0, 0.0, 1.0, 1.0, 0.0, 0.0, 1.0, 0.0, 1.0, 0.0, 1.0, 1.0, 1.0, 0.0, 1.0, 0.0, + 0.0, 1.0, 1.0, 1.0, 0.0, 1.0, 1.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, + ]; + + let make_attributes = |is_bin: bool| Attributes { point: vec![ Attribute::DataArray(DataArray { name: String::from("sample_scalars"), @@ -543,12 +550,11 @@ fn cube_complex_test() -> Result { Attribute::DataArray(DataArray { name: String::from("my_table"), elem: ElementType::LookupTable, - data: vec![ - 0.0f32, 0.0, 0.0, 1.0, 1.0, 0.0, 0.0, 1.0, 0.0, 1.0, 0.0, 1.0, 1.0, 1.0, 0.0, - 1.0, 0.0, 0.0, 1.0, 1.0, 1.0, 0.0, 1.0, 1.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, - 1.0, - ] - .into(), + data: if is_bin { + to_bin(my_table.clone()).into() + } else { + my_table.clone().into() + }, }), ], cell: vec![ @@ -585,6 +591,9 @@ fn cube_complex_test() -> Result { }, ], }; + let mut attributes = make_attributes(false); + let mut attributes_bin = make_attributes(true); + let points: IOBuffer = vec![ 0.0, 0.0, 0.0, 1.0, 0.0, 0.0f32, 1.0, 1.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0, 1.0, 0.0, 1.0, 1.0, 1.0, 1.0, 0.0, 1.0, 1.0, @@ -605,6 +614,19 @@ fn cube_complex_test() -> Result { }), }; + let out1_bin = Vtk { + version: Version::new_legacy(2, 0), + byte_order: ByteOrder::BigEndian, + title: String::from("Cube example"), + file_path: None, + data: DataSet::inline(PolyDataPiece { + points: points.clone(), + polys: polys.clone(), + data: attributes_bin.clone(), + ..Default::default() + }), + }; + let in2 = include_bytes!("../assets/cube_complex_topo.vtk"); let verts = Some(VertexNumbers::Legacy { @@ -650,6 +672,8 @@ fn cube_complex_test() -> Result { }, ]; + attributes_bin.cell = attributes.cell.clone(); + let out2 = Vtk { data: DataSet::inline(PolyDataPiece { points: points.clone(), @@ -660,16 +684,28 @@ fn cube_complex_test() -> Result { }), ..out1.clone() }; + let out2_bin = Vtk { + data: DataSet::inline(PolyDataPiece { + points: points.clone(), + polys: polys.clone(), + verts: verts.clone(), + data: attributes_bin.clone(), + ..Default::default() + }), + ..out1.clone() + }; test!(parse_ne(in1) => ne(&out1)); test_b!(parse_ne(String::new().write_vtk_ne(out1.clone())?.as_bytes()) => ne(&out1)); - test_b!(parse_ne(Vec::::new().write_vtk_ne(out1.clone())?) => ne(&out1)); - test_b!(parse_le(Vec::::new().write_vtk_le(out1.clone())?) => le(&out1)); - test_b!(parse_be(Vec::::new().write_vtk_be(out1.clone())?) => out1); + test_b!(parse_ne(String::new().write_vtk_ne(out1_bin.clone())?.as_bytes()) => ne(&out1)); + test_b!(parse_ne(Vec::::new().write_vtk_ne(out1_bin.clone())?) => ne(&out1_bin)); + test_b!(parse_le(Vec::::new().write_vtk_le(out1_bin.clone())?) => le(&out1_bin)); + test_b!(parse_be(Vec::::new().write_vtk_be(out1_bin.clone())?) => out1_bin); test_b!(parse_ne(in2) => ne(&out2)); test_b!(parse_ne(String::new().write_vtk_ne(out2.clone())?.as_bytes()) => ne(&out2)); - test_b!(parse_le(Vec::::new().write_vtk_le(out2.clone())?) => le(&out2)); - test_b!(parse_be(Vec::::new().write_vtk_be(out2.clone())?) => out2); + test_b!(parse_ne(String::new().write_vtk_ne(out2_bin.clone())?.as_bytes()) => ne(&out2)); + test_b!(parse_le(Vec::::new().write_vtk_le(out2_bin.clone())?) => le(&out2_bin)); + test_b!(parse_be(Vec::::new().write_vtk_be(out2_bin.clone())?) => out2_bin); Ok(()) }