extxyz/write.rs
1use crate::{error::ExtxyzError, Result};
2use std::io::Write;
3
4use extxyz_types::{escape, Frame, Value};
5
6/// Writes a sequence of frames to the given writer.
7///
8/// Each frame is serialized using [`write_frame`] and written in order.
9///
10/// # Arguments
11/// * `w` - The output writer.
12/// * `frames` - An iterable sequence of [`Frame`] values.
13///
14/// # Returns
15/// The number of frames successfully written.
16///
17/// # Errors
18/// Returns an error if writing to the output fails or if serializing any
19/// frame fails.
20///
21/// # Notes
22/// This function consumes the provided iterator.
23///
24/// # Examples
25/// ```ignore
26/// let frames = vec![frame1, frame2, frame3];
27/// let written = write_frames(&mut writer, frames)?;
28/// assert_eq!(written, 3);
29/// ```
30pub fn write_frames<W, I>(w: &mut W, frames: I) -> Result<usize>
31where
32 W: Write,
33 I: IntoIterator<Item = Frame>,
34{
35 let mut count = 0;
36 for frame in frames {
37 write_frame(w, &frame)?;
38 count += 1;
39 }
40 Ok(count)
41}
42
43/// Writes a single frame in extended XYZ (extxyz) format.
44///
45/// This function serializes a [`Frame`] into the extxyz text format and writes
46/// it to the provided writer.
47/// The output includes a `Properties` field derived from the frame's
48/// array data, even if it was not explicitly present in the input. The `Lattice`
49/// field, if present, is written in column-major order following the extxyz
50/// specification.
51///
52/// # Parameters
53/// - `w`: A writer implementing [`Write`] to which the frame will be written.
54/// - `frame`: The [`Frame`] to serialize.
55///
56/// # Errors
57/// Returns an error if:
58/// - Writing to the underlying writer fails.
59/// - The `Lattice` field exists but is not a valid 3×3 integer or float matrix.
60/// - Any internal formatting or serialization step fails.
61///
62/// # Notes
63/// - Keys and values in the frame's metadata are escaped as required by the
64/// extxyz format.
65/// - The ordering of metadata fields follows the internal ordering of the frame.
66/// - The `Properties` field is not taken directly from metadata but inferred
67/// from the atomic data arrays.
68///
69/// # Examples
70/// ```ignore
71/// use std::fs::File;
72/// use extxyz::write_frame;
73//
74/// let mut file = File::create("output.xyz")?;
75/// write_frame(&mut file, &frame)?;
76/// # Ok::<(), Box<dyn std::error::Error>>(())
77/// ```
78pub fn write_frame<W>(w: &mut W, frame: &Frame) -> Result<()>
79where
80 W: Write,
81{
82 let natoms = frame.natoms();
83
84 writeln!(w, "{natoms}")?;
85
86 // info
87 let info = frame.info_orderd();
88 let mut iter = info.iter().peekable();
89 while let Some((k, v)) = iter.next() {
90 // the inner datastructure will store "Properties" as a key (if exist), but in the
91 // write function the Properties field is deduct from the arr.
92 // When read the xyz may not have "Properties" field, but write will always have it.
93 // XXX: therefore in the read (the parser, if I impl myself) need to validate the
94 // properties is conform with what provided.
95 if *k == "Properties" {
96 continue;
97 }
98
99 let s = escape(k);
100 write!(w, "{s}")?;
101 write!(w, "=")?;
102
103 // in extxyz c implementation, lattice treated different write in column-wise and use
104 // single space as spliter
105 if *k == "Lattice" {
106 match v {
107 Value::MatrixInteger(_, _) => {
108 write!(w, "{v}")?;
109 }
110 Value::MatrixFloat(_, _) => {
111 write!(w, "{v}")?;
112 }
113 _ => {
114 // this is unreachable if the inner dict is not create manually
115 return Err(ExtxyzError::InvalidValue(
116 "Lattice must be a 3x3 int/float matrix",
117 ));
118 }
119 }
120 } else {
121 write!(w, "{v}")?;
122 }
123
124 // only add a space if there is more to print in info array
125 if iter.peek().is_some() {
126 write!(w, " ")?;
127 }
128 }
129
130 // "Properties" deduct from the arrs
131 write!(w, "Properties=")?;
132
133 let mut s = String::new();
134 let arrs = frame.arrs_orderd();
135 let mut iter = arrs.iter().peekable();
136 while let Some((k, v)) = iter.next() {
137 s.push_str(k);
138 s.push(':');
139 match v {
140 Value::VecInteger(_, _) => s.push_str("I:1"),
141 Value::VecFloat(_, _) => s.push_str("R:1"),
142 Value::VecBool(_, _) => s.push_str("L:1"),
143 Value::VecText(_, _) => s.push_str("S:1"),
144 Value::MatrixInteger(_, shape) => s.push_str(format!("I:{}", shape.1).as_str()),
145 Value::MatrixFloat(_, shape) => s.push_str(format!("R:{}", shape.1).as_str()),
146 Value::MatrixBool(_, shape) => s.push_str(format!("L:{}", shape.1).as_str()),
147 Value::MatrixText(_, shape) => s.push_str(format!("S:{}", shape.1).as_str()),
148 _ => {
149 // this is unreachable if the inner dict is not create manually
150 return Err(ExtxyzError::InvalidValue(
151 "arrs can only be vector or matrix",
152 ));
153 }
154 }
155
156 if iter.peek().is_some() {
157 s.push(':');
158 }
159 }
160 write!(w, "{}", escape(&s))?;
161 writeln!(w)?;
162
163 // arrays
164 for i in 0..natoms {
165 let mut iter = arrs.iter().peekable();
166 while let Some((_, v)) = iter.next() {
167 let i = i as usize;
168
169 // In legacy the libAtoms/extxyz's c impl the output format to:
170 // #define INTEGER_FMT "%8d"
171 // #define FLOAT_FMT "%16.8f"
172 // #define STRING_FMT "%s"
173 // #define BOOL_FMT "%.1s"
174
175 // new format rule of arr printing is:
176 // - for text, if it <8, pad to width 8 and left align (backward compatible otherwise
177 // shitty libAtoms/extxyz throw segfault), if its len >8 padding len +2
178 // - for float, the single value trimed to .8 precision with width of 16.
179 // - interger %8d
180 // - bool .1s
181
182 // store the columns but write row by row
183 match v {
184 Value::VecInteger(items, _) => write!(w, "{}", items[i])?,
185 Value::VecFloat(items, _) => write!(w, "{:>16.8}", items[i])?,
186 Value::VecBool(items, _) => write!(w, "{}", items[i])?,
187 Value::VecText(items, _) => {
188 let s = &items[i];
189 let sl = (*s).len();
190 if sl > 5 {
191 write!(w, "{1:<0$}", sl + 2, items[i])?
192 } else {
193 write!(w, "{:<5}", items[i])?
194 }
195 }
196 Value::MatrixInteger(items, _) => {
197 let s = &items[i];
198 let indent = " ";
199 let s = s
200 .iter()
201 .map(|i| format!("{i}"))
202 .collect::<Vec<_>>()
203 .join(indent);
204 write!(w, "{s}")?;
205 }
206 Value::MatrixFloat(items, _) => {
207 let s = &items[i];
208 let indent = " ";
209 let s = s
210 .iter()
211 .map(|i| {
212 let s = format!("{:>16.8}", i);
213 s
214 })
215 .collect::<Vec<_>>()
216 .join(indent);
217 write!(w, "{s}")?;
218 }
219 Value::MatrixBool(items, _) => {
220 let s = &items[i];
221 let indent = " ";
222 let s = s
223 .iter()
224 .map(|i| format!("{i}"))
225 .collect::<Vec<_>>()
226 .join(indent);
227 write!(w, "{s}")?;
228 }
229 Value::MatrixText(items, _) => {
230 let s = &items[i];
231 let indent = " ";
232 let s = s
233 .iter()
234 .map(|i| format!("{i}"))
235 .collect::<Vec<_>>()
236 .join(indent);
237 write!(w, "{s}")?;
238 }
239 _ => {
240 // this is unreachable if the inner dict is not create manually
241 return Err(ExtxyzError::InvalidValue(
242 "arrs can only be vector or matrix",
243 ));
244 }
245 }
246
247 if iter.peek().is_some() {
248 write!(w, " ")?; // 1 enforced space
249 }
250 }
251
252 writeln!(w)?;
253 }
254 w.flush()?;
255
256 Ok(())
257}
258
259#[cfg(test)]
260mod tests {
261 use std::io::{BufWriter, Cursor};
262
263 use crate::read_frame;
264
265 use super::*;
266
267 trait FrameNewExample {
268 fn new_example() -> Frame;
269 }
270
271 impl FrameNewExample for Frame {
272 fn new_example() -> Frame {
273 let inp = r#"4
274key1=a key2=a/b key3=a@b key4="a@b"
275Mg -4.25650 3.79180 -2.54123
276C -1.15405 2.86652 -1.26699
277C -5.53758 3.70936 0.63504
278C -7.28250 4.71303 -3.82016
279"#;
280 let mut rd = Cursor::new(inp.as_bytes());
281 read_frame(&mut rd).unwrap()
282 }
283 }
284
285 // // For test printing purpose
286 // struct TFrame(Frame);
287 //
288 // impl std::fmt::Display for TFrame {
289 // fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
290 // let mut buf = Vec::new();
291 // write_frame(&mut buf, &self.0).map_err(|_| std::fmt::Error)?;
292 // let s = std::str::from_utf8(&buf).map_err(|_| std::fmt::Error)?;
293 // f.write_str(s)
294 // }
295 // }
296
297 #[test]
298 fn test_write_frame() {
299 // this is a round trip from a text -> Frame -> text
300 let frame = Frame::new_example();
301 let mut buf = Vec::new();
302 {
303 let mut w = BufWriter::new(&mut buf);
304 write_frame(&mut w, &frame).unwrap();
305 }
306
307 let s = String::from_utf8(buf).unwrap();
308 let expect = r#"4
309key1=a key2=a/b key3=a@b key4=a@b Properties=species:S:1:pos:R:3
310Mg -4.25650000 3.79180000 -2.54123000
311C -1.15405000 2.86652000 -1.26699000
312C -5.53758000 3.70936000 0.63504000
313C -7.28250000 4.71303000 -3.82016000
314"#;
315 assert_eq!(s, expect);
316 }
317
318 #[test]
319 fn test_write_frames_default() {
320 let inp = r#"4
321key1=a key2=a/b key3=a@b key4="a@b"
322Mg -4.25650 3.79180 -2.54123
323C -1.15405 2.86652 -1.26699
324C -5.53758 3.70936 0.63504
325C -7.28250 4.71303 -3.82016
326"#;
327 let rd = Cursor::new(inp.as_bytes());
328 let frame1 = read_frame(&mut rd.clone()).unwrap();
329 let frame2 = read_frame(&mut rd.clone()).unwrap();
330 let frames = vec![frame1, frame2];
331 let mut buf = Vec::new();
332 {
333 let mut w = BufWriter::new(&mut buf);
334 write_frames(&mut w, frames).unwrap();
335 }
336 let s = String::from_utf8(buf).unwrap();
337 let expect = r#"4
338key1=a key2=a/b key3=a@b key4=a@b Properties=species:S:1:pos:R:3
339Mg -4.25650000 3.79180000 -2.54123000
340C -1.15405000 2.86652000 -1.26699000
341C -5.53758000 3.70936000 0.63504000
342C -7.28250000 4.71303000 -3.82016000
3434
344key1=a key2=a/b key3=a@b key4=a@b Properties=species:S:1:pos:R:3
345Mg -4.25650000 3.79180000 -2.54123000
346C -1.15405000 2.86652000 -1.26699000
347C -5.53758000 3.70936000 0.63504000
348C -7.28250000 4.71303000 -3.82016000
349"#;
350 assert_eq!(s, expect);
351 }
352}