1#[cfg(test)]
2mod test;
3
4use crate::{
5 math::{
6 Derivative, Differentiable, Quantity, Scalar, Tensor, TensorTuple, TensorVec, TensorVector,
7 integrate::{
8 BogackiShampine, BogackiShampineFixedStep, DormandPrince, DormandPrinceFixedStep,
9 Explicit, ExplicitDynamics, ExplicitDynamicsCoupled, IntegrationError, Spectrum, Times,
10 Verner8, Verner8FixedStep, Verner9, Verner9FixedStep,
11 },
12 },
13 units::Time,
14};
15use std::fmt::Debug;
16
17type Stack<X, T> = TensorTuple<X, Derivative<X, T>>;
18type Stacks<X, T> = TensorVector<Stack<X, T>>;
19type Slopes<X, T> = TensorVector<Derivative<Stack<X, T>, T>>;
20type Coupled<X, Y, T> = TensorTuple<Stack<X, T>, Y>;
21type Coupleds<X, Y, T> = TensorVector<Coupled<X, Y, T>>;
22type CoupledSlopes<X, Y, T> = TensorVector<Derivative<Coupled<X, Y, T>, T>>;
23
24#[doc = include_str!("doc.md")]
25pub trait Stacked<T = Time>
26where
27 Self: Debug,
28{
29 #[doc = include_str!("../doc.md")]
30 fn integrate_stacked<X, UX, UV, UA>(
31 &self,
32 mut function: impl FnMut(
33 Quantity<T>,
34 &X,
35 &Derivative<X, T>,
36 ) -> Result<Derivative<Derivative<X, T>, T>, String>,
37 time: &[Quantity<T>],
38 initial_position: X,
39 initial_velocity: Derivative<X, T>,
40 ) -> Result<(Times<T>, UX, UV, UA), IntegrationError>
41 where
42 Self: Explicit<Stack<X, T>, Stacks<X, T>, Slopes<X, T>, T>,
43 X: Differentiable<T> + Tensor,
44 Derivative<X, T>: Differentiable<T> + Tensor,
45 UX: TensorVec<Item = X>,
46 UV: TensorVec<Item = Derivative<X, T>>,
47 UA: TensorVec<Item = Derivative<Derivative<X, T>, T>>,
48 {
49 let (times, stacks, slopes) = self.integrate(
50 |t, stack: &Stack<X, T>| {
51 let acceleration = function(t, &stack.0, &stack.1)?;
52 Ok(TensorTuple(stack.1.clone(), acceleration))
53 },
54 time,
55 TensorTuple(initial_position, initial_velocity),
56 )?;
57 Ok(unstack(times, &stacks, &slopes))
58 }
59 #[doc = include_str!("bounded.md")]
60 fn integrate_stacked_bounded<X, UX, UV, UA>(
61 &self,
62 mut function: impl FnMut(
63 Quantity<T>,
64 &X,
65 &Derivative<X, T>,
66 ) -> Result<Derivative<Derivative<X, T>, T>, String>,
67 mut bound: impl FnMut(Quantity<T>, &X, &Derivative<X, T>) -> Result<Quantity<T>, String>,
68 safety: Scalar,
69 time: &[Quantity<T>],
70 initial_position: X,
71 initial_velocity: Derivative<X, T>,
72 ) -> Result<(Times<T>, UX, UV, UA), IntegrationError>
73 where
74 Self: Explicit<Stack<X, T>, Stacks<X, T>, Slopes<X, T>, T>,
75 X: Differentiable<T> + Tensor,
76 Derivative<X, T>: Differentiable<T> + Tensor,
77 UX: TensorVec<Item = X>,
78 UV: TensorVec<Item = Derivative<X, T>>,
79 UA: TensorVec<Item = Derivative<Derivative<X, T>, T>>,
80 {
81 let (times, stacks, slopes) = self.integrate_bounded(
82 |t, stack: &Stack<X, T>| {
83 let acceleration = function(t, &stack.0, &stack.1)?;
84 Ok(TensorTuple(stack.1.clone(), acceleration))
85 },
86 |t, stack: &Stack<X, T>| Ok(Spectrum::Imaginary(bound(t, &stack.0, &stack.1)?)),
87 safety,
88 time,
89 TensorTuple(initial_position, initial_velocity),
90 )?;
91 Ok(unstack(times, &stacks, &slopes))
92 }
93 #[expect(clippy::type_complexity)]
94 fn integrate_stacked_coupled<X, Y, UX, UV, UA, UY, UR>(
95 &self,
96 mut function: impl FnMut(
97 Quantity<T>,
98 &X,
99 &Derivative<X, T>,
100 &Y,
101 )
102 -> Result<(Derivative<Derivative<X, T>, T>, Derivative<Y, T>), String>,
103 time: &[Quantity<T>],
104 initial_position: X,
105 initial_velocity: Derivative<X, T>,
106 initial_state: Y,
107 ) -> Result<(Times<T>, UX, UV, UA, UY, UR), IntegrationError>
108 where
109 Self: Explicit<Coupled<X, Y, T>, Coupleds<X, Y, T>, CoupledSlopes<X, Y, T>, T>,
110 X: Differentiable<T> + Tensor,
111 Derivative<X, T>: Differentiable<T> + Tensor,
112 Y: Differentiable<T> + Tensor,
113 UX: TensorVec<Item = X>,
114 UV: TensorVec<Item = Derivative<X, T>>,
115 UA: TensorVec<Item = Derivative<Derivative<X, T>, T>>,
116 UY: TensorVec<Item = Y>,
117 UR: TensorVec<Item = Derivative<Y, T>>,
118 {
119 let (times, coupleds, slopes) = self.integrate(
120 |t, coupled: &Coupled<X, Y, T>| {
121 let (acceleration, rate) = function(t, &coupled.0.0, &coupled.0.1, &coupled.1)?;
122 Ok(TensorTuple(
123 TensorTuple(coupled.0.1.clone(), acceleration),
124 rate,
125 ))
126 },
127 time,
128 TensorTuple(
129 TensorTuple(initial_position, initial_velocity),
130 initial_state,
131 ),
132 )?;
133 Ok(unstack_coupled(times, &coupleds, &slopes))
134 }
135 #[expect(clippy::type_complexity)]
136 #[expect(clippy::too_many_arguments)]
137 fn integrate_stacked_coupled_bounded<X, Y, UX, UV, UA, UY, UR>(
138 &self,
139 mut function: impl FnMut(
140 Quantity<T>,
141 &X,
142 &Derivative<X, T>,
143 &Y,
144 )
145 -> Result<(Derivative<Derivative<X, T>, T>, Derivative<Y, T>), String>,
146 mut bound: impl FnMut(Quantity<T>, &X, &Derivative<X, T>, &Y) -> Result<Spectrum<T>, String>,
147 safety: Scalar,
148 time: &[Quantity<T>],
149 initial_position: X,
150 initial_velocity: Derivative<X, T>,
151 initial_state: Y,
152 ) -> Result<(Times<T>, UX, UV, UA, UY, UR), IntegrationError>
153 where
154 Self: Explicit<Coupled<X, Y, T>, Coupleds<X, Y, T>, CoupledSlopes<X, Y, T>, T>,
155 X: Differentiable<T> + Tensor,
156 Derivative<X, T>: Differentiable<T> + Tensor,
157 Y: Differentiable<T> + Tensor,
158 UX: TensorVec<Item = X>,
159 UV: TensorVec<Item = Derivative<X, T>>,
160 UA: TensorVec<Item = Derivative<Derivative<X, T>, T>>,
161 UY: TensorVec<Item = Y>,
162 UR: TensorVec<Item = Derivative<Y, T>>,
163 {
164 let (times, coupleds, slopes) = self.integrate_bounded(
165 |t, coupled: &Coupled<X, Y, T>| {
166 let (acceleration, rate) = function(t, &coupled.0.0, &coupled.0.1, &coupled.1)?;
167 Ok(TensorTuple(
168 TensorTuple(coupled.0.1.clone(), acceleration),
169 rate,
170 ))
171 },
172 |t, coupled: &Coupled<X, Y, T>| bound(t, &coupled.0.0, &coupled.0.1, &coupled.1),
173 safety,
174 time,
175 TensorTuple(
176 TensorTuple(initial_position, initial_velocity),
177 initial_state,
178 ),
179 )?;
180 Ok(unstack_coupled(times, &coupleds, &slopes))
181 }
182}
183
184fn unstack_coupled<X, Y, UX, UV, UA, UY, UR, T>(
185 times: Times<T>,
186 coupleds: &Coupleds<X, Y, T>,
187 slopes: &CoupledSlopes<X, Y, T>,
188) -> (Times<T>, UX, UV, UA, UY, UR)
189where
190 X: Differentiable<T> + Tensor,
191 Derivative<X, T>: Differentiable<T> + Tensor,
192 Y: Differentiable<T> + Tensor,
193 UX: TensorVec<Item = X>,
194 UV: TensorVec<Item = Derivative<X, T>>,
195 UA: TensorVec<Item = Derivative<Derivative<X, T>, T>>,
196 UY: TensorVec<Item = Y>,
197 UR: TensorVec<Item = Derivative<Y, T>>,
198{
199 let mut positions = UX::new();
200 let mut velocities = UV::new();
201 let mut accelerations = UA::new();
202 let mut states = UY::new();
203 let mut rates = UR::new();
204 coupleds
205 .iter()
206 .zip(slopes.iter())
207 .for_each(|(coupled, slope)| {
208 positions.push(coupled.0.0.clone());
209 velocities.push(coupled.0.1.clone());
210 accelerations.push(slope.0.1.clone());
211 states.push(coupled.1.clone());
212 rates.push(slope.1.clone());
213 });
214 (times, positions, velocities, accelerations, states, rates)
215}
216
217fn unstack<X, UX, UV, UA, T>(
218 times: Times<T>,
219 stacks: &Stacks<X, T>,
220 slopes: &Slopes<X, T>,
221) -> (Times<T>, UX, UV, UA)
222where
223 X: Differentiable<T> + Tensor,
224 Derivative<X, T>: Differentiable<T> + Tensor,
225 UX: TensorVec<Item = X>,
226 UV: TensorVec<Item = Derivative<X, T>>,
227 UA: TensorVec<Item = Derivative<Derivative<X, T>, T>>,
228{
229 let mut positions = UX::new();
230 let mut velocities = UV::new();
231 let mut accelerations = UA::new();
232 stacks.iter().zip(slopes.iter()).for_each(|(stack, slope)| {
233 positions.push(stack.0.clone());
234 velocities.push(stack.1.clone());
235 accelerations.push(slope.1.clone());
236 });
237 (times, positions, velocities, accelerations)
238}
239
240impl<I, X, UX, UV, UA, T> ExplicitDynamics<X, UX, UV, UA, T> for I
241where
242 I: Stacked<T> + Explicit<Stack<X, T>, Stacks<X, T>, Slopes<X, T>, T>,
243 X: Differentiable<T> + Tensor,
244 Derivative<X, T>: Differentiable<T> + Tensor,
245 UX: TensorVec<Item = X>,
246 UV: TensorVec<Item = Derivative<X, T>>,
247 UA: TensorVec<Item = Derivative<Derivative<X, T>, T>>,
248{
249 fn integrate(
250 &self,
251 function: impl FnMut(
252 Quantity<T>,
253 &X,
254 &Derivative<X, T>,
255 ) -> Result<Derivative<Derivative<X, T>, T>, String>,
256 time: &[Quantity<T>],
257 initial_position: X,
258 initial_velocity: Derivative<X, T>,
259 ) -> Result<(Times<T>, UX, UV, UA), IntegrationError> {
260 self.integrate_stacked(function, time, initial_position, initial_velocity)
261 }
262 fn integrate_bounded(
263 &self,
264 function: impl FnMut(
265 Quantity<T>,
266 &X,
267 &Derivative<X, T>,
268 ) -> Result<Derivative<Derivative<X, T>, T>, String>,
269 bound: impl FnMut(Quantity<T>, &X, &Derivative<X, T>) -> Result<Quantity<T>, String>,
270 safety: Scalar,
271 time: &[Quantity<T>],
272 initial_position: X,
273 initial_velocity: Derivative<X, T>,
274 ) -> Result<(Times<T>, UX, UV, UA), IntegrationError> {
275 self.integrate_stacked_bounded(
276 function,
277 bound,
278 safety,
279 time,
280 initial_position,
281 initial_velocity,
282 )
283 }
284}
285
286impl<I, X, Y, UX, UV, UA, UY, UR, T> ExplicitDynamicsCoupled<X, Y, UX, UV, UA, UY, UR, T> for I
287where
288 I: Stacked<T> + Explicit<Coupled<X, Y, T>, Coupleds<X, Y, T>, CoupledSlopes<X, Y, T>, T>,
289 X: Differentiable<T> + Tensor,
290 Derivative<X, T>: Differentiable<T> + Tensor,
291 Y: Differentiable<T> + Tensor,
292 UX: TensorVec<Item = X>,
293 UV: TensorVec<Item = Derivative<X, T>>,
294 UA: TensorVec<Item = Derivative<Derivative<X, T>, T>>,
295 UY: TensorVec<Item = Y>,
296 UR: TensorVec<Item = Derivative<Y, T>>,
297{
298 fn integrate_coupled(
299 &self,
300 function: impl FnMut(
301 Quantity<T>,
302 &X,
303 &Derivative<X, T>,
304 &Y,
305 )
306 -> Result<(Derivative<Derivative<X, T>, T>, Derivative<Y, T>), String>,
307 time: &[Quantity<T>],
308 initial_position: X,
309 initial_velocity: Derivative<X, T>,
310 initial_state: Y,
311 ) -> Result<(Times<T>, UX, UV, UA, UY, UR), IntegrationError> {
312 self.integrate_stacked_coupled(
313 function,
314 time,
315 initial_position,
316 initial_velocity,
317 initial_state,
318 )
319 }
320 fn integrate_coupled_bounded(
321 &self,
322 function: impl FnMut(
323 Quantity<T>,
324 &X,
325 &Derivative<X, T>,
326 &Y,
327 )
328 -> Result<(Derivative<Derivative<X, T>, T>, Derivative<Y, T>), String>,
329 bound: impl FnMut(Quantity<T>, &X, &Derivative<X, T>, &Y) -> Result<Spectrum<T>, String>,
330 safety: Scalar,
331 time: &[Quantity<T>],
332 initial_position: X,
333 initial_velocity: Derivative<X, T>,
334 initial_state: Y,
335 ) -> Result<(Times<T>, UX, UV, UA, UY, UR), IntegrationError> {
336 self.integrate_stacked_coupled_bounded(
337 function,
338 bound,
339 safety,
340 time,
341 initial_position,
342 initial_velocity,
343 initial_state,
344 )
345 }
346}
347
348macro_rules! stacked {
349 ($($integrator: ty),+) => {
350 $(impl<T> Stacked<T> for $integrator {})+
351 };
352}
353
354stacked!(
355 BogackiShampine,
356 BogackiShampineFixedStep,
357 DormandPrince,
358 DormandPrinceFixedStep,
359 Verner8,
360 Verner8FixedStep,
361 Verner9,
362 Verner9FixedStep
363);