1use std::fmt::{Display, LowerExp};
2use std::ops::AddAssign;
3
4use crate::cff::CutCFFIndex;
5use crate::utils::{F, FloatLike, PrecisionUpgradable};
6use itertools::Itertools;
7use spenso::algebra::{algebraic_traits::RefZero, complex::Complex};
8use symbolica::domains::dual::{DualNumberStructure, HyperDual};
9
10pub(crate) fn new_constant<T: Clone + RefZero>(shape: &HyperDual<T>, value: &T) -> HyperDual<T> {
11 let mut new = shape.clone();
12 new.values[0] = value.clone();
13 let new_values_iter = new.values.iter_mut().skip(1);
14 for v in new_values_iter {
15 *v = value.ref_zero();
16 }
17 new
18}
19
20pub(crate) fn new_from_values<T: Clone>(shape: &HyperDual<T>, values: &[T]) -> HyperDual<T> {
21 let mut new = shape.clone();
22 new.values = values.to_vec();
23 new
24}
25
26pub(crate) fn simple_n_deriv_shape(num_derivatives: usize) -> Vec<Vec<usize>> {
27 mixed_derivative_shape(&[num_derivatives])
28}
29
30pub(crate) fn mixed_derivative_shape(max_orders: &[usize]) -> Vec<Vec<usize>> {
35 let mut shape = max_orders
36 .iter()
37 .map(|order| 0..=*order)
38 .multi_cartesian_product()
39 .collect_vec();
40 shape.sort_by(|left, right| {
41 let left_degree: usize = left.iter().sum();
42 let right_degree: usize = right.iter().sum();
43 left_degree.cmp(&right_degree).then_with(|| {
44 if left_degree == 1 {
45 right.cmp(left)
46 } else {
47 left.cmp(right)
48 }
49 })
50 });
51 shape
52}
53
54#[derive(Clone, Copy, Debug, PartialEq, Eq)]
55pub(crate) struct CutCFFVariableIndices {
56 pub lu_cut: Option<usize>,
57 pub left_threshold: Option<usize>,
58 pub right_threshold: Option<usize>,
59}
60
61pub(crate) fn variable_indices_from_cut_cff_index(
62 cut_cff_index: &CutCFFIndex,
63) -> CutCFFVariableIndices {
64 let mut next_index = 0;
65
66 let lu_cut = cut_cff_index
67 .lu_cut_order
68 .filter(|order| *order > 1)
69 .map(|_| {
70 let index = next_index;
71 next_index += 1;
72 index
73 });
74 let left_threshold = cut_cff_index
75 .left_threshold_order
76 .filter(|order| *order > 1)
77 .map(|_| {
78 let index = next_index;
79 next_index += 1;
80 index
81 });
82 let right_threshold = cut_cff_index
83 .right_threshold_order
84 .filter(|order| *order > 1)
85 .map(|_| {
86 let index = next_index;
87 next_index += 1;
88 index
89 });
90
91 CutCFFVariableIndices {
92 lu_cut,
93 left_threshold,
94 right_threshold,
95 }
96}
97
98pub(crate) fn shape_from_cut_cff_index(cut_cff_index: &CutCFFIndex) -> Option<Vec<Vec<usize>>> {
99 let max_derivative_shape = [
100 cut_cff_index.lu_cut_order,
101 cut_cff_index.left_threshold_order,
102 cut_cff_index.right_threshold_order,
103 ]
104 .into_iter()
105 .flatten()
106 .filter(|order| *order > 1)
107 .map(|order| order - 1)
108 .collect_vec();
109
110 (!max_derivative_shape.is_empty()).then(|| mixed_derivative_shape(&max_derivative_shape))
111}
112
113impl<T> PrecisionUpgradable for HyperDual<T>
114where
115 T: PrecisionUpgradable,
116 T::Higher: Default,
117 T::Lower: Default,
118{
119 type Higher = HyperDual<T::Higher>;
120 type Lower = HyperDual<T::Lower>;
121
122 fn higher(&self) -> Self::Higher {
123 let shape = self.get_shape();
124 let owned_shape = shape
125 .into_iter()
126 .map(|slice| slice.to_vec())
127 .collect::<Vec<_>>();
128
129 let mut new_hyperdual = HyperDual::<T::Higher>::new(owned_shape);
130
131 new_hyperdual
132 .values
133 .iter_mut()
134 .zip(self.values.iter())
135 .for_each(|(new_v, old_v)| {
136 *new_v = old_v.higher();
137 });
138
139 new_hyperdual
140 }
141
142 fn lower(&self) -> Self::Lower {
143 let shape = self.get_shape();
144 let owned_shape = shape
145 .into_iter()
146 .map(|slice| slice.to_vec())
147 .collect::<Vec<_>>();
148
149 let mut new_hyperdual = HyperDual::<T::Lower>::new(owned_shape);
150
151 new_hyperdual
152 .values
153 .iter_mut()
154 .zip(self.values.iter())
155 .for_each(|(new_v, old_v)| {
156 *new_v = old_v.lower();
157 });
158
159 new_hyperdual
160 }
161}
162
163#[derive(Clone, Debug)]
164pub enum DualOrNot<T> {
165 Dual(HyperDual<T>),
166 NonDual(T),
167}
168
169impl<T> DualOrNot<T> {
170 pub fn unwrap_real(self) -> T {
171 match self {
172 DualOrNot::Dual(_dual) => panic!("Cannot unwrap real value from Dual variant"),
173 DualOrNot::NonDual(value) => value,
174 }
175 }
176}
177
178impl<T: LowerExp> Display for DualOrNot<T> {
179 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
180 match self {
181 DualOrNot::Dual(dual) => write!(
182 f,
183 "[{}]",
184 dual.values
185 .iter()
186 .map(|v| format!("{:+16e}", v))
187 .collect::<Vec<_>>()
188 .join(", ")
189 ),
190 DualOrNot::NonDual(value) => write!(f, "{:+16e}", value),
191 }
192 }
193}
194
195impl<T> AddAssign for DualOrNot<T>
196where
197 T: AddAssign<T>,
198{
199 fn add_assign(&mut self, rhs: Self) {
200 match (self, rhs) {
201 (DualOrNot::Dual(x), DualOrNot::Dual(y)) => {
202 assert_eq!(x.values.len(), y.values.len());
203 for (lhs, rhs) in x.values.iter_mut().zip(y.values) {
204 *lhs += rhs;
205 }
206 }
207 (DualOrNot::NonDual(x), DualOrNot::NonDual(y)) => {
208 *x += y;
209 }
210 _ => panic!("Cannot add DualOrNot of different types"),
211 }
212 }
213}
214
215impl<T: Clone> DualOrNot<T> {
216 pub fn new_from_slice(shape: &Option<HyperDual<T>>, values: &[T]) -> Self {
217 match shape {
218 Some(dual_shape) => {
219 let hyperdual = new_from_values(dual_shape, values);
220 DualOrNot::Dual(hyperdual)
221 }
222 None => {
223 assert_eq!(values.len(), 1);
224 DualOrNot::NonDual(values[0].clone())
225 }
226 }
227 }
228}
229
230pub(crate) fn extract_t_derivatives_complex<T: FloatLike>(
232 dual: HyperDual<Complex<F<T>>>,
233) -> Vec<Complex<F<T>>> {
234 let mut result = Vec::with_capacity(dual.values.len());
235 let mut n_factorial = Complex::new_re(dual.values[0].re.one());
236
237 for (order, value) in dual.values.iter().enumerate() {
238 if order > 0 {
239 n_factorial = &n_factorial * Complex::new_re(value.re.from_usize(order));
240 }
241 result.push(value.clone() * &n_factorial);
242 }
243
244 result
245}
246
247pub(crate) fn extract_t_derivatives<T: FloatLike>(dual: HyperDual<F<T>>) -> Vec<F<T>> {
249 let mut result = Vec::with_capacity(dual.values.len());
250 let mut n_factorial = dual.values[0].one();
251
252 for (order, value) in dual.values.iter().enumerate() {
253 if order > 0 {
254 n_factorial = &n_factorial * value.from_usize(order);
255 }
256 result.push(value.clone() * &n_factorial);
257 }
258
259 result
260}
261
262pub(crate) fn dualize_dual_t_to_dual_r_t<T: FloatLike>(
263 t_dual: HyperDual<F<T>>,
264 target_shape: HyperDual<F<T>>,
265 variable: usize,
266) -> HyperDual<F<T>> {
267 let mut new_dual = new_constant(&target_shape, &t_dual.values[0]);
268 let n_variables_of_target_shape = target_shape.get_shape()[0].len();
269
270 debug_assert!(variable < n_variables_of_target_shape);
271
272 for (i, value) in t_dual.values.iter().enumerate().skip(1) {
273 let dual_shape_to_find = {
274 let mut shape = vec![0; n_variables_of_target_shape];
275 shape[variable] = i;
276 shape
277 };
278
279 #[allow(clippy::expect_fun_call)]
280 let index_of_derivative = target_shape
281 .get_shape()
282 .iter()
283 .position(|shape| shape == &dual_shape_to_find)
284 .expect(&format!(
285 "Could not find derivative shape: {:?} in {:?}",
286 dual_shape_to_find,
287 target_shape.get_shape()
288 ));
289
290 new_dual.values[index_of_derivative] = value.clone();
291 }
292
293 new_dual
294}
295
296pub(crate) fn extract_coefficient_t_duals<T: Clone + RefZero + Default>(
297 dual: &HyperDual<T>,
298 t_variable: usize,
299) -> (Vec<Vec<usize>>, Vec<HyperDual<T>>) {
300 let mixed_shape = dual
301 .get_shape()
302 .into_iter()
303 .map(|orders| orders.to_vec())
304 .collect_vec();
305 let n_variables = mixed_shape.first().map_or(0, Vec::len);
306
307 debug_assert!(t_variable < n_variables);
308
309 let max_t_order = mixed_shape
310 .iter()
311 .map(|orders| orders[t_variable])
312 .max()
313 .unwrap_or(0);
314 let t_dual_shape = HyperDual::new(simple_n_deriv_shape(max_t_order));
315
316 let mut coefficient_orders = Vec::<Vec<usize>>::new();
317 let mut coefficient_duals = Vec::<HyperDual<T>>::new();
318
319 for (orders, value) in mixed_shape.iter().zip(dual.values.iter()) {
320 let t_order = orders[t_variable];
321 let projected_orders = orders
322 .iter()
323 .enumerate()
324 .filter_map(|(index, order)| (index != t_variable).then_some(*order))
325 .collect_vec();
326
327 let coefficient_index = if let Some(existing_index) = coefficient_orders
328 .iter()
329 .position(|existing_orders| existing_orders == &projected_orders)
330 {
331 existing_index
332 } else {
333 coefficient_orders.push(projected_orders);
334 coefficient_duals.push(new_constant(&t_dual_shape, &value.ref_zero()));
335 coefficient_duals.len() - 1
336 };
337
338 coefficient_duals[coefficient_index].values[t_order] = value.clone();
339 }
340
341 (coefficient_orders, coefficient_duals)
342}
343
344#[cfg(test)]
345mod tests {
346 use super::*;
347
348 #[test]
349 fn mixed_derivative_shape_supports_constant_axes() {
350 assert_eq!(mixed_derivative_shape(&[]), vec![Vec::<usize>::new()]);
351 assert_eq!(mixed_derivative_shape(&[0, 0]), vec![vec![0, 0]]);
352 assert_eq!(simple_n_deriv_shape(0), vec![vec![0]]);
353 assert_eq!(
354 mixed_derivative_shape(&[0, 2, 0, 1]),
355 vec![
356 vec![0, 0, 0, 0],
357 vec![0, 1, 0, 0],
358 vec![0, 0, 0, 1],
359 vec![0, 1, 0, 1],
360 vec![0, 2, 0, 0],
361 vec![0, 2, 0, 1],
362 ]
363 );
364
365 for orders in [vec![], vec![0, 0], vec![0, 2, 0, 1]] {
366 let shape = mixed_derivative_shape(&orders);
367 let mut values = vec![F(0.0_f64); shape.len()];
368 values[0] = F(3.0_f64);
369 let constant = HyperDual::from_values(shape, values);
370 let square = constant.clone() * constant;
371 assert_eq!(square.values[0], F(9.0_f64));
372 assert!(square.values[1..].iter().all(|value| *value == F(0.0_f64)));
373 }
374 }
375
376 #[test]
377 fn mixed_derivative_shape_preserves_four_axis_polynomial_coefficients() {
378 let shape = mixed_derivative_shape(&[2, 1, 1, 2]);
379 let variables = [2.0_f64, 3.0_f64, 5.0_f64, 7.0_f64]
380 .into_iter()
381 .enumerate()
382 .map(|(axis, base)| {
383 let mut values = vec![F(0.0_f64); shape.len()];
384 values[0] = F(base);
385 let mut linear_order = vec![0; 4];
386 linear_order[axis] = 1;
387 let linear_index = shape
388 .iter()
389 .position(|order| *order == linear_order)
390 .unwrap();
391 values[linear_index] = F(1.0_f64);
392 HyperDual::from_values(shape.clone(), values)
393 })
394 .collect_vec();
395 let polynomial = variables[0].clone()
396 * &variables[0]
397 * &variables[1]
398 * &variables[2]
399 * &variables[3]
400 * &variables[3];
401
402 for (orders, expected) in [
405 ([0, 0, 0, 0], 2940.0_f64),
406 ([1, 0, 1, 1], 168.0_f64),
407 ([2, 1, 1, 2], 1.0_f64),
408 ] {
409 let index = shape
410 .iter()
411 .position(|order| order.as_slice() == orders)
412 .unwrap();
413 assert_eq!(polynomial.values[index], F(expected));
414 }
415 }
416
417 #[test]
418 fn dualize_dual_t_to_dual_r_t_uses_requested_target_variable() {
419 let t_dual = HyperDual::new(simple_n_deriv_shape(1));
420 let t_dual = new_from_values(&t_dual, &[F(3.0_f64), F(5.0_f64)]);
421
422 let target_shape = HyperDual::new(
423 shape_from_cut_cff_index(&CutCFFIndex {
424 left_threshold_order: Some(2),
425 right_threshold_order: None,
426 lu_cut_order: Some(2),
427 })
428 .unwrap(),
429 );
430
431 let dualized = dualize_dual_t_to_dual_r_t(t_dual, target_shape, 1);
432
433 assert_eq!(
434 dualized.values,
435 vec![F(3.0_f64), F(0.0_f64), F(5.0_f64), F(0.0_f64)]
436 );
437 }
438
439 #[test]
440 fn dualize_dual_t_to_dual_r_t_keeps_constant_inputs_constant() {
441 let t_dual = HyperDual::new(simple_n_deriv_shape(0));
442 let t_dual = new_from_values(&t_dual, &[F(7.0_f64)]);
443
444 let target_shape = HyperDual::new(
445 shape_from_cut_cff_index(&CutCFFIndex {
446 left_threshold_order: Some(2),
447 right_threshold_order: None,
448 lu_cut_order: Some(2),
449 })
450 .unwrap(),
451 );
452
453 let dualized = dualize_dual_t_to_dual_r_t(t_dual, target_shape, 1);
454
455 assert_eq!(
456 dualized.values,
457 vec![F(7.0_f64), F(0.0_f64), F(0.0_f64), F(0.0_f64)]
458 );
459 }
460
461 #[test]
462 fn extract_coefficient_t_duals_factorizes_single_threshold_mixed_dual() {
463 let mixed_shape = HyperDual::new(
464 shape_from_cut_cff_index(&CutCFFIndex {
465 left_threshold_order: Some(2),
466 right_threshold_order: None,
467 lu_cut_order: Some(2),
468 })
469 .unwrap(),
470 );
471 let mixed_dual = new_from_values(
472 &mixed_shape,
473 &[F(2.0_f64), F(3.0_f64), F(5.0_f64), F(7.0_f64)],
474 );
475
476 let (coefficient_orders, coefficient_duals) = extract_coefficient_t_duals(&mixed_dual, 0);
477
478 assert_eq!(coefficient_orders, vec![vec![0], vec![1]]);
479 assert_eq!(coefficient_duals[0].values, vec![F(2.0_f64), F(3.0_f64)]);
480 assert_eq!(coefficient_duals[1].values, vec![F(5.0_f64), F(7.0_f64)]);
481 }
482
483 #[test]
484 fn extract_coefficient_t_duals_factorizes_iterated_mixed_dual() {
485 let mixed_shape = HyperDual::new(
486 shape_from_cut_cff_index(&CutCFFIndex {
487 left_threshold_order: Some(2),
488 right_threshold_order: Some(2),
489 lu_cut_order: Some(2),
490 })
491 .unwrap(),
492 );
493 let mixed_dual = new_from_values(
494 &mixed_shape,
495 &[
496 F(1.0_f64),
497 F(2.0_f64),
498 F(3.0_f64),
499 F(4.0_f64),
500 F(5.0_f64),
501 F(6.0_f64),
502 F(7.0_f64),
503 F(8.0_f64),
504 ],
505 );
506
507 let (coefficient_orders, coefficient_duals) = extract_coefficient_t_duals(&mixed_dual, 0);
508
509 assert_eq!(
510 coefficient_orders,
511 vec![vec![0, 0], vec![1, 0], vec![0, 1], vec![1, 1]]
512 );
513 assert_eq!(coefficient_duals[0].values, vec![F(1.0_f64), F(2.0_f64)]);
514 assert_eq!(coefficient_duals[1].values, vec![F(3.0_f64), F(7.0_f64)]);
515 assert_eq!(coefficient_duals[2].values, vec![F(4.0_f64), F(6.0_f64)]);
516 assert_eq!(coefficient_duals[3].values, vec![F(5.0_f64), F(8.0_f64)]);
517 }
518
519 #[test]
520 fn shape_from_cut_cff_index_preserves_canonical_t_left_right_order() {
521 let three_axis_shape = shape_from_cut_cff_index(&CutCFFIndex {
522 left_threshold_order: Some(3),
523 right_threshold_order: Some(2),
524 lu_cut_order: Some(2),
525 })
526 .unwrap();
527
528 assert_eq!(
529 three_axis_shape,
530 vec![
531 vec![0, 0, 0],
532 vec![1, 0, 0],
533 vec![0, 1, 0],
534 vec![0, 0, 1],
535 vec![0, 1, 1],
536 vec![0, 2, 0],
537 vec![1, 0, 1],
538 vec![1, 1, 0],
539 vec![0, 2, 1],
540 vec![1, 1, 1],
541 vec![1, 2, 0],
542 vec![1, 2, 1],
543 ]
544 );
545
546 assert_eq!(
547 shape_from_cut_cff_index(&CutCFFIndex {
548 left_threshold_order: Some(3),
549 right_threshold_order: None,
550 lu_cut_order: Some(3),
551 }),
552 Some(vec![
553 vec![0, 0],
554 vec![1, 0],
555 vec![0, 1],
556 vec![0, 2],
557 vec![1, 1],
558 vec![2, 0],
559 vec![1, 2],
560 vec![2, 1],
561 vec![2, 2],
562 ])
563 );
564
565 assert_eq!(
566 shape_from_cut_cff_index(&CutCFFIndex {
567 left_threshold_order: Some(1),
568 right_threshold_order: Some(2),
569 lu_cut_order: Some(2),
570 }),
571 Some(vec![vec![0, 0], vec![1, 0], vec![0, 1], vec![1, 1]])
572 );
573
574 assert_eq!(
575 shape_from_cut_cff_index(&CutCFFIndex {
576 left_threshold_order: Some(2),
577 right_threshold_order: Some(1),
578 lu_cut_order: Some(1),
579 }),
580 Some(vec![vec![0], vec![1]])
581 );
582 }
583
584 #[test]
585 fn variable_indices_from_cut_cff_index_tracks_active_axes_without_reordering() {
586 assert_eq!(
587 variable_indices_from_cut_cff_index(&CutCFFIndex {
588 left_threshold_order: Some(2),
589 right_threshold_order: Some(2),
590 lu_cut_order: Some(2),
591 }),
592 CutCFFVariableIndices {
593 lu_cut: Some(0),
594 left_threshold: Some(1),
595 right_threshold: Some(2),
596 }
597 );
598
599 assert_eq!(
600 variable_indices_from_cut_cff_index(&CutCFFIndex {
601 left_threshold_order: Some(2),
602 right_threshold_order: Some(2),
603 lu_cut_order: Some(1),
604 }),
605 CutCFFVariableIndices {
606 lu_cut: None,
607 left_threshold: Some(0),
608 right_threshold: Some(1),
609 }
610 );
611
612 assert_eq!(
613 variable_indices_from_cut_cff_index(&CutCFFIndex {
614 left_threshold_order: Some(1),
615 right_threshold_order: Some(2),
616 lu_cut_order: Some(2),
617 }),
618 CutCFFVariableIndices {
619 lu_cut: Some(0),
620 left_threshold: None,
621 right_threshold: Some(1),
622 }
623 );
624 }
625}