aeon_tk/mesh/
function.rs

1#![allow(clippy::needless_range_loop)]
2
3use crate::IRef;
4use crate::mesh::Mesh;
5use crate::{
6    image::{ImageMut, ImageRef},
7    kernel::{NodeSpace, node_from_vertex},
8};
9use std::ops::Range;
10
11/// An interface for computing values, gradients, and hessians of fields.
12pub trait Engine<const N: usize> {
13    fn space(&self) -> &NodeSpace<N>;
14
15    fn num_nodes(&self) -> usize {
16        self.space().num_nodes()
17    }
18
19    fn vertex_size(&self) -> [usize; N] {
20        self.space().vertex_size()
21    }
22
23    fn node_range(&self) -> Range<usize>;
24
25    fn alloc<T: Default>(&self, len: usize) -> &mut [T];
26
27    fn position(&self, vertex: [usize; N]) -> [f64; N] {
28        self.space().position(node_from_vertex(vertex))
29    }
30
31    fn index_from_vertex(&self, vertex: [usize; N]) -> usize {
32        self.space().index_from_vertex(vertex)
33    }
34    fn min_spacing(&self) -> f64 {
35        let spacing = self.space().spacing();
36        spacing
37            .iter()
38            .min_by(|a, b| a.total_cmp(&b))
39            .cloned()
40            .unwrap_or(1.0)
41    }
42
43    fn value(&self, field: &[f64], vertex: [usize; N]) -> f64;
44    fn derivative(&self, field: &[f64], axis: usize, vertex: [usize; N]) -> f64;
45    fn second_derivative(&self, field: &[f64], axis: usize, vertex: [usize; N]) -> f64;
46    fn mixed_derivative(&self, field: &[f64], i: usize, j: usize, vertex: [usize; N]) -> f64;
47    fn dissipation(&self, field: &[f64], axis: usize, vertex: [usize; N]) -> f64;
48}
49
50impl<'a, const N: usize, E: Engine<N>> Engine<N> for IRef<'a, E> {
51    fn space(&self) -> &NodeSpace<N> {
52        self.0.space()
53    }
54
55    fn node_range(&self) -> Range<usize> {
56        self.0.node_range()
57    }
58
59    fn alloc<T: Default>(&self, len: usize) -> &mut [T] {
60        self.0.alloc(len)
61    }
62
63    fn value(&self, field: &[f64], vertex: [usize; N]) -> f64 {
64        self.0.value(field, vertex)
65    }
66
67    fn derivative(&self, field: &[f64], axis: usize, vertex: [usize; N]) -> f64 {
68        self.0.derivative(field, axis, vertex)
69    }
70
71    fn second_derivative(&self, field: &[f64], axis: usize, vertex: [usize; N]) -> f64 {
72        self.0.second_derivative(field, axis, vertex)
73    }
74
75    fn mixed_derivative(&self, field: &[f64], i: usize, j: usize, vertex: [usize; N]) -> f64 {
76        self.0.mixed_derivative(field, i, j, vertex)
77    }
78
79    fn dissipation(&self, field: &[f64], axis: usize, vertex: [usize; N]) -> f64 {
80        self.0.dissipation(field, axis, vertex)
81    }
82}
83
84/// A function maps one set of scalar fields to another.
85pub trait Function<const N: usize> {
86    type Error;
87
88    /// Action of the function on an individual finite different block.
89    fn evaluate(
90        &self,
91        engine: impl Engine<N>,
92        input: ImageRef,
93        output: ImageMut,
94    ) -> Result<(), Self::Error>;
95
96    /// An (optional) preprocessing step run immediately before a function is applied, after
97    /// boundary conditions have been filled.
98    fn preprocess(&mut self, _mesh: &mut Mesh<N>, _input: ImageMut) -> Result<(), Self::Error> {
99        Ok(())
100    }
101}
102
103/// A projection takes in a position and returns a system of values.
104pub trait Projection<const N: usize> {
105    fn project(&self, position: [f64; N]) -> f64;
106}
107
108/// A Gaussian projection centered on a given point, with an amplitude and a sigma.
109#[derive(Clone, Copy)]
110pub struct Gaussian<const N: usize> {
111    pub amplitude: f64,
112    pub sigma: [f64; N],
113    pub center: [f64; N],
114}
115
116impl<const N: usize> Projection<N> for Gaussian<N> {
117    fn project(&self, position: [f64; N]) -> f64 {
118        let offset: [_; N] =
119            core::array::from_fn(|axis| (position[axis] - self.center[axis]) / self.sigma[axis]);
120        let r2: f64 = offset.map(|v| v * v).iter().sum();
121        self.amplitude * (-r2).exp()
122    }
123}
124
125#[derive(Clone, Copy)]
126pub struct TanH<const N: usize> {
127    pub amplitude: f64,
128    pub sigma: f64,
129    pub center: [f64; N],
130}
131
132impl<const N: usize> Projection<N> for TanH<N> {
133    fn project(&self, position: [f64; N]) -> f64 {
134        let offset =
135            core::array::from_fn::<_, N, _>(|axis| (position[axis] - self.center[axis]).powi(2))
136                .iter()
137                .sum::<f64>()
138                .sqrt()
139                / self.sigma;
140        self.amplitude * offset.tanh()
141    }
142}
143
144pub struct FunctionBorrowMut<'a, I>(pub &'a mut I);
145
146impl<'a, const N: usize, F: Function<N>> Function<N> for FunctionBorrowMut<'a, F> {
147    type Error = F::Error;
148
149    fn preprocess(&mut self, mesh: &mut Mesh<N>, input: ImageMut<'_>) -> Result<(), Self::Error> {
150        self.0.preprocess(mesh, input)
151    }
152
153    fn evaluate(
154        &self,
155        engine: impl Engine<N>,
156        input: ImageRef<'_>,
157        output: ImageMut<'_>,
158    ) -> Result<(), Self::Error> {
159        self.0.evaluate(engine, input, output)
160    }
161}
aeon_tk/mesh/function.rs

aeon_tk/mesh/
function.rs
