1use super::traj_it::TrajIterator;
20use super::{ExportCfg, INTERPOLATION_SAMPLES, InterpolationSnafu};
21use super::{Interpolatable, TrajError};
22use crate::errors::{NyxError, StateError};
23use crate::io::InputOutputError;
24use crate::io::watermark::pq_writer;
25use crate::linalg::DefaultAllocator;
26use crate::linalg::allocator::Allocator;
27use crate::md::prelude::{GuidanceMode, StateParameter};
28use crate::md::trajectory::smooth_state_diff_in_place;
29use crate::time::{Duration, Epoch, TimeSeries, TimeUnits};
30use anise::analysis::AnalysisError;
31use anise::analysis::specs::StateSpecTrait;
32use anise::astro::orbit::Orbit;
33use anise::errors::PhysicsError;
34use anise::prelude::{Aberration, Almanac};
35use arrow::array::{Array, Float64Builder, StringBuilder};
36use arrow::datatypes::{DataType, Field, Schema};
37use arrow::record_batch::RecordBatch;
38use hifitime::TimeScale;
39use log::{info, warn};
40use parquet::arrow::ArrowWriter;
41use snafu::ResultExt;
42use std::collections::HashMap;
43use std::error::Error;
44use std::fmt;
45use std::fs::File;
46use std::iter::Iterator;
47use std::ops;
48use std::ops::Bound::{Excluded, Included, Unbounded};
49use std::path::{Path, PathBuf};
50use std::sync::Arc;
51
52#[derive(Clone, PartialEq)]
54pub struct Traj<S: Interpolatable>
55where
56 DefaultAllocator: Allocator<S::VecLength> + Allocator<S::Size> + Allocator<S::Size, S::Size>,
57{
58 pub name: Option<String>,
60 pub states: Vec<S>,
62}
63
64impl<S: Interpolatable> Traj<S>
65where
66 DefaultAllocator: Allocator<S::VecLength> + Allocator<S::Size> + Allocator<S::Size, S::Size>,
67{
68 pub fn new() -> Self {
69 Self {
70 name: None,
71 states: Vec::new(),
72 }
73 }
74 pub fn finalize(&mut self) {
76 self.states.sort_by_key(|a| a.epoch());
78 self.states.dedup_by(|a, b| a.epoch().eq(&b.epoch()));
80 }
81
82 pub fn at(&self, epoch: Epoch) -> Result<S, TrajError> {
84 if self.states.is_empty() {
85 return Err(TrajError::EmptyTrajectory { epoch });
86 } else if self.first().epoch() > epoch || self.last().epoch() < epoch {
87 return Err(TrajError::NoInterpolationData {
88 epoch,
89 start: self.start_epoch(),
90 end: self.end_epoch(),
91 });
92 }
93 match self
94 .states
95 .binary_search_by(|state| state.epoch().cmp(&epoch))
96 {
97 Ok(idx) => {
98 Ok(self.states[idx])
100 }
101 Err(idx) => {
102 if idx == 0 || idx >= self.states.len() {
103 return Err(TrajError::NoInterpolationData {
106 epoch,
107 start: self.start_epoch(),
108 end: self.end_epoch(),
109 });
110 }
111 let num_left = INTERPOLATION_SAMPLES / 2;
116
117 let mut first_idx = idx.saturating_sub(num_left);
119 let last_idx = self.states.len().min(first_idx + INTERPOLATION_SAMPLES);
120
121 if last_idx == self.states.len() {
123 first_idx = last_idx.saturating_sub(2 * num_left);
124 }
125
126 let mut states = Vec::with_capacity(last_idx - first_idx);
127 for idx in first_idx..last_idx {
128 states.push(self.states[idx]);
129 }
130
131 self.states[idx]
132 .interpolate(epoch, &states)
133 .context(InterpolationSnafu)
134 }
135 }
136 }
137
138 pub fn first(&self) -> &S {
140 self.states.first().unwrap()
142 }
143
144 pub fn last(&self) -> &S {
146 self.states.last().unwrap()
147 }
148
149 pub fn start_epoch(&self) -> Epoch {
150 self.first().epoch()
151 }
152
153 pub fn end_epoch(&self) -> Epoch {
154 self.last().epoch()
155 }
156
157 pub fn every(&self, step: Duration) -> TrajIterator<'_, S> {
159 self.every_between(step, self.first().epoch(), self.last().epoch())
160 }
161
162 pub fn every_between(&self, step: Duration, start: Epoch, end: Epoch) -> TrajIterator<'_, S> {
164 TrajIterator {
165 time_series: TimeSeries::inclusive(
166 start.max(self.first().epoch()),
167 end.min(self.last().epoch()),
168 step,
169 ),
170 traj: self,
171 }
172 }
173
174 pub fn filter_by_epoch<R: ops::RangeBounds<Epoch>>(mut self, bound: R) -> Self {
176 self.states = self
177 .states
178 .iter()
179 .copied()
180 .filter(|s| bound.contains(&s.epoch()))
181 .collect::<Vec<_>>();
182 self
183 }
184
185 pub fn filter_by_offset<R: ops::RangeBounds<Duration>>(self, bound: R) -> Self {
188 if self.states.is_empty() {
189 return self;
190 }
191 let start = match bound.start_bound() {
193 Unbounded => self.states.first().unwrap().epoch(),
194 Included(offset) | Excluded(offset) => self.states.first().unwrap().epoch() + *offset,
195 };
196
197 let end = match bound.end_bound() {
198 Unbounded => self.states.last().unwrap().epoch(),
199 Included(offset) | Excluded(offset) => self.states.first().unwrap().epoch() + *offset,
200 };
201
202 self.filter_by_epoch(start..=end)
203 }
204 pub fn to_parquet_simple<P: AsRef<Path>>(&self, path: P) -> Result<PathBuf, Box<dyn Error>> {
206 self.to_parquet(path, ExportCfg::default())
207 }
208
209 pub fn to_parquet_with_cfg<P: AsRef<Path>>(
211 &self,
212 path: P,
213 cfg: ExportCfg,
214 ) -> Result<PathBuf, Box<dyn Error>> {
215 self.to_parquet(path, cfg)
216 }
217
218 pub fn to_parquet_with_step<P: AsRef<Path>>(
220 &self,
221 path: P,
222 step: Duration,
223 ) -> Result<(), Box<dyn Error>> {
224 self.to_parquet_with_cfg(
225 path,
226 ExportCfg {
227 step: Some(step),
228 ..Default::default()
229 },
230 )?;
231
232 Ok(())
233 }
234
235 pub fn to_parquet<P: AsRef<Path>>(
237 &self,
238 path: P,
239 cfg: ExportCfg,
240 ) -> Result<PathBuf, Box<dyn Error>> {
241 let tick = Epoch::now().unwrap();
242 info!("Exporting trajectory to parquet file...");
243
244 let path_buf = cfg.actual_path(path);
246
247 let states = if cfg.start_epoch.is_some() || cfg.end_epoch.is_some() || cfg.step.is_some() {
249 let start = cfg.start_epoch.unwrap_or_else(|| self.first().epoch());
251 let end = cfg.end_epoch.unwrap_or_else(|| self.last().epoch());
252 let step = cfg.step.unwrap_or_else(|| 1.minutes());
253 self.every_between(step, start, end).collect::<Vec<S>>()
254 } else {
255 self.states.to_vec()
256 };
257
258 let mut hdrs = vec![Field::new("Epoch (UTC)", DataType::Utf8, false)];
260
261 let frame = self.states[0].frame();
262 let more_meta = Some(vec![(
263 "Frame".to_string(),
264 serde_dhall::serialize(&frame)
265 .static_type_annotation()
266 .to_string()
267 .map_err(|e| {
268 Box::new(InputOutputError::SerializeDhall {
269 what: format!("frame `{frame}`"),
270 err: e.to_string(),
271 })
272 })?,
273 )]);
274
275 let requested_fields = match cfg.fields {
276 Some(fields) => fields,
277 None => S::export_params(),
278 };
279
280 let mut fields = Vec::new();
281 let mut field_nullable = Vec::new();
282 for field in requested_fields {
283 let mut any_ok = false;
284 let mut any_err = false;
285 for state in &states {
286 if state.value(field).is_ok() {
287 any_ok = true;
288 } else {
289 any_err = true;
290 }
291 }
292
293 if any_ok {
294 fields.push(field);
295 field_nullable.push(any_err);
296 }
297 }
298
299 for (field, nullable) in fields.iter().zip(field_nullable.iter().copied()) {
300 hdrs.push(field.to_field(more_meta.clone()).with_nullable(nullable));
301 }
302
303 let schema = Arc::new(Schema::new(hdrs));
305 let mut record: Vec<Arc<dyn Array>> = Vec::new();
306
307 let mut utc_epoch = StringBuilder::new();
311 for s in &states {
312 utc_epoch.append_value(s.epoch().to_time_scale(TimeScale::UTC).to_isoformat());
313 }
314 record.push(Arc::new(utc_epoch.finish()));
315
316 for field in fields {
318 if field == StateParameter::GuidanceMode() {
319 let mut guid_mode = StringBuilder::new();
320 for s in &states {
321 match s.value(field) {
322 Ok(value) => {
323 guid_mode.append_value(format!("{:?}", GuidanceMode::from(value)));
324 }
325 Err(_) => guid_mode.append_null(),
326 }
327 }
328 record.push(Arc::new(guid_mode.finish()));
329 } else {
330 let mut data = Float64Builder::new();
331 for s in &states {
332 match s.value(field) {
333 Ok(value) => data.append_value(value),
334 Err(_) => data.append_null(),
335 };
336 }
337 record.push(Arc::new(data.finish()));
338 }
339 }
340
341 info!(
342 "Serialized {} states from {} to {}",
343 states.len(),
344 states.first().unwrap().epoch(),
345 states.last().unwrap().epoch()
346 );
347
348 let mut metadata = HashMap::new();
350 metadata.insert("Purpose".to_string(), "Trajectory data".to_string());
351 if let Some(add_meta) = cfg.metadata {
352 for (k, v) in add_meta {
353 metadata.insert(k, v);
354 }
355 }
356
357 let props = pq_writer(Some(metadata));
358
359 let file = File::create(&path_buf)?;
360 let mut writer = ArrowWriter::try_new(file, schema.clone(), props).unwrap();
361
362 let batch = RecordBatch::try_new(schema, record)?;
363 writer.write(&batch)?;
364 writer.close()?;
365
366 let tock_time = Epoch::now().unwrap() - tick;
368 info!(
369 "Trajectory written to {} in {tock_time}",
370 path_buf.display()
371 );
372 Ok(path_buf)
373 }
374
375 pub fn resample(&self, step: Duration) -> Result<Self, NyxError> {
378 if self.states.is_empty() {
379 return Err(NyxError::Trajectory {
380 source: TrajError::CreationError {
381 msg: "No trajectory to convert".to_string(),
382 },
383 });
384 }
385
386 let mut traj = Self::new();
387 for state in self.every(step) {
388 traj.states.push(state);
389 }
390
391 traj.finalize();
392
393 Ok(traj)
394 }
395
396 pub fn rebuild(&self, epochs: &[Epoch]) -> Result<Self, NyxError> {
399 if self.states.is_empty() {
400 return Err(NyxError::Trajectory {
401 source: TrajError::CreationError {
402 msg: "No trajectory to convert".to_string(),
403 },
404 });
405 }
406
407 let mut traj = Self::new();
408 for epoch in epochs {
409 traj.states.push(self.at(*epoch)?);
410 }
411
412 traj.finalize();
413
414 Ok(traj)
415 }
416
417 pub fn ric_diff_to_parquet<P: AsRef<Path>>(
422 &self,
423 other: &Self,
424 path: P,
425 cfg: ExportCfg,
426 ) -> Result<PathBuf, TrajError> {
427 let tick = Epoch::now().unwrap();
428 info!("Exporting trajectory to parquet file...");
429
430 let path_buf = cfg.actual_path(path);
432
433 let mut hdrs = vec![Field::new("Epoch (UTC)", DataType::Utf8, false)];
435
436 for coord in ["X", "Y", "Z"] {
438 let mut meta = HashMap::new();
439 meta.insert("unit".to_string(), "km".to_string());
440
441 let field = Field::new(
442 format!("Delta {coord} (RIC) (km)"),
443 DataType::Float64,
444 false,
445 )
446 .with_metadata(meta);
447
448 hdrs.push(field);
449 }
450
451 for coord in ["x", "y", "z"] {
452 let mut meta = HashMap::new();
453 meta.insert("unit".to_string(), "km/s".to_string());
454
455 let field = Field::new(
456 format!("Delta V{coord} (RIC) (km/s)"),
457 DataType::Float64,
458 false,
459 )
460 .with_metadata(meta);
461
462 hdrs.push(field);
463 }
464
465 let frame = self.states[0].frame();
466 let more_meta = Some(vec![(
467 "Frame".to_string(),
468 serde_dhall::serialize(&frame)
469 .static_type_annotation()
470 .to_string()
471 .unwrap_or(frame.to_string()),
472 )]);
473
474 let mut cfg = cfg;
475
476 let mut fields = match cfg.fields {
477 Some(fields) => fields,
478 None => S::export_params(),
479 };
480
481 fields.retain(|param| {
483 param != &StateParameter::GuidanceMode() && self.first().value(*param).is_ok()
484 });
485
486 for field in &fields {
487 hdrs.push(field.to_field(more_meta.clone()));
488 }
489
490 let schema = Arc::new(Schema::new(hdrs));
492 let mut record: Vec<Arc<dyn Array>> = Vec::new();
493
494 cfg.start_epoch = if self.first().epoch() > other.first().epoch() {
496 Some(self.first().epoch())
497 } else {
498 Some(other.first().epoch())
499 };
500
501 cfg.end_epoch = if self.last().epoch() > other.last().epoch() {
502 Some(other.last().epoch())
503 } else {
504 Some(self.last().epoch())
505 };
506
507 let step = cfg.step.unwrap_or_else(|| 1.minutes());
509 let self_states = self
510 .every_between(step, cfg.start_epoch.unwrap(), cfg.end_epoch.unwrap())
511 .collect::<Vec<S>>();
512
513 let other_states = other
514 .every_between(step, cfg.start_epoch.unwrap(), cfg.end_epoch.unwrap())
515 .collect::<Vec<S>>();
516
517 let mut ric_diff = Vec::with_capacity(other_states.len());
519 for (other_state, self_state) in other_states.iter().zip(self_states.iter()) {
520 let self_orbit = self_state.orbit();
521 let other_orbit = other_state.orbit();
522
523 let this_ric_diff = self_orbit
524 .ric_difference(&other_orbit)
525 .map_err(|source: PhysicsError| TrajError::TrajPhysics { source })?;
526
527 ric_diff.push(this_ric_diff);
528 }
529
530 smooth_state_diff_in_place(&mut ric_diff, if other_states.len() > 5 { 5 } else { 1 });
531
532 let mut utc_epoch = StringBuilder::new();
536 for s in &self_states {
537 utc_epoch.append_value(s.epoch().to_time_scale(TimeScale::UTC).to_isoformat());
538 }
539 record.push(Arc::new(utc_epoch.finish()));
540
541 for coord_no in 0..6 {
543 let mut data = Float64Builder::new();
544 for this_ric_dff in &ric_diff {
545 data.append_value(this_ric_dff.to_cartesian_pos_vel()[coord_no]);
546 }
547 record.push(Arc::new(data.finish()));
548 }
549
550 for field in fields {
552 let mut data = Float64Builder::new();
553 for (other_state, self_state) in other_states.iter().zip(self_states.iter()) {
554 let self_val =
555 self_state
556 .value(field)
557 .map_err(|err: StateError| TrajError::TrajGeneric {
558 err: err.to_string(),
559 })?;
560 let other_val =
561 other_state
562 .value(field)
563 .map_err(|err: StateError| TrajError::TrajGeneric {
564 err: err.to_string(),
565 })?;
566 data.append_value(self_val - other_val);
567 }
568
569 record.push(Arc::new(data.finish()));
570 }
571
572 info!("Serialized {} states differences", self_states.len());
573
574 let mut metadata = HashMap::new();
576 metadata.insert(
577 "Purpose".to_string(),
578 "Trajectory difference data".to_string(),
579 );
580 if let Some(add_meta) = cfg.metadata {
581 for (k, v) in add_meta {
582 metadata.insert(k, v);
583 }
584 }
585
586 let props = pq_writer(Some(metadata));
587
588 let file = File::create(&path_buf).map_err(|err| TrajError::TrajGeneric {
589 err: format!("{err:?}"),
590 })?;
591 let mut writer = ArrowWriter::try_new(file, schema.clone(), props).unwrap();
592
593 let batch = RecordBatch::try_new(schema, record).map_err(|err| TrajError::TrajGeneric {
594 err: format!("{err:?}"),
595 })?;
596 writer.write(&batch).map_err(|err| TrajError::TrajGeneric {
597 err: format!("{err:?}"),
598 })?;
599 writer.close().map_err(|err| TrajError::TrajGeneric {
600 err: format!("{err:?}"),
601 })?;
602
603 let tock_time = Epoch::now().unwrap() - tick;
605 info!(
606 "Trajectory written to {} in {tock_time}",
607 path_buf.display()
608 );
609 Ok(path_buf)
610 }
611}
612
613impl<S: Interpolatable> ops::Add for Traj<S>
614where
615 DefaultAllocator: Allocator<S::VecLength> + Allocator<S::Size> + Allocator<S::Size, S::Size>,
616{
617 type Output = Result<Traj<S>, NyxError>;
618
619 fn add(self, other: Traj<S>) -> Self::Output {
621 &self + &other
622 }
623}
624
625impl<S: Interpolatable> ops::Add<&Traj<S>> for &Traj<S>
626where
627 DefaultAllocator: Allocator<S::VecLength> + Allocator<S::Size> + Allocator<S::Size, S::Size>,
628{
629 type Output = Result<Traj<S>, NyxError>;
630
631 fn add(self, other: &Traj<S>) -> Self::Output {
633 if self.first().frame() != other.first().frame() {
634 Err(NyxError::Trajectory {
635 source: TrajError::CreationError {
636 msg: format!(
637 "Frame mismatch in add operation: {} != {}",
638 self.first().frame(),
639 other.first().frame()
640 ),
641 },
642 })
643 } else {
644 if self.last().epoch() < other.first().epoch() {
645 let gap = other.first().epoch() - self.last().epoch();
646 warn!(
647 "Resulting merged trajectory will have a time-gap of {} starting at {}",
648 gap,
649 self.last().epoch()
650 );
651 }
652
653 let mut me = self.clone();
654 for state in &other
656 .states
657 .iter()
658 .copied()
659 .filter(|s| s.epoch() > self.last().epoch())
660 .collect::<Vec<S>>()
661 {
662 me.states.push(*state);
663 }
664 me.finalize();
665
666 Ok(me)
667 }
668 }
669}
670
671impl<S: Interpolatable> ops::AddAssign<&Traj<S>> for Traj<S>
672where
673 DefaultAllocator: Allocator<S::VecLength> + Allocator<S::Size> + Allocator<S::Size, S::Size>,
674{
675 fn add_assign(&mut self, rhs: &Self) {
681 *self = (self.clone() + rhs.clone()).unwrap();
682 }
683}
684
685impl<S: Interpolatable> fmt::Display for Traj<S>
686where
687 DefaultAllocator: Allocator<S::VecLength> + Allocator<S::Size> + Allocator<S::Size, S::Size>,
688{
689 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
690 if self.states.is_empty() {
691 write!(f, "Empty Trajectory!")
692 } else {
693 let dur = self.last().epoch() - self.first().epoch();
694 write!(
695 f,
696 "Trajectory {}in {} from {} to {} ({}, or {:.3} s) [{} states]",
697 match &self.name {
698 Some(name) => format!("of {name} "),
699 None => String::new(),
700 },
701 self.first().frame(),
702 self.first().epoch(),
703 self.last().epoch(),
704 dur,
705 dur.to_seconds(),
706 self.states.len()
707 )
708 }
709 }
710}
711
712impl<S: Interpolatable> fmt::Debug for Traj<S>
713where
714 DefaultAllocator: Allocator<S::VecLength> + Allocator<S::Size> + Allocator<S::Size, S::Size>,
715{
716 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
717 write!(f, "{self}",)
718 }
719}
720
721impl<S: Interpolatable> Default for Traj<S>
722where
723 DefaultAllocator: Allocator<S::VecLength> + Allocator<S::Size> + Allocator<S::Size, S::Size>,
724{
725 fn default() -> Self {
726 Self::new()
727 }
728}
729
730impl<S: Interpolatable> StateSpecTrait for Traj<S>
731where
732 DefaultAllocator: Allocator<S::VecLength> + Allocator<S::Size> + Allocator<S::Size, S::Size>,
733{
734 fn ab_corr(&self) -> Option<Aberration> {
735 None
736 }
737
738 fn evaluate(&self, epoch: Epoch, _almanac: &Almanac) -> Result<Orbit, AnalysisError> {
739 self.at(epoch)
740 .map(|state| state.orbit())
741 .map_err(|e| AnalysisError::GenericAnalysisError {
742 err: format!("{e}"),
743 })
744 }
745}