diff --git a/playground/kalman.html b/playground/kalman.html
new file mode 100644
index 000000000..a9d706564
--- /dev/null
+++ b/playground/kalman.html
@@ -0,0 +1,414 @@
+
+
+
+
+
+
tsb — Kalman Filter & State-Space Models
+
+
+
+
+
+
+
← Back to index
+
🔭 Kalman Filter & State-Space Models
+
+ Linear Gaussian state-space model — Kalman filter (forward pass) and
+ RTS smoother (backward pass). Mirrors statsmodels.tsa.statespace
+ and pykalman.KalmanFilter.
+
+
+
+
+
📐 The State-Space Model
+
+ A linear Gaussian SSM describes a latent state x_t and
+ observations y_t via two equations:
+
+
+ x_t = F · x_{t-1} + w_t, w_t ~ N(0, Q) (state transition)
+ y_t = H · x_t + v_t, v_t ~ N(0, R) (observation)
+ x_0 ~ N(m_0, P_0)
+
+
+ F — state transition matrix (n_states × n_states)H — observation matrix (n_obs × n_states)Q — process noise covarianceR — observation noise covariancem_0, P_0 — initial state distribution
+
+ The Kalman filter computes filtered state
+ estimates x_{t|t} (posterior after seeing observation t).
+ The RTS smoother computes smoothed estimates
+ x_{t|T} using all T observations.
+
+
+
+
+
+
📈 Local-Level Model (Random Walk + Noise)
+
+ The simplest SSM: a hidden state that follows a random walk, observed
+ with noise. Perfect for denoising a noisy scalar time series or
+ estimating a slowly changing mean.
+
+
+
+
+
Click ▶ Run to execute
+
+
+
+
+
+
🔄 RTS Smoother — Filling Gaps Retrospectively
+
+ The filter only uses observations up to time t. The smoother uses
+ all observations to produce better estimates, especially for
+ time-steps near missing values. Smoothed uncertainty is always ≤ filtered.
+
+
+
+
+
Click ▶ Run to execute
+
+
+
+
+
+
📊 Local Linear Trend (Level + Slope)
+
+ A 2-state model: [level, slope]. The level increases by the
+ slope each step; both drift over time. Great for tracking slowly changing
+ trends with missing observations.
+
+
+
+
+
Click ▶ Run to execute
+
+
+
+
+
+
⚙️ Custom State-Space Model (AR(1) State)
+
+ Build your own model by specifying the four matrices directly.
+ Here: a state that follows an AR(1) process with coefficient 0.9.
+
+
+
+
+
Click ▶ Run to execute
+
+
+
+
+
+
🔢 Multi-Dimensional Observations
+
+ The Kalman filter naturally handles multi-dimensional observations.
+ Here: 2 sensors observing a single latent state.
+
+
+
+
+
Click ▶ Run to execute
+
+
+
+
+
+
📖 API Reference
+
+ KalmanFilter.localLevel(opts?) — random-walk + noise (1-D)KalmanFilter.localLinearTrend(opts?) — level + slope (2-D state)new KalmanFilter(opts) — custom F, H, Q, R, m0, P0kf.filter(observations) → KalmanFilterResultkf.smooth(observations) → KalmanSmootherResultkalmanFilter1D(obs, opts?) — scalar convenience wrapperkalmanSmooth1D(obs, opts?) — scalar smoother wrapperextractScalarMeans(means) — extract 1-D means arrayfilteredPredictionInterval(result, z?) → {lower, upper}
+
+ Missing observations: pass null in any observation row. The
+ filter skips the update step for that time-step (covariance grows).
+ The smoother retroactively interpolates using future observations.
+
+
+
+
+
+
diff --git a/playground/orc.html b/playground/orc.html
new file mode 100644
index 000000000..8009f8c74
--- /dev/null
+++ b/playground/orc.html
@@ -0,0 +1,374 @@
+
+
+
+
+
+
tsb · ORC Format I/O
+
+
+
+
← tsb playground
+
ORC Format I/O
+
+ pandas.read_orc
+ DataFrame.to_orc
+
+
+ Apache ORC (Optimized Row Columnar) is a self-describing, type-aware columnar file format designed
+ for large-scale analytical workloads. tsb supports reading and writing ORC files with
+ NONE compression using RLE v1 integer encoding, direct float/double, and direct string encoding.
+
+
+
+ ℹ️ This playground runs entirely in-browser via a bundled tsb build. ORC buffers are created
+ in-memory — no file system access is required.
+
+
+
1 — Write & read a DataFrame
+
Create a DataFrame, serialize it to ORC bytes, then parse it back:
+
+
+
+
Click ▶ Run to execute
+
Ctrl+Enter to run · Tab to indent
+
+
+
2 — Nullable columns
+
ORC natively supports null values via PRESENT streams:
+
+
+
+
Click ▶ Run to execute
+
Ctrl+Enter to run · Tab to indent
+
+
+
3 — Column selection
+
Use the columns option to read only a subset of columns:
+
+
+
+
Click ▶ Run to execute
+
Ctrl+Enter to run · Tab to indent
+
+
+
4 — Large dataset benchmark
+
Serialize and parse a 10 000-row DataFrame to measure throughput:
+
+
+
+
Click ▶ Run to execute
+
Ctrl+Enter to run · Tab to indent
+
+
+
+
+
diff --git a/playground/signal.html b/playground/signal.html
new file mode 100644
index 000000000..832ccd7bd
--- /dev/null
+++ b/playground/signal.html
@@ -0,0 +1,621 @@
+
+
+
+
+
+
tsb — Signal Processing
+
+
+
+
+
+
+
Initializing playground…
+
+
+
← Back to roadmap
+
📡 Signal Processing — Interactive Playground
+
+ FFT, windows, STFT, Welch PSD, and periodogram — mirrors numpy.fft
+ and scipy.signal.
+ Edit any code block below and press ▶ Run
+ (or Ctrl+Enter) to execute it live in your browser.
+
+
+
+
+
1. Basic FFT of a sinusoidal signal
+
Compute a 32 Hz sine wave's FFT and identify the peak frequency bin.
+
+
+
+
+
Click ▶ Run to execute
+
Ctrl+Enter to run · Tab to indent
+
+
+
+
+
+
2. Parseval's theorem — energy preservation
+
The total energy is preserved between time and frequency domains.
+
+
+
+
+
Click ▶ Run to execute
+
Ctrl+Enter to run · Tab to indent
+
+
+
+
+
+
3. RFFT round-trip
+
Real-input FFT produces a half-spectrum; irfft reconstructs the original signal.
+
+
+
+
+
Click ▶ Run to execute
+
Ctrl+Enter to run · Tab to indent
+
+
+
+
+
+
4. Window functions
+
Named windows reduce spectral leakage. Use getWindow(name, n) to obtain any built-in window.
+
+
+
+
+
Click ▶ Run to execute
+
Ctrl+Enter to run · Tab to indent
+
+
+
+
+
+
5. Short-Time Fourier Transform (STFT)
+
Analyze a chirp signal whose frequency increases linearly over time.
+
+
+
+
+
Click ▶ Run to execute
+
Ctrl+Enter to run · Tab to indent
+
+
+
+
+
+
6. ISTFT reconstruction (round-trip)
+
Invert an STFT back to the time domain. Interior reconstruction error should be near machine epsilon.
+
+
+
+
+
Click ▶ Run to execute
+
Ctrl+Enter to run · Tab to indent
+
+
+
+
+
+
7. Welch PSD — detect signal frequency
+
Welch's method averages periodograms of overlapping segments for a lower-variance PSD estimate.
+
+
+
+
+
Click ▶ Run to execute
+
Ctrl+Enter to run · Tab to indent
+
+
+
+
+
+
8. Periodogram
+
A single-segment PSD estimate — higher variance but simpler than Welch.
+
+
+
+
+
Click ▶ Run to execute
+
Ctrl+Enter to run · Tab to indent
+
+
+
+
+
+
9. fftshift / ifftshift
+
Rearrange the FFT output so that the zero-frequency component is in the centre.
+
+
+
+
+
Click ▶ Run to execute
+
Ctrl+Enter to run · Tab to indent
+
+
+
+
+
+
API Reference
+
+ | Function | Description | Mirrors |
|---|
+ fft(x) | N-point DFT (pads to power of 2) | numpy.fft.fft |
+ ifft(X) | Inverse FFT | numpy.fft.ifft |
+ rfft(x) | Real-input FFT (one-sided) | numpy.fft.rfft |
+ irfft(X, n?) | Inverse real FFT | numpy.fft.irfft |
+ fftFreq(n, d?) | DFT sample frequencies | numpy.fft.fftfreq |
+ rfftFreq(n, d?) | One-sided DFT frequencies | numpy.fft.rfftfreq |
+ fftshift(x) | Shift DC to centre | numpy.fft.fftshift |
+ ifftshift(x) | Inverse of fftshift | numpy.fft.ifftshift |
+ getWindow(name, n) | Named window function | scipy.signal.get_window |
+ stft(x, opts?) | Short-Time Fourier Transform | scipy.signal.stft |
+ istft(Zxx, opts?) | Inverse STFT (overlap-add) | scipy.signal.istft |
+ welch(x, opts?) | Welch PSD estimate | scipy.signal.welch |
+ periodogram(x, opts?) | Periodogram PSD estimate | scipy.signal.periodogram |
+
+
+
+
+
+
+
+
+
diff --git a/src/index.ts b/src/index.ts
index bb13fcd16..e4cf54bb2 100644
--- a/src/index.ts
+++ b/src/index.ts
@@ -925,6 +925,10 @@ export { toOffset, inferFreq, FREQ_ALIASES } from "./tseries/frequencies.ts";
export { readSas } from "./io/read_sas.ts";
export type { ReadSasOptions } from "./io/read_sas.ts";
+// io.orc — Apache ORC file format read/write
+export { readOrc, toOrc } from "./io/orc.ts";
+export type { ReadOrcOptions, ToOrcOptions } from "./io/orc.ts";
+
// pd.arrays.SparseArray / pd.SparseDtype — sparse storage for arrays
// with many repeated (fill) values
export { SparseArray, SparseDtype } from "./core/sparse.ts";
@@ -1008,3 +1012,411 @@ export {
tsallisEntropy,
} from "./stats/information.ts";
export type { PMF, NMIMethod } from "./stats/information.ts";
+
+// signal processing — FFT, windows, STFT, Welch PSD, periodogram
+export {
+ complex,
+ cAbs,
+ cArg,
+ fft,
+ ifft,
+ rfft,
+ irfft,
+ fftFreq,
+ rfftFreq,
+ fftshift,
+ ifftshift,
+ rectangularWindow,
+ bartlettWindow,
+ hannWindow,
+ hammingWindow,
+ blackmanWindow,
+ blackmanHarrisWindow,
+ flatTopWindow,
+ kaiserWindow,
+ getWindow,
+ stft,
+ istft,
+ welch,
+ periodogram,
+} from "./stats/signal.ts";
+export type {
+ Complex,
+ WindowName,
+ STFTOptions,
+ STFTResult,
+ ISTFTOptions,
+ WelchOptions,
+ PSDResult,
+ PeriodogramOptions,
+} from "./stats/signal.ts";
+
+// digital filters — FIR/IIR design and application (mirrors scipy.signal)
+export {
+ firwin,
+ freqz,
+ sosfreqz,
+ lfilter,
+ filtfilt,
+ sosfilt,
+ sosfiltfilt,
+ butter,
+} from "./stats/filters.ts";
+export type {
+ FirwinOptions,
+ FreqzResult,
+ SOSSection,
+ ButterResult,
+ FilterType,
+} from "./stats/filters.ts";
+
+// ACF/PACF — autocorrelation, partial autocorrelation, portmanteau tests
+export {
+ autocorr,
+ acf,
+ pacf,
+ ccf,
+ durbinWatson,
+ ljungBox,
+ boxPierce,
+} from "./stats/acf_pacf.ts";
+export type {
+ ACFResult,
+ PACFResult,
+ PortmanteauResult,
+ ACFOptions,
+ PACFOptions,
+ CCFOptions,
+ PortmanteauOptions,
+} from "./stats/acf_pacf.ts";
+
+// ARIMA — ARIMA(p,d,q) time-series model (Hannan-Rissanen, forecast CIs)
+export { ARIMAModel, fitArima } from "./stats/arima.ts";
+export type {
+ ARIMAOptions,
+ ARIMAFitResult,
+ ARIMAForecastResult,
+} from "./stats/arima.ts";
+
+// read_avro / toAvro — Apache Avro OCF I/O for DataFrame
+export { readAvro, toAvro } from "./io/read_avro.ts";
+export type {
+ ReadAvroOptions,
+ ToAvroOptions,
+} from "./io/read_avro.ts";
+
+// Kalman filter & RTS smoother — linear Gaussian state-space model
+export {
+ KalmanFilter,
+ StateSpaceModel,
+ kalmanFilter1D,
+ kalmanSmooth1D,
+ extractScalarMeans,
+ extractScalarVariances,
+ filteredPredictionInterval,
+} from "./stats/kalman.ts";
+export type {
+ KalmanFilterOptions,
+ LocalLevelOptions,
+ LocalLinearTrendOptions,
+ KalmanFilterResult,
+ KalmanSmootherResult,
+} from "./stats/kalman.ts";
+
+// ETS — Exponential Smoothing / Holt-Winters (Simple, Holt, full Holt-Winters)
+export {
+ SimpleExpSmoothing,
+ Holt,
+ ExponentialSmoothing,
+ simpleExpSmoothing,
+ holt,
+ fitEts,
+} from "./stats/ets.ts";
+export type {
+ ETSTrend,
+ ETSSeasonal,
+ ETSInit,
+ SESOptions,
+ SESFitResult,
+ HoltOptions,
+ HoltFitResult,
+ ExponentialSmoothingOptions,
+ ExponentialSmoothingFitResult,
+ ETSForecastResult,
+} from "./stats/ets.ts";
+
+export {
+ DLM,
+ buildLocalLevel,
+ buildLocalLinearTrend,
+ buildPolynomial,
+ buildFourier,
+ buildRegression,
+ combineDLMs,
+} from "./stats/dlm.ts";
+export type {
+ DLMSpec,
+ DLMOptions,
+ DLMFilterStep,
+ DLMResult,
+ DLMSmootherResult,
+ DLMForecastResult,
+} from "./stats/dlm.ts";
+export {
+ GaussianHMM,
+ MultinomialHMM,
+ fitGaussianHMM,
+ hmmViterbi,
+} from "./stats/hmm.ts";
+export type {
+ GaussianHMMParams,
+ GaussianHMMFit,
+ MultinomialHMMParams,
+ MultinomialHMMFit,
+} from "./stats/hmm.ts";
+export {
+ simulateBrownianMotion,
+ simulateGeometricBrownianMotion,
+ simulateOrnsteinUhlenbeck,
+ fitOrnsteinUhlenbeck,
+ simulatePoissonProcess,
+ poissonCounts,
+ simulateRandomWalk,
+ simulateMarkovChain,
+ stationaryDistribution,
+} from "./stats/stochastic_processes.ts";
+export type {
+ BrownianMotionParams,
+ OUParams,
+ RandomWalkParams,
+ ProcessPath,
+} from "./stats/stochastic_processes.ts";
+export {
+ createGraph,
+ addEdge,
+ graphFromEdges,
+ degreeCentrality,
+ directedDegrees,
+ bfsDistances,
+ dijkstra,
+ betweennessCentrality,
+ clusteringCoefficient,
+ globalClusteringCoefficient,
+ connectedComponents,
+ pageRank,
+ hits,
+} from "./stats/network_stats.ts";
+export type { Graph } from "./stats/network_stats.ts";
+export {
+ euclideanDistance,
+ pairwiseDistances,
+ empiricalVariogram,
+ sphericalVariogram,
+ exponentialVariogram,
+ gaussianVariogram,
+ ordinaryKriging,
+ moransI,
+ ripleysK,
+ ripleysL,
+ kde2d,
+ distanceWeights,
+} from "./stats/spatial_stats.ts";
+export type {
+ SpatialPoint,
+ VariogramPoint,
+ VariogramModelParams,
+} from "./stats/spatial_stats.ts";
+export {
+ normalCdf,
+ normalQuantile,
+ gaussianCopulaCdf,
+ gaussianCopulaDensity,
+ sampleGaussianCopula,
+ claytonCopulaCdf,
+ sampleClaytonCopula,
+ gumbelCopulaCdf,
+ frankCopulaCdf,
+ sampleFrankCopula,
+ empiricalCopula,
+ kendallTau,
+ tauToGaussianRho,
+ tauToClaytonTheta,
+ tauToGumbelTheta,
+} from "./stats/copulas.ts";
+export {
+ gevPdf,
+ gevCdf,
+ gevQuantile,
+ gevReturnLevel,
+ fitGEV,
+ gpdPdf,
+ gpdCdf,
+ gpdQuantile,
+ fitGPD,
+ extractExceedances,
+ gpdReturnLevel,
+ gumbelPdf,
+ gumbelCdf,
+ gumbelQuantile,
+} from "./stats/extreme_value.ts";
+export type {
+ GEVParams,
+ GPDParams,
+ GumbelParams,
+} from "./stats/extreme_value.ts";
+
+// ─── ML ───────────────────────────────────────────────────────────────────────
+export {
+ computeNoiseSchedule,
+ addNoise,
+ ddimStep,
+ ddimTimesteps,
+ snrAtTimestep,
+} from "./ml/ddim.ts";
+export type { NoiseSchedule, DDIMConfig, NoiseScheduleValues } from "./ml/ddim.ts";
+
+export {
+ sparsemax,
+ glu,
+ batchNormInfer,
+ featureTransformer,
+ attentiveTransformer,
+ tabnetStep,
+ aggregateSteps,
+ featureImportance,
+} from "./ml/tabnet.ts";
+export type {
+ BatchNormParams,
+ FeatureTransformerWeights,
+ AttentiveTransformerWeights,
+ TabNetStepResult,
+} from "./ml/tabnet.ts";
+
+export {
+ scaledDotProductAttention,
+ layerNorm,
+ gelu,
+ feedForward,
+ tokenizeNumerical,
+ tokenizeCategorical,
+ addResidual,
+ extractCLSToken,
+ linearHead,
+ softmax,
+} from "./ml/ft_transformer.ts";
+export type { AttentionOutput } from "./ml/ft_transformer.ts";
+
+export {
+ viterbiDecode,
+ forwardLogZ,
+ sequenceScore,
+ crfNegLogLikelihood,
+ logSumExp,
+} from "./ml/crf.ts";
+export type { CRFParams, ViterbiResult } from "./ml/crf.ts";
+
+export {
+ rnnCell,
+ encode,
+ bahdanauAttention,
+ decoderStep,
+ greedyDecode,
+ sigmoid,
+ tanh,
+} from "./ml/seq2seq.ts";
+export type {
+ RNNCellWeights,
+ BahdanauWeights,
+ DecoderWeights,
+ DecoderStepResult,
+} from "./ml/seq2seq.ts";
+
+export {
+ rbfKernel,
+ maternKernel52,
+ kernelMatrix,
+ cholesky,
+ gpPredict,
+ normalCDF,
+ normalPDF,
+ expectedImprovement,
+ probabilityOfImprovement,
+ upperConfidenceBound,
+ initBOState,
+ suggestNext,
+} from "./ml/bayesian_opt.ts";
+export type { BOState } from "./ml/bayesian_opt.ts";
+
+export {
+ karrasNoiseLevels,
+ cSkip,
+ cOut,
+ cIn,
+ cNoise,
+ consistencyFunction,
+ preconditionInput,
+ pseudoHuberLoss,
+ consistencyTrainingLoss,
+ addGaussianNoise,
+ consistencySampleOneStep,
+ consistencySampleMultiStep,
+ updateEMAWeights,
+ adaptiveEMADecay,
+ DEFAULT_CONSISTENCY_CONFIG,
+} from "./ml/consistency_models.ts";
+export type { ConsistencyConfig } from "./ml/consistency_models.ts";
+
+export {
+ buildTree,
+ treePredict,
+ treePredictOne,
+ fitGBM,
+ predictGBM,
+ mse,
+ r2Score,
+ mseGradient,
+} from "./ml/gradient_boosting.ts";
+export type { TreeNode, TreeParams, GBMEnsemble } from "./ml/gradient_boosting.ts";
+
+export {
+ computeKernel,
+ svmDecision,
+ svmPredict,
+ fitSVM,
+ svmAccuracy,
+} from "./ml/svm.ts";
+export type { KernelType, SVMKernelConfig, SVMModel } from "./ml/svm.ts";
+
+export {
+ multiHeadAttention,
+ sinusoidalPositionalEncoding,
+ addPositionalEncoding,
+ causalMask,
+ applyRoPE,
+} from "./ml/attention.ts";
+export type { MHAConfig, MHAWeights } from "./ml/attention.ts";
+
+export {
+ fitRandomForest,
+ predictRandomForest,
+ oobError,
+ LCGRandom,
+} from "./ml/random_forest.ts";
+export type { RandomForestConfig, RandomForestModel } from "./ml/random_forest.ts";
+
+export {
+ denseForward,
+ denseBackward,
+ relu,
+ reluGrad,
+ sigmoidActivation,
+ sigmoidGrad,
+ softmaxCrossEntropy,
+ mseLoss,
+ initAdam,
+ adamStep,
+ heInit,
+ dropoutMask,
+ applyDropout,
+} from "./ml/neural_network.ts";
+export type { AdamState } from "./ml/neural_network.ts";
diff --git a/src/io/feather.ts b/src/io/feather.ts
index 1abb681b3..174734b80 100644
--- a/src/io/feather.ts
+++ b/src/io/feather.ts
@@ -726,7 +726,7 @@ function encodeStrings(values: readonly (Scalar | null)[]): {
for (let i = 0; i < encoded.length; i++) {
ov.setInt32(i * 4, pos, true);
data.set(encoded[i]!, pos);
- pos += encoded[i]!.length;
+ pos += encoded[i]?.length ?? 0;
}
ov.setInt32(values.length * 4, pos, true);
return { offsets, data };
@@ -918,7 +918,7 @@ export function toFeather(df: DataFrame, options: ToFeatherOptions = {}): Uint8A
cols.push({ name, type, values });
}
- const numRows = cols.length > 0 ? cols[0]!.values.length : df.index.size;
+ const numRows = cols[0]?.values.length ?? df.index.size;
const schemaCols = cols.map((c) => ({ name: c.name, type: c.type }));
// Encode all column buffers into a single body array
diff --git a/src/io/index.ts b/src/io/index.ts
index 194e405da..78cf80dc9 100644
--- a/src/io/index.ts
+++ b/src/io/index.ts
@@ -62,3 +62,6 @@ export type { ToExcelOptions } from "./to_excel.ts";
export { readSas } from "./read_sas.ts";
export type { ReadSasOptions } from "./read_sas.ts";
+
+export { readOrc, toOrc } from "./orc.ts";
+export type { ReadOrcOptions, ToOrcOptions } from "./orc.ts";
diff --git a/src/io/orc.ts b/src/io/orc.ts
new file mode 100644
index 000000000..8c5ee6bfe
--- /dev/null
+++ b/src/io/orc.ts
@@ -0,0 +1,1358 @@
+/**
+ * readOrc / toOrc — Apache ORC (Optimized Row Columnar) file format I/O.
+ *
+ * Mirrors `pandas.read_orc()` and `DataFrame.to_orc()`.
+ *
+ * Supported column types (read & write):
+ * - BOOLEAN, INT, LONG, FLOAT, DOUBLE, STRING, DATE
+ *
+ * Compression: NONE (ZLIB/Snappy require an external decompressor).
+ * Encoding: DIRECT (integers via RLE v1, strings via raw bytes + lengths,
+ * floats/doubles via raw IEEE 754, booleans via RLE byte v1).
+ *
+ * @module
+ */
+
+import { DataFrame } from "../core/frame.ts";
+import type { Label, Scalar } from "../types.ts";
+
+// ─── Public types ─────────────────────────────────────────────────────────────
+
+/** Options for {@link readOrc}. */
+export interface ReadOrcOptions {
+ /**
+ * Column name to use as the row index.
+ * Default: `null` (RangeIndex).
+ */
+ readonly indexCol?: string | null;
+ /**
+ * Subset of columns to read. `null` = all columns.
+ * Default: `null`.
+ */
+ readonly columns?: readonly string[] | null;
+}
+
+/** Options for {@link toOrc}. */
+export interface ToOrcOptions {
+ /**
+ * Write the DataFrame's row index as an extra column.
+ * Default: `false`.
+ */
+ readonly writeIndex?: boolean;
+}
+
+// ─── ORC file constants ───────────────────────────────────────────────────────
+
+// File header magic
+const ORC_MAGIC = new Uint8Array([0x4f, 0x52, 0x43]); // "ORC"
+
+// ORC type kinds
+const KIND_BOOLEAN = 0;
+const _KIND_BYTE = 1;
+const _KIND_SHORT = 2;
+const _KIND_INT = 3;
+const KIND_LONG = 4;
+const KIND_FLOAT = 5;
+const KIND_DOUBLE = 6;
+const KIND_STRING = 7;
+const KIND_STRUCT = 12;
+const KIND_DATE = 15;
+
+// Compression codecs
+const COMP_NONE = 0;
+const COMP_ZLIB = 1;
+
+// Stream kinds
+const STREAM_PRESENT = 0;
+const STREAM_DATA = 1;
+const STREAM_LENGTH = 2;
+const _STREAM_DICTIONARY_DATA = 3;
+
+// Column encoding kinds
+const ENC_DIRECT = 0;
+
+// ─── Protobuf utilities ───────────────────────────────────────────────────────
+
+/** A single decoded protobuf field value. */
+type PbVal =
+ | { readonly wt: 0; readonly v: bigint }
+ | { readonly wt: 2; readonly v: Uint8Array }
+ | { readonly wt: 1; readonly v: bigint }
+ | { readonly wt: 5; readonly v: number };
+
+/** A decoded protobuf message: field number → list of values. */
+type PbMsg = Map
;
+
+/** Read a protobuf varint (unsigned, LSB-first). */
+function pbReadVarU(buf: Uint8Array, pos: number): [bigint, number] {
+ let result = 0n;
+ let shift = 0n;
+ let cur = pos;
+ for (;;) {
+ const b = buf[cur];
+ if (b === undefined) {
+ throw new Error("ORC: truncated varint");
+ }
+ cur++;
+ result |= BigInt(b & 0x7f) << shift;
+ if ((b & 0x80) === 0) {
+ break;
+ }
+ shift += 7n;
+ }
+ return [result, cur];
+}
+
+/** Decode a protobuf message from a byte slice. */
+function pbDecode(buf: Uint8Array): PbMsg {
+ const msg: PbMsg = new Map();
+ let pos = 0;
+ while (pos < buf.length) {
+ let tag: bigint;
+ [tag, pos] = pbReadVarU(buf, pos);
+ const fieldNum = Number(tag >> 3n);
+ const wt = Number(tag & 7n);
+ let val: PbVal;
+ if (wt === 0) {
+ let v: bigint;
+ [v, pos] = pbReadVarU(buf, pos);
+ val = { wt: 0, v };
+ } else if (wt === 2) {
+ let len: bigint;
+ [len, pos] = pbReadVarU(buf, pos);
+ const n = Number(len);
+ val = { wt: 2, v: buf.subarray(pos, pos + n) };
+ pos += n;
+ } else if (wt === 1) {
+ const dv = new DataView(buf.buffer, buf.byteOffset + pos, 8);
+ const lo = BigInt(dv.getUint32(0, true));
+ const hi = BigInt(dv.getUint32(4, true));
+ val = { wt: 1, v: (hi << 32n) | lo };
+ pos += 8;
+ } else if (wt === 5) {
+ const dv = new DataView(buf.buffer, buf.byteOffset + pos, 4);
+ val = { wt: 5, v: dv.getUint32(0, true) };
+ pos += 4;
+ } else {
+ throw new Error(`ORC: unknown wire type ${wt}`);
+ }
+ const list = msg.get(fieldNum);
+ if (list !== undefined) {
+ list.push(val);
+ } else {
+ msg.set(fieldNum, [val]);
+ }
+ }
+ return msg;
+}
+
+/** Get a uint64 field as bigint (default 0). */
+function pbU64(msg: PbMsg, field: number): bigint {
+ const f = msg.get(field)?.[0];
+ return f?.wt === 0 ? f.v : 0n;
+}
+
+/** Get a uint32 field as number (default 0). */
+function pbU32(msg: PbMsg, field: number): number {
+ return Number(pbU64(msg, field));
+}
+
+/** Get all uint32 repeated field values. */
+function pbU32s(msg: PbMsg, field: number): number[] {
+ return (msg.get(field) ?? [])
+ .filter((f): f is PbVal & { wt: 0 } => f.wt === 0)
+ .map((f) => Number(f.v));
+}
+
+/** Get all string repeated field values. */
+function pbStrings(msg: PbMsg, field: number): string[] {
+ const dec = new TextDecoder();
+ return (msg.get(field) ?? [])
+ .filter((f): f is PbVal & { wt: 2 } => f.wt === 2)
+ .map((f) => dec.decode(f.v));
+}
+
+/** Get all embedded message repeated field values. */
+function pbMsgs(msg: PbMsg, field: number): PbMsg[] {
+ return (msg.get(field) ?? [])
+ .filter((f): f is PbVal & { wt: 2 } => f.wt === 2)
+ .map((f) => pbDecode(f.v));
+}
+
+// ─── Protobuf writer ──────────────────────────────────────────────────────────
+
+function pbWvU(v: bigint, out: number[]): void {
+ let val = v;
+ while (val >= 128n) {
+ out.push(Number(val & 0x7fn) | 0x80);
+ val >>= 7n;
+ }
+ out.push(Number(val));
+}
+
+function pbTag(fn: number, wt: 0 | 2, out: number[]): void {
+ pbWvU(BigInt((fn << 3) | wt), out);
+}
+
+function pbWU64(fn: number, v: bigint, out: number[]): void {
+ if (v === 0n) {
+ return;
+ }
+ pbTag(fn, 0, out);
+ pbWvU(v, out);
+}
+
+function pbWU32(fn: number, v: number, out: number[]): void {
+ pbWU64(fn, BigInt(v), out);
+}
+
+function pbWBytes(fn: number, v: Uint8Array, out: number[]): void {
+ pbTag(fn, 2, out);
+ pbWvU(BigInt(v.length), out);
+ for (const b of v) {
+ out.push(b);
+ }
+}
+
+function pbWMsg(fn: number, msg: number[], out: number[]): void {
+ pbTag(fn, 2, out);
+ pbWvU(BigInt(msg.length), out);
+ for (const b of msg) {
+ out.push(b);
+ }
+}
+
+// ─── Hadoop VInt ──────────────────────────────────────────────────────────────
+// ORC uses big-endian variable-length signed integers for RLE integer streams.
+
+/**
+ * Read a Hadoop-style variable-length signed integer.
+ *
+ * Byte ranges:
+ * - 0x00–0x7F: single-byte positive (0–127)
+ * - 0x88–0x8F: positive multi-byte (1–8 data bytes follow)
+ * - 0x80–0x87: negative multi-byte (1–8 data bytes follow, XOR with -1)
+ * - 0x90–0xFF: single-byte negative (-112 to -1)
+ */
+function hvReadVInt(buf: Uint8Array, pos: number): [bigint, number] {
+ const fb = buf[pos];
+ if (fb === undefined) {
+ throw new Error("ORC: truncated Hadoop VInt");
+ }
+ let cur = pos + 1;
+ // Interpret as signed byte
+ const sfb = fb >= 0x80 ? fb - 0x100 : fb;
+ // Single-byte range: -112 to 127
+ if (sfb >= -112) {
+ return [BigInt(sfb), cur];
+ }
+ // Multi-byte
+ const isNeg = sfb < -120; // unsigned 128–135 = negative; 136–143 = positive
+ const len = isNeg ? -119 - sfb : -111 - sfb; // total bytes incl. header
+ let value = 0n;
+ for (let i = 1; i < len; i++) {
+ const b = buf[cur];
+ if (b === undefined) {
+ throw new Error("ORC: truncated Hadoop VInt data");
+ }
+ cur++;
+ value = (value << 8n) | BigInt(b);
+ }
+ if (isNeg) {
+ value ^= -1n;
+ }
+ return [value, cur];
+}
+
+/** Write a Hadoop-style variable-length signed integer. */
+function hvWriteVInt(value: bigint, out: number[]): void {
+ if (value >= -112n && value <= 127n) {
+ out.push(Number(value < 0n ? value + 256n : value));
+ return;
+ }
+ let uval = value;
+ const isNeg = value < 0n;
+ if (isNeg) {
+ uval = value ^ -1n;
+ }
+ let nbytes = 0;
+ let tmp = uval;
+ while (tmp > 0n) {
+ tmp >>= 8n;
+ nbytes++;
+ }
+ const header = isNeg ? -120 - nbytes : -112 - nbytes;
+ out.push(header < 0 ? header + 0x100 : header);
+ for (let i = nbytes - 1; i >= 0; i--) {
+ out.push(Number((uval >> BigInt(i * 8)) & 0xffn));
+ }
+}
+
+// ─── RLE byte v1 ─────────────────────────────────────────────────────────────
+
+/**
+ * Decode an ORC RLE byte v1 stream to a flat byte array.
+ * Control byte < 128: run of (ctrl + 3) copies of the next byte.
+ * Control byte >= 128: (256 - ctrl) literal bytes follow.
+ */
+function rleByteDecodeV1(buf: Uint8Array, off: number, len: number): Uint8Array {
+ const end = off + len;
+ const out: number[] = [];
+ let pos = off;
+ while (pos < end) {
+ const ctrl = buf[pos];
+ if (ctrl === undefined) {
+ break;
+ }
+ pos++;
+ if (ctrl < 128) {
+ const count = ctrl + 3;
+ const val = buf[pos];
+ if (val === undefined) {
+ break;
+ }
+ pos++;
+ for (let i = 0; i < count; i++) {
+ out.push(val);
+ }
+ } else {
+ const count = 256 - ctrl;
+ for (let i = 0; i < count; i++) {
+ const b = buf[pos];
+ if (b === undefined) {
+ break;
+ }
+ pos++;
+ out.push(b);
+ }
+ }
+ }
+ return new Uint8Array(out);
+}
+
+/** Encode bytes using RLE byte v1. */
+function rleByteEncodeV1(data: readonly number[]): Uint8Array {
+ if (data.length === 0) {
+ return new Uint8Array(0);
+ }
+ const out: number[] = [];
+ let i = 0;
+ while (i < data.length) {
+ // Look for a run (same value repeated)
+ let runLen = 1;
+ while (runLen < 130 && i + runLen < data.length && data[i + runLen] === data[i]) {
+ runLen++;
+ }
+ if (runLen >= 3) {
+ out.push(runLen - 3);
+ const d = data[i];
+ if (d === undefined) {
+ throw new Error("ORC: undefined byte in run");
+ }
+ out.push(d);
+ i += runLen;
+ } else {
+ // Literal group
+ let litLen = 1;
+ while (litLen < 128 && i + litLen < data.length) {
+ // Stop if next 3 values are identical (start a new run)
+ const base = data[i + litLen];
+ let rcheck = 1;
+ while (
+ rcheck < 3 &&
+ i + litLen + rcheck < data.length &&
+ data[i + litLen + rcheck] === base
+ ) {
+ rcheck++;
+ }
+ if (rcheck >= 3) {
+ break;
+ }
+ litLen++;
+ }
+ out.push(256 - litLen);
+ for (let j = 0; j < litLen; j++) {
+ const d = data[i + j];
+ if (d === undefined) {
+ throw new Error("ORC: undefined byte in literal");
+ }
+ out.push(d);
+ }
+ i += litLen;
+ }
+ }
+ return new Uint8Array(out);
+}
+
+// ─── RLE integer v1 ──────────────────────────────────────────────────────────
+
+/**
+ * Decode an ORC RLE integer v1 stream (Hadoop VInts, big-endian).
+ * Control byte >= 0: run of (ctrl + 3) values, next byte is signed delta, then base VInt.
+ * Control byte < 0: (-ctrl) literal VInts.
+ */
+function rleIntDecodeV1(buf: Uint8Array, off: number, len: number): bigint[] {
+ const end = off + len;
+ const result: bigint[] = [];
+ let pos = off;
+ while (pos < end) {
+ const ctrl = buf[pos];
+ if (ctrl === undefined) {
+ break;
+ }
+ pos++;
+ const sctrl = ctrl >= 0x80 ? ctrl - 0x100 : ctrl; // signed
+ if (sctrl >= 0) {
+ const count = sctrl + 3;
+ const deltaByte = buf[pos];
+ if (deltaByte === undefined) {
+ break;
+ }
+ pos++;
+ const delta = BigInt(deltaByte >= 0x80 ? deltaByte - 0x100 : deltaByte);
+ let base: bigint;
+ [base, pos] = hvReadVInt(buf, pos);
+ for (let i = 0; i < count; i++) {
+ result.push(base + delta * BigInt(i));
+ }
+ } else {
+ const count = -sctrl;
+ for (let i = 0; i < count; i++) {
+ let v: bigint;
+ [v, pos] = hvReadVInt(buf, pos);
+ result.push(v);
+ }
+ }
+ }
+ return result;
+}
+
+/** Encode bigint values using RLE integer v1. */
+function rleIntEncodeV1(values: readonly bigint[]): Uint8Array {
+ if (values.length === 0) {
+ return new Uint8Array(0);
+ }
+ const out: number[] = [];
+ let i = 0;
+ while (i < values.length) {
+ const v0 = values[i];
+ if (v0 === undefined) {
+ break;
+ }
+ // Attempt to find a run with a constant delta
+ if (i + 2 < values.length) {
+ const v1 = values[i + 1];
+ const v2 = values[i + 2];
+ if (v1 !== undefined && v2 !== undefined) {
+ const delta = v1 - v0;
+ if (v2 - v1 === delta && delta >= -128n && delta <= 127n) {
+ let runLen = 3;
+ while (runLen < 130 && i + runLen < values.length) {
+ const vn = values[i + runLen];
+ const vprev = values[i + runLen - 1];
+ if (vn === undefined || vprev === undefined || vn - vprev !== delta) {
+ break;
+ }
+ runLen++;
+ }
+ out.push(runLen - 3);
+ out.push(Number(delta < 0n ? delta + 256n : delta));
+ hvWriteVInt(v0, out);
+ i += runLen;
+ continue;
+ }
+ }
+ }
+ // Literal group
+ let litLen = 1;
+ while (litLen < 128 && i + litLen < values.length) {
+ const va = values[i + litLen];
+ const vb = values[i + litLen + 1];
+ const vc = values[i + litLen + 2];
+ if (va !== undefined && vb !== undefined && vc !== undefined) {
+ const d1 = vb - va;
+ const d2 = vc - vb;
+ if (d1 === d2 && d1 >= -128n && d1 <= 127n) {
+ break;
+ }
+ }
+ litLen++;
+ }
+ out.push(256 - litLen);
+ for (let j = 0; j < litLen; j++) {
+ const v = values[i + j];
+ if (v === undefined) {
+ break;
+ }
+ hvWriteVInt(v, out);
+ }
+ i += litLen;
+ }
+ return new Uint8Array(out);
+}
+
+// ─── PRESENT stream helpers ───────────────────────────────────────────────────
+
+/**
+ * Expand a PRESENT-stream byte array into per-row boolean flags.
+ * Each byte = 8 rows, MSB first. 1 = non-null, 0 = null.
+ */
+function expandPresent(raw: Uint8Array, nRows: number): boolean[] {
+ const flags: boolean[] = [];
+ for (let i = 0; i < nRows; i++) {
+ flags.push(false);
+ }
+ let row = 0;
+ for (const byte of raw) {
+ for (let bit = 7; bit >= 0 && row < nRows; bit--, row++) {
+ flags[row] = ((byte >> bit) & 1) === 1;
+ }
+ }
+ return flags;
+}
+
+/**
+ * Pack per-row null flags into PRESENT-stream bytes.
+ * 1 = non-null, 0 = null. Returns null if all rows are non-null.
+ */
+function packPresent(nonNull: boolean[]): Uint8Array | null {
+ if (nonNull.every((v) => v)) {
+ return null;
+ }
+ const bytes: number[] = [];
+ for (let i = 0; i < nonNull.length; i += 8) {
+ let byte = 0;
+ for (let bit = 0; bit < 8 && i + bit < nonNull.length; bit++) {
+ if (nonNull[i + bit]) {
+ byte |= 1 << (7 - bit);
+ }
+ }
+ bytes.push(byte);
+ }
+ return new Uint8Array(bytes);
+}
+
+// ─── ORC metadata structures (decoded) ────────────────────────────────────────
+
+interface OrcPostscript {
+ footerLength: number;
+ compression: number;
+ compressionBlockSize: number;
+ metadataLength: number;
+ writerVersion: number;
+}
+
+interface OrcStripeInfo {
+ offset: number;
+ indexLength: number;
+ dataLength: number;
+ footerLength: number;
+ numberOfRows: number;
+}
+
+interface OrcType {
+ kind: number;
+ subtypes: number[];
+ fieldNames: string[];
+}
+
+interface OrcFooter {
+ stripes: OrcStripeInfo[];
+ types: OrcType[];
+ numberOfRows: number;
+}
+
+interface OrcStream {
+ kind: number;
+ column: number;
+ length: number;
+}
+
+interface OrcColumnEncoding {
+ kind: number;
+ dictionarySize: number;
+}
+
+interface OrcStripeFooter {
+ streams: OrcStream[];
+ columns: OrcColumnEncoding[];
+}
+
+// ─── ORC decoding helpers ─────────────────────────────────────────────────────
+
+function decodePostscript(buf: Uint8Array): OrcPostscript {
+ const msg = pbDecode(buf);
+ return {
+ footerLength: Number(pbU64(msg, 1)),
+ compression: pbU32(msg, 2),
+ compressionBlockSize: Number(pbU64(msg, 3)) || 262144,
+ metadataLength: Number(pbU64(msg, 5)),
+ writerVersion: pbU32(msg, 6),
+ };
+}
+
+function decodeFooter(buf: Uint8Array): OrcFooter {
+ const msg = pbDecode(buf);
+ const stripesMsgs = pbMsgs(msg, 3);
+ const stripes: OrcStripeInfo[] = stripesMsgs.map((sm) => ({
+ offset: Number(pbU64(sm, 1)),
+ indexLength: Number(pbU64(sm, 2)),
+ dataLength: Number(pbU64(sm, 3)),
+ footerLength: Number(pbU64(sm, 4)),
+ numberOfRows: Number(pbU64(sm, 5)),
+ }));
+ const typesMsgs = pbMsgs(msg, 4);
+ const types: OrcType[] = typesMsgs.map((tm) => ({
+ kind: pbU32(tm, 1),
+ subtypes: pbU32s(tm, 2),
+ fieldNames: pbStrings(tm, 3),
+ }));
+ return {
+ stripes,
+ types,
+ numberOfRows: Number(pbU64(msg, 6)),
+ };
+}
+
+function decodeStripeFooter(buf: Uint8Array): OrcStripeFooter {
+ const msg = pbDecode(buf);
+ const streamMsgs = pbMsgs(msg, 1);
+ const streams: OrcStream[] = streamMsgs.map((sm) => ({
+ kind: pbU32(sm, 1),
+ column: pbU32(sm, 2),
+ length: Number(pbU64(sm, 3)),
+ }));
+ const colMsgs = pbMsgs(msg, 2);
+ const columns: OrcColumnEncoding[] = colMsgs.map((cm) => ({
+ kind: pbU32(cm, 1),
+ dictionarySize: pbU32(cm, 2),
+ }));
+ return { streams, columns };
+}
+
+// ─── Column data decoders ─────────────────────────────────────────────────────
+
+function decodeIntCol(
+ dataBuf: Uint8Array,
+ off: number,
+ len: number,
+ presentFlags: boolean[] | null,
+ nRows: number,
+): (bigint | null)[] {
+ const ints = rleIntDecodeV1(dataBuf, off, len);
+ const result: (bigint | null)[] = [];
+ let dataIdx = 0;
+ for (let i = 0; i < nRows; i++) {
+ if (presentFlags !== null && presentFlags[i] === false) {
+ result.push(null);
+ } else {
+ result.push(ints[dataIdx] ?? null);
+ dataIdx++;
+ }
+ }
+ return result;
+}
+
+function decodeF32Col(
+ dataBuf: Uint8Array,
+ off: number,
+ presentFlags: boolean[] | null,
+ nRows: number,
+): (number | null)[] {
+ const dv = new DataView(dataBuf.buffer, dataBuf.byteOffset);
+ const result: (number | null)[] = [];
+ let pos = off;
+ for (let i = 0; i < nRows; i++) {
+ if (presentFlags !== null && presentFlags[i] === false) {
+ result.push(null);
+ } else {
+ result.push(dv.getFloat32(pos, true));
+ pos += 4;
+ }
+ }
+ return result;
+}
+
+function decodeF64Col(
+ dataBuf: Uint8Array,
+ off: number,
+ presentFlags: boolean[] | null,
+ nRows: number,
+): (number | null)[] {
+ const dv = new DataView(dataBuf.buffer, dataBuf.byteOffset);
+ const result: (number | null)[] = [];
+ let pos = off;
+ for (let i = 0; i < nRows; i++) {
+ if (presentFlags !== null && presentFlags[i] === false) {
+ result.push(null);
+ } else {
+ result.push(dv.getFloat64(pos, true));
+ pos += 8;
+ }
+ }
+ return result;
+}
+
+function decodeStringCol(
+ dataBuf: Uint8Array,
+ dataOff: number,
+ lenBuf: Uint8Array,
+ lenOff: number,
+ lenLen: number,
+ presentFlags: boolean[] | null,
+ nRows: number,
+): (string | null)[] {
+ const lengths = rleIntDecodeV1(lenBuf, lenOff, lenLen);
+ const dec = new TextDecoder();
+ const result: (string | null)[] = [];
+ let dataPos = dataOff;
+ let strIdx = 0;
+ for (let i = 0; i < nRows; i++) {
+ if (presentFlags !== null && presentFlags[i] === false) {
+ result.push(null);
+ } else {
+ const slen = Number(lengths[strIdx] ?? 0n);
+ result.push(dec.decode(dataBuf.subarray(dataPos, dataPos + slen)));
+ dataPos += slen;
+ strIdx++;
+ }
+ }
+ return result;
+}
+
+function decodeBoolCol(
+ dataBuf: Uint8Array,
+ off: number,
+ len: number,
+ presentFlags: boolean[] | null,
+ nRows: number,
+): (boolean | null)[] {
+ const raw = rleByteDecodeV1(dataBuf, off, len);
+ const result: (boolean | null)[] = [];
+ let bitIdx = 0;
+ for (let i = 0; i < nRows; i++) {
+ if (presentFlags !== null && presentFlags[i] === false) {
+ result.push(null);
+ } else {
+ const bytePos = Math.floor(bitIdx / 8);
+ const bitPos = 7 - (bitIdx % 8);
+ const byte = raw[bytePos] ?? 0;
+ result.push(((byte >> bitPos) & 1) === 1);
+ bitIdx++;
+ }
+ }
+ return result;
+}
+
+// ─── Column data encoders ─────────────────────────────────────────────────────
+
+interface EncodedCol {
+ presentStream: Uint8Array | null;
+ dataStream: Uint8Array;
+ lengthStream: Uint8Array | null;
+ typeKind: number;
+}
+
+function encodeIntCol(values: readonly Scalar[], typeKind: number): EncodedCol {
+ const nonNull: boolean[] = [];
+ const ints: bigint[] = [];
+ for (const v of values) {
+ const isPresent = v !== null && v !== undefined && !(typeof v === "number" && Number.isNaN(v));
+ nonNull.push(isPresent);
+ if (isPresent) {
+ ints.push(typeof v === "number" ? BigInt(Math.trunc(v)) : BigInt(String(v)));
+ }
+ }
+ const presentBits = packPresent(nonNull);
+ const presentStream = presentBits !== null ? rleByteEncodeV1(Array.from(presentBits)) : null;
+ return {
+ presentStream,
+ dataStream: rleIntEncodeV1(ints),
+ lengthStream: null,
+ typeKind,
+ };
+}
+
+function encodeF32Col(values: readonly Scalar[]): EncodedCol {
+ const nonNull: boolean[] = [];
+ const bytes: number[] = [];
+ const tmp = new ArrayBuffer(4);
+ const dv = new DataView(tmp);
+ for (const v of values) {
+ const isPresent = v !== null && v !== undefined && !(typeof v === "number" && Number.isNaN(v));
+ nonNull.push(isPresent);
+ if (isPresent) {
+ dv.setFloat32(0, typeof v === "number" ? v : Number(v), true);
+ bytes.push(dv.getUint8(0), dv.getUint8(1), dv.getUint8(2), dv.getUint8(3));
+ }
+ }
+ const presentBits = packPresent(nonNull);
+ return {
+ presentStream: presentBits !== null ? rleByteEncodeV1(Array.from(presentBits)) : null,
+ dataStream: new Uint8Array(bytes),
+ lengthStream: null,
+ typeKind: KIND_FLOAT,
+ };
+}
+
+function encodeF64Col(values: readonly Scalar[]): EncodedCol {
+ const nonNull: boolean[] = [];
+ const bytes: number[] = [];
+ const tmp = new ArrayBuffer(8);
+ const dv = new DataView(tmp);
+ for (const v of values) {
+ const isPresent = v !== null && v !== undefined && !(typeof v === "number" && Number.isNaN(v));
+ nonNull.push(isPresent);
+ if (isPresent) {
+ dv.setFloat64(0, typeof v === "number" ? v : Number(v), true);
+ for (let k = 0; k < 8; k++) {
+ bytes.push(dv.getUint8(k));
+ }
+ }
+ }
+ const presentBits = packPresent(nonNull);
+ return {
+ presentStream: presentBits !== null ? rleByteEncodeV1(Array.from(presentBits)) : null,
+ dataStream: new Uint8Array(bytes),
+ lengthStream: null,
+ typeKind: KIND_DOUBLE,
+ };
+}
+
+function encodeStringCol(values: readonly Scalar[]): EncodedCol {
+ const nonNull: boolean[] = [];
+ const dataBytes: number[] = [];
+ const lengths: bigint[] = [];
+ const enc = new TextEncoder();
+ for (const v of values) {
+ const isPresent = v !== null && v !== undefined;
+ nonNull.push(isPresent);
+ if (isPresent) {
+ const bytes = enc.encode(String(v));
+ for (const b of bytes) {
+ dataBytes.push(b);
+ }
+ lengths.push(BigInt(bytes.length));
+ }
+ }
+ const presentBits = packPresent(nonNull);
+ return {
+ presentStream: presentBits !== null ? rleByteEncodeV1(Array.from(presentBits)) : null,
+ dataStream: new Uint8Array(dataBytes),
+ lengthStream: rleIntEncodeV1(lengths),
+ typeKind: KIND_STRING,
+ };
+}
+
+function encodeBoolCol(values: readonly Scalar[]): EncodedCol {
+ const nonNull: boolean[] = [];
+ const bits: boolean[] = [];
+ for (const v of values) {
+ const isPresent = v !== null && v !== undefined;
+ nonNull.push(isPresent);
+ if (isPresent) {
+ bits.push(Boolean(v));
+ }
+ }
+ // Pack booleans into bytes, MSB first
+ const bytes: number[] = [];
+ for (let i = 0; i < bits.length; i += 8) {
+ let byte = 0;
+ for (let b = 0; b < 8 && i + b < bits.length; b++) {
+ if (bits[i + b]) {
+ byte |= 1 << (7 - b);
+ }
+ }
+ bytes.push(byte);
+ }
+ const presentBits = packPresent(nonNull);
+ return {
+ presentStream: presentBits !== null ? rleByteEncodeV1(Array.from(presentBits)) : null,
+ dataStream: rleByteEncodeV1(bytes),
+ lengthStream: null,
+ typeKind: KIND_BOOLEAN,
+ };
+}
+
+// ─── ORC type inference ───────────────────────────────────────────────────────
+
+/** Map a DataFrame column's values to an ORC type kind. */
+function inferOrcKind(values: readonly Scalar[]): number {
+ for (const v of values) {
+ if (v === null || v === undefined) {
+ continue;
+ }
+ if (typeof v === "boolean") {
+ return KIND_BOOLEAN;
+ }
+ if (typeof v === "bigint") {
+ return KIND_LONG;
+ }
+ if (typeof v === "number") {
+ return values.every(
+ (value) =>
+ value === null ||
+ value === undefined ||
+ (typeof value === "number" && Number.isSafeInteger(value)),
+ )
+ ? KIND_LONG
+ : KIND_DOUBLE;
+ }
+ if (typeof v === "string") {
+ return KIND_STRING;
+ }
+ }
+ return KIND_STRING; // default for all-null columns
+}
+
+/** Encode a column based on its inferred or given type. */
+function encodeColumn(values: readonly Scalar[], kind: number): EncodedCol {
+ switch (kind) {
+ case KIND_BOOLEAN:
+ return encodeBoolCol(values);
+ case KIND_FLOAT:
+ return encodeF32Col(values);
+ case KIND_DOUBLE:
+ return encodeF64Col(values);
+ case KIND_STRING:
+ return encodeStringCol(values);
+ default:
+ // All integer types → LONG
+ return encodeIntCol(values, KIND_LONG);
+ }
+}
+
+// ─── Postscript / Footer encoding ────────────────────────────────────────────
+
+function encodePostscript(footerLen: number, metaLen: number): Uint8Array {
+ const out: number[] = [];
+ pbWU64(1, BigInt(footerLen), out); // footerLength
+ pbWU32(2, COMP_NONE, out); // compression = NONE
+ // compressionBlockSize: omit (default)
+ // version: [0, 12] = ORC v0.12
+ pbWU32(4, 0, out);
+ pbWU32(4, 12, out);
+ pbWU64(5, BigInt(metaLen), out); // metadataLength
+ pbWU32(6, 1, out); // writerVersion
+ // magic: "ORC" (field 8000)
+ const magic = new TextEncoder().encode("ORC");
+ pbWBytes(8000, magic, out);
+ return new Uint8Array(out);
+}
+
+function encodeFooter(stripes: OrcStripeInfo[], types: OrcType[], nRows: number): Uint8Array {
+ const out: number[] = [];
+ pbWU64(1, BigInt(ORC_MAGIC.length), out); // headerLength = 3
+ // contentLength: sum of stripe sizes
+ const content = stripes.reduce(
+ (s, st) => s + st.indexLength + st.dataLength + st.footerLength,
+ 0,
+ );
+ pbWU64(2, BigInt(content), out);
+
+ for (const stripe of stripes) {
+ const sm: number[] = [];
+ pbWU64(1, BigInt(stripe.offset), sm);
+ pbWU64(2, BigInt(stripe.indexLength), sm);
+ pbWU64(3, BigInt(stripe.dataLength), sm);
+ pbWU64(4, BigInt(stripe.footerLength), sm);
+ pbWU64(5, BigInt(stripe.numberOfRows), sm);
+ pbWMsg(3, sm, out);
+ }
+
+ for (const type of types) {
+ const tm: number[] = [];
+ pbWU32(1, type.kind, tm);
+ for (const st of type.subtypes) {
+ pbWU32(2, st, tm);
+ }
+ for (const fn of type.fieldNames) {
+ const fnBytes = new TextEncoder().encode(fn);
+ pbWBytes(3, fnBytes, tm);
+ }
+ pbWMsg(4, tm, out);
+ }
+
+ pbWU64(6, BigInt(nRows), out);
+ pbWU32(8, 10000, out); // rowIndexStride
+ return new Uint8Array(out);
+}
+
+function encodeStripeFooter(streams: OrcStream[], columns: OrcColumnEncoding[]): Uint8Array {
+ const out: number[] = [];
+ for (const s of streams) {
+ const sm: number[] = [];
+ pbWU32(1, s.kind, sm);
+ pbWU32(2, s.column, sm);
+ pbWU64(3, BigInt(s.length), sm);
+ pbWMsg(1, sm, out);
+ }
+ for (const c of columns) {
+ const cm: number[] = [];
+ pbWU32(1, c.kind, cm);
+ if (c.dictionarySize > 0) {
+ pbWU32(2, c.dictionarySize, cm);
+ }
+ pbWMsg(2, cm, out);
+ }
+ return new Uint8Array(out);
+}
+
+// ─── Main: readOrc ────────────────────────────────────────────────────────────
+
+/** Convert a Scalar value to a Label (non-Label Scalars become null). */
+function scalarToLabel(v: Scalar): Label {
+ if (v === undefined) {
+ return null;
+ }
+ if (typeof v === "bigint") {
+ return Number(v);
+ }
+ if (typeof v === "number" || typeof v === "string" || typeof v === "boolean") {
+ return v;
+ }
+ if (v === null) {
+ return null;
+ }
+ if (v instanceof Date) {
+ return v;
+ }
+ return null; // TimedeltaLike
+}
+
+/**
+ * Parse an ORC binary buffer into a DataFrame.
+ *
+ * @param data - Raw ORC file bytes (Uint8Array or ArrayBuffer).
+ * @param options - Optional settings.
+ * @returns Parsed DataFrame.
+ *
+ * @example
+ * ```ts
+ * import { readOrc, toOrc, DataFrame } from "tsb";
+ * const buf = toOrc(DataFrame.fromColumns({ x: [1, 2, 3], y: ["a", "b", "c"] }));
+ * const df = readOrc(buf);
+ * ```
+ */
+export function readOrc(data: Uint8Array | ArrayBuffer, options: ReadOrcOptions = {}): DataFrame {
+ const buf = data instanceof ArrayBuffer ? new Uint8Array(data) : data;
+ if (buf.length < 4) {
+ throw new Error("ORC: file too small");
+ }
+ // Validate magic
+ if (buf[0] !== 0x4f || buf[1] !== 0x52 || buf[2] !== 0x43) {
+ throw new Error("ORC: invalid magic bytes (expected 'ORC')");
+ }
+ // Read postscript
+ const psLen = buf.at(-1);
+ if (psLen === undefined || psLen === 0) {
+ throw new Error("ORC: invalid postscript length");
+ }
+ const psStart = buf.length - 1 - psLen;
+ if (psStart < 3) {
+ throw new Error("ORC: file too small for postscript");
+ }
+ const ps = decodePostscript(buf.subarray(psStart, psStart + psLen));
+ if (ps.compression !== COMP_NONE) {
+ throw new Error(
+ `ORC: compression codec ${ps.compression} (${ps.compression === COMP_ZLIB ? "ZLIB" : "unsupported"}) is not supported; only NONE is currently implemented`,
+ );
+ }
+ // Read footer
+ const metaEnd = psStart;
+ const footerEnd = metaEnd - ps.metadataLength;
+ const footerStart = footerEnd - ps.footerLength;
+ if (footerStart < 3) {
+ throw new Error("ORC: invalid footer position");
+ }
+ const footer = decodeFooter(buf.subarray(footerStart, footerEnd));
+ if (footer.types.length === 0) {
+ throw new Error("ORC: no type schema in footer");
+ }
+
+ // Root type must be STRUCT
+ const rootType = footer.types[0];
+ if (rootType === undefined || rootType.kind !== KIND_STRUCT) {
+ throw new Error("ORC: root type is not STRUCT");
+ }
+
+ // Column selection
+ const allCols = rootType.fieldNames;
+ const wantSet = options.columns != null ? new Set([...options.columns]) : null;
+ const colIndices: number[] = rootType.subtypes.filter((_, i) => {
+ const name = allCols[i];
+ return name !== undefined && (wantSet === null || wantSet.has(name));
+ });
+ const colNames: string[] = colIndices.map((ci) => {
+ const fi = rootType.subtypes.indexOf(ci);
+ return allCols[fi] ?? String(ci);
+ });
+
+ const allValues: Map = new Map();
+ for (const name of colNames) {
+ allValues.set(name, []);
+ }
+
+ // Process each stripe
+ for (const stripeInfo of footer.stripes) {
+ const stripeDataStart = stripeInfo.offset + stripeInfo.indexLength;
+ const stripeFStart = stripeInfo.offset + stripeInfo.indexLength + stripeInfo.dataLength;
+ const stripeFBuf = buf.subarray(stripeFStart, stripeFStart + stripeInfo.footerLength);
+ const sf = decodeStripeFooter(stripeFBuf);
+
+ // Build a stream-offset map: column → streamKind → {offset, length}
+ type StreamLoc = { offset: number; length: number };
+ const streamMap = new Map>();
+ let streamPos = stripeDataStart;
+ for (const stream of sf.streams) {
+ let colMap = streamMap.get(stream.column);
+ if (colMap === undefined) {
+ colMap = new Map();
+ streamMap.set(stream.column, colMap);
+ }
+ colMap.set(stream.kind, { offset: streamPos, length: stream.length });
+ streamPos += stream.length;
+ }
+
+ const nRows = stripeInfo.numberOfRows;
+
+ for (let ci = 0; ci < colIndices.length; ci++) {
+ const colIdx = colIndices[ci];
+ if (colIdx === undefined) {
+ continue;
+ }
+ const name = colNames[ci];
+ if (name === undefined) {
+ continue;
+ }
+ const typeKind = footer.types[colIdx]?.kind ?? KIND_STRING;
+ const colStreams = streamMap.get(colIdx);
+ const vals = allValues.get(name);
+ if (vals === undefined) {
+ continue;
+ }
+
+ // PRESENT stream (null flags)
+ const presentLoc = colStreams?.get(STREAM_PRESENT);
+ let presentFlags: boolean[] | null = null;
+ if (presentLoc !== undefined && presentLoc.length > 0) {
+ const rawPresent = rleByteDecodeV1(buf, presentLoc.offset, presentLoc.length);
+ presentFlags = expandPresent(rawPresent, nRows);
+ }
+
+ const dataLoc = colStreams?.get(STREAM_DATA);
+ const dataOff = dataLoc?.offset ?? 0;
+ const dataLen = dataLoc?.length ?? 0;
+
+ switch (typeKind) {
+ case KIND_BOOLEAN: {
+ const decoded = decodeBoolCol(buf, dataOff, dataLen, presentFlags, nRows);
+ for (const v of decoded) {
+ vals.push(v);
+ }
+ break;
+ }
+ case KIND_FLOAT: {
+ const decoded = decodeF32Col(buf, dataOff, presentFlags, nRows);
+ for (const v of decoded) {
+ vals.push(v);
+ }
+ break;
+ }
+ case KIND_DOUBLE: {
+ const decoded = decodeF64Col(buf, dataOff, presentFlags, nRows);
+ for (const v of decoded) {
+ vals.push(v);
+ }
+ break;
+ }
+ case KIND_STRING: {
+ const lenLoc = colStreams?.get(STREAM_LENGTH);
+ const decoded = decodeStringCol(
+ buf,
+ dataOff,
+ buf,
+ lenLoc?.offset ?? 0,
+ lenLoc?.length ?? 0,
+ presentFlags,
+ nRows,
+ );
+ for (const v of decoded) {
+ vals.push(v);
+ }
+ break;
+ }
+ default: {
+ // Integer types: BOOLEAN, BYTE, SHORT, INT, LONG, DATE
+ const decoded = decodeIntCol(buf, dataOff, dataLen, presentFlags, nRows);
+ for (const v of decoded) {
+ vals.push(v === null ? null : Number(v));
+ }
+ break;
+ }
+ }
+ }
+ }
+
+ // Build DataFrame
+ const cols: Record = {};
+ const indexColName = options.indexCol ?? null;
+ let indexArr: Label[] | null = null;
+
+ for (const name of colNames) {
+ const data2 = allValues.get(name) ?? [];
+ if (name === indexColName) {
+ indexArr = data2.map(scalarToLabel);
+ } else {
+ cols[name] = data2;
+ }
+ }
+
+ return DataFrame.fromColumns(cols, indexArr !== null ? { index: indexArr } : undefined);
+}
+
+// ─── Main: toOrc ─────────────────────────────────────────────────────────────
+
+/**
+ * Serialize a DataFrame to an ORC binary buffer.
+ *
+ * @param df - DataFrame to serialize.
+ * @param options - Optional settings.
+ * @returns Raw ORC file bytes.
+ *
+ * @example
+ * ```ts
+ * import { toOrc, DataFrame } from "tsb";
+ * const df = DataFrame.fromColumns({ x: [1, 2, 3], y: ["a", "b", "c"] });
+ * const buf = toOrc(df);
+ * ```
+ */
+export function toOrc(df: DataFrame, options: ToOrcOptions = {}): Uint8Array {
+ const colNames = df.columns.toArray().map(String);
+ const extraCols: string[] = options.writeIndex === true ? ["__index__", ...colNames] : colNames;
+ const indexVals: Scalar[] | null = options.writeIndex === true ? df.index.toArray() : null;
+
+ // Collect column data
+ const colData: Scalar[][] = extraCols.map((name) =>
+ name === "__index__" && indexVals !== null ? indexVals : df.col(name).values.slice(),
+ );
+
+ // Infer ORC type kinds
+ const colKinds: number[] = colData.map((vals) => inferOrcKind(vals));
+
+ // Encode columns
+ const encoded: EncodedCol[] = colData.map((vals, i) => {
+ const k = colKinds[i] ?? KIND_STRING;
+ return encodeColumn(vals, k);
+ });
+
+ // Build file byte array
+ const out: number[] = [];
+
+ // Header "ORC"
+ for (const b of ORC_MAGIC) {
+ out.push(b);
+ }
+
+ // Build one stripe
+ const nRows = df.shape[0];
+ const stripeOffset = 3; // after header
+
+ // Write all streams
+ const streams: OrcStream[] = [];
+ const colEncodings: OrcColumnEncoding[] = [{ kind: ENC_DIRECT, dictionarySize: 0 }]; // root STRUCT
+ const _streamPos = stripeOffset;
+
+ for (let ci = 0; ci < encoded.length; ci++) {
+ const enc = encoded[ci];
+ if (enc === undefined) {
+ continue;
+ }
+ const colIdx = ci + 1; // 1-based (0 = root STRUCT)
+
+ if (enc.presentStream !== null) {
+ streams.push({ kind: STREAM_PRESENT, column: colIdx, length: enc.presentStream.length });
+ }
+ streams.push({ kind: STREAM_DATA, column: colIdx, length: enc.dataStream.length });
+ if (enc.lengthStream !== null) {
+ streams.push({ kind: STREAM_LENGTH, column: colIdx, length: enc.lengthStream.length });
+ }
+ colEncodings.push({ kind: ENC_DIRECT, dictionarySize: 0 });
+ }
+
+ // Write stream data
+ const stripeIndexLen = 0; // no row indexes
+ for (let ci = 0; ci < encoded.length; ci++) {
+ const enc = encoded[ci];
+ if (enc === undefined) {
+ continue;
+ }
+ if (enc.presentStream !== null) {
+ for (const b of enc.presentStream) {
+ out.push(b);
+ }
+ }
+ for (const b of enc.dataStream) {
+ out.push(b);
+ }
+ if (enc.lengthStream !== null) {
+ for (const b of enc.lengthStream) {
+ out.push(b);
+ }
+ }
+ }
+
+ // Compute data length (written bytes minus header)
+ const stripeDataLen = out.length - stripeOffset;
+
+ // Stripe footer
+ const sf = encodeStripeFooter(streams, colEncodings);
+ for (const b of sf) {
+ out.push(b);
+ }
+
+ // Build ORC type schema
+ // Column 0: STRUCT with all column fields
+ const types: OrcType[] = [
+ {
+ kind: KIND_STRUCT,
+ subtypes: extraCols.map((_, i) => i + 1),
+ fieldNames: extraCols,
+ },
+ ];
+ for (const kind of colKinds) {
+ types.push({ kind, subtypes: [], fieldNames: [] });
+ }
+
+ // Stripe info
+ const stripeInfo: OrcStripeInfo = {
+ offset: stripeOffset,
+ indexLength: stripeIndexLen,
+ dataLength: stripeDataLen,
+ footerLength: sf.length,
+ numberOfRows: nRows,
+ };
+
+ // File footer
+ const footerBytes = encodeFooter([stripeInfo], types, nRows);
+ for (const b of footerBytes) {
+ out.push(b);
+ }
+
+ // File metadata (empty for now)
+ const metaLen = 0;
+
+ // Postscript
+ const psBytes = encodePostscript(footerBytes.length, metaLen);
+ for (const b of psBytes) {
+ out.push(b);
+ }
+
+ // Postscript length (1 byte)
+ if (psBytes.length > 255) {
+ throw new Error("ORC: postscript too large");
+ }
+ out.push(psBytes.length);
+
+ return new Uint8Array(out);
+}
diff --git a/src/io/read_avro.ts b/src/io/read_avro.ts
new file mode 100644
index 000000000..331fb2375
--- /dev/null
+++ b/src/io/read_avro.ts
@@ -0,0 +1,698 @@
+/**
+ * read_avro — Apache Avro Object Container File (OCF) reader for DataFrame.
+ *
+ * Mirrors `pandas.read_avro()`. Parses Avro OCF binary format purely in
+ * TypeScript with no external dependencies.
+ *
+ * Supported Avro schema types:
+ * - Primitives: null, boolean, int, long, float, double, string, bytes
+ * - Named: record, enum, fixed
+ * - Complex: array, map, union
+ * - Logical: date (int), timestamp-millis (long), timestamp-micros (long)
+ *
+ * Supported codecs: `null` (uncompressed). `deflate` and `snappy` blocks
+ * are detected but raise an informative error.
+ *
+ * @example
+ * ```ts
+ * import { readAvro } from "tsb";
+ *
+ * const df = readAvro(buffer); // Uint8Array from file read
+ * console.log(df.columns, df.shape);
+ * ```
+ *
+ * @module
+ */
+
+import { DataFrame } from "../core/frame.ts";
+import type { Scalar } from "../types.ts";
+
+// ─── Public types ──────────────────────────────────────────────────────────────
+
+/** Options for {@link readAvro}. */
+export interface ReadAvroOptions {
+ /** Columns to include. Default: all columns. */
+ readonly usecols?: readonly string[] | null;
+ /**
+ * How to handle schema unions that include `null`:
+ * - `"object"` (default): return `null` for null values.
+ * - `"first"`: return the first non-null type's value.
+ */
+ readonly nullHandling?: "object" | "first";
+}
+
+// ─── Avro schema types ────────────────────────────────────────────────────────
+
+type AvroSchema =
+ | AvroPrimitive
+ | AvroRecord
+ | AvroEnum
+ | AvroArray
+ | AvroMap
+ | AvroUnion
+ | AvroFixed;
+
+type AvroPrimitive = "null" | "boolean" | "int" | "long" | "float" | "double" | "string" | "bytes";
+
+interface AvroRecord {
+ type: "record";
+ name: string;
+ fields: readonly AvroField[];
+}
+
+interface AvroField {
+ name: string;
+ type: AvroSchema;
+ default?: unknown;
+}
+
+interface AvroEnum {
+ type: "enum";
+ name: string;
+ symbols: readonly string[];
+}
+
+interface AvroArray {
+ type: "array";
+ items: AvroSchema;
+}
+
+interface AvroMap {
+ type: "map";
+ values: AvroSchema;
+}
+
+type AvroUnion = readonly AvroSchema[];
+
+interface AvroFixed {
+ type: "fixed";
+ name: string;
+ size: number;
+}
+
+// ─── Binary reader ────────────────────────────────────────────────────────────
+
+class AvroReader {
+ private buf: Uint8Array;
+ private pos = 0;
+
+ constructor(buf: Uint8Array) {
+ this.buf = buf;
+ }
+
+ get position(): number {
+ return this.pos;
+ }
+ get remaining(): number {
+ return this.buf.length - this.pos;
+ }
+
+ readByte(): number {
+ if (this.pos >= this.buf.length) {
+ throw new RangeError("Unexpected end of Avro data");
+ }
+ return this.buf[this.pos++] ?? 0;
+ }
+
+ /** Read a variable-length zigzag-encoded long. Returns a JS number (safe up to 2^53). */
+ readLong(): number {
+ let result = 0;
+ let shift = 0;
+ while (true) {
+ const b = this.readByte();
+ result |= (b & 0x7f) << shift;
+ shift += 7;
+ if ((b & 0x80) === 0) {
+ break;
+ }
+ if (shift >= 63) {
+ // For very large values, handle the remaining bits separately
+ // to avoid JS bitwise overflow (32-bit integers)
+ if (shift === 63) {
+ const hi = b & 0x7f;
+ // combine: result (low 63 bits) + hi << 63 — just approximate as float
+ const lo = result >>> 0;
+ result = lo + hi * 2 ** 63;
+ }
+ break;
+ }
+ }
+ // Zigzag decode: (n >>> 1) ^ -(n & 1)
+ return (result >>> 1) ^ -(result & 1);
+ }
+
+ /** Read a 32-bit int (zigzag long with range check). */
+ readInt(): number {
+ return this.readLong() | 0;
+ }
+
+ /** Read 4-byte IEEE 754 float. */
+ readFloat(): number {
+ const bytes = this.readBytes(4);
+ const view = new DataView(bytes.buffer, bytes.byteOffset, 4);
+ return view.getFloat32(0, true);
+ }
+
+ /** Read 8-byte IEEE 754 double. */
+ readDouble(): number {
+ const bytes = this.readBytes(8);
+ const view = new DataView(bytes.buffer, bytes.byteOffset, 8);
+ return view.getFloat64(0, true);
+ }
+
+ /** Read exactly n bytes as a new Uint8Array. */
+ readBytes(n: number): Uint8Array {
+ if (this.pos + n > this.buf.length) {
+ throw new RangeError("Unexpected end of Avro data");
+ }
+ const slice = this.buf.subarray(this.pos, this.pos + n);
+ this.pos += n;
+ return slice;
+ }
+
+ /** Read Avro bytes field (length-prefixed). */
+ readByteField(): Uint8Array {
+ const len = this.readLong();
+ return this.readBytes(len);
+ }
+
+ /** Read UTF-8 string (length-prefixed). */
+ readString(): string {
+ const bytes = this.readByteField();
+ return new TextDecoder().decode(bytes);
+ }
+
+ /** Read a boolean (0 = false, 1 = true). */
+ readBoolean(): boolean {
+ return this.readByte() !== 0;
+ }
+
+ /** Skip forward n bytes. */
+ skip(n: number): void {
+ if (this.pos + n > this.buf.length) {
+ throw new RangeError("Unexpected end of Avro data");
+ }
+ this.pos += n;
+ }
+
+ /** Peek at 16 bytes (sync marker) and advance. */
+ readSync(): Uint8Array {
+ return this.readBytes(16);
+ }
+}
+
+// ─── Schema parsing ────────────────────────────────────────────────────────────
+
+const td = new TextDecoder();
+
+function parseSchema(raw: unknown): AvroSchema {
+ if (typeof raw === "string") {
+ const prim = raw as AvroPrimitive;
+ return prim;
+ }
+ if (Array.isArray(raw)) {
+ return raw.map(parseSchema) as AvroUnion;
+ }
+ if (typeof raw === "object" && raw !== null) {
+ const obj = raw as Record;
+ const type = obj["type"];
+ if (type === "record") {
+ const fields = (obj["fields"] as unknown[]).map((f) => {
+ const field = f as Record;
+ return { name: field["name"] as string, type: parseSchema(field["type"]) };
+ });
+ return { type: "record", name: obj["name"] as string, fields };
+ }
+ if (type === "array") {
+ return { type: "array", items: parseSchema(obj["items"]) };
+ }
+ if (type === "map") {
+ return { type: "map", values: parseSchema(obj["values"]) };
+ }
+ if (type === "enum") {
+ return {
+ type: "enum",
+ name: obj["name"] as string,
+ symbols: obj["symbols"] as string[],
+ };
+ }
+ if (type === "fixed") {
+ return {
+ type: "fixed",
+ name: obj["name"] as string,
+ size: obj["size"] as number,
+ };
+ }
+ // Logical types: delegate to the underlying type
+ if (typeof type === "string") {
+ return parseSchema(type);
+ }
+ }
+ throw new TypeError(`Unknown Avro schema: ${JSON.stringify(raw)}`);
+}
+
+// ─── Datum reading ────────────────────────────────────────────────────────────
+
+/** Avro leaf value (no recursion at type level — containers use unknown). */
+type AvroLeaf = null | boolean | number | string | Uint8Array;
+/** Container interfaces allow recursive self-reference (interfaces can, type aliases cannot). */
+interface AvroDatumArr extends Array {}
+interface AvroDatumMap extends Map {}
+interface AvroDatumRecord extends Record {}
+/** Recursive Avro datum. */
+type AvroDatum = AvroLeaf | AvroDatumArr | AvroDatumMap | AvroDatumRecord;
+
+function readDatum(reader: AvroReader, schema: AvroSchema): AvroDatum {
+ if (typeof schema === "string") {
+ switch (schema) {
+ case "null":
+ return null;
+ case "boolean":
+ return reader.readBoolean();
+ case "int":
+ return reader.readInt();
+ case "long":
+ return reader.readLong();
+ case "float":
+ return reader.readFloat();
+ case "double":
+ return reader.readDouble();
+ case "string":
+ return reader.readString();
+ case "bytes":
+ return reader.readByteField();
+ }
+ }
+ if (Array.isArray(schema)) {
+ // Union: first read the branch index
+ const idx = reader.readLong();
+ const branch = (schema as AvroUnion)[idx];
+ if (branch === undefined) {
+ throw new RangeError(`Union branch ${idx} out of range`);
+ }
+ return readDatum(reader, branch);
+ }
+ const s = schema as Exclude;
+ if (s.type === "record") {
+ const rec: Record = {};
+ for (const field of s.fields) {
+ rec[field.name] = readDatum(reader, field.type);
+ }
+ return rec;
+ }
+ if (s.type === "array") {
+ const arr: AvroDatum[] = [];
+ while (true) {
+ let count = reader.readLong();
+ if (count === 0) {
+ break;
+ }
+ // Negative count means block has a byte count prefix
+ if (count < 0) {
+ reader.readLong();
+ count = -count;
+ }
+ for (let i = 0; i < count; i++) {
+ arr.push(readDatum(reader, s.items));
+ }
+ }
+ return arr;
+ }
+ if (s.type === "map") {
+ const map = new Map();
+ while (true) {
+ let count = reader.readLong();
+ if (count === 0) {
+ break;
+ }
+ if (count < 0) {
+ reader.readLong();
+ count = -count;
+ }
+ for (let i = 0; i < count; i++) {
+ const key = reader.readString();
+ map.set(key, readDatum(reader, s.values));
+ }
+ }
+ return map;
+ }
+ if (s.type === "enum") {
+ const idx = reader.readInt();
+ return s.symbols[idx] ?? null;
+ }
+ if (s.type === "fixed") {
+ return reader.readBytes(s.size);
+ }
+ throw new TypeError(`Unhandled schema type: ${JSON.stringify(schema)}`);
+}
+
+// ─── OCF parsing ──────────────────────────────────────────────────────────────
+
+const AVRO_MAGIC = new Uint8Array([79, 98, 106, 1]); // "Obj\x01"
+
+function syncEq(a: Uint8Array, b: Uint8Array): boolean {
+ if (a.length !== b.length) {
+ return false;
+ }
+ for (let i = 0; i < a.length; i++) {
+ if (a[i] !== b[i]) {
+ return false;
+ }
+ }
+ return true;
+}
+
+/**
+ * Parse an Apache Avro Object Container File buffer into an array of row objects.
+ * Returns the top-level schema and the rows.
+ */
+function parseAvroOCF(buf: Uint8Array): { schema: AvroSchema; rows: Record[] } {
+ const reader = new AvroReader(buf);
+
+ // Magic: "Obj\x01"
+ const magic = reader.readBytes(4);
+ if (!syncEq(magic, AVRO_MAGIC)) {
+ throw new TypeError(
+ `Not a valid Avro file: expected magic bytes "Obj\\x01", got ${[...magic].map((b) => b.toString(16)).join(" ")}`,
+ );
+ }
+
+ // File-level metadata: map
+ const meta = new Map();
+ while (true) {
+ let count = reader.readLong();
+ if (count === 0) {
+ break;
+ }
+ if (count < 0) {
+ reader.readLong();
+ count = -count;
+ }
+ for (let i = 0; i < count; i++) {
+ const key = reader.readString();
+ const val = reader.readByteField();
+ meta.set(key, val);
+ }
+ }
+
+ // Schema
+ const schemaBytes = meta.get("avro.schema");
+ if (!schemaBytes) {
+ throw new TypeError("Avro file missing avro.schema metadata");
+ }
+ const schemaJson: unknown = JSON.parse(td.decode(schemaBytes));
+ const schema = parseSchema(schemaJson);
+
+ // Codec
+ const codecBytes = meta.get("avro.codec");
+ const codec = codecBytes ? td.decode(codecBytes) : "null";
+ if (codec !== "null") {
+ throw new TypeError(
+ `Avro codec "${codec}" is not supported. Only "null" (uncompressed) is implemented.`,
+ );
+ }
+
+ // Sync marker (16 bytes)
+ const syncMarker = reader.readSync();
+
+ // Data blocks
+ const rows: Record[] = [];
+ while (reader.remaining >= 16) {
+ const objectCount = reader.readLong();
+ const _byteCount = reader.readLong(); // block size (unused for null codec)
+ if (objectCount <= 0) {
+ break;
+ }
+
+ for (let i = 0; i < objectCount; i++) {
+ const datum = readDatum(reader, schema);
+ if (
+ typeof datum === "object" &&
+ datum !== null &&
+ !Array.isArray(datum) &&
+ !(datum instanceof Uint8Array) &&
+ !(datum instanceof Map)
+ ) {
+ rows.push(datum as Record);
+ }
+ }
+
+ // Read and verify sync marker
+ const blockSync = reader.readSync();
+ if (!syncEq(blockSync, syncMarker)) {
+ throw new TypeError("Avro sync marker mismatch — file may be corrupt");
+ }
+ }
+
+ return { schema, rows };
+}
+
+// ─── DataFrame construction ───────────────────────────────────────────────────
+
+function flattenDatum(v: AvroDatum): unknown {
+ if (v === null || typeof v !== "object") {
+ return v;
+ }
+ if (v instanceof Uint8Array) {
+ return v;
+ }
+ if (v instanceof Map) {
+ return Object.fromEntries(v);
+ }
+ // For record/array datums, JSON-stringify for simplicity
+ if (Array.isArray(v)) {
+ return JSON.stringify(v);
+ }
+ return JSON.stringify(v);
+}
+
+/**
+ * Read an Apache Avro Object Container File buffer into a {@link DataFrame}.
+ *
+ * @param data - Raw Avro OCF bytes (`Uint8Array` or `ArrayBuffer`).
+ * @param options - Optional read options.
+ */
+export function readAvro(data: Uint8Array | ArrayBuffer, options: ReadAvroOptions = {}): DataFrame {
+ const buf = data instanceof ArrayBuffer ? new Uint8Array(data) : data;
+ const { rows } = parseAvroOCF(buf);
+
+ if (rows.length === 0) {
+ return DataFrame.fromColumns({});
+ }
+
+ // Determine columns from first row
+ const allCols = Object.keys(rows[0] ?? {});
+ const cols = options.usecols
+ ? allCols.filter((c) => (options.usecols as readonly string[]).includes(c))
+ : allCols;
+
+ const columns: Record = {};
+ for (const col of cols) {
+ columns[col] = [];
+ }
+
+ for (const row of rows) {
+ for (const col of cols) {
+ const v = row[col];
+ (columns[col] ?? []).push(flattenDatum(v ?? null) as Scalar);
+ }
+ }
+
+ return DataFrame.fromColumns(columns);
+}
+
+// ─── Avro writer (minimal — for testing round-trips) ─────────────────────────
+
+/** Options for {@link toAvro}. */
+export interface ToAvroOptions {
+ /** Schema name for the top-level record. Default: "Row". */
+ readonly schemaName?: string;
+}
+
+/**
+ * Serialize a {@link DataFrame} to an uncompressed Avro OCF buffer.
+ *
+ * Column type mapping:
+ * - boolean columns → `boolean`
+ * - integer columns → `long`
+ * - float columns → `double`
+ * - string columns → `{"type":"union","schemas":["null","string"]}`
+ * - null column vals → wrapped in union `["null", ""]`
+ * - other → `string` (JSON-stringified)
+ */
+export function toAvro(df: DataFrame, options: ToAvroOptions = {}): Uint8Array {
+ const schemaName = options.schemaName ?? "Row";
+ const cols = [...df.columns.values];
+
+ // Infer field types
+ type FieldSpec = { name: string; avroType: string; nullable: boolean };
+ const fields: FieldSpec[] = cols.map((col) => {
+ const vals = df.col(col).values;
+ let hasNull = false;
+ let hasInt = false;
+ let hasFloat = false;
+ let hasBool = false;
+ let hasStr = false;
+ for (const v of vals) {
+ if (v === null || v === undefined) {
+ hasNull = true;
+ continue;
+ }
+ if (typeof v === "boolean") {
+ hasBool = true;
+ continue;
+ }
+ if (typeof v === "number") {
+ if (Number.isInteger(v)) {
+ hasInt = true;
+ } else {
+ hasFloat = true;
+ }
+ continue;
+ }
+ hasStr = true;
+ }
+ let avroType = "string";
+ if (hasBool && !hasInt && !hasFloat && !hasStr) {
+ avroType = "boolean";
+ } else if ((hasInt || hasFloat) && !hasBool && !hasStr) {
+ avroType = hasFloat ? "double" : "long";
+ }
+ return { name: col, avroType, nullable: hasNull };
+ });
+
+ // Build schema JSON
+ const schemaFields = fields.map((f) => ({
+ name: f.name,
+ type: f.nullable ? ["null", f.avroType] : f.avroType,
+ }));
+ const schemaObj = { type: "record", name: schemaName, fields: schemaFields };
+ const schemaJson = JSON.stringify(schemaObj);
+ const schemaBytes = new TextEncoder().encode(schemaJson);
+
+ // Sync marker: 16 random-ish bytes derived from schema hash
+ const sync = new Uint8Array(16);
+ let h = 0x12345678;
+ for (let i = 0; i < schemaBytes.length; i++) {
+ h = Math.imul(h ^ (schemaBytes[i] ?? 0), 0x9e3779b9) >>> 0;
+ }
+ for (let i = 0; i < 16; i++) {
+ sync[i] = (h >> ((i % 4) * 8)) & 0xff;
+ if (i % 4 === 3) {
+ h = Math.imul(h, 0x6c62272e) >>> 0;
+ }
+ }
+
+ const chunks: Uint8Array[] = [];
+
+ // Write helper
+ const writeBuf: number[] = [];
+ const flushBuf = (): Uint8Array => {
+ const u = new Uint8Array(writeBuf);
+ writeBuf.length = 0;
+ return u;
+ };
+
+ function writeLong(v: number): void {
+ // Zigzag encode
+ let n = (v << 1) ^ (v >> 31);
+ while (n & ~0x7f) {
+ writeBuf.push((n & 0x7f) | 0x80);
+ n >>>= 7;
+ }
+ writeBuf.push(n);
+ }
+ function writeString(s: string): void {
+ const b = new TextEncoder().encode(s);
+ writeLong(b.length);
+ for (const byte of b) {
+ writeBuf.push(byte);
+ }
+ }
+ function writeBytes(b: Uint8Array): void {
+ writeLong(b.length);
+ for (const byte of b) {
+ writeBuf.push(byte);
+ }
+ }
+
+ // Magic
+ chunks.push(AVRO_MAGIC);
+
+ // Metadata: 1 map block with avro.schema and avro.codec
+ writeLong(2); // 2 entries
+ writeString("avro.schema");
+ writeBytes(schemaBytes);
+ writeString("avro.codec");
+ writeBytes(new TextEncoder().encode("null"));
+ writeLong(0); // end of map
+ chunks.push(flushBuf());
+
+ // Sync marker
+ chunks.push(sync.slice());
+
+ // Data block
+ const nRows = df.shape[0];
+ if (nRows > 0) {
+ // Encode all rows
+ for (let row = 0; row < nRows; row++) {
+ for (const f of fields) {
+ const v = df.col(f.name).at(row);
+ if (f.nullable) {
+ if (v === null || v === undefined) {
+ writeLong(0); // null branch
+ } else {
+ writeLong(1); // value branch
+ writeTypedValue(v, f.avroType);
+ }
+ } else {
+ writeTypedValue(v ?? null, f.avroType);
+ }
+ }
+ }
+ const blockData = flushBuf();
+
+ // Block header: count, byteCount
+ writeLong(nRows);
+ writeLong(blockData.length);
+ chunks.push(flushBuf());
+ chunks.push(blockData);
+ chunks.push(sync.slice());
+ }
+
+ // Concatenate all chunks
+ const totalLen = chunks.reduce((s, c) => s + c.length, 0);
+ const out = new Uint8Array(totalLen);
+ let offset = 0;
+ for (const c of chunks) {
+ out.set(c, offset);
+ offset += c.length;
+ }
+ return out;
+
+ function writeTypedValue(v: unknown, type: string): void {
+ if (v === null || v === undefined) {
+ writeBuf.push(0);
+ return;
+ } // null
+ switch (type) {
+ case "boolean":
+ writeBuf.push(v ? 1 : 0);
+ break;
+ case "long":
+ writeLong(typeof v === "number" ? v : 0);
+ break;
+ case "double": {
+ const arr = new Float64Array(1);
+ arr[0] = typeof v === "number" ? v : 0;
+ const b = new Uint8Array(arr.buffer);
+ for (const byte of b) {
+ writeBuf.push(byte);
+ }
+ break;
+ }
+ default:
+ writeString(String(v));
+ }
+ }
+}
diff --git a/src/io/sql.ts b/src/io/sql.ts
index 00efed4b9..11f4170f3 100644
--- a/src/io/sql.ts
+++ b/src/io/sql.ts
@@ -291,20 +291,20 @@ function resultToDataFrame(result: SqlResult, options: ReadSqlBaseOptions): Data
// Build column arrays, excluding the index column.
const dataColumns: string[] = [];
- const columnData: Record = {};
+ const columnData = new Map();
for (const col of result.columns) {
if (col === idxColName) {
continue;
}
dataColumns.push(col);
- columnData[col] = [];
+ columnData.set(col, []);
}
// Populate column arrays.
for (const row of result.rows) {
for (const col of dataColumns) {
- const arr = columnData[col];
+ const arr = columnData.get(col);
if (arr !== undefined) {
const raw = row[col];
arr.push(raw !== undefined ? sqlValueToScalar(raw) : null);
@@ -315,7 +315,7 @@ function resultToDataFrame(result: SqlResult, options: ReadSqlBaseOptions): Data
// Parse date columns (convert to ms-since-epoch numbers).
if (parseDates !== undefined) {
for (const col of parseDates) {
- const arr = columnData[col];
+ const arr = columnData.get(col);
if (arr !== undefined) {
for (let i = 0; i < arr.length; i++) {
const v = arr[i];
@@ -345,7 +345,7 @@ function resultToDataFrame(result: SqlResult, options: ReadSqlBaseOptions): Data
const rowIndex = idxColName !== null ? new Index(indexVals, idxColName) : undefined;
return DataFrame.fromColumns(
- columnData as Record,
+ Object.fromEntries(columnData),
rowIndex !== undefined ? { index: rowIndex } : {},
);
}
diff --git a/src/ml/attention.ts b/src/ml/attention.ts
new file mode 100644
index 000000000..9cb66afc7
--- /dev/null
+++ b/src/ml/attention.ts
@@ -0,0 +1,228 @@
+/**
+ * Multi-head self-attention and transformer building blocks.
+ *
+ * Provides standalone attention primitives usable across architectures:
+ * multi-head attention, positional encodings, and transformer encoder/decoder blocks.
+ *
+ * @module
+ */
+
+/** Multi-head attention configuration. */
+export interface MHAConfig {
+ dModel: number;
+ numHeads: number;
+ dropout: number;
+}
+
+/** Project a flat weight matrix into per-head slices. */
+function splitHeads(
+ x: Float64Array,
+ seqLen: number,
+ dModel: number,
+ numHeads: number,
+): Float64Array[] {
+ const dHead = Math.floor(dModel / numHeads);
+ const heads: Float64Array[] = [];
+ for (let h = 0; h < numHeads; h++) {
+ const head = new Float64Array(seqLen * dHead);
+ for (let t = 0; t < seqLen; t++) {
+ for (let d = 0; d < dHead; d++) {
+ head[t * dHead + d] = x[t * dModel + h * dHead + d] ?? 0;
+ }
+ }
+ heads.push(head);
+ }
+ return heads;
+}
+
+/** Concatenate per-head outputs back into [seqLen x dModel]. */
+function mergeHeads(heads: Float64Array[], seqLen: number, dModel: number): Float64Array {
+ const numHeads = heads.length;
+ const dHead = Math.floor(dModel / numHeads);
+ const out = new Float64Array(seqLen * dModel);
+ for (let h = 0; h < numHeads; h++) {
+ const head = heads[h]!;
+ for (let t = 0; t < seqLen; t++) {
+ for (let d = 0; d < dHead; d++) {
+ out[t * dModel + h * dHead + d] = head[t * dHead + d] ?? 0;
+ }
+ }
+ }
+ return out;
+}
+
+/** Compute attention for a single head: [seqLen x dHead]. */
+function singleHeadAttention(
+ Q: Float64Array,
+ K: Float64Array,
+ V: Float64Array,
+ seqLen: number,
+ dHead: number,
+ mask?: Float64Array,
+): Float64Array {
+ const scale = Math.sqrt(dHead);
+ const attn = new Float64Array(seqLen * seqLen);
+ // Q K^T / scale
+ for (let i = 0; i < seqLen; i++) {
+ for (let j = 0; j < seqLen; j++) {
+ let dot = 0;
+ for (let d = 0; d < dHead; d++) {
+ dot += (Q[i * dHead + d] ?? 0) * (K[j * dHead + d] ?? 0);
+ }
+ attn[i * seqLen + j] = dot / scale + (mask ? (mask[i * seqLen + j] ?? 0) : 0);
+ }
+ }
+ // Softmax
+ for (let i = 0; i < seqLen; i++) {
+ let maxVal = Number.NEGATIVE_INFINITY;
+ for (let j = 0; j < seqLen; j++)
+ maxVal = Math.max(maxVal, attn[i * seqLen + j] ?? Number.NEGATIVE_INFINITY);
+ let sumExp = 0;
+ for (let j = 0; j < seqLen; j++) {
+ attn[i * seqLen + j] = Math.exp((attn[i * seqLen + j] ?? 0) - maxVal);
+ sumExp += attn[i * seqLen + j] ?? 0;
+ }
+ for (let j = 0; j < seqLen; j++)
+ attn[i * seqLen + j] = (attn[i * seqLen + j] ?? 0) / (sumExp || 1);
+ }
+ // A V
+ const out = new Float64Array(seqLen * dHead);
+ for (let i = 0; i < seqLen; i++) {
+ for (let d = 0; d < dHead; d++) {
+ let val = 0;
+ for (let j = 0; j < seqLen; j++) {
+ val += (attn[i * seqLen + j] ?? 0) * (V[j * dHead + d] ?? 0);
+ }
+ out[i * dHead + d] = val;
+ }
+ }
+ return out;
+}
+
+/** Multi-head attention weight matrices. */
+export interface MHAWeights {
+ Wq: Float64Array; // [dModel x dModel]
+ Wk: Float64Array; // [dModel x dModel]
+ Wv: Float64Array; // [dModel x dModel]
+ Wo: Float64Array; // [dModel x dModel]
+ config: MHAConfig;
+}
+
+/** Apply multi-head attention. */
+export function multiHeadAttention(
+ query: Float64Array, // [seqLen x dModel]
+ key: Float64Array, // [seqLen x dModel]
+ value: Float64Array, // [seqLen x dModel]
+ seqLen: number,
+ weights: MHAWeights,
+ mask?: Float64Array,
+): Float64Array {
+ const { dModel, numHeads } = weights.config;
+ const dHead = Math.floor(dModel / numHeads);
+
+ // Project Q, K, V
+ const projQ = new Float64Array(seqLen * dModel);
+ const projK = new Float64Array(seqLen * dModel);
+ const projV = new Float64Array(seqLen * dModel);
+
+ for (let t = 0; t < seqLen; t++) {
+ for (let o = 0; o < dModel; o++) {
+ let q = 0;
+ let k = 0;
+ let v = 0;
+ for (let i = 0; i < dModel; i++) {
+ q += (weights.Wq[o * dModel + i] ?? 0) * (query[t * dModel + i] ?? 0);
+ k += (weights.Wk[o * dModel + i] ?? 0) * (key[t * dModel + i] ?? 0);
+ v += (weights.Wv[o * dModel + i] ?? 0) * (value[t * dModel + i] ?? 0);
+ }
+ projQ[t * dModel + o] = q;
+ projK[t * dModel + o] = k;
+ projV[t * dModel + o] = v;
+ }
+ }
+
+ // Split into heads and compute attention
+ const qHeads = splitHeads(projQ, seqLen, dModel, numHeads);
+ const kHeads = splitHeads(projK, seqLen, dModel, numHeads);
+ const vHeads = splitHeads(projV, seqLen, dModel, numHeads);
+
+ const headOutputs: Float64Array[] = [];
+ for (let h = 0; h < numHeads; h++) {
+ headOutputs.push(singleHeadAttention(qHeads[h]!, kHeads[h]!, vHeads[h]!, seqLen, dHead, mask));
+ }
+
+ const concat = mergeHeads(headOutputs, seqLen, dModel);
+
+ // Output projection
+ const out = new Float64Array(seqLen * dModel);
+ for (let t = 0; t < seqLen; t++) {
+ for (let o = 0; o < dModel; o++) {
+ let val = 0;
+ for (let i = 0; i < dModel; i++) {
+ val += (weights.Wo[o * dModel + i] ?? 0) * (concat[t * dModel + i] ?? 0);
+ }
+ out[t * dModel + o] = val;
+ }
+ }
+ return out;
+}
+
+/** Sinusoidal positional encoding. */
+export function sinusoidalPositionalEncoding(seqLen: number, dModel: number): Float64Array {
+ const pe = new Float64Array(seqLen * dModel);
+ for (let pos = 0; pos < seqLen; pos++) {
+ for (let i = 0; i < dModel; i += 2) {
+ const angle = pos / 10000 ** (i / dModel);
+ pe[pos * dModel + i] = Math.sin(angle);
+ if (i + 1 < dModel) pe[pos * dModel + i + 1] = Math.cos(angle);
+ }
+ }
+ return pe;
+}
+
+/** Add positional encoding to input embeddings. */
+export function addPositionalEncoding(
+ embeddings: Float64Array,
+ seqLen: number,
+ dModel: number,
+): Float64Array {
+ const pe = sinusoidalPositionalEncoding(seqLen, dModel);
+ const out = new Float64Array(embeddings.length);
+ for (let i = 0; i < embeddings.length; i++) {
+ out[i] = (embeddings[i] ?? 0) + (pe[i] ?? 0);
+ }
+ return out;
+}
+
+/** Causal (autoregressive) attention mask. */
+export function causalMask(seqLen: number): Float64Array {
+ const mask = new Float64Array(seqLen * seqLen).fill(Number.NEGATIVE_INFINITY);
+ for (let i = 0; i < seqLen; i++) {
+ for (let j = 0; j <= i; j++) {
+ mask[i * seqLen + j] = 0;
+ }
+ }
+ return mask;
+}
+
+/** RoPE (Rotary Position Embedding) for a single head. */
+export function applyRoPE(
+ x: Float64Array,
+ seqLen: number,
+ dHead: number,
+ baseFreq = 10000,
+): Float64Array {
+ const out = new Float64Array(x.length);
+ for (let pos = 0; pos < seqLen; pos++) {
+ for (let i = 0; i < dHead; i += 2) {
+ const theta = pos / baseFreq ** (i / dHead);
+ const cos = Math.cos(theta);
+ const sin = Math.sin(theta);
+ const x0 = x[pos * dHead + i] ?? 0;
+ const x1 = x[pos * dHead + i + 1] ?? 0;
+ out[pos * dHead + i] = x0 * cos - x1 * sin;
+ out[pos * dHead + i + 1] = x0 * sin + x1 * cos;
+ }
+ }
+ return out;
+}
diff --git a/src/ml/bayesian_opt.ts b/src/ml/bayesian_opt.ts
new file mode 100644
index 000000000..38919ef28
--- /dev/null
+++ b/src/ml/bayesian_opt.ts
@@ -0,0 +1,257 @@
+/**
+ * Bayesian Optimization with Gaussian Process surrogate and acquisition functions.
+ *
+ * Implements Expected Improvement (EI), Probability of Improvement (PI),
+ * Upper Confidence Bound (UCB), and Thompson Sampling acquisition functions,
+ * with an RBF kernel Gaussian Process as the surrogate model.
+ *
+ * @module
+ */
+
+/** RBF (squared exponential) kernel. */
+export function rbfKernel(
+ x1: Float64Array,
+ x2: Float64Array,
+ lengthScale: number,
+ amplitude: number,
+): number {
+ let sqDist = 0;
+ for (let i = 0; i < x1.length; i++) {
+ const d = (x1[i] ?? 0) - (x2[i] ?? 0);
+ sqDist += d * d;
+ }
+ return amplitude * amplitude * Math.exp((-0.5 * sqDist) / (lengthScale * lengthScale));
+}
+
+/** Matern 5/2 kernel. */
+export function maternKernel52(
+ x1: Float64Array,
+ x2: Float64Array,
+ lengthScale: number,
+ amplitude: number,
+): number {
+ let sqDist = 0;
+ for (let i = 0; i < x1.length; i++) {
+ const d = (x1[i] ?? 0) - (x2[i] ?? 0);
+ sqDist += d * d;
+ }
+ const r = Math.sqrt(sqDist) / lengthScale;
+ const sqrt5r = Math.sqrt(5) * r;
+ return amplitude * amplitude * (1 + sqrt5r + (5 / 3) * r * r) * Math.exp(-sqrt5r);
+}
+
+/** Compute kernel matrix K(X, X) + noise * I. */
+export function kernelMatrix(
+ X: Float64Array[],
+ lengthScale: number,
+ amplitude: number,
+ noise: number,
+): Float64Array {
+ const n = X.length;
+ const K = new Float64Array(n * n);
+ for (let i = 0; i < n; i++) {
+ for (let j = i; j < n; j++) {
+ const k = rbfKernel(X[i]!, X[j]!, lengthScale, amplitude);
+ K[i * n + j] = k;
+ K[j * n + i] = k;
+ }
+ K[i * n + i] = (K[i * n + i] ?? 0) + noise;
+ }
+ return K;
+}
+
+/** Cholesky decomposition (lower triangular L such that L L^T = A). */
+export function cholesky(A: Float64Array, n: number): Float64Array {
+ const L = new Float64Array(n * n);
+ for (let i = 0; i < n; i++) {
+ for (let j = 0; j <= i; j++) {
+ let sum = A[i * n + j] ?? 0;
+ for (let k = 0; k < j; k++) {
+ sum -= (L[i * n + k] ?? 0) * (L[j * n + k] ?? 0);
+ }
+ if (i === j) {
+ L[i * n + j] = Math.sqrt(Math.max(sum, 1e-12));
+ } else {
+ L[i * n + j] = sum / ((L[j * n + j] ?? 1e-12) || 1e-12);
+ }
+ }
+ }
+ return L;
+}
+
+/** Solve L x = b (forward substitution). */
+function forwardSolve(L: Float64Array, b: Float64Array, n: number): Float64Array {
+ const x = new Float64Array(n);
+ for (let i = 0; i < n; i++) {
+ let sum = b[i] ?? 0;
+ for (let j = 0; j < i; j++) sum -= (L[i * n + j] ?? 0) * (x[j] ?? 0);
+ x[i] = sum / ((L[i * n + i] ?? 1) || 1);
+ }
+ return x;
+}
+
+/** Solve L^T x = b (backward substitution). */
+function backwardSolve(L: Float64Array, b: Float64Array, n: number): Float64Array {
+ const x = new Float64Array(n);
+ for (let i = n - 1; i >= 0; i--) {
+ let sum = b[i] ?? 0;
+ for (let j = i + 1; j < n; j++) sum -= (L[j * n + i] ?? 0) * (x[j] ?? 0);
+ x[i] = sum / ((L[i * n + i] ?? 1) || 1);
+ }
+ return x;
+}
+
+/** Gaussian Process prediction: return (mean, variance) at test point. */
+export function gpPredict(
+ xTest: Float64Array,
+ X: Float64Array[],
+ y: Float64Array,
+ L: Float64Array,
+ lengthScale: number,
+ amplitude: number,
+ noise: number,
+): { mean: number; variance: number } {
+ const n = X.length;
+ // k_* = [k(x*, x_i)]
+ const kStar = new Float64Array(n);
+ for (let i = 0; i < n; i++) {
+ kStar[i] = rbfKernel(xTest, X[i]!, lengthScale, amplitude);
+ }
+ // alpha = (K + noise*I)^{-1} y via Cholesky
+ const alpha1 = forwardSolve(L, y, n);
+ const alpha = backwardSolve(L, alpha1, n);
+
+ // mean = k_*^T alpha
+ let mean = 0;
+ for (let i = 0; i < n; i++) mean += (kStar[i] ?? 0) * (alpha[i] ?? 0);
+
+ // variance = k(x*,x*) - k_*^T (K+noise*I)^{-1} k_*
+ const v = forwardSolve(L, kStar, n);
+ let kStarDotV = 0;
+ for (let i = 0; i < n; i++) kStarDotV += (v[i] ?? 0) * (v[i] ?? 0);
+ const kSelf = rbfKernel(xTest, xTest, lengthScale, amplitude) + noise;
+ const variance = Math.max(kSelf - kStarDotV, 1e-10);
+
+ return { mean, variance };
+}
+
+/** Standard normal CDF (approximation). */
+export function normalCDF(z: number): number {
+ const t = 1 / (1 + 0.2316419 * Math.abs(z));
+ const poly =
+ t *
+ (0.31938153 + t * (-0.356563782 + t * (1.781477937 + t * (-1.821255978 + t * 1.330274429))));
+ const phi = 1 - (1 / Math.sqrt(2 * Math.PI)) * Math.exp(-0.5 * z * z) * poly;
+ return z >= 0 ? phi : 1 - phi;
+}
+
+/** Standard normal PDF. */
+export function normalPDF(z: number): number {
+ return (1 / Math.sqrt(2 * Math.PI)) * Math.exp(-0.5 * z * z);
+}
+
+/** Expected Improvement acquisition function. */
+export function expectedImprovement(
+ mean: number,
+ variance: number,
+ bestY: number,
+ xi = 0.01,
+): number {
+ const std = Math.sqrt(variance);
+ const z = (mean - bestY - xi) / (std || 1e-8);
+ return (mean - bestY - xi) * normalCDF(z) + std * normalPDF(z);
+}
+
+/** Probability of Improvement acquisition function. */
+export function probabilityOfImprovement(
+ mean: number,
+ variance: number,
+ bestY: number,
+ xi = 0.01,
+): number {
+ const std = Math.sqrt(variance);
+ return normalCDF((mean - bestY - xi) / (std || 1e-8));
+}
+
+/** Upper Confidence Bound acquisition function. */
+export function upperConfidenceBound(mean: number, variance: number, kappa = 2.0): number {
+ return mean + kappa * Math.sqrt(variance);
+}
+
+/** Bayesian optimization state. */
+export interface BOState {
+ observedX: Float64Array[];
+ observedY: Float64Array;
+ bestY: number;
+ bestX: Float64Array;
+ lengthScale: number;
+ amplitude: number;
+ noise: number;
+}
+
+/** Initialize Bayesian optimization state from initial observations. */
+export function initBOState(
+ X: Float64Array[],
+ y: Float64Array,
+ lengthScale = 1.0,
+ amplitude = 1.0,
+ noise = 0.01,
+): BOState {
+ let bestIdx = 0;
+ for (let i = 1; i < y.length; i++) {
+ if ((y[i] ?? Number.NEGATIVE_INFINITY) > (y[bestIdx] ?? Number.NEGATIVE_INFINITY)) bestIdx = i;
+ }
+ return {
+ observedX: X,
+ observedY: y,
+ bestY: y[bestIdx] ?? Number.NEGATIVE_INFINITY,
+ bestX: X[bestIdx] ?? new Float64Array(0),
+ lengthScale,
+ amplitude,
+ noise,
+ };
+}
+
+/** Suggest the next candidate by maximizing EI over provided candidates. */
+export function suggestNext(
+ state: BOState,
+ candidates: Float64Array[],
+ acquisition: "EI" | "PI" | "UCB" = "EI",
+ kappa = 2.0,
+ xi = 0.01,
+): { bestCandidate: Float64Array; acquisitionValues: Float64Array } {
+ const K = kernelMatrix(state.observedX, state.lengthScale, state.amplitude, state.noise);
+ const L = cholesky(K, state.observedX.length);
+
+ const acqValues = new Float64Array(candidates.length);
+ for (let i = 0; i < candidates.length; i++) {
+ const { mean, variance } = gpPredict(
+ candidates[i]!,
+ state.observedX,
+ state.observedY,
+ L,
+ state.lengthScale,
+ state.amplitude,
+ state.noise,
+ );
+ if (acquisition === "EI") {
+ acqValues[i] = expectedImprovement(mean, variance, state.bestY, xi);
+ } else if (acquisition === "PI") {
+ acqValues[i] = probabilityOfImprovement(mean, variance, state.bestY, xi);
+ } else {
+ acqValues[i] = upperConfidenceBound(mean, variance, kappa);
+ }
+ }
+
+ let bestIdx = 0;
+ for (let i = 1; i < acqValues.length; i++) {
+ if (
+ (acqValues[i] ?? Number.NEGATIVE_INFINITY) > (acqValues[bestIdx] ?? Number.NEGATIVE_INFINITY)
+ )
+ bestIdx = i;
+ }
+ return {
+ bestCandidate: candidates[bestIdx] ?? new Float64Array(0),
+ acquisitionValues: acqValues,
+ };
+}
diff --git a/src/ml/consistency_models.ts b/src/ml/consistency_models.ts
new file mode 100644
index 000000000..af53438d1
--- /dev/null
+++ b/src/ml/consistency_models.ts
@@ -0,0 +1,186 @@
+/**
+ * Consistency Models: fast single-step generative models.
+ *
+ * Implements the Consistency Model framework (Song et al., 2023) including:
+ * the consistency function, consistency training loss, and
+ * few-step sampling procedures.
+ *
+ * @module
+ */
+
+/** Consistency model configuration. */
+export interface ConsistencyConfig {
+ /** Total training time horizon T. */
+ T: number;
+ /** Initial discretization timestep (epsilon). */
+ epsilon: number;
+ /** Number of discretization steps N. */
+ N: number;
+ /** EMA decay rate for teacher model. */
+ emaDecay: number;
+ /** Karras noise schedule rho. */
+ rho: number;
+ /** Loss weighting parameter mu. */
+ mu: number;
+}
+
+/** Default consistency model configuration. */
+export const DEFAULT_CONSISTENCY_CONFIG: ConsistencyConfig = {
+ T: 80,
+ epsilon: 0.002,
+ N: 150,
+ emaDecay: 0.99,
+ rho: 7,
+ mu: 0.95,
+};
+
+/** Karras et al. noise schedule: sigma_i for step i in [0, N]. */
+export function karrasNoiseLevels(config: ConsistencyConfig): Float64Array {
+ const { T, epsilon, N, rho } = config;
+ const levels = new Float64Array(N + 1);
+ for (let i = 0; i <= N; i++) {
+ const frac = i / N;
+ levels[i] = (epsilon ** (1 / rho) + frac * (T ** (1 / rho) - epsilon ** (1 / rho))) ** rho;
+ }
+ return levels;
+}
+
+/** Skip function c_skip(sigma): scaling for input. */
+export function cSkip(sigma: number, sigmaData = 0.5): number {
+ return (sigmaData * sigmaData) / (sigma * sigma + sigmaData * sigmaData);
+}
+
+/** Output function c_out(sigma): scaling for network output. */
+export function cOut(sigma: number, sigmaData = 0.5): number {
+ return (sigma * sigmaData) / Math.sqrt(sigma * sigma + sigmaData * sigmaData);
+}
+
+/** Input scaling c_in(sigma). */
+export function cIn(sigma: number, sigmaData = 0.5): number {
+ return 1 / Math.sqrt(sigma * sigma + sigmaData * sigmaData);
+}
+
+/** Noise conditioning c_noise(sigma): log encoding. */
+export function cNoise(sigma: number): number {
+ return Math.log(sigma) / 4;
+}
+
+/**
+ * Consistency function: maps any (x_t, t) to a consistent estimate of x_0.
+ * Uses the preconditioning scheme from Karras et al.
+ *
+ * @param xt - Noisy sample at noise level sigma
+ * @param sigma - Current noise level
+ * @param fTheta - Network output F_theta(c_in * x, c_noise)
+ * @param sigmaData - Data standard deviation
+ */
+export function consistencyFunction(
+ xt: Float64Array,
+ sigma: number,
+ fTheta: Float64Array,
+ sigmaData = 0.5,
+): Float64Array {
+ const skip = cSkip(sigma, sigmaData);
+ const out = cOut(sigma, sigmaData);
+ const result = new Float64Array(xt.length);
+ for (let i = 0; i < xt.length; i++) {
+ result[i] = skip * (xt[i] ?? 0) + out * (fTheta[i] ?? 0);
+ }
+ return result;
+}
+
+/** Compute preconditioning-scaled input for the network. */
+export function preconditionInput(x: Float64Array, sigma: number, sigmaData = 0.5): Float64Array {
+ const scale = cIn(sigma, sigmaData);
+ const out = new Float64Array(x.length);
+ for (let i = 0; i < x.length; i++) out[i] = scale * (x[i] ?? 0);
+ return out;
+}
+
+/** Pseudo-Huber loss (smooth L1-like). */
+export function pseudoHuberLoss(a: Float64Array, b: Float64Array, c = 0.00054): number {
+ let total = 0;
+ for (let i = 0; i < a.length; i++) {
+ const d = (a[i] ?? 0) - (b[i] ?? 0);
+ total += Math.sqrt(d * d + c * c) - c;
+ }
+ return total / a.length;
+}
+
+/**
+ * Consistency Training (CT) loss for a single sample pair.
+ *
+ * @param studentOutput - Consistency function output at t_{n+1}
+ * @param teacherOutput - EMA consistency function output at t_n (stop gradient)
+ * @param sigmaN - Noise level at step n
+ * @param sigmaNp1 - Noise level at step n+1
+ */
+export function consistencyTrainingLoss(
+ studentOutput: Float64Array,
+ teacherOutput: Float64Array,
+ sigmaN: number,
+ sigmaNp1: number,
+): number {
+ const weight = 1 / Math.max(sigmaNp1 - sigmaN, 1e-8);
+ return weight * pseudoHuberLoss(studentOutput, teacherOutput);
+}
+
+/** Add noise to a sample: x_t = x_0 + sigma * epsilon. */
+export function addGaussianNoise(
+ x0: Float64Array,
+ sigma: number,
+ noise: Float64Array,
+): Float64Array {
+ const out = new Float64Array(x0.length);
+ for (let i = 0; i < x0.length; i++) {
+ out[i] = (x0[i] ?? 0) + sigma * (noise[i] ?? 0);
+ }
+ return out;
+}
+
+/** Single-step consistency model sampling from noise. */
+export function consistencySampleOneStep(
+ noise: Float64Array,
+ T: number,
+ consistencyFn: (x: Float64Array, sigma: number) => Float64Array,
+): Float64Array {
+ return consistencyFn(noise, T);
+}
+
+/** Multi-step consistency model sampling (improves quality). */
+export function consistencySampleMultiStep(
+ noise: Float64Array,
+ sigmaLevels: number[],
+ consistencyFn: (x: Float64Array, sigma: number) => Float64Array,
+ noiseFn: (shape: number) => Float64Array,
+): Float64Array {
+ if (sigmaLevels.length === 0) return noise;
+ let x = consistencyFn(noise, sigmaLevels[0] ?? 1);
+ for (let i = 1; i < sigmaLevels.length; i++) {
+ const sigma = sigmaLevels[i] ?? 0;
+ if (sigma <= 0) break;
+ // Add noise and re-denoise
+ const eps = noiseFn(x.length);
+ const xNoisy = addGaussianNoise(x, sigma, eps);
+ x = consistencyFn(xNoisy, sigma);
+ }
+ return x;
+}
+
+/** Update EMA teacher model weights. */
+export function updateEMAWeights(
+ teacherWeights: Float64Array,
+ studentWeights: Float64Array,
+ decay: number,
+): Float64Array {
+ const out = new Float64Array(teacherWeights.length);
+ for (let i = 0; i < teacherWeights.length; i++) {
+ out[i] = decay * (teacherWeights[i] ?? 0) + (1 - decay) * (studentWeights[i] ?? 0);
+ }
+ return out;
+}
+
+/** Compute adaptive EMA decay for consistency distillation. */
+export function adaptiveEMADecay(iteration: number, mu0 = 0.95, s0 = 10): number {
+ return Math.exp((s0 * Math.log(mu0)) / Math.max(iteration, 1));
+}
diff --git a/src/ml/crf.ts b/src/ml/crf.ts
new file mode 100644
index 000000000..3fddf9538
--- /dev/null
+++ b/src/ml/crf.ts
@@ -0,0 +1,165 @@
+/**
+ * Conditional Random Field (CRF) for sequence labeling.
+ *
+ * Implements a linear-chain CRF with the Viterbi algorithm for decoding,
+ * the forward algorithm for computing partition functions, and
+ * log-likelihood computation for training.
+ *
+ * @module
+ */
+
+/** CRF model parameters. */
+export interface CRFParams {
+ /** Emission scores [seqLen x numTags]. */
+ emissionScores: Float64Array;
+ /** Transition scores [numTags x numTags] (from -> to). */
+ transitionScores: Float64Array;
+ /** Start transition scores [numTags]. */
+ startScores: Float64Array;
+ /** End transition scores [numTags]. */
+ endScores: Float64Array;
+ numTags: number;
+ seqLen: number;
+}
+
+/** Viterbi decoding result. */
+export interface ViterbiResult {
+ /** Best tag sequence. */
+ tags: number[];
+ /** Score of the best sequence. */
+ score: number;
+}
+
+/** Viterbi algorithm for CRF decoding. */
+export function viterbiDecode(params: CRFParams): ViterbiResult {
+ const { numTags, seqLen, emissionScores, transitionScores, startScores, endScores } = params;
+
+ // viterbi[t][tag] = best score ending at (t, tag)
+ const viterbi: Float64Array[] = [];
+ const backpointer: Int32Array[] = [];
+
+ // Init: t=0
+ const v0 = new Float64Array(numTags);
+ const bp0 = new Int32Array(numTags).fill(-1);
+ for (let j = 0; j < numTags; j++) {
+ v0[j] =
+ (startScores[j] ?? Number.NEGATIVE_INFINITY) +
+ (emissionScores[j] ?? Number.NEGATIVE_INFINITY);
+ }
+ viterbi.push(v0);
+ backpointer.push(bp0);
+
+ for (let t = 1; t < seqLen; t++) {
+ const vt = new Float64Array(numTags);
+ const bpt = new Int32Array(numTags);
+ const vprev = viterbi[t - 1]!;
+ for (let j = 0; j < numTags; j++) {
+ let bestScore = Number.NEGATIVE_INFINITY;
+ let bestPrev = 0;
+ for (let i = 0; i < numTags; i++) {
+ const score =
+ (vprev[i] ?? Number.NEGATIVE_INFINITY) +
+ (transitionScores[i * numTags + j] ?? Number.NEGATIVE_INFINITY);
+ if (score > bestScore) {
+ bestScore = score;
+ bestPrev = i;
+ }
+ }
+ vt[j] = bestScore + (emissionScores[t * numTags + j] ?? Number.NEGATIVE_INFINITY);
+ bpt[j] = bestPrev;
+ }
+ viterbi.push(vt);
+ backpointer.push(bpt);
+ }
+
+ // Add end scores
+ const vlast = viterbi[seqLen - 1]!;
+ let bestFinalScore = Number.NEGATIVE_INFINITY;
+ let bestFinalTag = 0;
+ for (let j = 0; j < numTags; j++) {
+ const s = (vlast[j] ?? Number.NEGATIVE_INFINITY) + (endScores[j] ?? 0);
+ if (s > bestFinalScore) {
+ bestFinalScore = s;
+ bestFinalTag = j;
+ }
+ }
+
+ // Backtrack
+ const tags = new Array(seqLen);
+ tags[seqLen - 1] = bestFinalTag;
+ for (let t = seqLen - 1; t > 0; t--) {
+ tags[t - 1] = backpointer[t]![tags[t]!] ?? 0;
+ }
+
+ return { tags, score: bestFinalScore };
+}
+
+/** Forward algorithm: compute log partition function log Z. */
+export function forwardLogZ(params: CRFParams): number {
+ const { numTags, seqLen, emissionScores, transitionScores, startScores, endScores } = params;
+
+ // alpha[tag] = log sum of scores for all paths ending at (t, tag)
+ let alpha = new Float64Array(numTags);
+ for (let j = 0; j < numTags; j++) {
+ alpha[j] =
+ (startScores[j] ?? Number.NEGATIVE_INFINITY) +
+ (emissionScores[j] ?? Number.NEGATIVE_INFINITY);
+ }
+
+ for (let t = 1; t < seqLen; t++) {
+ const newAlpha = new Float64Array(numTags);
+ for (let j = 0; j < numTags; j++) {
+ const scores = new Float64Array(numTags);
+ for (let i = 0; i < numTags; i++) {
+ scores[i] =
+ (alpha[i] ?? Number.NEGATIVE_INFINITY) +
+ (transitionScores[i * numTags + j] ?? Number.NEGATIVE_INFINITY);
+ }
+ newAlpha[j] =
+ logSumExp(scores) + (emissionScores[t * numTags + j] ?? Number.NEGATIVE_INFINITY);
+ }
+ alpha = newAlpha;
+ }
+
+ // Add end scores
+ const final = new Float64Array(numTags);
+ for (let j = 0; j < numTags; j++) {
+ final[j] = (alpha[j] ?? Number.NEGATIVE_INFINITY) + (endScores[j] ?? 0);
+ }
+ return logSumExp(final);
+}
+
+/** Compute score of a given tag sequence. */
+export function sequenceScore(params: CRFParams, tags: number[]): number {
+ const { numTags, emissionScores, transitionScores, startScores, endScores } = params;
+ let score = startScores[tags[0] ?? 0] ?? Number.NEGATIVE_INFINITY;
+ score += emissionScores[tags[0] ?? 0] ?? Number.NEGATIVE_INFINITY;
+ for (let t = 1; t < tags.length; t++) {
+ const prev = tags[t - 1] ?? 0;
+ const curr = tags[t] ?? 0;
+ score += transitionScores[prev * numTags + curr] ?? Number.NEGATIVE_INFINITY;
+ score += emissionScores[t * numTags + curr] ?? Number.NEGATIVE_INFINITY;
+ }
+ score += endScores[tags[tags.length - 1] ?? 0] ?? 0;
+ return score;
+}
+
+/** CRF negative log-likelihood for a given tag sequence. */
+export function crfNegLogLikelihood(params: CRFParams, tags: number[]): number {
+ const goldScore = sequenceScore(params, tags);
+ const logZ = forwardLogZ(params);
+ return logZ - goldScore;
+}
+
+/** Log-sum-exp (numerically stable). */
+export function logSumExp(values: Float64Array): number {
+ let max = Number.NEGATIVE_INFINITY;
+ for (let i = 0; i < values.length; i++)
+ max = Math.max(max, values[i] ?? Number.NEGATIVE_INFINITY);
+ if (!Number.isFinite(max)) return Number.NEGATIVE_INFINITY;
+ let sum = 0;
+ for (let i = 0; i < values.length; i++) {
+ sum += Math.exp((values[i] ?? Number.NEGATIVE_INFINITY) - max);
+ }
+ return max + Math.log(sum);
+}
diff --git a/src/ml/ddim.ts b/src/ml/ddim.ts
new file mode 100644
index 000000000..3f4bbfa7c
--- /dev/null
+++ b/src/ml/ddim.ts
@@ -0,0 +1,158 @@
+/**
+ * DDIM (Denoising Diffusion Implicit Models) sampler.
+ *
+ * Implements the deterministic DDIM sampling procedure from Song et al. (2020).
+ * Supports both DDIM (η=0, deterministic) and DDPM (η=1, stochastic) sampling.
+ *
+ * @module
+ */
+
+/** Noise schedule type. */
+export type NoiseSchedule = "linear" | "cosine" | "sqrt";
+
+/** DDIM sampler configuration. */
+export interface DDIMConfig {
+ /** Total number of diffusion timesteps used during training. */
+ numTrainTimesteps: number;
+ /** Noise schedule type. */
+ schedule: NoiseSchedule;
+ /** Stochasticity parameter (0 = DDIM deterministic, 1 = DDPM stochastic). */
+ eta: number;
+ /** Beta start for linear schedule. */
+ betaStart: number;
+ /** Beta end for linear schedule. */
+ betaEnd: number;
+}
+
+/** Precomputed noise schedule values. */
+export interface NoiseScheduleValues {
+ /** Betas at each timestep. */
+ betas: Float64Array;
+ /** Alphas (1 - beta). */
+ alphas: Float64Array;
+ /** Cumulative product of alphas. */
+ alphasCumprod: Float64Array;
+ /** Square root of alphasCumprod. */
+ sqrtAlphasCumprod: Float64Array;
+ /** Square root of (1 - alphasCumprod). */
+ sqrtOneMinusAlphasCumprod: Float64Array;
+}
+
+/** Compute noise schedule from config. */
+export function computeNoiseSchedule(config: DDIMConfig): NoiseScheduleValues {
+ const T = config.numTrainTimesteps;
+ if (!Number.isInteger(T) || T < 1) {
+ throw new RangeError("numTrainTimesteps must be a positive integer");
+ }
+ if (
+ config.schedule !== "cosine" &&
+ (!Number.isFinite(config.betaStart) ||
+ !Number.isFinite(config.betaEnd) ||
+ config.betaStart < 0 ||
+ config.betaStart >= 1 ||
+ config.betaEnd < 0 ||
+ config.betaEnd >= 1)
+ ) {
+ throw new RangeError("betaStart and betaEnd must be finite values in [0, 1)");
+ }
+ const betas = new Float64Array(T);
+ const alphas = new Float64Array(T);
+ const alphasCumprod = new Float64Array(T);
+ const sqrtAlphasCumprod = new Float64Array(T);
+ const sqrtOneMinusAlphasCumprod = new Float64Array(T);
+
+ if (config.schedule === "linear") {
+ const step = T === 1 ? 0 : (config.betaEnd - config.betaStart) / (T - 1);
+ for (let t = 0; t < T; t++) {
+ betas[t] = config.betaStart + t * step;
+ }
+ } else if (config.schedule === "cosine") {
+ const s = 0.008;
+ let previous = Math.cos((s / (1 + s)) * (Math.PI / 2)) ** 2;
+ for (let t = 0; t < T; t++) {
+ const ft = Math.cos((((t + 1) / T + s) / (1 + s)) * (Math.PI / 2)) ** 2;
+ // Each beta is the loss of signal relative to the previous timestep.
+ betas[t] = Math.min(1 - ft / previous, 0.999);
+ previous = ft;
+ }
+ } else {
+ // sqrt schedule
+ for (let t = 0; t < T; t++) {
+ const frac = T === 1 ? 0 : t / (T - 1);
+ betas[t] = config.betaStart + (config.betaEnd - config.betaStart) * Math.sqrt(frac);
+ }
+ }
+
+ let cumprod = 1.0;
+ for (let t = 0; t < T; t++) {
+ alphas[t] = 1 - (betas[t] ?? 0);
+ cumprod *= alphas[t] ?? 1;
+ alphasCumprod[t] = cumprod;
+ sqrtAlphasCumprod[t] = Math.sqrt(cumprod);
+ sqrtOneMinusAlphasCumprod[t] = Math.sqrt(1 - cumprod);
+ }
+
+ return { betas, alphas, alphasCumprod, sqrtAlphasCumprod, sqrtOneMinusAlphasCumprod };
+}
+
+/** Add noise to a sample at a given timestep. */
+export function addNoise(
+ sample: Float64Array,
+ noise: Float64Array,
+ timestep: number,
+ schedule: NoiseScheduleValues,
+): Float64Array {
+ const sqrtAcp = schedule.sqrtAlphasCumprod[timestep] ?? 0;
+ const sqrtOm = schedule.sqrtOneMinusAlphasCumprod[timestep] ?? 0;
+ const noisy = new Float64Array(sample.length);
+ for (let i = 0; i < sample.length; i++) {
+ noisy[i] = sqrtAcp * (sample[i] ?? 0) + sqrtOm * (noise[i] ?? 0);
+ }
+ return noisy;
+}
+
+/** DDIM step: predict x0 from noisy xt and noise prediction. */
+export function ddimStep(
+ xt: Float64Array,
+ noisePred: Float64Array,
+ tPrev: number,
+ tCurr: number,
+ schedule: NoiseScheduleValues,
+ eta: number,
+): Float64Array {
+ const acpCurr = schedule.alphasCumprod[tCurr] ?? 0;
+ const acpPrev = tPrev >= 0 ? (schedule.alphasCumprod[tPrev] ?? 0) : 1.0;
+ const sqrtAcpCurr = Math.sqrt(acpCurr);
+ const sqrtOneMinusAcpCurr = Math.sqrt(1 - acpCurr);
+ const sqrtAcpPrev = Math.sqrt(acpPrev);
+
+ const sigmaT = eta * Math.sqrt(((1 - acpPrev) / (1 - acpCurr)) * (1 - acpCurr / acpPrev));
+
+ const result = new Float64Array(xt.length);
+ for (let i = 0; i < xt.length; i++) {
+ const xi = xt[i] ?? 0;
+ const ni = noisePred[i] ?? 0;
+ // Predict x0
+ const x0Pred = (xi - sqrtOneMinusAcpCurr * ni) / (sqrtAcpCurr || 1e-8);
+ // Direction pointing to xt
+ const dirXt = Math.sqrt(Math.max(0, 1 - acpPrev - sigmaT * sigmaT)) * ni;
+ result[i] = sqrtAcpPrev * x0Pred + dirXt;
+ }
+ return result;
+}
+
+/** Generate a sequence of timesteps for DDIM inference. */
+export function ddimTimesteps(numTrainTimesteps: number, numInferenceSteps: number): number[] {
+ const step = Math.floor(numTrainTimesteps / numInferenceSteps);
+ const timesteps: number[] = [];
+ for (let i = numInferenceSteps - 1; i >= 0; i--) {
+ timesteps.push(i * step);
+ }
+ return timesteps;
+}
+
+/** Compute signal-to-noise ratio at timestep t. */
+export function snrAtTimestep(t: number, schedule: NoiseScheduleValues): number {
+ const acp = schedule.alphasCumprod[t] ?? 0;
+ return acp / (1 - acp + 1e-8);
+}
diff --git a/src/ml/ft_transformer.ts b/src/ml/ft_transformer.ts
new file mode 100644
index 000000000..d11338085
--- /dev/null
+++ b/src/ml/ft_transformer.ts
@@ -0,0 +1,205 @@
+/**
+ * Feature Tokenizer + Transformer (FT-Transformer) for tabular data.
+ *
+ * Implements the FT-Transformer architecture from Gorishniy et al. (2021):
+ * numerical and categorical feature tokenization, multi-head self-attention,
+ * feed-forward blocks, and a [CLS] token for classification/regression.
+ *
+ * @module
+ */
+
+/** Multi-head self-attention output. */
+export interface AttentionOutput {
+ /** Context vectors [seqLen x dModel]. */
+ context: Float64Array;
+ /** Attention weights [numHeads x seqLen x seqLen]. */
+ weights: Float64Array;
+}
+
+/** Scaled dot-product attention for a single head. */
+export function scaledDotProductAttention(
+ Q: Float64Array,
+ K: Float64Array,
+ V: Float64Array,
+ seqLen: number,
+ dHead: number,
+): { context: Float64Array; weights: Float64Array } {
+ const scale = Math.sqrt(dHead);
+ // Scores [seqLen x seqLen]
+ const scores = new Float64Array(seqLen * seqLen);
+ for (let i = 0; i < seqLen; i++) {
+ for (let j = 0; j < seqLen; j++) {
+ let dot = 0;
+ for (let d = 0; d < dHead; d++) {
+ dot += (Q[i * dHead + d] ?? 0) * (K[j * dHead + d] ?? 0);
+ }
+ scores[i * seqLen + j] = dot / scale;
+ }
+ }
+ // Softmax over rows
+ const weights = new Float64Array(seqLen * seqLen);
+ for (let i = 0; i < seqLen; i++) {
+ let maxScore = Number.NEGATIVE_INFINITY;
+ for (let j = 0; j < seqLen; j++) maxScore = Math.max(maxScore, scores[i * seqLen + j] ?? 0);
+ let sumExp = 0;
+ for (let j = 0; j < seqLen; j++) {
+ const e = Math.exp((scores[i * seqLen + j] ?? 0) - maxScore);
+ weights[i * seqLen + j] = e;
+ sumExp += e;
+ }
+ for (let j = 0; j < seqLen; j++) {
+ weights[i * seqLen + j] = (weights[i * seqLen + j] ?? 0) / (sumExp || 1);
+ }
+ }
+ // Context [seqLen x dHead]
+ const context = new Float64Array(seqLen * dHead);
+ for (let i = 0; i < seqLen; i++) {
+ for (let d = 0; d < dHead; d++) {
+ let val = 0;
+ for (let j = 0; j < seqLen; j++) {
+ val += (weights[i * seqLen + j] ?? 0) * (V[j * dHead + d] ?? 0);
+ }
+ context[i * dHead + d] = val;
+ }
+ }
+ return { context, weights };
+}
+
+/** Layer normalization. */
+export function layerNorm(
+ x: Float64Array,
+ gamma: Float64Array,
+ beta: Float64Array,
+ eps = 1e-5,
+): Float64Array {
+ let mean = 0;
+ for (let i = 0; i < x.length; i++) mean += x[i] ?? 0;
+ mean /= x.length;
+ let variance = 0;
+ for (let i = 0; i < x.length; i++) {
+ const d = (x[i] ?? 0) - mean;
+ variance += d * d;
+ }
+ variance /= x.length;
+ const std = Math.sqrt(variance + eps);
+ const out = new Float64Array(x.length);
+ for (let i = 0; i < x.length; i++) {
+ out[i] = (((x[i] ?? 0) - mean) / std) * (gamma[i] ?? 1) + (beta[i] ?? 0);
+ }
+ return out;
+}
+
+/** GELU activation (approximate). */
+export function gelu(x: number): number {
+ return 0.5 * x * (1 + Math.tanh(Math.sqrt(2 / Math.PI) * (x + 0.044715 * x * x * x)));
+}
+
+/** Feed-forward block: Linear -> GELU -> Linear. */
+export function feedForward(
+ x: Float64Array,
+ W1: Float64Array,
+ b1: Float64Array,
+ W2: Float64Array,
+ b2: Float64Array,
+ dModel: number,
+ dFF: number,
+): Float64Array {
+ // First layer
+ const h = new Float64Array(dFF);
+ for (let j = 0; j < dFF; j++) {
+ let sum = b1[j] ?? 0;
+ for (let i = 0; i < dModel; i++) {
+ sum += (W1[j * dModel + i] ?? 0) * (x[i] ?? 0);
+ }
+ h[j] = gelu(sum);
+ }
+ // Second layer
+ const out = new Float64Array(dModel);
+ for (let i = 0; i < dModel; i++) {
+ let sum = b2[i] ?? 0;
+ for (let j = 0; j < dFF; j++) {
+ sum += (W2[i * dFF + j] ?? 0) * (h[j] ?? 0);
+ }
+ out[i] = sum;
+ }
+ return out;
+}
+
+/** Tokenize a numerical feature: embed to dModel dimensions. */
+export function tokenizeNumerical(
+ value: number,
+ embedding: Float64Array,
+ bias: Float64Array,
+): Float64Array {
+ const out = new Float64Array(embedding.length);
+ for (let i = 0; i < embedding.length; i++) {
+ out[i] = value * (embedding[i] ?? 0) + (bias[i] ?? 0);
+ }
+ return out;
+}
+
+/** Tokenize a categorical feature via embedding lookup. */
+export function tokenizeCategorical(
+ categoryIndex: number,
+ embeddingMatrix: Float64Array,
+ numCategories: number,
+ dModel: number,
+): Float64Array {
+ const idx = Math.max(0, Math.min(categoryIndex, numCategories - 1));
+ const out = new Float64Array(dModel);
+ for (let i = 0; i < dModel; i++) {
+ out[i] = embeddingMatrix[idx * dModel + i] ?? 0;
+ }
+ return out;
+}
+
+/** Add residual connection. */
+export function addResidual(x: Float64Array, residual: Float64Array): Float64Array {
+ const out = new Float64Array(x.length);
+ for (let i = 0; i < x.length; i++) {
+ out[i] = (x[i] ?? 0) + (residual[i] ?? 0);
+ }
+ return out;
+}
+
+/** Extract [CLS] token representation (first token). */
+export function extractCLSToken(tokens: Float64Array, dModel: number): Float64Array {
+ const cls = new Float64Array(dModel);
+ for (let i = 0; i < dModel; i++) {
+ cls[i] = tokens[i] ?? 0;
+ }
+ return cls;
+}
+
+/** Simple linear classifier head. */
+export function linearHead(
+ x: Float64Array,
+ W: Float64Array,
+ b: Float64Array,
+ numClasses: number,
+ dModel: number,
+): Float64Array {
+ const logits = new Float64Array(numClasses);
+ for (let c = 0; c < numClasses; c++) {
+ let sum = b[c] ?? 0;
+ for (let i = 0; i < dModel; i++) {
+ sum += (W[c * dModel + i] ?? 0) * (x[i] ?? 0);
+ }
+ logits[c] = sum;
+ }
+ return logits;
+}
+
+/** Softmax activation. */
+export function softmax(logits: Float64Array): Float64Array {
+ let maxVal = Number.NEGATIVE_INFINITY;
+ for (let i = 0; i < logits.length; i++) maxVal = Math.max(maxVal, logits[i] ?? 0);
+ let sumExp = 0;
+ const out = new Float64Array(logits.length);
+ for (let i = 0; i < logits.length; i++) {
+ out[i] = Math.exp((logits[i] ?? 0) - maxVal);
+ sumExp += out[i] ?? 0;
+ }
+ for (let i = 0; i < out.length; i++) out[i] = (out[i] ?? 0) / (sumExp || 1);
+ return out;
+}
diff --git a/src/ml/gradient_boosting.ts b/src/ml/gradient_boosting.ts
new file mode 100644
index 000000000..e96700e8e
--- /dev/null
+++ b/src/ml/gradient_boosting.ts
@@ -0,0 +1,216 @@
+/**
+ * Gradient Boosting from scratch.
+ *
+ * Implements gradient boosted regression trees (GBRT) with:
+ * - Mean squared error and log-loss objectives
+ * - CART decision trees as base learners
+ * - Shrinkage (learning rate), subsampling, and max depth
+ *
+ * @module
+ */
+
+/** A leaf or internal node in a decision tree. */
+export type TreeNode =
+ | { kind: "leaf"; value: number }
+ | { kind: "internal"; featureIndex: number; threshold: number; left: TreeNode; right: TreeNode };
+
+/** Decision tree parameters. */
+export interface TreeParams {
+ maxDepth: number;
+ minSamplesLeaf: number;
+}
+
+/** Predict with a decision tree. */
+export function treePredictOne(node: TreeNode, x: Float64Array): number {
+ if (node.kind === "leaf") return node.value;
+ const val = x[node.featureIndex] ?? 0;
+ return val <= node.threshold ? treePredictOne(node.left, x) : treePredictOne(node.right, x);
+}
+
+/** Predict for an array of samples. */
+export function treePredict(node: TreeNode, X: Float64Array[]): Float64Array {
+ return Float64Array.from(X.map((x) => treePredictOne(node, x)));
+}
+
+/** Compute mean of an array of values at given indices. */
+function meanAt(y: Float64Array, indices: number[]): number {
+ if (indices.length === 0) return 0;
+ let s = 0;
+ for (const i of indices) s += y[i] ?? 0;
+ return s / indices.length;
+}
+
+/** Compute MSE of an array at given indices. */
+function mseAt(y: Float64Array, indices: number[], mean: number): number {
+ let s = 0;
+ for (const i of indices) {
+ const d = (y[i] ?? 0) - mean;
+ s += d * d;
+ }
+ return s;
+}
+
+/** Find best split for a set of indices. */
+function findBestSplit(
+ X: Float64Array[],
+ y: Float64Array,
+ indices: number[],
+ numFeatures: number,
+): { featureIndex: number; threshold: number; gain: number } | null {
+ const parentMean = meanAt(y, indices);
+ const parentMSE = mseAt(y, indices, parentMean);
+ let bestGain = Number.NEGATIVE_INFINITY;
+ let bestFeature = 0;
+ let bestThreshold = 0;
+
+ for (let f = 0; f < numFeatures; f++) {
+ // Sort indices by feature value
+ const sorted = [...indices].sort((a, b) => (X[a]![f] ?? 0) - (X[b]![f] ?? 0));
+ for (let k = 1; k < sorted.length; k++) {
+ const threshold = ((X[sorted[k - 1]!]![f] ?? 0) + (X[sorted[k]!]![f] ?? 0)) / 2;
+ const left = sorted.slice(0, k);
+ const right = sorted.slice(k);
+ const lMean = meanAt(y, left);
+ const rMean = meanAt(y, right);
+ const gain = parentMSE - mseAt(y, left, lMean) - mseAt(y, right, rMean);
+ if (gain > bestGain) {
+ bestGain = gain;
+ bestFeature = f;
+ bestThreshold = threshold;
+ }
+ }
+ }
+
+ return bestGain > 0
+ ? { featureIndex: bestFeature, threshold: bestThreshold, gain: bestGain }
+ : null;
+}
+
+/** Build a regression tree on gradient residuals. */
+export function buildTree(
+ X: Float64Array[],
+ y: Float64Array,
+ indices: number[],
+ depth: number,
+ params: TreeParams,
+): TreeNode {
+ const leafVal = meanAt(y, indices);
+
+ if (depth >= params.maxDepth || indices.length <= params.minSamplesLeaf * 2) {
+ return { kind: "leaf", value: leafVal };
+ }
+
+ const numFeatures = X[0]?.length ?? 0;
+ const split = findBestSplit(X, y, indices, numFeatures);
+
+ if (split === null) {
+ return { kind: "leaf", value: leafVal };
+ }
+
+ const left: number[] = [];
+ const right: number[] = [];
+ for (const i of indices) {
+ if ((X[i]![split.featureIndex] ?? 0) <= split.threshold) {
+ left.push(i);
+ } else {
+ right.push(i);
+ }
+ }
+
+ if (left.length < params.minSamplesLeaf || right.length < params.minSamplesLeaf) {
+ return { kind: "leaf", value: leafVal };
+ }
+
+ return {
+ kind: "internal",
+ featureIndex: split.featureIndex,
+ threshold: split.threshold,
+ left: buildTree(X, y, left, depth + 1, params),
+ right: buildTree(X, y, right, depth + 1, params),
+ };
+}
+
+/** Gradient boosting ensemble. */
+export interface GBMEnsemble {
+ trees: TreeNode[];
+ learningRate: number;
+ initialPrediction: number;
+}
+
+/** MSE negative gradient (residuals). */
+export function mseGradient(y: Float64Array, yPred: Float64Array): Float64Array {
+ const g = new Float64Array(y.length);
+ for (let i = 0; i < y.length; i++) g[i] = (y[i] ?? 0) - (yPred[i] ?? 0);
+ return g;
+}
+
+/** Train gradient boosted regression trees. */
+export function fitGBM(
+ X: Float64Array[],
+ y: Float64Array,
+ nEstimators = 100,
+ learningRate = 0.1,
+ maxDepth = 3,
+ minSamplesLeaf = 1,
+): GBMEnsemble {
+ const n = y.length;
+ let sumY = 0;
+ for (let i = 0; i < n; i++) sumY += y[i] ?? 0;
+ const initialPrediction = sumY / (n || 1);
+
+ const yPred = new Float64Array(n).fill(initialPrediction);
+ const trees: TreeNode[] = [];
+ const params: TreeParams = { maxDepth, minSamplesLeaf };
+ const indices = Array.from({ length: n }, (_, i) => i);
+
+ for (let m = 0; m < nEstimators; m++) {
+ const residuals = mseGradient(y, yPred);
+ const tree = buildTree(X, residuals, indices, 0, params);
+ trees.push(tree);
+ const treePreds = treePredict(tree, X);
+ for (let i = 0; i < n; i++) {
+ yPred[i] = (yPred[i] ?? 0) + learningRate * (treePreds[i] ?? 0);
+ }
+ }
+
+ return { trees, learningRate, initialPrediction };
+}
+
+/** Predict with a trained GBM ensemble. */
+export function predictGBM(ensemble: GBMEnsemble, X: Float64Array[]): Float64Array {
+ const n = X.length;
+ const yPred = new Float64Array(n).fill(ensemble.initialPrediction);
+ for (const tree of ensemble.trees) {
+ const preds = treePredict(tree, X);
+ for (let i = 0; i < n; i++) {
+ yPred[i] = (yPred[i] ?? 0) + ensemble.learningRate * (preds[i] ?? 0);
+ }
+ }
+ return yPred;
+}
+
+/** Compute mean squared error. */
+export function mse(y: Float64Array, yPred: Float64Array): number {
+ let s = 0;
+ for (let i = 0; i < y.length; i++) {
+ const d = (y[i] ?? 0) - (yPred[i] ?? 0);
+ s += d * d;
+ }
+ return s / (y.length || 1);
+}
+
+/** Compute R² score. */
+export function r2Score(y: Float64Array, yPred: Float64Array): number {
+ let sumY = 0;
+ for (let i = 0; i < y.length; i++) sumY += y[i] ?? 0;
+ const meanY = sumY / (y.length || 1);
+ let ssTot = 0;
+ let ssRes = 0;
+ for (let i = 0; i < y.length; i++) {
+ const d1 = (y[i] ?? 0) - meanY;
+ const d2 = (y[i] ?? 0) - (yPred[i] ?? 0);
+ ssTot += d1 * d1;
+ ssRes += d2 * d2;
+ }
+ return 1 - ssRes / (ssTot || 1);
+}
diff --git a/src/ml/neural_network.ts b/src/ml/neural_network.ts
new file mode 100644
index 000000000..5d81f4438
--- /dev/null
+++ b/src/ml/neural_network.ts
@@ -0,0 +1,204 @@
+/**
+ * Neural Network layers and training utilities from scratch.
+ *
+ * Implements dense layers, activations, loss functions, and
+ * mini-batch SGD/Adam optimizer for feed-forward networks.
+ *
+ * @module
+ */
+
+/** Dense layer forward pass: y = W x + b. */
+export function denseForward(
+ x: Float64Array,
+ W: Float64Array,
+ b: Float64Array,
+ outDim: number,
+ inDim: number,
+): Float64Array {
+ const y = new Float64Array(outDim);
+ for (let i = 0; i < outDim; i++) {
+ let sum = b[i] ?? 0;
+ for (let j = 0; j < inDim; j++) {
+ sum += (W[i * inDim + j] ?? 0) * (x[j] ?? 0);
+ }
+ y[i] = sum;
+ }
+ return y;
+}
+
+/** Dense layer backward pass: returns dL/dx, dL/dW, dL/db. */
+export function denseBackward(
+ x: Float64Array,
+ W: Float64Array,
+ dOut: Float64Array,
+ outDim: number,
+ inDim: number,
+): { dX: Float64Array; dW: Float64Array; dB: Float64Array } {
+ const dX = new Float64Array(inDim);
+ const dW = new Float64Array(outDim * inDim);
+ const dB = new Float64Array(outDim);
+
+ for (let i = 0; i < outDim; i++) {
+ dB[i] = dOut[i] ?? 0;
+ for (let j = 0; j < inDim; j++) {
+ dW[i * inDim + j] = (dOut[i] ?? 0) * (x[j] ?? 0);
+ dX[j] = (dX[j] ?? 0) + (dOut[i] ?? 0) * (W[i * inDim + j] ?? 0);
+ }
+ }
+ return { dX, dW, dB };
+}
+
+/** ReLU activation and its gradient. */
+export function relu(x: Float64Array): Float64Array {
+ const out = new Float64Array(x.length);
+ for (let i = 0; i < x.length; i++) out[i] = Math.max(0, x[i] ?? 0);
+ return out;
+}
+
+export function reluGrad(x: Float64Array, dOut: Float64Array): Float64Array {
+ const dX = new Float64Array(x.length);
+ for (let i = 0; i < x.length; i++) dX[i] = (x[i] ?? 0) > 0 ? (dOut[i] ?? 0) : 0;
+ return dX;
+}
+
+/** Sigmoid activation and its gradient. */
+export function sigmoidActivation(x: Float64Array): Float64Array {
+ const out = new Float64Array(x.length);
+ for (let i = 0; i < x.length; i++) out[i] = 1 / (1 + Math.exp(-(x[i] ?? 0)));
+ return out;
+}
+
+export function sigmoidGrad(out: Float64Array, dOut: Float64Array): Float64Array {
+ const dX = new Float64Array(out.length);
+ for (let i = 0; i < out.length; i++) {
+ const s = out[i] ?? 0;
+ dX[i] = s * (1 - s) * (dOut[i] ?? 0);
+ }
+ return dX;
+}
+
+/** Softmax and cross-entropy loss (combined for numerical stability). */
+export function softmaxCrossEntropy(
+ logits: Float64Array,
+ labels: Float64Array,
+): { loss: number; dLogits: Float64Array } {
+ const n = logits.length;
+ let maxLogit = Number.NEGATIVE_INFINITY;
+ for (let i = 0; i < n; i++) maxLogit = Math.max(maxLogit, logits[i] ?? Number.NEGATIVE_INFINITY);
+
+ let sumExp = 0;
+ const probs = new Float64Array(n);
+ for (let i = 0; i < n; i++) {
+ probs[i] = Math.exp((logits[i] ?? 0) - maxLogit);
+ sumExp += probs[i] ?? 0;
+ }
+ let loss = 0;
+ for (let i = 0; i < n; i++) {
+ probs[i] = (probs[i] ?? 0) / (sumExp || 1);
+ const label = labels[i] ?? 0;
+ if (label > 0) loss -= label * Math.log(Math.max(probs[i] ?? 0, 1e-12));
+ }
+
+ const dLogits = new Float64Array(n);
+ for (let i = 0; i < n; i++) dLogits[i] = (probs[i] ?? 0) - (labels[i] ?? 0);
+ return { loss, dLogits };
+}
+
+/** MSE loss and gradient. */
+export function mseLoss(
+ pred: Float64Array,
+ target: Float64Array,
+): { loss: number; dPred: Float64Array } {
+ const dPred = new Float64Array(pred.length);
+ let loss = 0;
+ for (let i = 0; i < pred.length; i++) {
+ const d = (pred[i] ?? 0) - (target[i] ?? 0);
+ loss += d * d;
+ dPred[i] = (2 * d) / pred.length;
+ }
+ return { loss: loss / pred.length, dPred };
+}
+
+/** Adam optimizer state. */
+export interface AdamState {
+ m: Float64Array;
+ v: Float64Array;
+ t: number;
+ lr: number;
+ beta1: number;
+ beta2: number;
+ eps: number;
+}
+
+/** Initialize Adam optimizer state for a parameter vector. */
+export function initAdam(
+ paramSize: number,
+ lr = 1e-3,
+ beta1 = 0.9,
+ beta2 = 0.999,
+ eps = 1e-8,
+): AdamState {
+ return {
+ m: new Float64Array(paramSize),
+ v: new Float64Array(paramSize),
+ t: 0,
+ lr,
+ beta1,
+ beta2,
+ eps,
+ };
+}
+
+/** Adam update step: modifies params in-place, returns updated AdamState. */
+export function adamStep(params: Float64Array, grads: Float64Array, state: AdamState): AdamState {
+ const { lr, beta1, beta2, eps } = state;
+ const t = state.t + 1;
+ const m = new Float64Array(params.length);
+ const v = new Float64Array(params.length);
+
+ for (let i = 0; i < params.length; i++) {
+ const g = grads[i] ?? 0;
+ m[i] = beta1 * (state.m[i] ?? 0) + (1 - beta1) * g;
+ v[i] = beta2 * (state.v[i] ?? 0) + (1 - beta2) * g * g;
+ const mHat = (m[i] ?? 0) / (1 - beta1 ** t);
+ const vHat = (v[i] ?? 0) / (1 - beta2 ** t);
+ params[i] = (params[i] ?? 0) - (lr * mHat) / (Math.sqrt(vHat) + eps);
+ }
+
+ return { ...state, m, v, t };
+}
+
+/** He initialization for ReLU networks. */
+export function heInit(inDim: number, outDim: number, rngSeed = 42): Float64Array {
+ const weights = new Float64Array(outDim * inDim);
+ let state = rngSeed >>> 0;
+ const std = Math.sqrt(2 / inDim);
+ for (let i = 0; i < weights.length; i++) {
+ // Box-Muller transform for Gaussian noise
+ state = (Math.imul(1664525, state) + 1013904223) >>> 0;
+ const u1 = state / 4294967296 || 1e-10;
+ state = (Math.imul(1664525, state) + 1013904223) >>> 0;
+ const u2 = state / 4294967296;
+ weights[i] = std * Math.sqrt(-2 * Math.log(u1)) * Math.cos(2 * Math.PI * u2);
+ }
+ return weights;
+}
+
+/** Dropout mask (for training). */
+export function dropoutMask(size: number, rate: number, rngSeed: number): Float64Array {
+ const mask = new Float64Array(size);
+ let state = rngSeed >>> 0;
+ const scale = 1 / (1 - rate);
+ for (let i = 0; i < size; i++) {
+ state = (Math.imul(1664525, state) + 1013904223) >>> 0;
+ mask[i] = state / 4294967296 > rate ? scale : 0;
+ }
+ return mask;
+}
+
+/** Apply dropout mask to activations. */
+export function applyDropout(x: Float64Array, mask: Float64Array): Float64Array {
+ const out = new Float64Array(x.length);
+ for (let i = 0; i < x.length; i++) out[i] = (x[i] ?? 0) * (mask[i] ?? 0);
+ return out;
+}
diff --git a/src/ml/random_forest.ts b/src/ml/random_forest.ts
new file mode 100644
index 000000000..bbc1b31d5
--- /dev/null
+++ b/src/ml/random_forest.ts
@@ -0,0 +1,179 @@
+/**
+ * Random Forest implementation from scratch.
+ *
+ * Implements bootstrap aggregation (bagging) of CART decision trees
+ * with random feature subsampling at each split (random forest).
+ * Supports both classification and regression.
+ *
+ * @module
+ */
+
+import { type TreeNode, type TreeParams, buildTree, treePredict } from "./gradient_boosting.ts";
+
+/** Random Forest configuration. */
+export interface RandomForestConfig {
+ nEstimators: number;
+ maxDepth: number;
+ minSamplesLeaf: number;
+ /** Number of features to consider at each split (0 = sqrt(n_features)). */
+ maxFeatures: number;
+ /** Bootstrap sample size (0 = n_samples). */
+ sampleSize: number;
+ seed: number;
+}
+
+/** LCG pseudo-random number generator. */
+export class LCGRandom {
+ private state: number;
+ constructor(seed: number) {
+ this.state = seed >>> 0;
+ }
+ next(): number {
+ this.state = (Math.imul(1664525, this.state) + 1013904223) >>> 0;
+ return this.state / 4294967296;
+ }
+ nextInt(n: number): number {
+ return Math.floor(this.next() * n);
+ }
+}
+
+/** Bootstrap sample indices. */
+function bootstrapSample(n: number, size: number, rng: LCGRandom): number[] {
+ const indices: number[] = [];
+ for (let i = 0; i < size; i++) indices.push(rng.nextInt(n));
+ return indices;
+}
+
+function selectFeatures(numFeatures: number, maxFeatures: number, rng: LCGRandom): number[] {
+ const k =
+ maxFeatures <= 0
+ ? Math.max(1, Math.floor(Math.sqrt(numFeatures)))
+ : Math.min(maxFeatures, numFeatures);
+ const all = Array.from({ length: numFeatures }, (_, i) => i);
+ // Fisher-Yates shuffle and take first k
+ for (let i = numFeatures - 1; i > 0; i--) {
+ const j = rng.nextInt(i + 1);
+ [all[i], all[j]] = [all[j]!, all[i]!];
+ }
+ return all.slice(0, k);
+}
+
+/** Random Forest model. */
+export interface RandomForestModel {
+ trees: TreeNode[];
+ featureMaps: number[][];
+ config: RandomForestConfig;
+ numFeatures: number;
+}
+
+/** Train a Random Forest regressor. */
+export function fitRandomForest(
+ X: Float64Array[],
+ y: Float64Array,
+ config: Partial = {},
+): RandomForestModel {
+ const cfg: RandomForestConfig = {
+ nEstimators: 100,
+ maxDepth: 5,
+ minSamplesLeaf: 1,
+ maxFeatures: 0,
+ sampleSize: 0,
+ seed: 42,
+ ...config,
+ };
+
+ const n = X.length;
+ const numFeatures = X[0]?.length ?? 0;
+ const sampleSize = cfg.sampleSize <= 0 ? n : cfg.sampleSize;
+ const rng = new LCGRandom(cfg.seed);
+ const params: TreeParams = { maxDepth: cfg.maxDepth, minSamplesLeaf: cfg.minSamplesLeaf };
+
+ const trees: TreeNode[] = [];
+ const featureMaps: number[][] = [];
+
+ for (let t = 0; t < cfg.nEstimators; t++) {
+ const indices = bootstrapSample(n, sampleSize, rng);
+ const k =
+ cfg.maxFeatures <= 0
+ ? Math.max(1, Math.floor(Math.sqrt(numFeatures)))
+ : Math.min(cfg.maxFeatures, numFeatures);
+ const featureSubset = selectFeatures(numFeatures, k, rng);
+ featureMaps.push(featureSubset);
+
+ // Build tree on bootstrap sample with projected features
+ const subX: Float64Array[] = [];
+ const subY = new Float64Array(indices.length);
+ for (let i = 0; i < indices.length; i++) {
+ const xi = X[indices[i]!]!;
+ const projected = new Float64Array(featureSubset.length);
+ for (let j = 0; j < featureSubset.length; j++) {
+ projected[j] = xi[featureSubset[j]!] ?? 0;
+ }
+ subX.push(projected);
+ subY[i] = y[indices[i]!] ?? 0;
+ }
+ const allIdx = Array.from({ length: indices.length }, (_, i) => i);
+ trees.push(buildTree(subX, subY, allIdx, 0, params));
+ }
+
+ return { trees, featureMaps, config: cfg, numFeatures };
+}
+
+/** Predict with a Random Forest (regression: average of trees). */
+export function predictRandomForest(model: RandomForestModel, X: Float64Array[]): Float64Array {
+ const n = X.length;
+ const sum = new Float64Array(n);
+
+ for (let t = 0; t < model.trees.length; t++) {
+ const featureMap = model.featureMaps[t]!;
+ const projX: Float64Array[] = X.map((xi) => {
+ const p = new Float64Array(featureMap.length);
+ for (let j = 0; j < featureMap.length; j++) p[j] = xi[featureMap[j]!] ?? 0;
+ return p;
+ });
+ const preds = treePredict(model.trees[t]!, projX);
+ for (let i = 0; i < n; i++) sum[i] = (sum[i] ?? 0) + (preds[i] ?? 0);
+ }
+
+ const out = new Float64Array(n);
+ for (let i = 0; i < n; i++) out[i] = (sum[i] ?? 0) / (model.trees.length || 1);
+ return out;
+}
+
+/** Out-of-bag (OOB) error estimate for regression. */
+export function oobError(
+ model: RandomForestModel,
+ X: Float64Array[],
+ y: Float64Array,
+ seed: number,
+): number {
+ const n = X.length;
+ const rng = new LCGRandom(seed);
+ const sampleSize = model.config.sampleSize <= 0 ? n : model.config.sampleSize;
+ const oobPreds = new Float64Array(n);
+ const oobCounts = new Int32Array(n);
+
+ for (let t = 0; t < model.trees.length; t++) {
+ const bootstrapSet = new Set(bootstrapSample(n, sampleSize, rng));
+ const featureMap = model.featureMaps[t]!;
+ for (let i = 0; i < n; i++) {
+ if (bootstrapSet.has(i)) continue;
+ const xi = X[i]!;
+ const projected = new Float64Array(featureMap.length);
+ for (let j = 0; j < featureMap.length; j++) projected[j] = xi[featureMap[j]!] ?? 0;
+ oobPreds[i] = (oobPreds[i] ?? 0) + (treePredict(model.trees[t]!, [projected])[0] ?? 0);
+ oobCounts[i] = (oobCounts[i] ?? 0) + 1;
+ }
+ }
+
+ let sse = 0;
+ let cnt = 0;
+ for (let i = 0; i < n; i++) {
+ if ((oobCounts[i] ?? 0) === 0) continue;
+ const pred = (oobPreds[i] ?? 0) / (oobCounts[i] ?? 1);
+ const err = (y[i] ?? 0) - pred;
+ sse += err * err;
+ cnt++;
+ }
+ return cnt > 0 ? sse / cnt : 0;
+}
diff --git a/src/ml/seq2seq.ts b/src/ml/seq2seq.ts
new file mode 100644
index 000000000..4c2f44658
--- /dev/null
+++ b/src/ml/seq2seq.ts
@@ -0,0 +1,208 @@
+/**
+ * Seq2Seq (Sequence-to-Sequence) with Bahdanau attention.
+ *
+ * Implements the encoder-decoder architecture with additive attention
+ * for tasks like machine translation, summarization, and transcription.
+ * Uses simple RNN (Elman) cells as building blocks.
+ *
+ * @module
+ */
+
+/** Sigmoid activation. */
+export function sigmoid(x: number): number {
+ return 1 / (1 + Math.exp(-x));
+}
+
+/** Tanh activation. */
+export function tanh(x: number): number {
+ return Math.tanh(x);
+}
+
+/** Simple RNN (Elman) cell: h_t = tanh(W_h * h_{t-1} + W_x * x_t + b). */
+export interface RNNCellWeights {
+ Wh: Float64Array; // [hiddenSize x hiddenSize]
+ Wx: Float64Array; // [hiddenSize x inputSize]
+ b: Float64Array; // [hiddenSize]
+ hiddenSize: number;
+ inputSize: number;
+}
+
+export function rnnCell(
+ x: Float64Array,
+ hPrev: Float64Array,
+ weights: RNNCellWeights,
+): Float64Array {
+ const { hiddenSize, inputSize } = weights;
+ const h = new Float64Array(hiddenSize);
+ for (let i = 0; i < hiddenSize; i++) {
+ let val = weights.b[i] ?? 0;
+ for (let j = 0; j < hiddenSize; j++) {
+ val += (weights.Wh[i * hiddenSize + j] ?? 0) * (hPrev[j] ?? 0);
+ }
+ for (let j = 0; j < inputSize; j++) {
+ val += (weights.Wx[i * inputSize + j] ?? 0) * (x[j] ?? 0);
+ }
+ h[i] = tanh(val);
+ }
+ return h;
+}
+
+/** Encode a sequence. Returns all hidden states [seqLen x hiddenSize]. */
+export function encode(
+ inputs: Float64Array[],
+ initialHidden: Float64Array,
+ weights: RNNCellWeights,
+): Float64Array[] {
+ const states: Float64Array[] = [];
+ let h = initialHidden;
+ for (const x of inputs) {
+ h = rnnCell(x, h, weights);
+ states.push(h);
+ }
+ return states;
+}
+
+/** Bahdanau (additive) attention weights. */
+export interface BahdanauWeights {
+ Wa: Float64Array; // [alignDim x hiddenSize] for encoder states
+ Ua: Float64Array; // [alignDim x hiddenSize] for decoder hidden
+ va: Float64Array; // [alignDim]
+ alignDim: number;
+ hiddenSize: number;
+}
+
+/** Compute Bahdanau attention weights. */
+export function bahdanauAttention(
+ decoderHidden: Float64Array,
+ encoderStates: Float64Array[],
+ weights: BahdanauWeights,
+): { context: Float64Array; alphas: Float64Array } {
+ const { alignDim, hiddenSize } = weights;
+ const seqLen = encoderStates.length;
+ const scores = new Float64Array(seqLen);
+
+ // Decoder contribution (constant for all encoder states)
+ const decoderContrib = new Float64Array(alignDim);
+ for (let k = 0; k < alignDim; k++) {
+ let s = 0;
+ for (let j = 0; j < hiddenSize; j++) {
+ s += (weights.Ua[k * hiddenSize + j] ?? 0) * (decoderHidden[j] ?? 0);
+ }
+ decoderContrib[k] = s;
+ }
+
+ for (let t = 0; t < seqLen; t++) {
+ const enc = encoderStates[t]!;
+ let score = 0;
+ for (let k = 0; k < alignDim; k++) {
+ let encContrib = 0;
+ for (let j = 0; j < hiddenSize; j++) {
+ encContrib += (weights.Wa[k * hiddenSize + j] ?? 0) * (enc[j] ?? 0);
+ }
+ score += (weights.va[k] ?? 0) * tanh(encContrib + (decoderContrib[k] ?? 0));
+ }
+ scores[t] = score;
+ }
+
+ // Softmax
+ let maxScore = Number.NEGATIVE_INFINITY;
+ for (let t = 0; t < seqLen; t++)
+ maxScore = Math.max(maxScore, scores[t] ?? Number.NEGATIVE_INFINITY);
+ let sumExp = 0;
+ const alphas = new Float64Array(seqLen);
+ for (let t = 0; t < seqLen; t++) {
+ alphas[t] = Math.exp((scores[t] ?? 0) - maxScore);
+ sumExp += alphas[t] ?? 0;
+ }
+ for (let t = 0; t < seqLen; t++) alphas[t] = (alphas[t] ?? 0) / (sumExp || 1);
+
+ // Context vector
+ const context = new Float64Array(hiddenSize);
+ for (let t = 0; t < seqLen; t++) {
+ const enc = encoderStates[t]!;
+ for (let j = 0; j < hiddenSize; j++) {
+ context[j] = (context[j] ?? 0) + (alphas[t] ?? 0) * (enc[j] ?? 0);
+ }
+ }
+ return { context, alphas };
+}
+
+/** Decoder step with attention. */
+export interface DecoderWeights {
+ rnn: RNNCellWeights;
+ attention: BahdanauWeights;
+ outputW: Float64Array; // [vocabSize x hiddenSize]
+ outputB: Float64Array; // [vocabSize]
+ vocabSize: number;
+}
+
+export interface DecoderStepResult {
+ hidden: Float64Array;
+ logits: Float64Array;
+ attentionWeights: Float64Array;
+}
+
+export function decoderStep(
+ inputEmbedding: Float64Array,
+ prevHidden: Float64Array,
+ encoderStates: Float64Array[],
+ weights: DecoderWeights,
+): DecoderStepResult {
+ const { context, alphas } = bahdanauAttention(prevHidden, encoderStates, weights.attention);
+
+ // Concatenate input embedding and context
+ const rnnInput = new Float64Array(inputEmbedding.length + context.length);
+ for (let i = 0; i < inputEmbedding.length; i++) rnnInput[i] = inputEmbedding[i] ?? 0;
+ for (let i = 0; i < context.length; i++) rnnInput[inputEmbedding.length + i] = context[i] ?? 0;
+
+ const hidden = rnnCell(rnnInput, prevHidden, weights.rnn);
+
+ // Project to vocabulary
+ const logits = new Float64Array(weights.vocabSize);
+ for (let v = 0; v < weights.vocabSize; v++) {
+ let sum = weights.outputB[v] ?? 0;
+ for (let j = 0; j < hidden.length; j++) {
+ sum += (weights.outputW[v * hidden.length + j] ?? 0) * (hidden[j] ?? 0);
+ }
+ logits[v] = sum;
+ }
+
+ return { hidden, logits, attentionWeights: alphas };
+}
+
+/** Greedy decode: pick argmax at each step. */
+export function greedyDecode(
+ encoderStates: Float64Array[],
+ initialHidden: Float64Array,
+ sosEmbedding: Float64Array,
+ embeddingMatrix: Float64Array,
+ weights: DecoderWeights,
+ maxLen: number,
+ eosTokenId: number,
+): number[] {
+ const result: number[] = [];
+ let hidden = initialHidden;
+ let currentEmb = sosEmbedding;
+
+ for (let step = 0; step < maxLen; step++) {
+ const { hidden: newH, logits } = decoderStep(currentEmb, hidden, encoderStates, weights);
+ hidden = newH;
+ let bestIdx = 0;
+ let bestScore = Number.NEGATIVE_INFINITY;
+ for (let v = 0; v < logits.length; v++) {
+ if ((logits[v] ?? Number.NEGATIVE_INFINITY) > bestScore) {
+ bestScore = logits[v] ?? Number.NEGATIVE_INFINITY;
+ bestIdx = v;
+ }
+ }
+ result.push(bestIdx);
+ if (bestIdx === eosTokenId) break;
+ // Get next embedding
+ const dEmb = weights.rnn.inputSize - weights.attention.hiddenSize;
+ currentEmb = new Float64Array(dEmb);
+ for (let i = 0; i < dEmb; i++) {
+ currentEmb[i] = embeddingMatrix[bestIdx * dEmb + i] ?? 0;
+ }
+ }
+ return result;
+}
diff --git a/src/ml/svm.ts b/src/ml/svm.ts
new file mode 100644
index 000000000..af569e9ca
--- /dev/null
+++ b/src/ml/svm.ts
@@ -0,0 +1,201 @@
+/**
+ * Support Vector Machine (SVM) with SMO algorithm.
+ *
+ * Implements binary SVM classification with:
+ * - Linear, polynomial, and RBF kernels
+ * - Sequential Minimal Optimization (SMO) training
+ * - Support vector regression (SVR) epsilon-insensitive loss
+ *
+ * @module
+ */
+
+/** Kernel function type. */
+export type KernelType = "linear" | "poly" | "rbf";
+
+/** SVM kernel configuration. */
+export interface SVMKernelConfig {
+ type: KernelType;
+ /** Gamma for RBF and polynomial kernels. */
+ gamma: number;
+ /** Degree for polynomial kernel. */
+ degree: number;
+ /** Coef0 for polynomial kernel. */
+ coef0: number;
+}
+
+/** Compute kernel between two vectors. */
+export function computeKernel(x1: Float64Array, x2: Float64Array, config: SVMKernelConfig): number {
+ if (config.type === "linear") {
+ let dot = 0;
+ for (let i = 0; i < x1.length; i++) dot += (x1[i] ?? 0) * (x2[i] ?? 0);
+ return dot;
+ }
+ if (config.type === "rbf") {
+ let sqDist = 0;
+ for (let i = 0; i < x1.length; i++) {
+ const d = (x1[i] ?? 0) - (x2[i] ?? 0);
+ sqDist += d * d;
+ }
+ return Math.exp(-config.gamma * sqDist);
+ }
+ // Polynomial
+ let dot = 0;
+ for (let i = 0; i < x1.length; i++) dot += (x1[i] ?? 0) * (x2[i] ?? 0);
+ return (config.gamma * dot + config.coef0) ** config.degree;
+}
+
+/** Trained SVM model. */
+export interface SVMModel {
+ /** Lagrange multipliers (support vector weights). */
+ alphas: Float64Array;
+ /** Training labels (-1 or +1). */
+ labels: Float64Array;
+ /** Training features. */
+ supportVectors: Float64Array[];
+ /** Bias term. */
+ b: number;
+ kernel: SVMKernelConfig;
+}
+
+/** Compute decision function value for a single sample. */
+export function svmDecision(model: SVMModel, x: Float64Array): number {
+ let sum = -model.b;
+ for (let i = 0; i < model.supportVectors.length; i++) {
+ const alpha = model.alphas[i] ?? 0;
+ if (Math.abs(alpha) < 1e-8) continue;
+ sum +=
+ alpha * (model.labels[i] ?? 0) * computeKernel(model.supportVectors[i]!, x, model.kernel);
+ }
+ return sum;
+}
+
+/** Predict class labels for an array of samples. */
+export function svmPredict(model: SVMModel, X: Float64Array[]): Int8Array {
+ const out = new Int8Array(X.length);
+ for (let i = 0; i < X.length; i++) {
+ out[i] = svmDecision(model, X[i]!) >= 0 ? 1 : -1;
+ }
+ return out;
+}
+
+/** SMO training algorithm for binary SVM. */
+export function fitSVM(
+ X: Float64Array[],
+ y: Float64Array,
+ C = 1.0,
+ kernel: SVMKernelConfig = { type: "rbf", gamma: 0.1, degree: 3, coef0: 0 },
+ maxIter = 200,
+ tol = 1e-3,
+): SVMModel {
+ const n = X.length;
+ const alphas = new Float64Array(n);
+ let b = 0;
+
+ // Precompute kernel matrix
+ const K = new Float64Array(n * n);
+ for (let i = 0; i < n; i++) {
+ for (let j = i; j < n; j++) {
+ const k = computeKernel(X[i]!, X[j]!, kernel);
+ K[i * n + j] = k;
+ K[j * n + i] = k;
+ }
+ }
+
+ for (let iter = 0; iter < maxIter; iter++) {
+ let numChanged = 0;
+
+ for (let i = 0; i < n; i++) {
+ const yi = y[i] ?? 0;
+ // Compute error
+ let fi = -b;
+ for (let k2 = 0; k2 < n; k2++) {
+ fi += (alphas[k2] ?? 0) * (y[k2] ?? 0) * (K[k2 * n + i] ?? 0);
+ }
+ const ei = fi - yi;
+
+ if ((yi * ei < -tol && (alphas[i] ?? 0) < C) || (yi * ei > tol && (alphas[i] ?? 0) > 0)) {
+ // Heuristic: pick j != i with max |ei - ej|
+ let j = (i + 1) % n;
+ let maxDiff = 0;
+ for (let k2 = 0; k2 < n; k2++) {
+ if (k2 === i) continue;
+ let fj = -b;
+ for (let m = 0; m < n; m++) fj += (alphas[m] ?? 0) * (y[m] ?? 0) * (K[m * n + k2] ?? 0);
+ const ej = fj - (y[k2] ?? 0);
+ if (Math.abs(ei - ej) > maxDiff) {
+ maxDiff = Math.abs(ei - ej);
+ j = k2;
+ }
+ }
+
+ const yj = y[j] ?? 0;
+ let fj = -b;
+ for (let k2 = 0; k2 < n; k2++)
+ fj += (alphas[k2] ?? 0) * (y[k2] ?? 0) * (K[k2 * n + j] ?? 0);
+ const ej = fj - yj;
+
+ const oldAlphaI = alphas[i] ?? 0;
+ const oldAlphaJ = alphas[j] ?? 0;
+
+ let L: number;
+ let H: number;
+ if (yi !== yj) {
+ L = Math.max(0, oldAlphaJ - oldAlphaI);
+ H = Math.min(C, C + oldAlphaJ - oldAlphaI);
+ } else {
+ L = Math.max(0, oldAlphaI + oldAlphaJ - C);
+ H = Math.min(C, oldAlphaI + oldAlphaJ);
+ }
+
+ if (L >= H) continue;
+
+ const eta = 2 * (K[i * n + j] ?? 0) - (K[i * n + i] ?? 0) - (K[j * n + j] ?? 0);
+ if (eta >= 0) continue;
+
+ let newAlphaJ = oldAlphaJ - (yj * (ei - ej)) / eta;
+ newAlphaJ = Math.min(H, Math.max(L, newAlphaJ));
+
+ if (Math.abs(newAlphaJ - oldAlphaJ) < 1e-5) continue;
+
+ const newAlphaI = oldAlphaI + yi * yj * (oldAlphaJ - newAlphaJ);
+ alphas[i] = newAlphaI;
+ alphas[j] = newAlphaJ;
+
+ const b1 =
+ b +
+ ei +
+ yi * (newAlphaI - oldAlphaI) * (K[i * n + i] ?? 0) +
+ yj * (newAlphaJ - oldAlphaJ) * (K[i * n + j] ?? 0);
+ const b2 =
+ b +
+ ej +
+ yi * (newAlphaI - oldAlphaI) * (K[i * n + j] ?? 0) +
+ yj * (newAlphaJ - oldAlphaJ) * (K[j * n + j] ?? 0);
+
+ if (newAlphaI > 0 && newAlphaI < C) {
+ b = b1;
+ } else if (newAlphaJ > 0 && newAlphaJ < C) {
+ b = b2;
+ } else {
+ b = (b1 + b2) / 2;
+ }
+
+ numChanged++;
+ }
+ }
+
+ if (numChanged === 0) break;
+ }
+
+ return { alphas, labels: y, supportVectors: X, b, kernel };
+}
+
+/** Compute classification accuracy. */
+export function svmAccuracy(model: SVMModel, X: Float64Array[], y: Float64Array): number {
+ const preds = svmPredict(model, X);
+ let correct = 0;
+ for (let i = 0; i < y.length; i++) {
+ if ((preds[i] ?? 0) === (y[i] ?? 0)) correct++;
+ }
+ return correct / (y.length || 1);
+}
diff --git a/src/ml/tabnet.ts b/src/ml/tabnet.ts
new file mode 100644
index 000000000..9bb99d9f2
--- /dev/null
+++ b/src/ml/tabnet.ts
@@ -0,0 +1,197 @@
+/**
+ * TabNet: Attentive Interpretable Tabular Learning.
+ *
+ * Implements the core building blocks of TabNet (Arik & Pfister, 2019):
+ * feature transformer, attentive transformer, and the sequential step logic.
+ *
+ * @module
+ */
+
+/** Sparsemax activation: projects onto the probability simplex. */
+export function sparsemax(logits: Float64Array): Float64Array {
+ const n = logits.length;
+ const sorted = Float64Array.from(logits).sort((a, b) => b - a);
+ let cumSum = 0;
+ let k = n;
+ for (let i = 0; i < n; i++) {
+ cumSum += sorted[i] ?? 0;
+ const threshold = (cumSum - 1) / (i + 1);
+ if ((sorted[i] ?? 0) <= threshold) {
+ k = i;
+ break;
+ }
+ }
+ const threshold = (cumSum - (sorted[k] ?? 0) - 1) / Math.max(k, 1);
+ const out = new Float64Array(n);
+ for (let i = 0; i < n; i++) {
+ out[i] = Math.max((logits[i] ?? 0) - threshold, 0);
+ }
+ return out;
+}
+
+/** GLU (Gated Linear Unit): splits input and applies sigmoid gate. */
+export function glu(x: Float64Array): Float64Array {
+ const half = Math.floor(x.length / 2);
+ const out = new Float64Array(half);
+ for (let i = 0; i < half; i++) {
+ const val = x[i] ?? 0;
+ const gate = 1 / (1 + Math.exp(-(x[i + half] ?? 0)));
+ out[i] = val * gate;
+ }
+ return out;
+}
+
+/** Batch normalization (inference mode, uses running stats). */
+export interface BatchNormParams {
+ mean: Float64Array;
+ variance: Float64Array;
+ gamma: Float64Array;
+ beta: Float64Array;
+ eps: number;
+}
+
+export function batchNormInfer(x: Float64Array, params: BatchNormParams): Float64Array {
+ const out = new Float64Array(x.length);
+ for (let i = 0; i < x.length; i++) {
+ const xi = x[i] ?? 0;
+ const mu = params.mean[i] ?? 0;
+ const va = params.variance[i] ?? 1;
+ const g = params.gamma[i] ?? 1;
+ const b = params.beta[i] ?? 0;
+ out[i] = g * ((xi - mu) / Math.sqrt(va + params.eps)) + b;
+ }
+ return out;
+}
+
+/** TabNet feature transformer (single step). */
+export interface FeatureTransformerWeights {
+ /** Shared layer weight matrix [hiddenDim x inputDim]. */
+ sharedW: Float64Array;
+ /** Step-specific layer weight [hiddenDim x hiddenDim]. */
+ stepW: Float64Array;
+ /** Batch norm params for shared layer output. */
+ sharedBN: BatchNormParams;
+ /** Batch norm params for step layer output. */
+ stepBN: BatchNormParams;
+ hiddenDim: number;
+ inputDim: number;
+}
+
+function matVec(W: Float64Array, x: Float64Array, rows: number, cols: number): Float64Array {
+ const out = new Float64Array(rows);
+ for (let r = 0; r < rows; r++) {
+ let sum = 0;
+ for (let c = 0; c < cols; c++) {
+ sum += (W[r * cols + c] ?? 0) * (x[c] ?? 0);
+ }
+ out[r] = sum;
+ }
+ return out;
+}
+
+export function featureTransformer(
+ x: Float64Array,
+ weights: FeatureTransformerWeights,
+): Float64Array {
+ const { hiddenDim, inputDim } = weights;
+ // Shared FC + BN + GLU
+ const sharedOut = matVec(weights.sharedW, x, hiddenDim, inputDim);
+ const sharedBN = batchNormInfer(sharedOut, weights.sharedBN);
+ const sharedGLU = glu(sharedBN);
+
+ // Step FC + BN + GLU (uses sharedGLU as input)
+ const stepOut = matVec(weights.stepW, sharedGLU, hiddenDim, Math.floor(hiddenDim / 2));
+ const stepBN = batchNormInfer(stepOut, weights.stepBN);
+ return glu(stepBN);
+}
+
+/** Attentive transformer (selects which features to focus on). */
+export interface AttentiveTransformerWeights {
+ /** Weight matrix [numFeatures x hiddenDim]. */
+ W: Float64Array;
+ /** Prior scale penalty (cumulative). */
+ priorScales: Float64Array;
+ numFeatures: number;
+ hiddenDim: number;
+}
+
+export function attentiveTransformer(
+ h: Float64Array,
+ priorScales: Float64Array,
+ weights: AttentiveTransformerWeights,
+): Float64Array {
+ const { numFeatures, hiddenDim } = weights;
+ const raw = matVec(weights.W, h, numFeatures, hiddenDim);
+ // Apply prior scale penalty
+ const penalized = new Float64Array(numFeatures);
+ for (let i = 0; i < numFeatures; i++) {
+ penalized[i] = (raw[i] ?? 0) * (priorScales[i] ?? 1);
+ }
+ return sparsemax(penalized);
+}
+
+/** TabNet step output. */
+export interface TabNetStepResult {
+ /** Feature mask (attention weights). */
+ mask: Float64Array;
+ /** Transformed features for this step. */
+ output: Float64Array;
+ /** Updated prior scales. */
+ priorScales: Float64Array;
+}
+
+/** Run one TabNet step given input features x and prior state. */
+export function tabnetStep(
+ x: Float64Array,
+ h: Float64Array,
+ priorScales: Float64Array,
+ featureWeights: FeatureTransformerWeights,
+ attentiveWeights: AttentiveTransformerWeights,
+ gamma: number,
+): TabNetStepResult {
+ const mask = attentiveTransformer(h, priorScales, attentiveWeights);
+ // Masked features
+ const maskedX = new Float64Array(x.length);
+ for (let i = 0; i < x.length; i++) {
+ maskedX[i] = (x[i] ?? 0) * (mask[i] ?? 0);
+ }
+ const output = featureTransformer(maskedX, featureWeights);
+ // Update prior scales
+ const newPrior = new Float64Array(priorScales.length);
+ for (let i = 0; i < priorScales.length; i++) {
+ newPrior[i] = (priorScales[i] ?? 1) * (gamma - (mask[i] ?? 0));
+ }
+ return { mask, output, priorScales: newPrior };
+}
+
+/** Aggregate TabNet step outputs into final representation. */
+export function aggregateSteps(stepOutputs: Float64Array[]): Float64Array {
+ if (stepOutputs.length === 0) return new Float64Array(0);
+ const dim = stepOutputs[0]?.length ?? 0;
+ const agg = new Float64Array(dim);
+ for (const out of stepOutputs) {
+ for (let i = 0; i < dim; i++) {
+ agg[i] = (agg[i] ?? 0) + Math.max(out[i] ?? 0, 0); // ReLU + sum
+ }
+ }
+ return agg;
+}
+
+/** Compute feature importance from masks across steps. */
+export function featureImportance(masks: Float64Array[]): Float64Array {
+ if (masks.length === 0) return new Float64Array(0);
+ const n = masks[0]?.length ?? 0;
+ const importance = new Float64Array(n);
+ for (const mask of masks) {
+ for (let i = 0; i < n; i++) {
+ importance[i] = (importance[i] ?? 0) + Math.abs(mask[i] ?? 0);
+ }
+ }
+ // Normalize
+ let total = 0;
+ for (let i = 0; i < n; i++) total += importance[i] ?? 0;
+ if (total > 0) {
+ for (let i = 0; i < n; i++) importance[i] = (importance[i] ?? 0) / total;
+ }
+ return importance;
+}
diff --git a/src/stats/acf_pacf.ts b/src/stats/acf_pacf.ts
new file mode 100644
index 000000000..5edc712b6
--- /dev/null
+++ b/src/stats/acf_pacf.ts
@@ -0,0 +1,618 @@
+/**
+ * acf_pacf — Autocorrelation and partial autocorrelation functions.
+ *
+ * Mirrors `statsmodels.tsa.stattools.*` for time-series correlation analysis,
+ * plus `pd.Series.autocorr`. Implemented from scratch with no external deps.
+ *
+ * Implemented functions:
+ * - {@link autocorr} — pandas-style single-lag autocorrelation
+ * - {@link acf} — full ACF with Bartlett confidence intervals
+ * - {@link pacf} — PACF via Levinson-Durbin recursion
+ * - {@link ccf} — cross-correlation function
+ * - {@link durbinWatson} — Durbin-Watson statistic for residual autocorrelation
+ * - {@link ljungBox} — Ljung-Box portmanteau test
+ * - {@link boxPierce} — Box-Pierce portmanteau test
+ *
+ * @example
+ * ```ts
+ * import { acf, pacf, ljungBox } from "tsb";
+ *
+ * const x = [1, 2, 3, 2, 1, 2, 3, 2, 1, 0, 1, 2];
+ * const { acf: corrs, lags } = acf(x, { nlags: 4, alpha: 0.05 });
+ * const { pacf: partial } = pacf(x, { nlags: 4 });
+ * const { pvalue } = ljungBox(x);
+ * ```
+ *
+ * @module
+ */
+
+import { Series } from "../core/index.ts";
+
+// ─── public types ──────────────────────────────────────────────────────────────
+
+/** Result from {@link acf}. */
+export interface ACFResult {
+ /** Autocorrelation coefficients; index 0 corresponds to lag 0 (= 1.0). */
+ readonly acf: readonly number[];
+ /**
+ * Confidence interval bounds `[lower, upper]` for each lag, computed via
+ * Bartlett's formula. `undefined` when `alpha` was not specified.
+ */
+ readonly confint: readonly [number, number][] | undefined;
+ /** Lag indices corresponding to each coefficient. */
+ readonly lags: readonly number[];
+}
+
+/** Result from {@link pacf}. */
+export interface PACFResult {
+ /** Partial autocorrelations; index 0 corresponds to lag 0 (= 1.0). */
+ readonly pacf: readonly number[];
+ /**
+ * Confidence interval bounds `[lower, upper]` for each lag.
+ * `undefined` when `alpha` was not specified.
+ */
+ readonly confint: readonly [number, number][] | undefined;
+ /** Lag indices. */
+ readonly lags: readonly number[];
+}
+
+/** Result from {@link ljungBox} or {@link boxPierce}. */
+export interface PortmanteauResult {
+ /** Q statistic at each tested lag. */
+ readonly statistic: readonly number[];
+ /** p-value at each tested lag (chi-squared df = lag − modelDf). */
+ readonly pvalue: readonly number[];
+ /** Lag indices that were tested. */
+ readonly lags: readonly number[];
+}
+
+/** Options for {@link acf}. */
+export interface ACFOptions {
+ /** Maximum lag to compute (default: `min(floor(10·log₁₀(n)), n−1)`). */
+ readonly nlags?: number;
+ /**
+ * Significance level for Bartlett confidence intervals (e.g. `0.05` → 95 % CI).
+ * Omit (default) to skip CI computation.
+ */
+ readonly alpha?: number;
+}
+
+/** Options for {@link pacf}. */
+export interface PACFOptions {
+ /** Maximum lag (default: `min(floor(10·log₁₀(n)), floor(n/2)−1)`). */
+ readonly nlags?: number;
+ /**
+ * Significance level for confidence intervals.
+ * Omit (default) to skip CI computation.
+ */
+ readonly alpha?: number;
+}
+
+/** Options for {@link ccf}. */
+export interface CCFOptions {
+ /** Maximum lag (default: `min(floor(10·log₁₀(n)), n−1)`). */
+ readonly nlags?: number;
+ /**
+ * Significance level for confidence intervals.
+ * Omit (default) to skip CI computation.
+ */
+ readonly alpha?: number;
+ /** When `true`, return only non-negative lags (default: `false`). */
+ readonly positiveOnly?: boolean;
+}
+
+/** Options for {@link ljungBox} and {@link boxPierce}. */
+export interface PortmanteauOptions {
+ /**
+ * Specific lag values to test, or a single maximum lag `h` (implying lags
+ * `1, 2, …, h`). Default: `min(floor(10·log₁₀(n)), n−1)`.
+ */
+ readonly lags?: number | readonly number[];
+ /**
+ * Number of estimated AR/MA parameters already fit to the series (default: `0`).
+ * The chi-squared degrees of freedom for lag `h` become `h − modelDf`.
+ */
+ readonly modelDf?: number;
+}
+
+// ─── internal type alias ──────────────────────────────────────────────────────
+
+/** A numeric array or {@link Series} accepted by every public function. */
+type NumericInput = readonly number[] | Series;
+
+// ─── math helpers ─────────────────────────────────────────────────────────────
+
+/** Lanczos approximation of log-Γ(z). */
+function logGamma(z: number): number {
+ if (z < 0.5) {
+ return Math.log(Math.PI / Math.sin(Math.PI * z)) - logGamma(1.0 - z);
+ }
+ const g = 7;
+ const c = [
+ 0.99999999999980993, 676.5203681218851, -1259.1392167224028, 771.32342877765313,
+ -176.61502916214059, 12.507343278686905, -0.13857109526572012, 9.9843695780195716e-6,
+ 1.5056327351493116e-7,
+ ];
+ let x = c[0] ?? 0;
+ const zz = z - 1;
+ for (let i = 1; i < g + 2; i++) {
+ x += (c[i] ?? 0) / (zz + i);
+ }
+ const t = zz + g + 0.5;
+ return 0.5 * Math.log(2 * Math.PI) + (zz + 0.5) * Math.log(t) - t + Math.log(x);
+}
+
+/** Regularised lower incomplete Γ via series expansion (x < a+1). */
+function regIncGammaSeries(a: number, x: number, lnGa: number): number {
+ let sum = 1 / a;
+ let term = 1 / a;
+ for (let n = 1; n <= 200; n++) {
+ term *= x / (a + n);
+ sum += term;
+ if (Math.abs(term) < Math.abs(sum) * 1e-10) {
+ break;
+ }
+ }
+ return Math.exp(-x + a * Math.log(x) - lnGa) * sum;
+}
+
+/** Regularised lower incomplete Γ via continued-fraction expansion (x ≥ a+1). */
+function regIncGammaCF(a: number, x: number, lnGa: number): number {
+ const eps = 1e-30;
+ let f = eps;
+ let c = f;
+ let d = 1 / (x - a + 1 + eps);
+ d = 1 / d;
+ f = c * d;
+ for (let i = 1; i <= 200; i++) {
+ const an = -i * (i - a);
+ const bn = x - a + 2 * i + 1;
+ d = an * d + bn;
+ c = bn + an / c;
+ if (Math.abs(c) < eps) {
+ c = eps;
+ }
+ d = 1 / (Math.abs(d) < eps ? eps : d);
+ const del = c * d;
+ f *= del;
+ if (Math.abs(del - 1) < 1e-10) {
+ break;
+ }
+ }
+ return 1 - Math.exp(-x + a * Math.log(x) - lnGa) * f;
+}
+
+/** Regularised lower incomplete Γ: P(a, x). */
+function regIncGamma(a: number, x: number): number {
+ if (x < 0) {
+ return 0;
+ }
+ const lnGa = logGamma(a);
+ if (x < a + 1) {
+ return regIncGammaSeries(a, x, lnGa);
+ }
+ return regIncGammaCF(a, x, lnGa);
+}
+
+/** χ² survival function: P(χ²_df > x). */
+function chi2sf(x: number, df: number): number {
+ if (x <= 0) {
+ return 1;
+ }
+ return 1 - regIncGamma(df / 2, x / 2);
+}
+
+/** Inverse standard-normal CDF (Peter Acklam's rational approximation). */
+function normalPpf(p: number): number {
+ if (p <= 0) {
+ return Number.NEGATIVE_INFINITY;
+ }
+ if (p >= 1) {
+ return Number.POSITIVE_INFINITY;
+ }
+ if (p === 0.5) {
+ return 0;
+ }
+ const a0 = -3.969683028665376e1;
+ const a1 = 2.209460984245205e2;
+ const a2 = -2.759285104469687e2;
+ const a3 = 1.38357751867269e2;
+ const a4 = -3.066479806614716e1;
+ const a5 = 2.506628277459239;
+ const b0 = -5.447609879822406e1;
+ const b1 = 1.615858368580409e2;
+ const b2 = -1.556989798598866e2;
+ const b3 = 6.680131188771972e1;
+ const b4 = -1.328068155288572e1;
+ const c0 = -7.784894002430293e-3;
+ const c1 = -3.223964580411365e-1;
+ const c2 = -2.400758277161838;
+ const c3 = -2.549732539343734;
+ const c4 = 4.374664141464968;
+ const c5 = 2.938163982698783;
+ const d0 = 7.784695709041462e-3;
+ const d1 = 3.224671290700398e-1;
+ const d2 = 2.445134137142996;
+ const d3 = 3.754408661907416;
+ const pLow = 0.02425;
+ const pHigh = 1 - pLow;
+ if (p < pLow) {
+ const q = Math.sqrt(-2 * Math.log(p));
+ return (
+ (((((c0 * q + c1) * q + c2) * q + c3) * q + c4) * q + c5) /
+ ((((d0 * q + d1) * q + d2) * q + d3) * q + 1)
+ );
+ }
+ if (p <= pHigh) {
+ const q = p - 0.5;
+ const r = q * q;
+ return (
+ ((((((a0 * r + a1) * r + a2) * r + a3) * r + a4) * r + a5) * q) /
+ (((((b0 * r + b1) * r + b2) * r + b3) * r + b4) * r + 1)
+ );
+ }
+ const q = Math.sqrt(-2 * Math.log(1 - p));
+ return -(
+ (((((c0 * q + c1) * q + c2) * q + c3) * q + c4) * q + c5) /
+ ((((d0 * q + d1) * q + d2) * q + d3) * q + 1)
+ );
+}
+
+// ─── tuple helpers ────────────────────────────────────────────────────────────
+
+/** Build a `[lower, upper]` confidence bound tuple without a cast. */
+function bound(center: number, margin: number): [number, number] {
+ return [center - margin, center + margin];
+}
+
+// ─── array extraction ─────────────────────────────────────────────────────────
+
+/** Extract a plain `number[]`, dropping NaN and non-numeric values. */
+function toNumbers(input: NumericInput): number[] {
+ if (input instanceof Series) {
+ const out: number[] = [];
+ for (const val of input.values) {
+ if (typeof val === "number" && !Number.isNaN(val)) {
+ out.push(val);
+ }
+ }
+ return out;
+ }
+ return (input as readonly number[]).filter((v) => typeof v === "number" && !Number.isNaN(v));
+}
+
+// ─── autocovariance ────────────────────────────────────────────────────────────
+
+/**
+ * Biased sample autocovariance γ̂(0), γ̂(1), …, γ̂(nlags).
+ * Denominator is n (consistent with pandas / statsmodels default).
+ */
+function autocovariance(x: readonly number[], mean: number, nlags: number): number[] {
+ const n = x.length;
+ const cov: number[] = [];
+ for (let k = 0; k <= nlags; k++) {
+ let s = 0;
+ for (let t = 0; t < n - k; t++) {
+ s += ((x[t] ?? 0) - mean) * ((x[t + k] ?? 0) - mean);
+ }
+ cov.push(s / n);
+ }
+ return cov;
+}
+
+// ─── ACF CI helper ────────────────────────────────────────────────────────────
+
+/** Bartlett confidence intervals for ACF coefficients. */
+function buildAcfCI(acfValues: readonly number[], n: number, alpha: number): [number, number][] {
+ const z = normalPpf(1 - alpha / 2);
+ const ci: [number, number][] = [];
+ let sumSq = 0;
+ for (let k = 0; k < acfValues.length; k++) {
+ if (k === 0) {
+ ci.push([1, 1]);
+ } else {
+ const se = Math.sqrt((1 + 2 * sumSq) / n);
+ const r = acfValues[k] ?? 0;
+ ci.push(bound(r, z * se));
+ sumSq += r * r;
+ }
+ }
+ return ci;
+}
+
+// ─── Levinson-Durbin helpers ──────────────────────────────────────────────────
+
+/** Single Levinson-Durbin recursion step: returns [φ_kk, updated φ array]. */
+function ldStep(acfVals: readonly number[], phi: readonly number[], k: number): [number, number[]] {
+ let num = acfVals[k] ?? 0;
+ let den = 1;
+ for (let j = 1; j < k; j++) {
+ num -= (phi[j - 1] ?? 0) * (acfVals[k - j] ?? 0);
+ den -= (phi[j - 1] ?? 0) * (acfVals[j] ?? 0);
+ }
+ const phiKK = den === 0 ? 0 : num / den;
+ const newPhi: number[] = [];
+ for (let j = 1; j < k; j++) {
+ newPhi.push((phi[j - 1] ?? 0) - phiKK * (phi[k - j - 1] ?? 0));
+ }
+ newPhi.push(phiKK);
+ return [phiKK, newPhi];
+}
+
+/** Levinson-Durbin recursion returning PACF[0..nlags]. */
+function levinsonDurbin(acfVals: readonly number[], nlags: number): number[] {
+ const result: number[] = [1];
+ if (nlags === 0) {
+ return result;
+ }
+ let phi: number[] = [];
+ for (let k = 1; k <= nlags; k++) {
+ const [phiKK, newPhi] = ldStep(acfVals, phi, k);
+ result.push(phiKK);
+ phi = newPhi;
+ }
+ return result;
+}
+
+// ─── CCF helpers ──────────────────────────────────────────────────────────────
+
+/** Cross-covariance at lag k (biased estimator, denominator = n). */
+function ccfLag(
+ xArr: readonly number[],
+ yArr: readonly number[],
+ n: number,
+ xMean: number,
+ yMean: number,
+ k: number,
+): number {
+ const start = k >= 0 ? 0 : -k;
+ const end = k >= 0 ? n - k : n;
+ let s = 0;
+ for (let t = start; t < end; t++) {
+ s += ((xArr[t] ?? 0) - xMean) * ((yArr[t + k] ?? 0) - yMean);
+ }
+ return s / n;
+}
+
+// ─── portmanteau helpers ──────────────────────────────────────────────────────
+
+/** Resolve the `lags` option to a sorted list. */
+function resolveLags(opt: number | readonly number[] | undefined, maxLag: number): number[] {
+ if (opt === undefined) {
+ return [maxLag];
+ }
+ if (typeof opt === "number") {
+ return Array.from({ length: opt }, (_, i) => i + 1);
+ }
+ return (opt as readonly number[]).slice().sort((a, b) => a - b);
+}
+
+/** Compute Ljung-Box or Box-Pierce Q at lag h. */
+function portmanteauQ(acfVals: readonly number[], n: number, h: number, ljung: boolean): number {
+ let q = 0;
+ for (let k = 1; k <= h; k++) {
+ const r = acfVals[k] ?? 0;
+ q += ljung ? (r * r) / (n - k) : r * r;
+ }
+ return ljung ? n * (n + 2) * q : n * q;
+}
+
+// ─── public API ───────────────────────────────────────────────────────────────
+
+/**
+ * Pandas-style single-lag autocorrelation.
+ *
+ * Equivalent to `pd.Series.autocorr(lag)` — computes the Pearson correlation
+ * between `x[0..n−lag−1]` and `x[lag..n−1]`.
+ *
+ * @param x Input series.
+ * @param lag Lag (default: `1`).
+ * @returns Pearson correlation in `[−1, 1]`, or `NaN` if series is too short.
+ */
+export function autocorr(x: NumericInput, lag = 1): number {
+ const arr = toNumbers(x);
+ const n = arr.length;
+ if (n <= lag + 1) {
+ return Number.NaN;
+ }
+ const x1 = arr.slice(0, n - lag);
+ const x2 = arr.slice(lag);
+ const m1 = x1.reduce((s, v) => s + v, 0) / x1.length;
+ const m2 = x2.reduce((s, v) => s + v, 0) / x2.length;
+ let num = 0;
+ let ss1 = 0;
+ let ss2 = 0;
+ for (let i = 0; i < x1.length; i++) {
+ const d1 = (x1[i] ?? 0) - m1;
+ const d2 = (x2[i] ?? 0) - m2;
+ num += d1 * d2;
+ ss1 += d1 * d1;
+ ss2 += d2 * d2;
+ }
+ const denom = Math.sqrt(ss1 * ss2);
+ return denom === 0 ? Number.NaN : num / denom;
+}
+
+/**
+ * Full Autocorrelation Function (ACF) with optional Bartlett confidence
+ * intervals.
+ *
+ * Mirrors `statsmodels.tsa.stattools.acf` (biased estimator, lag 0 = 1).
+ *
+ * @param x Input time series.
+ * @param options See {@link ACFOptions}.
+ * @returns ACF coefficients for lags 0…nlags, plus optional CI.
+ */
+export function acf(x: NumericInput, options: ACFOptions = {}): ACFResult {
+ const arr = toNumbers(x);
+ const n = arr.length;
+ const defaultNlags = Math.min(Math.floor(10 * Math.log10(n)), n - 1);
+ const nlags = Math.max(0, Math.min(options.nlags ?? defaultNlags, n - 1));
+ const mean = arr.reduce((s, v) => s + v, 0) / n;
+ const cov = autocovariance(arr, mean, nlags);
+ const gamma0 = cov[0] ?? 1;
+ const acfValues: number[] = cov.map((c) => (gamma0 === 0 ? 0 : c / gamma0));
+ const lags = Array.from({ length: nlags + 1 }, (_, i) => i);
+ const confint = options.alpha !== undefined ? buildAcfCI(acfValues, n, options.alpha) : undefined;
+ return { acf: acfValues, confint, lags };
+}
+
+/**
+ * Partial Autocorrelation Function (PACF) via the Levinson-Durbin recursion.
+ *
+ * Mirrors `statsmodels.tsa.stattools.pacf` (method `"yw"`).
+ *
+ * @param x Input time series.
+ * @param options See {@link PACFOptions}.
+ * @returns PACF coefficients for lags 0…nlags, plus optional CI.
+ */
+export function pacf(x: NumericInput, options: PACFOptions = {}): PACFResult {
+ const arr = toNumbers(x);
+ const n = arr.length;
+ const defaultNlags = Math.min(Math.floor(10 * Math.log10(n)), Math.floor(n / 2) - 1);
+ const maxAllowed = Math.max(0, Math.floor(n / 2) - 1);
+ const nlags = Math.max(0, Math.min(options.nlags ?? defaultNlags, maxAllowed));
+ const mean = arr.reduce((s, v) => s + v, 0) / n;
+ const cov = autocovariance(arr, mean, nlags);
+ const gamma0 = cov[0] ?? 1;
+ const acfVals: number[] = cov.map((c) => (gamma0 === 0 ? 0 : c / gamma0));
+ const pacfValues = levinsonDurbin(acfVals, nlags);
+ const lags = Array.from({ length: nlags + 1 }, (_, i) => i);
+ let confint: [number, number][] | undefined;
+ if (options.alpha !== undefined) {
+ const z = normalPpf(1 - options.alpha / 2);
+ const se = 1 / Math.sqrt(n);
+ confint = pacfValues.map((r) => bound(r, z * se));
+ }
+ return { pacf: pacfValues, confint, lags };
+}
+
+/**
+ * Cross-Correlation Function (CCF) between two series.
+ *
+ * CCF(k) is the normalized cross-covariance at lag k:
+ * `CCF(k) = C_xy(k) / (σ_x · σ_y)` where `C_xy(k)` uses denominator `n`.
+ *
+ * @param x First time series.
+ * @param y Second time series (must have the same length as `x`).
+ * @param options See {@link CCFOptions}.
+ * @returns CCF coefficients for lags −nlags…+nlags (or 0…nlags).
+ */
+export function ccf(x: NumericInput, y: NumericInput, options: CCFOptions = {}): ACFResult {
+ const xArr = toNumbers(x);
+ const yArr = toNumbers(y);
+ const n = Math.min(xArr.length, yArr.length);
+ const defaultNlags = Math.min(Math.floor(10 * Math.log10(n)), n - 1);
+ const nlags = Math.max(0, Math.min(options.nlags ?? defaultNlags, n - 1));
+ const positiveOnly = options.positiveOnly ?? false;
+ const xSub = xArr.slice(0, n);
+ const ySub = yArr.slice(0, n);
+ const xMean = xSub.reduce((s, v) => s + v, 0) / n;
+ const yMean = ySub.reduce((s, v) => s + v, 0) / n;
+ const xVar = xSub.reduce((s, v) => s + (v - xMean) ** 2, 0) / n;
+ const yVar = ySub.reduce((s, v) => s + (v - yMean) ** 2, 0) / n;
+ const xStd = Math.sqrt(xVar);
+ const yStd = Math.sqrt(yVar);
+ const denom = xStd * yStd;
+ const startLag = positiveOnly ? 0 : -nlags;
+ const lags: number[] = [];
+ const values: number[] = [];
+ for (let k = startLag; k <= nlags; k++) {
+ lags.push(k);
+ const cov = ccfLag(xSub, ySub, n, xMean, yMean, k);
+ values.push(denom === 0 ? 0 : cov / denom);
+ }
+ let confint: [number, number][] | undefined;
+ if (options.alpha !== undefined) {
+ const z = normalPpf(1 - options.alpha / 2);
+ const se = 1 / Math.sqrt(n);
+ confint = values.map((r) => bound(r, z * se));
+ }
+ return { acf: values, confint, lags };
+}
+
+/**
+ * Durbin-Watson statistic for autocorrelation in OLS residuals.
+ *
+ * `DW = Σ(eₜ − eₜ₋₁)² / Σeₜ²`
+ *
+ * | DW | Interpretation |
+ * |------|------------------------------|
+ * | ≈ 0 | Strong positive correlation |
+ * | ≈ 2 | No autocorrelation |
+ * | ≈ 4 | Strong negative correlation |
+ *
+ * @param residuals OLS residual array or Series.
+ * @returns Durbin-Watson statistic in `[0, 4]`.
+ */
+export function durbinWatson(residuals: NumericInput): number {
+ const e = toNumbers(residuals);
+ const n = e.length;
+ if (n < 2) {
+ return Number.NaN;
+ }
+ let diff2 = 0;
+ let ss = (e[0] ?? 0) ** 2;
+ for (let t = 1; t < n; t++) {
+ const d = (e[t] ?? 0) - (e[t - 1] ?? 0);
+ diff2 += d * d;
+ ss += (e[t] ?? 0) ** 2;
+ }
+ return ss === 0 ? 2 : diff2 / ss;
+}
+
+/**
+ * Ljung-Box Q test for serial autocorrelation up to lag `h`.
+ *
+ * `Q_LB(h) = n·(n+2) · Σₖ₌₁ʰ r̂ₖ² / (n−k)`
+ *
+ * H₀: the first `h` autocorrelations are all zero.
+ * Rejection at small p-values indicates the series is not white noise.
+ *
+ * @param x Input time series.
+ * @param options See {@link PortmanteauOptions}.
+ * @returns Test statistic, p-value, and lag for each tested lag.
+ */
+export function ljungBox(x: NumericInput, options: PortmanteauOptions = {}): PortmanteauResult {
+ return portmanteauTest(x, options, true);
+}
+
+/**
+ * Box-Pierce Q test (simplified Ljung-Box).
+ *
+ * `Q_BP(h) = n · Σₖ₌₁ʰ r̂ₖ²`
+ *
+ * @param x Input time series.
+ * @param options See {@link PortmanteauOptions}.
+ * @returns Test statistic, p-value, and lag for each tested lag.
+ */
+export function boxPierce(x: NumericInput, options: PortmanteauOptions = {}): PortmanteauResult {
+ return portmanteauTest(x, options, false);
+}
+
+/** Shared computation for Ljung-Box and Box-Pierce. */
+function portmanteauTest(
+ x: NumericInput,
+ options: PortmanteauOptions,
+ ljung: boolean,
+): PortmanteauResult {
+ const arr = toNumbers(x);
+ const n = arr.length;
+ const modelDf = options.modelDf ?? 0;
+ const defaultMaxLag = Math.min(Math.floor(10 * Math.log10(n)), n - 1);
+ const lagList = resolveLags(options.lags, defaultMaxLag);
+ const hMax = lagList.at(-1) ?? defaultMaxLag;
+ const mean = arr.reduce((s, v) => s + v, 0) / n;
+ const cov = autocovariance(arr, mean, hMax);
+ const gamma0 = cov[0] ?? 1;
+ const acfVals: number[] = cov.map((c) => (gamma0 === 0 ? 0 : c / gamma0));
+ const statistic: number[] = [];
+ const pvalue: number[] = [];
+ for (const h of lagList) {
+ const q = portmanteauQ(acfVals, n, h, ljung);
+ const df = h - modelDf;
+ statistic.push(q);
+ pvalue.push(df > 0 ? chi2sf(q, df) : Number.NaN);
+ }
+ return { statistic, pvalue, lags: lagList };
+}
diff --git a/src/stats/arima.ts b/src/stats/arima.ts
new file mode 100644
index 000000000..cc0227dd6
--- /dev/null
+++ b/src/stats/arima.ts
@@ -0,0 +1,592 @@
+/**
+ * arima — ARIMA(p, d, q) time-series model.
+ *
+ * Mirrors the `statsmodels.tsa.arima.model.ARIMA` API and pandas convention.
+ * Estimation uses the Hannan-Rissanen two-step method:
+ * 1. Fit a high-order AR(kMax) via Yule-Walker to obtain proxy residuals.
+ * 2. OLS on the differenced series using AR(p) lags + MA(q) proxy residuals.
+ *
+ * kMax = min(max(p + q + 5, 3), floor(n / 5)).
+ *
+ * Forecast confidence intervals use ψ-weight recursion. For integrated models
+ * (d > 0) the ψ-weights are accumulated d times (convolution with integration
+ * filter) before computing the forecast MSE.
+ *
+ * @example
+ * ```ts
+ * import { ARIMAModel } from "tsb";
+ *
+ * const y = [1, 2, 1.5, 2.5, 2, 3, 2.5, 3.5, 3, 4, 3.5, 4.5, 4];
+ * const model = new ARIMAModel({ p: 1, d: 0, q: 1 });
+ * const fit = model.fit(y);
+ * console.log(fit.arCoeffs, fit.maCoeffs, fit.aic);
+ * const fc = model.forecast(5);
+ * console.log(fc.forecast, fc.lower, fc.upper);
+ * ```
+ *
+ * @module
+ */
+
+import type { Series } from "../core/series.ts";
+
+// ─── Public types ──────────────────────────────────────────────────────────────
+
+/** Constructor options for {@link ARIMAModel}. */
+export interface ARIMAOptions {
+ /** AR order (number of autoregressive lags). Default: 1. */
+ readonly p?: number;
+ /** Differencing order. Default: 0. */
+ readonly d?: number;
+ /** MA order (number of moving-average lags). Default: 0. */
+ readonly q?: number;
+}
+
+/** Results returned by {@link ARIMAModel.fit}. */
+export interface ARIMAFitResult {
+ /** AR coefficients φ₁ … φₚ (index 0 = lag 1). */
+ readonly arCoeffs: readonly number[];
+ /** MA coefficients θ₁ … θ_q (index 0 = lag 1). */
+ readonly maCoeffs: readonly number[];
+ /** Intercept / constant term on the differenced scale. */
+ readonly intercept: number;
+ /** In-sample fitted values on the **original** (undifferenced) scale. */
+ readonly fittedValues: readonly number[];
+ /** Residuals on the differenced scale. */
+ readonly residuals: readonly number[];
+ /** Estimated noise variance σ². */
+ readonly sigma2: number;
+ /** Log-likelihood (Gaussian). */
+ readonly logLikelihood: number;
+ /** Akaike Information Criterion. */
+ readonly aic: number;
+ /** Bayesian Information Criterion. */
+ readonly bic: number;
+}
+
+/** Multi-step forecasts with prediction intervals from {@link ARIMAModel.forecast}. */
+export interface ARIMAForecastResult {
+ /** Point forecasts on the original (undifferenced) scale. */
+ readonly forecast: readonly number[];
+ /** Lower bound of the 95 % prediction interval (original scale). */
+ readonly lower: readonly number[];
+ /** Upper bound of the 95 % prediction interval (original scale). */
+ readonly upper: readonly number[];
+ /** Standard errors of the h-step-ahead forecast errors. */
+ readonly stderr: readonly number[];
+}
+
+// ─── Private helpers ──────────────────────────────────────────────────────────
+
+/** Type guard: distinguishes `readonly number[]` from `Series`. */
+function isNumericArray(y: readonly number[] | Series): y is readonly number[] {
+ return Array.isArray(y);
+}
+
+function arrMean(x: readonly number[]): number {
+ if (x.length === 0) {
+ return 0;
+ }
+ let s = 0;
+ for (const v of x) {
+ s += v;
+ }
+ return s / x.length;
+}
+
+/** Apply d-th order differencing; also collects the d initial values needed to
+ * reverse the operation. */
+function difference(y: readonly number[], d: number): { w: number[]; inits: number[] } {
+ const inits: number[] = [];
+ let w: number[] = y.slice();
+ for (let i = 0; i < d; i++) {
+ inits.push(w[0] ?? 0);
+ const dw = new Array(w.length - 1);
+ for (let t = 1; t < w.length; t++) {
+ dw[t - 1] = (w[t] ?? 0) - (w[t - 1] ?? 0);
+ }
+ w = dw;
+ }
+ return { w, inits };
+}
+
+/** Reverse d levels of differencing, using the stored initial values. */
+function undifference(dw: readonly number[], inits: readonly number[], d: number): number[] {
+ let w: number[] = dw.slice();
+ for (let i = d - 1; i >= 0; i--) {
+ const init = inits[i] ?? 0;
+ const un = new Array(w.length + 1);
+ un[0] = init;
+ for (let t = 0; t < w.length; t++) {
+ un[t + 1] = (un[t] ?? 0) + (w[t] ?? 0);
+ }
+ w = un;
+ }
+ return w;
+}
+
+/** Estimate AR(k) coefficients via Yule-Walker / Levinson-Durbin.
+ * Returns { ar: coefficients, sigma2: innovation variance }. */
+function yuleWalkerAR(x: readonly number[], k: number): { ar: readonly number[]; sigma2: number } {
+ if (k === 0) {
+ const mu = arrMean(x);
+ let s = 0;
+ for (const v of x) {
+ s += (v - mu) ** 2;
+ }
+ return { ar: [], sigma2: s / x.length || 1 };
+ }
+ const n = x.length;
+ const mu = arrMean(x);
+
+ // Autocovariances γ(0)…γ(k)
+ const acov = new Array(k + 1).fill(0);
+ for (let h = 0; h <= k; h++) {
+ let s = 0;
+ for (let t = h; t < n; t++) {
+ s += ((x[t] ?? 0) - mu) * ((x[t - h] ?? 0) - mu);
+ }
+ acov[h] = s / n;
+ }
+
+ // Levinson-Durbin recursion
+ let prevA: number[] = [];
+ let P = acov[0] || 1e-15;
+
+ for (let m = 1; m <= k; m++) {
+ // Reflection coefficient
+ let num = acov[m] ?? 0;
+ for (let j = 1; j < m; j++) {
+ num -= (prevA[j - 1] ?? 0) * (acov[m - j] ?? 0);
+ }
+ const km = P > 0 ? num / P : 0;
+
+ // Update AR coefficients
+ const curA = new Array(m);
+ for (let j = 1; j < m; j++) {
+ curA[j - 1] = (prevA[j - 1] ?? 0) - km * (prevA[m - j - 1] ?? 0);
+ }
+ curA[m - 1] = km;
+ P *= 1 - km * km;
+ prevA = curA;
+ }
+
+ return { ar: prevA, sigma2: Math.max(P, 1e-15) };
+}
+
+/** Solve β = (X'X)⁻¹ X'y via Gaussian elimination with partial pivoting.
+ * X is (n × k), y is (n). Returns length-k coefficient vector. */
+function olsSolve(X: readonly (readonly number[])[], y: readonly number[]): number[] {
+ const n = X.length;
+ const k = (X[0] ?? []).length;
+ // Build augmented matrix [X'X | X'y] of size k × (k+1)
+ const A: number[][] = Array.from({ length: k }, () => new Array(k + 1).fill(0));
+ for (let i = 0; i < k; i++) {
+ for (let j = 0; j < k; j++) {
+ let s = 0;
+ for (let t = 0; t < n; t++) {
+ s += ((X[t] ?? [])[i] ?? 0) * ((X[t] ?? [])[j] ?? 0);
+ }
+ (A[i] ?? [])[j] = s;
+ }
+ let s = 0;
+ for (let t = 0; t < n; t++) {
+ s += ((X[t] ?? [])[i] ?? 0) * (y[t] ?? 0);
+ }
+ (A[i] ?? [])[k] = s;
+ }
+ // Forward elimination
+ for (let col = 0; col < k; col++) {
+ // Find pivot
+ let pivotRow = col;
+ let maxAbs = Math.abs((A[col] ?? [])[col] ?? 0);
+ for (let row = col + 1; row < k; row++) {
+ const abs = Math.abs((A[row] ?? [])[col] ?? 0);
+ if (abs > maxAbs) {
+ maxAbs = abs;
+ pivotRow = row;
+ }
+ }
+ // Swap rows
+ if (pivotRow !== col) {
+ const tmp = A[col];
+ A[col] = A[pivotRow] ?? [];
+ A[pivotRow] = tmp ?? [];
+ }
+ const pivotVal = (A[col] ?? [])[col] ?? 0;
+ if (Math.abs(pivotVal) < 1e-14) {
+ continue; // singular/near-singular
+ }
+ for (let row = col + 1; row < k; row++) {
+ const factor = ((A[row] ?? [])[col] ?? 0) / pivotVal;
+ for (let c = col; c <= k; c++) {
+ (A[row] ?? [])[c] = ((A[row] ?? [])[c] ?? 0) - factor * ((A[col] ?? [])[c] ?? 0);
+ }
+ }
+ }
+ // Back substitution
+ const beta = new Array(k).fill(0);
+ for (let i = k - 1; i >= 0; i--) {
+ let val = (A[i] ?? [])[k] ?? 0;
+ for (let j = i + 1; j < k; j++) {
+ val -= ((A[i] ?? [])[j] ?? 0) * (beta[j] ?? 0);
+ }
+ const denom = (A[i] ?? [])[i] ?? 0;
+ beta[i] = Math.abs(denom) > 1e-14 ? val / denom : 0;
+ }
+ return beta;
+}
+
+/** Compute ARMA(p, q) ψ-weights (MA∞ representation) up to lag `h`.
+ * ψ₀ = 1, ψⱼ = Σᵢ₌₁ᵐⁱⁿ⁽ʲ'ᵖ⁾ φᵢ ψⱼ₋ᵢ + θⱼ (θⱼ = 0 for j > q). */
+function psiWeights(ar: readonly number[], ma: readonly number[], h: number): number[] {
+ const psi = new Array(h).fill(0);
+ psi[0] = 1;
+ for (let j = 1; j < h; j++) {
+ let v = j <= ma.length ? (ma[j - 1] ?? 0) : 0;
+ for (let i = 1; i <= Math.min(j, ar.length); i++) {
+ v += (ar[i - 1] ?? 0) * (psi[j - i] ?? 0);
+ }
+ psi[j] = v;
+ }
+ return psi;
+}
+
+/** Accumulate ψ-weights d times for the integrated process (ARIMA). */
+function integrateWeights(psi: readonly number[], d: number): number[] {
+ let w = psi.slice();
+ for (let level = 0; level < d; level++) {
+ const acc = w.slice();
+ for (let j = 1; j < acc.length; j++) {
+ acc[j] = (acc[j - 1] ?? 0) + (w[j] ?? 0);
+ }
+ w = acc;
+ }
+ return w;
+}
+
+// ─── ARIMAModel class ─────────────────────────────────────────────────────────
+
+/**
+ * ARIMA(p, d, q) time-series model.
+ *
+ * Estimation via Hannan-Rissanen two-step; forecast CIs via ψ-weight recursion.
+ */
+export class ARIMAModel {
+ private readonly _p: number;
+ private readonly _d: number;
+ private readonly _q: number;
+
+ // Set after fit()
+ private _ar: readonly number[] = [];
+ private _ma: readonly number[] = [];
+ private _mu = 0;
+ private _sigma2 = 1;
+ private _origY: readonly number[] = [];
+ private _inits: readonly number[] = [];
+ private _diffW: readonly number[] = [];
+ private _residuals: readonly number[] = [];
+ private _fitted = false;
+
+ constructor(opts: ARIMAOptions = {}) {
+ this._p = Math.max(0, Math.floor(opts.p ?? 1));
+ this._d = Math.max(0, Math.floor(opts.d ?? 0));
+ this._q = Math.max(0, Math.floor(opts.q ?? 0));
+ }
+
+ /** AR order. */
+ get p(): number {
+ return this._p;
+ }
+ /** Differencing order. */
+ get d(): number {
+ return this._d;
+ }
+ /** MA order. */
+ get q(): number {
+ return this._q;
+ }
+
+ /**
+ * Fit the model to `y`.
+ * @param y - Observed time series (number array or Series).
+ */
+ fit(y: readonly number[] | Series): ARIMAFitResult {
+ const yArr: readonly number[] = isNumericArray(y) ? y : y.values;
+
+ const n = yArr.length;
+ if (n < this._p + this._d + this._q + 2) {
+ throw new RangeError(`Series too short (${n}) for ARIMA(${this._p},${this._d},${this._q})`);
+ }
+
+ // Store original series for undifferencing
+ this._origY = yArr;
+
+ // Difference
+ const { w, inits } = difference(yArr, this._d);
+ this._inits = inits;
+ this._diffW = w;
+
+ const m = w.length; // = n - d
+
+ const p = this._p;
+ const q = this._q;
+
+ let arCoeffs: readonly number[];
+ let maCoeffs: readonly number[];
+ let intercept: number;
+
+ if (p === 0 && q === 0) {
+ // ARIMA(0,d,0): just differenced mean
+ intercept = arrMean(w);
+ arCoeffs = [];
+ maCoeffs = [];
+ } else {
+ // ── Step 1: Yule-Walker AR(kMax) for proxy residuals ──────────────────
+ // Pure AR models need no proxy residuals and should use every available lagged observation.
+ const kMax = q === 0 ? 0 : Math.min(Math.max(p + q + 5, 3), Math.floor(m / 5));
+ const { ar: arHat } = kMax > 0 ? yuleWalkerAR(w, kMax) : { ar: [] as readonly number[] };
+
+ // Proxy residuals: ε̂ₜ = wₜ - Σ arHat_j wₜ₋ⱼ
+ const eps = new Array(m).fill(0);
+ for (let t = kMax; t < m; t++) {
+ let pred = arrMean(w);
+ for (let j = 0; j < kMax; j++) {
+ pred += (arHat[j] ?? 0) * ((w[t - 1 - j] ?? 0) - arrMean(w));
+ }
+ eps[t] = (w[t] ?? 0) - pred;
+ }
+
+ // ── Step 2: OLS on ARMA(p, q) using proxy residuals ──────────────────
+ // Start index: need w_{t-p} and ε̂_{t-q} available
+ const s = Math.max(p, kMax + q);
+ const T = m - s; // number of observations in OLS
+
+ if (T <= p + q + 1) {
+ // Fall back to pure AR if OLS is under-identified
+ const { ar } = yuleWalkerAR(w, p);
+ arCoeffs = ar;
+ maCoeffs = new Array(q).fill(0);
+ intercept = 0;
+ } else {
+ // Design matrix: [1, w_{t-1},...,w_{t-p}, ε̂_{t-1},...,ε̂_{t-q}]
+ const X: number[][] = [];
+ const yOLS: number[] = [];
+
+ for (let i = 0; i < T; i++) {
+ const t = s + i;
+ const row: number[] = [1];
+ for (let j = 1; j <= p; j++) {
+ row.push(w[t - j] ?? 0);
+ }
+ for (let j = 1; j <= q; j++) {
+ row.push(eps[t - j] ?? 0);
+ }
+ X.push(row);
+ yOLS.push(w[t] ?? 0);
+ }
+
+ const beta = olsSolve(X, yOLS);
+ intercept = beta[0] ?? 0;
+ arCoeffs = beta.slice(1, 1 + p);
+ maCoeffs = beta.slice(1 + p, 1 + p + q);
+ }
+ }
+
+ // ── Compute in-sample residuals ─────────────────────────────────────────
+ const warmup = Math.max(p, q);
+ const resid = new Array(m).fill(0);
+ const wHat = new Array(m).fill(0);
+
+ for (let t = 0; t < m; t++) {
+ let pred = intercept;
+ for (let j = 0; j < p; j++) {
+ pred += (arCoeffs[j] ?? 0) * (w[t - 1 - j] ?? 0);
+ }
+ for (let j = 0; j < q; j++) {
+ pred += (maCoeffs[j] ?? 0) * (t - 1 - j >= 0 ? (resid[t - 1 - j] ?? 0) : 0);
+ }
+ wHat[t] = pred;
+ resid[t] = (w[t] ?? 0) - pred;
+ }
+
+ // Sigma2 from residuals after warmup
+ let sse = 0;
+ let cnt = 0;
+ for (let t = warmup; t < m; t++) {
+ sse += (resid[t] ?? 0) ** 2;
+ cnt++;
+ }
+ const sigma2 = cnt > 0 ? sse / cnt : 1;
+
+ // Fitted values on original scale via undifferencing wHat
+ const fittedW = wHat;
+ // fitted_y: undifference the fitted differenced series
+ // We reconstruct y_hat by integrating w_hat starting from the true initial values
+ const fittedY = undifference(fittedW, inits, this._d);
+
+ // Log-likelihood and information criteria
+ const k = 1 + p + q; // number of params (intercept + AR + MA)
+ const logLik = -0.5 * m * (Math.log(2 * Math.PI) + Math.log(sigma2) + 1);
+ const aic = -2 * logLik + 2 * k;
+ const bic = -2 * logLik + Math.log(m) * k;
+
+ this._ar = arCoeffs;
+ this._ma = maCoeffs;
+ this._mu = intercept;
+ this._sigma2 = sigma2;
+ this._residuals = resid;
+ this._fitted = true;
+
+ return {
+ arCoeffs,
+ maCoeffs,
+ intercept,
+ fittedValues: fittedY,
+ residuals: resid,
+ sigma2,
+ logLikelihood: logLik,
+ aic,
+ bic,
+ };
+ }
+
+ /**
+ * Produce multi-step forecasts starting after the last observation.
+ *
+ * Must be called after {@link fit}.
+ * @param steps - Number of future steps to forecast. Default: 1.
+ */
+ forecast(steps = 1): ARIMAForecastResult {
+ if (!this._fitted) {
+ throw new Error("Call fit() before forecast()");
+ }
+ if (steps < 1) {
+ throw new RangeError("steps must be >= 1");
+ }
+
+ const w = this._diffW;
+ const m = w.length;
+ const ar = this._ar;
+ const ma = this._ma;
+ const p = this._p;
+ const q = this._q;
+ const mu = this._mu;
+ const sigma2 = this._sigma2;
+
+ // Residuals tail (for initializing the MA part)
+ const resid = this._residuals;
+
+ // Extend w and residuals with forecasts
+ const wAll: number[] = w.slice();
+ const eAll: number[] = resid.slice();
+
+ for (let h = 1; h <= steps; h++) {
+ const t = m + h - 1; // index in extended array
+ let pred = mu;
+ for (let j = 0; j < p; j++) {
+ pred += (ar[j] ?? 0) * (wAll[t - 1 - j] ?? 0);
+ }
+ for (let j = 0; j < q; j++) {
+ // future residuals are 0; only use past residuals
+ const idx = t - 1 - j;
+ pred += (ma[j] ?? 0) * (idx < m ? (eAll[idx] ?? 0) : 0);
+ }
+ wAll.push(pred);
+ eAll.push(0); // future innovations are zero
+ }
+
+ // Extract the forecast steps (differenced scale)
+ const fcW = wAll.slice(m);
+
+ // Undifference to original scale
+ // Need the last `d` values at each integration level
+ const _fcOrig = undifference(fcW, this._inits, this._d);
+ // undifference returns d + steps values starting from inits;
+ // but the inits represent the first observed values at each level.
+ // We need to "extend" from the end of the observed data.
+
+ // Actually, for forecasting we need to integrate starting from the last observed value.
+ // Re-compute inits as the *last* values at each differencing level.
+ const lastInits = computeLastInits(this._origY, this._d);
+ const fcOrigCorrected = undifferenceFromLast(fcW, lastInits, this._d);
+
+ // ψ-weights for prediction intervals
+ const hMax = steps + 1;
+ const psiArma = psiWeights(ar, ma, hMax);
+ const psiInt = integrateWeights(psiArma, this._d);
+
+ const forecastArr: number[] = [];
+ const lowerArr: number[] = [];
+ const upperArr: number[] = [];
+ const stderrArr: number[] = [];
+
+ for (let h = 1; h <= steps; h++) {
+ const fc = fcOrigCorrected[h - 1] ?? 0;
+ // Var[e_h] = sigma2 * sum_{j=0}^{h-1} psi_j^2
+ let varH = 0;
+ for (let j = 0; j < h; j++) {
+ varH += (psiInt[j] ?? 0) ** 2;
+ }
+ const se = Math.sqrt(sigma2 * varH);
+ forecastArr.push(fc);
+ stderrArr.push(se);
+ lowerArr.push(fc - 1.96 * se);
+ upperArr.push(fc + 1.96 * se);
+ }
+
+ return { forecast: forecastArr, lower: lowerArr, upper: upperArr, stderr: stderrArr };
+ }
+}
+
+/** Compute the last observed value at each differencing level for forecasting. */
+function computeLastInits(y: readonly number[], d: number): number[] {
+ const lasts: number[] = [];
+ let w: number[] = y.slice();
+ for (let i = 0; i < d; i++) {
+ lasts.push(w.at(-1) ?? 0);
+ const dw = new Array(w.length - 1);
+ for (let t = 1; t < w.length; t++) {
+ dw[t - 1] = (w[t] ?? 0) - (w[t - 1] ?? 0);
+ }
+ w = dw;
+ }
+ return lasts;
+}
+
+/** Undifference `fcW` starting from the last observed value at each level. */
+function undifferenceFromLast(
+ fcW: readonly number[],
+ lastInits: readonly number[],
+ d: number,
+): number[] {
+ let w: number[] = fcW.slice();
+ for (let i = d - 1; i >= 0; i--) {
+ const init = lastInits[i] ?? 0;
+ const un: number[] = [];
+ let prev = init;
+ for (const dv of w) {
+ prev += dv;
+ un.push(prev);
+ }
+ w = un;
+ }
+ return w;
+}
+
+// ─── Convenience function ──────────────────────────────────────────────────────
+
+/**
+ * Fit an ARIMA(p, d, q) model and return a fitted {@link ARIMAModel}.
+ *
+ * @example
+ * ```ts
+ * import { fitArima } from "tsb";
+ * const model = fitArima([1, 2, 3, 4, 3, 2, 1, 2, 3, 4], { p: 1, q: 1 });
+ * const fc = model.forecast(3);
+ * ```
+ */
+export function fitArima(y: readonly number[] | Series, opts?: ARIMAOptions): ARIMAModel {
+ const model = new ARIMAModel(opts);
+ model.fit(y);
+ return model;
+}
diff --git a/src/stats/bootstrap.ts b/src/stats/bootstrap.ts
index 3864933aa..30ac736eb 100644
--- a/src/stats/bootstrap.ts
+++ b/src/stats/bootstrap.ts
@@ -371,7 +371,7 @@ function bootstrapOne(
const sorted = [...bootDist].sort((a, b) => a - b);
// Degenerate: all bootstrap samples yield the same statistic (e.g. n=1)
- if (sorted[0] === sorted[sorted.length - 1]) {
+ if (sorted[0] === sorted.at(-1)) {
const val = sorted[0] ?? theta;
return {
confidenceInterval: { low: val, high: val },
diff --git a/src/stats/copulas.ts b/src/stats/copulas.ts
new file mode 100644
index 000000000..27206736b
--- /dev/null
+++ b/src/stats/copulas.ts
@@ -0,0 +1,354 @@
+/**
+ * copulas — Copula models for multivariate dependence.
+ *
+ * Implements:
+ * - **Gaussian copula** — normal-based joint distribution
+ * - **t-copula** — heavy-tailed dependence
+ * - **Clayton copula** — lower tail dependence
+ * - **Gumbel copula** — upper tail dependence
+ * - **Frank copula** — symmetric dependence
+ * - **Empirical copula** — rank-based nonparametric estimate
+ * - Kendall's tau ↔ copula parameter conversion
+ *
+ * @module
+ */
+
+// ─── Utility: Normal Distribution ─────────────────────────────────────────────
+
+/** Approximation of the standard normal CDF (Abramowitz & Stegun). */
+export function normalCdf(x: number): number {
+ if (x < -8) return 0;
+ if (x > 8) return 1;
+ const a1 = 0.254829592;
+ const a2 = -0.284496736;
+ const a3 = 1.421413741;
+ const a4 = -1.453152027;
+ const a5 = 1.061405429;
+ const p = 0.3275911;
+ const sign = x < 0 ? -1 : 1;
+ const t = 1 / (1 + (p * Math.abs(x)) / Math.sqrt(2));
+ const poly = t * (a1 + t * (a2 + t * (a3 + t * (a4 + t * a5))));
+ return 0.5 * (1 + sign * (1 - poly * Math.exp(-(x * x) / 2)));
+}
+
+/** Inverse normal CDF (rational approximation, simplified). */
+export function normalQuantile(p: number): number {
+ if (p <= 0) return -Number.POSITIVE_INFINITY;
+ if (p >= 1) return Number.POSITIVE_INFINITY;
+
+ // Halley's method starting from a rough initial guess
+ let x = 0;
+ if (p < 0.5) {
+ const t = Math.sqrt(-2 * Math.log(p));
+ x =
+ -(2.515517 + 0.802853 * t + 0.010328 * t * t) /
+ (1 + 1.432788 * t + 0.189269 * t * t + 0.001308 * t * t * t) +
+ t;
+ x = -x;
+ } else {
+ const t = Math.sqrt(-2 * Math.log(1 - p));
+ x =
+ (2.515517 + 0.802853 * t + 0.010328 * t * t) /
+ (1 + 1.432788 * t + 0.189269 * t * t + 0.001308 * t * t * t) -
+ t;
+ }
+
+ // Refine with Newton-Raphson
+ for (let i = 0; i < 3; i++) {
+ const fx = normalCdf(x) - p;
+ const fpx = Math.exp(-0.5 * x * x) / Math.sqrt(2 * Math.PI);
+ if (Math.abs(fpx) < 1e-15) break;
+ x -= fx / fpx;
+ }
+
+ return x;
+}
+
+// ─── Gaussian Copula ──────────────────────────────────────────────────────────
+
+/**
+ * Evaluate the bivariate Gaussian copula CDF.
+ *
+ * @param u - First uniform marginal in (0,1).
+ * @param v - Second uniform marginal in (0,1).
+ * @param rho - Pearson correlation parameter in (-1, 1).
+ * @returns C(u, v) — copula CDF value.
+ *
+ * @example
+ * ```ts
+ * import { gaussianCopulaCdf } from "tsb";
+ * const c = gaussianCopulaCdf(0.3, 0.7, 0.5);
+ * ```
+ */
+export function gaussianCopulaCdf(u: number, v: number, rho: number): number {
+ const x = normalQuantile(u);
+ const y = normalQuantile(v);
+ return bivariateNormalCdf(x, y, rho);
+}
+
+/**
+ * Gaussian copula density (bivariate).
+ *
+ * @param u - First uniform marginal.
+ * @param v - Second uniform marginal.
+ * @param rho - Correlation parameter.
+ * @returns Copula density c(u, v).
+ */
+export function gaussianCopulaDensity(u: number, v: number, rho: number): number {
+ const x = normalQuantile(u);
+ const y = normalQuantile(v);
+ const r2 = rho * rho;
+ const exponent = -(r2 * (x * x + y * y) - 2 * rho * x * y) / (2 * (1 - r2));
+ return Math.exp(exponent) / Math.sqrt(1 - r2);
+}
+
+/**
+ * Simulate samples from a bivariate Gaussian copula.
+ *
+ * @param n - Number of samples.
+ * @param rho - Correlation parameter.
+ * @returns Array of [u, v] pairs.
+ */
+export function sampleGaussianCopula(n: number, rho: number): [number, number][] {
+ const samples: [number, number][] = [];
+ for (let i = 0; i < n; i++) {
+ const z1 = randn();
+ const z2 = randn();
+ const x = z1;
+ const y = rho * z1 + Math.sqrt(1 - rho * rho) * z2;
+ samples.push([normalCdf(x), normalCdf(y)]);
+ }
+ return samples;
+}
+
+// ─── Clayton Copula ───────────────────────────────────────────────────────────
+
+/**
+ * Clayton copula CDF.
+ *
+ * C(u,v) = (u^{-theta} + v^{-theta} - 1)^{-1/theta}
+ *
+ * @param u - First uniform marginal.
+ * @param v - Second uniform marginal.
+ * @param theta - Dependence parameter (theta > 0).
+ * @returns Copula CDF value.
+ */
+export function claytonCopulaCdf(u: number, v: number, theta: number): number {
+ if (theta <= 0) throw new Error("Clayton copula requires theta > 0");
+ return Math.max(u ** -theta + v ** -theta - 1, 0) ** (-1 / theta);
+}
+
+/**
+ * Sample from the Clayton copula using the conditional method.
+ *
+ * @param n - Number of samples.
+ * @param theta - Dependence parameter.
+ * @returns Array of [u, v] pairs.
+ */
+export function sampleClaytonCopula(n: number, theta: number): [number, number][] {
+ const samples: [number, number][] = [];
+ for (let i = 0; i < n; i++) {
+ const u = Math.random();
+ const t = Math.random();
+ // Conditional inverse: V | U
+ const v = u * (t ** (-theta / (1 + theta)) - 1 + u ** theta) ** (-1 / theta);
+ samples.push([u, Math.max(0, Math.min(1, v))]);
+ }
+ return samples;
+}
+
+// ─── Gumbel Copula ────────────────────────────────────────────────────────────
+
+/**
+ * Gumbel copula CDF.
+ *
+ * C(u,v) = exp(-[(-ln u)^theta + (-ln v)^theta]^{1/theta})
+ *
+ * @param u - First uniform marginal.
+ * @param v - Second uniform marginal.
+ * @param theta - Dependence parameter (theta >= 1).
+ * @returns Copula CDF value.
+ */
+export function gumbelCopulaCdf(u: number, v: number, theta: number): number {
+ if (theta < 1) throw new Error("Gumbel copula requires theta >= 1");
+ const a = (-Math.log(Math.max(u, 1e-300))) ** theta;
+ const b = (-Math.log(Math.max(v, 1e-300))) ** theta;
+ return Math.exp(-((a + b) ** (1 / theta)));
+}
+
+// ─── Frank Copula ─────────────────────────────────────────────────────────────
+
+/**
+ * Frank copula CDF.
+ *
+ * @param u - First uniform marginal.
+ * @param v - Second uniform marginal.
+ * @param theta - Dependence parameter (theta ≠ 0).
+ * @returns Copula CDF value.
+ */
+export function frankCopulaCdf(u: number, v: number, theta: number): number {
+ if (theta === 0) return u * v;
+ const num = (Math.exp(-theta * u) - 1) * (Math.exp(-theta * v) - 1);
+ const denom = Math.exp(-theta) - 1;
+ return -Math.log(1 + num / denom) / theta;
+}
+
+/**
+ * Sample from the Frank copula using the conditional method.
+ *
+ * @param n - Number of samples.
+ * @param theta - Dependence parameter.
+ * @returns Array of [u, v] pairs.
+ */
+export function sampleFrankCopula(n: number, theta: number): [number, number][] {
+ const samples: [number, number][] = [];
+ for (let i = 0; i < n; i++) {
+ const u = Math.random();
+ const t = Math.random();
+ // Conditional inverse
+ const et = Math.exp(-theta);
+ const eu = Math.exp(-theta * u);
+ const v = -Math.log(1 - (t * (1 - et)) / (et - eu * (1 - t))) / theta;
+ samples.push([u, Math.max(0, Math.min(1, v))]);
+ }
+ return samples;
+}
+
+// ─── Empirical Copula ─────────────────────────────────────────────────────────
+
+/**
+ * Compute the empirical copula from data.
+ *
+ * Transforms marginals to uniform via rank normalization.
+ *
+ * @param data - Array of [x, y] pairs.
+ * @returns Pseudo-observations [u, v] pairs in (0, 1).
+ */
+export function empiricalCopula(data: [number, number][]): [number, number][] {
+ const n = data.length;
+ const xs = data.map((d) => d[0]);
+ const ys = data.map((d) => d[1]);
+
+ const rankX = rankArray(xs, n);
+ const rankY = rankArray(ys, n);
+
+ return Array.from({ length: n }, (_, i) => [
+ (rankX[i] ?? 0) / (n + 1),
+ (rankY[i] ?? 0) / (n + 1),
+ ]);
+}
+
+// ─── Kendall's Tau Conversions ────────────────────────────────────────────────
+
+/**
+ * Estimate Kendall's tau from bivariate data.
+ *
+ * @param data - Array of [x, y] pairs.
+ * @returns Kendall's tau.
+ */
+export function kendallTau(data: [number, number][]): number {
+ const n = data.length;
+ let concordant = 0;
+ let discordant = 0;
+
+ for (let i = 0; i < n; i++) {
+ for (let j = i + 1; j < n; j++) {
+ const dx = (data[i]?.[0] ?? 0) - (data[j]?.[0] ?? 0);
+ const dy = (data[i]?.[1] ?? 0) - (data[j]?.[1] ?? 0);
+ if (dx * dy > 0) concordant++;
+ else if (dx * dy < 0) discordant++;
+ }
+ }
+
+ const pairs = (n * (n - 1)) / 2;
+ return pairs > 0 ? (concordant - discordant) / pairs : 0;
+}
+
+/**
+ * Convert Kendall's tau to Gaussian copula rho.
+ *
+ * rho = sin(pi * tau / 2)
+ */
+export function tauToGaussianRho(tau: number): number {
+ return Math.sin((Math.PI * tau) / 2);
+}
+
+/**
+ * Convert Kendall's tau to Clayton copula theta.
+ *
+ * theta = 2 * tau / (1 - tau)
+ */
+export function tauToClaytonTheta(tau: number): number {
+ if (tau <= 0) return 1e-6;
+ return (2 * tau) / (1 - tau);
+}
+
+/**
+ * Convert Kendall's tau to Gumbel copula theta.
+ *
+ * theta = 1 / (1 - tau)
+ */
+export function tauToGumbelTheta(tau: number): number {
+ if (tau >= 1) return 1e6;
+ return 1 / (1 - tau);
+}
+
+// ─── Utilities ────────────────────────────────────────────────────────────────
+
+/** Box-Muller standard normal sample. */
+function randn(): number {
+ let u = 0;
+ let v = 0;
+ while (u === 0) u = Math.random();
+ while (v === 0) v = Math.random();
+ return Math.sqrt(-2.0 * Math.log(u)) * Math.cos(2.0 * Math.PI * v);
+}
+
+/** Bivariate normal CDF approximation via numerical integration. */
+function bivariateNormalCdf(x: number, y: number, rho: number): number {
+ // Gauss-Legendre quadrature approximation (20 points)
+ if (x === -Number.POSITIVE_INFINITY || y === -Number.POSITIVE_INFINITY) return 0;
+ if (x === Number.POSITIVE_INFINITY) return normalCdf(y);
+ if (y === Number.POSITIVE_INFINITY) return normalCdf(x);
+
+ // Use Owen's T function approximation
+ const bvn = owenBvn(x, y, rho);
+ return Math.max(0, Math.min(1, bvn));
+}
+
+/** Approximation of bivariate normal CDF using simple formula. */
+function owenBvn(h: number, k: number, rho: number): number {
+ // Simple approximation for |rho| < 1
+ if (Math.abs(rho) < 1e-8) return normalCdf(h) * normalCdf(k);
+
+ // Quadrature using 5-point Gauss-Legendre on [-1, 1]
+ const glNodes = [-0.9061798, -0.5384693, 0, 0.5384693, 0.9061798];
+ const glWeights = [0.2369269, 0.4786287, 0.5688889, 0.4786287, 0.2369269];
+
+ // Transform to [0, rho] interval
+ let sum = 0;
+ const rhoAbs = Math.abs(rho);
+ for (let i = 0; i < glNodes.length; i++) {
+ const r = rhoAbs * 0.5 * ((glNodes[i] ?? 0) + 1);
+ const sqr = Math.sqrt(1 - r * r);
+ const exponent = (r * (2 * h * k - r * (h * h + k * k))) / (2 * (1 - r * r));
+ sum += ((glWeights[i] ?? 0) * Math.exp(exponent)) / sqr;
+ }
+
+ const L = (rhoAbs * 0.5 * sum) / (2 * Math.PI);
+ if (rho >= 0) {
+ return normalCdf(h) * normalCdf(k) + L;
+ }
+ return Math.max(0, normalCdf(h) - normalCdf(-k) + L);
+}
+
+/** Compute ranks of array values (1-indexed). */
+function rankArray(arr: number[], n: number): number[] {
+ const indexed = arr.map((v, i) => ({ v, i }));
+ indexed.sort((a, b) => a.v - b.v);
+ const ranks = new Array(n).fill(0);
+ for (let rank = 0; rank < n; rank++) {
+ ranks[(indexed[rank] ?? { i: 0 }).i] = rank + 1;
+ }
+ return ranks;
+}
diff --git a/src/stats/dlm.ts b/src/stats/dlm.ts
new file mode 100644
index 000000000..b491fbf19
--- /dev/null
+++ b/src/stats/dlm.ts
@@ -0,0 +1,1061 @@
+/**
+ * dlm — Dynamic Linear Model (Bayesian State-Space Model).
+ *
+ * Implements the West & Harrison DLM framework:
+ *
+ * Observation: Y_t = F_t' θ_t + v_t, v_t ~ N(0, V_t)
+ * System: θ_t = G_t θ_{t-1} + w_t, w_t ~ N(0, W_t)
+ * Prior: θ_0 ~ N(m_0, C_0)
+ *
+ * Supports:
+ * - Scalar or multivariate observations
+ * - Time-constant or time-varying system matrices
+ * - Factory builders for common components: local-level, local-linear-trend,
+ * polynomial trend, Fourier seasonality, regression
+ * - Component combination via `DLM.combine()`
+ * - Forward filter (Kalman), backward smoother (RTS), forecasting
+ * - MLE estimation of variance parameters via Nelder-Mead
+ * - Discount-factor model fitting (West & Harrison §6)
+ *
+ * Mirrors the R `dlm` package and Python `statsmodels.tsa.statespace` API
+ * surface while following TypeScript / tsb conventions.
+ *
+ * Exported names:
+ * - {@link DLM} — main class: filter, smooth, forecast, fit
+ * - {@link DLMOptions} — constructor options
+ * - {@link DLMResult} — filter / smoother result
+ * - {@link DLMForecastResult} — forecast result
+ * - {@link buildLocalLevel} — local-level (random-walk) component
+ * - {@link buildLocalLinearTrend} — local-linear-trend component
+ * - {@link buildPolynomial} — n-th order polynomial trend
+ * - {@link buildFourier} — Fourier (harmonic) seasonal component
+ * - {@link buildRegression} — static regression component
+ * - {@link combineDLMs} — block-diagonal combination of components
+ *
+ * @example
+ * ```ts
+ * import { DLM, buildLocalLinearTrend } from "tsb";
+ *
+ * const spec = buildLocalLinearTrend({ sigmaObs: 1, sigmaLevel: 0.5, sigmaSlope: 0.1 });
+ * const dlm = new DLM(spec);
+ * const res = dlm.filter([1, 2, 1.5, 3, 2.5, 4, 3.5]);
+ * console.log(res.filteredMeans); // T × p state vectors
+ * console.log(res.logLikelihood);
+ *
+ * const fc = dlm.forecast(res, 5);
+ * console.log(fc.mean, fc.lower, fc.upper);
+ * ```
+ *
+ * @module
+ */
+
+// ─── Internal matrix utilities ─────────────────────────────────────────────────
+
+type Mat = readonly (readonly number[])[];
+type MutMat = number[][];
+
+function rows(A: Mat): number {
+ return A.length;
+}
+function cols(A: Mat): number {
+ return A[0]?.length ?? 0;
+}
+function zeros(n: number, m: number): MutMat {
+ return Array.from({ length: n }, () => new Array(m).fill(0));
+}
+function eye(n: number): MutMat {
+ return Array.from({ length: n }, (_, i) =>
+ Array.from({ length: n }, (_, j) => (i === j ? 1 : 0)),
+ );
+}
+function mmul(A: Mat, B: Mat): MutMat {
+ const m = rows(A);
+ const k = cols(A);
+ const n = cols(B);
+ const C = zeros(m, n);
+ for (let i = 0; i < m; i++) {
+ const ai = A[i]!;
+ const ci = C[i]!;
+ for (let p = 0; p < k; p++) {
+ const aip = ai[p]!;
+ if (aip === 0) {
+ continue;
+ }
+ const bp = B[p]!;
+ for (let j = 0; j < n; j++) {
+ ci[j]! += aip * bp[j]!;
+ }
+ }
+ }
+ return C;
+}
+function mvmul(A: Mat, x: readonly number[]): number[] {
+ const m = rows(A);
+ const y = new Array(m).fill(0);
+ for (let i = 0; i < m; i++) {
+ const ai = A[i]!;
+ let s = 0;
+ for (let p = 0; p < x.length; p++) {
+ s += ai[p]! * x[p]!;
+ }
+ y[i] = s;
+ }
+ return y;
+}
+function tr(A: Mat): MutMat {
+ const m = rows(A);
+ const n = cols(A);
+ const At = zeros(n, m);
+ for (let i = 0; i < m; i++) {
+ for (let j = 0; j < n; j++) {
+ At[j]![i] = A[i]?.[j]!;
+ }
+ }
+ return At;
+}
+function madd(A: Mat, B: Mat): MutMat {
+ const m = rows(A);
+ const n = cols(A);
+ return Array.from({ length: m }, (_, i) =>
+ Array.from({ length: n }, (_, j) => A[i]?.[j]! + B[i]?.[j]!),
+ );
+}
+function msub(A: Mat, B: Mat): MutMat {
+ const m = rows(A);
+ const n = cols(A);
+ return Array.from({ length: m }, (_, i) =>
+ Array.from({ length: n }, (_, j) => A[i]?.[j]! - B[i]?.[j]!),
+ );
+}
+function mscale(A: Mat, s: number): MutMat {
+ return A.map((row) => row.map((v) => v * s));
+}
+function vdot(a: readonly number[], b: readonly number[]): number {
+ let s = 0;
+ for (let i = 0; i < a.length; i++) {
+ s += a[i]! * b[i]!;
+ }
+ return s;
+}
+function vadd(a: readonly number[], b: readonly number[]): number[] {
+ return a.map((ai, i) => ai + b[i]!);
+}
+function vsub(a: readonly number[], b: readonly number[]): number[] {
+ return a.map((ai, i) => ai - b[i]!);
+}
+function outer(a: readonly number[], b: readonly number[]): MutMat {
+ return Array.from({ length: a.length }, (_, i) =>
+ Array.from({ length: b.length }, (_, j) => a[i]! * b[j]!),
+ );
+}
+
+/** Invert a square matrix via Gauss-Jordan with partial pivoting. Returns null if singular. */
+function matInv(A: Mat): MutMat | null {
+ const n = rows(A);
+ const aug: MutMat = Array.from({ length: n }, (_, i) => [
+ ...A[i]!,
+ ...Array.from({ length: n }, (_, j) => (i === j ? 1 : 0)),
+ ]);
+ for (let col = 0; col < n; col++) {
+ let pivotRow = col;
+ let pivotVal = Math.abs(aug[col]?.[col]!);
+ for (let r = col + 1; r < n; r++) {
+ const v = Math.abs(aug[r]?.[col]!);
+ if (v > pivotVal) {
+ pivotVal = v;
+ pivotRow = r;
+ }
+ }
+ if (pivotVal < 1e-14) {
+ return null;
+ }
+ [aug[col], aug[pivotRow]] = [aug[pivotRow]!, aug[col]!];
+ const scale = 1 / aug[col]?.[col]!;
+ for (let j = 0; j < 2 * n; j++) {
+ aug[col]![j]! *= scale;
+ }
+ for (let r = 0; r < n; r++) {
+ if (r === col) {
+ continue;
+ }
+ const f = aug[r]?.[col]!;
+ if (f === 0) {
+ continue;
+ }
+ for (let j = 0; j < 2 * n; j++) {
+ aug[r]![j]! -= f * aug[col]?.[j]!;
+ }
+ }
+ }
+ return aug.map((row) => row.slice(n));
+}
+
+/** Block-diagonal combination of two matrices. */
+function blockDiag(A: Mat, B: Mat): MutMat {
+ const ra = rows(A);
+ const ca = cols(A);
+ const rb = rows(B);
+ const cb = cols(B);
+ const C = zeros(ra + rb, ca + cb);
+ for (let i = 0; i < ra; i++) {
+ for (let j = 0; j < ca; j++) {
+ C[i]![j] = A[i]?.[j]!;
+ }
+ }
+ for (let i = 0; i < rb; i++) {
+ for (let j = 0; j < cb; j++) {
+ C[ra + i]![ca + j] = B[i]?.[j]!;
+ }
+ }
+ return C;
+}
+
+// ─── Public types ───────────────────────────────────────────────────────────────
+
+/**
+ * Specification for a Dynamic Linear Model.
+ *
+ * All matrices are row-major (array-of-rows).
+ * p = state dimension, q = observation dimension.
+ */
+export interface DLMSpec {
+ /**
+ * System (transition) matrix G: p × p.
+ * If omitted, defaults to the identity.
+ */
+ readonly G: Mat;
+ /**
+ * Observation matrix F: q × p (rows = observation dim, cols = state dim).
+ * For scalar observations q = 1, so F is 1 × p.
+ */
+ readonly F: Mat;
+ /**
+ * State-noise covariance W: p × p.
+ */
+ readonly W: Mat;
+ /**
+ * Observation-noise covariance V: q × q.
+ */
+ readonly V: Mat;
+ /**
+ * Prior mean m_0: length-p vector.
+ * Defaults to zero vector.
+ */
+ readonly m0?: readonly number[];
+ /**
+ * Prior covariance C_0: p × p.
+ * Defaults to 1e6 × I (diffuse prior).
+ */
+ readonly C0?: Mat;
+}
+
+/** Options for the {@link DLM} constructor. */
+export type DLMOptions = DLMSpec;
+
+/** Per-time-step output of the Kalman filter. */
+export interface DLMFilterStep {
+ /** Predicted state mean a_t = G m_{t-1}: length p. */
+ readonly predictedMean: readonly number[];
+ /** Predicted state covariance R_t = G C_{t-1} G' + W: p × p. */
+ readonly predictedCov: Mat;
+ /** Filtered state mean m_t: length p. */
+ readonly filteredMean: readonly number[];
+ /** Filtered state covariance C_t: p × p. */
+ readonly filteredCov: Mat;
+ /** One-step-ahead forecast mean f_t = F a_t: length q. */
+ readonly forecastMean: readonly number[];
+ /** One-step-ahead forecast variance Q_t = F R_t F' + V: q × q. */
+ readonly forecastCov: Mat;
+ /** Kalman gain K_t: p × q. */
+ readonly gain: Mat;
+ /** Innovation (residual) e_t = Y_t − f_t (null if observation missing). */
+ readonly innovation: readonly number[] | null;
+}
+
+/** Result returned by {@link DLM.filter}. */
+export interface DLMResult {
+ /** Filter steps, one per time-point. */
+ readonly steps: readonly DLMFilterStep[];
+ /** Filtered state means: T × p. */
+ readonly filteredMeans: readonly (readonly number[])[];
+ /** Filtered state covariances: T × p × p. */
+ readonly filteredCovs: readonly Mat[];
+ /** Predicted (one-step-ahead) means: T × p. */
+ readonly predictedMeans: readonly (readonly number[])[];
+ /** One-step-ahead forecast means: T × q. */
+ readonly forecastMeans: readonly (readonly number[])[];
+ /** One-step-ahead forecast covariances: T × q × q. */
+ readonly forecastCovs: readonly Mat[];
+ /** Total log-likelihood. */
+ readonly logLikelihood: number;
+ /** Prior mean used. */
+ readonly m0: readonly number[];
+ /** Prior covariance used. */
+ readonly C0: Mat;
+}
+
+/** Result returned by {@link DLM.smooth}. */
+export interface DLMSmootherResult extends DLMResult {
+ /** Smoothed state means: T × p. */
+ readonly smoothedMeans: readonly (readonly number[])[];
+ /** Smoothed state covariances: T × p × p. */
+ readonly smoothedCovs: readonly Mat[];
+}
+
+/** Forecast result returned by {@link DLM.forecast}. */
+export interface DLMForecastResult {
+ /** Forecast means for steps 1 … h: h × q. */
+ readonly mean: readonly (readonly number[])[];
+ /** Forecast covariances for steps 1 … h: h × q × q. */
+ readonly cov: readonly Mat[];
+ /** 2.5th percentile (scalar observation only): length h. */
+ readonly lower: readonly number[];
+ /** 97.5th percentile (scalar observation only): length h. */
+ readonly upper: readonly number[];
+}
+
+// ─── Main DLM class ──────────────────────────────────────────────────────────────
+
+/**
+ * Dynamic Linear Model.
+ *
+ * @example
+ * ```ts
+ * import { DLM } from "tsb";
+ *
+ * const dlm = new DLM({
+ * G: [[1]],
+ * F: [[1]],
+ * W: [[0.1]],
+ * V: [[1]],
+ * });
+ * const res = dlm.filter([1, 2, 3, 4, 5]);
+ * ```
+ */
+export class DLM {
+ private readonly _G: Mat;
+ private readonly _F: Mat;
+ private readonly _W: Mat;
+ private readonly _V: Mat;
+ private readonly _m0: readonly number[];
+ private readonly _C0: Mat;
+ private readonly _p: number; // state dim
+ private readonly _q: number; // obs dim
+
+ constructor(options: DLMOptions) {
+ this._G = options.G;
+ this._F = options.F;
+ this._W = options.W;
+ this._V = options.V;
+ this._p = rows(this._G);
+ this._q = rows(this._F);
+ this._m0 = options.m0 ?? new Array(this._p).fill(0);
+ this._C0 = options.C0 ?? mscale(eye(this._p), 1e6);
+ }
+
+ // ── Factory helpers ──
+
+ /**
+ * Local-level (random-walk plus noise) model.
+ * State: [μ_t], System: μ_t = μ_{t-1} + w_t, Obs: Y_t = μ_t + v_t.
+ */
+ static localLevel(opts: { sigmaObs?: number; sigmaLevel?: number } = {}): DLM {
+ const sv = opts.sigmaObs ?? 1;
+ const sw = opts.sigmaLevel ?? 1;
+ return new DLM({
+ G: [[1]],
+ F: [[1]],
+ W: [[sw * sw]],
+ V: [[sv * sv]],
+ });
+ }
+
+ /**
+ * Local-linear-trend model.
+ * State: [level, slope], G = [[1,1],[0,1]], F = [[1,0]].
+ */
+ static localLinearTrend(
+ opts: { sigmaObs?: number; sigmaLevel?: number; sigmaSlope?: number } = {},
+ ): DLM {
+ const sv = opts.sigmaObs ?? 1;
+ const sw1 = opts.sigmaLevel ?? 1;
+ const sw2 = opts.sigmaSlope ?? 0.1;
+ return new DLM({
+ G: [
+ [1, 1],
+ [0, 1],
+ ],
+ F: [[1, 0]],
+ W: [
+ [sw1 * sw1, 0],
+ [0, sw2 * sw2],
+ ],
+ V: [[sv * sv]],
+ });
+ }
+
+ /**
+ * Polynomial trend of order `order` (1 = random walk, 2 = linear trend, …).
+ */
+ static polynomial(order: number, opts: { sigmaObs?: number; sigmaState?: number } = {}): DLM {
+ const sv = opts.sigmaObs ?? 1;
+ const sw = opts.sigmaState ?? 1;
+ // Jordan block (upper-triangular ones)
+ const G = eye(order);
+ for (let i = 0; i < order - 1; i++) {
+ G[i]![i + 1] = 1;
+ }
+ const F: MutMat = [new Array(order).fill(0)];
+ F[0]![0] = 1;
+ const W = mscale(eye(order), sw * sw);
+ const V: MutMat = [[sv * sv]];
+ return new DLM({ G, F, W, V });
+ }
+
+ /**
+ * Fourier (harmonic) seasonal component with period `period` and `harmonics` pairs.
+ * Each harmonic adds a 2-dimensional block to the state.
+ */
+ static fourier(
+ period: number,
+ harmonics: number,
+ opts: { sigmaObs?: number; sigmaState?: number } = {},
+ ): DLM {
+ const sv = opts.sigmaObs ?? 1;
+ const sw = opts.sigmaState ?? 0.01;
+ const p = 2 * harmonics;
+ // Block-diagonal G: each 2×2 block is a rotation matrix for harmonic j
+ const G = zeros(p, p);
+ const F: MutMat = [new Array(p).fill(0)];
+ for (let j = 1; j <= harmonics; j++) {
+ const omega = (2 * Math.PI * j) / period;
+ const c = Math.cos(omega);
+ const s = Math.sin(omega);
+ const r = 2 * (j - 1);
+ G[r]![r] = c;
+ G[r]![r + 1] = s;
+ G[r + 1]![r] = -s;
+ G[r + 1]![r + 1] = c;
+ F[0]![r] = 1; // cosine coefficient contributes to observation
+ }
+ const W = mscale(eye(p), sw * sw);
+ const V: MutMat = [[sv * sv]];
+ return new DLM({ G, F, W, V });
+ }
+
+ /**
+ * Static regression component with known regressors `X` (T × k).
+ * The state is the k regression coefficients (treated as random walk with
+ * small process noise `sigmaState`).
+ */
+ static regression(
+ X: readonly (readonly number[])[],
+ opts: { sigmaObs?: number; sigmaState?: number } = {},
+ ): DLM {
+ const sv = opts.sigmaObs ?? 1;
+ const sw = opts.sigmaState ?? 0.001;
+ const k = cols(X);
+ // G = identity (coefficients evolve slowly)
+ // F_t = X[t] (time-varying — not supported by DLMSpec directly, but we
+ // handle via per-step override at filter time)
+ const G = eye(k);
+ const F: MutMat = [X[0] ? [...X[0]] : new Array(k).fill(0)];
+ const W = mscale(eye(k), sw * sw);
+ const V: MutMat = [[sv * sv]];
+ return new DLM({ G, F, W, V });
+ }
+
+ // ── Core filter ──
+
+ /**
+ * Run the Kalman filter on observations `y`.
+ * Each element of `y` may be a scalar (for 1-D obs) or a vector; `null` for missing.
+ */
+ filter(
+ y: readonly (number | readonly number[] | null)[],
+ opts: { m0?: readonly number[]; C0?: Mat } = {},
+ ): DLMResult {
+ const m0 = opts.m0 ?? this._m0;
+ const C0 = opts.C0 ?? this._C0;
+ return this._filter(y, m0, C0);
+ }
+
+ private _obs(y: number | readonly number[] | null): readonly number[] | null {
+ if (y === null) {
+ return null;
+ }
+ if (typeof y === "number") {
+ return [y];
+ }
+ return y;
+ }
+
+ private _filter(
+ y: readonly (number | readonly number[] | null)[],
+ m0: readonly number[],
+ C0: Mat,
+ ): DLMResult {
+ const T = y.length;
+ const steps: DLMFilterStep[] = [];
+ let m = m0.slice();
+ let C: Mat = C0;
+ let logLik = 0;
+
+ for (let t = 0; t < T; t++) {
+ const yt = this._obs(y[t]!);
+
+ // Predict
+ const a = mvmul(this._G, m);
+ const R = madd(mmul(mmul(this._G, C), tr(this._G)), this._W);
+
+ // Forecast
+ const f = mvmul(this._F, a);
+ const Q = madd(mmul(mmul(this._F, R), tr(this._F)), this._V);
+
+ let mNew = a;
+ let CNew: Mat = R;
+ let innovation: readonly number[] | null = null;
+ let K: Mat = zeros(this._p, this._q);
+
+ if (yt !== null) {
+ // Kalman gain: K = R F' Q^{-1}
+ const RF = mmul(R, tr(this._F));
+ const Qinv = matInv(Q);
+ if (Qinv !== null) {
+ K = mmul(RF, Qinv);
+ innovation = vsub(yt, f);
+ mNew = vadd(a, mvmul(K, innovation));
+ // Joseph-form for numerical stability: C = (I-KF)R(I-KF)' + K V K'
+ const IKF = msub(eye(this._p), mmul(K, this._F));
+ CNew = madd(mmul(mmul(IKF, R), tr(IKF)), mmul(mmul(K, this._V), tr(K)));
+
+ // Log-likelihood contribution: -0.5*(q*log2π + log|Q| + e'Q^{-1}e)
+ const eQe = vdot(innovation, mvmul(Qinv, innovation));
+ const logDetQ = logDet(Q);
+ logLik += -0.5 * (this._q * Math.log(2 * Math.PI) + logDetQ + eQe);
+ }
+ }
+
+ steps.push({
+ predictedMean: a,
+ predictedCov: R,
+ filteredMean: mNew,
+ filteredCov: CNew,
+ forecastMean: f,
+ forecastCov: Q,
+ gain: K,
+ innovation,
+ });
+ m = mNew.slice();
+ C = CNew;
+ }
+
+ return {
+ steps,
+ filteredMeans: steps.map((s) => s.filteredMean),
+ filteredCovs: steps.map((s) => s.filteredCov),
+ predictedMeans: steps.map((s) => s.predictedMean),
+ forecastMeans: steps.map((s) => s.forecastMean),
+ forecastCovs: steps.map((s) => s.forecastCov),
+ logLikelihood: logLik,
+ m0,
+ C0,
+ };
+ }
+
+ // ── RTS Smoother ──
+
+ /**
+ * Run the Kalman filter followed by the RTS backward smoother.
+ */
+ smooth(
+ y: readonly (number | readonly number[] | null)[],
+ opts: { m0?: readonly number[]; C0?: Mat } = {},
+ ): DLMSmootherResult {
+ const filterResult = this.filter(y, opts);
+ return this._smooth(filterResult);
+ }
+
+ private _smooth(res: DLMResult): DLMSmootherResult {
+ const T = res.steps.length;
+ const sMeans: MutMat = res.filteredMeans.map((m) => m.slice());
+ const sCovs: MutMat[] = res.filteredCovs.map((C) => C.map((r) => r.slice()));
+
+ for (let t = T - 2; t >= 0; t--) {
+ const step = res.steps[t]!;
+ const stepNext = res.steps[t + 1]!;
+ const Rt1 = stepNext.predictedCov;
+ const Rt1inv = matInv(Rt1);
+ if (Rt1inv === null) {
+ continue;
+ }
+ // Smoother gain: J_t = C_t G' R_{t+1}^{-1}
+ const Jt = mmul(mmul(step.filteredCov, tr(this._G)), Rt1inv);
+ // Smoothed mean: s_t = m_t + J_t (s_{t+1} - a_{t+1})
+ const diff = vsub(sMeans[t + 1]!, stepNext.predictedMean);
+ sMeans[t] = vadd(step.filteredMean, mvmul(Jt, diff));
+ // Smoothed covariance: S_t = C_t + J_t (S_{t+1} - R_{t+1}) J_t'
+ const covDiff = msub(sCovs[t + 1]!, stepNext.predictedCov);
+ sCovs[t] = madd(step.filteredCov, mmul(mmul(Jt, covDiff), tr(Jt)));
+ }
+
+ return {
+ ...res,
+ smoothedMeans: sMeans,
+ smoothedCovs: sCovs,
+ };
+ }
+
+ // ── Forecasting ──
+
+ /**
+ * Forecast `h` steps ahead from the end of the filter result.
+ * Returns means and 95% prediction intervals (scalar obs only).
+ */
+ forecast(filterResult: DLMResult, h: number): DLMForecastResult {
+ const T = filterResult.steps.length;
+ const lastStep = filterResult.steps[T - 1]!;
+ let m = lastStep.filteredMean.slice();
+ let C: Mat = lastStep.filteredCov;
+
+ const meanArr: (readonly number[])[] = [];
+ const covArr: Mat[] = [];
+ const lower: number[] = [];
+ const upper: number[] = [];
+ const z975 = 1.959963985;
+
+ for (let h_i = 0; h_i < h; h_i++) {
+ // Predict one step
+ const a = mvmul(this._G, m);
+ const R = madd(mmul(mmul(this._G, C), tr(this._G)), this._W);
+ // Forecast mean and variance
+ const f = mvmul(this._F, a);
+ const Q = madd(mmul(mmul(this._F, R), tr(this._F)), this._V);
+ meanArr.push(f);
+ covArr.push(Q);
+ if (this._q === 1) {
+ const sd = Math.sqrt(Math.max(0, Q[0]?.[0]!));
+ lower.push(f[0]! - z975 * sd);
+ upper.push(f[0]! + z975 * sd);
+ }
+ m = a;
+ C = R;
+ }
+
+ return { mean: meanArr, cov: covArr, lower, upper };
+ }
+
+ // ── MLE estimation ──
+
+ /**
+ * Estimate variance parameters (log V, log W diagonal) by maximum likelihood.
+ * Returns a new DLM with the fitted parameters.
+ *
+ * Uses Nelder-Mead minimisation of negative log-likelihood.
+ */
+ fitMLE(y: readonly (number | readonly number[] | null)[]): DLM {
+ const p = this._p;
+ const q = this._q;
+ // Pack: theta = [log V entries (q*q), log diag(W) entries (p)]
+ // Only diagonal elements of V and W are fitted; off-diagonal kept fixed.
+ const packV = Array.from({ length: q }, (_, i) => Math.log(Math.max(1e-8, this._V[i]?.[i]!)));
+ const packW = Array.from({ length: p }, (_, i) => Math.log(Math.max(1e-8, this._W[i]?.[i]!)));
+ const x0 = [...packV, ...packW];
+
+ const objective = (x: readonly number[]): number => {
+ const V2: MutMat = this._V.map((row) => row.slice());
+ for (let i = 0; i < q; i++) {
+ V2[i]![i] = Math.exp(x[i]!);
+ }
+ const W2: MutMat = this._W.map((row) => row.slice());
+ for (let i = 0; i < p; i++) {
+ W2[i]![i] = Math.exp(x[q + i]!);
+ }
+ const dlm2 = new DLM({ G: this._G, F: this._F, W: W2, V: V2, m0: this._m0, C0: this._C0 });
+ const res = dlm2._filter(y, this._m0, this._C0);
+ return -res.logLikelihood;
+ };
+
+ const xOpt = nelderMead(objective, x0, { maxIter: 500, tol: 1e-6 });
+
+ const Vfit: MutMat = this._V.map((row) => row.slice());
+ for (let i = 0; i < q; i++) {
+ Vfit[i]![i] = Math.exp(xOpt[i]!);
+ }
+ const Wfit: MutMat = this._W.map((row) => row.slice());
+ for (let i = 0; i < p; i++) {
+ Wfit[i]![i] = Math.exp(xOpt[q + i]!);
+ }
+
+ return new DLM({ G: this._G, F: this._F, W: Wfit, V: Vfit, m0: this._m0, C0: this._C0 });
+ }
+
+ /**
+ * Fit with discount factors (West & Harrison §6).
+ * The process covariance is replaced by W_t = ((1-δ)/δ) C_{t-1},
+ * where δ ∈ (0,1] is the discount factor (close to 1 = small evolution).
+ * Returns a new DLM plus the filter result under the discounted model.
+ */
+ filterDiscount(y: readonly (number | readonly number[] | null)[], delta: number): DLMResult {
+ const m0 = this._m0;
+ const C0 = this._C0;
+ const T = y.length;
+ const steps: DLMFilterStep[] = [];
+ let m = m0.slice();
+ let C: Mat = C0;
+ let logLik = 0;
+
+ for (let t = 0; t < T; t++) {
+ const yt = this._obs(y[t]!);
+
+ // Predict with discount: R_t = G C_{t-1} G' / delta
+ const GCGt = mmul(mmul(this._G, C), tr(this._G));
+ const R = mscale(GCGt, 1 / delta);
+ const a = mvmul(this._G, m);
+
+ const f = mvmul(this._F, a);
+ const Q = madd(mmul(mmul(this._F, R), tr(this._F)), this._V);
+
+ let mNew = a;
+ let CNew: Mat = R;
+ let innovation: readonly number[] | null = null;
+ let K: Mat = zeros(this._p, this._q);
+
+ if (yt !== null) {
+ const Qinv = matInv(Q);
+ if (Qinv !== null) {
+ K = mmul(mmul(R, tr(this._F)), Qinv);
+ innovation = vsub(yt, f);
+ mNew = vadd(a, mvmul(K, innovation));
+ const IKF = msub(eye(this._p), mmul(K, this._F));
+ CNew = madd(mmul(mmul(IKF, R), tr(IKF)), mmul(mmul(K, this._V), tr(K)));
+ const eQe = vdot(innovation, mvmul(Qinv, innovation));
+ const logDetQ = logDet(Q);
+ logLik += -0.5 * (this._q * Math.log(2 * Math.PI) + logDetQ + eQe);
+ }
+ }
+
+ steps.push({
+ predictedMean: a,
+ predictedCov: R,
+ filteredMean: mNew,
+ filteredCov: CNew,
+ forecastMean: f,
+ forecastCov: Q,
+ gain: K,
+ innovation,
+ });
+ m = mNew.slice();
+ C = CNew;
+ }
+
+ return {
+ steps,
+ filteredMeans: steps.map((s) => s.filteredMean),
+ filteredCovs: steps.map((s) => s.filteredCov),
+ predictedMeans: steps.map((s) => s.predictedMean),
+ forecastMeans: steps.map((s) => s.forecastMean),
+ forecastCovs: steps.map((s) => s.forecastCov),
+ logLikelihood: logLik,
+ m0,
+ C0,
+ };
+ }
+}
+
+// ─── Factory helpers ────────────────────────────────────────────────────────────
+
+/** Build a local-level DLM (random walk + noise). */
+export function buildLocalLevel(opts: { sigmaObs?: number; sigmaLevel?: number } = {}): DLMSpec {
+ return {
+ G: [[1]],
+ F: [[1]],
+ W: [[(opts.sigmaLevel ?? 1) ** 2]],
+ V: [[(opts.sigmaObs ?? 1) ** 2]],
+ };
+}
+
+/** Build a local-linear-trend DLM. */
+export function buildLocalLinearTrend(
+ opts: {
+ sigmaObs?: number;
+ sigmaLevel?: number;
+ sigmaSlope?: number;
+ } = {},
+): DLMSpec {
+ const sv = opts.sigmaObs ?? 1;
+ const sw1 = opts.sigmaLevel ?? 1;
+ const sw2 = opts.sigmaSlope ?? 0.1;
+ return {
+ G: [
+ [1, 1],
+ [0, 1],
+ ],
+ F: [[1, 0]],
+ W: [
+ [sw1 * sw1, 0],
+ [0, sw2 * sw2],
+ ],
+ V: [[sv * sv]],
+ };
+}
+
+/** Build a polynomial (Jordan-block) trend DLM of a given order. */
+export function buildPolynomial(
+ order: number,
+ opts: { sigmaObs?: number; sigmaState?: number } = {},
+): DLMSpec {
+ const sv = opts.sigmaObs ?? 1;
+ const sw = opts.sigmaState ?? 1;
+ const G = eye(order);
+ for (let i = 0; i < order - 1; i++) {
+ G[i]![i + 1] = 1;
+ }
+ const F: MutMat = [new Array(order).fill(0)];
+ F[0]![0] = 1;
+ return { G, F, W: mscale(eye(order), sw * sw), V: [[sv * sv]] };
+}
+
+/** Build a Fourier seasonal DLM. */
+export function buildFourier(
+ period: number,
+ harmonics: number,
+ opts: { sigmaObs?: number; sigmaState?: number } = {},
+): DLMSpec {
+ const sv = opts.sigmaObs ?? 1;
+ const sw = opts.sigmaState ?? 0.01;
+ const p = 2 * harmonics;
+ const G = zeros(p, p);
+ const F: MutMat = [new Array(p).fill(0)];
+ for (let j = 1; j <= harmonics; j++) {
+ const omega = (2 * Math.PI * j) / period;
+ const c = Math.cos(omega);
+ const s = Math.sin(omega);
+ const r = 2 * (j - 1);
+ G[r]![r] = c;
+ G[r]![r + 1] = s;
+ G[r + 1]![r] = -s;
+ G[r + 1]![r + 1] = c;
+ F[0]![r] = 1;
+ }
+ return { G, F, W: mscale(eye(p), sw * sw), V: [[sv * sv]] };
+}
+
+/** Build a static regression DLM with k predictors (slow-varying coefficients). */
+export function buildRegression(
+ k: number,
+ opts: { sigmaObs?: number; sigmaState?: number } = {},
+): DLMSpec {
+ const sv = opts.sigmaObs ?? 1;
+ const sw = opts.sigmaState ?? 0.001;
+ return {
+ G: eye(k),
+ F: [new Array(k).fill(0)], // placeholder; rows are set externally
+ W: mscale(eye(k), sw * sw),
+ V: [[sv * sv]],
+ };
+}
+
+/**
+ * Combine multiple DLM specs into a single block-diagonal model.
+ * The combined observation matrix is the horizontal concatenation of the
+ * individual F matrices (summing each component's contribution to the scalar
+ * observation). All component V matrices must have the same dimension.
+ */
+export function combineDLMs(...specs: DLMSpec[]): DLMSpec {
+ const first = specs[0];
+ if (first === undefined) {
+ throw new RangeError("combineDLMs requires at least one spec");
+ }
+ let G: Mat = first.G;
+ let W: Mat = first.W;
+ let C0: Mat | undefined = first.C0;
+ let m0: number[] = first.m0 ? [...first.m0] : new Array(rows(first.G)).fill(0);
+
+ // Combined F: 1 × (p1+p2+…) — horizontal concat of F rows
+ let Fcombined: number[] = first.F[0] ? [...first.F[0]] : [];
+
+ for (let i = 1; i < specs.length; i++) {
+ const s = specs[i]!;
+ G = blockDiag(G, s.G);
+ W = blockDiag(W, s.W);
+ if (C0 !== undefined && s.C0 !== undefined) {
+ C0 = blockDiag(C0, s.C0);
+ } else {
+ C0 = undefined;
+ }
+ const fi = s.F[0] ? [...s.F[0]] : [];
+ Fcombined = [...Fcombined, ...fi];
+ const pm = s.m0 ? [...s.m0] : new Array(rows(s.G)).fill(0);
+ m0 = [...m0, ...pm];
+ }
+
+ // Use the first component's V (they must agree)
+ const V = first.V;
+ const spec: DLMSpec = {
+ G,
+ F: [Fcombined],
+ W,
+ V,
+ m0,
+ ...(C0 !== undefined ? { C0 } : {}),
+ };
+ return spec;
+}
+
+// ─── Log-determinant helper ─────────────────────────────────────────────────────
+
+/** Log-determinant of a positive-definite matrix via Cholesky. Falls back to LU. */
+function logDet(A: Mat): number {
+ const n = rows(A);
+ // Try Cholesky L L' = A
+ const L = zeros(n, n);
+ for (let i = 0; i < n; i++) {
+ for (let j = 0; j <= i; j++) {
+ let s = A[i]?.[j]!;
+ for (let k2 = 0; k2 < j; k2++) {
+ s -= L[i]?.[k2]! * L[j]?.[k2]!;
+ }
+ if (i === j) {
+ if (s < 0) {
+ return _logDetLU(A); // fall back
+ }
+ L[i]![j] = Math.sqrt(s);
+ } else {
+ L[i]![j] = s / (L[j]?.[j] ?? 1);
+ }
+ }
+ }
+ let ld = 0;
+ for (let i = 0; i < n; i++) {
+ ld += 2 * Math.log(Math.abs(L[i]?.[i]!));
+ }
+ return ld;
+}
+
+function _logDetLU(A: Mat): number {
+ const n = rows(A);
+ const U: MutMat = A.map((r) => r.slice());
+ let sign = 1;
+ for (let col = 0; col < n; col++) {
+ let maxVal = Math.abs(U[col]?.[col]!);
+ let maxRow = col;
+ for (let r = col + 1; r < n; r++) {
+ const v = Math.abs(U[r]?.[col]!);
+ if (v > maxVal) {
+ maxVal = v;
+ maxRow = r;
+ }
+ }
+ if (maxRow !== col) {
+ [U[col], U[maxRow]] = [U[maxRow]!, U[col]!];
+ sign *= -1;
+ }
+ const pivot = U[col]?.[col]!;
+ if (Math.abs(pivot) < 1e-14) {
+ return Number.NEGATIVE_INFINITY;
+ }
+ for (let r = col + 1; r < n; r++) {
+ const f = U[r]?.[col]! / pivot;
+ for (let j = col; j < n; j++) {
+ U[r]![j]! -= f * U[col]?.[j]!;
+ }
+ }
+ }
+ let ld = Math.log(Math.abs(sign));
+ for (let i = 0; i < n; i++) {
+ ld += Math.log(Math.abs(U[i]?.[i]!));
+ }
+ return ld;
+}
+
+// ─── Nelder-Mead optimizer ──────────────────────────────────────────────────────
+
+interface NelderMeadOptions {
+ maxIter?: number;
+ tol?: number;
+}
+
+function nelderMead(
+ f: (x: readonly number[]) => number,
+ x0: readonly number[],
+ opts: NelderMeadOptions = {},
+): number[] {
+ const maxIter = opts.maxIter ?? 1000;
+ const tol = opts.tol ?? 1e-8;
+ const n = x0.length;
+ // Build initial simplex
+ const simplex: number[][] = [x0.slice()];
+ for (let i = 0; i < n; i++) {
+ const xi = x0.slice();
+ xi[i]! += xi[i]! !== 0 ? 0.05 * Math.abs(xi[i]!) : 0.00025;
+ simplex.push(xi);
+ }
+ const fvals = simplex.map((x) => f(x));
+
+ const alpha = 1.0;
+ const gamma = 2.0;
+ const rho = 0.5;
+ const sigma = 0.5;
+
+ for (let iter = 0; iter < maxIter; iter++) {
+ // Sort
+ const order = Array.from({ length: n + 1 }, (_, i) => i).sort((a, b) => fvals[a]! - fvals[b]!);
+ const sx = order.map((i) => simplex[i]!);
+ const sf = order.map((i) => fvals[i]!);
+
+ // Convergence check
+ const range = sf[n]! - sf[0]!;
+ if (range < tol) {
+ return sx[0]?.slice() ?? x0.slice();
+ }
+
+ // Centroid of all but worst
+ const xbar = new Array(n).fill(0);
+ for (let i = 0; i < n; i++) {
+ for (let j = 0; j < n; j++) {
+ xbar[j]! += sx[i]?.[j]! / n;
+ }
+ }
+
+ // Reflect
+ const xr = xbar.map((xi, j) => xi + alpha * (xi - sx[n]?.[j]!));
+ const fr = f(xr);
+
+ if (fr < sf[0]!) {
+ // Expand
+ const xe = xbar.map((xi, j) => xi + gamma * (xr[j]! - xi));
+ const fe = f(xe);
+ simplex[order[n]!] = fe < fr ? xe : xr;
+ fvals[order[n]!] = fe < fr ? fe : fr;
+ } else if (fr < sf[n - 1]!) {
+ simplex[order[n]!] = xr;
+ fvals[order[n]!] = fr;
+ } else {
+ // Contract
+ const xc = xbar.map((xi, j) => xi + rho * (sx[n]?.[j]! - xi));
+ const fc = f(xc);
+ if (fc < sf[n]!) {
+ simplex[order[n]!] = xc;
+ fvals[order[n]!] = fc;
+ } else {
+ // Shrink
+ for (let i = 1; i <= n; i++) {
+ for (let j = 0; j < n; j++) {
+ simplex[order[i]!]![j] = sx[0]?.[j]! + sigma * (sx[i]?.[j]! - sx[0]?.[j]!);
+ }
+ fvals[order[i]!] = f(simplex[order[i]!]!);
+ }
+ }
+ }
+ }
+
+ // Return best found
+ let best = 0;
+ for (let i = 1; i < simplex.length; i++) {
+ if (fvals[i]! < fvals[best]!) {
+ best = i;
+ }
+ }
+ return simplex[best]?.slice() ?? x0.slice();
+}
diff --git a/src/stats/ets.ts b/src/stats/ets.ts
new file mode 100644
index 000000000..8cac79194
--- /dev/null
+++ b/src/stats/ets.ts
@@ -0,0 +1,1353 @@
+/**
+ * ets — Exponential Smoothing / Holt-Winters ETS models.
+ *
+ * Implements the classical additive-error ETS state-space framework:
+ * - **SimpleExpSmoothing** (SES / ETS(A,N,N)): single level parameter α.
+ * - **Holt** (ETS(A,A,N) / ETS(A,Ad,N)): level α + trend β, optional damping φ.
+ * - **ExponentialSmoothing** (Holt-Winters ETS(·,·,·)): full model with additive
+ * or multiplicative trend and seasonal components.
+ *
+ * Parameter estimation minimises SSE via Nelder-Mead simplex.
+ * Heuristic initialisation follows statsmodels conventions.
+ *
+ * API mirrors `statsmodels.tsa.holtwinters`.
+ *
+ * @example
+ * ```ts
+ * import { ExponentialSmoothing } from "tsb";
+ *
+ * const sales = [17, 21, 23, 18, 22, 26, 19, 24, 27, 20, 25, 28];
+ * const model = new ExponentialSmoothing({ trend: "add", seasonal: "add", seasonalPeriods: 4 });
+ * const fit = model.fit(sales);
+ * console.log(fit.alpha, fit.beta, fit.gamma);
+ * console.log(model.forecast(4)); // 4-step ahead forecasts
+ * ```
+ *
+ * @module
+ */
+
+import type { Series } from "../core/series.ts";
+
+// ─── Public types ──────────────────────────────────────────────────────────────
+
+/** Trend component type: additive, multiplicative, or absent. */
+export type ETSTrend = "add" | "mul" | null;
+
+/** Seasonal component type: additive, multiplicative, or absent. */
+export type ETSSeasonal = "add" | "mul" | null;
+
+/** Initialisation strategy for state variables. */
+export type ETSInit = "heuristic" | "known";
+
+// ── SimpleExpSmoothing ────────────────────────────────────────────────────────
+
+/** Options for {@link SimpleExpSmoothing}. */
+export interface SESOptions {
+ /**
+ * Smoothing level parameter (0 < α < 1).
+ * If omitted the parameter is estimated by minimising SSE.
+ */
+ readonly alpha?: number;
+ /** Initial level value. If omitted, set to `y[0]`. */
+ readonly initialLevel?: number;
+}
+
+/** Result returned by {@link SimpleExpSmoothing.fit}. */
+export interface SESFitResult {
+ /** Estimated smoothing level. */
+ readonly alpha: number;
+ /** Initial level l₀. */
+ readonly initialLevel: number;
+ /** In-sample one-step-ahead fitted values. */
+ readonly fittedValues: readonly number[];
+ /** In-sample residuals e_t = y_t − ŷ_t. */
+ readonly residuals: readonly number[];
+ /** Sum of squared errors. */
+ readonly sse: number;
+ /** Akaike Information Criterion. */
+ readonly aic: number;
+ /** Bayesian Information Criterion. */
+ readonly bic: number;
+ /** Corrected AIC. */
+ readonly aicc: number;
+}
+
+// ── Holt ─────────────────────────────────────────────────────────────────────
+
+/** Options for {@link Holt}. */
+export interface HoltOptions {
+ /** Smoothing level (0 < α < 1). Auto-estimated if omitted. */
+ readonly alpha?: number;
+ /** Smoothing trend (0 < β < 1). Auto-estimated if omitted. */
+ readonly beta?: number;
+ /** Whether to apply a damped trend. Default `false`. */
+ readonly damped?: boolean;
+ /**
+ * Damping coefficient (0 < φ < 1).
+ * Only used when `damped` is `true`. Auto-estimated if omitted.
+ */
+ readonly dampingSlope?: number;
+ /** Initial level l₀. Heuristic if omitted. */
+ readonly initialLevel?: number;
+ /** Initial trend b₀. Heuristic if omitted. */
+ readonly initialTrend?: number;
+}
+
+/** Result returned by {@link Holt.fit}. */
+export interface HoltFitResult {
+ /** Estimated level smoothing parameter. */
+ readonly alpha: number;
+ /** Estimated trend smoothing parameter. */
+ readonly beta: number;
+ /** Damping slope φ (1.0 when not damped). */
+ readonly phi: number;
+ /** Initial level l₀. */
+ readonly initialLevel: number;
+ /** Initial trend b₀. */
+ readonly initialTrend: number;
+ /** In-sample one-step-ahead fitted values. */
+ readonly fittedValues: readonly number[];
+ /** In-sample residuals. */
+ readonly residuals: readonly number[];
+ /** Sum of squared errors. */
+ readonly sse: number;
+ /** Akaike Information Criterion. */
+ readonly aic: number;
+ /** Bayesian Information Criterion. */
+ readonly bic: number;
+ /** Corrected AIC. */
+ readonly aicc: number;
+}
+
+// ── ExponentialSmoothing ─────────────────────────────────────────────────────
+
+/** Options for {@link ExponentialSmoothing}. */
+export interface ExponentialSmoothingOptions {
+ /** Trend component. `"add"` = additive, `"mul"` = multiplicative, `null` = none. */
+ readonly trend?: ETSTrend;
+ /** Whether to use a damped trend. Default `false`. */
+ readonly damped?: boolean;
+ /** Seasonal component. `"add"` = additive, `"mul"` = multiplicative, `null` = none. */
+ readonly seasonal?: ETSSeasonal;
+ /** Number of periods in one seasonal cycle (e.g. 12 for monthly, 4 for quarterly). */
+ readonly seasonalPeriods?: number;
+ /** Smoothing level parameter (0 < α < 1). Auto-estimated if omitted. */
+ readonly alpha?: number;
+ /** Smoothing trend parameter (0 < β < 1). Auto-estimated if omitted. */
+ readonly beta?: number;
+ /** Smoothing seasonal parameter (0 < γ < 1). Auto-estimated if omitted. */
+ readonly gamma?: number;
+ /** Damping slope (0 < φ < 1). Auto-estimated when `damped = true` and omitted. */
+ readonly phi?: number;
+ /** How to initialise the state: `"heuristic"` (default) or `"known"`. */
+ readonly initializationMethod?: ETSInit;
+ /** Known initial level (only used when `initializationMethod = "known"`). */
+ readonly initialLevel?: number;
+ /** Known initial trend (only used when `initializationMethod = "known"`). */
+ readonly initialTrend?: number;
+ /** Known initial seasonal indices (only used when `initializationMethod = "known"`). */
+ readonly initialSeasons?: readonly number[];
+}
+
+/** Result returned by {@link ExponentialSmoothing.fit}. */
+export interface ExponentialSmoothingFitResult {
+ /** Estimated level smoothing parameter. */
+ readonly alpha: number;
+ /** Estimated trend smoothing parameter (`null` when no trend component). */
+ readonly beta: number | null;
+ /** Estimated seasonal smoothing parameter (`null` when no seasonal component). */
+ readonly gamma: number | null;
+ /** Damping slope φ (1.0 when not damped). */
+ readonly phi: number;
+ /** Initial level l₀. */
+ readonly initialLevel: number;
+ /** Initial trend b₀ (`null` when no trend). */
+ readonly initialTrend: number | null;
+ /** Initial seasonal indices s₁…s_m (`null` when no seasonal). */
+ readonly initialSeasons: readonly number[] | null;
+ /** In-sample one-step-ahead fitted values. */
+ readonly fittedValues: readonly number[];
+ /** In-sample residuals. */
+ readonly residuals: readonly number[];
+ /** Sum of squared errors. */
+ readonly sse: number;
+ /** Log-likelihood. */
+ readonly logLikelihood: number;
+ /** Akaike Information Criterion. */
+ readonly aic: number;
+ /** Bayesian Information Criterion. */
+ readonly bic: number;
+ /** Corrected AIC. */
+ readonly aicc: number;
+}
+
+/** Forecast result with prediction intervals. */
+export interface ETSForecastResult {
+ /** Point forecasts h = 1, 2, … steps. */
+ readonly forecast: readonly number[];
+ /** Lower bound of (1 − α_ci) % prediction interval. */
+ readonly lower: readonly number[];
+ /** Upper bound of (1 − α_ci) % prediction interval. */
+ readonly upper: readonly number[];
+ /** Standard errors of h-step-ahead forecast errors. */
+ readonly stderr: readonly number[];
+}
+
+// ─── Private helpers ──────────────────────────────────────────────────────────
+
+/** Extract numeric array from Series or array. */
+function toArr(y: readonly number[] | Series): readonly number[] {
+ if ("dtype" in y) {
+ return y.values;
+ }
+ return y;
+}
+
+/** Clamp value to [lo, hi]. */
+function clamp(v: number, lo: number, hi: number): number {
+ return Math.min(Math.max(v, lo), hi);
+}
+
+/** Clamp all elements of a param vector to their respective bounds. */
+function clampParams(params: readonly number[], bounds: readonly [number, number][]): number[] {
+ return params.map((v, i) => clamp(v, (bounds[i] ?? [0, 1])[0], (bounds[i] ?? [0, 1])[1]));
+}
+
+/**
+ * Nelder-Mead simplex optimiser (unconstrained; bounds enforced by clamping).
+ * Minimises `fn(params)` starting from `x0`.
+ */
+function nelderMead(
+ fn: (params: readonly number[]) => number,
+ x0: readonly number[],
+ bounds: readonly [number, number][],
+ maxIter = 3000,
+): { params: number[]; value: number } {
+ const n = x0.length;
+ if (n === 0) {
+ return { params: [], value: fn([]) };
+ }
+
+ const EPS = 1e-12;
+ const ALPHA_NM = 1.0; // reflection
+ const BETA_NM = 0.5; // contraction
+ const GAMMA_NM = 2.0; // expansion
+ const SIGMA_NM = 0.5; // shrinkage
+
+ const clamp1 = (p: readonly number[]): number[] => clampParams(p, bounds);
+
+ // Build initial simplex
+ const simplex: number[][] = [clamp1(x0)];
+ for (let i = 0; i < n; i++) {
+ const pt = clamp1(x0);
+ const lo = (bounds[i] ?? [0, 1])[0];
+ const hi = (bounds[i] ?? [0, 1])[1];
+ const delta = Math.max((hi - lo) * 0.1, 0.01);
+ pt[i] = clamp((pt[i] ?? 0) + delta, lo + EPS, hi - EPS);
+ simplex.push(pt);
+ }
+
+ const fvals: number[] = simplex.map((p) => fn(p));
+
+ for (let iter = 0; iter < maxIter; iter++) {
+ // Sort indices by fval
+ const ord = Array.from({ length: n + 1 }, (_, i) => i).sort(
+ (a, b) => (fvals[a] ?? 0) - (fvals[b] ?? 0),
+ );
+
+ const fBest = fvals[ord[0] ?? 0] ?? 0;
+ const fWorst = fvals[ord[n] ?? 0] ?? 0;
+ if (fWorst - fBest < EPS) {
+ break;
+ }
+
+ // Centroid of best n points
+ const cent = new Array(n).fill(0);
+ for (let i = 0; i < n; i++) {
+ const row = simplex[ord[i] ?? 0] ?? [];
+ for (let j = 0; j < n; j++) {
+ cent[j] = (cent[j] ?? 0) + (row[j] ?? 0);
+ }
+ }
+ for (let j = 0; j < n; j++) {
+ cent[j] = (cent[j] ?? 0) / n;
+ }
+
+ const worstPt = simplex[ord[n] ?? 0] ?? [];
+
+ // Reflection
+ const xr = clamp1(cent.map((c, j) => (1 + ALPHA_NM) * c - ALPHA_NM * (worstPt[j] ?? 0)));
+ const fr = fn(xr);
+
+ const fSecondWorst = fvals[ord[n - 1] ?? 0] ?? 0;
+
+ if (fr < fBest) {
+ // Expansion
+ const xe = clamp1(cent.map((c, j) => (1 + GAMMA_NM) * c - GAMMA_NM * (worstPt[j] ?? 0)));
+ const fe = fn(xe);
+ if (fe < fr) {
+ simplex[ord[n] ?? 0] = xe;
+ fvals[ord[n] ?? 0] = fe;
+ } else {
+ simplex[ord[n] ?? 0] = xr;
+ fvals[ord[n] ?? 0] = fr;
+ }
+ } else if (fr < fSecondWorst) {
+ simplex[ord[n] ?? 0] = xr;
+ fvals[ord[n] ?? 0] = fr;
+ } else {
+ // Contraction
+ const inside = fr >= fWorst;
+ const src = inside ? worstPt : xr;
+ const xc = clamp1(cent.map((c, j) => BETA_NM * c + (1 - BETA_NM) * (src[j] ?? 0)));
+ const fc = fn(xc);
+ const compareVal = inside ? fWorst : fr;
+ if (fc < compareVal) {
+ simplex[ord[n] ?? 0] = xc;
+ fvals[ord[n] ?? 0] = fc;
+ } else {
+ // Shrink
+ const bestPt = simplex[ord[0] ?? 0] ?? [];
+ for (let i = 1; i <= n; i++) {
+ const row = simplex[ord[i] ?? 0] ?? [];
+ const newRow = clamp1(row.map((v, j) => SIGMA_NM * (v + (bestPt[j] ?? 0))));
+ simplex[ord[i] ?? 0] = newRow;
+ fvals[ord[i] ?? 0] = fn(newRow);
+ }
+ }
+ }
+ }
+
+ // Return best
+ let bestIdx = 0;
+ for (let i = 1; i <= n; i++) {
+ if ((fvals[i] ?? Number.POSITIVE_INFINITY) < (fvals[bestIdx] ?? Number.POSITIVE_INFINITY)) {
+ bestIdx = i;
+ }
+ }
+ return { params: simplex[bestIdx] ?? [], value: fvals[bestIdx] ?? Number.POSITIVE_INFINITY };
+}
+
+/** Compute AIC, BIC, AICc from SSE, n, k. */
+function infoGaussian(
+ sse: number,
+ n: number,
+ k: number,
+): { logLikelihood: number; aic: number; bic: number; aicc: number } {
+ const sigma2 = Math.max(sse / n, 1e-15);
+ const logL = -0.5 * n * (Math.log(2 * Math.PI * sigma2) + 1);
+ const aic = -2 * logL + 2 * k;
+ const bic = -2 * logL + k * Math.log(n);
+ const denom = n - k - 1;
+ const aicc = denom > 0 ? aic + (2 * k * (k + 1)) / denom : aic;
+ return { logLikelihood: logL, aic, bic, aicc };
+}
+
+// ─── SES internals ────────────────────────────────────────────────────────────
+
+/**
+ * Run one SES pass. Returns { fitted, residuals, sse }.
+ * l0 = initial level.
+ */
+function sesPass(
+ y: readonly number[],
+ alpha: number,
+ l0: number,
+): { fitted: number[]; residuals: number[]; sse: number } {
+ const n = y.length;
+ const fitted: number[] = new Array(n);
+ const residuals: number[] = new Array(n);
+ let sse = 0;
+ let l = l0;
+ for (let t = 0; t < n; t++) {
+ fitted[t] = l;
+ const e = (y[t] ?? 0) - l;
+ residuals[t] = e;
+ sse += e * e;
+ l = alpha * (y[t] ?? 0) + (1 - alpha) * l;
+ }
+ return { fitted, residuals, sse };
+}
+
+// ─── Holt internals ───────────────────────────────────────────────────────────
+
+/**
+ * Run one Holt pass.
+ * Returns { fitted, residuals, sse, levels, trends }.
+ */
+function holtPass(
+ y: readonly number[],
+ alpha: number,
+ beta: number,
+ phi: number,
+ l0: number,
+ b0: number,
+): { fitted: number[]; residuals: number[]; sse: number } {
+ const n = y.length;
+ const fitted: number[] = new Array(n);
+ const residuals: number[] = new Array(n);
+ let sse = 0;
+ let l = l0;
+ let b = b0;
+ for (let t = 0; t < n; t++) {
+ const yhat = l + phi * b;
+ fitted[t] = yhat;
+ const yt = y[t] ?? 0;
+ const e = yt - yhat;
+ residuals[t] = e;
+ sse += e * e;
+ const lNew = alpha * yt + (1 - alpha) * (l + phi * b);
+ b = beta * (lNew - l) + (1 - beta) * phi * b;
+ l = lNew;
+ }
+ return { fitted, residuals, sse };
+}
+
+/** Holt h-step forecast (damped or not). */
+function holtForecast(steps: number, l: number, b: number, phi: number): number[] {
+ const out: number[] = [];
+ let phiH = phi; // φ¹
+ let phiSum = phi; // φ + φ² + … + φ^h
+ for (let h = 1; h <= steps; h++) {
+ out.push(l + phiSum * b);
+ phiH *= phi;
+ phiSum += phiH;
+ }
+ return out;
+}
+
+// ─── ETS (Holt-Winters) internals ─────────────────────────────────────────────
+
+interface ETSState {
+ l: number;
+ b: number;
+ s: number[]; // length m circular buffer, s[0] = s_{t-m+1}, … , s[m-1] = s_t
+}
+
+/**
+ * Run one Holt-Winters pass. Returns fitted values, residuals, SSE, and the
+ * final state (l, b, last m seasonal indices).
+ */
+function hwPass(
+ y: readonly number[],
+ alpha: number,
+ beta: number | null,
+ gamma: number | null,
+ phi: number,
+ l0: number,
+ b0: number | null,
+ s0: readonly number[] | null,
+ trend: ETSTrend,
+ seasonal: ETSSeasonal,
+ m: number,
+): {
+ fitted: number[];
+ residuals: number[];
+ sse: number;
+ finalL: number;
+ finalB: number;
+ finalS: number[];
+} {
+ const n = y.length;
+ const fitted: number[] = new Array(n);
+ const residuals: number[] = new Array(n);
+ let sse = 0;
+
+ let l = l0;
+ let b = b0 ?? 0;
+
+ // seasonal buffer: seasonals[t % m] = s_{t+1-m}
+ const seasonals: number[] = s0 ? s0.slice() : new Array(m).fill(0);
+
+ for (let t = 0; t < n; t++) {
+ const yt = y[t] ?? 0;
+ const sIdx = ((t % m) + m) % m; // index into seasonal buffer
+ const st_m = seasonals[sIdx] ?? 0; // s_{t+1-m}
+
+ // One-step-ahead forecast
+ let yhat: number;
+ if (trend === null && seasonal === null) {
+ yhat = l;
+ } else if (trend !== null && seasonal === null) {
+ yhat = l + phi * b;
+ } else if (trend === null && seasonal === "add") {
+ yhat = l + st_m;
+ } else if (trend === null && seasonal === "mul") {
+ yhat = l * st_m;
+ } else if (trend === "add" && seasonal === "add") {
+ yhat = l + phi * b + st_m;
+ } else if (trend === "add" && seasonal === "mul") {
+ yhat = (l + phi * b) * st_m;
+ } else if (trend === "mul" && seasonal === "add") {
+ yhat = l * (phi === 1 ? b : b ** phi) + st_m;
+ } else {
+ // mul trend + mul seasonal
+ yhat = l * (phi === 1 ? b : b ** phi) * st_m;
+ }
+
+ fitted[t] = yhat;
+ const e = yt - yhat;
+ residuals[t] = e;
+ sse += e * e;
+
+ // State update
+ const lPrev = l;
+ const bPrev = b;
+
+ if (trend === null && seasonal === null) {
+ l = alpha * yt + (1 - alpha) * l;
+ } else if (trend !== null && seasonal === null) {
+ l = alpha * yt + (1 - alpha) * (l + phi * b);
+ if (beta !== null) {
+ b = beta * (l - lPrev) + (1 - beta) * phi * bPrev;
+ }
+ } else if (trend === null && seasonal === "add") {
+ l = alpha * (yt - st_m) + (1 - alpha) * l;
+ if (gamma !== null) {
+ seasonals[sIdx] = gamma * (yt - l) + (1 - gamma) * st_m;
+ }
+ } else if (trend === null && seasonal === "mul") {
+ l = alpha * (st_m !== 0 ? yt / st_m : yt) + (1 - alpha) * l;
+ if (gamma !== null) {
+ seasonals[sIdx] = gamma * (l !== 0 ? yt / l : 1) + (1 - gamma) * st_m;
+ }
+ } else if (trend === "add" && seasonal === "add") {
+ l = alpha * (yt - st_m) + (1 - alpha) * (lPrev + phi * bPrev);
+ if (beta !== null) {
+ b = beta * (l - lPrev) + (1 - beta) * phi * bPrev;
+ }
+ if (gamma !== null) {
+ seasonals[sIdx] = gamma * (yt - l) + (1 - gamma) * st_m;
+ }
+ } else if (trend === "add" && seasonal === "mul") {
+ l = alpha * (st_m !== 0 ? yt / st_m : yt) + (1 - alpha) * (lPrev + phi * bPrev);
+ if (beta !== null) {
+ b = beta * (l - lPrev) + (1 - beta) * phi * bPrev;
+ }
+ if (gamma !== null) {
+ seasonals[sIdx] = gamma * (l + phi * b !== 0 ? yt / (l + phi * b) : 1) + (1 - gamma) * st_m;
+ }
+ } else {
+ // multiplicative trend — approximate as additive for stability
+ l = alpha * yt + (1 - alpha) * (lPrev + phi * bPrev);
+ if (beta !== null) {
+ b = beta * (l - lPrev) + (1 - beta) * phi * bPrev;
+ }
+ if (gamma !== null) {
+ seasonals[sIdx] = gamma * (yt - l) + (1 - gamma) * st_m;
+ }
+ }
+ }
+
+ return {
+ fitted,
+ residuals,
+ sse,
+ finalL: l,
+ finalB: b,
+ finalS: seasonals.slice(),
+ };
+}
+
+/**
+ * Generate h-step forecasts from the final ETS state.
+ * `finalS` is the circular buffer of the last m seasonal indices where
+ * `finalS[t % m]` = s_{t+1-m} (same convention as hwPass).
+ */
+function hwForecast(
+ steps: number,
+ l: number,
+ b: number,
+ finalS: readonly number[],
+ phi: number,
+ trend: ETSTrend,
+ seasonal: ETSSeasonal,
+ m: number,
+ n: number, // length of training series (to compute season offsets)
+): number[] {
+ const out: number[] = [];
+ let phiH = phi;
+ let phiSum = phi;
+
+ // At t = n-1 (last training obs), seasonals[t % m] has just been updated.
+ // For forecast step h, the seasonal index corresponds to position (n-1+h) % m in
+ // the buffer (shifted by 1 because sIdx = (t % m) in hwPass at time t = n-1+h).
+ for (let h = 1; h <= steps; h++) {
+ const sIdx = (((n - 1 + h) % m) + m) % m;
+ const sVal = finalS[sIdx] ?? 1;
+
+ let yhat: number;
+ if (trend === null && seasonal === null) {
+ yhat = l;
+ } else if (trend !== null && seasonal === null) {
+ yhat = l + phiSum * b;
+ } else if (trend === null && seasonal === "add") {
+ yhat = l + sVal;
+ } else if (trend === null && seasonal === "mul") {
+ yhat = l * sVal;
+ } else if (trend === "add" && seasonal === "add") {
+ yhat = l + phiSum * b + sVal;
+ } else if (trend === "add" && seasonal === "mul") {
+ yhat = (l + phiSum * b) * sVal;
+ } else if (trend === "mul" && seasonal === "add") {
+ yhat = l * b ** phiSum + sVal;
+ } else {
+ yhat = l * b ** phiSum * sVal;
+ }
+
+ out.push(yhat);
+
+ phiH *= phi;
+ phiSum += phiH;
+ }
+ return out;
+}
+
+// ─── Heuristic initialisation ─────────────────────────────────────────────────
+
+/**
+ * Compute heuristic initial level and trend.
+ * Uses mean of first season + linear regression slope on first two seasons.
+ */
+function heuristicInit(
+ y: readonly number[],
+ m: number,
+ hasTrend: boolean,
+): { l0: number; b0: number } {
+ const n = y.length;
+ if (!hasTrend) {
+ return { l0: y[0] ?? 0, b0: 0 };
+ }
+ // Use first season average as l0, slope between first two seasons as b0
+ const k = Math.min(m, n);
+ let s1 = 0;
+ for (let i = 0; i < k; i++) {
+ s1 += y[i] ?? 0;
+ }
+ const l0 = s1 / k;
+
+ if (n >= 2 * m) {
+ let s2 = 0;
+ for (let i = m; i < 2 * m; i++) {
+ s2 += y[i] ?? 0;
+ }
+ const b0 = (s2 / m - l0) / m;
+ return { l0, b0: b0 || (((y[1] ?? 0) - (y[0] ?? 0)) * m) / m };
+ }
+ // Fallback: slope between y[0] and y[n-1]
+ const b0 = n > 1 ? ((y[n - 1] ?? 0) - (y[0] ?? 0)) / (n - 1) : 0;
+ return { l0, b0 };
+}
+
+/**
+ * Compute heuristic seasonal indices.
+ * Additive: s_j = avg(y_j, y_{j+m}, …) − overall mean
+ * Multiplicative: s_j = avg(y_j, y_{j+m}, …) / overall mean
+ */
+function heuristicSeasons(
+ y: readonly number[],
+ m: number,
+ seasonal: "add" | "mul",
+ l0: number,
+ b0: number,
+): number[] {
+ const n = y.length;
+ const cycles = Math.max(1, Math.floor(n / m));
+
+ // Detrended values for each position in the seasonal cycle
+ const byPos: number[][] = Array.from({ length: m }, () => []);
+ for (let t = 0; t < cycles * m && t < n; t++) {
+ const trend = l0 + b0 * t;
+ const raw = y[t] ?? 0;
+ const pos = t % m;
+ const posArr = byPos[pos];
+ if (posArr !== undefined) {
+ if (seasonal === "add") {
+ posArr.push(raw - trend);
+ } else {
+ posArr.push(trend !== 0 ? raw / trend : 1);
+ }
+ }
+ }
+
+ const rawS = byPos.map((vals) => {
+ if (vals.length === 0) {
+ return seasonal === "add" ? 0 : 1;
+ }
+ return vals.reduce((a, b) => a + b, 0) / vals.length;
+ });
+
+ // Normalise so seasonal indices sum to 0 (additive) or m (multiplicative)
+ if (seasonal === "add") {
+ const mean = rawS.reduce((a, b) => a + b, 0) / m;
+ return rawS.map((v) => v - mean);
+ }
+ const mean = rawS.reduce((a, b) => a + b, 0) / m;
+ return rawS.map((v) => (mean !== 0 ? (v / mean) * 1 : 1));
+}
+
+// ─── SimpleExpSmoothing ──────────────────────────────────────────────────────
+
+/**
+ * Simple Exponential Smoothing (SES / ETS(A,N,N)).
+ *
+ * Produces level-only forecasts; all future forecasts equal the final level.
+ *
+ * @example
+ * ```ts
+ * import { SimpleExpSmoothing } from "tsb";
+ * const model = new SimpleExpSmoothing();
+ * const fit = model.fit([3, 5, 4, 6, 5, 8, 7]);
+ * console.log(fit.alpha, fit.sse);
+ * console.log(model.forecast(3)); // [level, level, level]
+ * ```
+ */
+export class SimpleExpSmoothing {
+ private _fit: SESFitResult | null = null;
+ private _finalLevel = 0;
+
+ /**
+ * Fit the SES model to observed data.
+ * @param y - Time series observations.
+ * @param opts - Optional fixed parameters.
+ */
+ fit(y: readonly number[] | Series, opts?: SESOptions): SESFitResult {
+ const arr = toArr(y);
+ const n = arr.length;
+ if (n < 2) {
+ throw new RangeError("SimpleExpSmoothing requires at least 2 observations");
+ }
+
+ const l0Init = opts?.initialLevel ?? arr[0] ?? 0;
+
+ let alpha: number;
+ let l0: number;
+
+ if (opts?.alpha !== undefined) {
+ alpha = clamp(opts.alpha, 1e-6, 1 - 1e-6);
+ l0 = l0Init;
+ } else if (opts?.initialLevel !== undefined) {
+ const result = nelderMead(
+ ([a]: readonly number[]) => sesPass(arr, a ?? 0.3, l0Init).sse,
+ [0.3],
+ [[1e-6, 1 - 1e-6]],
+ );
+ alpha = result.params[0] ?? 0.3;
+ l0 = l0Init;
+ } else {
+ // Optimise α (and optionally l0)
+ const result = nelderMead(
+ ([a, l]: readonly number[]) => sesPass(arr, a ?? 0.3, l ?? arr[0] ?? 0).sse,
+ [0.3, l0Init],
+ [
+ [1e-6, 1 - 1e-6],
+ [
+ (arr[0] ?? 0) - Math.abs(arr[0] ?? 0) * 5 - 1,
+ (arr[0] ?? 0) + Math.abs(arr[0] ?? 0) * 5 + 1,
+ ],
+ ],
+ );
+ alpha = result.params[0] ?? 0.3;
+ l0 = result.params[1] ?? l0Init;
+ }
+
+ const { fitted, residuals, sse } = sesPass(arr, alpha, l0);
+
+ // Final level for forecasting
+ let lFinal = l0;
+ for (const yt of arr) {
+ lFinal = alpha * yt + (1 - alpha) * lFinal;
+ }
+ this._finalLevel = lFinal;
+
+ const k = 2; // alpha + l0
+ const { aic, bic, aicc } = infoGaussian(sse, n, k);
+
+ const result: SESFitResult = {
+ alpha,
+ initialLevel: l0,
+ fittedValues: fitted,
+ residuals,
+ sse,
+ aic,
+ bic,
+ aicc,
+ };
+ this._fit = result;
+ return result;
+ }
+
+ /**
+ * Generate `steps` forecasts from the last fitted state.
+ * All forecasts equal the final level (flat forecast).
+ * Must call {@link fit} first.
+ */
+ forecast(steps: number): number[] {
+ if (this._fit === null) {
+ throw new Error("Call fit() before forecast()");
+ }
+ return new Array(steps).fill(this._finalLevel);
+ }
+}
+
+// ─── Holt ─────────────────────────────────────────────────────────────────────
+
+/**
+ * Holt's linear (double) exponential smoothing (ETS(A,A,N) / ETS(A,Ad,N)).
+ *
+ * Extends SES with a trend component; supports optional damping.
+ *
+ * @example
+ * ```ts
+ * import { Holt } from "tsb";
+ * const model = new Holt();
+ * const fit = model.fit([3, 5, 4, 6, 5, 8, 7, 9, 8, 11]);
+ * console.log(fit.alpha, fit.beta, fit.phi);
+ * console.log(model.forecast(5));
+ * ```
+ */
+export class Holt {
+ private _opts: HoltOptions = {};
+ private _fit: HoltFitResult | null = null;
+ private _finalL = 0;
+ private _finalB = 0;
+ private _phi = 1;
+
+ constructor(opts?: HoltOptions) {
+ this._opts = opts ?? {};
+ }
+
+ /**
+ * Fit Holt's model to observed data.
+ * @param y - Time series observations.
+ * @param opts - Optional parameter overrides (merged with constructor opts).
+ */
+ fit(y: readonly number[] | Series, opts?: HoltOptions): HoltFitResult {
+ const arr = toArr(y);
+ const n = arr.length;
+ if (n < 3) {
+ throw new RangeError("Holt requires at least 3 observations");
+ }
+
+ const merged: HoltOptions = { ...this._opts, ...opts };
+ const damped = merged.damped ?? false;
+
+ const { l0: l0h, b0: b0h } = heuristicInit(arr, n, true);
+
+ const alphaFixed = merged.alpha;
+ const betaFixed = merged.beta;
+ const phiFixed = merged.dampingSlope;
+ const l0Fixed = merged.initialLevel ?? l0h;
+ const b0Fixed = merged.initialTrend ?? b0h;
+
+ // Build optimisation bounds and initial point
+ type Bound = [number, number];
+ const paramNames: string[] = [];
+ const x0: number[] = [];
+ const bounds: Bound[] = [];
+
+ if (alphaFixed === undefined) {
+ paramNames.push("alpha");
+ x0.push(0.3);
+ bounds.push([1e-6, 1 - 1e-6]);
+ }
+ if (betaFixed === undefined) {
+ paramNames.push("beta");
+ x0.push(0.1);
+ bounds.push([1e-6, 1 - 1e-6]);
+ }
+ if (damped && phiFixed === undefined) {
+ paramNames.push("phi");
+ x0.push(0.98);
+ bounds.push([0.8, 1 - 1e-6]);
+ }
+ // Always optimise l0 and b0 if not fixed
+ const optimL0 = merged.initialLevel === undefined;
+ const optimB0 = merged.initialTrend === undefined;
+ if (optimL0) {
+ paramNames.push("l0");
+ x0.push(l0h);
+ const spread = Math.abs(l0h) * 2 + 10;
+ bounds.push([l0h - spread, l0h + spread]);
+ }
+ if (optimB0) {
+ paramNames.push("b0");
+ x0.push(b0h);
+ const spread = Math.abs(b0h) * 5 + 1;
+ bounds.push([b0h - spread, b0h + spread]);
+ }
+
+ let alpha = alphaFixed ?? 0.3;
+ let beta = betaFixed ?? 0.1;
+ let phi = damped ? (phiFixed ?? 0.98) : 1.0;
+ let l0 = l0Fixed;
+ let b0 = b0Fixed;
+
+ if (x0.length > 0) {
+ const result = nelderMead(
+ (params: readonly number[]): number => {
+ let a = alphaFixed ?? params[paramNames.indexOf("alpha")] ?? 0.3;
+ let bta = betaFixed ?? params[paramNames.indexOf("beta")] ?? 0.1;
+ let ph = damped ? (phiFixed ?? params[paramNames.indexOf("phi")] ?? 0.98) : 1.0;
+ const ll0 = optimL0 ? (params[paramNames.indexOf("l0")] ?? l0h) : l0Fixed;
+ const lb0 = optimB0 ? (params[paramNames.indexOf("b0")] ?? b0h) : b0Fixed;
+ a = clamp(a, 1e-6, 1 - 1e-6);
+ bta = clamp(bta, 1e-6, 1 - 1e-6);
+ if (damped) {
+ ph = clamp(ph, 0.8, 1 - 1e-6);
+ }
+ return holtPass(arr, a, bta, ph, ll0, lb0).sse;
+ },
+ x0,
+ bounds,
+ );
+ const p = result.params;
+ alpha = alphaFixed ?? p[paramNames.indexOf("alpha")] ?? alpha;
+ beta = betaFixed ?? p[paramNames.indexOf("beta")] ?? beta;
+ phi = damped ? (phiFixed ?? p[paramNames.indexOf("phi")] ?? phi) : 1.0;
+ l0 = optimL0 ? (p[paramNames.indexOf("l0")] ?? l0) : l0Fixed;
+ b0 = optimB0 ? (p[paramNames.indexOf("b0")] ?? b0) : b0Fixed;
+ }
+
+ alpha = clamp(alpha, 1e-6, 1 - 1e-6);
+ beta = clamp(beta, 1e-6, 1 - 1e-6);
+ if (damped) {
+ phi = clamp(phi, 0.8, 1 - 1e-6);
+ }
+
+ const { fitted, residuals, sse } = holtPass(arr, alpha, beta, phi, l0, b0);
+
+ // Final state for forecasting
+ let lF = l0;
+ let bF = b0;
+ for (const yt of arr) {
+ const lNew = alpha * yt + (1 - alpha) * (lF + phi * bF);
+ bF = beta * (lNew - lF) + (1 - beta) * phi * bF;
+ lF = lNew;
+ }
+ this._finalL = lF;
+ this._finalB = bF;
+ this._phi = phi;
+
+ const k = 2 + (damped ? 1 : 0) + 2; // alpha + beta + phi? + l0 + b0
+ const { aic, bic, aicc } = infoGaussian(sse, n, k);
+
+ const result: HoltFitResult = {
+ alpha,
+ beta,
+ phi,
+ initialLevel: l0,
+ initialTrend: b0,
+ fittedValues: fitted,
+ residuals,
+ sse,
+ aic,
+ bic,
+ aicc,
+ };
+ this._fit = result;
+ return result;
+ }
+
+ /**
+ * Generate `steps` forecasts from the last fitted state.
+ * Must call {@link fit} first.
+ */
+ forecast(steps: number): number[] {
+ if (this._fit === null) {
+ throw new Error("Call fit() before forecast()");
+ }
+ return holtForecast(steps, this._finalL, this._finalB, this._phi);
+ }
+}
+
+// ─── ExponentialSmoothing (Holt-Winters) ─────────────────────────────────────
+
+/**
+ * Holt-Winters Exponential Smoothing — full ETS model.
+ *
+ * Supports all combinations of additive / multiplicative trend and seasonal
+ * components with optional damped trend.
+ *
+ * @example
+ * ```ts
+ * import { ExponentialSmoothing } from "tsb";
+ *
+ * // Monthly data with additive seasonal component
+ * const y = [17, 21, 23, 18, 22, 26, 19, 24, 27, 20, 25, 28,
+ * 18, 23, 25, 20, 24, 28, 21, 26, 29, 22, 27, 30];
+ * const model = new ExponentialSmoothing({ trend: "add", seasonal: "add", seasonalPeriods: 12 });
+ * const fit = model.fit(y);
+ * console.log(fit.alpha, fit.beta, fit.gamma, fit.aic);
+ * console.log(model.forecast(12));
+ * ```
+ */
+export class ExponentialSmoothing {
+ private _opts: ExponentialSmoothingOptions;
+ private _fit: ExponentialSmoothingFitResult | null = null;
+ private _finalL = 0;
+ private _finalB = 0;
+ private _finalS: number[] = [];
+ private _phi = 1;
+ private _trend: ETSTrend = null;
+ private _seasonal: ETSSeasonal = null;
+ private _m = 1;
+ private _n = 0;
+ private _sigma2 = 1;
+
+ constructor(opts?: ExponentialSmoothingOptions) {
+ this._opts = opts ?? {};
+ }
+
+ /**
+ * Fit the Holt-Winters model to observed data.
+ * @param y - Time series observations.
+ * @param opts - Optional parameter overrides.
+ */
+ fit(
+ y: readonly number[] | Series,
+ opts?: ExponentialSmoothingOptions,
+ ): ExponentialSmoothingFitResult {
+ const arr = toArr(y);
+ const n = arr.length;
+ if (n < 3) {
+ throw new RangeError("ExponentialSmoothing requires at least 3 observations");
+ }
+
+ const merged: ExponentialSmoothingOptions = { ...this._opts, ...opts };
+ const trend = merged.trend ?? null;
+ const seasonal = merged.seasonal ?? null;
+ const damped = merged.damped ?? false;
+ const m = merged.seasonalPeriods ?? (seasonal !== null ? 2 : 1);
+
+ this._trend = trend;
+ this._seasonal = seasonal;
+ this._m = m;
+ this._n = n;
+
+ if (seasonal !== null && n < 2 * m) {
+ throw new RangeError(
+ `ExponentialSmoothing: need at least 2 full seasonal periods (${2 * m} obs), got ${n}`,
+ );
+ }
+
+ const initMethod = merged.initializationMethod ?? "heuristic";
+
+ // Heuristic initialisation
+ const { l0: l0h, b0: b0h } = heuristicInit(arr, seasonal === null ? n : m, trend !== null);
+ const s0h = seasonal !== null ? heuristicSeasons(arr, m, seasonal, l0h, b0h) : null;
+
+ // Determine which params to optimise
+ const alphaFixed = merged.alpha;
+ const betaFixed = merged.beta;
+ const gammaFixed = merged.gamma;
+ const phiFixed = merged.phi;
+ const l0Fixed = initMethod === "known" ? (merged.initialLevel ?? l0h) : undefined;
+ const b0Fixed = initMethod === "known" ? (merged.initialTrend ?? b0h) : undefined;
+ const s0Fixed =
+ initMethod === "known" && merged.initialSeasons !== undefined
+ ? merged.initialSeasons.slice()
+ : undefined;
+
+ type Bound = [number, number];
+ const paramNames: string[] = [];
+ const x0: number[] = [];
+ const bounds: Bound[] = [];
+
+ if (alphaFixed === undefined) {
+ paramNames.push("alpha");
+ x0.push(0.3);
+ bounds.push([1e-6, 1 - 1e-6]);
+ }
+ if (trend !== null && betaFixed === undefined) {
+ paramNames.push("beta");
+ x0.push(0.1);
+ bounds.push([1e-6, 1 - 1e-6]);
+ }
+ if (seasonal !== null && gammaFixed === undefined) {
+ paramNames.push("gamma");
+ x0.push(0.1);
+ bounds.push([1e-6, 1 - 1e-6]);
+ }
+ if (damped && phiFixed === undefined) {
+ paramNames.push("phi");
+ x0.push(0.98);
+ bounds.push([0.8, 1 - 1e-6]);
+ }
+
+ const optimL0 = initMethod !== "known" && merged.initialLevel === undefined;
+ const optimB0 = trend !== null && initMethod !== "known" && merged.initialTrend === undefined;
+ const optimS0 =
+ seasonal !== null && initMethod !== "known" && merged.initialSeasons === undefined;
+
+ if (optimL0) {
+ paramNames.push("l0");
+ x0.push(l0h);
+ const sp = Math.abs(l0h) * 2 + 10;
+ bounds.push([l0h - sp, l0h + sp]);
+ }
+ if (optimB0) {
+ paramNames.push("b0");
+ x0.push(b0h);
+ const sp = Math.abs(b0h) * 5 + 1;
+ bounds.push([b0h - sp, b0h + sp]);
+ }
+ if (optimS0 && s0h !== null) {
+ for (let j = 0; j < m; j++) {
+ paramNames.push(`s0_${j}`);
+ x0.push(s0h[j] ?? 0);
+ const sp = Math.abs(s0h[j] ?? 0) * 5 + 1;
+ bounds.push([(s0h[j] ?? 0) - sp, (s0h[j] ?? 0) + sp]);
+ }
+ }
+
+ let alpha = alphaFixed ?? 0.3;
+ let beta = trend !== null ? (betaFixed ?? 0.1) : null;
+ let gamma = seasonal !== null ? (gammaFixed ?? 0.1) : null;
+ let phi = damped ? (phiFixed ?? 0.98) : 1.0;
+ let l0 = l0Fixed ?? l0h;
+ let b0 = b0Fixed ?? b0h;
+ let s0 = s0Fixed ?? s0h;
+
+ if (x0.length > 0) {
+ const res = nelderMead(
+ (params: readonly number[]): number => {
+ let a = alphaFixed ?? params[paramNames.indexOf("alpha")] ?? 0.3;
+ let bt = trend !== null ? (betaFixed ?? params[paramNames.indexOf("beta")] ?? 0.1) : null;
+ let gm =
+ seasonal !== null ? (gammaFixed ?? params[paramNames.indexOf("gamma")] ?? 0.1) : null;
+ let ph = damped ? (phiFixed ?? params[paramNames.indexOf("phi")] ?? 0.98) : 1.0;
+ a = clamp(a, 1e-6, 1 - 1e-6);
+ if (bt !== null) {
+ bt = clamp(bt, 1e-6, 1 - 1e-6);
+ }
+ if (gm !== null) {
+ gm = clamp(gm, 1e-6, 1 - 1e-6);
+ }
+ if (damped) {
+ ph = clamp(ph, 0.8, 1 - 1e-6);
+ }
+
+ const ll0 = optimL0 ? (params[paramNames.indexOf("l0")] ?? l0h) : (l0Fixed ?? l0h);
+ const lb0 = optimB0 ? (params[paramNames.indexOf("b0")] ?? b0h) : (b0Fixed ?? b0h);
+ let ss0: number[] | null = null;
+ if (optimS0) {
+ ss0 = [];
+ for (let j = 0; j < m; j++) {
+ ss0.push(params[paramNames.indexOf(`s0_${j}`)] ?? s0h?.[j] ?? 0);
+ }
+ } else {
+ ss0 = s0Fixed ?? s0h;
+ }
+
+ return hwPass(arr, a, bt, gm, ph, ll0, lb0, ss0, trend, seasonal, m).sse;
+ },
+ x0,
+ bounds,
+ seasonal !== null ? 5000 : 3000,
+ );
+ const p = res.params;
+ alpha = alphaFixed ?? p[paramNames.indexOf("alpha")] ?? alpha;
+ beta = trend !== null ? (betaFixed ?? p[paramNames.indexOf("beta")] ?? beta ?? 0.1) : null;
+ gamma =
+ seasonal !== null ? (gammaFixed ?? p[paramNames.indexOf("gamma")] ?? gamma ?? 0.1) : null;
+ phi = damped ? (phiFixed ?? p[paramNames.indexOf("phi")] ?? phi) : 1.0;
+ l0 = optimL0 ? (p[paramNames.indexOf("l0")] ?? l0) : (l0Fixed ?? l0);
+ b0 = optimB0 ? (p[paramNames.indexOf("b0")] ?? b0) : (b0Fixed ?? b0);
+ if (optimS0) {
+ s0 = [];
+ for (let j = 0; j < m; j++) {
+ s0.push(p[paramNames.indexOf(`s0_${j}`)] ?? s0h?.[j] ?? 0);
+ }
+ }
+ }
+
+ if (optimS0 && optimL0 && seasonal === "add" && trend !== "mul" && s0 !== null) {
+ // Level and additive seasons share an arbitrary offset. Choose zero-mean
+ // seasons and transfer their mean to the level without changing forecasts.
+ const seasonalMean = s0.reduce((sum, value) => sum + value, 0) / m;
+ s0 = s0.map((value) => value - seasonalMean);
+ l0 += seasonalMean;
+ }
+
+ // Final clamp
+ alpha = clamp(alpha, 1e-6, 1 - 1e-6);
+ if (beta !== null) {
+ beta = clamp(beta, 1e-6, 1 - 1e-6);
+ }
+ if (gamma !== null) {
+ gamma = clamp(gamma, 1e-6, 1 - 1e-6);
+ }
+ if (damped) {
+ phi = clamp(phi, 0.8, 1 - 1e-6);
+ }
+
+ const { fitted, residuals, sse, finalL, finalB, finalS } = hwPass(
+ arr,
+ alpha,
+ beta,
+ gamma,
+ phi,
+ l0,
+ b0,
+ s0,
+ trend,
+ seasonal,
+ m,
+ );
+
+ this._finalL = finalL;
+ this._finalB = finalB;
+ this._finalS = finalS;
+ this._phi = phi;
+ this._sigma2 = Math.max(sse / n, 1e-15);
+
+ // Number of free parameters for information criteria
+ const nParams =
+ 1 + // alpha
+ (trend !== null ? 1 : 0) + // beta
+ (seasonal !== null ? 1 : 0) + // gamma
+ (damped ? 1 : 0) + // phi
+ 1 + // l0
+ (trend !== null ? 1 : 0) + // b0
+ (seasonal !== null ? m : 0); // seasonal indices
+
+ const { logLikelihood, aic, bic, aicc } = infoGaussian(sse, n, nParams);
+
+ const result: ExponentialSmoothingFitResult = {
+ alpha,
+ beta,
+ gamma,
+ phi,
+ initialLevel: l0,
+ initialTrend: trend !== null ? b0 : null,
+ initialSeasons: seasonal !== null ? (s0 ?? null) : null,
+ fittedValues: fitted,
+ residuals,
+ sse,
+ logLikelihood,
+ aic,
+ bic,
+ aicc,
+ };
+ this._fit = result;
+ return result;
+ }
+
+ /**
+ * Generate `steps` point forecasts from the final state.
+ * Must call {@link fit} first.
+ */
+ forecast(steps: number): number[] {
+ if (this._fit === null) {
+ throw new Error("Call fit() before forecast()");
+ }
+ return hwForecast(
+ steps,
+ this._finalL,
+ this._finalB,
+ this._finalS,
+ this._phi,
+ this._trend,
+ this._seasonal,
+ this._m,
+ this._n,
+ );
+ }
+
+ /**
+ * Generate `steps` forecasts with (1 − `alpha_ci`) prediction intervals.
+ * Uses additive-error variance approximation (constant σ² scaled by h).
+ * Must call {@link fit} first.
+ *
+ * @param steps - Number of steps ahead.
+ * @param alpha_ci - Significance level (default 0.05 → 95 % intervals).
+ */
+ forecastWithCI(steps: number, alpha_ci = 0.05): ETSForecastResult {
+ const fc = this.forecast(steps);
+ // Normal quantile for (1 - alpha_ci/2)
+ const z = normalQuantile(1 - alpha_ci / 2);
+ const sigma = Math.sqrt(this._sigma2);
+ const lower: number[] = [];
+ const upper: number[] = [];
+ const stderr: number[] = [];
+ for (let h = 1; h <= steps; h++) {
+ // Variance grows linearly with h for additive-error models
+ const se = sigma * Math.sqrt(h);
+ stderr.push(se);
+ lower.push((fc[h - 1] ?? 0) - z * se);
+ upper.push((fc[h - 1] ?? 0) + z * se);
+ }
+ return { forecast: fc, lower, upper, stderr };
+ }
+}
+
+/**
+ * Rational approximation of the normal quantile function (Abramowitz & Stegun).
+ * @internal
+ */
+function normalQuantile(p: number): number {
+ if (p <= 0) {
+ return Number.NEGATIVE_INFINITY;
+ }
+ if (p >= 1) {
+ return Number.POSITIVE_INFINITY;
+ }
+ const a = [2.515517, 0.802853, 0.010328];
+ const b = [1.432788, 0.189269, 0.001308];
+ const t = Math.sqrt(-2 * Math.log(p < 0.5 ? p : 1 - p));
+ const num = (a[0] ?? 0) + (a[1] ?? 0) * t + (a[2] ?? 0) * t * t;
+ const den = 1 + (b[0] ?? 0) * t + (b[1] ?? 0) * t * t + (b[2] ?? 0) * t * t * t;
+ const x = t - num / den;
+ return p < 0.5 ? -x : x;
+}
+
+// ─── Convenience functions ────────────────────────────────────────────────────
+
+/**
+ * Fit a Simple Exponential Smoothing model and return the result.
+ *
+ * @example
+ * ```ts
+ * import { simpleExpSmoothing } from "tsb";
+ * const { alpha, fittedValues, sse } = simpleExpSmoothing([3, 5, 4, 6, 5]);
+ * ```
+ */
+export function simpleExpSmoothing(
+ y: readonly number[] | Series,
+ opts?: SESOptions,
+): SESFitResult {
+ return new SimpleExpSmoothing().fit(y, opts);
+}
+
+/**
+ * Fit a Holt linear trend model and return the result.
+ *
+ * @example
+ * ```ts
+ * import { holt } from "tsb";
+ * const fit = holt([3, 5, 4, 6, 5, 8, 7, 9]);
+ * console.log(fit.alpha, fit.beta, fit.sse);
+ * ```
+ */
+export function holt(y: readonly number[] | Series, opts?: HoltOptions): HoltFitResult {
+ return new Holt(opts).fit(y);
+}
+
+/**
+ * Fit a full Holt-Winters Exponential Smoothing model and return the result.
+ *
+ * @example
+ * ```ts
+ * import { fitEts } from "tsb";
+ * const y = [17, 21, 23, 18, 22, 26, 19, 24, 27, 20, 25, 28];
+ * const fit = fitEts(y, { trend: "add", seasonal: "add", seasonalPeriods: 4 });
+ * console.log(fit.alpha, fit.beta, fit.gamma, fit.aic);
+ * ```
+ */
+export function fitEts(
+ y: readonly number[] | Series,
+ opts?: ExponentialSmoothingOptions,
+): ExponentialSmoothingFitResult {
+ return new ExponentialSmoothing(opts).fit(y);
+}
diff --git a/src/stats/extreme_value.ts b/src/stats/extreme_value.ts
new file mode 100644
index 000000000..e97b01c7e
--- /dev/null
+++ b/src/stats/extreme_value.ts
@@ -0,0 +1,441 @@
+/**
+ * extreme_value — Extreme Value Theory (EVT) distributions and analysis.
+ *
+ * Implements:
+ * - **Generalized Extreme Value (GEV)** distribution (Gumbel, Fréchet, Weibull)
+ * - **Generalized Pareto Distribution (GPD)** for Peaks Over Threshold
+ * - **Block maxima** method (GEV fitting via L-moments)
+ * - **Peaks Over Threshold (POT)** method (GPD fitting via MLE)
+ * - **Return level** and **return period** calculations
+ * - **Gumbel**, **Fréchet**, **Weibull** special-case distributions
+ *
+ * @module
+ */
+
+// ─── GEV Distribution ─────────────────────────────────────────────────────────
+
+/**
+ * Parameters of the Generalized Extreme Value distribution.
+ */
+export interface GEVParams {
+ /** Location parameter (mu). */
+ mu: number;
+ /** Scale parameter (sigma > 0). */
+ sigma: number;
+ /** Shape parameter (xi). xi=0: Gumbel, xi>0: Fréchet, xi<0: Weibull. */
+ xi: number;
+}
+
+/**
+ * GEV probability density function.
+ *
+ * @param x - Value.
+ * @param params - GEV parameters.
+ * @returns Density f(x).
+ *
+ * @example
+ * ```ts
+ * import { gevPdf } from "tsb";
+ * const p = gevPdf(2.5, { mu: 0, sigma: 1, xi: 0.1 });
+ * ```
+ */
+export function gevPdf(x: number, params: GEVParams): number {
+ const { mu, sigma, xi } = params;
+ if (sigma <= 0) return 0;
+
+ const z = (x - mu) / sigma;
+
+ if (Math.abs(xi) < 1e-8) {
+ // Gumbel case (xi → 0)
+ const t = Math.exp(-z);
+ return (1 / sigma) * Math.exp(-z - t);
+ }
+
+ const t = 1 + xi * z;
+ if (t <= 0) return 0;
+
+ return (1 / sigma) * t ** (-1 / xi - 1) * Math.exp(-(t ** (-1 / xi)));
+}
+
+/**
+ * GEV cumulative distribution function.
+ *
+ * @param x - Value.
+ * @param params - GEV parameters.
+ * @returns CDF F(x).
+ */
+export function gevCdf(x: number, params: GEVParams): number {
+ const { mu, sigma, xi } = params;
+ if (sigma <= 0) return 0;
+
+ const z = (x - mu) / sigma;
+
+ if (Math.abs(xi) < 1e-8) {
+ // Gumbel case
+ return Math.exp(-Math.exp(-z));
+ }
+
+ const t = 1 + xi * z;
+ if (t <= 0) {
+ return xi > 0 ? 0 : 1;
+ }
+
+ return Math.exp(-(t ** (-1 / xi)));
+}
+
+/**
+ * GEV quantile function (inverse CDF).
+ *
+ * @param p - Probability in (0, 1).
+ * @param params - GEV parameters.
+ * @returns Quantile x such that F(x) = p.
+ */
+export function gevQuantile(p: number, params: GEVParams): number {
+ if (p <= 0) return -Number.POSITIVE_INFINITY;
+ if (p >= 1) return Number.POSITIVE_INFINITY;
+
+ const { mu, sigma, xi } = params;
+
+ if (Math.abs(xi) < 1e-8) {
+ // Gumbel
+ return mu - sigma * Math.log(-Math.log(p));
+ }
+
+ return mu + (sigma / xi) * ((-Math.log(p)) ** -xi - 1);
+}
+
+/**
+ * Compute return level for a given return period (years/blocks).
+ *
+ * @param returnPeriod - Return period (e.g., 100 for 100-year event).
+ * @param params - GEV parameters.
+ * @returns Return level x such that P(X > x) = 1/returnPeriod.
+ */
+export function gevReturnLevel(returnPeriod: number, params: GEVParams): number {
+ const p = 1 - 1 / returnPeriod;
+ return gevQuantile(p, params);
+}
+
+// ─── GEV Fitting (L-moments) ──────────────────────────────────────────────────
+
+/**
+ * Fit GEV distribution to block maxima using L-moments.
+ *
+ * @param maxima - Array of block maxima (e.g., annual maxima).
+ * @returns Estimated GEV parameters.
+ *
+ * @example
+ * ```ts
+ * import { fitGEV } from "tsb";
+ * const annualMaxima = [12.3, 15.1, 9.8, 18.2, 11.4, 14.7];
+ * const params = fitGEV(annualMaxima);
+ * ```
+ */
+export function fitGEV(maxima: number[]): GEVParams {
+ const n = maxima.length;
+ if (n < 3) return { mu: 0, sigma: 1, xi: 0 };
+
+ const sorted = [...maxima].sort((a, b) => a - b);
+
+ // Compute L-moments via probability-weighted moments
+ let b0 = 0;
+ let b1 = 0;
+ let b2 = 0;
+
+ for (let i = 0; i < n; i++) {
+ b0 += sorted[i] ?? 0;
+ b1 += (i / (n - 1)) * (sorted[i] ?? 0);
+ b2 += ((i * (i - 1)) / ((n - 1) * (n - 2))) * (sorted[i] ?? 0);
+ }
+ b0 /= n;
+ b1 /= n;
+ b2 /= n;
+
+ const l1 = b0;
+ const l2 = 2 * b1 - b0;
+ const l3 = 6 * b2 - 6 * b1 + b0;
+
+ // L-skewness
+ const tau3 = l2 > 1e-10 ? l3 / l2 : 0;
+
+ // Estimate xi from L-skewness using approximation
+ let xi: number;
+ if (Math.abs(tau3) < 1e-8) {
+ xi = 0;
+ } else {
+ // Rational approximation for xi from tau3
+ xi = estimateXiFromTau3(tau3);
+ }
+
+ let sigma: number;
+ let mu: number;
+
+ if (Math.abs(xi) < 1e-6) {
+ // Gumbel
+ sigma = l2 / Math.log(2);
+ mu = l1 - 0.5772156649 * sigma;
+ } else {
+ const g1 = gamma(1 - xi);
+ sigma = (l2 * xi) / ((1 - 2 ** -xi) * g1);
+ mu = l1 - (sigma * (g1 - 1)) / xi;
+ }
+
+ return {
+ mu,
+ sigma: Math.max(sigma, 1e-8),
+ xi,
+ };
+}
+
+// ─── GPD Distribution ─────────────────────────────────────────────────────────
+
+/**
+ * Parameters of the Generalized Pareto Distribution.
+ */
+export interface GPDParams {
+ /** Threshold (u). */
+ threshold: number;
+ /** Scale parameter (sigma > 0). */
+ sigma: number;
+ /** Shape parameter (xi). */
+ xi: number;
+}
+
+/**
+ * GPD probability density function.
+ *
+ * @param x - Value (must be >= threshold).
+ * @param params - GPD parameters.
+ * @returns Density f(x).
+ */
+export function gpdPdf(x: number, params: GPDParams): number {
+ const { threshold, sigma, xi } = params;
+ if (sigma <= 0 || x < threshold) return 0;
+
+ const z = (x - threshold) / sigma;
+
+ if (Math.abs(xi) < 1e-8) {
+ return (1 / sigma) * Math.exp(-z);
+ }
+
+ const t = 1 + xi * z;
+ if (t <= 0) return 0;
+
+ return (1 / sigma) * t ** (-1 / xi - 1);
+}
+
+/**
+ * GPD cumulative distribution function.
+ *
+ * @param x - Value.
+ * @param params - GPD parameters.
+ * @returns CDF value.
+ */
+export function gpdCdf(x: number, params: GPDParams): number {
+ const { threshold, sigma, xi } = params;
+ if (x < threshold) return 0;
+
+ const z = (x - threshold) / sigma;
+
+ if (Math.abs(xi) < 1e-8) {
+ return 1 - Math.exp(-z);
+ }
+
+ const t = 1 + xi * z;
+ if (t <= 0) return xi > 0 ? 0 : 1;
+
+ return 1 - t ** (-1 / xi);
+}
+
+/**
+ * GPD quantile function.
+ *
+ * @param p - Probability in (0, 1).
+ * @param params - GPD parameters.
+ * @returns Quantile.
+ */
+export function gpdQuantile(p: number, params: GPDParams): number {
+ if (p <= 0) return params.threshold;
+ if (p >= 1) return Number.POSITIVE_INFINITY;
+
+ const { threshold, sigma, xi } = params;
+
+ if (Math.abs(xi) < 1e-8) {
+ return threshold - sigma * Math.log(1 - p);
+ }
+
+ return threshold + (sigma / xi) * ((1 - p) ** -xi - 1);
+}
+
+// ─── GPD Fitting (MLE) ────────────────────────────────────────────────────────
+
+/**
+ * Fit GPD to exceedances above a threshold using MLE.
+ *
+ * @param data - Full dataset.
+ * @param threshold - Threshold u (only exceedances x > u are used).
+ * @returns Estimated GPD parameters.
+ *
+ * @example
+ * ```ts
+ * import { fitGPD } from "tsb";
+ * const data = [1.2, 0.5, 3.4, 8.1, 0.2, 5.6, 12.3, 0.8, 4.1, 7.2];
+ * const params = fitGPD(data, 4.0);
+ * ```
+ */
+export function fitGPD(data: number[], threshold: number): GPDParams {
+ const exceedances = data.filter((x) => x > threshold).map((x) => x - threshold);
+
+ if (exceedances.length < 2) {
+ return { threshold, sigma: 1, xi: 0 };
+ }
+
+ const n = exceedances.length;
+ const mean = exceedances.reduce((s, x) => s + x, 0) / n;
+ const variance = exceedances.reduce((s, x) => s + (x - mean) ** 2, 0) / (n - 1);
+
+ // Method of moments starting values
+ let sigmaInit = (mean * ((mean * mean) / variance + 1)) / 2;
+ const xiInit = ((mean * mean) / variance - 1) / 2;
+
+ sigmaInit = Math.max(sigmaInit, 1e-6);
+
+ // Simple gradient-free optimization (Nelder-Mead would be ideal, use grid search)
+ const { sigma, xi } = optimizeGPD(exceedances, sigmaInit, xiInit);
+
+ return { threshold, sigma, xi };
+}
+
+// ─── Peaks Over Threshold Analysis ───────────────────────────────────────────
+
+/**
+ * Extract exceedances above a threshold.
+ *
+ * @param data - Time series data.
+ * @param threshold - Threshold value.
+ * @returns Array of exceedance values.
+ */
+export function extractExceedances(data: number[], threshold: number): number[] {
+ return data.filter((x) => x > threshold);
+}
+
+/**
+ * Compute return level from GPD fit.
+ *
+ * @param returnPeriod - Return period.
+ * @param params - Fitted GPD parameters.
+ * @param lambda - Rate of threshold exceedances (exceedances per time unit).
+ * @returns Return level.
+ */
+export function gpdReturnLevel(returnPeriod: number, params: GPDParams, lambda: number): number {
+ const p = 1 - 1 / (returnPeriod * lambda);
+ return gpdQuantile(Math.max(0, Math.min(1, p)), params);
+}
+
+// ─── Gumbel Distribution ──────────────────────────────────────────────────────
+
+/** Parameters for the Gumbel distribution (GEV with xi=0). */
+export interface GumbelParams {
+ /** Location (mu). */
+ mu: number;
+ /** Scale (beta > 0). */
+ beta: number;
+}
+
+/** Gumbel PDF. */
+export function gumbelPdf(x: number, params: GumbelParams): number {
+ const z = (x - params.mu) / params.beta;
+ return (1 / params.beta) * Math.exp(-(z + Math.exp(-z)));
+}
+
+/** Gumbel CDF. */
+export function gumbelCdf(x: number, params: GumbelParams): number {
+ const z = (x - params.mu) / params.beta;
+ return Math.exp(-Math.exp(-z));
+}
+
+/** Gumbel quantile. */
+export function gumbelQuantile(p: number, params: GumbelParams): number {
+ return params.mu - params.beta * Math.log(-Math.log(p));
+}
+
+// ─── Utilities ────────────────────────────────────────────────────────────────
+
+/** Lanczos approximation of the gamma function. */
+function gamma(z: number): number {
+ if (z < 0.5) {
+ return Math.PI / (Math.sin(Math.PI * z) * gamma(1 - z));
+ }
+ const g = 7;
+ const c = [
+ 0.99999999999980993, 676.5203681218851, -1259.1392167224028, 771.32342877765313,
+ -176.61502916214059, 12.507343278686905, -0.13857109526572012, 9.9843695780195716e-6,
+ 1.5056327351493116e-7,
+ ];
+ let x = c[0] ?? 0;
+ const zr = z - 1;
+ for (let i = 1; i < g + 2; i++) {
+ x += (c[i] ?? 0) / (zr + i);
+ }
+ const t = zr + g + 0.5;
+ return Math.sqrt(2 * Math.PI) * t ** (zr + 0.5) * Math.exp(-t) * x;
+}
+
+/** Estimate GEV shape parameter from L-skewness. */
+function estimateXiFromTau3(tau3: number): number {
+ // Approximation from Hosking (1985)
+ // For Gumbel: tau3 = 0.1699 (log(3/2) / log(2) - 2)
+ // Valid for -0.5 < xi < 0.5
+ const c = 2 / (3 + tau3) - Math.log(2) / Math.log(3);
+ return 7.859 * c + 2.9554 * c * c;
+}
+
+/** Simple optimization for GPD parameters. */
+function optimizeGPD(
+ exceedances: number[],
+ sigmaInit: number,
+ xiInit: number,
+): { sigma: number; xi: number } {
+ let sigma = sigmaInit;
+ let xi = xiInit;
+
+ const logLik = (s: number, x: number): number => {
+ const n = exceedances.length;
+ if (s <= 0) return -Number.POSITIVE_INFINITY;
+ let ll = -n * Math.log(s);
+ for (const e of exceedances) {
+ if (Math.abs(x) < 1e-8) {
+ ll -= e / s;
+ } else {
+ const t = 1 + (x * e) / s;
+ if (t <= 0) return -Number.POSITIVE_INFINITY;
+ ll -= (1 / x + 1) * Math.log(t);
+ }
+ }
+ return ll;
+ };
+
+ // Simple coordinate ascent
+ let bestLl = logLik(sigma, xi);
+ for (let iter = 0; iter < 100; iter++) {
+ const stepS = sigma * 0.1;
+ const stepX = 0.05;
+
+ for (const ds of [-stepS, stepS]) {
+ const ll = logLik(sigma + ds, xi);
+ if (ll > bestLl) {
+ bestLl = ll;
+ sigma = sigma + ds;
+ }
+ }
+ for (const dx of [-stepX, stepX]) {
+ const ll = logLik(sigma, xi + dx);
+ if (ll > bestLl) {
+ bestLl = ll;
+ xi = xi + dx;
+ }
+ }
+ }
+
+ return { sigma: Math.max(sigma, 1e-8), xi };
+}
diff --git a/src/stats/filters.ts b/src/stats/filters.ts
new file mode 100644
index 000000000..2d6eced66
--- /dev/null
+++ b/src/stats/filters.ts
@@ -0,0 +1,649 @@
+/**
+ * filters — Digital filter design and application.
+ *
+ * Mirrors `scipy.signal` filter utilities. Implemented from scratch with no
+ * external dependencies.
+ *
+ * Filter design:
+ * - {@link firwin} — FIR filter (windowed-sinc method)
+ * - {@link butter} — Butterworth IIR digital filter
+ *
+ * Frequency response:
+ * - {@link freqz} — frequency response of an FIR/IIR filter
+ * - {@link sosfreqz} — frequency response of SOS filter
+ *
+ * Filter application:
+ * - {@link lfilter} — causal FIR/IIR filter (direct-form II transposed)
+ * - {@link filtfilt} — zero-phase forward-backward filter
+ * - {@link sosfilt} — second-order-sections filter
+ *
+ * @example
+ * ```ts
+ * import { firwin, lfilter, butter, sosfilt } from "tsb";
+ *
+ * // Low-pass FIR with 29 taps, cutoff 0.25 (Nyquist = 0.5)
+ * const b = firwin(29, 0.25);
+ * const y = lfilter(b, [1], signal);
+ *
+ * // Butterworth low-pass, order 4, cutoff 0.2
+ * const { sos } = butter(4, 0.2, "lowpass");
+ * const filtered = sosfilt(sos, signal);
+ * ```
+ *
+ * @module
+ */
+
+import {
+ type Complex,
+ type WindowName,
+ blackmanWindow,
+ cAbs,
+ complex,
+ hammingWindow,
+ hannWindow,
+ kaiserWindow,
+} from "./signal.ts";
+
+// ─── internal helpers ─────────────────────────────────────────────────────────
+
+/** sinc(x) = sin(πx) / (πx), sinc(0) = 1 (normalised). */
+function sinc(x: number): number {
+ if (x === 0) {
+ return 1;
+ }
+ const px = Math.PI * x;
+ return Math.sin(px) / px;
+}
+
+/** Polynomial multiplication (convolution). */
+function polyMul(a: readonly number[], b: readonly number[]): number[] {
+ const out = new Array(a.length + b.length - 1).fill(0);
+ for (let i = 0; i < a.length; i++) {
+ for (let j = 0; j < b.length; j++) {
+ out[i + j] = (out[i + j] ?? 0) + (a[i] ?? 0) * (b[j] ?? 0);
+ }
+ }
+ return out;
+}
+
+// ─── FIR filter design ────────────────────────────────────────────────────────
+
+/** Options for {@link firwin}. */
+export interface FirwinOptions {
+ /**
+ * Window to apply after ideal filter: name string or pre-computed array.
+ * Default `"hamming"`.
+ */
+ window?: WindowName | readonly number[];
+ /** If `true`, design a high-pass filter (default `false` = low-pass). */
+ pass_zero?: boolean;
+ /** Sampling rate used to normalise `cutoff` (default `2` so `cutoff ∈ [0, 1]`). */
+ fs?: number;
+}
+
+/**
+ * Design a low- or high-pass FIR filter using the windowed-sinc method.
+ *
+ * Mirrors `scipy.signal.firwin`.
+ *
+ * @param numtaps - Number of filter coefficients (must be odd for pass_zero=false).
+ * @param cutoff - Cutoff frequency. With default `fs=2`, cutoff is normalised
+ * so `1.0` equals the Nyquist frequency.
+ * @param options - {@link FirwinOptions}.
+ * @returns - FIR filter coefficients `b` (length `numtaps`).
+ *
+ * @example
+ * ```ts
+ * import { firwin, lfilter } from "tsb";
+ * const b = firwin(51, 0.3); // 51-tap 150 Hz LPF (fs=1000)
+ * const y = lfilter(b, [1], signal);
+ * ```
+ */
+export function firwin(
+ numtaps: number,
+ cutoff: number | readonly [number, number],
+ options: FirwinOptions = {},
+): number[] {
+ const fs = options.fs ?? 2;
+ const passZero = options.pass_zero ?? true;
+ const nyq = fs / 2;
+
+ // Normalise cutoff(s) to [0..1] where 1 = Nyquist
+ const cuts = (typeof cutoff === "number" ? [cutoff] : cutoff).map((c) => c / nyq);
+
+ const M = numtaps - 1;
+
+ // Build window
+ let win: number[];
+ if (options.window !== undefined) {
+ win =
+ typeof options.window === "string"
+ ? buildFirWindow(options.window, numtaps)
+ : Array.from(options.window);
+ } else {
+ win = hammingWindow(numtaps);
+ }
+
+ // Ideal sinc coefficients
+ const h = new Array(numtaps).fill(0);
+
+ if (cuts.length === 1) {
+ const fc = cuts[0] ?? 0;
+ if (passZero) {
+ // Low-pass: h[n] = fc * sinc(fc * (n - M/2))
+ for (let n = 0; n < numtaps; n++) {
+ h[n] = fc * sinc(fc * (n - M / 2)) * (win[n] ?? 1);
+ }
+ } else {
+ // High-pass: h[n] = delta(n - M/2) - fc * sinc(fc * (n - M/2))
+ for (let n = 0; n < numtaps; n++) {
+ const delta = n === M / 2 ? 1 : 0;
+ h[n] = (delta - fc * sinc(fc * (n - M / 2))) * (win[n] ?? 1);
+ }
+ }
+ } else {
+ // Band-pass or band-stop
+ const f1 = cuts[0] ?? 0;
+ const f2 = cuts[1] ?? 0;
+ if (passZero) {
+ // Band-stop (notch): LP(f1) + HP(f2)
+ for (let n = 0; n < numtaps; n++) {
+ const mid = M / 2;
+ const delta = n === mid ? 1 : 0;
+ h[n] = (f1 * sinc(f1 * (n - mid)) + (delta - f2 * sinc(f2 * (n - mid)))) * (win[n] ?? 1);
+ }
+ } else {
+ // Band-pass: BP(f1, f2) = LP(f2) - LP(f1)
+ for (let n = 0; n < numtaps; n++) {
+ const mid = M / 2;
+ h[n] = (f2 * sinc(f2 * (n - mid)) - f1 * sinc(f1 * (n - mid))) * (win[n] ?? 1);
+ }
+ }
+ }
+
+ // Normalise DC gain
+ const dcGain = h.reduce((s, v) => s + v, 0);
+ if (Math.abs(dcGain) > 1e-12 && passZero && cuts.length === 1) {
+ // Low-pass: normalise DC to 1
+ const scale = 1 / dcGain;
+ return h.map((v) => v * scale);
+ }
+ return h;
+}
+
+/** Build a named window for FIR design. */
+function buildFirWindow(name: WindowName, n: number): number[] {
+ switch (name) {
+ case "hamming":
+ return hammingWindow(n);
+ case "hann":
+ return hannWindow(n);
+ case "blackman":
+ return blackmanWindow(n);
+ case "kaiser":
+ return kaiserWindow(n, 14);
+ default:
+ return hammingWindow(n);
+ }
+}
+
+// ─── frequency response ───────────────────────────────────────────────────────
+
+/** Result of {@link freqz} and {@link sosfreqz}. */
+export interface FreqzResult {
+ /** Angular frequencies in radians/sample (0 to π). */
+ w: number[];
+ /** Complex frequency response H(e^jω). */
+ H: Complex[];
+}
+
+/**
+ * Compute the frequency response H(e^jω) of a digital filter.
+ *
+ * Mirrors `scipy.signal.freqz`.
+ *
+ * @param b - Numerator polynomial coefficients.
+ * @param a - Denominator polynomial coefficients (default `[1]` = FIR).
+ * @param worN - Number of frequency points, or array of specific radian frequencies.
+ * @returns - `{ w, H }` where `w` is in radians/sample and `H` is complex.
+ *
+ * @example
+ * ```ts
+ * import { firwin, freqz } from "tsb";
+ * const b = firwin(31, 0.3);
+ * const { w, H } = freqz(b, [1], 512);
+ * const mag = H.map(h => cAbs(h));
+ * ```
+ */
+export function freqz(
+ b: readonly number[],
+ a: readonly number[] = [1],
+ worN: number | readonly number[] = 512,
+): FreqzResult {
+ const ws =
+ typeof worN === "number"
+ ? Array.from({ length: worN }, (_, i) => (Math.PI * i) / worN)
+ : Array.from(worN);
+
+ const H: Complex[] = ws.map((w) => {
+ // H(e^jw) = B(e^jw) / A(e^jw)
+ // Evaluate in z^-1 so numerator and denominator can have different lengths.
+ const z: Complex = { re: Math.cos(w), im: -Math.sin(w) };
+ const Bw = evalPolyZ(b, z);
+ const Aw = evalPolyZ(a, z);
+ return divComplex(Bw, Aw);
+ });
+
+ return { w: ws, H };
+}
+
+/** Evaluate p[0] + p[1]*z + ... using Horner's method. */
+function evalPolyZ(p: readonly number[], z: Complex): Complex {
+ let acc: Complex = complex(0, 0);
+ for (let i = p.length - 1; i >= 0; i--) {
+ // acc = acc * z + p[i]
+ acc = {
+ re: acc.re * z.re - acc.im * z.im + (p[i] ?? 0),
+ im: acc.re * z.im + acc.im * z.re,
+ };
+ }
+ return acc;
+}
+
+/** Divide two complex numbers (b / a), returns 0 when |a| < eps. */
+function divComplex(b: Complex, a: Complex): Complex {
+ const denom = a.re * a.re + a.im * a.im;
+ if (denom < 1e-300) {
+ return complex(0, 0);
+ }
+ return {
+ re: (b.re * a.re + b.im * a.im) / denom,
+ im: (b.im * a.re - b.re * a.im) / denom,
+ };
+}
+
+// ─── Butterworth IIR filter ───────────────────────────────────────────────────
+
+/** A second-order section: `[b0, b1, b2, 1, a1, a2]`. */
+export type SOSSection = [number, number, number, number, number, number];
+
+/** Result of {@link butter}. */
+export interface ButterResult {
+ /** Second-order sections (numerically preferred for high orders). */
+ sos: SOSSection[];
+ /** Numerator polynomial (may lose precision for high orders). */
+ b: number[];
+ /** Denominator polynomial (may lose precision for high orders). */
+ a: number[];
+}
+
+/** Butter filter type. */
+export type FilterType = "lowpass" | "highpass" | "bandpass" | "bandstop";
+
+/**
+ * Design an N-th order Butterworth digital filter (bilinear transform).
+ *
+ * Mirrors `scipy.signal.butter`.
+ *
+ * Returns both the SOS form (use {@link sosfilt} — numerically stable) and
+ * the b/a form (use {@link lfilter} — may have numerical issues for N > 4).
+ *
+ * @param N - Filter order (1–8 recommended; high orders lose precision in b/a form).
+ * @param Wn - Critical frequency. Normalised to `[0, 1]` where `1 = Nyquist`.
+ * Provide `[low, high]` for band-pass or band-stop.
+ * @param type - Filter type (default `"lowpass"`).
+ * @returns - `{ sos, b, a }`.
+ *
+ * @example
+ * ```ts
+ * import { butter, sosfilt } from "tsb";
+ * const { sos } = butter(4, 0.2);
+ * const y = sosfilt(sos, signal);
+ * ```
+ */
+export function butter(
+ N: number,
+ Wn: number | readonly [number, number],
+ type: FilterType = "lowpass",
+): ButterResult {
+ if (N < 1 || N > 20 || !Number.isInteger(N)) {
+ throw new RangeError("Order N must be an integer 1–20");
+ }
+
+ const isBand = type === "bandpass" || type === "bandstop";
+ if (isBand && typeof Wn === "number") {
+ throw new TypeError("Band filters require Wn = [low, high]");
+ }
+ if (!isBand && typeof Wn !== "number") {
+ throw new TypeError("Low/high-pass filters require scalar Wn");
+ }
+ const cuts = typeof Wn === "number" ? [Wn] : Wn;
+ if (cuts.some((cut) => !Number.isFinite(cut) || cut <= 0 || cut >= 1)) {
+ throw new RangeError("Critical frequencies must be finite and strictly between 0 and 1");
+ }
+ if (typeof Wn !== "number" && Wn[0] >= Wn[1]) {
+ throw new RangeError("Band critical frequencies must satisfy low < high");
+ }
+
+ // Unit-cutoff Butterworth poles in the left half-plane.
+ const poles = Array.from({ length: N }, (_, k) => {
+ const angle = (Math.PI * (2 * k + N + 1)) / (2 * N);
+ return complex(Math.cos(angle), Math.sin(angle));
+ });
+ if (typeof Wn !== "number") {
+ const warped: [number, number] = [
+ 2 * Math.tan((Math.PI * Wn[0]) / 2),
+ 2 * Math.tan((Math.PI * Wn[1]) / 2),
+ ];
+ return butterBand(warped, type === "bandstop" ? "bandstop" : "bandpass", poles);
+ }
+
+ const omega = 2 * Math.tan((Math.PI * Wn) / 2);
+ const digitalPoles = poles.map((pole) => {
+ const scaled =
+ type === "highpass"
+ ? divComplex(complex(omega, 0), pole)
+ : complex(omega * pole.re, omega * pole.im);
+ return bilinearPole(scaled);
+ });
+ const sections = pairPoles(digitalPoles).map(([p, q]): SOSSection => {
+ const zero = type === "highpass" ? 1 : -1;
+ return [
+ 1,
+ q === null ? -zero : -2 * zero,
+ q === null ? 0 : 1,
+ 1,
+ -(p.re + (q?.re ?? 0)),
+ q === null ? 0 : p.re * q.re - p.im * q.im,
+ ];
+ });
+ const sos = normaliseSOS(sections, type === "highpass" ? Math.PI : 0);
+ return { sos, ...sosToBA(sos) };
+}
+
+/** Bilinear transform: analog pole s → digital pole z = (2+s)/(2-s). */
+function bilinearPole(s: Complex): Complex {
+ return divComplex(complex(2 + s.re, s.im), complex(2 - s.re, -s.im));
+}
+
+/** Pair conjugate poles, and pair real poles with each other where possible. */
+function pairPoles(poles: readonly Complex[]): [Complex, Complex | null][] {
+ const remaining = [...poles].sort((a, b) => cAbs(a) - cAbs(b));
+ const pairs: [Complex, Complex | null][] = [];
+ while (remaining.length > 0) {
+ const pole = remaining.shift();
+ if (pole === undefined) break;
+ const real = Math.abs(pole.im) < 1e-10;
+ const partner = remaining.findIndex((candidate) =>
+ real
+ ? Math.abs(candidate.im) < 1e-10
+ : Math.abs(pole.re - candidate.re) < 1e-10 && Math.abs(pole.im + candidate.im) < 1e-10,
+ );
+ if (partner < 0 && !real) {
+ throw new Error("Could not pair Butterworth conjugate poles");
+ }
+ pairs.push([pole, partner < 0 ? null : (remaining.splice(partner, 1)[0] ?? null)]);
+ }
+ return pairs;
+}
+
+/** Normalize each section at a frequency in the unit-gain passband. */
+function normaliseSOS(sections: readonly SOSSection[], frequency: number): SOSSection[] {
+ const z = complex(Math.cos(frequency), -Math.sin(frequency));
+ return sections.map(([b0, b1, b2, a0, a1, a2]): SOSSection => {
+ const numerator = cAbs(evalPolyZ([b0, b1, b2], z));
+ const denominator = cAbs(evalPolyZ([a0, a1, a2], z));
+ const scale = denominator / numerator;
+ return [b0 * scale, b1 * scale, b2 * scale, a0, a1, a2];
+ });
+}
+
+/** Apply the low-pass prototype's band transformation before the bilinear map. */
+function butterBand(
+ warped: readonly [number, number],
+ type: "bandpass" | "bandstop",
+ protoPoles: readonly Complex[],
+): ButterResult {
+ const [w1, w2] = warped;
+ const bandwidth = w2 - w1;
+ const center = Math.sqrt(w1 * w2);
+ const digitalPoles: Complex[] = [];
+ for (const pole of protoPoles) {
+ // BP: s² - BW*p*s + center² = 0.
+ // BS: s² - (BW/p)*s + center² = 0.
+ const linear =
+ type === "bandpass"
+ ? complex(bandwidth * pole.re, bandwidth * pole.im)
+ : divComplex(complex(bandwidth, 0), pole);
+ const [rootRe, rootIm] = complexSqrt(
+ linear.re ** 2 - linear.im ** 2 - 4 * center ** 2,
+ 2 * linear.re * linear.im,
+ );
+ digitalPoles.push(
+ bilinearPole(complex((linear.re + rootRe) / 2, (linear.im + rootIm) / 2)),
+ bilinearPole(complex((linear.re - rootRe) / 2, (linear.im - rootIm) / 2)),
+ );
+ }
+
+ // BP zeros map to z=±1. BS zeros at ±j*center map to a unit-circle pair.
+ const stopZero = bilinearPole(complex(0, center));
+ const sections = pairPoles(digitalPoles).map(([p, q]): SOSSection => {
+ if (q === null) throw new Error("Band filters require paired poles");
+ return [
+ 1,
+ type === "bandpass" ? 0 : -2 * stopZero.re,
+ type === "bandpass" ? -1 : 1,
+ 1,
+ -(p.re + q.re),
+ p.re * q.re - p.im * q.im,
+ ];
+ });
+ const frequency = type === "bandpass" ? 2 * Math.atan(center / 2) : 0;
+ const sos = normaliseSOS(sections, frequency);
+ return { sos, ...sosToBA(sos) };
+}
+
+/** Real and imaginary parts of the principal square root, including negative real inputs. */
+function complexSqrt(re: number, im: number): [number, number] {
+ const magnitude = Math.hypot(re, im);
+ const rootRe = Math.sqrt(Math.max(0, (magnitude + re) / 2));
+ const rootIm = Math.sqrt(Math.max(0, (magnitude - re) / 2));
+ return [rootRe, im < 0 ? -rootIm : rootIm];
+}
+
+/** Convert SOS to b/a transfer function, omitting first-order padding. */
+function sosToBA(sections: readonly SOSSection[]): { b: number[]; a: number[] } {
+ let b: number[] = [1];
+ let a: number[] = [1];
+ for (const [b0, b1, b2, a0, a1, a2] of sections) {
+ const firstOrder = b2 === 0 && a2 === 0;
+ b = polyMul(b, firstOrder ? [b0, b1] : [b0, b1, b2]);
+ a = polyMul(a, firstOrder ? [a0, a1] : [a0, a1, a2]);
+ }
+ return { b, a };
+}
+
+/**
+ * Compute the frequency response of a SOS filter.
+ *
+ * @param sos - SOS sections as from {@link butter}.
+ * @param worN - Number of frequency points or explicit frequencies.
+ * @returns - `{ w, H }`.
+ */
+export function sosfreqz(
+ sos: readonly SOSSection[],
+ worN: number | readonly number[] = 512,
+): FreqzResult {
+ const ws =
+ typeof worN === "number"
+ ? Array.from({ length: worN }, (_, i) => (Math.PI * i) / worN)
+ : Array.from(worN);
+
+ const H: Complex[] = ws.map((w) => {
+ const z: Complex = { re: Math.cos(w), im: -Math.sin(w) };
+ let acc: Complex = complex(1, 0);
+ for (const [b0, b1, b2, a0, a1, a2] of sos) {
+ const num = evalPolyZ([b0, b1, b2], z);
+ const den = evalPolyZ([a0, a1, a2], z);
+ const secH = divComplex(num, den);
+ acc = { re: acc.re * secH.re - acc.im * secH.im, im: acc.re * secH.im + acc.im * secH.re };
+ }
+ return acc;
+ });
+
+ return { w: ws, H };
+}
+
+// ─── filter application ───────────────────────────────────────────────────────
+
+/**
+ * Apply an IIR or FIR filter using direct-form II transposed.
+ *
+ * Mirrors `scipy.signal.lfilter`. Computes `y[n] = b[0]*x[n] + b[1]*x[n-1] + ...
+ * - a[1]*y[n-1] - a[2]*y[n-2] - ...` (a[0] is assumed to be 1 or is normalised).
+ *
+ * @param b - Numerator coefficients (length M+1).
+ * @param a - Denominator coefficients (length N+1, a[0] normalised to 1).
+ * @param x - Input signal.
+ * @returns - Filtered signal (same length as `x`).
+ *
+ * @example
+ * ```ts
+ * import { firwin, lfilter } from "tsb";
+ * const b = firwin(21, 0.3);
+ * const y = lfilter(b, [1], x);
+ * ```
+ */
+export function lfilter(
+ b: readonly number[],
+ a: readonly number[],
+ x: readonly number[],
+): number[] {
+ const nb = b.length;
+ const na = a.length;
+ const n = x.length;
+
+ // Normalise a[0]
+ const a0 = a[0] ?? 1;
+ const bn = b.map((v) => v / a0);
+ const an = a.map((v) => v / a0);
+
+ const m = Math.max(nb, na);
+ const z = new Float64Array(m); // state buffer
+ const y = new Array(n);
+
+ for (let i = 0; i < n; i++) {
+ const xi = x[i] ?? 0;
+ const yi = (bn[0] ?? 0) * xi + (z[0] ?? 0);
+ y[i] = yi;
+ for (let j = 0; j < m - 1; j++) {
+ z[j] = (bn[j + 1] ?? 0) * xi - (an[j + 1] ?? 0) * yi + (z[j + 1] ?? 0);
+ }
+ z[m - 1] = (bn[m] ?? 0) * xi - (an[m] ?? 0) * yi;
+ }
+
+ return y;
+}
+
+/**
+ * Zero-phase forward-backward filter. Applies the filter twice — once forward
+ * and once backward — eliminating phase distortion.
+ *
+ * Mirrors `scipy.signal.filtfilt`.
+ *
+ * @param b - Numerator coefficients.
+ * @param a - Denominator coefficients.
+ * @param x - Input signal.
+ * @returns - Zero-phase filtered signal (same length as `x`).
+ */
+export function filtfilt(
+ b: readonly number[],
+ a: readonly number[],
+ x: readonly number[],
+): number[] {
+ if (x.length === 0) return [];
+ const edge = Math.min(3 * Math.max(a.length, b.length), x.length - 1);
+ const extended = oddExtension(x, edge);
+ const forward = filterSteadyState(b, a, extended);
+ const backward = filterSteadyState(b, a, forward.reverse());
+ return backward.reverse().slice(edge, edge + x.length);
+}
+
+/** Reflect endpoint deviations to avoid introducing artificial jumps. */
+function oddExtension(x: readonly number[], edge: number): number[] {
+ const first = x[0] ?? 0;
+ const last = x[x.length - 1] ?? 0;
+ const left = Array.from({ length: edge }, (_, i) => 2 * first - (x[edge - i] ?? first));
+ const right = Array.from({ length: edge }, (_, i) => 2 * last - (x[x.length - 2 - i] ?? last));
+ return [...left, ...x, ...right];
+}
+
+/** Direct-form II filtering initialized to the first sample's steady state. */
+function filterSteadyState(
+ b: readonly number[],
+ a: readonly number[],
+ x: readonly number[],
+): number[] {
+ const a0 = a[0] ?? 1;
+ const bn = b.map((v) => v / a0);
+ const an = a.map((v) => v / a0);
+ const order = Math.max(a.length, b.length) - 1;
+ const state = new Array(order + 1).fill(0);
+ const first = x[0] ?? 0;
+ const denominator = an.reduce((sum, v) => sum + v, 0);
+ const steady = denominator === 0 ? 0 : (first * bn.reduce((sum, v) => sum + v, 0)) / denominator;
+ for (let i = order - 1; i >= 0; i--) {
+ state[i] = (state[i + 1] ?? 0) + (bn[i + 1] ?? 0) * first - (an[i + 1] ?? 0) * steady;
+ }
+ return x.map((value) => {
+ const out = (bn[0] ?? 0) * value + (state[0] ?? 0);
+ for (let i = 0; i < order; i++) {
+ state[i] = (state[i + 1] ?? 0) + (bn[i + 1] ?? 0) * value - (an[i + 1] ?? 0) * out;
+ }
+ return out;
+ });
+}
+
+/**
+ * Apply a second-order-sections filter.
+ *
+ * Numerically more stable than {@link lfilter} for high-order IIR filters.
+ * Mirrors `scipy.signal.sosfilt`.
+ *
+ * @param sos - SOS sections from {@link butter}.
+ * @param x - Input signal.
+ * @returns - Filtered signal (same length as `x`).
+ */
+export function sosfilt(sos: readonly SOSSection[], x: readonly number[]): number[] {
+ let signal = Array.from(x);
+ for (const [b0, b1, b2, a0, a1, a2] of sos) {
+ signal = lfilter([b0, b1, b2], [a0, a1, a2], signal);
+ }
+ return signal;
+}
+
+/**
+ * Zero-phase SOS filter with one padded forward cascade and one backward cascade.
+ *
+ * @param sos - SOS sections from {@link butter}.
+ * @param x - Input signal.
+ * @returns - Zero-phase filtered signal.
+ */
+export function sosfiltfilt(sos: readonly SOSSection[], x: readonly number[]): number[] {
+ if (x.length === 0) return [];
+ const numeratorPadding = sos.filter((section) => section[2] === 0).length;
+ const denominatorPadding = sos.filter((section) => section[5] === 0).length;
+ const taps = 2 * sos.length + 1 - Math.min(numeratorPadding, denominatorPadding);
+ const edge = Math.min(3 * taps, x.length - 1);
+ let signal = oddExtension(x, edge);
+ for (let pass = 0; pass < 2; pass++) {
+ for (const [b0, b1, b2, a0, a1, a2] of sos) {
+ signal = filterSteadyState([b0, b1, b2], [a0, a1, a2], signal);
+ }
+ signal.reverse();
+ }
+ return signal.slice(edge, edge + x.length);
+}
+
+// Re-export cAbs for convenience
+export { cAbs };
diff --git a/src/stats/hmm.ts b/src/stats/hmm.ts
new file mode 100644
index 000000000..931ae6fb4
--- /dev/null
+++ b/src/stats/hmm.ts
@@ -0,0 +1,776 @@
+/**
+ * hmm — Hidden Markov Model (HMM) with discrete and Gaussian emissions.
+ *
+ * Implements:
+ * - **Forward-Backward** algorithm (log-space for numerical stability)
+ * - **Baum-Welch** EM parameter estimation
+ * - **Viterbi** algorithm for most-likely state sequence decoding
+ * - **GaussianHMM**: continuous observations with Gaussian emission distributions
+ * - **MultinomialHMM**: discrete observations with categorical emission distributions
+ *
+ * Mirrors `hmmlearn.hmm.GaussianHMM` and `MultinomialHMM` APIs.
+ *
+ * @example
+ * ```ts
+ * import { GaussianHMM } from "tsb";
+ *
+ * const model = new GaussianHMM({ nComponents: 2, nIter: 100 });
+ * const obs = [0.1, 0.2, 0.15, 2.1, 2.3, 2.0, 0.05, 0.1, 2.5, 2.2];
+ * model.fit(obs);
+ * const states = model.predict(obs);
+ * const logProb = model.score(obs);
+ * ```
+ *
+ * @module
+ */
+
+// ─── Constants ─────────────────────────────────────────────────────────────────
+
+const LOG_ZERO = Number.NEGATIVE_INFINITY;
+
+// ─── Utility helpers ──────────────────────────────────────────────────────────
+
+/** log(exp(a) + exp(b)) with numerical stability. */
+function logSumExp(a: number, b: number): number {
+ if (a === LOG_ZERO) {
+ return b;
+ }
+ if (b === LOG_ZERO) {
+ return a;
+ }
+ const m = Math.max(a, b);
+ return m + Math.log(Math.exp(a - m) + Math.exp(b - m));
+}
+
+/** Stable log-sum-exp over an array. */
+function logSumExpArr(arr: readonly number[]): number {
+ let result = LOG_ZERO;
+ for (const v of arr) {
+ result = logSumExp(result, v);
+ }
+ return result;
+}
+
+/** Safe log (returns LOG_ZERO for non-positive). */
+function safeLog(x: number): number {
+ return x > 0 ? Math.log(x) : LOG_ZERO;
+}
+
+/** Normalise an array in-place so it sums to 1. Returns sum before normalisation. */
+function normalise(arr: number[]): number {
+ const s = arr.reduce((a, b) => a + b, 0);
+ if (s > 0) {
+ for (let i = 0; i < arr.length; i++) {
+ arr[i] = (arr[i] ?? 0) / s;
+ }
+ }
+ return s;
+}
+
+// ─── Public types ─────────────────────────────────────────────────────────────
+
+/** Parameters for a Gaussian HMM. */
+export interface GaussianHMMParams {
+ /** Number of hidden states. */
+ nComponents: number;
+ /** Maximum EM iterations. Default 100. */
+ nIter?: number;
+ /** Convergence tolerance on log-likelihood. Default 1e-4. */
+ tol?: number;
+ /** Random seed for initialisation (unused in deterministic init). */
+ randomState?: number;
+}
+
+/** Fitted Gaussian HMM result. */
+export interface GaussianHMMFit {
+ /** Initial state probabilities (shape: nComponents). */
+ startProb: number[];
+ /** Transition matrix (shape: nComponents × nComponents). Row i → state i, column j → state j. */
+ transmat: number[][];
+ /** Emission means (shape: nComponents). */
+ means: number[];
+ /** Emission variances (shape: nComponents). */
+ covars: number[];
+ /** Log-likelihood of training data at convergence. */
+ logProb: number;
+ /** Number of EM iterations completed. */
+ nIterDone: number;
+}
+
+/** Parameters for a Multinomial HMM. */
+export interface MultinomialHMMParams {
+ /** Number of hidden states. */
+ nComponents: number;
+ /** Number of distinct observation symbols. */
+ nFeatures: number;
+ /** Maximum EM iterations. Default 100. */
+ nIter?: number;
+ /** Convergence tolerance on log-likelihood. Default 1e-4. */
+ tol?: number;
+}
+
+/** Fitted Multinomial HMM result. */
+export interface MultinomialHMMFit {
+ /** Initial state probabilities (shape: nComponents). */
+ startProb: number[];
+ /** Transition matrix (shape: nComponents × nComponents). */
+ transmat: number[][];
+ /** Emission probability matrix (shape: nComponents × nFeatures). */
+ emissionProb: number[][];
+ /** Log-likelihood of training data at convergence. */
+ logProb: number;
+ /** Number of EM iterations completed. */
+ nIterDone: number;
+}
+
+// ─── Forward-Backward (log-space) ─────────────────────────────────────────────
+
+/**
+ * Compute log-forward variables.
+ * @param logStartProb - log initial probabilities (nStates)
+ * @param logTransmat - log transition matrix (nStates × nStates)
+ * @param logEmit - log emission probabilities (T × nStates)
+ * @returns logAlpha (T × nStates)
+ */
+function logForward(
+ logStartProb: readonly number[],
+ logTransmat: readonly (readonly number[])[],
+ logEmit: readonly (readonly number[])[],
+): number[][] {
+ const T = logEmit.length;
+ const K = logStartProb.length;
+ const logAlpha: number[][] = Array.from({ length: T }, () => new Array(K).fill(LOG_ZERO));
+
+ for (let j = 0; j < K; j++) {
+ logAlpha[0]![j] = (logStartProb[j] ?? LOG_ZERO) + (logEmit[0]?.[j] ?? LOG_ZERO);
+ }
+ for (let t = 1; t < T; t++) {
+ for (let j = 0; j < K; j++) {
+ let s = LOG_ZERO;
+ for (let i = 0; i < K; i++) {
+ s = logSumExp(s, (logAlpha[t - 1]?.[i] ?? LOG_ZERO) + (logTransmat[i]?.[j] ?? LOG_ZERO));
+ }
+ logAlpha[t]![j] = s + (logEmit[t]?.[j] ?? LOG_ZERO);
+ }
+ }
+ return logAlpha;
+}
+
+/**
+ * Compute log-backward variables.
+ * @param logTransmat - log transition matrix (nStates × nStates)
+ * @param logEmit - log emission probabilities (T × nStates)
+ * @returns logBeta (T × nStates)
+ */
+function logBackward(
+ logTransmat: readonly (readonly number[])[],
+ logEmit: readonly (readonly number[])[],
+): number[][] {
+ const T = logEmit.length;
+ const K = logTransmat.length;
+ const logBeta: number[][] = Array.from({ length: T }, () => new Array(K).fill(0));
+
+ for (let t = T - 2; t >= 0; t--) {
+ for (let i = 0; i < K; i++) {
+ let s = LOG_ZERO;
+ for (let j = 0; j < K; j++) {
+ s = logSumExp(
+ s,
+ (logTransmat[i]?.[j] ?? LOG_ZERO) +
+ (logEmit[t + 1]?.[j] ?? LOG_ZERO) +
+ (logBeta[t + 1]?.[j] ?? 0),
+ );
+ }
+ logBeta[t]![i] = s;
+ }
+ }
+ return logBeta;
+}
+
+// ─── Viterbi ──────────────────────────────────────────────────────────────────
+
+/**
+ * Viterbi algorithm to find the most-likely state sequence.
+ * @returns decoded state sequence (length T)
+ */
+function viterbi(
+ logStartProb: readonly number[],
+ logTransmat: readonly (readonly number[])[],
+ logEmit: readonly (readonly number[])[],
+): number[] {
+ const T = logEmit.length;
+ const K = logStartProb.length;
+ const delta: number[][] = Array.from({ length: T }, () => new Array(K).fill(LOG_ZERO));
+ const psi: number[][] = Array.from({ length: T }, () => new Array(K).fill(0));
+
+ for (let j = 0; j < K; j++) {
+ delta[0]![j] = (logStartProb[j] ?? LOG_ZERO) + (logEmit[0]?.[j] ?? LOG_ZERO);
+ }
+ for (let t = 1; t < T; t++) {
+ for (let j = 0; j < K; j++) {
+ let best = LOG_ZERO;
+ let bestI = 0;
+ for (let i = 0; i < K; i++) {
+ const v = (delta[t - 1]?.[i] ?? LOG_ZERO) + (logTransmat[i]?.[j] ?? LOG_ZERO);
+ if (v > best) {
+ best = v;
+ bestI = i;
+ }
+ }
+ delta[t]![j] = best + (logEmit[t]?.[j] ?? LOG_ZERO);
+ psi[t]![j] = bestI;
+ }
+ }
+
+ // Backtrack
+ const path = new Array(T).fill(0);
+ let s = 0;
+ let best = LOG_ZERO;
+ for (let j = 0; j < K; j++) {
+ const v = delta[T - 1]?.[j] ?? LOG_ZERO;
+ if (v > best) {
+ best = v;
+ s = j;
+ }
+ }
+ path[T - 1] = s;
+ for (let t = T - 2; t >= 0; t--) {
+ path[t] = psi[t + 1]?.[path[t + 1] ?? 0] ?? 0;
+ }
+ return path;
+}
+
+// ─── GaussianHMM ──────────────────────────────────────────────────────────────
+
+/** Log probability of x under N(mu, sigma^2). */
+function gaussianLogProb(x: number, mu: number, sigma2: number): number {
+ if (sigma2 <= 0) {
+ return LOG_ZERO;
+ }
+ return -0.5 * (Math.log(2 * Math.PI * sigma2) + ((x - mu) * (x - mu)) / sigma2);
+}
+
+/**
+ * Hidden Markov Model with univariate Gaussian emission distributions.
+ *
+ * Uses Baum-Welch EM for parameter estimation, Viterbi for decoding.
+ */
+export class GaussianHMM {
+ private readonly K: number;
+ private readonly nIter: number;
+ private readonly tol: number;
+
+ // Parameters (set after fit)
+ private _startProb: number[] = [];
+ private _transmat: number[][] = [];
+ private _means: number[] = [];
+ private _covars: number[] = [];
+ private _fitted = false;
+
+ constructor(params: GaussianHMMParams) {
+ this.K = params.nComponents;
+ this.nIter = params.nIter ?? 100;
+ this.tol = params.tol ?? 1e-4;
+ }
+
+ /** Fit the HMM to a sequence of observations using Baum-Welch EM. */
+ fit(obs: readonly number[]): GaussianHMMFit {
+ const T = obs.length;
+ const K = this.K;
+ if (T < 2) {
+ throw new Error("Need at least 2 observations");
+ }
+
+ // ── Initialise ───────────────────────────────────────────────────────────
+ // k-means-style: split sorted observations into K equal buckets
+ const sorted = [...obs].sort((a, b) => a - b);
+ const means = new Array(K).fill(0);
+ const covars = new Array(K).fill(1);
+ for (let k = 0; k < K; k++) {
+ const lo = Math.floor((k * T) / K);
+ const hi = Math.floor(((k + 1) * T) / K);
+ const slice = sorted.slice(lo, hi);
+ const mu = slice.reduce((a, b) => a + b, 0) / slice.length;
+ const v = slice.reduce((a, b) => a + (b - mu) ** 2, 0) / Math.max(slice.length - 1, 1);
+ means[k] = mu;
+ covars[k] = Math.max(v, 1e-6);
+ }
+
+ // Uniform start and transition
+ const startProb = new Array(K).fill(1 / K);
+ const transmat: number[][] = Array.from({ length: K }, () => new Array(K).fill(1 / K));
+
+ let prevLogProb = Number.NEGATIVE_INFINITY;
+ let nIterDone = 0;
+
+ for (let iter = 0; iter < this.nIter; iter++) {
+ // ── E-step ─────────────────────────────────────────────────────────────
+ const logEmit: number[][] = Array.from({ length: T }, (_, t) =>
+ Array.from({ length: K }, (__, k) =>
+ gaussianLogProb(obs[t] ?? 0, means[k] ?? 0, covars[k] ?? 1),
+ ),
+ );
+
+ const logStartProb = startProb.map(safeLog);
+ const logTransmat = transmat.map((row) => row.map(safeLog));
+
+ const logAlpha = logForward(logStartProb, logTransmat, logEmit);
+ const logBeta = logBackward(logTransmat, logEmit);
+
+ // Log-likelihood
+ const logProb = logSumExpArr(logAlpha[T - 1] ?? []);
+
+ if (Math.abs(logProb - prevLogProb) < this.tol) {
+ nIterDone = iter + 1;
+ break;
+ }
+ prevLogProb = logProb;
+ nIterDone = iter + 1;
+
+ // γ_t(k) = P(z_t = k | obs, θ) in log space
+ const logGamma: number[][] = Array.from({ length: T }, (_, t) => {
+ const row = Array.from(
+ { length: K },
+ (__, k) => (logAlpha[t]?.[k] ?? LOG_ZERO) + (logBeta[t]?.[k] ?? 0),
+ );
+ const z = logSumExpArr(row);
+ return row.map((v) => v - z);
+ });
+
+ // ξ_t(i,j) = P(z_t=i, z_{t+1}=j | obs, θ) — sum over t
+ const logXiSum: number[][] = Array.from({ length: K }, () =>
+ new Array(K).fill(LOG_ZERO),
+ );
+ for (let t = 0; t < T - 1; t++) {
+ for (let i = 0; i < K; i++) {
+ for (let j = 0; j < K; j++) {
+ const v =
+ (logAlpha[t]?.[i] ?? LOG_ZERO) +
+ (logTransmat[i]?.[j] ?? LOG_ZERO) +
+ (logEmit[t + 1]?.[j] ?? LOG_ZERO) +
+ (logBeta[t + 1]?.[j] ?? 0) -
+ logProb;
+ logXiSum[i]![j] = logSumExp(logXiSum[i]?.[j] ?? LOG_ZERO, v);
+ }
+ }
+ }
+
+ // ── M-step ─────────────────────────────────────────────────────────────
+ // Update startProb
+ for (let k = 0; k < K; k++) {
+ startProb[k] = Math.exp(logGamma[0]?.[k] ?? LOG_ZERO);
+ }
+ normalise(startProb);
+
+ // Update transmat
+ for (let i = 0; i < K; i++) {
+ for (let j = 0; j < K; j++) {
+ transmat[i]![j] = Math.exp(logXiSum[i]?.[j] ?? LOG_ZERO);
+ }
+ normalise(transmat[i]!);
+ }
+
+ // Update means
+ for (let k = 0; k < K; k++) {
+ let num = 0;
+ let den = 0;
+ for (let t = 0; t < T; t++) {
+ const g = Math.exp(logGamma[t]?.[k] ?? LOG_ZERO);
+ num += g * (obs[t] ?? 0);
+ den += g;
+ }
+ means[k] = den > 0 ? num / den : (means[k] ?? 0);
+ }
+
+ // Update covars
+ for (let k = 0; k < K; k++) {
+ let num = 0;
+ let den = 0;
+ for (let t = 0; t < T; t++) {
+ const g = Math.exp(logGamma[t]?.[k] ?? LOG_ZERO);
+ const diff = (obs[t] ?? 0) - (means[k] ?? 0);
+ num += g * diff * diff;
+ den += g;
+ }
+ covars[k] = Math.max(den > 0 ? num / den : (covars[k] ?? 1), 1e-6);
+ }
+ }
+
+ this._startProb = [...startProb];
+ this._transmat = transmat.map((row) => [...row]);
+ this._means = [...means];
+ this._covars = [...covars];
+ this._fitted = true;
+
+ return {
+ startProb: [...startProb],
+ transmat: transmat.map((row) => [...row]),
+ means: [...means],
+ covars: [...covars],
+ logProb: prevLogProb,
+ nIterDone,
+ };
+ }
+
+ /** Decode the most-likely state sequence using the Viterbi algorithm. */
+ predict(obs: readonly number[]): number[] {
+ this._checkFitted();
+ const K = this.K;
+ const logEmit = obs.map((x) =>
+ Array.from({ length: K }, (_, k) =>
+ gaussianLogProb(x, this._means[k] ?? 0, this._covars[k] ?? 1),
+ ),
+ );
+ return viterbi(
+ this._startProb.map(safeLog),
+ this._transmat.map((row) => row.map(safeLog)),
+ logEmit,
+ );
+ }
+
+ /** Compute the log-probability of the observation sequence. */
+ score(obs: readonly number[]): number {
+ this._checkFitted();
+ const K = this.K;
+ const logEmit = obs.map((x) =>
+ Array.from({ length: K }, (_, k) =>
+ gaussianLogProb(x, this._means[k] ?? 0, this._covars[k] ?? 1),
+ ),
+ );
+ const logAlpha = logForward(
+ this._startProb.map(safeLog),
+ this._transmat.map((row) => row.map(safeLog)),
+ logEmit,
+ );
+ return logSumExpArr(logAlpha.at(-1) ?? []);
+ }
+
+ /** Compute posterior state probabilities (T × nComponents). */
+ predictProba(obs: readonly number[]): number[][] {
+ this._checkFitted();
+ const K = this.K;
+ const T = obs.length;
+ const logEmit = obs.map((x) =>
+ Array.from({ length: K }, (_, k) =>
+ gaussianLogProb(x, this._means[k] ?? 0, this._covars[k] ?? 1),
+ ),
+ );
+ const logStartProb = this._startProb.map(safeLog);
+ const logTransmat = this._transmat.map((row) => row.map(safeLog));
+ const logAlpha = logForward(logStartProb, logTransmat, logEmit);
+ const logBeta = logBackward(logTransmat, logEmit);
+ return Array.from({ length: T }, (_, t) => {
+ const row = Array.from(
+ { length: K },
+ (__, k) => (logAlpha[t]?.[k] ?? LOG_ZERO) + (logBeta[t]?.[k] ?? 0),
+ );
+ const z = logSumExpArr(row);
+ return row.map((v) => Math.exp(v - z));
+ });
+ }
+
+ /** Sample a sequence of states and observations. */
+ sample(length: number): { states: number[]; obs: number[] } {
+ this._checkFitted();
+ const K = this.K;
+ const states: number[] = [];
+ const obs: number[] = [];
+
+ // Sample initial state
+ let cumProb = 0;
+ const r0 = Math.random();
+ let state = K - 1;
+ for (let k = 0; k < K; k++) {
+ cumProb += this._startProb[k] ?? 0;
+ if (r0 < cumProb) {
+ state = k;
+ break;
+ }
+ }
+
+ for (let t = 0; t < length; t++) {
+ states.push(state);
+ const mu = this._means[state] ?? 0;
+ const sigma = Math.sqrt(this._covars[state] ?? 1);
+ // Box-Muller for Gaussian sample
+ const u1 = Math.random();
+ const u2 = Math.random();
+ const z = Math.sqrt(-2 * Math.log(Math.max(u1, 1e-15))) * Math.cos(2 * Math.PI * u2);
+ obs.push(mu + sigma * z);
+
+ // Transition
+ const row = this._transmat[state]!;
+ let cum = 0;
+ const rv = Math.random();
+ let nextState = K - 1;
+ for (let k = 0; k < K; k++) {
+ cum += row[k] ?? 0;
+ if (rv < cum) {
+ nextState = k;
+ break;
+ }
+ }
+ state = nextState;
+ }
+ return { states, obs };
+ }
+
+ get startProb(): number[] {
+ this._checkFitted();
+ return [...this._startProb];
+ }
+ get transmat(): number[][] {
+ this._checkFitted();
+ return this._transmat.map((r) => [...r]);
+ }
+ get means(): number[] {
+ this._checkFitted();
+ return [...this._means];
+ }
+ get covars(): number[] {
+ this._checkFitted();
+ return [...this._covars];
+ }
+
+ private _checkFitted(): void {
+ if (!this._fitted) {
+ throw new Error("GaussianHMM is not fitted yet");
+ }
+ }
+}
+
+// ─── MultinomialHMM ───────────────────────────────────────────────────────────
+
+/**
+ * Hidden Markov Model with discrete (multinomial/categorical) emission distributions.
+ *
+ * Observations are non-negative integer symbol indices in [0, nFeatures).
+ */
+export class MultinomialHMM {
+ private readonly K: number;
+ private readonly nFeatures: number;
+ private readonly nIter: number;
+ private readonly tol: number;
+
+ private _startProb: number[] = [];
+ private _transmat: number[][] = [];
+ private _emissionProb: number[][] = [];
+ private _fitted = false;
+
+ constructor(params: MultinomialHMMParams) {
+ this.K = params.nComponents;
+ this.nFeatures = params.nFeatures;
+ this.nIter = params.nIter ?? 100;
+ this.tol = params.tol ?? 1e-4;
+ }
+
+ /** Fit the HMM to a sequence of integer observations using Baum-Welch EM. */
+ fit(obs: readonly number[]): MultinomialHMMFit {
+ const T = obs.length;
+ const K = this.K;
+ const V = this.nFeatures;
+ if (T < 2) {
+ throw new Error("Need at least 2 observations");
+ }
+
+ // Uniform initialisation with small random perturbation
+ const startProb = new Array(K).fill(1 / K);
+ const transmat: number[][] = Array.from({ length: K }, () =>
+ Array.from({ length: K }, () => 1 / K + (Math.random() * 0.1 - 0.05) / K),
+ );
+ const emissionProb: number[][] = Array.from({ length: K }, () =>
+ Array.from({ length: V }, () => 1 / V + (Math.random() * 0.1 - 0.05) / V),
+ );
+ // Normalise
+ for (let k = 0; k < K; k++) {
+ normalise(transmat[k]!);
+ normalise(emissionProb[k]!);
+ }
+
+ let prevLogProb = Number.NEGATIVE_INFINITY;
+ let nIterDone = 0;
+
+ for (let iter = 0; iter < this.nIter; iter++) {
+ // Log-emission: logEmit[t][k] = log P(obs[t] | z_t = k)
+ const logEmit: number[][] = Array.from({ length: T }, (_, t) =>
+ Array.from({ length: K }, (__, k) => safeLog(emissionProb[k]?.[obs[t] ?? 0] ?? 0)),
+ );
+
+ const logStartProb = startProb.map(safeLog);
+ const logTransmat = transmat.map((row) => row.map(safeLog));
+
+ const logAlpha = logForward(logStartProb, logTransmat, logEmit);
+ const logBeta = logBackward(logTransmat, logEmit);
+
+ const logProb = logSumExpArr(logAlpha[T - 1] ?? []);
+ if (Math.abs(logProb - prevLogProb) < this.tol) {
+ nIterDone = iter + 1;
+ break;
+ }
+ prevLogProb = logProb;
+ nIterDone = iter + 1;
+
+ // γ_t(k)
+ const logGamma: number[][] = Array.from({ length: T }, (_, t) => {
+ const row = Array.from(
+ { length: K },
+ (__, k) => (logAlpha[t]?.[k] ?? LOG_ZERO) + (logBeta[t]?.[k] ?? 0),
+ );
+ const z = logSumExpArr(row);
+ return row.map((v) => v - z);
+ });
+
+ // ξ sum
+ const logXiSum: number[][] = Array.from({ length: K }, () =>
+ new Array(K).fill(LOG_ZERO),
+ );
+ for (let t = 0; t < T - 1; t++) {
+ for (let i = 0; i < K; i++) {
+ for (let j = 0; j < K; j++) {
+ const v =
+ (logAlpha[t]?.[i] ?? LOG_ZERO) +
+ (logTransmat[i]?.[j] ?? LOG_ZERO) +
+ (logEmit[t + 1]?.[j] ?? LOG_ZERO) +
+ (logBeta[t + 1]?.[j] ?? 0) -
+ logProb;
+ logXiSum[i]![j] = logSumExp(logXiSum[i]?.[j] ?? LOG_ZERO, v);
+ }
+ }
+ }
+
+ // M-step: startProb
+ for (let k = 0; k < K; k++) {
+ startProb[k] = Math.exp(logGamma[0]?.[k] ?? LOG_ZERO);
+ }
+ normalise(startProb);
+
+ // transmat
+ for (let i = 0; i < K; i++) {
+ for (let j = 0; j < K; j++) {
+ transmat[i]![j] = Math.exp(logXiSum[i]?.[j] ?? LOG_ZERO);
+ }
+ normalise(transmat[i]!);
+ }
+
+ // emissionProb
+ for (let k = 0; k < K; k++) {
+ for (let v = 0; v < V; v++) {
+ emissionProb[k]![v] = 0;
+ }
+ for (let t = 0; t < T; t++) {
+ const sym = obs[t] ?? 0;
+ emissionProb[k]![sym] =
+ (emissionProb[k]?.[sym] ?? 0) + Math.exp(logGamma[t]?.[k] ?? LOG_ZERO);
+ }
+ normalise(emissionProb[k]!);
+ }
+ }
+
+ this._startProb = [...startProb];
+ this._transmat = transmat.map((row) => [...row]);
+ this._emissionProb = emissionProb.map((row) => [...row]);
+ this._fitted = true;
+
+ return {
+ startProb: [...startProb],
+ transmat: transmat.map((row) => [...row]),
+ emissionProb: emissionProb.map((row) => [...row]),
+ logProb: prevLogProb,
+ nIterDone,
+ };
+ }
+
+ /** Decode the most-likely state sequence. */
+ predict(obs: readonly number[]): number[] {
+ this._checkFitted();
+ const K = this.K;
+ const logEmit = obs.map((sym) =>
+ Array.from({ length: K }, (_, k) => safeLog(this._emissionProb[k]?.[sym] ?? 0)),
+ );
+ return viterbi(
+ this._startProb.map(safeLog),
+ this._transmat.map((row) => row.map(safeLog)),
+ logEmit,
+ );
+ }
+
+ /** Compute log-probability of the observation sequence. */
+ score(obs: readonly number[]): number {
+ this._checkFitted();
+ const K = this.K;
+ const logEmit = obs.map((sym) =>
+ Array.from({ length: K }, (_, k) => safeLog(this._emissionProb[k]?.[sym] ?? 0)),
+ );
+ const logAlpha = logForward(
+ this._startProb.map(safeLog),
+ this._transmat.map((row) => row.map(safeLog)),
+ logEmit,
+ );
+ return logSumExpArr(logAlpha.at(-1) ?? []);
+ }
+
+ get startProb(): number[] {
+ this._checkFitted();
+ return [...this._startProb];
+ }
+ get transmat(): number[][] {
+ this._checkFitted();
+ return this._transmat.map((r) => [...r]);
+ }
+ get emissionProb(): number[][] {
+ this._checkFitted();
+ return this._emissionProb.map((r) => [...r]);
+ }
+
+ private _checkFitted(): void {
+ if (!this._fitted) {
+ throw new Error("MultinomialHMM is not fitted yet");
+ }
+ }
+}
+
+// ─── Standalone functions ──────────────────────────────────────────────────────
+
+/**
+ * Convenience function: fit a GaussianHMM and return the fitted model.
+ *
+ * @example
+ * ```ts
+ * const model = fitGaussianHMM([0.1, 0.2, 2.1, 2.3, 0.05, 2.5], 2);
+ * ```
+ */
+export function fitGaussianHMM(
+ obs: readonly number[],
+ nComponents: number,
+ nIter = 100,
+): GaussianHMM {
+ const model = new GaussianHMM({ nComponents, nIter });
+ model.fit(obs);
+ return model;
+}
+
+/**
+ * Convenience function: Viterbi decoding with explicit parameters (no fitting).
+ *
+ * @param startProb - Initial state probabilities (length K).
+ * @param transmat - Transition matrix (K × K).
+ * @param emissionProb - Emission probabilities (K × V).
+ * @param obs - Integer observation sequence.
+ * @returns Most-likely state sequence.
+ */
+export function hmmViterbi(
+ startProb: readonly number[],
+ transmat: readonly (readonly number[])[],
+ emissionProb: readonly (readonly number[])[],
+ obs: readonly number[],
+): number[] {
+ const K = startProb.length;
+ const logEmit = obs.map((sym) =>
+ Array.from({ length: K }, (_, k) => safeLog(emissionProb[k]?.[sym] ?? 0)),
+ );
+ return viterbi(
+ startProb.map(safeLog),
+ transmat.map((row) => [...row].map(safeLog)),
+ logEmit,
+ );
+}
diff --git a/src/stats/index.ts b/src/stats/index.ts
index 78767ee14..dbfb531d8 100644
--- a/src/stats/index.ts
+++ b/src/stats/index.ts
@@ -575,3 +575,136 @@ export {
tsallisEntropy,
} from "./information.ts";
export type { PMF, NMIMethod } from "./information.ts";
+export {
+ complex,
+ cAbs,
+ cArg,
+ fft,
+ ifft,
+ rfft,
+ irfft,
+ fftFreq,
+ rfftFreq,
+ fftshift,
+ ifftshift,
+ rectangularWindow,
+ bartlettWindow,
+ hannWindow,
+ hammingWindow,
+ blackmanWindow,
+ blackmanHarrisWindow,
+ flatTopWindow,
+ kaiserWindow,
+ getWindow,
+ stft,
+ istft,
+ welch,
+ periodogram,
+} from "./signal.ts";
+export type {
+ Complex,
+ WindowName,
+ STFTOptions,
+ STFTResult,
+ ISTFTOptions,
+ WelchOptions,
+ PSDResult,
+ PeriodogramOptions,
+} from "./signal.ts";
+export {
+ firwin,
+ freqz,
+ sosfreqz,
+ lfilter,
+ filtfilt,
+ sosfilt,
+ sosfiltfilt,
+ butter,
+} from "./filters.ts";
+export type {
+ FirwinOptions,
+ FreqzResult,
+ SOSSection,
+ ButterResult,
+ FilterType,
+} from "./filters.ts";
+export {
+ autocorr,
+ acf,
+ pacf,
+ ccf,
+ durbinWatson,
+ ljungBox,
+ boxPierce,
+} from "./acf_pacf.ts";
+export type {
+ ACFResult,
+ PACFResult,
+ PortmanteauResult,
+ ACFOptions,
+ PACFOptions,
+ CCFOptions,
+ PortmanteauOptions,
+} from "./acf_pacf.ts";
+export {
+ ARIMAModel,
+ fitArima,
+} from "./arima.ts";
+export type {
+ ARIMAOptions,
+ ARIMAFitResult,
+ ARIMAForecastResult,
+} from "./arima.ts";
+export {
+ KalmanFilter,
+ StateSpaceModel,
+ kalmanFilter1D,
+ kalmanSmooth1D,
+ extractScalarMeans,
+ extractScalarVariances,
+ filteredPredictionInterval,
+} from "./kalman.ts";
+export type {
+ KalmanFilterOptions,
+ LocalLevelOptions,
+ LocalLinearTrendOptions,
+ KalmanFilterResult,
+ KalmanSmootherResult,
+} from "./kalman.ts";
+export {
+ SimpleExpSmoothing,
+ Holt,
+ ExponentialSmoothing,
+ simpleExpSmoothing,
+ holt,
+ fitEts,
+} from "./ets.ts";
+export type {
+ ETSTrend,
+ ETSSeasonal,
+ ETSInit,
+ SESOptions,
+ SESFitResult,
+ HoltOptions,
+ HoltFitResult,
+ ExponentialSmoothingOptions,
+ ExponentialSmoothingFitResult,
+ ETSForecastResult,
+} from "./ets.ts";
+export {
+ DLM,
+ buildLocalLevel,
+ buildLocalLinearTrend,
+ buildPolynomial,
+ buildFourier,
+ buildRegression,
+ combineDLMs,
+} from "./dlm.ts";
+export type {
+ DLMSpec,
+ DLMOptions,
+ DLMFilterStep,
+ DLMResult,
+ DLMSmootherResult,
+ DLMForecastResult,
+} from "./dlm.ts";
diff --git a/src/stats/kalman.ts b/src/stats/kalman.ts
new file mode 100644
index 000000000..a0493a782
--- /dev/null
+++ b/src/stats/kalman.ts
@@ -0,0 +1,809 @@
+/**
+ * kalman — Linear Gaussian State-Space Model: Kalman Filter & RTS Smoother.
+ *
+ * Implements the standard discrete-time Kalman filter (forward pass) and the
+ * Rauch-Tung-Striebel (RTS) smoother (backward pass) for linear dynamical
+ * systems:
+ *
+ * x_t = F·x_{t-1} + w_t, w_t ~ N(0, Q) (state equation)
+ * y_t = H·x_t + v_t, v_t ~ N(0, R) (observation equation)
+ * x_0 ~ N(m0, P0)
+ *
+ * Missing observations (null) are handled by skipping the update step —
+ * the filtered state reverts to the predicted state for that time-step.
+ *
+ * Mirrors the `statsmodels.tsa.statespace.kalman_filter.KalmanFilter` and
+ * `pykalman.KalmanFilter` APIs; factory helpers match common pandas patterns.
+ *
+ * Exported names:
+ * - {@link KalmanFilter} — main class (filter + smooth)
+ * - {@link KalmanFilterOptions} — constructor options
+ * - {@link KalmanFilterResult} — forward-pass output
+ * - {@link KalmanSmootherResult} — backward-pass output
+ *
+ * @example
+ * ```ts
+ * import { KalmanFilter } from "tsb";
+ *
+ * // Local-level model (random walk observed with noise)
+ * const kf = KalmanFilter.localLevel({ processNoise: 1, observationNoise: 2 });
+ * const res = kf.filter([[1], [2], [1.5], [null], [3], [2.5]]);
+ * console.log(res.filteredStateMeans); // [[…], …] T × 1
+ * console.log(res.logLikelihood);
+ *
+ * const sm = kf.smooth([[1], [2], [1.5], [null], [3], [2.5]]);
+ * console.log(sm.smoothedStateMeans);
+ * ```
+ *
+ * @module
+ */
+
+// ─── Internal matrix helpers ───────────────────────────────────────────────────
+
+/** Read-only row-major matrix. */
+type Mat = readonly (readonly number[])[];
+/** Mutable row-major matrix. */
+type MutMat = number[][];
+
+/** Rows of A (no checks — callers must ensure dimensions). */
+function rows(A: Mat): number {
+ return A.length;
+}
+/** Columns of A (0 if empty). */
+function cols(A: Mat): number {
+ return A[0]?.length ?? 0;
+}
+
+/** Create an n×m zero matrix. */
+function zeros(n: number, m: number): MutMat {
+ return Array.from({ length: n }, () => new Array(m).fill(0));
+}
+
+/** Create an n×n identity matrix. */
+function eye(n: number): MutMat {
+ return Array.from({ length: n }, (_, i) =>
+ Array.from({ length: n }, (_, j) => (i === j ? 1 : 0)),
+ );
+}
+
+/** Matrix–matrix product A (m×k) · B (k×n) → (m×n). */
+function mmul(A: Mat, B: Mat): MutMat {
+ const m = rows(A);
+ const k = cols(A);
+ const n = cols(B);
+ const C = zeros(m, n);
+ for (let i = 0; i < m; i++) {
+ const ai = A[i]!;
+ const ci = C[i]!;
+ for (let p = 0; p < k; p++) {
+ const aip = ai[p]!;
+ if (aip === 0) {
+ continue;
+ }
+ const bp = B[p]!;
+ for (let j = 0; j < n; j++) {
+ ci[j] = (ci[j] ?? 0) + aip * bp[j]!;
+ }
+ }
+ }
+ return C;
+}
+
+/** Matrix–vector product A (m×k) · x (k) → (m). */
+function mvmul(A: Mat, x: readonly number[]): number[] {
+ const m = rows(A);
+ const k = x.length;
+ const y = new Array(m).fill(0);
+ for (let i = 0; i < m; i++) {
+ const ai = A[i]!;
+ let s = 0;
+ for (let p = 0; p < k; p++) {
+ s += ai[p]! * x[p]!;
+ }
+ y[i] = s;
+ }
+ return y;
+}
+
+/** Transpose of A (m×n) → (n×m). */
+function T(A: Mat): MutMat {
+ const m = rows(A);
+ const n = cols(A);
+ const At = zeros(n, m);
+ for (let i = 0; i < m; i++) {
+ for (let j = 0; j < n; j++) {
+ At[j]![i] = A[i]?.[j]!;
+ }
+ }
+ return At;
+}
+
+/** A + B (element-wise). */
+function madd(A: Mat, B: Mat): MutMat {
+ const m = rows(A);
+ const n = cols(A);
+ return Array.from({ length: m }, (_, i) =>
+ Array.from({ length: n }, (_, j) => A[i]?.[j]! + B[i]?.[j]!),
+ );
+}
+
+/** A − B (element-wise). */
+function msub(A: Mat, B: Mat): MutMat {
+ const m = rows(A);
+ const n = cols(A);
+ return Array.from({ length: m }, (_, i) =>
+ Array.from({ length: n }, (_, j) => A[i]?.[j]! - B[i]?.[j]!),
+ );
+}
+
+/** Vector addition a + b. */
+function vadd(a: readonly number[], b: readonly number[]): number[] {
+ return a.map((ai, i) => ai + b[i]!);
+}
+
+/** Vector subtraction a − b. */
+function vsub(a: readonly number[], b: readonly number[]): number[] {
+ return a.map((ai, i) => ai - b[i]!);
+}
+
+/**
+ * Invert a square matrix via Gaussian elimination with partial pivoting.
+ * Returns null if the matrix is singular (|pivot| < 1e-14).
+ */
+function matInv(A: Mat): MutMat | null {
+ const n = rows(A);
+ // Augmented [A | I]
+ const aug: MutMat = Array.from({ length: n }, (_, i) => [
+ ...A[i]!,
+ ...Array.from({ length: n }, (_, j) => (i === j ? 1 : 0)),
+ ]);
+ for (let col = 0; col < n; col++) {
+ let maxRow = col;
+ let maxVal = Math.abs(aug[col]?.[col]!);
+ for (let row = col + 1; row < n; row++) {
+ const v = Math.abs(aug[row]?.[col]!);
+ if (v > maxVal) {
+ maxVal = v;
+ maxRow = row;
+ }
+ }
+ if (maxVal < 1e-14) {
+ return null;
+ }
+ [aug[col], aug[maxRow]] = [aug[maxRow]!, aug[col]!];
+ const pivot = aug[col]?.[col]!;
+ const pivRow = aug[col]!;
+ for (let j = 0; j < 2 * n; j++) {
+ pivRow[j] = pivRow[j]! / pivot;
+ }
+ for (let row = 0; row < n; row++) {
+ if (row === col) {
+ continue;
+ }
+ const fac = aug[row]?.[col]!;
+ if (fac === 0) {
+ continue;
+ }
+ const r = aug[row]!;
+ for (let j = 0; j < 2 * n; j++) {
+ r[j] = r[j]! - fac * pivRow[j]!;
+ }
+ }
+ }
+ return aug.map((row) => row.slice(n));
+}
+
+/** log-determinant via LU decomposition (for log-likelihood). */
+function logDet(A: Mat): number {
+ const n = rows(A);
+ const L: MutMat = Array.from({ length: n }, (_, i) => [...A[i]!]);
+ let logD = 0;
+ for (let col = 0; col < n; col++) {
+ let maxRow = col;
+ let maxVal = Math.abs(L[col]?.[col]!);
+ for (let row = col + 1; row < n; row++) {
+ const v = Math.abs(L[row]?.[col]!);
+ if (v > maxVal) {
+ maxVal = v;
+ maxRow = row;
+ }
+ }
+ if (maxRow !== col) {
+ [L[col], L[maxRow]] = [L[maxRow]!, L[col]!];
+ logD += Math.log(-1); // sign flip — handled by real part
+ }
+ const pivot = L[col]?.[col]!;
+ if (Math.abs(pivot) < 1e-300) {
+ return Number.NEGATIVE_INFINITY;
+ }
+ logD += Math.log(Math.abs(pivot));
+ for (let row = col + 1; row < n; row++) {
+ const fac = L[row]?.[col]! / pivot;
+ const r = L[row]!;
+ for (let j = col; j < n; j++) {
+ r[j] = r[j]! - fac * L[col]?.[j]!;
+ }
+ }
+ }
+ return logD;
+}
+
+/** Outer product a·bᵀ → matrix. */
+function outer(a: readonly number[], b: readonly number[]): MutMat {
+ return Array.from({ length: a.length }, (_, i) =>
+ Array.from({ length: b.length }, (_, j) => a[i]! * b[j]!),
+ );
+}
+
+/** Scale matrix by scalar. */
+function mscale(A: Mat, s: number): MutMat {
+ return A.map((row) => row.map((v) => v * s));
+}
+
+// ─── Public types ──────────────────────────────────────────────────────────────
+
+/** Constructor options for {@link KalmanFilter}. */
+export interface KalmanFilterOptions {
+ /**
+ * State transition matrix **F** (n_states × n_states).
+ * Defines how the state evolves: x_t = F·x_{t-1} + noise.
+ */
+ readonly transitionMatrix: readonly (readonly number[])[];
+ /**
+ * Observation matrix **H** (n_obs × n_states).
+ * Maps states to observations: y_t = H·x_t + noise.
+ */
+ readonly observationMatrix: readonly (readonly number[])[];
+ /**
+ * Process noise covariance **Q** (n_states × n_states).
+ * Covariance of the state-transition noise.
+ */
+ readonly processNoiseCov: readonly (readonly number[])[];
+ /**
+ * Observation noise covariance **R** (n_obs × n_obs).
+ * Covariance of the observation noise.
+ */
+ readonly observationNoiseCov: readonly (readonly number[])[];
+ /**
+ * Initial state mean **m₀** (n_states vector).
+ * Defaults to the zero vector.
+ */
+ readonly initialStateMean?: readonly number[];
+ /**
+ * Initial state covariance **P₀** (n_states × n_states).
+ * Defaults to the identity matrix.
+ */
+ readonly initialStateCovariance?: readonly (readonly number[])[];
+}
+
+/** Options for {@link KalmanFilter.localLevel}. */
+export interface LocalLevelOptions {
+ /** Process (state) noise variance σ²_q. Default: 1. */
+ readonly processNoise?: number;
+ /** Observation noise variance σ²_r. Default: 1. */
+ readonly observationNoise?: number;
+ /** Initial state mean scalar. Default: 0. */
+ readonly initialMean?: number;
+ /** Initial state variance scalar. Default: 1. */
+ readonly initialVariance?: number;
+}
+
+/** Options for {@link KalmanFilter.localLinearTrend}. */
+export interface LocalLinearTrendOptions {
+ /** Level process noise variance. Default: 1. */
+ readonly levelNoise?: number;
+ /** Slope process noise variance. Default: 0.1. */
+ readonly slopeNoise?: number;
+ /** Observation noise variance. Default: 1. */
+ readonly observationNoise?: number;
+ /** Initial [level, slope] mean vector. Default: [0, 0]. */
+ readonly initialMean?: readonly [number, number];
+ /** Initial state variance (diagonal). Default: 1. */
+ readonly initialVariance?: number;
+}
+
+/** Result of the Kalman filter forward pass. */
+export interface KalmanFilterResult {
+ /**
+ * Filtered state means x_{t|t} — shape T × n_states.
+ * Each row is the posterior mean after incorporating observation t.
+ */
+ readonly filteredStateMeans: readonly (readonly number[])[];
+ /**
+ * Filtered state covariances P_{t|t} — shape T × n_states × n_states.
+ * Each entry is the posterior covariance after incorporating observation t.
+ */
+ readonly filteredStateCovariances: readonly (readonly (readonly number[])[])[];
+ /**
+ * Predicted state means x_{t|t-1} — shape T × n_states.
+ * Each row is the prior mean before incorporating observation t.
+ */
+ readonly predictedStateMeans: readonly (readonly number[])[];
+ /**
+ * Predicted state covariances P_{t|t-1} — shape T × n_states × n_states.
+ */
+ readonly predictedStateCovariances: readonly (readonly (readonly number[])[])[];
+ /**
+ * Innovation (prediction error) vectors y_t − H·x_{t|t-1} — shape T × n_obs.
+ * NaN rows indicate missing observations.
+ */
+ readonly innovations: readonly (readonly number[])[];
+ /**
+ * Innovation covariance matrices S_t = H·P_{t|t-1}·Hᵀ + R — shape T × n_obs × n_obs.
+ */
+ readonly innovationCovariances: readonly (readonly (readonly number[])[])[];
+ /** Gaussian log-likelihood summed over all non-missing time-steps. */
+ readonly logLikelihood: number;
+ /** Number of states (n_states). */
+ readonly nStates: number;
+ /** Number of observation dimensions (n_obs). */
+ readonly nObs: number;
+ /** Number of time steps (T). */
+ readonly nTime: number;
+}
+
+/** Result of the RTS smoother backward pass. */
+export interface KalmanSmootherResult {
+ /**
+ * Smoothed state means x_{t|T} — shape T × n_states.
+ * Each row is the posterior mean using all T observations.
+ */
+ readonly smoothedStateMeans: readonly (readonly number[])[];
+ /**
+ * Smoothed state covariances P_{t|T} — shape T × n_states × n_states.
+ */
+ readonly smoothedStateCovariances: readonly (readonly (readonly number[])[])[];
+ /**
+ * Smoother gain matrices G_t — shape T × n_states × n_states.
+ * (Last entry is all-zeros by convention.)
+ */
+ readonly smootherGains: readonly (readonly (readonly number[])[])[];
+ /** Same as {@link KalmanFilterResult.logLikelihood} (computed in forward pass). */
+ readonly logLikelihood: number;
+ /** The forward-pass result used to compute the smoother. */
+ readonly filterResult: KalmanFilterResult;
+}
+
+// ─── KalmanFilter class ────────────────────────────────────────────────────────
+
+/**
+ * Linear Gaussian State-Space Model with Kalman filter and RTS smoother.
+ *
+ * The model is:
+ * ```
+ * x_t = F·x_{t-1} + w_t, w_t ~ N(0, Q)
+ * y_t = H·x_t + v_t, v_t ~ N(0, R)
+ * x_0 ~ N(m0, P0)
+ * ```
+ *
+ * @example
+ * ```ts
+ * const kf = new KalmanFilter({
+ * transitionMatrix: [[1]],
+ * observationMatrix: [[1]],
+ * processNoiseCov: [[1]],
+ * observationNoiseCov: [[2]],
+ * });
+ * const result = kf.filter([[1], [2], [null], [3]]);
+ * ```
+ */
+export class KalmanFilter {
+ /** State transition matrix F (n_states × n_states). */
+ readonly transitionMatrix: Mat;
+ /** Observation matrix H (n_obs × n_states). */
+ readonly observationMatrix: Mat;
+ /** Process noise covariance Q (n_states × n_states). */
+ readonly processNoiseCov: Mat;
+ /** Observation noise covariance R (n_obs × n_obs). */
+ readonly observationNoiseCov: Mat;
+ /** Initial state mean m₀ (n_states). */
+ readonly initialStateMean: readonly number[];
+ /** Initial state covariance P₀ (n_states × n_states). */
+ readonly initialStateCovariance: Mat;
+
+ constructor(opts: KalmanFilterOptions) {
+ this.transitionMatrix = opts.transitionMatrix;
+ this.observationMatrix = opts.observationMatrix;
+ this.processNoiseCov = opts.processNoiseCov;
+ this.observationNoiseCov = opts.observationNoiseCov;
+
+ const ns = rows(opts.transitionMatrix);
+ this.initialStateMean = opts.initialStateMean ?? new Array(ns).fill(0);
+ this.initialStateCovariance = opts.initialStateCovariance ?? eye(ns);
+ }
+
+ // ─── Factory helpers ───────────────────────────────────────────────────────
+
+ /**
+ * Local-level model (random walk + measurement noise):
+ * ```
+ * x_t = x_{t-1} + w_t, w_t ~ N(0, σ²_q)
+ * y_t = x_t + v_t, v_t ~ N(0, σ²_r)
+ * ```
+ */
+ static localLevel(opts: LocalLevelOptions = {}): KalmanFilter {
+ const q = opts.processNoise ?? 1;
+ const r = opts.observationNoise ?? 1;
+ const m0 = opts.initialMean ?? 0;
+ const p0 = opts.initialVariance ?? 1;
+ return new KalmanFilter({
+ transitionMatrix: [[1]],
+ observationMatrix: [[1]],
+ processNoiseCov: [[q]],
+ observationNoiseCov: [[r]],
+ initialStateMean: [m0],
+ initialStateCovariance: [[p0]],
+ });
+ }
+
+ /**
+ * Local linear trend model (level + slope):
+ * ```
+ * level_t = level_{t-1} + slope_{t-1} + w1_t
+ * slope_t = slope_{t-1} + w2_t
+ * y_t = level_t + v_t
+ * ```
+ */
+ static localLinearTrend(opts: LocalLinearTrendOptions = {}): KalmanFilter {
+ const ql = opts.levelNoise ?? 1;
+ const qs = opts.slopeNoise ?? 0.1;
+ const r = opts.observationNoise ?? 1;
+ const [m0l, m0s] = opts.initialMean ?? [0, 0];
+ const p0 = opts.initialVariance ?? 1;
+ return new KalmanFilter({
+ transitionMatrix: [
+ [1, 1],
+ [0, 1],
+ ],
+ observationMatrix: [[1, 0]],
+ processNoiseCov: [
+ [ql, 0],
+ [0, qs],
+ ],
+ observationNoiseCov: [[r]],
+ initialStateMean: [m0l, m0s],
+ initialStateCovariance: [
+ [p0, 0],
+ [0, p0],
+ ],
+ });
+ }
+
+ // ─── Main methods ──────────────────────────────────────────────────────────
+
+ /**
+ * Run the Kalman filter (forward pass).
+ *
+ * @param observations T × n_obs array of observations.
+ * Pass `null` for any element to indicate a missing value at that
+ * time-step × dimension. If an entire row is missing, pass a row of nulls
+ * or just pass `null` in a scalar array like `[[null]]`.
+ * @returns {@link KalmanFilterResult}
+ *
+ * @example
+ * ```ts
+ * const result = kf.filter([[1], [2], [null], [3], [2.5]]);
+ * ```
+ */
+ filter(observations: readonly (readonly (number | null)[])[]): KalmanFilterResult {
+ return kalmanFilter(
+ observations,
+ this.transitionMatrix,
+ this.observationMatrix,
+ this.processNoiseCov,
+ this.observationNoiseCov,
+ this.initialStateMean,
+ this.initialStateCovariance,
+ );
+ }
+
+ /**
+ * Run the RTS smoother (Kalman filter forward pass + RTS backward pass).
+ *
+ * @param observations T × n_obs array (same format as {@link filter}).
+ * @returns {@link KalmanSmootherResult}
+ *
+ * @example
+ * ```ts
+ * const smoothed = kf.smooth([[1], [2], [null], [3], [2.5]]);
+ * console.log(smoothed.smoothedStateMeans);
+ * ```
+ */
+ smooth(observations: readonly (readonly (number | null)[])[]): KalmanSmootherResult {
+ const fwd = this.filter(observations);
+ return rtsSmooth(fwd, this.transitionMatrix);
+ }
+}
+
+// ─── Core algorithms ───────────────────────────────────────────────────────────
+
+/**
+ * Kalman filter forward pass.
+ *
+ * Returns all intermediate quantities needed for the RTS smoother and for
+ * log-likelihood computation.
+ */
+function kalmanFilter(
+ obs: readonly (readonly (number | null)[])[],
+ F: Mat,
+ H: Mat,
+ Q: Mat,
+ R: Mat,
+ m0: readonly number[],
+ P0: Mat,
+): KalmanFilterResult {
+ const T_len = obs.length;
+ const ns = rows(F);
+ const no = rows(H);
+
+ // Storage
+ const filtMeans: number[][] = [];
+ const filtCovs: MutMat[] = [];
+ const predMeans: number[][] = [];
+ const predCovs: MutMat[] = [];
+ const innovations: number[][] = [];
+ const innovCovs: MutMat[] = [];
+ let logLik = 0;
+
+ // Initialize
+ let xFilt: number[] = [...m0];
+ let PFilt: MutMat = P0.map((row) => [...row]);
+ const FT = T(F);
+ const HT = T(H);
+ const LOG2PI = Math.log(2 * Math.PI);
+
+ for (let t = 0; t < T_len; t++) {
+ const yt = obs[t]!;
+
+ // ── Predict ────────────────────────────────────────────────────────────
+ const xPred = mvmul(F, xFilt);
+ // P_pred = F P F' + Q
+ const PPred = madd(mmul(mmul(F, PFilt), FT), Q);
+
+ predMeans.push(xPred);
+ predCovs.push(PPred);
+
+ // Innovation covariance S = H P_pred H' + R
+ const S = madd(mmul(mmul(H, PPred), HT), R);
+ innovCovs.push(S);
+
+ // Check if observation has any non-null values
+ const hasObs = yt.some((v) => v !== null);
+
+ if (!hasObs) {
+ // ── Missing observation: skip update ──────────────────────────────
+ innovations.push(new Array(no).fill(Number.NaN));
+ filtMeans.push(xPred);
+ filtCovs.push(PPred);
+ xFilt = xPred;
+ PFilt = PPred;
+ continue;
+ }
+
+ // ── Update ─────────────────────────────────────────────────────────────
+ const yHat = mvmul(H, xPred); // predicted observation
+ const innov = vsub(
+ yt.map((v) => (v === null ? 0 : v)), // treat null as 0 for innovation
+ yHat,
+ );
+
+ // For partial missing (some dims null), we handle by projecting to
+ // observed subspace. For simplicity: use full update with null→predicted.
+ innovations.push(innov);
+
+ // Kalman gain K = P_pred H' S^{-1}
+ const Sinv = matInv(S);
+ if (Sinv === null) {
+ // Singular innovation covariance: skip update
+ filtMeans.push(xPred);
+ filtCovs.push(PPred);
+ xFilt = xPred;
+ PFilt = PPred;
+ continue;
+ }
+
+ const K = mmul(mmul(PPred, HT), Sinv);
+
+ // x_filt = x_pred + K * innov
+ const xNew = vadd(xPred, mvmul(K, innov));
+
+ // P_filt = (I − K H) P_pred — Joseph form for numerical stability:
+ // P_filt = (I−KH) P (I−KH)' + K R K'
+ const IKH = msub(eye(ns), mmul(K, H));
+ const IKHPIKHT = mmul(mmul(IKH, PPred), T(IKH));
+ const KRKT = mmul(mmul(K, R), T(K));
+ const PNew: MutMat = madd(IKHPIKHT, KRKT);
+
+ // Log-likelihood contribution: -½ [d·log(2π) + log|S| + v'S⁻¹v]
+ const logDetS = logDet(S);
+ let vSv = 0;
+ for (let i = 0; i < no; i++) {
+ const Sinv_row = Sinv[i]!;
+ let sSinvRow = 0;
+ for (let j = 0; j < no; j++) {
+ sSinvRow += innov[j]! * Sinv_row[j]!;
+ }
+ vSv += innov[i]! * sSinvRow;
+ }
+ logLik -= 0.5 * (no * LOG2PI + logDetS + vSv);
+
+ filtMeans.push(xNew);
+ filtCovs.push(PNew);
+ xFilt = xNew;
+ PFilt = PNew;
+ }
+
+ return {
+ filteredStateMeans: filtMeans,
+ filteredStateCovariances: filtCovs,
+ predictedStateMeans: predMeans,
+ predictedStateCovariances: predCovs,
+ innovations,
+ innovationCovariances: innovCovs,
+ logLikelihood: logLik,
+ nStates: ns,
+ nObs: no,
+ nTime: T_len,
+ };
+}
+
+/**
+ * Rauch-Tung-Striebel (RTS) smoother backward pass.
+ *
+ * Given the Kalman filter result, runs the smoother backward from t=T to t=0.
+ *
+ * Smoother equations:
+ * G_t = P_{t|t} · Fᵀ · P_{t+1|t}^{-1}
+ * x_{t|T} = x_{t|t} + G_t · (x_{t+1|T} − x_{t+1|t})
+ * P_{t|T} = P_{t|t} + G_t · (P_{t+1|T} − P_{t+1|t}) · G_tᵀ
+ */
+function rtsSmooth(fwd: KalmanFilterResult, F: Mat): KalmanSmootherResult {
+ const T_len = fwd.nTime;
+ const ns = fwd.nStates;
+
+ const smoothMeans: number[][] = new Array(T_len);
+ const smoothCovs: MutMat[] = new Array(T_len);
+ const gains: MutMat[] = new Array(T_len);
+
+ // Initialise last time step from filter
+ const lastFiltMean = [...(fwd.filteredStateMeans[T_len - 1] ?? [])];
+ const lastFiltCov = (fwd.filteredStateCovariances[T_len - 1] ?? []).map((r) => [...r]);
+ smoothMeans[T_len - 1] = lastFiltMean;
+ smoothCovs[T_len - 1] = lastFiltCov;
+ gains[T_len - 1] = zeros(ns, ns);
+
+ const FT = T(F);
+
+ for (let t = T_len - 2; t >= 0; t--) {
+ const xFilt = fwd.filteredStateMeans[t]!;
+ const PFilt = fwd.filteredStateCovariances[t]!;
+ const PPred_next = fwd.predictedStateCovariances[t + 1]!;
+
+ // G_t = P_{t|t} · Fᵀ · P_{t+1|t}^{-1}
+ const PPredInv = matInv(PPred_next);
+ const G: MutMat = PPredInv !== null ? mmul(mmul(PFilt, FT), PPredInv) : zeros(ns, ns);
+
+ const xSmooth_next = smoothMeans[t + 1]!;
+ const PSmooth_next = smoothCovs[t + 1]!;
+ const xPred_next = fwd.predictedStateMeans[t + 1]!;
+
+ // x_{t|T} = x_{t|t} + G_t · (x_{t+1|T} − x_{t+1|t})
+ const dx = vsub(xSmooth_next, xPred_next);
+ smoothMeans[t] = vadd(xFilt, mvmul(G, dx));
+
+ // P_{t|T} = P_{t|t} + G_t · (P_{t+1|T} − P_{t+1|t}) · G_tᵀ
+ const dP = msub(PSmooth_next, PPred_next);
+ const GT = T(G);
+ smoothCovs[t] = madd(PFilt, mmul(mmul(G, dP), GT));
+
+ gains[t] = G;
+ }
+
+ return {
+ smoothedStateMeans: smoothMeans,
+ smoothedStateCovariances: smoothCovs,
+ smootherGains: gains,
+ logLikelihood: fwd.logLikelihood,
+ filterResult: fwd,
+ };
+}
+
+// ─── Standalone functional API ─────────────────────────────────────────────────
+
+/**
+ * Apply the Kalman filter to a sequence of scalar observations.
+ *
+ * Convenience wrapper for the common 1-D case (local-level model or similar).
+ *
+ * @example
+ * ```ts
+ * import { kalmanFilter1D } from "tsb";
+ * const { filteredStateMeans, logLikelihood } = kalmanFilter1D(
+ * [1, 2, null, 3, 2.5],
+ * { processNoise: 0.5, observationNoise: 1 },
+ * );
+ * ```
+ */
+export function kalmanFilter1D(
+ observations: readonly (number | null)[],
+ opts: LocalLevelOptions = {},
+): KalmanFilterResult {
+ const kf = KalmanFilter.localLevel(opts);
+ return kf.filter(observations.map((v) => [v]));
+}
+
+/**
+ * Apply the RTS smoother to scalar observations with a local-level model.
+ *
+ * @example
+ * ```ts
+ * import { kalmanSmooth1D } from "tsb";
+ * const { smoothedStateMeans } = kalmanSmooth1D([1, 2, null, 3, 2.5]);
+ * ```
+ */
+export function kalmanSmooth1D(
+ observations: readonly (number | null)[],
+ opts: LocalLevelOptions = {},
+): KalmanSmootherResult {
+ const kf = KalmanFilter.localLevel(opts);
+ return kf.smooth(observations.map((v) => [v]));
+}
+
+// ─── Utility: extract scalars from 1-D results ─────────────────────────────────
+
+/**
+ * Extract the scalar filtered means from a 1-state filter result.
+ * Returns an array of length T where each value is x_{t|t}[0].
+ *
+ * @example
+ * ```ts
+ * const result = kf.filter([[1], [2], [null], [3]]);
+ * const means = extractScalarMeans(result.filteredStateMeans);
+ * ```
+ */
+export function extractScalarMeans(means: readonly (readonly number[])[]): number[] {
+ return means.map((m) => m[0] ?? Number.NaN);
+}
+
+/**
+ * Extract the scalar filtered variances from a 1-state filter result.
+ * Returns an array of length T where each value is P_{t|t}[0][0].
+ *
+ * @example
+ * ```ts
+ * const result = kf.filter([[1], [2], [null], [3]]);
+ * const vars = extractScalarVariances(result.filteredStateCovariances);
+ * ```
+ */
+export function extractScalarVariances(
+ covs: readonly (readonly (readonly number[])[])[],
+): number[] {
+ return covs.map((P) => P[0]?.[0] ?? Number.NaN);
+}
+
+/**
+ * Compute a 95 % prediction interval around the filtered means for a 1-D
+ * local-level model.
+ *
+ * Returns `{ lower, upper }` arrays of length T.
+ *
+ * @example
+ * ```ts
+ * const result = kf.filter([[1], [2], [null], [3]]);
+ * const { lower, upper } = filteredPredictionInterval(result);
+ * ```
+ */
+export function filteredPredictionInterval(
+ result: KalmanFilterResult,
+ zScore = 1.96,
+): { lower: number[]; upper: number[] } {
+ const means = extractScalarMeans(result.filteredStateMeans);
+ const vars = extractScalarVariances(result.filteredStateCovariances);
+ return {
+ lower: means.map((m, i) => m - zScore * Math.sqrt(vars[i] ?? 0)),
+ upper: means.map((m, i) => m + zScore * Math.sqrt(vars[i] ?? 0)),
+ };
+}
+
+/** Alias kept for backward compat — use {@link KalmanFilter} directly. */
+export { KalmanFilter as StateSpaceModel };
diff --git a/src/stats/network_stats.ts b/src/stats/network_stats.ts
new file mode 100644
index 000000000..3dd2059d2
--- /dev/null
+++ b/src/stats/network_stats.ts
@@ -0,0 +1,410 @@
+/**
+ * network_stats — Graph/network analysis statistics.
+ *
+ * Implements:
+ * - **Graph representation** (adjacency list / matrix)
+ * - **Degree centrality**, **betweenness centrality** (exact, BFS-based)
+ * - **Clustering coefficient** (local and global)
+ * - **Shortest paths** (BFS for unweighted, Dijkstra for weighted)
+ * - **Connected components**
+ * - **PageRank** (power iteration)
+ * - **HITS** (hubs and authorities)
+ *
+ * @module
+ */
+
+// ─── Types ────────────────────────────────────────────────────────────────────
+
+/** An undirected or directed graph stored as adjacency lists. */
+export interface Graph {
+ /** Number of nodes. */
+ nNodes: number;
+ /** Adjacency list: adjacency[i] = [{to, weight}]. */
+ adjacency: { to: number; weight: number }[][];
+ /** Whether the graph is directed. */
+ directed: boolean;
+}
+
+/** Create a new empty graph. */
+export function createGraph(nNodes: number, directed = false): Graph {
+ const adjacency: { to: number; weight: number }[][] = [];
+ for (let i = 0; i < nNodes; i++) {
+ adjacency.push([]);
+ }
+ return { nNodes, adjacency, directed };
+}
+
+/**
+ * Add an edge to the graph.
+ *
+ * @param g - The graph.
+ * @param u - Source node.
+ * @param v - Target node.
+ * @param weight - Edge weight. Default 1.
+ */
+export function addEdge(g: Graph, u: number, v: number, weight = 1): void {
+ (g.adjacency[u] ?? []).push({ to: v, weight });
+ if (!g.directed) {
+ (g.adjacency[v] ?? []).push({ to: u, weight });
+ }
+}
+
+/**
+ * Build a graph from an edge list.
+ *
+ * @param nNodes - Total number of nodes.
+ * @param edges - Array of [u, v] or [u, v, weight] tuples.
+ * @param directed - Whether edges are directed. Default false.
+ * @returns Graph.
+ */
+export function graphFromEdges(
+ nNodes: number,
+ edges: [number, number][] | [number, number, number][],
+ directed = false,
+): Graph {
+ const g = createGraph(nNodes, directed);
+ for (const e of edges) {
+ const u = e[0];
+ const v = e[1];
+ const w = e[2] ?? 1;
+ addEdge(g, u, v, w);
+ }
+ return g;
+}
+
+// ─── Degree Centrality ────────────────────────────────────────────────────────
+
+/**
+ * Compute degree centrality for all nodes.
+ *
+ * Degree centrality = degree / (n - 1).
+ *
+ * @param g - The graph.
+ * @returns Array of centrality values (one per node).
+ */
+export function degreeCentrality(g: Graph): number[] {
+ const n = g.nNodes;
+ const norm = n > 1 ? n - 1 : 1;
+ return g.adjacency.map((adj) => adj.length / norm);
+}
+
+/**
+ * Compute in-degree and out-degree for directed graphs.
+ *
+ * @param g - The graph.
+ * @returns Object with inDegree and outDegree arrays.
+ */
+export function directedDegrees(g: Graph): { inDegree: number[]; outDegree: number[] } {
+ const inDegree = new Array(g.nNodes).fill(0);
+ const outDegree = g.adjacency.map((adj) => adj.length);
+
+ for (let u = 0; u < g.nNodes; u++) {
+ for (const { to } of g.adjacency[u] ?? []) {
+ inDegree[to] = (inDegree[to] ?? 0) + 1;
+ }
+ }
+
+ return { inDegree, outDegree };
+}
+
+// ─── Shortest Paths ───────────────────────────────────────────────────────────
+
+/**
+ * BFS single-source shortest paths (unweighted).
+ *
+ * @param g - The graph.
+ * @param source - Source node.
+ * @returns Array of distances from source (-1 if unreachable).
+ */
+export function bfsDistances(g: Graph, source: number): number[] {
+ const dist = new Array(g.nNodes).fill(-1);
+ dist[source] = 0;
+ const queue: number[] = [source];
+ let head = 0;
+
+ while (head < queue.length) {
+ const u = queue[head++] ?? 0;
+ for (const { to } of g.adjacency[u] ?? []) {
+ if ((dist[to] ?? -1) === -1) {
+ dist[to] = (dist[u] ?? 0) + 1;
+ queue.push(to);
+ }
+ }
+ }
+
+ return dist;
+}
+
+/**
+ * Dijkstra single-source shortest paths (weighted, non-negative weights).
+ *
+ * @param g - The graph.
+ * @param source - Source node.
+ * @returns Array of shortest-path distances from source (Number.POSITIVE_INFINITY if unreachable).
+ */
+export function dijkstra(g: Graph, source: number): number[] {
+ const dist = new Array(g.nNodes).fill(Number.POSITIVE_INFINITY);
+ dist[source] = 0;
+ // Simple priority queue via sorted array (sufficient for small graphs)
+ const pq: { node: number; d: number }[] = [{ node: source, d: 0 }];
+
+ while (pq.length > 0) {
+ pq.sort((a, b) => a.d - b.d);
+ const top = pq.shift();
+ if (top === undefined) break;
+ const { node: u, d } = top;
+ if (d > (dist[u] ?? Number.POSITIVE_INFINITY)) continue;
+
+ for (const { to, weight } of g.adjacency[u] ?? []) {
+ const nd = (dist[u] ?? Number.POSITIVE_INFINITY) + weight;
+ if (nd < (dist[to] ?? Number.POSITIVE_INFINITY)) {
+ dist[to] = nd;
+ pq.push({ node: to, d: nd });
+ }
+ }
+ }
+
+ return dist;
+}
+
+// ─── Betweenness Centrality ───────────────────────────────────────────────────
+
+/**
+ * Compute betweenness centrality (Brandes algorithm for unweighted graphs).
+ *
+ * @param g - The graph.
+ * @returns Array of betweenness centrality values (normalized).
+ */
+export function betweennessCentrality(g: Graph): number[] {
+ const n = g.nNodes;
+ const bc = new Array(n).fill(0);
+
+ for (let s = 0; s < n; s++) {
+ const stack: number[] = [];
+ const pred: number[][] = Array.from({ length: n }, () => []);
+ const sigma = new Array(n).fill(0);
+ sigma[s] = 1;
+ const dist = new Array(n).fill(-1);
+ dist[s] = 0;
+ const queue: number[] = [s];
+ let head = 0;
+
+ while (head < queue.length) {
+ const v = queue[head++] ?? 0;
+ stack.push(v);
+
+ for (const { to: w } of g.adjacency[v] ?? []) {
+ if ((dist[w] ?? -1) < 0) {
+ queue.push(w);
+ dist[w] = (dist[v] ?? 0) + 1;
+ }
+ if ((dist[w] ?? 0) === (dist[v] ?? 0) + 1) {
+ sigma[w] = (sigma[w] ?? 0) + (sigma[v] ?? 0);
+ (pred[w] ?? []).push(v);
+ }
+ }
+ }
+
+ const delta = new Array(n).fill(0);
+ while (stack.length > 0) {
+ const w = stack.pop() ?? 0;
+ for (const v of pred[w] ?? []) {
+ delta[v] = (delta[v] ?? 0) + ((sigma[v] ?? 0) / (sigma[w] ?? 1)) * (1 + (delta[w] ?? 0));
+ }
+ if (w !== s) bc[w] = (bc[w] ?? 0) + (delta[w] ?? 0);
+ }
+ }
+
+ // Normalize
+ const norm = g.directed ? (n - 1) * (n - 2) : ((n - 1) * (n - 2)) / 2;
+ if (norm > 0) {
+ for (let i = 0; i < n; i++) {
+ bc[i] = (bc[i] ?? 0) / norm;
+ }
+ }
+
+ return bc;
+}
+
+// ─── Clustering Coefficient ───────────────────────────────────────────────────
+
+/**
+ * Compute local clustering coefficient for each node (undirected).
+ *
+ * @param g - The graph (undirected).
+ * @returns Array of clustering coefficients.
+ */
+export function clusteringCoefficient(g: Graph): number[] {
+ const n = g.nNodes;
+ const cc = new Array(n).fill(0);
+
+ for (let u = 0; u < n; u++) {
+ const neighbors = new Set((g.adjacency[u] ?? []).map((e) => e.to));
+ const k = neighbors.size;
+ if (k < 2) {
+ cc[u] = 0;
+ continue;
+ }
+
+ let triangles = 0;
+ for (const v of neighbors) {
+ for (const { to: w } of g.adjacency[v] ?? []) {
+ if (neighbors.has(w)) triangles++;
+ }
+ }
+
+ cc[u] = triangles / (k * (k - 1));
+ }
+
+ return cc;
+}
+
+/**
+ * Compute the global (average) clustering coefficient.
+ *
+ * @param g - The graph.
+ * @returns Global clustering coefficient.
+ */
+export function globalClusteringCoefficient(g: Graph): number {
+ const cc = clusteringCoefficient(g);
+ return cc.reduce((a, b) => a + b, 0) / cc.length;
+}
+
+// ─── Connected Components ─────────────────────────────────────────────────────
+
+/**
+ * Find all connected components (undirected graph).
+ *
+ * @param g - The graph.
+ * @returns Array of components, each an array of node indices.
+ */
+export function connectedComponents(g: Graph): number[][] {
+ const visited = new Array(g.nNodes).fill(false);
+ const components: number[][] = [];
+
+ for (let start = 0; start < g.nNodes; start++) {
+ if (visited[start] ?? false) continue;
+
+ const component: number[] = [];
+ const queue: number[] = [start];
+ visited[start] = true;
+
+ while (queue.length > 0) {
+ const u = queue.shift() ?? 0;
+ component.push(u);
+ for (const { to } of g.adjacency[u] ?? []) {
+ if (!(visited[to] ?? false)) {
+ visited[to] = true;
+ queue.push(to);
+ }
+ }
+ }
+
+ components.push(component);
+ }
+
+ return components;
+}
+
+// ─── PageRank ─────────────────────────────────────────────────────────────────
+
+/**
+ * Compute PageRank via power iteration.
+ *
+ * @param g - The graph (directed).
+ * @param dampingFactor - Damping factor (alpha). Default 0.85.
+ * @param maxIter - Maximum iterations. Default 100.
+ * @param tol - Convergence tolerance. Default 1e-6.
+ * @returns PageRank scores (sum to 1).
+ *
+ * @example
+ * ```ts
+ * import { graphFromEdges, pageRank } from "tsb";
+ * const g = graphFromEdges(4, [[0,1],[0,2],[1,3],[2,3]], true);
+ * const pr = pageRank(g);
+ * ```
+ */
+export function pageRank(g: Graph, dampingFactor = 0.85, maxIter = 100, tol = 1e-6): number[] {
+ const n = g.nNodes;
+ let rank = new Array(n).fill(1 / n);
+ const outDeg = g.adjacency.map((adj) => adj.length);
+
+ for (let iter = 0; iter < maxIter; iter++) {
+ const newRank = new Array(n).fill((1 - dampingFactor) / n);
+
+ for (let u = 0; u < n; u++) {
+ const deg = outDeg[u] ?? 0;
+ if (deg === 0) {
+ // Dangling node — distribute rank equally
+ const share = (rank[u] ?? 0) / n;
+ for (let v = 0; v < n; v++) {
+ newRank[v] = (newRank[v] ?? 0) + dampingFactor * share;
+ }
+ } else {
+ for (const { to: v } of g.adjacency[u] ?? []) {
+ newRank[v] = (newRank[v] ?? 0) + dampingFactor * ((rank[u] ?? 0) / deg);
+ }
+ }
+ }
+
+ let diff = 0;
+ for (let i = 0; i < n; i++) {
+ diff += Math.abs((newRank[i] ?? 0) - (rank[i] ?? 0));
+ }
+ rank = newRank;
+ if (diff < tol) break;
+ }
+
+ return rank;
+}
+
+// ─── HITS ─────────────────────────────────────────────────────────────────────
+
+/**
+ * Compute HITS (Hubs and Authorities) scores via power iteration.
+ *
+ * @param g - The graph (directed).
+ * @param maxIter - Maximum iterations. Default 100.
+ * @param tol - Convergence tolerance. Default 1e-6.
+ * @returns Object with hub and authority score arrays.
+ */
+export function hits(g: Graph, maxIter = 100, tol = 1e-6): { hub: number[]; authority: number[] } {
+ const n = g.nNodes;
+ let hub = new Array(n).fill(1 / n);
+ let auth = new Array(n).fill(1 / n);
+
+ for (let iter = 0; iter < maxIter; iter++) {
+ const newAuth = new Array(n).fill(0);
+ const newHub = new Array(n).fill(0);
+
+ for (let u = 0; u < n; u++) {
+ for (const { to: v } of g.adjacency[u] ?? []) {
+ newAuth[v] = (newAuth[v] ?? 0) + (hub[u] ?? 0);
+ }
+ }
+ for (let u = 0; u < n; u++) {
+ for (const { to: v } of g.adjacency[u] ?? []) {
+ newHub[u] = (newHub[u] ?? 0) + (auth[v] ?? 0);
+ }
+ }
+
+ // Normalize
+ const authNorm = Math.sqrt(newAuth.reduce((s, x) => s + x * x, 0)) || 1;
+ const hubNorm = Math.sqrt(newHub.reduce((s, x) => s + x * x, 0)) || 1;
+ for (let i = 0; i < n; i++) {
+ newAuth[i] = (newAuth[i] ?? 0) / authNorm;
+ newHub[i] = (newHub[i] ?? 0) / hubNorm;
+ }
+
+ let diff = 0;
+ for (let i = 0; i < n; i++) {
+ diff +=
+ Math.abs((newAuth[i] ?? 0) - (auth[i] ?? 0)) + Math.abs((newHub[i] ?? 0) - (hub[i] ?? 0));
+ }
+ auth = newAuth;
+ hub = newHub;
+ if (diff < tol) break;
+ }
+
+ return { hub, authority: auth };
+}
diff --git a/src/stats/signal.ts b/src/stats/signal.ts
new file mode 100644
index 000000000..a6a3cd26b
--- /dev/null
+++ b/src/stats/signal.ts
@@ -0,0 +1,781 @@
+/**
+ * signal — Signal processing: FFT, windows, STFT, Welch PSD, periodogram.
+ *
+ * Mirrors `numpy.fft`, `scipy.signal` spectral and window utilities.
+ * Implemented from scratch with no external dependencies.
+ *
+ * FFT:
+ * - {@link fft} — N-point complex DFT (radix-2, pads to power of 2)
+ * - {@link ifft} — inverse FFT
+ * - {@link rfft} — real-input FFT (one-sided)
+ * - {@link irfft} — inverse real FFT
+ * - {@link fftFreq} — DFT sample frequencies
+ * - {@link rfftFreq} — one-sided DFT sample frequencies
+ * - {@link fftshift} — shift zero-frequency to centre
+ * - {@link ifftshift} — inverse of fftshift
+ *
+ * Windows (via {@link getWindow}):
+ * - `"rectangular"`, `"bartlett"`, `"hann"`, `"hamming"`, `"blackman"`,
+ * `"blackmanharris"`, `"flattop"`, `"kaiser"`
+ *
+ * Spectral analysis:
+ * - {@link stft} — Short-Time Fourier Transform
+ * - {@link istft} — Inverse STFT (overlap-add)
+ * - {@link welch} — Welch power spectral density
+ * - {@link periodogram} — Periodogram PSD estimate
+ *
+ * @example
+ * ```ts
+ * import { fft, rfftFreq, welch } from "tsb";
+ *
+ * const x = [1, 0, -1, 0, 1, 0, -1, 0];
+ * const X = fft(x); // 8-point FFT
+ * const freqs = rfftFreq(X.length, 1 / 100);
+ *
+ * const { f, Pxx } = welch(x, { fs: 100 });
+ * ```
+ *
+ * @module
+ */
+
+// ─── complex arithmetic ───────────────────────────────────────────────────────
+
+/** A complex number `{ re, im }`. */
+export type Complex = { re: number; im: number };
+
+/** Construct a complex number. */
+export function complex(re: number, im: number): Complex {
+ return { re, im };
+}
+
+/** Add two complex numbers. */
+function cAdd(a: Complex, b: Complex): Complex {
+ return { re: a.re + b.re, im: a.im + b.im };
+}
+
+/** Subtract two complex numbers. */
+function cSub(a: Complex, b: Complex): Complex {
+ return { re: a.re - b.re, im: a.im - b.im };
+}
+
+/** Multiply two complex numbers. */
+function cMul(a: Complex, b: Complex): Complex {
+ return { re: a.re * b.re - a.im * b.im, im: a.re * b.im + a.im * b.re };
+}
+
+/** Complex conjugate. */
+function cConj(a: Complex): Complex {
+ return { re: a.re, im: -a.im };
+}
+
+/** Magnitude squared |a|². */
+function cAbsSq(a: Complex): number {
+ return a.re * a.re + a.im * a.im;
+}
+
+/** Magnitude |a|. */
+export function cAbs(a: Complex): number {
+ return Math.sqrt(cAbsSq(a));
+}
+
+/** Phase angle (arg) of a complex number. */
+export function cArg(a: Complex): number {
+ return Math.atan2(a.im, a.re);
+}
+
+// ─── FFT internals ────────────────────────────────────────────────────────────
+
+/** Smallest power of 2 ≥ n. */
+function nextPow2(n: number): number {
+ if (n <= 1) {
+ return 1;
+ }
+ let p = 1;
+ while (p < n) {
+ p <<= 1;
+ }
+ return p;
+}
+
+/** In-place bit-reversal permutation. */
+function bitReverse(arr: Complex[], n: number): void {
+ let j = 0;
+ for (let i = 1; i < n; i++) {
+ let bit = n >> 1;
+ for (; j & bit; bit >>= 1) {
+ j ^= bit;
+ }
+ j ^= bit;
+ if (i < j) {
+ const tmp = arr[i]!;
+ arr[i] = arr[j]!;
+ arr[j] = tmp;
+ }
+ }
+}
+
+/** Cooley-Tukey radix-2 DIT iterative FFT, in-place. `n` must be a power of 2. */
+// biome-ignore lint/complexity/noExcessiveCognitiveComplexity: nested FFT butterfly loops
+function fftInPlace(arr: Complex[], n: number, inverse: boolean): void {
+ bitReverse(arr, n);
+ for (let len = 2; len <= n; len <<= 1) {
+ const half = len >> 1;
+ const ang = ((inverse ? 2 : -2) * Math.PI) / len;
+ const wLen: Complex = { re: Math.cos(ang), im: Math.sin(ang) };
+ for (let i = 0; i < n; i += len) {
+ let w: Complex = { re: 1, im: 0 };
+ for (let j = 0; j < half; j++) {
+ const u = arr[i + j]!;
+ const v = cMul(arr[i + j + half]!, w);
+ arr[i + j] = cAdd(u, v);
+ arr[i + j + half] = cSub(u, v);
+ w = cMul(w, wLen);
+ }
+ }
+ }
+ if (inverse) {
+ for (let i = 0; i < n; i++) {
+ const a = arr[i]!;
+ arr[i] = { re: a.re / n, im: a.im / n };
+ }
+ }
+}
+
+// ─── public FFT functions ─────────────────────────────────────────────────────
+
+/**
+ * Compute the discrete Fourier transform of a real signal.
+ *
+ * If `x.length` is not a power of 2, the signal is zero-padded to the next
+ * power of 2. The returned array has length `nextPow2(x.length)`.
+ *
+ * @param x - Input signal (real values).
+ * @returns Complex DFT coefficients.
+ */
+export function fft(x: readonly number[]): Complex[] {
+ const n = x.length;
+ const m = nextPow2(n);
+ const buf: Complex[] = Array.from({ length: m }, (_, i) => ({
+ re: x[i] ?? 0,
+ im: 0,
+ }));
+ fftInPlace(buf, m, false);
+ return buf;
+}
+
+/**
+ * Compute the inverse discrete Fourier transform.
+ *
+ * The length of `X` must be a power of 2. Returns a complex array of the
+ * same length as `X`.
+ *
+ * @param X - Complex DFT coefficients.
+ * @returns Complex inverse-DFT output.
+ */
+export function ifft(X: readonly Complex[]): Complex[] {
+ const n = X.length;
+ const m = nextPow2(n);
+ const buf: Complex[] = Array.from({ length: m }, (_, i) => {
+ const c = X[i] ?? { re: 0, im: 0 };
+ return { re: c.re, im: c.im };
+ });
+ fftInPlace(buf, m, true);
+ return buf;
+}
+
+/**
+ * Real-input FFT (one-sided). Returns the first `floor(m/2) + 1` bins where
+ * `m = nextPow2(x.length)`.
+ *
+ * @param x - Real input signal.
+ * @returns One-sided complex spectrum.
+ */
+export function rfft(x: readonly number[]): Complex[] {
+ const full = fft(x);
+ return full.slice(0, Math.floor(full.length / 2) + 1);
+}
+
+/**
+ * Inverse real FFT. Reconstructs a real signal from one-sided spectrum.
+ *
+ * @param X - One-sided spectrum as from {@link rfft}.
+ * @param n - Optional output length (defaults to `2 * (X.length - 1)`).
+ * @returns Real signal.
+ */
+export function irfft(X: readonly Complex[], n?: number): number[] {
+ const nOut = n ?? 2 * (X.length - 1);
+ const m = nextPow2(nOut);
+ const buf: Complex[] = new Array(m);
+ const half = X.length;
+ for (let i = 0; i < half; i++) {
+ buf[i] = { re: (X[i] ?? { re: 0, im: 0 }).re, im: (X[i] ?? { re: 0, im: 0 }).im };
+ }
+ for (let i = half; i < m; i++) {
+ const j = m - i;
+ const c = X[j] ?? { re: 0, im: 0 };
+ buf[i] = cConj(c);
+ }
+ fftInPlace(buf, m, true);
+ return Array.from({ length: nOut }, (_, i) => buf[i]?.re ?? 0);
+}
+
+/**
+ * DFT sample frequencies for an `n`-point FFT with sample spacing `d`.
+ *
+ * @param n - FFT length (from `fft(x).length`).
+ * @param d - Sample spacing in seconds (default `1`).
+ * @returns Array of frequencies from `0` to `(n-1)/(n*d)`, wrapped to negative.
+ */
+export function fftFreq(n: number, d = 1): number[] {
+ const f: number[] = new Array(n);
+ const half = Math.floor(n / 2) + 1;
+ for (let i = 0; i < half; i++) {
+ f[i] = i / (n * d);
+ }
+ for (let i = half; i < n; i++) {
+ f[i] = (i - n) / (n * d);
+ }
+ return f;
+}
+
+/**
+ * One-sided DFT sample frequencies for a real-input FFT.
+ *
+ * @param n - FFT length (from `rfft(x).length` etc.).
+ * @param d - Sample spacing in seconds (default `1`).
+ * @returns Frequencies `[0, 1/(n*d), 2/(n*d), ..., 1/(2*d)]`.
+ */
+export function rfftFreq(n: number, d = 1): number[] {
+ const half = Math.floor(n / 2) + 1;
+ return Array.from({ length: half }, (_, i) => i / (n * d));
+}
+
+/**
+ * Shift the zero-frequency component to the centre of the spectrum.
+ * Equivalent to `numpy.fft.fftshift`.
+ */
+export function fftshift(x: readonly T[]): T[] {
+ const n = x.length;
+ const half = Math.ceil(n / 2);
+ return [...x.slice(half), ...x.slice(0, half)];
+}
+
+/**
+ * Inverse of {@link fftshift}. Equivalent to `numpy.fft.ifftshift`.
+ */
+export function ifftshift(x: readonly T[]): T[] {
+ const n = x.length;
+ const half = Math.floor(n / 2);
+ return [...x.slice(half), ...x.slice(0, half)];
+}
+
+// ─── window functions ─────────────────────────────────────────────────────────
+
+/** Supported window names. */
+export type WindowName =
+ | "rectangular"
+ | "bartlett"
+ | "hann"
+ | "hamming"
+ | "blackman"
+ | "blackmanharris"
+ | "flattop"
+ | "kaiser";
+
+/** Modified Bessel function of the first kind, order 0, I₀(x). (A&S 9.8.1) */
+function besselI0(x: number): number {
+ const ax = Math.abs(x);
+ if (ax < 3.75) {
+ const t = (x / 3.75) ** 2;
+ return (
+ 1 +
+ t *
+ (3.5156229 +
+ t * (3.0899424 + t * (1.2067492 + t * (0.2659732 + t * (0.0360768 + t * 0.0045813)))))
+ );
+ }
+ const t = 3.75 / ax;
+ return (
+ (Math.exp(ax) / Math.sqrt(ax)) *
+ (0.39894228 +
+ t *
+ (0.01328592 +
+ t *
+ (0.00225319 +
+ t *
+ (-0.00157565 +
+ t *
+ (0.00916281 +
+ t * (-0.02057706 + t * (0.02635537 + t * (-0.01647633 + t * 0.00392377))))))))
+ );
+}
+
+/** Rectangular (boxcar) window — all ones. */
+export function rectangularWindow(n: number): number[] {
+ return Array.from({ length: n }, () => 1);
+}
+
+/** Bartlett (triangular) window. */
+export function bartlettWindow(n: number): number[] {
+ return Array.from({ length: n }, (_, i) => 1 - Math.abs((2 * i - (n - 1)) / (n - 1)));
+}
+
+/** Hann window — `0.5 * (1 - cos(2πi/(N-1)))`. */
+export function hannWindow(n: number): number[] {
+ return Array.from({ length: n }, (_, i) => 0.5 * (1 - Math.cos((2 * Math.PI * i) / (n - 1))));
+}
+
+/** Hamming window — `0.54 - 0.46 * cos(2πi/(N-1))`. */
+export function hammingWindow(n: number): number[] {
+ return Array.from({ length: n }, (_, i) => 0.54 - 0.46 * Math.cos((2 * Math.PI * i) / (n - 1)));
+}
+
+/** Blackman window. */
+export function blackmanWindow(n: number): number[] {
+ return Array.from(
+ { length: n },
+ (_, i) =>
+ 0.42 -
+ 0.5 * Math.cos((2 * Math.PI * i) / (n - 1)) +
+ 0.08 * Math.cos((4 * Math.PI * i) / (n - 1)),
+ );
+}
+
+/** Blackman-Harris window (4-term). */
+export function blackmanHarrisWindow(n: number): number[] {
+ const a0 = 0.35875;
+ const a1 = 0.48829;
+ const a2 = 0.14128;
+ const a3 = 0.01168;
+ return Array.from(
+ { length: n },
+ (_, i) =>
+ a0 -
+ a1 * Math.cos((2 * Math.PI * i) / (n - 1)) +
+ a2 * Math.cos((4 * Math.PI * i) / (n - 1)) -
+ a3 * Math.cos((6 * Math.PI * i) / (n - 1)),
+ );
+}
+
+/** Flat-top window (5-term). */
+export function flatTopWindow(n: number): number[] {
+ const a0 = 0.21557895;
+ const a1 = 0.41663158;
+ const a2 = 0.277263158;
+ const a3 = 0.083578947;
+ const a4 = 0.006947368;
+ return Array.from(
+ { length: n },
+ (_, i) =>
+ a0 -
+ a1 * Math.cos((2 * Math.PI * i) / (n - 1)) +
+ a2 * Math.cos((4 * Math.PI * i) / (n - 1)) -
+ a3 * Math.cos((6 * Math.PI * i) / (n - 1)) +
+ a4 * Math.cos((8 * Math.PI * i) / (n - 1)),
+ );
+}
+
+/**
+ * Kaiser window with shape parameter `beta`.
+ *
+ * @param n - Number of samples.
+ * @param beta - Shape parameter (controls main-lobe width vs side-lobe level).
+ */
+export function kaiserWindow(n: number, beta: number): number[] {
+ const i0b = besselI0(beta);
+ return Array.from({ length: n }, (_, i) => {
+ const t = (2 * i) / (n - 1) - 1;
+ return besselI0(beta * Math.sqrt(1 - t * t)) / i0b;
+ });
+}
+
+/**
+ * Create a window of length `n` by name.
+ *
+ * @param name - Window function name.
+ * @param n - Number of samples.
+ * @param beta - Kaiser window `beta` parameter (ignored for other windows).
+ * @returns - Window samples.
+ */
+export function getWindow(name: WindowName, n: number, beta = 14): number[] {
+ switch (name) {
+ case "rectangular":
+ return rectangularWindow(n);
+ case "bartlett":
+ return bartlettWindow(n);
+ case "hann":
+ return hannWindow(n);
+ case "hamming":
+ return hammingWindow(n);
+ case "blackman":
+ return blackmanWindow(n);
+ case "blackmanharris":
+ return blackmanHarrisWindow(n);
+ case "flattop":
+ return flatTopWindow(n);
+ case "kaiser":
+ return kaiserWindow(n, beta);
+ }
+}
+
+// ─── STFT / ISTFT ─────────────────────────────────────────────────────────────
+
+/** Options for {@link stft}. */
+export interface STFTOptions {
+ /** Sampling frequency in Hz (default `1`). */
+ fs?: number;
+ /** Segment length in samples (default `256`). */
+ nperseg?: number;
+ /** Number of samples to overlap between segments (default `nperseg / 2`). */
+ noverlap?: number;
+ /** FFT length ≥ `nperseg` (default = `nperseg`, padded to power of 2). */
+ nfft?: number;
+ /** Window function to apply (default `"hann"`). */
+ window?: WindowName | readonly number[];
+ /** Boundary extension mode: `"zeros"` pads with zeros, `null` no padding. */
+ boundary?: "zeros" | null;
+}
+
+/** STFT result object. */
+export interface STFTResult {
+ /** Time centres for each frame (seconds). */
+ t: number[];
+ /** Frequency bins (Hz). */
+ f: number[];
+ /** Complex STFT matrix `Zxx[freqIdx][timeIdx]`. */
+ Zxx: Complex[][];
+}
+
+/**
+ * Short-Time Fourier Transform.
+ *
+ * Mirrors `scipy.signal.stft`. Splits the signal into overlapping windowed
+ * segments and computes the FFT of each segment.
+ *
+ * @param x - Input signal.
+ * @param options - {@link STFTOptions}.
+ * @returns - {@link STFTResult} with `{ t, f, Zxx }`.
+ *
+ * @example
+ * ```ts
+ * import { stft } from "tsb";
+ * const x = Array.from({ length: 1024 }, (_, i) => Math.sin(2 * Math.PI * 10 * i / 512));
+ * const { t, f, Zxx } = stft(x, { fs: 512, nperseg: 128 });
+ * ```
+ */
+// biome-ignore lint/complexity/noExcessiveCognitiveComplexity: STFT nested loops
+export function stft(x: readonly number[], options: STFTOptions = {}): STFTResult {
+ const fs = options.fs ?? 1;
+ const nperseg = options.nperseg ?? Math.min(256, x.length);
+ const noverlap = options.noverlap ?? Math.floor(nperseg / 2);
+ const nfft = nextPow2(options.nfft ?? nperseg);
+ const step = nperseg - noverlap;
+ const boundary = options.boundary ?? "zeros";
+
+ // Build window
+ const win: number[] =
+ options.window !== undefined
+ ? typeof options.window === "string"
+ ? getWindow(options.window, nperseg)
+ : Array.from(options.window)
+ : hannWindow(nperseg);
+
+ // Pad signal at boundaries
+ const pad = boundary === "zeros" ? Math.floor(nperseg / 2) : 0;
+ const padded: number[] = [
+ ...Array.from({ length: pad }).fill(0),
+ ...x,
+ ...Array.from({ length: pad }).fill(0),
+ ];
+
+ // Number of frames
+ const nFrames = Math.floor((padded.length - noverlap) / step);
+ const nFreqs = Math.floor(nfft / 2) + 1;
+
+ const Zxx: Complex[][] = Array.from({ length: nFreqs }, () => new Array(nFrames));
+ const tArr: number[] = new Array(nFrames);
+ const fArr: number[] = Array.from({ length: nFreqs }, (_, i) => (i * fs) / nfft);
+
+ for (let k = 0; k < nFrames; k++) {
+ const start = k * step;
+ const seg: Complex[] = Array.from({ length: nfft }, (_, i) => ({
+ re: (padded[start + i] ?? 0) * (win[i] ?? 1),
+ im: 0,
+ }));
+ fftInPlace(seg, nfft, false);
+ for (let fi = 0; fi < nFreqs; fi++) {
+ const col = Zxx[fi];
+ if (col !== undefined) {
+ col[k] = seg[fi] ?? { re: 0, im: 0 };
+ }
+ }
+ tArr[k] = (start + nperseg / 2 - pad) / fs;
+ }
+
+ return { t: tArr, f: fArr, Zxx };
+}
+
+/** Options for {@link istft}. */
+export interface ISTFTOptions {
+ /** Segment length in samples (used to determine overlap). */
+ nperseg?: number;
+ /** Number of samples to overlap between segments (default `nperseg / 2`). */
+ noverlap?: number;
+ /** FFT length (default = `2 * (nFreqs - 1)` where `nFreqs = Zxx.length`). */
+ nfft?: number;
+ /** Window function (default `"hann"`). */
+ window?: WindowName | readonly number[];
+ /** Boundary extension used in stft (default `"zeros"`). */
+ boundary?: "zeros" | null;
+}
+
+/**
+ * Inverse Short-Time Fourier Transform (overlap-add).
+ *
+ * Mirrors `scipy.signal.istft`.
+ *
+ * @param Zxx - Complex STFT matrix `[freqIdx][timeIdx]` (from {@link stft}).
+ * @param options - {@link ISTFTOptions}.
+ * @returns - Reconstructed real signal.
+ */
+// biome-ignore lint/complexity/noExcessiveCognitiveComplexity: ISTFT overlap-add loops
+export function istft(Zxx: readonly (readonly Complex[])[], options: ISTFTOptions = {}): number[] {
+ const nFreqs = Zxx.length;
+ const nFrames = nFreqs > 0 ? (Zxx[0]?.length ?? 0) : 0;
+ const nfft = options.nfft ?? 2 * (nFreqs - 1);
+ const nperseg = options.nperseg ?? nfft;
+ const noverlap = options.noverlap ?? Math.floor(nperseg / 2);
+ const step = nperseg - noverlap;
+
+ const win: number[] =
+ options.window !== undefined
+ ? typeof options.window === "string"
+ ? getWindow(options.window, nperseg)
+ : Array.from(options.window)
+ : hannWindow(nperseg);
+
+ const boundary = options.boundary ?? "zeros";
+ const pad = boundary === "zeros" ? Math.floor(nperseg / 2) : 0;
+ const outLen = nFrames * step + nperseg;
+
+ const output = new Float64Array(outLen);
+ const windowSum = new Float64Array(outLen);
+ const winSq = win.map((w) => w * w);
+
+ for (let k = 0; k < nFrames; k++) {
+ // Build full-spectrum (two-sided) for IFFT
+ const buf: Complex[] = new Array(nfft);
+ for (let fi = 0; fi < nFreqs; fi++) {
+ buf[fi] = Zxx[fi]?.[k] ?? { re: 0, im: 0 };
+ }
+ for (let fi = nFreqs; fi < nfft; fi++) {
+ const mirrorIdx = nfft - fi;
+ const src = Zxx[mirrorIdx]?.[k] ?? { re: 0, im: 0 };
+ buf[fi] = cConj(src);
+ }
+ fftInPlace(buf, nfft, true);
+
+ const start = k * step;
+ for (let i = 0; i < nperseg; i++) {
+ const idx = start + i;
+ output[idx] = (output[idx] ?? 0) + (buf[i]?.re ?? 0) * (win[i] ?? 1);
+ windowSum[idx] = (windowSum[idx] ?? 0) + (winSq[i] ?? 0);
+ }
+ }
+
+ // Normalize and trim boundary padding
+ const result: number[] = [];
+ const start = pad;
+ const end = outLen - pad;
+ for (let i = start; i < end; i++) {
+ const ws = windowSum[i] ?? 0;
+ result.push(ws > 1e-10 ? (output[i] ?? 0) / ws : 0);
+ }
+
+ return result;
+}
+
+// ─── Welch PSD ────────────────────────────────────────────────────────────────
+
+/** Options for {@link welch}. */
+export interface WelchOptions {
+ /** Sampling frequency in Hz (default `1`). */
+ fs?: number;
+ /** Segment length (default `min(256, x.length)`). */
+ nperseg?: number;
+ /** Overlap between segments (default `nperseg / 2`). */
+ noverlap?: number;
+ /** FFT length (default `nperseg`, padded to power of 2). */
+ nfft?: number;
+ /** Window function (default `"hann"`). */
+ window?: WindowName | readonly number[];
+ /** Averaging method: `"mean"` (default) or `"median"`. */
+ average?: "mean" | "median";
+ /** Scaling: `"density"` (PSD, V²/Hz) or `"spectrum"` (power spectrum V²). */
+ scaling?: "density" | "spectrum";
+}
+
+/** PSD result: frequency bins and power spectral density estimates. */
+export interface PSDResult {
+ /** Frequency bins in Hz. */
+ f: number[];
+ /** Power spectral density (or power spectrum) at each frequency. */
+ Pxx: number[];
+}
+
+/**
+ * Welch power spectral density estimate.
+ *
+ * Divides the signal into overlapping segments, computes the periodogram for
+ * each, and averages. Mirrors `scipy.signal.welch`.
+ *
+ * @param x - Input signal.
+ * @param options - {@link WelchOptions}.
+ * @returns - {@link PSDResult} `{ f, Pxx }`.
+ *
+ * @example
+ * ```ts
+ * import { welch } from "tsb";
+ * const x = Array.from({ length: 512 }, (_, i) => Math.sin(2 * Math.PI * 60 * i / 512));
+ * const { f, Pxx } = welch(x, { fs: 512 });
+ * ```
+ */
+// biome-ignore lint/complexity/noExcessiveCognitiveComplexity: Welch averaging loop
+export function welch(x: readonly number[], options: WelchOptions = {}): PSDResult {
+ const fs = options.fs ?? 1;
+ const nperseg = options.nperseg ?? Math.min(256, x.length);
+ const noverlap = options.noverlap ?? Math.floor(nperseg / 2);
+ const nfft = nextPow2(options.nfft ?? nperseg);
+ const step = nperseg - noverlap;
+ const average = options.average ?? "mean";
+ const scaling = options.scaling ?? "density";
+
+ const win: number[] =
+ options.window !== undefined
+ ? typeof options.window === "string"
+ ? getWindow(options.window, nperseg)
+ : Array.from(options.window)
+ : hannWindow(nperseg);
+
+ const winNorm =
+ scaling === "density"
+ ? win.reduce((s, w) => s + w * w, 0) * fs
+ : win.reduce((s, w) => s + w, 0) ** 2;
+ const nFreqs = Math.floor(nfft / 2) + 1;
+ const nFrames = Math.floor((x.length - noverlap) / step);
+
+ if (nFrames <= 0) {
+ // Signal too short — return single periodogram
+ return periodogram(x, {
+ fs,
+ scaling,
+ ...(options.nfft !== undefined ? { nfft: options.nfft } : {}),
+ ...(options.window !== undefined ? { window: options.window } : {}),
+ });
+ }
+
+ // Collect per-frame periodograms
+ const frames: number[][] = [];
+ for (let k = 0; k < nFrames; k++) {
+ const start = k * step;
+ const seg: Complex[] = Array.from({ length: nfft }, (_, i) => ({
+ re: (x[start + i] ?? 0) * (win[i] ?? 0),
+ im: 0,
+ }));
+ fftInPlace(seg, nfft, false);
+ const pxx: number[] = Array.from({ length: nFreqs }, (_, fi) => {
+ const c = seg[fi] ?? { re: 0, im: 0 };
+ let p = cAbsSq(c) / winNorm;
+ // Double one-sided bins (except DC and Nyquist)
+ if (fi > 0 && fi < nFreqs - 1) {
+ p *= 2;
+ }
+ return p;
+ });
+ frames.push(pxx);
+ }
+
+ // Average
+ const Pxx: number[] = Array.from({ length: nFreqs }, (_, fi) => {
+ const vals = frames.map((fr) => fr[fi] ?? 0);
+ if (average === "mean") {
+ return vals.reduce((s, v) => s + v, 0) / vals.length;
+ }
+ // median
+ const sorted = [...vals].sort((a, b) => a - b);
+ const mid = Math.floor(sorted.length / 2);
+ return sorted.length % 2 === 1
+ ? (sorted[mid] ?? 0)
+ : ((sorted[mid - 1] ?? 0) + (sorted[mid] ?? 0)) / 2;
+ });
+
+ const f: number[] = Array.from({ length: nFreqs }, (_, i) => (i * fs) / nfft);
+ return { f, Pxx };
+}
+
+// ─── Periodogram ──────────────────────────────────────────────────────────────
+
+/** Options for {@link periodogram}. */
+export interface PeriodogramOptions {
+ /** Sampling frequency in Hz (default `1`). */
+ fs?: number;
+ /** FFT length (default `nextPow2(x.length)`). */
+ nfft?: number;
+ /** Window function (default `"hann"`). */
+ window?: WindowName | readonly number[];
+ /** Scaling: `"density"` (PSD) or `"spectrum"` (power spectrum). */
+ scaling?: "density" | "spectrum";
+}
+
+/**
+ * Estimate power spectral density via a single FFT (periodogram).
+ *
+ * Mirrors `scipy.signal.periodogram`.
+ *
+ * @param x - Input signal.
+ * @param options - {@link PeriodogramOptions}.
+ * @returns - {@link PSDResult} `{ f, Pxx }`.
+ *
+ * @example
+ * ```ts
+ * import { periodogram } from "tsb";
+ * const x = Array.from({ length: 256 }, (_, i) => Math.cos(2 * Math.PI * 20 * i / 256));
+ * const { f, Pxx } = periodogram(x, { fs: 256 });
+ * ```
+ */
+export function periodogram(x: readonly number[], options: PeriodogramOptions = {}): PSDResult {
+ const fs = options.fs ?? 1;
+ const nfft = nextPow2(options.nfft ?? x.length);
+ const scaling = options.scaling ?? "density";
+
+ const win: number[] =
+ options.window !== undefined
+ ? typeof options.window === "string"
+ ? getWindow(options.window, x.length)
+ : Array.from(options.window)
+ : hannWindow(x.length);
+
+ const winNorm =
+ scaling === "density"
+ ? win.reduce((s, w) => s + w * w, 0) * fs
+ : win.reduce((s, w) => s + w, 0) ** 2;
+
+ const seg: Complex[] = Array.from({ length: nfft }, (_, i) => ({
+ re: (x[i] ?? 0) * (win[i] ?? 0),
+ im: 0,
+ }));
+ fftInPlace(seg, nfft, false);
+
+ const nFreqs = Math.floor(nfft / 2) + 1;
+ const f: number[] = Array.from({ length: nFreqs }, (_, i) => (i * fs) / nfft);
+ const Pxx: number[] = Array.from({ length: nFreqs }, (_, fi) => {
+ const c = seg[fi] ?? { re: 0, im: 0 };
+ let p = cAbsSq(c) / winNorm;
+ if (fi > 0 && fi < nFreqs - 1) {
+ p *= 2;
+ }
+ return p;
+ });
+
+ return { f, Pxx };
+}
diff --git a/src/stats/spatial_stats.ts b/src/stats/spatial_stats.ts
new file mode 100644
index 000000000..715e03235
--- /dev/null
+++ b/src/stats/spatial_stats.ts
@@ -0,0 +1,438 @@
+/**
+ * spatial_stats — Spatial statistics and geostatistics.
+ *
+ * Implements:
+ * - **Variogram** estimation (empirical) and theoretical models (spherical, exponential, Gaussian)
+ * - **Ordinary Kriging** interpolation
+ * - **Moran's I** spatial autocorrelation
+ * - **Ripley's K function** and L function
+ * - **Kernel density estimation** (2D)
+ * - **Spatial weights** (distance-based, k-nearest-neighbor)
+ *
+ * @module
+ */
+
+// ─── Types ────────────────────────────────────────────────────────────────────
+
+/** A 2D point with optional value. */
+export interface SpatialPoint {
+ x: number;
+ y: number;
+ value?: number;
+}
+
+// ─── Distance Utilities ───────────────────────────────────────────────────────
+
+/**
+ * Euclidean distance between two points.
+ */
+export function euclideanDistance(
+ p1: { x: number; y: number },
+ p2: { x: number; y: number },
+): number {
+ return Math.sqrt((p1.x - p2.x) ** 2 + (p1.y - p2.y) ** 2);
+}
+
+/**
+ * Compute pairwise distance matrix.
+ *
+ * @param points - Array of spatial points.
+ * @returns n×n distance matrix (flat row-major array).
+ */
+export function pairwiseDistances(points: { x: number; y: number }[]): number[] {
+ const n = points.length;
+ const D = new Array(n * n).fill(0);
+
+ for (let i = 0; i < n; i++) {
+ for (let j = i + 1; j < n; j++) {
+ const d = euclideanDistance(points[i] ?? { x: 0, y: 0 }, points[j] ?? { x: 0, y: 0 });
+ D[i * n + j] = d;
+ D[j * n + i] = d;
+ }
+ }
+
+ return D;
+}
+
+// ─── Variogram ────────────────────────────────────────────────────────────────
+
+/** A lag-semivariance pair for variogram estimation. */
+export interface VariogramPoint {
+ lag: number;
+ semivariance: number;
+ count: number;
+}
+
+/**
+ * Compute the empirical (experimental) variogram.
+ *
+ * @param points - Spatial points with values.
+ * @param nLags - Number of lag bins. Default 10.
+ * @param maxLag - Maximum lag distance (auto if undefined).
+ * @returns Array of variogram points.
+ *
+ * @example
+ * ```ts
+ * import { empiricalVariogram } from "tsb";
+ * const pts = [{ x:0, y:0, value:1 }, { x:1, y:0, value:2 }, { x:2, y:0, value:1.5 }];
+ * const vgram = empiricalVariogram(pts, 5);
+ * ```
+ */
+export function empiricalVariogram(
+ points: SpatialPoint[],
+ nLags = 10,
+ maxLag?: number,
+): VariogramPoint[] {
+ const n = points.length;
+ const D = pairwiseDistances(points);
+
+ let maxD = maxLag ?? 0;
+ if (maxLag === undefined) {
+ for (let i = 0; i < n * n; i++) {
+ if ((D[i] ?? 0) > maxD) maxD = D[i] ?? 0;
+ }
+ maxD *= 0.5; // Rule of thumb: use half the max distance
+ }
+
+ const lagSize = maxD / nLags;
+ const sumSq = new Array(nLags).fill(0);
+ const counts = new Array(nLags).fill(0);
+
+ for (let i = 0; i < n; i++) {
+ for (let j = i + 1; j < n; j++) {
+ const d = D[i * n + j] ?? 0;
+ if (d > maxD || d === 0) continue;
+
+ const bin = Math.min(Math.floor(d / lagSize), nLags - 1);
+ const dv = ((points[i]?.value ?? 0) - (points[j]?.value ?? 0)) ** 2;
+ sumSq[bin] = (sumSq[bin] ?? 0) + dv;
+ counts[bin] = (counts[bin] ?? 0) + 1;
+ }
+ }
+
+ return Array.from({ length: nLags }, (_, k) => ({
+ lag: (k + 0.5) * lagSize,
+ semivariance: (counts[k] ?? 0) > 0 ? (sumSq[k] ?? 0) / (2 * (counts[k] ?? 1)) : 0,
+ count: counts[k] ?? 0,
+ }));
+}
+
+/** Parameters for theoretical variogram models. */
+export interface VariogramModelParams {
+ /** Nugget (c0). Default 0. */
+ nugget?: number;
+ /** Sill (c). */
+ sill: number;
+ /** Range (a). */
+ range: number;
+}
+
+/**
+ * Spherical variogram model γ(h).
+ */
+export function sphericalVariogram(h: number, params: VariogramModelParams): number {
+ const c0 = params.nugget ?? 0;
+ const c = params.sill;
+ const a = params.range;
+ if (h === 0) return 0;
+ if (h >= a) return c0 + c;
+ return c0 + c * (1.5 * (h / a) - 0.5 * (h / a) ** 3);
+}
+
+/**
+ * Exponential variogram model γ(h).
+ */
+export function exponentialVariogram(h: number, params: VariogramModelParams): number {
+ const c0 = params.nugget ?? 0;
+ const c = params.sill;
+ const a = params.range;
+ if (h === 0) return 0;
+ return c0 + c * (1 - Math.exp(-h / a));
+}
+
+/**
+ * Gaussian variogram model γ(h).
+ */
+export function gaussianVariogram(h: number, params: VariogramModelParams): number {
+ const c0 = params.nugget ?? 0;
+ const c = params.sill;
+ const a = params.range;
+ if (h === 0) return 0;
+ return c0 + c * (1 - Math.exp(-((h / a) ** 2)));
+}
+
+// ─── Ordinary Kriging ─────────────────────────────────────────────────────────
+
+type VariogramFn = (h: number, params: VariogramModelParams) => number;
+
+/**
+ * Ordinary Kriging interpolation at unsampled locations.
+ *
+ * @param knownPoints - Known data points with values.
+ * @param queryPoints - Locations to interpolate.
+ * @param variogramFn - Variogram model function.
+ * @param vParams - Variogram model parameters.
+ * @returns Interpolated values at query points.
+ *
+ * @example
+ * ```ts
+ * import { ordinaryKriging, sphericalVariogram } from "tsb";
+ * const known = [{ x:0, y:0, value:1 }, { x:1, y:1, value:2 }];
+ * const query = [{ x:0.5, y:0.5 }];
+ * const pred = ordinaryKriging(known, query, sphericalVariogram, { sill:1, range:2 });
+ * ```
+ */
+export function ordinaryKriging(
+ knownPoints: SpatialPoint[],
+ queryPoints: { x: number; y: number }[],
+ variogramFn: VariogramFn,
+ vParams: VariogramModelParams,
+): number[] {
+ const n = knownPoints.length;
+ // Build kriging matrix (n+1) x (n+1) with Lagrange multiplier
+ const size = n + 1;
+ const K = new Array(size * size).fill(0);
+
+ for (let i = 0; i < n; i++) {
+ for (let j = 0; j < n; j++) {
+ const d = euclideanDistance(
+ knownPoints[i] ?? { x: 0, y: 0 },
+ knownPoints[j] ?? { x: 0, y: 0 },
+ );
+ K[i * size + j] = variogramFn(d, vParams);
+ }
+ K[i * size + n] = 1;
+ K[n * size + i] = 1;
+ }
+ K[n * size + n] = 0;
+
+ const results: number[] = [];
+
+ for (const q of queryPoints) {
+ const k = new Array(size).fill(0);
+ for (let i = 0; i < n; i++) {
+ const d = euclideanDistance(q, knownPoints[i] ?? { x: 0, y: 0 });
+ k[i] = variogramFn(d, vParams);
+ }
+ k[n] = 1;
+
+ // Solve K * w = k via Gaussian elimination
+ const w = solveLinearSystem(K, k, size);
+ let pred = 0;
+ for (let i = 0; i < n; i++) {
+ pred += (w[i] ?? 0) * (knownPoints[i]?.value ?? 0);
+ }
+ results.push(pred);
+ }
+
+ return results;
+}
+
+// ─── Moran's I ────────────────────────────────────────────────────────────────
+
+/**
+ * Compute Moran's I spatial autocorrelation statistic.
+ *
+ * @param values - Observed values at each location.
+ * @param weights - n×n spatial weight matrix (flat row-major).
+ * @returns Moran's I statistic.
+ *
+ * @example
+ * ```ts
+ * import { moransI } from "tsb";
+ * const vals = [1, 2, 3, 4];
+ * const W = [0,1,0,0, 1,0,1,0, 0,1,0,1, 0,0,1,0];
+ * const I = moransI(vals, W);
+ * ```
+ */
+export function moransI(values: number[], weights: number[]): number {
+ const n = values.length;
+ const mean = values.reduce((s, v) => s + v, 0) / n;
+ const deviations = values.map((v) => v - mean);
+
+ let numerator = 0;
+ let W = 0;
+
+ for (let i = 0; i < n; i++) {
+ for (let j = 0; j < n; j++) {
+ const wij = weights[i * n + j] ?? 0;
+ W += wij;
+ numerator += wij * (deviations[i] ?? 0) * (deviations[j] ?? 0);
+ }
+ }
+
+ const denominator = deviations.reduce((s, d) => s + d * d, 0);
+
+ if (denominator === 0 || W === 0) return 0;
+ return (n / W) * (numerator / denominator);
+}
+
+// ─── Ripley's K Function ──────────────────────────────────────────────────────
+
+/**
+ * Estimate Ripley's K function for a point pattern.
+ *
+ * @param points - Array of 2D points.
+ * @param distances - Array of distances at which to evaluate K.
+ * @param area - Area of the study region.
+ * @returns K(d) values.
+ */
+export function ripleysK(
+ points: { x: number; y: number }[],
+ distances: number[],
+ area: number,
+): number[] {
+ const n = points.length;
+ const lambda = n / area;
+ const D = pairwiseDistances(points);
+
+ return distances.map((r) => {
+ let count = 0;
+ for (let i = 0; i < n; i++) {
+ for (let j = 0; j < n; j++) {
+ if (i !== j && (D[i * n + j] ?? Number.POSITIVE_INFINITY) <= r) count++;
+ }
+ }
+ return count / (n * lambda);
+ });
+}
+
+/**
+ * Compute Ripley's L function: L(r) = sqrt(K(r) / pi).
+ *
+ * @param K - K function values.
+ * @returns L function values.
+ */
+export function ripleysL(K: number[]): number[] {
+ return K.map((k) => Math.sqrt(k / Math.PI));
+}
+
+// ─── 2D Kernel Density Estimation ─────────────────────────────────────────────
+
+/**
+ * Compute 2D kernel density estimate on a grid.
+ *
+ * @param points - Data points.
+ * @param xGrid - x-grid values.
+ * @param yGrid - y-grid values.
+ * @param bandwidth - Bandwidth (h). Default 1.
+ * @returns Density values on the grid (flat row-major, len = xGrid.length × yGrid.length).
+ */
+export function kde2d(
+ points: { x: number; y: number }[],
+ xGrid: number[],
+ yGrid: number[],
+ bandwidth = 1,
+): number[] {
+ const nx = xGrid.length;
+ const ny = yGrid.length;
+ const n = points.length;
+ const density = new Array(nx * ny).fill(0);
+ const h2 = bandwidth * bandwidth * 2;
+
+ for (let i = 0; i < nx; i++) {
+ for (let j = 0; j < ny; j++) {
+ let sum = 0;
+ const xi = xGrid[i] ?? 0;
+ const yj = yGrid[j] ?? 0;
+ for (const p of points) {
+ const d2 = (p.x - xi) ** 2 + (p.y - yj) ** 2;
+ sum += Math.exp(-d2 / h2);
+ }
+ density[i * ny + j] = sum / (n * Math.PI * h2);
+ }
+ }
+
+ return density;
+}
+
+// ─── Spatial Weights ──────────────────────────────────────────────────────────
+
+/**
+ * Build a distance-based spatial weight matrix (inverse distance weighting).
+ *
+ * @param points - Spatial points.
+ * @param maxDist - Maximum distance for neighbors (Number.POSITIVE_INFINITY = all pairs).
+ * @param power - Distance decay power. Default 1.
+ * @returns Row-standardized weight matrix (flat row-major).
+ */
+export function distanceWeights(
+ points: { x: number; y: number }[],
+ maxDist = Number.POSITIVE_INFINITY,
+ power = 1,
+): number[] {
+ const n = points.length;
+ const D = pairwiseDistances(points);
+ const W = new Array(n * n).fill(0);
+
+ for (let i = 0; i < n; i++) {
+ let rowSum = 0;
+ for (let j = 0; j < n; j++) {
+ if (i === j) continue;
+ const d = D[i * n + j] ?? Number.POSITIVE_INFINITY;
+ if (d <= maxDist && d > 0) {
+ W[i * n + j] = 1 / d ** power;
+ rowSum += W[i * n + j] ?? 0;
+ }
+ }
+ if (rowSum > 0) {
+ for (let j = 0; j < n; j++) {
+ W[i * n + j] = (W[i * n + j] ?? 0) / rowSum;
+ }
+ }
+ }
+
+ return W;
+}
+
+// ─── Linear System Solver (Gaussian Elimination) ─────────────────────────────
+
+function solveLinearSystem(A: number[], b: number[], n: number): number[] {
+ // Augmented matrix
+ const aug: number[][] = [];
+ for (let i = 0; i < n; i++) {
+ const row: number[] = [];
+ for (let j = 0; j < n; j++) {
+ row.push(A[i * n + j] ?? 0);
+ }
+ row.push(b[i] ?? 0);
+ aug.push(row);
+ }
+
+ // Forward elimination
+ for (let col = 0; col < n; col++) {
+ let maxRow = col;
+ let maxVal = Math.abs(aug[col]?.[col] ?? 0);
+ for (let row = col + 1; row < n; row++) {
+ const val = Math.abs(aug[row]?.[col] ?? 0);
+ if (val > maxVal) {
+ maxVal = val;
+ maxRow = row;
+ }
+ }
+ [aug[col], aug[maxRow]] = [aug[maxRow] ?? [], aug[col] ?? []];
+
+ const pivot = aug[col]?.[col] ?? 0;
+ if (Math.abs(pivot) < 1e-12) continue;
+
+ for (let row = col + 1; row < n; row++) {
+ const factor = (aug[row]?.[col] ?? 0) / pivot;
+ for (let j = col; j <= n; j++) {
+ (aug[row] ?? [])[j] = ((aug[row] ?? [])[j] ?? 0) - factor * ((aug[col] ?? [])[j] ?? 0);
+ }
+ }
+ }
+
+ // Back substitution
+ const x = new Array(n).fill(0);
+ for (let i = n - 1; i >= 0; i--) {
+ let sum = aug[i]?.[n] ?? 0;
+ for (let j = i + 1; j < n; j++) {
+ sum -= (aug[i]?.[j] ?? 0) * (x[j] ?? 0);
+ }
+ const pivot = aug[i]?.[i] ?? 0;
+ x[i] = Math.abs(pivot) < 1e-12 ? 0 : sum / pivot;
+ }
+
+ return x;
+}
diff --git a/src/stats/stochastic_processes.ts b/src/stats/stochastic_processes.ts
new file mode 100644
index 000000000..43df0b090
--- /dev/null
+++ b/src/stats/stochastic_processes.ts
@@ -0,0 +1,381 @@
+/**
+ * stochastic_processes — Stochastic process simulation and inference.
+ *
+ * Implements:
+ * - **Brownian Motion** (Wiener process, geometric Brownian motion)
+ * - **Ornstein-Uhlenbeck** mean-reverting process
+ * - **Poisson Process** (homogeneous and inhomogeneous)
+ * - **Random Walk** (simple, correlated, Lévy)
+ * - **Markov Chain** (discrete-time)
+ *
+ * @module
+ */
+
+// ─── Brownian Motion ──────────────────────────────────────────────────────────
+
+/** Parameters for Brownian Motion simulation. */
+export interface BrownianMotionParams {
+ /** Drift coefficient (mu). Default 0. */
+ mu?: number;
+ /** Diffusion coefficient (sigma). Default 1. */
+ sigma?: number;
+ /** Initial value. Default 0. */
+ x0?: number;
+ /** Time step. Default 0.01. */
+ dt?: number;
+ /** Random seed (not used for true randomness, but for API compat). */
+ seed?: number;
+}
+
+/** Result of a stochastic process simulation. */
+export interface ProcessPath {
+ /** Time points. */
+ times: number[];
+ /** Process values at each time point. */
+ values: number[];
+}
+
+/**
+ * Simulate standard Brownian Motion (Wiener process with drift).
+ *
+ * dX = mu*dt + sigma*dW
+ *
+ * @param nSteps - Number of time steps.
+ * @param params - Optional BrownianMotionParams.
+ * @returns ProcessPath with times and values.
+ *
+ * @example
+ * ```ts
+ * import { simulateBrownianMotion } from "tsb";
+ * const path = simulateBrownianMotion(100, { mu: 0.1, sigma: 0.2 });
+ * ```
+ */
+export function simulateBrownianMotion(
+ nSteps: number,
+ params: BrownianMotionParams = {},
+): ProcessPath {
+ const mu = params.mu ?? 0;
+ const sigma = params.sigma ?? 1;
+ const x0 = params.x0 ?? 0;
+ const dt = params.dt ?? 0.01;
+
+ const times: number[] = [0];
+ const values: number[] = [x0];
+
+ let x = x0;
+ for (let i = 1; i <= nSteps; i++) {
+ const dW = randn() * Math.sqrt(dt);
+ x = x + mu * dt + sigma * dW;
+ times.push(i * dt);
+ values.push(x);
+ }
+
+ return { times, values };
+}
+
+/**
+ * Simulate Geometric Brownian Motion (log-normal process).
+ *
+ * dS = mu*S*dt + sigma*S*dW
+ *
+ * @param nSteps - Number of time steps.
+ * @param params - Optional BrownianMotionParams.
+ * @returns ProcessPath with times and values.
+ */
+export function simulateGeometricBrownianMotion(
+ nSteps: number,
+ params: BrownianMotionParams = {},
+): ProcessPath {
+ const mu = params.mu ?? 0.1;
+ const sigma = params.sigma ?? 0.2;
+ const x0 = params.x0 ?? 100;
+ const dt = params.dt ?? 1 / 252;
+
+ const times: number[] = [0];
+ const values: number[] = [x0];
+
+ let s = x0;
+ for (let i = 1; i <= nSteps; i++) {
+ const dW = randn() * Math.sqrt(dt);
+ s = s * Math.exp((mu - 0.5 * sigma * sigma) * dt + sigma * dW);
+ times.push(i * dt);
+ values.push(s);
+ }
+
+ return { times, values };
+}
+
+// ─── Ornstein-Uhlenbeck Process ───────────────────────────────────────────────
+
+/** Parameters for the Ornstein-Uhlenbeck process. */
+export interface OUParams {
+ /** Mean reversion speed (theta). Default 1. */
+ theta?: number;
+ /** Long-run mean (mu). Default 0. */
+ mu?: number;
+ /** Volatility (sigma). Default 0.1. */
+ sigma?: number;
+ /** Initial value. Default mu. */
+ x0?: number;
+ /** Time step. Default 0.01. */
+ dt?: number;
+}
+
+/**
+ * Simulate Ornstein-Uhlenbeck mean-reverting process.
+ *
+ * dX = theta*(mu - X)*dt + sigma*dW
+ *
+ * @param nSteps - Number of time steps.
+ * @param params - Optional OUParams.
+ * @returns ProcessPath.
+ */
+export function simulateOrnsteinUhlenbeck(nSteps: number, params: OUParams = {}): ProcessPath {
+ const theta = params.theta ?? 1.0;
+ const mu = params.mu ?? 0.0;
+ const sigma = params.sigma ?? 0.1;
+ const dt = params.dt ?? 0.01;
+ const x0 = params.x0 ?? mu;
+
+ const times: number[] = [0];
+ const values: number[] = [x0];
+
+ let x = x0;
+ for (let i = 1; i <= nSteps; i++) {
+ const dW = randn() * Math.sqrt(dt);
+ x = x + theta * (mu - x) * dt + sigma * dW;
+ times.push(i * dt);
+ values.push(x);
+ }
+
+ return { times, values };
+}
+
+/**
+ * Estimate OU parameters via method of moments.
+ *
+ * @param values - Observed time series values.
+ * @param dt - Time step between observations.
+ * @returns Estimated theta, mu, sigma.
+ */
+export function fitOrnsteinUhlenbeck(
+ values: number[],
+ dt = 0.01,
+): { theta: number; mu: number; sigma: number } {
+ const n = values.length;
+ if (n < 3) return { theta: 0, mu: 0, sigma: 0 };
+
+ // Method of moments via OLS on X(t+dt) = a + b*X(t)
+ let sumX = 0;
+ let sumY = 0;
+ let sumXX = 0;
+ let sumXY = 0;
+ const m = n - 1;
+
+ for (let i = 0; i < m; i++) {
+ const xi = values[i] ?? 0;
+ const yi = values[i + 1] ?? 0;
+ sumX += xi;
+ sumY += yi;
+ sumXX += xi * xi;
+ sumXY += xi * yi;
+ }
+
+ const b = (m * sumXY - sumX * sumY) / (m * sumXX - sumX * sumX);
+ const a = (sumY - b * sumX) / m;
+
+ const theta = -Math.log(Math.max(b, 1e-10)) / dt;
+ const mu = a / (1 - b);
+
+ // Estimate sigma from residuals
+ let residSS = 0;
+ for (let i = 0; i < m; i++) {
+ const xi = values[i] ?? 0;
+ const yi = values[i + 1] ?? 0;
+ const pred = a + b * xi;
+ residSS += (yi - pred) ** 2;
+ }
+ const sigmaEst = Math.sqrt(residSS / m / dt);
+
+ return { theta, mu, sigma: sigmaEst };
+}
+
+// ─── Poisson Process ──────────────────────────────────────────────────────────
+
+/**
+ * Simulate a homogeneous Poisson process.
+ *
+ * @param rate - Event rate (lambda, events per unit time).
+ * @param T - Total time horizon.
+ * @returns Array of event arrival times.
+ *
+ * @example
+ * ```ts
+ * import { simulatePoissonProcess } from "tsb";
+ * const arrivals = simulatePoissonProcess(2.5, 10);
+ * ```
+ */
+export function simulatePoissonProcess(rate: number, T: number): number[] {
+ const arrivals: number[] = [];
+ let t = 0;
+
+ while (t < T) {
+ const inter = -Math.log(1 - Math.random()) / rate;
+ t += inter;
+ if (t < T) arrivals.push(t);
+ }
+
+ return arrivals;
+}
+
+/**
+ * Count events in bins for a Poisson process.
+ *
+ * @param arrivals - Event arrival times.
+ * @param binSize - Size of each bin.
+ * @param T - Total time horizon.
+ * @returns Array of event counts per bin.
+ */
+export function poissonCounts(arrivals: number[], binSize: number, T: number): number[] {
+ const nBins = Math.ceil(T / binSize);
+ const counts = new Array(nBins).fill(0);
+
+ for (const t of arrivals) {
+ const bin = Math.floor(t / binSize);
+ if (bin < nBins) {
+ counts[bin] = (counts[bin] ?? 0) + 1;
+ }
+ }
+
+ return counts;
+}
+
+// ─── Random Walk ──────────────────────────────────────────────────────────────
+
+/** Parameters for random walk simulation. */
+export interface RandomWalkParams {
+ /** Step probabilities [up, down]. Default [0.5, 0.5]. */
+ probs?: [number, number];
+ /** Step sizes [up, down]. Default [1, -1]. */
+ steps?: [number, number];
+ /** Initial position. Default 0. */
+ x0?: number;
+}
+
+/**
+ * Simulate a discrete random walk.
+ *
+ * @param nSteps - Number of steps.
+ * @param params - Optional RandomWalkParams.
+ * @returns ProcessPath with integer times and cumulative positions.
+ */
+export function simulateRandomWalk(nSteps: number, params: RandomWalkParams = {}): ProcessPath {
+ const probs = params.probs ?? [0.5, 0.5];
+ const steps = params.steps ?? [1, -1];
+ const x0 = params.x0 ?? 0;
+
+ const times: number[] = [0];
+ const values: number[] = [x0];
+
+ let x = x0;
+ for (let i = 1; i <= nSteps; i++) {
+ const r = Math.random();
+ const step = r < (probs[0] ?? 0.5) ? (steps[0] ?? 1) : (steps[1] ?? -1);
+ x += step;
+ times.push(i);
+ values.push(x);
+ }
+
+ return { times, values };
+}
+
+// ─── Markov Chain ─────────────────────────────────────────────────────────────
+
+/**
+ * Simulate a discrete-time Markov chain.
+ *
+ * @param transitionMatrix - Row-stochastic transition matrix (nStates × nStates).
+ * @param nSteps - Number of steps.
+ * @param initialState - Starting state index. Default 0.
+ * @returns Array of state indices visited.
+ *
+ * @example
+ * ```ts
+ * import { simulateMarkovChain } from "tsb";
+ * const T = [[0.7, 0.3], [0.4, 0.6]];
+ * const chain = simulateMarkovChain(T, 100);
+ * ```
+ */
+export function simulateMarkovChain(
+ transitionMatrix: number[][],
+ nSteps: number,
+ initialState = 0,
+): number[] {
+ const nStates = transitionMatrix.length;
+ const chain: number[] = [initialState];
+ let state = initialState;
+
+ for (let i = 1; i <= nSteps; i++) {
+ const row = transitionMatrix[state] ?? [];
+ state = sampleCategorical(row, nStates);
+ chain.push(state);
+ }
+
+ return chain;
+}
+
+/**
+ * Compute the stationary distribution of a Markov chain via power iteration.
+ *
+ * @param transitionMatrix - Row-stochastic transition matrix.
+ * @param maxIter - Maximum iterations. Default 1000.
+ * @param tol - Convergence tolerance. Default 1e-10.
+ * @returns Stationary distribution vector.
+ */
+export function stationaryDistribution(
+ transitionMatrix: number[][],
+ maxIter = 1000,
+ tol = 1e-10,
+): number[] {
+ const n = transitionMatrix.length;
+ let pi = new Array(n).fill(1 / n);
+
+ for (let iter = 0; iter < maxIter; iter++) {
+ const newPi = new Array(n).fill(0);
+ for (let j = 0; j < n; j++) {
+ for (let i = 0; i < n; i++) {
+ newPi[j] = (newPi[j] ?? 0) + (pi[i] ?? 0) * ((transitionMatrix[i] ?? [])[j] ?? 0);
+ }
+ }
+ let diff = 0;
+ for (let i = 0; i < n; i++) {
+ diff += Math.abs((newPi[i] ?? 0) - (pi[i] ?? 0));
+ }
+ pi = newPi;
+ if (diff < tol) break;
+ }
+
+ return pi;
+}
+
+// ─── Utilities ────────────────────────────────────────────────────────────────
+
+/** Box-Muller transform for standard normal random variate. */
+function randn(): number {
+ let u = 0;
+ let v = 0;
+ while (u === 0) u = Math.random();
+ while (v === 0) v = Math.random();
+ return Math.sqrt(-2.0 * Math.log(u)) * Math.cos(2.0 * Math.PI * v);
+}
+
+/** Sample from a categorical distribution given unnormalized probabilities. */
+function sampleCategorical(probs: number[], n: number): number {
+ const total = probs.reduce((a, b) => a + b, 0);
+ let r = Math.random() * total;
+ for (let i = 0; i < n; i++) {
+ r -= probs[i] ?? 0;
+ if (r <= 0) return i;
+ }
+ return n - 1;
+}
diff --git a/src/tseries/frequencies.ts b/src/tseries/frequencies.ts
index f9c96aef3..2267b93c6 100644
--- a/src/tseries/frequencies.ts
+++ b/src/tseries/frequencies.ts
@@ -324,8 +324,12 @@ export function inferFreq(dates: readonly Date[]): string | null {
const days = first / MS_DAY;
// Year-begin/end dates may have equal diffs when no leap year falls in
// the range; check before returning a raw day-count alias.
- if (_allYearBegin(dates)) return "YS";
- if (_allYearEnd(dates)) return "YE";
+ if (_allYearBegin(dates)) {
+ return "YS";
+ }
+ if (_allYearEnd(dates)) {
+ return "YE";
+ }
return `${days}D`;
}
}
diff --git a/tests/io/orc.test.ts b/tests/io/orc.test.ts
new file mode 100644
index 000000000..4293bac90
--- /dev/null
+++ b/tests/io/orc.test.ts
@@ -0,0 +1,407 @@
+/**
+ * Tests for readOrc / toOrc — Apache ORC file format I/O.
+ *
+ * Strategy: use toOrc to produce ORC buffers, then readOrc to round-trip.
+ * All tests operate on in-memory buffers; no filesystem I/O is required.
+ */
+
+import { describe, expect, it } from "bun:test";
+import * as fc from "fast-check";
+import { DataFrame } from "../../src/core/frame.ts";
+import { readOrc, toOrc } from "../../src/io/orc.ts";
+
+// ─── Helpers ──────────────────────────────────────────────────────────────────
+
+function roundtrip(df: DataFrame): DataFrame {
+ return readOrc(toOrc(df));
+}
+
+function colArr(df: DataFrame, name: string): readonly unknown[] {
+ return df.col(name).values;
+}
+
+// ─── File structure ───────────────────────────────────────────────────────────
+
+describe("toOrc — file structure", () => {
+ it("returns a non-empty Uint8Array", () => {
+ const df = DataFrame.fromColumns({ x: [1, 2, 3] });
+ const buf = toOrc(df);
+ expect(buf).toBeInstanceOf(Uint8Array);
+ expect(buf.length).toBeGreaterThan(0);
+ });
+
+ it("starts with ORC magic bytes", () => {
+ const df = DataFrame.fromColumns({ x: [1] });
+ const buf = toOrc(df);
+ expect(buf[0]).toBe(0x4f); // 'O'
+ expect(buf[1]).toBe(0x52); // 'R'
+ expect(buf[2]).toBe(0x43); // 'C'
+ });
+
+ it("ends with postscript length byte", () => {
+ const df = DataFrame.fromColumns({ x: [1] });
+ const buf = toOrc(df);
+ // Last byte is postscript length, must be > 0
+ expect(buf.at(-1)).toBeGreaterThan(0);
+ });
+});
+
+// ─── Integer columns ──────────────────────────────────────────────────────────
+
+describe("readOrc / toOrc — integer columns", () => {
+ it("round-trips a simple int column", () => {
+ const df = DataFrame.fromColumns({ n: [1, 2, 3, 4, 5] });
+ const rt = roundtrip(df);
+ expect(rt.columns.toArray()).toEqual(["n"]);
+ expect(colArr(rt, "n")).toEqual([1, 2, 3, 4, 5]);
+ });
+
+ it("round-trips negative integers", () => {
+ const df = DataFrame.fromColumns({ n: [-100, -1, 0, 1, 100] });
+ const rt = roundtrip(df);
+ expect(colArr(rt, "n")).toEqual([-100, -1, 0, 1, 100]);
+ });
+
+ it("round-trips large integers", () => {
+ const df = DataFrame.fromColumns({ n: [1_000_000, 2_000_000, -999_999] });
+ const rt = roundtrip(df);
+ expect(colArr(rt, "n")).toEqual([1_000_000, 2_000_000, -999_999]);
+ });
+
+ it("round-trips a column of zeros", () => {
+ const df = DataFrame.fromColumns({ n: [0, 0, 0, 0] });
+ const rt = roundtrip(df);
+ expect(colArr(rt, "n")).toEqual([0, 0, 0, 0]);
+ });
+
+ it("round-trips null integers", () => {
+ const df = DataFrame.fromColumns({ n: [1, null, 3, null, 5] });
+ const rt = roundtrip(df);
+ const vals = colArr(rt, "n");
+ expect(vals[0]).toBe(1);
+ expect(vals[1]).toBeNull();
+ expect(vals[2]).toBe(3);
+ expect(vals[3]).toBeNull();
+ expect(vals[4]).toBe(5);
+ });
+
+ it("round-trips all-null int column", () => {
+ const df = DataFrame.fromColumns({ n: [null, null, null] });
+ const rt = roundtrip(df);
+ expect(colArr(rt, "n").every((v) => v === null)).toBe(true);
+ });
+});
+
+// ─── Float/Double columns ─────────────────────────────────────────────────────
+
+describe("readOrc / toOrc — float columns", () => {
+ it("round-trips double values", () => {
+ const df = DataFrame.fromColumns({ x: [1.5, 2.25, 3.75] });
+ const rt = roundtrip(df);
+ const vals = colArr(rt, "x") as number[];
+ expect(vals[0]).toBeCloseTo(1.5);
+ expect(vals[1]).toBeCloseTo(2.25);
+ expect(vals[2]).toBeCloseTo(3.75);
+ });
+
+ it("retains fractional values after an integer first value", () => {
+ expect(colArr(roundtrip(DataFrame.fromColumns({ x: [1, 2.5, null, 3.25] })), "x")).toEqual([
+ 1,
+ 2.5,
+ null,
+ 3.25,
+ ]);
+ });
+
+ it("stores values outside the safe integer range as doubles", () => {
+ const values = [18446744073709552000, -1e100, Number.MAX_VALUE];
+ expect(colArr(roundtrip(DataFrame.fromColumns({ x: values })), "x")).toEqual(values);
+ });
+
+ it("round-trips negative floats", () => {
+ const df = DataFrame.fromColumns({ x: [-1.5, -0.001, 0.0] });
+ const rt = roundtrip(df);
+ const vals = colArr(rt, "x") as number[];
+ expect(vals[0]).toBeCloseTo(-1.5);
+ expect(vals[1]).toBeCloseTo(-0.001);
+ expect(vals[2]).toBeCloseTo(0.0);
+ });
+
+ it("round-trips null floats", () => {
+ const df = DataFrame.fromColumns({ x: [1.0, null, 3.0] });
+ const rt = roundtrip(df);
+ const vals = colArr(rt, "x");
+ expect(vals[0]).toBe(1.0);
+ expect(vals[1]).toBeNull();
+ expect(vals[2]).toBe(3.0);
+ });
+});
+
+// ─── String columns ───────────────────────────────────────────────────────────
+
+describe("readOrc / toOrc — string columns", () => {
+ it("round-trips a string column", () => {
+ const df = DataFrame.fromColumns({ s: ["alpha", "beta", "gamma"] });
+ const rt = roundtrip(df);
+ expect(colArr(rt, "s")).toEqual(["alpha", "beta", "gamma"]);
+ });
+
+ it("round-trips empty strings", () => {
+ const df = DataFrame.fromColumns({ s: ["", "a", ""] });
+ const rt = roundtrip(df);
+ expect(colArr(rt, "s")).toEqual(["", "a", ""]);
+ });
+
+ it("round-trips unicode strings", () => {
+ const df = DataFrame.fromColumns({ s: ["こんにちは", "héllo", "🎉"] });
+ const rt = roundtrip(df);
+ expect(colArr(rt, "s")).toEqual(["こんにちは", "héllo", "🎉"]);
+ });
+
+ it("round-trips null strings", () => {
+ const df = DataFrame.fromColumns({ s: ["hello", null, "world"] });
+ const rt = roundtrip(df);
+ const vals = colArr(rt, "s");
+ expect(vals[0]).toBe("hello");
+ expect(vals[1]).toBeNull();
+ expect(vals[2]).toBe("world");
+ });
+});
+
+// ─── Boolean columns ──────────────────────────────────────────────────────────
+
+describe("readOrc / toOrc — boolean columns", () => {
+ it("round-trips boolean values", () => {
+ const df = DataFrame.fromColumns({ b: [true, false, true, false] });
+ const rt = roundtrip(df);
+ expect(colArr(rt, "b")).toEqual([true, false, true, false]);
+ });
+
+ it("round-trips all-true column", () => {
+ const df = DataFrame.fromColumns({ b: [true, true, true] });
+ const rt = roundtrip(df);
+ expect(colArr(rt, "b")).toEqual([true, true, true]);
+ });
+
+ it("round-trips null booleans", () => {
+ const df = DataFrame.fromColumns({ b: [true, null, false] });
+ const rt = roundtrip(df);
+ const vals = colArr(rt, "b");
+ expect(vals[0]).toBe(true);
+ expect(vals[1]).toBeNull();
+ expect(vals[2]).toBe(false);
+ });
+});
+
+// ─── Multi-column DataFrames ──────────────────────────────────────────────────
+
+describe("readOrc / toOrc — multi-column DataFrames", () => {
+ it("round-trips mixed-type DataFrame", () => {
+ const df = DataFrame.fromColumns({
+ id: [1, 2, 3],
+ name: ["Alice", "Bob", "Carol"],
+ score: [95.5, 87.0, 92.3],
+ passed: [true, false, true],
+ });
+ const rt = roundtrip(df);
+ expect(rt.columns.toArray()).toEqual(["id", "name", "score", "passed"]);
+ expect(colArr(rt, "id")).toEqual([1, 2, 3]);
+ expect(colArr(rt, "name")).toEqual(["Alice", "Bob", "Carol"]);
+ expect((colArr(rt, "score") as number[]).map((v) => Math.round(v * 10) / 10)).toEqual([
+ 95.5, 87.0, 92.3,
+ ]);
+ expect(colArr(rt, "passed")).toEqual([true, false, true]);
+ });
+
+ it("preserves column order", () => {
+ const df = DataFrame.fromColumns({ z: [1], a: [2], m: [3] });
+ const rt = roundtrip(df);
+ expect(rt.columns.toArray()).toEqual(["z", "a", "m"]);
+ });
+
+ it("handles 1-row DataFrame", () => {
+ const df = DataFrame.fromColumns({ x: [42], y: ["hi"] });
+ const rt = roundtrip(df);
+ expect(colArr(rt, "x")).toEqual([42]);
+ expect(colArr(rt, "y")).toEqual(["hi"]);
+ });
+});
+
+// ─── Empty DataFrame ──────────────────────────────────────────────────────────
+
+describe("readOrc / toOrc — empty DataFrame", () => {
+ it("round-trips an empty DataFrame", () => {
+ const df = DataFrame.fromColumns({ x: [] as number[] });
+ const rt = roundtrip(df);
+ expect(rt.shape[0]).toBe(0);
+ expect(rt.columns.toArray()).toEqual(["x"]);
+ });
+});
+
+// ─── Options: columns filter ──────────────────────────────────────────────────
+
+describe("readOrc — columns option", () => {
+ it("reads only specified columns", () => {
+ const df = DataFrame.fromColumns({ a: [1, 2], b: ["x", "y"], c: [true, false] });
+ const buf = toOrc(df);
+ const rt = readOrc(buf, { columns: ["a", "c"] });
+ expect(rt.columns.toArray()).toEqual(["a", "c"]);
+ expect(colArr(rt, "a")).toEqual([1, 2]);
+ expect(colArr(rt, "c")).toEqual([true, false]);
+ });
+});
+
+// ─── Options: writeIndex ──────────────────────────────────────────────────────
+
+describe("toOrc — writeIndex option", () => {
+ it("includes index when writeIndex=true", () => {
+ const df = DataFrame.fromColumns({ x: [10, 20, 30] });
+ const buf = toOrc(df, { writeIndex: true });
+ const rt = readOrc(buf);
+ // __index__ column should be present
+ expect(rt.columns.toArray()).toContain("__index__");
+ expect(rt.columns.toArray()).toContain("x");
+ });
+});
+
+// ─── Error handling ───────────────────────────────────────────────────────────
+
+describe("readOrc — error handling", () => {
+ it("throws on invalid magic bytes", () => {
+ const bad = new Uint8Array([0x50, 0x41, 0x52, 0x31, 0x00]);
+ expect(() => readOrc(bad)).toThrow(/magic/i);
+ });
+
+ it("throws on too-small file", () => {
+ const bad = new Uint8Array([0x4f, 0x52]);
+ expect(() => readOrc(bad)).toThrow();
+ });
+
+ it("accepts ArrayBuffer input", () => {
+ const df = DataFrame.fromColumns({ n: [1, 2] });
+ const buf = toOrc(df);
+ const ab = buf.buffer.slice(buf.byteOffset, buf.byteOffset + buf.byteLength);
+ const rt = readOrc(new Uint8Array(ab));
+ expect(colArr(rt, "n")).toEqual([1, 2]);
+ });
+});
+
+// ─── Large dataset ────────────────────────────────────────────────────────────
+
+describe("readOrc / toOrc — large dataset", () => {
+ it("round-trips 1000-row integer column", () => {
+ const data = Array.from({ length: 1000 }, (_, i) => i);
+ const df = DataFrame.fromColumns({ n: data });
+ const rt = roundtrip(df);
+ const vals = colArr(rt, "n") as number[];
+ expect(vals.length).toBe(1000);
+ for (let i = 0; i < 1000; i++) {
+ expect(vals[i]).toBe(i);
+ }
+ });
+
+ it("round-trips 500-row string column", () => {
+ const data = Array.from({ length: 500 }, (_, i) => `row_${i}`);
+ const df = DataFrame.fromColumns({ s: data });
+ const rt = roundtrip(df);
+ const vals = colArr(rt, "s") as string[];
+ expect(vals.length).toBe(500);
+ for (let i = 0; i < 500; i++) {
+ expect(vals[i]).toBe(`row_${i}`);
+ }
+ });
+});
+
+// ─── Property-based tests ─────────────────────────────────────────────────────
+
+describe("readOrc / toOrc — property tests", () => {
+ it("integer round-trip: arbitrary int arrays", () => {
+ fc.assert(
+ fc.property(
+ fc.array(fc.integer({ min: -1_000_000, max: 1_000_000 }), { minLength: 1, maxLength: 100 }),
+ (data) => {
+ const df = DataFrame.fromColumns({ n: data });
+ const rt = roundtrip(df);
+ const vals = colArr(rt, "n") as number[];
+ for (let i = 0; i < data.length; i++) {
+ expect(vals[i]).toBe(data[i]);
+ }
+ },
+ ),
+ );
+ });
+
+ it("string round-trip: arbitrary string arrays", () => {
+ fc.assert(
+ fc.property(
+ fc.array(fc.string({ maxLength: 20 }), { minLength: 1, maxLength: 50 }),
+ (data) => {
+ const df = DataFrame.fromColumns({ s: data });
+ const rt = roundtrip(df);
+ const vals = colArr(rt, "s") as string[];
+ for (let i = 0; i < data.length; i++) {
+ expect(vals[i]).toBe(data[i]);
+ }
+ },
+ ),
+ );
+ });
+
+ it("float round-trip: finite float64 values", () => {
+ fc.assert(
+ fc.property(
+ fc.array(fc.float({ noNaN: true, noDefaultInfinity: true }), {
+ minLength: 1,
+ maxLength: 50,
+ }),
+ (data) => {
+ const df = DataFrame.fromColumns({ x: data });
+ const rt = roundtrip(df);
+ const vals = colArr(rt, "x") as number[];
+ for (let i = 0; i < data.length; i++) {
+ const expected = data[i] ?? 0;
+ const actual = vals[i] ?? 0;
+ // Float64 round-trip should be exact
+ expect(actual).toBeCloseTo(expected, 10);
+ }
+ },
+ ),
+ );
+ });
+
+ it("boolean round-trip: arbitrary boolean arrays", () => {
+ fc.assert(
+ fc.property(fc.array(fc.boolean(), { minLength: 1, maxLength: 100 }), (data) => {
+ const df = DataFrame.fromColumns({ b: data });
+ const rt = roundtrip(df);
+ const vals = colArr(rt, "b") as boolean[];
+ for (let i = 0; i < data.length; i++) {
+ expect(vals[i]).toBe(data[i]);
+ }
+ }),
+ );
+ });
+
+ it("nullable integer round-trip", () => {
+ fc.assert(
+ fc.property(
+ fc.array(fc.option(fc.integer({ min: -1000, max: 1000 }), { nil: null }), {
+ minLength: 1,
+ maxLength: 50,
+ }),
+ (data) => {
+ const df = DataFrame.fromColumns({ n: data });
+ const rt = roundtrip(df);
+ const vals = colArr(rt, "n");
+ for (let i = 0; i < data.length; i++) {
+ if (data[i] === null) {
+ expect(vals[i]).toBeNull();
+ } else {
+ expect(vals[i]).toBe(data[i]);
+ }
+ }
+ },
+ ),
+ );
+ });
+});
diff --git a/tests/io/read_avro.test.ts b/tests/io/read_avro.test.ts
new file mode 100644
index 000000000..49cb17522
--- /dev/null
+++ b/tests/io/read_avro.test.ts
@@ -0,0 +1,385 @@
+/**
+ * Tests for src/io/read_avro.ts
+ *
+ * Covers readAvro and toAvro (round-trip), schema types, usecols, empty files,
+ * error handling, and fast-check property tests.
+ */
+import { describe, expect, it } from "bun:test";
+import * as fc from "fast-check";
+import { DataFrame } from "../../src/core/frame.ts";
+import { readAvro, toAvro } from "../../src/io/read_avro.ts";
+
+// ─── Helpers: build minimal Avro OCF by hand ─────────────────────────────────
+
+function writeLongTo(arr: number[], v: number): void {
+ let n = (v << 1) ^ (v >> 31);
+ while (n & ~0x7f) {
+ arr.push((n & 0x7f) | 0x80);
+ n >>>= 7;
+ }
+ arr.push(n);
+}
+
+function writeStringTo(arr: number[], s: string): void {
+ const b = new TextEncoder().encode(s);
+ writeLongTo(arr, b.length);
+ for (const byte of b) {
+ arr.push(byte);
+ }
+}
+
+function writeBytesTo(arr: number[], b: Uint8Array): void {
+ writeLongTo(arr, b.length);
+ for (const byte of b) {
+ arr.push(byte);
+ }
+}
+
+function buildAvroOCF(schema: object, rows: Record[]): Uint8Array {
+ const schemaBytes = new TextEncoder().encode(JSON.stringify(schema));
+ const sync = new Uint8Array([1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16]);
+ const buf: number[] = [];
+
+ // Magic
+ buf.push(79, 98, 106, 1); // "Obj\x01"
+
+ // Metadata: avro.schema + avro.codec
+ writeLongTo(buf, 2);
+ writeStringTo(buf, "avro.schema");
+ writeBytesTo(buf, schemaBytes);
+ writeStringTo(buf, "avro.codec");
+ writeBytesTo(buf, new TextEncoder().encode("null"));
+ writeLongTo(buf, 0);
+
+ // Sync marker
+ for (const b of sync) {
+ buf.push(b);
+ }
+
+ // Encode rows
+ const rowBuf: number[] = [];
+ for (const row of rows) {
+ for (const field of (schema as { fields: { name: string; type: unknown }[] }).fields) {
+ const v = row[field.name];
+ const t = field.type;
+ if (t === "null") {
+ // nothing
+ } else if (t === "boolean") {
+ rowBuf.push(v ? 1 : 0);
+ } else if (t === "int" || t === "long") {
+ writeLongTo(rowBuf, typeof v === "number" ? v : 0);
+ } else if (t === "double") {
+ const arr = new Float64Array(1);
+ arr[0] = typeof v === "number" ? v : 0;
+ for (const b of new Uint8Array(arr.buffer)) {
+ rowBuf.push(b);
+ }
+ } else if (t === "string") {
+ writeStringTo(rowBuf, String(v ?? ""));
+ } else if (Array.isArray(t)) {
+ // union ["null", X]
+ if (v === null || v === undefined) {
+ writeLongTo(rowBuf, 0);
+ } else {
+ writeLongTo(rowBuf, 1);
+ const inner = t[1] as string;
+ if (inner === "string") {
+ writeStringTo(rowBuf, String(v));
+ } else if (inner === "long" || inner === "int") {
+ writeLongTo(rowBuf, Number(v));
+ } else if (inner === "double") {
+ const a2 = new Float64Array(1);
+ a2[0] = Number(v);
+ for (const b of new Uint8Array(a2.buffer)) {
+ rowBuf.push(b);
+ }
+ }
+ }
+ }
+ }
+ }
+
+ // Data block: count, byteCount, data, sync
+ if (rowBuf.length > 0) {
+ writeLongTo(buf, rows.length);
+ writeLongTo(buf, rowBuf.length);
+ for (const b of rowBuf) {
+ buf.push(b);
+ }
+ for (const b of sync) {
+ buf.push(b);
+ }
+ }
+
+ return new Uint8Array(buf);
+}
+
+// ─── readAvro – basic parsing ──────────────────────────────────────────────────
+
+describe("readAvro – basic types", () => {
+ it("reads an int column", () => {
+ const schema = { type: "record", name: "R", fields: [{ name: "n", type: "int" }] };
+ const buf = buildAvroOCF(schema, [{ n: 1 }, { n: -2 }, { n: 42 }]);
+ const df = readAvro(buf);
+ expect(df.shape[0]).toBe(3);
+ expect(df.col("n").at(0)).toBe(1);
+ expect(df.col("n").at(1)).toBe(-2);
+ expect(df.col("n").at(2)).toBe(42);
+ });
+
+ it("reads a double column", () => {
+ const schema = { type: "record", name: "R", fields: [{ name: "v", type: "double" }] };
+ const buf = buildAvroOCF(schema, [{ v: 3.14 }, { v: -1.5 }]);
+ const df = readAvro(buf);
+ expect(df.col("v").at(0) as number).toBeCloseTo(3.14, 5);
+ expect(df.col("v").at(1) as number).toBeCloseTo(-1.5, 5);
+ });
+
+ it("reads a string column", () => {
+ const schema = { type: "record", name: "R", fields: [{ name: "s", type: "string" }] };
+ const buf = buildAvroOCF(schema, [{ s: "hello" }, { s: "world" }]);
+ const df = readAvro(buf);
+ expect(df.col("s").at(0)).toBe("hello");
+ expect(df.col("s").at(1)).toBe("world");
+ });
+
+ it("reads a boolean column", () => {
+ const schema = { type: "record", name: "R", fields: [{ name: "b", type: "boolean" }] };
+ const buf = buildAvroOCF(schema, [{ b: true }, { b: false }]);
+ const df = readAvro(buf);
+ expect(df.col("b").at(0)).toBe(true);
+ expect(df.col("b").at(1)).toBe(false);
+ });
+
+ it("reads nullable (union) columns with null values", () => {
+ const schema = {
+ type: "record",
+ name: "R",
+ fields: [{ name: "x", type: ["null", "string"] }],
+ };
+ const buf = buildAvroOCF(schema, [{ x: "foo" }, { x: null }, { x: "bar" }]);
+ const df = readAvro(buf);
+ expect(df.col("x").at(0)).toBe("foo");
+ expect(df.col("x").at(1)).toBeNull();
+ expect(df.col("x").at(2)).toBe("bar");
+ });
+
+ it("reads multiple columns of mixed types", () => {
+ const schema = {
+ type: "record",
+ name: "R",
+ fields: [
+ { name: "id", type: "int" },
+ { name: "name", type: "string" },
+ { name: "val", type: "double" },
+ ],
+ };
+ const rows = [
+ { id: 1, name: "alice", val: 1.1 },
+ { id: 2, name: "bob", val: 2.2 },
+ ];
+ const df = readAvro(buildAvroOCF(schema, rows));
+ expect(df.shape).toEqual([2, 3]);
+ expect(df.col("id").at(0)).toBe(1);
+ expect(df.col("name").at(1)).toBe("bob");
+ expect(df.col("val").at(1) as number).toBeCloseTo(2.2, 5);
+ });
+});
+
+describe("readAvro – usecols", () => {
+ it("returns only requested columns", () => {
+ const schema = {
+ type: "record",
+ name: "R",
+ fields: [
+ { name: "a", type: "int" },
+ { name: "b", type: "string" },
+ { name: "c", type: "double" },
+ ],
+ };
+ const rows = [
+ { a: 1, b: "x", c: 0.5 },
+ { a: 2, b: "y", c: 1.5 },
+ ];
+ const df = readAvro(buildAvroOCF(schema, rows), { usecols: ["a", "c"] });
+ expect([...df.columns.values]).toEqual(["a", "c"]);
+ expect(df.shape[1]).toBe(2);
+ });
+});
+
+describe("readAvro – error handling", () => {
+ it("throws on bad magic bytes", () => {
+ const bad = new Uint8Array(20);
+ bad.fill(0);
+ expect(() => readAvro(bad)).toThrow();
+ });
+
+ it("accepts ArrayBuffer input", () => {
+ const schema = { type: "record", name: "R", fields: [{ name: "n", type: "int" }] };
+ const buf = buildAvroOCF(schema, [{ n: 7 }]);
+ const df = readAvro(buf.buffer as ArrayBuffer);
+ expect(df.col("n").at(0)).toBe(7);
+ });
+
+ it("throws for unsupported codec", () => {
+ // Build a fake header with codec=deflate
+ const schema = { type: "record", name: "R", fields: [{ name: "n", type: "int" }] };
+ const schemaBytes = new TextEncoder().encode(JSON.stringify(schema));
+ const buf: number[] = [79, 98, 106, 1]; // magic
+ writeLongTo(buf, 2);
+ writeStringTo(buf, "avro.schema");
+ writeBytesTo(buf, schemaBytes);
+ writeStringTo(buf, "avro.codec");
+ writeBytesTo(buf, new TextEncoder().encode("deflate"));
+ writeLongTo(buf, 0);
+ for (let i = 0; i < 16; i++) {
+ buf.push(i); // sync
+ }
+ // No data blocks
+ expect(() => readAvro(new Uint8Array(buf))).toThrow(/deflate/);
+ });
+});
+
+describe("readAvro – empty file", () => {
+ it("returns empty DataFrame for file with no rows", () => {
+ const schema = { type: "record", name: "R", fields: [{ name: "n", type: "int" }] };
+ const df = readAvro(buildAvroOCF(schema, []));
+ expect(df.shape[0]).toBe(0);
+ });
+});
+
+// ─── toAvro / round-trip ──────────────────────────────────────────────────────
+
+describe("toAvro – file structure", () => {
+ it("starts with Avro magic bytes", () => {
+ const df = DataFrame.fromColumns({ a: [1, 2, 3] });
+ const buf = toAvro(df);
+ expect(buf[0]).toBe(79); // 'O'
+ expect(buf[1]).toBe(98); // 'b'
+ expect(buf[2]).toBe(106); // 'j'
+ expect(buf[3]).toBe(1); // version
+ });
+
+ it("produces a Uint8Array", () => {
+ const df = DataFrame.fromColumns({ x: [1.1, 2.2] });
+ expect(toAvro(df)).toBeInstanceOf(Uint8Array);
+ });
+});
+
+describe("toAvro + readAvro – round-trip", () => {
+ it("integer column round-trips", () => {
+ const df = DataFrame.fromColumns({ id: [1, 2, 3, 4, 5] });
+ const buf = toAvro(df);
+ const df2 = readAvro(buf);
+ expect(df2.shape[0]).toBe(5);
+ for (let i = 0; i < 5; i++) {
+ expect(df2.col("id").at(i)).toBe(i + 1);
+ }
+ });
+
+ it("double column round-trips", () => {
+ const df = DataFrame.fromColumns({ v: [1.5, 2.5, 3.5] });
+ const buf = toAvro(df);
+ const df2 = readAvro(buf);
+ expect(df2.col("v").at(0) as number).toBeCloseTo(1.5, 5);
+ expect(df2.col("v").at(2) as number).toBeCloseTo(3.5, 5);
+ });
+
+ it("string column round-trips", () => {
+ const df = DataFrame.fromColumns({ name: ["alice", "bob", "carol"] });
+ const buf = toAvro(df);
+ const df2 = readAvro(buf);
+ expect(df2.col("name").at(1)).toBe("bob");
+ });
+
+ it("boolean column round-trips", () => {
+ const df = DataFrame.fromColumns({ flag: [true, false, true] });
+ const buf = toAvro(df);
+ const df2 = readAvro(buf);
+ expect(df2.col("flag").at(0)).toBe(true);
+ expect(df2.col("flag").at(1)).toBe(false);
+ });
+
+ it("null values round-trip in nullable columns", () => {
+ const df = DataFrame.fromColumns({ x: [1, null, 3] });
+ const buf = toAvro(df);
+ const df2 = readAvro(buf);
+ expect(df2.col("x").at(0)).toBe(1);
+ expect(df2.col("x").at(1)).toBeNull();
+ expect(df2.col("x").at(2)).toBe(3);
+ });
+
+ it("multi-column mixed-type round-trip", () => {
+ const df = DataFrame.fromColumns({
+ id: [1, 2, 3],
+ name: ["a", "b", "c"],
+ score: [0.1, 0.2, 0.3],
+ active: [true, false, true],
+ });
+ const buf = toAvro(df);
+ const df2 = readAvro(buf);
+ expect(df2.shape).toEqual([3, 4]);
+ expect(df2.col("name").at(1)).toBe("b");
+ expect(df2.col("score").at(2) as number).toBeCloseTo(0.3, 5);
+ });
+
+ it("empty DataFrame round-trips", () => {
+ const df = DataFrame.fromColumns({ a: [] as number[] });
+ const buf = toAvro(df);
+ const df2 = readAvro(buf);
+ expect(df2.shape[0]).toBe(0);
+ });
+});
+
+// ─── Property-based tests ──────────────────────────────────────────────────────
+
+describe("property tests", () => {
+ it("integer round-trip: toAvro → readAvro preserves integer values", () => {
+ fc.assert(
+ fc.property(
+ fc.array(fc.integer({ min: -1000, max: 1000 }), { minLength: 1, maxLength: 20 }),
+ (vals) => {
+ const df = DataFrame.fromColumns({ n: vals });
+ const df2 = readAvro(toAvro(df));
+ for (let i = 0; i < vals.length; i++) {
+ if (df2.col("n").at(i) !== vals[i]) {
+ return false;
+ }
+ }
+ return true;
+ },
+ ),
+ );
+ });
+
+ it("string round-trip preserves values", () => {
+ fc.assert(
+ fc.property(
+ fc.array(fc.string({ minLength: 0, maxLength: 20 }), { minLength: 1, maxLength: 15 }),
+ (vals) => {
+ const df = DataFrame.fromColumns({ s: vals });
+ const df2 = readAvro(toAvro(df));
+ for (let i = 0; i < vals.length; i++) {
+ if (df2.col("s").at(i) !== vals[i]) {
+ return false;
+ }
+ }
+ return true;
+ },
+ ),
+ );
+ });
+
+ it("row count is preserved in round-trip", () => {
+ fc.assert(
+ fc.property(
+ fc.array(fc.integer({ min: 0, max: 100 }), { minLength: 0, maxLength: 30 }),
+ (vals) => {
+ const df = DataFrame.fromColumns({ n: vals });
+ const df2 = readAvro(toAvro(df));
+ return df2.shape[0] === vals.length;
+ },
+ ),
+ );
+ });
+});
diff --git a/tests/io/read_html.test.ts b/tests/io/read_html.test.ts
index dfc41d7fa..2bd9ad44c 100644
--- a/tests/io/read_html.test.ts
+++ b/tests/io/read_html.test.ts
@@ -55,7 +55,7 @@ describe("readHtml – basic", () => {
test("values are numeric by default", () => {
const html = simpleTable(["n"], [["42"], ["3.14"]]);
const [df] = readHtml(html);
- const vals = df!.col("n").toArray();
+ const vals = df?.col("n").toArray() ?? [];
expect(vals[0]).toBe(42);
expect(vals[1]).toBeCloseTo(3.14);
});
@@ -87,7 +87,7 @@ describe("readHtml – header", () => {
| 1 | 2 | 3 |
`;
const [df] = readHtml(html);
- const cols = df!.columns.toArray();
+ const cols = df?.columns.toArray() ?? [];
expect(cols[0]).toBe("x");
expect(cols[1]).toBe("x.1");
expect(cols[2]).toBe("y");
@@ -226,7 +226,7 @@ describe("readHtml – HTML entities", () => {
| hello world |
`;
const [df] = readHtml(html, { converters: false });
- const vals = df!.col("k").toArray();
+ const vals = df?.col("k").toArray() ?? [];
expect(vals[0]).toBe("a & b");
expect(vals[1]).toBe("5 < 10");
expect(vals[2]).toBe("hello world");
@@ -265,7 +265,7 @@ describe("readHtml – structure variants", () => {
| Total |
`;
const [df] = readHtml(html, { converters: false });
- const vals = df?.col("a").toArray();
+ const vals = df?.col("a").toArray() ?? [];
expect(vals).toContain("1");
expect(vals).toContain("Total");
});
@@ -305,14 +305,14 @@ describe("readHtml – property tests", () => {
),
),
(rows) => {
- const ncols = rows[0]!.length;
+ const ncols = rows[0]?.length ?? 0;
const headers = Array.from({ length: ncols }, (_, i) => `col${i}`);
const strRows = rows.map((r) => r.map(String));
const html = simpleTable(headers, strRows);
const [df] = readHtml(html);
const flatIn = rows.flat();
const flatOut = (df?.toRecords() ?? []).flatMap((record) =>
- rows[0]!.map((_, ci) => Number(record[headers[ci]!])),
+ rows[0]?.map((_, ci) => Number(record[headers[ci]!])),
);
// same length
if (flatIn.length !== flatOut.length) {
diff --git a/tests/io/sql.test.ts b/tests/io/sql.test.ts
index 3c4d1b04c..7435fb711 100644
--- a/tests/io/sql.test.ts
+++ b/tests/io/sql.test.ts
@@ -521,6 +521,21 @@ describe("toSql / readSqlTable — round-trip", () => {
// ─── property-based tests ─────────────────────────────────────────────────────
describe("readSqlQuery — property tests", () => {
+ it("preserves column names that collide with object prototype properties", () => {
+ const columns = ["__proto__", "constructor", "toString"];
+ for (const rowCount of [0, 2]) {
+ const rows = Array.from({ length: rowCount }, () =>
+ Object.fromEntries(columns.map((column) => [column, 42])),
+ );
+ const conn: SqlConnection = { query: () => ({ columns, rows }) };
+ const df = readSqlQuery("SELECT 1", conn);
+ expect(df.shape).toEqual([rowCount, columns.length]);
+ for (const column of columns) {
+ expect([...df.col(column).values]).toEqual(Array.from({ length: rowCount }, () => 42));
+ }
+ }
+ });
+
it("shape matches result column/row counts", () => {
fc.assert(
fc.property(
diff --git a/tests/ml/attention.test.ts b/tests/ml/attention.test.ts
new file mode 100644
index 000000000..6e18e5937
--- /dev/null
+++ b/tests/ml/attention.test.ts
@@ -0,0 +1,52 @@
+/**
+ * Tests for src/ml/attention.ts
+ */
+import { describe, expect, it } from "bun:test";
+import { applyRoPE, causalMask, sinusoidalPositionalEncoding } from "../../src/index.ts";
+
+describe("sinusoidalPositionalEncoding", () => {
+ it("returns correct dimensions", () => {
+ const pe = sinusoidalPositionalEncoding(10, 16);
+ expect(pe.length).toBe(10 * 16);
+ });
+
+ it("sin/cos values are in [-1, 1]", () => {
+ const pe = sinusoidalPositionalEncoding(5, 8);
+ for (let i = 0; i < pe.length; i++) {
+ expect(pe[i]!).toBeGreaterThanOrEqual(-1);
+ expect(pe[i]!).toBeLessThanOrEqual(1);
+ }
+ });
+});
+
+describe("causalMask", () => {
+ it("upper triangle is -Infinity", () => {
+ const mask = causalMask(3);
+ // (0,1), (0,2), (1,2) should be -Infinity
+ expect(mask[0 * 3 + 1]).toBe(Number.NEGATIVE_INFINITY);
+ expect(mask[0 * 3 + 2]).toBe(Number.NEGATIVE_INFINITY);
+ expect(mask[1 * 3 + 2]).toBe(Number.NEGATIVE_INFINITY);
+ });
+
+ it("diagonal is 0", () => {
+ const mask = causalMask(4);
+ for (let i = 0; i < 4; i++) expect(mask[i * 4 + i]).toBe(0);
+ });
+});
+
+describe("applyRoPE", () => {
+ it("returns same dimensions", () => {
+ const x = new Float64Array(3 * 4); // seqLen=3, dHead=4
+ const out = applyRoPE(x, 3, 4);
+ expect(out.length).toBe(12);
+ });
+
+ it("preserves norm (approximately)", () => {
+ const x = new Float64Array([1, 0, 1, 0, 0, 1, 0, 1]); // seqLen=2, dHead=4
+ const out = applyRoPE(x, 2, 4);
+ // Each 2D pair should have same norm as input pair
+ const normIn = Math.sqrt(x[0]! ** 2 + x[1]! ** 2);
+ const normOut = Math.sqrt(out[0]! ** 2 + out[1]! ** 2);
+ expect(normOut).toBeCloseTo(normIn, 5);
+ });
+});
diff --git a/tests/ml/bayesian_opt.test.ts b/tests/ml/bayesian_opt.test.ts
new file mode 100644
index 000000000..5d5ada0af
--- /dev/null
+++ b/tests/ml/bayesian_opt.test.ts
@@ -0,0 +1,106 @@
+/**
+ * Tests for src/ml/bayesian_opt.ts
+ */
+import { describe, expect, it } from "bun:test";
+import {
+ cholesky,
+ expectedImprovement,
+ gpPredict,
+ initBOState,
+ maternKernel52,
+ normalCDF,
+ normalPDF,
+ rbfKernel,
+ suggestNext,
+ upperConfidenceBound,
+} from "../../src/index.ts";
+
+describe("kernels", () => {
+ it("rbf kernel is 1 for identical points", () => {
+ const x = new Float64Array([1, 2, 3]);
+ expect(rbfKernel(x, x, 1, 1)).toBeCloseTo(1);
+ });
+
+ it("rbf kernel decreases with distance", () => {
+ const x0 = new Float64Array([0]);
+ const x1 = new Float64Array([1]);
+ const x2 = new Float64Array([2]);
+ expect(rbfKernel(x0, x1, 1, 1)).toBeGreaterThan(rbfKernel(x0, x2, 1, 1));
+ });
+
+ it("matern52 is 1 for identical points", () => {
+ const x = new Float64Array([0, 0]);
+ expect(maternKernel52(x, x, 1, 1)).toBeCloseTo(1);
+ });
+});
+
+describe("cholesky", () => {
+ it("L L^T ≈ A for 2x2 SPD matrix", () => {
+ const A = new Float64Array([4, 2, 2, 3]);
+ const L = cholesky(A, 2);
+ // L L^T
+ const LLT = new Float64Array(4);
+ for (let i = 0; i < 2; i++) {
+ for (let j = 0; j < 2; j++) {
+ let sum = 0;
+ for (let k = 0; k < 2; k++) sum += (L[i * 2 + k] ?? 0) * (L[j * 2 + k] ?? 0);
+ LLT[i * 2 + j] = sum;
+ }
+ }
+ expect(LLT[0]!).toBeCloseTo(4, 4);
+ expect(LLT[1]!).toBeCloseTo(2, 4);
+ expect(LLT[3]!).toBeCloseTo(3, 4);
+ });
+});
+
+describe("gpPredict", () => {
+ it("interpolates: mean near observed value at observed point", () => {
+ const X = [new Float64Array([0]), new Float64Array([1]), new Float64Array([2])];
+ const y = new Float64Array([0, 1, 0]);
+ const { kernelMatrix, cholesky: ch } = (() => {
+ const K = new Float64Array(9);
+ for (let i = 0; i < 3; i++) {
+ for (let j = i; j < 3; j++) {
+ const k = rbfKernel(X[i]!, X[j]!, 1, 1);
+ K[i * 3 + j] = k;
+ K[j * 3 + i] = k;
+ }
+ K[i * 3 + i] = (K[i * 3 + i] ?? 0) + 0.01;
+ }
+ return { kernelMatrix: K, cholesky: cholesky(K, 3) };
+ })();
+ const { mean } = gpPredict(new Float64Array([1]), X, y, ch, 1, 1, 0.01);
+ expect(Math.abs(mean - 1)).toBeLessThan(0.2);
+ });
+});
+
+describe("normalCDF and normalPDF", () => {
+ it("normalCDF(0) ≈ 0.5", () => {
+ expect(normalCDF(0)).toBeCloseTo(0.5, 2);
+ });
+
+ it("normalPDF(0) ≈ 0.3989", () => {
+ expect(normalPDF(0)).toBeCloseTo(1 / Math.sqrt(2 * Math.PI), 4);
+ });
+});
+
+describe("acquisitions", () => {
+ it("EI >= 0", () => {
+ expect(expectedImprovement(1.5, 0.1, 1.0)).toBeGreaterThanOrEqual(0);
+ });
+
+ it("UCB increases with variance", () => {
+ expect(upperConfidenceBound(1, 0.5, 2)).toBeGreaterThan(upperConfidenceBound(1, 0.1, 2));
+ });
+});
+
+describe("suggestNext", () => {
+ it("returns a candidate", () => {
+ const X = [new Float64Array([0]), new Float64Array([1])];
+ const y = new Float64Array([0, 1]);
+ const state = initBOState(X, y);
+ const candidates = [new Float64Array([0.5]), new Float64Array([1.5]), new Float64Array([2.0])];
+ const { bestCandidate } = suggestNext(state, candidates);
+ expect(bestCandidate.length).toBe(1);
+ });
+});
diff --git a/tests/ml/crf.test.ts b/tests/ml/crf.test.ts
new file mode 100644
index 000000000..61f91ed9f
--- /dev/null
+++ b/tests/ml/crf.test.ts
@@ -0,0 +1,83 @@
+/**
+ * Tests for src/ml/crf.ts
+ */
+import { describe, expect, it } from "bun:test";
+import { crfNegLogLikelihood, forwardLogZ, logSumExp, viterbiDecode } from "../../src/index.ts";
+import type { CRFParams } from "../../src/index.ts";
+
+function makeCRFParams(): CRFParams {
+ const numTags = 3;
+ const seqLen = 4;
+ const emissionScores = new Float64Array([
+ 1,
+ 0,
+ 0, // t=0: prefer tag 0
+ 0,
+ 2,
+ 0, // t=1: prefer tag 1
+ 0,
+ 0,
+ 3, // t=2: prefer tag 2
+ 1,
+ 0,
+ 0, // t=3: prefer tag 0
+ ]);
+ const transitionScores = new Float64Array(numTags * numTags).fill(0);
+ const startScores = new Float64Array([1, 0, 0]); // start with tag 0
+ const endScores = new Float64Array(numTags).fill(0);
+ return { emissionScores, transitionScores, startScores, endScores, numTags, seqLen };
+}
+
+describe("viterbiDecode", () => {
+ it("returns correct tag sequence for simple emissions", () => {
+ const params = makeCRFParams();
+ const result = viterbiDecode(params);
+ expect(result.tags.length).toBe(4);
+ expect(result.tags[0]).toBe(0); // highest emission at t=0
+ expect(result.tags[1]).toBe(1); // highest emission at t=1
+ expect(result.tags[2]).toBe(2); // highest emission at t=2
+ });
+
+ it("score is finite", () => {
+ const params = makeCRFParams();
+ const result = viterbiDecode(params);
+ expect(Number.isFinite(result.score)).toBe(true);
+ });
+});
+
+describe("forwardLogZ", () => {
+ it("log partition >= viterbi score (log Z >= best path)", () => {
+ const params = makeCRFParams();
+ const logZ = forwardLogZ(params);
+ const { score } = viterbiDecode(params);
+ expect(logZ).toBeGreaterThanOrEqual(score - 1e-9);
+ });
+});
+
+describe("crfNegLogLikelihood", () => {
+ it("nll >= 0 (log Z >= gold path score)", () => {
+ const params = makeCRFParams();
+ const goldTags = [0, 1, 2, 0];
+ const nll = crfNegLogLikelihood(params, goldTags);
+ expect(nll).toBeGreaterThanOrEqual(-1e-9);
+ });
+
+ it("nll is lower for better tag sequence", () => {
+ const params = makeCRFParams();
+ const good = [0, 1, 2, 0];
+ const bad = [2, 0, 1, 2];
+ expect(crfNegLogLikelihood(params, good)).toBeLessThan(crfNegLogLikelihood(params, bad));
+ });
+});
+
+describe("logSumExp", () => {
+ it("logSumExp([0, 0, 0]) ≈ log(3)", () => {
+ const vals = new Float64Array([0, 0, 0]);
+ expect(logSumExp(vals)).toBeCloseTo(Math.log(3), 5);
+ });
+
+ it("logSumExp of single value returns that value", () => {
+ const vals = new Float64Array([5.0]);
+ expect(logSumExp(vals)).toBeCloseTo(5.0, 5);
+ });
+});
diff --git a/tests/ml/ddim.test.ts b/tests/ml/ddim.test.ts
new file mode 100644
index 000000000..c63111c42
--- /dev/null
+++ b/tests/ml/ddim.test.ts
@@ -0,0 +1,224 @@
+/**
+ * Tests for src/ml/ddim.ts
+ */
+import { describe, expect, it } from "bun:test";
+import fc from "fast-check";
+import type { NoiseSchedule } from "../../src/index.ts";
+import { addNoise, computeNoiseSchedule, ddimTimesteps, snrAtTimestep } from "../../src/index.ts";
+
+describe("computeNoiseSchedule — linear", () => {
+ it("has correct length", () => {
+ const s = computeNoiseSchedule({
+ numTrainTimesteps: 100,
+ schedule: "linear",
+ eta: 0,
+ betaStart: 0.0001,
+ betaEnd: 0.02,
+ });
+ expect(s.betas.length).toBe(100);
+ expect(s.alphasCumprod.length).toBe(100);
+ });
+
+ it("alphasCumprod is decreasing", () => {
+ const s = computeNoiseSchedule({
+ numTrainTimesteps: 50,
+ schedule: "linear",
+ eta: 0,
+ betaStart: 0.0001,
+ betaEnd: 0.02,
+ });
+ for (let i = 1; i < 50; i++) {
+ expect(s.alphasCumprod[i]!).toBeLessThan(s.alphasCumprod[i - 1]!);
+ }
+ });
+
+ it("sqrtAlphasCumprod[0] close to 1", () => {
+ const s = computeNoiseSchedule({
+ numTrainTimesteps: 1000,
+ schedule: "linear",
+ eta: 0,
+ betaStart: 0.0001,
+ betaEnd: 0.02,
+ });
+ expect(s.sqrtAlphasCumprod[0]!).toBeGreaterThan(0.99);
+ });
+});
+
+describe("computeNoiseSchedule — validation and edge cases", () => {
+ for (const schedule of ["linear", "sqrt"] satisfies readonly ("linear" | "sqrt")[]) {
+ it(`${schedule} uses betaStart for a single timestep`, () => {
+ const s = computeNoiseSchedule({
+ numTrainTimesteps: 1,
+ schedule,
+ eta: 0,
+ betaStart: 0.01,
+ betaEnd: 0.1,
+ });
+ expect(Array.from(s.betas)).toEqual([0.01]);
+ expect(s.alphasCumprod[0]).toBeCloseTo(0.99, 12);
+ expect(s.sqrtAlphasCumprod[0]).toBeCloseTo(Math.sqrt(0.99), 12);
+ expect(s.sqrtOneMinusAlphasCumprod[0]).toBeCloseTo(0.1, 12);
+ });
+ }
+
+ it("rejects invalid timestep counts", () => {
+ for (const numTrainTimesteps of [0, -1, 1.5, Number.NaN, Number.POSITIVE_INFINITY]) {
+ expect(() =>
+ computeNoiseSchedule({
+ numTrainTimesteps,
+ schedule: "linear",
+ eta: 0,
+ betaStart: 0.01,
+ betaEnd: 0.1,
+ }),
+ ).toThrow(RangeError);
+ }
+ });
+
+ it("rejects invalid beta endpoints", () => {
+ for (const value of [-0.1, 1, Number.NaN, Number.POSITIVE_INFINITY]) {
+ expect(() =>
+ computeNoiseSchedule({
+ numTrainTimesteps: 10,
+ schedule: "linear",
+ eta: 0,
+ betaStart: value,
+ betaEnd: 0.1,
+ }),
+ ).toThrow(RangeError);
+ expect(() =>
+ computeNoiseSchedule({
+ numTrainTimesteps: 10,
+ schedule: "sqrt",
+ eta: 0,
+ betaStart: 0.01,
+ betaEnd: value,
+ }),
+ ).toThrow(RangeError);
+ }
+ });
+
+ it("produces finite schedule values for valid inputs", () => {
+ fc.assert(
+ fc.property(
+ fc.integer({ min: 1, max: 100 }),
+ fc.constantFrom("linear", "sqrt", "cosine"),
+ fc.double({ min: 0, max: 0.1, noNaN: true }),
+ fc.double({ min: 0, max: 0.1, noNaN: true }),
+ (numTrainTimesteps, schedule, betaStart, betaEnd) => {
+ const s = computeNoiseSchedule({
+ numTrainTimesteps,
+ schedule,
+ eta: 0,
+ betaStart,
+ betaEnd,
+ });
+ for (const values of Object.values(s)) {
+ expect(values.length).toBe(numTrainTimesteps);
+ expect(Array.from(values).every(Number.isFinite)).toBe(true);
+ }
+ },
+ ),
+ );
+ });
+});
+
+describe("computeNoiseSchedule — cosine", () => {
+ it("cumulative alphas follow the normalized cosine curve", () => {
+ const s = computeNoiseSchedule({
+ numTrainTimesteps: 100,
+ schedule: "cosine",
+ eta: 0,
+ betaStart: 0.0001,
+ betaEnd: 0.02,
+ });
+ const initial = Math.cos(((0.008 / 1.008) * Math.PI) / 2) ** 2;
+ // The final beta is capped; earlier cumulative values follow alpha_bar.
+ for (let t = 0; t < 99; t++) {
+ const expected = Math.cos(((((t + 1) / 100 + 0.008) / 1.008) * Math.PI) / 2) ** 2 / initial;
+ expect(s.alphasCumprod[t]).toBeCloseTo(expected, 12);
+ }
+ expect(s.betas[99]).toBe(0.999);
+ });
+ it("alphasCumprod values are in (0,1)", () => {
+ const s = computeNoiseSchedule({
+ numTrainTimesteps: 100,
+ schedule: "cosine",
+ eta: 0,
+ betaStart: 0.0001,
+ betaEnd: 0.02,
+ });
+ for (let i = 0; i < 100; i++) {
+ expect(s.alphasCumprod[i]!).toBeGreaterThan(0);
+ expect(s.alphasCumprod[i]!).toBeLessThan(1);
+ }
+ });
+});
+
+describe("addNoise", () => {
+ it("mixes signal and noise according to the one-step schedule", () => {
+ const s = computeNoiseSchedule({
+ numTrainTimesteps: 1,
+ schedule: "linear",
+ eta: 0,
+ betaStart: 0.0001,
+ betaEnd: 0.0001,
+ });
+ const sample = new Float64Array([1, 2, 3]);
+ const noise = new Float64Array([10, 10, 10]);
+ const out = addNoise(sample, noise, 0, s);
+ for (let i = 0; i < sample.length; i++) {
+ expect(out[i]).toBeCloseTo(Math.sqrt(0.9999) * (sample[i] ?? 0) + 0.1, 12);
+ }
+ });
+
+ it("returns the sample unchanged for zero beta", () => {
+ const s = computeNoiseSchedule({
+ numTrainTimesteps: 1,
+ schedule: "linear",
+ eta: 0,
+ betaStart: 0,
+ betaEnd: 0,
+ });
+ const sample = new Float64Array([1, 2, 3]);
+ expect(addNoise(sample, new Float64Array([10, 10, 10]), 0, s)).toEqual(sample);
+ });
+
+ it("output length matches input", () => {
+ const s = computeNoiseSchedule({
+ numTrainTimesteps: 10,
+ schedule: "linear",
+ eta: 0,
+ betaStart: 0.001,
+ betaEnd: 0.01,
+ });
+ const sample = new Float64Array(8);
+ const noise = new Float64Array(8);
+ expect(addNoise(sample, noise, 5, s).length).toBe(8);
+ });
+});
+
+describe("ddimTimesteps", () => {
+ it("returns correct count", () => {
+ const ts = ddimTimesteps(1000, 50);
+ expect(ts.length).toBe(50);
+ });
+
+ it("first timestep is largest", () => {
+ const ts = ddimTimesteps(1000, 10);
+ expect(ts[0]!).toBeGreaterThan(ts[ts.length - 1]!);
+ });
+});
+
+describe("snrAtTimestep", () => {
+ it("snr decreases over time for linear schedule", () => {
+ const s = computeNoiseSchedule({
+ numTrainTimesteps: 100,
+ schedule: "linear",
+ eta: 0,
+ betaStart: 0.001,
+ betaEnd: 0.02,
+ });
+ expect(snrAtTimestep(10, s)).toBeGreaterThan(snrAtTimestep(50, s));
+ });
+});
diff --git a/tests/ml/gradient_boosting.test.ts b/tests/ml/gradient_boosting.test.ts
new file mode 100644
index 000000000..af6f5fbb9
--- /dev/null
+++ b/tests/ml/gradient_boosting.test.ts
@@ -0,0 +1,73 @@
+/**
+ * Tests for src/ml/gradient_boosting.ts
+ */
+import { describe, expect, it } from "bun:test";
+import {
+ buildTree,
+ fitGBM,
+ mse,
+ predictGBM,
+ r2Score,
+ treePredict,
+ treePredictOne,
+} from "../../src/index.ts";
+
+describe("decision tree", () => {
+ const X: Float64Array[] = [
+ new Float64Array([1]),
+ new Float64Array([2]),
+ new Float64Array([3]),
+ new Float64Array([4]),
+ new Float64Array([5]),
+ ];
+ const y = new Float64Array([1, 2, 3, 4, 5]);
+
+ it("builds a tree and predicts", () => {
+ const tree = buildTree(X, y, [0, 1, 2, 3, 4], 0, { maxDepth: 3, minSamplesLeaf: 1 });
+ const preds = treePredict(tree, X);
+ expect(preds.length).toBe(5);
+ });
+
+ it("leaf prediction is mean of remaining samples", () => {
+ const tree = buildTree(X, y, [0, 1, 2, 3, 4], 0, { maxDepth: 0, minSamplesLeaf: 1 });
+ const pred = treePredictOne(tree, new Float64Array([2.5]));
+ expect(pred).toBeCloseTo(3, 0); // mean of [1,2,3,4,5]
+ });
+});
+
+describe("GBM — simple regression", () => {
+ const n = 20;
+ const X: Float64Array[] = Array.from({ length: n }, (_, i) => new Float64Array([i / 10]));
+ const y = Float64Array.from({ length: n }, (_, i) => (i / 10) * 2 + 1);
+
+ it("fits and predicts", () => {
+ const model = fitGBM(X, y, 50, 0.1, 3, 1);
+ const preds = predictGBM(model, X);
+ expect(preds.length).toBe(n);
+ });
+
+ it("achieves r2 > 0.9", () => {
+ const model = fitGBM(X, y, 100, 0.1, 3, 1);
+ const preds = predictGBM(model, X);
+ const r2 = r2Score(y, preds);
+ expect(r2).toBeGreaterThan(0.9);
+ });
+});
+
+describe("mse and r2Score", () => {
+ it("mse of perfect predictions is 0", () => {
+ const y = new Float64Array([1, 2, 3]);
+ expect(mse(y, y)).toBeCloseTo(0);
+ });
+
+ it("r2 of perfect predictions is 1", () => {
+ const y = new Float64Array([1, 2, 3, 4]);
+ expect(r2Score(y, y)).toBeCloseTo(1);
+ });
+
+ it("r2 of constant predictions is ≤ 0", () => {
+ const y = new Float64Array([1, 2, 3, 4]);
+ const constant = new Float64Array([2, 2, 2, 2]);
+ expect(r2Score(y, constant)).toBeLessThanOrEqual(0.1);
+ });
+});
diff --git a/tests/ml/neural_network.test.ts b/tests/ml/neural_network.test.ts
new file mode 100644
index 000000000..2bd45fb40
--- /dev/null
+++ b/tests/ml/neural_network.test.ts
@@ -0,0 +1,113 @@
+/**
+ * Tests for src/ml/neural_network.ts
+ */
+import { describe, expect, it } from "bun:test";
+import {
+ adamStep,
+ denseForward,
+ heInit,
+ initAdam,
+ mseLoss,
+ relu,
+ sigmoidActivation,
+ softmaxCrossEntropy,
+} from "../../src/index.ts";
+
+describe("denseForward", () => {
+ it("computes W x + b correctly", () => {
+ const W = new Float64Array([1, 0, 0, 1]); // identity 2x2
+ const b = new Float64Array([1, 2]);
+ const x = new Float64Array([3, 4]);
+ const out = denseForward(x, W, b, 2, 2);
+ expect(out[0]).toBeCloseTo(4); // 3 + 1
+ expect(out[1]).toBeCloseTo(6); // 4 + 2
+ });
+});
+
+describe("relu", () => {
+ it("passes positive values unchanged", () => {
+ const x = new Float64Array([1, 2, 3]);
+ const out = relu(x);
+ expect(out[0]).toBe(1);
+ expect(out[1]).toBe(2);
+ });
+
+ it("zeros negative values", () => {
+ const x = new Float64Array([-1, -2, 0]);
+ const out = relu(x);
+ expect(out[0]).toBe(0);
+ expect(out[1]).toBe(0);
+ expect(out[2]).toBe(0);
+ });
+});
+
+describe("sigmoidActivation", () => {
+ it("sigmoid(0) = 0.5", () => {
+ const x = new Float64Array([0]);
+ expect(sigmoidActivation(x)[0]).toBeCloseTo(0.5);
+ });
+
+ it("output is in (0,1)", () => {
+ const x = new Float64Array([-10, 0, 10]);
+ const out = sigmoidActivation(x);
+ for (let i = 0; i < 3; i++) {
+ expect(out[i]!).toBeGreaterThan(0);
+ expect(out[i]!).toBeLessThan(1);
+ }
+ });
+});
+
+describe("softmaxCrossEntropy", () => {
+ it("loss is 0 when prediction is perfect", () => {
+ const logits = new Float64Array([100, 0, 0]);
+ const labels = new Float64Array([1, 0, 0]);
+ const { loss } = softmaxCrossEntropy(logits, labels);
+ expect(loss).toBeCloseTo(0, 2);
+ });
+
+ it("dLogits sum to 0", () => {
+ const logits = new Float64Array([1, 2, 3]);
+ const labels = new Float64Array([0, 1, 0]);
+ const { dLogits } = softmaxCrossEntropy(logits, labels);
+ let sum = 0;
+ for (let i = 0; i < 3; i++) sum += dLogits[i]!;
+ expect(sum).toBeCloseTo(0, 5);
+ });
+});
+
+describe("mseLoss", () => {
+ it("loss is 0 for perfect prediction", () => {
+ const y = new Float64Array([1, 2, 3]);
+ const { loss } = mseLoss(y, y);
+ expect(loss).toBeCloseTo(0);
+ });
+});
+
+describe("adamStep", () => {
+ it("matches the bias-corrected update for a constant unit gradient", () => {
+ const params = new Float64Array([1.0]);
+ const grads = new Float64Array([1.0]);
+ let state = initAdam(1);
+ for (let i = 0; i < 100; i++) {
+ state = adamStep(params, grads, state);
+ }
+ // Both bias-corrected moments are 1, so epsilon makes each update
+ // slightly smaller than the learning rate.
+ expect(params[0]).toBeCloseTo(1 - (100 * 0.001) / (1 + 1e-8), 12);
+ expect(state.t).toBe(100);
+ });
+});
+
+describe("heInit", () => {
+ it("returns correct size", () => {
+ const w = heInit(10, 5);
+ expect(w.length).toBe(50);
+ });
+
+ it("mean is roughly 0", () => {
+ const w = heInit(100, 100);
+ let sum = 0;
+ for (let i = 0; i < w.length; i++) sum += w[i]!;
+ expect(Math.abs(sum / w.length)).toBeLessThan(0.1);
+ });
+});
diff --git a/tests/ml/random_forest.test.ts b/tests/ml/random_forest.test.ts
new file mode 100644
index 000000000..a1ebef184
--- /dev/null
+++ b/tests/ml/random_forest.test.ts
@@ -0,0 +1,47 @@
+/**
+ * Tests for src/ml/random_forest.ts
+ */
+import { describe, expect, it } from "bun:test";
+import { LCGRandom, fitRandomForest, predictRandomForest, r2Score } from "../../src/index.ts";
+
+describe("LCGRandom", () => {
+ it("returns values in [0, 1)", () => {
+ const rng = new LCGRandom(42);
+ for (let i = 0; i < 100; i++) {
+ const v = rng.next();
+ expect(v).toBeGreaterThanOrEqual(0);
+ expect(v).toBeLessThan(1);
+ }
+ });
+
+ it("nextInt returns values in [0, n)", () => {
+ const rng = new LCGRandom(12345);
+ for (let i = 0; i < 50; i++) {
+ const v = rng.nextInt(10);
+ expect(v).toBeGreaterThanOrEqual(0);
+ expect(v).toBeLessThan(10);
+ }
+ });
+});
+
+describe("fitRandomForest — regression", () => {
+ const n = 30;
+ const X: Float64Array[] = Array.from(
+ { length: n },
+ (_, i) => new Float64Array([i / 10, (i % 5) / 5]),
+ );
+ const y = Float64Array.from({ length: n }, (_, i) => (i / 10) * 2 + (i % 5) / 5);
+
+ it("trains and predicts without error", () => {
+ const model = fitRandomForest(X, y, { nEstimators: 20, maxDepth: 3, seed: 1 });
+ const preds = predictRandomForest(model, X);
+ expect(preds.length).toBe(n);
+ });
+
+ it("achieves r2 > 0.7 on training set", () => {
+ const model = fitRandomForest(X, y, { nEstimators: 50, maxDepth: 5, seed: 42 });
+ const preds = predictRandomForest(model, X);
+ const r2 = r2Score(y, preds);
+ expect(r2).toBeGreaterThan(0.7);
+ });
+});
diff --git a/tests/ml/svm.test.ts b/tests/ml/svm.test.ts
new file mode 100644
index 000000000..112f64a7e
--- /dev/null
+++ b/tests/ml/svm.test.ts
@@ -0,0 +1,43 @@
+/**
+ * Tests for src/ml/svm.ts
+ */
+import { describe, expect, it } from "bun:test";
+import { computeKernel, fitSVM, svmAccuracy, svmPredict } from "../../src/index.ts";
+
+describe("computeKernel", () => {
+ it("linear kernel is dot product", () => {
+ const x1 = new Float64Array([1, 2]);
+ const x2 = new Float64Array([3, 4]);
+ const cfg = { type: "linear" as const, gamma: 1, degree: 2, coef0: 0 };
+ expect(computeKernel(x1, x2, cfg)).toBeCloseTo(11);
+ });
+
+ it("rbf kernel is 1 for identical points", () => {
+ const x = new Float64Array([1, 2, 3]);
+ const cfg = { type: "rbf" as const, gamma: 0.5, degree: 2, coef0: 0 };
+ expect(computeKernel(x, x, cfg)).toBeCloseTo(1);
+ });
+});
+
+describe("fitSVM — linearly separable", () => {
+ const X: Float64Array[] = [
+ new Float64Array([-2]),
+ new Float64Array([-1]),
+ new Float64Array([1]),
+ new Float64Array([2]),
+ ];
+ const y = new Float64Array([-1, -1, 1, 1]);
+
+ it("achieves 100% accuracy on training set", () => {
+ const model = fitSVM(X, y, 1.0, { type: "linear", gamma: 1, degree: 2, coef0: 0 }, 100);
+ const acc = svmAccuracy(model, X, y);
+ expect(acc).toBeCloseTo(1.0, 1);
+ });
+
+ it("predicts correct labels", () => {
+ const model = fitSVM(X, y, 1.0, { type: "linear", gamma: 1, degree: 2, coef0: 0 }, 100);
+ const preds = svmPredict(model, [new Float64Array([-1.5]), new Float64Array([1.5])]);
+ expect(preds[0]).toBe(-1);
+ expect(preds[1]).toBe(1);
+ });
+});
diff --git a/tests/stats/acf_pacf.test.ts b/tests/stats/acf_pacf.test.ts
new file mode 100644
index 000000000..411b6e00e
--- /dev/null
+++ b/tests/stats/acf_pacf.test.ts
@@ -0,0 +1,508 @@
+/**
+ * Tests for src/stats/acf_pacf.ts
+ *
+ * Covers autocorr, acf, pacf, ccf, durbinWatson, ljungBox, boxPierce.
+ * Numerical references cross-checked against statsmodels / scipy.
+ */
+import { describe, expect, it } from "bun:test";
+import fc from "fast-check";
+import {
+ Series,
+ acf,
+ autocorr,
+ boxPierce,
+ ccf,
+ durbinWatson,
+ ljungBox,
+ pacf,
+} from "../../src/index.ts";
+
+// ─── helpers ──────────────────────────────────────────────────────────────────
+
+function round(v: number, dp = 6): number {
+ const f = 10 ** dp;
+ return Math.round(v * f) / f;
+}
+
+// AR(1) process: x[t] = phi * x[t-1] + noise (deterministic, no noise)
+function ar1(phi: number, n: number): number[] {
+ const xs: number[] = [1];
+ for (let i = 1; i < n; i++) {
+ xs.push(phi * (xs[i - 1] ?? 0));
+ }
+ return xs;
+}
+
+// White noise from a simple LCG seed
+function lcgNoise(n: number, seed = 42): number[] {
+ let s = seed;
+ const out: number[] = [];
+ for (let i = 0; i < n; i++) {
+ s = (s * 1664525 + 1013904223) & 0x7fffffff;
+ out.push(s / 0x7fffffff - 0.5);
+ }
+ return out;
+}
+
+// ─── autocorr ────────────────────────────────────────────────────────────────
+
+describe("autocorr", () => {
+ it("returns 1.0 at lag 0", () => {
+ const x = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10];
+ expect(autocorr(x, 0)).toBeCloseTo(1.0, 5);
+ });
+
+ it("returns 1.0 for perfectly correlated shifted copies (linear series)", () => {
+ // x = [1,2,...,10]; lag=1 gives almost perfect correlation
+ const x = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10];
+ expect(autocorr(x, 1)).toBeCloseTo(1.0, 2);
+ });
+
+ it("returns -1 for alternating ±1 series at lag 1", () => {
+ const x = [1, -1, 1, -1, 1, -1, 1, -1, 1, -1];
+ expect(autocorr(x, 1)).toBeCloseTo(-1.0, 5);
+ });
+
+ it("returns NaN for series too short for given lag", () => {
+ expect(Number.isNaN(autocorr([1, 2], 2))).toBe(true);
+ });
+
+ it("accepts Series input", () => {
+ const s = new Series({ data: [1, 2, 3, 4, 5, 6] });
+ const arr = [1, 2, 3, 4, 5, 6];
+ expect(autocorr(s, 1)).toBeCloseTo(autocorr(arr, 1), 8);
+ });
+
+ it("property: |autocorr| ≤ 1 for any series", () => {
+ fc.assert(
+ fc.property(
+ fc.array(fc.float({ noNaN: true, noDefaultInfinity: true }), {
+ minLength: 5,
+ maxLength: 30,
+ }),
+ (xs) => {
+ const r = autocorr(xs, 1);
+ if (!Number.isNaN(r)) {
+ expect(Math.abs(r)).toBeLessThanOrEqual(1 + 1e-9);
+ }
+ },
+ ),
+ { numRuns: 200 },
+ );
+ });
+});
+
+// ─── acf ─────────────────────────────────────────────────────────────────────
+
+describe("acf", () => {
+ it("lag-0 coefficient is always 1.0", () => {
+ const x = [3, 1, 4, 1, 5, 9, 2, 6];
+ const result = acf(x);
+ expect(result.acf[0]).toBe(1.0);
+ });
+
+ it("lags array starts at 0", () => {
+ const x = [1, 2, 3, 4, 5, 6, 7, 8];
+ const result = acf(x, { nlags: 3 });
+ expect(result.lags).toEqual([0, 1, 2, 3]);
+ });
+
+ it("respects nlags parameter", () => {
+ const x = Array.from({ length: 20 }, (_, i) => i);
+ const result = acf(x, { nlags: 5 });
+ expect(result.acf.length).toBe(6); // lags 0..5
+ });
+
+ it("linear series matches finite-sample autocovariance normalization", () => {
+ const x = Array.from({ length: 20 }, (_, i) => i);
+ const result = acf(x, { nlags: 5 });
+ // statsmodels 0.15.0: acf(np.arange(20), nlags=5, fft=False).
+ const expected = [
+ 1, 0.85, 0.7015037593984962, 0.556015037593985, 0.4150375939849624, 0.2800751879699248,
+ ];
+ expected.forEach((value, k) => expect(result.acf[k]).toBeCloseTo(value, 12));
+ });
+
+ it("alternating series has negative ACF at odd lags", () => {
+ const x = Array.from({ length: 20 }, (_, i) => (i % 2 === 0 ? 1 : -1));
+ const result = acf(x, { nlags: 3 });
+ expect(result.acf[1]).toBeLessThan(0);
+ expect(result.acf[2]).toBeGreaterThan(0); // lag 2 positive
+ });
+
+ it("returns CI when alpha is specified", () => {
+ const x = lcgNoise(50);
+ const result = acf(x, { nlags: 5, alpha: 0.05 });
+ expect(result.confint).toBeDefined();
+ expect(result.confint?.length).toBe(6);
+ // lag-0 CI is always [1, 1]
+ const ci0 = result.confint?.[0];
+ expect(ci0?.[0]).toBe(1);
+ expect(ci0?.[1]).toBe(1);
+ });
+
+ it("CI bounds are ordered [lower, upper] for lags ≥ 1", () => {
+ const x = lcgNoise(40);
+ const result = acf(x, { nlags: 4, alpha: 0.05 });
+ for (let k = 1; k <= 4; k++) {
+ const ci = result.confint?.[k];
+ if (ci !== undefined) {
+ expect(ci[0]).toBeLessThanOrEqual(ci[1]);
+ }
+ }
+ });
+
+ it("no CI when alpha is omitted", () => {
+ const x = [1, 2, 3, 4, 5];
+ const result = acf(x);
+ expect(result.confint).toBeUndefined();
+ });
+
+ it("known AR(1) with φ=0.8 — ACF(1) ≈ 0.8", () => {
+ // For AR(1) with large n, ACF(k) ≈ φ^k
+ const x = ar1(0.8, 200);
+ const result = acf(x, { nlags: 3 });
+ // Expected: ACF(1) ≈ 0.8, ACF(2) ≈ 0.64, ACF(3) ≈ 0.512
+ expect(result.acf[1]).toBeGreaterThan(0.7);
+ expect(result.acf[2]).toBeGreaterThan(0.55);
+ });
+
+ it("property: ACF values are in [-1, 1]", () => {
+ fc.assert(
+ fc.property(
+ fc.array(fc.float({ noNaN: true, noDefaultInfinity: true }), {
+ minLength: 5,
+ maxLength: 40,
+ }),
+ (xs) => {
+ const result = acf(xs, { nlags: 3 });
+ for (const r of result.acf) {
+ expect(Math.abs(r)).toBeLessThanOrEqual(1 + 1e-9);
+ }
+ },
+ ),
+ { numRuns: 200 },
+ );
+ });
+});
+
+// ─── pacf ────────────────────────────────────────────────────────────────────
+
+describe("pacf", () => {
+ it("lag-0 PACF is always 1.0", () => {
+ const x = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10];
+ const result = pacf(x);
+ expect(result.pacf[0]).toBe(1.0);
+ });
+
+ it("lags array starts at 0", () => {
+ const x = Array.from({ length: 20 }, (_, i) => i);
+ const result = pacf(x, { nlags: 3 });
+ expect(result.lags).toEqual([0, 1, 2, 3]);
+ });
+
+ it("AR(1) φ=0.7: PACF[1] ≈ 0.7, PACF[k>1] ≈ 0", () => {
+ const noise = lcgNoise(200, 7);
+ const x: number[] = [noise[0] ?? 0];
+ for (let i = 1; i < 200; i++) {
+ x.push(0.7 * (x[i - 1] ?? 0) + (noise[i] ?? 0) * 0.2);
+ }
+ const result = pacf(x, { nlags: 4 });
+ // PACF[1] should be close to 0.7
+ expect(result.pacf[1]).toBeGreaterThan(0.55);
+ expect(result.pacf[1]).toBeLessThan(0.85);
+ // PACF[2..4] should be close to 0 for a true AR(1)
+ expect(Math.abs(result.pacf[2] ?? 0)).toBeLessThan(0.25);
+ expect(Math.abs(result.pacf[3] ?? 0)).toBeLessThan(0.25);
+ });
+
+ it("returns CI when alpha is specified", () => {
+ const x = lcgNoise(50);
+ const result = pacf(x, { nlags: 4, alpha: 0.05 });
+ expect(result.confint).toBeDefined();
+ expect(result.confint?.length).toBe(5);
+ });
+
+ it("CI bounds ordered [lower, upper]", () => {
+ const x = lcgNoise(40);
+ const result = pacf(x, { nlags: 4, alpha: 0.05 });
+ for (const ci of result.confint ?? []) {
+ expect(ci[0]).toBeLessThanOrEqual(ci[1]);
+ }
+ });
+
+ it("no CI when alpha is omitted", () => {
+ const x = [1, 2, 3, 4, 5];
+ expect(pacf(x).confint).toBeUndefined();
+ });
+
+ it("property: PACF[0] = 1 always", () => {
+ fc.assert(
+ fc.property(
+ fc.array(fc.float({ noNaN: true, noDefaultInfinity: true }), {
+ minLength: 6,
+ maxLength: 40,
+ }),
+ (xs) => {
+ const result = pacf(xs);
+ expect(result.pacf[0]).toBe(1.0);
+ },
+ ),
+ { numRuns: 200 },
+ );
+ });
+});
+
+// ─── ccf ─────────────────────────────────────────────────────────────────────
+
+describe("ccf", () => {
+ it("CCF of identical series at lag 0 is 1.0", () => {
+ const x = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10];
+ const result = ccf(x, x, { nlags: 3, positiveOnly: true });
+ expect(result.acf[0]).toBeCloseTo(1.0, 5);
+ });
+
+ it("detects a known lag: y = shift(x, 2)", () => {
+ const x = [0, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 0, 0];
+ const y = [1, 2, 3, 4, 5, 6, 7, 8, 9, 0, 0, 0, 0]; // x shifted left by 2
+ const result = ccf(x, y, { nlags: 4, positiveOnly: false });
+ // Maximum CCF should be near lag -2 (x leads y by 2)
+ // or equivalently, CCF(k=2) for y(t+2) vs x(t)
+ expect(result.lags.length).toBeGreaterThan(0);
+ });
+
+ it("returns CI when alpha specified", () => {
+ const x = lcgNoise(30);
+ const y = lcgNoise(30, 99);
+ const result = ccf(x, y, { nlags: 3, alpha: 0.05 });
+ expect(result.confint).toBeDefined();
+ });
+
+ it("positiveOnly returns only non-negative lags", () => {
+ const x = lcgNoise(20);
+ const y = lcgNoise(20, 11);
+ const result = ccf(x, y, { nlags: 3, positiveOnly: true });
+ for (const lag of result.lags) {
+ expect(lag).toBeGreaterThanOrEqual(0);
+ }
+ });
+
+ it("two-sided returns negative lags", () => {
+ const x = lcgNoise(20);
+ const y = lcgNoise(20, 22);
+ const result = ccf(x, y, { nlags: 3, positiveOnly: false });
+ expect(result.lags.some((l) => l < 0)).toBe(true);
+ });
+});
+
+// ─── durbinWatson ─────────────────────────────────────────────────────────────
+
+describe("durbinWatson", () => {
+ it("returns ~2 for random noise (no autocorrelation)", () => {
+ const e = lcgNoise(100);
+ const dw = durbinWatson(e);
+ expect(dw).toBeGreaterThan(1.5);
+ expect(dw).toBeLessThan(2.5);
+ });
+
+ it("returns ~0 for strongly positively autocorrelated residuals", () => {
+ // Residuals that are all the same sign (strongly autocorrelated)
+ const e = Array.from({ length: 20 }, (_, i) => 1 + i * 0.001);
+ const dw = durbinWatson(e);
+ expect(dw).toBeLessThan(0.1);
+ });
+
+ it("returns ~4 for alternating residuals (negative autocorrelation)", () => {
+ const e = Array.from({ length: 20 }, (_, i) => (i % 2 === 0 ? 1 : -1));
+ const dw = durbinWatson(e);
+ expect(dw).toBeCloseTo((4 * (e.length - 1)) / e.length, 12);
+ });
+
+ it("returns 2 for all-zero residuals", () => {
+ const e = Array.from({ length: 10 }, () => 0);
+ expect(durbinWatson(e)).toBe(2);
+ });
+
+ it("returns NaN for single-element input", () => {
+ expect(Number.isNaN(durbinWatson([1]))).toBe(true);
+ });
+
+ it("accepts Series input", () => {
+ const e = [1, -1, 1, -1, 1, -1, 1, -1];
+ const s = new Series({ data: e });
+ expect(durbinWatson(s)).toBeCloseTo(durbinWatson(e), 8);
+ });
+
+ it("property: DW ∈ [0, 4]", () => {
+ fc.assert(
+ fc.property(
+ fc.array(fc.float({ noNaN: true, noDefaultInfinity: true }), {
+ minLength: 2,
+ maxLength: 50,
+ }),
+ (xs) => {
+ const dw = durbinWatson(xs);
+ if (!Number.isNaN(dw)) {
+ expect(dw).toBeGreaterThanOrEqual(0 - 1e-9);
+ expect(dw).toBeLessThanOrEqual(4 + 1e-9);
+ }
+ },
+ ),
+ { numRuns: 300 },
+ );
+ });
+});
+
+// ─── ljungBox ────────────────────────────────────────────────────────────────
+
+describe("ljungBox", () => {
+ it("high p-value for white noise", () => {
+ const x = lcgNoise(100);
+ const result = ljungBox(x);
+ // White noise: should usually not reject H0
+ expect(result.pvalue[0]).toBeGreaterThan(0.0);
+ });
+
+ it("very low p-value for AR(1) process (structured autocorrelation)", () => {
+ const x = ar1(0.9, 100);
+ const result = ljungBox(x, { lags: [5] });
+ expect(result.pvalue[0]).toBeLessThan(0.01);
+ });
+
+ it("statistic is non-negative", () => {
+ const x = lcgNoise(50);
+ const result = ljungBox(x, { lags: [3, 5, 8] });
+ for (const q of result.statistic) {
+ expect(q).toBeGreaterThanOrEqual(0);
+ }
+ });
+
+ it("lags array matches requested lags", () => {
+ const x = lcgNoise(40);
+ const result = ljungBox(x, { lags: [1, 3, 6] });
+ expect(result.lags).toEqual([1, 3, 6]);
+ expect(result.statistic.length).toBe(3);
+ expect(result.pvalue.length).toBe(3);
+ });
+
+ it("when lags is a number h, returns h p-values for lags 1..h", () => {
+ const x = lcgNoise(30);
+ const result = ljungBox(x, { lags: 5 });
+ expect(result.lags).toEqual([1, 2, 3, 4, 5]);
+ });
+
+ it("pvalue is NaN when df ≤ 0 (modelDf ≥ lag)", () => {
+ const x = lcgNoise(30);
+ const result = ljungBox(x, { lags: [1], modelDf: 1 });
+ expect(Number.isNaN(result.pvalue[0])).toBe(true);
+ });
+
+ it("Ljung-Box Q > Box-Pierce Q for same data (finite-sample correction)", () => {
+ const x = ar1(0.5, 50);
+ const lb = ljungBox(x, { lags: [5] });
+ const bp = boxPierce(x, { lags: [5] });
+ // LB statistic ≥ BP statistic (LB has larger finite-sample correction)
+ expect(lb.statistic[0] ?? 0).toBeGreaterThan((bp.statistic[0] ?? 0) * 0.9);
+ });
+
+ it("property: statistic ≥ 0 for any series", () => {
+ fc.assert(
+ fc.property(
+ fc.array(fc.float({ noNaN: true, noDefaultInfinity: true }), {
+ minLength: 10,
+ maxLength: 50,
+ }),
+ (xs) => {
+ const result = ljungBox(xs, { lags: [3] });
+ expect(result.statistic[0] ?? 0).toBeGreaterThanOrEqual(0);
+ },
+ ),
+ { numRuns: 200 },
+ );
+ });
+});
+
+// ─── boxPierce ────────────────────────────────────────────────────────────────
+
+describe("boxPierce", () => {
+ it("statistic is non-negative", () => {
+ const x = lcgNoise(50);
+ const result = boxPierce(x, { lags: [4] });
+ expect(result.statistic[0] ?? 0).toBeGreaterThanOrEqual(0);
+ });
+
+ it("high p-value for white noise", () => {
+ const x = lcgNoise(200);
+ const result = boxPierce(x);
+ expect(result.pvalue[0] ?? 0).toBeGreaterThan(0);
+ });
+
+ it("very low p-value for strong autocorrelation", () => {
+ const x = ar1(0.95, 100);
+ const result = boxPierce(x, { lags: [10] });
+ expect(result.pvalue[0] ?? 0).toBeLessThan(0.001);
+ });
+
+ it("lags array matches requested lags", () => {
+ const x = lcgNoise(40);
+ const result = boxPierce(x, { lags: [2, 5] });
+ expect(result.lags).toEqual([2, 5]);
+ });
+
+ it("monotone Q: Q(h+1) ≥ Q(h) for series with autocorrelation", () => {
+ const x = ar1(0.7, 80);
+ const result = boxPierce(x, { lags: [1, 2, 3, 4, 5] });
+ for (let i = 1; i < result.statistic.length; i++) {
+ // Q is cumulative: adding one more lag adds r_k^2 ≥ 0
+ expect(result.statistic[i] ?? 0).toBeGreaterThanOrEqual(
+ (result.statistic[i - 1] ?? 0) - 1e-9,
+ );
+ }
+ });
+
+ it("known numerical check — constant series gives Q=0", () => {
+ // All ACF values at lag ≥ 1 are NaN or 0 for a constant series → Q = 0
+ const x = Array.from({ length: 10 }, () => 5);
+ const result = boxPierce(x, { lags: [3] });
+ expect(result.statistic[0] ?? 0).toBe(0);
+ });
+});
+
+// ─── known values (cross-checked against statsmodels) ────────────────────────
+
+describe("known values (statsmodels reference)", () => {
+ // statsmodels reference values:
+ // import statsmodels.tsa.stattools as sm
+ // x = [1, 2, 3, 2, 1, 2, 3, 2, 1, 2]
+ // sm.acf(x, nlags=4, fft=False)
+ // => [1.0, -0.0020408163265305938, -0.8, -0.006122448979591823, 0.5836734693877551]
+ const x = [1, 2, 3, 2, 1, 2, 3, 2, 1, 2];
+
+ it("ACF matches statsmodels for x=[1,2,3,2,1,2,3,2,1,2]", () => {
+ const result = acf(x, { nlags: 4 });
+ expect(round(result.acf[0] ?? 0, 4)).toBe(1.0);
+ expect(round(result.acf[1] ?? 0, 4)).toBe(round(-0.0020408163265305938, 4));
+ expect(round(result.acf[2] ?? 0, 4)).toBe(round(-0.8, 4));
+ });
+
+ it("Pearson autocorrelation uses separate shifted means", () => {
+ const acfVal = acf(x, { nlags: 1 }).acf[1] ?? 0;
+ // Pearson correlation centers each shifted slice separately; ACF centers
+ // the entire original sample. Their values need not have the same sign.
+ expect(autocorr(x, 1)).toBeCloseTo(0, 12);
+ expect(acfVal).toBeCloseTo(-0.0020408163265305938, 12);
+ });
+
+ it("Durbin-Watson for [1,-1,1,-1,...] is close to 4", () => {
+ const e = [1, -1, 1, -1, 1, -1, 1, -1, 1, -1];
+ expect(durbinWatson(e)).toBeCloseTo(3.6, 12);
+ });
+
+ it("Ljung-Box Q for x=[1,2,3,...,10], lag=1 is finite and positive", () => {
+ const xs = Array.from({ length: 10 }, (_, i) => i + 1);
+ const result = ljungBox(xs, { lags: [1] });
+ const q = result.statistic[0] ?? 0;
+ expect(q).toBeGreaterThan(0);
+ expect(Number.isFinite(q)).toBe(true);
+ });
+});
diff --git a/tests/stats/arima.test.ts b/tests/stats/arima.test.ts
new file mode 100644
index 000000000..57139eead
--- /dev/null
+++ b/tests/stats/arima.test.ts
@@ -0,0 +1,350 @@
+/**
+ * Tests for src/stats/arima.ts
+ *
+ * Covers ARIMA(p,d,q) estimation, in-sample fitted values, multi-step
+ * forecasting, prediction intervals, AIC/BIC, and edge cases.
+ * Numerical references cross-checked against statsmodels.
+ */
+import { describe, expect, it } from "bun:test";
+import * as fc from "fast-check";
+import { ARIMAModel, fitArima } from "../../src/index.ts";
+
+// ─── Helpers ───────────────────────────────────────────────────────────────────
+
+/** Generate a deterministic AR(1) series: x_t = phi * x_{t-1} + noise */
+function ar1Series(phi: number, n: number, noiseAmp = 0.1): number[] {
+ const xs: number[] = [1.0];
+ // LCG for reproducibility
+ let seed = 42;
+ const rand = (): number => {
+ seed = (seed * 1664525 + 1013904223) & 0x7fffffff;
+ return (seed / 0x7fffffff - 0.5) * 2 * noiseAmp;
+ };
+ for (let i = 1; i < n; i++) {
+ xs.push(phi * (xs[i - 1] ?? 0) + rand());
+ }
+ return xs;
+}
+
+/** Mean absolute error */
+function mae(a: readonly number[], b: readonly number[]): number {
+ let s = 0;
+ for (let i = 0; i < a.length; i++) {
+ s += Math.abs((a[i] ?? 0) - (b[i] ?? 0));
+ }
+ return s / a.length;
+}
+
+// ─── ARIMAModel construction ────────────────────────────────────────────────────
+
+describe("ARIMAModel construction", () => {
+ it("stores default p=1,d=0,q=0", () => {
+ const m = new ARIMAModel();
+ expect(m.p).toBe(1);
+ expect(m.d).toBe(0);
+ expect(m.q).toBe(0);
+ });
+
+ it("stores custom p,d,q", () => {
+ const m = new ARIMAModel({ p: 2, d: 1, q: 1 });
+ expect(m.p).toBe(2);
+ expect(m.d).toBe(1);
+ expect(m.q).toBe(1);
+ });
+
+ it("clamps negative orders to 0", () => {
+ const m = new ARIMAModel({ p: -1, d: -2, q: -3 });
+ expect(m.p).toBe(0);
+ expect(m.d).toBe(0);
+ expect(m.q).toBe(0);
+ });
+});
+
+// ─── fit() ─────────────────────────────────────────────────────────────────────
+
+describe("fit – AR(1)", () => {
+ const y = ar1Series(0.7, 200, 0.05);
+
+ it("returns arCoeffs of length p", () => {
+ const { arCoeffs } = new ARIMAModel({ p: 1, d: 0, q: 0 }).fit(y);
+ expect(arCoeffs.length).toBe(1);
+ });
+
+ it("AR(1) uses every available lagged observation and matches OLS", () => {
+ const { arCoeffs, intercept } = new ARIMAModel({ p: 1, d: 0, q: 0 }).fit(y);
+ // statsmodels 0.15.0: OLS(y[1:], add_constant(y[:-1])).fit().params.
+ expect(arCoeffs[0]).toBeCloseTo(0.6746405572472494, 10);
+ expect(intercept).toBeCloseTo(-0.00022833154176704845, 10);
+ });
+
+ it("fittedValues length equals series length", () => {
+ const result = new ARIMAModel({ p: 1, d: 0, q: 0 }).fit(y);
+ expect(result.fittedValues.length).toBe(y.length);
+ });
+
+ it("residuals length equals series length minus d", () => {
+ const result = new ARIMAModel({ p: 1, d: 0, q: 0 }).fit(y);
+ expect(result.residuals.length).toBe(y.length);
+ });
+
+ it("sigma2 is positive", () => {
+ const { sigma2 } = new ARIMAModel({ p: 1, d: 0, q: 0 }).fit(y);
+ expect(sigma2).toBeGreaterThan(0);
+ });
+
+ it("AIC < 0 for a well-fit model (negative log-likelihood dominates)", () => {
+ const { aic } = new ARIMAModel({ p: 1, d: 0, q: 0 }).fit(y);
+ expect(typeof aic).toBe("number");
+ expect(Number.isFinite(aic)).toBe(true);
+ });
+
+ it("BIC >= AIC (more penalty per parameter for n > 8)", () => {
+ const { aic, bic } = new ARIMAModel({ p: 1, d: 0, q: 0 }).fit(y);
+ expect(bic).toBeGreaterThanOrEqual(aic);
+ });
+});
+
+describe("fit – MA(1)", () => {
+ // MA(1): x_t = noise_t + theta * noise_{t-1}
+ it("fits MA(1) without error", () => {
+ const y2 = ar1Series(0.5, 100, 0.2);
+ const result = new ARIMAModel({ p: 0, d: 0, q: 1 }).fit(y2);
+ expect(result.maCoeffs.length).toBe(1);
+ });
+});
+
+describe("fit – ARMA(1,1)", () => {
+ it("fits ARMA(1,1) and returns finite AIC", () => {
+ const y3 = ar1Series(0.5, 150, 0.1);
+ const result = new ARIMAModel({ p: 1, d: 0, q: 1 }).fit(y3);
+ expect(result.arCoeffs.length).toBe(1);
+ expect(result.maCoeffs.length).toBe(1);
+ expect(Number.isFinite(result.aic)).toBe(true);
+ });
+});
+
+describe("fit – ARIMA(1,1,0)", () => {
+ // Random walk + drift
+ const rw: number[] = [100];
+ let s = 99;
+ for (let i = 1; i < 150; i++) {
+ s = (s * 1664525 + 1013904223) & 0x7fffffff;
+ rw.push(rw[i - 1]! + 0.5 + (s / 0x7fffffff - 0.5) * 2);
+ }
+
+ it("fittedValues length equals original n (not differenced n)", () => {
+ const result = new ARIMAModel({ p: 1, d: 1, q: 0 }).fit(rw);
+ expect(result.fittedValues.length).toBe(rw.length);
+ });
+
+ it("residuals length equals differenced series (n - d)", () => {
+ const result = new ARIMAModel({ p: 1, d: 1, q: 0 }).fit(rw);
+ expect(result.residuals.length).toBe(rw.length - 1);
+ });
+
+ it("AIC is finite", () => {
+ const result = new ARIMAModel({ p: 1, d: 1, q: 0 }).fit(rw);
+ expect(Number.isFinite(result.aic)).toBe(true);
+ });
+});
+
+describe("fit – ARIMA(0,0,0)", () => {
+ it("fits with just an intercept", () => {
+ const y4 = [1, 2, 3, 4, 3, 2, 1, 2, 3];
+ const result = new ARIMAModel({ p: 0, d: 0, q: 0 }).fit(y4);
+ expect(result.arCoeffs.length).toBe(0);
+ expect(result.maCoeffs.length).toBe(0);
+ expect(Number.isFinite(result.intercept)).toBe(true);
+ });
+});
+
+describe("fit – error on short series", () => {
+ it("throws RangeError if series too short", () => {
+ expect(() => new ARIMAModel({ p: 2, d: 1, q: 1 }).fit([1, 2, 3])).toThrow(RangeError);
+ });
+});
+
+// ─── forecast() ────────────────────────────────────────────────────────────────
+
+describe("forecast – AR(1)", () => {
+ const y5 = ar1Series(0.6, 100, 0.05);
+ const model = new ARIMAModel({ p: 1, d: 0, q: 0 });
+ model.fit(y5);
+
+ it("returns correct number of steps", () => {
+ const fc5 = model.forecast(5);
+ expect(fc5.forecast.length).toBe(5);
+ expect(fc5.lower.length).toBe(5);
+ expect(fc5.upper.length).toBe(5);
+ expect(fc5.stderr.length).toBe(5);
+ });
+
+ it("lower < forecast < upper for all steps", () => {
+ const fc3 = model.forecast(3);
+ for (let h = 0; h < 3; h++) {
+ expect(fc3.lower[h] ?? 0).toBeLessThan(fc3.forecast[h] ?? 0);
+ expect(fc3.upper[h] ?? 0).toBeGreaterThan(fc3.forecast[h] ?? 0);
+ }
+ });
+
+ it("stderr increases monotonically for AR(1) with |phi| < 1", () => {
+ const fc4 = model.forecast(4);
+ for (let h = 1; h < 4; h++) {
+ expect(fc4.stderr[h] ?? 0).toBeGreaterThanOrEqual(fc4.stderr[h - 1] ?? 0);
+ }
+ });
+
+ it("step-1 stderr ≈ sqrt(sigma2) for AR(1) (sigma * psi_0 = sigma)", () => {
+ const fitResult = model.fit(y5);
+ const fc1 = model.forecast(1);
+ expect(fc1.stderr[0] ?? 0).toBeCloseTo(Math.sqrt(fitResult.sigma2), 3);
+ });
+
+ it("default steps=1 works", () => {
+ expect(model.forecast().forecast.length).toBe(1);
+ });
+
+ it("throws if called before fit", () => {
+ const m2 = new ARIMAModel({ p: 1 });
+ expect(() => m2.forecast(1)).toThrow();
+ });
+
+ it("throws on steps < 1", () => {
+ expect(() => model.forecast(0)).toThrow(RangeError);
+ });
+});
+
+describe("forecast – ARIMA(0,1,0) (random walk)", () => {
+ const rw2: number[] = [10];
+ for (let i = 1; i < 80; i++) {
+ rw2.push(rw2[i - 1]! + 0.1);
+ }
+ const model2 = new ARIMAModel({ p: 0, d: 1, q: 0 });
+ model2.fit(rw2);
+
+ it("forecast[0] ≈ last observed + drift", () => {
+ const fc = model2.forecast(1);
+ expect(typeof fc.forecast[0]).toBe("number");
+ expect(Number.isFinite(fc.forecast[0] ?? Number.NaN)).toBe(true);
+ });
+
+ it("stderr grows with horizon for I(1)", () => {
+ const fc = model2.forecast(5);
+ expect(fc.stderr[4] ?? 0).toBeGreaterThan(fc.stderr[0] ?? 0);
+ });
+});
+
+// ─── fitArima convenience function ─────────────────────────────────────────────
+
+describe("fitArima", () => {
+ it("returns ARIMAModel with forecast method", () => {
+ const model3 = fitArima(ar1Series(0.5, 80, 0.1), { p: 1, q: 0 });
+ const fc = model3.forecast(3);
+ expect(fc.forecast.length).toBe(3);
+ });
+
+ it("accepts Series via duck-typing", async () => {
+ const { Series } = await import("../../src/index.ts");
+ const s = new Series({ data: [1, 2, 3, 4, 3, 2, 1, 2, 3, 4, 3, 2, 1, 2, 3, 4] });
+ const model4 = fitArima(s, { p: 1, d: 0, q: 0 });
+ expect(model4.p).toBe(1);
+ expect(model4.forecast(2).forecast.length).toBe(2);
+ });
+});
+
+// ─── Property-based tests ───────────────────────────────────────────────────────
+
+describe("property tests", () => {
+ it("fitted value count always equals input length", () => {
+ fc.assert(
+ fc.property(
+ fc.array(fc.float({ noNaN: true, noDefaultInfinity: true, min: -100, max: 100 }), {
+ minLength: 20,
+ maxLength: 60,
+ }),
+ fc.integer({ min: 0, max: 2 }),
+ fc.integer({ min: 0, max: 2 }),
+ (y, p, q) => {
+ const model5 = new ARIMAModel({ p, d: 0, q });
+ try {
+ const res = model5.fit(y);
+ return res.fittedValues.length === y.length;
+ } catch {
+ return true; // short series allowed to throw
+ }
+ },
+ ),
+ );
+ });
+
+ it("forecast intervals always satisfy lower <= forecast <= upper", () => {
+ fc.assert(
+ fc.property(
+ fc.array(fc.float({ noNaN: true, noDefaultInfinity: true, min: -50, max: 50 }), {
+ minLength: 20,
+ maxLength: 50,
+ }),
+ (y) => {
+ const model6 = new ARIMAModel({ p: 1, d: 0, q: 0 });
+ try {
+ model6.fit(y);
+ const fc2 = model6.forecast(3);
+ for (let h = 0; h < 3; h++) {
+ if ((fc2.lower[h] ?? 0) > (fc2.forecast[h] ?? 0) + 1e-6) {
+ return false;
+ }
+ if ((fc2.upper[h] ?? 0) < (fc2.forecast[h] ?? 0) - 1e-6) {
+ return false;
+ }
+ }
+ return true;
+ } catch {
+ return true;
+ }
+ },
+ ),
+ );
+ });
+
+ it("sigma2 is always positive after fit", () => {
+ fc.assert(
+ fc.property(
+ fc.array(fc.float({ noNaN: true, noDefaultInfinity: true, min: -100, max: 100 }), {
+ minLength: 15,
+ maxLength: 40,
+ }),
+ (y) => {
+ const model7 = new ARIMAModel({ p: 1, d: 0, q: 0 });
+ try {
+ const res = model7.fit(y);
+ return res.sigma2 > 0;
+ } catch {
+ return true;
+ }
+ },
+ ),
+ );
+ });
+});
+
+// ─── AIC / BIC ordering ─────────────────────────────────────────────────────────
+
+describe("information criteria", () => {
+ it("higher-order model has smaller AIC on a long series with signal", () => {
+ const y6 = ar1Series(0.8, 300, 0.1);
+ const m1 = new ARIMAModel({ p: 1, d: 0, q: 0 }).fit(y6);
+ const m2 = new ARIMAModel({ p: 2, d: 0, q: 0 }).fit(y6);
+ // AR(2) should not have dramatically worse AIC on an AR(1) series
+ expect(m2.aic).not.toBeNaN();
+ expect(m1.aic).not.toBeNaN();
+ });
+
+ it("fitted values MAE on AR(1) is less than naive mean", () => {
+ const y7 = ar1Series(0.85, 200, 0.05);
+ const result = new ARIMAModel({ p: 1, d: 0, q: 0 }).fit(y7);
+ const mean = y7.reduce((s, v) => s + v, 0) / y7.length;
+ const naiveMae = mae(y7, new Array(y7.length).fill(mean));
+ const modelMae = mae(y7, result.fittedValues);
+ expect(modelMae).toBeLessThan(naiveMae);
+ });
+});
diff --git a/tests/stats/dlm.test.ts b/tests/stats/dlm.test.ts
new file mode 100644
index 000000000..e3b1c4136
--- /dev/null
+++ b/tests/stats/dlm.test.ts
@@ -0,0 +1,446 @@
+/**
+ * Tests for dlm.ts — Dynamic Linear Model (State-Space).
+ */
+import { describe, expect, test } from "bun:test";
+import * as fc from "fast-check";
+import {
+ DLM,
+ buildFourier,
+ buildLocalLevel,
+ buildLocalLinearTrend,
+ buildPolynomial,
+ buildRegression,
+ combineDLMs,
+} from "../../src/stats/dlm.ts";
+
+// ─── helpers ───────────────────────────────────────────────────────────────────
+function approx(a: number, b: number, tol = 1e-6): boolean {
+ return Math.abs(a - b) <= tol * (1 + Math.abs(b));
+}
+function approxVec(a: readonly number[], b: readonly number[], tol = 1e-4): boolean {
+ return a.length === b.length && a.every((v, i) => approx(v, b[i]!, tol));
+}
+
+// ─── Local-level model ─────────────────────────────────────────────────────────
+describe("DLM — local-level", () => {
+ const y = [1, 2, 3, 2, 3, 4, 3, 4, 5, 4];
+
+ test("filter returns correct dimensions", () => {
+ const dlm = DLM.localLevel({ sigmaObs: 1, sigmaLevel: 0.5 });
+ const res = dlm.filter(y);
+ expect(res.steps.length).toBe(y.length);
+ expect(res.filteredMeans.length).toBe(y.length);
+ expect(res.filteredMeans[0]?.length).toBe(1);
+ expect(res.forecastMeans.length).toBe(y.length);
+ });
+
+ test("log-likelihood is finite and negative", () => {
+ const dlm = DLM.localLevel({ sigmaObs: 1, sigmaLevel: 0.5 });
+ const res = dlm.filter(y);
+ expect(Number.isFinite(res.logLikelihood)).toBe(true);
+ expect(res.logLikelihood).toBeLessThan(0);
+ });
+
+ test("filtered means track the signal", () => {
+ const dlm = DLM.localLevel({ sigmaObs: 1, sigmaLevel: 0.5 });
+ const res = dlm.filter(y);
+ // Filtered means should be between 0 and 6 for this data
+ for (const m of res.filteredMeans) {
+ expect(m[0]).toBeGreaterThan(0);
+ expect(m[0]).toBeLessThan(6);
+ }
+ });
+
+ test("filter with missing observations", () => {
+ const dlm = DLM.localLevel({ sigmaObs: 1, sigmaLevel: 0.5 });
+ const yMissing = [1, null, 3, null, 5];
+ const res = dlm.filter(yMissing);
+ expect(res.steps.length).toBe(5);
+ // Steps with null observations have null innovation
+ expect(res.steps[1]?.innovation).toBeNull();
+ expect(res.steps[3]?.innovation).toBeNull();
+ // Steps with observations have non-null innovation
+ expect(res.steps[0]?.innovation).not.toBeNull();
+ });
+
+ test("higher observation noise → filtered closer to prior", () => {
+ const dlmLow = DLM.localLevel({ sigmaObs: 0.01, sigmaLevel: 1 });
+ const dlmHigh = DLM.localLevel({ sigmaObs: 100, sigmaLevel: 1 });
+ const resLow = dlmLow.filter(y);
+ const resHigh = dlmHigh.filter(y);
+ // With low obs noise, filtered mean ≈ observation
+ expect(Math.abs(resLow.filteredMeans[0]?.[0]! - y[0]!)).toBeLessThan(0.1);
+ // With high obs noise, filtered mean stays near prior (zero init, drifts slowly)
+ // Just check the values are different
+ expect(resLow.filteredMeans[5]?.[0]).not.toBeCloseTo(resHigh.filteredMeans[5]?.[0]!, 1);
+ });
+});
+
+// ─── Local-linear-trend model ──────────────────────────────────────────────────
+describe("DLM — local-linear-trend", () => {
+ // Linearly increasing data with noise
+ const y = Array.from({ length: 20 }, (_, t) => t + 0.5 * (Math.sin(t) * 0.1));
+
+ test("filter dimensions", () => {
+ const dlm = DLM.localLinearTrend({ sigmaObs: 0.5, sigmaLevel: 0.1, sigmaSlope: 0.01 });
+ const res = dlm.filter(y);
+ expect(res.filteredMeans.length).toBe(y.length);
+ expect(res.filteredMeans[0]?.length).toBe(2); // [level, slope]
+ });
+
+ test("slope converges near 1 for linear data", () => {
+ const dlm = DLM.localLinearTrend({ sigmaObs: 0.1, sigmaLevel: 0.01, sigmaSlope: 0.01 });
+ const res = dlm.filter(y);
+ const lastSlope = res.filteredMeans[y.length - 1]?.[1]!;
+ expect(lastSlope).toBeGreaterThan(0.5);
+ expect(lastSlope).toBeLessThan(2);
+ });
+
+ test("log-likelihood improves with better sigma", () => {
+ const dlmGood = DLM.localLinearTrend({ sigmaObs: 0.5, sigmaLevel: 0.1, sigmaSlope: 0.01 });
+ const dlmBad = DLM.localLinearTrend({ sigmaObs: 10, sigmaLevel: 10, sigmaSlope: 10 });
+ expect(dlmGood.filter(y).logLikelihood).toBeGreaterThan(dlmBad.filter(y).logLikelihood);
+ });
+});
+
+// ─── Smoother ──────────────────────────────────────────────────────────────────
+describe("DLM — RTS smoother", () => {
+ const y = [1, 2, 1.5, 3, 2.5, 4, 3.5, 5, 4.5, 6];
+
+ test("smoothed means have same length as data", () => {
+ const dlm = DLM.localLevel({ sigmaObs: 1, sigmaLevel: 0.5 });
+ const res = dlm.smooth(y);
+ expect(res.smoothedMeans.length).toBe(y.length);
+ expect(res.smoothedCovs.length).toBe(y.length);
+ });
+
+ test("smoother log-likelihood equals filter log-likelihood", () => {
+ const dlm = DLM.localLevel({ sigmaObs: 1, sigmaLevel: 0.5 });
+ const filterRes = dlm.filter(y);
+ const smoothRes = dlm.smooth(y);
+ expect(smoothRes.logLikelihood).toBeCloseTo(filterRes.logLikelihood, 6);
+ });
+
+ test("smoothed covariances ≤ filtered covariances (trace)", () => {
+ const dlm = DLM.localLinearTrend({ sigmaObs: 1, sigmaLevel: 0.5, sigmaSlope: 0.1 });
+ const res = dlm.smooth(y);
+ for (let t = 0; t < y.length - 1; t++) {
+ const filtTrace = res.filteredCovs[t]?.[0]?.[0]! + (res.filteredCovs[t]?.[1]?.[1] ?? 0);
+ const smTrace = res.smoothedCovs[t]?.[0]?.[0]! + (res.smoothedCovs[t]?.[1]?.[1] ?? 0);
+ // Smoother should not increase uncertainty
+ expect(smTrace).toBeLessThanOrEqual(filtTrace + 1e-6);
+ }
+ });
+
+ test("last smoothed = last filtered", () => {
+ const dlm = DLM.localLevel({ sigmaObs: 1, sigmaLevel: 0.5 });
+ const res = dlm.smooth(y);
+ const T = y.length;
+ expect(res.smoothedMeans[T - 1]?.[0]).toBeCloseTo(res.filteredMeans[T - 1]?.[0]!, 6);
+ });
+});
+
+// ─── Forecasting ───────────────────────────────────────────────────────────────
+describe("DLM — forecasting", () => {
+ const y = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10];
+
+ test("forecast returns h steps", () => {
+ const dlm = DLM.localLevel({ sigmaObs: 0.5, sigmaLevel: 0.2 });
+ const res = dlm.filter(y);
+ const fc = dlm.forecast(res, 5);
+ expect(fc.mean.length).toBe(5);
+ expect(fc.lower.length).toBe(5);
+ expect(fc.upper.length).toBe(5);
+ });
+
+ test("forecast mean > lower, < upper", () => {
+ const dlm = DLM.localLevel({ sigmaObs: 0.5, sigmaLevel: 0.2 });
+ const res = dlm.filter(y);
+ const fc = dlm.forecast(res, 5);
+ for (let i = 0; i < 5; i++) {
+ expect(fc.mean[i]?.[0]).toBeGreaterThan(fc.lower[i]!);
+ expect(fc.mean[i]?.[0]).toBeLessThan(fc.upper[i]!);
+ }
+ });
+
+ test("prediction intervals widen over forecast horizon", () => {
+ const dlm = DLM.localLevel({ sigmaObs: 0.5, sigmaLevel: 0.2 });
+ const res = dlm.filter(y);
+ const fc = dlm.forecast(res, 5);
+ // Intervals should widen (or stay same) as h increases
+ for (let i = 1; i < 5; i++) {
+ const widthPrev = fc.upper[i - 1]! - fc.lower[i - 1]!;
+ const widthCurr = fc.upper[i]! - fc.lower[i]!;
+ expect(widthCurr).toBeGreaterThanOrEqual(widthPrev - 1e-10);
+ }
+ });
+
+ test("local-linear-trend forecast extrapolates trend", () => {
+ const dlm = DLM.localLinearTrend({ sigmaObs: 0.1, sigmaLevel: 0.01, sigmaSlope: 0.001 });
+ const res = dlm.filter(y);
+ const fc = dlm.forecast(res, 3);
+ // Should forecast above 10
+ expect(fc.mean[0]?.[0]).toBeGreaterThan(9);
+ expect(fc.mean[2]?.[0]).toBeGreaterThan(fc.mean[0]?.[0]!);
+ });
+});
+
+// ─── Polynomial DLM ────────────────────────────────────────────────────────────
+describe("buildPolynomial", () => {
+ test("order=1 is equivalent to local-level", () => {
+ const spec = buildPolynomial(1, { sigmaObs: 1, sigmaState: 0.5 });
+ expect(spec.G).toEqual([[1]]);
+ expect(spec.F).toEqual([[1]]);
+ });
+
+ test("order=2 has 2×2 G with upper-triangular 1", () => {
+ const spec = buildPolynomial(2);
+ expect(spec.G[0]).toEqual([1, 1]);
+ expect(spec.G[1]).toEqual([0, 1]);
+ });
+
+ test("filter runs for order=3", () => {
+ const spec = buildPolynomial(3, { sigmaObs: 1, sigmaState: 0.1 });
+ const dlm = new DLM(spec);
+ const res = dlm.filter([1, 2, 3, 4, 5]);
+ expect(res.filteredMeans.length).toBe(5);
+ expect(res.filteredMeans[0]?.length).toBe(3);
+ });
+});
+
+// ─── Fourier seasonal DLM ──────────────────────────────────────────────────────
+describe("buildFourier", () => {
+ test("state dimension is 2*harmonics", () => {
+ const spec = buildFourier(12, 3);
+ expect(spec.G.length).toBe(6);
+ expect(spec.F[0]?.length).toBe(6);
+ });
+
+ test("rotation matrix property: G G' = I (for each 2×2 block)", () => {
+ const spec = buildFourier(12, 2);
+ const G = spec.G;
+ // Check first block is rotation
+ const c = G[0]?.[0]!;
+ const s = G[1]?.[0]!;
+ expect(c * c + s * s).toBeCloseTo(1, 6);
+ });
+
+ test("filter runs without error on seasonal data", () => {
+ // Synthetic seasonal data: period=4, 2 cycles
+ const y = [1, 2, 3, 2, 1, 2, 3, 2];
+ const spec = buildFourier(4, 2, { sigmaObs: 0.5, sigmaState: 0.1 });
+ const dlm = new DLM(spec);
+ const res = dlm.filter(y);
+ expect(res.steps.length).toBe(8);
+ expect(Number.isFinite(res.logLikelihood)).toBe(true);
+ });
+});
+
+// ─── buildRegression ───────────────────────────────────────────────────────────
+describe("buildRegression", () => {
+ test("spec has correct dimensions", () => {
+ const spec = buildRegression(3, { sigmaObs: 1, sigmaState: 0.001 });
+ expect(spec.G.length).toBe(3);
+ expect(spec.G[0]?.length).toBe(3);
+ expect(spec.F[0]?.length).toBe(3);
+ });
+});
+
+// ─── combineDLMs ───────────────────────────────────────────────────────────────
+describe("combineDLMs", () => {
+ test("combines state dimensions", () => {
+ const ll = buildLocalLevel({ sigmaObs: 1, sigmaLevel: 0.5 });
+ const llt = buildLocalLinearTrend({ sigmaObs: 1, sigmaLevel: 0.5, sigmaSlope: 0.1 });
+ const combined = combineDLMs(ll, llt);
+ // ll has 1 state, llt has 2 → combined has 3
+ expect(combined.G.length).toBe(3);
+ expect(combined.F[0]?.length).toBe(3);
+ expect(combined.W.length).toBe(3);
+ });
+
+ test("single spec is identity", () => {
+ const ll = buildLocalLevel({ sigmaObs: 1, sigmaLevel: 0.5 });
+ const combined = combineDLMs(ll);
+ expect(combined.G).toEqual(ll.G);
+ });
+
+ test("combined DLM can be filtered", () => {
+ const ll = buildLocalLevel({ sigmaObs: 1, sigmaLevel: 0.5 });
+ const fourier = buildFourier(4, 1, { sigmaObs: 1, sigmaState: 0.1 });
+ const combined = combineDLMs(ll, fourier);
+ const dlm = new DLM(combined);
+ const y = [1, 2, 3, 2, 1, 2, 3, 2];
+ const res = dlm.filter(y);
+ expect(res.steps.length).toBe(8);
+ expect(Number.isFinite(res.logLikelihood)).toBe(true);
+ });
+
+ test("throws on empty input", () => {
+ expect(() => combineDLMs()).toThrow();
+ });
+});
+
+// ─── MLE fitting ──────────────────────────────────────────────────────────────
+describe("DLM.fitMLE", () => {
+ test("returns a DLM instance", () => {
+ const dlm = DLM.localLevel({ sigmaObs: 1, sigmaLevel: 1 });
+ const y = [1, 2, 1.5, 2.5, 3, 2, 3.5, 3, 4, 3.5];
+ const fitted = dlm.fitMLE(y);
+ expect(fitted).toBeInstanceOf(DLM);
+ });
+
+ test("fitted log-likelihood ≥ initial log-likelihood", () => {
+ const dlm = DLM.localLevel({ sigmaObs: 5, sigmaLevel: 5 });
+ const y = [1, 2, 1.5, 2.5, 3, 2, 3.5, 3, 4, 3.5];
+ const fitted = dlm.fitMLE(y);
+ const llInitial = dlm.filter(y).logLikelihood;
+ const llFitted = fitted.filter(y).logLikelihood;
+ expect(llFitted).toBeGreaterThanOrEqual(llInitial - 0.1);
+ });
+});
+
+// ─── Discount factor filter ───────────────────────────────────────────────────
+describe("DLM.filterDiscount", () => {
+ test("discount=1 matches standard filter closely", () => {
+ const dlm = DLM.localLevel({ sigmaObs: 1, sigmaLevel: 0.5 });
+ const y = [1, 2, 3, 4, 5];
+ const _stdRes = dlm.filter(y);
+ // With discount=1, R_t = G C G' (no extra growth) — different from W-augmented
+ const discRes = dlm.filterDiscount(y, 1.0);
+ expect(discRes.steps.length).toBe(y.length);
+ expect(Number.isFinite(discRes.logLikelihood)).toBe(true);
+ });
+
+ test("lower discount factor → wider prediction intervals", () => {
+ const dlm = DLM.localLevel({ sigmaObs: 1, sigmaLevel: 0 });
+ const y = [1, 2, 3, 4, 5, 6, 7, 8];
+ const res095 = dlm.filterDiscount(y, 0.95);
+ const res099 = dlm.filterDiscount(y, 0.99);
+ // delta=0.95 has more uncertainty → larger Q
+ const q095 = res095.steps[3]?.forecastCov[0]?.[0]!;
+ const q099 = res099.steps[3]?.forecastCov[0]?.[0]!;
+ expect(q095).toBeGreaterThanOrEqual(q099 - 1e-6);
+ });
+});
+
+// ─── Factory static methods ────────────────────────────────────────────────────
+describe("DLM factory methods", () => {
+ test("DLM.localLevel matches buildLocalLevel dimensions", () => {
+ const dlm1 = DLM.localLevel({ sigmaObs: 1, sigmaLevel: 0.5 });
+ const spec = buildLocalLevel({ sigmaObs: 1, sigmaLevel: 0.5 });
+ // Both have G=[[1]], F=[[1]]
+ const dlm2 = new DLM(spec);
+ const y = [1, 2, 3];
+ const r1 = dlm1.filter(y);
+ const r2 = dlm2.filter(y);
+ expect(r1.filteredMeans[0]?.[0]).toBeCloseTo(r2.filteredMeans[0]?.[0]!, 4);
+ });
+
+ test("DLM.fourier factory", () => {
+ const dlm = DLM.fourier(12, 3, { sigmaObs: 1, sigmaState: 0.01 });
+ expect(dlm).toBeInstanceOf(DLM);
+ const res = dlm.filter([1, 2, 3, 2, 1, 2, 3, 2, 1, 2, 3, 2]);
+ expect(res.steps.length).toBe(12);
+ });
+
+ test("DLM.polynomial(1) is same as localLevel", () => {
+ const dlm = DLM.polynomial(1, { sigmaObs: 1, sigmaState: 0.5 });
+ expect(dlm).toBeInstanceOf(DLM);
+ });
+});
+
+// ─── Edge cases ────────────────────────────────────────────────────────────────
+describe("DLM edge cases", () => {
+ test("single observation", () => {
+ const dlm = DLM.localLevel({ sigmaObs: 1, sigmaLevel: 0.5 });
+ const res = dlm.filter([3.14]);
+ expect(res.filteredMeans.length).toBe(1);
+ expect(Number.isFinite(res.logLikelihood)).toBe(true);
+ });
+
+ test("all missing observations", () => {
+ const dlm = DLM.localLevel({ sigmaObs: 1, sigmaLevel: 0.5 });
+ const res = dlm.filter([null, null, null]);
+ expect(res.steps.length).toBe(3);
+ // No log-likelihood contributions from missing obs
+ expect(res.logLikelihood).toBe(0);
+ });
+
+ test("constant series", () => {
+ const dlm = DLM.localLevel({ sigmaObs: 0.5, sigmaLevel: 0.1 });
+ const res = dlm.filter([5, 5, 5, 5, 5]);
+ // Filtered means should converge toward 5
+ const last = res.filteredMeans[4]?.[0]!;
+ expect(last).toBeGreaterThan(4);
+ expect(last).toBeLessThan(6);
+ });
+
+ test("custom prior m0 and C0", () => {
+ const dlm = DLM.localLevel({ sigmaObs: 1, sigmaLevel: 0.5 });
+ const _res1 = dlm.filter([1, 2, 3], { m0: [0], C0: [[1000]] });
+ const res2 = dlm.filter([1, 2, 3], { m0: [5], C0: [[0.01]] });
+ // Strong prior on m0=5 should keep filtered mean near 5 initially
+ expect(res2.filteredMeans[0]?.[0]).toBeGreaterThan(3);
+ });
+});
+
+// ─── Property-based tests ──────────────────────────────────────────────────────
+describe("DLM property tests", () => {
+ test("logLikelihood is finite for any finite observations", () => {
+ fc.assert(
+ fc.property(
+ fc.array(fc.float({ min: -100, max: 100, noNaN: true }), { minLength: 2, maxLength: 20 }),
+ fc.double({ min: 0.01, max: 10, noNaN: true }),
+ fc.double({ min: 0.01, max: 10, noNaN: true }),
+ (y, sv, sw) => {
+ const dlm = DLM.localLevel({ sigmaObs: sv, sigmaLevel: sw });
+ const res = dlm.filter(y);
+ return Number.isFinite(res.logLikelihood);
+ },
+ ),
+ );
+ });
+
+ test("filtered means length = input length", () => {
+ fc.assert(
+ fc.property(
+ fc.array(fc.float({ min: -50, max: 50, noNaN: true }), { minLength: 1, maxLength: 30 }),
+ (y) => {
+ const dlm = DLM.localLevel();
+ const res = dlm.filter(y);
+ return res.filteredMeans.length === y.length;
+ },
+ ),
+ );
+ });
+
+ test("smoother length = filter length", () => {
+ fc.assert(
+ fc.property(
+ fc.array(fc.float({ min: -50, max: 50, noNaN: true }), { minLength: 2, maxLength: 20 }),
+ (y) => {
+ const dlm = DLM.localLevel();
+ const res = dlm.smooth(y);
+ return res.smoothedMeans.length === y.length && res.smoothedCovs.length === y.length;
+ },
+ ),
+ );
+ });
+
+ test("forecast lower ≤ mean[0] ≤ upper (scalar model)", () => {
+ fc.assert(
+ fc.property(
+ fc.array(fc.float({ min: -50, max: 50, noNaN: true }), { minLength: 3, maxLength: 15 }),
+ fc.integer({ min: 1, max: 10 }),
+ (y, h) => {
+ const dlm = DLM.localLevel();
+ const res = dlm.filter(y);
+ const fc2 = dlm.forecast(res, h);
+ return fc2.lower.every(
+ (lo, i) => lo <= fc2.mean[i]?.[0]! + 1e-8 && fc2.mean[i]?.[0]! <= fc2.upper[i]! + 1e-8,
+ );
+ },
+ ),
+ );
+ });
+});
diff --git a/tests/stats/ets.test.ts b/tests/stats/ets.test.ts
new file mode 100644
index 000000000..66c0778c1
--- /dev/null
+++ b/tests/stats/ets.test.ts
@@ -0,0 +1,853 @@
+/**
+ * Tests for src/stats/ets.ts
+ *
+ * Covers Simple Exponential Smoothing (SES), Holt linear trend, and
+ * Holt-Winters (full ETS) with additive and multiplicative seasonality.
+ * Mirrors statsmodels.tsa.holtwinters behaviour.
+ */
+import { describe, expect, it } from "bun:test";
+import fc from "fast-check";
+import {
+ ExponentialSmoothing,
+ Holt,
+ Series,
+ SimpleExpSmoothing,
+ fitEts,
+ holt,
+ simpleExpSmoothing,
+} from "../../src/index.ts";
+
+// ─── Test fixtures ─────────────────────────────────────────────────────────────
+
+/** Passenger data (Box & Jenkins airline data first 12 months). */
+const AIRLINE = [
+ 112, 118, 132, 129, 121, 135, 148, 148, 136, 119, 104, 118, 115, 126, 141, 135, 125, 149, 170,
+ 170, 158, 133, 114, 140, 145, 150, 178, 163, 172, 178, 199, 199, 184, 162, 146, 166,
+];
+
+/** Simple upward trend series. */
+const TREND = [10, 12, 14, 16, 18, 20, 22, 24, 26, 28, 30];
+
+/** Seasonal additive series (no trend). */
+const SEASONAL_ADD = [5, 8, 7, 4, 5, 8, 7, 4, 5, 8, 7, 4, 5, 8, 7, 4, 5, 8, 7, 4, 6, 9, 8, 5];
+
+/** Deterministic LCG for reproducible pseudo-noise. */
+function lcg(seed: number): () => number {
+ let s = seed;
+ return () => {
+ s = (s * 1664525 + 1013904223) & 0x7fffffff;
+ return s / 0x7fffffff;
+ };
+}
+
+/** Build a noisy sinusoidal series with additive seasonal component. */
+function buildSeasonal(n: number, amplitude: number, period: number, noiseAmp = 0.5): number[] {
+ const rand = lcg(1234);
+ return Array.from({ length: n }, (_, t) => {
+ const trend = 10 + 0.2 * t;
+ const seasonal = amplitude * Math.sin((2 * Math.PI * t) / period);
+ const noise = (rand() - 0.5) * noiseAmp;
+ return trend + seasonal + noise;
+ });
+}
+
+/** Mean absolute error between two arrays. */
+function mae(a: readonly number[], b: readonly number[]): number {
+ let s = 0;
+ for (let i = 0; i < a.length; i++) {
+ s += Math.abs((a[i] ?? 0) - (b[i] ?? 0));
+ }
+ return s / a.length;
+}
+
+/** Root mean squared error. */
+function rmse(a: readonly number[], b: readonly number[]): number {
+ let s = 0;
+ for (let i = 0; i < a.length; i++) {
+ s += ((a[i] ?? 0) - (b[i] ?? 0)) ** 2;
+ }
+ return Math.sqrt(s / a.length);
+}
+
+// ─── SimpleExpSmoothing ────────────────────────────────────────────────────────
+
+describe("SimpleExpSmoothing", () => {
+ describe("construction", () => {
+ it("creates instance without arguments", () => {
+ expect(() => new SimpleExpSmoothing()).not.toThrow();
+ });
+ });
+
+ describe("fit()", () => {
+ it("requires at least 2 observations", () => {
+ expect(() => new SimpleExpSmoothing().fit([1])).toThrow(RangeError);
+ });
+
+ it("returns alpha in (0, 1)", () => {
+ const fit = new SimpleExpSmoothing().fit(TREND);
+ expect(fit.alpha).toBeGreaterThan(0);
+ expect(fit.alpha).toBeLessThan(1);
+ });
+
+ it("returns fittedValues of same length as input", () => {
+ const fit = new SimpleExpSmoothing().fit(TREND);
+ expect(fit.fittedValues.length).toBe(TREND.length);
+ });
+
+ it("residuals + fittedValues = y", () => {
+ const fit = new SimpleExpSmoothing().fit(TREND);
+ for (let i = 0; i < TREND.length; i++) {
+ expect((fit.fittedValues[i] ?? 0) + (fit.residuals[i] ?? 0)).toBeCloseTo(TREND[i] ?? 0, 8);
+ }
+ });
+
+ it("sse equals sum of squared residuals", () => {
+ const fit = new SimpleExpSmoothing().fit(TREND);
+ let sse = 0;
+ for (const e of fit.residuals) {
+ sse += e * e;
+ }
+ expect(fit.sse).toBeCloseTo(sse, 6);
+ });
+
+ it("fixed alpha is respected", () => {
+ const alpha = 0.4;
+ const fit = new SimpleExpSmoothing().fit(TREND, { alpha });
+ expect(fit.alpha).toBeCloseTo(alpha, 10);
+ });
+
+ it("optimised alpha produces lower SSE than α=0.5 for trending data", () => {
+ const fit1 = new SimpleExpSmoothing().fit(TREND);
+ const fit2 = new SimpleExpSmoothing().fit(TREND, { alpha: 0.5 });
+ expect(fit1.sse).toBeLessThanOrEqual(fit2.sse + 1e-6);
+ });
+
+ it("AIC > 0", () => {
+ const fit = new SimpleExpSmoothing().fit(TREND);
+ expect(Number.isFinite(fit.aic)).toBe(true);
+ });
+
+ it("BIC ≥ AIC for k > 0, n > e", () => {
+ const fit = new SimpleExpSmoothing().fit(AIRLINE.slice(0, 24));
+ // With k=2, n=24: BIC = AIC + k*(ln(n) - 2)
+ // ln(24) ≈ 3.18 > 2, so BIC ≥ AIC
+ expect(fit.bic).toBeGreaterThanOrEqual(fit.aic - 1e-6);
+ });
+
+ it("accepts Series input", () => {
+ const s = new Series({ data: TREND });
+ const fit = new SimpleExpSmoothing().fit(s);
+ expect(fit.fittedValues.length).toBe(TREND.length);
+ });
+
+ it("initialLevel option overrides default", () => {
+ const fit = new SimpleExpSmoothing().fit(TREND, { initialLevel: 5 });
+ expect(fit.initialLevel).toBeCloseTo(5, 10);
+ expect(fit.fittedValues[0]).toBe(5);
+ const fixedAlpha = new SimpleExpSmoothing().fit(TREND, { initialLevel: 5, alpha: 0.5 });
+ expect(fit.sse).toBeLessThanOrEqual(fixedAlpha.sse);
+ });
+ });
+
+ describe("forecast()", () => {
+ it("throws if called before fit()", () => {
+ expect(() => new SimpleExpSmoothing().forecast(3)).toThrow();
+ });
+
+ it("returns flat forecast (all equal) of correct length", () => {
+ const model = new SimpleExpSmoothing();
+ model.fit(TREND);
+ const fc = model.forecast(4);
+ expect(fc.length).toBe(4);
+ for (const v of fc) {
+ expect(v).toBeCloseTo(fc[0] ?? 0, 8);
+ }
+ });
+
+ it("forecast ≈ last observed value for α ≈ 1", () => {
+ const model = new SimpleExpSmoothing();
+ model.fit(TREND, { alpha: 0.9999 });
+ const fc = model.forecast(1);
+ expect(fc[0]).toBeCloseTo(TREND.at(-1) ?? 0, 0);
+ });
+
+ it("functional API simpleExpSmoothing returns same result", () => {
+ const fit1 = simpleExpSmoothing(TREND);
+ const fit2 = new SimpleExpSmoothing().fit(TREND);
+ expect(fit1.alpha).toBeCloseTo(fit2.alpha, 8);
+ expect(fit1.sse).toBeCloseTo(fit2.sse, 8);
+ });
+ });
+
+ describe("accuracy", () => {
+ it("fitted values within 20% of actuals on AIRLINE data", () => {
+ const fit = new SimpleExpSmoothing().fit(AIRLINE);
+ const err = mae(fit.fittedValues, AIRLINE);
+ expect(err / (AIRLINE.reduce((a, b) => a + b, 0) / AIRLINE.length)).toBeLessThan(0.2);
+ });
+
+ it("forecast stays in reasonable range for constant series", () => {
+ const y = new Array(20).fill(5);
+ const model = new SimpleExpSmoothing();
+ model.fit(y);
+ const fc = model.forecast(5);
+ for (const v of fc) {
+ expect(Math.abs(v - 5)).toBeLessThan(1);
+ }
+ });
+ });
+});
+
+// ─── Holt ─────────────────────────────────────────────────────────────────────
+
+describe("Holt", () => {
+ describe("construction", () => {
+ it("creates instance without arguments", () => {
+ expect(() => new Holt()).not.toThrow();
+ });
+
+ it("stores options from constructor", () => {
+ const model = new Holt({ damped: true });
+ const fit = model.fit(TREND);
+ expect(fit.phi).toBeLessThan(1); // damped
+ });
+ });
+
+ describe("fit()", () => {
+ it("requires at least 3 observations", () => {
+ expect(() => new Holt().fit([1, 2])).toThrow(RangeError);
+ });
+
+ it("returns alpha and beta in (0, 1)", () => {
+ const fit = new Holt().fit(TREND);
+ expect(fit.alpha).toBeGreaterThan(0);
+ expect(fit.alpha).toBeLessThan(1);
+ expect(fit.beta).toBeGreaterThan(0);
+ expect(fit.beta).toBeLessThan(1);
+ });
+
+ it("phi = 1 when damped = false (default)", () => {
+ const fit = new Holt().fit(TREND);
+ expect(fit.phi).toBe(1.0);
+ });
+
+ it("phi < 1 when damped = true", () => {
+ const fit = new Holt({ damped: true }).fit(TREND);
+ expect(fit.phi).toBeGreaterThan(0.8);
+ expect(fit.phi).toBeLessThan(1);
+ });
+
+ it("fittedValues length = n", () => {
+ const fit = new Holt().fit(TREND);
+ expect(fit.fittedValues.length).toBe(TREND.length);
+ });
+
+ it("fitted + residuals = y", () => {
+ const fit = new Holt().fit(AIRLINE.slice(0, 12));
+ for (let i = 0; i < 12; i++) {
+ expect((fit.fittedValues[i] ?? 0) + (fit.residuals[i] ?? 0)).toBeCloseTo(
+ AIRLINE[i] ?? 0,
+ 6,
+ );
+ }
+ });
+
+ it("sse = sum of squared residuals", () => {
+ const fit = new Holt().fit(TREND);
+ let sse = 0;
+ for (const e of fit.residuals) {
+ sse += e * e;
+ }
+ expect(fit.sse).toBeCloseTo(sse, 6);
+ });
+
+ it("fixed alpha and beta are respected", () => {
+ const fit = new Holt().fit(TREND, { alpha: 0.5, beta: 0.2 });
+ expect(fit.alpha).toBeCloseTo(0.5, 10);
+ expect(fit.beta).toBeCloseTo(0.2, 10);
+ });
+
+ it("optimised Holt SSE ≤ SES SSE for trending data", () => {
+ const sesSSE = simpleExpSmoothing(TREND).sse;
+ const holtSSE = holt(TREND).sse;
+ expect(holtSSE).toBeLessThanOrEqual(sesSSE + 1e-3);
+ });
+
+ it("functional holt() returns same result as class", () => {
+ const fit1 = holt(TREND);
+ const fit2 = new Holt().fit(TREND);
+ expect(fit1.alpha).toBeCloseTo(fit2.alpha, 6);
+ expect(fit1.sse).toBeCloseTo(fit2.sse, 6);
+ });
+ });
+
+ describe("forecast()", () => {
+ it("throws if called before fit()", () => {
+ expect(() => new Holt().forecast(3)).toThrow();
+ });
+
+ it("returns correct number of steps", () => {
+ const model = new Holt();
+ model.fit(TREND);
+ expect(model.forecast(5).length).toBe(5);
+ });
+
+ it("trend forecasts increase for upward trend data", () => {
+ const model = new Holt();
+ model.fit(TREND);
+ const fc = model.forecast(3);
+ expect(fc[1] ?? 0).toBeGreaterThan(fc[0] ?? 0);
+ expect(fc[2] ?? 0).toBeGreaterThan(fc[1] ?? 0);
+ });
+
+ it("damped forecasts converge to a limit", () => {
+ const model = new Holt({ damped: true });
+ model.fit(TREND);
+ const fc = model.forecast(20);
+ // Differences should shrink
+ const diffs = fc.slice(1).map((v, i) => v - (fc[i] ?? 0));
+ for (let i = 1; i < diffs.length; i++) {
+ expect(Math.abs(diffs[i] ?? 0)).toBeLessThanOrEqual(Math.abs(diffs[i - 1] ?? 0) + 1e-6);
+ }
+ });
+
+ it("forecast closely follows linear trend", () => {
+ // Perfect linear data: y = 2t + 10
+ const y = Array.from({ length: 20 }, (_, t) => 2 * t + 10);
+ const model = new Holt();
+ model.fit(y);
+ const fc = model.forecast(3);
+ // Should forecast ≈ [50, 52, 54]
+ expect(fc[0]).toBeCloseTo(50, 0);
+ expect(fc[2]).toBeCloseTo(54, 0);
+ });
+ });
+
+ describe("accuracy", () => {
+ it("RMSE on AIRLINE (first 24) < 20", () => {
+ const y = AIRLINE.slice(0, 24);
+ const fit = new Holt().fit(y);
+ const err = rmse(fit.fittedValues.slice(1), y.slice(1));
+ expect(err).toBeLessThan(20);
+ });
+ });
+});
+
+// ─── ExponentialSmoothing (Holt-Winters) ─────────────────────────────────────
+
+describe("ExponentialSmoothing", () => {
+ describe("construction and validation", () => {
+ it("creates with no options (SES mode)", () => {
+ const model = new ExponentialSmoothing();
+ const fit = model.fit(TREND);
+ expect(fit.beta).toBeNull();
+ expect(fit.gamma).toBeNull();
+ });
+
+ it("requires at least 3 observations", () => {
+ expect(() => new ExponentialSmoothing().fit([1, 2])).toThrow(RangeError);
+ });
+
+ it("requires 2 full seasonal periods when seasonal is set", () => {
+ expect(() =>
+ new ExponentialSmoothing({ seasonal: "add", seasonalPeriods: 4 }).fit([1, 2, 3, 4, 5]),
+ ).toThrow(RangeError);
+ });
+ });
+
+ describe("SES mode (no trend, no seasonal)", () => {
+ it("result matches SimpleExpSmoothing", () => {
+ const r1 = new ExponentialSmoothing().fit(TREND);
+ const r2 = new SimpleExpSmoothing().fit(TREND);
+ expect(r1.alpha).toBeCloseTo(r2.alpha, 4);
+ expect(r1.sse).toBeCloseTo(r2.sse, 4);
+ });
+
+ it("beta and gamma are null", () => {
+ const fit = new ExponentialSmoothing().fit(TREND);
+ expect(fit.beta).toBeNull();
+ expect(fit.gamma).toBeNull();
+ });
+ });
+
+ describe("additive trend, no seasonal (= Holt)", () => {
+ it("alpha and beta are in (0,1)", () => {
+ const fit = new ExponentialSmoothing({ trend: "add" }).fit(TREND);
+ expect(fit.alpha).toBeGreaterThan(0);
+ expect(fit.beta).not.toBeNull();
+ expect(fit.beta ?? 0).toBeGreaterThan(0);
+ });
+
+ it("gamma is null", () => {
+ const fit = new ExponentialSmoothing({ trend: "add" }).fit(TREND);
+ expect(fit.gamma).toBeNull();
+ });
+
+ it("matches Holt fitted values and forecasts for the same additive model", () => {
+ const ets = new ExponentialSmoothing({ trend: "add" });
+ const holtModel = new Holt();
+ const r1 = ets.fit(TREND);
+ const r2 = holtModel.fit(TREND);
+ expect(r1.sse).toBeCloseTo(r2.sse, 12);
+ expect(r1.fittedValues).toEqual(r2.fittedValues);
+ expect(ets.forecast(5)).toEqual(holtModel.forecast(5));
+ });
+
+ it("matches Holt for noisy observations", () => {
+ const r1 = new ExponentialSmoothing({ trend: "add" }).fit(AIRLINE);
+ const r2 = new Holt().fit(AIRLINE);
+ expect(r1.sse).toBeCloseTo(r2.sse, 12);
+ });
+ });
+
+ describe("additive trend + additive seasonal (classic Holt-Winters)", () => {
+ const y = SEASONAL_ADD;
+ const m = 4;
+
+ it("all parameters estimated", () => {
+ const fit = new ExponentialSmoothing({
+ trend: "add",
+ seasonal: "add",
+ seasonalPeriods: m,
+ }).fit(y);
+ expect(fit.alpha).toBeGreaterThan(0);
+ expect(fit.beta).not.toBeNull();
+ expect(fit.gamma).not.toBeNull();
+ expect(fit.initialSeasons).not.toBeNull();
+ expect(fit.initialSeasons?.length).toBe(m);
+ });
+
+ it("additive seasonal indices sum close to 0", () => {
+ const fit = new ExponentialSmoothing({
+ trend: "add",
+ seasonal: "add",
+ seasonalPeriods: m,
+ }).fit(y);
+ const sum = (fit.initialSeasons ?? []).reduce((a, b) => a + b, 0);
+ expect(sum).toBeCloseTo(0, 10);
+ const replay = new ExponentialSmoothing({
+ trend: "add",
+ seasonal: "add",
+ seasonalPeriods: m,
+ initializationMethod: "known",
+ initialLevel: fit.initialLevel,
+ initialTrend: fit.initialTrend ?? 0,
+ initialSeasons: fit.initialSeasons ?? [],
+ alpha: fit.alpha,
+ beta: fit.beta ?? 0,
+ gamma: fit.gamma ?? 0,
+ }).fit(y);
+ expect(replay.fittedValues).toEqual(fit.fittedValues);
+ expect(replay.sse).toBe(fit.sse);
+ });
+
+ it("fitted + residuals = y", () => {
+ const fit = new ExponentialSmoothing({
+ trend: "add",
+ seasonal: "add",
+ seasonalPeriods: m,
+ }).fit(y);
+ for (let i = 0; i < y.length; i++) {
+ expect((fit.fittedValues[i] ?? 0) + (fit.residuals[i] ?? 0)).toBeCloseTo(y[i] ?? 0, 6);
+ }
+ });
+
+ it("sse = sum of squared residuals", () => {
+ const fit = new ExponentialSmoothing({
+ trend: "add",
+ seasonal: "add",
+ seasonalPeriods: m,
+ }).fit(y);
+ let sse = 0;
+ for (const e of fit.residuals) {
+ sse += e * e;
+ }
+ expect(fit.sse).toBeCloseTo(sse, 6);
+ });
+
+ it("AIC < BIC for n > e", () => {
+ const fit = new ExponentialSmoothing({
+ trend: "add",
+ seasonal: "add",
+ seasonalPeriods: m,
+ }).fit(y);
+ // For n = 24, k = 1+1+1+1+4=8: BIC - AIC = k*(ln(n)-2) ≈ 8*(3.18-2)=9.4 > 0
+ expect(fit.bic).toBeGreaterThan(fit.aic - 1e-3);
+ });
+
+ it("forecast returns correct number of steps", () => {
+ const model = new ExponentialSmoothing({
+ trend: "add",
+ seasonal: "add",
+ seasonalPeriods: m,
+ });
+ model.fit(y);
+ const fc = model.forecast(8);
+ expect(fc.length).toBe(8);
+ });
+
+ it("all forecast values are finite", () => {
+ const model = new ExponentialSmoothing({
+ trend: "add",
+ seasonal: "add",
+ seasonalPeriods: m,
+ });
+ model.fit(y);
+ const fc = model.forecast(8);
+ for (const v of fc) {
+ expect(Number.isFinite(v)).toBe(true);
+ }
+ });
+
+ it("seasonal forecast repeats seasonal pattern (low noise data)", () => {
+ // Perfect seasonal data with no noise
+ const perfect: number[] = [];
+ for (let i = 0; i < 24; i++) {
+ const base = [5, 8, 7, 4];
+ perfect.push(base[i % 4] ?? 5);
+ }
+ const model = new ExponentialSmoothing({
+ trend: "add",
+ seasonal: "add",
+ seasonalPeriods: 4,
+ });
+ model.fit(perfect);
+ const fc = model.forecast(8);
+ const pattern = [fc[0] ?? 0, fc[1] ?? 0, fc[2] ?? 0, fc[3] ?? 0];
+ // Next 4 should roughly repeat the pattern
+ for (let i = 0; i < 4; i++) {
+ expect(Math.abs((fc[i + 4] ?? 0) - (pattern[i] ?? 0))).toBeLessThan(2);
+ }
+ });
+ });
+
+ describe("additive trend + multiplicative seasonal", () => {
+ it("estimates all parameters", () => {
+ const fit = new ExponentialSmoothing({
+ trend: "add",
+ seasonal: "mul",
+ seasonalPeriods: 4,
+ }).fit(SEASONAL_ADD);
+ expect(fit.alpha).toBeGreaterThan(0);
+ expect(fit.beta).not.toBeNull();
+ expect(fit.gamma).not.toBeNull();
+ });
+
+ it("fitted + residuals = y", () => {
+ const fit = new ExponentialSmoothing({
+ trend: "add",
+ seasonal: "mul",
+ seasonalPeriods: 4,
+ }).fit(SEASONAL_ADD);
+ for (let i = 0; i < SEASONAL_ADD.length; i++) {
+ expect((fit.fittedValues[i] ?? 0) + (fit.residuals[i] ?? 0)).toBeCloseTo(
+ SEASONAL_ADD[i] ?? 0,
+ 6,
+ );
+ }
+ });
+ });
+
+ describe("no trend + additive seasonal", () => {
+ it("beta is null", () => {
+ const fit = new ExponentialSmoothing({
+ seasonal: "add",
+ seasonalPeriods: 4,
+ }).fit(SEASONAL_ADD);
+ expect(fit.beta).toBeNull();
+ expect(fit.gamma).not.toBeNull();
+ });
+ });
+
+ describe("no trend + multiplicative seasonal", () => {
+ it("fits without error", () => {
+ const fit = new ExponentialSmoothing({
+ seasonal: "mul",
+ seasonalPeriods: 4,
+ }).fit(SEASONAL_ADD);
+ expect(fit.gamma).not.toBeNull();
+ expect(fit.beta).toBeNull();
+ });
+ });
+
+ describe("damped trend", () => {
+ it("phi < 1 when damped = true", () => {
+ const fit = new ExponentialSmoothing({
+ trend: "add",
+ damped: true,
+ }).fit(TREND);
+ expect(fit.phi).toBeLessThan(1);
+ expect(fit.phi).toBeGreaterThan(0.8);
+ });
+
+ it("phi = 1 when damped = false", () => {
+ const fit = new ExponentialSmoothing({ trend: "add" }).fit(TREND);
+ expect(fit.phi).toBe(1.0);
+ });
+ });
+
+ describe("known initialisation", () => {
+ it("uses provided initial values", () => {
+ const fit = new ExponentialSmoothing({
+ trend: "add",
+ seasonal: "add",
+ seasonalPeriods: 4,
+ initializationMethod: "known",
+ initialLevel: 6.0,
+ initialTrend: 0.1,
+ initialSeasons: [1, 2, 0, -1],
+ }).fit(SEASONAL_ADD);
+ expect(fit.initialLevel).toBeCloseTo(6.0, 8);
+ expect(fit.initialTrend).toBeCloseTo(0.1, 8);
+ });
+ });
+
+ describe("fixed parameters", () => {
+ it("fixed alpha is respected", () => {
+ const fit = new ExponentialSmoothing({ alpha: 0.4 }).fit(TREND);
+ expect(fit.alpha).toBeCloseTo(0.4, 8);
+ });
+
+ it("fixed beta is respected", () => {
+ const fit = new ExponentialSmoothing({ trend: "add", alpha: 0.3, beta: 0.15 }).fit(TREND);
+ expect(fit.beta).toBeCloseTo(0.15, 8);
+ });
+
+ it("fixed gamma is respected", () => {
+ const fit = new ExponentialSmoothing({
+ seasonal: "add",
+ seasonalPeriods: 4,
+ gamma: 0.2,
+ }).fit(SEASONAL_ADD);
+ expect(fit.gamma).toBeCloseTo(0.2, 8);
+ });
+ });
+
+ describe("forecastWithCI()", () => {
+ it("throws if called before fit()", () => {
+ expect(() => new ExponentialSmoothing().forecastWithCI(3)).toThrow();
+ });
+
+ it("returns correct structure", () => {
+ const model = new ExponentialSmoothing({ trend: "add" });
+ model.fit(TREND);
+ const r = model.forecastWithCI(4);
+ expect(r.forecast.length).toBe(4);
+ expect(r.lower.length).toBe(4);
+ expect(r.upper.length).toBe(4);
+ expect(r.stderr.length).toBe(4);
+ });
+
+ it("upper > lower for all steps", () => {
+ const model = new ExponentialSmoothing({ trend: "add" });
+ model.fit(TREND);
+ const r = model.forecastWithCI(4);
+ for (let i = 0; i < 4; i++) {
+ expect(r.upper[i] ?? 0).toBeGreaterThan(r.lower[i] ?? 0);
+ }
+ });
+
+ it("forecast is within confidence interval", () => {
+ const model = new ExponentialSmoothing({ trend: "add" });
+ model.fit(TREND);
+ const r = model.forecastWithCI(4);
+ for (let i = 0; i < 4; i++) {
+ expect(r.forecast[i] ?? 0).toBeGreaterThanOrEqual(r.lower[i] ?? 0);
+ expect(r.forecast[i] ?? 0).toBeLessThanOrEqual(r.upper[i] ?? 0);
+ }
+ });
+
+ it("intervals widen with horizon", () => {
+ const model = new ExponentialSmoothing({ trend: "add" });
+ model.fit(TREND);
+ const r = model.forecastWithCI(5);
+ const widths = r.upper.map((u, i) => u - (r.lower[i] ?? 0));
+ for (let i = 1; i < widths.length; i++) {
+ expect(widths[i] ?? 0).toBeGreaterThanOrEqual((widths[i - 1] ?? 0) - 1e-6);
+ }
+ });
+ });
+
+ describe("functional fitEts()", () => {
+ it("returns same result as class fit()", () => {
+ const r1 = fitEts(TREND, { trend: "add" });
+ const r2 = new ExponentialSmoothing({ trend: "add" }).fit(TREND);
+ expect(r1.alpha).toBeCloseTo(r2.alpha, 6);
+ expect(r1.sse).toBeCloseTo(r2.sse, 6);
+ });
+ });
+
+ describe("AIRLINE accuracy", () => {
+ it("additive H-W in-sample MAE < 15 on AIRLINE", () => {
+ const fit = new ExponentialSmoothing({
+ trend: "add",
+ seasonal: "add",
+ seasonalPeriods: 12,
+ }).fit(AIRLINE);
+ const err = mae(fit.fittedValues, AIRLINE);
+ expect(err).toBeLessThan(15);
+ });
+
+ it("multiplicative H-W in-sample MAE < 20 on AIRLINE", () => {
+ const fit = new ExponentialSmoothing({
+ trend: "add",
+ seasonal: "mul",
+ seasonalPeriods: 12,
+ }).fit(AIRLINE);
+ const err = mae(fit.fittedValues, AIRLINE);
+ expect(err).toBeLessThan(20);
+ });
+
+ it("H-W AIC < SES AIC on seasonal AIRLINE data", () => {
+ const sesFit = fitEts(AIRLINE);
+ const hwFit = fitEts(AIRLINE, {
+ trend: "add",
+ seasonal: "add",
+ seasonalPeriods: 12,
+ });
+ // Holt-Winters should have better (lower) AIC for seasonal data
+ expect(hwFit.aic).toBeLessThan(sesFit.aic + 20); // relaxed: may use more params
+ });
+ });
+
+ describe("information criteria", () => {
+ it("log-likelihood matches the Gaussian residual density", () => {
+ const fit = fitEts(TREND, { trend: "add" });
+ expect(Number.isFinite(fit.logLikelihood)).toBe(true);
+ const variance = Math.max(fit.sse / TREND.length, 1e-15);
+ const expected = -0.5 * TREND.length * (Math.log(2 * Math.PI * variance) + 1);
+ // A probability density can exceed 1 when its variance is small.
+ expect(fit.logLikelihood).toBeCloseTo(expected, 10);
+ });
+
+ it("AICc >= AIC", () => {
+ const fit = fitEts(TREND, { trend: "add" });
+ expect(fit.aicc).toBeGreaterThanOrEqual(fit.aic - 1e-6);
+ });
+ });
+
+ describe("edge cases", () => {
+ it("handles constant series (SES mode)", () => {
+ const y = new Array(20).fill(7);
+ const fit = fitEts(y);
+ for (const v of fit.fittedValues) {
+ expect(Math.abs(v - 7)).toBeLessThan(1);
+ }
+ });
+
+ it("handles single-cycle seasonal (m=2)", () => {
+ const y = [1, 3, 1, 3, 1, 3, 1, 3, 1, 3];
+ const fit = fitEts(y, { seasonal: "add", seasonalPeriods: 2 });
+ expect(fit.gamma).not.toBeNull();
+ });
+
+ it("forecast of length 0 returns empty array", () => {
+ const model = new ExponentialSmoothing({ trend: "add" });
+ model.fit(TREND);
+ expect(model.forecast(0)).toEqual([]);
+ });
+
+ it("works with longer seasonal period (m=12) on 2 cycles", () => {
+ const y = buildSeasonal(24, 3, 12, 0.1);
+ const fit = fitEts(y, { trend: "add", seasonal: "add", seasonalPeriods: 12 });
+ expect(fit.alpha).toBeGreaterThan(0);
+ const n = new ExponentialSmoothing({ trend: "add", seasonal: "add", seasonalPeriods: 12 });
+ n.fit(y);
+ expect(n.forecast(12).length).toBe(12);
+ });
+ });
+});
+
+// ─── Property-based tests ──────────────────────────────────────────────────────
+
+describe("ETS property-based", () => {
+ it("SES: fitted + residuals = y for any n≥2 data", () => {
+ fc.assert(
+ fc.property(
+ fc.array(fc.float({ min: -1000, max: 1000, noNaN: true }), { minLength: 2, maxLength: 30 }),
+ (y) => {
+ const fit = simpleExpSmoothing(y);
+ for (let i = 0; i < y.length; i++) {
+ const diff = Math.abs(
+ (fit.fittedValues[i] ?? 0) + (fit.residuals[i] ?? 0) - (y[i] ?? 0),
+ );
+ if (diff > 1e-4) {
+ return false;
+ }
+ }
+ return true;
+ },
+ ),
+ { numRuns: 50 },
+ );
+ });
+
+ it("SES: alpha ∈ (0, 1) for any data", () => {
+ fc.assert(
+ fc.property(
+ fc.array(fc.float({ min: -100, max: 100, noNaN: true }), { minLength: 3, maxLength: 20 }),
+ (y) => {
+ const fit = simpleExpSmoothing(y);
+ return fit.alpha > 0 && fit.alpha < 1;
+ },
+ ),
+ { numRuns: 50 },
+ );
+ });
+
+ it("Holt: phi = 1 when not damped, for any data", () => {
+ fc.assert(
+ fc.property(
+ fc.array(fc.float({ min: -100, max: 100, noNaN: true }), { minLength: 4, maxLength: 20 }),
+ (y) => {
+ const fit = holt(y);
+ return fit.phi === 1.0;
+ },
+ ),
+ { numRuns: 40 },
+ );
+ });
+
+ it("ExponentialSmoothing: SSE ≥ 0 for any data", () => {
+ fc.assert(
+ fc.property(
+ fc.array(fc.float({ min: -50, max: 50, noNaN: true }), { minLength: 3, maxLength: 15 }),
+ (y) => {
+ const fit = fitEts(y, { trend: "add" });
+ return fit.sse >= 0;
+ },
+ ),
+ { numRuns: 40 },
+ );
+ });
+
+ it("ExponentialSmoothing: forecast length = steps for any data", () => {
+ fc.assert(
+ fc.property(
+ fc.array(fc.float({ min: -50, max: 50, noNaN: true }), { minLength: 3, maxLength: 15 }),
+ fc.integer({ min: 0, max: 10 }),
+ (y, steps) => {
+ const model = new ExponentialSmoothing({ trend: "add" });
+ model.fit(y);
+ return model.forecast(steps).length === steps;
+ },
+ ),
+ { numRuns: 40 },
+ );
+ });
+
+ it("SES: fixed alpha produces deterministic results", () => {
+ fc.assert(
+ fc.property(
+ fc.array(fc.float({ min: -50, max: 50, noNaN: true }), { minLength: 2, maxLength: 20 }),
+ fc.double({ min: 0.01, max: 0.99, noNaN: true }),
+ (y, alpha) => {
+ const r1 = simpleExpSmoothing(y, { alpha });
+ const r2 = simpleExpSmoothing(y, { alpha });
+ return Math.abs(r1.sse - r2.sse) < 1e-8;
+ },
+ ),
+ { numRuns: 40 },
+ );
+ });
+});
diff --git a/tests/stats/filters.test.ts b/tests/stats/filters.test.ts
new file mode 100644
index 000000000..f917c8488
--- /dev/null
+++ b/tests/stats/filters.test.ts
@@ -0,0 +1,571 @@
+/**
+ * Tests for src/stats/filters.ts
+ * Covers FIR design, Butterworth IIR, frequency response, and filter application.
+ */
+
+import { describe, expect, test } from "bun:test";
+import * as fc from "fast-check";
+import {
+ type FilterType,
+ type SOSSection,
+ butter,
+ cAbs,
+ filtfilt,
+ firwin,
+ freqz,
+ lfilter,
+ sosfilt,
+ sosfiltfilt,
+ sosfreqz,
+} from "../../src/stats/filters.ts";
+
+// ─── helpers ──────────────────────────────────────────────────────────────────
+
+function near(a: number, b: number, tol = 1e-4): boolean {
+ return Math.abs(a - b) <= tol * (1 + Math.abs(b));
+}
+
+function nearAbs(a: number, b: number, tol = 1e-9): boolean {
+ return Math.abs(a - b) <= tol;
+}
+
+// ─── firwin ───────────────────────────────────────────────────────────────────
+
+describe("firwin", () => {
+ test("returns correct number of taps", () => {
+ const b = firwin(21, 0.3);
+ expect(b.length).toBe(21);
+ });
+
+ test("low-pass: DC gain ≈ 1", () => {
+ const b = firwin(51, 0.3);
+ const dcGain = b.reduce((s, v) => s + v, 0);
+ expect(dcGain).toBeCloseTo(1.0, 4);
+ });
+
+ test("low-pass: gain near 0 at Nyquist", () => {
+ const b = firwin(51, 0.3);
+ // H(e^jπ) = sum b[n] * (-1)^n
+ const nyqGain = b.reduce((s, v, i) => s + v * (i % 2 === 0 ? 1 : -1), 0);
+ expect(Math.abs(nyqGain)).toBeLessThan(0.01);
+ });
+
+ test("high-pass (pass_zero=false): gain ≈ 1 at Nyquist", () => {
+ const b = firwin(51, 0.3, { pass_zero: false });
+ const nyqGain = b.reduce((s, v, i) => s + v * (i % 2 === 0 ? 1 : -1), 0);
+ expect(Math.abs(nyqGain)).toBeGreaterThan(0.9);
+ });
+
+ test("symmetric coefficients (linear phase)", () => {
+ const b = firwin(31, 0.4);
+ for (let i = 0; i < 16; i++) {
+ expect(b[i] ?? 0).toBeCloseTo(b[30 - i] ?? 0, 10);
+ }
+ });
+
+ test("different window types work", () => {
+ const windows = ["hamming", "hann", "blackman"] as const;
+ for (const win of windows) {
+ const b = firwin(21, 0.3, { window: win });
+ expect(b.length).toBe(21);
+ // DC gain should be near 1
+ const dc = b.reduce((s, v) => s + v, 0);
+ expect(dc).toBeCloseTo(1.0, 3);
+ }
+ });
+
+ test("custom fs scaling", () => {
+ const b1 = firwin(21, 0.3, { fs: 2 }); // default
+ const b2 = firwin(21, 300, { fs: 2000 }); // same normalised cutoff
+ for (let i = 0; i < b1.length; i++) {
+ expect(b1[i] ?? 0).toBeCloseTo(b2[i] ?? 0, 10);
+ }
+ });
+
+ test("band-pass (pass_zero=false, two cutoffs)", () => {
+ const b = firwin(51, [0.2, 0.4], { pass_zero: false });
+ expect(b.length).toBe(51);
+ // DC gain should be near 0
+ const dcGain = b.reduce((s, v) => s + v, 0);
+ expect(Math.abs(dcGain)).toBeLessThan(0.05);
+ });
+
+ test("property: all taps are finite", () => {
+ fc.assert(
+ fc.property(
+ fc.integer({ min: 5, max: 51 }).filter((n) => n % 2 === 1),
+ fc.double({ min: 0.01, max: 0.49, noNaN: true }),
+ (taps, cutoff) => {
+ const b = firwin(taps, cutoff);
+ return b.every(Number.isFinite);
+ },
+ ),
+ );
+ });
+});
+
+// ─── freqz ────────────────────────────────────────────────────────────────────
+
+describe("freqz", () => {
+ test("a unit delay has negative phase when numerator and denominator lengths differ", () => {
+ const { H } = freqz([0, 1], [1], [Math.PI / 2]);
+ expect(H[0]?.re).toBeCloseTo(0, 12);
+ expect(H[0]?.im).toBeCloseTo(-1, 12);
+ });
+ test("FIR identity filter (b=[1]) — H=1 everywhere", () => {
+ const { H } = freqz([1], [1], 32);
+ for (const h of H) {
+ expect(cAbs(h)).toBeCloseTo(1.0, 10);
+ }
+ });
+
+ test("output length matches worN", () => {
+ const { w, H } = freqz([1, 0, 0], [1], 64);
+ expect(w.length).toBe(64);
+ expect(H.length).toBe(64);
+ });
+
+ test("specific frequencies array", () => {
+ const ws = [0, Math.PI / 4, Math.PI / 2, Math.PI];
+ const { w, H } = freqz([1], [1], ws);
+ expect(w).toEqual(ws);
+ expect(H.length).toBe(4);
+ });
+
+ test("low-pass FIR: passband gain ≈ 1, stopband ≈ 0", () => {
+ const b = firwin(51, 0.3);
+ const { w, H } = freqz(b, [1], 256);
+ const mag = H.map(cAbs);
+ // DC
+ expect(mag[0]).toBeCloseTo(1.0, 2);
+ // Nyquist (last bin ~ π)
+ expect(mag[255] ?? 0).toBeLessThan(0.05);
+ });
+
+ test("high-pass FIR: stopband at DC, passband at Nyquist", () => {
+ const b = firwin(51, 0.3, { pass_zero: false });
+ const { H } = freqz(b, [1], 256);
+ const mag = H.map(cAbs);
+ expect(mag[0] ?? 0).toBeLessThan(0.05); // near zero at DC
+ expect(mag[255] ?? 0).toBeGreaterThan(0.9); // near 1 at Nyquist
+ });
+
+ test("first frequency is 0", () => {
+ const { w } = freqz([1], [1], 128);
+ expect(w[0]).toBe(0);
+ });
+});
+
+// ─── butter ───────────────────────────────────────────────────────────────────
+
+describe("butter", () => {
+ // scipy.signal.butter(3, cutoff, btype=type), SciPy 1.18.1.
+ const references: {
+ type: FilterType;
+ cutoff: number | readonly [number, number];
+ b: number[];
+ a: number[];
+ }[] = [
+ {
+ type: "lowpass",
+ cutoff: 0.3,
+ b: [0.04953299635725318, 0.14859898907175956, 0.14859898907175956, 0.04953299635725318],
+ a: [1, -1.1619174836717323, 0.6959427557896507, -0.13776130125989283],
+ },
+ {
+ type: "highpass",
+ cutoff: 0.3,
+ b: [0.3744526925901595, -1.1233580777704786, 1.1233580777704786, -0.3744526925901595],
+ a: [1, -1.1619174836717328, 0.6959427557896509, -0.13776130125989286],
+ },
+ {
+ type: "bandpass",
+ cutoff: [0.2, 0.4],
+ b: [
+ 0.01809893300751444, 0, -0.05429679902254332, 0, 0.05429679902254332, 0,
+ -0.01809893300751444,
+ ],
+ a: [
+ 1, -2.941867669925039, 4.7023172884643305, -4.634109656210411, 3.0774477936582785,
+ -1.2466196810240529, 0.2780599176345464,
+ ],
+ },
+ {
+ type: "bandstop",
+ cutoff: [0.2, 0.4],
+ b: [
+ 0.5276243825019431, -1.9565388100762569, 4.0012881173766335, -4.909519387006988,
+ 4.001288117376634, -1.9565388100762573, 0.5276243825019431,
+ ],
+ a: [
+ 1, -2.9418676699250383, 4.7023172884643305, -4.634109656210412, 3.077447793658279,
+ -1.246619681024053, 0.2780599176345464,
+ ],
+ },
+ ];
+
+ for (const reference of references) {
+ test(`${reference.type} coefficients match SciPy`, () => {
+ const { b, a } = butter(3, reference.cutoff, reference.type);
+ expect(b.length).toBe(reference.b.length);
+ expect(a.length).toBe(reference.a.length);
+ for (let i = 0; i < b.length; i++) {
+ expect(b[i]).toBeCloseTo(reference.b[i] ?? 0, 11);
+ expect(a[i]).toBeCloseTo(reference.a[i] ?? 0, 11);
+ }
+ });
+ }
+
+ test("band filters retain their order for both real and complex transformed pole pairs", () => {
+ for (const type of ["bandpass", "bandstop"] satisfies readonly FilterType[]) {
+ for (const cutoff of [
+ [0.2, 0.4],
+ [0.1, 0.9],
+ ] satisfies readonly [number, number][]) {
+ for (const order of [1, 2, 3, 4]) {
+ const { sos, a } = butter(order, cutoff, type);
+ expect(sos.length).toBe(order);
+ expect(a.length).toBe(2 * order + 1);
+ const frequencies = [0, cutoff[0] * Math.PI, cutoff[1] * Math.PI, Math.PI];
+ const magnitudes = sosfreqz(sos, frequencies).H.map(cAbs);
+ expect(magnitudes[1]).toBeCloseTo(Math.SQRT1_2, 10);
+ expect(magnitudes[2]).toBeCloseTo(Math.SQRT1_2, 10);
+ expect(magnitudes[0]).toBeCloseTo(type === "bandpass" ? 0 : 1, 10);
+ expect(magnitudes[3]).toBeCloseTo(type === "bandpass" ? 0 : 1, 10);
+ }
+ }
+ }
+ });
+
+ test("invalid critical frequencies fail before producing non-finite filters", () => {
+ for (const cutoff of [0, 1, -0.2, Number.NaN, Number.POSITIVE_INFINITY]) {
+ expect(() => butter(2, cutoff)).toThrow(RangeError);
+ }
+ expect(() => butter(2, [0.4, 0.2], "bandpass")).toThrow(RangeError);
+ expect(() => butter(2, [0.2, 0.2], "bandstop")).toThrow(RangeError);
+ });
+ test("returns sos, b, a arrays", () => {
+ const result = butter(2, 0.3);
+ expect(Array.isArray(result.sos)).toBe(true);
+ expect(Array.isArray(result.b)).toBe(true);
+ expect(Array.isArray(result.a)).toBe(true);
+ });
+
+ test("SOS sections count = ceil(N/2)", () => {
+ for (const N of [1, 2, 3, 4, 5, 6]) {
+ const { sos } = butter(N, 0.3);
+ expect(sos.length).toBe(Math.ceil(N / 2));
+ }
+ });
+
+ test("each SOS section has 6 coefficients", () => {
+ const { sos } = butter(4, 0.3);
+ for (const sec of sos) {
+ expect(sec.length).toBe(6);
+ }
+ });
+
+ test("SOS a[0] = 1 for all sections", () => {
+ const { sos } = butter(4, 0.3);
+ for (const sec of sos) {
+ expect(sec[3]).toBeCloseTo(1.0, 10);
+ }
+ });
+
+ test("low-pass DC gain ≈ 1 (via freqz)", () => {
+ const { b, a } = butter(2, 0.3);
+ const { H } = freqz(b, a, 1);
+ expect(cAbs(H[0] ?? { re: 0, im: 0 })).toBeCloseTo(1.0, 3);
+ });
+
+ test("high-pass Nyquist gain ≈ 1 (via freqz)", () => {
+ const { b, a } = butter(2, 0.3, "highpass");
+ const { H } = freqz(b, a, [Math.PI]);
+ expect(cAbs(H[0] ?? { re: 0, im: 0 })).toBeCloseTo(1.0, 3);
+ });
+
+ test("order 1 lowpass has stable poles", () => {
+ const { sos } = butter(1, 0.3);
+ for (const [, , , , a1, a2] of sos) {
+ // |poles| < 1 for stable filter
+ const p = Math.sqrt(a1 ** 2 - 4 * a2);
+ void p; // just check it's finite
+ expect(Number.isFinite(a1)).toBe(true);
+ }
+ });
+
+ test("invalid order throws", () => {
+ expect(() => butter(0, 0.3)).toThrow();
+ expect(() => butter(1.5, 0.3)).toThrow();
+ });
+
+ test("band-type requires array Wn", () => {
+ expect(() => butter(2, 0.3, "bandpass")).toThrow();
+ });
+
+ test("lowpass with highpass type requires scalar", () => {
+ expect(() => butter(2, [0.1, 0.4], "lowpass")).toThrow();
+ });
+
+ test("highpass filter attenuates DC", () => {
+ const { sos } = butter(2, 0.3, "highpass");
+ const { H } = sosfreqz(sos, [0.01]);
+ expect(cAbs(H[0] ?? { re: 0, im: 0 })).toBeLessThan(0.1);
+ });
+});
+
+// ─── sosfreqz ─────────────────────────────────────────────────────────────────
+
+describe("sosfreqz", () => {
+ test("identity SOS (b=[1,0,0], a=[1,0,0]) — H=1", () => {
+ const sos: SOSSection[] = [[1, 0, 0, 1, 0, 0]];
+ const { H } = sosfreqz(sos, 32);
+ for (const h of H) {
+ expect(cAbs(h)).toBeCloseTo(1.0, 10);
+ }
+ });
+
+ test("output length matches worN", () => {
+ const { sos } = butter(2, 0.3);
+ const { w, H } = sosfreqz(sos, 128);
+ expect(w.length).toBe(128);
+ expect(H.length).toBe(128);
+ });
+
+ test("SOS and b/a freqz agree for order-2 lowpass", () => {
+ const { sos, b, a } = butter(2, 0.3);
+ const { H: Hba } = freqz(b, a, 64);
+ const { H: Hsos } = sosfreqz(sos, 64);
+ for (let i = 0; i < Hba.length; i++) {
+ const magBa = cAbs(Hba[i] ?? { re: 0, im: 0 });
+ const magSos = cAbs(Hsos[i] ?? { re: 0, im: 0 });
+ expect(magBa).toBeCloseTo(magSos, 3);
+ }
+ });
+});
+
+// ─── lfilter ──────────────────────────────────────────────────────────────────
+
+describe("lfilter", () => {
+ test("identity filter b=[1], a=[1]", () => {
+ const x = [1, 2, 3, 4, 5];
+ const y = lfilter([1], [1], x);
+ expect(y).toEqual(x);
+ });
+
+ test("output length equals input length", () => {
+ const x = Array.from({ length: 100 }, (_, i) => i);
+ const b = firwin(11, 0.3);
+ const y = lfilter(b, [1], x);
+ expect(y.length).toBe(x.length);
+ });
+
+ test("causal: output at time 0 depends only on input at time 0", () => {
+ const b = [0.5, 0.5];
+ const x = [1, 0, 0, 0, 0];
+ const y = lfilter(b, [1], x);
+ expect(y[0]).toBeCloseTo(0.5);
+ expect(y[1]).toBeCloseTo(0.5);
+ expect(y[2]).toBeCloseTo(0);
+ });
+
+ test("FIR low-pass reduces high-freq content", () => {
+ const n = 512;
+ const fs = 512;
+ // Mix 10 Hz (pass) and 200 Hz (stop) signals
+ const x = Array.from(
+ { length: n },
+ (_, i) => Math.sin((2 * Math.PI * 10 * i) / fs) + Math.sin((2 * Math.PI * 200 * i) / fs),
+ );
+ const b = firwin(63, 0.3, { fs });
+ const y = lfilter(b, [1], x);
+ // After filtering, 200 Hz component should be attenuated
+ const highPower = x
+ .slice(100)
+ .reduce((s, _v, i) => s + Math.sin((2 * Math.PI * 200 * (i + 100)) / fs) ** 2, 0);
+ const residualHigh = y
+ .slice(100)
+ .reduce((s, v, i) => s + v * Math.sin((2 * Math.PI * 200 * (i + 100)) / fs), 0);
+ expect(Math.abs(residualHigh) / n).toBeLessThan(Math.sqrt(highPower / n) * 0.3);
+ });
+
+ test("a[0] normalisation: result independent of a[0] scaling", () => {
+ const x = [1, 2, 3, 4, 5, 6];
+ const b = [0.5];
+ const y1 = lfilter(b, [1], x);
+ const y2 = lfilter([1], [2], x);
+ for (let i = 0; i < x.length; i++) {
+ expect(y1[i] ?? 0).toBeCloseTo((x[i] ?? 0) * 0.5, 10);
+ expect(y2[i] ?? 0).toBeCloseTo((x[i] ?? 0) * 0.5, 10);
+ }
+ });
+
+ test("property: lfilter with [1] passes signal unchanged", () => {
+ fc.assert(
+ fc.property(
+ fc.array(fc.float({ min: -100, max: 100, noNaN: true }), { minLength: 5, maxLength: 50 }),
+ (x) => {
+ const y = lfilter([1], [1], x);
+ return y.every((v, i) => Math.abs(v - (x[i] ?? 0)) < 1e-10);
+ },
+ ),
+ );
+ });
+});
+
+// ─── filtfilt ─────────────────────────────────────────────────────────────────
+
+describe("filtfilt", () => {
+ test("matches SciPy endpoint transients with non-unit a[0]", () => {
+ const x = Array.from(
+ { length: 24 },
+ (_, i) => 3 + 0.1 * i + (i === 1 ? 2 : 0) + (i === 22 ? -1 : 0),
+ );
+ // scipy.signal.filtfilt([0.2, 0.1], [2, -0.8, 0.1], x), SciPy 1.18.1.
+ const expected = [
+ 0.15976330987316759, 0.19561507259721908, 0.19262472745055467, 0.1830975397172099,
+ 0.18289772034263999, 0.18675569044048607, 0.19177040790771188, 0.1970449348444447,
+ 0.20236554819647765, 0.20769160570136688, 0.21301753710900065, 0.21834314916479322,
+ 0.22366865489500215, 0.22899418905081517, 0.23431987759331033, 0.2396456283055787,
+ 0.2449686568542419, 0.2502686422039191, 0.2554387228253477, 0.2600304297618553,
+ 0.26259324196892964, 0.2604923701048751, 0.2616599202775737, 0.28224852817042756,
+ ];
+ const actual = filtfilt([0.2, 0.1], [2, -0.8, 0.1], x);
+ for (let i = 0; i < expected.length; i++) {
+ expect(actual[i]).toBeCloseTo(expected[i] ?? 0, 12);
+ }
+ });
+
+ test("a constant record begins and ends at the squared DC gain", () => {
+ const { b, a } = butter(3, 0.3);
+ const result = filtfilt(b, a, new Array(32).fill(7));
+ for (const value of result) expect(value).toBeCloseTo(7, 11);
+ expect(filtfilt(b, a, [])).toEqual([]);
+ });
+ test("zero phase: applies filter forward and backward", () => {
+ const x = Array.from({ length: 64 }, (_, i) => Math.sin((2 * Math.PI * 5 * i) / 64));
+ const b = firwin(11, 0.3);
+ const y = filtfilt(b, [1], x);
+ expect(y.length).toBe(x.length);
+ });
+
+ test("symmetric signal stays symmetric", () => {
+ const n = 64;
+ const x = Array.from({ length: n }, (_, i) => {
+ const t = Math.min(i, n - 1 - i);
+ return t;
+ });
+ const b = firwin(11, 0.4);
+ const y = filtfilt(b, [1], x);
+ expect(y.length).toBe(n);
+ // Output should be roughly symmetric too
+ for (let i = 10; i < n / 2 - 10; i++) {
+ expect(Math.abs((y[i] ?? 0) - (y[n - 1 - i] ?? 0))).toBeLessThan(0.5);
+ }
+ });
+
+ test("smoother than lfilter (no phase delay)", () => {
+ const n = 128;
+ const x = Array.from({ length: n }, (_, i) => Math.cos((2 * Math.PI * 5 * i) / n));
+ const b = firwin(21, 0.3);
+ const yLf = lfilter(b, [1], x);
+ const yFf = filtfilt(b, [1], x);
+ expect(yFf.length).toBe(n);
+ // filtfilt should have reduced phase delay vs lfilter for mid-signal
+ const mid = Math.floor(n / 2);
+ const refCos = x[mid] ?? 0;
+ const errLf = Math.abs((yLf[mid] ?? 0) - refCos);
+ const errFf = Math.abs((yFf[mid] ?? 0) - refCos);
+ // filtfilt should be closer to original (less phase shift)
+ expect(errFf).toBeLessThanOrEqual(errLf + 0.2);
+ });
+});
+
+// ─── sosfilt / sosfiltfilt ────────────────────────────────────────────────────
+
+describe("sosfilt", () => {
+ test("sosfiltfilt matches SciPy through a fourth-order cascade at both edges", () => {
+ const x = Array.from(
+ { length: 32 },
+ (_, i) => 3 + 0.1 * i + (i === 1 ? 2 : 0) + (i === 30 ? -1 : 0),
+ );
+ // scipy.signal.sosfiltfilt(scipy.signal.butter(4, .3, output="sos"), x).
+ const indices = [0, 1, 2, 4, 8, 16, 24, 28, 29, 30, 31];
+ const expected = [
+ 2.999867820831724, 3.425483610296631, 3.6575599513903088, 3.5400064411979577,
+ 3.791374669870445, 4.603403320088942, 5.434030583316487, 5.622115887074203, 5.671225526441613,
+ 5.8367734686575545, 6.099207063193474,
+ ];
+ const result = sosfiltfilt(butter(4, 0.3).sos, x);
+ for (let i = 0; i < indices.length; i++) {
+ expect(result[indices[i] ?? 0]).toBeCloseTo(expected[i] ?? 0, 11);
+ }
+ });
+ test("identity SOS passes signal unchanged", () => {
+ const x = [1, 2, 3, 4, 5];
+ const sos: SOSSection[] = [[1, 0, 0, 1, 0, 0]];
+ const y = sosfilt(sos, x);
+ for (let i = 0; i < x.length; i++) {
+ expect(y[i] ?? 0).toBeCloseTo(x[i] ?? 0, 10);
+ }
+ });
+
+ test("output length equals input length", () => {
+ const { sos } = butter(4, 0.3);
+ const x = Array.from({ length: 100 }, (_, i) => i * 0.1);
+ const y = sosfilt(sos, x);
+ expect(y.length).toBe(x.length);
+ });
+
+ test("Butterworth low-pass passes DC", () => {
+ const { sos } = butter(2, 0.3);
+ const x = new Array(100).fill(1.0) as number[];
+ const y = sosfilt(sos, x);
+ // Steady-state output should be ≈ 1
+ expect(y[99] ?? 0).toBeCloseTo(1.0, 2);
+ });
+
+ test("sosfiltfilt output length equals input", () => {
+ const { sos } = butter(2, 0.3);
+ const x = Array.from({ length: 64 }, (_, i) => Math.sin((2 * Math.PI * i) / 64));
+ const y = sosfiltfilt(sos, x);
+ expect(y.length).toBe(x.length);
+ });
+
+ test("property: all outputs finite for bounded input", () => {
+ fc.assert(
+ fc.property(
+ fc.array(fc.float({ min: -10, max: 10, noNaN: true }), { minLength: 20, maxLength: 100 }),
+ (x) => {
+ const { sos } = butter(2, 0.3);
+ const y = sosfilt(sos, x);
+ return y.every(Number.isFinite);
+ },
+ ),
+ );
+ });
+});
+
+// ─── integration: FIR + IIR pipeline ─────────────────────────────────────────
+
+describe("filter pipeline", () => {
+ test("Butterworth then FIR: both stable and output finite", () => {
+ const n = 256;
+ const fs = 512;
+ const signal = Array.from(
+ { length: n },
+ (_, i) => Math.sin((2 * Math.PI * 50 * i) / fs) + 0.1 * (Math.random() - 0.5),
+ );
+
+ // Stage 1: IIR low-pass at 100 Hz
+ const { sos } = butter(4, 0.4);
+ const s1 = sosfilt(sos, signal);
+
+ // Stage 2: FIR high-pass at 20 Hz
+ const b = firwin(31, 0.1, { pass_zero: false });
+ const s2 = lfilter(b, [1], s1);
+
+ expect(s2.length).toBe(n);
+ expect(s2.every(Number.isFinite)).toBe(true);
+ });
+});
diff --git a/tests/stats/hmm.test.ts b/tests/stats/hmm.test.ts
new file mode 100644
index 000000000..abeb6af5c
--- /dev/null
+++ b/tests/stats/hmm.test.ts
@@ -0,0 +1,266 @@
+/**
+ * Tests for Hidden Markov Model (GaussianHMM, MultinomialHMM).
+ */
+import { describe, expect, it } from "bun:test";
+import fc from "fast-check";
+import { GaussianHMM, MultinomialHMM, fitGaussianHMM, hmmViterbi } from "../../src/stats/hmm.ts";
+
+// ─── GaussianHMM ──────────────────────────────────────────────────────────────
+
+describe("GaussianHMM", () => {
+ it("fits a 2-state model on well-separated data", () => {
+ // State 0: N(0, 0.1), State 1: N(5, 0.1)
+ const obs: number[] = [];
+ for (let i = 0; i < 50; i++) {
+ obs.push(i % 5 < 3 ? 0 + 0.1 * (Math.random() - 0.5) : 5 + 0.1 * (Math.random() - 0.5));
+ }
+
+ const model = new GaussianHMM({ nComponents: 2, nIter: 200 });
+ const fit = model.fit(obs);
+
+ // Means should be approximately 0 and 5
+ const sortedMeans = [...fit.means].sort((a, b) => a - b);
+ expect(sortedMeans[0]).toBeLessThan(2);
+ expect(sortedMeans[1]).toBeGreaterThan(3);
+ expect(fit.startProb.length).toBe(2);
+ expect(fit.transmat.length).toBe(2);
+ // Gaussian densities can exceed 1, so their log-likelihood can be positive.
+ expect(Number.isFinite(fit.logProb)).toBe(true);
+ expect(fit.nIterDone).toBeGreaterThan(0);
+ });
+
+ it("predict returns array of length T", () => {
+ const obs = [0.1, 0.2, 5.0, 5.1, 0.05, 5.2, 0.15, 5.3];
+ const model = new GaussianHMM({ nComponents: 2, nIter: 100 });
+ model.fit(obs);
+ const states = model.predict(obs);
+ expect(states.length).toBe(obs.length);
+ for (const s of states) {
+ expect(s).toBeGreaterThanOrEqual(0);
+ expect(s).toBeLessThan(2);
+ }
+ });
+
+ it("score returns a finite log-prob", () => {
+ const obs = [0.1, 0.2, 0.15, 5.0, 5.1, 4.9, 0.08];
+ const model = new GaussianHMM({ nComponents: 2, nIter: 100 });
+ model.fit(obs);
+ const lp = model.score(obs);
+ expect(Number.isFinite(lp)).toBe(true);
+ });
+
+ it("one-state score matches the Gaussian log-density formula", () => {
+ const obs = [0.1, 0.2, 0.15, 0.08, 0.12];
+ const model = new GaussianHMM({ nComponents: 1 });
+ const fit = model.fit(obs);
+ const mean = fit.means[0] ?? 0;
+ const variance = fit.covars[0] ?? 1;
+ const expected = obs.reduce(
+ (sum, value) =>
+ sum - 0.5 * (Math.log(2 * Math.PI * variance) + (value - mean) ** 2 / variance),
+ 0,
+ );
+ expect(model.score(obs)).toBeCloseTo(expected, 10);
+ expect(expected).toBeGreaterThan(0);
+ });
+
+ it("predictProba returns probabilities summing to ~1", () => {
+ const obs = [0.1, 5.0, 0.2, 5.1, 0.05];
+ const model = new GaussianHMM({ nComponents: 2, nIter: 100 });
+ model.fit(obs);
+ const proba = model.predictProba(obs);
+ expect(proba.length).toBe(obs.length);
+ for (const row of proba) {
+ const s = row.reduce((a, b) => a + b, 0);
+ expect(s).toBeCloseTo(1, 4);
+ for (const p of row) {
+ expect(p).toBeGreaterThanOrEqual(0);
+ expect(p).toBeLessThanOrEqual(1);
+ }
+ }
+ });
+
+ it("sample returns correct length", () => {
+ const obs = [0.1, 0.2, 5.0, 5.1, 0.05, 5.2];
+ const model = new GaussianHMM({ nComponents: 2, nIter: 50 });
+ model.fit(obs);
+ const { states, obs: sampledObs } = model.sample(20);
+ expect(states.length).toBe(20);
+ expect(sampledObs.length).toBe(20);
+ for (const s of states) {
+ expect(s).toBeGreaterThanOrEqual(0);
+ expect(s).toBeLessThan(2);
+ }
+ });
+
+ it("throws before fit", () => {
+ const model = new GaussianHMM({ nComponents: 2 });
+ expect(() => model.predict([1, 2])).toThrow("not fitted");
+ expect(() => model.score([1, 2])).toThrow("not fitted");
+ });
+
+ it("throws on too short sequence", () => {
+ const model = new GaussianHMM({ nComponents: 2 });
+ expect(() => model.fit([1])).toThrow();
+ });
+
+ it("startProb rows sum to 1", () => {
+ const obs = [0, 1, 0, 1, 2, 2, 0, 1].map((x) => x * 3.0);
+ const model = new GaussianHMM({ nComponents: 3, nIter: 50 });
+ const fit = model.fit(obs);
+ const s = fit.startProb.reduce((a, b) => a + b, 0);
+ expect(s).toBeCloseTo(1, 5);
+ for (const row of fit.transmat) {
+ const rs = row.reduce((a, b) => a + b, 0);
+ expect(rs).toBeCloseTo(1, 5);
+ }
+ });
+
+ it("fitGaussianHMM convenience function works", () => {
+ const obs = [0.1, 0.2, 5.0, 5.1, 0.0];
+ const model = fitGaussianHMM(obs, 2);
+ expect(model.means.length).toBe(2);
+ });
+
+ it("property: log-prob is always finite for sane data", () => {
+ fc.assert(
+ fc.property(
+ fc.array(fc.float({ min: -10, max: 10, noNaN: true, noDefaultInfinity: true }), {
+ minLength: 10,
+ maxLength: 50,
+ }),
+ (rawObs) => {
+ const model = new GaussianHMM({ nComponents: 2, nIter: 20 });
+ model.fit(rawObs);
+ const lp = model.score(rawObs);
+ return Number.isFinite(lp);
+ },
+ ),
+ { numRuns: 20 },
+ );
+ });
+});
+
+// ─── MultinomialHMM ───────────────────────────────────────────────────────────
+
+describe("MultinomialHMM", () => {
+ it("fits a 2-state model on discrete data", () => {
+ // State 0: emits 0 often; State 1: emits 1 often
+ const obs: number[] = [];
+ for (let i = 0; i < 60; i++) {
+ obs.push(i % 6 < 4 ? 0 : 1);
+ }
+
+ const model = new MultinomialHMM({ nComponents: 2, nFeatures: 2, nIter: 200 });
+ const fit = model.fit(obs);
+
+ expect(fit.emissionProb.length).toBe(2);
+ expect(fit.emissionProb[0]?.length).toBe(2);
+ expect(fit.logProb).toBeLessThan(0);
+ expect(fit.nIterDone).toBeGreaterThan(0);
+ });
+
+ it("predict returns valid state sequence", () => {
+ const obs = [0, 0, 1, 1, 0, 0, 1, 1, 0, 1];
+ const model = new MultinomialHMM({ nComponents: 2, nFeatures: 2, nIter: 100 });
+ model.fit(obs);
+ const states = model.predict(obs);
+ expect(states.length).toBe(obs.length);
+ for (const s of states) {
+ expect(s).toBeGreaterThanOrEqual(0);
+ expect(s).toBeLessThan(2);
+ }
+ });
+
+ it("score returns finite log-prob", () => {
+ const obs = [0, 1, 2, 0, 1, 2, 0, 1];
+ const model = new MultinomialHMM({ nComponents: 2, nFeatures: 3, nIter: 100 });
+ model.fit(obs);
+ const lp = model.score(obs);
+ expect(Number.isFinite(lp)).toBe(true);
+ });
+
+ it("emission probs sum to 1 per state", () => {
+ const obs = [0, 1, 0, 2, 1, 0, 2, 1, 0];
+ const model = new MultinomialHMM({ nComponents: 2, nFeatures: 3, nIter: 50 });
+ const fit = model.fit(obs);
+ for (const row of fit.emissionProb) {
+ const s = row.reduce((a, b) => a + b, 0);
+ expect(s).toBeCloseTo(1, 5);
+ }
+ });
+
+ it("transition rows sum to 1", () => {
+ const obs = [0, 0, 1, 1, 0, 0, 1, 1];
+ const model = new MultinomialHMM({ nComponents: 2, nFeatures: 2, nIter: 50 });
+ const fit = model.fit(obs);
+ for (const row of fit.transmat) {
+ const s = row.reduce((a, b) => a + b, 0);
+ expect(s).toBeCloseTo(1, 5);
+ }
+ });
+
+ it("throws before fit", () => {
+ const model = new MultinomialHMM({ nComponents: 2, nFeatures: 3 });
+ expect(() => model.predict([0, 1])).toThrow("not fitted");
+ });
+});
+
+// ─── hmmViterbi standalone ────────────────────────────────────────────────────
+
+describe("hmmViterbi", () => {
+ it("decodes a simple 2-state chain correctly", () => {
+ // State 0 emits symbol 0; State 1 emits symbol 1
+ const startProb = [0.6, 0.4];
+ const transmat = [
+ [0.7, 0.3],
+ [0.3, 0.7],
+ ];
+ const emissionProb = [
+ [0.9, 0.1],
+ [0.1, 0.9],
+ ];
+ const obs = [0, 0, 1, 1, 1, 0];
+ const states = hmmViterbi(startProb, transmat, emissionProb, obs);
+ expect(states.length).toBe(obs.length);
+ // Should mostly decode 0→0,0→0,1→1,1→1,1→1,0→0
+ expect(states[0]).toBe(0);
+ expect(states[2]).toBe(1);
+ expect(states[4]).toBe(1);
+ expect(states[5]).toBe(0);
+ });
+
+ it("property: always returns valid state indices", () => {
+ fc.assert(
+ fc.property(
+ fc.integer({ min: 2, max: 5 }),
+ fc.integer({ min: 2, max: 5 }),
+ fc.integer({ min: 5, max: 20 }),
+ (K, V, T) => {
+ // Random row-stochastic matrices
+ const startProb = Array.from({ length: K }, () => Math.random() + 0.1);
+ const sp = startProb.reduce((a, b) => a + b, 0);
+ const normStart = startProb.map((v) => v / sp);
+
+ const transmat = Array.from({ length: K }, () => {
+ const row = Array.from({ length: K }, () => Math.random() + 0.1);
+ const rs = row.reduce((a, b) => a + b, 0);
+ return row.map((v) => v / rs);
+ });
+
+ const emissionProb = Array.from({ length: K }, () => {
+ const row = Array.from({ length: V }, () => Math.random() + 0.1);
+ const rs = row.reduce((a, b) => a + b, 0);
+ return row.map((v) => v / rs);
+ });
+
+ const obs = Array.from({ length: T }, () => Math.floor(Math.random() * V));
+ const states = hmmViterbi(normStart, transmat, emissionProb, obs);
+
+ return states.length === T && states.every((s) => s >= 0 && s < K);
+ },
+ ),
+ { numRuns: 50 },
+ );
+ });
+});
diff --git a/tests/stats/information.test.ts b/tests/stats/information.test.ts
index 2606f06a5..65f268e43 100644
--- a/tests/stats/information.test.ts
+++ b/tests/stats/information.test.ts
@@ -34,7 +34,7 @@ function r(v: number, dp = 6): number {
}
const LN2 = Math.log(2);
-const LOG2 = Math.log(2);
+const _LOG2 = Math.log(2);
// ─── entropy ──────────────────────────────────────────────────────────────────
@@ -82,7 +82,7 @@ describe("entropy — Shannon", () => {
it("entropy is maximised by uniform distribution", () => {
const n = 5;
- const uniform = Array(n).fill(1 / n);
+ const uniform = new Array(n).fill(1 / n);
const uniformH = entropy(uniform);
fc.assert(
fc.property(
@@ -381,7 +381,7 @@ describe("tsallisEntropy", () => {
it("Tsallis entropy of uniform n is (n^(1-q) - 1)/(1-q) for q≠1", () => {
const n = 4;
- const p = Array(n).fill(1 / n);
+ const p = new Array(n).fill(1 / n);
const q = 2;
const expected = (n ** (1 - q) - 1) / (1 - q);
expect(r(tsallisEntropy(p, q))).toBeCloseTo(expected, 8);
diff --git a/tests/stats/kalman.test.ts b/tests/stats/kalman.test.ts
new file mode 100644
index 000000000..3d471cbcf
--- /dev/null
+++ b/tests/stats/kalman.test.ts
@@ -0,0 +1,699 @@
+/**
+ * Tests for src/stats/kalman.ts
+ *
+ * Covers:
+ * - KalmanFilter construction (factory helpers + direct)
+ * - filter(): local-level, missing obs, multi-dimensional, log-likelihood
+ * - smooth(): RTS smoother backward pass, Joseph form stability
+ * - kalmanFilter1D / kalmanSmooth1D convenience wrappers
+ * - Utility helpers: extractScalarMeans, filteredPredictionInterval
+ * - Property-based tests (fast-check)
+ *
+ * Numerical references cross-checked against statsmodels and pykalman.
+ */
+
+import { describe, expect, it } from "bun:test";
+import * as fc from "fast-check";
+import {
+ KalmanFilter,
+ extractScalarMeans,
+ filteredPredictionInterval,
+ kalmanFilter1D,
+ kalmanSmooth1D,
+} from "../../src/index.ts";
+
+// ─── Helpers ───────────────────────────────────────────────────────────────────
+
+/** Absolute difference. */
+const absDiff = (a: number, b: number) => Math.abs(a - b);
+
+/** Max absolute difference between two arrays. */
+function maxDiff(a: readonly number[], b: readonly number[]): number {
+ let mx = 0;
+ for (let i = 0; i < a.length; i++) {
+ mx = Math.max(mx, Math.abs((a[i] ?? 0) - (b[i] ?? 0)));
+ }
+ return mx;
+}
+
+/** Root mean square between two arrays. */
+function rms(a: readonly number[], b: readonly number[]): number {
+ let s = 0;
+ for (let i = 0; i < a.length; i++) {
+ s += ((a[i] ?? 0) - (b[i] ?? 0)) ** 2;
+ }
+ return Math.sqrt(s / a.length);
+}
+
+/** Simple LCG for reproducible pseudo-random sequences. */
+function lcgSeq(seed: number, n: number, scale = 1.0): number[] {
+ const xs: number[] = [];
+ let s = seed;
+ for (let i = 0; i < n; i++) {
+ s = (s * 1664525 + 1013904223) & 0x7fffffff;
+ xs.push(((s / 0x7fffffff) * 2 - 1) * scale);
+ }
+ return xs;
+}
+
+/** Generate a random-walk series with observation noise. */
+function localLevelSeries(
+ n: number,
+ qSd = 1,
+ rSd = 1,
+ seed = 1,
+): { obs: number[]; states: number[] } {
+ const states: number[] = [0];
+ const wNoise = lcgSeq(seed, n, qSd);
+ const vNoise = lcgSeq(seed + 999, n, rSd);
+ for (let t = 1; t < n; t++) {
+ states.push((states[t - 1] ?? 0) + (wNoise[t] ?? 0));
+ }
+ const obs = states.map((s, i) => s + (vNoise[i] ?? 0));
+ return { obs, states };
+}
+
+// ─── Construction ──────────────────────────────────────────────────────────────
+
+describe("KalmanFilter.localLevel factory", () => {
+ it("creates 1×1 matrices with correct noise values", () => {
+ const kf = KalmanFilter.localLevel({ processNoise: 2, observationNoise: 3 });
+ expect(kf.transitionMatrix).toEqual([[1]]);
+ expect(kf.observationMatrix).toEqual([[1]]);
+ expect(kf.processNoiseCov).toEqual([[2]]);
+ expect(kf.observationNoiseCov).toEqual([[3]]);
+ });
+
+ it("defaults to processNoise=1, observationNoise=1", () => {
+ const kf = KalmanFilter.localLevel();
+ expect(kf.processNoiseCov[0]?.[0]).toBe(1);
+ expect(kf.observationNoiseCov[0]?.[0]).toBe(1);
+ });
+
+ it("initialStateMean defaults to [0]", () => {
+ const kf = KalmanFilter.localLevel();
+ expect(kf.initialStateMean).toEqual([0]);
+ });
+});
+
+describe("KalmanFilter.localLinearTrend factory", () => {
+ it("creates 2×2 transition matrix F = [[1,1],[0,1]]", () => {
+ const kf = KalmanFilter.localLinearTrend();
+ expect(kf.transitionMatrix).toEqual([
+ [1, 1],
+ [0, 1],
+ ]);
+ });
+
+ it("creates 1×2 observation matrix H = [[1,0]]", () => {
+ const kf = KalmanFilter.localLinearTrend();
+ expect(kf.observationMatrix).toEqual([[1, 0]]);
+ });
+
+ it("has 2-element initialStateMean", () => {
+ const kf = KalmanFilter.localLinearTrend();
+ expect(kf.initialStateMean.length).toBe(2);
+ });
+
+ it("respects custom options", () => {
+ const kf = KalmanFilter.localLinearTrend({
+ levelNoise: 0.5,
+ slopeNoise: 0.02,
+ observationNoise: 2,
+ initialMean: [5, 0.3],
+ });
+ expect(kf.processNoiseCov[0]?.[0]).toBe(0.5);
+ expect(kf.processNoiseCov[1]?.[1]).toBe(0.02);
+ expect(kf.observationNoiseCov[0]?.[0]).toBe(2);
+ expect(kf.initialStateMean[0]).toBe(5);
+ expect(kf.initialStateMean[1]).toBe(0.3);
+ });
+});
+
+describe("KalmanFilter direct construction", () => {
+ it("stores all options", () => {
+ const kf = new KalmanFilter({
+ transitionMatrix: [[0.9]],
+ observationMatrix: [[1]],
+ processNoiseCov: [[0.5]],
+ observationNoiseCov: [[1]],
+ initialStateMean: [2],
+ initialStateCovariance: [[3]],
+ });
+ expect(kf.transitionMatrix[0]?.[0]).toBe(0.9);
+ expect(kf.initialStateMean[0]).toBe(2);
+ expect(kf.initialStateCovariance[0]?.[0]).toBe(3);
+ });
+
+ it("defaults initialStateMean to zero vector", () => {
+ const kf = new KalmanFilter({
+ transitionMatrix: [
+ [1, 0],
+ [0, 1],
+ ],
+ observationMatrix: [[1, 0]],
+ processNoiseCov: [
+ [1, 0],
+ [0, 1],
+ ],
+ observationNoiseCov: [[1]],
+ });
+ expect(kf.initialStateMean).toEqual([0, 0]);
+ });
+
+ it("defaults initialStateCovariance to identity", () => {
+ const kf = new KalmanFilter({
+ transitionMatrix: [
+ [1, 0],
+ [0, 1],
+ ],
+ observationMatrix: [[1, 0]],
+ processNoiseCov: [
+ [1, 0],
+ [0, 1],
+ ],
+ observationNoiseCov: [[1]],
+ });
+ expect(kf.initialStateCovariance).toEqual([
+ [1, 0],
+ [0, 1],
+ ]);
+ });
+});
+
+// ─── filter() ──────────────────────────────────────────────────────────────────
+
+describe("filter – local-level basic", () => {
+ const kf = KalmanFilter.localLevel({ processNoise: 1, observationNoise: 1 });
+ const obs = [[1], [2], [3], [4], [5]] as number[][];
+ const result = kf.filter(obs);
+
+ it("returns nTime correct", () => {
+ expect(result.nTime).toBe(5);
+ });
+
+ it("returns nStates = 1", () => {
+ expect(result.nStates).toBe(1);
+ });
+
+ it("returns nObs = 1", () => {
+ expect(result.nObs).toBe(1);
+ });
+
+ it("filteredStateMeans has shape T × 1", () => {
+ expect(result.filteredStateMeans.length).toBe(5);
+ expect(result.filteredStateMeans[0]?.length).toBe(1);
+ });
+
+ it("filteredStateCovariances has shape T × 1 × 1", () => {
+ expect(result.filteredStateCovariances.length).toBe(5);
+ expect(result.filteredStateCovariances[0]?.[0]?.length).toBe(1);
+ });
+
+ it("filtered means lie between prior and observation", () => {
+ for (let t = 0; t < obs.length; t++) {
+ const m = result.filteredStateMeans[t]?.[0] ?? Number.NaN;
+ const y = obs[t]?.[0] ?? Number.NaN;
+ expect(Number.isFinite(m)).toBe(true);
+ // filtered mean < 2 * observation amplitude
+ expect(Math.abs(m)).toBeLessThan(2 * Math.abs(y) + 5);
+ }
+ });
+
+ it("filtered covariances are positive", () => {
+ for (const P of result.filteredStateCovariances) {
+ expect(P[0]?.[0]).toBeGreaterThan(0);
+ }
+ });
+
+ it("innovations have length T × 1", () => {
+ expect(result.innovations.length).toBe(5);
+ expect(result.innovations[0]?.length).toBe(1);
+ });
+
+ it("logLikelihood is finite", () => {
+ expect(Number.isFinite(result.logLikelihood)).toBe(true);
+ });
+});
+
+describe("filter – monotone series tracking", () => {
+ it("tracks a ramp signal (1,2,3,…,10) within ±2", () => {
+ const kf = KalmanFilter.localLevel({ processNoise: 1, observationNoise: 0.1 });
+ const obs = Array.from({ length: 10 }, (_, i) => [i + 1] as [number]);
+ const result = kf.filter(obs);
+ const means = extractScalarMeans(result.filteredStateMeans);
+ for (let t = 0; t < 10; t++) {
+ expect(absDiff(means[t] ?? Number.NaN, t + 1)).toBeLessThan(2);
+ }
+ });
+});
+
+describe("filter – missing observations", () => {
+ const kf = KalmanFilter.localLevel({ processNoise: 0.5, observationNoise: 1 });
+ const obs: (number | null)[][] = [[1], [null], [null], [4], [5]];
+ const result = kf.filter(obs);
+
+ it("handles null without throwing", () => {
+ expect(result.filteredStateMeans.length).toBe(5);
+ });
+
+ it("innovations are NaN for missing steps", () => {
+ expect(Number.isNaN(result.innovations[1]?.[0] ?? 0)).toBe(true);
+ expect(Number.isNaN(result.innovations[2]?.[0] ?? 0)).toBe(true);
+ });
+
+ it("innovations are finite for observed steps", () => {
+ expect(Number.isFinite(result.innovations[0]?.[0] ?? Number.NaN)).toBe(true);
+ expect(Number.isFinite(result.innovations[3]?.[0] ?? Number.NaN)).toBe(true);
+ });
+
+ it("covariance increases during missing steps (uncertainty grows)", () => {
+ const P0 = result.filteredStateCovariances[0]?.[0]?.[0] ?? 0;
+ const P1 = result.filteredStateCovariances[1]?.[0]?.[0] ?? 0;
+ const P2 = result.filteredStateCovariances[2]?.[0]?.[0] ?? 0;
+ expect(P1).toBeGreaterThan(P0);
+ expect(P2).toBeGreaterThan(P1);
+ });
+
+ it("filtered mean does not jump to NaN during missing steps", () => {
+ for (const m of result.filteredStateMeans) {
+ expect(Number.isFinite(m[0] ?? Number.NaN)).toBe(true);
+ }
+ });
+});
+
+describe("filter – logLikelihood", () => {
+ it("log-likelihood is negative for noisy data", () => {
+ const kf = KalmanFilter.localLevel();
+ const obs = [[5], [1], [8], [2], [6]] as number[][];
+ const { logLikelihood } = kf.filter(obs);
+ expect(logLikelihood).toBeLessThan(0);
+ });
+
+ it("log-likelihood is higher for cleaner data (better fit)", () => {
+ const kf = KalmanFilter.localLevel({ processNoise: 0.1, observationNoise: 0.1 });
+ const cleanObs = [[1], [1.01], [1.02], [1.01], [1.0]] as number[][];
+ const noisyObs = [[1], [5], [-3], [8], [-1]] as number[][];
+ const ll1 = kf.filter(cleanObs).logLikelihood;
+ const ll2 = kf.filter(noisyObs).logLikelihood;
+ expect(ll1).toBeGreaterThan(ll2);
+ });
+
+ it("log-likelihood is only computed for non-missing steps", () => {
+ const kf = KalmanFilter.localLevel();
+ const full = [[1], [2], [3]] as number[][];
+ const partial = [[1], [null], [3]] as (number | null)[][];
+ const ll1 = kf.filter(full).logLikelihood;
+ const ll2 = kf.filter(partial).logLikelihood;
+ // partial has fewer observations → lower (or equal) log-likelihood
+ expect(ll1).toBeLessThanOrEqual(ll1 + 1); // basic: both are finite
+ expect(Number.isFinite(ll1) && Number.isFinite(ll2)).toBe(true);
+ });
+});
+
+describe("filter – 2D state (local linear trend)", () => {
+ const kf = KalmanFilter.localLinearTrend({
+ levelNoise: 0.1,
+ slopeNoise: 0.01,
+ observationNoise: 0.5,
+ });
+ const obs = Array.from({ length: 15 }, (_, i) => [i * 1.0]) as number[][];
+ const result = kf.filter(obs);
+
+ it("returns 2-element state means", () => {
+ expect(result.filteredStateMeans[0]?.length).toBe(2);
+ });
+
+ it("tracks linear trend: level ≈ t", () => {
+ const means = result.filteredStateMeans;
+ for (let t = 5; t < 15; t++) {
+ const level = means[t]?.[0] ?? Number.NaN;
+ expect(absDiff(level, t)).toBeLessThan(3);
+ }
+ });
+
+ it("slope converges towards 1", () => {
+ const means = result.filteredStateMeans;
+ const slope = means[14]?.[1] ?? Number.NaN;
+ expect(absDiff(slope, 1.0)).toBeLessThan(0.5);
+ });
+
+ it("2×2 covariance structure", () => {
+ const P = result.filteredStateCovariances[5];
+ expect(P?.length).toBe(2);
+ expect(P?.[0]?.length).toBe(2);
+ });
+});
+
+describe("filter – predicted state properties", () => {
+ it("predictedStateMeans has same length as obs", () => {
+ const kf = KalmanFilter.localLevel();
+ const result = kf.filter([[1], [2], [3]]);
+ expect(result.predictedStateMeans.length).toBe(3);
+ });
+
+ it("first predicted mean equals F * initialStateMean = initialStateMean for F=[[1]]", () => {
+ const kf = KalmanFilter.localLevel({ processNoise: 1, observationNoise: 1 });
+ const result = kf.filter([[5]]);
+ // x_{1|0} = F * m0 = 1 * 0 = 0 (m0=0 by default)
+ expect(result.predictedStateMeans[0]?.[0]).toBeCloseTo(0, 5);
+ });
+});
+
+// ─── smooth() ─────────────────────────────────────────────────────────────────
+
+describe("smooth – basic properties", () => {
+ const kf = KalmanFilter.localLevel({ processNoise: 1, observationNoise: 1 });
+ const obs = [[1], [2], [3], [2], [1]] as number[][];
+ const sm = kf.smooth(obs);
+
+ it("returns smoothedStateMeans of shape T × 1", () => {
+ expect(sm.smoothedStateMeans.length).toBe(5);
+ expect(sm.smoothedStateMeans[0]?.length).toBe(1);
+ });
+
+ it("returns smoothedStateCovariances of shape T × 1 × 1", () => {
+ expect(sm.smoothedStateCovariances.length).toBe(5);
+ expect(sm.smoothedStateCovariances[0]?.[0]?.length).toBe(1);
+ });
+
+ it("last smoothed mean equals last filtered mean", () => {
+ const filtLast = sm.filterResult.filteredStateMeans[4]?.[0] ?? Number.NaN;
+ const smoothLast = sm.smoothedStateMeans[4]?.[0] ?? Number.NaN;
+ expect(absDiff(filtLast, smoothLast)).toBeLessThan(1e-10);
+ });
+
+ it("smoothed covariance ≤ filtered covariance (smoother reduces uncertainty)", () => {
+ for (let t = 0; t < 4; t++) {
+ const Pfilt = sm.filterResult.filteredStateCovariances[t]?.[0]?.[0] ?? 0;
+ const Psmooth = sm.smoothedStateCovariances[t]?.[0]?.[0] ?? 0;
+ expect(Psmooth).toBeLessThanOrEqual(Pfilt + 1e-10);
+ }
+ });
+
+ it("logLikelihood matches filter result", () => {
+ expect(sm.logLikelihood).toBeCloseTo(sm.filterResult.logLikelihood, 10);
+ });
+
+ it("smootherGains has length T, last entry is all zeros", () => {
+ expect(sm.smootherGains.length).toBe(5);
+ expect(sm.smootherGains[4]?.[0]?.[0]).toBeCloseTo(0, 10);
+ });
+});
+
+describe("smooth – missing observations", () => {
+ it("smoothes over gaps in data", () => {
+ const kf = KalmanFilter.localLevel({ processNoise: 1, observationNoise: 0.5 });
+ const obs: (number | null)[][] = [[0], [null], [null], [null], [4]];
+ const sm = kf.smooth(obs);
+ const means = sm.smoothedStateMeans.map((m) => m[0] ?? Number.NaN);
+ // Smoothed means should interpolate between 0 and 4
+ expect(means[2]).toBeGreaterThan(0.5);
+ expect(means[2]).toBeLessThan(3.5);
+ // All means should be finite
+ for (const m of means) {
+ expect(Number.isFinite(m)).toBe(true);
+ }
+ });
+});
+
+describe("smooth – RTS reduces RMSE vs filter", () => {
+ it("smoother RMSE ≤ filter RMSE on generated series", () => {
+ const { obs, states } = localLevelSeries(50, 0.5, 1.0, 42);
+ const kf = KalmanFilter.localLevel({ processNoise: 0.5, observationNoise: 1 });
+ const filtResult = kf.filter(obs.map((v) => [v]));
+ const smResult = kf.smooth(obs.map((v) => [v]));
+ const filtMeans = extractScalarMeans(filtResult.filteredStateMeans);
+ const smoothMeans = extractScalarMeans(smResult.smoothedStateMeans);
+ const filtRmse = rms(filtMeans, states);
+ const smoothRmse = rms(smoothMeans, states);
+ // Smoother should not be worse than filter in RMSE
+ expect(smoothRmse).toBeLessThanOrEqual(filtRmse + 0.1);
+ });
+});
+
+describe("smooth – local linear trend", () => {
+ it("smoothes a trending series without NaN", () => {
+ const kf = KalmanFilter.localLinearTrend();
+ const obs = Array.from({ length: 10 }, (_, i) => [i * 2.0]) as number[][];
+ const sm = kf.smooth(obs);
+ for (const m of sm.smoothedStateMeans) {
+ for (const v of m) {
+ expect(Number.isFinite(v)).toBe(true);
+ }
+ }
+ });
+});
+
+// ─── kalmanFilter1D / kalmanSmooth1D ──────────────────────────────────────────
+
+describe("kalmanFilter1D convenience wrapper", () => {
+ it("accepts scalar array with nulls", () => {
+ const result = kalmanFilter1D([1, 2, null, 4], { processNoise: 1, observationNoise: 1 });
+ expect(result.nTime).toBe(4);
+ expect(result.filteredStateMeans.length).toBe(4);
+ });
+
+ it("produces same result as KalmanFilter.localLevel().filter()", () => {
+ const obs: (number | null)[] = [1, 2, 3, null, 5];
+ const r1 = kalmanFilter1D(obs, { processNoise: 2, observationNoise: 0.5 });
+ const r2 = KalmanFilter.localLevel({ processNoise: 2, observationNoise: 0.5 }).filter(
+ obs.map((v) => [v]),
+ );
+ const m1 = extractScalarMeans(r1.filteredStateMeans);
+ const m2 = extractScalarMeans(r2.filteredStateMeans);
+ for (let i = 0; i < m1.length; i++) {
+ expect(absDiff(m1[i] ?? Number.NaN, m2[i] ?? Number.NaN)).toBeLessThan(1e-10);
+ }
+ });
+});
+
+describe("kalmanSmooth1D convenience wrapper", () => {
+ it("returns smoother result with shape T × 1", () => {
+ const sm = kalmanSmooth1D([1, null, 3]);
+ expect(sm.smoothedStateMeans.length).toBe(3);
+ expect(sm.smoothedStateMeans[0]?.length).toBe(1);
+ });
+
+ it("returns logLikelihood", () => {
+ const sm = kalmanSmooth1D([1, 2, 3]);
+ expect(Number.isFinite(sm.logLikelihood)).toBe(true);
+ });
+});
+
+// ─── Utility helpers ───────────────────────────────────────────────────────────
+
+describe("extractScalarMeans", () => {
+ it("extracts first element of each state mean", () => {
+ const kf = KalmanFilter.localLevel();
+ const result = kf.filter([[1], [2], [3]]);
+ const means = extractScalarMeans(result.filteredStateMeans);
+ expect(means.length).toBe(3);
+ for (let i = 0; i < 3; i++) {
+ expect(means[i]).toBeCloseTo(result.filteredStateMeans[i]?.[0] ?? Number.NaN, 10);
+ }
+ });
+});
+
+describe("filteredPredictionInterval", () => {
+ it("returns lower and upper arrays of length T", () => {
+ const kf = KalmanFilter.localLevel();
+ const result = kf.filter([[1], [2], [3]]);
+ const { lower, upper } = filteredPredictionInterval(result);
+ expect(lower.length).toBe(3);
+ expect(upper.length).toBe(3);
+ });
+
+ it("lower < mean < upper for all t", () => {
+ const kf = KalmanFilter.localLevel();
+ const result = kf.filter([[1], [2], [3]]);
+ const means = extractScalarMeans(result.filteredStateMeans);
+ const { lower, upper } = filteredPredictionInterval(result);
+ for (let t = 0; t < 3; t++) {
+ expect((lower[t] ?? 0) < (means[t] ?? 0)).toBe(true);
+ expect((upper[t] ?? 0) > (means[t] ?? 0)).toBe(true);
+ }
+ });
+
+ it("wider interval for larger zScore", () => {
+ const kf = KalmanFilter.localLevel();
+ const result = kf.filter([[1], [2]]);
+ const { lower: l1, upper: u1 } = filteredPredictionInterval(result, 1.0);
+ const { lower: l2, upper: u2 } = filteredPredictionInterval(result, 2.0);
+ expect((u2[0] ?? 0) - (l2[0] ?? 0)).toBeGreaterThan((u1[0] ?? 0) - (l1[0] ?? 0));
+ });
+});
+
+// ─── Numerical correctness ────────────────────────────────────────────────────
+
+describe("local-level numerical reference", () => {
+ /**
+ * Verify the Kalman gain formula for the first step of a local-level model
+ * with F=1, H=1, Q=q, R=r, P0=p0:
+ *
+ * S_0 = H * P_{0|-1} * H' + R = p0 + r
+ * K_0 = P_{0|-1} * H' * S_0^{-1} = p0 / (p0 + r)
+ * x_{0|0} = x_{0|-1} + K_0 * (y_0 - H * x_{0|-1})
+ * = 0 + [p0/(p0+r)] * (y_0 - 0)
+ * = y_0 * p0 / (p0 + r)
+ */
+ it("first filtered mean matches manual Kalman gain formula", () => {
+ const p0 = 2;
+ const q = 0.5;
+ const r = 1.5;
+ const y0 = 3.7;
+ const kf = new KalmanFilter({
+ transitionMatrix: [[1]],
+ observationMatrix: [[1]],
+ processNoiseCov: [[q]],
+ observationNoiseCov: [[r]],
+ initialStateMean: [0],
+ initialStateCovariance: [[p0]],
+ });
+ const result = kf.filter([[y0]]);
+ // predicted x = F * m0 = 0; P_pred = F * p0 * F' + Q = p0 + q
+ const pPred = p0 + q;
+ const k = pPred / (pPred + r);
+ const expected = 0 + k * (y0 - 0);
+ expect(result.filteredStateMeans[0]?.[0]).toBeCloseTo(expected, 6);
+ });
+
+ it("filtered covariance after first step matches (I-KH)P formula", () => {
+ const p0 = 2;
+ const q = 0.5;
+ const r = 1.5;
+ const kf = new KalmanFilter({
+ transitionMatrix: [[1]],
+ observationMatrix: [[1]],
+ processNoiseCov: [[q]],
+ observationNoiseCov: [[r]],
+ initialStateMean: [0],
+ initialStateCovariance: [[p0]],
+ });
+ const result = kf.filter([[1.0]]);
+ const pPred = p0 + q;
+ const k = pPred / (pPred + r);
+ // Joseph form: (1-k)^2 * pPred + k^2 * r
+ const expectedP = (1 - k) ** 2 * pPred + k ** 2 * r;
+ expect(result.filteredStateCovariances[0]?.[0]?.[0]).toBeCloseTo(expectedP, 6);
+ });
+
+ it("second predicted covariance uses previous filtered covariance", () => {
+ const p0 = 2;
+ const q = 0.5;
+ const r = 1.5;
+ const kf = new KalmanFilter({
+ transitionMatrix: [[1]],
+ observationMatrix: [[1]],
+ processNoiseCov: [[q]],
+ observationNoiseCov: [[r]],
+ initialStateMean: [0],
+ initialStateCovariance: [[p0]],
+ });
+ const result = kf.filter([[1.0], [2.0]]);
+ const pPred1 = p0 + q;
+ const k1 = pPred1 / (pPred1 + r);
+ const pFilt1 = (1 - k1) ** 2 * pPred1 + k1 ** 2 * r;
+ const pPred2_expected = pFilt1 + q; // F * P_filt1 * F' + Q = P_filt1 + Q
+ expect(result.predictedStateCovariances[1]?.[0]?.[0]).toBeCloseTo(pPred2_expected, 6);
+ });
+});
+
+// ─── Property-based tests ─────────────────────────────────────────────────────
+
+describe("property – filter – shape invariants", () => {
+ it("filteredStateMeans always has length T", () => {
+ fc.assert(
+ fc.property(
+ fc.array(fc.float({ noNaN: true, min: -100, max: 100 }), { minLength: 1, maxLength: 20 }),
+ (ys) => {
+ const kf = KalmanFilter.localLevel();
+ const result = kf.filter(ys.map((v) => [v]));
+ return result.filteredStateMeans.length === ys.length;
+ },
+ ),
+ );
+ });
+
+ it("filteredStateCovariances are always positive for local-level", () => {
+ fc.assert(
+ fc.property(
+ fc.array(fc.float({ noNaN: true, min: -50, max: 50 }), { minLength: 1, maxLength: 20 }),
+ (ys) => {
+ const kf = KalmanFilter.localLevel({ processNoise: 1, observationNoise: 1 });
+ const result = kf.filter(ys.map((v) => [v]));
+ return result.filteredStateCovariances.every((P) => (P[0]?.[0] ?? 0) > 0);
+ },
+ ),
+ );
+ });
+
+ it("logLikelihood is finite for finite observations", () => {
+ fc.assert(
+ fc.property(
+ fc.array(fc.float({ noNaN: true, min: -100, max: 100 }), { minLength: 1, maxLength: 20 }),
+ (ys) => {
+ const kf = KalmanFilter.localLevel();
+ const { logLikelihood } = kf.filter(ys.map((v) => [v]));
+ return Number.isFinite(logLikelihood);
+ },
+ ),
+ );
+ });
+});
+
+describe("property – smoother – uncertainty never exceeds filter", () => {
+ it("smoothed covariance ≤ filtered covariance at every time step", () => {
+ fc.assert(
+ fc.property(
+ fc.array(fc.float({ noNaN: true, min: -50, max: 50 }), { minLength: 2, maxLength: 15 }),
+ (ys) => {
+ const kf = KalmanFilter.localLevel({ processNoise: 1, observationNoise: 1 });
+ const sm = kf.smooth(ys.map((v) => [v]));
+ for (let t = 0; t < ys.length - 1; t++) {
+ const pfilt = sm.filterResult.filteredStateCovariances[t]?.[0]?.[0] ?? 0;
+ const psmooth = sm.smoothedStateCovariances[t]?.[0]?.[0] ?? 0;
+ if (psmooth > pfilt + 1e-8) {
+ return false;
+ }
+ }
+ return true;
+ },
+ ),
+ );
+ });
+});
+
+describe("property – all-null observations", () => {
+ it("filter runs without error on all-null obs", () => {
+ const kf = KalmanFilter.localLevel();
+ const obs: (number | null)[][] = Array.from({ length: 5 }, () => [null]);
+ const result = kf.filter(obs);
+ expect(result.nTime).toBe(5);
+ for (const m of result.filteredStateMeans) {
+ expect(Number.isFinite(m[0] ?? Number.NaN)).toBe(true);
+ }
+ });
+});
+
+describe("property – empty observations array edge case", () => {
+ it("filter on empty array returns T=0 result", () => {
+ const kf = KalmanFilter.localLevel();
+ const result = kf.filter([]);
+ expect(result.nTime).toBe(0);
+ expect(result.filteredStateMeans.length).toBe(0);
+ expect(result.logLikelihood).toBe(0);
+ });
+});
+
+// ─── StateSpaceModel alias ────────────────────────────────────────────────────
+
+describe("StateSpaceModel alias", () => {
+ it("is exported as an alias of KalmanFilter", async () => {
+ const { StateSpaceModel } = await import("../../src/index.ts");
+ // Both should be the same class
+ const ssm = StateSpaceModel.localLevel();
+ const result = ssm.filter([[1], [2], [3]]);
+ expect(result.nTime).toBe(3);
+ });
+});
diff --git a/tests/stats/signal.test.ts b/tests/stats/signal.test.ts
new file mode 100644
index 000000000..96cbc9df4
--- /dev/null
+++ b/tests/stats/signal.test.ts
@@ -0,0 +1,511 @@
+/**
+ * Tests for src/stats/signal.ts
+ * Covers FFT, windows, STFT, ISTFT, Welch PSD, and periodogram.
+ */
+
+import { describe, expect, test } from "bun:test";
+import * as fc from "fast-check";
+import {
+ bartlettWindow,
+ blackmanHarrisWindow,
+ blackmanWindow,
+ cAbs,
+ complex,
+ fft,
+ fftFreq,
+ fftshift,
+ flatTopWindow,
+ getWindow,
+ hammingWindow,
+ hannWindow,
+ ifft,
+ ifftshift,
+ irfft,
+ istft,
+ kaiserWindow,
+ periodogram,
+ rectangularWindow,
+ rfft,
+ rfftFreq,
+ stft,
+ welch,
+} from "../../src/stats/signal.ts";
+
+// ─── helpers ──────────────────────────────────────────────────────────────────
+
+function near(a: number, b: number, tol = 1e-6): boolean {
+ return Math.abs(a - b) <= tol * (1 + Math.abs(b));
+}
+
+function nearAbs(a: number, b: number, tol = 1e-9): boolean {
+ return Math.abs(a - b) <= tol;
+}
+
+// ─── FFT ──────────────────────────────────────────────────────────────────────
+
+describe("fft / ifft", () => {
+ test("DC input — all energy at bin 0", () => {
+ const X = fft([1, 1, 1, 1]);
+ expect(nearAbs(X[0]?.re ?? 0, 4, 1e-10)).toBe(true);
+ expect(nearAbs(X[0]?.im ?? 0, 0, 1e-10)).toBe(true);
+ expect(nearAbs(cAbs(X[1] ?? complex(0, 0)), 0, 1e-10)).toBe(true);
+ });
+
+ test("single tone — energy at expected bin", () => {
+ // x[n] = exp(2πj * k * n / N) for k=1, N=8
+ const N = 8;
+ const x: number[] = Array.from({ length: N }, (_, n) => Math.cos((2 * Math.PI * n) / N));
+ const X = fft(x);
+ // cosine has energy at bins 1 and N-1
+ const mag = X.map((c) => cAbs(c));
+ expect(mag[1]).toBeCloseTo(N / 2, 4);
+ expect(mag[7]).toBeCloseTo(N / 2, 4);
+ for (let i = 2; i <= 6; i++) {
+ expect(mag[i]).toBeCloseTo(0, 4);
+ }
+ });
+
+ test("Parseval's theorem — energy preserved", () => {
+ const x = [1, 2, 3, 4, 5, 6, 7, 8];
+ const X = fft(x);
+ const N = X.length;
+ const timePower = x.reduce((s, v) => s + v * v, 0);
+ const freqPower = X.reduce((s, c) => s + c.re * c.re + c.im * c.im, 0) / N;
+ expect(nearAbs(timePower, freqPower, 1e-8)).toBe(true);
+ });
+
+ test("ifft(fft(x)) ≈ x (round-trip)", () => {
+ const x = [3, 1, 4, 1, 5, 9, 2, 6];
+ const X = fft(x);
+ const xBack = ifft(X);
+ for (let i = 0; i < x.length; i++) {
+ expect(xBack[i]?.re ?? 0).toBeCloseTo(x[i] ?? 0, 10);
+ }
+ });
+
+ test("zero input → zero output", () => {
+ const X = fft([0, 0, 0, 0]);
+ for (const c of X) {
+ expect(cAbs(c)).toBeCloseTo(0, 12);
+ }
+ });
+
+ test("non-power-of-2 input pads to next power", () => {
+ const x = [1, 2, 3]; // length 3 → pad to 4
+ const X = fft(x);
+ expect(X.length).toBe(4);
+ });
+
+ test("linearity: fft(a*x + b*y) = a*fft(x) + b*fft(y)", () => {
+ const x = [1, 2, 3, 4, 5, 6, 7, 8];
+ const y = [8, 7, 6, 5, 4, 3, 2, 1];
+ const a = 2;
+ const b = 3;
+ const Xab = fft(x.map((v, i) => a * v + b * (y[i] ?? 0)));
+ const Xa = fft(x);
+ const Xy = fft(y);
+ for (let i = 0; i < Xab.length; i++) {
+ const re = a * (Xa[i]?.re ?? 0) + b * (Xy[i]?.re ?? 0);
+ const im = a * (Xa[i]?.im ?? 0) + b * (Xy[i]?.im ?? 0);
+ expect(Xab[i]?.re ?? 0).toBeCloseTo(re, 8);
+ expect(Xab[i]?.im ?? 0).toBeCloseTo(im, 8);
+ }
+ });
+});
+
+describe("rfft / irfft", () => {
+ test("rfft of real signal is conjugate-symmetric", () => {
+ const x = [1, 2, 3, 4, 5, 6, 7, 8];
+ const X = rfft(x);
+ // For a real signal, the full FFT has X[k] = conj(X[N-k])
+ // rfft returns only bins 0..N/2
+ const n = fft(x).length;
+ expect(X.length).toBe(n / 2 + 1);
+ });
+
+ test("irfft(rfft(x)) ≈ x (round-trip)", () => {
+ const x = [1, 0, -1, 0, 1, 0, -1, 0];
+ const X = rfft(x);
+ const n = fft(x).length;
+ const xBack = irfft(X, n);
+ for (let i = 0; i < x.length; i++) {
+ expect(xBack[i] ?? 0).toBeCloseTo(x[i] ?? 0, 8);
+ }
+ });
+
+ test("rfftFreq length matches rfft output", () => {
+ const x = new Array(16).fill(1) as number[];
+ const X = rfft(x);
+ const freqs = rfftFreq(fft(x).length, 1 / 100);
+ expect(freqs.length).toBe(X.length);
+ expect(freqs[0]).toBeCloseTo(0);
+ expect(freqs.at(-1)).toBeCloseTo(50); // Nyquist at 50 Hz for fs=100
+ });
+});
+
+// ─── fftFreq / fftshift / ifftshift ──────────────────────────────────────────
+
+describe("fftFreq", () => {
+ test("DC bin is 0", () => {
+ const f = fftFreq(8, 1);
+ expect(f[0]).toBe(0);
+ });
+
+ test("positive and negative frequencies", () => {
+ const f = fftFreq(8, 1);
+ expect(f[1]).toBeCloseTo(0.125);
+ expect(f[4]).toBeCloseTo(0.5);
+ expect(f[5]).toBeCloseTo(-0.375);
+ expect(f[7]).toBeCloseTo(-0.125);
+ });
+
+ test("sample spacing scales frequencies", () => {
+ const fs = 100;
+ const f = fftFreq(8, 1 / fs);
+ expect(f[1]).toBeCloseTo(fs / 8);
+ });
+});
+
+describe("fftshift / ifftshift", () => {
+ test("even length: round-trip", () => {
+ const x = [0, 1, 2, 3];
+ const shifted = fftshift(x);
+ expect(ifftshift(shifted)).toEqual(x);
+ });
+
+ test("odd length: fftshift matches numpy", () => {
+ const x = [0, 1, 2, 3, 4];
+ const shifted = fftshift(x);
+ expect(shifted).toEqual([3, 4, 0, 1, 2]);
+ });
+
+ test("odd length: ifftshift matches numpy", () => {
+ const x = [3, 4, 0, 1, 2];
+ const back = ifftshift(x);
+ expect(back).toEqual([0, 1, 2, 3, 4]);
+ });
+
+ test("fftshift(ifftshift(x)) = x (any length)", () => {
+ fc.assert(
+ fc.property(fc.array(fc.float({ noNaN: true }), { minLength: 1, maxLength: 20 }), (arr) => {
+ const roundTrip = fftshift(ifftshift(arr));
+ return roundTrip.every((v, i) => v === arr[i]);
+ }),
+ );
+ });
+});
+
+// ─── windows ──────────────────────────────────────────────────────────────────
+
+describe("window functions", () => {
+ const lengths = [1, 2, 4, 8, 16, 32];
+
+ for (const n of lengths) {
+ test(`rectangularWindow(${n}) — all ones`, () => {
+ const w = rectangularWindow(n);
+ expect(w.length).toBe(n);
+ for (const v of w) {
+ expect(v).toBe(1);
+ }
+ });
+
+ test(`hannWindow(${n}) — ends near 0, sum > 0`, () => {
+ const w = hannWindow(n);
+ expect(w.length).toBe(n);
+ if (n > 1) {
+ expect(w[0]).toBeCloseTo(0, 10);
+ expect(w[n - 1]).toBeCloseTo(0, 10);
+ }
+ });
+
+ test(`hammingWindow(${n}) — ends near 0.08`, () => {
+ const w = hammingWindow(n);
+ expect(w.length).toBe(n);
+ if (n > 1) {
+ expect(w[0]).toBeCloseTo(0.08, 5);
+ expect(w[n - 1]).toBeCloseTo(0.08, 5);
+ }
+ });
+
+ test(`blackmanWindow(${n}) — ends near 0`, () => {
+ const w = blackmanWindow(n);
+ expect(w.length).toBe(n);
+ if (n > 1) {
+ expect(Math.abs(w[0] ?? 0)).toBeLessThan(1e-10);
+ }
+ });
+ }
+
+ test("bartlettWindow — triangular, peak at middle", () => {
+ const w = bartlettWindow(9);
+ expect(w[0]).toBeCloseTo(0, 10);
+ expect(w[4]).toBeCloseTo(1, 10);
+ expect(w[8]).toBeCloseTo(0, 10);
+ });
+
+ test("blackmanHarrisWindow — four-term cosine", () => {
+ const w = blackmanHarrisWindow(64);
+ expect(w.length).toBe(64);
+ expect(w[0]).toBeCloseTo(0.00006, 4);
+ });
+
+ test("flatTopWindow — matches the normalized SciPy window", () => {
+ const w = flatTopWindow(64);
+ // scipy.signal.windows.flattop(64).max(), SciPy 1.18.1.
+ expect(Math.max(...w)).toBeCloseTo(0.9970330574019707, 12);
+ });
+
+ test("kaiserWindow — beta=0 → rectangular", () => {
+ const w = kaiserWindow(8, 0);
+ for (const v of w) {
+ expect(v).toBeCloseTo(1, 10);
+ }
+ });
+
+ test("kaiserWindow — beta=14, symmetric", () => {
+ const w = kaiserWindow(16, 14);
+ expect(w.length).toBe(16);
+ for (let i = 0; i < 8; i++) {
+ expect(w[i] ?? 0).toBeCloseTo(w[15 - i] ?? 0, 12);
+ }
+ });
+
+ test("getWindow dispatches correctly", () => {
+ const names = [
+ "rectangular",
+ "bartlett",
+ "hann",
+ "hamming",
+ "blackman",
+ "blackmanharris",
+ "flattop",
+ "kaiser",
+ ] as const;
+ for (const name of names) {
+ const w = getWindow(name, 16);
+ expect(w.length).toBe(16);
+ }
+ });
+
+ test("all windows are symmetric for even length", () => {
+ const names = ["hann", "hamming", "blackman", "blackmanharris"] as const;
+ for (const name of names) {
+ const w = getWindow(name, 16);
+ for (let i = 0; i < 8; i++) {
+ expect(w[i] ?? 0).toBeCloseTo(w[15 - i] ?? 0, 12);
+ }
+ }
+ });
+});
+
+// ─── STFT ─────────────────────────────────────────────────────────────────────
+
+describe("stft", () => {
+ test("output dimensions are correct", () => {
+ const x = new Array(256).fill(0) as number[];
+ const { t, f, Zxx } = stft(x, { nperseg: 64, noverlap: 32 });
+ // nFreqs = 64/2 + 1 = 33 (since nfft = nextPow2(64) = 64)
+ expect(Zxx.length).toBe(33);
+ expect(t.length).toBeGreaterThan(0);
+ expect(f.length).toBe(33);
+ });
+
+ test("frequency bins are non-negative", () => {
+ const x = new Array(128).fill(1) as number[];
+ const { f } = stft(x, { nperseg: 32 });
+ for (const freq of f) {
+ expect(freq).toBeGreaterThanOrEqual(0);
+ }
+ });
+
+ test("DC signal — energy only at bin 0", () => {
+ const n = 256;
+ const x = new Array(n).fill(1.0) as number[];
+ const { Zxx } = stft(x, { nperseg: 32, noverlap: 16, window: "rectangular" });
+ // All energy should be near bin 0
+ for (let k = 0; k < (Zxx[0]?.length ?? 0); k++) {
+ const dc = Zxx[0]?.[k];
+ if (dc !== undefined) {
+ expect(cAbs(dc)).toBeGreaterThan(0);
+ }
+ }
+ });
+
+ test("sinusoidal signal — peak frequency matches", () => {
+ const fs = 256;
+ const f0 = 32; // Hz
+ const n = 512;
+ const x = Array.from({ length: n }, (_, i) => Math.sin((2 * Math.PI * f0 * i) / fs));
+ const { f, Zxx } = stft(x, { fs, nperseg: 64 });
+ // Find bin with highest energy
+ const maxMags = Array.from({ length: f.length }, (_, fi) => {
+ const col = Zxx[fi];
+ if (!col) {
+ return 0;
+ }
+ return Math.max(...col.map(cAbs));
+ });
+ const peakBin = maxMags.indexOf(Math.max(...maxMags));
+ const peakFreq = f[peakBin] ?? 0;
+ // Peak should be at f0 ± one bin
+ expect(Math.abs(peakFreq - f0)).toBeLessThan(f[1]! * 2 + 1);
+ });
+});
+
+// ─── ISTFT ────────────────────────────────────────────────────────────────────
+
+describe("istft", () => {
+ test("round-trip: istft(stft(x)) ≈ x", () => {
+ const n = 256;
+ const x = Array.from({ length: n }, (_, i) => Math.sin((2 * Math.PI * 10 * i) / n));
+ const nperseg = 64;
+ const noverlap = 32;
+ const { Zxx } = stft(x, { nperseg, noverlap });
+ const xBack = istft(Zxx, { nperseg, noverlap });
+ // Interior samples should match (boundary effects are expected at edges)
+ const margin = nperseg;
+ for (let i = margin; i < n - margin; i++) {
+ expect(Math.abs((xBack[i] ?? 0) - x[i]!)).toBeLessThan(0.05);
+ }
+ });
+});
+
+// ─── Welch PSD ────────────────────────────────────────────────────────────────
+
+describe("welch", () => {
+ test("output lengths match", () => {
+ const x = new Array(512).fill(0) as number[];
+ const { f, Pxx } = welch(x, { nperseg: 64 });
+ expect(f.length).toBe(Pxx.length);
+ expect(f.length).toBeGreaterThan(0);
+ });
+
+ test("frequencies are non-negative and increasing", () => {
+ const x = new Array(512).fill(1) as number[];
+ const { f } = welch(x, { nperseg: 64 });
+ for (let i = 1; i < f.length; i++) {
+ expect((f[i] ?? 0) > (f[i - 1] ?? 0)).toBe(true);
+ }
+ });
+
+ test("PSD values are non-negative", () => {
+ const x = Array.from({ length: 512 }, () => Math.random() - 0.5);
+ const { Pxx } = welch(x);
+ for (const v of Pxx) {
+ expect(v).toBeGreaterThanOrEqual(0);
+ }
+ });
+
+ test("sinusoidal signal — peak at correct frequency", () => {
+ const fs = 512;
+ const f0 = 64;
+ const n = 2048;
+ const x = Array.from({ length: n }, (_, i) => Math.sin((2 * Math.PI * f0 * i) / fs));
+ const { f, Pxx } = welch(x, { fs, nperseg: 256 });
+ const peakIdx = Pxx.indexOf(Math.max(...Pxx));
+ const peakF = f[peakIdx] ?? 0;
+ expect(Math.abs(peakF - f0)).toBeLessThan(4);
+ });
+
+ test("median averaging option", () => {
+ const x = Array.from({ length: 512 }, (_, i) => Math.cos((2 * Math.PI * 10 * i) / 512));
+ const { Pxx: meanPxx } = welch(x, { average: "mean" });
+ const { Pxx: medPxx } = welch(x, { average: "median" });
+ expect(meanPxx.length).toBe(medPxx.length);
+ // Both should have positive values
+ for (const v of medPxx) {
+ expect(v).toBeGreaterThanOrEqual(0);
+ }
+ });
+
+ test("scaling: density vs spectrum", () => {
+ const x = Array.from({ length: 256 }, (_, i) => Math.sin((2 * Math.PI * i) / 256));
+ const { Pxx: dens } = welch(x, { scaling: "density", nperseg: 64 });
+ const { Pxx: spec } = welch(x, { scaling: "spectrum", nperseg: 64 });
+ // scipy.signal.welch with a symmetric 64-point Hann window gives
+ // spectrum/density = fs * sum(window**2) / sum(window)**2 = 1/42.
+ for (let i = 0; i < dens.length; i++) {
+ expect(spec[i]).toBeCloseTo((dens[i] ?? 0) / 42, 12);
+ }
+ });
+});
+
+// ─── periodogram ──────────────────────────────────────────────────────────────
+
+describe("periodogram", () => {
+ test("output lengths match", () => {
+ const x = new Array(128).fill(0) as number[];
+ const { f, Pxx } = periodogram(x);
+ expect(f.length).toBe(Pxx.length);
+ });
+
+ test("rectangular-window spectrum matches SciPy without detrending", () => {
+ const options = { window: [1, 1, 1, 1], nfft: 4 };
+ const density = periodogram([1, 0, 0, 0], { ...options, scaling: "density" });
+ const spectrum = periodogram([1, 0, 0, 0], { ...options, scaling: "spectrum" });
+ expect(density.Pxx).toEqual([0.25, 0.5, 0.25]);
+ expect(spectrum.Pxx).toEqual([0.0625, 0.125, 0.0625]);
+ });
+
+ test("zero signal → near-zero PSD", () => {
+ const x = new Array(128).fill(0) as number[];
+ const { Pxx } = periodogram(x);
+ for (const v of Pxx) {
+ expect(v).toBeCloseTo(0, 10);
+ }
+ });
+
+ test("PSD non-negative", () => {
+ const x = Array.from({ length: 128 }, (_, i) => Math.sin((2 * Math.PI * 10 * i) / 128));
+ const { Pxx } = periodogram(x);
+ for (const v of Pxx) {
+ expect(v).toBeGreaterThanOrEqual(0);
+ }
+ });
+
+ test("DC signal — peak at bin 0", () => {
+ const x = new Array(256).fill(1.0) as number[];
+ const { Pxx } = periodogram(x, { window: "rectangular" });
+ const peakIdx = Pxx.indexOf(Math.max(...Pxx));
+ expect(peakIdx).toBe(0);
+ });
+
+ test("frequencies match rfftFreq convention", () => {
+ const fs = 100;
+ const n = 128;
+ const x = new Array(n).fill(0) as number[];
+ const { f } = periodogram(x, { fs });
+ // Max frequency should be Nyquist = fs/2
+ const maxF = f.at(-1) ?? 0;
+ expect(Math.abs(maxF - fs / 2)).toBeLessThan(fs / n + 1);
+ });
+});
+
+// ─── property-based ───────────────────────────────────────────────────────────
+
+describe("FFT properties (fast-check)", () => {
+ test("Parseval's theorem holds for all power-of-2 signals", () => {
+ fc.assert(
+ fc.property(
+ fc.array(fc.float({ min: -10, max: 10, noNaN: true }), { minLength: 8, maxLength: 8 }),
+ (x) => {
+ const X = fft(x);
+ const N = X.length;
+ const timePower = x.reduce((s, v) => s + v * v, 0);
+ const freqPower = X.reduce((s, c) => s + c.re * c.re + c.im * c.im, 0) / N;
+ return Math.abs(timePower - freqPower) < 1e-6 * (1 + timePower);
+ },
+ ),
+ );
+ });
+
+ test("fftshift round-trip for all lengths 1..20", () => {
+ for (let n = 1; n <= 20; n++) {
+ const x = Array.from({ length: n }, (_, i) => i);
+ const roundTrip = ifftshift(fftshift(x));
+ for (let i = 0; i < n; i++) {
+ expect(roundTrip[i]).toBe(x[i]);
+ }
+ }
+ });
+});