Skip to main content

infotheory/api/
types.rs

1//! Public type definitions for the spec-first API surface.
2
3use crate::coders::CoderType;
4use crate::error::{InfotheoryError, InfotheoryResult};
5use std::num::NonZeroUsize;
6use std::sync::Arc;
7
8/// Typed ZPAQ method specification.
9#[derive(Clone, Debug, Eq, PartialEq)]
10pub enum ZpaqMethodSpec {
11    /// Literal ZPAQ method string.
12    Literal {
13        /// ZPAQ method string accepted by the backend.
14        value: String,
15    },
16}
17
18impl ZpaqMethodSpec {
19    /// Construct a literal ZPAQ method specification.
20    pub fn literal(value: impl Into<String>) -> Self {
21        Self::Literal {
22            value: value.into(),
23        }
24    }
25
26    /// Return the canonical method string.
27    pub fn value(&self) -> &str {
28        match self {
29            Self::Literal { value } => value,
30        }
31    }
32}
33
34impl From<String> for ZpaqMethodSpec {
35    fn from(value: String) -> Self {
36        Self::literal(value)
37    }
38}
39
40impl From<&str> for ZpaqMethodSpec {
41    fn from(value: &str) -> Self {
42        Self::literal(value)
43    }
44}
45
46/// How generated symbols should update the model state.
47#[derive(Clone, Copy, Debug, PartialEq, Eq)]
48#[non_exhaustive]
49pub enum GenerationUpdateMode {
50    /// Keep adapting/fitting on generated bytes.
51    Adaptive,
52    /// Freeze fitted parameters/statistics and only advance conditioning state.
53    Frozen,
54}
55
56/// How to pick the next byte from the model distribution.
57#[derive(Clone, Copy, Debug, PartialEq, Eq)]
58#[non_exhaustive]
59pub enum GenerationStrategy {
60    /// Deterministic argmax over the next-byte distribution.
61    Greedy,
62    /// Seeded sampling from the next-byte distribution.
63    Sample,
64}
65
66/// Generation options shared by the library API and CLI.
67#[derive(Clone, Copy, Debug)]
68#[non_exhaustive]
69pub struct GenerationConfig {
70    /// Byte-selection strategy.
71    pub strategy: GenerationStrategy,
72    /// Whether generated bytes should keep adapting the model.
73    pub update_mode: GenerationUpdateMode,
74    /// RNG seed used by [`GenerationStrategy::Sample`].
75    pub seed: u64,
76    /// Softmax temperature for sampling. `<= 0` behaves like greedy.
77    pub temperature: f64,
78    /// Optional top-k truncation. `0` disables it.
79    pub top_k: usize,
80    /// Optional nucleus truncation. Values `>= 1.0` disable it.
81    pub top_p: f64,
82}
83
84impl Default for GenerationConfig {
85    fn default() -> Self {
86        Self::sampled_frozen(42)
87    }
88}
89
90impl GenerationConfig {
91    /// Deterministic frozen continuation.
92    pub const fn greedy_frozen() -> Self {
93        Self {
94            strategy: GenerationStrategy::Greedy,
95            update_mode: GenerationUpdateMode::Frozen,
96            seed: 0xD00D_F00D_CAFE_BABEu64,
97            temperature: 1.0,
98            top_k: 0,
99            top_p: 1.0,
100        }
101    }
102
103    /// Seeded frozen sampling from the model distribution.
104    pub const fn sampled_frozen(seed: u64) -> Self {
105        Self {
106            strategy: GenerationStrategy::Sample,
107            update_mode: GenerationUpdateMode::Frozen,
108            seed,
109            temperature: 1.0,
110            top_k: 0,
111            top_p: 1.0,
112        }
113    }
114}
115
116/// Core predictive model class used by the library.
117///
118/// `RateBackend` is the shared model class behind entropy-rate estimation,
119/// rate-coded compression, generation, and the world-model interface used by
120/// MC-AIXI/AIQI planners.
121#[derive(Clone)]
122#[non_exhaustive]
123pub enum RateBackend {
124    /// ROSA+ suffix-automaton estimator.
125    ///
126    /// `max_order < 0` enables ROSA's adaptive order selection over the full
127    /// suffix automaton; `max_order >= 0` caps the predictive context length.
128    RosaPlus {
129        /// Maximum context length used by the ROSA predictor.
130        ///
131        /// `< 0` means "adaptive / full SAM"; `>= 0` is a fixed cap.
132        max_order: i64,
133    },
134    /// Local contiguous match predictor.
135    Match {
136        /// Number of retained hash bits for suffix lookup.
137        hash_bits: usize,
138        /// Minimum repeat length required before predicting.
139        min_len: usize,
140        /// Maximum repeat length used for confidence scaling.
141        max_len: usize,
142        /// Residual probability mass left for non-match symbols.
143        base_mix: f64,
144        /// Confidence multiplier applied to short-match tapering.
145        confidence_scale: f64,
146    },
147    /// Sparse/gapped local match predictor.
148    SparseMatch {
149        /// Number of retained hash bits for spaced-suffix lookup.
150        hash_bits: usize,
151        /// Minimum spaced repeat length required before predicting.
152        min_len: usize,
153        /// Maximum spaced repeat length used for confidence scaling.
154        max_len: usize,
155        /// Minimum gap between matched bytes.
156        gap_min: usize,
157        /// Maximum gap between matched bytes.
158        gap_max: usize,
159        /// Residual probability mass left for non-match symbols.
160        base_mix: f64,
161        /// Confidence multiplier applied to short-match tapering.
162        confidence_scale: f64,
163    },
164    /// Pure-Rust bounded-memory PPMD-style model.
165    Ppmd {
166        /// Maximum context order.
167        order: usize,
168        /// Approximate memory budget in MiB.
169        memory_mb: usize,
170    },
171    /// Exact online Sequitur grammar backend with byte-level predictive readout.
172    Sequitur {
173        /// Maximum number of terminal bytes retained per grammar-derived context.
174        context_bytes: usize,
175    },
176    #[cfg(feature = "backend-mamba")]
177    /// Typed Mamba method specification.
178    MambaMethod {
179        /// Mamba method specification.
180        method: crate::mambazip::MethodSpec,
181    },
182    #[cfg(feature = "backend-rwkv")]
183    /// Typed RWKV7 method specification.
184    Rwkv7Method {
185        /// RWKV7 method specification.
186        method: crate::rwkvzip::MethodSpec,
187    },
188    /// ZPAQ compression-based rate model (streamable methods only).
189    Zpaq {
190        /// Typed ZPAQ method specification.
191        method: ZpaqMethodSpec,
192    },
193    /// Online mixture over `RateBackend` experts.
194    ///
195    /// `Bayes`, `Switching`, and `Convex` follow
196    /// "On Ensemble Techniques for AIXI Approximation"; `FadingBayes`,
197    /// `Mdl`, and `Neural` are repository extensions.
198    Mixture {
199        /// Mixture expert/runtime specification.
200        spec: Arc<MixtureSpec>,
201    },
202    /// Particle-latent filter ensemble.
203    Particle {
204        /// Particle filter specification.
205        spec: Arc<ParticleSpec>,
206    },
207    /// Calibrated wrapper over another bytewise backend.
208    Calibrated {
209        /// Calibration specification.
210        spec: Arc<CalibratedSpec>,
211    },
212    /// Action-Conditional CTW (single context tree).
213    Ctw {
214        /// Context tree depth.
215        depth: usize,
216    },
217    /// Factorized Action-Conditional CTW (k trees for k-bit percepts).
218    FacCtw {
219        /// Base context depth.
220        base_depth: usize,
221        /// Planner/percept bit cardinality for AIXI-style consumers.
222        ///
223        /// Rate-backend byte execution uses `encoding_bits` as the symbol
224        /// decomposition width. `num_percept_bits` is retained for controller
225        /// specs whose percept cardinality can differ from byte encoding width.
226        num_percept_bits: usize,
227        /// Encoding width in bits for rate execution.
228        ///
229        /// Valid compiled FAC-CTW specs require `1..=8`.
230        encoding_bits: usize,
231        /// Optional bit order for decomposing symbols.
232        ///
233        /// `None` keeps the compatibility default: 8-bit byte symbols use
234        /// MSB-first ordering, while non-byte widths retain the legacy
235        /// LSB-first FAC-CTW convention. `Some(true)` selects MSB-first
236        /// explicitly; `Some(false)` selects legacy LSB-first explicitly.
237        msb_first: Option<bool>,
238    },
239}
240
241/// Compression backend used by NCD/compression-size operations.
242#[derive(Clone)]
243#[non_exhaustive]
244pub enum CompressionBackend {
245    /// ZPAQ compressor with explicit method string.
246    Zpaq {
247        /// Typed ZPAQ method specification.
248        method: ZpaqMethodSpec,
249        /// Internal ZPAQ compression thread count (`1` => single-threaded).
250        threads: NonZeroUsize,
251    },
252    #[cfg(feature = "backend-rwkv")]
253    /// RWKV7 compressor configured by typed method specification.
254    Rwkv7 {
255        /// RWKV7 method specification.
256        method: crate::rwkvzip::MethodSpec,
257        /// Entropy coder used for coding model PDFs.
258        coder: CoderType,
259    },
260    /// Generic rate-coded compressor wrapping an arbitrary rate backend.
261    Rate {
262        /// Predictive rate backend.
263        rate_backend: RateBackend,
264        /// Entropy coder used for coding model PDFs.
265        coder: CoderType,
266        /// Framing mode for output payloads.
267        framing: crate::compression::FramingMode,
268    },
269}
270
271/// Shared maximum nesting depth for recursive mixture specifications.
272pub const MAX_MIXTURE_NESTING: usize = 8;
273
274/// Mixture policy kind for rate-backend mixtures.
275#[derive(Clone, Copy, Debug, Eq, PartialEq)]
276#[non_exhaustive]
277pub enum MixtureKind {
278    /// Standard Bayesian mixture with fixed expert weights.
279    Bayes,
280    /// Bayesian mixture with exponential weight decay.
281    FadingBayes,
282    /// Switching mixture using the fixed-share update from
283    /// "On Ensemble Techniques for AIXI Approximation".
284    Switching,
285    /// Online convex mixture with projected-simplex weight updates.
286    Convex,
287    /// MDL-style best-expert selector.
288    Mdl,
289    /// Bytewise neural logistic mixer (PAQ style adaptation).
290    Neural,
291}
292
293/// Adaptive schedule family for switching and convex mixtures.
294#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
295#[non_exhaustive]
296pub enum MixtureScheduleMode {
297    /// Use the implementation's default exposed parameterization.
298    ///
299    /// - `Switching`: constant switch rate `alpha`
300    /// - `Convex`: step size `alpha / sqrt(t)`
301    #[default]
302    Default,
303    /// Use the theorem schedule from
304    /// "On Ensemble Techniques for AIXI Approximation".
305    ///
306    /// - `Switching`: `alpha_t = 1 / t`
307    /// - `Convex`: `eta_t = epsilon / sqrt(t)` under this implementation's
308    ///   natural-log gradient, matching the bit-loss schedule analyzed in
309    ///   "On Ensemble Techniques for AIXI Approximation" after
310    ///   accounting for the `1 / ln(2)` factor in the base-2 gradient
311    ///
312    /// This preserves configured expert priors; exact theorem hypotheses for
313    /// switching still additionally require uniform priors.
314    Theorem,
315}
316
317/// Fixed context families for calibrated PDF wrappers.
318#[derive(Clone, Copy, Debug, Eq, PartialEq)]
319#[non_exhaustive]
320pub enum CalibrationContextKind {
321    /// Single global calibration row.
322    Global,
323    /// Previous-byte class only.
324    ByteClass,
325    /// Text-structure-aware context hash.
326    Text,
327    /// Repeat-aware context hash.
328    Repeat,
329    /// Joint text/repeat-aware context hash.
330    TextRepeat,
331}
332
333/// Configuration for an SSE-calibrated wrapper rate backend.
334#[derive(Clone)]
335#[non_exhaustive]
336pub struct CalibratedSpec {
337    /// Base backend whose bit probabilities are calibrated.
338    pub base: RateBackend,
339    /// Context family controlling SSE table row selection.
340    pub context: CalibrationContextKind,
341    /// Stretched-probability bins per row.
342    pub bins: usize,
343    /// Initial online learning fraction for observed-bit updates.
344    pub learning_rate: f64,
345    /// Symmetric stretched-logit clip used by SSE quantization.
346    pub bias_clip: f64,
347}
348
349impl CalibratedSpec {
350    /// Create a new SSE-calibrated spec with PAQ/ZPAQ-style defaults.
351    pub fn new(base: RateBackend, context: CalibrationContextKind) -> Self {
352        Self {
353            base,
354            context,
355            bins: 32,
356            learning_rate: 1.0 / 32.0,
357            bias_clip: 16.0,
358        }
359    }
360
361    /// Override the per-row SSE bin count.
362    #[must_use]
363    pub fn with_bins(mut self, bins: usize) -> Self {
364        self.bins = bins;
365        self
366    }
367
368    /// Override the initial online learning fraction.
369    #[must_use]
370    pub fn with_learning_rate(mut self, learning_rate: f64) -> Self {
371        self.learning_rate = learning_rate;
372        self
373    }
374
375    /// Override the symmetric stretched-logit clip.
376    #[must_use]
377    pub fn with_bias_clip(mut self, bias_clip: f64) -> Self {
378        self.bias_clip = bias_clip;
379        self
380    }
381}
382
383/// Expert specification for mixture backends.
384///
385/// Algorithm-specific configuration (such as ROSA's `max_order`) lives inside
386/// the expert's [`RateBackend`] variant; the expert spec itself only carries
387/// mixture-level metadata.
388#[derive(Clone)]
389#[non_exhaustive]
390pub struct MixtureExpertSpec {
391    /// Optional expert display name.
392    pub name: Option<String>,
393    /// Log prior weight (natural log). Uniform priors can be `0.0`.
394    pub log_prior: f64,
395    /// Underlying backend for this expert.
396    pub backend: RateBackend,
397}
398
399impl MixtureExpertSpec {
400    /// Create a new expert spec with a default 0.0 prior and no length limits.
401    pub fn new(backend: RateBackend) -> Self {
402        Self {
403            name: None,
404            log_prior: 0.0,
405            backend,
406        }
407    }
408
409    /// Set the human-readable expert name.
410    #[must_use]
411    pub fn with_name(mut self, name: impl Into<String>) -> Self {
412        self.name = Some(name.into());
413        self
414    }
415
416    /// Clear any previously set expert name.
417    #[must_use]
418    pub fn without_name(mut self) -> Self {
419        self.name = None;
420        self
421    }
422
423    /// Set the log-prior weight (natural log) for this expert.
424    #[must_use]
425    pub fn with_log_prior(mut self, log_prior: f64) -> Self {
426        self.log_prior = log_prior;
427        self
428    }
429}
430
431/// Mixture specification for rate-backend mixtures.
432#[derive(Clone)]
433#[non_exhaustive]
434pub struct MixtureSpec {
435    /// Mixture policy.
436    pub kind: MixtureKind,
437    /// Adaptive schedule family for supported mixture kinds.
438    pub schedule: MixtureScheduleMode,
439    /// Shared scalar parameter: switch rate for `Switching`, step-size scale for `Convex`,
440    /// learning rate for `Neural`, and generic alpha for the remaining families.
441    ///
442    /// In theorem mode for `Switching` and `Convex`, this field is retained for API
443    /// compatibility but is not used by the update schedule.
444    pub alpha: f64,
445    /// Decay factor for fading Bayes mixtures.
446    pub decay: Option<f64>,
447    /// Expert list.
448    pub experts: Vec<MixtureExpertSpec>,
449}
450
451/// Configuration for a particle-latent filter ensemble rate backend.
452#[derive(Clone, Debug)]
453#[non_exhaustive]
454pub struct ParticleSpec {
455    /// Number of particles in the ensemble.
456    pub num_particles: usize,
457    /// Context window length for rolling byte context.
458    pub context_window: usize,
459    /// Number of latent update unroll steps per byte.
460    pub unroll_steps: usize,
461    /// Number of latent cells per particle.
462    pub num_cells: usize,
463    /// Dimensionality of each latent cell.
464    pub cell_dim: usize,
465    /// Number of discrete rules for soft routing.
466    pub num_rules: usize,
467    /// Hidden dimension for the selector MLP.
468    pub selector_hidden: usize,
469    /// Hidden dimension for each rule MLP.
470    pub rule_hidden: usize,
471    /// Dimension of per-rule noise input (ignored when deterministic).
472    pub noise_dim: usize,
473    /// Whether to use fully deterministic execution (no RNG).
474    pub deterministic: bool,
475    /// Whether to inject noise into rule inputs (ignored when deterministic).
476    pub enable_noise: bool,
477    /// Base scale for deterministic hash-noise injected into rule inputs.
478    pub noise_scale: f64,
479    /// Number of steps over which injected noise linearly anneals to zero.
480    pub noise_anneal_steps: usize,
481    /// Learning rate for readout layer SGD.
482    pub learning_rate_readout: f64,
483    /// Learning rate for selector MLP SGD.
484    pub learning_rate_selector: f64,
485    /// Learning rate for rule MLP SGD.
486    pub learning_rate_rule: f64,
487    /// Truncated backpropagation-through-time depth (number of recent steps).
488    pub bptt_depth: usize,
489    /// Momentum coefficient for selector/rule online updates (in [0, 1)).
490    pub optimizer_momentum: f64,
491    /// Gradient clipping threshold (max abs value per element).
492    pub grad_clip: f64,
493    /// Latent cell state clipping threshold (max abs value per element).
494    pub state_clip: f64,
495    /// Forgetting factor for particle log-weights (0 = no forgetting).
496    pub forget_lambda: f64,
497    /// Effective sample size ratio threshold for resampling (in (0, 1]).
498    pub resample_threshold: f64,
499    /// Fraction of particles to mutate after resampling (in [0, 1]).
500    pub mutate_fraction: f64,
501    /// Scale of hash-noise perturbation applied during mutation.
502    pub mutate_scale: f64,
503    /// Whether mutation also perturbs model parameters (state is always mutated).
504    pub mutate_model_params: bool,
505    /// Diagnostics print interval in steps (0 disables particle diagnostics logs).
506    pub diagnostics_interval: usize,
507    /// Minimum probability floor for numerical stability.
508    pub min_prob: f64,
509    /// Master seed for deterministic initialization and mutation.
510    pub seed: u64,
511}
512
513impl RateBackend {
514    /// Returns the current build's implicit default rate backend.
515    pub fn try_default() -> InfotheoryResult<Self> {
516        crate::runtime::first_enabled_default_rate_backend_spec().ok_or_else(|| {
517            let mut requested = Vec::new();
518            for descriptor in crate::runtime::RATE_BACKEND_REGISTRY {
519                if let Some(feature) = descriptor.feature
520                    && !requested.contains(&feature)
521                {
522                    requested.push(feature);
523                }
524            }
525            InfotheoryError::unsupported(format!(
526                "no default rate backend is available in this build; enable one of: {}",
527                requested.join(", ")
528            ))
529        })
530    }
531
532    /// Stable internal backend identity.
533    pub(crate) fn kind(&self) -> crate::runtime::RateBackendKind {
534        match self {
535            RateBackend::RosaPlus { .. } => crate::runtime::RateBackendKind::RosaPlus,
536            RateBackend::Match { .. } => crate::runtime::RateBackendKind::Match,
537            RateBackend::SparseMatch { .. } => crate::runtime::RateBackendKind::SparseMatch,
538            RateBackend::Ppmd { .. } => crate::runtime::RateBackendKind::Ppmd,
539            RateBackend::Sequitur { .. } => crate::runtime::RateBackendKind::Sequitur,
540            #[cfg(feature = "backend-mamba")]
541            RateBackend::MambaMethod { .. } => crate::runtime::RateBackendKind::Mamba,
542            #[cfg(feature = "backend-rwkv")]
543            RateBackend::Rwkv7Method { .. } => crate::runtime::RateBackendKind::Rwkv7,
544            RateBackend::Zpaq { .. } => crate::runtime::RateBackendKind::Zpaq,
545            RateBackend::Mixture { .. } => crate::runtime::RateBackendKind::Mixture,
546            RateBackend::Particle { .. } => crate::runtime::RateBackendKind::Particle,
547            RateBackend::Calibrated { .. } => crate::runtime::RateBackendKind::Calibrated,
548            RateBackend::Ctw { .. } => crate::runtime::RateBackendKind::Ctw,
549            RateBackend::FacCtw { .. } => crate::runtime::RateBackendKind::FacCtw,
550        }
551    }
552
553    pub(crate) fn descriptor(
554        &self,
555    ) -> Result<&'static crate::runtime::RateBackendDescriptor, String> {
556        crate::runtime::describe_rate_backend_kind(self.kind())
557    }
558
559    /// Validate and canonicalize this backend in the provided compilation environment.
560    pub fn validate_in(
561        &self,
562        env: &crate::spec::SpecEnvironment,
563    ) -> crate::spec::SpecResult<crate::spec::ValidatedRateBackend> {
564        crate::spec::core::validate_rate_backend_in(self, env)
565    }
566
567    /// Validate and canonicalize this backend using the default environment.
568    pub fn validate(&self) -> crate::spec::SpecResult<crate::spec::ValidatedRateBackend> {
569        self.validate_in(&crate::spec::SpecEnvironment::default())
570    }
571
572    /// Validate and compile this backend in the provided compilation environment.
573    pub fn compile_in(
574        &self,
575        env: &crate::spec::SpecEnvironment,
576    ) -> crate::spec::SpecResult<crate::spec::CompiledRateBackend> {
577        self.validate_in(env)?.compile()
578    }
579
580    /// Validate and compile this backend using the default environment.
581    pub fn compile(&self) -> crate::spec::SpecResult<crate::spec::CompiledRateBackend> {
582        self.validate()?.compile()
583    }
584}
585
586impl CompressionBackend {
587    /// Construct a ZPAQ compression backend with default internal threading (`threads = 1`).
588    pub fn zpaq(method: impl Into<ZpaqMethodSpec>) -> Self {
589        Self::Zpaq {
590            method: method.into(),
591            threads: NonZeroUsize::MIN,
592        }
593    }
594
595    /// Construct a ZPAQ compression backend with explicit internal thread count.
596    pub fn zpaq_with_threads(method: impl Into<ZpaqMethodSpec>, threads: NonZeroUsize) -> Self {
597        Self::Zpaq {
598            method: method.into(),
599            threads,
600        }
601    }
602
603    /// Returns the current build's implicit default compression backend.
604    pub fn try_default() -> InfotheoryResult<Self> {
605        if crate::runtime::COMPRESSION_BACKEND_REGISTRY
606            .iter()
607            .any(|descriptor| {
608                descriptor.enabled
609                    && descriptor.kind == crate::runtime::CompressionBackendKind::Zpaq
610            })
611        {
612            return Ok(Self::zpaq("5"));
613        }
614
615        Ok(CompressionBackend::Rate {
616            rate_backend: RateBackend::try_default()?,
617            coder: crate::coders::CoderType::AC,
618            framing: crate::compression::FramingMode::Raw,
619        })
620    }
621
622    /// Stable internal backend identity.
623    pub(crate) fn kind(&self) -> crate::runtime::CompressionBackendKind {
624        match self {
625            CompressionBackend::Zpaq { .. } => crate::runtime::CompressionBackendKind::Zpaq,
626            #[cfg(feature = "backend-rwkv")]
627            CompressionBackend::Rwkv7 { .. } => crate::runtime::CompressionBackendKind::Rwkv7,
628            CompressionBackend::Rate { coder, .. } => match coder {
629                crate::coders::CoderType::AC => crate::runtime::CompressionBackendKind::RateAc,
630                crate::coders::CoderType::RANS => crate::runtime::CompressionBackendKind::RateRans,
631            },
632        }
633    }
634
635    /// Canonical registry descriptor for this compression backend.
636    pub(crate) fn descriptor(
637        &self,
638    ) -> Result<&'static crate::runtime::CompressionBackendDescriptor, String> {
639        crate::runtime::describe_compression_backend_kind(self.kind())
640    }
641
642    /// Validate and canonicalize this backend in the provided compilation environment.
643    pub fn validate_in(
644        &self,
645        env: &crate::spec::SpecEnvironment,
646    ) -> crate::spec::SpecResult<crate::spec::ValidatedCompressionBackend> {
647        crate::spec::core::validate_compression_backend_in(self, env)
648    }
649
650    /// Validate and canonicalize this backend using the default environment.
651    pub fn validate(&self) -> crate::spec::SpecResult<crate::spec::ValidatedCompressionBackend> {
652        self.validate_in(&crate::spec::SpecEnvironment::default())
653    }
654
655    /// Validate and compile this backend in the provided compilation environment.
656    pub fn compile_in(
657        &self,
658        env: &crate::spec::SpecEnvironment,
659    ) -> crate::spec::SpecResult<crate::spec::CompiledCompressionBackend> {
660        self.validate_in(env)?.compile()
661    }
662
663    /// Validate and compile this backend using the default environment.
664    pub fn compile(&self) -> crate::spec::SpecResult<crate::spec::CompiledCompressionBackend> {
665        self.validate()?.compile()
666    }
667}
668/// Parse a mixture kind name with the shared alias table used across CLI, Python, and WASM.
669pub fn parse_mixture_kind_name(kind: &str) -> Result<MixtureKind, String> {
670    match kind.trim().to_ascii_lowercase().as_str() {
671        "bayes" => Ok(MixtureKind::Bayes),
672        "fading-bayes" => Ok(MixtureKind::FadingBayes),
673        "switching" => Ok(MixtureKind::Switching),
674        "convex" => Ok(MixtureKind::Convex),
675        "mdl" => Ok(MixtureKind::Mdl),
676        "neural" => Ok(MixtureKind::Neural),
677        other => Err(format!("unknown mixture kind '{other}'")),
678    }
679}
680
681/// Parse a mixture schedule mode with the shared alias table used across CLI, Python, and WASM.
682pub fn parse_mixture_schedule_name(schedule: &str) -> Result<MixtureScheduleMode, String> {
683    match schedule.trim().to_ascii_lowercase().as_str() {
684        "" | "default" => Ok(MixtureScheduleMode::Default),
685        "theorem" => Ok(MixtureScheduleMode::Theorem),
686        other => Err(format!("unknown mixture schedule '{other}'")),
687    }
688}
689
690impl MixtureSpec {
691    /// Build a mixture specification from kind and expert list.
692    pub fn new(kind: MixtureKind, experts: Vec<MixtureExpertSpec>) -> Self {
693        Self {
694            kind,
695            schedule: MixtureScheduleMode::Default,
696            alpha: 0.01,
697            decay: None,
698            experts,
699        }
700    }
701
702    /// Set the schedule family.
703    pub fn with_schedule(mut self, schedule: MixtureScheduleMode) -> Self {
704        self.schedule = schedule;
705        self
706    }
707
708    /// Set the family-specific alpha parameter.
709    pub fn with_alpha(mut self, alpha: f64) -> Self {
710        self.alpha = alpha;
711        self
712    }
713
714    /// Set fading decay factor.
715    pub fn with_decay(mut self, decay: f64) -> Self {
716        self.decay = Some(decay);
717        self
718    }
719
720    /// Validate the mixture configuration before building runtime state.
721    pub fn validate(&self) -> InfotheoryResult<()> {
722        validate_mixture_spec_with_depth(self, MAX_MIXTURE_NESTING)
723            .map_err(InfotheoryError::invalid_backend_config)
724    }
725
726    /// Convert to executable expert configs for runtime mixture evaluation.
727    pub fn build_experts(&self) -> Vec<crate::mixture::ExpertConfig> {
728        self.experts
729            .iter()
730            .map(|spec| {
731                crate::mixture::ExpertConfig::from_rate_backend(
732                    spec.name.clone(),
733                    spec.log_prior,
734                    spec.backend.clone(),
735                )
736            })
737            .collect()
738    }
739}
740
741fn validate_mixture_spec_with_depth(spec: &MixtureSpec, depth: usize) -> Result<(), String> {
742    if depth == 0 {
743        return Err("mixture spec nesting too deep".to_string());
744    }
745    validate_mixture_spec_shallow(spec)?;
746    for (index, expert) in spec.experts.iter().enumerate() {
747        validate_rate_backend_with_depth(&expert.backend, depth - 1).map_err(|err| {
748            if let Some(name) = expert.name.as_deref() {
749                format!("mixture expert '{name}' invalid: {err}")
750            } else {
751                format!("mixture expert #{} invalid: {err}", index + 1)
752            }
753        })?;
754    }
755    Ok(())
756}
757
758fn validate_mixture_spec_shallow(spec: &MixtureSpec) -> Result<(), String> {
759    if spec.experts.is_empty() {
760        return Err("mixture spec must include at least one expert".to_string());
761    }
762    if !spec.alpha.is_finite() {
763        return Err("mixture alpha must be finite".to_string());
764    }
765    if spec
766        .experts
767        .iter()
768        .any(|expert| !expert.log_prior.is_finite())
769    {
770        return Err("mixture expert log_prior must be finite".to_string());
771    }
772    if let Some(decay) = spec.decay
773        && !(decay.is_finite() && decay > 0.0 && decay < 1.0)
774    {
775        return Err("mixture decay must be in (0, 1)".to_string());
776    }
777    if matches!(spec.kind, MixtureKind::FadingBayes) && spec.decay.is_none() {
778        return Err("fading Bayes mixture requires decay".to_string());
779    }
780    if spec.schedule != MixtureScheduleMode::Default
781        && !matches!(spec.kind, MixtureKind::Switching | MixtureKind::Convex)
782    {
783        return Err(
784            "mixture schedule is only supported for switching and convex mixtures".to_string(),
785        );
786    }
787    match (spec.kind, spec.schedule) {
788        (MixtureKind::Switching, MixtureScheduleMode::Default) => {
789            if !(0.0..=1.0).contains(&spec.alpha) {
790                return Err("switching mixture alpha must be in [0, 1]".to_string());
791            }
792        }
793        (MixtureKind::Convex, MixtureScheduleMode::Default)
794        | (MixtureKind::Neural, MixtureScheduleMode::Default) => {
795            if spec.alpha <= 0.0 {
796                return Err("mixture alpha must be > 0".to_string());
797            }
798        }
799        (MixtureKind::Neural, MixtureScheduleMode::Theorem) => unreachable!(),
800        _ => {}
801    }
802    Ok(())
803}
804
805fn validate_rate_backend_with_depth(backend: &RateBackend, depth: usize) -> Result<(), String> {
806    let descriptor = crate::runtime::try_describe_rate_backend(backend)
807        .map_err(|err| format!("{err} (while validating rate backend)"))?;
808    if !descriptor.enabled {
809        let Some(feature) = descriptor.feature else {
810            return Err(format!(
811                "internal backend registry mismatch: disabled backend '{}' is missing required feature metadata",
812                descriptor.canonical
813            ));
814        };
815        return Err(format!(
816            "backend '{}' requires infotheory feature '{}'",
817            descriptor.canonical, feature
818        ));
819    }
820    match backend {
821        RateBackend::Sequitur { context_bytes } => {
822            if *context_bytes < 2 {
823                Err("sequitur context_bytes must be >= 2".to_string())
824            } else {
825                Ok(())
826            }
827        }
828        RateBackend::Mixture { spec } => validate_mixture_spec_with_depth(spec.as_ref(), depth),
829        RateBackend::Particle { spec } => spec.validate().map_err(|err| err.to_string()),
830        RateBackend::Calibrated { spec } => {
831            if depth == 0 {
832                return Err("calibrated spec nesting too deep".to_string());
833            }
834            validate_rate_backend_with_depth(&spec.base, depth - 1)
835                .map_err(|err| format!("calibrated base invalid: {err}"))
836        }
837        _ => Ok(()),
838    }
839}
840
841/// Validate a rate backend, including nested mixture/calibrated subgraphs.
842pub fn validate_rate_backend(backend: &RateBackend) -> InfotheoryResult<()> {
843    validate_rate_backend_with_depth(backend, MAX_MIXTURE_NESTING)
844        .map_err(InfotheoryError::invalid_backend_config)
845}
846
847/// Validate a compression backend, including nested rate-backed compressors.
848pub fn validate_compression_backend(backend: &CompressionBackend) -> InfotheoryResult<()> {
849    let descriptor = crate::runtime::try_describe_compression_backend(backend)
850        .map_err(InfotheoryError::invalid_backend_config)?;
851    if !descriptor.enabled {
852        let Some(feature) = descriptor.feature else {
853            return Err(InfotheoryError::invalid_backend_config(format!(
854                "internal backend registry mismatch: disabled compression backend '{}' is missing required feature metadata",
855                descriptor.canonical
856            )));
857        };
858        return Err(InfotheoryError::invalid_backend_config(format!(
859            "compression backend '{}' requires infotheory feature '{}'",
860            descriptor.canonical, feature
861        )));
862    }
863    if let CompressionBackend::Rate { rate_backend, .. } = backend {
864        validate_rate_backend(rate_backend).map_err(|err| {
865            InfotheoryError::invalid_backend_config(format!(
866                "rate-coded compression backend invalid: {err}"
867            ))
868        })?;
869    }
870    Ok(())
871}
872
873impl Default for ParticleSpec {
874    fn default() -> Self {
875        Self {
876            num_particles: 16,
877            context_window: 32,
878            unroll_steps: 2,
879            num_cells: 8,
880            cell_dim: 32,
881            num_rules: 4,
882            selector_hidden: 64,
883            rule_hidden: 64,
884            noise_dim: 8,
885            deterministic: true,
886            enable_noise: false,
887            noise_scale: 0.10,
888            noise_anneal_steps: 8192,
889            learning_rate_readout: 0.01,
890            learning_rate_selector: 1e-4,
891            learning_rate_rule: 3e-4,
892            bptt_depth: 3,
893            optimizer_momentum: 0.05,
894            grad_clip: 1.0,
895            state_clip: 8.0,
896            forget_lambda: 0.0,
897            resample_threshold: 0.5,
898            mutate_fraction: 0.1,
899            mutate_scale: 0.01,
900            mutate_model_params: false,
901            diagnostics_interval: 0,
902            min_prob: 2f64.powi(-24),
903            seed: 42,
904        }
905    }
906}
907
908impl ParticleSpec {
909    /// Validate all fields, returning an error message on failure.
910    pub fn validate(&self) -> InfotheoryResult<()> {
911        if self.num_particles == 0 {
912            return Err(InfotheoryError::invalid_backend_config(
913                "num_particles must be > 0",
914            ));
915        }
916        if self.context_window == 0 {
917            return Err(InfotheoryError::invalid_backend_config(
918                "context_window must be > 0",
919            ));
920        }
921        if self.unroll_steps == 0 {
922            return Err(InfotheoryError::invalid_backend_config(
923                "unroll_steps must be > 0",
924            ));
925        }
926        if self.num_cells == 0 {
927            return Err(InfotheoryError::invalid_backend_config(
928                "num_cells must be > 0",
929            ));
930        }
931        if self.cell_dim == 0 {
932            return Err(InfotheoryError::invalid_backend_config(
933                "cell_dim must be > 0",
934            ));
935        }
936        if self.num_rules == 0 {
937            return Err(InfotheoryError::invalid_backend_config(
938                "num_rules must be > 0",
939            ));
940        }
941        if self.selector_hidden == 0 {
942            return Err(InfotheoryError::invalid_backend_config(
943                "selector_hidden must be > 0",
944            ));
945        }
946        if self.rule_hidden == 0 {
947            return Err(InfotheoryError::invalid_backend_config(
948                "rule_hidden must be > 0",
949            ));
950        }
951        if !self.learning_rate_readout.is_finite() || self.learning_rate_readout < 0.0 {
952            return Err(InfotheoryError::invalid_backend_config(
953                "learning_rate_readout must be finite and non-negative",
954            ));
955        }
956        if !self.learning_rate_selector.is_finite() || self.learning_rate_selector < 0.0 {
957            return Err(InfotheoryError::invalid_backend_config(
958                "learning_rate_selector must be finite and non-negative",
959            ));
960        }
961        if !self.learning_rate_rule.is_finite() || self.learning_rate_rule < 0.0 {
962            return Err(InfotheoryError::invalid_backend_config(
963                "learning_rate_rule must be finite and non-negative",
964            ));
965        }
966        if !self.noise_scale.is_finite() || self.noise_scale < 0.0 {
967            return Err(InfotheoryError::invalid_backend_config(
968                "noise_scale must be finite and non-negative",
969            ));
970        }
971        if !self.optimizer_momentum.is_finite()
972            || self.optimizer_momentum < 0.0
973            || self.optimizer_momentum >= 1.0
974        {
975            return Err(InfotheoryError::invalid_backend_config(
976                "optimizer_momentum must be finite and in [0, 1)",
977            ));
978        }
979        if self.bptt_depth == 0 {
980            return Err(InfotheoryError::invalid_backend_config(
981                "bptt_depth must be > 0",
982            ));
983        }
984        if !(self.resample_threshold > 0.0 && self.resample_threshold <= 1.0) {
985            return Err(InfotheoryError::invalid_backend_config(
986                "resample_threshold must be in (0, 1]",
987            ));
988        }
989        if !(self.mutate_fraction >= 0.0 && self.mutate_fraction <= 1.0) {
990            return Err(InfotheoryError::invalid_backend_config(
991                "mutate_fraction must be in [0, 1]",
992            ));
993        }
994        if !(self.min_prob > 0.0 && self.min_prob < 0.5) {
995            return Err(InfotheoryError::invalid_backend_config(
996                "min_prob must be in (0, 0.5)",
997            ));
998        }
999        Ok(())
1000    }
1001}
1002
1003impl crate::spec::CanonicalJson for RateBackend {
1004    fn to_canonical_json_value(&self) -> crate::spec::SpecResult<serde_json::Value> {
1005        crate::spec::rate_backend_to_json_value(self)
1006    }
1007}
1008
1009impl crate::spec::CanonicalJson for CompressionBackend {
1010    fn to_canonical_json_value(&self) -> crate::spec::SpecResult<serde_json::Value> {
1011        crate::spec::compression_backend_to_json_value(self)
1012    }
1013}
1014
1015impl crate::spec::CanonicalJson for MixtureSpec {
1016    fn to_canonical_json_value(&self) -> crate::spec::SpecResult<serde_json::Value> {
1017        crate::spec::mixture_spec_to_json_value(self)
1018    }
1019}
1020
1021impl crate::spec::CanonicalJson for ParticleSpec {
1022    fn to_canonical_json_value(&self) -> crate::spec::SpecResult<serde_json::Value> {
1023        Ok(crate::spec::particle_spec_to_json_value(self))
1024    }
1025}
1026
1027impl crate::spec::CanonicalJson for CalibratedSpec {
1028    fn to_canonical_json_value(&self) -> crate::spec::SpecResult<serde_json::Value> {
1029        crate::spec::calibrated_spec_to_json_value(self)
1030    }
1031}
1032
1033#[cfg(test)]
1034mod tests {
1035    use super::*;
1036
1037    #[test]
1038    fn fading_bayes_decay_rejects_zero_and_one_boundaries() {
1039        let expert = MixtureExpertSpec {
1040            name: Some("placeholder".to_string()),
1041            log_prior: 0.0,
1042            backend: RateBackend::RosaPlus { max_order: -1 },
1043        };
1044
1045        for &decay in &[0.0, 1.0] {
1046            let spec =
1047                MixtureSpec::new(MixtureKind::FadingBayes, vec![expert.clone()]).with_decay(decay);
1048            let err = validate_mixture_spec_shallow(&spec).expect_err("boundary decay should fail");
1049            assert!(
1050                err.contains("(0, 1)"),
1051                "unexpected error for decay={decay}: {err}"
1052            );
1053        }
1054
1055        let valid = MixtureSpec::new(MixtureKind::FadingBayes, vec![expert]).with_decay(0.5);
1056        validate_mixture_spec_shallow(&valid).expect("strictly interior decay should validate");
1057    }
1058}