Skip to main content

powerio_pkg/
study.rs

1//! Cumulative study edits for `.pio.json` packages.
2
3use 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/// Additive study block stored on a package envelope.
15#[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/// One cumulative commit in a study block.
49#[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/// One study edit. Unknown edit kinds are preserved and rejected only when a
64/// caller tries to materialize them.
65#[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
263/// Apply commits `0..=commit_index` to a balanced model and return the updated
264/// model plus JSON Pointer paths touched by the edits.
265pub(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
312/// Dry run identity resolution for every known study edit.
313pub(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}