gam 0.3.107

Generalized penalized likelihood engine
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317

//! Closed-form and spectral zonal Wahba kernels on S².
//!
//! This module owns the scalar/SIMD kernel dispatch for intrinsic sphere
//! smooths. Callers in `basis` handle data validation, coordinate transforms,
//! and matrix assembly.

use super::polylog::{dilog_unit, trilog_unit};
use super::sphere_spectral::{
    pseudo_s2_truncated_coefficients, sobolev_s2_truncated_coefficients,
    sphere_truncated_spectral_derivative_eval, sphere_truncated_spectral_eval,
};
use super::{BasisError, SphereWahbaKernel};

/// Pseudo-spline Wahba kernel on S² (mgcv `makeR`-style closed form).
#[inline]
pub(crate) fn wahba_sphere_kernel_pseudo_from_cos(cos_gamma: f64, m: usize) -> f64 {
    let cg = cos_gamma.clamp(-1.0, 1.0);
    let z = (1.0 - cg).max(f64::EPSILON * 1.0e-4);
    let w = 0.5 * z;
    let c0 = w.sqrt();
    let a = (1.0 + 1.0 / c0).ln();
    let c = 2.0 * c0;
    let two_pi = 2.0 * std::f64::consts::PI;
    match m {
        1 => {
            let q1 = 2.0 * a * w - c + 1.0;
            (q1 - 0.5) / two_pi
        }
        2 => {
            let w2 = w * w;
            let q2 = a * (6.0 * w2 - 2.0 * w) - 3.0 * c * w + 3.0 * w + 0.5;
            (q2 / 2.0 - 1.0 / 6.0) / two_pi
        }
        3 => {
            let w2 = w * w;
            let w3 = w2 * w;
            let q3 = (a * (60.0 * w3 - 36.0 * w2) + 30.0 * w2 + c * (8.0 * w - 30.0 * w2)
                - 3.0 * w
                + 1.0)
                / 3.0;
            (q3 / 6.0 - 1.0 / 24.0) / two_pi
        }
        _ => {
            let w2 = w * w;
            let w3 = w2 * w;
            let w4 = w3 * w;
            let q4 = a * (70.0 * w4 - 60.0 * w3 + 6.0 * w2)
                + 35.0 * w3 * (1.0 - c)
                + c * 55.0 * w2 / 3.0
                - 12.5 * w2
                - w / 3.0
                + 0.25;
            (q4 / 24.0 - 1.0 / 120.0) / two_pi
        }
    }
}

/// Exact derivative `dK_m^{pseudo}/d(cos gamma)` of the pseudo-spline Wahba
/// kernel [`wahba_sphere_kernel_pseudo_from_cos`].
///
/// The forward kernel is a polynomial in `w = (1 - cos gamma)/2` with the
/// auxiliary terms `c0 = sqrt(w)`, `c = 2 c0`, and
/// `a = ln(1 + 1/c0)`. Differentiating in `w` and applying
/// `dw/d(cos gamma) = -1/2` gives the analytic `dK/d(cos gamma)` below.
#[inline]
pub(crate) fn wahba_sphere_kernel_pseudo_derivative_dcos(cos_gamma: f64, m: usize) -> f64 {
    let cg = cos_gamma.clamp(-1.0, 1.0);
    let z = (1.0 - cg).max(f64::EPSILON * 1.0e-4);
    let w = 0.5 * z;
    let c0 = w.sqrt();
    let a = (1.0 + 1.0 / c0).ln();
    let c = 2.0 * c0;
    let two_pi = 2.0 * std::f64::consts::PI;
    let da_dw = -1.0 / (2.0 * c0 * c0 * (c0 + 1.0));
    let dc_dw = 1.0 / c0;
    let dk_dw = match m {
        1 => {
            let dq1_dw = 2.0 * a + 2.0 * w * da_dw - dc_dw;
            dq1_dw / two_pi
        }
        2 => {
            let dq2_dw = da_dw * (6.0 * w * w - 2.0 * w) + a * (12.0 * w - 2.0)
                - 3.0 * (dc_dw * w + c)
                + 3.0;
            (dq2_dw / 2.0) / two_pi
        }
        3 => {
            let w2 = w * w;
            let w3 = w2 * w;
            let dinner_dw = da_dw * (60.0 * w3 - 36.0 * w2)
                + a * (180.0 * w2 - 72.0 * w)
                + 60.0 * w
                + (dc_dw * (8.0 * w - 30.0 * w2) + c * (8.0 - 60.0 * w))
                - 3.0;
            let dq3_dw = dinner_dw / 3.0;
            (dq3_dw / 6.0) / two_pi
        }
        _ => {
            let w2 = w * w;
            let w3 = w2 * w;
            let w4 = w3 * w;
            let dq4_dw = da_dw * (70.0 * w4 - 60.0 * w3 + 6.0 * w2)
                + a * (280.0 * w3 - 180.0 * w2 + 12.0 * w)
                + 35.0 * (3.0 * w2 * (1.0 - c) - w3 * dc_dw)
                + (55.0 / 3.0) * (dc_dw * w2 + c * 2.0 * w)
                - 25.0 * w
                - 1.0 / 3.0;
            (dq4_dw / 24.0) / two_pi
        }
    };
    dk_dw * (-0.5)
}

// ============================================================================
// Wahba/Sobolev kernel on S²
// ============================================================================
//
// `K_m^{Sobolev}(gamma) = (1/4pi) * sum_{l >= 1} (2l + 1)
// * [l(l + 1)]^{-m} * P_l(cos gamma)`.
//
// For `m in {1, 2, 3}` we use the closed-form expressions derived in
// Beatson & zu Castell, "Thinplate Splines on the Sphere", SIGMA 14 (2018)
// 083 (Section 6.2). For `m = 4`, we fall back to a truncated Legendre series.

/// Sobolev `K_m^{Sobolev}` reproducing kernel on S², closed-form for
/// `m in {1, 2, 3}` plus spectral series for `m = 4`.
#[inline]
pub(crate) fn wahba_sphere_kernel_sobolev_closed_form(cos_gamma: f64, m: usize) -> f64 {
    let cos_g = cos_gamma.clamp(-1.0, 1.0);
    let four_pi = 4.0 * std::f64::consts::PI;
    let pi2_6 = std::f64::consts::PI * std::f64::consts::PI / 6.0;
    let u = ((1.0 - cos_g) * 0.5).max(f64::EPSILON * 1.0e-4);
    let one_minus_u = (1.0 - u).max(f64::EPSILON * 1.0e-4);
    match m {
        1 => (-u.ln() - 1.0) / four_pi,
        2 => (dilog_unit(one_minus_u) + 1.0 - pi2_6) / four_pi,
        3 => {
            const ZETA3: f64 = 1.2020569031595942853997381615114499907649862923404988817922;
            let li3_u = trilog_unit(u);
            let li2_one_minus_u = dilog_unit(one_minus_u);
            let li2_u = dilog_unit(u);
            (-2.0 * li3_u - li2_one_minus_u + u.ln() * li2_u + 2.0 * ZETA3 + pi2_6 - 2.0) / four_pi
        }
        _ => wahba_sphere_kernel_sobolev_spectral(cos_g, m),
    }
}

/// Spectral Legendre-series evaluation of the Sobolev kernel
/// `K_m^{Sobolev}(gamma) = (1/4pi) sum_{l >= 1} (2l+1) *
/// [l(l+1)]^{-m} * P_l(cos gamma)`.
#[inline]
pub(crate) fn wahba_sphere_kernel_sobolev_spectral(cos_gamma: f64, m: usize) -> f64 {
    let l_max = match m {
        1 => 4096_usize,
        2 => 256,
        3 => 128,
        _ => 96,
    };
    let x = cos_gamma.clamp(-1.0, 1.0);
    let m_i = m as i32;
    let four_pi = 4.0 * std::f64::consts::PI;
    let mut p_l_minus_1 = 1.0_f64;
    let mut p_l = x;
    let mut sum = 3.0 * p_l / (four_pi * 2.0_f64.powi(m_i));
    for l in 1..l_max {
        let p_l_plus_1 =
            ((2 * l + 1) as f64 * x * p_l - (l as f64) * p_l_minus_1) / ((l + 1) as f64);
        let ell = (l + 1) as f64;
        let eigen = (ell * (ell + 1.0)).powi(m_i);
        let weight = (2.0 * ell + 1.0) / four_pi;
        sum += weight * p_l_plus_1 / eigen;
        p_l_minus_1 = p_l;
        p_l = p_l_plus_1;
    }
    sum
}

/// Evaluate the Wahba sphere reproducing kernel at a single `cos gamma`.
#[inline]
pub(crate) fn wahba_sphere_kernel_from_cos_kind(
    cos_gamma: f64,
    penalty_order: usize,
    kernel: SphereWahbaKernel,
) -> Result<f64, BasisError> {
    if !(1..=4).contains(&penalty_order) {
        crate::bail_invalid_basis!(
            "spherical spline penalty_order must be one of 1, 2, 3, 4; got {penalty_order}"
        );
    }
    let value = match kernel {
        SphereWahbaKernel::Sobolev => {
            wahba_sphere_kernel_sobolev_closed_form(cos_gamma, penalty_order)
        }
        SphereWahbaKernel::Pseudo => wahba_sphere_kernel_pseudo_from_cos(cos_gamma, penalty_order),
        SphereWahbaKernel::SobolevTruncated { lmax } => {
            let coeffs = sobolev_s2_truncated_coefficients(lmax as usize, penalty_order);
            sphere_truncated_spectral_eval(cos_gamma, &coeffs)
        }
        SphereWahbaKernel::PseudoTruncated { lmax } => {
            let coeffs = pseudo_s2_truncated_coefficients(lmax as usize, penalty_order);
            sphere_truncated_spectral_eval(cos_gamma, &coeffs)
        }
    };
    if !value.is_finite() {
        crate::bail_invalid_basis!("spherical spline kernel produced a non-finite value");
    }
    Ok(value)
}

/// SIMD lane-wise evaluation. Both Sobolev and pseudo-spline branches are
/// scalar-per-lane because the closed forms contain non-vector elementary and
/// polylogarithm calls.
#[inline]
pub(crate) fn wahba_sphere_kernel_from_cos_simd_kind(
    cos_gamma: wide::f64x4,
    penalty_order: usize,
    kernel: SphereWahbaKernel,
) -> wide::f64x4 {
    use wide::f64x4;
    if !(1..=4).contains(&penalty_order) {
        return f64x4::from(f64::NAN);
    }
    let cg = cos_gamma.fast_max(f64x4::from(-1.0)).fast_min(f64x4::ONE);
    let lanes = cg.to_array();
    let f = |x: f64| -> f64 {
        match kernel {
            SphereWahbaKernel::Sobolev => wahba_sphere_kernel_sobolev_closed_form(x, penalty_order),
            SphereWahbaKernel::Pseudo => wahba_sphere_kernel_pseudo_from_cos(x, penalty_order),
            SphereWahbaKernel::SobolevTruncated { lmax } => {
                let coeffs = sobolev_s2_truncated_coefficients(lmax as usize, penalty_order);
                sphere_truncated_spectral_eval(x, &coeffs)
            }
            SphereWahbaKernel::PseudoTruncated { lmax } => {
                let coeffs = pseudo_s2_truncated_coefficients(lmax as usize, penalty_order);
                sphere_truncated_spectral_eval(x, &coeffs)
            }
        }
    };
    f64x4::from(lanes.map(f))
}

/// Spectral derivative of the Sobolev sphere kernel w.r.t. `cos gamma`.
pub(crate) fn wahba_sphere_kernel_sobolev_derivative_dcos(x: f64, m: usize) -> f64 {
    const POLE_LIMIT_THRESHOLD: f64 = 1.0e-10;

    let l_max = match m {
        1 => 4096_usize,
        2 => 256,
        3 => 128,
        _ => 96,
    };
    let x = x.clamp(-1.0, 1.0);
    let m_i = m as i32;
    let four_pi = 4.0 * std::f64::consts::PI;
    if x.abs() > 1.0 - POLE_LIMIT_THRESHOLD {
        let pole = if x.is_sign_negative() {
            -1.0_f64
        } else {
            1.0_f64
        };
        let mut sum = 0.0_f64;
        for l in 1..=l_max {
            let ell = l as f64;
            let sign = if pole < 0.0 && l % 2 == 0 { -1.0 } else { 1.0 };
            let p_l_prime = 0.5 * ell * (ell + 1.0) * sign;
            let eigen = (ell * (ell + 1.0)).powi(m_i);
            let weight = (2.0 * ell + 1.0) / four_pi;
            sum += weight * p_l_prime / eigen;
        }
        return sum;
    }
    let one_minus_x2 = (1.0 - x * x).max(f64::EPSILON);
    let mut p_lm1 = 1.0_f64;
    let mut p_l = x;
    let mut l = 1_usize;
    let mut sum = 0.0_f64;
    loop {
        let p_l_prime = (l as f64) * (p_lm1 - x * p_l) / one_minus_x2;
        let ell = l as f64;
        let eigen = (ell * (ell + 1.0)).powi(m_i);
        let weight = (2.0 * ell + 1.0) / four_pi;
        sum += weight * p_l_prime / eigen;
        if l >= l_max {
            break;
        }
        let p_lp1 = ((2 * l + 1) as f64 * x * p_l - (l as f64) * p_lm1) / ((l + 1) as f64);
        p_lm1 = p_l;
        p_l = p_lp1;
        l += 1;
    }
    sum
}

/// Unified `dK/d(cos gamma)` for any [`SphereWahbaKernel`] kind.
#[inline]
pub(crate) fn wahba_sphere_kernel_derivative_dcos_kind(
    cos_gamma: f64,
    penalty_order: usize,
    kernel: SphereWahbaKernel,
) -> f64 {
    match kernel {
        SphereWahbaKernel::Sobolev => {
            wahba_sphere_kernel_sobolev_derivative_dcos(cos_gamma, penalty_order)
        }
        SphereWahbaKernel::Pseudo => {
            wahba_sphere_kernel_pseudo_derivative_dcos(cos_gamma, penalty_order)
        }
        SphereWahbaKernel::SobolevTruncated { lmax } => {
            let coeffs = sobolev_s2_truncated_coefficients(lmax as usize, penalty_order);
            sphere_truncated_spectral_derivative_eval(cos_gamma, &coeffs)
        }
        SphereWahbaKernel::PseudoTruncated { lmax } => {
            let coeffs = pseudo_s2_truncated_coefficients(lmax as usize, penalty_order);
            sphere_truncated_spectral_derivative_eval(cos_gamma, &coeffs)
        }
    }
}