Skip to content

Commit

Permalink
feat: Add --network-only argument to CLI convert command.
Browse files Browse the repository at this point in the history 
This uses the update pywr-v1-schema. The error handling in the
convert command is also improved.
  • Loading branch information
@jetuk
jetuk committed Apr 26, 2024
1 parent 90e42b3 commit b5b14ee
Show file tree
Hide file tree
Showing 3 changed files with 163 additions and 86 deletions.
2 changes: 1 addition & 1 deletion Cargo.toml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -46,6 +46,6 @@ pyo3-log = "0.10"
tracing = { version = "0.1", features = ["log"] }
csv = "1.1"
hdf5 = { git = "https://github.com/aldanor/hdf5-rust.git", package = "hdf5", features = ["static", "zlib"] }
pywr-v1-schema = { git = "https://github.com/pywr/pywr-schema/", tag = "v0.12.0", package = "pywr-schema" }
pywr-v1-schema = { git = "https://github.com/pywr/pywr-schema/", tag = "v0.13.0", package = "pywr-schema" }
chrono = { version = "0.4.34" }
schemars = { version = "0.8.16", features = ["chrono"] }
117 changes: 86 additions & 31 deletions pywr-cli/src/main.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@ mod tracing;

use crate::tracing::setup_tracing;
use ::tracing::info;
use anyhow::{Context, Result};
use anyhow::{bail, Context, Result};
use clap::{Parser, Subcommand, ValueEnum};
#[cfg(feature = "ipm-ocl")]
use pywr_core::solvers::{ClIpmF32Solver, ClIpmF64Solver, ClIpmSolverSettings};
Expand All @@ -12,7 +12,8 @@ use pywr_core::solvers::{HighsSolver, HighsSolverSettings};
#[cfg(feature = "ipm-simd")]
use pywr_core::solvers::{SimdIpmF64Solver, SimdIpmSolverSettings};
use pywr_core::test_utils::make_random_model;
use pywr_schema::model::{PywrModel, PywrMultiNetworkModel};
use pywr_schema::model::{PywrModel, PywrMultiNetworkModel, PywrNetwork};
use pywr_schema::ConversionError;
use rand::SeedableRng;
use rand_chacha::ChaCha8Rng;
use schemars::schema_for;
Expand Down Expand Up @@ -62,10 +63,16 @@ struct Cli {
enum Commands {
Convert {
/// Path to Pywr v1.x JSON.
model: PathBuf,
input: PathBuf,
/// Path to output Pywr v2 JSON.
// TODO support printing to stdout?
output: PathBuf,
/// Stop if there is an error converting the model.
#[arg(short, long, default_value_t = false)]
stop_on_error: bool,
/// Convert only the network schema.
#[arg(short, long, default_value_t = false)]
network_only: bool,
},

Run {
Expand Down Expand Up @@ -122,7 +129,12 @@ fn main() -> Result<()> {

match &cli.command {
Some(command) => match command {
Commands::Convert { model, stop_on_error } => convert(model, *stop_on_error),
Commands::Convert {
input,
output,
stop_on_error,
network_only,
} => convert(input, output, *stop_on_error, *network_only)?,
Commands::Run {
model,
solver,
Expand Down Expand Up @@ -153,51 +165,94 @@ fn main() -> Result<()> {
Ok(())
}

fn convert(path: &Path, stop_on_error: bool) {
if path.is_dir() {
for entry in path.read_dir().expect("read_dir call failed").flatten() {
fn convert(in_path: &Path, out_path: &Path, stop_on_error: bool, network_only: bool) -> Result<()> {
if in_path.is_dir() {
if !out_path.is_dir() {
bail!("Output path must be an existing directory when input path is a directory");
}

for entry in in_path
.read_dir()
.with_context(|| format!("Failed to read directory: {:?}", in_path))?
.flatten()
{
let path = entry.path();
if path.is_file()
&& (path.extension().unwrap() == "json")
&& (!path.file_stem().unwrap().to_str().unwrap().contains("_v2"))
{
v1_to_v2(&path, stop_on_error);

if path.is_file() {
if let Some(ext) = path.extension() {
if ext == "json" {
let out_fn = out_path.join(
path.file_name()
.with_context(|| "Failed to determine output filename.".to_string())?,
);

v1_to_v2(&path, &out_fn, stop_on_error, network_only)?;
}
}
}
}
} else {
v1_to_v2(path, stop_on_error);
if out_path.is_dir() {
bail!("Output path must be a file when input path is a file");
}

v1_to_v2(in_path, out_path, stop_on_error, network_only)?;
}

Ok(())
}

fn v1_to_v2(path: &Path, stop_on_error: bool) {
info!("Model: {}", path.display());
fn v1_to_v2(in_path: &Path, out_path: &Path, stop_on_error: bool, network_only: bool) -> Result<()> {
info!("Converting file: {}", in_path.display());

let data = std::fs::read_to_string(path).unwrap();
// Load the v1 schema
let schema: pywr_v1_schema::PywrModel = serde_json::from_str(data.as_str()).unwrap();
// Convert to v2 schema and collect any errors
let (schema_v2, errors) = PywrModel::from_v1(schema);
let data = std::fs::read_to_string(in_path).with_context(|| format!("Failed to read file: {:?}", in_path))?;

if network_only {
let schema: pywr_v1_schema::PywrNetwork = serde_json::from_str(data.as_str())
.with_context(|| format!("Failed deserialise Pywr v1 network file: {:?}", in_path))?;
// Convert to v2 schema and collect any errors
let (schema_v2, errors) = PywrNetwork::from_v1(schema);

handle_conversion_errors(&errors, stop_on_error)?;

std::fs::write(
out_path,
serde_json::to_string_pretty(&schema_v2).with_context(|| "Failed serialise Pywr v2 network".to_string())?,
)
.with_context(|| format!("Failed to write file: {:?}", out_path))?;
} else {
// Load the v1 schema
let schema: pywr_v1_schema::PywrModel = serde_json::from_str(data.as_str())
.with_context(|| format!("Failed deserialise Pywr v1 model file: {:?}", in_path))?;
// Convert to v2 schema and collect any errors
let (schema_v2, errors) = PywrModel::from_v1(schema);

handle_conversion_errors(&errors, stop_on_error)?;

std::fs::write(
out_path,
serde_json::to_string_pretty(&schema_v2).with_context(|| "Failed serialise Pywr v2 model".to_string())?,
)
.with_context(|| format!("Failed to write file: {:?}", out_path))?;
}

Ok(())
}

fn handle_conversion_errors(errors: &[ConversionError], stop_on_error: bool) -> Result<()> {
if !errors.is_empty() {
info!("Model converted with {} errors:", errors.len());
info!("File converted with {} errors:", errors.len());
for error in errors {
info!(" {}", error);
}
if stop_on_error {
return;
bail!("File conversion failed with at-least one error!");
}
} else {
info!("Model converted with zero errors!");
info!("File converted with zero errors!");
}

// There must be a better way to do this!!
let mut new_file_name = path.file_stem().unwrap().to_os_string();
new_file_name.push("_v2");
let mut new_file_name = PathBuf::from(new_file_name);
new_file_name.set_extension("json");
let new_file_pth = path.parent().unwrap().join(new_file_name);

std::fs::write(new_file_pth, serde_json::to_string_pretty(&schema_v2).unwrap()).unwrap();
Ok(())
}

fn run(path: &Path, solver: &Solver, data_path: Option<&Path>, output_path: Option<&Path>) {
Expand Down
130 changes: 76 additions & 54 deletions pywr-schema/src/model.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -176,6 +176,80 @@ impl PywrNetwork {
Ok(serde_json::from_str(data.as_str())?)
}

/// Convert a v1 network to a v2 network.
///
/// This function is used to convert a v1 model to a v2 model. The conversion is not always
/// possible and may result in errors. The errors are returned as a vector of [`ConversionError`]s.
/// alongside the (partially) converted model. This may result in a model that will not
/// function as expected. The user should check the errors and the converted model to ensure
/// that the conversion has been successful.
pub fn from_v1(v1: pywr_v1_schema::PywrNetwork) -> (Self, Vec<ConversionError>) {
let mut errors = Vec::new();

// Extract nodes and any timeseries data from the v1 nodes
let nodes_and_ts: Vec<NodeAndTimeseries> = match v1.nodes {
Some(nodes) => nodes
.into_iter()
.filter_map(|n| match n.try_into() {
Ok(n) => Some(n),
Err(e) => {
errors.push(e);
None
}
})
.collect::<Vec<_>>(),
None => Vec::new(),
};

let mut ts_data = nodes_and_ts
.iter()
.filter_map(|n| n.timeseries.clone())
.flatten()
.collect::<Vec<_>>();

let nodes = nodes_and_ts.into_iter().map(|n| n.node).collect::<Vec<_>>();

let edges = match v1.edges {
Some(edges) => edges.into_iter().map(|e| e.into()).collect(),
None => Vec::new(),
};

let parameters = if let Some(v1_parameters) = v1.parameters {
let mut unnamed_count: usize = 0;
let (parameters, param_ts_data) =
convert_parameter_v1_to_v2(v1_parameters, &mut unnamed_count, &mut errors);
ts_data.extend(param_ts_data);
Some(parameters)
} else {
None
};

let timeseries = if !ts_data.is_empty() {
let ts = convert_from_v1_data(ts_data, &v1.tables, &mut errors);
Some(ts)
} else {
None
};

// TODO convert v1 tables!
let tables = None;
let outputs = None;
let metric_sets = None;

(
Self {
nodes,
edges,
parameters,
tables,
timeseries,
metric_sets,
outputs,
},
errors,
)
}

pub fn get_node_by_name(&self, name: &str) -> Option<&Node> {
self.nodes.iter().find(|n| n.name() == name)
}
Expand Down Expand Up @@ -457,60 +531,8 @@ impl PywrModel {

let timestepper = v1.timestepper.into();

// Extract nodes and any timeseries data from the v1 nodes
let nodes_and_ts: Vec<NodeAndTimeseries> = v1
.nodes
.clone()
.into_iter()
.filter_map(|n| match n.try_into() {
Ok(n) => Some(n),
Err(e) => {
errors.push(e);
None
}
})
.collect::<Vec<_>>();

let mut ts_data = nodes_and_ts
.iter()
.filter_map(|n| n.timeseries.clone())
.flatten()
.collect::<Vec<_>>();

let nodes = nodes_and_ts.into_iter().map(|n| n.node).collect::<Vec<_>>();

let edges = v1.edges.into_iter().map(|e| e.into()).collect();

let parameters = if let Some(v1_parameters) = v1.parameters {
let mut unnamed_count: usize = 0;
let (parameters, param_ts_data) =
convert_parameter_v1_to_v2(v1_parameters, &mut unnamed_count, &mut errors);
ts_data.extend(param_ts_data);
Some(parameters)
} else {
None
};

let timeseries = if !ts_data.is_empty() {
let ts = convert_from_v1_data(ts_data, &v1.tables, &mut errors);
Some(ts)
} else {
None
};

// TODO convert v1 tables!
let tables = None;
let outputs = None;
let metric_sets = None;
let network = PywrNetwork {
nodes,
edges,
parameters,
tables,
timeseries,
metric_sets,
outputs,
};
let (network, network_errors) = PywrNetwork::from_v1(v1.network);
errors.extend(network_errors);

(
Self {
Expand Down

0 comments on commit b5b14ee

Please sign in to comment.