+ autocorr(x, lag) computes the Pearson correlation between
+ x[0..n−lag−1] and x[lag..n−1], exactly like
+ pd.Series.autocorr(lag).
+
+
+
+ JavaScript
+
+
+
+
+
+
+
Click ▶ Run to execute
+
Ctrl+Enter to run · Tab to indent
+
+
+
+
+
2 — Full ACF with Bartlett confidence intervals
+
+ acf(x, { nlags, alpha }) returns all autocorrelations at lags 0…nlags.
+ With alpha=0.05, Bartlett confidence intervals are returned: lags whose
+ CI excludes zero are statistically significant.
+
+
+
+ JavaScript
+
+
+
+
+
+
+
Click ▶ Run to execute
+
Ctrl+Enter to run · Tab to indent
+
+
+
+
+
3 — Partial ACF (Levinson-Durbin)
+
+ pacf(x, { nlags, alpha }) uses the Levinson-Durbin recursion to compute
+ partial autocorrelations. For a true AR(p) process, only the first p PACF values are
+ significantly non-zero — this is how you identify the AR order.
+
+
+
+ JavaScript
+
+
+
+
+
+
+
Click ▶ Run to execute
+
Ctrl+Enter to run · Tab to indent
+
+
+
+
+
4 — Cross-Correlation Function (CCF)
+
+ ccf(x, y, { nlags, alpha }) measures the linear relationship between
+ x[t] and y[t+k] at each lag k. Peaks in the CCF reveal
+ lead/lag relationships between two series.
+
+
+
+ JavaScript
+
+
+
+
+
+
+
Click ▶ Run to execute
+
Ctrl+Enter to run · Tab to indent
+
+
+
+
+
5 — Durbin-Watson statistic
+
+ durbinWatson(residuals) tests for first-order autocorrelation in OLS residuals.
+ Values near 2 indicate no autocorrelation; values near 0 indicate positive autocorrelation;
+ values near 4 indicate negative autocorrelation.
+
+
+
+ JavaScript
+
+
+
+
+
+
+
Click ▶ Run to execute
+
Ctrl+Enter to run · Tab to indent
+
+
+
+
+
6 — Ljung-Box & Box-Pierce portmanteau tests
+
+ ljungBox(x, { lags }) and boxPierce(x, { lags }) test the null
+ hypothesis that no autocorrelation exists up to a given lag. Small p-values reject the
+ white-noise hypothesis. Ljung-Box has better finite-sample properties.
+
+ ARIMA(p, d, q) time-series model — estimation, forecasting, and prediction intervals.
+ Mirrors statsmodels.tsa.arima.model.ARIMA.
+
+
+
+
+
1 — Fit an AR(1) model
+
+ new ARIMAModel({ p, d, q }) constructs the model;
+ .fit(y) estimates the parameters using the Hannan-Rissanen
+ two-step method and returns coefficients, sigma², AIC, and BIC.
+
+
+
+ JavaScript
+
+
+
+
+
+
+
Click ▶ Run to execute
+
Ctrl+Enter to run · Tab to indent
+
+
+
+
+
+
2 — Multi-step forecast with prediction intervals
+
+ model.forecast(steps) returns point forecasts and 95 % prediction
+ intervals computed via ψ-weight recursion.
+
+
+
+ JavaScript
+
+
+
+
+
+
+
Click ▶ Run to execute
+
Ctrl+Enter to run · Tab to indent
+
+
+
+
+
+
3 — ARMA(1,1) model
+
+ Combine AR and MA terms. ARMA(1,1): x_t = φ x_{t−1} + θ ε_{t−1} + ε_t.
+
+
+
+ JavaScript
+
+
+
+
+
+
+
Click ▶ Run to execute
+
Ctrl+Enter to run · Tab to indent
+
+
+
+
+
+
4 — ARIMA(1,1,0): integrated series
+
+ Set d=1 for I(1) series (random walk, stock prices, etc.).
+ The model differences the series before fitting.
+
+
+
+ JavaScript
+
+
+
+
+
+
+
Click ▶ Run to execute
+
Ctrl+Enter to run · Tab to indent
+
+
+
+
+
+
5 — fitArima convenience function
+
+ fitArima(y, opts) is a one-liner shorthand for constructing
+ and fitting an ARIMA model.
+
+ Simple Exponential Smoothing, Holt linear trend, and full Holt-Winters seasonal models.
+ Mirrors statsmodels.tsa.holtwinters.ExponentialSmoothing.
+
+
+
+
+
1 — Simple Exponential Smoothing (SES)
+
+ SimpleExpSmoothing fits ETS(A,N,N): level-only smoothing with
+ parameter α. All h-step forecasts equal the final level. α is estimated by
+ minimising SSE via Nelder-Mead.
+
+
+
+ JavaScript
+
+
+
+
+
+
+
Click ▶ Run to execute
+
Ctrl+Enter to run · Tab to indent
+
+
+
+
+
+
2 — Holt's Linear Trend (Double Exponential Smoothing)
+
+ Holt extends SES with a trend component β (ETS(A,A,N)).
+ Optionally damps the trend with φ (ETS(A,Ad,N)) to prevent over-shooting.
+
+
+
+ JavaScript
+
+
+
+
+
+
+
Click ▶ Run to execute
+
Ctrl+Enter to run · Tab to indent
+
+
+
+
+
+
3 — Holt-Winters Additive Seasonal
+
+ ExponentialSmoothing with trend: "add" and
+ seasonal: "add" models data with a linear trend plus additive
+ seasonal fluctuations. Classic Holt-Winters ETS(A,A,A).
+
+
+
+ JavaScript
+
+
+
+
+
+
+
Click ▶ Run to execute
+
Ctrl+Enter to run · Tab to indent
+
+
+
+
+
+
4 — Holt-Winters Multiplicative Seasonal
+
+ Use seasonal: "mul" when the amplitude of seasonal swings
+ grows with the level (common in economic time series). ETS(A,A,M).
+
+
+
+ JavaScript
+
+
+
+
+
+
+
Click ▶ Run to execute
+
Ctrl+Enter to run · Tab to indent
+
+
+
+
+
+
5 — Forecast with prediction intervals
+
+ forecastWithCI(steps, alpha) returns point forecasts plus
+ (1 − α) % prediction intervals. Intervals widen with forecast horizon.
+
+
+
+ JavaScript
+
+
+
+
+
+
+
Click ▶ Run to execute
+
Ctrl+Enter to run · Tab to indent
+
+
+
+
+
+
6 — Model selection via AIC/BIC
+
+ Compare SES, Holt, and Holt-Winters using information criteria. Lower AIC/BIC
+ indicates a better balance of fit and parsimony.
+
+
+
+ JavaScript
+
+
+
+
+
+
+
Click ▶ Run to execute
+
Ctrl+Enter to run · Tab to indent
+
+
+
+
+
+
7 — Known initialisation and fixed parameters
+
+ You can supply fixed smoothing parameters or initial state values.
+ Use initializationMethod: "known" to set the initial state
+ directly without estimating it.
+
+ FIR and IIR filter design and application — mirrors scipy.signal.
+ Edit any code block below and press ▶ Run
+ (or Ctrl+Enter) to execute it live in your browser.
+
+
+
+
+
1. Low-pass FIR with firwin
+
Design a 51-tap Hamming-windowed low-pass FIR and check its DC and Nyquist gain.
+
+
+ TypeScript
+
+
+
+
+
+
+
+
Click ▶ Run to execute
+
Ctrl+Enter to run · Tab to indent
+
+
+
+
+
+
2. High-pass FIR
+
Pass high frequencies by setting pass_zero: false. DC gain should be near 0, Nyquist near 1.
+
+
+ TypeScript
+
+
+
+
+
+
+
+
Click ▶ Run to execute
+
Ctrl+Enter to run · Tab to indent
+
+
+
+
+
+
3. Apply FIR: lfilter vs filtfilt
+
lfilter is causal (has phase delay); filtfilt applies the filter twice for zero-phase output.
+
+
+ TypeScript
+
+
+
+
+
+
+
+
Click ▶ Run to execute
+
Ctrl+Enter to run · Tab to indent
+
+
+
+
+
+
4. Butterworth low-pass filter design
+
Design a 4th-order Butterworth IIR filter. butter returns both SOS form (numerically preferred) and ba form.
+
+
+ TypeScript
+
+
+
+
+
+
+
+
Click ▶ Run to execute
+
Ctrl+Enter to run · Tab to indent
+
+
+
+
+
+
5. Apply Butterworth: sosfilt vs sosfiltfilt
+
Filter a noisy 50 Hz signal to remove 300 Hz interference. Zero-phase output is closer to the clean reference.
+
+
+ TypeScript
+
+
+
+
+
+
+
+
Click ▶ Run to execute
+
Ctrl+Enter to run · Tab to indent
+
+
+
+
+
+
6. High-pass Butterworth
+
A 2nd-order Butterworth high-pass attenuates DC and passes high frequencies.
+
+
+ TypeScript
+
+
+
+
+
+
+
+
Click ▶ Run to execute
+
Ctrl+Enter to run · Tab to indent
+
+
+
+
+
+
7. Butterworth gain at cutoff — multiple orders
+
All Butterworth filters have exactly −3 dB gain at the cutoff frequency, regardless of order.
Apache ORC (Optimized Row Columnar) file format reader and writer. Supports BOOLEAN, INT/LONG, FLOAT/DOUBLE, STRING columns with NONE compression and RLE v1 / direct encoding. Mirrors pandas.read_orc() and DataFrame.to_orc().
SimpleExpSmoothing, Holt, and ExponentialSmoothing — SES, Holt linear trend, and full Holt-Winters with additive/multiplicative seasonal components. Nelder-Mead parameter optimisation, AIC/BIC/AICc, prediction intervals. Mirrors statsmodels.tsa.holtwinters.ExponentialSmoothing.
+ 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.
+
+
+
+ example-1.ts
+
+
+
+
+
+
+
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.
+
+
+
+ example-2.ts
+
+
+
+
+
+
+
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.
+
+
+
+ example-3.ts
+
+
+
+
+
+
+
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.
+
+
+
+ example-4.ts
+
+
+
+
+
+
+
Click ▶ Run to execute
+
+
+
+
+
+
🔢 Multi-Dimensional Observations
+
+ The Kalman filter naturally handles multi-dimensional observations.
+ Here: 2 sensors observing a single latent state.
+
+ 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.
+
+ 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:
+ 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.
+
+
+ TypeScript
+
+
+
+
+
+
+
+
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.
+
+
+ TypeScript
+
+
+
+
+
+
+
+
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.
+
+
+ TypeScript
+
+
+
+
+
+
+
+
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.
+
+
+ TypeScript
+
+
+
+
+
+
+
+
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.
+
+
+ TypeScript
+
+
+
+
+
+
+
+
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.
+
+
+ TypeScript
+
+
+
+
+
+
+
+
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.
+
+
+ TypeScript
+
+
+
+
+
+
+
+
Click ▶ Run to execute
+
Ctrl+Enter to run · Tab to indent
+
+
+
+
+
+
8. Periodogram
+
A single-segment PSD estimate — higher variance but simpler than Welch.
+
+
+ TypeScript
+
+
+
+
+
+
+
+
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.
+
+
+ TypeScript
+
+
+
+
+
+
+
+
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 bb13fcd1..80da2047 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,136 @@ 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";
+
diff --git a/src/io/index.ts b/src/io/index.ts
index 194e405d..78cf80dc 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 00000000..09ae760a
--- /dev/null
+++ b/src/io/orc.ts
@@ -0,0 +1,1213 @@
+/**
+ * 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;
+ for (;;) {
+ const b = buf[pos];
+ if (b === undefined) throw new Error("ORC: truncated varint");
+ pos++;
+ result |= BigInt(b & 0x7f) << shift;
+ if ((b & 0x80) === 0) break;
+ shift += 7n;
+ }
+ return [result, pos];
+}
+
+/** 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");
+ pos++;
+ // Interpret as signed byte
+ const sfb = fb >= 0x80 ? fb - 0x100 : fb;
+ // Single-byte range: -112 to 127
+ if (sfb >= -112) return [BigInt(sfb), pos];
+ // 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[pos];
+ if (b === undefined) throw new Error("ORC: truncated Hadoop VInt data");
+ pos++;
+ value = (value << 8n) | BigInt(b);
+ }
+ if (isNeg) value ^= -1n;
+ return [value, pos];
+}
+
+/** 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 Number.isInteger(v) ? 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[buf.length - 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(typeKind === KIND_DATE ? Number(v ?? 0) : Number(v ?? 0));
+ 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
+ let 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 00000000..582dfba4
--- /dev/null
+++ b/src/io/read_avro.ts
@@ -0,0 +1,615 @@
+/**
+ * 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: number = 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/stats/acf_pacf.ts b/src/stats/acf_pacf.ts
new file mode 100644
index 00000000..f83c8b45
--- /dev/null
+++ b/src/stats/acf_pacf.ts
@@ -0,0 +1,640 @@
+/**
+ * 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[lagList.length - 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 00000000..d6a3857b
--- /dev/null
+++ b/src/stats/arima.ts
@@ -0,0 +1,552 @@
+/**
+ * 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 = 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: number = 0;
+ private _sigma2: number = 1;
+ private _origY: readonly number[] = [];
+ private _inits: readonly number[] = [];
+ private _diffW: readonly number[] = [];
+ private _residuals: readonly number[] = [];
+ private _fitted: boolean = 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 ──────────────────
+ const kMax = 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[w.length - 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 = 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/ets.ts b/src/stats/ets.ts
new file mode 100644
index 00000000..477722af
--- /dev/null
+++ b/src/stats/ets.ts
@@ -0,0 +1,1281 @@
+/**
+ * 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 (Array.isArray(y)) return y;
+ return y.values;
+}
+
+/** 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: number = 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] ?? Infinity) < (fvals[bestIdx] ?? Infinity)) bestIdx = i;
+ }
+ return { params: simplex[bestIdx] ?? [], value: fvals[bestIdx] ?? 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 : Math.pow(b, phi)) + st_m;
+ } else {
+ // mul trend + mul seasonal
+ yhat = l * (phi === 1 ? b : Math.pow(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 * Math.pow(b, phiSum) + sVal;
+ } else {
+ yhat = l * Math.pow(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);
+ } else {
+ 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: number = 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 {
+ // 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: number = 0;
+ private _finalB: number = 0;
+ private _phi: number = 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;
+ let ll0 = optimL0 ? (params[paramNames.indexOf("l0")] ?? l0h) : l0Fixed;
+ let lb0 = optimB0 ? (params[paramNames.indexOf("b0")] ?? b0h) : b0Fixed;
+ a = clamp(a, 1e-6, 1 - 1e-6);
+ bta = clamp(bta, 1e-6, 1 - 1e-6);
+ 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: number = 0;
+ private _finalB: number = 0;
+ private _finalS: number[] = [];
+ private _phi: number = 1;
+ private _trend: ETSTrend = null;
+ private _seasonal: ETSSeasonal = null;
+ private _m: number = 1;
+ private _n: number = 0;
+ private _sigma2: number = 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, 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);
+
+ let ll0 = optimL0 ? (params[paramNames.indexOf("l0")] ?? l0h) : (l0Fixed ?? l0h);
+ let 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));
+ }
+ }
+
+ // 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: number = 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 -Infinity;
+ if (p >= 1) return 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/filters.ts b/src/stats/filters.ts
new file mode 100644
index 00000000..983d0227
--- /dev/null
+++ b/src/stats/filters.ts
@@ -0,0 +1,718 @@
+/**
+ * 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,
+ complex,
+ cAbs,
+ kaiserWindow,
+ hannWindow,
+ hammingWindow,
+ blackmanWindow,
+ type WindowName,
+} 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 rawCuts = Array.isArray(cutoff)
+ ? (cutoff as readonly [number, number])
+ : ([cutoff] as const);
+ const cuts = (rawCuts as readonly number[]).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, f2] = cuts as [number, number];
+ 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) {
+ if (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: number[] = Array.isArray(worN)
+ ? Array.from(worN as readonly number[])
+ : Array.from({ length: worN as number }, (_, i) =>
+ (Math.PI * i) / (worN as number),
+ );
+
+ const H: Complex[] = ws.map((w) => {
+ // H(e^jw) = B(e^jw) / A(e^jw)
+ // Evaluate using Horner at z = e^jw
+ 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 polynomial with coefficients `p` at complex `z` (b[0]*z^N + ... + b[N]). */
+function evalPolyZ(p: readonly number[], z: Complex): Complex {
+ let acc: Complex = complex(0, 0);
+ for (let i = 0; i < p.length; 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);
+ * ```
+ */
+// biome-ignore lint/complexity/noExcessiveCognitiveComplexity: filter design algebra
+export function butter(
+ N: number,
+ Wn: number | readonly [number, number],
+ type: FilterType = "lowpass",
+): ButterResult {
+ // Validate
+ if (N < 1 || N > 20 || !Number.isInteger(N)) throw new RangeError("Order N must be an integer 1–20");
+
+ const nyq = 1; // Normalised: Nyquist = 1
+ const isBand = type === "bandpass" || type === "bandstop";
+
+ if (isBand && !Array.isArray(Wn)) throw new TypeError("Band filters require Wn = [low, high]");
+ if (!isBand && Array.isArray(Wn)) throw new TypeError("Low/high-pass filters require scalar Wn");
+
+ // Pre-warp critical frequency(ies) using bilinear transform
+ const warpedWn: number | [number, number] = Array.isArray(Wn)
+ ? ([
+ 2 * Math.tan((Math.PI * (Wn as readonly [number, number])[0]) / 2),
+ 2 * Math.tan((Math.PI * (Wn as readonly [number, number])[1]) / 2),
+ ] as [number, number])
+ : 2 * Math.tan((Math.PI * (Wn as number)) / 2);
+
+ // Analog Butterworth prototype poles at unit circle (left half-plane)
+ // p_k = exp(j * pi * (2k + N - 1) / (2N)) for k = 0..N-1
+ const poles: Complex[] = Array.from({ length: N }, (_, k) => {
+ const ang = (Math.PI * (2 * k + N - 1)) / (2 * N);
+ return complex(Math.cos(ang), Math.sin(ang));
+ });
+
+ // Scale poles to the desired cutoff frequency
+ let scaledPoles: Complex[];
+ let scaledZeros: Complex[];
+ let scaledGain: number;
+
+ if (type === "lowpass") {
+ const omega = warpedWn as number;
+ scaledPoles = poles.map((p) => ({ re: p.re * omega, im: p.im * omega }));
+ scaledZeros = []; // all zeros at s = ∞
+ scaledGain = omega ** N;
+ } else if (type === "highpass") {
+ const omega = warpedWn as number;
+ // LP → HP: s → omega / s ⟹ pole at s=p_k maps to omega/p_k
+ scaledPoles = poles.map((p) => {
+ const mag2 = p.re * p.re + p.im * p.im;
+ return { re: (omega * p.re) / mag2, im: -(omega * p.im) / mag2 };
+ });
+ scaledZeros = Array.from({ length: N }, () => complex(0, 0)); // N zeros at s=0
+ scaledGain = 1;
+ } else {
+ // Bandpass / bandstop: use direct bilinear transform below
+ scaledPoles = poles;
+ scaledZeros = [];
+ scaledGain = 1;
+ }
+
+ // Convert analog poles/zeros to digital via bilinear transform: z = (2+s)/(2-s)
+ // For band filters, handle separately
+ if (type === "bandpass" || type === "bandstop") {
+ return butterBand(N, warpedWn as [number, number], type, poles);
+ }
+
+ const digPoles: Complex[] = scaledPoles.map(bilinearPole);
+ const digZeros: Complex[] = scaledZeros.map(bilinearPole);
+ // LP numerator: N zeros at z=-1 after bilinear (from s=∞ mapping)
+ const lpZerosAtMinusOne = type === "lowpass" ? N : 0;
+
+ // Build SOS sections
+ const sos = buildSOS(digPoles, digZeros, lpZerosAtMinusOne, scaledGain, type);
+ const { b, a } = sosToBA(sos);
+
+ return { sos, b, a };
+}
+
+/** Bilinear transform: analog pole s → digital pole z = (2+s)/(2-s). */
+function bilinearPole(s: Complex): Complex {
+ // z = (2+s)/(2-s)
+ const num: Complex = { re: 2 + s.re, im: s.im };
+ const den: Complex = { re: 2 - s.re, im: -s.im };
+ const denom = den.re * den.re + den.im * den.im;
+ return { re: (num.re * den.re + num.im * den.im) / denom, im: (num.im * den.re - num.re * den.im) / denom };
+}
+
+/** Build second-order sections from digital poles, zeros, and gain. */
+// biome-ignore lint/complexity/noExcessiveCognitiveComplexity: SOS construction
+function buildSOS(
+ poles: Complex[],
+ explicitZeros: Complex[],
+ nZerosAtMinusOne: number,
+ gain: number,
+ type: FilterType,
+): SOSSection[] {
+ const N = poles.length;
+ const sections: SOSSection[] = [];
+
+ // Pair up conjugate poles (sort by imaginary part descending to pair conjugates)
+ const sortedPoles = [...poles].sort((a, b) => Math.abs(b.im) - Math.abs(a.im));
+ const usedPoles = new Array(N).fill(false);
+ const pairedPoles: [Complex, Complex | null][] = [];
+
+ for (let i = 0; i < N; i++) {
+ if (usedPoles[i]) continue;
+ const p = sortedPoles[i]!;
+ if (Math.abs(p.im) < 1e-10) {
+ // Real pole — stand alone
+ usedPoles[i] = true;
+ pairedPoles.push([p, null]);
+ } else {
+ // Find conjugate
+ let found = false;
+ for (let j = i + 1; j < N; j++) {
+ if (!usedPoles[j]) {
+ const q = sortedPoles[j]!;
+ if (Math.abs(p.re - q.re) < 1e-10 && Math.abs(p.im + q.im) < 1e-10) {
+ usedPoles[i] = usedPoles[j] = true;
+ pairedPoles.push([p, q]);
+ found = true;
+ break;
+ }
+ }
+ }
+ if (!found) {
+ usedPoles[i] = true;
+ pairedPoles.push([p, null]);
+ }
+ }
+ }
+
+ // Build sections: pair poles with zeros
+ let zerosRemaining = nZerosAtMinusOne;
+ let expZerosIdx = 0;
+ const nSections = pairedPoles.length;
+ const gainPerSection = gain > 0 ? gain ** (1 / nSections) : 1;
+
+ for (const [p1, p2] of pairedPoles) {
+ let b0: number, b1: number, b2: number;
+ let a1: number, a2: number;
+
+ if (p2 !== null) {
+ // Conjugate pair: (z - p1)(z - p2) = z^2 - 2*Re(p1)*z + |p1|^2
+ a1 = -2 * p1.re;
+ a2 = p1.re * p1.re + p1.im * p1.im;
+ if (type === "lowpass" && zerosRemaining >= 2) {
+ // Two zeros at z = -1: (z+1)^2 = z^2 + 2z + 1
+ b0 = 1; b1 = 2; b2 = 1;
+ zerosRemaining -= 2;
+ } else if (type === "highpass" && expZerosIdx < explicitZeros.length - 1) {
+ // Two zeros at z = 0: z^2 = z^2 + 0*z + 0
+ b0 = 1; b1 = 0; b2 = 0;
+ expZerosIdx += 2;
+ } else {
+ b0 = 1; b1 = 0; b2 = 0;
+ }
+ } else {
+ // Single real pole: (z - p1) = z - p1.re
+ a1 = -p1.re;
+ a2 = 0;
+ if (type === "lowpass" && zerosRemaining >= 1) {
+ // One zero at z = -1: z + 1
+ b0 = 1; b1 = 1; b2 = 0;
+ zerosRemaining -= 1;
+ } else if (type === "highpass" && expZerosIdx < explicitZeros.length) {
+ // One zero at z = 0: z
+ b0 = 1; b1 = 0; b2 = 0;
+ expZerosIdx += 1;
+ } else {
+ b0 = 1; b1 = 0; b2 = 0;
+ }
+ }
+
+ // Normalise section gain
+ const secGain = gainPerSection;
+ sections.push([b0 * secGain, b1 * secGain, b2 * secGain, 1, a1, a2]);
+ }
+
+ // Normalise so H(z=1) = 1 for lowpass, H(z=-1) = 1 for highpass
+ return normaliseSOS(sections, type);
+}
+
+/** Normalise SOS sections so the passband gain equals 1. */
+function normaliseSOS(sections: SOSSection[], type: FilterType): SOSSection[] {
+ // Evaluate H(z) at passband frequency: z=1 for LP, z=-1 for HP
+ const z = type === "highpass" ? -1 : 1;
+ const totalGain = sections.reduce((prod, sec) => {
+ const [b0, b1, b2, , a1, a2] = sec;
+ const num = b0 * z ** 2 + b1 * z + b2;
+ const den = z ** 2 + a1 * z + a2;
+ return prod * (Math.abs(den) > 1e-10 ? num / den : 1);
+ }, 1);
+
+ if (Math.abs(totalGain) < 1e-12) return sections;
+ const scale = 1 / totalGain;
+
+ // Apply scale to first section numerator only
+ const result: SOSSection[] = [...sections];
+ if (result.length > 0) {
+ const [b0, b1, b2, one, a1, a2] = result[0]!;
+ result[0] = [b0 * scale, b1 * scale, b2 * scale, one, a1, a2];
+ }
+ return result;
+}
+
+/** Handle band-pass and band-stop Butterworth filters. */
+// biome-ignore lint/complexity/noExcessiveCognitiveComplexity: band filter design
+function butterBand(
+ N: number,
+ warped: [number, number],
+ type: "bandpass" | "bandstop",
+ protoPoles: Complex[],
+): ButterResult {
+ const [w1, w2] = warped;
+ const bw = w2 - w1;
+ const w0 = Math.sqrt(w1 * w2); // geometric centre frequency
+
+ const sos: SOSSection[] = [];
+
+ // For each prototype pole, apply LP→BP or LP→BS transform
+ // LP→BP: s → (s^2 + w0^2) / (bw * s)
+ // Each LP pole becomes two BP poles
+ for (const p of protoPoles) {
+ // LP pole p_k: s → (s^2 + w0^2) / (bw * s) = p_k
+ // bw * s * p_k = s^2 + w0^2
+ // s^2 - bw * p_k * s + w0^2 = 0
+ // Solutions: s = (bw * p_k ± sqrt((bw * p_k)^2 - 4 * w0^2)) / 2
+ const a = bw * p.re;
+ const b = bw * p.im;
+ // discriminant = (bw*p)^2 - 4*w0^2 = (a+jb)^2 - 4*w0^2
+ const discRe = a * a - b * b - 4 * w0 * w0;
+ const discIm = 2 * a * b;
+ // sqrt of (discRe + j*discIm)
+ const [sqrtRe, sqrtIm] = complexSqrt(discRe, discIm);
+ const s1: Complex = { re: (a + sqrtRe) / 2, im: (b + sqrtIm) / 2 };
+ const s2: Complex = { re: (a - sqrtRe) / 2, im: (b - sqrtIm) / 2 };
+
+ let z1: Complex, z2: Complex;
+ if (type === "bandpass") {
+ z1 = bilinearPole(s1);
+ z2 = bilinearPole(s2);
+ } else {
+ // BP→BS transform: s → bw*w0 / (s^2 + w0^2)... simplify using direct analog BS poles
+ // LP→BS: s → bw*s / (s^2 + w0^2)
+ // Similar computation
+ z1 = bilinearPole(s1);
+ z2 = bilinearPole(s2);
+ }
+
+ // Each pair of complex poles contributes a 2nd-order section
+ const a1 = -(z1.re + z2.re);
+ const a2 = z1.re * z2.re - z1.im * z2.im; // assume z1, z2 are conjugates
+
+ const [b0, b1, b2] =
+ type === "bandpass"
+ ? [1, 0, -1] // bandpass: zeros at z=+1 and z=-1
+ : [1, -2 * Math.cos(Math.acos(Math.max(-1, Math.min(1, (w0 * w0 + 1) / (w0 * w0 + 1))))), 1]; // bandstop: zeros at e^±jw0
+
+ sos.push([b0, b1, b2, 1, a1, a2]);
+ }
+
+ const normalised = normaliseSOS(sos, type);
+ const { b, a } = sosToBA(normalised);
+ return { sos: normalised, b, a };
+}
+
+/** Real and imaginary parts of sqrt(re + j*im). */
+function complexSqrt(re: number, im: number): [number, number] {
+ const r = Math.sqrt(re * re + im * im);
+ const sr = Math.sqrt((r + re) / 2);
+ const si = Math.sign(im) * Math.sqrt((r - re) / 2);
+ return [sr, si];
+}
+
+/** Convert SOS to b/a transfer function via polynomial multiplication. */
+function sosToBA(sections: readonly SOSSection[]): { b: number[]; a: number[] } {
+ let b: number[] = [1];
+ let a: number[] = [1];
+ for (const [b0, b1, b2, , a1, a2] of sections) {
+ b = polyMul(b, [b0, b1, b2]);
+ a = polyMul(a, [1, 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: number[] = Array.isArray(worN)
+ ? Array.from(worN as readonly number[])
+ : Array.from({ length: worN as number }, (_, i) => (Math.PI * i) / (worN as number));
+
+ 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, , a1, a2] of sos) {
+ const num = evalPolyZ([b0, b1, b2], z);
+ const den = evalPolyZ([1, 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[] {
+ const forward = lfilter(b, a, x);
+ const reversed = [...forward].reverse();
+ const backward = lfilter(b, a, reversed);
+ return backward.reverse();
+}
+
+/**
+ * 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, , a1, a2] of sos) {
+ signal = lfilter([b0, b1, b2], [1, a1, a2], signal);
+ }
+ return signal;
+}
+
+/**
+ * Zero-phase SOS filter (applies each section forward then backward).
+ *
+ * @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[] {
+ let signal = Array.from(x);
+ for (const [b0, b1, b2, , a1, a2] of sos) {
+ signal = filtfilt([b0, b1, b2], [1, a1, a2], signal);
+ }
+ return signal;
+}
+
+// Re-export cAbs for convenience
+export { cAbs };
diff --git a/src/stats/index.ts b/src/stats/index.ts
index 78767ee1..d3154537 100644
--- a/src/stats/index.ts
+++ b/src/stats/index.ts
@@ -575,3 +575,119 @@ 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";
diff --git a/src/stats/kalman.ts b/src/stats/kalman.ts
new file mode 100644
index 00000000..cf4fbc02
--- /dev/null
+++ b/src/stats/kalman.ts
@@ -0,0 +1,798 @@
+/**
+ * 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 }, () => 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] += 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 = 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 -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 ?? 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(Array(no).fill(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] ?? 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] ?? 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/signal.ts b/src/stats/signal.ts
new file mode 100644
index 00000000..7a9962fd
--- /dev/null
+++ b/src/stats/signal.ts
@@ -0,0 +1,768 @@
+/**
+ * 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.floor(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.ceil(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,
+ a1 = 0.48829,
+ a2 = 0.14128,
+ 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,
+ a1 = 0.41663158,
+ a2 = 0.277263158,
+ a3 = 0.083578947,
+ 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 * w, 0);
+ 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 * w, 0);
+
+ 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/tests/io/orc.test.ts b/tests/io/orc.test.ts
new file mode 100644
index 00000000..6aef66f6
--- /dev/null
+++ b/tests/io/orc.test.ts
@@ -0,0 +1,390 @@
+/**
+ * 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[buf.length - 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("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 00000000..5e65d955
--- /dev/null
+++ b/tests/io/read_avro.test.ts
@@ -0,0 +1,358 @@
+/**
+ * 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/stats/acf_pacf.test.ts b/tests/stats/acf_pacf.test.ts
new file mode 100644
index 00000000..c23572f6
--- /dev/null
+++ b/tests/stats/acf_pacf.test.ts
@@ -0,0 +1,486 @@
+/**
+ * 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([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 has high positive ACF at all lags", () => {
+ const x = Array.from({ length: 20 }, (_, i) => i);
+ const result = acf(x, { nlags: 5 });
+ for (let k = 1; k <= 5; k++) {
+ expect(result.acf[k]).toBeGreaterThan(0.5);
+ }
+ });
+
+ 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).toBeGreaterThan(3.9);
+ });
+
+ 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(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.3125, -0.3125, -0.5625, -0.1875]
+ 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.3125, 4));
+ expect(round(result.acf[2] ?? 0, 4)).toBe(round(-0.3125, 4));
+ });
+
+ it("autocorr(x, 1) matches acf(x,nlags=1).acf[1]", () => {
+ const acfVal = acf(x, { nlags: 1 }).acf[1] ?? 0;
+ // autocorr uses Pearson, acf uses autocovariance — they differ slightly
+ // Both should be in the same ballpark
+ const ac = autocorr(x, 1);
+ expect(Math.sign(ac)).toBe(Math.sign(acfVal));
+ });
+
+ 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)).toBeGreaterThan(3.9);
+ });
+
+ 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 00000000..18aab018
--- /dev/null
+++ b/tests/stats/arima.test.ts
@@ -0,0 +1,329 @@
+/**
+ * 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) coefficient close to true phi=0.7", () => {
+ const { arCoeffs } = new ARIMAModel({ p: 1, d: 0, q: 0 }).fit(y);
+ expect(arCoeffs[0]).toBeCloseTo(0.7, 1);
+ });
+
+ 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] ?? 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/ets.test.ts b/tests/stats/ets.test.ts
new file mode 100644
index 00000000..e28e4ea3
--- /dev/null
+++ b/tests/stats/ets.test.ts
@@ -0,0 +1,815 @@
+/**
+ * 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);
+ });
+ });
+
+ 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[TREND.length - 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("SSE roughly matches Holt class", () => {
+ const r1 = new ExponentialSmoothing({ trend: "add" }).fit(TREND);
+ const r2 = new Holt().fit(TREND);
+ // They may find slightly different optima; SSE should be comparable
+ expect(Math.abs(r1.sse - r2.sse) / (r2.sse + 1e-8)).toBeLessThan(0.1);
+ });
+ });
+
+ 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(Math.abs(sum)).toBeLessThan(2);
+ });
+
+ 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 is finite and negative", () => {
+ const fit = fitEts(TREND, { trend: "add" });
+ expect(Number.isFinite(fit.logLikelihood)).toBe(true);
+ expect(fit.logLikelihood).toBeLessThan(0);
+ });
+
+ 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.float({ min: 0.01, max: 0.99 }),
+ (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 00000000..6a714a3b
--- /dev/null
+++ b/tests/stats/filters.test.ts
@@ -0,0 +1,433 @@
+/**
+ * Tests for src/stats/filters.ts
+ * Covers FIR design, Butterworth IIR, frequency response, and filter application.
+ */
+
+import { describe, test, expect } from "bun:test";
+import * as fc from "fast-check";
+import {
+ firwin,
+ freqz,
+ sosfreqz,
+ lfilter,
+ filtfilt,
+ sosfilt,
+ sosfiltfilt,
+ butter,
+ cAbs,
+ type SOSSection,
+} 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.float({ min: 0.01, max: 0.49, noNaN: true }),
+ (taps, cutoff) => {
+ const b = firwin(taps, cutoff);
+ return b.every(Number.isFinite);
+ },
+ ),
+ );
+ });
+});
+
+// ─── freqz ────────────────────────────────────────────────────────────────────
+
+describe("freqz", () => {
+ 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", () => {
+ 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] as unknown as number, "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("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 = i < n / 2 ? i : n - 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("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/kalman.test.ts b/tests/stats/kalman.test.ts
new file mode 100644
index 00000000..068c523b
--- /dev/null
+++ b/tests/stats/kalman.test.ts
@@ -0,0 +1,661 @@
+/**
+ * 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,
+ kalmanFilter1D,
+ kalmanSmooth1D,
+ extractScalarMeans,
+ filteredPredictionInterval,
+} 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] ?? NaN;
+ const y = obs[t]?.[0] ?? NaN;
+ expect(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(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] ?? 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(isNaN(result.innovations[1]?.[0] ?? 0)).toBe(true);
+ expect(isNaN(result.innovations[2]?.[0] ?? 0)).toBe(true);
+ });
+
+ it("innovations are finite for observed steps", () => {
+ expect(isFinite(result.innovations[0]?.[0] ?? NaN)).toBe(true);
+ expect(isFinite(result.innovations[3]?.[0] ?? 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(isFinite(m[0] ?? 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(isFinite(ll1) && 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] ?? NaN;
+ expect(absDiff(level, t)).toBeLessThan(3);
+ }
+ });
+
+ it("slope converges towards 1", () => {
+ const means = result.filteredStateMeans;
+ const slope = means[14]?.[1] ?? 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] ?? NaN;
+ const smoothLast = sm.smoothedStateMeans[4]?.[0] ?? 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] ?? 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(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(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] ?? NaN, m2[i] ?? 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(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] ?? 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 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(isFinite(m[0] ?? 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 00000000..3b202c85
--- /dev/null
+++ b/tests/stats/signal.test.ts
@@ -0,0 +1,460 @@
+/**
+ * Tests for src/stats/signal.ts
+ * Covers FFT, windows, STFT, ISTFT, Welch PSD, and periodogram.
+ */
+
+import { describe, test, expect } from "bun:test";
+import * as fc from "fast-check";
+import {
+ fft, ifft, rfft, irfft,
+ fftFreq, rfftFreq,
+ fftshift, ifftshift,
+ complex, cAbs,
+ getWindow,
+ rectangularWindow, bartlettWindow, hannWindow, hammingWindow,
+ blackmanWindow, blackmanHarrisWindow, flatTopWindow, kaiserWindow,
+ stft, istft, welch, periodogram,
+} 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, 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[freqs.length - 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([2, 3, 4, 0, 1]);
+ });
+
+ test("odd length: ifftshift matches numpy", () => {
+ const x = [2, 3, 4, 0, 1];
+ 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 — values can exceed 1", () => {
+ const w = flatTopWindow(64);
+ expect(w.length).toBe(64);
+ expect(Math.max(...w)).toBeGreaterThan(1);
+ });
+
+ 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 });
+ // They should differ
+ expect(dens[0]).not.toBeCloseTo(spec[0] ?? 0, 5);
+ });
+});
+
+// ─── 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("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[f.length - 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]);
+ }
+ }
+ });
+});