Skip to main content

gamlss_core/
param.rs

1use std::{marker::PhantomData, ops::Range};
2
3use crate::{DesignMatrix, LinearPredictorBlock, ModelError, PredictorBlock};
4
5/// Helper for assigning sequential offsets to typed parameter block tuples.
6///
7/// This is the safe construction path for ordinary models: create each
8/// [`ParameterBlock`] with any placeholder offset, then call
9/// `ParameterBlocks::new((...))` to lay the tuple out from zero. Low-level
10/// constructors that accept explicit offsets remain available for advanced
11/// layouts and integration code.
12#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
13pub struct ParameterBlocks;
14
15impl ParameterBlocks {
16    /// Assigns sequential offsets starting at zero.
17    ///
18    /// # Panics
19    ///
20    /// Panics if the sequential layout does not fit in `usize`. Use
21    /// [`Self::try_new`] when block sizes may come from unchecked external
22    /// input.
23    #[allow(clippy::new_ret_no_self)]
24    #[must_use]
25    #[inline]
26    pub fn new<Blocks>(blocks: Blocks) -> Blocks
27    where
28        Blocks: AssignParameterOffsets,
29    {
30        Self::with_start(0, blocks)
31    }
32
33    /// Assigns sequential offsets starting at `start`.
34    ///
35    /// # Panics
36    ///
37    /// Panics if the sequential layout does not fit in `usize`. Use
38    /// [`Self::try_with_start`] when block sizes may come from unchecked
39    /// external input.
40    #[must_use]
41    #[inline]
42    pub fn with_start<Blocks>(start: usize, blocks: Blocks) -> Blocks
43    where
44        Blocks: AssignParameterOffsets,
45    {
46        blocks.assign_offsets(start)
47    }
48
49    /// Assigns sequential offsets starting at zero.
50    ///
51    /// # Errors
52    ///
53    /// Returns [`ModelError::BlockRangeOverflow`] if any assigned block range
54    /// would not fit in `usize`.
55    #[allow(clippy::new_ret_no_self)]
56    #[inline]
57    pub fn try_new<Blocks>(blocks: Blocks) -> Result<Blocks, ModelError>
58    where
59        Blocks: TryAssignParameterOffsets,
60    {
61        Self::try_with_start(0, blocks)
62    }
63
64    /// Assigns sequential offsets starting at `start`.
65    ///
66    /// # Errors
67    ///
68    /// Returns [`ModelError::BlockRangeOverflow`] if any assigned block range
69    /// would not fit in `usize`.
70    #[inline]
71    pub fn try_with_start<Blocks>(start: usize, blocks: Blocks) -> Result<Blocks, ModelError>
72    where
73        Blocks: TryAssignParameterOffsets,
74    {
75        blocks.try_assign_offsets(start)
76    }
77}
78
79/// Marker for the location parameter `mu`.
80#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
81pub struct Mu;
82
83impl ParameterName for Mu {
84    const NAME: &'static str = "mu";
85}
86
87/// Marker for a mathematical mean parameter.
88#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
89pub struct Mean;
90
91impl ParameterName for Mean {
92    const NAME: &'static str = "mean";
93}
94
95/// Marker for a median parameter.
96#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
97pub struct Median;
98
99impl ParameterName for Median {
100    const NAME: &'static str = "median";
101}
102
103/// Marker for a component mean in mixture or zero-adjusted models.
104#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
105pub struct ComponentMean;
106
107impl ParameterName for ComponentMean {
108    const NAME: &'static str = "component_mean";
109}
110
111/// Marker for an unconditional total mean in mixture or zero-adjusted models.
112#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
113pub struct TotalMean;
114
115impl ParameterName for TotalMean {
116    const NAME: &'static str = "total_mean";
117}
118
119/// Marker for the scale parameter `sigma`.
120#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
121pub struct Sigma;
122
123impl ParameterName for Sigma {
124    const NAME: &'static str = "sigma";
125}
126
127/// Marker for a coefficient of variation parameter.
128#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
129pub struct Cv;
130
131impl ParameterName for Cv {
132    const NAME: &'static str = "cv";
133}
134
135/// Marker for a log-standard-deviation parameter.
136#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
137pub struct LogSd;
138
139impl ParameterName for LogSd {
140    const NAME: &'static str = "log_sd";
141}
142
143/// Marker for the location parameter of a log-scale distribution.
144#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
145pub struct LogLocation;
146
147impl ParameterName for LogLocation {
148    const NAME: &'static str = "log_location";
149}
150
151/// Marker for the third GAMLSS parameter `nu`.
152#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
153pub struct Nu;
154
155impl ParameterName for Nu {
156    const NAME: &'static str = "nu";
157}
158
159/// Marker for the fourth GAMLSS parameter `tau`.
160#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
161pub struct Tau;
162
163impl ParameterName for Tau {
164    const NAME: &'static str = "tau";
165}
166
167/// Marker for the rate parameter of a distribution.
168#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
169pub struct Rate;
170
171impl ParameterName for Rate {
172    const NAME: &'static str = "rate";
173}
174
175/// Marker for a dispersion parameter.
176#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
177pub struct Dispersion;
178
179impl ParameterName for Dispersion {
180    const NAME: &'static str = "dispersion";
181}
182
183/// Marker for the shape parameter of a distribution.
184#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
185pub struct Shape;
186
187impl ParameterName for Shape {
188    const NAME: &'static str = "shape";
189}
190
191/// Marker for a negative-binomial size parameter.
192#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
193pub struct Size;
194
195impl ParameterName for Size {
196    const NAME: &'static str = "size";
197}
198
199/// Marker for the scale parameter of a distribution.
200#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
201pub struct Scale;
202
203impl ParameterName for Scale {
204    const NAME: &'static str = "scale";
205}
206
207/// Marker for a probability parameter.
208#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
209pub struct Probability;
210
211impl ParameterName for Probability {
212    const NAME: &'static str = "probability";
213}
214
215/// Marker for a zero-mass or zero-inflation probability parameter.
216#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
217pub struct ZeroProbability;
218
219impl ParameterName for ZeroProbability {
220    const NAME: &'static str = "zero_probability";
221}
222
223/// Marker for a one-mass probability parameter.
224#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
225pub struct OneProbability;
226
227impl ParameterName for OneProbability {
228    const NAME: &'static str = "one_probability";
229}
230
231/// Marker for a power parameter.
232#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
233pub struct Power;
234
235impl ParameterName for Power {
236    const NAME: &'static str = "power";
237}
238
239/// Marker for the precision parameter of a distribution.
240///
241/// Used for mean/precision parameterizations, e.g. the beta distribution, where
242/// `precision > 0` controls the concentration around the mean.
243#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
244pub struct Precision;
245
246impl ParameterName for Precision {
247    const NAME: &'static str = "precision";
248}
249
250/// Typed coefficient block for a single distribution parameter.
251///
252/// `P` specifies the parameter role, `L` specifies the link function, `X` holds
253/// the predictor block, and `Penalty` adds regularization. The block stores its
254/// coefficient range within the common beta vector; use [`Self::range`] and
255/// [`Self::len`] to inspect that layout.
256#[derive(Debug, Clone, PartialEq, Eq)]
257pub struct ParameterBlock<P, L, X, Penalty> {
258    x: X,
259    penalty: Penalty,
260    offset: usize,
261    len: usize,
262    marker: PhantomData<(P, L)>,
263}
264
265impl<P, L, X, Penalty> ParameterBlock<P, L, X, Penalty>
266where
267    X: PredictorBlock,
268{
269    /// Creates a block, taking `len` from `x.nparams()`.
270    #[must_use]
271    #[inline]
272    pub fn new(x: X, penalty: Penalty, offset: usize) -> Self {
273        let len = x.nparams();
274        Self::from_len(x, penalty, offset, len)
275    }
276
277    /// Creates a block from a generic predictor.
278    ///
279    /// This is a synonym for [`Self::new`], kept for code where the explicit
280    /// `predictor` word makes the call more readable.
281    #[must_use]
282    #[inline]
283    pub fn from_predictor(x: X, penalty: Penalty, offset: usize) -> Self {
284        Self::new(x, penalty, offset)
285    }
286}
287
288impl<P, L, X, Penalty> ParameterBlock<P, L, LinearPredictorBlock<X>, Penalty>
289where
290    X: DesignMatrix,
291{
292    /// Creates a linear block from a design matrix.
293    #[must_use]
294    #[inline]
295    pub fn linear(x: X, penalty: Penalty, offset: usize) -> Self {
296        Self::new(LinearPredictorBlock::new(x), penalty, offset)
297    }
298}
299
300impl<P, L, X, Penalty> ParameterBlock<P, L, X, Penalty> {
301    #[inline]
302    const fn from_len(x: X, penalty: Penalty, offset: usize, len: usize) -> Self {
303        Self {
304            x,
305            penalty,
306            offset,
307            len,
308            marker: PhantomData,
309        }
310    }
311
312    /// Returns a copy of the block with a new offset.
313    #[must_use]
314    #[inline]
315    pub const fn with_offset(mut self, offset: usize) -> Self {
316        self.offset = offset;
317        self
318    }
319
320    /// Predictor block.
321    #[must_use]
322    #[inline]
323    pub const fn x(&self) -> &X {
324        &self.x
325    }
326
327    /// Penalty applied to the block's coefficients.
328    #[must_use]
329    #[inline]
330    pub const fn penalty(&self) -> &Penalty {
331        &self.penalty
332    }
333
334    #[must_use]
335    #[inline]
336    pub(crate) const fn offset(&self) -> usize {
337        self.offset
338    }
339
340    /// Coefficient range of the block in the common beta vector.
341    ///
342    /// # Panics
343    ///
344    /// Panics if `offset + len` overflows. Use [`Self::try_range`] when the
345    /// offset may come from unchecked external input.
346    #[must_use]
347    #[inline]
348    pub const fn range(&self) -> Range<usize> {
349        self.offset..self.end()
350    }
351
352    /// Index immediately after the last coefficient of the block.
353    ///
354    /// # Panics
355    ///
356    /// Panics if `offset + len` overflows. Use [`Self::try_range`] for
357    /// recoverable validation.
358    #[must_use]
359    #[inline]
360    pub const fn end(&self) -> usize {
361        self.offset
362            .checked_add(self.len)
363            .expect("parameter block range end must fit in usize")
364    }
365
366    /// Number of coefficients in the block.
367    #[must_use]
368    #[inline]
369    pub const fn len(&self) -> usize {
370        self.len
371    }
372
373    /// `true` if the block contains no coefficients.
374    #[must_use]
375    #[inline]
376    pub const fn is_empty(&self) -> bool {
377        self.len == 0
378    }
379}
380
381impl<P, L, X, Penalty> ParameterBlock<P, L, X, Penalty>
382where
383    P: ParameterName,
384{
385    /// Validates and returns the block's coefficient range.
386    ///
387    /// # Errors
388    ///
389    /// Returns [`ModelError::BlockRangeOverflow`] if `offset + len` does not fit
390    /// in `usize`.
391    #[inline]
392    pub fn try_range(&self) -> Result<Range<usize>, ModelError> {
393        let end = self
394            .offset
395            .checked_add(self.len)
396            .ok_or(ModelError::BlockRangeOverflow {
397                parameter: P::NAME,
398                offset: self.offset,
399                len: self.len,
400            })?;
401        Ok(self.offset..end)
402    }
403}
404
405/// Stable public name for a distribution parameter marker.
406pub trait ParameterName {
407    /// Name used in parameter layouts and unpacked coefficient views.
408    const NAME: &'static str;
409}
410
411/// Tuple contract implemented for typed parameter block tuples up to arity 8.
412pub trait AssignParameterOffsets: Sized {
413    /// Returns `self` with sequential offsets starting at `start`.
414    #[must_use]
415    fn assign_offsets(self, start: usize) -> Self;
416}
417
418/// Fallible tuple contract for assigning typed parameter block offsets.
419pub trait TryAssignParameterOffsets: Sized {
420    /// Returns `self` with sequential offsets starting at `start`.
421    ///
422    /// # Errors
423    ///
424    /// Returns [`ModelError::BlockRangeOverflow`] if a block range would not
425    /// fit in `usize`.
426    fn try_assign_offsets(self, start: usize) -> Result<Self, ModelError>;
427}
428
429macro_rules! impl_assign_offsets {
430    (
431        types = ($($block:ident),+);
432        vars = ($($var:ident),+)
433    ) => {
434        impl<$($block,)+> AssignParameterOffsets for ($($block,)+)
435        where
436            $($block: OffsetAssignable,)+
437        {
438            #[inline]
439            fn assign_offsets(self, start: usize) -> Self {
440                let ($($var,)+) = self;
441                let mut offset = start;
442                $(
443                    let $var = $var.with_assigned_offset(offset);
444                    let len = $var.assigned_len();
445                    offset = offset
446                        .checked_add(len)
447                        .expect("parameter block layout must fit in usize");
448                )+
449                let _ = offset;
450                ($($var,)+)
451            }
452        }
453
454        impl<$($block,)+> TryAssignParameterOffsets for ($($block,)+)
455        where
456            $($block: TryOffsetAssignable,)+
457        {
458            #[inline]
459            fn try_assign_offsets(self, start: usize) -> Result<Self, ModelError> {
460                let ($($var,)+) = self;
461                let mut offset = start;
462                $(
463                    let $var = $var.with_assigned_offset(offset);
464                    let assigned_offset = $var.assigned_offset();
465                    let assigned_len = $var.assigned_len();
466                    offset = offset.checked_add(assigned_len).ok_or(
467                        ModelError::BlockRangeOverflow {
468                            parameter: $var.assigned_name(),
469                            offset: assigned_offset,
470                            len: assigned_len,
471                        },
472                    )?;
473                )+
474                let _ = offset;
475                Ok(($($var,)+))
476            }
477        }
478    };
479}
480
481impl<P, L, X, Penalty> OffsetAssignable for ParameterBlock<P, L, X, Penalty> {
482    fn with_assigned_offset(self, offset: usize) -> Self {
483        self.with_offset(offset)
484    }
485
486    fn assigned_len(&self) -> usize {
487        self.len()
488    }
489}
490
491impl<P, L, X, Penalty> TryOffsetAssignable for ParameterBlock<P, L, X, Penalty>
492where
493    P: ParameterName,
494{
495    fn with_assigned_offset(self, offset: usize) -> Self {
496        self.with_offset(offset)
497    }
498
499    fn assigned_offset(&self) -> usize {
500        self.offset()
501    }
502
503    fn assigned_len(&self) -> usize {
504        self.len()
505    }
506
507    fn assigned_name(&self) -> &'static str {
508        P::NAME
509    }
510}
511
512trait OffsetAssignable: Sized {
513    fn with_assigned_offset(self, offset: usize) -> Self;
514    fn assigned_len(&self) -> usize;
515}
516
517trait TryOffsetAssignable: Sized {
518    fn with_assigned_offset(self, offset: usize) -> Self;
519    fn assigned_offset(&self) -> usize;
520    fn assigned_len(&self) -> usize;
521    fn assigned_name(&self) -> &'static str;
522}
523
524impl_assign_offsets!(types = (B1); vars = (b1));
525impl_assign_offsets!(types = (B1, B2); vars = (b1, b2));
526impl_assign_offsets!(types = (B1, B2, B3); vars = (b1, b2, b3));
527impl_assign_offsets!(types = (B1, B2, B3, B4); vars = (b1, b2, b3, b4));
528impl_assign_offsets!(types = (B1, B2, B3, B4, B5); vars = (b1, b2, b3, b4, b5));
529impl_assign_offsets!(types = (B1, B2, B3, B4, B5, B6); vars = (b1, b2, b3, b4, b5, b6));
530impl_assign_offsets!(
531    types = (B1, B2, B3, B4, B5, B6, B7);
532    vars = (b1, b2, b3, b4, b5, b6, b7)
533);
534impl_assign_offsets!(
535    types = (B1, B2, B3, B4, B5, B6, B7, B8);
536    vars = (b1, b2, b3, b4, b5, b6, b7, b8)
537);
538
539#[cfg(test)]
540mod tests {
541    use crate::{DenseDesign, Identity, LinearPredictorBlock, NoPenalty};
542
543    use super::{
544        Mu, Nu, ParameterBlock, ParameterBlocks, Precision, Rate, Scale, Shape, Sigma, Tau,
545    };
546
547    #[test]
548    fn parameter_blocks_assign_offsets_for_one_block() {
549        let mu = ParameterBlock::<Mu, Identity, _, _>::linear(
550            DenseDesign::from_rows(&[[1.0, 2.0]]),
551            NoPenalty,
552            99,
553        );
554
555        let (mu,) = ParameterBlocks::new((mu,));
556
557        assert_eq!(mu.range(), 0..2);
558    }
559
560    #[test]
561    fn parameter_blocks_assign_offsets_for_two_blocks() {
562        let mu = ParameterBlock::<Mu, Identity, _, _>::linear(
563            DenseDesign::from_rows(&[[1.0, 2.0]]),
564            NoPenalty,
565            99,
566        );
567        let sigma = ParameterBlock::<Sigma, Identity, _, _>::linear(
568            DenseDesign::from_rows(&[[1.0, 2.0, 3.0]]),
569            NoPenalty,
570            99,
571        );
572
573        let (mu, sigma) = ParameterBlocks::new((mu, sigma));
574
575        assert_eq!(mu.range(), 0..2);
576        assert_eq!(sigma.range(), 2..5);
577    }
578
579    #[test]
580    fn parameter_blocks_assign_offsets_for_eight_blocks_with_start() {
581        let blocks = (
582            intercept_block::<Mu>(),
583            intercept_block::<Sigma>(),
584            intercept_block::<Nu>(),
585            intercept_block::<Tau>(),
586            intercept_block::<Shape>(),
587            intercept_block::<Scale>(),
588            intercept_block::<Rate>(),
589            intercept_block::<Precision>(),
590        );
591
592        let (b1, b2, b3, b4, b5, b6, b7, b8) = ParameterBlocks::with_start(10, blocks);
593
594        assert_eq!(b1.range(), 10..11);
595        assert_eq!(b2.range(), 11..12);
596        assert_eq!(b3.range(), 12..13);
597        assert_eq!(b4.range(), 13..14);
598        assert_eq!(b5.range(), 14..15);
599        assert_eq!(b6.range(), 15..16);
600        assert_eq!(b7.range(), 16..17);
601        assert_eq!(b8.range(), 17..18);
602    }
603
604    #[test]
605    fn parameter_block_try_range_reports_overflow() {
606        let block = ParameterBlock::<Mu, Identity, _, _>::linear(
607            DenseDesign::from_rows(&[[1.0, 2.0]]),
608            NoPenalty,
609            usize::MAX,
610        );
611
612        assert_eq!(
613            block.try_range().unwrap_err(),
614            crate::ModelError::BlockRangeOverflow {
615                parameter: "mu",
616                offset: usize::MAX,
617                len: 2,
618            }
619        );
620    }
621
622    #[test]
623    fn parameter_blocks_try_with_start_reports_layout_overflow() {
624        let mu = ParameterBlock::<Mu, Identity, _, _>::linear(
625            DenseDesign::from_rows(&[[1.0, 2.0]]),
626            NoPenalty,
627            99,
628        );
629
630        assert_eq!(
631            ParameterBlocks::try_with_start(usize::MAX, (mu,)).unwrap_err(),
632            crate::ModelError::BlockRangeOverflow {
633                parameter: "mu",
634                offset: usize::MAX,
635                len: 2,
636            }
637        );
638    }
639
640    fn intercept_block<P>()
641    -> ParameterBlock<P, Identity, LinearPredictorBlock<DenseDesign>, NoPenalty> {
642        ParameterBlock::linear(DenseDesign::intercept(1), NoPenalty, 99)
643    }
644}