1use std::{marker::PhantomData, ops::Range};
2
3use crate::{DesignMatrix, LinearPredictorBlock, ModelError, PredictorBlock};
4
5#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
13pub struct ParameterBlocks;
14
15impl ParameterBlocks {
16 #[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 #[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 #[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 #[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#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
81pub struct Mu;
82
83impl ParameterName for Mu {
84 const NAME: &'static str = "mu";
85}
86
87#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
89pub struct Mean;
90
91impl ParameterName for Mean {
92 const NAME: &'static str = "mean";
93}
94
95#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
97pub struct Median;
98
99impl ParameterName for Median {
100 const NAME: &'static str = "median";
101}
102
103#[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#[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#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
121pub struct Sigma;
122
123impl ParameterName for Sigma {
124 const NAME: &'static str = "sigma";
125}
126
127#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
129pub struct Cv;
130
131impl ParameterName for Cv {
132 const NAME: &'static str = "cv";
133}
134
135#[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#[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#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
153pub struct Nu;
154
155impl ParameterName for Nu {
156 const NAME: &'static str = "nu";
157}
158
159#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
161pub struct Tau;
162
163impl ParameterName for Tau {
164 const NAME: &'static str = "tau";
165}
166
167#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
169pub struct Rate;
170
171impl ParameterName for Rate {
172 const NAME: &'static str = "rate";
173}
174
175#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
177pub struct Dispersion;
178
179impl ParameterName for Dispersion {
180 const NAME: &'static str = "dispersion";
181}
182
183#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
185pub struct Shape;
186
187impl ParameterName for Shape {
188 const NAME: &'static str = "shape";
189}
190
191#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
193pub struct Size;
194
195impl ParameterName for Size {
196 const NAME: &'static str = "size";
197}
198
199#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
201pub struct Scale;
202
203impl ParameterName for Scale {
204 const NAME: &'static str = "scale";
205}
206
207#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
209pub struct Probability;
210
211impl ParameterName for Probability {
212 const NAME: &'static str = "probability";
213}
214
215#[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#[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#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
233pub struct Power;
234
235impl ParameterName for Power {
236 const NAME: &'static str = "power";
237}
238
239#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
244pub struct Precision;
245
246impl ParameterName for Precision {
247 const NAME: &'static str = "precision";
248}
249
250#[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 #[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 #[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 #[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 #[must_use]
314 #[inline]
315 pub const fn with_offset(mut self, offset: usize) -> Self {
316 self.offset = offset;
317 self
318 }
319
320 #[must_use]
322 #[inline]
323 pub const fn x(&self) -> &X {
324 &self.x
325 }
326
327 #[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 #[must_use]
347 #[inline]
348 pub const fn range(&self) -> Range<usize> {
349 self.offset..self.end()
350 }
351
352 #[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 #[must_use]
368 #[inline]
369 pub const fn len(&self) -> usize {
370 self.len
371 }
372
373 #[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 #[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
405pub trait ParameterName {
407 const NAME: &'static str;
409}
410
411pub trait AssignParameterOffsets: Sized {
413 #[must_use]
415 fn assign_offsets(self, start: usize) -> Self;
416}
417
418pub trait TryAssignParameterOffsets: Sized {
420 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}