Skip to content

Commit

Permalink
Add support for Ascii and Binary color scalars
Browse files Browse the repository at this point in the history
This is implemented in the writer and parser.
  • Loading branch information
elrnv committed Sep 6, 2024
1 parent 01a4aaf commit d1bce11
Show file tree
Hide file tree
Showing 3 changed files with 166 additions and 16 deletions.
10 changes: 8 additions & 2 deletions src/parser.rs
Original file line number Diff line number Diff line change
Expand Up @@ -290,8 +290,14 @@ impl<BO: ByteOrder + 'static> VtkParser<BO> {
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::U8, 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((
Expand Down
114 changes: 112 additions & 2 deletions src/writer.rs
Original file line number Diff line number Diff line change
Expand Up @@ -176,6 +176,8 @@ mod write_vtk_impl {
DataSet(DataSetError),
NewLine,

InvalidColorScalarsType,

/// Generic formatting error originating from [`std::fmt::Error`].
FormatError,
/// Generic IO error originating from [`std::io::Error`].
Expand All @@ -191,6 +193,10 @@ mod write_vtk_impl {
Error::Header(header_err) => write!(f, "Header: {}", header_err),
Error::DataSet(data_set_err) => write!(f, "Data set: {}", data_set_err),
Error::NewLine => write!(f, "New line"),
Error::InvalidColorScalarsType => write!(
f,
"Invalid color scalars type, must be unsigned char or float"
),
Error::FormatError => write!(f, "Format error"),
Error::IOError(kind) => write!(f, "IO Error: {:?}", kind),
}
Expand Down Expand Up @@ -247,6 +253,9 @@ mod write_vtk_impl {
data: Vec<T>,
) -> Result;
fn write_buf<BO: ByteOrder>(&mut self, data: IOBuffer) -> Result;
fn write_color_scalars_buf<BO: ByteOrder>(&mut self, buf: IOBuffer) -> Result {
self.write_buf::<BO>(buf)
}

fn write_attributes<BO: ByteOrder>(
&mut self,
Expand Down Expand Up @@ -296,7 +305,7 @@ mod write_vtk_impl {
))
},
)?;
self.write_buf::<BO>(data).map_err(|e| {
self.write_color_scalars_buf::<BO>(data).map_err(|e| {
Error::Attribute(AttributeError::ColorScalars(EntryPart::Data(
e.into(),
)))
Expand All @@ -308,7 +317,7 @@ mod write_vtk_impl {
Error::Attribute(AttributeError::LookupTable(EntryPart::Header))
},
)?;
self.write_buf::<BO>(data).map_err(|e| {
self.write_color_scalars_buf::<BO>(data).map_err(|e| {
Error::Attribute(AttributeError::LookupTable(EntryPart::Data(
e.into(),
)))
Expand Down Expand Up @@ -844,6 +853,54 @@ mod write_vtk_impl {
let buf = IOBuffer::from(data);
self.write_buf::<BO>(buf)
}
fn write_color_scalars_buf<BO: ByteOrder>(&mut self, buf: IOBuffer) -> Result {
fn write_buf_impl<T, W, F>(vec: Vec<T>, 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<T: Into<f64>>(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 * 255 as u8)?,
IOBuffer::U8(v) => write_buf_impl(v, &mut self.0, |x| x as u8)?,
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<BO: ByteOrder>(&mut self, buf: IOBuffer) -> Result {
fn write_buf_impl<T, W, E>(vec: Vec<T>, writer: &mut W, elem_writer: E) -> Result
where
Expand Down Expand Up @@ -910,6 +967,9 @@ mod write_vtk_impl {
fn write_buf<BO: ByteOrder>(&mut self, buf: IOBuffer) -> Result {
BinaryWriter(self).write_buf::<BO>(buf)
}
fn write_color_scalars_buf<BO: ByteOrder>(&mut self, buf: IOBuffer) -> Result {
BinaryWriter(self).write_color_scalars_buf::<BO>(buf)
}
}

impl<W: std::fmt::Write> WriteVtkImpl for AsciiWriter<W> {
Expand Down Expand Up @@ -950,6 +1010,56 @@ mod write_vtk_impl {
writeln!(&mut self.0, "{}", data)?;
Ok(())
}
fn write_color_scalars_buf<BO: ByteOrder>(&mut self, data: IOBuffer) -> Result {
fn write_buf_impl<T, W, F>(vec: Vec<T>, 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<T: Into<i64>>(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 {
Expand Down
58 changes: 46 additions & 12 deletions tests/legacy.rs
Original file line number Diff line number Diff line change
Expand Up @@ -530,7 +530,14 @@ fn cube_complex_test() -> Result {
],
});

let mut attributes = Attributes {
let to_bin = |v: Vec<f32>| -> Vec<u8> { 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"),
Expand All @@ -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![
Expand Down Expand Up @@ -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,
Expand All @@ -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 {
Expand Down Expand Up @@ -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(),
Expand All @@ -660,16 +684,26 @@ 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::<u8>::new().write_vtk_ne(out1.clone())?) => ne(&out1));
test_b!(parse_le(Vec::<u8>::new().write_vtk_le(out1.clone())?) => le(&out1));
test_b!(parse_be(Vec::<u8>::new().write_vtk_be(out1.clone())?) => out1);
test_b!(parse_ne(Vec::<u8>::new().write_vtk_ne(out1_bin.clone())?) => ne(&out1_bin));
test_b!(parse_le(Vec::<u8>::new().write_vtk_le(out1_bin.clone())?) => le(&out1_bin));
test_b!(parse_be(Vec::<u8>::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::<u8>::new().write_vtk_le(out2.clone())?) => le(&out2));
test_b!(parse_be(Vec::<u8>::new().write_vtk_be(out2.clone())?) => out2);
test_b!(parse_le(Vec::<u8>::new().write_vtk_le(out2_bin.clone())?) => le(&out2_bin));
test_b!(parse_be(Vec::<u8>::new().write_vtk_be(out2_bin.clone())?) => out2_bin);
Ok(())
}

Expand Down

0 comments on commit d1bce11

Please sign in to comment.