nyx_space/od/process/solution/
import.rs

1/*
2    Nyx, blazing fast astrodynamics
3    Copyright (C) 2018-onwards Christopher Rabotin <christopher.rabotin@gmail.com>
4
5    This program is free software: you can redistribute it and/or modify
6    it under the terms of the GNU Affero General Public License as published
7    by the Free Software Foundation, either version 3 of the License, or
8    (at your option) any later version.
9
10    This program is distributed in the hope that it will be useful,
11    but WITHOUT ANY WARRANTY; without even the implied warranty of
12    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
13    GNU Affero General Public License for more details.
14
15    You should have received a copy of the GNU Affero General Public License
16    along with this program.  If not, see <https://www.gnu.org/licenses/>.
17*/
18
19// Add this code within the impl block for ODSolution,
20// potentially in a new file like `src/od/process/solution/import.rs`
21// and ensure necessary imports are present.
22
23use crate::io::{ArrowSnafu, InputOutputError, MissingDataSnafu, ParquetSnafu, StdIOSnafu};
24use crate::linalg::allocator::Allocator;
25use crate::linalg::{Const, DefaultAllocator, DimName, OMatrix, OVector, SMatrix};
26use crate::od::estimate::*;
27use crate::od::msr::MeasurementType;
28use crate::od::*;
29use crate::Spacecraft;
30use anise::frames::Frame;
31use anise::prelude::Orbit;
32use anise::structure::spacecraft::{DragData, Mass, SRPData};
33use arrow::array::RecordBatchReader;
34use arrow::array::{Array, BooleanArray, Float64Array, StringArray};
35use hifitime::Epoch;
36use indexmap::IndexSet;
37use log::{info, warn};
38use msr::sensitivity::TrackerSensitivity;
39use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;
40use snafu::prelude::*;
41use std::collections::{BTreeMap, HashMap};
42use std::fs::File;
43use std::path::Path;
44use std::str::FromStr;
45use std::sync::Arc;
46
47use super::ODSolution; // Needed for StateType bounds
48
49// --- Function Definition ---
50
51impl<MsrSize, Trk> ODSolution<Spacecraft, KfEstimate<Spacecraft>, MsrSize, Trk>
52where
53    MsrSize: DimName + std::fmt::Debug + Clone, // Added Debug+Clone for error messages/vec construction
54    Trk: TrackerSensitivity<Spacecraft, Spacecraft> + Clone, // Added Clone for devices
55    // Bounds needed for KfEstimate and Spacecraft
56    DefaultAllocator: Allocator<MsrSize>
57        + Allocator<MsrSize, <Spacecraft as State>::Size>
58        + Allocator<Const<1>, MsrSize>
59        + Allocator<<Spacecraft as State>::Size>
60        + Allocator<<Spacecraft as State>::Size, <Spacecraft as State>::Size>
61        + Allocator<MsrSize, MsrSize>
62        + Allocator<<Spacecraft as State>::Size, MsrSize>
63        + Allocator<<Spacecraft as State>::VecLength>,
64    <DefaultAllocator as Allocator<<Spacecraft as State>::VecLength>>::Buffer<f64>: Send,
65    <DefaultAllocator as Allocator<<Spacecraft as State>::Size>>::Buffer<f64>: Copy,
66    <DefaultAllocator as Allocator<<Spacecraft as State>::Size, <Spacecraft as State>::Size>>::Buffer<f64>: Copy,
67
68{
69    /// Loads an OD solution from a Parquet file created by `ODSolution::to_parquet`.
70    ///
71    /// The `devices` map must be provided by the caller as it contains potentially complex
72    /// state (like Almanac references) that isn't serialized in the Parquet file.
73    ///
74    /// Note: This function currently assumes the StateType is `Spacecraft` and the
75    /// estimate type is `KfEstimate<Spacecraft>`.
76    pub fn from_parquet<P: AsRef<Path>>(
77        path: P,
78        devices: BTreeMap<String, Trk>,
79    ) -> Result<Self, InputOutputError> {
80
81
82     let file = File::open(&path).context(StdIOSnafu {
83          action: "opening OD solution file",
84      })?;
85
86      let builder = ParquetRecordBatchReaderBuilder::try_new(file).unwrap();
87
88      let mut metadata = HashMap::new();
89      // Build the custom metadata
90      if let Some(file_metadata) = builder.metadata().file_metadata().key_value_metadata() {
91          for key_value in file_metadata {
92              if !key_value.key.starts_with("ARROW:") {
93                  metadata.insert(
94                      key_value.key.clone(),
95                      key_value.value.clone().unwrap_or("[unset]".to_string()),
96                  );
97              }
98          }
99      }
100
101      // Check the schema
102      let mut has_epoch = false; // Required
103      let mut frame = None;
104      let mut srp_area_m2 = None;
105      let mut drag_area_m2 = None;
106
107      let schema = builder.schema().clone();
108
109      let reader = builder.build().context(ParquetSnafu {
110          action: "building Parquet reader for OD results",
111      })?;
112
113      for field in &reader.schema().fields {
114          if field.name().as_str() == "Epoch (UTC)" {
115              has_epoch = true;
116          } else {
117            if let Some(frame_info) = field.metadata().get("Frame") {
118                // Frame is expected to be serialized as Dhall.
119                match serde_dhall::from_str(frame_info).parse::<Frame>() {
120                    Err(e) => {
121                        return Err(InputOutputError::ParseDhall {
122                            data: frame_info.to_string(),
123                            err: format!("{e}"),
124                        })
125                    }
126                    Ok(deser_frame) => frame = Some(deser_frame),
127                };
128            }
129            if let Some(info) = field.metadata().get("SRP Area (m2)") {
130                srp_area_m2 = Some(info.parse::<f64>().unwrap_or(0.0));
131            }
132            if let Some(info) = field.metadata().get("Drag Area (m2)"){
133                drag_area_m2 = Some(info.parse::<f64>().unwrap_or(0.0));
134            }
135        }
136      }
137
138      ensure!(
139          has_epoch,
140          MissingDataSnafu {
141              which: "Epoch (UTC)"
142          }
143      );
144
145      ensure!(
146          frame.is_some(),
147          MissingDataSnafu {
148              which: "Frame in metadata"
149          }
150      );
151
152
153        let mut estimates: Vec<KfEstimate<Spacecraft>> = Vec::new();
154        let mut residuals: Vec<Option<Residual<MsrSize>>> = Vec::new();
155        let mut gains: Vec<Option<OMatrix<f64, <Spacecraft as State>::Size, MsrSize>>> = Vec::new();
156        let mut filter_smoother_ratios: Vec<Option<OVector<f64, <Spacecraft as State>::Size>>> =
157            Vec::new();
158        let mut measurement_types_found = IndexSet::new();
159
160        let state_size = <Spacecraft as State>::Size::DIM;
161
162        // State item names used in column naming
163        let state_items = ["X", "Y", "Z", "Vx", "Vy", "Vz", "Cr", "Cd", "Mass"];
164        let mut cov_units = vec![];
165
166        for i in 0..state_items.len() {
167            for j in i..state_items.len() {
168                let cov_unit = if i < 3 {
169                    if j < 3 {
170                        "km^2"
171                    } else if (3..6).contains(&j) {
172                        "km^2/s"
173                    } else if j == 8 {
174                        "km*kg"
175                    } else {
176                        "km"
177                    }
178                } else if (3..6).contains(&i) {
179                    if (3..6).contains(&j) {
180                        "km^2/s^2"
181                    } else if j == 8 {
182                        "km/s*kg"
183                    } else {
184                        "km/s"
185                    }
186                } else if i == 8 || j == 8 {
187                    "kg^2"
188                } else {
189                    "unitless"
190                };
191
192                cov_units.push(cov_unit);
193            }
194        }
195
196        // --- Pre-parse Measurement Types from Schema ---
197        // Infer measurement types from residual column names
198        for field in schema.fields() {
199             if let Some(msr_type_str) = field.name().strip_prefix("Prefit residual: ") {
200                 if let Some(bracket_pos) = msr_type_str.find(" (") {
201                     let type_name = &msr_type_str[..bracket_pos];
202                     if let Ok(msr_type) = MeasurementType::from_str(type_name) {
203                          measurement_types_found.insert(msr_type);
204                     } else {
205                         warn!("Could not parse measurement type from column: {}", field.name());
206                     }
207                 }
208             }
209        }
210        if measurement_types_found.is_empty() {
211             warn!("Could not automatically detect any measurement types from Parquet column names. Residuals may not be loaded correctly.");
212        } else {
213             info!("Detected measurement types: {measurement_types_found:?}");
214        }
215
216
217
218        // while let Some(record_batch) = reader.next() {
219        for record_batch in reader {
220            let batch = record_batch.context(ArrowSnafu {
221                action: "reading record batch from OD results",
222            })?;
223
224            let num_rows = batch.num_rows();
225
226            // --- Helper to get column data ---
227            let get_col = |name: &str| -> Result<Arc<dyn Array>, InputOutputError> {
228                batch
229                    .column_by_name(name)
230                    .ok_or_else(|| InputOutputError::MissingData {
231                        which: format!("column '{name}' in OD results"),
232                    })
233                    .cloned() // Clone the Arc<dyn Array>
234            };
235
236            // --- Extract Columns (handle potential errors) ---
237
238            let epoch_col = get_col("Epoch (UTC)")?
239                .as_any()
240                .downcast_ref::<StringArray>()
241                .ok_or_else(|| InputOutputError::ArrowError {
242                     action: "downcasting Epoch column",
243                     source: arrow::error::ArrowError::CastError("Could not cast Epoch to StringArray".to_string()),
244                 })?.clone(); // Clone the concrete array
245
246            // State component columns
247            let x_col = get_col("x (km)")?.as_any().downcast_ref::<Float64Array>().ok_or_else(|| InputOutputError::ArrowError{action: "downcasting X", source: arrow::error::ArrowError::CastError("".to_string())})?.clone();
248            let y_col = get_col("y (km)")?.as_any().downcast_ref::<Float64Array>().ok_or_else(|| InputOutputError::ArrowError{action: "downcasting Y", source: arrow::error::ArrowError::CastError("".to_string())})?.clone();
249            let z_col = get_col("z (km)")?.as_any().downcast_ref::<Float64Array>().ok_or_else(|| InputOutputError::ArrowError{action: "downcasting Z", source: arrow::error::ArrowError::CastError("".to_string())})?.clone();
250            let vx_col = get_col("vx (km/s)")?.as_any().downcast_ref::<Float64Array>().ok_or_else(|| InputOutputError::ArrowError{action: "downcasting VX", source: arrow::error::ArrowError::CastError("".to_string())})?.clone();
251            let vy_col = get_col("vy (km/s)")?.as_any().downcast_ref::<Float64Array>().ok_or_else(|| InputOutputError::ArrowError{action: "downcasting VY", source: arrow::error::ArrowError::CastError("".to_string())})?.clone();
252            let vz_col = get_col("vz (km/s)")?.as_any().downcast_ref::<Float64Array>().ok_or_else(|| InputOutputError::ArrowError{action: "downcasting VZ", source: arrow::error::ArrowError::CastError("".to_string())})?.clone();
253
254            let cr_col = get_col("cr")?.as_any().downcast_ref::<Float64Array>().ok_or_else(|| InputOutputError::ArrowError{action: "downcasting Cr (unitless)", source: arrow::error::ArrowError::CastError("".to_string())})?.clone();
255            let cd_col = get_col("cd")?.as_any().downcast_ref::<Float64Array>().ok_or_else(|| InputOutputError::ArrowError{action: "downcasting Cd (unitless)", source: arrow::error::ArrowError::CastError("".to_string())})?.clone();
256
257            let dry_mass_col = get_col("dry_mass (kg)")?.as_any().downcast_ref::<Float64Array>().ok_or_else(|| InputOutputError::ArrowError{action: "downcasting dry_mass (kg)", source: arrow::error::ArrowError::CastError("".to_string())})?.clone();
258            let prop_mass_col = get_col("prop_mass (kg)")?.as_any().downcast_ref::<Float64Array>().ok_or_else(|| InputOutputError::ArrowError{action: "downcasting prop_mass (kg)", source: arrow::error::ArrowError::CastError("".to_string())})?.clone();
259
260            // Covariance columns (store them for iteration)
261            let mut cov_cols = Vec::new();
262            for i in 0..state_size {
263                for j in i..state_size {
264                     // Column names need frame and units, which were part of the export but hard to reconstruct perfectly here.
265                     // We'll guess the base name format. Robust parsing would require metadata storage.
266                     let base_name = format!("Covariance {}*{}", state_items[i], state_items[j]);
267                     // Find the actual column name (it has frame/units appended)
268                     let col_name = schema.fields().iter()
269                         .find(|f| f.name().starts_with(&base_name))
270                         .map(|f| f.name().as_str())
271                         .ok_or_else(|| InputOutputError::ParquetError {
272                              action: "seeking covariance column",
273                              source: parquet::errors::ParquetError::General("Column not found".to_string()),
274                          })?;
275                     cov_cols.push(get_col(col_name)?.as_any().downcast_ref::<Float64Array>().ok_or_else(|| InputOutputError::ArrowError{action: "downcasting covariance column", source: arrow::error::ArrowError::CastError("".to_string())})?.clone());
276                }
277            }
278
279
280            // Residual related columns
281            let rejected_col = get_col("Residual Rejected").ok().and_then(|arr| arr.as_any().downcast_ref::<BooleanArray>().cloned());
282            let tracker_col = get_col("Tracker").ok().and_then(|arr| arr.as_any().downcast_ref::<StringArray>().cloned());
283            let ratio_col = get_col("Residual ratio").ok().and_then(|arr| arr.as_any().downcast_ref::<Float64Array>().cloned());
284
285            let mut residual_data_cols: HashMap<MeasurementType, BTreeMap<String, Float64Array>> = HashMap::new();
286            for msr_type in &measurement_types_found {
287                 let mut type_cols = BTreeMap::new();
288                 let prefixes = ["Prefit residual", "Postfit residual", "Measurement noise", "Real observation", "Computed observation"];
289                 for prefix in prefixes {
290                      // Again, guessing column name format
291                      let base_name = format!("{}: {:?} ({})", prefix, msr_type, msr_type.unit());
292                      if let Ok(col) = get_col(&base_name) {
293                            if let Some(arr) = col.as_any().downcast_ref::<Float64Array>() {
294                                 type_cols.insert(prefix.to_string(), arr.clone());
295                            }
296                      }
297                 }
298                 residual_data_cols.insert(*msr_type, type_cols);
299            }
300
301            // Gain columns (store for iteration)
302            let mut gain_cols: Vec<Option<Float64Array>> = Vec::new();
303            let mut gain_available = true;
304            for i in 0..state_size {
305                for f in &measurement_types_found {
306                    // Guessing format - needs robust parsing or metadata
307                    let base_name = format!(
308                        "Gain {}*{f:?} ({}*{})",
309                        state_items[i],
310                        cov_units[i],
311                        f.unit()
312                    );
313                    let col_name = schema.fields().iter()
314                         .find(|f| f.name().starts_with(&base_name))
315                         .map(|f| f.name().as_str());
316
317                    if let Some(name) = col_name {
318                          if let Ok(col) = get_col(name) {
319                               gain_cols.push(col.as_any().downcast_ref::<Float64Array>().cloned());
320                          } else {
321                               gain_cols.push(None); // Column missing
322                               gain_available = false;
323                          }
324                    } else {
325                          // If *any* gain column is missing, assume no gains were stored (e.g., smoother run)
326                          gain_available = false;
327                          break; // No need to check further gain columns
328                    }
329                }
330                if !gain_available { break; }
331            }
332            if !gain_available { gain_cols.clear(); } // Ensure empty if incomplete
333
334
335            // FSR columns (store for iteration)
336            let mut fsr_cols: Vec<Option<Float64Array>> = Vec::new();
337            let mut fsr_available = true;
338            // for i in 0..state_size {
339            for state_item in state_items.iter().take(state_size) {
340                 // Guessing format
341                 let base_name = format!("Filter-smoother ratio {state_item}");
342                 let col_name = schema.fields().iter()
343                     .find(|f| f.name().starts_with(&base_name))
344                     .map(|f| f.name().as_str());
345
346                 if let Some(name) = col_name {
347                      if let Ok(col) = get_col(name) {
348                            fsr_cols.push(col.as_any().downcast_ref::<Float64Array>().cloned());
349                      } else {
350                            fsr_cols.push(None);
351                            fsr_available = false;
352                      }
353                 } else {
354                      fsr_available = false;
355                      break;
356                 }
357            }
358             if !fsr_available { fsr_cols.clear(); }
359
360
361            // --- Iterate through rows in the batch ---
362            for i in 0..num_rows {
363
364                let epoch = Epoch::from_gregorian_str(epoch_col.value(i)).map_err(|e| {
365                    InputOutputError::Inconsistency {
366                        msg: format!("{e} when parsing epoch"),
367                    }
368                })?;
369
370                // Reconstruct spacecraft
371                let nominal_state = Spacecraft::builder().orbit(
372                    Orbit::cartesian(
373                        x_col.value(i),
374                        y_col.value(i),
375                        z_col.value(i),
376                        vx_col.value(i),
377                        vy_col.value(i),
378                        vz_col.value(i),
379                        epoch,
380                        frame.expect("somehow frame isn't set")
381                    )).mass(
382                        Mass::from_dry_and_prop_masses(
383                            dry_mass_col.value(i),
384                            prop_mass_col.value(i))
385                    ).srp(SRPData {
386                        area_m2: srp_area_m2.expect("somehow srp area isn't set"),
387                        coeff_reflectivity: cr_col.value(i)
388                    }).drag(DragData{
389                        area_m2: drag_area_m2.expect("somehow dragarea isn't set"),
390                        coeff_drag: cd_col.value(i)
391                    }).build();
392
393                // Reconstruct Covariance
394                let mut covar = SMatrix::<f64, 9, 9>::zeros();
395                let mut cov_col_idx = 0;
396                for row in 0..state_size {
397                    for col in row..state_size {
398                        let val = cov_cols[cov_col_idx].value(i);
399                        covar[(row, col)] = val;
400                        if row != col {
401                            covar[(col, row)] = val; // Symmetric
402                        }
403                        cov_col_idx += 1;
404                    }
405                }
406
407                // Reconstruct KfEstimate
408                let estimate = KfEstimate {
409                    nominal_state,
410                    state_deviation: OVector::<f64, Const<9>>::zeros(), // Deviation not stored
411                    covar,
412                    covar_bar: covar, // Not stored, use covar
413                    stm: OMatrix::<f64, Const<9>, Const<9>>::identity(), // Not stored
414                    predicted: false, // Not stored
415                };
416                estimates.push(estimate);
417
418                // Reconstruct Residual (if applicable)
419                let is_rejected_opt = rejected_col.as_ref().and_then(|col| if col.is_valid(i) { Some(col.value(i)) } else { None });
420                let tracker_opt = tracker_col.as_ref().and_then(|col| if col.is_valid(i) { Some(col.value(i).to_string()) } else { None });
421                let ratio_opt = ratio_col.as_ref().and_then(|col| if col.is_valid(i) { Some(col.value(i)) } else { None });
422
423                let current_residual: Option<Residual<MsrSize>> = if let (Some(is_rejected), Some(tracker), Some(ratio)) = (is_rejected_opt, tracker_opt.clone(), ratio_opt) {
424                     // It's a measurement update
425                     let mut prefit_vec = OVector::<f64, MsrSize>::zeros();
426                     let mut postfit_vec = OVector::<f64, MsrSize>::zeros();
427                     let mut noise_vec = OVector::<f64, MsrSize>::zeros();
428                     let mut real_obs_vec = OVector::<f64, MsrSize>::zeros();
429                     let mut comp_obs_vec = OVector::<f64, MsrSize>::zeros();
430                     let mut current_msr_types = IndexSet::with_capacity(MsrSize::DIM);
431
432                     let mut msr_idx = 0;
433                     for (msr_type, type_cols) in &residual_data_cols {
434                           if msr_idx >= MsrSize::DIM { break; } // Should not happen if MsrSize matches data
435
436                           // Check if data exists for this type *at this row*
437                           let prefit_val = type_cols.get("Prefit residual").and_then(|col| if col.is_valid(i) { Some(col.value(i)) } else { None });
438                           let postfit_val = type_cols.get("Postfit residual").and_then(|col| if col.is_valid(i) { Some(col.value(i)) } else { None });
439                           let noise_val = type_cols.get("Measurement noise").and_then(|col| if col.is_valid(i) { Some(col.value(i)) } else { None });
440                           let real_val = type_cols.get("Real observation").and_then(|col| if col.is_valid(i) { Some(col.value(i)) } else { None });
441                           let comp_val = type_cols.get("Computed observation").and_then(|col| if col.is_valid(i) { Some(col.value(i)) } else { None });
442
443                           // Only include if *at least one* value is present for this type in this row
444                           if prefit_val.is_some() || postfit_val.is_some() || noise_val.is_some() || real_val.is_some() || comp_val.is_some() {
445                                prefit_vec[msr_idx] = prefit_val.unwrap_or(f64::NAN); // Or handle differently
446                                postfit_vec[msr_idx] = postfit_val.unwrap_or(f64::NAN);
447                                noise_vec[msr_idx] = noise_val.unwrap_or(f64::NAN);
448                                real_obs_vec[msr_idx] = real_val.unwrap_or(f64::NAN);
449                                comp_obs_vec[msr_idx] = comp_val.unwrap_or(f64::NAN);
450                                current_msr_types.insert(*msr_type);
451                                msr_idx += 1;
452                           }
453                     }
454
455                     // Resize vectors if fewer than MsrSize types were found for this row
456                     // This part is tricky and depends on how multi-type residuals were stored.
457                     // Assuming vectors should always have MsrSize, filled potentially with NaN.
458
459                     let resid = Residual {
460                          epoch,
461                          prefit: prefit_vec,
462                          postfit: postfit_vec,
463                          tracker_msr_noise: noise_vec,
464                          ratio,
465                          real_obs: real_obs_vec,
466                          computed_obs: comp_obs_vec,
467                          msr_types: current_msr_types, // Store types found for this row
468                          rejected: is_rejected,
469                          tracker: Some(tracker),
470                     };
471                     Some(resid)
472
473                } else {
474                    // Not all parts of a residual were present, assume it was a time update
475                    None
476                };
477                residuals.push(current_residual);
478
479
480                // Reconstruct Gain (if available)
481                let current_gain: Option<OMatrix<f64, <Spacecraft as State>::Size, MsrSize>> = if gain_available && !gain_cols.is_empty() {
482                     let mut gain_mat = OMatrix::<f64, <Spacecraft as State>::Size, MsrSize>::zeros();
483                     let mut all_valid = true;
484                     let mut col_idx = 0;
485                     'gain_outer: for row in 0..state_size {
486                          for col in 0..MsrSize::DIM {
487                               if let Some(gain_col) = &gain_cols[col_idx] {
488                                    if gain_col.is_valid(i) {
489                                         gain_mat[(row, col)] = gain_col.value(i);
490                                    } else {
491                                         all_valid = false;
492                                         break 'gain_outer; // Null found, entire matrix is None
493                                    }
494                               } else {
495                                    // Should not happen if gain_available is true, but safeguard
496                                    all_valid = false;
497                                    break 'gain_outer;
498                               }
499                               col_idx += 1;
500                          }
501                     }
502                     if all_valid { Some(gain_mat) } else { None }
503                } else { None };
504                gains.push(current_gain);
505
506                // Reconstruct FSR (if available)
507                let current_fsr: Option<OVector<f64, <Spacecraft as State>::Size>> = if fsr_available && !fsr_cols.is_empty() {
508                      let mut fsr_vec = OVector::<f64, <Spacecraft as State>::Size>::zeros();
509                      let mut all_valid = true;
510                      for k in 0..state_size {
511                           if let Some(fsr_col) = &fsr_cols[k] {
512                                if fsr_col.is_valid(i) {
513                                     fsr_vec[k] = fsr_col.value(i);
514                                } else {
515                                     all_valid = false;
516                                     break;
517                                }
518                           } else {
519                                all_valid = false;
520                                break;
521                           }
522                      }
523                      if all_valid { Some(fsr_vec) } else { None }
524                 } else { None };
525                 filter_smoother_ratios.push(current_fsr);
526
527            } // End row loop
528        } // End batch loop
529
530
531        // --- Final Construction ---
532        Ok(ODSolution {
533            estimates,
534            residuals,
535            gains,
536            filter_smoother_ratios,
537            devices, // Provided by user
538            measurement_types: measurement_types_found, // Determined from columns
539        })
540    }
541}
nyx_space/od/process/solution/import.rs

nyx_space/od/process/solution/
import.rs
