Skip to main content

powerio_matrix/io/
sensitivity.rs

1//! Streamed Matrix Market output for DC sensitivity matrices.
2
3use std::fs::File;
4use std::io::{BufReader, BufWriter, Write};
5use std::path::{Path, PathBuf};
6use std::time::{SystemTime, UNIX_EPOCH};
7
8use crate::Result;
9use crate::indexed::IndexedNetwork;
10use crate::matrix::sensitivity::for_each_ptdf_lodf_entry;
11use crate::matrix::{SensitivityMetadata, SensitivityOptions};
12
13/// Write PTDF and LODF Matrix Market files from the option based sensitivity
14/// path and return the metadata for the same entries. The iterative solver path
15/// streams retained coordinates through temp files, so it does not keep the
16/// full sparse output in memory.
17pub fn write_sensitivity_mtx_with_options(
18    case: &IndexedNetwork,
19    options: &SensitivityOptions,
20    ptdf_path: impl AsRef<Path>,
21    lodf_path: impl AsRef<Path>,
22) -> Result<SensitivityMetadata> {
23    let mut ptdf = CoordinateMtxWriter::new(ptdf_path.as_ref())?;
24    let mut lodf = CoordinateMtxWriter::new(lodf_path.as_ref())?;
25
26    let metadata = match for_each_ptdf_lodf_entry(
27        case,
28        options,
29        |row, col, value| ptdf.write_entry(row, col, value),
30        |row, col, value| lodf.write_entry(row, col, value),
31    ) {
32        Ok(metadata) => metadata,
33        Err(err) => {
34            ptdf.cleanup();
35            lodf.cleanup();
36            return Err(err);
37        }
38    };
39
40    ptdf.finish(metadata.ptdf.rows, metadata.ptdf.cols)?;
41    lodf.finish(metadata.lodf.rows, metadata.lodf.cols)?;
42    Ok(metadata)
43}
44
45struct CoordinateMtxWriter {
46    target_path: PathBuf,
47    body_path: PathBuf,
48    final_tmp_path: PathBuf,
49    body: Option<BufWriter<File>>,
50    nnz: usize,
51}
52
53impl CoordinateMtxWriter {
54    fn new(target_path: &Path) -> Result<Self> {
55        let body_path = temp_path(target_path, "body");
56        let final_tmp_path = temp_path(target_path, "final");
57        let body = BufWriter::new(File::create(&body_path)?);
58        Ok(Self {
59            target_path: target_path.to_path_buf(),
60            body_path,
61            final_tmp_path,
62            body: Some(body),
63            nnz: 0,
64        })
65    }
66
67    fn write_entry(&mut self, row: usize, col: usize, value: f64) -> Result<()> {
68        if value == 0.0 {
69            return Ok(());
70        }
71        let body = self
72            .body
73            .as_mut()
74            .expect("coordinate writer body is open before finish");
75        writeln!(body, "{} {} {:.16e}", row + 1, col + 1, value)?;
76        self.nnz += 1;
77        Ok(())
78    }
79
80    fn finish(mut self, rows: usize, cols: usize) -> Result<()> {
81        if let Some(mut body) = self.body.take() {
82            body.flush()?;
83        }
84
85        let mut out = BufWriter::new(File::create(&self.final_tmp_path)?);
86        writeln!(out, "%%MatrixMarket matrix coordinate real general")?;
87        writeln!(out, "% written by powerio")?;
88        writeln!(out, "{rows} {cols} {}", self.nnz)?;
89        let mut body = BufReader::new(File::open(&self.body_path)?);
90        std::io::copy(&mut body, &mut out)?;
91        out.flush()?;
92
93        std::fs::rename(&self.final_tmp_path, &self.target_path)?;
94        let _ = std::fs::remove_file(&self.body_path);
95        Ok(())
96    }
97
98    fn cleanup(&mut self) {
99        if let Some(mut body) = self.body.take() {
100            let _ = body.flush();
101        }
102        let _ = std::fs::remove_file(&self.body_path);
103        let _ = std::fs::remove_file(&self.final_tmp_path);
104    }
105}
106
107fn temp_path(target_path: &Path, suffix: &str) -> PathBuf {
108    let pid = std::process::id();
109    let nanos = SystemTime::now()
110        .duration_since(UNIX_EPOCH)
111        .map_or(0, |duration| duration.as_nanos());
112    let name = target_path
113        .file_name()
114        .map_or_else(|| "matrix".into(), |name| name.to_string_lossy());
115    target_path.with_file_name(format!(".{name}.{pid}.{nanos}.{suffix}.tmp"))
116}