1use std::collections::{BTreeMap, BTreeSet, HashMap};
4
5use serde::{Deserialize, Serialize, de::Error as _};
6use serde_json::{Map, Value, json};
7
8use crate::model::ModelPayload;
9use crate::operating::{
10 ElementRef, ElementUpdate, IdentityIndex, apply_update_fields, json_error, payload_key,
11 resolve_update, resolve_update_row, validate_update_fields_survived,
12};
13
14#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
16#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))]
17#[non_exhaustive]
18pub struct StudyBlock {
19 #[serde(default, skip_serializing_if = "Option::is_none")]
20 pub label: Option<String>,
21 #[serde(default, skip_serializing_if = "Option::is_none")]
22 pub author: Option<String>,
23 #[serde(default, skip_serializing_if = "Option::is_none")]
24 pub created_at: Option<String>,
25 #[serde(default, skip_serializing_if = "Option::is_none")]
26 pub base_operating_point: Option<usize>,
27 #[serde(default, skip_serializing_if = "Vec::is_empty")]
28 pub commits: Vec<StudyCommit>,
29 #[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
30 pub app: BTreeMap<String, Value>,
31 #[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
32 pub metadata: BTreeMap<String, Value>,
33}
34
35impl StudyBlock {
36 #[must_use]
37 pub fn is_empty(&self) -> bool {
38 self.label.is_none()
39 && self.author.is_none()
40 && self.created_at.is_none()
41 && self.base_operating_point.is_none()
42 && self.commits.is_empty()
43 && self.app.is_empty()
44 && self.metadata.is_empty()
45 }
46}
47
48#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
50#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))]
51#[non_exhaustive]
52pub struct StudyCommit {
53 #[serde(default, skip_serializing_if = "Option::is_none")]
54 pub label: Option<String>,
55 #[serde(default, skip_serializing_if = "Option::is_none")]
56 pub created_at: Option<String>,
57 #[serde(default, skip_serializing_if = "Vec::is_empty")]
58 pub edits: Vec<StudyEdit>,
59 #[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
60 pub metadata: BTreeMap<String, Value>,
61}
62
63#[derive(Clone, Debug, PartialEq)]
66#[non_exhaustive]
67pub enum StudyEdit {
68 DemandDelta {
69 bus: ElementRef,
70 p_mw: f64,
71 q_mvar: Option<f64>,
72 },
73 RatingDelta {
74 branch: ElementRef,
75 delta_mw: f64,
76 },
77 SetFields {
78 update: ElementUpdate,
79 },
80 Unknown {
81 kind: String,
82 value: Value,
83 },
84}
85
86impl StudyEdit {
87 #[must_use]
88 pub fn kind(&self) -> &str {
89 match self {
90 Self::DemandDelta { .. } => "demand_delta",
91 Self::RatingDelta { .. } => "rating_delta",
92 Self::SetFields { .. } => "set_fields",
93 Self::Unknown { kind, .. } => kind,
94 }
95 }
96}
97
98#[cfg(feature = "schema")]
99impl schemars::JsonSchema for StudyEdit {
100 fn schema_name() -> std::borrow::Cow<'static, str> {
101 "StudyEdit".into()
102 }
103
104 fn json_schema(generator: &mut schemars::SchemaGenerator) -> schemars::Schema {
105 let element_ref = generator.subschema_for::<ElementRef>().to_value();
106 let element_update = generator.subschema_for::<ElementUpdate>().to_value();
107 let schema = json!({
108 "type": "object",
109 "oneOf": [
110 {
111 "type": "object",
112 "required": ["kind", "bus", "p_mw"],
113 "properties": {
114 "kind": { "const": "demand_delta" },
115 "bus": element_ref,
116 "p_mw": { "type": "number" },
117 "q_mvar": { "type": "number" }
118 }
119 },
120 {
121 "type": "object",
122 "required": ["kind", "branch", "delta_mw"],
123 "properties": {
124 "kind": { "const": "rating_delta" },
125 "branch": element_ref,
126 "delta_mw": { "type": "number" }
127 }
128 },
129 {
130 "type": "object",
131 "required": ["kind", "update"],
132 "properties": {
133 "kind": { "const": "set_fields" },
134 "update": element_update
135 }
136 },
137 {
138 "type": "object",
139 "required": ["kind"],
140 "properties": {
141 "kind": {
142 "type": "string",
143 "not": {
144 "enum": ["demand_delta", "rating_delta", "set_fields"]
145 }
146 }
147 },
148 "additionalProperties": true
149 }
150 ]
151 });
152 schemars::Schema::try_from(schema).expect("study edit schema is an object")
153 }
154}
155
156impl Serialize for StudyEdit {
157 fn serialize<S: serde::Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
158 match self {
159 Self::DemandDelta { bus, p_mw, q_mvar } => {
160 #[derive(Serialize)]
161 struct Wire<'a> {
162 kind: &'static str,
163 bus: &'a ElementRef,
164 p_mw: f64,
165 #[serde(skip_serializing_if = "Option::is_none")]
166 q_mvar: Option<f64>,
167 }
168 Wire {
169 kind: "demand_delta",
170 bus,
171 p_mw: *p_mw,
172 q_mvar: *q_mvar,
173 }
174 .serialize(serializer)
175 }
176 Self::RatingDelta { branch, delta_mw } => {
177 #[derive(Serialize)]
178 struct Wire<'a> {
179 kind: &'static str,
180 branch: &'a ElementRef,
181 delta_mw: f64,
182 }
183 Wire {
184 kind: "rating_delta",
185 branch,
186 delta_mw: *delta_mw,
187 }
188 .serialize(serializer)
189 }
190 Self::SetFields { update } => {
191 #[derive(Serialize)]
192 struct Wire<'a> {
193 kind: &'static str,
194 update: &'a ElementUpdate,
195 }
196 Wire {
197 kind: "set_fields",
198 update,
199 }
200 .serialize(serializer)
201 }
202 Self::Unknown { value, .. } => value.serialize(serializer),
203 }
204 }
205}
206
207impl<'de> Deserialize<'de> for StudyEdit {
208 fn deserialize<D: serde::Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
209 let value = Value::deserialize(deserializer)?;
210 let object = value
211 .as_object()
212 .ok_or_else(|| serde::de::Error::custom("study edit must be an object"))?;
213 let kind = object
214 .get("kind")
215 .and_then(Value::as_str)
216 .ok_or_else(|| serde::de::Error::custom("study edit needs string `kind`"))?;
217 match kind {
218 "demand_delta" => {
219 #[derive(Deserialize)]
220 struct Wire {
221 bus: ElementRef,
222 p_mw: f64,
223 #[serde(default)]
224 q_mvar: Option<f64>,
225 }
226 let wire = serde_json::from_value::<Wire>(value).map_err(D::Error::custom)?;
227 Ok(Self::DemandDelta {
228 bus: wire.bus,
229 p_mw: wire.p_mw,
230 q_mvar: wire.q_mvar,
231 })
232 }
233 "rating_delta" => {
234 #[derive(Deserialize)]
235 struct Wire {
236 branch: ElementRef,
237 delta_mw: f64,
238 }
239 let wire = serde_json::from_value::<Wire>(value).map_err(D::Error::custom)?;
240 Ok(Self::RatingDelta {
241 branch: wire.branch,
242 delta_mw: wire.delta_mw,
243 })
244 }
245 "set_fields" => {
246 #[derive(Deserialize)]
247 struct Wire {
248 update: ElementUpdate,
249 }
250 let wire = serde_json::from_value::<Wire>(value).map_err(D::Error::custom)?;
251 Ok(Self::SetFields {
252 update: wire.update,
253 })
254 }
255 other => Ok(Self::Unknown {
256 kind: other.to_owned(),
257 value,
258 }),
259 }
260 }
261}
262
263pub(crate) fn apply_study_to_model(
266 model: &ModelPayload,
267 study: &StudyBlock,
268 commit_index: usize,
269) -> serde_json::Result<(ModelPayload, BTreeSet<String>)> {
270 if !matches!(model, ModelPayload::Balanced { .. }) {
271 return Err(json_error(
272 "STUDY.WRONG_MODEL_KIND: study materialization requires a balanced package",
273 ));
274 }
275 if study.commits.get(commit_index).is_none() {
276 return Err(json_error(format!(
277 "package has no study commit {commit_index}"
278 )));
279 }
280
281 let mut value = serde_json::to_value(model)?;
282 let payload_key = payload_key(model);
283 let payload = value
284 .as_object_mut()
285 .and_then(|root| root.get_mut(payload_key))
286 .and_then(Value::as_object_mut)
287 .ok_or_else(|| json_error(format!("model payload missing `{payload_key}` object")))?;
288
289 let mut indexes = HashMap::new();
290 let mut updated_paths = BTreeSet::new();
291 let mut set_field_updates = Vec::new();
292 let mut set_field_rows = Vec::new();
293 let mut context = StudyApplyContext {
294 payload,
295 payload_key,
296 indexes: &mut indexes,
297 updated_paths: &mut updated_paths,
298 set_field_updates: &mut set_field_updates,
299 set_field_rows: &mut set_field_rows,
300 };
301 for (commit_pos, commit) in study.commits.iter().take(commit_index + 1).enumerate() {
302 for (edit_pos, edit) in commit.edits.iter().enumerate() {
303 context.apply_edit(edit, commit_pos, edit_pos)?;
304 }
305 }
306
307 let updated = serde_json::from_value(value)?;
308 validate_update_fields_survived(&updated, &set_field_updates, &set_field_rows)?;
309 Ok((updated, updated_paths))
310}
311
312pub(crate) fn check_study_identities(
314 model: &ModelPayload,
315 study: &StudyBlock,
316) -> Vec<(usize, usize, String)> {
317 let payload_key = payload_key(model);
318 let payload = match serde_json::to_value(model) {
319 Ok(Value::Object(mut root)) => match root.remove(payload_key) {
320 Some(Value::Object(payload)) => payload,
321 _ => {
322 return vec![(
323 0,
324 0,
325 format!("model payload missing `{payload_key}` object"),
326 )];
327 }
328 },
329 _ => return vec![(0, 0, "model payload did not serialize to object".to_owned())],
330 };
331
332 let mut indexes = HashMap::new();
333 let mut findings = Vec::new();
334 for (commit_pos, commit) in study.commits.iter().enumerate() {
335 for (edit_pos, edit) in commit.edits.iter().enumerate() {
336 let result = match edit {
337 StudyEdit::DemandDelta { bus, .. } => {
338 resolve_update_row(&payload, &mut indexes, bus).map(|_| ())
339 }
340 StudyEdit::RatingDelta { branch, .. } => {
341 resolve_update_row(&payload, &mut indexes, branch).map(|_| ())
342 }
343 StudyEdit::SetFields { update } => {
344 resolve_update(&payload, &mut indexes, update).map(|_| ())
345 }
346 StudyEdit::Unknown { .. } => Ok(()),
347 };
348 if let Err(message) = result {
349 findings.push((commit_pos, edit_pos, message));
350 }
351 }
352 }
353 findings
354}
355
356struct StudyApplyContext<'a> {
357 payload: &'a mut Map<String, Value>,
358 payload_key: &'a str,
359 indexes: &'a mut HashMap<String, IdentityIndex>,
360 updated_paths: &'a mut BTreeSet<String>,
361 set_field_updates: &'a mut Vec<ElementUpdate>,
362 set_field_rows: &'a mut Vec<usize>,
363}
364
365impl StudyApplyContext<'_> {
366 fn apply_edit(
367 &mut self,
368 edit: &StudyEdit,
369 commit_pos: usize,
370 edit_pos: usize,
371 ) -> serde_json::Result<()> {
372 match edit {
373 StudyEdit::DemandDelta { bus, p_mw, q_mvar } => {
374 let bus_row =
375 resolve_update_row(self.payload, self.indexes, bus).map_err(json_error)?;
376 let touched = apply_demand_delta(self.payload, bus_row, *p_mw, *q_mvar)?;
377 for path in touched {
378 self.updated_paths
379 .insert(format!("/model/{}/{path}", self.payload_key));
380 }
381 self.indexes.remove("loads");
382 }
383 StudyEdit::RatingDelta { branch, delta_mw } => {
384 let branch_row =
385 resolve_update_row(self.payload, self.indexes, branch).map_err(json_error)?;
386 let branch = row_object_mut(self.payload, "branches", branch_row)?;
387 let old = number_field(branch, "rate_a", "branch", branch_row)?;
388 branch.insert("rate_a".to_owned(), json!(old + delta_mw));
389 self.updated_paths.insert(format!(
390 "/model/{}/branches/{branch_row}/rate_a",
391 self.payload_key
392 ));
393 }
394 StudyEdit::SetFields { update } => {
395 let row = resolve_update(self.payload, self.indexes, update).map_err(json_error)?;
396 apply_update_fields(self.payload, &update.element.table, row, &update.fields)?;
397 for field in update.fields.keys() {
398 self.updated_paths.insert(format!(
399 "/model/{}/{}/{row}/{field}",
400 self.payload_key, update.element.table
401 ));
402 }
403 self.set_field_updates.push(update.clone());
404 self.set_field_rows.push(row);
405 }
406 StudyEdit::Unknown { kind, .. } => {
407 return Err(json_error(format!(
408 "STUDY.UNKNOWN_EDIT_KIND: study commit {commit_pos} edit {edit_pos} has unsupported kind `{kind}`"
409 )));
410 }
411 }
412 Ok(())
413 }
414}
415
416fn apply_demand_delta(
417 payload: &mut Map<String, Value>,
418 bus_row: usize,
419 p_delta: f64,
420 q_delta: Option<f64>,
421) -> serde_json::Result<Vec<String>> {
422 let (bus_id, bus_uid) = {
423 let bus = row_object(payload, "buses", bus_row)?;
424 let bus_id = bus
425 .get("id")
426 .and_then(Value::as_u64)
427 .ok_or_else(|| json_error(format!("bus row {bus_row} has no numeric `id`")))?;
428 let bus_uid = bus
429 .get("uid")
430 .and_then(Value::as_str)
431 .map_or_else(|| format!("buses:{bus_row}"), str::to_owned);
432 (bus_id, bus_uid)
433 };
434
435 let loads = payload
436 .get_mut("loads")
437 .and_then(Value::as_array_mut)
438 .ok_or_else(|| json_error("balanced payload has no `loads` array"))?;
439 let mut rows = Vec::new();
440 let mut total_p = 0.0;
441 let mut total_q = 0.0;
442 for (row, load) in loads.iter().enumerate() {
443 let Some(load) = load.as_object() else {
444 continue;
445 };
446 if load.get("bus").and_then(Value::as_u64) != Some(bus_id) {
447 continue;
448 }
449 if !load
450 .get("in_service")
451 .and_then(Value::as_bool)
452 .unwrap_or(true)
453 {
454 continue;
455 }
456 let p = load.get("p").and_then(Value::as_f64).unwrap_or(0.0);
457 let q = load.get("q").and_then(Value::as_f64).unwrap_or(0.0);
458 rows.push((row, p, q));
459 total_p += p;
460 total_q += q;
461 }
462
463 if rows.is_empty() || total_p == 0.0 {
464 let q = q_delta.unwrap_or(0.0);
465 loads.push(json!({
466 "bus": bus_id,
467 "p": p_delta,
468 "q": q,
469 "voltage_model": null,
470 "in_service": true,
471 "uid": format!("study:load:{bus_uid}"),
472 "extras": {
473 "study": {
474 "synthetic": true,
475 "source": "demand_delta"
476 }
477 }
478 }));
479 let row = loads.len() - 1;
480 return Ok(vec![
481 format!("loads/{row}/p"),
482 format!("loads/{row}/q"),
483 format!("loads/{row}/uid"),
484 format!("loads/{row}/extras"),
485 ]);
486 }
487
488 let q_delta = q_delta.unwrap_or_else(|| p_delta * total_q / total_p);
489 let mut touched = Vec::new();
490 for (row, p, q) in rows {
491 let p_share = p / total_p;
492 let q_share = if total_q == 0.0 { p_share } else { q / total_q };
493 let load = loads
494 .get_mut(row)
495 .and_then(Value::as_object_mut)
496 .ok_or_else(|| json_error(format!("load row {row} disappeared")))?;
497 load.insert("p".to_owned(), json!(p + p_delta * p_share));
498 load.insert("q".to_owned(), json!(q + q_delta * q_share));
499 touched.push(format!("loads/{row}/p"));
500 touched.push(format!("loads/{row}/q"));
501 }
502 Ok(touched)
503}
504
505fn row_object<'a>(
506 payload: &'a Map<String, Value>,
507 table_name: &str,
508 row: usize,
509) -> serde_json::Result<&'a Map<String, Value>> {
510 payload
511 .get(table_name)
512 .and_then(Value::as_array)
513 .and_then(|table| table.get(row))
514 .and_then(Value::as_object)
515 .ok_or_else(|| json_error(format!("table `{table_name}` has no object row {row}")))
516}
517
518fn row_object_mut<'a>(
519 payload: &'a mut Map<String, Value>,
520 table_name: &str,
521 row: usize,
522) -> serde_json::Result<&'a mut Map<String, Value>> {
523 payload
524 .get_mut(table_name)
525 .and_then(Value::as_array_mut)
526 .and_then(|table| table.get_mut(row))
527 .and_then(Value::as_object_mut)
528 .ok_or_else(|| json_error(format!("table `{table_name}` has no object row {row}")))
529}
530
531fn number_field(
532 object: &Map<String, Value>,
533 field: &str,
534 label: &str,
535 row: usize,
536) -> serde_json::Result<f64> {
537 object
538 .get(field)
539 .and_then(Value::as_f64)
540 .ok_or_else(|| json_error(format!("{label} row {row} has no numeric `{field}` field")))
541}