rs_route 0.2.0

A cli tool for muskingum-cunge routing in ngiab
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240

use crate::io::results::SimulationResults;
use anyhow::{Context, Result};
use chrono::NaiveDateTime;
use netcdf::{self, FileMut};
use std::path::PathBuf;
use std::sync::{Arc, Mutex};

pub fn init_netcdf_output(
    output_dir: PathBuf,
    filename: &str,
    _num_flowpaths: usize,
    timesteps: Vec<f64>,
    reference_time: &NaiveDateTime,
) -> Result<Arc<Mutex<FileMut>>> {
    // Create NetCDF file
    let mut file = netcdf::create(output_dir.join(filename))
        .with_context(|| format!("Failed to create NetCDF file: {}", filename))?;

    // Add dimensions
    file.add_dimension("feature_id", 0)
        .context("Failed to add feature_id dimension")?;
    file.add_dimension("time", timesteps.len())
        .context("Failed to add time dimension")?;

    // Add variables
    // Time variable
    let mut time_var = file
        .add_variable::<f64>("time", &["time"])
        .context("Failed to add time variable")?;
    time_var.put_attribute("_FillValue", -9999.0)?;
    time_var.put_attribute("long_name", "valid output time")?;
    time_var.put_attribute("standard_name", "time")?;
    time_var.put_attribute(
        "units",
        format!(
            "seconds since {}",
            reference_time.format("%Y-%m-%d %H:%M:%S")
        ),
    )?;
    time_var.put_attribute("missing_value", -9999.0)?;
    time_var
        .put_values(&timesteps, ..)
        .context("Failed to write time values")?;

    // Feature ID variable
    let mut feature_var = file
        .add_variable::<i64>("feature_id", &["feature_id"])
        .context("Failed to add feature_id variable")?;
    feature_var.put_attribute("long_name", "Segment ID")?;

    // Flow variable
    let mut flow_var = file
        .add_variable::<f32>("flow", &["feature_id", "time"])
        .context("Failed to add flow variable")?;
    flow_var.put_attribute("_FillValue", -9999.0f32)?;
    flow_var.put_attribute("long_name", "Flow")?;
    flow_var.put_attribute("units", "m3 s-1")?;
    flow_var.put_attribute("missing_value", -9999.0f32)?;

    // Velocity variable
    let mut velocity_var = file
        .add_variable::<f32>("velocity", &["feature_id", "time"])
        .context("Failed to add velocity variable")?;
    velocity_var.put_attribute("_FillValue", -9999.0f32)?;
    velocity_var.put_attribute("long_name", "Velocity")?;
    velocity_var.put_attribute("units", "m/s")?;
    velocity_var.put_attribute("missing_value", -9999.0f32)?;

    // Depth variable
    let mut depth_var = file
        .add_variable::<f32>("depth", &["feature_id", "time"])
        .context("Failed to add depth variable")?;
    depth_var.put_attribute("_FillValue", -9999.0f32)?;
    depth_var.put_attribute("long_name", "Depth")?;
    depth_var.put_attribute("units", "m")?;
    depth_var.put_attribute("missing_value", -9999.0f32)?;

    // Global attributes
    file.add_attribute("TITLE", "OUTPUT FROM RS-ROUTE")?;
    file.add_attribute(
        "file_reference_time",
        reference_time.format("%Y-%m-%d_%H:%M:%S").to_string(),
    )?;
    file.add_attribute("code_version", "")?;

    // Additional expected variables
    let _ = file.add_variable::<f32>("type", &["feature_id"])?;
    let _ = file.add_variable::<f32>("nudge", &["feature_id"])?;

    Ok(Arc::new(Mutex::new(file)))
}

pub fn write_batch(
    output_file: &Arc<Mutex<FileMut>>,
    batch: &[Arc<SimulationResults>],
) -> Result<()> {
    let mut file = output_file
        .lock()
        .map_err(|e| anyhow::anyhow!("Failed to acquire NetCDF file lock: {}", e))?;

    // Get current index once
    let feature_var = file
        .variable("feature_id")
        .ok_or_else(|| anyhow::anyhow!("feature_id variable not found"))?;
    let start_idx = feature_var.len();

    // Prepare all data arrays
    let mut all_feature_ids = Vec::with_capacity(batch.len());
    let mut all_flows = Vec::new();
    let mut all_velocities = Vec::new();
    let mut all_depths = Vec::new();

    let expected_timesteps = file
        .dimension("time")
        .ok_or_else(|| anyhow::anyhow!("time dimension not found"))?
        .len();

    for results in batch {
        all_feature_ids.push(results.feature_id);

        // Downsample data
        let actual_timesteps = results.flow_data.len();
        let downsampling = actual_timesteps / expected_timesteps;

        for i in (downsampling - 1..actual_timesteps).step_by(downsampling) {
            all_flows.push(results.flow_data[i]);
            all_velocities.push(results.velocity_data[i]);
            all_depths.push(results.depth_data[i]);
        }
    }

    // Write all feature IDs at once
    let mut feature_var = file
        .variable_mut("feature_id")
        .ok_or_else(|| anyhow::anyhow!("feature_id variable not found"))?;
    feature_var
        .put_values(&all_feature_ids, start_idx..)
        .context("Failed to write feature_ids")?;

    // Reshape and write flow data
    let flow_2d: Vec<Vec<f32>> = all_flows
        .chunks(expected_timesteps)
        .map(|chunk| chunk.to_vec())
        .collect();

    let mut flow_var = file
        .variable_mut("flow")
        .ok_or_else(|| anyhow::anyhow!("flow variable not found"))?;
    for (i, row) in flow_2d.iter().enumerate() {
        flow_var.put_values(row, (start_idx + i, ..))?;
    }

    // Similar for velocity and depth
    let velocity_2d: Vec<Vec<f32>> = all_velocities
        .chunks(expected_timesteps)
        .map(|chunk| chunk.to_vec())
        .collect();

    let mut velocity_var = file
        .variable_mut("velocity")
        .ok_or_else(|| anyhow::anyhow!("velocity variable not found"))?;
    for (i, row) in velocity_2d.iter().enumerate() {
        velocity_var.put_values(row, (start_idx + i, ..))?;
    }

    let depth_2d: Vec<Vec<f32>> = all_depths
        .chunks(expected_timesteps)
        .map(|chunk| chunk.to_vec())
        .collect();

    let mut depth_var = file
        .variable_mut("depth")
        .ok_or_else(|| anyhow::anyhow!("depth variable not found"))?;
    for (i, row) in depth_2d.iter().enumerate() {
        depth_var.put_values(row, (start_idx + i, ..))?;
    }

    Ok(())
}

// Function to write results to NetCDF
pub fn write_output(
    output_file: &Arc<Mutex<FileMut>>,
    results: &Arc<SimulationResults>,
) -> Result<()> {
    // Get lock on file
    let mut file = output_file
        .lock()
        .map_err(|e| anyhow::anyhow!("Failed to acquire NetCDF file lock: {}", e))?;

    // figure out the downsampling that needs to be done
    let expected_timesteps = file
        .dimension("time")
        .ok_or_else(|| anyhow::anyhow!("time dimension not found"))?;
    let actual_timesteps = results.flow_data.len();
    let downsampling = actual_timesteps / expected_timesteps.len();
    let mut downsampled_flow_data = Vec::with_capacity(expected_timesteps.len());
    let mut downsampled_velocity_data = Vec::with_capacity(expected_timesteps.len());
    let mut downsampled_depth_data = Vec::with_capacity(expected_timesteps.len());
    for i in (downsampling - 1..actual_timesteps).step_by(downsampling) {
        downsampled_flow_data.push(results.flow_data[i]);
        downsampled_velocity_data.push(results.velocity_data[i]);
        downsampled_depth_data.push(results.depth_data[i]);
    }

    // Get feature variable
    let mut feature_var = file
        .variable_mut("feature_id")
        .ok_or_else(|| anyhow::anyhow!("feature_id variable not found"))?;
    let fidx = feature_var.len();
    feature_var
        .put_value(results.feature_id, fidx)
        .context("Failed to write feature_id")?;

    // Flow variable
    let mut flow_var = file
        .variable_mut("flow")
        .ok_or_else(|| anyhow::anyhow!("flow variable not found"))?;
    flow_var
        .put_values(&downsampled_flow_data, (fidx, ..))
        .context("Failed to write flow data")?;

    // Velocity variable
    let mut velocity_var = file
        .variable_mut("velocity")
        .ok_or_else(|| anyhow::anyhow!("velocity variable not found"))?;
    velocity_var
        .put_values(&downsampled_velocity_data, (fidx, ..))
        .context("Failed to write velocity data")?;

    // Depth variable
    let mut depth_var = file
        .variable_mut("depth")
        .ok_or_else(|| anyhow::anyhow!("depth variable not found"))?;
    depth_var
        .put_values(&downsampled_depth_data, (fidx, ..))
        .context("Failed to write depth data")?;

    Ok(())
}