1use crate::error::{InterpolateError, InterpolateResult};
26
27#[derive(Debug, Clone)]
36pub struct PdeOperator {
37 pub terms: Vec<(f64, Vec<usize>)>,
39 pub dim: usize,
41}
42
43impl PdeOperator {
44 pub fn laplacian(dim: usize) -> Self {
49 let terms = (0..dim)
50 .map(|i| {
51 let mut order = vec![0usize; dim];
52 order[i] = 2;
53 (1.0_f64, order)
54 })
55 .collect();
56 Self { terms, dim }
57 }
58
59 pub fn advection_1d(speed: f64) -> Self {
61 Self {
62 terms: vec![(speed, vec![1])],
63 dim: 1,
64 }
65 }
66
67 pub fn custom(terms: Vec<(f64, Vec<usize>)>, dim: usize) -> InterpolateResult<Self> {
76 for (_, ref order) in &terms {
77 if order.len() != dim {
78 return Err(InterpolateError::InvalidInput {
79 message: format!(
80 "Multi-index length {} does not match dim {}",
81 order.len(),
82 dim
83 ),
84 });
85 }
86 }
87 Ok(Self { terms, dim })
88 }
89
90 pub fn apply_fd(&self, center: &[f64], f_at: impl Fn(&[f64]) -> f64, h: f64) -> f64 {
115 assert_eq!(
116 center.len(),
117 self.dim,
118 "center length must equal operator dim"
119 );
120 let f = &f_at;
121 self.terms
122 .iter()
123 .map(|(coeff, order)| coeff * apply_term_fd(center, order, f, h))
124 .sum()
125 }
126
127 pub fn try_apply_fd(
130 &self,
131 center: &[f64],
132 f_at: impl Fn(&[f64]) -> f64,
133 h: f64,
134 ) -> InterpolateResult<f64> {
135 if center.len() != self.dim {
136 return Err(InterpolateError::DimensionMismatch(format!(
137 "center has {} components, operator has dim {}",
138 center.len(),
139 self.dim
140 )));
141 }
142 let f = &f_at;
143 let mut total = 0.0_f64;
144 for (coeff, order) in &self.terms {
145 let val = try_apply_term_fd(center, order, f, h)?;
146 total += coeff * val;
147 }
148 Ok(total)
149 }
150}
151
152fn apply_term_fd(center: &[f64], order: &[usize], f: &impl Fn(&[f64]) -> f64, h: f64) -> f64 {
159 let total_order: usize = order.iter().sum();
160
161 match total_order {
162 0 => f(center),
163 1 => {
164 let dim = order.iter().position(|&o| o == 1).unwrap_or(0);
166 central_diff_1(center, dim, f, h)
167 }
168 2 => {
169 let active: Vec<usize> = order
170 .iter()
171 .enumerate()
172 .filter(|(_, &o)| o > 0)
173 .map(|(i, _)| i)
174 .collect();
175 if active.len() == 1 {
176 central_diff_2(center, active[0], f, h)
178 } else {
179 central_diff_mixed(center, active[0], active[1], f, h)
181 }
182 }
183 n => apply_higher_order(center, order, f, h, n),
185 }
186}
187
188fn try_apply_term_fd(
190 center: &[f64],
191 order: &[usize],
192 f: &impl Fn(&[f64]) -> f64,
193 h: f64,
194) -> InterpolateResult<f64> {
195 let total: usize = order.iter().sum();
196 if total >= 4 {
197 return Err(InterpolateError::NotImplemented(format!(
198 "Finite-difference stencil not implemented for total derivative order {total}"
199 )));
200 }
201 Ok(apply_term_fd(center, order, f, h))
202}
203
204fn central_diff_1(center: &[f64], dim: usize, f: &impl Fn(&[f64]) -> f64, h: f64) -> f64 {
206 let mut cp = center.to_vec();
207 let mut cm = center.to_vec();
208 cp[dim] += h;
209 cm[dim] -= h;
210 (f(&cp) - f(&cm)) / (2.0 * h)
211}
212
213fn central_diff_2(center: &[f64], dim: usize, f: &impl Fn(&[f64]) -> f64, h: f64) -> f64 {
215 let mut cp = center.to_vec();
216 let mut cm = center.to_vec();
217 cp[dim] += h;
218 cm[dim] -= h;
219 (f(&cp) - 2.0 * f(center) + f(&cm)) / (h * h)
220}
221
222fn central_diff_mixed(
225 center: &[f64],
226 d0: usize,
227 d1: usize,
228 f: &impl Fn(&[f64]) -> f64,
229 h: f64,
230) -> f64 {
231 let mut pp = center.to_vec();
232 let mut pm = center.to_vec();
233 let mut mp = center.to_vec();
234 let mut mm = center.to_vec();
235 pp[d0] += h;
236 pp[d1] += h;
237 pm[d0] += h;
238 pm[d1] -= h;
239 mp[d0] -= h;
240 mp[d1] += h;
241 mm[d0] -= h;
242 mm[d1] -= h;
243 (f(&pp) - f(&pm) - f(&mp) + f(&mm)) / (4.0 * h * h)
244}
245
246fn apply_higher_order(
250 center: &[f64],
251 order: &[usize],
252 f: &impl Fn(&[f64]) -> f64,
253 h: f64,
254 _total: usize,
255) -> f64 {
256 let mut ops: Vec<(usize, u8)> = Vec::new();
258 for (d, &o) in order.iter().enumerate() {
259 let n2 = o / 2;
260 let n1 = o % 2;
261 for _ in 0..n2 {
262 ops.push((d, 2));
263 }
264 for _ in 0..n1 {
265 ops.push((d, 1));
266 }
267 }
268 apply_ops_recursively(center, &ops, f, h)
270}
271
272fn apply_ops_recursively(
273 center: &[f64],
274 ops: &[(usize, u8)],
275 f: &impl Fn(&[f64]) -> f64,
276 h: f64,
277) -> f64 {
278 if ops.is_empty() {
279 return f(center);
280 }
281 let (d, degree) = ops[0];
282 let rest = &ops[1..];
283
284 let inner_fn = move |x: &[f64]| apply_ops_recursively(x, rest, f, h);
285
286 if degree == 1 {
287 central_diff_1(center, d, &inner_fn, h)
288 } else {
289 central_diff_2(center, d, &inner_fn, h)
290 }
291}
292
293#[non_exhaustive]
299#[derive(Debug, Clone, PartialEq)]
300pub enum RbfKernel {
301 ThinPlateSpline,
303 Multiquadric,
305 InverseMultiquadric,
307 Gaussian,
309}
310
311impl RbfKernel {
312 pub fn eval(&self, r: f64, eps: f64) -> f64 {
314 match self {
315 RbfKernel::ThinPlateSpline => {
316 if r < 1e-300 {
317 0.0
318 } else {
319 r * r * r.ln()
320 }
321 }
322 RbfKernel::Multiquadric => (1.0 + (eps * r) * (eps * r)).sqrt(),
323 RbfKernel::InverseMultiquadric => 1.0 / (1.0 + (eps * r) * (eps * r)).sqrt(),
324 RbfKernel::Gaussian => (-(eps * r) * (eps * r)).exp(),
325 }
326 }
327}
328
329#[derive(Debug, Clone)]
331pub struct PhysicsInformedRbfConfig {
332 pub pde_weight: f64,
334 pub n_collocation: usize,
337 pub kernel: RbfKernel,
339 pub epsilon: f64,
341 pub ridge: f64,
343}
344
345impl Default for PhysicsInformedRbfConfig {
346 fn default() -> Self {
347 Self {
348 pde_weight: 1.0,
349 n_collocation: 20,
350 kernel: RbfKernel::Multiquadric,
351 epsilon: 1.0,
352 ridge: 1e-10,
353 }
354 }
355}
356
357#[derive(Debug, Clone)]
377pub struct PhysicsInformedRbf {
378 config: PhysicsInformedRbfConfig,
379 centers: Vec<Vec<f64>>,
381 coeffs: Vec<f64>,
383 collocation_pts: Vec<Vec<f64>>,
385 operator: PdeOperator,
387}
388
389impl PhysicsInformedRbf {
390 fn dist(a: &[f64], b: &[f64]) -> f64 {
393 a.iter()
394 .zip(b.iter())
395 .map(|(&ai, &bi)| (ai - bi) * (ai - bi))
396 .sum::<f64>()
397 .sqrt()
398 }
399
400 fn rbf_matrix(centers: &[Vec<f64>], kernel: &RbfKernel, eps: f64) -> Vec<Vec<f64>> {
403 let n = centers.len();
404 let mut phi = vec![vec![0.0f64; n]; n];
405 for i in 0..n {
406 for j in 0..n {
407 let r = Self::dist(¢ers[i], ¢ers[j]);
408 phi[i][j] = kernel.eval(r, eps);
409 }
410 }
411 phi
412 }
413
414 fn pde_operator_row(
418 c: &[f64],
419 centers: &[Vec<f64>],
420 kernel: &RbfKernel,
421 eps: f64,
422 op: &PdeOperator,
423 h: f64,
424 ) -> Vec<f64> {
425 centers
426 .iter()
427 .map(|xi| {
428 let phi_j = |x: &[f64]| {
429 let r = Self::dist(x, xi);
430 kernel.eval(r, eps)
431 };
432 op.apply_fd(c, phi_j, h)
433 })
434 .collect()
435 }
436
437 fn mat_vec(a: &[Vec<f64>], x: &[f64]) -> Vec<f64> {
440 a.iter()
441 .map(|row| row.iter().zip(x.iter()).map(|(&a, &b)| a * b).sum())
442 .collect()
443 }
444
445 fn gram(a: &[Vec<f64>]) -> Vec<Vec<f64>> {
447 let n = if a.is_empty() { 0 } else { a[0].len() };
448 let mut g = vec![vec![0.0f64; n]; n];
449 for row in a {
450 for i in 0..n {
451 for j in 0..n {
452 g[i][j] += row[i] * row[j];
453 }
454 }
455 }
456 g
457 }
458
459 fn at_vec(a: &[Vec<f64>], v: &[f64]) -> Vec<f64> {
461 let n = if a.is_empty() { 0 } else { a[0].len() };
462 let mut out = vec![0.0f64; n];
463 for (row, &vi) in a.iter().zip(v.iter()) {
464 for j in 0..n {
465 out[j] += row[j] * vi;
466 }
467 }
468 out
469 }
470
471 fn cholesky_solve(a: &[Vec<f64>], b: &[f64]) -> InterpolateResult<Vec<f64>> {
474 use crate::random_features::cholesky_solve as rf_chol;
475 rf_chol(a, b)
476 }
477
478 fn make_collocation(data: &[Vec<f64>], n_collocation: usize, seed: u64) -> Vec<Vec<f64>> {
482 if data.is_empty() || n_collocation == 0 {
483 return Vec::new();
484 }
485 let dim = data[0].len();
486 let mut mins = vec![f64::INFINITY; dim];
488 let mut maxs = vec![f64::NEG_INFINITY; dim];
489 for pt in data {
490 for (d, &v) in pt.iter().enumerate() {
491 if v < mins[d] {
492 mins[d] = v;
493 }
494 if v > maxs[d] {
495 maxs[d] = v;
496 }
497 }
498 }
499 for d in 0..dim {
501 let range = (maxs[d] - mins[d]).max(1e-12);
502 mins[d] += 0.05 * range;
503 maxs[d] -= 0.05 * range;
504 }
505 let mut state = seed.wrapping_add(1);
507 let next = |s: &mut u64| -> f64 {
508 *s = s
509 .wrapping_mul(6_364_136_223_846_793_005)
510 .wrapping_add(1_442_695_040_888_963_407);
511 (*s >> 11) as f64 / (1u64 << 53) as f64
512 };
513 (0..n_collocation)
514 .map(|_| {
515 (0..dim)
516 .map(|d| mins[d] + next(&mut state) * (maxs[d] - mins[d]))
517 .collect()
518 })
519 .collect()
520 }
521
522 pub fn fit(
533 points: &[Vec<f64>],
534 values: &[f64],
535 operator: PdeOperator,
536 rhs_fn: impl Fn(&[f64]) -> f64,
537 config: PhysicsInformedRbfConfig,
538 ) -> InterpolateResult<Self> {
539 if points.is_empty() {
540 return Err(InterpolateError::InsufficientData(
541 "No data points provided".to_string(),
542 ));
543 }
544 if points.len() != values.len() {
545 return Err(InterpolateError::DimensionMismatch(format!(
546 "points ({}) and values ({}) have different lengths",
547 points.len(),
548 values.len()
549 )));
550 }
551 let n = points.len();
552 let kernel = &config.kernel;
553 let eps = config.epsilon;
554 let lambda = config.pde_weight;
555 let h_fd = 1e-4;
557
558 let phi = Self::rbf_matrix(points, kernel, eps);
560
561 let colloc = Self::make_collocation(points, config.n_collocation, 42);
563 let m = colloc.len();
564
565 let mut l_mat: Vec<Vec<f64>> = Vec::with_capacity(m);
566 for c in &colloc {
567 let row = Self::pde_operator_row(c, points, kernel, eps, &operator, h_fd);
568 l_mat.push(row);
569 }
570
571 let g: Vec<f64> = colloc.iter().map(|c| rhs_fn(c)).collect();
573
574 let phi_t_phi = Self::gram(&phi);
576 let l_t_l = Self::gram(&l_mat);
577
578 let mut lhs = vec![vec![0.0f64; n]; n];
579 for i in 0..n {
580 for j in 0..n {
581 lhs[i][j] = phi_t_phi[i][j] + lambda * l_t_l[i][j];
582 if i == j {
583 lhs[i][j] += config.ridge;
584 }
585 }
586 }
587
588 let phi_t_y = Self::at_vec(&phi, values);
589 let l_t_g = Self::at_vec(&l_mat, &g);
590
591 let mut rhs_vec = vec![0.0f64; n];
592 for i in 0..n {
593 rhs_vec[i] = phi_t_y[i] + lambda * l_t_g[i];
594 }
595
596 let coeffs = Self::cholesky_solve(&lhs, &rhs_vec)?;
597
598 Ok(Self {
599 config,
600 centers: points.to_vec(),
601 coeffs,
602 collocation_pts: colloc,
603 operator,
604 })
605 }
606
607 pub fn eval(&self, x: &[f64]) -> f64 {
609 self.centers
610 .iter()
611 .zip(self.coeffs.iter())
612 .map(|(xi, &ai)| {
613 let r = Self::dist(x, xi);
614 ai * self.config.kernel.eval(r, self.config.epsilon)
615 })
616 .sum()
617 }
618
619 pub fn eval_batch(&self, points: &[Vec<f64>]) -> Vec<f64> {
621 points.iter().map(|x| self.eval(x)).collect()
622 }
623
624 pub fn pde_residual(&self, x: &[f64], rhs_fn: impl Fn(&[f64]) -> f64) -> f64 {
628 let h = 1e-4;
629 let f_fn = |pt: &[f64]| self.eval(pt);
630 let lf = self.operator.apply_fd(x, f_fn, h);
631 lf - rhs_fn(x)
632 }
633
634 pub fn collocation_pts(&self) -> &[Vec<f64>] {
636 &self.collocation_pts
637 }
638
639 pub fn n_centers(&self) -> usize {
641 self.centers.len()
642 }
643}
644
645#[cfg(test)]
650mod tests {
651 use super::*;
652 use std::f64::consts::PI;
653
654 #[test]
657 fn test_laplacian_1d() {
658 let lap = PdeOperator::laplacian(1);
660 let f = |x: &[f64]| x[0] * x[0];
661 let val = lap.apply_fd(&[1.5], f, 1e-4);
662 assert!(
663 (val - 2.0).abs() < 1e-5,
664 "1D Laplacian of x² should be 2, got {val}"
665 );
666 }
667
668 #[test]
669 fn test_laplacian_2d() {
670 let lap = PdeOperator::laplacian(2);
672 let f = |x: &[f64]| x[0] * x[0] + x[1] * x[1];
673 let val = lap.apply_fd(&[1.0, 2.0], f, 1e-4);
674 assert!(
675 (val - 4.0).abs() < 1e-5,
676 "2D Laplacian of x²+y² should be 4, got {val}"
677 );
678 }
679
680 #[test]
681 fn test_advection_1d() {
682 let speed = 2.0;
684 let op = PdeOperator::advection_1d(speed);
685 let x0 = PI / 4.0;
686 let f = |x: &[f64]| x[0].sin();
687 let val = op.apply_fd(&[x0], f, 1e-5);
688 let expected = speed * x0.cos();
689 assert!(
690 (val - expected).abs() < 1e-4,
691 "Advection stencil: got {val}, expected {expected}"
692 );
693 }
694
695 #[test]
696 fn test_custom_operator() {
697 let op = PdeOperator::custom(vec![(3.0, vec![1, 0])], 2).expect("custom op");
699 let f = |x: &[f64]| x[0] + x[1];
700 let val = op.apply_fd(&[0.5, 0.5], f, 1e-5);
701 assert!((val - 3.0).abs() < 1e-4, "Custom op value {val}");
702 }
703
704 #[test]
705 fn test_custom_operator_wrong_dim() {
706 let result = PdeOperator::custom(vec![(1.0, vec![1, 0])], 3);
707 assert!(
708 result.is_err(),
709 "Should fail when multi-index dim != operator dim"
710 );
711 }
712
713 #[test]
714 fn test_try_apply_fd_dimension_check() {
715 let lap = PdeOperator::laplacian(2);
716 let f = |x: &[f64]| x[0];
717 let result = lap.try_apply_fd(&[1.0], f, 1e-4);
718 assert!(result.is_err(), "Should error on wrong center length");
719 }
720
721 #[test]
724 fn test_pifr_fit_and_eval() {
725 let pts: Vec<Vec<f64>> = (0..5)
727 .flat_map(|i| (0..5).map(move |j| vec![i as f64 * 0.25, j as f64 * 0.25]))
728 .collect();
729 let vals: Vec<f64> = pts.iter().map(|p| p[0] * p[0] - p[1] * p[1]).collect();
730
731 let op = PdeOperator::laplacian(2);
732 let rhs_fn = |_x: &[f64]| 0.0_f64; let config = PhysicsInformedRbfConfig {
734 pde_weight: 1.0,
735 n_collocation: 10,
736 kernel: RbfKernel::Multiquadric,
737 epsilon: 2.0,
738 ridge: 1e-8,
739 };
740
741 let interp =
742 PhysicsInformedRbf::fit(&pts, &vals, op, rhs_fn, config).expect("fit should succeed");
743
744 for (pt, &v) in pts.iter().zip(vals.iter()) {
746 let pred = interp.eval(pt);
747 assert!(
748 (pred - v).abs() < 0.1,
749 "Training error too large at {:?}: pred={pred:.4}, true={v:.4}",
750 pt
751 );
752 }
753 }
754
755 #[test]
756 fn test_pifr_laplacian_pde_residual_small() {
757 let pts: Vec<Vec<f64>> = (0..5)
760 .flat_map(|i| (0..5).map(move |j| vec![i as f64 * 0.25, j as f64 * 0.25]))
761 .collect();
762 let vals: Vec<f64> = pts.iter().map(|p| p[0] * p[0] - p[1] * p[1]).collect();
763
764 let op = PdeOperator::laplacian(2);
765 let rhs_fn = |_: &[f64]| 0.0f64;
766 let config = PhysicsInformedRbfConfig {
767 pde_weight: 50.0,
768 n_collocation: 20,
769 kernel: RbfKernel::Gaussian,
770 epsilon: 2.0,
771 ridge: 1e-9,
772 };
773
774 let interp = PhysicsInformedRbf::fit(&pts, &vals, op.clone(), rhs_fn, config).expect("fit");
775
776 let res = interp.pde_residual(&[0.3, 0.3], |_| 0.0);
780 assert!(
781 res.abs() < 10.0,
782 "PDE residual too large (> 10.0): {res:.4} — penalty should reduce it"
783 );
784 }
785
786 #[test]
787 fn test_pifr_eval_batch() {
788 let pts: Vec<Vec<f64>> = (0..5).map(|i| vec![i as f64 * 0.5]).collect();
789 let vals: Vec<f64> = pts.iter().map(|p| p[0] * p[0]).collect();
790
791 let op = PdeOperator::advection_1d(0.0);
792 let config = PhysicsInformedRbfConfig {
793 pde_weight: 0.01,
794 n_collocation: 5,
795 kernel: RbfKernel::Multiquadric,
796 epsilon: 1.0,
797 ridge: 1e-8,
798 };
799
800 let interp = PhysicsInformedRbf::fit(&pts, &vals, op, |_| 0.0, config).expect("fit 1D");
801
802 let batch_pts: Vec<Vec<f64>> = vec![vec![0.25], vec![0.75], vec![1.25]];
803 let results = interp.eval_batch(&batch_pts);
804 assert_eq!(results.len(), 3, "batch eval length");
805 for v in &results {
806 assert!(v.is_finite(), "batch eval should be finite");
807 }
808 }
809
810 #[test]
811 fn test_pifr_n_centers() {
812 let pts: Vec<Vec<f64>> = (0..6).map(|i| vec![i as f64 * 0.2]).collect();
813 let vals: Vec<f64> = pts.iter().map(|p| p[0]).collect();
814 let op = PdeOperator::laplacian(1);
815 let config = PhysicsInformedRbfConfig::default();
816 let interp = PhysicsInformedRbf::fit(&pts, &vals, op, |_| 0.0, config).expect("fit");
817 assert_eq!(interp.n_centers(), 6);
818 }
819
820 #[test]
821 fn test_pifr_error_empty_points() {
822 let op = PdeOperator::laplacian(1);
823 let result = PhysicsInformedRbf::fit(&[], &[], op, |_| 0.0, Default::default());
824 assert!(result.is_err());
825 }
826
827 #[test]
828 fn test_pifr_error_length_mismatch() {
829 let pts: Vec<Vec<f64>> = (0..5).map(|i| vec![i as f64]).collect();
830 let vals: Vec<f64> = vec![0.0; 3];
831 let op = PdeOperator::laplacian(1);
832 let result = PhysicsInformedRbf::fit(&pts, &vals, op, |_| 0.0, Default::default());
833 assert!(result.is_err());
834 }
835
836 #[test]
837 fn test_rbf_kernel_variants() {
838 let kernels = [
839 RbfKernel::ThinPlateSpline,
840 RbfKernel::Multiquadric,
841 RbfKernel::InverseMultiquadric,
842 RbfKernel::Gaussian,
843 ];
844 for kernel in &kernels {
845 let v = kernel.eval(1.0, 1.0);
846 assert!(
847 v.is_finite() && v >= 0.0,
848 "kernel {:?} returned {v}",
849 kernel
850 );
851 if *kernel == RbfKernel::ThinPlateSpline {
853 assert_eq!(kernel.eval(0.0, 1.0), 0.0);
854 }
855 }
856 }
857}