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

nyx_space/od/process/solution/
import.rs
