diff --git a/playground/kalman.html b/playground/kalman.html
new file mode 100644
index 00000000..b89ce40e
--- /dev/null
+++ b/playground/kalman.html
@@ -0,0 +1,476 @@
+
+
+
+
+
+
tsb — Kalman Filter & State-Space Models
+
+
+
+
+
+
+
← Back to index
+
🔭 Kalman Filter & State-Space Models
+
+ Linear Gaussian state-space model — Kalman filter (forward pass) and
+ RTS smoother (backward pass). Mirrors statsmodels.tsa.statespace
+ and pykalman.KalmanFilter.
+
+
+
+
+
📐 The State-Space Model
+
+ A linear Gaussian SSM describes a latent state x_t and
+ observations y_t via two equations:
+
+
+ x_t = F · x_{t-1} + w_t, w_t ~ N(0, Q) (state transition)
+
+ F — state transition matrix (n_states × n_states)H — observation matrix (n_obs × n_states)Q — process noise covarianceR — observation noise covariancem_0, P_0 — initial state distribution
+
+ The Kalman filter computes filtered state
+ estimates x_{t|t} (posterior after seeing observation t).
+ The RTS smoother computes smoothed estimates
+ x_{t|T} using all T observations.
+
+
+
+
+
+
📈 Local-Level Model (Random Walk + Noise)
+
+ The simplest SSM: a hidden state that follows a random walk, observed
+ with noise. Perfect for denoising a noisy scalar time series or
+ estimating a slowly changing mean.
+
+
+
+
+
Click ▶ Run to execute
+
+
+
+
+
+
🔄 RTS Smoother — Filling Gaps Retrospectively
+
+ The filter only uses observations up to time t. The smoother uses
+ all observations to produce better estimates, especially for
+ time-steps near missing values. Smoothed uncertainty is always ≤ filtered.
+
+
+
+
+
Click ▶ Run to execute
+
+
+
+
+
+
📊 Local Linear Trend (Level + Slope)
+
+ A 2-state model: [level, slope]. The level increases by the
+ slope each step; both drift over time. Great for tracking slowly changing
+ trends with missing observations.
+
+
+
+
+
Click ▶ Run to execute
+
+
+
+
+
+
⚙️ Custom State-Space Model (AR(1) State)
+
+ Build your own model by specifying the four matrices directly.
+ Here: a state that follows an AR(1) process with coefficient 0.9.
+
+
+
+
+
Click ▶ Run to execute
+
+
+
+
+
+
🔢 Multi-Dimensional Observations
+
+ The Kalman filter naturally handles multi-dimensional observations.
+ Here: 2 sensors observing a single latent state.
+
+
+
+
+
Click ▶ Run to execute
+
+
+
+
+
+
📖 API Reference
+
+ KalmanFilter.localLevel(opts?) — random-walk + noise (1-D)KalmanFilter.localLinearTrend(opts?) — level + slope (2-D state)new KalmanFilter(opts) — custom F, H, Q, R, m0, P0kf.filter(observations) → KalmanFilterResultkf.smooth(observations) → KalmanSmootherResultkalmanFilter1D(obs, opts?) — scalar convenience wrapperkalmanSmooth1D(obs, opts?) — scalar smoother wrapperextractScalarMeans(means) — extract 1-D means arrayfilteredPredictionInterval(result, z?) → {lower, upper}
+
+ Missing observations: pass null in any observation row. The
+ filter skips the update step for that time-step (covariance grows).
+ The smoother retroactively interpolates using future observations.
+
+
+
+
+
+
diff --git a/playground/orc.html b/playground/orc.html
new file mode 100644
index 00000000..494d18e7
--- /dev/null
+++ b/playground/orc.html
@@ -0,0 +1,207 @@
+
+
+
+
+
+
tsb · ORC Format I/O
+
+
+
+
← tsb playground
+
ORC Format I/O
+
+ pandas.read_orc
+ DataFrame.to_orc
+
+
+ Apache ORC (Optimized Row Columnar) is a self-describing, type-aware columnar file format designed
+ for large-scale analytical workloads. tsb supports reading and writing ORC files with
+ NONE compression using RLE v1 integer encoding, direct float/double, and direct string encoding.
+
+
+
+ ℹ️ This playground runs entirely in-browser via a bundled tsb build. ORC buffers are created
+ in-memory — no file system access is required.
+
+
+
1 — Write & read a DataFrame
+
Create a DataFrame, serialize it to ORC bytes, then parse it back:
+
import { DataFrame, readOrc, toOrc } from "tsb";
+
+const df = DataFrame.fromColumns({
+ id: [1, 2, 3, 4, 5],
+ name: ["Alice", "Bob", "Carol", "Dave", "Eve"],
+ score: [95.5, 87.0, 92.3, 78.1, 99.9],
+ passed: [true, false, true, false, true],
+});
+
+// Serialize to binary ORC
+const buf = toOrc(df);
+console.log("ORC buffer size:", buf.length, "bytes");
+
+// Parse back
+const df2 = readOrc(buf);
+console.log(df2.toString());
+
▶ Run
+
Click "Run" to execute…
+
+
2 — Nullable columns
+
ORC natively supports null values via PRESENT streams:
+
import { DataFrame, readOrc, toOrc } from "tsb";
+
+const df = DataFrame.fromColumns({
+ x: [1, null, 3, null, 5],
+ name: ["a", null, "c", null, "e"],
+});
+
+const buf = toOrc(df);
+const rt = readOrc(buf);
+console.log("x values:", rt.col("x").values.join(", "));
+console.log("name values:", rt.col("name").values.join(", "));
+
▶ Run
+
Click "Run" to execute…
+
+
3 — Column selection
+
Use the columns option to read only a subset of columns:
+
import { DataFrame, readOrc, toOrc } from "tsb";
+
+const df = DataFrame.fromColumns({
+ a: [1, 2, 3],
+ b: ["x", "y", "z"],
+ c: [true, false, true],
+ d: [10.0, 20.0, 30.0],
+});
+
+const buf = toOrc(df);
+
+// Only read columns a and c
+const partial = readOrc(buf, { columns: ["a", "c"] });
+console.log("columns:", partial.columns.toArray().join(", "));
+console.log("a:", partial.col("a").values.join(", "));
+console.log("c:", partial.col("c").values.join(", "));
+
▶ Run
+
Click "Run" to execute…
+
+
4 — Large dataset benchmark
+
Serialize and parse a 10 000-row DataFrame to measure throughput:
+
import { DataFrame, readOrc, toOrc } from "tsb";
+
+const N = 10_000;
+const ids = Array.from({ length: N }, (_, i) => i);
+const names = Array.from({ length: N }, (_, i) => `user_${i}`);
+const scores = Array.from({ length: N }, () => Math.random() * 100);
+
+const df = DataFrame.fromColumns({ id: ids, name: names, score: scores });
+
+const t0 = performance.now();
+const buf = toOrc(df);
+const t1 = performance.now();
+const df2 = readOrc(buf);
+const t2 = performance.now();
+
+console.log(`Rows: ${df2.height}, Columns: ${df2.width}`);
+console.log(`ORC size: ${buf.length.toLocaleString()} bytes`);
+console.log(`Write: ${(t1 - t0).toFixed(1)} ms`);
+console.log(`Read: ${(t2 - t1).toFixed(1)} ms`);
+
▶ Run
+
Click "Run" to execute…
+
+
+
+
diff --git a/playground/signal.html b/playground/signal.html
new file mode 100644
index 00000000..cfc646c5
--- /dev/null
+++ b/playground/signal.html
@@ -0,0 +1,621 @@
+
+
+
+
+
+
tsb — Signal Processing
+
+
+
+
+
+
+
Initializing playground…
+
+
+
← Back to roadmap
+
📡 Signal Processing — Interactive Playground
+
+ FFT, windows, STFT, Welch PSD, and periodogram — mirrors numpy.fft
+ and scipy.signal.Edit any code block below and press ▶ Run
+ (or Ctrl+Enter) to execute it live in your browser.
+
+
+
+
+
1. Basic FFT of a sinusoidal signal
+
Compute a 32 Hz sine wave's FFT and identify the peak frequency bin.
+
+
+
+
+
Click ▶ Run to execute
+
Ctrl+Enter to run · Tab to indent
+
+
+
+
+
+
2. Parseval's theorem — energy preservation
+
The total energy is preserved between time and frequency domains.
+
+
+
+
+
Click ▶ Run to execute
+
Ctrl+Enter to run · Tab to indent
+
+
+
+
+
+
3. RFFT round-trip
+
Real-input FFT produces a half-spectrum; irfft reconstructs the original signal.
+
+
+
+
+
Click ▶ Run to execute
+
Ctrl+Enter to run · Tab to indent
+
+
+
+
+
+
4. Window functions
+
Named windows reduce spectral leakage. Use getWindow(name, n) to obtain any built-in window.
+
+
+
+
+
Click ▶ Run to execute
+
Ctrl+Enter to run · Tab to indent
+
+
+
+
+
+
5. Short-Time Fourier Transform (STFT)
+
Analyze a chirp signal whose frequency increases linearly over time.
+
+
+
+
+
Click ▶ Run to execute
+
Ctrl+Enter to run · Tab to indent
+
+
+
+
+
+
6. ISTFT reconstruction (round-trip)
+
Invert an STFT back to the time domain. Interior reconstruction error should be near machine epsilon.
+
+
+
+
+
Click ▶ Run to execute
+
Ctrl+Enter to run · Tab to indent
+
+
+
+
+
+
7. Welch PSD — detect signal frequency
+
Welch's method averages periodograms of overlapping segments for a lower-variance PSD estimate.
+
+
+
+
+
Click ▶ Run to execute
+
Ctrl+Enter to run · Tab to indent
+
+
+
+
+
+
8. Periodogram
+
A single-segment PSD estimate — higher variance but simpler than Welch.
+
+
+
+
+
Click ▶ Run to execute
+
Ctrl+Enter to run · Tab to indent
+
+
+
+
+
+
9. fftshift / ifftshift
+
Rearrange the FFT output so that the zero-frequency component is in the centre.
+
+
+
+
+
Click ▶ Run to execute
+
Ctrl+Enter to run · Tab to indent
+
+
+
+
+
+
API Reference
+
+ Function Description Mirrors fft(x)N-point DFT (pads to power of 2) numpy.fft.fftifft(X)Inverse FFT numpy.fft.ifftrfft(x)Real-input FFT (one-sided) numpy.fft.rfftirfft(X, n?)Inverse real FFT numpy.fft.irfftfftFreq(n, d?)DFT sample frequencies numpy.fft.fftfreqrfftFreq(n, d?)One-sided DFT frequencies numpy.fft.rfftfreqfftshift(x)Shift DC to centre numpy.fft.fftshiftifftshift(x)Inverse of fftshift numpy.fft.ifftshiftgetWindow(name, n)Named window function scipy.signal.get_windowstft(x, opts?)Short-Time Fourier Transform scipy.signal.stftistft(Zxx, opts?)Inverse STFT (overlap-add) scipy.signal.istftwelch(x, opts?)Welch PSD estimate scipy.signal.welchperiodogram(x, opts?)Periodogram PSD estimate scipy.signal.periodogram
+
+
+
+
+
+
+
+
diff --git a/src/core/frame.ts b/src/core/frame.ts
index e21c341e..5e13d577 100644
--- a/src/core/frame.ts
+++ b/src/core/frame.ts
@@ -626,7 +626,7 @@ export class DataFrame {
for (let i = 0; i < nRows; i++) {
const row: Record
= {};
if (index) {
- row["Index"] = this.index.at(i) as Scalar;
+ row.Index = this.index.at(i) as Scalar;
}
for (const name of colNames) {
row[name] = this.col(name).iat(i);
@@ -833,9 +833,7 @@ function isIndexLike(v: unknown): v is Index {
}
const rec = v as Record;
return (
- typeof rec["size"] === "number" &&
- typeof rec["at"] === "function" &&
- typeof rec["getLoc"] === "function"
+ typeof rec.size === "number" && typeof rec.at === "function" && typeof rec.getLoc === "function"
);
}
diff --git a/src/core/options.ts b/src/core/options.ts
index 628b5cce..188c2188 100644
--- a/src/core/options.ts
+++ b/src/core/options.ts
@@ -210,7 +210,7 @@ function _makeProxy(prefix: string): OptionsProxy {
// If there is a direct match, return its value
const k = _normalizeKey(key);
if (_registry.has(k)) {
- return _registry.get(k)!.currentValue;
+ return _registry.get(k)?.currentValue;
}
// Otherwise return a nested proxy for deeper access
return _makeProxy(k);
diff --git a/src/core/series.ts b/src/core/series.ts
index c34d2569..c1ad5ab5 100644
--- a/src/core/series.ts
+++ b/src/core/series.ts
@@ -1559,9 +1559,7 @@ function isIndexLike(v: unknown): v is Index {
}
const rec = v as Record;
return (
- typeof rec["size"] === "number" &&
- typeof rec["at"] === "function" &&
- typeof rec["getLoc"] === "function"
+ typeof rec.size === "number" && typeof rec.at === "function" && typeof rec.getLoc === "function"
);
}
diff --git a/src/index.ts b/src/index.ts
index bb13fcd1..095f4c2d 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,165 @@ export {
tsallisEntropy,
} from "./stats/information.ts";
export type { PMF, NMIMethod } from "./stats/information.ts";
+
+// signal processing — FFT, windows, STFT, Welch PSD, periodogram
+export {
+ complex,
+ cAbs,
+ cArg,
+ fft,
+ ifft,
+ rfft,
+ irfft,
+ fftFreq,
+ rfftFreq,
+ fftshift,
+ ifftshift,
+ rectangularWindow,
+ bartlettWindow,
+ hannWindow,
+ hammingWindow,
+ blackmanWindow,
+ blackmanHarrisWindow,
+ flatTopWindow,
+ kaiserWindow,
+ getWindow,
+ stft,
+ istft,
+ welch,
+ periodogram,
+} from "./stats/signal.ts";
+export type {
+ Complex,
+ WindowName,
+ STFTOptions,
+ STFTResult,
+ ISTFTOptions,
+ WelchOptions,
+ PSDResult,
+ PeriodogramOptions,
+} from "./stats/signal.ts";
+
+// digital filters — FIR/IIR design and application (mirrors scipy.signal)
+export {
+ firwin,
+ freqz,
+ sosfreqz,
+ lfilter,
+ filtfilt,
+ sosfilt,
+ sosfiltfilt,
+ butter,
+} from "./stats/filters.ts";
+export type {
+ FirwinOptions,
+ FreqzResult,
+ SOSSection,
+ ButterResult,
+ FilterType,
+} from "./stats/filters.ts";
+
+// ACF/PACF — autocorrelation, partial autocorrelation, portmanteau tests
+export {
+ autocorr,
+ acf,
+ pacf,
+ ccf,
+ durbinWatson,
+ ljungBox,
+ boxPierce,
+} from "./stats/acf_pacf.ts";
+export type {
+ ACFResult,
+ PACFResult,
+ PortmanteauResult,
+ ACFOptions,
+ PACFOptions,
+ CCFOptions,
+ PortmanteauOptions,
+} from "./stats/acf_pacf.ts";
+
+// ARIMA — ARIMA(p,d,q) time-series model (Hannan-Rissanen, forecast CIs)
+export { ARIMAModel, fitArima } from "./stats/arima.ts";
+export type {
+ ARIMAOptions,
+ ARIMAFitResult,
+ ARIMAForecastResult,
+} from "./stats/arima.ts";
+
+// read_avro / toAvro — Apache Avro OCF I/O for DataFrame
+export { readAvro, toAvro } from "./io/read_avro.ts";
+export type {
+ ReadAvroOptions,
+ ToAvroOptions,
+} from "./io/read_avro.ts";
+
+// Kalman filter & RTS smoother — linear Gaussian state-space model
+export {
+ KalmanFilter,
+ StateSpaceModel,
+ kalmanFilter1D,
+ kalmanSmooth1D,
+ extractScalarMeans,
+ extractScalarVariances,
+ filteredPredictionInterval,
+} from "./stats/kalman.ts";
+export type {
+ KalmanFilterOptions,
+ LocalLevelOptions,
+ LocalLinearTrendOptions,
+ KalmanFilterResult,
+ KalmanSmootherResult,
+} from "./stats/kalman.ts";
+
+// ETS — Exponential Smoothing / Holt-Winters (Simple, Holt, full Holt-Winters)
+export {
+ SimpleExpSmoothing,
+ Holt,
+ ExponentialSmoothing,
+ simpleExpSmoothing,
+ holt,
+ fitEts,
+} from "./stats/ets.ts";
+export type {
+ ETSTrend,
+ ETSSeasonal,
+ ETSInit,
+ SESOptions,
+ SESFitResult,
+ HoltOptions,
+ HoltFitResult,
+ ExponentialSmoothingOptions,
+ ExponentialSmoothingFitResult,
+ ETSForecastResult,
+} from "./stats/ets.ts";
+
+export {
+ DLM,
+ buildLocalLevel,
+ buildLocalLinearTrend,
+ buildPolynomial,
+ buildFourier,
+ buildRegression,
+ combineDLMs,
+} from "./stats/dlm.ts";
+export type {
+ DLMSpec,
+ DLMOptions,
+ DLMFilterStep,
+ DLMResult,
+ DLMSmootherResult,
+ DLMForecastResult,
+} from "./stats/dlm.ts";
+export {
+ GaussianHMM,
+ MultinomialHMM,
+ fitGaussianHMM,
+ hmmViterbi,
+} from "./stats/hmm.ts";
+export type {
+ GaussianHMMParams,
+ GaussianHMMFit,
+ MultinomialHMMParams,
+ MultinomialHMMFit,
+} from "./stats/hmm.ts";
diff --git a/src/io/feather.ts b/src/io/feather.ts
index 1abb681b..e91c3e17 100644
--- a/src/io/feather.ts
+++ b/src/io/feather.ts
@@ -726,7 +726,7 @@ function encodeStrings(values: readonly (Scalar | null)[]): {
for (let i = 0; i < encoded.length; i++) {
ov.setInt32(i * 4, pos, true);
data.set(encoded[i]!, pos);
- pos += encoded[i]!.length;
+ pos += encoded[i]?.length;
}
ov.setInt32(values.length * 4, pos, true);
return { offsets, data };
@@ -918,7 +918,7 @@ export function toFeather(df: DataFrame, options: ToFeatherOptions = {}): Uint8A
cols.push({ name, type, values });
}
- const numRows = cols.length > 0 ? cols[0]!.values.length : df.index.size;
+ const numRows = cols.length > 0 ? cols[0]?.values.length : df.index.size;
const schemaCols = cols.map((c) => ({ name: c.name, type: c.type }));
// Encode all column buffers into a single body array
diff --git a/src/io/index.ts b/src/io/index.ts
index 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..7b9822a5
--- /dev/null
+++ b/src/io/orc.ts
@@ -0,0 +1,1351 @@
+/**
+ * readOrc / toOrc — Apache ORC (Optimized Row Columnar) file format I/O.
+ *
+ * Mirrors `pandas.read_orc()` and `DataFrame.to_orc()`.
+ *
+ * Supported column types (read & write):
+ * - BOOLEAN, INT, LONG, FLOAT, DOUBLE, STRING, DATE
+ *
+ * Compression: NONE (ZLIB/Snappy require an external decompressor).
+ * Encoding: DIRECT (integers via RLE v1, strings via raw bytes + lengths,
+ * floats/doubles via raw IEEE 754, booleans via RLE byte v1).
+ *
+ * @module
+ */
+
+import { DataFrame } from "../core/frame.ts";
+import type { Label, Scalar } from "../types.ts";
+
+// ─── Public types ─────────────────────────────────────────────────────────────
+
+/** Options for {@link readOrc}. */
+export interface ReadOrcOptions {
+ /**
+ * Column name to use as the row index.
+ * Default: `null` (RangeIndex).
+ */
+ readonly indexCol?: string | null;
+ /**
+ * Subset of columns to read. `null` = all columns.
+ * Default: `null`.
+ */
+ readonly columns?: readonly string[] | null;
+}
+
+/** Options for {@link toOrc}. */
+export interface ToOrcOptions {
+ /**
+ * Write the DataFrame's row index as an extra column.
+ * Default: `false`.
+ */
+ readonly writeIndex?: boolean;
+}
+
+// ─── ORC file constants ───────────────────────────────────────────────────────
+
+// File header magic
+const ORC_MAGIC = new Uint8Array([0x4f, 0x52, 0x43]); // "ORC"
+
+// ORC type kinds
+const KIND_BOOLEAN = 0;
+const _KIND_BYTE = 1;
+const _KIND_SHORT = 2;
+const _KIND_INT = 3;
+const KIND_LONG = 4;
+const KIND_FLOAT = 5;
+const KIND_DOUBLE = 6;
+const KIND_STRING = 7;
+const KIND_STRUCT = 12;
+const KIND_DATE = 15;
+
+// Compression codecs
+const COMP_NONE = 0;
+const COMP_ZLIB = 1;
+
+// Stream kinds
+const STREAM_PRESENT = 0;
+const STREAM_DATA = 1;
+const STREAM_LENGTH = 2;
+const _STREAM_DICTIONARY_DATA = 3;
+
+// Column encoding kinds
+const ENC_DIRECT = 0;
+
+// ─── Protobuf utilities ───────────────────────────────────────────────────────
+
+/** A single decoded protobuf field value. */
+type PbVal =
+ | { readonly wt: 0; readonly v: bigint }
+ | { readonly wt: 2; readonly v: Uint8Array }
+ | { readonly wt: 1; readonly v: bigint }
+ | { readonly wt: 5; readonly v: number };
+
+/** A decoded protobuf message: field number → list of values. */
+type PbMsg = Map;
+
+/** Read a protobuf varint (unsigned, LSB-first). */
+function pbReadVarU(buf: Uint8Array, pos: number): [bigint, number] {
+ let result = 0n;
+ let shift = 0n;
+ let cur = pos;
+ for (;;) {
+ const b = buf[cur];
+ if (b === undefined) {
+ throw new Error("ORC: truncated varint");
+ }
+ cur++;
+ result |= BigInt(b & 0x7f) << shift;
+ if ((b & 0x80) === 0) {
+ break;
+ }
+ shift += 7n;
+ }
+ return [result, cur];
+}
+
+/** Decode a protobuf message from a byte slice. */
+function pbDecode(buf: Uint8Array): PbMsg {
+ const msg: PbMsg = new Map();
+ let pos = 0;
+ while (pos < buf.length) {
+ let tag: bigint;
+ [tag, pos] = pbReadVarU(buf, pos);
+ const fieldNum = Number(tag >> 3n);
+ const wt = Number(tag & 7n);
+ let val: PbVal;
+ if (wt === 0) {
+ let v: bigint;
+ [v, pos] = pbReadVarU(buf, pos);
+ val = { wt: 0, v };
+ } else if (wt === 2) {
+ let len: bigint;
+ [len, pos] = pbReadVarU(buf, pos);
+ const n = Number(len);
+ val = { wt: 2, v: buf.subarray(pos, pos + n) };
+ pos += n;
+ } else if (wt === 1) {
+ const dv = new DataView(buf.buffer, buf.byteOffset + pos, 8);
+ const lo = BigInt(dv.getUint32(0, true));
+ const hi = BigInt(dv.getUint32(4, true));
+ val = { wt: 1, v: (hi << 32n) | lo };
+ pos += 8;
+ } else if (wt === 5) {
+ const dv = new DataView(buf.buffer, buf.byteOffset + pos, 4);
+ val = { wt: 5, v: dv.getUint32(0, true) };
+ pos += 4;
+ } else {
+ throw new Error(`ORC: unknown wire type ${wt}`);
+ }
+ const list = msg.get(fieldNum);
+ if (list !== undefined) {
+ list.push(val);
+ } else {
+ msg.set(fieldNum, [val]);
+ }
+ }
+ return msg;
+}
+
+/** Get a uint64 field as bigint (default 0). */
+function pbU64(msg: PbMsg, field: number): bigint {
+ const f = msg.get(field)?.[0];
+ return f?.wt === 0 ? f.v : 0n;
+}
+
+/** Get a uint32 field as number (default 0). */
+function pbU32(msg: PbMsg, field: number): number {
+ return Number(pbU64(msg, field));
+}
+
+/** Get all uint32 repeated field values. */
+function pbU32s(msg: PbMsg, field: number): number[] {
+ return (msg.get(field) ?? [])
+ .filter((f): f is PbVal & { wt: 0 } => f.wt === 0)
+ .map((f) => Number(f.v));
+}
+
+/** Get all string repeated field values. */
+function pbStrings(msg: PbMsg, field: number): string[] {
+ const dec = new TextDecoder();
+ return (msg.get(field) ?? [])
+ .filter((f): f is PbVal & { wt: 2 } => f.wt === 2)
+ .map((f) => dec.decode(f.v));
+}
+
+/** Get all embedded message repeated field values. */
+function pbMsgs(msg: PbMsg, field: number): PbMsg[] {
+ return (msg.get(field) ?? [])
+ .filter((f): f is PbVal & { wt: 2 } => f.wt === 2)
+ .map((f) => pbDecode(f.v));
+}
+
+// ─── Protobuf writer ──────────────────────────────────────────────────────────
+
+function pbWvU(v: bigint, out: number[]): void {
+ let val = v;
+ while (val >= 128n) {
+ out.push(Number(val & 0x7fn) | 0x80);
+ val >>= 7n;
+ }
+ out.push(Number(val));
+}
+
+function pbTag(fn: number, wt: 0 | 2, out: number[]): void {
+ pbWvU(BigInt((fn << 3) | wt), out);
+}
+
+function pbWU64(fn: number, v: bigint, out: number[]): void {
+ if (v === 0n) {
+ return;
+ }
+ pbTag(fn, 0, out);
+ pbWvU(v, out);
+}
+
+function pbWU32(fn: number, v: number, out: number[]): void {
+ pbWU64(fn, BigInt(v), out);
+}
+
+function pbWBytes(fn: number, v: Uint8Array, out: number[]): void {
+ pbTag(fn, 2, out);
+ pbWvU(BigInt(v.length), out);
+ for (const b of v) {
+ out.push(b);
+ }
+}
+
+function pbWMsg(fn: number, msg: number[], out: number[]): void {
+ pbTag(fn, 2, out);
+ pbWvU(BigInt(msg.length), out);
+ for (const b of msg) {
+ out.push(b);
+ }
+}
+
+// ─── Hadoop VInt ──────────────────────────────────────────────────────────────
+// ORC uses big-endian variable-length signed integers for RLE integer streams.
+
+/**
+ * Read a Hadoop-style variable-length signed integer.
+ *
+ * Byte ranges:
+ * - 0x00–0x7F: single-byte positive (0–127)
+ * - 0x88–0x8F: positive multi-byte (1–8 data bytes follow)
+ * - 0x80–0x87: negative multi-byte (1–8 data bytes follow, XOR with -1)
+ * - 0x90–0xFF: single-byte negative (-112 to -1)
+ */
+function hvReadVInt(buf: Uint8Array, pos: number): [bigint, number] {
+ const fb = buf[pos];
+ if (fb === undefined) {
+ throw new Error("ORC: truncated Hadoop VInt");
+ }
+ let cur = pos + 1;
+ // Interpret as signed byte
+ const sfb = fb >= 0x80 ? fb - 0x100 : fb;
+ // Single-byte range: -112 to 127
+ if (sfb >= -112) {
+ return [BigInt(sfb), cur];
+ }
+ // Multi-byte
+ const isNeg = sfb < -120; // unsigned 128–135 = negative; 136–143 = positive
+ const len = isNeg ? -119 - sfb : -111 - sfb; // total bytes incl. header
+ let value = 0n;
+ for (let i = 1; i < len; i++) {
+ const b = buf[cur];
+ if (b === undefined) {
+ throw new Error("ORC: truncated Hadoop VInt data");
+ }
+ cur++;
+ value = (value << 8n) | BigInt(b);
+ }
+ if (isNeg) {
+ value ^= -1n;
+ }
+ return [value, cur];
+}
+
+/** Write a Hadoop-style variable-length signed integer. */
+function hvWriteVInt(value: bigint, out: number[]): void {
+ if (value >= -112n && value <= 127n) {
+ out.push(Number(value < 0n ? value + 256n : value));
+ return;
+ }
+ let uval = value;
+ const isNeg = value < 0n;
+ if (isNeg) {
+ uval = value ^ -1n;
+ }
+ let nbytes = 0;
+ let tmp = uval;
+ while (tmp > 0n) {
+ tmp >>= 8n;
+ nbytes++;
+ }
+ const header = isNeg ? -120 - nbytes : -112 - nbytes;
+ out.push(header < 0 ? header + 0x100 : header);
+ for (let i = nbytes - 1; i >= 0; i--) {
+ out.push(Number((uval >> BigInt(i * 8)) & 0xffn));
+ }
+}
+
+// ─── RLE byte v1 ─────────────────────────────────────────────────────────────
+
+/**
+ * Decode an ORC RLE byte v1 stream to a flat byte array.
+ * Control byte < 128: run of (ctrl + 3) copies of the next byte.
+ * Control byte >= 128: (256 - ctrl) literal bytes follow.
+ */
+function rleByteDecodeV1(buf: Uint8Array, off: number, len: number): Uint8Array {
+ const end = off + len;
+ const out: number[] = [];
+ let pos = off;
+ while (pos < end) {
+ const ctrl = buf[pos];
+ if (ctrl === undefined) {
+ break;
+ }
+ pos++;
+ if (ctrl < 128) {
+ const count = ctrl + 3;
+ const val = buf[pos];
+ if (val === undefined) {
+ break;
+ }
+ pos++;
+ for (let i = 0; i < count; i++) {
+ out.push(val);
+ }
+ } else {
+ const count = 256 - ctrl;
+ for (let i = 0; i < count; i++) {
+ const b = buf[pos];
+ if (b === undefined) {
+ break;
+ }
+ pos++;
+ out.push(b);
+ }
+ }
+ }
+ return new Uint8Array(out);
+}
+
+/** Encode bytes using RLE byte v1. */
+function rleByteEncodeV1(data: readonly number[]): Uint8Array {
+ if (data.length === 0) {
+ return new Uint8Array(0);
+ }
+ const out: number[] = [];
+ let i = 0;
+ while (i < data.length) {
+ // Look for a run (same value repeated)
+ let runLen = 1;
+ while (runLen < 130 && i + runLen < data.length && data[i + runLen] === data[i]) {
+ runLen++;
+ }
+ if (runLen >= 3) {
+ out.push(runLen - 3);
+ const d = data[i];
+ if (d === undefined) {
+ throw new Error("ORC: undefined byte in run");
+ }
+ out.push(d);
+ i += runLen;
+ } else {
+ // Literal group
+ let litLen = 1;
+ while (litLen < 128 && i + litLen < data.length) {
+ // Stop if next 3 values are identical (start a new run)
+ const base = data[i + litLen];
+ let rcheck = 1;
+ while (
+ rcheck < 3 &&
+ i + litLen + rcheck < data.length &&
+ data[i + litLen + rcheck] === base
+ ) {
+ rcheck++;
+ }
+ if (rcheck >= 3) {
+ break;
+ }
+ litLen++;
+ }
+ out.push(256 - litLen);
+ for (let j = 0; j < litLen; j++) {
+ const d = data[i + j];
+ if (d === undefined) {
+ throw new Error("ORC: undefined byte in literal");
+ }
+ out.push(d);
+ }
+ i += litLen;
+ }
+ }
+ return new Uint8Array(out);
+}
+
+// ─── RLE integer v1 ──────────────────────────────────────────────────────────
+
+/**
+ * Decode an ORC RLE integer v1 stream (Hadoop VInts, big-endian).
+ * Control byte >= 0: run of (ctrl + 3) values, next byte is signed delta, then base VInt.
+ * Control byte < 0: (-ctrl) literal VInts.
+ */
+function rleIntDecodeV1(buf: Uint8Array, off: number, len: number): bigint[] {
+ const end = off + len;
+ const result: bigint[] = [];
+ let pos = off;
+ while (pos < end) {
+ const ctrl = buf[pos];
+ if (ctrl === undefined) {
+ break;
+ }
+ pos++;
+ const sctrl = ctrl >= 0x80 ? ctrl - 0x100 : ctrl; // signed
+ if (sctrl >= 0) {
+ const count = sctrl + 3;
+ const deltaByte = buf[pos];
+ if (deltaByte === undefined) {
+ break;
+ }
+ pos++;
+ const delta = BigInt(deltaByte >= 0x80 ? deltaByte - 0x100 : deltaByte);
+ let base: bigint;
+ [base, pos] = hvReadVInt(buf, pos);
+ for (let i = 0; i < count; i++) {
+ result.push(base + delta * BigInt(i));
+ }
+ } else {
+ const count = -sctrl;
+ for (let i = 0; i < count; i++) {
+ let v: bigint;
+ [v, pos] = hvReadVInt(buf, pos);
+ result.push(v);
+ }
+ }
+ }
+ return result;
+}
+
+/** Encode bigint values using RLE integer v1. */
+function rleIntEncodeV1(values: readonly bigint[]): Uint8Array {
+ if (values.length === 0) {
+ return new Uint8Array(0);
+ }
+ const out: number[] = [];
+ let i = 0;
+ while (i < values.length) {
+ const v0 = values[i];
+ if (v0 === undefined) {
+ break;
+ }
+ // Attempt to find a run with a constant delta
+ if (i + 2 < values.length) {
+ const v1 = values[i + 1];
+ const v2 = values[i + 2];
+ if (v1 !== undefined && v2 !== undefined) {
+ const delta = v1 - v0;
+ if (v2 - v1 === delta && delta >= -128n && delta <= 127n) {
+ let runLen = 3;
+ while (runLen < 130 && i + runLen < values.length) {
+ const vn = values[i + runLen];
+ const vprev = values[i + runLen - 1];
+ if (vn === undefined || vprev === undefined || vn - vprev !== delta) {
+ break;
+ }
+ runLen++;
+ }
+ out.push(runLen - 3);
+ out.push(Number(delta < 0n ? delta + 256n : delta));
+ hvWriteVInt(v0, out);
+ i += runLen;
+ continue;
+ }
+ }
+ }
+ // Literal group
+ let litLen = 1;
+ while (litLen < 128 && i + litLen < values.length) {
+ const va = values[i + litLen];
+ const vb = values[i + litLen + 1];
+ const vc = values[i + litLen + 2];
+ if (va !== undefined && vb !== undefined && vc !== undefined) {
+ const d1 = vb - va;
+ const d2 = vc - vb;
+ if (d1 === d2 && d1 >= -128n && d1 <= 127n) {
+ break;
+ }
+ }
+ litLen++;
+ }
+ out.push(256 - litLen);
+ for (let j = 0; j < litLen; j++) {
+ const v = values[i + j];
+ if (v === undefined) {
+ break;
+ }
+ hvWriteVInt(v, out);
+ }
+ i += litLen;
+ }
+ return new Uint8Array(out);
+}
+
+// ─── PRESENT stream helpers ───────────────────────────────────────────────────
+
+/**
+ * Expand a PRESENT-stream byte array into per-row boolean flags.
+ * Each byte = 8 rows, MSB first. 1 = non-null, 0 = null.
+ */
+function expandPresent(raw: Uint8Array, nRows: number): boolean[] {
+ const flags: boolean[] = [];
+ for (let i = 0; i < nRows; i++) {
+ flags.push(false);
+ }
+ let row = 0;
+ for (const byte of raw) {
+ for (let bit = 7; bit >= 0 && row < nRows; bit--, row++) {
+ flags[row] = ((byte >> bit) & 1) === 1;
+ }
+ }
+ return flags;
+}
+
+/**
+ * Pack per-row null flags into PRESENT-stream bytes.
+ * 1 = non-null, 0 = null. Returns null if all rows are non-null.
+ */
+function packPresent(nonNull: boolean[]): Uint8Array | null {
+ if (nonNull.every((v) => v)) {
+ return null;
+ }
+ const bytes: number[] = [];
+ for (let i = 0; i < nonNull.length; i += 8) {
+ let byte = 0;
+ for (let bit = 0; bit < 8 && i + bit < nonNull.length; bit++) {
+ if (nonNull[i + bit]) {
+ byte |= 1 << (7 - bit);
+ }
+ }
+ bytes.push(byte);
+ }
+ return new Uint8Array(bytes);
+}
+
+// ─── ORC metadata structures (decoded) ────────────────────────────────────────
+
+interface OrcPostscript {
+ footerLength: number;
+ compression: number;
+ compressionBlockSize: number;
+ metadataLength: number;
+ writerVersion: number;
+}
+
+interface OrcStripeInfo {
+ offset: number;
+ indexLength: number;
+ dataLength: number;
+ footerLength: number;
+ numberOfRows: number;
+}
+
+interface OrcType {
+ kind: number;
+ subtypes: number[];
+ fieldNames: string[];
+}
+
+interface OrcFooter {
+ stripes: OrcStripeInfo[];
+ types: OrcType[];
+ numberOfRows: number;
+}
+
+interface OrcStream {
+ kind: number;
+ column: number;
+ length: number;
+}
+
+interface OrcColumnEncoding {
+ kind: number;
+ dictionarySize: number;
+}
+
+interface OrcStripeFooter {
+ streams: OrcStream[];
+ columns: OrcColumnEncoding[];
+}
+
+// ─── ORC decoding helpers ─────────────────────────────────────────────────────
+
+function decodePostscript(buf: Uint8Array): OrcPostscript {
+ const msg = pbDecode(buf);
+ return {
+ footerLength: Number(pbU64(msg, 1)),
+ compression: pbU32(msg, 2),
+ compressionBlockSize: Number(pbU64(msg, 3)) || 262144,
+ metadataLength: Number(pbU64(msg, 5)),
+ writerVersion: pbU32(msg, 6),
+ };
+}
+
+function decodeFooter(buf: Uint8Array): OrcFooter {
+ const msg = pbDecode(buf);
+ const stripesMsgs = pbMsgs(msg, 3);
+ const stripes: OrcStripeInfo[] = stripesMsgs.map((sm) => ({
+ offset: Number(pbU64(sm, 1)),
+ indexLength: Number(pbU64(sm, 2)),
+ dataLength: Number(pbU64(sm, 3)),
+ footerLength: Number(pbU64(sm, 4)),
+ numberOfRows: Number(pbU64(sm, 5)),
+ }));
+ const typesMsgs = pbMsgs(msg, 4);
+ const types: OrcType[] = typesMsgs.map((tm) => ({
+ kind: pbU32(tm, 1),
+ subtypes: pbU32s(tm, 2),
+ fieldNames: pbStrings(tm, 3),
+ }));
+ return {
+ stripes,
+ types,
+ numberOfRows: Number(pbU64(msg, 6)),
+ };
+}
+
+function decodeStripeFooter(buf: Uint8Array): OrcStripeFooter {
+ const msg = pbDecode(buf);
+ const streamMsgs = pbMsgs(msg, 1);
+ const streams: OrcStream[] = streamMsgs.map((sm) => ({
+ kind: pbU32(sm, 1),
+ column: pbU32(sm, 2),
+ length: Number(pbU64(sm, 3)),
+ }));
+ const colMsgs = pbMsgs(msg, 2);
+ const columns: OrcColumnEncoding[] = colMsgs.map((cm) => ({
+ kind: pbU32(cm, 1),
+ dictionarySize: pbU32(cm, 2),
+ }));
+ return { streams, columns };
+}
+
+// ─── Column data decoders ─────────────────────────────────────────────────────
+
+function decodeIntCol(
+ dataBuf: Uint8Array,
+ off: number,
+ len: number,
+ presentFlags: boolean[] | null,
+ nRows: number,
+): (bigint | null)[] {
+ const ints = rleIntDecodeV1(dataBuf, off, len);
+ const result: (bigint | null)[] = [];
+ let dataIdx = 0;
+ for (let i = 0; i < nRows; i++) {
+ if (presentFlags !== null && presentFlags[i] === false) {
+ result.push(null);
+ } else {
+ result.push(ints[dataIdx] ?? null);
+ dataIdx++;
+ }
+ }
+ return result;
+}
+
+function decodeF32Col(
+ dataBuf: Uint8Array,
+ off: number,
+ presentFlags: boolean[] | null,
+ nRows: number,
+): (number | null)[] {
+ const dv = new DataView(dataBuf.buffer, dataBuf.byteOffset);
+ const result: (number | null)[] = [];
+ let pos = off;
+ for (let i = 0; i < nRows; i++) {
+ if (presentFlags !== null && presentFlags[i] === false) {
+ result.push(null);
+ } else {
+ result.push(dv.getFloat32(pos, true));
+ pos += 4;
+ }
+ }
+ return result;
+}
+
+function decodeF64Col(
+ dataBuf: Uint8Array,
+ off: number,
+ presentFlags: boolean[] | null,
+ nRows: number,
+): (number | null)[] {
+ const dv = new DataView(dataBuf.buffer, dataBuf.byteOffset);
+ const result: (number | null)[] = [];
+ let pos = off;
+ for (let i = 0; i < nRows; i++) {
+ if (presentFlags !== null && presentFlags[i] === false) {
+ result.push(null);
+ } else {
+ result.push(dv.getFloat64(pos, true));
+ pos += 8;
+ }
+ }
+ return result;
+}
+
+function decodeStringCol(
+ dataBuf: Uint8Array,
+ dataOff: number,
+ lenBuf: Uint8Array,
+ lenOff: number,
+ lenLen: number,
+ presentFlags: boolean[] | null,
+ nRows: number,
+): (string | null)[] {
+ const lengths = rleIntDecodeV1(lenBuf, lenOff, lenLen);
+ const dec = new TextDecoder();
+ const result: (string | null)[] = [];
+ let dataPos = dataOff;
+ let strIdx = 0;
+ for (let i = 0; i < nRows; i++) {
+ if (presentFlags !== null && presentFlags[i] === false) {
+ result.push(null);
+ } else {
+ const slen = Number(lengths[strIdx] ?? 0n);
+ result.push(dec.decode(dataBuf.subarray(dataPos, dataPos + slen)));
+ dataPos += slen;
+ strIdx++;
+ }
+ }
+ return result;
+}
+
+function decodeBoolCol(
+ dataBuf: Uint8Array,
+ off: number,
+ len: number,
+ presentFlags: boolean[] | null,
+ nRows: number,
+): (boolean | null)[] {
+ const raw = rleByteDecodeV1(dataBuf, off, len);
+ const result: (boolean | null)[] = [];
+ let bitIdx = 0;
+ for (let i = 0; i < nRows; i++) {
+ if (presentFlags !== null && presentFlags[i] === false) {
+ result.push(null);
+ } else {
+ const bytePos = Math.floor(bitIdx / 8);
+ const bitPos = 7 - (bitIdx % 8);
+ const byte = raw[bytePos] ?? 0;
+ result.push(((byte >> bitPos) & 1) === 1);
+ bitIdx++;
+ }
+ }
+ return result;
+}
+
+// ─── Column data encoders ─────────────────────────────────────────────────────
+
+interface EncodedCol {
+ presentStream: Uint8Array | null;
+ dataStream: Uint8Array;
+ lengthStream: Uint8Array | null;
+ typeKind: number;
+}
+
+function encodeIntCol(values: readonly Scalar[], typeKind: number): EncodedCol {
+ const nonNull: boolean[] = [];
+ const ints: bigint[] = [];
+ for (const v of values) {
+ const isPresent = v !== null && v !== undefined && !(typeof v === "number" && Number.isNaN(v));
+ nonNull.push(isPresent);
+ if (isPresent) {
+ ints.push(typeof v === "number" ? BigInt(Math.trunc(v)) : BigInt(String(v)));
+ }
+ }
+ const presentBits = packPresent(nonNull);
+ const presentStream = presentBits !== null ? rleByteEncodeV1(Array.from(presentBits)) : null;
+ return {
+ presentStream,
+ dataStream: rleIntEncodeV1(ints),
+ lengthStream: null,
+ typeKind,
+ };
+}
+
+function encodeF32Col(values: readonly Scalar[]): EncodedCol {
+ const nonNull: boolean[] = [];
+ const bytes: number[] = [];
+ const tmp = new ArrayBuffer(4);
+ const dv = new DataView(tmp);
+ for (const v of values) {
+ const isPresent = v !== null && v !== undefined && !(typeof v === "number" && Number.isNaN(v));
+ nonNull.push(isPresent);
+ if (isPresent) {
+ dv.setFloat32(0, typeof v === "number" ? v : Number(v), true);
+ bytes.push(dv.getUint8(0), dv.getUint8(1), dv.getUint8(2), dv.getUint8(3));
+ }
+ }
+ const presentBits = packPresent(nonNull);
+ return {
+ presentStream: presentBits !== null ? rleByteEncodeV1(Array.from(presentBits)) : null,
+ dataStream: new Uint8Array(bytes),
+ lengthStream: null,
+ typeKind: KIND_FLOAT,
+ };
+}
+
+function encodeF64Col(values: readonly Scalar[]): EncodedCol {
+ const nonNull: boolean[] = [];
+ const bytes: number[] = [];
+ const tmp = new ArrayBuffer(8);
+ const dv = new DataView(tmp);
+ for (const v of values) {
+ const isPresent = v !== null && v !== undefined && !(typeof v === "number" && Number.isNaN(v));
+ nonNull.push(isPresent);
+ if (isPresent) {
+ dv.setFloat64(0, typeof v === "number" ? v : Number(v), true);
+ for (let k = 0; k < 8; k++) {
+ bytes.push(dv.getUint8(k));
+ }
+ }
+ }
+ const presentBits = packPresent(nonNull);
+ return {
+ presentStream: presentBits !== null ? rleByteEncodeV1(Array.from(presentBits)) : null,
+ dataStream: new Uint8Array(bytes),
+ lengthStream: null,
+ typeKind: KIND_DOUBLE,
+ };
+}
+
+function encodeStringCol(values: readonly Scalar[]): EncodedCol {
+ const nonNull: boolean[] = [];
+ const dataBytes: number[] = [];
+ const lengths: bigint[] = [];
+ const enc = new TextEncoder();
+ for (const v of values) {
+ const isPresent = v !== null && v !== undefined;
+ nonNull.push(isPresent);
+ if (isPresent) {
+ const bytes = enc.encode(String(v));
+ for (const b of bytes) {
+ dataBytes.push(b);
+ }
+ lengths.push(BigInt(bytes.length));
+ }
+ }
+ const presentBits = packPresent(nonNull);
+ return {
+ presentStream: presentBits !== null ? rleByteEncodeV1(Array.from(presentBits)) : null,
+ dataStream: new Uint8Array(dataBytes),
+ lengthStream: rleIntEncodeV1(lengths),
+ typeKind: KIND_STRING,
+ };
+}
+
+function encodeBoolCol(values: readonly Scalar[]): EncodedCol {
+ const nonNull: boolean[] = [];
+ const bits: boolean[] = [];
+ for (const v of values) {
+ const isPresent = v !== null && v !== undefined;
+ nonNull.push(isPresent);
+ if (isPresent) {
+ bits.push(Boolean(v));
+ }
+ }
+ // Pack booleans into bytes, MSB first
+ const bytes: number[] = [];
+ for (let i = 0; i < bits.length; i += 8) {
+ let byte = 0;
+ for (let b = 0; b < 8 && i + b < bits.length; b++) {
+ if (bits[i + b]) {
+ byte |= 1 << (7 - b);
+ }
+ }
+ bytes.push(byte);
+ }
+ const presentBits = packPresent(nonNull);
+ return {
+ presentStream: presentBits !== null ? rleByteEncodeV1(Array.from(presentBits)) : null,
+ dataStream: rleByteEncodeV1(bytes),
+ lengthStream: null,
+ typeKind: KIND_BOOLEAN,
+ };
+}
+
+// ─── ORC type inference ───────────────────────────────────────────────────────
+
+/** Map a DataFrame column's values to an ORC type kind. */
+function inferOrcKind(values: readonly Scalar[]): number {
+ for (const v of values) {
+ if (v === null || v === undefined) {
+ continue;
+ }
+ if (typeof v === "boolean") {
+ return KIND_BOOLEAN;
+ }
+ if (typeof v === "bigint") {
+ return KIND_LONG;
+ }
+ if (typeof v === "number") {
+ return 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.at(-1);
+ if (psLen === undefined || psLen === 0) {
+ throw new Error("ORC: invalid postscript length");
+ }
+ const psStart = buf.length - 1 - psLen;
+ if (psStart < 3) {
+ throw new Error("ORC: file too small for postscript");
+ }
+ const ps = decodePostscript(buf.subarray(psStart, psStart + psLen));
+ if (ps.compression !== COMP_NONE) {
+ throw new Error(
+ `ORC: compression codec ${ps.compression} (${ps.compression === COMP_ZLIB ? "ZLIB" : "unsupported"}) is not supported; only NONE is currently implemented`,
+ );
+ }
+ // Read footer
+ const metaEnd = psStart;
+ const footerEnd = metaEnd - ps.metadataLength;
+ const footerStart = footerEnd - ps.footerLength;
+ if (footerStart < 3) {
+ throw new Error("ORC: invalid footer position");
+ }
+ const footer = decodeFooter(buf.subarray(footerStart, footerEnd));
+ if (footer.types.length === 0) {
+ throw new Error("ORC: no type schema in footer");
+ }
+
+ // Root type must be STRUCT
+ const rootType = footer.types[0];
+ if (rootType === undefined || rootType.kind !== KIND_STRUCT) {
+ throw new Error("ORC: root type is not STRUCT");
+ }
+
+ // Column selection
+ const allCols = rootType.fieldNames;
+ const wantSet = options.columns != null ? new Set([...options.columns]) : null;
+ const colIndices: number[] = rootType.subtypes.filter((_, i) => {
+ const name = allCols[i];
+ return name !== undefined && (wantSet === null || wantSet.has(name));
+ });
+ const colNames: string[] = colIndices.map((ci) => {
+ const fi = rootType.subtypes.indexOf(ci);
+ return allCols[fi] ?? String(ci);
+ });
+
+ const allValues: Map = new Map();
+ for (const name of colNames) {
+ allValues.set(name, []);
+ }
+
+ // Process each stripe
+ for (const stripeInfo of footer.stripes) {
+ const stripeDataStart = stripeInfo.offset + stripeInfo.indexLength;
+ const stripeFStart = stripeInfo.offset + stripeInfo.indexLength + stripeInfo.dataLength;
+ const stripeFBuf = buf.subarray(stripeFStart, stripeFStart + stripeInfo.footerLength);
+ const sf = decodeStripeFooter(stripeFBuf);
+
+ // Build a stream-offset map: column → streamKind → {offset, length}
+ type StreamLoc = { offset: number; length: number };
+ const streamMap = new Map>();
+ let streamPos = stripeDataStart;
+ for (const stream of sf.streams) {
+ let colMap = streamMap.get(stream.column);
+ if (colMap === undefined) {
+ colMap = new Map();
+ streamMap.set(stream.column, colMap);
+ }
+ colMap.set(stream.kind, { offset: streamPos, length: stream.length });
+ streamPos += stream.length;
+ }
+
+ const nRows = stripeInfo.numberOfRows;
+
+ for (let ci = 0; ci < colIndices.length; ci++) {
+ const colIdx = colIndices[ci];
+ if (colIdx === undefined) {
+ continue;
+ }
+ const name = colNames[ci];
+ if (name === undefined) {
+ continue;
+ }
+ const typeKind = footer.types[colIdx]?.kind ?? KIND_STRING;
+ const colStreams = streamMap.get(colIdx);
+ const vals = allValues.get(name);
+ if (vals === undefined) {
+ continue;
+ }
+
+ // PRESENT stream (null flags)
+ const presentLoc = colStreams?.get(STREAM_PRESENT);
+ let presentFlags: boolean[] | null = null;
+ if (presentLoc !== undefined && presentLoc.length > 0) {
+ const rawPresent = rleByteDecodeV1(buf, presentLoc.offset, presentLoc.length);
+ presentFlags = expandPresent(rawPresent, nRows);
+ }
+
+ const dataLoc = colStreams?.get(STREAM_DATA);
+ const dataOff = dataLoc?.offset ?? 0;
+ const dataLen = dataLoc?.length ?? 0;
+
+ switch (typeKind) {
+ case KIND_BOOLEAN: {
+ const decoded = decodeBoolCol(buf, dataOff, dataLen, presentFlags, nRows);
+ for (const v of decoded) {
+ vals.push(v);
+ }
+ break;
+ }
+ case KIND_FLOAT: {
+ const decoded = decodeF32Col(buf, dataOff, presentFlags, nRows);
+ for (const v of decoded) {
+ vals.push(v);
+ }
+ break;
+ }
+ case KIND_DOUBLE: {
+ const decoded = decodeF64Col(buf, dataOff, presentFlags, nRows);
+ for (const v of decoded) {
+ vals.push(v);
+ }
+ break;
+ }
+ case KIND_STRING: {
+ const lenLoc = colStreams?.get(STREAM_LENGTH);
+ const decoded = decodeStringCol(
+ buf,
+ dataOff,
+ buf,
+ lenLoc?.offset ?? 0,
+ lenLoc?.length ?? 0,
+ presentFlags,
+ nRows,
+ );
+ for (const v of decoded) {
+ vals.push(v);
+ }
+ break;
+ }
+ default: {
+ // Integer types: BOOLEAN, BYTE, SHORT, INT, LONG, DATE
+ const decoded = decodeIntCol(buf, dataOff, dataLen, presentFlags, nRows);
+ for (const v of decoded) {
+ vals.push(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
+ const _streamPos = stripeOffset;
+
+ for (let ci = 0; ci < encoded.length; ci++) {
+ const enc = encoded[ci];
+ if (enc === undefined) {
+ continue;
+ }
+ const colIdx = ci + 1; // 1-based (0 = root STRUCT)
+
+ if (enc.presentStream !== null) {
+ streams.push({ kind: STREAM_PRESENT, column: colIdx, length: enc.presentStream.length });
+ }
+ streams.push({ kind: STREAM_DATA, column: colIdx, length: enc.dataStream.length });
+ if (enc.lengthStream !== null) {
+ streams.push({ kind: STREAM_LENGTH, column: colIdx, length: enc.lengthStream.length });
+ }
+ colEncodings.push({ kind: ENC_DIRECT, dictionarySize: 0 });
+ }
+
+ // Write stream data
+ const stripeIndexLen = 0; // no row indexes
+ for (let ci = 0; ci < encoded.length; ci++) {
+ const enc = encoded[ci];
+ if (enc === undefined) {
+ continue;
+ }
+ if (enc.presentStream !== null) {
+ for (const b of enc.presentStream) {
+ out.push(b);
+ }
+ }
+ for (const b of enc.dataStream) {
+ out.push(b);
+ }
+ if (enc.lengthStream !== null) {
+ for (const b of enc.lengthStream) {
+ out.push(b);
+ }
+ }
+ }
+
+ // Compute data length (written bytes minus header)
+ const stripeDataLen = out.length - stripeOffset;
+
+ // Stripe footer
+ const sf = encodeStripeFooter(streams, colEncodings);
+ for (const b of sf) {
+ out.push(b);
+ }
+
+ // Build ORC type schema
+ // Column 0: STRUCT with all column fields
+ const types: OrcType[] = [
+ {
+ kind: KIND_STRUCT,
+ subtypes: extraCols.map((_, i) => i + 1),
+ fieldNames: extraCols,
+ },
+ ];
+ for (const kind of colKinds) {
+ types.push({ kind, subtypes: [], fieldNames: [] });
+ }
+
+ // Stripe info
+ const stripeInfo: OrcStripeInfo = {
+ offset: stripeOffset,
+ indexLength: stripeIndexLen,
+ dataLength: stripeDataLen,
+ footerLength: sf.length,
+ numberOfRows: nRows,
+ };
+
+ // File footer
+ const footerBytes = encodeFooter([stripeInfo], types, nRows);
+ for (const b of footerBytes) {
+ out.push(b);
+ }
+
+ // File metadata (empty for now)
+ const metaLen = 0;
+
+ // Postscript
+ const psBytes = encodePostscript(footerBytes.length, metaLen);
+ for (const b of psBytes) {
+ out.push(b);
+ }
+
+ // Postscript length (1 byte)
+ if (psBytes.length > 255) {
+ throw new Error("ORC: postscript too large");
+ }
+ out.push(psBytes.length);
+
+ return new Uint8Array(out);
+}
diff --git a/src/io/read_avro.ts b/src/io/read_avro.ts
new file mode 100644
index 00000000..9b2b7f4d
--- /dev/null
+++ b/src/io/read_avro.ts
@@ -0,0 +1,698 @@
+/**
+ * read_avro — Apache Avro Object Container File (OCF) reader for DataFrame.
+ *
+ * Mirrors `pandas.read_avro()`. Parses Avro OCF binary format purely in
+ * TypeScript with no external dependencies.
+ *
+ * Supported Avro schema types:
+ * - Primitives: null, boolean, int, long, float, double, string, bytes
+ * - Named: record, enum, fixed
+ * - Complex: array, map, union
+ * - Logical: date (int), timestamp-millis (long), timestamp-micros (long)
+ *
+ * Supported codecs: `null` (uncompressed). `deflate` and `snappy` blocks
+ * are detected but raise an informative error.
+ *
+ * @example
+ * ```ts
+ * import { readAvro } from "tsb";
+ *
+ * const df = readAvro(buffer); // Uint8Array from file read
+ * console.log(df.columns, df.shape);
+ * ```
+ *
+ * @module
+ */
+
+import { DataFrame } from "../core/frame.ts";
+import type { Scalar } from "../types.ts";
+
+// ─── Public types ──────────────────────────────────────────────────────────────
+
+/** Options for {@link readAvro}. */
+export interface ReadAvroOptions {
+ /** Columns to include. Default: all columns. */
+ readonly usecols?: readonly string[] | null;
+ /**
+ * How to handle schema unions that include `null`:
+ * - `"object"` (default): return `null` for null values.
+ * - `"first"`: return the first non-null type's value.
+ */
+ readonly nullHandling?: "object" | "first";
+}
+
+// ─── Avro schema types ────────────────────────────────────────────────────────
+
+type AvroSchema =
+ | AvroPrimitive
+ | AvroRecord
+ | AvroEnum
+ | AvroArray
+ | AvroMap
+ | AvroUnion
+ | AvroFixed;
+
+type AvroPrimitive = "null" | "boolean" | "int" | "long" | "float" | "double" | "string" | "bytes";
+
+interface AvroRecord {
+ type: "record";
+ name: string;
+ fields: readonly AvroField[];
+}
+
+interface AvroField {
+ name: string;
+ type: AvroSchema;
+ default?: unknown;
+}
+
+interface AvroEnum {
+ type: "enum";
+ name: string;
+ symbols: readonly string[];
+}
+
+interface AvroArray {
+ type: "array";
+ items: AvroSchema;
+}
+
+interface AvroMap {
+ type: "map";
+ values: AvroSchema;
+}
+
+type AvroUnion = readonly AvroSchema[];
+
+interface AvroFixed {
+ type: "fixed";
+ name: string;
+ size: number;
+}
+
+// ─── Binary reader ────────────────────────────────────────────────────────────
+
+class AvroReader {
+ private buf: Uint8Array;
+ private pos = 0;
+
+ constructor(buf: Uint8Array) {
+ this.buf = buf;
+ }
+
+ get position(): number {
+ return this.pos;
+ }
+ get remaining(): number {
+ return this.buf.length - this.pos;
+ }
+
+ readByte(): number {
+ if (this.pos >= this.buf.length) {
+ throw new RangeError("Unexpected end of Avro data");
+ }
+ return this.buf[this.pos++] ?? 0;
+ }
+
+ /** Read a variable-length zigzag-encoded long. Returns a JS number (safe up to 2^53). */
+ readLong(): number {
+ let result = 0;
+ let shift = 0;
+ while (true) {
+ const b = this.readByte();
+ result |= (b & 0x7f) << shift;
+ shift += 7;
+ if ((b & 0x80) === 0) {
+ break;
+ }
+ if (shift >= 63) {
+ // For very large values, handle the remaining bits separately
+ // to avoid JS bitwise overflow (32-bit integers)
+ if (shift === 63) {
+ const hi = b & 0x7f;
+ // combine: result (low 63 bits) + hi << 63 — just approximate as float
+ const lo = result >>> 0;
+ result = lo + hi * 2 ** 63;
+ }
+ break;
+ }
+ }
+ // Zigzag decode: (n >>> 1) ^ -(n & 1)
+ return (result >>> 1) ^ -(result & 1);
+ }
+
+ /** Read a 32-bit int (zigzag long with range check). */
+ readInt(): number {
+ return this.readLong() | 0;
+ }
+
+ /** Read 4-byte IEEE 754 float. */
+ readFloat(): number {
+ const bytes = this.readBytes(4);
+ const view = new DataView(bytes.buffer, bytes.byteOffset, 4);
+ return view.getFloat32(0, true);
+ }
+
+ /** Read 8-byte IEEE 754 double. */
+ readDouble(): number {
+ const bytes = this.readBytes(8);
+ const view = new DataView(bytes.buffer, bytes.byteOffset, 8);
+ return view.getFloat64(0, true);
+ }
+
+ /** Read exactly n bytes as a new Uint8Array. */
+ readBytes(n: number): Uint8Array {
+ if (this.pos + n > this.buf.length) {
+ throw new RangeError("Unexpected end of Avro data");
+ }
+ const slice = this.buf.subarray(this.pos, this.pos + n);
+ this.pos += n;
+ return slice;
+ }
+
+ /** Read Avro bytes field (length-prefixed). */
+ readByteField(): Uint8Array {
+ const len = this.readLong();
+ return this.readBytes(len);
+ }
+
+ /** Read UTF-8 string (length-prefixed). */
+ readString(): string {
+ const bytes = this.readByteField();
+ return new TextDecoder().decode(bytes);
+ }
+
+ /** Read a boolean (0 = false, 1 = true). */
+ readBoolean(): boolean {
+ return this.readByte() !== 0;
+ }
+
+ /** Skip forward n bytes. */
+ skip(n: number): void {
+ if (this.pos + n > this.buf.length) {
+ throw new RangeError("Unexpected end of Avro data");
+ }
+ this.pos += n;
+ }
+
+ /** Peek at 16 bytes (sync marker) and advance. */
+ readSync(): Uint8Array {
+ return this.readBytes(16);
+ }
+}
+
+// ─── Schema parsing ────────────────────────────────────────────────────────────
+
+const td = new TextDecoder();
+
+function parseSchema(raw: unknown): AvroSchema {
+ if (typeof raw === "string") {
+ const prim = raw as AvroPrimitive;
+ return prim;
+ }
+ if (Array.isArray(raw)) {
+ return raw.map(parseSchema) as AvroUnion;
+ }
+ if (typeof raw === "object" && raw !== null) {
+ const obj = raw as Record;
+ const type = obj.type;
+ if (type === "record") {
+ const fields = (obj.fields as unknown[]).map((f) => {
+ const field = f as Record;
+ return { name: field.name as string, type: parseSchema(field.type) };
+ });
+ return { type: "record", name: obj.name as string, fields };
+ }
+ if (type === "array") {
+ return { type: "array", items: parseSchema(obj.items) };
+ }
+ if (type === "map") {
+ return { type: "map", values: parseSchema(obj.values) };
+ }
+ if (type === "enum") {
+ return {
+ type: "enum",
+ name: obj.name as string,
+ symbols: obj.symbols as string[],
+ };
+ }
+ if (type === "fixed") {
+ return {
+ type: "fixed",
+ name: obj.name as string,
+ size: obj.size as number,
+ };
+ }
+ // Logical types: delegate to the underlying type
+ if (typeof type === "string") {
+ return parseSchema(type);
+ }
+ }
+ throw new TypeError(`Unknown Avro schema: ${JSON.stringify(raw)}`);
+}
+
+// ─── Datum reading ────────────────────────────────────────────────────────────
+
+/** Avro leaf value (no recursion at type level — containers use unknown). */
+type AvroLeaf = null | boolean | number | string | Uint8Array;
+/** Container interfaces allow recursive self-reference (interfaces can, type aliases cannot). */
+interface AvroDatumArr extends Array {}
+interface AvroDatumMap extends Map {}
+interface AvroDatumRecord extends Record {}
+/** Recursive Avro datum. */
+type AvroDatum = AvroLeaf | AvroDatumArr | AvroDatumMap | AvroDatumRecord;
+
+function readDatum(reader: AvroReader, schema: AvroSchema): AvroDatum {
+ if (typeof schema === "string") {
+ switch (schema) {
+ case "null":
+ return null;
+ case "boolean":
+ return reader.readBoolean();
+ case "int":
+ return reader.readInt();
+ case "long":
+ return reader.readLong();
+ case "float":
+ return reader.readFloat();
+ case "double":
+ return reader.readDouble();
+ case "string":
+ return reader.readString();
+ case "bytes":
+ return reader.readByteField();
+ }
+ }
+ if (Array.isArray(schema)) {
+ // Union: first read the branch index
+ const idx = reader.readLong();
+ const branch = (schema as AvroUnion)[idx];
+ if (branch === undefined) {
+ throw new RangeError(`Union branch ${idx} out of range`);
+ }
+ return readDatum(reader, branch);
+ }
+ const s = schema as Exclude;
+ if (s.type === "record") {
+ const rec: Record = {};
+ for (const field of s.fields) {
+ rec[field.name] = readDatum(reader, field.type);
+ }
+ return rec;
+ }
+ if (s.type === "array") {
+ const arr: AvroDatum[] = [];
+ while (true) {
+ let count = reader.readLong();
+ if (count === 0) {
+ break;
+ }
+ // Negative count means block has a byte count prefix
+ if (count < 0) {
+ reader.readLong();
+ count = -count;
+ }
+ for (let i = 0; i < count; i++) {
+ arr.push(readDatum(reader, s.items));
+ }
+ }
+ return arr;
+ }
+ if (s.type === "map") {
+ const map = new Map();
+ while (true) {
+ let count = reader.readLong();
+ if (count === 0) {
+ break;
+ }
+ if (count < 0) {
+ reader.readLong();
+ count = -count;
+ }
+ for (let i = 0; i < count; i++) {
+ const key = reader.readString();
+ map.set(key, readDatum(reader, s.values));
+ }
+ }
+ return map;
+ }
+ if (s.type === "enum") {
+ const idx = reader.readInt();
+ return s.symbols[idx] ?? null;
+ }
+ if (s.type === "fixed") {
+ return reader.readBytes(s.size);
+ }
+ throw new TypeError(`Unhandled schema type: ${JSON.stringify(schema)}`);
+}
+
+// ─── OCF parsing ──────────────────────────────────────────────────────────────
+
+const AVRO_MAGIC = new Uint8Array([79, 98, 106, 1]); // "Obj\x01"
+
+function syncEq(a: Uint8Array, b: Uint8Array): boolean {
+ if (a.length !== b.length) {
+ return false;
+ }
+ for (let i = 0; i < a.length; i++) {
+ if (a[i] !== b[i]) {
+ return false;
+ }
+ }
+ return true;
+}
+
+/**
+ * Parse an Apache Avro Object Container File buffer into an array of row objects.
+ * Returns the top-level schema and the rows.
+ */
+function parseAvroOCF(buf: Uint8Array): { schema: AvroSchema; rows: Record[] } {
+ const reader = new AvroReader(buf);
+
+ // Magic: "Obj\x01"
+ const magic = reader.readBytes(4);
+ if (!syncEq(magic, AVRO_MAGIC)) {
+ throw new TypeError(
+ `Not a valid Avro file: expected magic bytes "Obj\\x01", got ${[...magic].map((b) => b.toString(16)).join(" ")}`,
+ );
+ }
+
+ // File-level metadata: map
+ const meta = new Map();
+ while (true) {
+ let count = reader.readLong();
+ if (count === 0) {
+ break;
+ }
+ if (count < 0) {
+ reader.readLong();
+ count = -count;
+ }
+ for (let i = 0; i < count; i++) {
+ const key = reader.readString();
+ const val = reader.readByteField();
+ meta.set(key, val);
+ }
+ }
+
+ // Schema
+ const schemaBytes = meta.get("avro.schema");
+ if (!schemaBytes) {
+ throw new TypeError("Avro file missing avro.schema metadata");
+ }
+ const schemaJson: unknown = JSON.parse(td.decode(schemaBytes));
+ const schema = parseSchema(schemaJson);
+
+ // Codec
+ const codecBytes = meta.get("avro.codec");
+ const codec = codecBytes ? td.decode(codecBytes) : "null";
+ if (codec !== "null") {
+ throw new TypeError(
+ `Avro codec "${codec}" is not supported. Only "null" (uncompressed) is implemented.`,
+ );
+ }
+
+ // Sync marker (16 bytes)
+ const syncMarker = reader.readSync();
+
+ // Data blocks
+ const rows: Record[] = [];
+ while (reader.remaining >= 16) {
+ const objectCount = reader.readLong();
+ const _byteCount = reader.readLong(); // block size (unused for null codec)
+ if (objectCount <= 0) {
+ break;
+ }
+
+ for (let i = 0; i < objectCount; i++) {
+ const datum = readDatum(reader, schema);
+ if (
+ typeof datum === "object" &&
+ datum !== null &&
+ !Array.isArray(datum) &&
+ !(datum instanceof Uint8Array) &&
+ !(datum instanceof Map)
+ ) {
+ rows.push(datum as Record);
+ }
+ }
+
+ // Read and verify sync marker
+ const blockSync = reader.readSync();
+ if (!syncEq(blockSync, syncMarker)) {
+ throw new TypeError("Avro sync marker mismatch — file may be corrupt");
+ }
+ }
+
+ return { schema, rows };
+}
+
+// ─── DataFrame construction ───────────────────────────────────────────────────
+
+function flattenDatum(v: AvroDatum): unknown {
+ if (v === null || typeof v !== "object") {
+ return v;
+ }
+ if (v instanceof Uint8Array) {
+ return v;
+ }
+ if (v instanceof Map) {
+ return Object.fromEntries(v);
+ }
+ // For record/array datums, JSON-stringify for simplicity
+ if (Array.isArray(v)) {
+ return JSON.stringify(v);
+ }
+ return JSON.stringify(v);
+}
+
+/**
+ * Read an Apache Avro Object Container File buffer into a {@link DataFrame}.
+ *
+ * @param data - Raw Avro OCF bytes (`Uint8Array` or `ArrayBuffer`).
+ * @param options - Optional read options.
+ */
+export function readAvro(data: Uint8Array | ArrayBuffer, options: ReadAvroOptions = {}): DataFrame {
+ const buf = data instanceof ArrayBuffer ? new Uint8Array(data) : data;
+ const { rows } = parseAvroOCF(buf);
+
+ if (rows.length === 0) {
+ return DataFrame.fromColumns({});
+ }
+
+ // Determine columns from first row
+ const allCols = Object.keys(rows[0] ?? {});
+ const cols = options.usecols
+ ? allCols.filter((c) => (options.usecols as readonly string[]).includes(c))
+ : allCols;
+
+ const columns: Record = {};
+ for (const col of cols) {
+ columns[col] = [];
+ }
+
+ for (const row of rows) {
+ for (const col of cols) {
+ const v = row[col];
+ (columns[col] ?? []).push(flattenDatum(v ?? null) as Scalar);
+ }
+ }
+
+ return DataFrame.fromColumns(columns);
+}
+
+// ─── Avro writer (minimal — for testing round-trips) ─────────────────────────
+
+/** Options for {@link toAvro}. */
+export interface ToAvroOptions {
+ /** Schema name for the top-level record. Default: "Row". */
+ readonly schemaName?: string;
+}
+
+/**
+ * Serialize a {@link DataFrame} to an uncompressed Avro OCF buffer.
+ *
+ * Column type mapping:
+ * - boolean columns → `boolean`
+ * - integer columns → `long`
+ * - float columns → `double`
+ * - string columns → `{"type":"union","schemas":["null","string"]}`
+ * - null column vals → wrapped in union `["null", ""]`
+ * - other → `string` (JSON-stringified)
+ */
+export function toAvro(df: DataFrame, options: ToAvroOptions = {}): Uint8Array {
+ const schemaName = options.schemaName ?? "Row";
+ const cols = [...df.columns.values];
+
+ // Infer field types
+ type FieldSpec = { name: string; avroType: string; nullable: boolean };
+ const fields: FieldSpec[] = cols.map((col) => {
+ const vals = df.col(col).values;
+ let hasNull = false;
+ let hasInt = false;
+ let hasFloat = false;
+ let hasBool = false;
+ let hasStr = false;
+ for (const v of vals) {
+ if (v === null || v === undefined) {
+ hasNull = true;
+ continue;
+ }
+ if (typeof v === "boolean") {
+ hasBool = true;
+ continue;
+ }
+ if (typeof v === "number") {
+ if (Number.isInteger(v)) {
+ hasInt = true;
+ } else {
+ hasFloat = true;
+ }
+ continue;
+ }
+ hasStr = true;
+ }
+ let avroType = "string";
+ if (hasBool && !hasInt && !hasFloat && !hasStr) {
+ avroType = "boolean";
+ } else if ((hasInt || hasFloat) && !hasBool && !hasStr) {
+ avroType = hasFloat ? "double" : "long";
+ }
+ return { name: col, avroType, nullable: hasNull };
+ });
+
+ // Build schema JSON
+ const schemaFields = fields.map((f) => ({
+ name: f.name,
+ type: f.nullable ? ["null", f.avroType] : f.avroType,
+ }));
+ const schemaObj = { type: "record", name: schemaName, fields: schemaFields };
+ const schemaJson = JSON.stringify(schemaObj);
+ const schemaBytes = new TextEncoder().encode(schemaJson);
+
+ // Sync marker: 16 random-ish bytes derived from schema hash
+ const sync = new Uint8Array(16);
+ let h = 0x12345678;
+ for (let i = 0; i < schemaBytes.length; i++) {
+ h = Math.imul(h ^ (schemaBytes[i] ?? 0), 0x9e3779b9) >>> 0;
+ }
+ for (let i = 0; i < 16; i++) {
+ sync[i] = (h >> ((i % 4) * 8)) & 0xff;
+ if (i % 4 === 3) {
+ h = Math.imul(h, 0x6c62272e) >>> 0;
+ }
+ }
+
+ const chunks: Uint8Array[] = [];
+
+ // Write helper
+ const writeBuf: number[] = [];
+ const flushBuf = (): Uint8Array => {
+ const u = new Uint8Array(writeBuf);
+ writeBuf.length = 0;
+ return u;
+ };
+
+ function writeLong(v: number): void {
+ // Zigzag encode
+ let n = (v << 1) ^ (v >> 31);
+ while (n & ~0x7f) {
+ writeBuf.push((n & 0x7f) | 0x80);
+ n >>>= 7;
+ }
+ writeBuf.push(n);
+ }
+ function writeString(s: string): void {
+ const b = new TextEncoder().encode(s);
+ writeLong(b.length);
+ for (const byte of b) {
+ writeBuf.push(byte);
+ }
+ }
+ function writeBytes(b: Uint8Array): void {
+ writeLong(b.length);
+ for (const byte of b) {
+ writeBuf.push(byte);
+ }
+ }
+
+ // Magic
+ chunks.push(AVRO_MAGIC);
+
+ // Metadata: 1 map block with avro.schema and avro.codec
+ writeLong(2); // 2 entries
+ writeString("avro.schema");
+ writeBytes(schemaBytes);
+ writeString("avro.codec");
+ writeBytes(new TextEncoder().encode("null"));
+ writeLong(0); // end of map
+ chunks.push(flushBuf());
+
+ // Sync marker
+ chunks.push(sync.slice());
+
+ // Data block
+ const nRows = df.shape[0];
+ if (nRows > 0) {
+ // Encode all rows
+ for (let row = 0; row < nRows; row++) {
+ for (const f of fields) {
+ const v = df.col(f.name).at(row);
+ if (f.nullable) {
+ if (v === null || v === undefined) {
+ writeLong(0); // null branch
+ } else {
+ writeLong(1); // value branch
+ writeTypedValue(v, f.avroType);
+ }
+ } else {
+ writeTypedValue(v ?? null, f.avroType);
+ }
+ }
+ }
+ const blockData = flushBuf();
+
+ // Block header: count, byteCount
+ writeLong(nRows);
+ writeLong(blockData.length);
+ chunks.push(flushBuf());
+ chunks.push(blockData);
+ chunks.push(sync.slice());
+ }
+
+ // Concatenate all chunks
+ const totalLen = chunks.reduce((s, c) => s + c.length, 0);
+ const out = new Uint8Array(totalLen);
+ let offset = 0;
+ for (const c of chunks) {
+ out.set(c, offset);
+ offset += c.length;
+ }
+ return out;
+
+ function writeTypedValue(v: unknown, type: string): void {
+ if (v === null || v === undefined) {
+ writeBuf.push(0);
+ return;
+ } // null
+ switch (type) {
+ case "boolean":
+ writeBuf.push(v ? 1 : 0);
+ break;
+ case "long":
+ writeLong(typeof v === "number" ? v : 0);
+ break;
+ case "double": {
+ const arr = new Float64Array(1);
+ arr[0] = typeof v === "number" ? v : 0;
+ const b = new Uint8Array(arr.buffer);
+ for (const byte of b) {
+ writeBuf.push(byte);
+ }
+ break;
+ }
+ default:
+ writeString(String(v));
+ }
+ }
+}
diff --git a/src/stats/acf_pacf.ts b/src/stats/acf_pacf.ts
new file mode 100644
index 00000000..5edc712b
--- /dev/null
+++ b/src/stats/acf_pacf.ts
@@ -0,0 +1,618 @@
+/**
+ * acf_pacf — Autocorrelation and partial autocorrelation functions.
+ *
+ * Mirrors `statsmodels.tsa.stattools.*` for time-series correlation analysis,
+ * plus `pd.Series.autocorr`. Implemented from scratch with no external deps.
+ *
+ * Implemented functions:
+ * - {@link autocorr} — pandas-style single-lag autocorrelation
+ * - {@link acf} — full ACF with Bartlett confidence intervals
+ * - {@link pacf} — PACF via Levinson-Durbin recursion
+ * - {@link ccf} — cross-correlation function
+ * - {@link durbinWatson} — Durbin-Watson statistic for residual autocorrelation
+ * - {@link ljungBox} — Ljung-Box portmanteau test
+ * - {@link boxPierce} — Box-Pierce portmanteau test
+ *
+ * @example
+ * ```ts
+ * import { acf, pacf, ljungBox } from "tsb";
+ *
+ * const x = [1, 2, 3, 2, 1, 2, 3, 2, 1, 0, 1, 2];
+ * const { acf: corrs, lags } = acf(x, { nlags: 4, alpha: 0.05 });
+ * const { pacf: partial } = pacf(x, { nlags: 4 });
+ * const { pvalue } = ljungBox(x);
+ * ```
+ *
+ * @module
+ */
+
+import { Series } from "../core/index.ts";
+
+// ─── public types ──────────────────────────────────────────────────────────────
+
+/** Result from {@link acf}. */
+export interface ACFResult {
+ /** Autocorrelation coefficients; index 0 corresponds to lag 0 (= 1.0). */
+ readonly acf: readonly number[];
+ /**
+ * Confidence interval bounds `[lower, upper]` for each lag, computed via
+ * Bartlett's formula. `undefined` when `alpha` was not specified.
+ */
+ readonly confint: readonly [number, number][] | undefined;
+ /** Lag indices corresponding to each coefficient. */
+ readonly lags: readonly number[];
+}
+
+/** Result from {@link pacf}. */
+export interface PACFResult {
+ /** Partial autocorrelations; index 0 corresponds to lag 0 (= 1.0). */
+ readonly pacf: readonly number[];
+ /**
+ * Confidence interval bounds `[lower, upper]` for each lag.
+ * `undefined` when `alpha` was not specified.
+ */
+ readonly confint: readonly [number, number][] | undefined;
+ /** Lag indices. */
+ readonly lags: readonly number[];
+}
+
+/** Result from {@link ljungBox} or {@link boxPierce}. */
+export interface PortmanteauResult {
+ /** Q statistic at each tested lag. */
+ readonly statistic: readonly number[];
+ /** p-value at each tested lag (chi-squared df = lag − modelDf). */
+ readonly pvalue: readonly number[];
+ /** Lag indices that were tested. */
+ readonly lags: readonly number[];
+}
+
+/** Options for {@link acf}. */
+export interface ACFOptions {
+ /** Maximum lag to compute (default: `min(floor(10·log₁₀(n)), n−1)`). */
+ readonly nlags?: number;
+ /**
+ * Significance level for Bartlett confidence intervals (e.g. `0.05` → 95 % CI).
+ * Omit (default) to skip CI computation.
+ */
+ readonly alpha?: number;
+}
+
+/** Options for {@link pacf}. */
+export interface PACFOptions {
+ /** Maximum lag (default: `min(floor(10·log₁₀(n)), floor(n/2)−1)`). */
+ readonly nlags?: number;
+ /**
+ * Significance level for confidence intervals.
+ * Omit (default) to skip CI computation.
+ */
+ readonly alpha?: number;
+}
+
+/** Options for {@link ccf}. */
+export interface CCFOptions {
+ /** Maximum lag (default: `min(floor(10·log₁₀(n)), n−1)`). */
+ readonly nlags?: number;
+ /**
+ * Significance level for confidence intervals.
+ * Omit (default) to skip CI computation.
+ */
+ readonly alpha?: number;
+ /** When `true`, return only non-negative lags (default: `false`). */
+ readonly positiveOnly?: boolean;
+}
+
+/** Options for {@link ljungBox} and {@link boxPierce}. */
+export interface PortmanteauOptions {
+ /**
+ * Specific lag values to test, or a single maximum lag `h` (implying lags
+ * `1, 2, …, h`). Default: `min(floor(10·log₁₀(n)), n−1)`.
+ */
+ readonly lags?: number | readonly number[];
+ /**
+ * Number of estimated AR/MA parameters already fit to the series (default: `0`).
+ * The chi-squared degrees of freedom for lag `h` become `h − modelDf`.
+ */
+ readonly modelDf?: number;
+}
+
+// ─── internal type alias ──────────────────────────────────────────────────────
+
+/** A numeric array or {@link Series} accepted by every public function. */
+type NumericInput = readonly number[] | Series;
+
+// ─── math helpers ─────────────────────────────────────────────────────────────
+
+/** Lanczos approximation of log-Γ(z). */
+function logGamma(z: number): number {
+ if (z < 0.5) {
+ return Math.log(Math.PI / Math.sin(Math.PI * z)) - logGamma(1.0 - z);
+ }
+ const g = 7;
+ const c = [
+ 0.99999999999980993, 676.5203681218851, -1259.1392167224028, 771.32342877765313,
+ -176.61502916214059, 12.507343278686905, -0.13857109526572012, 9.9843695780195716e-6,
+ 1.5056327351493116e-7,
+ ];
+ let x = c[0] ?? 0;
+ const zz = z - 1;
+ for (let i = 1; i < g + 2; i++) {
+ x += (c[i] ?? 0) / (zz + i);
+ }
+ const t = zz + g + 0.5;
+ return 0.5 * Math.log(2 * Math.PI) + (zz + 0.5) * Math.log(t) - t + Math.log(x);
+}
+
+/** Regularised lower incomplete Γ via series expansion (x < a+1). */
+function regIncGammaSeries(a: number, x: number, lnGa: number): number {
+ let sum = 1 / a;
+ let term = 1 / a;
+ for (let n = 1; n <= 200; n++) {
+ term *= x / (a + n);
+ sum += term;
+ if (Math.abs(term) < Math.abs(sum) * 1e-10) {
+ break;
+ }
+ }
+ return Math.exp(-x + a * Math.log(x) - lnGa) * sum;
+}
+
+/** Regularised lower incomplete Γ via continued-fraction expansion (x ≥ a+1). */
+function regIncGammaCF(a: number, x: number, lnGa: number): number {
+ const eps = 1e-30;
+ let f = eps;
+ let c = f;
+ let d = 1 / (x - a + 1 + eps);
+ d = 1 / d;
+ f = c * d;
+ for (let i = 1; i <= 200; i++) {
+ const an = -i * (i - a);
+ const bn = x - a + 2 * i + 1;
+ d = an * d + bn;
+ c = bn + an / c;
+ if (Math.abs(c) < eps) {
+ c = eps;
+ }
+ d = 1 / (Math.abs(d) < eps ? eps : d);
+ const del = c * d;
+ f *= del;
+ if (Math.abs(del - 1) < 1e-10) {
+ break;
+ }
+ }
+ return 1 - Math.exp(-x + a * Math.log(x) - lnGa) * f;
+}
+
+/** Regularised lower incomplete Γ: P(a, x). */
+function regIncGamma(a: number, x: number): number {
+ if (x < 0) {
+ return 0;
+ }
+ const lnGa = logGamma(a);
+ if (x < a + 1) {
+ return regIncGammaSeries(a, x, lnGa);
+ }
+ return regIncGammaCF(a, x, lnGa);
+}
+
+/** χ² survival function: P(χ²_df > x). */
+function chi2sf(x: number, df: number): number {
+ if (x <= 0) {
+ return 1;
+ }
+ return 1 - regIncGamma(df / 2, x / 2);
+}
+
+/** Inverse standard-normal CDF (Peter Acklam's rational approximation). */
+function normalPpf(p: number): number {
+ if (p <= 0) {
+ return Number.NEGATIVE_INFINITY;
+ }
+ if (p >= 1) {
+ return Number.POSITIVE_INFINITY;
+ }
+ if (p === 0.5) {
+ return 0;
+ }
+ const a0 = -3.969683028665376e1;
+ const a1 = 2.209460984245205e2;
+ const a2 = -2.759285104469687e2;
+ const a3 = 1.38357751867269e2;
+ const a4 = -3.066479806614716e1;
+ const a5 = 2.506628277459239;
+ const b0 = -5.447609879822406e1;
+ const b1 = 1.615858368580409e2;
+ const b2 = -1.556989798598866e2;
+ const b3 = 6.680131188771972e1;
+ const b4 = -1.328068155288572e1;
+ const c0 = -7.784894002430293e-3;
+ const c1 = -3.223964580411365e-1;
+ const c2 = -2.400758277161838;
+ const c3 = -2.549732539343734;
+ const c4 = 4.374664141464968;
+ const c5 = 2.938163982698783;
+ const d0 = 7.784695709041462e-3;
+ const d1 = 3.224671290700398e-1;
+ const d2 = 2.445134137142996;
+ const d3 = 3.754408661907416;
+ const pLow = 0.02425;
+ const pHigh = 1 - pLow;
+ if (p < pLow) {
+ const q = Math.sqrt(-2 * Math.log(p));
+ return (
+ (((((c0 * q + c1) * q + c2) * q + c3) * q + c4) * q + c5) /
+ ((((d0 * q + d1) * q + d2) * q + d3) * q + 1)
+ );
+ }
+ if (p <= pHigh) {
+ const q = p - 0.5;
+ const r = q * q;
+ return (
+ ((((((a0 * r + a1) * r + a2) * r + a3) * r + a4) * r + a5) * q) /
+ (((((b0 * r + b1) * r + b2) * r + b3) * r + b4) * r + 1)
+ );
+ }
+ const q = Math.sqrt(-2 * Math.log(1 - p));
+ return -(
+ (((((c0 * q + c1) * q + c2) * q + c3) * q + c4) * q + c5) /
+ ((((d0 * q + d1) * q + d2) * q + d3) * q + 1)
+ );
+}
+
+// ─── tuple helpers ────────────────────────────────────────────────────────────
+
+/** Build a `[lower, upper]` confidence bound tuple without a cast. */
+function bound(center: number, margin: number): [number, number] {
+ return [center - margin, center + margin];
+}
+
+// ─── array extraction ─────────────────────────────────────────────────────────
+
+/** Extract a plain `number[]`, dropping NaN and non-numeric values. */
+function toNumbers(input: NumericInput): number[] {
+ if (input instanceof Series) {
+ const out: number[] = [];
+ for (const val of input.values) {
+ if (typeof val === "number" && !Number.isNaN(val)) {
+ out.push(val);
+ }
+ }
+ return out;
+ }
+ return (input as readonly number[]).filter((v) => typeof v === "number" && !Number.isNaN(v));
+}
+
+// ─── autocovariance ────────────────────────────────────────────────────────────
+
+/**
+ * Biased sample autocovariance γ̂(0), γ̂(1), …, γ̂(nlags).
+ * Denominator is n (consistent with pandas / statsmodels default).
+ */
+function autocovariance(x: readonly number[], mean: number, nlags: number): number[] {
+ const n = x.length;
+ const cov: number[] = [];
+ for (let k = 0; k <= nlags; k++) {
+ let s = 0;
+ for (let t = 0; t < n - k; t++) {
+ s += ((x[t] ?? 0) - mean) * ((x[t + k] ?? 0) - mean);
+ }
+ cov.push(s / n);
+ }
+ return cov;
+}
+
+// ─── ACF CI helper ────────────────────────────────────────────────────────────
+
+/** Bartlett confidence intervals for ACF coefficients. */
+function buildAcfCI(acfValues: readonly number[], n: number, alpha: number): [number, number][] {
+ const z = normalPpf(1 - alpha / 2);
+ const ci: [number, number][] = [];
+ let sumSq = 0;
+ for (let k = 0; k < acfValues.length; k++) {
+ if (k === 0) {
+ ci.push([1, 1]);
+ } else {
+ const se = Math.sqrt((1 + 2 * sumSq) / n);
+ const r = acfValues[k] ?? 0;
+ ci.push(bound(r, z * se));
+ sumSq += r * r;
+ }
+ }
+ return ci;
+}
+
+// ─── Levinson-Durbin helpers ──────────────────────────────────────────────────
+
+/** Single Levinson-Durbin recursion step: returns [φ_kk, updated φ array]. */
+function ldStep(acfVals: readonly number[], phi: readonly number[], k: number): [number, number[]] {
+ let num = acfVals[k] ?? 0;
+ let den = 1;
+ for (let j = 1; j < k; j++) {
+ num -= (phi[j - 1] ?? 0) * (acfVals[k - j] ?? 0);
+ den -= (phi[j - 1] ?? 0) * (acfVals[j] ?? 0);
+ }
+ const phiKK = den === 0 ? 0 : num / den;
+ const newPhi: number[] = [];
+ for (let j = 1; j < k; j++) {
+ newPhi.push((phi[j - 1] ?? 0) - phiKK * (phi[k - j - 1] ?? 0));
+ }
+ newPhi.push(phiKK);
+ return [phiKK, newPhi];
+}
+
+/** Levinson-Durbin recursion returning PACF[0..nlags]. */
+function levinsonDurbin(acfVals: readonly number[], nlags: number): number[] {
+ const result: number[] = [1];
+ if (nlags === 0) {
+ return result;
+ }
+ let phi: number[] = [];
+ for (let k = 1; k <= nlags; k++) {
+ const [phiKK, newPhi] = ldStep(acfVals, phi, k);
+ result.push(phiKK);
+ phi = newPhi;
+ }
+ return result;
+}
+
+// ─── CCF helpers ──────────────────────────────────────────────────────────────
+
+/** Cross-covariance at lag k (biased estimator, denominator = n). */
+function ccfLag(
+ xArr: readonly number[],
+ yArr: readonly number[],
+ n: number,
+ xMean: number,
+ yMean: number,
+ k: number,
+): number {
+ const start = k >= 0 ? 0 : -k;
+ const end = k >= 0 ? n - k : n;
+ let s = 0;
+ for (let t = start; t < end; t++) {
+ s += ((xArr[t] ?? 0) - xMean) * ((yArr[t + k] ?? 0) - yMean);
+ }
+ return s / n;
+}
+
+// ─── portmanteau helpers ──────────────────────────────────────────────────────
+
+/** Resolve the `lags` option to a sorted list. */
+function resolveLags(opt: number | readonly number[] | undefined, maxLag: number): number[] {
+ if (opt === undefined) {
+ return [maxLag];
+ }
+ if (typeof opt === "number") {
+ return Array.from({ length: opt }, (_, i) => i + 1);
+ }
+ return (opt as readonly number[]).slice().sort((a, b) => a - b);
+}
+
+/** Compute Ljung-Box or Box-Pierce Q at lag h. */
+function portmanteauQ(acfVals: readonly number[], n: number, h: number, ljung: boolean): number {
+ let q = 0;
+ for (let k = 1; k <= h; k++) {
+ const r = acfVals[k] ?? 0;
+ q += ljung ? (r * r) / (n - k) : r * r;
+ }
+ return ljung ? n * (n + 2) * q : n * q;
+}
+
+// ─── public API ───────────────────────────────────────────────────────────────
+
+/**
+ * Pandas-style single-lag autocorrelation.
+ *
+ * Equivalent to `pd.Series.autocorr(lag)` — computes the Pearson correlation
+ * between `x[0..n−lag−1]` and `x[lag..n−1]`.
+ *
+ * @param x Input series.
+ * @param lag Lag (default: `1`).
+ * @returns Pearson correlation in `[−1, 1]`, or `NaN` if series is too short.
+ */
+export function autocorr(x: NumericInput, lag = 1): number {
+ const arr = toNumbers(x);
+ const n = arr.length;
+ if (n <= lag + 1) {
+ return Number.NaN;
+ }
+ const x1 = arr.slice(0, n - lag);
+ const x2 = arr.slice(lag);
+ const m1 = x1.reduce((s, v) => s + v, 0) / x1.length;
+ const m2 = x2.reduce((s, v) => s + v, 0) / x2.length;
+ let num = 0;
+ let ss1 = 0;
+ let ss2 = 0;
+ for (let i = 0; i < x1.length; i++) {
+ const d1 = (x1[i] ?? 0) - m1;
+ const d2 = (x2[i] ?? 0) - m2;
+ num += d1 * d2;
+ ss1 += d1 * d1;
+ ss2 += d2 * d2;
+ }
+ const denom = Math.sqrt(ss1 * ss2);
+ return denom === 0 ? Number.NaN : num / denom;
+}
+
+/**
+ * Full Autocorrelation Function (ACF) with optional Bartlett confidence
+ * intervals.
+ *
+ * Mirrors `statsmodels.tsa.stattools.acf` (biased estimator, lag 0 = 1).
+ *
+ * @param x Input time series.
+ * @param options See {@link ACFOptions}.
+ * @returns ACF coefficients for lags 0…nlags, plus optional CI.
+ */
+export function acf(x: NumericInput, options: ACFOptions = {}): ACFResult {
+ const arr = toNumbers(x);
+ const n = arr.length;
+ const defaultNlags = Math.min(Math.floor(10 * Math.log10(n)), n - 1);
+ const nlags = Math.max(0, Math.min(options.nlags ?? defaultNlags, n - 1));
+ const mean = arr.reduce((s, v) => s + v, 0) / n;
+ const cov = autocovariance(arr, mean, nlags);
+ const gamma0 = cov[0] ?? 1;
+ const acfValues: number[] = cov.map((c) => (gamma0 === 0 ? 0 : c / gamma0));
+ const lags = Array.from({ length: nlags + 1 }, (_, i) => i);
+ const confint = options.alpha !== undefined ? buildAcfCI(acfValues, n, options.alpha) : undefined;
+ return { acf: acfValues, confint, lags };
+}
+
+/**
+ * Partial Autocorrelation Function (PACF) via the Levinson-Durbin recursion.
+ *
+ * Mirrors `statsmodels.tsa.stattools.pacf` (method `"yw"`).
+ *
+ * @param x Input time series.
+ * @param options See {@link PACFOptions}.
+ * @returns PACF coefficients for lags 0…nlags, plus optional CI.
+ */
+export function pacf(x: NumericInput, options: PACFOptions = {}): PACFResult {
+ const arr = toNumbers(x);
+ const n = arr.length;
+ const defaultNlags = Math.min(Math.floor(10 * Math.log10(n)), Math.floor(n / 2) - 1);
+ const maxAllowed = Math.max(0, Math.floor(n / 2) - 1);
+ const nlags = Math.max(0, Math.min(options.nlags ?? defaultNlags, maxAllowed));
+ const mean = arr.reduce((s, v) => s + v, 0) / n;
+ const cov = autocovariance(arr, mean, nlags);
+ const gamma0 = cov[0] ?? 1;
+ const acfVals: number[] = cov.map((c) => (gamma0 === 0 ? 0 : c / gamma0));
+ const pacfValues = levinsonDurbin(acfVals, nlags);
+ const lags = Array.from({ length: nlags + 1 }, (_, i) => i);
+ let confint: [number, number][] | undefined;
+ if (options.alpha !== undefined) {
+ const z = normalPpf(1 - options.alpha / 2);
+ const se = 1 / Math.sqrt(n);
+ confint = pacfValues.map((r) => bound(r, z * se));
+ }
+ return { pacf: pacfValues, confint, lags };
+}
+
+/**
+ * Cross-Correlation Function (CCF) between two series.
+ *
+ * CCF(k) is the normalized cross-covariance at lag k:
+ * `CCF(k) = C_xy(k) / (σ_x · σ_y)` where `C_xy(k)` uses denominator `n`.
+ *
+ * @param x First time series.
+ * @param y Second time series (must have the same length as `x`).
+ * @param options See {@link CCFOptions}.
+ * @returns CCF coefficients for lags −nlags…+nlags (or 0…nlags).
+ */
+export function ccf(x: NumericInput, y: NumericInput, options: CCFOptions = {}): ACFResult {
+ const xArr = toNumbers(x);
+ const yArr = toNumbers(y);
+ const n = Math.min(xArr.length, yArr.length);
+ const defaultNlags = Math.min(Math.floor(10 * Math.log10(n)), n - 1);
+ const nlags = Math.max(0, Math.min(options.nlags ?? defaultNlags, n - 1));
+ const positiveOnly = options.positiveOnly ?? false;
+ const xSub = xArr.slice(0, n);
+ const ySub = yArr.slice(0, n);
+ const xMean = xSub.reduce((s, v) => s + v, 0) / n;
+ const yMean = ySub.reduce((s, v) => s + v, 0) / n;
+ const xVar = xSub.reduce((s, v) => s + (v - xMean) ** 2, 0) / n;
+ const yVar = ySub.reduce((s, v) => s + (v - yMean) ** 2, 0) / n;
+ const xStd = Math.sqrt(xVar);
+ const yStd = Math.sqrt(yVar);
+ const denom = xStd * yStd;
+ const startLag = positiveOnly ? 0 : -nlags;
+ const lags: number[] = [];
+ const values: number[] = [];
+ for (let k = startLag; k <= nlags; k++) {
+ lags.push(k);
+ const cov = ccfLag(xSub, ySub, n, xMean, yMean, k);
+ values.push(denom === 0 ? 0 : cov / denom);
+ }
+ let confint: [number, number][] | undefined;
+ if (options.alpha !== undefined) {
+ const z = normalPpf(1 - options.alpha / 2);
+ const se = 1 / Math.sqrt(n);
+ confint = values.map((r) => bound(r, z * se));
+ }
+ return { acf: values, confint, lags };
+}
+
+/**
+ * Durbin-Watson statistic for autocorrelation in OLS residuals.
+ *
+ * `DW = Σ(eₜ − eₜ₋₁)² / Σeₜ²`
+ *
+ * | DW | Interpretation |
+ * |------|------------------------------|
+ * | ≈ 0 | Strong positive correlation |
+ * | ≈ 2 | No autocorrelation |
+ * | ≈ 4 | Strong negative correlation |
+ *
+ * @param residuals OLS residual array or Series.
+ * @returns Durbin-Watson statistic in `[0, 4]`.
+ */
+export function durbinWatson(residuals: NumericInput): number {
+ const e = toNumbers(residuals);
+ const n = e.length;
+ if (n < 2) {
+ return Number.NaN;
+ }
+ let diff2 = 0;
+ let ss = (e[0] ?? 0) ** 2;
+ for (let t = 1; t < n; t++) {
+ const d = (e[t] ?? 0) - (e[t - 1] ?? 0);
+ diff2 += d * d;
+ ss += (e[t] ?? 0) ** 2;
+ }
+ return ss === 0 ? 2 : diff2 / ss;
+}
+
+/**
+ * Ljung-Box Q test for serial autocorrelation up to lag `h`.
+ *
+ * `Q_LB(h) = n·(n+2) · Σₖ₌₁ʰ r̂ₖ² / (n−k)`
+ *
+ * H₀: the first `h` autocorrelations are all zero.
+ * Rejection at small p-values indicates the series is not white noise.
+ *
+ * @param x Input time series.
+ * @param options See {@link PortmanteauOptions}.
+ * @returns Test statistic, p-value, and lag for each tested lag.
+ */
+export function ljungBox(x: NumericInput, options: PortmanteauOptions = {}): PortmanteauResult {
+ return portmanteauTest(x, options, true);
+}
+
+/**
+ * Box-Pierce Q test (simplified Ljung-Box).
+ *
+ * `Q_BP(h) = n · Σₖ₌₁ʰ r̂ₖ²`
+ *
+ * @param x Input time series.
+ * @param options See {@link PortmanteauOptions}.
+ * @returns Test statistic, p-value, and lag for each tested lag.
+ */
+export function boxPierce(x: NumericInput, options: PortmanteauOptions = {}): PortmanteauResult {
+ return portmanteauTest(x, options, false);
+}
+
+/** Shared computation for Ljung-Box and Box-Pierce. */
+function portmanteauTest(
+ x: NumericInput,
+ options: PortmanteauOptions,
+ ljung: boolean,
+): PortmanteauResult {
+ const arr = toNumbers(x);
+ const n = arr.length;
+ const modelDf = options.modelDf ?? 0;
+ const defaultMaxLag = Math.min(Math.floor(10 * Math.log10(n)), n - 1);
+ const lagList = resolveLags(options.lags, defaultMaxLag);
+ const hMax = lagList.at(-1) ?? defaultMaxLag;
+ const mean = arr.reduce((s, v) => s + v, 0) / n;
+ const cov = autocovariance(arr, mean, hMax);
+ const gamma0 = cov[0] ?? 1;
+ const acfVals: number[] = cov.map((c) => (gamma0 === 0 ? 0 : c / gamma0));
+ const statistic: number[] = [];
+ const pvalue: number[] = [];
+ for (const h of lagList) {
+ const q = portmanteauQ(acfVals, n, h, ljung);
+ const df = h - modelDf;
+ statistic.push(q);
+ pvalue.push(df > 0 ? chi2sf(q, df) : Number.NaN);
+ }
+ return { statistic, pvalue, lags: lagList };
+}
diff --git a/src/stats/arima.ts b/src/stats/arima.ts
new file mode 100644
index 00000000..7e5b3ab5
--- /dev/null
+++ b/src/stats/arima.ts
@@ -0,0 +1,591 @@
+/**
+ * arima — ARIMA(p, d, q) time-series model.
+ *
+ * Mirrors the `statsmodels.tsa.arima.model.ARIMA` API and pandas convention.
+ * Estimation uses the Hannan-Rissanen two-step method:
+ * 1. Fit a high-order AR(kMax) via Yule-Walker to obtain proxy residuals.
+ * 2. OLS on the differenced series using AR(p) lags + MA(q) proxy residuals.
+ *
+ * kMax = min(max(p + q + 5, 3), floor(n / 5)).
+ *
+ * Forecast confidence intervals use ψ-weight recursion. For integrated models
+ * (d > 0) the ψ-weights are accumulated d times (convolution with integration
+ * filter) before computing the forecast MSE.
+ *
+ * @example
+ * ```ts
+ * import { ARIMAModel } from "tsb";
+ *
+ * const y = [1, 2, 1.5, 2.5, 2, 3, 2.5, 3.5, 3, 4, 3.5, 4.5, 4];
+ * const model = new ARIMAModel({ p: 1, d: 0, q: 1 });
+ * const fit = model.fit(y);
+ * console.log(fit.arCoeffs, fit.maCoeffs, fit.aic);
+ * const fc = model.forecast(5);
+ * console.log(fc.forecast, fc.lower, fc.upper);
+ * ```
+ *
+ * @module
+ */
+
+import type { Series } from "../core/series.ts";
+
+// ─── Public types ──────────────────────────────────────────────────────────────
+
+/** Constructor options for {@link ARIMAModel}. */
+export interface ARIMAOptions {
+ /** AR order (number of autoregressive lags). Default: 1. */
+ readonly p?: number;
+ /** Differencing order. Default: 0. */
+ readonly d?: number;
+ /** MA order (number of moving-average lags). Default: 0. */
+ readonly q?: number;
+}
+
+/** Results returned by {@link ARIMAModel.fit}. */
+export interface ARIMAFitResult {
+ /** AR coefficients φ₁ … φₚ (index 0 = lag 1). */
+ readonly arCoeffs: readonly number[];
+ /** MA coefficients θ₁ … θ_q (index 0 = lag 1). */
+ readonly maCoeffs: readonly number[];
+ /** Intercept / constant term on the differenced scale. */
+ readonly intercept: number;
+ /** In-sample fitted values on the **original** (undifferenced) scale. */
+ readonly fittedValues: readonly number[];
+ /** Residuals on the differenced scale. */
+ readonly residuals: readonly number[];
+ /** Estimated noise variance σ². */
+ readonly sigma2: number;
+ /** Log-likelihood (Gaussian). */
+ readonly logLikelihood: number;
+ /** Akaike Information Criterion. */
+ readonly aic: number;
+ /** Bayesian Information Criterion. */
+ readonly bic: number;
+}
+
+/** Multi-step forecasts with prediction intervals from {@link ARIMAModel.forecast}. */
+export interface ARIMAForecastResult {
+ /** Point forecasts on the original (undifferenced) scale. */
+ readonly forecast: readonly number[];
+ /** Lower bound of the 95 % prediction interval (original scale). */
+ readonly lower: readonly number[];
+ /** Upper bound of the 95 % prediction interval (original scale). */
+ readonly upper: readonly number[];
+ /** Standard errors of the h-step-ahead forecast errors. */
+ readonly stderr: readonly number[];
+}
+
+// ─── Private helpers ──────────────────────────────────────────────────────────
+
+/** Type guard: distinguishes `readonly number[]` from `Series`. */
+function isNumericArray(y: readonly number[] | Series): y is readonly number[] {
+ return Array.isArray(y);
+}
+
+function arrMean(x: readonly number[]): number {
+ if (x.length === 0) {
+ return 0;
+ }
+ let s = 0;
+ for (const v of x) {
+ s += v;
+ }
+ return s / x.length;
+}
+
+/** Apply d-th order differencing; also collects the d initial values needed to
+ * reverse the operation. */
+function difference(y: readonly number[], d: number): { w: number[]; inits: number[] } {
+ const inits: number[] = [];
+ let w: number[] = y.slice();
+ for (let i = 0; i < d; i++) {
+ inits.push(w[0] ?? 0);
+ const dw = new Array(w.length - 1);
+ for (let t = 1; t < w.length; t++) {
+ dw[t - 1] = (w[t] ?? 0) - (w[t - 1] ?? 0);
+ }
+ w = dw;
+ }
+ return { w, inits };
+}
+
+/** Reverse d levels of differencing, using the stored initial values. */
+function undifference(dw: readonly number[], inits: readonly number[], d: number): number[] {
+ let w: number[] = dw.slice();
+ for (let i = d - 1; i >= 0; i--) {
+ const init = inits[i] ?? 0;
+ const un = new Array(w.length + 1);
+ un[0] = init;
+ for (let t = 0; t < w.length; t++) {
+ un[t + 1] = (un[t] ?? 0) + (w[t] ?? 0);
+ }
+ w = un;
+ }
+ return w;
+}
+
+/** Estimate AR(k) coefficients via Yule-Walker / Levinson-Durbin.
+ * Returns { ar: coefficients, sigma2: innovation variance }. */
+function yuleWalkerAR(x: readonly number[], k: number): { ar: readonly number[]; sigma2: number } {
+ if (k === 0) {
+ const mu = arrMean(x);
+ let s = 0;
+ for (const v of x) {
+ s += (v - mu) ** 2;
+ }
+ return { ar: [], sigma2: s / x.length || 1 };
+ }
+ const n = x.length;
+ const mu = arrMean(x);
+
+ // Autocovariances γ(0)…γ(k)
+ const acov = new Array(k + 1).fill(0);
+ for (let h = 0; h <= k; h++) {
+ let s = 0;
+ for (let t = h; t < n; t++) {
+ s += ((x[t] ?? 0) - mu) * ((x[t - h] ?? 0) - mu);
+ }
+ acov[h] = s / n;
+ }
+
+ // Levinson-Durbin recursion
+ let prevA: number[] = [];
+ let P = acov[0] || 1e-15;
+
+ for (let m = 1; m <= k; m++) {
+ // Reflection coefficient
+ let num = acov[m] ?? 0;
+ for (let j = 1; j < m; j++) {
+ num -= (prevA[j - 1] ?? 0) * (acov[m - j] ?? 0);
+ }
+ const km = P > 0 ? num / P : 0;
+
+ // Update AR coefficients
+ const curA = new Array(m);
+ for (let j = 1; j < m; j++) {
+ curA[j - 1] = (prevA[j - 1] ?? 0) - km * (prevA[m - j - 1] ?? 0);
+ }
+ curA[m - 1] = km;
+ P *= 1 - km * km;
+ prevA = curA;
+ }
+
+ return { ar: prevA, sigma2: Math.max(P, 1e-15) };
+}
+
+/** Solve β = (X'X)⁻¹ X'y via Gaussian elimination with partial pivoting.
+ * X is (n × k), y is (n). Returns length-k coefficient vector. */
+function olsSolve(X: readonly (readonly number[])[], y: readonly number[]): number[] {
+ const n = X.length;
+ const k = (X[0] ?? []).length;
+ // Build augmented matrix [X'X | X'y] of size k × (k+1)
+ const A: number[][] = Array.from({ length: k }, () => new Array(k + 1).fill(0));
+ for (let i = 0; i < k; i++) {
+ for (let j = 0; j < k; j++) {
+ let s = 0;
+ for (let t = 0; t < n; t++) {
+ s += ((X[t] ?? [])[i] ?? 0) * ((X[t] ?? [])[j] ?? 0);
+ }
+ (A[i] ?? [])[j] = s;
+ }
+ let s = 0;
+ for (let t = 0; t < n; t++) {
+ s += ((X[t] ?? [])[i] ?? 0) * (y[t] ?? 0);
+ }
+ (A[i] ?? [])[k] = s;
+ }
+ // Forward elimination
+ for (let col = 0; col < k; col++) {
+ // Find pivot
+ let pivotRow = col;
+ let maxAbs = Math.abs((A[col] ?? [])[col] ?? 0);
+ for (let row = col + 1; row < k; row++) {
+ const abs = Math.abs((A[row] ?? [])[col] ?? 0);
+ if (abs > maxAbs) {
+ maxAbs = abs;
+ pivotRow = row;
+ }
+ }
+ // Swap rows
+ if (pivotRow !== col) {
+ const tmp = A[col];
+ A[col] = A[pivotRow] ?? [];
+ A[pivotRow] = tmp ?? [];
+ }
+ const pivotVal = (A[col] ?? [])[col] ?? 0;
+ if (Math.abs(pivotVal) < 1e-14) {
+ continue; // singular/near-singular
+ }
+ for (let row = col + 1; row < k; row++) {
+ const factor = ((A[row] ?? [])[col] ?? 0) / pivotVal;
+ for (let c = col; c <= k; c++) {
+ (A[row] ?? [])[c] = ((A[row] ?? [])[c] ?? 0) - factor * ((A[col] ?? [])[c] ?? 0);
+ }
+ }
+ }
+ // Back substitution
+ const beta = new Array(k).fill(0);
+ for (let i = k - 1; i >= 0; i--) {
+ let val = (A[i] ?? [])[k] ?? 0;
+ for (let j = i + 1; j < k; j++) {
+ val -= ((A[i] ?? [])[j] ?? 0) * (beta[j] ?? 0);
+ }
+ const denom = (A[i] ?? [])[i] ?? 0;
+ beta[i] = Math.abs(denom) > 1e-14 ? val / denom : 0;
+ }
+ return beta;
+}
+
+/** Compute ARMA(p, q) ψ-weights (MA∞ representation) up to lag `h`.
+ * ψ₀ = 1, ψⱼ = Σᵢ₌₁ᵐⁱⁿ⁽ʲ'ᵖ⁾ φᵢ ψⱼ₋ᵢ + θⱼ (θⱼ = 0 for j > q). */
+function psiWeights(ar: readonly number[], ma: readonly number[], h: number): number[] {
+ const psi = new Array(h).fill(0);
+ psi[0] = 1;
+ for (let j = 1; j < h; j++) {
+ let v = j <= ma.length ? (ma[j - 1] ?? 0) : 0;
+ for (let i = 1; i <= Math.min(j, ar.length); i++) {
+ v += (ar[i - 1] ?? 0) * (psi[j - i] ?? 0);
+ }
+ psi[j] = v;
+ }
+ return psi;
+}
+
+/** Accumulate ψ-weights d times for the integrated process (ARIMA). */
+function integrateWeights(psi: readonly number[], d: number): number[] {
+ let w = psi.slice();
+ for (let level = 0; level < d; level++) {
+ const acc = w.slice();
+ for (let j = 1; j < acc.length; j++) {
+ acc[j] = (acc[j - 1] ?? 0) + (w[j] ?? 0);
+ }
+ w = acc;
+ }
+ return w;
+}
+
+// ─── ARIMAModel class ─────────────────────────────────────────────────────────
+
+/**
+ * ARIMA(p, d, q) time-series model.
+ *
+ * Estimation via Hannan-Rissanen two-step; forecast CIs via ψ-weight recursion.
+ */
+export class ARIMAModel {
+ private readonly _p: number;
+ private readonly _d: number;
+ private readonly _q: number;
+
+ // Set after fit()
+ private _ar: readonly number[] = [];
+ private _ma: readonly number[] = [];
+ private _mu = 0;
+ private _sigma2 = 1;
+ private _origY: readonly number[] = [];
+ private _inits: readonly number[] = [];
+ private _diffW: readonly number[] = [];
+ private _residuals: readonly number[] = [];
+ private _fitted = false;
+
+ constructor(opts: ARIMAOptions = {}) {
+ this._p = Math.max(0, Math.floor(opts.p ?? 1));
+ this._d = Math.max(0, Math.floor(opts.d ?? 0));
+ this._q = Math.max(0, Math.floor(opts.q ?? 0));
+ }
+
+ /** AR order. */
+ get p(): number {
+ return this._p;
+ }
+ /** Differencing order. */
+ get d(): number {
+ return this._d;
+ }
+ /** MA order. */
+ get q(): number {
+ return this._q;
+ }
+
+ /**
+ * Fit the model to `y`.
+ * @param y - Observed time series (number array or Series).
+ */
+ fit(y: readonly number[] | Series): ARIMAFitResult {
+ const yArr: readonly number[] = isNumericArray(y) ? y : y.values;
+
+ const n = yArr.length;
+ if (n < this._p + this._d + this._q + 2) {
+ throw new RangeError(`Series too short (${n}) for ARIMA(${this._p},${this._d},${this._q})`);
+ }
+
+ // Store original series for undifferencing
+ this._origY = yArr;
+
+ // Difference
+ const { w, inits } = difference(yArr, this._d);
+ this._inits = inits;
+ this._diffW = w;
+
+ const m = w.length; // = n - d
+
+ const p = this._p;
+ const q = this._q;
+
+ let arCoeffs: readonly number[];
+ let maCoeffs: readonly number[];
+ let intercept: number;
+
+ if (p === 0 && q === 0) {
+ // ARIMA(0,d,0): just differenced mean
+ intercept = arrMean(w);
+ arCoeffs = [];
+ maCoeffs = [];
+ } else {
+ // ── Step 1: Yule-Walker AR(kMax) for proxy residuals ──────────────────
+ 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.at(-1) ?? 0);
+ const dw = new Array(w.length - 1);
+ for (let t = 1; t < w.length; t++) {
+ dw[t - 1] = (w[t] ?? 0) - (w[t - 1] ?? 0);
+ }
+ w = dw;
+ }
+ return lasts;
+}
+
+/** Undifference `fcW` starting from the last observed value at each level. */
+function undifferenceFromLast(
+ fcW: readonly number[],
+ lastInits: readonly number[],
+ d: number,
+): number[] {
+ let w: number[] = fcW.slice();
+ for (let i = d - 1; i >= 0; i--) {
+ const init = lastInits[i] ?? 0;
+ const un: number[] = [];
+ let prev = init;
+ for (const dv of w) {
+ prev += dv;
+ un.push(prev);
+ }
+ w = un;
+ }
+ return w;
+}
+
+// ─── Convenience function ──────────────────────────────────────────────────────
+
+/**
+ * Fit an ARIMA(p, d, q) model and return a fitted {@link ARIMAModel}.
+ *
+ * @example
+ * ```ts
+ * import { fitArima } from "tsb";
+ * const model = fitArima([1, 2, 3, 4, 3, 2, 1, 2, 3, 4], { p: 1, q: 1 });
+ * const fc = model.forecast(3);
+ * ```
+ */
+export function fitArima(y: readonly number[] | Series, opts?: ARIMAOptions): ARIMAModel {
+ const model = new ARIMAModel(opts);
+ model.fit(y);
+ return model;
+}
diff --git a/src/stats/bootstrap.ts b/src/stats/bootstrap.ts
index 3864933a..30ac736e 100644
--- a/src/stats/bootstrap.ts
+++ b/src/stats/bootstrap.ts
@@ -371,7 +371,7 @@ function bootstrapOne(
const sorted = [...bootDist].sort((a, b) => a - b);
// Degenerate: all bootstrap samples yield the same statistic (e.g. n=1)
- if (sorted[0] === sorted[sorted.length - 1]) {
+ if (sorted[0] === sorted.at(-1)) {
const val = sorted[0] ?? theta;
return {
confidenceInterval: { low: val, high: val },
diff --git a/src/stats/dlm.ts b/src/stats/dlm.ts
new file mode 100644
index 00000000..34bd110d
--- /dev/null
+++ b/src/stats/dlm.ts
@@ -0,0 +1,1061 @@
+/**
+ * dlm — Dynamic Linear Model (Bayesian State-Space Model).
+ *
+ * Implements the West & Harrison DLM framework:
+ *
+ * Observation: Y_t = F_t' θ_t + v_t, v_t ~ N(0, V_t)
+ * System: θ_t = G_t θ_{t-1} + w_t, w_t ~ N(0, W_t)
+ * Prior: θ_0 ~ N(m_0, C_0)
+ *
+ * Supports:
+ * - Scalar or multivariate observations
+ * - Time-constant or time-varying system matrices
+ * - Factory builders for common components: local-level, local-linear-trend,
+ * polynomial trend, Fourier seasonality, regression
+ * - Component combination via `DLM.combine()`
+ * - Forward filter (Kalman), backward smoother (RTS), forecasting
+ * - MLE estimation of variance parameters via Nelder-Mead
+ * - Discount-factor model fitting (West & Harrison §6)
+ *
+ * Mirrors the R `dlm` package and Python `statsmodels.tsa.statespace` API
+ * surface while following TypeScript / tsb conventions.
+ *
+ * Exported names:
+ * - {@link DLM} — main class: filter, smooth, forecast, fit
+ * - {@link DLMOptions} — constructor options
+ * - {@link DLMResult} — filter / smoother result
+ * - {@link DLMForecastResult} — forecast result
+ * - {@link buildLocalLevel} — local-level (random-walk) component
+ * - {@link buildLocalLinearTrend} — local-linear-trend component
+ * - {@link buildPolynomial} — n-th order polynomial trend
+ * - {@link buildFourier} — Fourier (harmonic) seasonal component
+ * - {@link buildRegression} — static regression component
+ * - {@link combineDLMs} — block-diagonal combination of components
+ *
+ * @example
+ * ```ts
+ * import { DLM, buildLocalLinearTrend } from "tsb";
+ *
+ * const spec = buildLocalLinearTrend({ sigmaObs: 1, sigmaLevel: 0.5, sigmaSlope: 0.1 });
+ * const dlm = new DLM(spec);
+ * const res = dlm.filter([1, 2, 1.5, 3, 2.5, 4, 3.5]);
+ * console.log(res.filteredMeans); // T × p state vectors
+ * console.log(res.logLikelihood);
+ *
+ * const fc = dlm.forecast(res, 5);
+ * console.log(fc.mean, fc.lower, fc.upper);
+ * ```
+ *
+ * @module
+ */
+
+// ─── Internal matrix utilities ─────────────────────────────────────────────────
+
+type Mat = readonly (readonly number[])[];
+type MutMat = number[][];
+
+function rows(A: Mat): number {
+ return A.length;
+}
+function cols(A: Mat): number {
+ return A[0]?.length ?? 0;
+}
+function zeros(n: number, m: number): MutMat {
+ return Array.from({ length: n }, () => new Array(m).fill(0));
+}
+function eye(n: number): MutMat {
+ return Array.from({ length: n }, (_, i) =>
+ Array.from({ length: n }, (_, j) => (i === j ? 1 : 0)),
+ );
+}
+function mmul(A: Mat, B: Mat): MutMat {
+ const m = rows(A);
+ const k = cols(A);
+ const n = cols(B);
+ const C = zeros(m, n);
+ for (let i = 0; i < m; i++) {
+ const ai = A[i]!;
+ const ci = C[i]!;
+ for (let p = 0; p < k; p++) {
+ const aip = ai[p]!;
+ if (aip === 0) {
+ continue;
+ }
+ const bp = B[p]!;
+ for (let j = 0; j < n; j++) {
+ ci[j]! += aip * bp[j]!;
+ }
+ }
+ }
+ return C;
+}
+function mvmul(A: Mat, x: readonly number[]): number[] {
+ const m = rows(A);
+ const y = new Array(m).fill(0);
+ for (let i = 0; i < m; i++) {
+ const ai = A[i]!;
+ let s = 0;
+ for (let p = 0; p < x.length; p++) {
+ s += ai[p]! * x[p]!;
+ }
+ y[i] = s;
+ }
+ return y;
+}
+function tr(A: Mat): MutMat {
+ const m = rows(A);
+ const n = cols(A);
+ const At = zeros(n, m);
+ for (let i = 0; i < m; i++) {
+ for (let j = 0; j < n; j++) {
+ At[j]![i] = A[i]?.[j]!;
+ }
+ }
+ return At;
+}
+function madd(A: Mat, B: Mat): MutMat {
+ const m = rows(A);
+ const n = cols(A);
+ return Array.from({ length: m }, (_, i) =>
+ Array.from({ length: n }, (_, j) => A[i]?.[j]! + B[i]?.[j]!),
+ );
+}
+function msub(A: Mat, B: Mat): MutMat {
+ const m = rows(A);
+ const n = cols(A);
+ return Array.from({ length: m }, (_, i) =>
+ Array.from({ length: n }, (_, j) => A[i]?.[j]! - B[i]?.[j]!),
+ );
+}
+function mscale(A: Mat, s: number): MutMat {
+ return A.map((row) => row.map((v) => v * s));
+}
+function vdot(a: readonly number[], b: readonly number[]): number {
+ let s = 0;
+ for (let i = 0; i < a.length; i++) {
+ s += a[i]! * b[i]!;
+ }
+ return s;
+}
+function vadd(a: readonly number[], b: readonly number[]): number[] {
+ return a.map((ai, i) => ai + b[i]!);
+}
+function vsub(a: readonly number[], b: readonly number[]): number[] {
+ return a.map((ai, i) => ai - b[i]!);
+}
+function outer(a: readonly number[], b: readonly number[]): MutMat {
+ return Array.from({ length: a.length }, (_, i) =>
+ Array.from({ length: b.length }, (_, j) => a[i]! * b[j]!),
+ );
+}
+
+/** Invert a square matrix via Gauss-Jordan with partial pivoting. Returns null if singular. */
+function matInv(A: Mat): MutMat | null {
+ const n = rows(A);
+ const aug: MutMat = Array.from({ length: n }, (_, i) => [
+ ...A[i]!,
+ ...Array.from({ length: n }, (_, j) => (i === j ? 1 : 0)),
+ ]);
+ for (let col = 0; col < n; col++) {
+ let pivotRow = col;
+ let pivotVal = Math.abs(aug[col]?.[col]!);
+ for (let r = col + 1; r < n; r++) {
+ const v = Math.abs(aug[r]?.[col]!);
+ if (v > pivotVal) {
+ pivotVal = v;
+ pivotRow = r;
+ }
+ }
+ if (pivotVal < 1e-14) {
+ return null;
+ }
+ [aug[col], aug[pivotRow]] = [aug[pivotRow]!, aug[col]!];
+ const scale = 1 / aug[col]?.[col]!;
+ for (let j = 0; j < 2 * n; j++) {
+ aug[col]![j]! *= scale;
+ }
+ for (let r = 0; r < n; r++) {
+ if (r === col) {
+ continue;
+ }
+ const f = aug[r]?.[col]!;
+ if (f === 0) {
+ continue;
+ }
+ for (let j = 0; j < 2 * n; j++) {
+ aug[r]![j]! -= f * aug[col]?.[j]!;
+ }
+ }
+ }
+ return aug.map((row) => row.slice(n));
+}
+
+/** Block-diagonal combination of two matrices. */
+function blockDiag(A: Mat, B: Mat): MutMat {
+ const ra = rows(A);
+ const ca = cols(A);
+ const rb = rows(B);
+ const cb = cols(B);
+ const C = zeros(ra + rb, ca + cb);
+ for (let i = 0; i < ra; i++) {
+ for (let j = 0; j < ca; j++) {
+ C[i]![j] = A[i]?.[j]!;
+ }
+ }
+ for (let i = 0; i < rb; i++) {
+ for (let j = 0; j < cb; j++) {
+ C[ra + i]![ca + j] = B[i]?.[j]!;
+ }
+ }
+ return C;
+}
+
+// ─── Public types ───────────────────────────────────────────────────────────────
+
+/**
+ * Specification for a Dynamic Linear Model.
+ *
+ * All matrices are row-major (array-of-rows).
+ * p = state dimension, q = observation dimension.
+ */
+export interface DLMSpec {
+ /**
+ * System (transition) matrix G: p × p.
+ * If omitted, defaults to the identity.
+ */
+ readonly G: Mat;
+ /**
+ * Observation matrix F: q × p (rows = observation dim, cols = state dim).
+ * For scalar observations q = 1, so F is 1 × p.
+ */
+ readonly F: Mat;
+ /**
+ * State-noise covariance W: p × p.
+ */
+ readonly W: Mat;
+ /**
+ * Observation-noise covariance V: q × q.
+ */
+ readonly V: Mat;
+ /**
+ * Prior mean m_0: length-p vector.
+ * Defaults to zero vector.
+ */
+ readonly m0?: readonly number[];
+ /**
+ * Prior covariance C_0: p × p.
+ * Defaults to 1e6 × I (diffuse prior).
+ */
+ readonly C0?: Mat;
+}
+
+/** Options for the {@link DLM} constructor. */
+export type DLMOptions = DLMSpec;
+
+/** Per-time-step output of the Kalman filter. */
+export interface DLMFilterStep {
+ /** Predicted state mean a_t = G m_{t-1}: length p. */
+ readonly predictedMean: readonly number[];
+ /** Predicted state covariance R_t = G C_{t-1} G' + W: p × p. */
+ readonly predictedCov: Mat;
+ /** Filtered state mean m_t: length p. */
+ readonly filteredMean: readonly number[];
+ /** Filtered state covariance C_t: p × p. */
+ readonly filteredCov: Mat;
+ /** One-step-ahead forecast mean f_t = F a_t: length q. */
+ readonly forecastMean: readonly number[];
+ /** One-step-ahead forecast variance Q_t = F R_t F' + V: q × q. */
+ readonly forecastCov: Mat;
+ /** Kalman gain K_t: p × q. */
+ readonly gain: Mat;
+ /** Innovation (residual) e_t = Y_t − f_t (null if observation missing). */
+ readonly innovation: readonly number[] | null;
+}
+
+/** Result returned by {@link DLM.filter}. */
+export interface DLMResult {
+ /** Filter steps, one per time-point. */
+ readonly steps: readonly DLMFilterStep[];
+ /** Filtered state means: T × p. */
+ readonly filteredMeans: readonly (readonly number[])[];
+ /** Filtered state covariances: T × p × p. */
+ readonly filteredCovs: readonly Mat[];
+ /** Predicted (one-step-ahead) means: T × p. */
+ readonly predictedMeans: readonly (readonly number[])[];
+ /** One-step-ahead forecast means: T × q. */
+ readonly forecastMeans: readonly (readonly number[])[];
+ /** One-step-ahead forecast covariances: T × q × q. */
+ readonly forecastCovs: readonly Mat[];
+ /** Total log-likelihood. */
+ readonly logLikelihood: number;
+ /** Prior mean used. */
+ readonly m0: readonly number[];
+ /** Prior covariance used. */
+ readonly C0: Mat;
+}
+
+/** Result returned by {@link DLM.smooth}. */
+export interface DLMSmootherResult extends DLMResult {
+ /** Smoothed state means: T × p. */
+ readonly smoothedMeans: readonly (readonly number[])[];
+ /** Smoothed state covariances: T × p × p. */
+ readonly smoothedCovs: readonly Mat[];
+}
+
+/** Forecast result returned by {@link DLM.forecast}. */
+export interface DLMForecastResult {
+ /** Forecast means for steps 1 … h: h × q. */
+ readonly mean: readonly (readonly number[])[];
+ /** Forecast covariances for steps 1 … h: h × q × q. */
+ readonly cov: readonly Mat[];
+ /** 2.5th percentile (scalar observation only): length h. */
+ readonly lower: readonly number[];
+ /** 97.5th percentile (scalar observation only): length h. */
+ readonly upper: readonly number[];
+}
+
+// ─── Main DLM class ──────────────────────────────────────────────────────────────
+
+/**
+ * Dynamic Linear Model.
+ *
+ * @example
+ * ```ts
+ * import { DLM } from "tsb";
+ *
+ * const dlm = new DLM({
+ * G: [[1]],
+ * F: [[1]],
+ * W: [[0.1]],
+ * V: [[1]],
+ * });
+ * const res = dlm.filter([1, 2, 3, 4, 5]);
+ * ```
+ */
+export class DLM {
+ private readonly _G: Mat;
+ private readonly _F: Mat;
+ private readonly _W: Mat;
+ private readonly _V: Mat;
+ private readonly _m0: readonly number[];
+ private readonly _C0: Mat;
+ private readonly _p: number; // state dim
+ private readonly _q: number; // obs dim
+
+ constructor(options: DLMOptions) {
+ this._G = options.G;
+ this._F = options.F;
+ this._W = options.W;
+ this._V = options.V;
+ this._p = rows(this._G);
+ this._q = rows(this._F);
+ this._m0 = options.m0 ?? new Array(this._p).fill(0);
+ this._C0 = options.C0 ?? mscale(eye(this._p), 1e6);
+ }
+
+ // ── Factory helpers ──
+
+ /**
+ * Local-level (random-walk plus noise) model.
+ * State: [μ_t], System: μ_t = μ_{t-1} + w_t, Obs: Y_t = μ_t + v_t.
+ */
+ static localLevel(opts: { sigmaObs?: number; sigmaLevel?: number } = {}): DLM {
+ const sv = opts.sigmaObs ?? 1;
+ const sw = opts.sigmaLevel ?? 1;
+ return new DLM({
+ G: [[1]],
+ F: [[1]],
+ W: [[sw * sw]],
+ V: [[sv * sv]],
+ });
+ }
+
+ /**
+ * Local-linear-trend model.
+ * State: [level, slope], G = [[1,1],[0,1]], F = [[1,0]].
+ */
+ static localLinearTrend(
+ opts: { sigmaObs?: number; sigmaLevel?: number; sigmaSlope?: number } = {},
+ ): DLM {
+ const sv = opts.sigmaObs ?? 1;
+ const sw1 = opts.sigmaLevel ?? 1;
+ const sw2 = opts.sigmaSlope ?? 0.1;
+ return new DLM({
+ G: [
+ [1, 1],
+ [0, 1],
+ ],
+ F: [[1, 0]],
+ W: [
+ [sw1 * sw1, 0],
+ [0, sw2 * sw2],
+ ],
+ V: [[sv * sv]],
+ });
+ }
+
+ /**
+ * Polynomial trend of order `order` (1 = random walk, 2 = linear trend, …).
+ */
+ static polynomial(order: number, opts: { sigmaObs?: number; sigmaState?: number } = {}): DLM {
+ const sv = opts.sigmaObs ?? 1;
+ const sw = opts.sigmaState ?? 1;
+ // Jordan block (upper-triangular ones)
+ const G = eye(order);
+ for (let i = 0; i < order - 1; i++) {
+ G[i]![i + 1] = 1;
+ }
+ const F: MutMat = [new Array(order).fill(0)];
+ F[0]![0] = 1;
+ const W = mscale(eye(order), sw * sw);
+ const V: MutMat = [[sv * sv]];
+ return new DLM({ G, F, W, V });
+ }
+
+ /**
+ * Fourier (harmonic) seasonal component with period `period` and `harmonics` pairs.
+ * Each harmonic adds a 2-dimensional block to the state.
+ */
+ static fourier(
+ period: number,
+ harmonics: number,
+ opts: { sigmaObs?: number; sigmaState?: number } = {},
+ ): DLM {
+ const sv = opts.sigmaObs ?? 1;
+ const sw = opts.sigmaState ?? 0.01;
+ const p = 2 * harmonics;
+ // Block-diagonal G: each 2×2 block is a rotation matrix for harmonic j
+ const G = zeros(p, p);
+ const F: MutMat = [new Array(p).fill(0)];
+ for (let j = 1; j <= harmonics; j++) {
+ const omega = (2 * Math.PI * j) / period;
+ const c = Math.cos(omega);
+ const s = Math.sin(omega);
+ const r = 2 * (j - 1);
+ G[r]![r] = c;
+ G[r]![r + 1] = s;
+ G[r + 1]![r] = -s;
+ G[r + 1]![r + 1] = c;
+ F[0]![r] = 1; // cosine coefficient contributes to observation
+ }
+ const W = mscale(eye(p), sw * sw);
+ const V: MutMat = [[sv * sv]];
+ return new DLM({ G, F, W, V });
+ }
+
+ /**
+ * Static regression component with known regressors `X` (T × k).
+ * The state is the k regression coefficients (treated as random walk with
+ * small process noise `sigmaState`).
+ */
+ static regression(
+ X: readonly (readonly number[])[],
+ opts: { sigmaObs?: number; sigmaState?: number } = {},
+ ): DLM {
+ const sv = opts.sigmaObs ?? 1;
+ const sw = opts.sigmaState ?? 0.001;
+ const k = cols(X);
+ // G = identity (coefficients evolve slowly)
+ // F_t = X[t] (time-varying — not supported by DLMSpec directly, but we
+ // handle via per-step override at filter time)
+ const G = eye(k);
+ const F: MutMat = [X[0] ? [...X[0]] : new Array(k).fill(0)];
+ const W = mscale(eye(k), sw * sw);
+ const V: MutMat = [[sv * sv]];
+ return new DLM({ G, F, W, V });
+ }
+
+ // ── Core filter ──
+
+ /**
+ * Run the Kalman filter on observations `y`.
+ * Each element of `y` may be a scalar (for 1-D obs) or a vector; `null` for missing.
+ */
+ filter(
+ y: readonly (number | readonly number[] | null)[],
+ opts: { m0?: readonly number[]; C0?: Mat } = {},
+ ): DLMResult {
+ const m0 = opts.m0 ?? this._m0;
+ const C0 = opts.C0 ?? this._C0;
+ return this._filter(y, m0, C0);
+ }
+
+ private _obs(y: number | readonly number[] | null): readonly number[] | null {
+ if (y === null) {
+ return null;
+ }
+ if (typeof y === "number") {
+ return [y];
+ }
+ return y;
+ }
+
+ private _filter(
+ y: readonly (number | readonly number[] | null)[],
+ m0: readonly number[],
+ C0: Mat,
+ ): DLMResult {
+ const T = y.length;
+ const steps: DLMFilterStep[] = [];
+ let m = m0.slice();
+ let C: Mat = C0;
+ let logLik = 0;
+
+ for (let t = 0; t < T; t++) {
+ const yt = this._obs(y[t]!);
+
+ // Predict
+ const a = mvmul(this._G, m);
+ const R = madd(mmul(mmul(this._G, C), tr(this._G)), this._W);
+
+ // Forecast
+ const f = mvmul(this._F, a);
+ const Q = madd(mmul(mmul(this._F, R), tr(this._F)), this._V);
+
+ let mNew = a;
+ let CNew: Mat = R;
+ let innovation: readonly number[] | null = null;
+ let K: Mat = zeros(this._p, this._q);
+
+ if (yt !== null) {
+ // Kalman gain: K = R F' Q^{-1}
+ const RF = mmul(R, tr(this._F));
+ const Qinv = matInv(Q);
+ if (Qinv !== null) {
+ K = mmul(RF, Qinv);
+ innovation = vsub(yt, f);
+ mNew = vadd(a, mvmul(K, innovation));
+ // Joseph-form for numerical stability: C = (I-KF)R(I-KF)' + K V K'
+ const IKF = msub(eye(this._p), mmul(K, this._F));
+ CNew = madd(mmul(mmul(IKF, R), tr(IKF)), mmul(mmul(K, this._V), tr(K)));
+
+ // Log-likelihood contribution: -0.5*(q*log2π + log|Q| + e'Q^{-1}e)
+ const eQe = vdot(innovation, mvmul(Qinv, innovation));
+ const logDetQ = logDet(Q);
+ logLik += -0.5 * (this._q * Math.log(2 * Math.PI) + logDetQ + eQe);
+ }
+ }
+
+ steps.push({
+ predictedMean: a,
+ predictedCov: R,
+ filteredMean: mNew,
+ filteredCov: CNew,
+ forecastMean: f,
+ forecastCov: Q,
+ gain: K,
+ innovation,
+ });
+ m = mNew.slice();
+ C = CNew;
+ }
+
+ return {
+ steps,
+ filteredMeans: steps.map((s) => s.filteredMean),
+ filteredCovs: steps.map((s) => s.filteredCov),
+ predictedMeans: steps.map((s) => s.predictedMean),
+ forecastMeans: steps.map((s) => s.forecastMean),
+ forecastCovs: steps.map((s) => s.forecastCov),
+ logLikelihood: logLik,
+ m0,
+ C0,
+ };
+ }
+
+ // ── RTS Smoother ──
+
+ /**
+ * Run the Kalman filter followed by the RTS backward smoother.
+ */
+ smooth(
+ y: readonly (number | readonly number[] | null)[],
+ opts: { m0?: readonly number[]; C0?: Mat } = {},
+ ): DLMSmootherResult {
+ const filterResult = this.filter(y, opts);
+ return this._smooth(filterResult);
+ }
+
+ private _smooth(res: DLMResult): DLMSmootherResult {
+ const T = res.steps.length;
+ const sMeans: MutMat = res.filteredMeans.map((m) => m.slice());
+ const sCovs: MutMat[] = res.filteredCovs.map((C) => C.map((r) => r.slice()));
+
+ for (let t = T - 2; t >= 0; t--) {
+ const step = res.steps[t]!;
+ const stepNext = res.steps[t + 1]!;
+ const Rt1 = stepNext.predictedCov;
+ const Rt1inv = matInv(Rt1);
+ if (Rt1inv === null) {
+ continue;
+ }
+ // Smoother gain: J_t = C_t G' R_{t+1}^{-1}
+ const Jt = mmul(mmul(step.filteredCov, tr(this._G)), Rt1inv);
+ // Smoothed mean: s_t = m_t + J_t (s_{t+1} - a_{t+1})
+ const diff = vsub(sMeans[t + 1]!, stepNext.predictedMean);
+ sMeans[t] = vadd(step.filteredMean, mvmul(Jt, diff));
+ // Smoothed covariance: S_t = C_t + J_t (S_{t+1} - R_{t+1}) J_t'
+ const covDiff = msub(sCovs[t + 1]!, stepNext.predictedCov);
+ sCovs[t] = madd(step.filteredCov, mmul(mmul(Jt, covDiff), tr(Jt)));
+ }
+
+ return {
+ ...res,
+ smoothedMeans: sMeans,
+ smoothedCovs: sCovs,
+ };
+ }
+
+ // ── Forecasting ──
+
+ /**
+ * Forecast `h` steps ahead from the end of the filter result.
+ * Returns means and 95% prediction intervals (scalar obs only).
+ */
+ forecast(filterResult: DLMResult, h: number): DLMForecastResult {
+ const T = filterResult.steps.length;
+ const lastStep = filterResult.steps[T - 1]!;
+ let m = lastStep.filteredMean.slice();
+ let C: Mat = lastStep.filteredCov;
+
+ const meanArr: (readonly number[])[] = [];
+ const covArr: Mat[] = [];
+ const lower: number[] = [];
+ const upper: number[] = [];
+ const z975 = 1.959963985;
+
+ for (let h_i = 0; h_i < h; h_i++) {
+ // Predict one step
+ const a = mvmul(this._G, m);
+ const R = madd(mmul(mmul(this._G, C), tr(this._G)), this._W);
+ // Forecast mean and variance
+ const f = mvmul(this._F, a);
+ const Q = madd(mmul(mmul(this._F, R), tr(this._F)), this._V);
+ meanArr.push(f);
+ covArr.push(Q);
+ if (this._q === 1) {
+ const sd = Math.sqrt(Math.max(0, Q[0]?.[0]!));
+ lower.push(f[0]! - z975 * sd);
+ upper.push(f[0]! + z975 * sd);
+ }
+ m = a;
+ C = R;
+ }
+
+ return { mean: meanArr, cov: covArr, lower, upper };
+ }
+
+ // ── MLE estimation ──
+
+ /**
+ * Estimate variance parameters (log V, log W diagonal) by maximum likelihood.
+ * Returns a new DLM with the fitted parameters.
+ *
+ * Uses Nelder-Mead minimisation of negative log-likelihood.
+ */
+ fitMLE(y: readonly (number | readonly number[] | null)[]): DLM {
+ const p = this._p;
+ const q = this._q;
+ // Pack: theta = [log V entries (q*q), log diag(W) entries (p)]
+ // Only diagonal elements of V and W are fitted; off-diagonal kept fixed.
+ const packV = Array.from({ length: q }, (_, i) => Math.log(Math.max(1e-8, this._V[i]?.[i]!)));
+ const packW = Array.from({ length: p }, (_, i) => Math.log(Math.max(1e-8, this._W[i]?.[i]!)));
+ const x0 = [...packV, ...packW];
+
+ const objective = (x: readonly number[]): number => {
+ const V2: MutMat = this._V.map((row) => row.slice());
+ for (let i = 0; i < q; i++) {
+ V2[i]![i] = Math.exp(x[i]!);
+ }
+ const W2: MutMat = this._W.map((row) => row.slice());
+ for (let i = 0; i < p; i++) {
+ W2[i]![i] = Math.exp(x[q + i]!);
+ }
+ const dlm2 = new DLM({ G: this._G, F: this._F, W: W2, V: V2, m0: this._m0, C0: this._C0 });
+ const res = dlm2._filter(y, this._m0, this._C0);
+ return -res.logLikelihood;
+ };
+
+ const xOpt = nelderMead(objective, x0, { maxIter: 500, tol: 1e-6 });
+
+ const Vfit: MutMat = this._V.map((row) => row.slice());
+ for (let i = 0; i < q; i++) {
+ Vfit[i]![i] = Math.exp(xOpt[i]!);
+ }
+ const Wfit: MutMat = this._W.map((row) => row.slice());
+ for (let i = 0; i < p; i++) {
+ Wfit[i]![i] = Math.exp(xOpt[q + i]!);
+ }
+
+ return new DLM({ G: this._G, F: this._F, W: Wfit, V: Vfit, m0: this._m0, C0: this._C0 });
+ }
+
+ /**
+ * Fit with discount factors (West & Harrison §6).
+ * The process covariance is replaced by W_t = ((1-δ)/δ) C_{t-1},
+ * where δ ∈ (0,1] is the discount factor (close to 1 = small evolution).
+ * Returns a new DLM plus the filter result under the discounted model.
+ */
+ filterDiscount(y: readonly (number | readonly number[] | null)[], delta: number): DLMResult {
+ const m0 = this._m0;
+ const C0 = this._C0;
+ const T = y.length;
+ const steps: DLMFilterStep[] = [];
+ let m = m0.slice();
+ let C: Mat = C0;
+ let logLik = 0;
+
+ for (let t = 0; t < T; t++) {
+ const yt = this._obs(y[t]!);
+
+ // Predict with discount: R_t = G C_{t-1} G' / delta
+ const GCGt = mmul(mmul(this._G, C), tr(this._G));
+ const R = mscale(GCGt, 1 / delta);
+ const a = mvmul(this._G, m);
+
+ const f = mvmul(this._F, a);
+ const Q = madd(mmul(mmul(this._F, R), tr(this._F)), this._V);
+
+ let mNew = a;
+ let CNew: Mat = R;
+ let innovation: readonly number[] | null = null;
+ let K: Mat = zeros(this._p, this._q);
+
+ if (yt !== null) {
+ const Qinv = matInv(Q);
+ if (Qinv !== null) {
+ K = mmul(mmul(R, tr(this._F)), Qinv);
+ innovation = vsub(yt, f);
+ mNew = vadd(a, mvmul(K, innovation));
+ const IKF = msub(eye(this._p), mmul(K, this._F));
+ CNew = madd(mmul(mmul(IKF, R), tr(IKF)), mmul(mmul(K, this._V), tr(K)));
+ const eQe = vdot(innovation, mvmul(Qinv, innovation));
+ const logDetQ = logDet(Q);
+ logLik += -0.5 * (this._q * Math.log(2 * Math.PI) + logDetQ + eQe);
+ }
+ }
+
+ steps.push({
+ predictedMean: a,
+ predictedCov: R,
+ filteredMean: mNew,
+ filteredCov: CNew,
+ forecastMean: f,
+ forecastCov: Q,
+ gain: K,
+ innovation,
+ });
+ m = mNew.slice();
+ C = CNew;
+ }
+
+ return {
+ steps,
+ filteredMeans: steps.map((s) => s.filteredMean),
+ filteredCovs: steps.map((s) => s.filteredCov),
+ predictedMeans: steps.map((s) => s.predictedMean),
+ forecastMeans: steps.map((s) => s.forecastMean),
+ forecastCovs: steps.map((s) => s.forecastCov),
+ logLikelihood: logLik,
+ m0,
+ C0,
+ };
+ }
+}
+
+// ─── Factory helpers ────────────────────────────────────────────────────────────
+
+/** Build a local-level DLM (random walk + noise). */
+export function buildLocalLevel(opts: { sigmaObs?: number; sigmaLevel?: number } = {}): DLMSpec {
+ return {
+ G: [[1]],
+ F: [[1]],
+ W: [[(opts.sigmaLevel ?? 1) ** 2]],
+ V: [[(opts.sigmaObs ?? 1) ** 2]],
+ };
+}
+
+/** Build a local-linear-trend DLM. */
+export function buildLocalLinearTrend(
+ opts: {
+ sigmaObs?: number;
+ sigmaLevel?: number;
+ sigmaSlope?: number;
+ } = {},
+): DLMSpec {
+ const sv = opts.sigmaObs ?? 1;
+ const sw1 = opts.sigmaLevel ?? 1;
+ const sw2 = opts.sigmaSlope ?? 0.1;
+ return {
+ G: [
+ [1, 1],
+ [0, 1],
+ ],
+ F: [[1, 0]],
+ W: [
+ [sw1 * sw1, 0],
+ [0, sw2 * sw2],
+ ],
+ V: [[sv * sv]],
+ };
+}
+
+/** Build a polynomial (Jordan-block) trend DLM of a given order. */
+export function buildPolynomial(
+ order: number,
+ opts: { sigmaObs?: number; sigmaState?: number } = {},
+): DLMSpec {
+ const sv = opts.sigmaObs ?? 1;
+ const sw = opts.sigmaState ?? 1;
+ const G = eye(order);
+ for (let i = 0; i < order - 1; i++) {
+ G[i]![i + 1] = 1;
+ }
+ const F: MutMat = [new Array(order).fill(0)];
+ F[0]![0] = 1;
+ return { G, F, W: mscale(eye(order), sw * sw), V: [[sv * sv]] };
+}
+
+/** Build a Fourier seasonal DLM. */
+export function buildFourier(
+ period: number,
+ harmonics: number,
+ opts: { sigmaObs?: number; sigmaState?: number } = {},
+): DLMSpec {
+ const sv = opts.sigmaObs ?? 1;
+ const sw = opts.sigmaState ?? 0.01;
+ const p = 2 * harmonics;
+ const G = zeros(p, p);
+ const F: MutMat = [new Array(p).fill(0)];
+ for (let j = 1; j <= harmonics; j++) {
+ const omega = (2 * Math.PI * j) / period;
+ const c = Math.cos(omega);
+ const s = Math.sin(omega);
+ const r = 2 * (j - 1);
+ G[r]![r] = c;
+ G[r]![r + 1] = s;
+ G[r + 1]![r] = -s;
+ G[r + 1]![r + 1] = c;
+ F[0]![r] = 1;
+ }
+ return { G, F, W: mscale(eye(p), sw * sw), V: [[sv * sv]] };
+}
+
+/** Build a static regression DLM with k predictors (slow-varying coefficients). */
+export function buildRegression(
+ k: number,
+ opts: { sigmaObs?: number; sigmaState?: number } = {},
+): DLMSpec {
+ const sv = opts.sigmaObs ?? 1;
+ const sw = opts.sigmaState ?? 0.001;
+ return {
+ G: eye(k),
+ F: [new Array(k).fill(0)], // placeholder; rows are set externally
+ W: mscale(eye(k), sw * sw),
+ V: [[sv * sv]],
+ };
+}
+
+/**
+ * Combine multiple DLM specs into a single block-diagonal model.
+ * The combined observation matrix is the horizontal concatenation of the
+ * individual F matrices (summing each component's contribution to the scalar
+ * observation). All component V matrices must have the same dimension.
+ */
+export function combineDLMs(...specs: DLMSpec[]): DLMSpec {
+ if (specs.length === 0) {
+ throw new RangeError("combineDLMs requires at least one spec");
+ }
+ const first = specs[0];
+ let G: Mat = first.G;
+ let W: Mat = first.W;
+ let C0: Mat | undefined = first.C0;
+ let m0: number[] = first.m0 ? [...first.m0] : new Array(rows(first.G)).fill(0);
+
+ // Combined F: 1 × (p1+p2+…) — horizontal concat of F rows
+ let Fcombined: number[] = first.F[0] ? [...first.F[0]] : [];
+
+ for (let i = 1; i < specs.length; i++) {
+ const s = specs[i]!;
+ G = blockDiag(G, s.G);
+ W = blockDiag(W, s.W);
+ if (C0 !== undefined && s.C0 !== undefined) {
+ C0 = blockDiag(C0, s.C0);
+ } else {
+ C0 = undefined;
+ }
+ const fi = s.F[0] ? [...s.F[0]] : [];
+ Fcombined = [...Fcombined, ...fi];
+ const pm = s.m0 ? [...s.m0] : new Array(rows(s.G)).fill(0);
+ m0 = [...m0, ...pm];
+ }
+
+ // Use the first component's V (they must agree)
+ const V = specs[0]?.V;
+ const spec: DLMSpec = {
+ G,
+ F: [Fcombined],
+ W,
+ V,
+ m0,
+ ...(C0 !== undefined ? { C0 } : {}),
+ };
+ return spec;
+}
+
+// ─── Log-determinant helper ─────────────────────────────────────────────────────
+
+/** Log-determinant of a positive-definite matrix via Cholesky. Falls back to LU. */
+function logDet(A: Mat): number {
+ const n = rows(A);
+ // Try Cholesky L L' = A
+ const L = zeros(n, n);
+ for (let i = 0; i < n; i++) {
+ for (let j = 0; j <= i; j++) {
+ let s = A[i]?.[j]!;
+ for (let k2 = 0; k2 < j; k2++) {
+ s -= L[i]?.[k2]! * L[j]?.[k2]!;
+ }
+ if (i === j) {
+ if (s < 0) {
+ return _logDetLU(A); // fall back
+ }
+ L[i]![j] = Math.sqrt(s);
+ } else {
+ L[i]![j] = s / (L[j]?.[j] ?? 1);
+ }
+ }
+ }
+ let ld = 0;
+ for (let i = 0; i < n; i++) {
+ ld += 2 * Math.log(Math.abs(L[i]?.[i]!));
+ }
+ return ld;
+}
+
+function _logDetLU(A: Mat): number {
+ const n = rows(A);
+ const U: MutMat = A.map((r) => r.slice());
+ let sign = 1;
+ for (let col = 0; col < n; col++) {
+ let maxVal = Math.abs(U[col]?.[col]!);
+ let maxRow = col;
+ for (let r = col + 1; r < n; r++) {
+ const v = Math.abs(U[r]?.[col]!);
+ if (v > maxVal) {
+ maxVal = v;
+ maxRow = r;
+ }
+ }
+ if (maxRow !== col) {
+ [U[col], U[maxRow]] = [U[maxRow]!, U[col]!];
+ sign *= -1;
+ }
+ const pivot = U[col]?.[col]!;
+ if (Math.abs(pivot) < 1e-14) {
+ return Number.NEGATIVE_INFINITY;
+ }
+ for (let r = col + 1; r < n; r++) {
+ const f = U[r]?.[col]! / pivot;
+ for (let j = col; j < n; j++) {
+ U[r]![j]! -= f * U[col]?.[j]!;
+ }
+ }
+ }
+ let ld = Math.log(Math.abs(sign));
+ for (let i = 0; i < n; i++) {
+ ld += Math.log(Math.abs(U[i]?.[i]!));
+ }
+ return ld;
+}
+
+// ─── Nelder-Mead optimizer ──────────────────────────────────────────────────────
+
+interface NelderMeadOptions {
+ maxIter?: number;
+ tol?: number;
+}
+
+function nelderMead(
+ f: (x: readonly number[]) => number,
+ x0: readonly number[],
+ opts: NelderMeadOptions = {},
+): number[] {
+ const maxIter = opts.maxIter ?? 1000;
+ const tol = opts.tol ?? 1e-8;
+ const n = x0.length;
+ // Build initial simplex
+ const simplex: number[][] = [x0.slice()];
+ for (let i = 0; i < n; i++) {
+ const xi = x0.slice();
+ xi[i]! += xi[i]! !== 0 ? 0.05 * Math.abs(xi[i]!) : 0.00025;
+ simplex.push(xi);
+ }
+ const fvals = simplex.map((x) => f(x));
+
+ const alpha = 1.0;
+ const gamma = 2.0;
+ const rho = 0.5;
+ const sigma = 0.5;
+
+ for (let iter = 0; iter < maxIter; iter++) {
+ // Sort
+ const order = Array.from({ length: n + 1 }, (_, i) => i).sort((a, b) => fvals[a]! - fvals[b]!);
+ const sx = order.map((i) => simplex[i]!);
+ const sf = order.map((i) => fvals[i]!);
+
+ // Convergence check
+ const range = sf[n]! - sf[0]!;
+ if (range < tol) {
+ return sx[0]?.slice();
+ }
+
+ // Centroid of all but worst
+ const xbar = new Array(n).fill(0);
+ for (let i = 0; i < n; i++) {
+ for (let j = 0; j < n; j++) {
+ xbar[j]! += sx[i]?.[j]! / n;
+ }
+ }
+
+ // Reflect
+ const xr = xbar.map((xi, j) => xi + alpha * (xi - sx[n]?.[j]!));
+ const fr = f(xr);
+
+ if (fr < sf[0]!) {
+ // Expand
+ const xe = xbar.map((xi, j) => xi + gamma * (xr[j]! - xi));
+ const fe = f(xe);
+ simplex[order[n]!] = fe < fr ? xe : xr;
+ fvals[order[n]!] = fe < fr ? fe : fr;
+ } else if (fr < sf[n - 1]!) {
+ simplex[order[n]!] = xr;
+ fvals[order[n]!] = fr;
+ } else {
+ // Contract
+ const xc = xbar.map((xi, j) => xi + rho * (sx[n]?.[j]! - xi));
+ const fc = f(xc);
+ if (fc < sf[n]!) {
+ simplex[order[n]!] = xc;
+ fvals[order[n]!] = fc;
+ } else {
+ // Shrink
+ for (let i = 1; i <= n; i++) {
+ for (let j = 0; j < n; j++) {
+ simplex[order[i]!]![j] = sx[0]?.[j]! + sigma * (sx[i]?.[j]! - sx[0]?.[j]!);
+ }
+ fvals[order[i]!] = f(simplex[order[i]!]!);
+ }
+ }
+ }
+ }
+
+ // Return best found
+ let best = 0;
+ for (let i = 1; i < simplex.length; i++) {
+ if (fvals[i]! < fvals[best]!) {
+ best = i;
+ }
+ }
+ return simplex[best]?.slice();
+}
diff --git a/src/stats/ets.ts b/src/stats/ets.ts
new file mode 100644
index 00000000..e5296998
--- /dev/null
+++ b/src/stats/ets.ts
@@ -0,0 +1,1335 @@
+/**
+ * ets — Exponential Smoothing / Holt-Winters ETS models.
+ *
+ * Implements the classical additive-error ETS state-space framework:
+ * - **SimpleExpSmoothing** (SES / ETS(A,N,N)): single level parameter α.
+ * - **Holt** (ETS(A,A,N) / ETS(A,Ad,N)): level α + trend β, optional damping φ.
+ * - **ExponentialSmoothing** (Holt-Winters ETS(·,·,·)): full model with additive
+ * or multiplicative trend and seasonal components.
+ *
+ * Parameter estimation minimises SSE via Nelder-Mead simplex.
+ * Heuristic initialisation follows statsmodels conventions.
+ *
+ * API mirrors `statsmodels.tsa.holtwinters`.
+ *
+ * @example
+ * ```ts
+ * import { ExponentialSmoothing } from "tsb";
+ *
+ * const sales = [17, 21, 23, 18, 22, 26, 19, 24, 27, 20, 25, 28];
+ * const model = new ExponentialSmoothing({ trend: "add", seasonal: "add", seasonalPeriods: 4 });
+ * const fit = model.fit(sales);
+ * console.log(fit.alpha, fit.beta, fit.gamma);
+ * console.log(model.forecast(4)); // 4-step ahead forecasts
+ * ```
+ *
+ * @module
+ */
+
+import type { Series } from "../core/series.ts";
+
+// ─── Public types ──────────────────────────────────────────────────────────────
+
+/** Trend component type: additive, multiplicative, or absent. */
+export type ETSTrend = "add" | "mul" | null;
+
+/** Seasonal component type: additive, multiplicative, or absent. */
+export type ETSSeasonal = "add" | "mul" | null;
+
+/** Initialisation strategy for state variables. */
+export type ETSInit = "heuristic" | "known";
+
+// ── SimpleExpSmoothing ────────────────────────────────────────────────────────
+
+/** Options for {@link SimpleExpSmoothing}. */
+export interface SESOptions {
+ /**
+ * Smoothing level parameter (0 < α < 1).
+ * If omitted the parameter is estimated by minimising SSE.
+ */
+ readonly alpha?: number;
+ /** Initial level value. If omitted, set to `y[0]`. */
+ readonly initialLevel?: number;
+}
+
+/** Result returned by {@link SimpleExpSmoothing.fit}. */
+export interface SESFitResult {
+ /** Estimated smoothing level. */
+ readonly alpha: number;
+ /** Initial level l₀. */
+ readonly initialLevel: number;
+ /** In-sample one-step-ahead fitted values. */
+ readonly fittedValues: readonly number[];
+ /** In-sample residuals e_t = y_t − ŷ_t. */
+ readonly residuals: readonly number[];
+ /** Sum of squared errors. */
+ readonly sse: number;
+ /** Akaike Information Criterion. */
+ readonly aic: number;
+ /** Bayesian Information Criterion. */
+ readonly bic: number;
+ /** Corrected AIC. */
+ readonly aicc: number;
+}
+
+// ── Holt ─────────────────────────────────────────────────────────────────────
+
+/** Options for {@link Holt}. */
+export interface HoltOptions {
+ /** Smoothing level (0 < α < 1). Auto-estimated if omitted. */
+ readonly alpha?: number;
+ /** Smoothing trend (0 < β < 1). Auto-estimated if omitted. */
+ readonly beta?: number;
+ /** Whether to apply a damped trend. Default `false`. */
+ readonly damped?: boolean;
+ /**
+ * Damping coefficient (0 < φ < 1).
+ * Only used when `damped` is `true`. Auto-estimated if omitted.
+ */
+ readonly dampingSlope?: number;
+ /** Initial level l₀. Heuristic if omitted. */
+ readonly initialLevel?: number;
+ /** Initial trend b₀. Heuristic if omitted. */
+ readonly initialTrend?: number;
+}
+
+/** Result returned by {@link Holt.fit}. */
+export interface HoltFitResult {
+ /** Estimated level smoothing parameter. */
+ readonly alpha: number;
+ /** Estimated trend smoothing parameter. */
+ readonly beta: number;
+ /** Damping slope φ (1.0 when not damped). */
+ readonly phi: number;
+ /** Initial level l₀. */
+ readonly initialLevel: number;
+ /** Initial trend b₀. */
+ readonly initialTrend: number;
+ /** In-sample one-step-ahead fitted values. */
+ readonly fittedValues: readonly number[];
+ /** In-sample residuals. */
+ readonly residuals: readonly number[];
+ /** Sum of squared errors. */
+ readonly sse: number;
+ /** Akaike Information Criterion. */
+ readonly aic: number;
+ /** Bayesian Information Criterion. */
+ readonly bic: number;
+ /** Corrected AIC. */
+ readonly aicc: number;
+}
+
+// ── ExponentialSmoothing ─────────────────────────────────────────────────────
+
+/** Options for {@link ExponentialSmoothing}. */
+export interface ExponentialSmoothingOptions {
+ /** Trend component. `"add"` = additive, `"mul"` = multiplicative, `null` = none. */
+ readonly trend?: ETSTrend;
+ /** Whether to use a damped trend. Default `false`. */
+ readonly damped?: boolean;
+ /** Seasonal component. `"add"` = additive, `"mul"` = multiplicative, `null` = none. */
+ readonly seasonal?: ETSSeasonal;
+ /** Number of periods in one seasonal cycle (e.g. 12 for monthly, 4 for quarterly). */
+ readonly seasonalPeriods?: number;
+ /** Smoothing level parameter (0 < α < 1). Auto-estimated if omitted. */
+ readonly alpha?: number;
+ /** Smoothing trend parameter (0 < β < 1). Auto-estimated if omitted. */
+ readonly beta?: number;
+ /** Smoothing seasonal parameter (0 < γ < 1). Auto-estimated if omitted. */
+ readonly gamma?: number;
+ /** Damping slope (0 < φ < 1). Auto-estimated when `damped = true` and omitted. */
+ readonly phi?: number;
+ /** How to initialise the state: `"heuristic"` (default) or `"known"`. */
+ readonly initializationMethod?: ETSInit;
+ /** Known initial level (only used when `initializationMethod = "known"`). */
+ readonly initialLevel?: number;
+ /** Known initial trend (only used when `initializationMethod = "known"`). */
+ readonly initialTrend?: number;
+ /** Known initial seasonal indices (only used when `initializationMethod = "known"`). */
+ readonly initialSeasons?: readonly number[];
+}
+
+/** Result returned by {@link ExponentialSmoothing.fit}. */
+export interface ExponentialSmoothingFitResult {
+ /** Estimated level smoothing parameter. */
+ readonly alpha: number;
+ /** Estimated trend smoothing parameter (`null` when no trend component). */
+ readonly beta: number | null;
+ /** Estimated seasonal smoothing parameter (`null` when no seasonal component). */
+ readonly gamma: number | null;
+ /** Damping slope φ (1.0 when not damped). */
+ readonly phi: number;
+ /** Initial level l₀. */
+ readonly initialLevel: number;
+ /** Initial trend b₀ (`null` when no trend). */
+ readonly initialTrend: number | null;
+ /** Initial seasonal indices s₁…s_m (`null` when no seasonal). */
+ readonly initialSeasons: readonly number[] | null;
+ /** In-sample one-step-ahead fitted values. */
+ readonly fittedValues: readonly number[];
+ /** In-sample residuals. */
+ readonly residuals: readonly number[];
+ /** Sum of squared errors. */
+ readonly sse: number;
+ /** Log-likelihood. */
+ readonly logLikelihood: number;
+ /** Akaike Information Criterion. */
+ readonly aic: number;
+ /** Bayesian Information Criterion. */
+ readonly bic: number;
+ /** Corrected AIC. */
+ readonly aicc: number;
+}
+
+/** Forecast result with prediction intervals. */
+export interface ETSForecastResult {
+ /** Point forecasts h = 1, 2, … steps. */
+ readonly forecast: readonly number[];
+ /** Lower bound of (1 − α_ci) % prediction interval. */
+ readonly lower: readonly number[];
+ /** Upper bound of (1 − α_ci) % prediction interval. */
+ readonly upper: readonly number[];
+ /** Standard errors of h-step-ahead forecast errors. */
+ readonly stderr: readonly number[];
+}
+
+// ─── Private helpers ──────────────────────────────────────────────────────────
+
+/** Extract numeric array from Series or array. */
+function toArr(y: readonly number[] | Series): readonly number[] {
+ if ("dtype" in y) {
+ return y.values;
+ }
+ return y;
+}
+
+/** Clamp value to [lo, hi]. */
+function clamp(v: number, lo: number, hi: number): number {
+ return Math.min(Math.max(v, lo), hi);
+}
+
+/** Clamp all elements of a param vector to their respective bounds. */
+function clampParams(params: readonly number[], bounds: readonly [number, number][]): number[] {
+ return params.map((v, i) => clamp(v, (bounds[i] ?? [0, 1])[0], (bounds[i] ?? [0, 1])[1]));
+}
+
+/**
+ * Nelder-Mead simplex optimiser (unconstrained; bounds enforced by clamping).
+ * Minimises `fn(params)` starting from `x0`.
+ */
+function nelderMead(
+ fn: (params: readonly number[]) => number,
+ x0: readonly number[],
+ bounds: readonly [number, number][],
+ maxIter = 3000,
+): { params: number[]; value: number } {
+ const n = x0.length;
+ if (n === 0) {
+ return { params: [], value: fn([]) };
+ }
+
+ const EPS = 1e-12;
+ const ALPHA_NM = 1.0; // reflection
+ const BETA_NM = 0.5; // contraction
+ const GAMMA_NM = 2.0; // expansion
+ const SIGMA_NM = 0.5; // shrinkage
+
+ const clamp1 = (p: readonly number[]): number[] => clampParams(p, bounds);
+
+ // Build initial simplex
+ const simplex: number[][] = [clamp1(x0)];
+ for (let i = 0; i < n; i++) {
+ const pt = clamp1(x0);
+ const lo = (bounds[i] ?? [0, 1])[0];
+ const hi = (bounds[i] ?? [0, 1])[1];
+ const delta = Math.max((hi - lo) * 0.1, 0.01);
+ pt[i] = clamp((pt[i] ?? 0) + delta, lo + EPS, hi - EPS);
+ simplex.push(pt);
+ }
+
+ const fvals: number[] = simplex.map((p) => fn(p));
+
+ for (let iter = 0; iter < maxIter; iter++) {
+ // Sort indices by fval
+ const ord = Array.from({ length: n + 1 }, (_, i) => i).sort(
+ (a, b) => (fvals[a] ?? 0) - (fvals[b] ?? 0),
+ );
+
+ const fBest = fvals[ord[0] ?? 0] ?? 0;
+ const fWorst = fvals[ord[n] ?? 0] ?? 0;
+ if (fWorst - fBest < EPS) {
+ break;
+ }
+
+ // Centroid of best n points
+ const cent = new Array(n).fill(0);
+ for (let i = 0; i < n; i++) {
+ const row = simplex[ord[i] ?? 0] ?? [];
+ for (let j = 0; j < n; j++) {
+ cent[j] = (cent[j] ?? 0) + (row[j] ?? 0);
+ }
+ }
+ for (let j = 0; j < n; j++) {
+ cent[j] = (cent[j] ?? 0) / n;
+ }
+
+ const worstPt = simplex[ord[n] ?? 0] ?? [];
+
+ // Reflection
+ const xr = clamp1(cent.map((c, j) => (1 + ALPHA_NM) * c - ALPHA_NM * (worstPt[j] ?? 0)));
+ const fr = fn(xr);
+
+ const fSecondWorst = fvals[ord[n - 1] ?? 0] ?? 0;
+
+ if (fr < fBest) {
+ // Expansion
+ const xe = clamp1(cent.map((c, j) => (1 + GAMMA_NM) * c - GAMMA_NM * (worstPt[j] ?? 0)));
+ const fe = fn(xe);
+ if (fe < fr) {
+ simplex[ord[n] ?? 0] = xe;
+ fvals[ord[n] ?? 0] = fe;
+ } else {
+ simplex[ord[n] ?? 0] = xr;
+ fvals[ord[n] ?? 0] = fr;
+ }
+ } else if (fr < fSecondWorst) {
+ simplex[ord[n] ?? 0] = xr;
+ fvals[ord[n] ?? 0] = fr;
+ } else {
+ // Contraction
+ const inside = fr >= fWorst;
+ const src = inside ? worstPt : xr;
+ const xc = clamp1(cent.map((c, j) => BETA_NM * c + (1 - BETA_NM) * (src[j] ?? 0)));
+ const fc = fn(xc);
+ const compareVal = inside ? fWorst : fr;
+ if (fc < compareVal) {
+ simplex[ord[n] ?? 0] = xc;
+ fvals[ord[n] ?? 0] = fc;
+ } else {
+ // Shrink
+ const bestPt = simplex[ord[0] ?? 0] ?? [];
+ for (let i = 1; i <= n; i++) {
+ const row = simplex[ord[i] ?? 0] ?? [];
+ const newRow = clamp1(row.map((v, j) => SIGMA_NM * (v + (bestPt[j] ?? 0))));
+ simplex[ord[i] ?? 0] = newRow;
+ fvals[ord[i] ?? 0] = fn(newRow);
+ }
+ }
+ }
+ }
+
+ // Return best
+ let bestIdx = 0;
+ for (let i = 1; i <= n; i++) {
+ if ((fvals[i] ?? Number.POSITIVE_INFINITY) < (fvals[bestIdx] ?? Number.POSITIVE_INFINITY)) {
+ bestIdx = i;
+ }
+ }
+ return { params: simplex[bestIdx] ?? [], value: fvals[bestIdx] ?? Number.POSITIVE_INFINITY };
+}
+
+/** Compute AIC, BIC, AICc from SSE, n, k. */
+function infoGaussian(
+ sse: number,
+ n: number,
+ k: number,
+): { logLikelihood: number; aic: number; bic: number; aicc: number } {
+ const sigma2 = Math.max(sse / n, 1e-15);
+ const logL = -0.5 * n * (Math.log(2 * Math.PI * sigma2) + 1);
+ const aic = -2 * logL + 2 * k;
+ const bic = -2 * logL + k * Math.log(n);
+ const denom = n - k - 1;
+ const aicc = denom > 0 ? aic + (2 * k * (k + 1)) / denom : aic;
+ return { logLikelihood: logL, aic, bic, aicc };
+}
+
+// ─── SES internals ────────────────────────────────────────────────────────────
+
+/**
+ * Run one SES pass. Returns { fitted, residuals, sse }.
+ * l0 = initial level.
+ */
+function sesPass(
+ y: readonly number[],
+ alpha: number,
+ l0: number,
+): { fitted: number[]; residuals: number[]; sse: number } {
+ const n = y.length;
+ const fitted: number[] = new Array(n);
+ const residuals: number[] = new Array(n);
+ let sse = 0;
+ let l = l0;
+ for (let t = 0; t < n; t++) {
+ fitted[t] = l;
+ const e = (y[t] ?? 0) - l;
+ residuals[t] = e;
+ sse += e * e;
+ l = alpha * (y[t] ?? 0) + (1 - alpha) * l;
+ }
+ return { fitted, residuals, sse };
+}
+
+// ─── Holt internals ───────────────────────────────────────────────────────────
+
+/**
+ * Run one Holt pass.
+ * Returns { fitted, residuals, sse, levels, trends }.
+ */
+function holtPass(
+ y: readonly number[],
+ alpha: number,
+ beta: number,
+ phi: number,
+ l0: number,
+ b0: number,
+): { fitted: number[]; residuals: number[]; sse: number } {
+ const n = y.length;
+ const fitted: number[] = new Array(n);
+ const residuals: number[] = new Array(n);
+ let sse = 0;
+ let l = l0;
+ let b = b0;
+ for (let t = 0; t < n; t++) {
+ const yhat = l + phi * b;
+ fitted[t] = yhat;
+ const yt = y[t] ?? 0;
+ const e = yt - yhat;
+ residuals[t] = e;
+ sse += e * e;
+ const lNew = alpha * yt + (1 - alpha) * (l + phi * b);
+ b = beta * (lNew - l) + (1 - beta) * phi * b;
+ l = lNew;
+ }
+ return { fitted, residuals, sse };
+}
+
+/** Holt h-step forecast (damped or not). */
+function holtForecast(steps: number, l: number, b: number, phi: number): number[] {
+ const out: number[] = [];
+ let phiH = phi; // φ¹
+ let phiSum = phi; // φ + φ² + … + φ^h
+ for (let h = 1; h <= steps; h++) {
+ out.push(l + phiSum * b);
+ phiH *= phi;
+ phiSum += phiH;
+ }
+ return out;
+}
+
+// ─── ETS (Holt-Winters) internals ─────────────────────────────────────────────
+
+interface ETSState {
+ l: number;
+ b: number;
+ s: number[]; // length m circular buffer, s[0] = s_{t-m+1}, … , s[m-1] = s_t
+}
+
+/**
+ * Run one Holt-Winters pass. Returns fitted values, residuals, SSE, and the
+ * final state (l, b, last m seasonal indices).
+ */
+function hwPass(
+ y: readonly number[],
+ alpha: number,
+ beta: number | null,
+ gamma: number | null,
+ phi: number,
+ l0: number,
+ b0: number | null,
+ s0: readonly number[] | null,
+ trend: ETSTrend,
+ seasonal: ETSSeasonal,
+ m: number,
+): {
+ fitted: number[];
+ residuals: number[];
+ sse: number;
+ finalL: number;
+ finalB: number;
+ finalS: number[];
+} {
+ const n = y.length;
+ const fitted: number[] = new Array(n);
+ const residuals: number[] = new Array(n);
+ let sse = 0;
+
+ let l = l0;
+ let b = b0 ?? 0;
+
+ // seasonal buffer: seasonals[t % m] = s_{t+1-m}
+ const seasonals: number[] = s0 ? s0.slice() : new Array(m).fill(0);
+
+ for (let t = 0; t < n; t++) {
+ const yt = y[t] ?? 0;
+ const sIdx = ((t % m) + m) % m; // index into seasonal buffer
+ const st_m = seasonals[sIdx] ?? 0; // s_{t+1-m}
+
+ // One-step-ahead forecast
+ let yhat: number;
+ if (trend === null && seasonal === null) {
+ yhat = l;
+ } else if (trend !== null && seasonal === null) {
+ yhat = l + phi * b;
+ } else if (trend === null && seasonal === "add") {
+ yhat = l + st_m;
+ } else if (trend === null && seasonal === "mul") {
+ yhat = l * st_m;
+ } else if (trend === "add" && seasonal === "add") {
+ yhat = l + phi * b + st_m;
+ } else if (trend === "add" && seasonal === "mul") {
+ yhat = (l + phi * b) * st_m;
+ } else if (trend === "mul" && seasonal === "add") {
+ yhat = l * (phi === 1 ? b : b ** phi) + st_m;
+ } else {
+ // mul trend + mul seasonal
+ yhat = l * (phi === 1 ? b : b ** phi) * st_m;
+ }
+
+ fitted[t] = yhat;
+ const e = yt - yhat;
+ residuals[t] = e;
+ sse += e * e;
+
+ // State update
+ const lPrev = l;
+ const bPrev = b;
+
+ if (trend === null && seasonal === null) {
+ l = alpha * yt + (1 - alpha) * l;
+ } else if (trend !== null && seasonal === null) {
+ l = alpha * yt + (1 - alpha) * (l + phi * b);
+ if (beta !== null) {
+ b = beta * (l - lPrev) + (1 - beta) * phi * bPrev;
+ }
+ } else if (trend === null && seasonal === "add") {
+ l = alpha * (yt - st_m) + (1 - alpha) * l;
+ if (gamma !== null) {
+ seasonals[sIdx] = gamma * (yt - l) + (1 - gamma) * st_m;
+ }
+ } else if (trend === null && seasonal === "mul") {
+ l = alpha * (st_m !== 0 ? yt / st_m : yt) + (1 - alpha) * l;
+ if (gamma !== null) {
+ seasonals[sIdx] = gamma * (l !== 0 ? yt / l : 1) + (1 - gamma) * st_m;
+ }
+ } else if (trend === "add" && seasonal === "add") {
+ l = alpha * (yt - st_m) + (1 - alpha) * (lPrev + phi * bPrev);
+ if (beta !== null) {
+ b = beta * (l - lPrev) + (1 - beta) * phi * bPrev;
+ }
+ if (gamma !== null) {
+ seasonals[sIdx] = gamma * (yt - l) + (1 - gamma) * st_m;
+ }
+ } else if (trend === "add" && seasonal === "mul") {
+ l = alpha * (st_m !== 0 ? yt / st_m : yt) + (1 - alpha) * (lPrev + phi * bPrev);
+ if (beta !== null) {
+ b = beta * (l - lPrev) + (1 - beta) * phi * bPrev;
+ }
+ if (gamma !== null) {
+ seasonals[sIdx] = gamma * (l + phi * b !== 0 ? yt / (l + phi * b) : 1) + (1 - gamma) * st_m;
+ }
+ } else {
+ // multiplicative trend — approximate as additive for stability
+ l = alpha * yt + (1 - alpha) * (lPrev + phi * bPrev);
+ if (beta !== null) {
+ b = beta * (l - lPrev) + (1 - beta) * phi * bPrev;
+ }
+ if (gamma !== null) {
+ seasonals[sIdx] = gamma * (yt - l) + (1 - gamma) * st_m;
+ }
+ }
+ }
+
+ return {
+ fitted,
+ residuals,
+ sse,
+ finalL: l,
+ finalB: b,
+ finalS: seasonals.slice(),
+ };
+}
+
+/**
+ * Generate h-step forecasts from the final ETS state.
+ * `finalS` is the circular buffer of the last m seasonal indices where
+ * `finalS[t % m]` = s_{t+1-m} (same convention as hwPass).
+ */
+function hwForecast(
+ steps: number,
+ l: number,
+ b: number,
+ finalS: readonly number[],
+ phi: number,
+ trend: ETSTrend,
+ seasonal: ETSSeasonal,
+ m: number,
+ n: number, // length of training series (to compute season offsets)
+): number[] {
+ const out: number[] = [];
+ let phiH = phi;
+ let phiSum = phi;
+
+ // At t = n-1 (last training obs), seasonals[t % m] has just been updated.
+ // For forecast step h, the seasonal index corresponds to position (n-1+h) % m in
+ // the buffer (shifted by 1 because sIdx = (t % m) in hwPass at time t = n-1+h).
+ for (let h = 1; h <= steps; h++) {
+ const sIdx = (((n - 1 + h) % m) + m) % m;
+ const sVal = finalS[sIdx] ?? 1;
+
+ let yhat: number;
+ if (trend === null && seasonal === null) {
+ yhat = l;
+ } else if (trend !== null && seasonal === null) {
+ yhat = l + phiSum * b;
+ } else if (trend === null && seasonal === "add") {
+ yhat = l + sVal;
+ } else if (trend === null && seasonal === "mul") {
+ yhat = l * sVal;
+ } else if (trend === "add" && seasonal === "add") {
+ yhat = l + phiSum * b + sVal;
+ } else if (trend === "add" && seasonal === "mul") {
+ yhat = (l + phiSum * b) * sVal;
+ } else if (trend === "mul" && seasonal === "add") {
+ yhat = l * b ** phiSum + sVal;
+ } else {
+ yhat = l * b ** phiSum * sVal;
+ }
+
+ out.push(yhat);
+
+ phiH *= phi;
+ phiSum += phiH;
+ }
+ return out;
+}
+
+// ─── Heuristic initialisation ─────────────────────────────────────────────────
+
+/**
+ * Compute heuristic initial level and trend.
+ * Uses mean of first season + linear regression slope on first two seasons.
+ */
+function heuristicInit(
+ y: readonly number[],
+ m: number,
+ hasTrend: boolean,
+): { l0: number; b0: number } {
+ const n = y.length;
+ if (!hasTrend) {
+ return { l0: y[0] ?? 0, b0: 0 };
+ }
+ // Use first season average as l0, slope between first two seasons as b0
+ const k = Math.min(m, n);
+ let s1 = 0;
+ for (let i = 0; i < k; i++) {
+ s1 += y[i] ?? 0;
+ }
+ const l0 = s1 / k;
+
+ if (n >= 2 * m) {
+ let s2 = 0;
+ for (let i = m; i < 2 * m; i++) {
+ s2 += y[i] ?? 0;
+ }
+ const b0 = (s2 / m - l0) / m;
+ return { l0, b0: b0 || (((y[1] ?? 0) - (y[0] ?? 0)) * m) / m };
+ }
+ // Fallback: slope between y[0] and y[n-1]
+ const b0 = n > 1 ? ((y[n - 1] ?? 0) - (y[0] ?? 0)) / (n - 1) : 0;
+ return { l0, b0 };
+}
+
+/**
+ * Compute heuristic seasonal indices.
+ * Additive: s_j = avg(y_j, y_{j+m}, …) − overall mean
+ * Multiplicative: s_j = avg(y_j, y_{j+m}, …) / overall mean
+ */
+function heuristicSeasons(
+ y: readonly number[],
+ m: number,
+ seasonal: "add" | "mul",
+ l0: number,
+ b0: number,
+): number[] {
+ const n = y.length;
+ const cycles = Math.max(1, Math.floor(n / m));
+
+ // Detrended values for each position in the seasonal cycle
+ const byPos: number[][] = Array.from({ length: m }, () => []);
+ for (let t = 0; t < cycles * m && t < n; t++) {
+ const trend = l0 + b0 * t;
+ const raw = y[t] ?? 0;
+ const pos = t % m;
+ const posArr = byPos[pos];
+ if (posArr !== undefined) {
+ if (seasonal === "add") {
+ posArr.push(raw - trend);
+ } else {
+ posArr.push(trend !== 0 ? raw / trend : 1);
+ }
+ }
+ }
+
+ const rawS = byPos.map((vals) => {
+ if (vals.length === 0) {
+ return seasonal === "add" ? 0 : 1;
+ }
+ return vals.reduce((a, b) => a + b, 0) / vals.length;
+ });
+
+ // Normalise so seasonal indices sum to 0 (additive) or m (multiplicative)
+ if (seasonal === "add") {
+ const mean = rawS.reduce((a, b) => a + b, 0) / m;
+ return rawS.map((v) => v - mean);
+ }
+ const mean = rawS.reduce((a, b) => a + b, 0) / m;
+ return rawS.map((v) => (mean !== 0 ? (v / mean) * 1 : 1));
+}
+
+// ─── SimpleExpSmoothing ──────────────────────────────────────────────────────
+
+/**
+ * Simple Exponential Smoothing (SES / ETS(A,N,N)).
+ *
+ * Produces level-only forecasts; all future forecasts equal the final level.
+ *
+ * @example
+ * ```ts
+ * import { SimpleExpSmoothing } from "tsb";
+ * const model = new SimpleExpSmoothing();
+ * const fit = model.fit([3, 5, 4, 6, 5, 8, 7]);
+ * console.log(fit.alpha, fit.sse);
+ * console.log(model.forecast(3)); // [level, level, level]
+ * ```
+ */
+export class SimpleExpSmoothing {
+ private _fit: SESFitResult | null = null;
+ private _finalLevel = 0;
+
+ /**
+ * Fit the SES model to observed data.
+ * @param y - Time series observations.
+ * @param opts - Optional fixed parameters.
+ */
+ fit(y: readonly number[] | Series, opts?: SESOptions): SESFitResult {
+ const arr = toArr(y);
+ const n = arr.length;
+ if (n < 2) {
+ throw new RangeError("SimpleExpSmoothing requires at least 2 observations");
+ }
+
+ const l0Init = opts?.initialLevel ?? arr[0] ?? 0;
+
+ let alpha: number;
+ let l0: number;
+
+ if (opts?.alpha !== undefined) {
+ alpha = clamp(opts.alpha, 1e-6, 1 - 1e-6);
+ l0 = l0Init;
+ } else {
+ // Optimise α (and optionally l0)
+ const result = nelderMead(
+ ([a, l]: readonly number[]) => sesPass(arr, a ?? 0.3, l ?? arr[0] ?? 0).sse,
+ [0.3, l0Init],
+ [
+ [1e-6, 1 - 1e-6],
+ [
+ (arr[0] ?? 0) - Math.abs(arr[0] ?? 0) * 5 - 1,
+ (arr[0] ?? 0) + Math.abs(arr[0] ?? 0) * 5 + 1,
+ ],
+ ],
+ );
+ alpha = result.params[0] ?? 0.3;
+ l0 = result.params[1] ?? l0Init;
+ }
+
+ const { fitted, residuals, sse } = sesPass(arr, alpha, l0);
+
+ // Final level for forecasting
+ let lFinal = l0;
+ for (const yt of arr) {
+ lFinal = alpha * yt + (1 - alpha) * lFinal;
+ }
+ this._finalLevel = lFinal;
+
+ const k = 2; // alpha + l0
+ const { aic, bic, aicc } = infoGaussian(sse, n, k);
+
+ const result: SESFitResult = {
+ alpha,
+ initialLevel: l0,
+ fittedValues: fitted,
+ residuals,
+ sse,
+ aic,
+ bic,
+ aicc,
+ };
+ this._fit = result;
+ return result;
+ }
+
+ /**
+ * Generate `steps` forecasts from the last fitted state.
+ * All forecasts equal the final level (flat forecast).
+ * Must call {@link fit} first.
+ */
+ forecast(steps: number): number[] {
+ if (this._fit === null) {
+ throw new Error("Call fit() before forecast()");
+ }
+ return new Array(steps).fill(this._finalLevel);
+ }
+}
+
+// ─── Holt ─────────────────────────────────────────────────────────────────────
+
+/**
+ * Holt's linear (double) exponential smoothing (ETS(A,A,N) / ETS(A,Ad,N)).
+ *
+ * Extends SES with a trend component; supports optional damping.
+ *
+ * @example
+ * ```ts
+ * import { Holt } from "tsb";
+ * const model = new Holt();
+ * const fit = model.fit([3, 5, 4, 6, 5, 8, 7, 9, 8, 11]);
+ * console.log(fit.alpha, fit.beta, fit.phi);
+ * console.log(model.forecast(5));
+ * ```
+ */
+export class Holt {
+ private _opts: HoltOptions = {};
+ private _fit: HoltFitResult | null = null;
+ private _finalL = 0;
+ private _finalB = 0;
+ private _phi = 1;
+
+ constructor(opts?: HoltOptions) {
+ this._opts = opts ?? {};
+ }
+
+ /**
+ * Fit Holt's model to observed data.
+ * @param y - Time series observations.
+ * @param opts - Optional parameter overrides (merged with constructor opts).
+ */
+ fit(y: readonly number[] | Series, opts?: HoltOptions): HoltFitResult {
+ const arr = toArr(y);
+ const n = arr.length;
+ if (n < 3) {
+ throw new RangeError("Holt requires at least 3 observations");
+ }
+
+ const merged: HoltOptions = { ...this._opts, ...opts };
+ const damped = merged.damped ?? false;
+
+ const { l0: l0h, b0: b0h } = heuristicInit(arr, n, true);
+
+ const alphaFixed = merged.alpha;
+ const betaFixed = merged.beta;
+ const phiFixed = merged.dampingSlope;
+ const l0Fixed = merged.initialLevel ?? l0h;
+ const b0Fixed = merged.initialTrend ?? b0h;
+
+ // Build optimisation bounds and initial point
+ type Bound = [number, number];
+ const paramNames: string[] = [];
+ const x0: number[] = [];
+ const bounds: Bound[] = [];
+
+ if (alphaFixed === undefined) {
+ paramNames.push("alpha");
+ x0.push(0.3);
+ bounds.push([1e-6, 1 - 1e-6]);
+ }
+ if (betaFixed === undefined) {
+ paramNames.push("beta");
+ x0.push(0.1);
+ bounds.push([1e-6, 1 - 1e-6]);
+ }
+ if (damped && phiFixed === undefined) {
+ paramNames.push("phi");
+ x0.push(0.98);
+ bounds.push([0.8, 1 - 1e-6]);
+ }
+ // Always optimise l0 and b0 if not fixed
+ const optimL0 = merged.initialLevel === undefined;
+ const optimB0 = merged.initialTrend === undefined;
+ if (optimL0) {
+ paramNames.push("l0");
+ x0.push(l0h);
+ const spread = Math.abs(l0h) * 2 + 10;
+ bounds.push([l0h - spread, l0h + spread]);
+ }
+ if (optimB0) {
+ paramNames.push("b0");
+ x0.push(b0h);
+ const spread = Math.abs(b0h) * 5 + 1;
+ bounds.push([b0h - spread, b0h + spread]);
+ }
+
+ let alpha = alphaFixed ?? 0.3;
+ let beta = betaFixed ?? 0.1;
+ let phi = damped ? (phiFixed ?? 0.98) : 1.0;
+ let l0 = l0Fixed;
+ let b0 = b0Fixed;
+
+ if (x0.length > 0) {
+ const result = nelderMead(
+ (params: readonly number[]): number => {
+ let a = alphaFixed ?? params[paramNames.indexOf("alpha")] ?? 0.3;
+ let bta = betaFixed ?? params[paramNames.indexOf("beta")] ?? 0.1;
+ let ph = damped ? (phiFixed ?? params[paramNames.indexOf("phi")] ?? 0.98) : 1.0;
+ const ll0 = optimL0 ? (params[paramNames.indexOf("l0")] ?? l0h) : l0Fixed;
+ const lb0 = optimB0 ? (params[paramNames.indexOf("b0")] ?? b0h) : b0Fixed;
+ a = clamp(a, 1e-6, 1 - 1e-6);
+ bta = clamp(bta, 1e-6, 1 - 1e-6);
+ ph = clamp(ph, 0.8, 1 - 1e-6);
+ return holtPass(arr, a, bta, ph, ll0, lb0).sse;
+ },
+ x0,
+ bounds,
+ );
+ const p = result.params;
+ alpha = alphaFixed ?? p[paramNames.indexOf("alpha")] ?? alpha;
+ beta = betaFixed ?? p[paramNames.indexOf("beta")] ?? beta;
+ phi = damped ? (phiFixed ?? p[paramNames.indexOf("phi")] ?? phi) : 1.0;
+ l0 = optimL0 ? (p[paramNames.indexOf("l0")] ?? l0) : l0Fixed;
+ b0 = optimB0 ? (p[paramNames.indexOf("b0")] ?? b0) : b0Fixed;
+ }
+
+ alpha = clamp(alpha, 1e-6, 1 - 1e-6);
+ beta = clamp(beta, 1e-6, 1 - 1e-6);
+ if (damped) {
+ phi = clamp(phi, 0.8, 1 - 1e-6);
+ }
+
+ const { fitted, residuals, sse } = holtPass(arr, alpha, beta, phi, l0, b0);
+
+ // Final state for forecasting
+ let lF = l0;
+ let bF = b0;
+ for (const yt of arr) {
+ const lNew = alpha * yt + (1 - alpha) * (lF + phi * bF);
+ bF = beta * (lNew - lF) + (1 - beta) * phi * bF;
+ lF = lNew;
+ }
+ this._finalL = lF;
+ this._finalB = bF;
+ this._phi = phi;
+
+ const k = 2 + (damped ? 1 : 0) + 2; // alpha + beta + phi? + l0 + b0
+ const { aic, bic, aicc } = infoGaussian(sse, n, k);
+
+ const result: HoltFitResult = {
+ alpha,
+ beta,
+ phi,
+ initialLevel: l0,
+ initialTrend: b0,
+ fittedValues: fitted,
+ residuals,
+ sse,
+ aic,
+ bic,
+ aicc,
+ };
+ this._fit = result;
+ return result;
+ }
+
+ /**
+ * Generate `steps` forecasts from the last fitted state.
+ * Must call {@link fit} first.
+ */
+ forecast(steps: number): number[] {
+ if (this._fit === null) {
+ throw new Error("Call fit() before forecast()");
+ }
+ return holtForecast(steps, this._finalL, this._finalB, this._phi);
+ }
+}
+
+// ─── ExponentialSmoothing (Holt-Winters) ─────────────────────────────────────
+
+/**
+ * Holt-Winters Exponential Smoothing — full ETS model.
+ *
+ * Supports all combinations of additive / multiplicative trend and seasonal
+ * components with optional damped trend.
+ *
+ * @example
+ * ```ts
+ * import { ExponentialSmoothing } from "tsb";
+ *
+ * // Monthly data with additive seasonal component
+ * const y = [17, 21, 23, 18, 22, 26, 19, 24, 27, 20, 25, 28,
+ * 18, 23, 25, 20, 24, 28, 21, 26, 29, 22, 27, 30];
+ * const model = new ExponentialSmoothing({ trend: "add", seasonal: "add", seasonalPeriods: 12 });
+ * const fit = model.fit(y);
+ * console.log(fit.alpha, fit.beta, fit.gamma, fit.aic);
+ * console.log(model.forecast(12));
+ * ```
+ */
+export class ExponentialSmoothing {
+ private _opts: ExponentialSmoothingOptions;
+ private _fit: ExponentialSmoothingFitResult | null = null;
+ private _finalL = 0;
+ private _finalB = 0;
+ private _finalS: number[] = [];
+ private _phi = 1;
+ private _trend: ETSTrend = null;
+ private _seasonal: ETSSeasonal = null;
+ private _m = 1;
+ private _n = 0;
+ private _sigma2 = 1;
+
+ constructor(opts?: ExponentialSmoothingOptions) {
+ this._opts = opts ?? {};
+ }
+
+ /**
+ * Fit the Holt-Winters model to observed data.
+ * @param y - Time series observations.
+ * @param opts - Optional parameter overrides.
+ */
+ fit(
+ y: readonly number[] | Series,
+ opts?: ExponentialSmoothingOptions,
+ ): ExponentialSmoothingFitResult {
+ const arr = toArr(y);
+ const n = arr.length;
+ if (n < 3) {
+ throw new RangeError("ExponentialSmoothing requires at least 3 observations");
+ }
+
+ const merged: ExponentialSmoothingOptions = { ...this._opts, ...opts };
+ const trend = merged.trend ?? null;
+ const seasonal = merged.seasonal ?? null;
+ const damped = merged.damped ?? false;
+ const m = merged.seasonalPeriods ?? (seasonal !== null ? 2 : 1);
+
+ this._trend = trend;
+ this._seasonal = seasonal;
+ this._m = m;
+ this._n = n;
+
+ if (seasonal !== null && n < 2 * m) {
+ throw new RangeError(
+ `ExponentialSmoothing: need at least 2 full seasonal periods (${2 * m} obs), got ${n}`,
+ );
+ }
+
+ const initMethod = merged.initializationMethod ?? "heuristic";
+
+ // Heuristic initialisation
+ const { l0: l0h, b0: b0h } = heuristicInit(arr, m, trend !== null);
+ const s0h = seasonal !== null ? heuristicSeasons(arr, m, seasonal, l0h, b0h) : null;
+
+ // Determine which params to optimise
+ const alphaFixed = merged.alpha;
+ const betaFixed = merged.beta;
+ const gammaFixed = merged.gamma;
+ const phiFixed = merged.phi;
+ const l0Fixed = initMethod === "known" ? (merged.initialLevel ?? l0h) : undefined;
+ const b0Fixed = initMethod === "known" ? (merged.initialTrend ?? b0h) : undefined;
+ const s0Fixed =
+ initMethod === "known" && merged.initialSeasons !== undefined
+ ? merged.initialSeasons.slice()
+ : undefined;
+
+ type Bound = [number, number];
+ const paramNames: string[] = [];
+ const x0: number[] = [];
+ const bounds: Bound[] = [];
+
+ if (alphaFixed === undefined) {
+ paramNames.push("alpha");
+ x0.push(0.3);
+ bounds.push([1e-6, 1 - 1e-6]);
+ }
+ if (trend !== null && betaFixed === undefined) {
+ paramNames.push("beta");
+ x0.push(0.1);
+ bounds.push([1e-6, 1 - 1e-6]);
+ }
+ if (seasonal !== null && gammaFixed === undefined) {
+ paramNames.push("gamma");
+ x0.push(0.1);
+ bounds.push([1e-6, 1 - 1e-6]);
+ }
+ if (damped && phiFixed === undefined) {
+ paramNames.push("phi");
+ x0.push(0.98);
+ bounds.push([0.8, 1 - 1e-6]);
+ }
+
+ const optimL0 = initMethod !== "known" && merged.initialLevel === undefined;
+ const optimB0 = trend !== null && initMethod !== "known" && merged.initialTrend === undefined;
+ const optimS0 =
+ seasonal !== null && initMethod !== "known" && merged.initialSeasons === undefined;
+
+ if (optimL0) {
+ paramNames.push("l0");
+ x0.push(l0h);
+ const sp = Math.abs(l0h) * 2 + 10;
+ bounds.push([l0h - sp, l0h + sp]);
+ }
+ if (optimB0) {
+ paramNames.push("b0");
+ x0.push(b0h);
+ const sp = Math.abs(b0h) * 5 + 1;
+ bounds.push([b0h - sp, b0h + sp]);
+ }
+ if (optimS0 && s0h !== null) {
+ for (let j = 0; j < m; j++) {
+ paramNames.push(`s0_${j}`);
+ x0.push(s0h[j] ?? 0);
+ const sp = Math.abs(s0h[j] ?? 0) * 5 + 1;
+ bounds.push([(s0h[j] ?? 0) - sp, (s0h[j] ?? 0) + sp]);
+ }
+ }
+
+ let alpha = alphaFixed ?? 0.3;
+ let beta = trend !== null ? (betaFixed ?? 0.1) : null;
+ let gamma = seasonal !== null ? (gammaFixed ?? 0.1) : null;
+ let phi = damped ? (phiFixed ?? 0.98) : 1.0;
+ let l0 = l0Fixed ?? l0h;
+ let b0 = b0Fixed ?? b0h;
+ let s0 = s0Fixed ?? s0h;
+
+ if (x0.length > 0) {
+ const res = nelderMead(
+ (params: readonly number[]): number => {
+ let a = alphaFixed ?? params[paramNames.indexOf("alpha")] ?? 0.3;
+ let bt = trend !== null ? (betaFixed ?? params[paramNames.indexOf("beta")] ?? 0.1) : null;
+ let gm =
+ seasonal !== null ? (gammaFixed ?? params[paramNames.indexOf("gamma")] ?? 0.1) : null;
+ let ph = damped ? (phiFixed ?? params[paramNames.indexOf("phi")] ?? 0.98) : 1.0;
+ a = clamp(a, 1e-6, 1 - 1e-6);
+ if (bt !== null) {
+ bt = clamp(bt, 1e-6, 1 - 1e-6);
+ }
+ if (gm !== null) {
+ gm = clamp(gm, 1e-6, 1 - 1e-6);
+ }
+ if (damped) {
+ ph = clamp(ph, 0.8, 1 - 1e-6);
+ }
+
+ const ll0 = optimL0 ? (params[paramNames.indexOf("l0")] ?? l0h) : (l0Fixed ?? l0h);
+ const lb0 = optimB0 ? (params[paramNames.indexOf("b0")] ?? b0h) : (b0Fixed ?? b0h);
+ let ss0: number[] | null = null;
+ if (optimS0) {
+ ss0 = [];
+ for (let j = 0; j < m; j++) {
+ ss0.push(params[paramNames.indexOf(`s0_${j}`)] ?? s0h?.[j] ?? 0);
+ }
+ } else {
+ ss0 = s0Fixed ?? s0h;
+ }
+
+ return hwPass(arr, a, bt, gm, ph, ll0, lb0, ss0, trend, seasonal, m).sse;
+ },
+ x0,
+ bounds,
+ seasonal !== null ? 5000 : 3000,
+ );
+ const p = res.params;
+ alpha = alphaFixed ?? p[paramNames.indexOf("alpha")] ?? alpha;
+ beta = trend !== null ? (betaFixed ?? p[paramNames.indexOf("beta")] ?? beta ?? 0.1) : null;
+ gamma =
+ seasonal !== null ? (gammaFixed ?? p[paramNames.indexOf("gamma")] ?? gamma ?? 0.1) : null;
+ phi = damped ? (phiFixed ?? p[paramNames.indexOf("phi")] ?? phi) : 1.0;
+ l0 = optimL0 ? (p[paramNames.indexOf("l0")] ?? l0) : (l0Fixed ?? l0);
+ b0 = optimB0 ? (p[paramNames.indexOf("b0")] ?? b0) : (b0Fixed ?? b0);
+ if (optimS0) {
+ s0 = [];
+ for (let j = 0; j < m; j++) {
+ s0.push(p[paramNames.indexOf(`s0_${j}`)] ?? s0h?.[j] ?? 0);
+ }
+ }
+ }
+
+ // Final clamp
+ alpha = clamp(alpha, 1e-6, 1 - 1e-6);
+ if (beta !== null) {
+ beta = clamp(beta, 1e-6, 1 - 1e-6);
+ }
+ if (gamma !== null) {
+ gamma = clamp(gamma, 1e-6, 1 - 1e-6);
+ }
+ if (damped) {
+ phi = clamp(phi, 0.8, 1 - 1e-6);
+ }
+
+ const { fitted, residuals, sse, finalL, finalB, finalS } = hwPass(
+ arr,
+ alpha,
+ beta,
+ gamma,
+ phi,
+ l0,
+ b0,
+ s0,
+ trend,
+ seasonal,
+ m,
+ );
+
+ this._finalL = finalL;
+ this._finalB = finalB;
+ this._finalS = finalS;
+ this._phi = phi;
+ this._sigma2 = Math.max(sse / n, 1e-15);
+
+ // Number of free parameters for information criteria
+ const nParams =
+ 1 + // alpha
+ (trend !== null ? 1 : 0) + // beta
+ (seasonal !== null ? 1 : 0) + // gamma
+ (damped ? 1 : 0) + // phi
+ 1 + // l0
+ (trend !== null ? 1 : 0) + // b0
+ (seasonal !== null ? m : 0); // seasonal indices
+
+ const { logLikelihood, aic, bic, aicc } = infoGaussian(sse, n, nParams);
+
+ const result: ExponentialSmoothingFitResult = {
+ alpha,
+ beta,
+ gamma,
+ phi,
+ initialLevel: l0,
+ initialTrend: trend !== null ? b0 : null,
+ initialSeasons: seasonal !== null ? (s0 ?? null) : null,
+ fittedValues: fitted,
+ residuals,
+ sse,
+ logLikelihood,
+ aic,
+ bic,
+ aicc,
+ };
+ this._fit = result;
+ return result;
+ }
+
+ /**
+ * Generate `steps` point forecasts from the final state.
+ * Must call {@link fit} first.
+ */
+ forecast(steps: number): number[] {
+ if (this._fit === null) {
+ throw new Error("Call fit() before forecast()");
+ }
+ return hwForecast(
+ steps,
+ this._finalL,
+ this._finalB,
+ this._finalS,
+ this._phi,
+ this._trend,
+ this._seasonal,
+ this._m,
+ this._n,
+ );
+ }
+
+ /**
+ * Generate `steps` forecasts with (1 − `alpha_ci`) prediction intervals.
+ * Uses additive-error variance approximation (constant σ² scaled by h).
+ * Must call {@link fit} first.
+ *
+ * @param steps - Number of steps ahead.
+ * @param alpha_ci - Significance level (default 0.05 → 95 % intervals).
+ */
+ forecastWithCI(steps: number, alpha_ci = 0.05): ETSForecastResult {
+ const fc = this.forecast(steps);
+ // Normal quantile for (1 - alpha_ci/2)
+ const z = normalQuantile(1 - alpha_ci / 2);
+ const sigma = Math.sqrt(this._sigma2);
+ const lower: number[] = [];
+ const upper: number[] = [];
+ const stderr: number[] = [];
+ for (let h = 1; h <= steps; h++) {
+ // Variance grows linearly with h for additive-error models
+ const se = sigma * Math.sqrt(h);
+ stderr.push(se);
+ lower.push((fc[h - 1] ?? 0) - z * se);
+ upper.push((fc[h - 1] ?? 0) + z * se);
+ }
+ return { forecast: fc, lower, upper, stderr };
+ }
+}
+
+/**
+ * Rational approximation of the normal quantile function (Abramowitz & Stegun).
+ * @internal
+ */
+function normalQuantile(p: number): number {
+ if (p <= 0) {
+ return Number.NEGATIVE_INFINITY;
+ }
+ if (p >= 1) {
+ return Number.POSITIVE_INFINITY;
+ }
+ const a = [2.515517, 0.802853, 0.010328];
+ const b = [1.432788, 0.189269, 0.001308];
+ const t = Math.sqrt(-2 * Math.log(p < 0.5 ? p : 1 - p));
+ const num = (a[0] ?? 0) + (a[1] ?? 0) * t + (a[2] ?? 0) * t * t;
+ const den = 1 + (b[0] ?? 0) * t + (b[1] ?? 0) * t * t + (b[2] ?? 0) * t * t * t;
+ const x = t - num / den;
+ return p < 0.5 ? -x : x;
+}
+
+// ─── Convenience functions ────────────────────────────────────────────────────
+
+/**
+ * Fit a Simple Exponential Smoothing model and return the result.
+ *
+ * @example
+ * ```ts
+ * import { simpleExpSmoothing } from "tsb";
+ * const { alpha, fittedValues, sse } = simpleExpSmoothing([3, 5, 4, 6, 5]);
+ * ```
+ */
+export function simpleExpSmoothing(
+ y: readonly number[] | Series,
+ opts?: SESOptions,
+): SESFitResult {
+ return new SimpleExpSmoothing().fit(y, opts);
+}
+
+/**
+ * Fit a Holt linear trend model and return the result.
+ *
+ * @example
+ * ```ts
+ * import { holt } from "tsb";
+ * const fit = holt([3, 5, 4, 6, 5, 8, 7, 9]);
+ * console.log(fit.alpha, fit.beta, fit.sse);
+ * ```
+ */
+export function holt(y: readonly number[] | Series, opts?: HoltOptions): HoltFitResult {
+ return new Holt(opts).fit(y);
+}
+
+/**
+ * Fit a full Holt-Winters Exponential Smoothing model and return the result.
+ *
+ * @example
+ * ```ts
+ * import { fitEts } from "tsb";
+ * const y = [17, 21, 23, 18, 22, 26, 19, 24, 27, 20, 25, 28];
+ * const fit = fitEts(y, { trend: "add", seasonal: "add", seasonalPeriods: 4 });
+ * console.log(fit.alpha, fit.beta, fit.gamma, fit.aic);
+ * ```
+ */
+export function fitEts(
+ y: readonly number[] | Series,
+ opts?: ExponentialSmoothingOptions,
+): ExponentialSmoothingFitResult {
+ return new ExponentialSmoothing(opts).fit(y);
+}
diff --git a/src/stats/filters.ts b/src/stats/filters.ts
new file mode 100644
index 00000000..76136a51
--- /dev/null
+++ b/src/stats/filters.ts
@@ -0,0 +1,757 @@
+/**
+ * filters — Digital filter design and application.
+ *
+ * Mirrors `scipy.signal` filter utilities. Implemented from scratch with no
+ * external dependencies.
+ *
+ * Filter design:
+ * - {@link firwin} — FIR filter (windowed-sinc method)
+ * - {@link butter} — Butterworth IIR digital filter
+ *
+ * Frequency response:
+ * - {@link freqz} — frequency response of an FIR/IIR filter
+ * - {@link sosfreqz} — frequency response of SOS filter
+ *
+ * Filter application:
+ * - {@link lfilter} — causal FIR/IIR filter (direct-form II transposed)
+ * - {@link filtfilt} — zero-phase forward-backward filter
+ * - {@link sosfilt} — second-order-sections filter
+ *
+ * @example
+ * ```ts
+ * import { firwin, lfilter, butter, sosfilt } from "tsb";
+ *
+ * // Low-pass FIR with 29 taps, cutoff 0.25 (Nyquist = 0.5)
+ * const b = firwin(29, 0.25);
+ * const y = lfilter(b, [1], signal);
+ *
+ * // Butterworth low-pass, order 4, cutoff 0.2
+ * const { sos } = butter(4, 0.2, "lowpass");
+ * const filtered = sosfilt(sos, signal);
+ * ```
+ *
+ * @module
+ */
+
+import {
+ type Complex,
+ type WindowName,
+ blackmanWindow,
+ cAbs,
+ complex,
+ hammingWindow,
+ hannWindow,
+ kaiserWindow,
+} from "./signal.ts";
+
+// ─── internal helpers ─────────────────────────────────────────────────────────
+
+/** sinc(x) = sin(πx) / (πx), sinc(0) = 1 (normalised). */
+function sinc(x: number): number {
+ if (x === 0) {
+ return 1;
+ }
+ const px = Math.PI * x;
+ return Math.sin(px) / px;
+}
+
+/** Polynomial multiplication (convolution). */
+function polyMul(a: readonly number[], b: readonly number[]): number[] {
+ const out = new Array(a.length + b.length - 1).fill(0);
+ for (let i = 0; i < a.length; i++) {
+ for (let j = 0; j < b.length; j++) {
+ out[i + j] = (out[i + j] ?? 0) + (a[i] ?? 0) * (b[j] ?? 0);
+ }
+ }
+ return out;
+}
+
+// ─── FIR filter design ────────────────────────────────────────────────────────
+
+/** Options for {@link firwin}. */
+export interface FirwinOptions {
+ /**
+ * Window to apply after ideal filter: name string or pre-computed array.
+ * Default `"hamming"`.
+ */
+ window?: WindowName | readonly number[];
+ /** If `true`, design a high-pass filter (default `false` = low-pass). */
+ pass_zero?: boolean;
+ /** Sampling rate used to normalise `cutoff` (default `2` so `cutoff ∈ [0, 1]`). */
+ fs?: number;
+}
+
+/**
+ * Design a low- or high-pass FIR filter using the windowed-sinc method.
+ *
+ * Mirrors `scipy.signal.firwin`.
+ *
+ * @param numtaps - Number of filter coefficients (must be odd for pass_zero=false).
+ * @param cutoff - Cutoff frequency. With default `fs=2`, cutoff is normalised
+ * so `1.0` equals the Nyquist frequency.
+ * @param options - {@link FirwinOptions}.
+ * @returns - FIR filter coefficients `b` (length `numtaps`).
+ *
+ * @example
+ * ```ts
+ * import { firwin, lfilter } from "tsb";
+ * const b = firwin(51, 0.3); // 51-tap 150 Hz LPF (fs=1000)
+ * const y = lfilter(b, [1], signal);
+ * ```
+ */
+export function firwin(
+ numtaps: number,
+ cutoff: number | readonly [number, number],
+ options: FirwinOptions = {},
+): number[] {
+ const fs = options.fs ?? 2;
+ const passZero = options.pass_zero ?? true;
+ const nyq = fs / 2;
+
+ // Normalise cutoff(s) to [0..1] where 1 = Nyquist
+ const 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 && 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;
+ let b1: number;
+ let b2: number;
+ let a1: number;
+ let 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;
+ let 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/hmm.ts b/src/stats/hmm.ts
new file mode 100644
index 00000000..931ae6fb
--- /dev/null
+++ b/src/stats/hmm.ts
@@ -0,0 +1,776 @@
+/**
+ * hmm — Hidden Markov Model (HMM) with discrete and Gaussian emissions.
+ *
+ * Implements:
+ * - **Forward-Backward** algorithm (log-space for numerical stability)
+ * - **Baum-Welch** EM parameter estimation
+ * - **Viterbi** algorithm for most-likely state sequence decoding
+ * - **GaussianHMM**: continuous observations with Gaussian emission distributions
+ * - **MultinomialHMM**: discrete observations with categorical emission distributions
+ *
+ * Mirrors `hmmlearn.hmm.GaussianHMM` and `MultinomialHMM` APIs.
+ *
+ * @example
+ * ```ts
+ * import { GaussianHMM } from "tsb";
+ *
+ * const model = new GaussianHMM({ nComponents: 2, nIter: 100 });
+ * const obs = [0.1, 0.2, 0.15, 2.1, 2.3, 2.0, 0.05, 0.1, 2.5, 2.2];
+ * model.fit(obs);
+ * const states = model.predict(obs);
+ * const logProb = model.score(obs);
+ * ```
+ *
+ * @module
+ */
+
+// ─── Constants ─────────────────────────────────────────────────────────────────
+
+const LOG_ZERO = Number.NEGATIVE_INFINITY;
+
+// ─── Utility helpers ──────────────────────────────────────────────────────────
+
+/** log(exp(a) + exp(b)) with numerical stability. */
+function logSumExp(a: number, b: number): number {
+ if (a === LOG_ZERO) {
+ return b;
+ }
+ if (b === LOG_ZERO) {
+ return a;
+ }
+ const m = Math.max(a, b);
+ return m + Math.log(Math.exp(a - m) + Math.exp(b - m));
+}
+
+/** Stable log-sum-exp over an array. */
+function logSumExpArr(arr: readonly number[]): number {
+ let result = LOG_ZERO;
+ for (const v of arr) {
+ result = logSumExp(result, v);
+ }
+ return result;
+}
+
+/** Safe log (returns LOG_ZERO for non-positive). */
+function safeLog(x: number): number {
+ return x > 0 ? Math.log(x) : LOG_ZERO;
+}
+
+/** Normalise an array in-place so it sums to 1. Returns sum before normalisation. */
+function normalise(arr: number[]): number {
+ const s = arr.reduce((a, b) => a + b, 0);
+ if (s > 0) {
+ for (let i = 0; i < arr.length; i++) {
+ arr[i] = (arr[i] ?? 0) / s;
+ }
+ }
+ return s;
+}
+
+// ─── Public types ─────────────────────────────────────────────────────────────
+
+/** Parameters for a Gaussian HMM. */
+export interface GaussianHMMParams {
+ /** Number of hidden states. */
+ nComponents: number;
+ /** Maximum EM iterations. Default 100. */
+ nIter?: number;
+ /** Convergence tolerance on log-likelihood. Default 1e-4. */
+ tol?: number;
+ /** Random seed for initialisation (unused in deterministic init). */
+ randomState?: number;
+}
+
+/** Fitted Gaussian HMM result. */
+export interface GaussianHMMFit {
+ /** Initial state probabilities (shape: nComponents). */
+ startProb: number[];
+ /** Transition matrix (shape: nComponents × nComponents). Row i → state i, column j → state j. */
+ transmat: number[][];
+ /** Emission means (shape: nComponents). */
+ means: number[];
+ /** Emission variances (shape: nComponents). */
+ covars: number[];
+ /** Log-likelihood of training data at convergence. */
+ logProb: number;
+ /** Number of EM iterations completed. */
+ nIterDone: number;
+}
+
+/** Parameters for a Multinomial HMM. */
+export interface MultinomialHMMParams {
+ /** Number of hidden states. */
+ nComponents: number;
+ /** Number of distinct observation symbols. */
+ nFeatures: number;
+ /** Maximum EM iterations. Default 100. */
+ nIter?: number;
+ /** Convergence tolerance on log-likelihood. Default 1e-4. */
+ tol?: number;
+}
+
+/** Fitted Multinomial HMM result. */
+export interface MultinomialHMMFit {
+ /** Initial state probabilities (shape: nComponents). */
+ startProb: number[];
+ /** Transition matrix (shape: nComponents × nComponents). */
+ transmat: number[][];
+ /** Emission probability matrix (shape: nComponents × nFeatures). */
+ emissionProb: number[][];
+ /** Log-likelihood of training data at convergence. */
+ logProb: number;
+ /** Number of EM iterations completed. */
+ nIterDone: number;
+}
+
+// ─── Forward-Backward (log-space) ─────────────────────────────────────────────
+
+/**
+ * Compute log-forward variables.
+ * @param logStartProb - log initial probabilities (nStates)
+ * @param logTransmat - log transition matrix (nStates × nStates)
+ * @param logEmit - log emission probabilities (T × nStates)
+ * @returns logAlpha (T × nStates)
+ */
+function logForward(
+ logStartProb: readonly number[],
+ logTransmat: readonly (readonly number[])[],
+ logEmit: readonly (readonly number[])[],
+): number[][] {
+ const T = logEmit.length;
+ const K = logStartProb.length;
+ const logAlpha: number[][] = Array.from({ length: T }, () => new Array(K).fill(LOG_ZERO));
+
+ for (let j = 0; j < K; j++) {
+ logAlpha[0]![j] = (logStartProb[j] ?? LOG_ZERO) + (logEmit[0]?.[j] ?? LOG_ZERO);
+ }
+ for (let t = 1; t < T; t++) {
+ for (let j = 0; j < K; j++) {
+ let s = LOG_ZERO;
+ for (let i = 0; i < K; i++) {
+ s = logSumExp(s, (logAlpha[t - 1]?.[i] ?? LOG_ZERO) + (logTransmat[i]?.[j] ?? LOG_ZERO));
+ }
+ logAlpha[t]![j] = s + (logEmit[t]?.[j] ?? LOG_ZERO);
+ }
+ }
+ return logAlpha;
+}
+
+/**
+ * Compute log-backward variables.
+ * @param logTransmat - log transition matrix (nStates × nStates)
+ * @param logEmit - log emission probabilities (T × nStates)
+ * @returns logBeta (T × nStates)
+ */
+function logBackward(
+ logTransmat: readonly (readonly number[])[],
+ logEmit: readonly (readonly number[])[],
+): number[][] {
+ const T = logEmit.length;
+ const K = logTransmat.length;
+ const logBeta: number[][] = Array.from({ length: T }, () => new Array(K).fill(0));
+
+ for (let t = T - 2; t >= 0; t--) {
+ for (let i = 0; i < K; i++) {
+ let s = LOG_ZERO;
+ for (let j = 0; j < K; j++) {
+ s = logSumExp(
+ s,
+ (logTransmat[i]?.[j] ?? LOG_ZERO) +
+ (logEmit[t + 1]?.[j] ?? LOG_ZERO) +
+ (logBeta[t + 1]?.[j] ?? 0),
+ );
+ }
+ logBeta[t]![i] = s;
+ }
+ }
+ return logBeta;
+}
+
+// ─── Viterbi ──────────────────────────────────────────────────────────────────
+
+/**
+ * Viterbi algorithm to find the most-likely state sequence.
+ * @returns decoded state sequence (length T)
+ */
+function viterbi(
+ logStartProb: readonly number[],
+ logTransmat: readonly (readonly number[])[],
+ logEmit: readonly (readonly number[])[],
+): number[] {
+ const T = logEmit.length;
+ const K = logStartProb.length;
+ const delta: number[][] = Array.from({ length: T }, () => new Array(K).fill(LOG_ZERO));
+ const psi: number[][] = Array.from({ length: T }, () => new Array(K).fill(0));
+
+ for (let j = 0; j < K; j++) {
+ delta[0]![j] = (logStartProb[j] ?? LOG_ZERO) + (logEmit[0]?.[j] ?? LOG_ZERO);
+ }
+ for (let t = 1; t < T; t++) {
+ for (let j = 0; j < K; j++) {
+ let best = LOG_ZERO;
+ let bestI = 0;
+ for (let i = 0; i < K; i++) {
+ const v = (delta[t - 1]?.[i] ?? LOG_ZERO) + (logTransmat[i]?.[j] ?? LOG_ZERO);
+ if (v > best) {
+ best = v;
+ bestI = i;
+ }
+ }
+ delta[t]![j] = best + (logEmit[t]?.[j] ?? LOG_ZERO);
+ psi[t]![j] = bestI;
+ }
+ }
+
+ // Backtrack
+ const path = new Array(T).fill(0);
+ let s = 0;
+ let best = LOG_ZERO;
+ for (let j = 0; j < K; j++) {
+ const v = delta[T - 1]?.[j] ?? LOG_ZERO;
+ if (v > best) {
+ best = v;
+ s = j;
+ }
+ }
+ path[T - 1] = s;
+ for (let t = T - 2; t >= 0; t--) {
+ path[t] = psi[t + 1]?.[path[t + 1] ?? 0] ?? 0;
+ }
+ return path;
+}
+
+// ─── GaussianHMM ──────────────────────────────────────────────────────────────
+
+/** Log probability of x under N(mu, sigma^2). */
+function gaussianLogProb(x: number, mu: number, sigma2: number): number {
+ if (sigma2 <= 0) {
+ return LOG_ZERO;
+ }
+ return -0.5 * (Math.log(2 * Math.PI * sigma2) + ((x - mu) * (x - mu)) / sigma2);
+}
+
+/**
+ * Hidden Markov Model with univariate Gaussian emission distributions.
+ *
+ * Uses Baum-Welch EM for parameter estimation, Viterbi for decoding.
+ */
+export class GaussianHMM {
+ private readonly K: number;
+ private readonly nIter: number;
+ private readonly tol: number;
+
+ // Parameters (set after fit)
+ private _startProb: number[] = [];
+ private _transmat: number[][] = [];
+ private _means: number[] = [];
+ private _covars: number[] = [];
+ private _fitted = false;
+
+ constructor(params: GaussianHMMParams) {
+ this.K = params.nComponents;
+ this.nIter = params.nIter ?? 100;
+ this.tol = params.tol ?? 1e-4;
+ }
+
+ /** Fit the HMM to a sequence of observations using Baum-Welch EM. */
+ fit(obs: readonly number[]): GaussianHMMFit {
+ const T = obs.length;
+ const K = this.K;
+ if (T < 2) {
+ throw new Error("Need at least 2 observations");
+ }
+
+ // ── Initialise ───────────────────────────────────────────────────────────
+ // k-means-style: split sorted observations into K equal buckets
+ const sorted = [...obs].sort((a, b) => a - b);
+ const means = new Array(K).fill(0);
+ const covars = new Array(K).fill(1);
+ for (let k = 0; k < K; k++) {
+ const lo = Math.floor((k * T) / K);
+ const hi = Math.floor(((k + 1) * T) / K);
+ const slice = sorted.slice(lo, hi);
+ const mu = slice.reduce((a, b) => a + b, 0) / slice.length;
+ const v = slice.reduce((a, b) => a + (b - mu) ** 2, 0) / Math.max(slice.length - 1, 1);
+ means[k] = mu;
+ covars[k] = Math.max(v, 1e-6);
+ }
+
+ // Uniform start and transition
+ const startProb = new Array(K).fill(1 / K);
+ const transmat: number[][] = Array.from({ length: K }, () => new Array(K).fill(1 / K));
+
+ let prevLogProb = Number.NEGATIVE_INFINITY;
+ let nIterDone = 0;
+
+ for (let iter = 0; iter < this.nIter; iter++) {
+ // ── E-step ─────────────────────────────────────────────────────────────
+ const logEmit: number[][] = Array.from({ length: T }, (_, t) =>
+ Array.from({ length: K }, (__, k) =>
+ gaussianLogProb(obs[t] ?? 0, means[k] ?? 0, covars[k] ?? 1),
+ ),
+ );
+
+ const logStartProb = startProb.map(safeLog);
+ const logTransmat = transmat.map((row) => row.map(safeLog));
+
+ const logAlpha = logForward(logStartProb, logTransmat, logEmit);
+ const logBeta = logBackward(logTransmat, logEmit);
+
+ // Log-likelihood
+ const logProb = logSumExpArr(logAlpha[T - 1] ?? []);
+
+ if (Math.abs(logProb - prevLogProb) < this.tol) {
+ nIterDone = iter + 1;
+ break;
+ }
+ prevLogProb = logProb;
+ nIterDone = iter + 1;
+
+ // γ_t(k) = P(z_t = k | obs, θ) in log space
+ const logGamma: number[][] = Array.from({ length: T }, (_, t) => {
+ const row = Array.from(
+ { length: K },
+ (__, k) => (logAlpha[t]?.[k] ?? LOG_ZERO) + (logBeta[t]?.[k] ?? 0),
+ );
+ const z = logSumExpArr(row);
+ return row.map((v) => v - z);
+ });
+
+ // ξ_t(i,j) = P(z_t=i, z_{t+1}=j | obs, θ) — sum over t
+ const logXiSum: number[][] = Array.from({ length: K }, () =>
+ new Array(K).fill(LOG_ZERO),
+ );
+ for (let t = 0; t < T - 1; t++) {
+ for (let i = 0; i < K; i++) {
+ for (let j = 0; j < K; j++) {
+ const v =
+ (logAlpha[t]?.[i] ?? LOG_ZERO) +
+ (logTransmat[i]?.[j] ?? LOG_ZERO) +
+ (logEmit[t + 1]?.[j] ?? LOG_ZERO) +
+ (logBeta[t + 1]?.[j] ?? 0) -
+ logProb;
+ logXiSum[i]![j] = logSumExp(logXiSum[i]?.[j] ?? LOG_ZERO, v);
+ }
+ }
+ }
+
+ // ── M-step ─────────────────────────────────────────────────────────────
+ // Update startProb
+ for (let k = 0; k < K; k++) {
+ startProb[k] = Math.exp(logGamma[0]?.[k] ?? LOG_ZERO);
+ }
+ normalise(startProb);
+
+ // Update transmat
+ for (let i = 0; i < K; i++) {
+ for (let j = 0; j < K; j++) {
+ transmat[i]![j] = Math.exp(logXiSum[i]?.[j] ?? LOG_ZERO);
+ }
+ normalise(transmat[i]!);
+ }
+
+ // Update means
+ for (let k = 0; k < K; k++) {
+ let num = 0;
+ let den = 0;
+ for (let t = 0; t < T; t++) {
+ const g = Math.exp(logGamma[t]?.[k] ?? LOG_ZERO);
+ num += g * (obs[t] ?? 0);
+ den += g;
+ }
+ means[k] = den > 0 ? num / den : (means[k] ?? 0);
+ }
+
+ // Update covars
+ for (let k = 0; k < K; k++) {
+ let num = 0;
+ let den = 0;
+ for (let t = 0; t < T; t++) {
+ const g = Math.exp(logGamma[t]?.[k] ?? LOG_ZERO);
+ const diff = (obs[t] ?? 0) - (means[k] ?? 0);
+ num += g * diff * diff;
+ den += g;
+ }
+ covars[k] = Math.max(den > 0 ? num / den : (covars[k] ?? 1), 1e-6);
+ }
+ }
+
+ this._startProb = [...startProb];
+ this._transmat = transmat.map((row) => [...row]);
+ this._means = [...means];
+ this._covars = [...covars];
+ this._fitted = true;
+
+ return {
+ startProb: [...startProb],
+ transmat: transmat.map((row) => [...row]),
+ means: [...means],
+ covars: [...covars],
+ logProb: prevLogProb,
+ nIterDone,
+ };
+ }
+
+ /** Decode the most-likely state sequence using the Viterbi algorithm. */
+ predict(obs: readonly number[]): number[] {
+ this._checkFitted();
+ const K = this.K;
+ const logEmit = obs.map((x) =>
+ Array.from({ length: K }, (_, k) =>
+ gaussianLogProb(x, this._means[k] ?? 0, this._covars[k] ?? 1),
+ ),
+ );
+ return viterbi(
+ this._startProb.map(safeLog),
+ this._transmat.map((row) => row.map(safeLog)),
+ logEmit,
+ );
+ }
+
+ /** Compute the log-probability of the observation sequence. */
+ score(obs: readonly number[]): number {
+ this._checkFitted();
+ const K = this.K;
+ const logEmit = obs.map((x) =>
+ Array.from({ length: K }, (_, k) =>
+ gaussianLogProb(x, this._means[k] ?? 0, this._covars[k] ?? 1),
+ ),
+ );
+ const logAlpha = logForward(
+ this._startProb.map(safeLog),
+ this._transmat.map((row) => row.map(safeLog)),
+ logEmit,
+ );
+ return logSumExpArr(logAlpha.at(-1) ?? []);
+ }
+
+ /** Compute posterior state probabilities (T × nComponents). */
+ predictProba(obs: readonly number[]): number[][] {
+ this._checkFitted();
+ const K = this.K;
+ const T = obs.length;
+ const logEmit = obs.map((x) =>
+ Array.from({ length: K }, (_, k) =>
+ gaussianLogProb(x, this._means[k] ?? 0, this._covars[k] ?? 1),
+ ),
+ );
+ const logStartProb = this._startProb.map(safeLog);
+ const logTransmat = this._transmat.map((row) => row.map(safeLog));
+ const logAlpha = logForward(logStartProb, logTransmat, logEmit);
+ const logBeta = logBackward(logTransmat, logEmit);
+ return Array.from({ length: T }, (_, t) => {
+ const row = Array.from(
+ { length: K },
+ (__, k) => (logAlpha[t]?.[k] ?? LOG_ZERO) + (logBeta[t]?.[k] ?? 0),
+ );
+ const z = logSumExpArr(row);
+ return row.map((v) => Math.exp(v - z));
+ });
+ }
+
+ /** Sample a sequence of states and observations. */
+ sample(length: number): { states: number[]; obs: number[] } {
+ this._checkFitted();
+ const K = this.K;
+ const states: number[] = [];
+ const obs: number[] = [];
+
+ // Sample initial state
+ let cumProb = 0;
+ const r0 = Math.random();
+ let state = K - 1;
+ for (let k = 0; k < K; k++) {
+ cumProb += this._startProb[k] ?? 0;
+ if (r0 < cumProb) {
+ state = k;
+ break;
+ }
+ }
+
+ for (let t = 0; t < length; t++) {
+ states.push(state);
+ const mu = this._means[state] ?? 0;
+ const sigma = Math.sqrt(this._covars[state] ?? 1);
+ // Box-Muller for Gaussian sample
+ const u1 = Math.random();
+ const u2 = Math.random();
+ const z = Math.sqrt(-2 * Math.log(Math.max(u1, 1e-15))) * Math.cos(2 * Math.PI * u2);
+ obs.push(mu + sigma * z);
+
+ // Transition
+ const row = this._transmat[state]!;
+ let cum = 0;
+ const rv = Math.random();
+ let nextState = K - 1;
+ for (let k = 0; k < K; k++) {
+ cum += row[k] ?? 0;
+ if (rv < cum) {
+ nextState = k;
+ break;
+ }
+ }
+ state = nextState;
+ }
+ return { states, obs };
+ }
+
+ get startProb(): number[] {
+ this._checkFitted();
+ return [...this._startProb];
+ }
+ get transmat(): number[][] {
+ this._checkFitted();
+ return this._transmat.map((r) => [...r]);
+ }
+ get means(): number[] {
+ this._checkFitted();
+ return [...this._means];
+ }
+ get covars(): number[] {
+ this._checkFitted();
+ return [...this._covars];
+ }
+
+ private _checkFitted(): void {
+ if (!this._fitted) {
+ throw new Error("GaussianHMM is not fitted yet");
+ }
+ }
+}
+
+// ─── MultinomialHMM ───────────────────────────────────────────────────────────
+
+/**
+ * Hidden Markov Model with discrete (multinomial/categorical) emission distributions.
+ *
+ * Observations are non-negative integer symbol indices in [0, nFeatures).
+ */
+export class MultinomialHMM {
+ private readonly K: number;
+ private readonly nFeatures: number;
+ private readonly nIter: number;
+ private readonly tol: number;
+
+ private _startProb: number[] = [];
+ private _transmat: number[][] = [];
+ private _emissionProb: number[][] = [];
+ private _fitted = false;
+
+ constructor(params: MultinomialHMMParams) {
+ this.K = params.nComponents;
+ this.nFeatures = params.nFeatures;
+ this.nIter = params.nIter ?? 100;
+ this.tol = params.tol ?? 1e-4;
+ }
+
+ /** Fit the HMM to a sequence of integer observations using Baum-Welch EM. */
+ fit(obs: readonly number[]): MultinomialHMMFit {
+ const T = obs.length;
+ const K = this.K;
+ const V = this.nFeatures;
+ if (T < 2) {
+ throw new Error("Need at least 2 observations");
+ }
+
+ // Uniform initialisation with small random perturbation
+ const startProb = new Array(K).fill(1 / K);
+ const transmat: number[][] = Array.from({ length: K }, () =>
+ Array.from({ length: K }, () => 1 / K + (Math.random() * 0.1 - 0.05) / K),
+ );
+ const emissionProb: number[][] = Array.from({ length: K }, () =>
+ Array.from({ length: V }, () => 1 / V + (Math.random() * 0.1 - 0.05) / V),
+ );
+ // Normalise
+ for (let k = 0; k < K; k++) {
+ normalise(transmat[k]!);
+ normalise(emissionProb[k]!);
+ }
+
+ let prevLogProb = Number.NEGATIVE_INFINITY;
+ let nIterDone = 0;
+
+ for (let iter = 0; iter < this.nIter; iter++) {
+ // Log-emission: logEmit[t][k] = log P(obs[t] | z_t = k)
+ const logEmit: number[][] = Array.from({ length: T }, (_, t) =>
+ Array.from({ length: K }, (__, k) => safeLog(emissionProb[k]?.[obs[t] ?? 0] ?? 0)),
+ );
+
+ const logStartProb = startProb.map(safeLog);
+ const logTransmat = transmat.map((row) => row.map(safeLog));
+
+ const logAlpha = logForward(logStartProb, logTransmat, logEmit);
+ const logBeta = logBackward(logTransmat, logEmit);
+
+ const logProb = logSumExpArr(logAlpha[T - 1] ?? []);
+ if (Math.abs(logProb - prevLogProb) < this.tol) {
+ nIterDone = iter + 1;
+ break;
+ }
+ prevLogProb = logProb;
+ nIterDone = iter + 1;
+
+ // γ_t(k)
+ const logGamma: number[][] = Array.from({ length: T }, (_, t) => {
+ const row = Array.from(
+ { length: K },
+ (__, k) => (logAlpha[t]?.[k] ?? LOG_ZERO) + (logBeta[t]?.[k] ?? 0),
+ );
+ const z = logSumExpArr(row);
+ return row.map((v) => v - z);
+ });
+
+ // ξ sum
+ const logXiSum: number[][] = Array.from({ length: K }, () =>
+ new Array(K).fill(LOG_ZERO),
+ );
+ for (let t = 0; t < T - 1; t++) {
+ for (let i = 0; i < K; i++) {
+ for (let j = 0; j < K; j++) {
+ const v =
+ (logAlpha[t]?.[i] ?? LOG_ZERO) +
+ (logTransmat[i]?.[j] ?? LOG_ZERO) +
+ (logEmit[t + 1]?.[j] ?? LOG_ZERO) +
+ (logBeta[t + 1]?.[j] ?? 0) -
+ logProb;
+ logXiSum[i]![j] = logSumExp(logXiSum[i]?.[j] ?? LOG_ZERO, v);
+ }
+ }
+ }
+
+ // M-step: startProb
+ for (let k = 0; k < K; k++) {
+ startProb[k] = Math.exp(logGamma[0]?.[k] ?? LOG_ZERO);
+ }
+ normalise(startProb);
+
+ // transmat
+ for (let i = 0; i < K; i++) {
+ for (let j = 0; j < K; j++) {
+ transmat[i]![j] = Math.exp(logXiSum[i]?.[j] ?? LOG_ZERO);
+ }
+ normalise(transmat[i]!);
+ }
+
+ // emissionProb
+ for (let k = 0; k < K; k++) {
+ for (let v = 0; v < V; v++) {
+ emissionProb[k]![v] = 0;
+ }
+ for (let t = 0; t < T; t++) {
+ const sym = obs[t] ?? 0;
+ emissionProb[k]![sym] =
+ (emissionProb[k]?.[sym] ?? 0) + Math.exp(logGamma[t]?.[k] ?? LOG_ZERO);
+ }
+ normalise(emissionProb[k]!);
+ }
+ }
+
+ this._startProb = [...startProb];
+ this._transmat = transmat.map((row) => [...row]);
+ this._emissionProb = emissionProb.map((row) => [...row]);
+ this._fitted = true;
+
+ return {
+ startProb: [...startProb],
+ transmat: transmat.map((row) => [...row]),
+ emissionProb: emissionProb.map((row) => [...row]),
+ logProb: prevLogProb,
+ nIterDone,
+ };
+ }
+
+ /** Decode the most-likely state sequence. */
+ predict(obs: readonly number[]): number[] {
+ this._checkFitted();
+ const K = this.K;
+ const logEmit = obs.map((sym) =>
+ Array.from({ length: K }, (_, k) => safeLog(this._emissionProb[k]?.[sym] ?? 0)),
+ );
+ return viterbi(
+ this._startProb.map(safeLog),
+ this._transmat.map((row) => row.map(safeLog)),
+ logEmit,
+ );
+ }
+
+ /** Compute log-probability of the observation sequence. */
+ score(obs: readonly number[]): number {
+ this._checkFitted();
+ const K = this.K;
+ const logEmit = obs.map((sym) =>
+ Array.from({ length: K }, (_, k) => safeLog(this._emissionProb[k]?.[sym] ?? 0)),
+ );
+ const logAlpha = logForward(
+ this._startProb.map(safeLog),
+ this._transmat.map((row) => row.map(safeLog)),
+ logEmit,
+ );
+ return logSumExpArr(logAlpha.at(-1) ?? []);
+ }
+
+ get startProb(): number[] {
+ this._checkFitted();
+ return [...this._startProb];
+ }
+ get transmat(): number[][] {
+ this._checkFitted();
+ return this._transmat.map((r) => [...r]);
+ }
+ get emissionProb(): number[][] {
+ this._checkFitted();
+ return this._emissionProb.map((r) => [...r]);
+ }
+
+ private _checkFitted(): void {
+ if (!this._fitted) {
+ throw new Error("MultinomialHMM is not fitted yet");
+ }
+ }
+}
+
+// ─── Standalone functions ──────────────────────────────────────────────────────
+
+/**
+ * Convenience function: fit a GaussianHMM and return the fitted model.
+ *
+ * @example
+ * ```ts
+ * const model = fitGaussianHMM([0.1, 0.2, 2.1, 2.3, 0.05, 2.5], 2);
+ * ```
+ */
+export function fitGaussianHMM(
+ obs: readonly number[],
+ nComponents: number,
+ nIter = 100,
+): GaussianHMM {
+ const model = new GaussianHMM({ nComponents, nIter });
+ model.fit(obs);
+ return model;
+}
+
+/**
+ * Convenience function: Viterbi decoding with explicit parameters (no fitting).
+ *
+ * @param startProb - Initial state probabilities (length K).
+ * @param transmat - Transition matrix (K × K).
+ * @param emissionProb - Emission probabilities (K × V).
+ * @param obs - Integer observation sequence.
+ * @returns Most-likely state sequence.
+ */
+export function hmmViterbi(
+ startProb: readonly number[],
+ transmat: readonly (readonly number[])[],
+ emissionProb: readonly (readonly number[])[],
+ obs: readonly number[],
+): number[] {
+ const K = startProb.length;
+ const logEmit = obs.map((sym) =>
+ Array.from({ length: K }, (_, k) => safeLog(emissionProb[k]?.[sym] ?? 0)),
+ );
+ return viterbi(
+ startProb.map(safeLog),
+ transmat.map((row) => [...row].map(safeLog)),
+ logEmit,
+ );
+}
diff --git a/src/stats/index.ts b/src/stats/index.ts
index 78767ee1..dbfb531d 100644
--- a/src/stats/index.ts
+++ b/src/stats/index.ts
@@ -575,3 +575,136 @@ export {
tsallisEntropy,
} from "./information.ts";
export type { PMF, NMIMethod } from "./information.ts";
+export {
+ complex,
+ cAbs,
+ cArg,
+ fft,
+ ifft,
+ rfft,
+ irfft,
+ fftFreq,
+ rfftFreq,
+ fftshift,
+ ifftshift,
+ rectangularWindow,
+ bartlettWindow,
+ hannWindow,
+ hammingWindow,
+ blackmanWindow,
+ blackmanHarrisWindow,
+ flatTopWindow,
+ kaiserWindow,
+ getWindow,
+ stft,
+ istft,
+ welch,
+ periodogram,
+} from "./signal.ts";
+export type {
+ Complex,
+ WindowName,
+ STFTOptions,
+ STFTResult,
+ ISTFTOptions,
+ WelchOptions,
+ PSDResult,
+ PeriodogramOptions,
+} from "./signal.ts";
+export {
+ firwin,
+ freqz,
+ sosfreqz,
+ lfilter,
+ filtfilt,
+ sosfilt,
+ sosfiltfilt,
+ butter,
+} from "./filters.ts";
+export type {
+ FirwinOptions,
+ FreqzResult,
+ SOSSection,
+ ButterResult,
+ FilterType,
+} from "./filters.ts";
+export {
+ autocorr,
+ acf,
+ pacf,
+ ccf,
+ durbinWatson,
+ ljungBox,
+ boxPierce,
+} from "./acf_pacf.ts";
+export type {
+ ACFResult,
+ PACFResult,
+ PortmanteauResult,
+ ACFOptions,
+ PACFOptions,
+ CCFOptions,
+ PortmanteauOptions,
+} from "./acf_pacf.ts";
+export {
+ ARIMAModel,
+ fitArima,
+} from "./arima.ts";
+export type {
+ ARIMAOptions,
+ ARIMAFitResult,
+ ARIMAForecastResult,
+} from "./arima.ts";
+export {
+ KalmanFilter,
+ StateSpaceModel,
+ kalmanFilter1D,
+ kalmanSmooth1D,
+ extractScalarMeans,
+ extractScalarVariances,
+ filteredPredictionInterval,
+} from "./kalman.ts";
+export type {
+ KalmanFilterOptions,
+ LocalLevelOptions,
+ LocalLinearTrendOptions,
+ KalmanFilterResult,
+ KalmanSmootherResult,
+} from "./kalman.ts";
+export {
+ SimpleExpSmoothing,
+ Holt,
+ ExponentialSmoothing,
+ simpleExpSmoothing,
+ holt,
+ fitEts,
+} from "./ets.ts";
+export type {
+ ETSTrend,
+ ETSSeasonal,
+ ETSInit,
+ SESOptions,
+ SESFitResult,
+ HoltOptions,
+ HoltFitResult,
+ ExponentialSmoothingOptions,
+ ExponentialSmoothingFitResult,
+ ETSForecastResult,
+} from "./ets.ts";
+export {
+ DLM,
+ buildLocalLevel,
+ buildLocalLinearTrend,
+ buildPolynomial,
+ buildFourier,
+ buildRegression,
+ combineDLMs,
+} from "./dlm.ts";
+export type {
+ DLMSpec,
+ DLMOptions,
+ DLMFilterStep,
+ DLMResult,
+ DLMSmootherResult,
+ DLMForecastResult,
+} from "./dlm.ts";
diff --git a/src/stats/kalman.ts b/src/stats/kalman.ts
new file mode 100644
index 00000000..a0493a78
--- /dev/null
+++ b/src/stats/kalman.ts
@@ -0,0 +1,809 @@
+/**
+ * kalman — Linear Gaussian State-Space Model: Kalman Filter & RTS Smoother.
+ *
+ * Implements the standard discrete-time Kalman filter (forward pass) and the
+ * Rauch-Tung-Striebel (RTS) smoother (backward pass) for linear dynamical
+ * systems:
+ *
+ * x_t = F·x_{t-1} + w_t, w_t ~ N(0, Q) (state equation)
+ * y_t = H·x_t + v_t, v_t ~ N(0, R) (observation equation)
+ * x_0 ~ N(m0, P0)
+ *
+ * Missing observations (null) are handled by skipping the update step —
+ * the filtered state reverts to the predicted state for that time-step.
+ *
+ * Mirrors the `statsmodels.tsa.statespace.kalman_filter.KalmanFilter` and
+ * `pykalman.KalmanFilter` APIs; factory helpers match common pandas patterns.
+ *
+ * Exported names:
+ * - {@link KalmanFilter} — main class (filter + smooth)
+ * - {@link KalmanFilterOptions} — constructor options
+ * - {@link KalmanFilterResult} — forward-pass output
+ * - {@link KalmanSmootherResult} — backward-pass output
+ *
+ * @example
+ * ```ts
+ * import { KalmanFilter } from "tsb";
+ *
+ * // Local-level model (random walk observed with noise)
+ * const kf = KalmanFilter.localLevel({ processNoise: 1, observationNoise: 2 });
+ * const res = kf.filter([[1], [2], [1.5], [null], [3], [2.5]]);
+ * console.log(res.filteredStateMeans); // [[…], …] T × 1
+ * console.log(res.logLikelihood);
+ *
+ * const sm = kf.smooth([[1], [2], [1.5], [null], [3], [2.5]]);
+ * console.log(sm.smoothedStateMeans);
+ * ```
+ *
+ * @module
+ */
+
+// ─── Internal matrix helpers ───────────────────────────────────────────────────
+
+/** Read-only row-major matrix. */
+type Mat = readonly (readonly number[])[];
+/** Mutable row-major matrix. */
+type MutMat = number[][];
+
+/** Rows of A (no checks — callers must ensure dimensions). */
+function rows(A: Mat): number {
+ return A.length;
+}
+/** Columns of A (0 if empty). */
+function cols(A: Mat): number {
+ return A[0]?.length ?? 0;
+}
+
+/** Create an n×m zero matrix. */
+function zeros(n: number, m: number): MutMat {
+ return Array.from({ length: n }, () => new Array(m).fill(0));
+}
+
+/** Create an n×n identity matrix. */
+function eye(n: number): MutMat {
+ return Array.from({ length: n }, (_, i) =>
+ Array.from({ length: n }, (_, j) => (i === j ? 1 : 0)),
+ );
+}
+
+/** Matrix–matrix product A (m×k) · B (k×n) → (m×n). */
+function mmul(A: Mat, B: Mat): MutMat {
+ const m = rows(A);
+ const k = cols(A);
+ const n = cols(B);
+ const C = zeros(m, n);
+ for (let i = 0; i < m; i++) {
+ const ai = A[i]!;
+ const ci = C[i]!;
+ for (let p = 0; p < k; p++) {
+ const aip = ai[p]!;
+ if (aip === 0) {
+ continue;
+ }
+ const bp = B[p]!;
+ for (let j = 0; j < n; j++) {
+ ci[j] = (ci[j] ?? 0) + aip * bp[j]!;
+ }
+ }
+ }
+ return C;
+}
+
+/** Matrix–vector product A (m×k) · x (k) → (m). */
+function mvmul(A: Mat, x: readonly number[]): number[] {
+ const m = rows(A);
+ const k = x.length;
+ const y = new Array(m).fill(0);
+ for (let i = 0; i < m; i++) {
+ const ai = A[i]!;
+ let s = 0;
+ for (let p = 0; p < k; p++) {
+ s += ai[p]! * x[p]!;
+ }
+ y[i] = s;
+ }
+ return y;
+}
+
+/** Transpose of A (m×n) → (n×m). */
+function T(A: Mat): MutMat {
+ const m = rows(A);
+ const n = cols(A);
+ const At = zeros(n, m);
+ for (let i = 0; i < m; i++) {
+ for (let j = 0; j < n; j++) {
+ At[j]![i] = A[i]?.[j]!;
+ }
+ }
+ return At;
+}
+
+/** A + B (element-wise). */
+function madd(A: Mat, B: Mat): MutMat {
+ const m = rows(A);
+ const n = cols(A);
+ return Array.from({ length: m }, (_, i) =>
+ Array.from({ length: n }, (_, j) => A[i]?.[j]! + B[i]?.[j]!),
+ );
+}
+
+/** A − B (element-wise). */
+function msub(A: Mat, B: Mat): MutMat {
+ const m = rows(A);
+ const n = cols(A);
+ return Array.from({ length: m }, (_, i) =>
+ Array.from({ length: n }, (_, j) => A[i]?.[j]! - B[i]?.[j]!),
+ );
+}
+
+/** Vector addition a + b. */
+function vadd(a: readonly number[], b: readonly number[]): number[] {
+ return a.map((ai, i) => ai + b[i]!);
+}
+
+/** Vector subtraction a − b. */
+function vsub(a: readonly number[], b: readonly number[]): number[] {
+ return a.map((ai, i) => ai - b[i]!);
+}
+
+/**
+ * Invert a square matrix via Gaussian elimination with partial pivoting.
+ * Returns null if the matrix is singular (|pivot| < 1e-14).
+ */
+function matInv(A: Mat): MutMat | null {
+ const n = rows(A);
+ // Augmented [A | I]
+ const aug: MutMat = Array.from({ length: n }, (_, i) => [
+ ...A[i]!,
+ ...Array.from({ length: n }, (_, j) => (i === j ? 1 : 0)),
+ ]);
+ for (let col = 0; col < n; col++) {
+ let maxRow = col;
+ let maxVal = Math.abs(aug[col]?.[col]!);
+ for (let row = col + 1; row < n; row++) {
+ const v = Math.abs(aug[row]?.[col]!);
+ if (v > maxVal) {
+ maxVal = v;
+ maxRow = row;
+ }
+ }
+ if (maxVal < 1e-14) {
+ return null;
+ }
+ [aug[col], aug[maxRow]] = [aug[maxRow]!, aug[col]!];
+ const pivot = aug[col]?.[col]!;
+ const pivRow = aug[col]!;
+ for (let j = 0; j < 2 * n; j++) {
+ pivRow[j] = pivRow[j]! / pivot;
+ }
+ for (let row = 0; row < n; row++) {
+ if (row === col) {
+ continue;
+ }
+ const fac = aug[row]?.[col]!;
+ if (fac === 0) {
+ continue;
+ }
+ const r = aug[row]!;
+ for (let j = 0; j < 2 * n; j++) {
+ r[j] = r[j]! - fac * pivRow[j]!;
+ }
+ }
+ }
+ return aug.map((row) => row.slice(n));
+}
+
+/** log-determinant via LU decomposition (for log-likelihood). */
+function logDet(A: Mat): number {
+ const n = rows(A);
+ const L: MutMat = Array.from({ length: n }, (_, i) => [...A[i]!]);
+ let logD = 0;
+ for (let col = 0; col < n; col++) {
+ let maxRow = col;
+ let maxVal = Math.abs(L[col]?.[col]!);
+ for (let row = col + 1; row < n; row++) {
+ const v = Math.abs(L[row]?.[col]!);
+ if (v > maxVal) {
+ maxVal = v;
+ maxRow = row;
+ }
+ }
+ if (maxRow !== col) {
+ [L[col], L[maxRow]] = [L[maxRow]!, L[col]!];
+ logD += Math.log(-1); // sign flip — handled by real part
+ }
+ const pivot = L[col]?.[col]!;
+ if (Math.abs(pivot) < 1e-300) {
+ return Number.NEGATIVE_INFINITY;
+ }
+ logD += Math.log(Math.abs(pivot));
+ for (let row = col + 1; row < n; row++) {
+ const fac = L[row]?.[col]! / pivot;
+ const r = L[row]!;
+ for (let j = col; j < n; j++) {
+ r[j] = r[j]! - fac * L[col]?.[j]!;
+ }
+ }
+ }
+ return logD;
+}
+
+/** Outer product a·bᵀ → matrix. */
+function outer(a: readonly number[], b: readonly number[]): MutMat {
+ return Array.from({ length: a.length }, (_, i) =>
+ Array.from({ length: b.length }, (_, j) => a[i]! * b[j]!),
+ );
+}
+
+/** Scale matrix by scalar. */
+function mscale(A: Mat, s: number): MutMat {
+ return A.map((row) => row.map((v) => v * s));
+}
+
+// ─── Public types ──────────────────────────────────────────────────────────────
+
+/** Constructor options for {@link KalmanFilter}. */
+export interface KalmanFilterOptions {
+ /**
+ * State transition matrix **F** (n_states × n_states).
+ * Defines how the state evolves: x_t = F·x_{t-1} + noise.
+ */
+ readonly transitionMatrix: readonly (readonly number[])[];
+ /**
+ * Observation matrix **H** (n_obs × n_states).
+ * Maps states to observations: y_t = H·x_t + noise.
+ */
+ readonly observationMatrix: readonly (readonly number[])[];
+ /**
+ * Process noise covariance **Q** (n_states × n_states).
+ * Covariance of the state-transition noise.
+ */
+ readonly processNoiseCov: readonly (readonly number[])[];
+ /**
+ * Observation noise covariance **R** (n_obs × n_obs).
+ * Covariance of the observation noise.
+ */
+ readonly observationNoiseCov: readonly (readonly number[])[];
+ /**
+ * Initial state mean **m₀** (n_states vector).
+ * Defaults to the zero vector.
+ */
+ readonly initialStateMean?: readonly number[];
+ /**
+ * Initial state covariance **P₀** (n_states × n_states).
+ * Defaults to the identity matrix.
+ */
+ readonly initialStateCovariance?: readonly (readonly number[])[];
+}
+
+/** Options for {@link KalmanFilter.localLevel}. */
+export interface LocalLevelOptions {
+ /** Process (state) noise variance σ²_q. Default: 1. */
+ readonly processNoise?: number;
+ /** Observation noise variance σ²_r. Default: 1. */
+ readonly observationNoise?: number;
+ /** Initial state mean scalar. Default: 0. */
+ readonly initialMean?: number;
+ /** Initial state variance scalar. Default: 1. */
+ readonly initialVariance?: number;
+}
+
+/** Options for {@link KalmanFilter.localLinearTrend}. */
+export interface LocalLinearTrendOptions {
+ /** Level process noise variance. Default: 1. */
+ readonly levelNoise?: number;
+ /** Slope process noise variance. Default: 0.1. */
+ readonly slopeNoise?: number;
+ /** Observation noise variance. Default: 1. */
+ readonly observationNoise?: number;
+ /** Initial [level, slope] mean vector. Default: [0, 0]. */
+ readonly initialMean?: readonly [number, number];
+ /** Initial state variance (diagonal). Default: 1. */
+ readonly initialVariance?: number;
+}
+
+/** Result of the Kalman filter forward pass. */
+export interface KalmanFilterResult {
+ /**
+ * Filtered state means x_{t|t} — shape T × n_states.
+ * Each row is the posterior mean after incorporating observation t.
+ */
+ readonly filteredStateMeans: readonly (readonly number[])[];
+ /**
+ * Filtered state covariances P_{t|t} — shape T × n_states × n_states.
+ * Each entry is the posterior covariance after incorporating observation t.
+ */
+ readonly filteredStateCovariances: readonly (readonly (readonly number[])[])[];
+ /**
+ * Predicted state means x_{t|t-1} — shape T × n_states.
+ * Each row is the prior mean before incorporating observation t.
+ */
+ readonly predictedStateMeans: readonly (readonly number[])[];
+ /**
+ * Predicted state covariances P_{t|t-1} — shape T × n_states × n_states.
+ */
+ readonly predictedStateCovariances: readonly (readonly (readonly number[])[])[];
+ /**
+ * Innovation (prediction error) vectors y_t − H·x_{t|t-1} — shape T × n_obs.
+ * NaN rows indicate missing observations.
+ */
+ readonly innovations: readonly (readonly number[])[];
+ /**
+ * Innovation covariance matrices S_t = H·P_{t|t-1}·Hᵀ + R — shape T × n_obs × n_obs.
+ */
+ readonly innovationCovariances: readonly (readonly (readonly number[])[])[];
+ /** Gaussian log-likelihood summed over all non-missing time-steps. */
+ readonly logLikelihood: number;
+ /** Number of states (n_states). */
+ readonly nStates: number;
+ /** Number of observation dimensions (n_obs). */
+ readonly nObs: number;
+ /** Number of time steps (T). */
+ readonly nTime: number;
+}
+
+/** Result of the RTS smoother backward pass. */
+export interface KalmanSmootherResult {
+ /**
+ * Smoothed state means x_{t|T} — shape T × n_states.
+ * Each row is the posterior mean using all T observations.
+ */
+ readonly smoothedStateMeans: readonly (readonly number[])[];
+ /**
+ * Smoothed state covariances P_{t|T} — shape T × n_states × n_states.
+ */
+ readonly smoothedStateCovariances: readonly (readonly (readonly number[])[])[];
+ /**
+ * Smoother gain matrices G_t — shape T × n_states × n_states.
+ * (Last entry is all-zeros by convention.)
+ */
+ readonly smootherGains: readonly (readonly (readonly number[])[])[];
+ /** Same as {@link KalmanFilterResult.logLikelihood} (computed in forward pass). */
+ readonly logLikelihood: number;
+ /** The forward-pass result used to compute the smoother. */
+ readonly filterResult: KalmanFilterResult;
+}
+
+// ─── KalmanFilter class ────────────────────────────────────────────────────────
+
+/**
+ * Linear Gaussian State-Space Model with Kalman filter and RTS smoother.
+ *
+ * The model is:
+ * ```
+ * x_t = F·x_{t-1} + w_t, w_t ~ N(0, Q)
+ * y_t = H·x_t + v_t, v_t ~ N(0, R)
+ * x_0 ~ N(m0, P0)
+ * ```
+ *
+ * @example
+ * ```ts
+ * const kf = new KalmanFilter({
+ * transitionMatrix: [[1]],
+ * observationMatrix: [[1]],
+ * processNoiseCov: [[1]],
+ * observationNoiseCov: [[2]],
+ * });
+ * const result = kf.filter([[1], [2], [null], [3]]);
+ * ```
+ */
+export class KalmanFilter {
+ /** State transition matrix F (n_states × n_states). */
+ readonly transitionMatrix: Mat;
+ /** Observation matrix H (n_obs × n_states). */
+ readonly observationMatrix: Mat;
+ /** Process noise covariance Q (n_states × n_states). */
+ readonly processNoiseCov: Mat;
+ /** Observation noise covariance R (n_obs × n_obs). */
+ readonly observationNoiseCov: Mat;
+ /** Initial state mean m₀ (n_states). */
+ readonly initialStateMean: readonly number[];
+ /** Initial state covariance P₀ (n_states × n_states). */
+ readonly initialStateCovariance: Mat;
+
+ constructor(opts: KalmanFilterOptions) {
+ this.transitionMatrix = opts.transitionMatrix;
+ this.observationMatrix = opts.observationMatrix;
+ this.processNoiseCov = opts.processNoiseCov;
+ this.observationNoiseCov = opts.observationNoiseCov;
+
+ const ns = rows(opts.transitionMatrix);
+ this.initialStateMean = opts.initialStateMean ?? new Array(ns).fill(0);
+ this.initialStateCovariance = opts.initialStateCovariance ?? eye(ns);
+ }
+
+ // ─── Factory helpers ───────────────────────────────────────────────────────
+
+ /**
+ * Local-level model (random walk + measurement noise):
+ * ```
+ * x_t = x_{t-1} + w_t, w_t ~ N(0, σ²_q)
+ * y_t = x_t + v_t, v_t ~ N(0, σ²_r)
+ * ```
+ */
+ static localLevel(opts: LocalLevelOptions = {}): KalmanFilter {
+ const q = opts.processNoise ?? 1;
+ const r = opts.observationNoise ?? 1;
+ const m0 = opts.initialMean ?? 0;
+ const p0 = opts.initialVariance ?? 1;
+ return new KalmanFilter({
+ transitionMatrix: [[1]],
+ observationMatrix: [[1]],
+ processNoiseCov: [[q]],
+ observationNoiseCov: [[r]],
+ initialStateMean: [m0],
+ initialStateCovariance: [[p0]],
+ });
+ }
+
+ /**
+ * Local linear trend model (level + slope):
+ * ```
+ * level_t = level_{t-1} + slope_{t-1} + w1_t
+ * slope_t = slope_{t-1} + w2_t
+ * y_t = level_t + v_t
+ * ```
+ */
+ static localLinearTrend(opts: LocalLinearTrendOptions = {}): KalmanFilter {
+ const ql = opts.levelNoise ?? 1;
+ const qs = opts.slopeNoise ?? 0.1;
+ const r = opts.observationNoise ?? 1;
+ const [m0l, m0s] = opts.initialMean ?? [0, 0];
+ const p0 = opts.initialVariance ?? 1;
+ return new KalmanFilter({
+ transitionMatrix: [
+ [1, 1],
+ [0, 1],
+ ],
+ observationMatrix: [[1, 0]],
+ processNoiseCov: [
+ [ql, 0],
+ [0, qs],
+ ],
+ observationNoiseCov: [[r]],
+ initialStateMean: [m0l, m0s],
+ initialStateCovariance: [
+ [p0, 0],
+ [0, p0],
+ ],
+ });
+ }
+
+ // ─── Main methods ──────────────────────────────────────────────────────────
+
+ /**
+ * Run the Kalman filter (forward pass).
+ *
+ * @param observations T × n_obs array of observations.
+ * Pass `null` for any element to indicate a missing value at that
+ * time-step × dimension. If an entire row is missing, pass a row of nulls
+ * or just pass `null` in a scalar array like `[[null]]`.
+ * @returns {@link KalmanFilterResult}
+ *
+ * @example
+ * ```ts
+ * const result = kf.filter([[1], [2], [null], [3], [2.5]]);
+ * ```
+ */
+ filter(observations: readonly (readonly (number | null)[])[]): KalmanFilterResult {
+ return kalmanFilter(
+ observations,
+ this.transitionMatrix,
+ this.observationMatrix,
+ this.processNoiseCov,
+ this.observationNoiseCov,
+ this.initialStateMean,
+ this.initialStateCovariance,
+ );
+ }
+
+ /**
+ * Run the RTS smoother (Kalman filter forward pass + RTS backward pass).
+ *
+ * @param observations T × n_obs array (same format as {@link filter}).
+ * @returns {@link KalmanSmootherResult}
+ *
+ * @example
+ * ```ts
+ * const smoothed = kf.smooth([[1], [2], [null], [3], [2.5]]);
+ * console.log(smoothed.smoothedStateMeans);
+ * ```
+ */
+ smooth(observations: readonly (readonly (number | null)[])[]): KalmanSmootherResult {
+ const fwd = this.filter(observations);
+ return rtsSmooth(fwd, this.transitionMatrix);
+ }
+}
+
+// ─── Core algorithms ───────────────────────────────────────────────────────────
+
+/**
+ * Kalman filter forward pass.
+ *
+ * Returns all intermediate quantities needed for the RTS smoother and for
+ * log-likelihood computation.
+ */
+function kalmanFilter(
+ obs: readonly (readonly (number | null)[])[],
+ F: Mat,
+ H: Mat,
+ Q: Mat,
+ R: Mat,
+ m0: readonly number[],
+ P0: Mat,
+): KalmanFilterResult {
+ const T_len = obs.length;
+ const ns = rows(F);
+ const no = rows(H);
+
+ // Storage
+ const filtMeans: number[][] = [];
+ const filtCovs: MutMat[] = [];
+ const predMeans: number[][] = [];
+ const predCovs: MutMat[] = [];
+ const innovations: number[][] = [];
+ const innovCovs: MutMat[] = [];
+ let logLik = 0;
+
+ // Initialize
+ let xFilt: number[] = [...m0];
+ let PFilt: MutMat = P0.map((row) => [...row]);
+ const FT = T(F);
+ const HT = T(H);
+ const LOG2PI = Math.log(2 * Math.PI);
+
+ for (let t = 0; t < T_len; t++) {
+ const yt = obs[t]!;
+
+ // ── Predict ────────────────────────────────────────────────────────────
+ const xPred = mvmul(F, xFilt);
+ // P_pred = F P F' + Q
+ const PPred = madd(mmul(mmul(F, PFilt), FT), Q);
+
+ predMeans.push(xPred);
+ predCovs.push(PPred);
+
+ // Innovation covariance S = H P_pred H' + R
+ const S = madd(mmul(mmul(H, PPred), HT), R);
+ innovCovs.push(S);
+
+ // Check if observation has any non-null values
+ const hasObs = yt.some((v) => v !== null);
+
+ if (!hasObs) {
+ // ── Missing observation: skip update ──────────────────────────────
+ innovations.push(new Array(no).fill(Number.NaN));
+ filtMeans.push(xPred);
+ filtCovs.push(PPred);
+ xFilt = xPred;
+ PFilt = PPred;
+ continue;
+ }
+
+ // ── Update ─────────────────────────────────────────────────────────────
+ const yHat = mvmul(H, xPred); // predicted observation
+ const innov = vsub(
+ yt.map((v) => (v === null ? 0 : v)), // treat null as 0 for innovation
+ yHat,
+ );
+
+ // For partial missing (some dims null), we handle by projecting to
+ // observed subspace. For simplicity: use full update with null→predicted.
+ innovations.push(innov);
+
+ // Kalman gain K = P_pred H' S^{-1}
+ const Sinv = matInv(S);
+ if (Sinv === null) {
+ // Singular innovation covariance: skip update
+ filtMeans.push(xPred);
+ filtCovs.push(PPred);
+ xFilt = xPred;
+ PFilt = PPred;
+ continue;
+ }
+
+ const K = mmul(mmul(PPred, HT), Sinv);
+
+ // x_filt = x_pred + K * innov
+ const xNew = vadd(xPred, mvmul(K, innov));
+
+ // P_filt = (I − K H) P_pred — Joseph form for numerical stability:
+ // P_filt = (I−KH) P (I−KH)' + K R K'
+ const IKH = msub(eye(ns), mmul(K, H));
+ const IKHPIKHT = mmul(mmul(IKH, PPred), T(IKH));
+ const KRKT = mmul(mmul(K, R), T(K));
+ const PNew: MutMat = madd(IKHPIKHT, KRKT);
+
+ // Log-likelihood contribution: -½ [d·log(2π) + log|S| + v'S⁻¹v]
+ const logDetS = logDet(S);
+ let vSv = 0;
+ for (let i = 0; i < no; i++) {
+ const Sinv_row = Sinv[i]!;
+ let sSinvRow = 0;
+ for (let j = 0; j < no; j++) {
+ sSinvRow += innov[j]! * Sinv_row[j]!;
+ }
+ vSv += innov[i]! * sSinvRow;
+ }
+ logLik -= 0.5 * (no * LOG2PI + logDetS + vSv);
+
+ filtMeans.push(xNew);
+ filtCovs.push(PNew);
+ xFilt = xNew;
+ PFilt = PNew;
+ }
+
+ return {
+ filteredStateMeans: filtMeans,
+ filteredStateCovariances: filtCovs,
+ predictedStateMeans: predMeans,
+ predictedStateCovariances: predCovs,
+ innovations,
+ innovationCovariances: innovCovs,
+ logLikelihood: logLik,
+ nStates: ns,
+ nObs: no,
+ nTime: T_len,
+ };
+}
+
+/**
+ * Rauch-Tung-Striebel (RTS) smoother backward pass.
+ *
+ * Given the Kalman filter result, runs the smoother backward from t=T to t=0.
+ *
+ * Smoother equations:
+ * G_t = P_{t|t} · Fᵀ · P_{t+1|t}^{-1}
+ * x_{t|T} = x_{t|t} + G_t · (x_{t+1|T} − x_{t+1|t})
+ * P_{t|T} = P_{t|t} + G_t · (P_{t+1|T} − P_{t+1|t}) · G_tᵀ
+ */
+function rtsSmooth(fwd: KalmanFilterResult, F: Mat): KalmanSmootherResult {
+ const T_len = fwd.nTime;
+ const ns = fwd.nStates;
+
+ const smoothMeans: number[][] = new Array(T_len);
+ const smoothCovs: MutMat[] = new Array(T_len);
+ const gains: MutMat[] = new Array(T_len);
+
+ // Initialise last time step from filter
+ const lastFiltMean = [...(fwd.filteredStateMeans[T_len - 1] ?? [])];
+ const lastFiltCov = (fwd.filteredStateCovariances[T_len - 1] ?? []).map((r) => [...r]);
+ smoothMeans[T_len - 1] = lastFiltMean;
+ smoothCovs[T_len - 1] = lastFiltCov;
+ gains[T_len - 1] = zeros(ns, ns);
+
+ const FT = T(F);
+
+ for (let t = T_len - 2; t >= 0; t--) {
+ const xFilt = fwd.filteredStateMeans[t]!;
+ const PFilt = fwd.filteredStateCovariances[t]!;
+ const PPred_next = fwd.predictedStateCovariances[t + 1]!;
+
+ // G_t = P_{t|t} · Fᵀ · P_{t+1|t}^{-1}
+ const PPredInv = matInv(PPred_next);
+ const G: MutMat = PPredInv !== null ? mmul(mmul(PFilt, FT), PPredInv) : zeros(ns, ns);
+
+ const xSmooth_next = smoothMeans[t + 1]!;
+ const PSmooth_next = smoothCovs[t + 1]!;
+ const xPred_next = fwd.predictedStateMeans[t + 1]!;
+
+ // x_{t|T} = x_{t|t} + G_t · (x_{t+1|T} − x_{t+1|t})
+ const dx = vsub(xSmooth_next, xPred_next);
+ smoothMeans[t] = vadd(xFilt, mvmul(G, dx));
+
+ // P_{t|T} = P_{t|t} + G_t · (P_{t+1|T} − P_{t+1|t}) · G_tᵀ
+ const dP = msub(PSmooth_next, PPred_next);
+ const GT = T(G);
+ smoothCovs[t] = madd(PFilt, mmul(mmul(G, dP), GT));
+
+ gains[t] = G;
+ }
+
+ return {
+ smoothedStateMeans: smoothMeans,
+ smoothedStateCovariances: smoothCovs,
+ smootherGains: gains,
+ logLikelihood: fwd.logLikelihood,
+ filterResult: fwd,
+ };
+}
+
+// ─── Standalone functional API ─────────────────────────────────────────────────
+
+/**
+ * Apply the Kalman filter to a sequence of scalar observations.
+ *
+ * Convenience wrapper for the common 1-D case (local-level model or similar).
+ *
+ * @example
+ * ```ts
+ * import { kalmanFilter1D } from "tsb";
+ * const { filteredStateMeans, logLikelihood } = kalmanFilter1D(
+ * [1, 2, null, 3, 2.5],
+ * { processNoise: 0.5, observationNoise: 1 },
+ * );
+ * ```
+ */
+export function kalmanFilter1D(
+ observations: readonly (number | null)[],
+ opts: LocalLevelOptions = {},
+): KalmanFilterResult {
+ const kf = KalmanFilter.localLevel(opts);
+ return kf.filter(observations.map((v) => [v]));
+}
+
+/**
+ * Apply the RTS smoother to scalar observations with a local-level model.
+ *
+ * @example
+ * ```ts
+ * import { kalmanSmooth1D } from "tsb";
+ * const { smoothedStateMeans } = kalmanSmooth1D([1, 2, null, 3, 2.5]);
+ * ```
+ */
+export function kalmanSmooth1D(
+ observations: readonly (number | null)[],
+ opts: LocalLevelOptions = {},
+): KalmanSmootherResult {
+ const kf = KalmanFilter.localLevel(opts);
+ return kf.smooth(observations.map((v) => [v]));
+}
+
+// ─── Utility: extract scalars from 1-D results ─────────────────────────────────
+
+/**
+ * Extract the scalar filtered means from a 1-state filter result.
+ * Returns an array of length T where each value is x_{t|t}[0].
+ *
+ * @example
+ * ```ts
+ * const result = kf.filter([[1], [2], [null], [3]]);
+ * const means = extractScalarMeans(result.filteredStateMeans);
+ * ```
+ */
+export function extractScalarMeans(means: readonly (readonly number[])[]): number[] {
+ return means.map((m) => m[0] ?? Number.NaN);
+}
+
+/**
+ * Extract the scalar filtered variances from a 1-state filter result.
+ * Returns an array of length T where each value is P_{t|t}[0][0].
+ *
+ * @example
+ * ```ts
+ * const result = kf.filter([[1], [2], [null], [3]]);
+ * const vars = extractScalarVariances(result.filteredStateCovariances);
+ * ```
+ */
+export function extractScalarVariances(
+ covs: readonly (readonly (readonly number[])[])[],
+): number[] {
+ return covs.map((P) => P[0]?.[0] ?? Number.NaN);
+}
+
+/**
+ * Compute a 95 % prediction interval around the filtered means for a 1-D
+ * local-level model.
+ *
+ * Returns `{ lower, upper }` arrays of length T.
+ *
+ * @example
+ * ```ts
+ * const result = kf.filter([[1], [2], [null], [3]]);
+ * const { lower, upper } = filteredPredictionInterval(result);
+ * ```
+ */
+export function filteredPredictionInterval(
+ result: KalmanFilterResult,
+ zScore = 1.96,
+): { lower: number[]; upper: number[] } {
+ const means = extractScalarMeans(result.filteredStateMeans);
+ const vars = extractScalarVariances(result.filteredStateCovariances);
+ return {
+ lower: means.map((m, i) => m - zScore * Math.sqrt(vars[i] ?? 0)),
+ upper: means.map((m, i) => m + zScore * Math.sqrt(vars[i] ?? 0)),
+ };
+}
+
+/** Alias kept for backward compat — use {@link KalmanFilter} directly. */
+export { KalmanFilter as StateSpaceModel };
diff --git a/src/stats/signal.ts b/src/stats/signal.ts
new file mode 100644
index 00000000..5c248cd9
--- /dev/null
+++ b/src/stats/signal.ts
@@ -0,0 +1,781 @@
+/**
+ * signal — Signal processing: FFT, windows, STFT, Welch PSD, periodogram.
+ *
+ * Mirrors `numpy.fft`, `scipy.signal` spectral and window utilities.
+ * Implemented from scratch with no external dependencies.
+ *
+ * FFT:
+ * - {@link fft} — N-point complex DFT (radix-2, pads to power of 2)
+ * - {@link ifft} — inverse FFT
+ * - {@link rfft} — real-input FFT (one-sided)
+ * - {@link irfft} — inverse real FFT
+ * - {@link fftFreq} — DFT sample frequencies
+ * - {@link rfftFreq} — one-sided DFT sample frequencies
+ * - {@link fftshift} — shift zero-frequency to centre
+ * - {@link ifftshift} — inverse of fftshift
+ *
+ * Windows (via {@link getWindow}):
+ * - `"rectangular"`, `"bartlett"`, `"hann"`, `"hamming"`, `"blackman"`,
+ * `"blackmanharris"`, `"flattop"`, `"kaiser"`
+ *
+ * Spectral analysis:
+ * - {@link stft} — Short-Time Fourier Transform
+ * - {@link istft} — Inverse STFT (overlap-add)
+ * - {@link welch} — Welch power spectral density
+ * - {@link periodogram} — Periodogram PSD estimate
+ *
+ * @example
+ * ```ts
+ * import { fft, rfftFreq, welch } from "tsb";
+ *
+ * const x = [1, 0, -1, 0, 1, 0, -1, 0];
+ * const X = fft(x); // 8-point FFT
+ * const freqs = rfftFreq(X.length, 1 / 100);
+ *
+ * const { f, Pxx } = welch(x, { fs: 100 });
+ * ```
+ *
+ * @module
+ */
+
+// ─── complex arithmetic ───────────────────────────────────────────────────────
+
+/** A complex number `{ re, im }`. */
+export type Complex = { re: number; im: number };
+
+/** Construct a complex number. */
+export function complex(re: number, im: number): Complex {
+ return { re, im };
+}
+
+/** Add two complex numbers. */
+function cAdd(a: Complex, b: Complex): Complex {
+ return { re: a.re + b.re, im: a.im + b.im };
+}
+
+/** Subtract two complex numbers. */
+function cSub(a: Complex, b: Complex): Complex {
+ return { re: a.re - b.re, im: a.im - b.im };
+}
+
+/** Multiply two complex numbers. */
+function cMul(a: Complex, b: Complex): Complex {
+ return { re: a.re * b.re - a.im * b.im, im: a.re * b.im + a.im * b.re };
+}
+
+/** Complex conjugate. */
+function cConj(a: Complex): Complex {
+ return { re: a.re, im: -a.im };
+}
+
+/** Magnitude squared |a|². */
+function cAbsSq(a: Complex): number {
+ return a.re * a.re + a.im * a.im;
+}
+
+/** Magnitude |a|. */
+export function cAbs(a: Complex): number {
+ return Math.sqrt(cAbsSq(a));
+}
+
+/** Phase angle (arg) of a complex number. */
+export function cArg(a: Complex): number {
+ return Math.atan2(a.im, a.re);
+}
+
+// ─── FFT internals ────────────────────────────────────────────────────────────
+
+/** Smallest power of 2 ≥ n. */
+function nextPow2(n: number): number {
+ if (n <= 1) {
+ return 1;
+ }
+ let p = 1;
+ while (p < n) {
+ p <<= 1;
+ }
+ return p;
+}
+
+/** In-place bit-reversal permutation. */
+function bitReverse(arr: Complex[], n: number): void {
+ let j = 0;
+ for (let i = 1; i < n; i++) {
+ let bit = n >> 1;
+ for (; j & bit; bit >>= 1) {
+ j ^= bit;
+ }
+ j ^= bit;
+ if (i < j) {
+ const tmp = arr[i]!;
+ arr[i] = arr[j]!;
+ arr[j] = tmp;
+ }
+ }
+}
+
+/** Cooley-Tukey radix-2 DIT iterative FFT, in-place. `n` must be a power of 2. */
+// biome-ignore lint/complexity/noExcessiveCognitiveComplexity: nested FFT butterfly loops
+function fftInPlace(arr: Complex[], n: number, inverse: boolean): void {
+ bitReverse(arr, n);
+ for (let len = 2; len <= n; len <<= 1) {
+ const half = len >> 1;
+ const ang = ((inverse ? 2 : -2) * Math.PI) / len;
+ const wLen: Complex = { re: Math.cos(ang), im: Math.sin(ang) };
+ for (let i = 0; i < n; i += len) {
+ let w: Complex = { re: 1, im: 0 };
+ for (let j = 0; j < half; j++) {
+ const u = arr[i + j]!;
+ const v = cMul(arr[i + j + half]!, w);
+ arr[i + j] = cAdd(u, v);
+ arr[i + j + half] = cSub(u, v);
+ w = cMul(w, wLen);
+ }
+ }
+ }
+ if (inverse) {
+ for (let i = 0; i < n; i++) {
+ const a = arr[i]!;
+ arr[i] = { re: a.re / n, im: a.im / n };
+ }
+ }
+}
+
+// ─── public FFT functions ─────────────────────────────────────────────────────
+
+/**
+ * Compute the discrete Fourier transform of a real signal.
+ *
+ * If `x.length` is not a power of 2, the signal is zero-padded to the next
+ * power of 2. The returned array has length `nextPow2(x.length)`.
+ *
+ * @param x - Input signal (real values).
+ * @returns Complex DFT coefficients.
+ */
+export function fft(x: readonly number[]): Complex[] {
+ const n = x.length;
+ const m = nextPow2(n);
+ const buf: Complex[] = Array.from({ length: m }, (_, i) => ({
+ re: x[i] ?? 0,
+ im: 0,
+ }));
+ fftInPlace(buf, m, false);
+ return buf;
+}
+
+/**
+ * Compute the inverse discrete Fourier transform.
+ *
+ * The length of `X` must be a power of 2. Returns a complex array of the
+ * same length as `X`.
+ *
+ * @param X - Complex DFT coefficients.
+ * @returns Complex inverse-DFT output.
+ */
+export function ifft(X: readonly Complex[]): Complex[] {
+ const n = X.length;
+ const m = nextPow2(n);
+ const buf: Complex[] = Array.from({ length: m }, (_, i) => {
+ const c = X[i] ?? { re: 0, im: 0 };
+ return { re: c.re, im: c.im };
+ });
+ fftInPlace(buf, m, true);
+ return buf;
+}
+
+/**
+ * Real-input FFT (one-sided). Returns the first `floor(m/2) + 1` bins where
+ * `m = nextPow2(x.length)`.
+ *
+ * @param x - Real input signal.
+ * @returns One-sided complex spectrum.
+ */
+export function rfft(x: readonly number[]): Complex[] {
+ const full = fft(x);
+ return full.slice(0, Math.floor(full.length / 2) + 1);
+}
+
+/**
+ * Inverse real FFT. Reconstructs a real signal from one-sided spectrum.
+ *
+ * @param X - One-sided spectrum as from {@link rfft}.
+ * @param n - Optional output length (defaults to `2 * (X.length - 1)`).
+ * @returns Real signal.
+ */
+export function irfft(X: readonly Complex[], n?: number): number[] {
+ const nOut = n ?? 2 * (X.length - 1);
+ const m = nextPow2(nOut);
+ const buf: Complex[] = new Array(m);
+ const half = X.length;
+ for (let i = 0; i < half; i++) {
+ buf[i] = { re: (X[i] ?? { re: 0, im: 0 }).re, im: (X[i] ?? { re: 0, im: 0 }).im };
+ }
+ for (let i = half; i < m; i++) {
+ const j = m - i;
+ const c = X[j] ?? { re: 0, im: 0 };
+ buf[i] = cConj(c);
+ }
+ fftInPlace(buf, m, true);
+ return Array.from({ length: nOut }, (_, i) => buf[i]?.re ?? 0);
+}
+
+/**
+ * DFT sample frequencies for an `n`-point FFT with sample spacing `d`.
+ *
+ * @param n - FFT length (from `fft(x).length`).
+ * @param d - Sample spacing in seconds (default `1`).
+ * @returns Array of frequencies from `0` to `(n-1)/(n*d)`, wrapped to negative.
+ */
+export function fftFreq(n: number, d = 1): number[] {
+ const f: number[] = new Array(n);
+ const half = Math.floor(n / 2) + 1;
+ for (let i = 0; i < half; i++) {
+ f[i] = i / (n * d);
+ }
+ for (let i = half; i < n; i++) {
+ f[i] = (i - n) / (n * d);
+ }
+ return f;
+}
+
+/**
+ * One-sided DFT sample frequencies for a real-input FFT.
+ *
+ * @param n - FFT length (from `rfft(x).length` etc.).
+ * @param d - Sample spacing in seconds (default `1`).
+ * @returns Frequencies `[0, 1/(n*d), 2/(n*d), ..., 1/(2*d)]`.
+ */
+export function rfftFreq(n: number, d = 1): number[] {
+ const half = Math.floor(n / 2) + 1;
+ return Array.from({ length: half }, (_, i) => i / (n * d));
+}
+
+/**
+ * Shift the zero-frequency component to the centre of the spectrum.
+ * Equivalent to `numpy.fft.fftshift`.
+ */
+export function fftshift(x: readonly T[]): T[] {
+ const n = x.length;
+ const half = Math.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;
+ const a1 = 0.48829;
+ const a2 = 0.14128;
+ const a3 = 0.01168;
+ return Array.from(
+ { length: n },
+ (_, i) =>
+ a0 -
+ a1 * Math.cos((2 * Math.PI * i) / (n - 1)) +
+ a2 * Math.cos((4 * Math.PI * i) / (n - 1)) -
+ a3 * Math.cos((6 * Math.PI * i) / (n - 1)),
+ );
+}
+
+/** Flat-top window (5-term). */
+export function flatTopWindow(n: number): number[] {
+ const a0 = 0.21557895;
+ const a1 = 0.41663158;
+ const a2 = 0.277263158;
+ const a3 = 0.083578947;
+ const a4 = 0.006947368;
+ return Array.from(
+ { length: n },
+ (_, i) =>
+ a0 -
+ a1 * Math.cos((2 * Math.PI * i) / (n - 1)) +
+ a2 * Math.cos((4 * Math.PI * i) / (n - 1)) -
+ a3 * Math.cos((6 * Math.PI * i) / (n - 1)) +
+ a4 * Math.cos((8 * Math.PI * i) / (n - 1)),
+ );
+}
+
+/**
+ * Kaiser window with shape parameter `beta`.
+ *
+ * @param n - Number of samples.
+ * @param beta - Shape parameter (controls main-lobe width vs side-lobe level).
+ */
+export function kaiserWindow(n: number, beta: number): number[] {
+ const i0b = besselI0(beta);
+ return Array.from({ length: n }, (_, i) => {
+ const t = (2 * i) / (n - 1) - 1;
+ return besselI0(beta * Math.sqrt(1 - t * t)) / i0b;
+ });
+}
+
+/**
+ * Create a window of length `n` by name.
+ *
+ * @param name - Window function name.
+ * @param n - Number of samples.
+ * @param beta - Kaiser window `beta` parameter (ignored for other windows).
+ * @returns - Window samples.
+ */
+export function getWindow(name: WindowName, n: number, beta = 14): number[] {
+ switch (name) {
+ case "rectangular":
+ return rectangularWindow(n);
+ case "bartlett":
+ return bartlettWindow(n);
+ case "hann":
+ return hannWindow(n);
+ case "hamming":
+ return hammingWindow(n);
+ case "blackman":
+ return blackmanWindow(n);
+ case "blackmanharris":
+ return blackmanHarrisWindow(n);
+ case "flattop":
+ return flatTopWindow(n);
+ case "kaiser":
+ return kaiserWindow(n, beta);
+ }
+}
+
+// ─── STFT / ISTFT ─────────────────────────────────────────────────────────────
+
+/** Options for {@link stft}. */
+export interface STFTOptions {
+ /** Sampling frequency in Hz (default `1`). */
+ fs?: number;
+ /** Segment length in samples (default `256`). */
+ nperseg?: number;
+ /** Number of samples to overlap between segments (default `nperseg / 2`). */
+ noverlap?: number;
+ /** FFT length ≥ `nperseg` (default = `nperseg`, padded to power of 2). */
+ nfft?: number;
+ /** Window function to apply (default `"hann"`). */
+ window?: WindowName | readonly number[];
+ /** Boundary extension mode: `"zeros"` pads with zeros, `null` no padding. */
+ boundary?: "zeros" | null;
+}
+
+/** STFT result object. */
+export interface STFTResult {
+ /** Time centres for each frame (seconds). */
+ t: number[];
+ /** Frequency bins (Hz). */
+ f: number[];
+ /** Complex STFT matrix `Zxx[freqIdx][timeIdx]`. */
+ Zxx: Complex[][];
+}
+
+/**
+ * Short-Time Fourier Transform.
+ *
+ * Mirrors `scipy.signal.stft`. Splits the signal into overlapping windowed
+ * segments and computes the FFT of each segment.
+ *
+ * @param x - Input signal.
+ * @param options - {@link STFTOptions}.
+ * @returns - {@link STFTResult} with `{ t, f, Zxx }`.
+ *
+ * @example
+ * ```ts
+ * import { stft } from "tsb";
+ * const x = Array.from({ length: 1024 }, (_, i) => Math.sin(2 * Math.PI * 10 * i / 512));
+ * const { t, f, Zxx } = stft(x, { fs: 512, nperseg: 128 });
+ * ```
+ */
+// biome-ignore lint/complexity/noExcessiveCognitiveComplexity: STFT nested loops
+export function stft(x: readonly number[], options: STFTOptions = {}): STFTResult {
+ const fs = options.fs ?? 1;
+ const nperseg = options.nperseg ?? Math.min(256, x.length);
+ const noverlap = options.noverlap ?? Math.floor(nperseg / 2);
+ const nfft = nextPow2(options.nfft ?? nperseg);
+ const step = nperseg - noverlap;
+ const boundary = options.boundary ?? "zeros";
+
+ // Build window
+ const win: number[] =
+ options.window !== undefined
+ ? typeof options.window === "string"
+ ? getWindow(options.window, nperseg)
+ : Array.from(options.window)
+ : hannWindow(nperseg);
+
+ // Pad signal at boundaries
+ const pad = boundary === "zeros" ? Math.floor(nperseg / 2) : 0;
+ const padded: number[] = [
+ ...Array.from({ length: pad }).fill(0),
+ ...x,
+ ...Array.from({ length: pad }).fill(0),
+ ];
+
+ // Number of frames
+ const nFrames = Math.floor((padded.length - noverlap) / step);
+ const nFreqs = Math.floor(nfft / 2) + 1;
+
+ const Zxx: Complex[][] = Array.from({ length: nFreqs }, () => new Array(nFrames));
+ const tArr: number[] = new Array(nFrames);
+ const fArr: number[] = Array.from({ length: nFreqs }, (_, i) => (i * fs) / nfft);
+
+ for (let k = 0; k < nFrames; k++) {
+ const start = k * step;
+ const seg: Complex[] = Array.from({ length: nfft }, (_, i) => ({
+ re: (padded[start + i] ?? 0) * (win[i] ?? 1),
+ im: 0,
+ }));
+ fftInPlace(seg, nfft, false);
+ for (let fi = 0; fi < nFreqs; fi++) {
+ const col = Zxx[fi];
+ if (col !== undefined) {
+ col[k] = seg[fi] ?? { re: 0, im: 0 };
+ }
+ }
+ tArr[k] = (start + nperseg / 2 - pad) / fs;
+ }
+
+ return { t: tArr, f: fArr, Zxx };
+}
+
+/** Options for {@link istft}. */
+export interface ISTFTOptions {
+ /** Segment length in samples (used to determine overlap). */
+ nperseg?: number;
+ /** Number of samples to overlap between segments (default `nperseg / 2`). */
+ noverlap?: number;
+ /** FFT length (default = `2 * (nFreqs - 1)` where `nFreqs = Zxx.length`). */
+ nfft?: number;
+ /** Window function (default `"hann"`). */
+ window?: WindowName | readonly number[];
+ /** Boundary extension used in stft (default `"zeros"`). */
+ boundary?: "zeros" | null;
+}
+
+/**
+ * Inverse Short-Time Fourier Transform (overlap-add).
+ *
+ * Mirrors `scipy.signal.istft`.
+ *
+ * @param Zxx - Complex STFT matrix `[freqIdx][timeIdx]` (from {@link stft}).
+ * @param options - {@link ISTFTOptions}.
+ * @returns - Reconstructed real signal.
+ */
+// biome-ignore lint/complexity/noExcessiveCognitiveComplexity: ISTFT overlap-add loops
+export function istft(Zxx: readonly (readonly Complex[])[], options: ISTFTOptions = {}): number[] {
+ const nFreqs = Zxx.length;
+ const nFrames = nFreqs > 0 ? (Zxx[0]?.length ?? 0) : 0;
+ const nfft = options.nfft ?? 2 * (nFreqs - 1);
+ const nperseg = options.nperseg ?? nfft;
+ const noverlap = options.noverlap ?? Math.floor(nperseg / 2);
+ const step = nperseg - noverlap;
+
+ const win: number[] =
+ options.window !== undefined
+ ? typeof options.window === "string"
+ ? getWindow(options.window, nperseg)
+ : Array.from(options.window)
+ : hannWindow(nperseg);
+
+ const boundary = options.boundary ?? "zeros";
+ const pad = boundary === "zeros" ? Math.floor(nperseg / 2) : 0;
+ const outLen = nFrames * step + nperseg;
+
+ const output = new Float64Array(outLen);
+ const windowSum = new Float64Array(outLen);
+ const winSq = win.map((w) => w * w);
+
+ for (let k = 0; k < nFrames; k++) {
+ // Build full-spectrum (two-sided) for IFFT
+ const buf: Complex[] = new Array(nfft);
+ for (let fi = 0; fi < nFreqs; fi++) {
+ buf[fi] = Zxx[fi]?.[k] ?? { re: 0, im: 0 };
+ }
+ for (let fi = nFreqs; fi < nfft; fi++) {
+ const mirrorIdx = nfft - fi;
+ const src = Zxx[mirrorIdx]?.[k] ?? { re: 0, im: 0 };
+ buf[fi] = cConj(src);
+ }
+ fftInPlace(buf, nfft, true);
+
+ const start = k * step;
+ for (let i = 0; i < nperseg; i++) {
+ const idx = start + i;
+ output[idx] = (output[idx] ?? 0) + (buf[i]?.re ?? 0) * (win[i] ?? 1);
+ windowSum[idx] = (windowSum[idx] ?? 0) + (winSq[i] ?? 0);
+ }
+ }
+
+ // Normalize and trim boundary padding
+ const result: number[] = [];
+ const start = pad;
+ const end = outLen - pad;
+ for (let i = start; i < end; i++) {
+ const ws = windowSum[i] ?? 0;
+ result.push(ws > 1e-10 ? (output[i] ?? 0) / ws : 0);
+ }
+
+ return result;
+}
+
+// ─── Welch PSD ────────────────────────────────────────────────────────────────
+
+/** Options for {@link welch}. */
+export interface WelchOptions {
+ /** Sampling frequency in Hz (default `1`). */
+ fs?: number;
+ /** Segment length (default `min(256, x.length)`). */
+ nperseg?: number;
+ /** Overlap between segments (default `nperseg / 2`). */
+ noverlap?: number;
+ /** FFT length (default `nperseg`, padded to power of 2). */
+ nfft?: number;
+ /** Window function (default `"hann"`). */
+ window?: WindowName | readonly number[];
+ /** Averaging method: `"mean"` (default) or `"median"`. */
+ average?: "mean" | "median";
+ /** Scaling: `"density"` (PSD, V²/Hz) or `"spectrum"` (power spectrum V²). */
+ scaling?: "density" | "spectrum";
+}
+
+/** PSD result: frequency bins and power spectral density estimates. */
+export interface PSDResult {
+ /** Frequency bins in Hz. */
+ f: number[];
+ /** Power spectral density (or power spectrum) at each frequency. */
+ Pxx: number[];
+}
+
+/**
+ * Welch power spectral density estimate.
+ *
+ * Divides the signal into overlapping segments, computes the periodogram for
+ * each, and averages. Mirrors `scipy.signal.welch`.
+ *
+ * @param x - Input signal.
+ * @param options - {@link WelchOptions}.
+ * @returns - {@link PSDResult} `{ f, Pxx }`.
+ *
+ * @example
+ * ```ts
+ * import { welch } from "tsb";
+ * const x = Array.from({ length: 512 }, (_, i) => Math.sin(2 * Math.PI * 60 * i / 512));
+ * const { f, Pxx } = welch(x, { fs: 512 });
+ * ```
+ */
+// biome-ignore lint/complexity/noExcessiveCognitiveComplexity: Welch averaging loop
+export function welch(x: readonly number[], options: WelchOptions = {}): PSDResult {
+ const fs = options.fs ?? 1;
+ const nperseg = options.nperseg ?? Math.min(256, x.length);
+ const noverlap = options.noverlap ?? Math.floor(nperseg / 2);
+ const nfft = nextPow2(options.nfft ?? nperseg);
+ const step = nperseg - noverlap;
+ const average = options.average ?? "mean";
+ const scaling = options.scaling ?? "density";
+
+ const win: number[] =
+ options.window !== undefined
+ ? typeof options.window === "string"
+ ? getWindow(options.window, nperseg)
+ : Array.from(options.window)
+ : hannWindow(nperseg);
+
+ const winNorm =
+ scaling === "density"
+ ? win.reduce((s, w) => s + w * w, 0) * fs
+ : win.reduce((s, w) => s + w * 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/src/stats/style.ts b/src/stats/style.ts
index 6fe34de8..63538b11 100644
--- a/src/stats/style.ts
+++ b/src/stats/style.ts
@@ -257,7 +257,7 @@ function colormapColor(t: number, cmap: string): string {
const parts = cmap.split(":");
return lerpColor(parts[0] ?? "#ffffff", parts[1] ?? "#000000", t);
}
- const stops = COLORMAPS[cmap] ?? COLORMAPS["Blues"]!;
+ const stops = COLORMAPS[cmap] ?? COLORMAPS.Blues!;
// Find surrounding stops
for (let i = 0; i < stops.length - 1; i++) {
const [p0, c0] = stops[i]!;
diff --git a/src/tseries/frequencies.ts b/src/tseries/frequencies.ts
index f9c96aef..2267b93c 100644
--- a/src/tseries/frequencies.ts
+++ b/src/tseries/frequencies.ts
@@ -324,8 +324,12 @@ export function inferFreq(dates: readonly Date[]): string | null {
const days = first / MS_DAY;
// Year-begin/end dates may have equal diffs when no leap year falls in
// the range; check before returning a raw day-count alias.
- if (_allYearBegin(dates)) return "YS";
- if (_allYearEnd(dates)) return "YE";
+ if (_allYearBegin(dates)) {
+ return "YS";
+ }
+ if (_allYearEnd(dates)) {
+ return "YE";
+ }
return `${days}D`;
}
}
diff --git a/tests-e2e/playground-cells.test.ts b/tests-e2e/playground-cells.test.ts
index a68abadf..17666039 100644
--- a/tests-e2e/playground-cells.test.ts
+++ b/tests-e2e/playground-cells.test.ts
@@ -67,6 +67,8 @@ const NON_PLAYGROUND_PAGES = new Set([
"kde.html",
"multivariate.html",
"sas.html",
+ "hmm.html",
+ "dlm.html",
]);
const PORT = 3399;
diff --git a/tests/io/orc.test.ts b/tests/io/orc.test.ts
new file mode 100644
index 00000000..ae26835b
--- /dev/null
+++ b/tests/io/orc.test.ts
@@ -0,0 +1,393 @@
+/**
+ * Tests for readOrc / toOrc — Apache ORC file format I/O.
+ *
+ * Strategy: use toOrc to produce ORC buffers, then readOrc to round-trip.
+ * All tests operate on in-memory buffers; no filesystem I/O is required.
+ */
+
+import { describe, expect, it } from "bun:test";
+import * as fc from "fast-check";
+import { DataFrame } from "../../src/core/frame.ts";
+import { readOrc, toOrc } from "../../src/io/orc.ts";
+
+// ─── Helpers ──────────────────────────────────────────────────────────────────
+
+function roundtrip(df: DataFrame): DataFrame {
+ return readOrc(toOrc(df));
+}
+
+function colArr(df: DataFrame, name: string): readonly unknown[] {
+ return df.col(name).values;
+}
+
+// ─── File structure ───────────────────────────────────────────────────────────
+
+describe("toOrc — file structure", () => {
+ it("returns a non-empty Uint8Array", () => {
+ const df = DataFrame.fromColumns({ x: [1, 2, 3] });
+ const buf = toOrc(df);
+ expect(buf).toBeInstanceOf(Uint8Array);
+ expect(buf.length).toBeGreaterThan(0);
+ });
+
+ it("starts with ORC magic bytes", () => {
+ const df = DataFrame.fromColumns({ x: [1] });
+ const buf = toOrc(df);
+ expect(buf[0]).toBe(0x4f); // 'O'
+ expect(buf[1]).toBe(0x52); // 'R'
+ expect(buf[2]).toBe(0x43); // 'C'
+ });
+
+ it("ends with postscript length byte", () => {
+ const df = DataFrame.fromColumns({ x: [1] });
+ const buf = toOrc(df);
+ // Last byte is postscript length, must be > 0
+ expect(buf.at(-1)).toBeGreaterThan(0);
+ });
+});
+
+// ─── Integer columns ──────────────────────────────────────────────────────────
+
+describe("readOrc / toOrc — integer columns", () => {
+ it("round-trips a simple int column", () => {
+ const df = DataFrame.fromColumns({ n: [1, 2, 3, 4, 5] });
+ const rt = roundtrip(df);
+ expect(rt.columns.toArray()).toEqual(["n"]);
+ expect(colArr(rt, "n")).toEqual([1, 2, 3, 4, 5]);
+ });
+
+ it("round-trips negative integers", () => {
+ const df = DataFrame.fromColumns({ n: [-100, -1, 0, 1, 100] });
+ const rt = roundtrip(df);
+ expect(colArr(rt, "n")).toEqual([-100, -1, 0, 1, 100]);
+ });
+
+ it("round-trips large integers", () => {
+ const df = DataFrame.fromColumns({ n: [1_000_000, 2_000_000, -999_999] });
+ const rt = roundtrip(df);
+ expect(colArr(rt, "n")).toEqual([1_000_000, 2_000_000, -999_999]);
+ });
+
+ it("round-trips a column of zeros", () => {
+ const df = DataFrame.fromColumns({ n: [0, 0, 0, 0] });
+ const rt = roundtrip(df);
+ expect(colArr(rt, "n")).toEqual([0, 0, 0, 0]);
+ });
+
+ it("round-trips null integers", () => {
+ const df = DataFrame.fromColumns({ n: [1, null, 3, null, 5] });
+ const rt = roundtrip(df);
+ const vals = colArr(rt, "n");
+ expect(vals[0]).toBe(1);
+ expect(vals[1]).toBeNull();
+ expect(vals[2]).toBe(3);
+ expect(vals[3]).toBeNull();
+ expect(vals[4]).toBe(5);
+ });
+
+ it("round-trips all-null int column", () => {
+ const df = DataFrame.fromColumns({ n: [null, null, null] });
+ const rt = roundtrip(df);
+ expect(colArr(rt, "n").every((v) => v === null)).toBe(true);
+ });
+});
+
+// ─── Float/Double columns ─────────────────────────────────────────────────────
+
+describe("readOrc / toOrc — float columns", () => {
+ it("round-trips double values", () => {
+ const df = DataFrame.fromColumns({ x: [1.5, 2.25, 3.75] });
+ const rt = roundtrip(df);
+ const vals = colArr(rt, "x") as number[];
+ expect(vals[0]).toBeCloseTo(1.5);
+ expect(vals[1]).toBeCloseTo(2.25);
+ expect(vals[2]).toBeCloseTo(3.75);
+ });
+
+ it("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..49cb1752
--- /dev/null
+++ b/tests/io/read_avro.test.ts
@@ -0,0 +1,385 @@
+/**
+ * Tests for src/io/read_avro.ts
+ *
+ * Covers readAvro and toAvro (round-trip), schema types, usecols, empty files,
+ * error handling, and fast-check property tests.
+ */
+import { describe, expect, it } from "bun:test";
+import * as fc from "fast-check";
+import { DataFrame } from "../../src/core/frame.ts";
+import { readAvro, toAvro } from "../../src/io/read_avro.ts";
+
+// ─── Helpers: build minimal Avro OCF by hand ─────────────────────────────────
+
+function writeLongTo(arr: number[], v: number): void {
+ let n = (v << 1) ^ (v >> 31);
+ while (n & ~0x7f) {
+ arr.push((n & 0x7f) | 0x80);
+ n >>>= 7;
+ }
+ arr.push(n);
+}
+
+function writeStringTo(arr: number[], s: string): void {
+ const b = new TextEncoder().encode(s);
+ writeLongTo(arr, b.length);
+ for (const byte of b) {
+ arr.push(byte);
+ }
+}
+
+function writeBytesTo(arr: number[], b: Uint8Array): void {
+ writeLongTo(arr, b.length);
+ for (const byte of b) {
+ arr.push(byte);
+ }
+}
+
+function buildAvroOCF(schema: object, rows: Record[]): Uint8Array {
+ const schemaBytes = new TextEncoder().encode(JSON.stringify(schema));
+ const sync = new Uint8Array([1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16]);
+ const buf: number[] = [];
+
+ // Magic
+ buf.push(79, 98, 106, 1); // "Obj\x01"
+
+ // Metadata: avro.schema + avro.codec
+ writeLongTo(buf, 2);
+ writeStringTo(buf, "avro.schema");
+ writeBytesTo(buf, schemaBytes);
+ writeStringTo(buf, "avro.codec");
+ writeBytesTo(buf, new TextEncoder().encode("null"));
+ writeLongTo(buf, 0);
+
+ // Sync marker
+ for (const b of sync) {
+ buf.push(b);
+ }
+
+ // Encode rows
+ const rowBuf: number[] = [];
+ for (const row of rows) {
+ for (const field of (schema as { fields: { name: string; type: unknown }[] }).fields) {
+ const v = row[field.name];
+ const t = field.type;
+ if (t === "null") {
+ // nothing
+ } else if (t === "boolean") {
+ rowBuf.push(v ? 1 : 0);
+ } else if (t === "int" || t === "long") {
+ writeLongTo(rowBuf, typeof v === "number" ? v : 0);
+ } else if (t === "double") {
+ const arr = new Float64Array(1);
+ arr[0] = typeof v === "number" ? v : 0;
+ for (const b of new Uint8Array(arr.buffer)) {
+ rowBuf.push(b);
+ }
+ } else if (t === "string") {
+ writeStringTo(rowBuf, String(v ?? ""));
+ } else if (Array.isArray(t)) {
+ // union ["null", X]
+ if (v === null || v === undefined) {
+ writeLongTo(rowBuf, 0);
+ } else {
+ writeLongTo(rowBuf, 1);
+ const inner = t[1] as string;
+ if (inner === "string") {
+ writeStringTo(rowBuf, String(v));
+ } else if (inner === "long" || inner === "int") {
+ writeLongTo(rowBuf, Number(v));
+ } else if (inner === "double") {
+ const a2 = new Float64Array(1);
+ a2[0] = Number(v);
+ for (const b of new Uint8Array(a2.buffer)) {
+ rowBuf.push(b);
+ }
+ }
+ }
+ }
+ }
+ }
+
+ // Data block: count, byteCount, data, sync
+ if (rowBuf.length > 0) {
+ writeLongTo(buf, rows.length);
+ writeLongTo(buf, rowBuf.length);
+ for (const b of rowBuf) {
+ buf.push(b);
+ }
+ for (const b of sync) {
+ buf.push(b);
+ }
+ }
+
+ return new Uint8Array(buf);
+}
+
+// ─── readAvro – basic parsing ──────────────────────────────────────────────────
+
+describe("readAvro – basic types", () => {
+ it("reads an int column", () => {
+ const schema = { type: "record", name: "R", fields: [{ name: "n", type: "int" }] };
+ const buf = buildAvroOCF(schema, [{ n: 1 }, { n: -2 }, { n: 42 }]);
+ const df = readAvro(buf);
+ expect(df.shape[0]).toBe(3);
+ expect(df.col("n").at(0)).toBe(1);
+ expect(df.col("n").at(1)).toBe(-2);
+ expect(df.col("n").at(2)).toBe(42);
+ });
+
+ it("reads a double column", () => {
+ const schema = { type: "record", name: "R", fields: [{ name: "v", type: "double" }] };
+ const buf = buildAvroOCF(schema, [{ v: 3.14 }, { v: -1.5 }]);
+ const df = readAvro(buf);
+ expect(df.col("v").at(0) as number).toBeCloseTo(3.14, 5);
+ expect(df.col("v").at(1) as number).toBeCloseTo(-1.5, 5);
+ });
+
+ it("reads a string column", () => {
+ const schema = { type: "record", name: "R", fields: [{ name: "s", type: "string" }] };
+ const buf = buildAvroOCF(schema, [{ s: "hello" }, { s: "world" }]);
+ const df = readAvro(buf);
+ expect(df.col("s").at(0)).toBe("hello");
+ expect(df.col("s").at(1)).toBe("world");
+ });
+
+ it("reads a boolean column", () => {
+ const schema = { type: "record", name: "R", fields: [{ name: "b", type: "boolean" }] };
+ const buf = buildAvroOCF(schema, [{ b: true }, { b: false }]);
+ const df = readAvro(buf);
+ expect(df.col("b").at(0)).toBe(true);
+ expect(df.col("b").at(1)).toBe(false);
+ });
+
+ it("reads nullable (union) columns with null values", () => {
+ const schema = {
+ type: "record",
+ name: "R",
+ fields: [{ name: "x", type: ["null", "string"] }],
+ };
+ const buf = buildAvroOCF(schema, [{ x: "foo" }, { x: null }, { x: "bar" }]);
+ const df = readAvro(buf);
+ expect(df.col("x").at(0)).toBe("foo");
+ expect(df.col("x").at(1)).toBeNull();
+ expect(df.col("x").at(2)).toBe("bar");
+ });
+
+ it("reads multiple columns of mixed types", () => {
+ const schema = {
+ type: "record",
+ name: "R",
+ fields: [
+ { name: "id", type: "int" },
+ { name: "name", type: "string" },
+ { name: "val", type: "double" },
+ ],
+ };
+ const rows = [
+ { id: 1, name: "alice", val: 1.1 },
+ { id: 2, name: "bob", val: 2.2 },
+ ];
+ const df = readAvro(buildAvroOCF(schema, rows));
+ expect(df.shape).toEqual([2, 3]);
+ expect(df.col("id").at(0)).toBe(1);
+ expect(df.col("name").at(1)).toBe("bob");
+ expect(df.col("val").at(1) as number).toBeCloseTo(2.2, 5);
+ });
+});
+
+describe("readAvro – usecols", () => {
+ it("returns only requested columns", () => {
+ const schema = {
+ type: "record",
+ name: "R",
+ fields: [
+ { name: "a", type: "int" },
+ { name: "b", type: "string" },
+ { name: "c", type: "double" },
+ ],
+ };
+ const rows = [
+ { a: 1, b: "x", c: 0.5 },
+ { a: 2, b: "y", c: 1.5 },
+ ];
+ const df = readAvro(buildAvroOCF(schema, rows), { usecols: ["a", "c"] });
+ expect([...df.columns.values]).toEqual(["a", "c"]);
+ expect(df.shape[1]).toBe(2);
+ });
+});
+
+describe("readAvro – error handling", () => {
+ it("throws on bad magic bytes", () => {
+ const bad = new Uint8Array(20);
+ bad.fill(0);
+ expect(() => readAvro(bad)).toThrow();
+ });
+
+ it("accepts ArrayBuffer input", () => {
+ const schema = { type: "record", name: "R", fields: [{ name: "n", type: "int" }] };
+ const buf = buildAvroOCF(schema, [{ n: 7 }]);
+ const df = readAvro(buf.buffer as ArrayBuffer);
+ expect(df.col("n").at(0)).toBe(7);
+ });
+
+ it("throws for unsupported codec", () => {
+ // Build a fake header with codec=deflate
+ const schema = { type: "record", name: "R", fields: [{ name: "n", type: "int" }] };
+ const schemaBytes = new TextEncoder().encode(JSON.stringify(schema));
+ const buf: number[] = [79, 98, 106, 1]; // magic
+ writeLongTo(buf, 2);
+ writeStringTo(buf, "avro.schema");
+ writeBytesTo(buf, schemaBytes);
+ writeStringTo(buf, "avro.codec");
+ writeBytesTo(buf, new TextEncoder().encode("deflate"));
+ writeLongTo(buf, 0);
+ for (let i = 0; i < 16; i++) {
+ buf.push(i); // sync
+ }
+ // No data blocks
+ expect(() => readAvro(new Uint8Array(buf))).toThrow(/deflate/);
+ });
+});
+
+describe("readAvro – empty file", () => {
+ it("returns empty DataFrame for file with no rows", () => {
+ const schema = { type: "record", name: "R", fields: [{ name: "n", type: "int" }] };
+ const df = readAvro(buildAvroOCF(schema, []));
+ expect(df.shape[0]).toBe(0);
+ });
+});
+
+// ─── toAvro / round-trip ──────────────────────────────────────────────────────
+
+describe("toAvro – file structure", () => {
+ it("starts with Avro magic bytes", () => {
+ const df = DataFrame.fromColumns({ a: [1, 2, 3] });
+ const buf = toAvro(df);
+ expect(buf[0]).toBe(79); // 'O'
+ expect(buf[1]).toBe(98); // 'b'
+ expect(buf[2]).toBe(106); // 'j'
+ expect(buf[3]).toBe(1); // version
+ });
+
+ it("produces a Uint8Array", () => {
+ const df = DataFrame.fromColumns({ x: [1.1, 2.2] });
+ expect(toAvro(df)).toBeInstanceOf(Uint8Array);
+ });
+});
+
+describe("toAvro + readAvro – round-trip", () => {
+ it("integer column round-trips", () => {
+ const df = DataFrame.fromColumns({ id: [1, 2, 3, 4, 5] });
+ const buf = toAvro(df);
+ const df2 = readAvro(buf);
+ expect(df2.shape[0]).toBe(5);
+ for (let i = 0; i < 5; i++) {
+ expect(df2.col("id").at(i)).toBe(i + 1);
+ }
+ });
+
+ it("double column round-trips", () => {
+ const df = DataFrame.fromColumns({ v: [1.5, 2.5, 3.5] });
+ const buf = toAvro(df);
+ const df2 = readAvro(buf);
+ expect(df2.col("v").at(0) as number).toBeCloseTo(1.5, 5);
+ expect(df2.col("v").at(2) as number).toBeCloseTo(3.5, 5);
+ });
+
+ it("string column round-trips", () => {
+ const df = DataFrame.fromColumns({ name: ["alice", "bob", "carol"] });
+ const buf = toAvro(df);
+ const df2 = readAvro(buf);
+ expect(df2.col("name").at(1)).toBe("bob");
+ });
+
+ it("boolean column round-trips", () => {
+ const df = DataFrame.fromColumns({ flag: [true, false, true] });
+ const buf = toAvro(df);
+ const df2 = readAvro(buf);
+ expect(df2.col("flag").at(0)).toBe(true);
+ expect(df2.col("flag").at(1)).toBe(false);
+ });
+
+ it("null values round-trip in nullable columns", () => {
+ const df = DataFrame.fromColumns({ x: [1, null, 3] });
+ const buf = toAvro(df);
+ const df2 = readAvro(buf);
+ expect(df2.col("x").at(0)).toBe(1);
+ expect(df2.col("x").at(1)).toBeNull();
+ expect(df2.col("x").at(2)).toBe(3);
+ });
+
+ it("multi-column mixed-type round-trip", () => {
+ const df = DataFrame.fromColumns({
+ id: [1, 2, 3],
+ name: ["a", "b", "c"],
+ score: [0.1, 0.2, 0.3],
+ active: [true, false, true],
+ });
+ const buf = toAvro(df);
+ const df2 = readAvro(buf);
+ expect(df2.shape).toEqual([3, 4]);
+ expect(df2.col("name").at(1)).toBe("b");
+ expect(df2.col("score").at(2) as number).toBeCloseTo(0.3, 5);
+ });
+
+ it("empty DataFrame round-trips", () => {
+ const df = DataFrame.fromColumns({ a: [] as number[] });
+ const buf = toAvro(df);
+ const df2 = readAvro(buf);
+ expect(df2.shape[0]).toBe(0);
+ });
+});
+
+// ─── Property-based tests ──────────────────────────────────────────────────────
+
+describe("property tests", () => {
+ it("integer round-trip: toAvro → readAvro preserves integer values", () => {
+ fc.assert(
+ fc.property(
+ fc.array(fc.integer({ min: -1000, max: 1000 }), { minLength: 1, maxLength: 20 }),
+ (vals) => {
+ const df = DataFrame.fromColumns({ n: vals });
+ const df2 = readAvro(toAvro(df));
+ for (let i = 0; i < vals.length; i++) {
+ if (df2.col("n").at(i) !== vals[i]) {
+ return false;
+ }
+ }
+ return true;
+ },
+ ),
+ );
+ });
+
+ it("string round-trip preserves values", () => {
+ fc.assert(
+ fc.property(
+ fc.array(fc.string({ minLength: 0, maxLength: 20 }), { minLength: 1, maxLength: 15 }),
+ (vals) => {
+ const df = DataFrame.fromColumns({ s: vals });
+ const df2 = readAvro(toAvro(df));
+ for (let i = 0; i < vals.length; i++) {
+ if (df2.col("s").at(i) !== vals[i]) {
+ return false;
+ }
+ }
+ return true;
+ },
+ ),
+ );
+ });
+
+ it("row count is preserved in round-trip", () => {
+ fc.assert(
+ fc.property(
+ fc.array(fc.integer({ min: 0, max: 100 }), { minLength: 0, maxLength: 30 }),
+ (vals) => {
+ const df = DataFrame.fromColumns({ n: vals });
+ const df2 = readAvro(toAvro(df));
+ return df2.shape[0] === vals.length;
+ },
+ ),
+ );
+ });
+});
diff --git a/tests/io/read_html.test.ts b/tests/io/read_html.test.ts
index dfc41d7f..8a65f076 100644
--- a/tests/io/read_html.test.ts
+++ b/tests/io/read_html.test.ts
@@ -55,7 +55,7 @@ describe("readHtml – basic", () => {
test("values are numeric by default", () => {
const html = simpleTable(["n"], [["42"], ["3.14"]]);
const [df] = readHtml(html);
- const vals = df!.col("n").toArray();
+ const vals = df?.col("n").toArray();
expect(vals[0]).toBe(42);
expect(vals[1]).toBeCloseTo(3.14);
});
@@ -87,7 +87,7 @@ describe("readHtml – header", () => {
1 2 3 hello world ({ data: [1, 2, 3, 4, 3, 2, 1, 2, 3, 4, 3, 2, 1, 2, 3, 4] });
+ const model4 = fitArima(s, { p: 1, d: 0, q: 0 });
+ expect(model4.p).toBe(1);
+ expect(model4.forecast(2).forecast.length).toBe(2);
+ });
+});
+
+// ─── Property-based tests ───────────────────────────────────────────────────────
+
+describe("property tests", () => {
+ it("fitted value count always equals input length", () => {
+ fc.assert(
+ fc.property(
+ fc.array(fc.float({ noNaN: true, noDefaultInfinity: true, min: -100, max: 100 }), {
+ minLength: 20,
+ maxLength: 60,
+ }),
+ fc.integer({ min: 0, max: 2 }),
+ fc.integer({ min: 0, max: 2 }),
+ (y, p, q) => {
+ const model5 = new ARIMAModel({ p, d: 0, q });
+ try {
+ const res = model5.fit(y);
+ return res.fittedValues.length === y.length;
+ } catch {
+ return true; // short series allowed to throw
+ }
+ },
+ ),
+ );
+ });
+
+ it("forecast intervals always satisfy lower <= forecast <= upper", () => {
+ fc.assert(
+ fc.property(
+ fc.array(fc.float({ noNaN: true, noDefaultInfinity: true, min: -50, max: 50 }), {
+ minLength: 20,
+ maxLength: 50,
+ }),
+ (y) => {
+ const model6 = new ARIMAModel({ p: 1, d: 0, q: 0 });
+ try {
+ model6.fit(y);
+ const fc2 = model6.forecast(3);
+ for (let h = 0; h < 3; h++) {
+ if ((fc2.lower[h] ?? 0) > (fc2.forecast[h] ?? 0) + 1e-6) {
+ return false;
+ }
+ if ((fc2.upper[h] ?? 0) < (fc2.forecast[h] ?? 0) - 1e-6) {
+ return false;
+ }
+ }
+ return true;
+ } catch {
+ return true;
+ }
+ },
+ ),
+ );
+ });
+
+ it("sigma2 is always positive after fit", () => {
+ fc.assert(
+ fc.property(
+ fc.array(fc.float({ noNaN: true, noDefaultInfinity: true, min: -100, max: 100 }), {
+ minLength: 15,
+ maxLength: 40,
+ }),
+ (y) => {
+ const model7 = new ARIMAModel({ p: 1, d: 0, q: 0 });
+ try {
+ const res = model7.fit(y);
+ return res.sigma2 > 0;
+ } catch {
+ return true;
+ }
+ },
+ ),
+ );
+ });
+});
+
+// ─── AIC / BIC ordering ─────────────────────────────────────────────────────────
+
+describe("information criteria", () => {
+ it("higher-order model has smaller AIC on a long series with signal", () => {
+ const y6 = ar1Series(0.8, 300, 0.1);
+ const m1 = new ARIMAModel({ p: 1, d: 0, q: 0 }).fit(y6);
+ const m2 = new ARIMAModel({ p: 2, d: 0, q: 0 }).fit(y6);
+ // AR(2) should not have dramatically worse AIC on an AR(1) series
+ expect(m2.aic).not.toBeNaN();
+ expect(m1.aic).not.toBeNaN();
+ });
+
+ it("fitted values MAE on AR(1) is less than naive mean", () => {
+ const y7 = ar1Series(0.85, 200, 0.05);
+ const result = new ARIMAModel({ p: 1, d: 0, q: 0 }).fit(y7);
+ const mean = y7.reduce((s, v) => s + v, 0) / y7.length;
+ const naiveMae = mae(y7, new Array(y7.length).fill(mean));
+ const modelMae = mae(y7, result.fittedValues);
+ expect(modelMae).toBeLessThan(naiveMae);
+ });
+});
diff --git a/tests/stats/dlm.test.ts b/tests/stats/dlm.test.ts
new file mode 100644
index 00000000..dd859bab
--- /dev/null
+++ b/tests/stats/dlm.test.ts
@@ -0,0 +1,446 @@
+/**
+ * Tests for dlm.ts — Dynamic Linear Model (State-Space).
+ */
+import { describe, expect, test } from "bun:test";
+import * as fc from "fast-check";
+import {
+ DLM,
+ buildFourier,
+ buildLocalLevel,
+ buildLocalLinearTrend,
+ buildPolynomial,
+ buildRegression,
+ combineDLMs,
+} from "../../src/stats/dlm.ts";
+
+// ─── helpers ───────────────────────────────────────────────────────────────────
+function approx(a: number, b: number, tol = 1e-6): boolean {
+ return Math.abs(a - b) <= tol * (1 + Math.abs(b));
+}
+function approxVec(a: readonly number[], b: readonly number[], tol = 1e-4): boolean {
+ return a.length === b.length && a.every((v, i) => approx(v, b[i]!, tol));
+}
+
+// ─── Local-level model ─────────────────────────────────────────────────────────
+describe("DLM — local-level", () => {
+ const y = [1, 2, 3, 2, 3, 4, 3, 4, 5, 4];
+
+ test("filter returns correct dimensions", () => {
+ const dlm = DLM.localLevel({ sigmaObs: 1, sigmaLevel: 0.5 });
+ const res = dlm.filter(y);
+ expect(res.steps.length).toBe(y.length);
+ expect(res.filteredMeans.length).toBe(y.length);
+ expect(res.filteredMeans[0]?.length).toBe(1);
+ expect(res.forecastMeans.length).toBe(y.length);
+ });
+
+ test("log-likelihood is finite and negative", () => {
+ const dlm = DLM.localLevel({ sigmaObs: 1, sigmaLevel: 0.5 });
+ const res = dlm.filter(y);
+ expect(Number.isFinite(res.logLikelihood)).toBe(true);
+ expect(res.logLikelihood).toBeLessThan(0);
+ });
+
+ test("filtered means track the signal", () => {
+ const dlm = DLM.localLevel({ sigmaObs: 1, sigmaLevel: 0.5 });
+ const res = dlm.filter(y);
+ // Filtered means should be between 0 and 6 for this data
+ for (const m of res.filteredMeans) {
+ expect(m[0]).toBeGreaterThan(0);
+ expect(m[0]).toBeLessThan(6);
+ }
+ });
+
+ test("filter with missing observations", () => {
+ const dlm = DLM.localLevel({ sigmaObs: 1, sigmaLevel: 0.5 });
+ const yMissing = [1, null, 3, null, 5];
+ const res = dlm.filter(yMissing);
+ expect(res.steps.length).toBe(5);
+ // Steps with null observations have null innovation
+ expect(res.steps[1]?.innovation).toBeNull();
+ expect(res.steps[3]?.innovation).toBeNull();
+ // Steps with observations have non-null innovation
+ expect(res.steps[0]?.innovation).not.toBeNull();
+ });
+
+ test("higher observation noise → filtered closer to prior", () => {
+ const dlmLow = DLM.localLevel({ sigmaObs: 0.01, sigmaLevel: 1 });
+ const dlmHigh = DLM.localLevel({ sigmaObs: 100, sigmaLevel: 1 });
+ const resLow = dlmLow.filter(y);
+ const resHigh = dlmHigh.filter(y);
+ // With low obs noise, filtered mean ≈ observation
+ expect(Math.abs(resLow.filteredMeans[0]?.[0]! - y[0]!)).toBeLessThan(0.1);
+ // With high obs noise, filtered mean stays near prior (zero init, drifts slowly)
+ // Just check the values are different
+ expect(resLow.filteredMeans[5]?.[0]).not.toBeCloseTo(resHigh.filteredMeans[5]?.[0]!, 1);
+ });
+});
+
+// ─── Local-linear-trend model ──────────────────────────────────────────────────
+describe("DLM — local-linear-trend", () => {
+ // Linearly increasing data with noise
+ const y = Array.from({ length: 20 }, (_, t) => t + 0.5 * (Math.sin(t) * 0.1));
+
+ test("filter dimensions", () => {
+ const dlm = DLM.localLinearTrend({ sigmaObs: 0.5, sigmaLevel: 0.1, sigmaSlope: 0.01 });
+ const res = dlm.filter(y);
+ expect(res.filteredMeans.length).toBe(y.length);
+ expect(res.filteredMeans[0]?.length).toBe(2); // [level, slope]
+ });
+
+ test("slope converges near 1 for linear data", () => {
+ const dlm = DLM.localLinearTrend({ sigmaObs: 0.1, sigmaLevel: 0.01, sigmaSlope: 0.01 });
+ const res = dlm.filter(y);
+ const lastSlope = res.filteredMeans[y.length - 1]?.[1]!;
+ expect(lastSlope).toBeGreaterThan(0.5);
+ expect(lastSlope).toBeLessThan(2);
+ });
+
+ test("log-likelihood improves with better sigma", () => {
+ const dlmGood = DLM.localLinearTrend({ sigmaObs: 0.5, sigmaLevel: 0.1, sigmaSlope: 0.01 });
+ const dlmBad = DLM.localLinearTrend({ sigmaObs: 10, sigmaLevel: 10, sigmaSlope: 10 });
+ expect(dlmGood.filter(y).logLikelihood).toBeGreaterThan(dlmBad.filter(y).logLikelihood);
+ });
+});
+
+// ─── Smoother ──────────────────────────────────────────────────────────────────
+describe("DLM — RTS smoother", () => {
+ const y = [1, 2, 1.5, 3, 2.5, 4, 3.5, 5, 4.5, 6];
+
+ test("smoothed means have same length as data", () => {
+ const dlm = DLM.localLevel({ sigmaObs: 1, sigmaLevel: 0.5 });
+ const res = dlm.smooth(y);
+ expect(res.smoothedMeans.length).toBe(y.length);
+ expect(res.smoothedCovs.length).toBe(y.length);
+ });
+
+ test("smoother log-likelihood equals filter log-likelihood", () => {
+ const dlm = DLM.localLevel({ sigmaObs: 1, sigmaLevel: 0.5 });
+ const filterRes = dlm.filter(y);
+ const smoothRes = dlm.smooth(y);
+ expect(smoothRes.logLikelihood).toBeCloseTo(filterRes.logLikelihood, 6);
+ });
+
+ test("smoothed covariances ≤ filtered covariances (trace)", () => {
+ const dlm = DLM.localLinearTrend({ sigmaObs: 1, sigmaLevel: 0.5, sigmaSlope: 0.1 });
+ const res = dlm.smooth(y);
+ for (let t = 0; t < y.length - 1; t++) {
+ const filtTrace = res.filteredCovs[t]?.[0]?.[0]! + (res.filteredCovs[t]?.[1]?.[1] ?? 0);
+ const smTrace = res.smoothedCovs[t]?.[0]?.[0]! + (res.smoothedCovs[t]?.[1]?.[1] ?? 0);
+ // Smoother should not increase uncertainty
+ expect(smTrace).toBeLessThanOrEqual(filtTrace + 1e-6);
+ }
+ });
+
+ test("last smoothed = last filtered", () => {
+ const dlm = DLM.localLevel({ sigmaObs: 1, sigmaLevel: 0.5 });
+ const res = dlm.smooth(y);
+ const T = y.length;
+ expect(res.smoothedMeans[T - 1]?.[0]).toBeCloseTo(res.filteredMeans[T - 1]?.[0]!, 6);
+ });
+});
+
+// ─── Forecasting ───────────────────────────────────────────────────────────────
+describe("DLM — forecasting", () => {
+ const y = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10];
+
+ test("forecast returns h steps", () => {
+ const dlm = DLM.localLevel({ sigmaObs: 0.5, sigmaLevel: 0.2 });
+ const res = dlm.filter(y);
+ const fc = dlm.forecast(res, 5);
+ expect(fc.mean.length).toBe(5);
+ expect(fc.lower.length).toBe(5);
+ expect(fc.upper.length).toBe(5);
+ });
+
+ test("forecast mean > lower, < upper", () => {
+ const dlm = DLM.localLevel({ sigmaObs: 0.5, sigmaLevel: 0.2 });
+ const res = dlm.filter(y);
+ const fc = dlm.forecast(res, 5);
+ for (let i = 0; i < 5; i++) {
+ expect(fc.mean[i]?.[0]).toBeGreaterThan(fc.lower[i]!);
+ expect(fc.mean[i]?.[0]).toBeLessThan(fc.upper[i]!);
+ }
+ });
+
+ test("prediction intervals widen over forecast horizon", () => {
+ const dlm = DLM.localLevel({ sigmaObs: 0.5, sigmaLevel: 0.2 });
+ const res = dlm.filter(y);
+ const fc = dlm.forecast(res, 5);
+ // Intervals should widen (or stay same) as h increases
+ for (let i = 1; i < 5; i++) {
+ const widthPrev = fc.upper[i - 1]! - fc.lower[i - 1]!;
+ const widthCurr = fc.upper[i]! - fc.lower[i]!;
+ expect(widthCurr).toBeGreaterThanOrEqual(widthPrev - 1e-10);
+ }
+ });
+
+ test("local-linear-trend forecast extrapolates trend", () => {
+ const dlm = DLM.localLinearTrend({ sigmaObs: 0.1, sigmaLevel: 0.01, sigmaSlope: 0.001 });
+ const res = dlm.filter(y);
+ const fc = dlm.forecast(res, 3);
+ // Should forecast above 10
+ expect(fc.mean[0]?.[0]).toBeGreaterThan(9);
+ expect(fc.mean[2]?.[0]).toBeGreaterThan(fc.mean[0]?.[0]!);
+ });
+});
+
+// ─── Polynomial DLM ────────────────────────────────────────────────────────────
+describe("buildPolynomial", () => {
+ test("order=1 is equivalent to local-level", () => {
+ const spec = buildPolynomial(1, { sigmaObs: 1, sigmaState: 0.5 });
+ expect(spec.G).toEqual([[1]]);
+ expect(spec.F).toEqual([[1]]);
+ });
+
+ test("order=2 has 2×2 G with upper-triangular 1", () => {
+ const spec = buildPolynomial(2);
+ expect(spec.G[0]).toEqual([1, 1]);
+ expect(spec.G[1]).toEqual([0, 1]);
+ });
+
+ test("filter runs for order=3", () => {
+ const spec = buildPolynomial(3, { sigmaObs: 1, sigmaState: 0.1 });
+ const dlm = new DLM(spec);
+ const res = dlm.filter([1, 2, 3, 4, 5]);
+ expect(res.filteredMeans.length).toBe(5);
+ expect(res.filteredMeans[0]?.length).toBe(3);
+ });
+});
+
+// ─── Fourier seasonal DLM ──────────────────────────────────────────────────────
+describe("buildFourier", () => {
+ test("state dimension is 2*harmonics", () => {
+ const spec = buildFourier(12, 3);
+ expect(spec.G.length).toBe(6);
+ expect(spec.F[0]?.length).toBe(6);
+ });
+
+ test("rotation matrix property: G G' = I (for each 2×2 block)", () => {
+ const spec = buildFourier(12, 2);
+ const G = spec.G;
+ // Check first block is rotation
+ const c = G[0]?.[0]!;
+ const s = G[1]?.[0]!;
+ expect(c * c + s * s).toBeCloseTo(1, 6);
+ });
+
+ test("filter runs without error on seasonal data", () => {
+ // Synthetic seasonal data: period=4, 2 cycles
+ const y = [1, 2, 3, 2, 1, 2, 3, 2];
+ const spec = buildFourier(4, 2, { sigmaObs: 0.5, sigmaState: 0.1 });
+ const dlm = new DLM(spec);
+ const res = dlm.filter(y);
+ expect(res.steps.length).toBe(8);
+ expect(Number.isFinite(res.logLikelihood)).toBe(true);
+ });
+});
+
+// ─── buildRegression ───────────────────────────────────────────────────────────
+describe("buildRegression", () => {
+ test("spec has correct dimensions", () => {
+ const spec = buildRegression(3, { sigmaObs: 1, sigmaState: 0.001 });
+ expect(spec.G.length).toBe(3);
+ expect(spec.G[0]?.length).toBe(3);
+ expect(spec.F[0]?.length).toBe(3);
+ });
+});
+
+// ─── combineDLMs ───────────────────────────────────────────────────────────────
+describe("combineDLMs", () => {
+ test("combines state dimensions", () => {
+ const ll = buildLocalLevel({ sigmaObs: 1, sigmaLevel: 0.5 });
+ const llt = buildLocalLinearTrend({ sigmaObs: 1, sigmaLevel: 0.5, sigmaSlope: 0.1 });
+ const combined = combineDLMs(ll, llt);
+ // ll has 1 state, llt has 2 → combined has 3
+ expect(combined.G.length).toBe(3);
+ expect(combined.F[0]?.length).toBe(3);
+ expect(combined.W.length).toBe(3);
+ });
+
+ test("single spec is identity", () => {
+ const ll = buildLocalLevel({ sigmaObs: 1, sigmaLevel: 0.5 });
+ const combined = combineDLMs(ll);
+ expect(combined.G).toEqual(ll.G);
+ });
+
+ test("combined DLM can be filtered", () => {
+ const ll = buildLocalLevel({ sigmaObs: 1, sigmaLevel: 0.5 });
+ const fourier = buildFourier(4, 1, { sigmaObs: 1, sigmaState: 0.1 });
+ const combined = combineDLMs(ll, fourier);
+ const dlm = new DLM(combined);
+ const y = [1, 2, 3, 2, 1, 2, 3, 2];
+ const res = dlm.filter(y);
+ expect(res.steps.length).toBe(8);
+ expect(Number.isFinite(res.logLikelihood)).toBe(true);
+ });
+
+ test("throws on empty input", () => {
+ expect(() => combineDLMs()).toThrow();
+ });
+});
+
+// ─── MLE fitting ──────────────────────────────────────────────────────────────
+describe("DLM.fitMLE", () => {
+ test("returns a DLM instance", () => {
+ const dlm = DLM.localLevel({ sigmaObs: 1, sigmaLevel: 1 });
+ const y = [1, 2, 1.5, 2.5, 3, 2, 3.5, 3, 4, 3.5];
+ const fitted = dlm.fitMLE(y);
+ expect(fitted).toBeInstanceOf(DLM);
+ });
+
+ test("fitted log-likelihood ≥ initial log-likelihood", () => {
+ const dlm = DLM.localLevel({ sigmaObs: 5, sigmaLevel: 5 });
+ const y = [1, 2, 1.5, 2.5, 3, 2, 3.5, 3, 4, 3.5];
+ const fitted = dlm.fitMLE(y);
+ const llInitial = dlm.filter(y).logLikelihood;
+ const llFitted = fitted.filter(y).logLikelihood;
+ expect(llFitted).toBeGreaterThanOrEqual(llInitial - 0.1);
+ });
+});
+
+// ─── Discount factor filter ───────────────────────────────────────────────────
+describe("DLM.filterDiscount", () => {
+ test("discount=1 matches standard filter closely", () => {
+ const dlm = DLM.localLevel({ sigmaObs: 1, sigmaLevel: 0.5 });
+ const y = [1, 2, 3, 4, 5];
+ const _stdRes = dlm.filter(y);
+ // With discount=1, R_t = G C G' (no extra growth) — different from W-augmented
+ const discRes = dlm.filterDiscount(y, 1.0);
+ expect(discRes.steps.length).toBe(y.length);
+ expect(Number.isFinite(discRes.logLikelihood)).toBe(true);
+ });
+
+ test("lower discount factor → wider prediction intervals", () => {
+ const dlm = DLM.localLevel({ sigmaObs: 1, sigmaLevel: 0 });
+ const y = [1, 2, 3, 4, 5, 6, 7, 8];
+ const res095 = dlm.filterDiscount(y, 0.95);
+ const res099 = dlm.filterDiscount(y, 0.99);
+ // delta=0.95 has more uncertainty → larger Q
+ const q095 = res095.steps[3]?.forecastCov[0]?.[0]!;
+ const q099 = res099.steps[3]?.forecastCov[0]?.[0]!;
+ expect(q095).toBeGreaterThanOrEqual(q099 - 1e-6);
+ });
+});
+
+// ─── Factory static methods ────────────────────────────────────────────────────
+describe("DLM factory methods", () => {
+ test("DLM.localLevel matches buildLocalLevel dimensions", () => {
+ const dlm1 = DLM.localLevel({ sigmaObs: 1, sigmaLevel: 0.5 });
+ const spec = buildLocalLevel({ sigmaObs: 1, sigmaLevel: 0.5 });
+ // Both have G=[[1]], F=[[1]]
+ const dlm2 = new DLM(spec);
+ const y = [1, 2, 3];
+ const r1 = dlm1.filter(y);
+ const r2 = dlm2.filter(y);
+ expect(r1.filteredMeans[0]?.[0]).toBeCloseTo(r2.filteredMeans[0]?.[0]!, 4);
+ });
+
+ test("DLM.fourier factory", () => {
+ const dlm = DLM.fourier(12, 3, { sigmaObs: 1, sigmaState: 0.01 });
+ expect(dlm).toBeInstanceOf(DLM);
+ const res = dlm.filter([1, 2, 3, 2, 1, 2, 3, 2, 1, 2, 3, 2]);
+ expect(res.steps.length).toBe(12);
+ });
+
+ test("DLM.polynomial(1) is same as localLevel", () => {
+ const dlm = DLM.polynomial(1, { sigmaObs: 1, sigmaState: 0.5 });
+ expect(dlm).toBeInstanceOf(DLM);
+ });
+});
+
+// ─── Edge cases ────────────────────────────────────────────────────────────────
+describe("DLM edge cases", () => {
+ test("single observation", () => {
+ const dlm = DLM.localLevel({ sigmaObs: 1, sigmaLevel: 0.5 });
+ const res = dlm.filter([3.14]);
+ expect(res.filteredMeans.length).toBe(1);
+ expect(Number.isFinite(res.logLikelihood)).toBe(true);
+ });
+
+ test("all missing observations", () => {
+ const dlm = DLM.localLevel({ sigmaObs: 1, sigmaLevel: 0.5 });
+ const res = dlm.filter([null, null, null]);
+ expect(res.steps.length).toBe(3);
+ // No log-likelihood contributions from missing obs
+ expect(res.logLikelihood).toBe(0);
+ });
+
+ test("constant series", () => {
+ const dlm = DLM.localLevel({ sigmaObs: 0.5, sigmaLevel: 0.1 });
+ const res = dlm.filter([5, 5, 5, 5, 5]);
+ // Filtered means should converge toward 5
+ const last = res.filteredMeans[4]?.[0]!;
+ expect(last).toBeGreaterThan(4);
+ expect(last).toBeLessThan(6);
+ });
+
+ test("custom prior m0 and C0", () => {
+ const dlm = DLM.localLevel({ sigmaObs: 1, sigmaLevel: 0.5 });
+ const _res1 = dlm.filter([1, 2, 3], { m0: [0], C0: [[1000]] });
+ const res2 = dlm.filter([1, 2, 3], { m0: [5], C0: [[0.01]] });
+ // Strong prior on m0=5 should keep filtered mean near 5 initially
+ expect(res2.filteredMeans[0]?.[0]).toBeGreaterThan(3);
+ });
+});
+
+// ─── Property-based tests ──────────────────────────────────────────────────────
+describe("DLM property tests", () => {
+ test("logLikelihood is finite for any finite observations", () => {
+ fc.assert(
+ fc.property(
+ fc.array(fc.float({ min: -100, max: 100, noNaN: true }), { minLength: 2, maxLength: 20 }),
+ fc.float({ min: 0.01, max: 10, noNaN: true }),
+ fc.float({ min: 0.01, max: 10, noNaN: true }),
+ (y, sv, sw) => {
+ const dlm = DLM.localLevel({ sigmaObs: sv, sigmaLevel: sw });
+ const res = dlm.filter(y);
+ return Number.isFinite(res.logLikelihood);
+ },
+ ),
+ );
+ });
+
+ test("filtered means length = input length", () => {
+ fc.assert(
+ fc.property(
+ fc.array(fc.float({ min: -50, max: 50, noNaN: true }), { minLength: 1, maxLength: 30 }),
+ (y) => {
+ const dlm = DLM.localLevel();
+ const res = dlm.filter(y);
+ return res.filteredMeans.length === y.length;
+ },
+ ),
+ );
+ });
+
+ test("smoother length = filter length", () => {
+ fc.assert(
+ fc.property(
+ fc.array(fc.float({ min: -50, max: 50, noNaN: true }), { minLength: 2, maxLength: 20 }),
+ (y) => {
+ const dlm = DLM.localLevel();
+ const res = dlm.smooth(y);
+ return res.smoothedMeans.length === y.length && res.smoothedCovs.length === y.length;
+ },
+ ),
+ );
+ });
+
+ test("forecast lower ≤ mean[0] ≤ upper (scalar model)", () => {
+ fc.assert(
+ fc.property(
+ fc.array(fc.float({ min: -50, max: 50, noNaN: true }), { minLength: 3, maxLength: 15 }),
+ fc.integer({ min: 1, max: 10 }),
+ (y, h) => {
+ const dlm = DLM.localLevel();
+ const res = dlm.filter(y);
+ const fc2 = dlm.forecast(res, h);
+ return fc2.lower.every(
+ (lo, i) => lo <= fc2.mean[i]?.[0]! + 1e-8 && fc2.mean[i]?.[0]! <= fc2.upper[i]! + 1e-8,
+ );
+ },
+ ),
+ );
+ });
+});
diff --git a/tests/stats/ets.test.ts b/tests/stats/ets.test.ts
new file mode 100644
index 00000000..c3d64270
--- /dev/null
+++ b/tests/stats/ets.test.ts
@@ -0,0 +1,824 @@
+/**
+ * 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.at(-1) ?? 0, 0);
+ });
+
+ it("functional API simpleExpSmoothing returns same result", () => {
+ const fit1 = simpleExpSmoothing(TREND);
+ const fit2 = new SimpleExpSmoothing().fit(TREND);
+ expect(fit1.alpha).toBeCloseTo(fit2.alpha, 8);
+ expect(fit1.sse).toBeCloseTo(fit2.sse, 8);
+ });
+ });
+
+ describe("accuracy", () => {
+ it("fitted values within 20% of actuals on AIRLINE data", () => {
+ const fit = new SimpleExpSmoothing().fit(AIRLINE);
+ const err = mae(fit.fittedValues, AIRLINE);
+ expect(err / (AIRLINE.reduce((a, b) => a + b, 0) / AIRLINE.length)).toBeLessThan(0.2);
+ });
+
+ it("forecast stays in reasonable range for constant series", () => {
+ const y = new Array(20).fill(5);
+ const model = new SimpleExpSmoothing();
+ model.fit(y);
+ const fc = model.forecast(5);
+ for (const v of fc) {
+ expect(Math.abs(v - 5)).toBeLessThan(1);
+ }
+ });
+ });
+});
+
+// ─── Holt ─────────────────────────────────────────────────────────────────────
+
+describe("Holt", () => {
+ describe("construction", () => {
+ it("creates instance without arguments", () => {
+ expect(() => new Holt()).not.toThrow();
+ });
+
+ it("stores options from constructor", () => {
+ const model = new Holt({ damped: true });
+ const fit = model.fit(TREND);
+ expect(fit.phi).toBeLessThan(1); // damped
+ });
+ });
+
+ describe("fit()", () => {
+ it("requires at least 3 observations", () => {
+ expect(() => new Holt().fit([1, 2])).toThrow(RangeError);
+ });
+
+ it("returns alpha and beta in (0, 1)", () => {
+ const fit = new Holt().fit(TREND);
+ expect(fit.alpha).toBeGreaterThan(0);
+ expect(fit.alpha).toBeLessThan(1);
+ expect(fit.beta).toBeGreaterThan(0);
+ expect(fit.beta).toBeLessThan(1);
+ });
+
+ it("phi = 1 when damped = false (default)", () => {
+ const fit = new Holt().fit(TREND);
+ expect(fit.phi).toBe(1.0);
+ });
+
+ it("phi < 1 when damped = true", () => {
+ const fit = new Holt({ damped: true }).fit(TREND);
+ expect(fit.phi).toBeGreaterThan(0.8);
+ expect(fit.phi).toBeLessThan(1);
+ });
+
+ it("fittedValues length = n", () => {
+ const fit = new Holt().fit(TREND);
+ expect(fit.fittedValues.length).toBe(TREND.length);
+ });
+
+ it("fitted + residuals = y", () => {
+ const fit = new Holt().fit(AIRLINE.slice(0, 12));
+ for (let i = 0; i < 12; i++) {
+ expect((fit.fittedValues[i] ?? 0) + (fit.residuals[i] ?? 0)).toBeCloseTo(
+ AIRLINE[i] ?? 0,
+ 6,
+ );
+ }
+ });
+
+ it("sse = sum of squared residuals", () => {
+ const fit = new Holt().fit(TREND);
+ let sse = 0;
+ for (const e of fit.residuals) {
+ sse += e * e;
+ }
+ expect(fit.sse).toBeCloseTo(sse, 6);
+ });
+
+ it("fixed alpha and beta are respected", () => {
+ const fit = new Holt().fit(TREND, { alpha: 0.5, beta: 0.2 });
+ expect(fit.alpha).toBeCloseTo(0.5, 10);
+ expect(fit.beta).toBeCloseTo(0.2, 10);
+ });
+
+ it("optimised Holt SSE ≤ SES SSE for trending data", () => {
+ const sesSSE = simpleExpSmoothing(TREND).sse;
+ const holtSSE = holt(TREND).sse;
+ expect(holtSSE).toBeLessThanOrEqual(sesSSE + 1e-3);
+ });
+
+ it("functional holt() returns same result as class", () => {
+ const fit1 = holt(TREND);
+ const fit2 = new Holt().fit(TREND);
+ expect(fit1.alpha).toBeCloseTo(fit2.alpha, 6);
+ expect(fit1.sse).toBeCloseTo(fit2.sse, 6);
+ });
+ });
+
+ describe("forecast()", () => {
+ it("throws if called before fit()", () => {
+ expect(() => new Holt().forecast(3)).toThrow();
+ });
+
+ it("returns correct number of steps", () => {
+ const model = new Holt();
+ model.fit(TREND);
+ expect(model.forecast(5).length).toBe(5);
+ });
+
+ it("trend forecasts increase for upward trend data", () => {
+ const model = new Holt();
+ model.fit(TREND);
+ const fc = model.forecast(3);
+ expect(fc[1] ?? 0).toBeGreaterThan(fc[0] ?? 0);
+ expect(fc[2] ?? 0).toBeGreaterThan(fc[1] ?? 0);
+ });
+
+ it("damped forecasts converge to a limit", () => {
+ const model = new Holt({ damped: true });
+ model.fit(TREND);
+ const fc = model.forecast(20);
+ // Differences should shrink
+ const diffs = fc.slice(1).map((v, i) => v - (fc[i] ?? 0));
+ for (let i = 1; i < diffs.length; i++) {
+ expect(Math.abs(diffs[i] ?? 0)).toBeLessThanOrEqual(Math.abs(diffs[i - 1] ?? 0) + 1e-6);
+ }
+ });
+
+ it("forecast closely follows linear trend", () => {
+ // Perfect linear data: y = 2t + 10
+ const y = Array.from({ length: 20 }, (_, t) => 2 * t + 10);
+ const model = new Holt();
+ model.fit(y);
+ const fc = model.forecast(3);
+ // Should forecast ≈ [50, 52, 54]
+ expect(fc[0]).toBeCloseTo(50, 0);
+ expect(fc[2]).toBeCloseTo(54, 0);
+ });
+ });
+
+ describe("accuracy", () => {
+ it("RMSE on AIRLINE (first 24) < 20", () => {
+ const y = AIRLINE.slice(0, 24);
+ const fit = new Holt().fit(y);
+ const err = rmse(fit.fittedValues.slice(1), y.slice(1));
+ expect(err).toBeLessThan(20);
+ });
+ });
+});
+
+// ─── ExponentialSmoothing (Holt-Winters) ─────────────────────────────────────
+
+describe("ExponentialSmoothing", () => {
+ describe("construction and validation", () => {
+ it("creates with no options (SES mode)", () => {
+ const model = new ExponentialSmoothing();
+ const fit = model.fit(TREND);
+ expect(fit.beta).toBeNull();
+ expect(fit.gamma).toBeNull();
+ });
+
+ it("requires at least 3 observations", () => {
+ expect(() => new ExponentialSmoothing().fit([1, 2])).toThrow(RangeError);
+ });
+
+ it("requires 2 full seasonal periods when seasonal is set", () => {
+ expect(() =>
+ new ExponentialSmoothing({ seasonal: "add", seasonalPeriods: 4 }).fit([1, 2, 3, 4, 5]),
+ ).toThrow(RangeError);
+ });
+ });
+
+ describe("SES mode (no trend, no seasonal)", () => {
+ it("result matches SimpleExpSmoothing", () => {
+ const r1 = new ExponentialSmoothing().fit(TREND);
+ const r2 = new SimpleExpSmoothing().fit(TREND);
+ expect(r1.alpha).toBeCloseTo(r2.alpha, 4);
+ expect(r1.sse).toBeCloseTo(r2.sse, 4);
+ });
+
+ it("beta and gamma are null", () => {
+ const fit = new ExponentialSmoothing().fit(TREND);
+ expect(fit.beta).toBeNull();
+ expect(fit.gamma).toBeNull();
+ });
+ });
+
+ describe("additive trend, no seasonal (= Holt)", () => {
+ it("alpha and beta are in (0,1)", () => {
+ const fit = new ExponentialSmoothing({ trend: "add" }).fit(TREND);
+ expect(fit.alpha).toBeGreaterThan(0);
+ expect(fit.beta).not.toBeNull();
+ expect(fit.beta ?? 0).toBeGreaterThan(0);
+ });
+
+ it("gamma is null", () => {
+ const fit = new ExponentialSmoothing({ trend: "add" }).fit(TREND);
+ expect(fit.gamma).toBeNull();
+ });
+
+ it("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..09d84507
--- /dev/null
+++ b/tests/stats/filters.test.ts
@@ -0,0 +1,437 @@
+/**
+ * Tests for src/stats/filters.ts
+ * Covers FIR design, Butterworth IIR, frequency response, and filter application.
+ */
+
+import { describe, expect, test } from "bun:test";
+import * as fc from "fast-check";
+import {
+ type SOSSection,
+ butter,
+ cAbs,
+ filtfilt,
+ firwin,
+ freqz,
+ lfilter,
+ sosfilt,
+ sosfiltfilt,
+ sosfreqz,
+} from "../../src/stats/filters.ts";
+
+// ─── helpers ──────────────────────────────────────────────────────────────────
+
+function near(a: number, b: number, tol = 1e-4): boolean {
+ return Math.abs(a - b) <= tol * (1 + Math.abs(b));
+}
+
+function nearAbs(a: number, b: number, tol = 1e-9): boolean {
+ return Math.abs(a - b) <= tol;
+}
+
+// ─── firwin ───────────────────────────────────────────────────────────────────
+
+describe("firwin", () => {
+ test("returns correct number of taps", () => {
+ const b = firwin(21, 0.3);
+ expect(b.length).toBe(21);
+ });
+
+ test("low-pass: DC gain ≈ 1", () => {
+ const b = firwin(51, 0.3);
+ const dcGain = b.reduce((s, v) => s + v, 0);
+ expect(dcGain).toBeCloseTo(1.0, 4);
+ });
+
+ test("low-pass: gain near 0 at Nyquist", () => {
+ const b = firwin(51, 0.3);
+ // H(e^jπ) = sum b[n] * (-1)^n
+ const nyqGain = b.reduce((s, v, i) => s + v * (i % 2 === 0 ? 1 : -1), 0);
+ expect(Math.abs(nyqGain)).toBeLessThan(0.01);
+ });
+
+ test("high-pass (pass_zero=false): gain ≈ 1 at Nyquist", () => {
+ const b = firwin(51, 0.3, { pass_zero: false });
+ const nyqGain = b.reduce((s, v, i) => s + v * (i % 2 === 0 ? 1 : -1), 0);
+ expect(Math.abs(nyqGain)).toBeGreaterThan(0.9);
+ });
+
+ test("symmetric coefficients (linear phase)", () => {
+ const b = firwin(31, 0.4);
+ for (let i = 0; i < 16; i++) {
+ expect(b[i] ?? 0).toBeCloseTo(b[30 - i] ?? 0, 10);
+ }
+ });
+
+ test("different window types work", () => {
+ const windows = ["hamming", "hann", "blackman"] as const;
+ for (const win of windows) {
+ const b = firwin(21, 0.3, { window: win });
+ expect(b.length).toBe(21);
+ // DC gain should be near 1
+ const dc = b.reduce((s, v) => s + v, 0);
+ expect(dc).toBeCloseTo(1.0, 3);
+ }
+ });
+
+ test("custom fs scaling", () => {
+ const b1 = firwin(21, 0.3, { fs: 2 }); // default
+ const b2 = firwin(21, 300, { fs: 2000 }); // same normalised cutoff
+ for (let i = 0; i < b1.length; i++) {
+ expect(b1[i] ?? 0).toBeCloseTo(b2[i] ?? 0, 10);
+ }
+ });
+
+ test("band-pass (pass_zero=false, two cutoffs)", () => {
+ const b = firwin(51, [0.2, 0.4], { pass_zero: false });
+ expect(b.length).toBe(51);
+ // DC gain should be near 0
+ const dcGain = b.reduce((s, v) => s + v, 0);
+ expect(Math.abs(dcGain)).toBeLessThan(0.05);
+ });
+
+ test("property: all taps are finite", () => {
+ fc.assert(
+ fc.property(
+ fc.integer({ min: 5, max: 51 }).filter((n) => n % 2 === 1),
+ fc.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/hmm.test.ts b/tests/stats/hmm.test.ts
new file mode 100644
index 00000000..0133556c
--- /dev/null
+++ b/tests/stats/hmm.test.ts
@@ -0,0 +1,252 @@
+/**
+ * Tests for Hidden Markov Model (GaussianHMM, MultinomialHMM).
+ */
+import { describe, expect, it } from "bun:test";
+import fc from "fast-check";
+import { GaussianHMM, MultinomialHMM, fitGaussianHMM, hmmViterbi } from "../../src/stats/hmm.ts";
+
+// ─── GaussianHMM ──────────────────────────────────────────────────────────────
+
+describe("GaussianHMM", () => {
+ it("fits a 2-state model on well-separated data", () => {
+ // State 0: N(0, 0.1), State 1: N(5, 0.1)
+ const obs: number[] = [];
+ for (let i = 0; i < 50; i++) {
+ obs.push(i % 5 < 3 ? 0 + 0.1 * (Math.random() - 0.5) : 5 + 0.1 * (Math.random() - 0.5));
+ }
+
+ const model = new GaussianHMM({ nComponents: 2, nIter: 200 });
+ const fit = model.fit(obs);
+
+ // Means should be approximately 0 and 5
+ const sortedMeans = [...fit.means].sort((a, b) => a - b);
+ expect(sortedMeans[0]).toBeLessThan(2);
+ expect(sortedMeans[1]).toBeGreaterThan(3);
+ expect(fit.startProb.length).toBe(2);
+ expect(fit.transmat.length).toBe(2);
+ expect(fit.logProb).toBeLessThan(0); // log-prob is negative
+ expect(fit.nIterDone).toBeGreaterThan(0);
+ });
+
+ it("predict returns array of length T", () => {
+ const obs = [0.1, 0.2, 5.0, 5.1, 0.05, 5.2, 0.15, 5.3];
+ const model = new GaussianHMM({ nComponents: 2, nIter: 100 });
+ model.fit(obs);
+ const states = model.predict(obs);
+ expect(states.length).toBe(obs.length);
+ for (const s of states) {
+ expect(s).toBeGreaterThanOrEqual(0);
+ expect(s).toBeLessThan(2);
+ }
+ });
+
+ it("score returns a finite log-prob", () => {
+ const obs = [0.1, 0.2, 0.15, 5.0, 5.1, 4.9, 0.08];
+ const model = new GaussianHMM({ nComponents: 2, nIter: 100 });
+ model.fit(obs);
+ const lp = model.score(obs);
+ expect(Number.isFinite(lp)).toBe(true);
+ expect(lp).toBeLessThan(0);
+ });
+
+ it("predictProba returns probabilities summing to ~1", () => {
+ const obs = [0.1, 5.0, 0.2, 5.1, 0.05];
+ const model = new GaussianHMM({ nComponents: 2, nIter: 100 });
+ model.fit(obs);
+ const proba = model.predictProba(obs);
+ expect(proba.length).toBe(obs.length);
+ for (const row of proba) {
+ const s = row.reduce((a, b) => a + b, 0);
+ expect(s).toBeCloseTo(1, 4);
+ for (const p of row) {
+ expect(p).toBeGreaterThanOrEqual(0);
+ expect(p).toBeLessThanOrEqual(1);
+ }
+ }
+ });
+
+ it("sample returns correct length", () => {
+ const obs = [0.1, 0.2, 5.0, 5.1, 0.05, 5.2];
+ const model = new GaussianHMM({ nComponents: 2, nIter: 50 });
+ model.fit(obs);
+ const { states, obs: sampledObs } = model.sample(20);
+ expect(states.length).toBe(20);
+ expect(sampledObs.length).toBe(20);
+ for (const s of states) {
+ expect(s).toBeGreaterThanOrEqual(0);
+ expect(s).toBeLessThan(2);
+ }
+ });
+
+ it("throws before fit", () => {
+ const model = new GaussianHMM({ nComponents: 2 });
+ expect(() => model.predict([1, 2])).toThrow("not fitted");
+ expect(() => model.score([1, 2])).toThrow("not fitted");
+ });
+
+ it("throws on too short sequence", () => {
+ const model = new GaussianHMM({ nComponents: 2 });
+ expect(() => model.fit([1])).toThrow();
+ });
+
+ it("startProb rows sum to 1", () => {
+ const obs = [0, 1, 0, 1, 2, 2, 0, 1].map((x) => x * 3.0);
+ const model = new GaussianHMM({ nComponents: 3, nIter: 50 });
+ const fit = model.fit(obs);
+ const s = fit.startProb.reduce((a, b) => a + b, 0);
+ expect(s).toBeCloseTo(1, 5);
+ for (const row of fit.transmat) {
+ const rs = row.reduce((a, b) => a + b, 0);
+ expect(rs).toBeCloseTo(1, 5);
+ }
+ });
+
+ it("fitGaussianHMM convenience function works", () => {
+ const obs = [0.1, 0.2, 5.0, 5.1, 0.0];
+ const model = fitGaussianHMM(obs, 2);
+ expect(model.means.length).toBe(2);
+ });
+
+ it("property: log-prob is always finite for sane data", () => {
+ fc.assert(
+ fc.property(
+ fc.array(fc.float({ min: -10, max: 10, noNaN: true, noDefaultInfinity: true }), {
+ minLength: 10,
+ maxLength: 50,
+ }),
+ (rawObs) => {
+ const model = new GaussianHMM({ nComponents: 2, nIter: 20 });
+ model.fit(rawObs);
+ const lp = model.score(rawObs);
+ return Number.isFinite(lp);
+ },
+ ),
+ { numRuns: 20 },
+ );
+ });
+});
+
+// ─── MultinomialHMM ───────────────────────────────────────────────────────────
+
+describe("MultinomialHMM", () => {
+ it("fits a 2-state model on discrete data", () => {
+ // State 0: emits 0 often; State 1: emits 1 often
+ const obs: number[] = [];
+ for (let i = 0; i < 60; i++) {
+ obs.push(i % 6 < 4 ? 0 : 1);
+ }
+
+ const model = new MultinomialHMM({ nComponents: 2, nFeatures: 2, nIter: 200 });
+ const fit = model.fit(obs);
+
+ expect(fit.emissionProb.length).toBe(2);
+ expect(fit.emissionProb[0]?.length).toBe(2);
+ expect(fit.logProb).toBeLessThan(0);
+ expect(fit.nIterDone).toBeGreaterThan(0);
+ });
+
+ it("predict returns valid state sequence", () => {
+ const obs = [0, 0, 1, 1, 0, 0, 1, 1, 0, 1];
+ const model = new MultinomialHMM({ nComponents: 2, nFeatures: 2, nIter: 100 });
+ model.fit(obs);
+ const states = model.predict(obs);
+ expect(states.length).toBe(obs.length);
+ for (const s of states) {
+ expect(s).toBeGreaterThanOrEqual(0);
+ expect(s).toBeLessThan(2);
+ }
+ });
+
+ it("score returns finite log-prob", () => {
+ const obs = [0, 1, 2, 0, 1, 2, 0, 1];
+ const model = new MultinomialHMM({ nComponents: 2, nFeatures: 3, nIter: 100 });
+ model.fit(obs);
+ const lp = model.score(obs);
+ expect(Number.isFinite(lp)).toBe(true);
+ expect(lp).toBeLessThan(0);
+ });
+
+ it("emission probs sum to 1 per state", () => {
+ const obs = [0, 1, 0, 2, 1, 0, 2, 1, 0];
+ const model = new MultinomialHMM({ nComponents: 2, nFeatures: 3, nIter: 50 });
+ const fit = model.fit(obs);
+ for (const row of fit.emissionProb) {
+ const s = row.reduce((a, b) => a + b, 0);
+ expect(s).toBeCloseTo(1, 5);
+ }
+ });
+
+ it("transition rows sum to 1", () => {
+ const obs = [0, 0, 1, 1, 0, 0, 1, 1];
+ const model = new MultinomialHMM({ nComponents: 2, nFeatures: 2, nIter: 50 });
+ const fit = model.fit(obs);
+ for (const row of fit.transmat) {
+ const s = row.reduce((a, b) => a + b, 0);
+ expect(s).toBeCloseTo(1, 5);
+ }
+ });
+
+ it("throws before fit", () => {
+ const model = new MultinomialHMM({ nComponents: 2, nFeatures: 3 });
+ expect(() => model.predict([0, 1])).toThrow("not fitted");
+ });
+});
+
+// ─── hmmViterbi standalone ────────────────────────────────────────────────────
+
+describe("hmmViterbi", () => {
+ it("decodes a simple 2-state chain correctly", () => {
+ // State 0 emits symbol 0; State 1 emits symbol 1
+ const startProb = [0.6, 0.4];
+ const transmat = [
+ [0.7, 0.3],
+ [0.3, 0.7],
+ ];
+ const emissionProb = [
+ [0.9, 0.1],
+ [0.1, 0.9],
+ ];
+ const obs = [0, 0, 1, 1, 1, 0];
+ const states = hmmViterbi(startProb, transmat, emissionProb, obs);
+ expect(states.length).toBe(obs.length);
+ // Should mostly decode 0→0,0→0,1→1,1→1,1→1,0→0
+ expect(states[0]).toBe(0);
+ expect(states[2]).toBe(1);
+ expect(states[4]).toBe(1);
+ expect(states[5]).toBe(0);
+ });
+
+ it("property: always returns valid state indices", () => {
+ fc.assert(
+ fc.property(
+ fc.integer({ min: 2, max: 5 }),
+ fc.integer({ min: 2, max: 5 }),
+ fc.integer({ min: 5, max: 20 }),
+ (K, V, T) => {
+ // Random row-stochastic matrices
+ const startProb = Array.from({ length: K }, () => Math.random() + 0.1);
+ const sp = startProb.reduce((a, b) => a + b, 0);
+ const normStart = startProb.map((v) => v / sp);
+
+ const transmat = Array.from({ length: K }, () => {
+ const row = Array.from({ length: K }, () => Math.random() + 0.1);
+ const rs = row.reduce((a, b) => a + b, 0);
+ return row.map((v) => v / rs);
+ });
+
+ const emissionProb = Array.from({ length: K }, () => {
+ const row = Array.from({ length: V }, () => Math.random() + 0.1);
+ const rs = row.reduce((a, b) => a + b, 0);
+ return row.map((v) => v / rs);
+ });
+
+ const obs = Array.from({ length: T }, () => Math.floor(Math.random() * V));
+ const states = hmmViterbi(normStart, transmat, emissionProb, obs);
+
+ return states.length === T && states.every((s) => s >= 0 && s < K);
+ },
+ ),
+ { numRuns: 50 },
+ );
+ });
+});
diff --git a/tests/stats/information.test.ts b/tests/stats/information.test.ts
index 2606f06a..65f268e4 100644
--- a/tests/stats/information.test.ts
+++ b/tests/stats/information.test.ts
@@ -34,7 +34,7 @@ function r(v: number, dp = 6): number {
}
const LN2 = Math.log(2);
-const LOG2 = Math.log(2);
+const _LOG2 = Math.log(2);
// ─── entropy ──────────────────────────────────────────────────────────────────
@@ -82,7 +82,7 @@ describe("entropy — Shannon", () => {
it("entropy is maximised by uniform distribution", () => {
const n = 5;
- const uniform = Array(n).fill(1 / n);
+ const uniform = new Array(n).fill(1 / n);
const uniformH = entropy(uniform);
fc.assert(
fc.property(
@@ -381,7 +381,7 @@ describe("tsallisEntropy", () => {
it("Tsallis entropy of uniform n is (n^(1-q) - 1)/(1-q) for q≠1", () => {
const n = 4;
- const p = Array(n).fill(1 / n);
+ const p = new Array(n).fill(1 / n);
const q = 2;
const expected = (n ** (1 - q) - 1) / (1 - q);
expect(r(tsallisEntropy(p, q))).toBeCloseTo(expected, 8);
diff --git a/tests/stats/kalman.test.ts b/tests/stats/kalman.test.ts
new file mode 100644
index 00000000..3d471cbc
--- /dev/null
+++ b/tests/stats/kalman.test.ts
@@ -0,0 +1,699 @@
+/**
+ * Tests for src/stats/kalman.ts
+ *
+ * Covers:
+ * - KalmanFilter construction (factory helpers + direct)
+ * - filter(): local-level, missing obs, multi-dimensional, log-likelihood
+ * - smooth(): RTS smoother backward pass, Joseph form stability
+ * - kalmanFilter1D / kalmanSmooth1D convenience wrappers
+ * - Utility helpers: extractScalarMeans, filteredPredictionInterval
+ * - Property-based tests (fast-check)
+ *
+ * Numerical references cross-checked against statsmodels and pykalman.
+ */
+
+import { describe, expect, it } from "bun:test";
+import * as fc from "fast-check";
+import {
+ KalmanFilter,
+ extractScalarMeans,
+ filteredPredictionInterval,
+ kalmanFilter1D,
+ kalmanSmooth1D,
+} from "../../src/index.ts";
+
+// ─── Helpers ───────────────────────────────────────────────────────────────────
+
+/** Absolute difference. */
+const absDiff = (a: number, b: number) => Math.abs(a - b);
+
+/** Max absolute difference between two arrays. */
+function maxDiff(a: readonly number[], b: readonly number[]): number {
+ let mx = 0;
+ for (let i = 0; i < a.length; i++) {
+ mx = Math.max(mx, Math.abs((a[i] ?? 0) - (b[i] ?? 0)));
+ }
+ return mx;
+}
+
+/** Root mean square between two arrays. */
+function rms(a: readonly number[], b: readonly number[]): number {
+ let s = 0;
+ for (let i = 0; i < a.length; i++) {
+ s += ((a[i] ?? 0) - (b[i] ?? 0)) ** 2;
+ }
+ return Math.sqrt(s / a.length);
+}
+
+/** Simple LCG for reproducible pseudo-random sequences. */
+function lcgSeq(seed: number, n: number, scale = 1.0): number[] {
+ const xs: number[] = [];
+ let s = seed;
+ for (let i = 0; i < n; i++) {
+ s = (s * 1664525 + 1013904223) & 0x7fffffff;
+ xs.push(((s / 0x7fffffff) * 2 - 1) * scale);
+ }
+ return xs;
+}
+
+/** Generate a random-walk series with observation noise. */
+function localLevelSeries(
+ n: number,
+ qSd = 1,
+ rSd = 1,
+ seed = 1,
+): { obs: number[]; states: number[] } {
+ const states: number[] = [0];
+ const wNoise = lcgSeq(seed, n, qSd);
+ const vNoise = lcgSeq(seed + 999, n, rSd);
+ for (let t = 1; t < n; t++) {
+ states.push((states[t - 1] ?? 0) + (wNoise[t] ?? 0));
+ }
+ const obs = states.map((s, i) => s + (vNoise[i] ?? 0));
+ return { obs, states };
+}
+
+// ─── Construction ──────────────────────────────────────────────────────────────
+
+describe("KalmanFilter.localLevel factory", () => {
+ it("creates 1×1 matrices with correct noise values", () => {
+ const kf = KalmanFilter.localLevel({ processNoise: 2, observationNoise: 3 });
+ expect(kf.transitionMatrix).toEqual([[1]]);
+ expect(kf.observationMatrix).toEqual([[1]]);
+ expect(kf.processNoiseCov).toEqual([[2]]);
+ expect(kf.observationNoiseCov).toEqual([[3]]);
+ });
+
+ it("defaults to processNoise=1, observationNoise=1", () => {
+ const kf = KalmanFilter.localLevel();
+ expect(kf.processNoiseCov[0]?.[0]).toBe(1);
+ expect(kf.observationNoiseCov[0]?.[0]).toBe(1);
+ });
+
+ it("initialStateMean defaults to [0]", () => {
+ const kf = KalmanFilter.localLevel();
+ expect(kf.initialStateMean).toEqual([0]);
+ });
+});
+
+describe("KalmanFilter.localLinearTrend factory", () => {
+ it("creates 2×2 transition matrix F = [[1,1],[0,1]]", () => {
+ const kf = KalmanFilter.localLinearTrend();
+ expect(kf.transitionMatrix).toEqual([
+ [1, 1],
+ [0, 1],
+ ]);
+ });
+
+ it("creates 1×2 observation matrix H = [[1,0]]", () => {
+ const kf = KalmanFilter.localLinearTrend();
+ expect(kf.observationMatrix).toEqual([[1, 0]]);
+ });
+
+ it("has 2-element initialStateMean", () => {
+ const kf = KalmanFilter.localLinearTrend();
+ expect(kf.initialStateMean.length).toBe(2);
+ });
+
+ it("respects custom options", () => {
+ const kf = KalmanFilter.localLinearTrend({
+ levelNoise: 0.5,
+ slopeNoise: 0.02,
+ observationNoise: 2,
+ initialMean: [5, 0.3],
+ });
+ expect(kf.processNoiseCov[0]?.[0]).toBe(0.5);
+ expect(kf.processNoiseCov[1]?.[1]).toBe(0.02);
+ expect(kf.observationNoiseCov[0]?.[0]).toBe(2);
+ expect(kf.initialStateMean[0]).toBe(5);
+ expect(kf.initialStateMean[1]).toBe(0.3);
+ });
+});
+
+describe("KalmanFilter direct construction", () => {
+ it("stores all options", () => {
+ const kf = new KalmanFilter({
+ transitionMatrix: [[0.9]],
+ observationMatrix: [[1]],
+ processNoiseCov: [[0.5]],
+ observationNoiseCov: [[1]],
+ initialStateMean: [2],
+ initialStateCovariance: [[3]],
+ });
+ expect(kf.transitionMatrix[0]?.[0]).toBe(0.9);
+ expect(kf.initialStateMean[0]).toBe(2);
+ expect(kf.initialStateCovariance[0]?.[0]).toBe(3);
+ });
+
+ it("defaults initialStateMean to zero vector", () => {
+ const kf = new KalmanFilter({
+ transitionMatrix: [
+ [1, 0],
+ [0, 1],
+ ],
+ observationMatrix: [[1, 0]],
+ processNoiseCov: [
+ [1, 0],
+ [0, 1],
+ ],
+ observationNoiseCov: [[1]],
+ });
+ expect(kf.initialStateMean).toEqual([0, 0]);
+ });
+
+ it("defaults initialStateCovariance to identity", () => {
+ const kf = new KalmanFilter({
+ transitionMatrix: [
+ [1, 0],
+ [0, 1],
+ ],
+ observationMatrix: [[1, 0]],
+ processNoiseCov: [
+ [1, 0],
+ [0, 1],
+ ],
+ observationNoiseCov: [[1]],
+ });
+ expect(kf.initialStateCovariance).toEqual([
+ [1, 0],
+ [0, 1],
+ ]);
+ });
+});
+
+// ─── filter() ──────────────────────────────────────────────────────────────────
+
+describe("filter – local-level basic", () => {
+ const kf = KalmanFilter.localLevel({ processNoise: 1, observationNoise: 1 });
+ const obs = [[1], [2], [3], [4], [5]] as number[][];
+ const result = kf.filter(obs);
+
+ it("returns nTime correct", () => {
+ expect(result.nTime).toBe(5);
+ });
+
+ it("returns nStates = 1", () => {
+ expect(result.nStates).toBe(1);
+ });
+
+ it("returns nObs = 1", () => {
+ expect(result.nObs).toBe(1);
+ });
+
+ it("filteredStateMeans has shape T × 1", () => {
+ expect(result.filteredStateMeans.length).toBe(5);
+ expect(result.filteredStateMeans[0]?.length).toBe(1);
+ });
+
+ it("filteredStateCovariances has shape T × 1 × 1", () => {
+ expect(result.filteredStateCovariances.length).toBe(5);
+ expect(result.filteredStateCovariances[0]?.[0]?.length).toBe(1);
+ });
+
+ it("filtered means lie between prior and observation", () => {
+ for (let t = 0; t < obs.length; t++) {
+ const m = result.filteredStateMeans[t]?.[0] ?? Number.NaN;
+ const y = obs[t]?.[0] ?? Number.NaN;
+ expect(Number.isFinite(m)).toBe(true);
+ // filtered mean < 2 * observation amplitude
+ expect(Math.abs(m)).toBeLessThan(2 * Math.abs(y) + 5);
+ }
+ });
+
+ it("filtered covariances are positive", () => {
+ for (const P of result.filteredStateCovariances) {
+ expect(P[0]?.[0]).toBeGreaterThan(0);
+ }
+ });
+
+ it("innovations have length T × 1", () => {
+ expect(result.innovations.length).toBe(5);
+ expect(result.innovations[0]?.length).toBe(1);
+ });
+
+ it("logLikelihood is finite", () => {
+ expect(Number.isFinite(result.logLikelihood)).toBe(true);
+ });
+});
+
+describe("filter – monotone series tracking", () => {
+ it("tracks a ramp signal (1,2,3,…,10) within ±2", () => {
+ const kf = KalmanFilter.localLevel({ processNoise: 1, observationNoise: 0.1 });
+ const obs = Array.from({ length: 10 }, (_, i) => [i + 1] as [number]);
+ const result = kf.filter(obs);
+ const means = extractScalarMeans(result.filteredStateMeans);
+ for (let t = 0; t < 10; t++) {
+ expect(absDiff(means[t] ?? Number.NaN, t + 1)).toBeLessThan(2);
+ }
+ });
+});
+
+describe("filter – missing observations", () => {
+ const kf = KalmanFilter.localLevel({ processNoise: 0.5, observationNoise: 1 });
+ const obs: (number | null)[][] = [[1], [null], [null], [4], [5]];
+ const result = kf.filter(obs);
+
+ it("handles null without throwing", () => {
+ expect(result.filteredStateMeans.length).toBe(5);
+ });
+
+ it("innovations are NaN for missing steps", () => {
+ expect(Number.isNaN(result.innovations[1]?.[0] ?? 0)).toBe(true);
+ expect(Number.isNaN(result.innovations[2]?.[0] ?? 0)).toBe(true);
+ });
+
+ it("innovations are finite for observed steps", () => {
+ expect(Number.isFinite(result.innovations[0]?.[0] ?? Number.NaN)).toBe(true);
+ expect(Number.isFinite(result.innovations[3]?.[0] ?? Number.NaN)).toBe(true);
+ });
+
+ it("covariance increases during missing steps (uncertainty grows)", () => {
+ const P0 = result.filteredStateCovariances[0]?.[0]?.[0] ?? 0;
+ const P1 = result.filteredStateCovariances[1]?.[0]?.[0] ?? 0;
+ const P2 = result.filteredStateCovariances[2]?.[0]?.[0] ?? 0;
+ expect(P1).toBeGreaterThan(P0);
+ expect(P2).toBeGreaterThan(P1);
+ });
+
+ it("filtered mean does not jump to NaN during missing steps", () => {
+ for (const m of result.filteredStateMeans) {
+ expect(Number.isFinite(m[0] ?? Number.NaN)).toBe(true);
+ }
+ });
+});
+
+describe("filter – logLikelihood", () => {
+ it("log-likelihood is negative for noisy data", () => {
+ const kf = KalmanFilter.localLevel();
+ const obs = [[5], [1], [8], [2], [6]] as number[][];
+ const { logLikelihood } = kf.filter(obs);
+ expect(logLikelihood).toBeLessThan(0);
+ });
+
+ it("log-likelihood is higher for cleaner data (better fit)", () => {
+ const kf = KalmanFilter.localLevel({ processNoise: 0.1, observationNoise: 0.1 });
+ const cleanObs = [[1], [1.01], [1.02], [1.01], [1.0]] as number[][];
+ const noisyObs = [[1], [5], [-3], [8], [-1]] as number[][];
+ const ll1 = kf.filter(cleanObs).logLikelihood;
+ const ll2 = kf.filter(noisyObs).logLikelihood;
+ expect(ll1).toBeGreaterThan(ll2);
+ });
+
+ it("log-likelihood is only computed for non-missing steps", () => {
+ const kf = KalmanFilter.localLevel();
+ const full = [[1], [2], [3]] as number[][];
+ const partial = [[1], [null], [3]] as (number | null)[][];
+ const ll1 = kf.filter(full).logLikelihood;
+ const ll2 = kf.filter(partial).logLikelihood;
+ // partial has fewer observations → lower (or equal) log-likelihood
+ expect(ll1).toBeLessThanOrEqual(ll1 + 1); // basic: both are finite
+ expect(Number.isFinite(ll1) && Number.isFinite(ll2)).toBe(true);
+ });
+});
+
+describe("filter – 2D state (local linear trend)", () => {
+ const kf = KalmanFilter.localLinearTrend({
+ levelNoise: 0.1,
+ slopeNoise: 0.01,
+ observationNoise: 0.5,
+ });
+ const obs = Array.from({ length: 15 }, (_, i) => [i * 1.0]) as number[][];
+ const result = kf.filter(obs);
+
+ it("returns 2-element state means", () => {
+ expect(result.filteredStateMeans[0]?.length).toBe(2);
+ });
+
+ it("tracks linear trend: level ≈ t", () => {
+ const means = result.filteredStateMeans;
+ for (let t = 5; t < 15; t++) {
+ const level = means[t]?.[0] ?? Number.NaN;
+ expect(absDiff(level, t)).toBeLessThan(3);
+ }
+ });
+
+ it("slope converges towards 1", () => {
+ const means = result.filteredStateMeans;
+ const slope = means[14]?.[1] ?? Number.NaN;
+ expect(absDiff(slope, 1.0)).toBeLessThan(0.5);
+ });
+
+ it("2×2 covariance structure", () => {
+ const P = result.filteredStateCovariances[5];
+ expect(P?.length).toBe(2);
+ expect(P?.[0]?.length).toBe(2);
+ });
+});
+
+describe("filter – predicted state properties", () => {
+ it("predictedStateMeans has same length as obs", () => {
+ const kf = KalmanFilter.localLevel();
+ const result = kf.filter([[1], [2], [3]]);
+ expect(result.predictedStateMeans.length).toBe(3);
+ });
+
+ it("first predicted mean equals F * initialStateMean = initialStateMean for F=[[1]]", () => {
+ const kf = KalmanFilter.localLevel({ processNoise: 1, observationNoise: 1 });
+ const result = kf.filter([[5]]);
+ // x_{1|0} = F * m0 = 1 * 0 = 0 (m0=0 by default)
+ expect(result.predictedStateMeans[0]?.[0]).toBeCloseTo(0, 5);
+ });
+});
+
+// ─── smooth() ─────────────────────────────────────────────────────────────────
+
+describe("smooth – basic properties", () => {
+ const kf = KalmanFilter.localLevel({ processNoise: 1, observationNoise: 1 });
+ const obs = [[1], [2], [3], [2], [1]] as number[][];
+ const sm = kf.smooth(obs);
+
+ it("returns smoothedStateMeans of shape T × 1", () => {
+ expect(sm.smoothedStateMeans.length).toBe(5);
+ expect(sm.smoothedStateMeans[0]?.length).toBe(1);
+ });
+
+ it("returns smoothedStateCovariances of shape T × 1 × 1", () => {
+ expect(sm.smoothedStateCovariances.length).toBe(5);
+ expect(sm.smoothedStateCovariances[0]?.[0]?.length).toBe(1);
+ });
+
+ it("last smoothed mean equals last filtered mean", () => {
+ const filtLast = sm.filterResult.filteredStateMeans[4]?.[0] ?? Number.NaN;
+ const smoothLast = sm.smoothedStateMeans[4]?.[0] ?? Number.NaN;
+ expect(absDiff(filtLast, smoothLast)).toBeLessThan(1e-10);
+ });
+
+ it("smoothed covariance ≤ filtered covariance (smoother reduces uncertainty)", () => {
+ for (let t = 0; t < 4; t++) {
+ const Pfilt = sm.filterResult.filteredStateCovariances[t]?.[0]?.[0] ?? 0;
+ const Psmooth = sm.smoothedStateCovariances[t]?.[0]?.[0] ?? 0;
+ expect(Psmooth).toBeLessThanOrEqual(Pfilt + 1e-10);
+ }
+ });
+
+ it("logLikelihood matches filter result", () => {
+ expect(sm.logLikelihood).toBeCloseTo(sm.filterResult.logLikelihood, 10);
+ });
+
+ it("smootherGains has length T, last entry is all zeros", () => {
+ expect(sm.smootherGains.length).toBe(5);
+ expect(sm.smootherGains[4]?.[0]?.[0]).toBeCloseTo(0, 10);
+ });
+});
+
+describe("smooth – missing observations", () => {
+ it("smoothes over gaps in data", () => {
+ const kf = KalmanFilter.localLevel({ processNoise: 1, observationNoise: 0.5 });
+ const obs: (number | null)[][] = [[0], [null], [null], [null], [4]];
+ const sm = kf.smooth(obs);
+ const means = sm.smoothedStateMeans.map((m) => m[0] ?? Number.NaN);
+ // Smoothed means should interpolate between 0 and 4
+ expect(means[2]).toBeGreaterThan(0.5);
+ expect(means[2]).toBeLessThan(3.5);
+ // All means should be finite
+ for (const m of means) {
+ expect(Number.isFinite(m)).toBe(true);
+ }
+ });
+});
+
+describe("smooth – RTS reduces RMSE vs filter", () => {
+ it("smoother RMSE ≤ filter RMSE on generated series", () => {
+ const { obs, states } = localLevelSeries(50, 0.5, 1.0, 42);
+ const kf = KalmanFilter.localLevel({ processNoise: 0.5, observationNoise: 1 });
+ const filtResult = kf.filter(obs.map((v) => [v]));
+ const smResult = kf.smooth(obs.map((v) => [v]));
+ const filtMeans = extractScalarMeans(filtResult.filteredStateMeans);
+ const smoothMeans = extractScalarMeans(smResult.smoothedStateMeans);
+ const filtRmse = rms(filtMeans, states);
+ const smoothRmse = rms(smoothMeans, states);
+ // Smoother should not be worse than filter in RMSE
+ expect(smoothRmse).toBeLessThanOrEqual(filtRmse + 0.1);
+ });
+});
+
+describe("smooth – local linear trend", () => {
+ it("smoothes a trending series without NaN", () => {
+ const kf = KalmanFilter.localLinearTrend();
+ const obs = Array.from({ length: 10 }, (_, i) => [i * 2.0]) as number[][];
+ const sm = kf.smooth(obs);
+ for (const m of sm.smoothedStateMeans) {
+ for (const v of m) {
+ expect(Number.isFinite(v)).toBe(true);
+ }
+ }
+ });
+});
+
+// ─── kalmanFilter1D / kalmanSmooth1D ──────────────────────────────────────────
+
+describe("kalmanFilter1D convenience wrapper", () => {
+ it("accepts scalar array with nulls", () => {
+ const result = kalmanFilter1D([1, 2, null, 4], { processNoise: 1, observationNoise: 1 });
+ expect(result.nTime).toBe(4);
+ expect(result.filteredStateMeans.length).toBe(4);
+ });
+
+ it("produces same result as KalmanFilter.localLevel().filter()", () => {
+ const obs: (number | null)[] = [1, 2, 3, null, 5];
+ const r1 = kalmanFilter1D(obs, { processNoise: 2, observationNoise: 0.5 });
+ const r2 = KalmanFilter.localLevel({ processNoise: 2, observationNoise: 0.5 }).filter(
+ obs.map((v) => [v]),
+ );
+ const m1 = extractScalarMeans(r1.filteredStateMeans);
+ const m2 = extractScalarMeans(r2.filteredStateMeans);
+ for (let i = 0; i < m1.length; i++) {
+ expect(absDiff(m1[i] ?? Number.NaN, m2[i] ?? Number.NaN)).toBeLessThan(1e-10);
+ }
+ });
+});
+
+describe("kalmanSmooth1D convenience wrapper", () => {
+ it("returns smoother result with shape T × 1", () => {
+ const sm = kalmanSmooth1D([1, null, 3]);
+ expect(sm.smoothedStateMeans.length).toBe(3);
+ expect(sm.smoothedStateMeans[0]?.length).toBe(1);
+ });
+
+ it("returns logLikelihood", () => {
+ const sm = kalmanSmooth1D([1, 2, 3]);
+ expect(Number.isFinite(sm.logLikelihood)).toBe(true);
+ });
+});
+
+// ─── Utility helpers ───────────────────────────────────────────────────────────
+
+describe("extractScalarMeans", () => {
+ it("extracts first element of each state mean", () => {
+ const kf = KalmanFilter.localLevel();
+ const result = kf.filter([[1], [2], [3]]);
+ const means = extractScalarMeans(result.filteredStateMeans);
+ expect(means.length).toBe(3);
+ for (let i = 0; i < 3; i++) {
+ expect(means[i]).toBeCloseTo(result.filteredStateMeans[i]?.[0] ?? Number.NaN, 10);
+ }
+ });
+});
+
+describe("filteredPredictionInterval", () => {
+ it("returns lower and upper arrays of length T", () => {
+ const kf = KalmanFilter.localLevel();
+ const result = kf.filter([[1], [2], [3]]);
+ const { lower, upper } = filteredPredictionInterval(result);
+ expect(lower.length).toBe(3);
+ expect(upper.length).toBe(3);
+ });
+
+ it("lower < mean < upper for all t", () => {
+ const kf = KalmanFilter.localLevel();
+ const result = kf.filter([[1], [2], [3]]);
+ const means = extractScalarMeans(result.filteredStateMeans);
+ const { lower, upper } = filteredPredictionInterval(result);
+ for (let t = 0; t < 3; t++) {
+ expect((lower[t] ?? 0) < (means[t] ?? 0)).toBe(true);
+ expect((upper[t] ?? 0) > (means[t] ?? 0)).toBe(true);
+ }
+ });
+
+ it("wider interval for larger zScore", () => {
+ const kf = KalmanFilter.localLevel();
+ const result = kf.filter([[1], [2]]);
+ const { lower: l1, upper: u1 } = filteredPredictionInterval(result, 1.0);
+ const { lower: l2, upper: u2 } = filteredPredictionInterval(result, 2.0);
+ expect((u2[0] ?? 0) - (l2[0] ?? 0)).toBeGreaterThan((u1[0] ?? 0) - (l1[0] ?? 0));
+ });
+});
+
+// ─── Numerical correctness ────────────────────────────────────────────────────
+
+describe("local-level numerical reference", () => {
+ /**
+ * Verify the Kalman gain formula for the first step of a local-level model
+ * with F=1, H=1, Q=q, R=r, P0=p0:
+ *
+ * S_0 = H * P_{0|-1} * H' + R = p0 + r
+ * K_0 = P_{0|-1} * H' * S_0^{-1} = p0 / (p0 + r)
+ * x_{0|0} = x_{0|-1} + K_0 * (y_0 - H * x_{0|-1})
+ * = 0 + [p0/(p0+r)] * (y_0 - 0)
+ * = y_0 * p0 / (p0 + r)
+ */
+ it("first filtered mean matches manual Kalman gain formula", () => {
+ const p0 = 2;
+ const q = 0.5;
+ const r = 1.5;
+ const y0 = 3.7;
+ const kf = new KalmanFilter({
+ transitionMatrix: [[1]],
+ observationMatrix: [[1]],
+ processNoiseCov: [[q]],
+ observationNoiseCov: [[r]],
+ initialStateMean: [0],
+ initialStateCovariance: [[p0]],
+ });
+ const result = kf.filter([[y0]]);
+ // predicted x = F * m0 = 0; P_pred = F * p0 * F' + Q = p0 + q
+ const pPred = p0 + q;
+ const k = pPred / (pPred + r);
+ const expected = 0 + k * (y0 - 0);
+ expect(result.filteredStateMeans[0]?.[0]).toBeCloseTo(expected, 6);
+ });
+
+ it("filtered covariance after first step matches (I-KH)P formula", () => {
+ const p0 = 2;
+ const q = 0.5;
+ const r = 1.5;
+ const kf = new KalmanFilter({
+ transitionMatrix: [[1]],
+ observationMatrix: [[1]],
+ processNoiseCov: [[q]],
+ observationNoiseCov: [[r]],
+ initialStateMean: [0],
+ initialStateCovariance: [[p0]],
+ });
+ const result = kf.filter([[1.0]]);
+ const pPred = p0 + q;
+ const k = pPred / (pPred + r);
+ // Joseph form: (1-k)^2 * pPred + k^2 * r
+ const expectedP = (1 - k) ** 2 * pPred + k ** 2 * r;
+ expect(result.filteredStateCovariances[0]?.[0]?.[0]).toBeCloseTo(expectedP, 6);
+ });
+
+ it("second predicted covariance uses previous filtered covariance", () => {
+ const p0 = 2;
+ const q = 0.5;
+ const r = 1.5;
+ const kf = new KalmanFilter({
+ transitionMatrix: [[1]],
+ observationMatrix: [[1]],
+ processNoiseCov: [[q]],
+ observationNoiseCov: [[r]],
+ initialStateMean: [0],
+ initialStateCovariance: [[p0]],
+ });
+ const result = kf.filter([[1.0], [2.0]]);
+ const pPred1 = p0 + q;
+ const k1 = pPred1 / (pPred1 + r);
+ const pFilt1 = (1 - k1) ** 2 * pPred1 + k1 ** 2 * r;
+ const pPred2_expected = pFilt1 + q; // F * P_filt1 * F' + Q = P_filt1 + Q
+ expect(result.predictedStateCovariances[1]?.[0]?.[0]).toBeCloseTo(pPred2_expected, 6);
+ });
+});
+
+// ─── Property-based tests ─────────────────────────────────────────────────────
+
+describe("property – filter – shape invariants", () => {
+ it("filteredStateMeans always has length T", () => {
+ fc.assert(
+ fc.property(
+ fc.array(fc.float({ noNaN: true, min: -100, max: 100 }), { minLength: 1, maxLength: 20 }),
+ (ys) => {
+ const kf = KalmanFilter.localLevel();
+ const result = kf.filter(ys.map((v) => [v]));
+ return result.filteredStateMeans.length === ys.length;
+ },
+ ),
+ );
+ });
+
+ it("filteredStateCovariances are always positive for local-level", () => {
+ fc.assert(
+ fc.property(
+ fc.array(fc.float({ noNaN: true, min: -50, max: 50 }), { minLength: 1, maxLength: 20 }),
+ (ys) => {
+ const kf = KalmanFilter.localLevel({ processNoise: 1, observationNoise: 1 });
+ const result = kf.filter(ys.map((v) => [v]));
+ return result.filteredStateCovariances.every((P) => (P[0]?.[0] ?? 0) > 0);
+ },
+ ),
+ );
+ });
+
+ it("logLikelihood is finite for finite observations", () => {
+ fc.assert(
+ fc.property(
+ fc.array(fc.float({ noNaN: true, min: -100, max: 100 }), { minLength: 1, maxLength: 20 }),
+ (ys) => {
+ const kf = KalmanFilter.localLevel();
+ const { logLikelihood } = kf.filter(ys.map((v) => [v]));
+ return Number.isFinite(logLikelihood);
+ },
+ ),
+ );
+ });
+});
+
+describe("property – smoother – uncertainty never exceeds filter", () => {
+ it("smoothed covariance ≤ filtered covariance at every time step", () => {
+ fc.assert(
+ fc.property(
+ fc.array(fc.float({ noNaN: true, min: -50, max: 50 }), { minLength: 2, maxLength: 15 }),
+ (ys) => {
+ const kf = KalmanFilter.localLevel({ processNoise: 1, observationNoise: 1 });
+ const sm = kf.smooth(ys.map((v) => [v]));
+ for (let t = 0; t < ys.length - 1; t++) {
+ const pfilt = sm.filterResult.filteredStateCovariances[t]?.[0]?.[0] ?? 0;
+ const psmooth = sm.smoothedStateCovariances[t]?.[0]?.[0] ?? 0;
+ if (psmooth > pfilt + 1e-8) {
+ return false;
+ }
+ }
+ return true;
+ },
+ ),
+ );
+ });
+});
+
+describe("property – all-null observations", () => {
+ it("filter runs without error on all-null obs", () => {
+ const kf = KalmanFilter.localLevel();
+ const obs: (number | null)[][] = Array.from({ length: 5 }, () => [null]);
+ const result = kf.filter(obs);
+ expect(result.nTime).toBe(5);
+ for (const m of result.filteredStateMeans) {
+ expect(Number.isFinite(m[0] ?? Number.NaN)).toBe(true);
+ }
+ });
+});
+
+describe("property – empty observations array edge case", () => {
+ it("filter on empty array returns T=0 result", () => {
+ const kf = KalmanFilter.localLevel();
+ const result = kf.filter([]);
+ expect(result.nTime).toBe(0);
+ expect(result.filteredStateMeans.length).toBe(0);
+ expect(result.logLikelihood).toBe(0);
+ });
+});
+
+// ─── StateSpaceModel alias ────────────────────────────────────────────────────
+
+describe("StateSpaceModel alias", () => {
+ it("is exported as an alias of KalmanFilter", async () => {
+ const { StateSpaceModel } = await import("../../src/index.ts");
+ // Both should be the same class
+ const ssm = StateSpaceModel.localLevel();
+ const result = ssm.filter([[1], [2], [3]]);
+ expect(result.nTime).toBe(3);
+ });
+});
diff --git a/tests/stats/signal.test.ts b/tests/stats/signal.test.ts
new file mode 100644
index 00000000..5938d36d
--- /dev/null
+++ b/tests/stats/signal.test.ts
@@ -0,0 +1,500 @@
+/**
+ * Tests for src/stats/signal.ts
+ * Covers FFT, windows, STFT, ISTFT, Welch PSD, and periodogram.
+ */
+
+import { describe, expect, test } from "bun:test";
+import * as fc from "fast-check";
+import {
+ bartlettWindow,
+ blackmanHarrisWindow,
+ blackmanWindow,
+ cAbs,
+ complex,
+ fft,
+ fftFreq,
+ fftshift,
+ flatTopWindow,
+ getWindow,
+ hammingWindow,
+ hannWindow,
+ ifft,
+ ifftshift,
+ irfft,
+ istft,
+ kaiserWindow,
+ periodogram,
+ rectangularWindow,
+ rfft,
+ rfftFreq,
+ stft,
+ welch,
+} from "../../src/stats/signal.ts";
+
+// ─── helpers ──────────────────────────────────────────────────────────────────
+
+function near(a: number, b: number, tol = 1e-6): boolean {
+ return Math.abs(a - b) <= tol * (1 + Math.abs(b));
+}
+
+function nearAbs(a: number, b: number, tol = 1e-9): boolean {
+ return Math.abs(a - b) <= tol;
+}
+
+// ─── FFT ──────────────────────────────────────────────────────────────────────
+
+describe("fft / ifft", () => {
+ test("DC input — all energy at bin 0", () => {
+ const X = fft([1, 1, 1, 1]);
+ expect(nearAbs(X[0]?.re ?? 0, 4, 1e-10)).toBe(true);
+ expect(nearAbs(X[0]?.im ?? 0, 0, 1e-10)).toBe(true);
+ expect(nearAbs(cAbs(X[1] ?? complex(0, 0)), 0, 1e-10)).toBe(true);
+ });
+
+ test("single tone — energy at expected bin", () => {
+ // x[n] = exp(2πj * k * n / N) for k=1, N=8
+ const N = 8;
+ const x: number[] = Array.from({ length: N }, (_, n) => Math.cos((2 * Math.PI * n) / N));
+ const X = fft(x);
+ // cosine has energy at bins 1 and N-1
+ const mag = X.map((c) => cAbs(c));
+ expect(mag[1]).toBeCloseTo(N / 2, 4);
+ expect(mag[7]).toBeCloseTo(N / 2, 4);
+ for (let i = 2; i <= 6; i++) {
+ expect(mag[i]).toBeCloseTo(0, 4);
+ }
+ });
+
+ test("Parseval's theorem — energy preserved", () => {
+ const x = [1, 2, 3, 4, 5, 6, 7, 8];
+ const X = fft(x);
+ const N = X.length;
+ const timePower = x.reduce((s, v) => s + v * v, 0);
+ const freqPower = X.reduce((s, c) => s + c.re * c.re + c.im * c.im, 0) / N;
+ expect(nearAbs(timePower, freqPower, 1e-8)).toBe(true);
+ });
+
+ test("ifft(fft(x)) ≈ x (round-trip)", () => {
+ const x = [3, 1, 4, 1, 5, 9, 2, 6];
+ const X = fft(x);
+ const xBack = ifft(X);
+ for (let i = 0; i < x.length; i++) {
+ expect(xBack[i]?.re ?? 0).toBeCloseTo(x[i] ?? 0, 10);
+ }
+ });
+
+ test("zero input → zero output", () => {
+ const X = fft([0, 0, 0, 0]);
+ for (const c of X) {
+ expect(cAbs(c)).toBeCloseTo(0, 12);
+ }
+ });
+
+ test("non-power-of-2 input pads to next power", () => {
+ const x = [1, 2, 3]; // length 3 → pad to 4
+ const X = fft(x);
+ expect(X.length).toBe(4);
+ });
+
+ test("linearity: fft(a*x + b*y) = a*fft(x) + b*fft(y)", () => {
+ const x = [1, 2, 3, 4, 5, 6, 7, 8];
+ const y = [8, 7, 6, 5, 4, 3, 2, 1];
+ const a = 2;
+ const b = 3;
+ const Xab = fft(x.map((v, i) => a * v + b * (y[i] ?? 0)));
+ const Xa = fft(x);
+ const Xy = fft(y);
+ for (let i = 0; i < Xab.length; i++) {
+ const re = a * (Xa[i]?.re ?? 0) + b * (Xy[i]?.re ?? 0);
+ const im = a * (Xa[i]?.im ?? 0) + b * (Xy[i]?.im ?? 0);
+ expect(Xab[i]?.re ?? 0).toBeCloseTo(re, 8);
+ expect(Xab[i]?.im ?? 0).toBeCloseTo(im, 8);
+ }
+ });
+});
+
+describe("rfft / irfft", () => {
+ test("rfft of real signal is conjugate-symmetric", () => {
+ const x = [1, 2, 3, 4, 5, 6, 7, 8];
+ const X = rfft(x);
+ // For a real signal, the full FFT has X[k] = conj(X[N-k])
+ // rfft returns only bins 0..N/2
+ const n = fft(x).length;
+ expect(X.length).toBe(n / 2 + 1);
+ });
+
+ test("irfft(rfft(x)) ≈ x (round-trip)", () => {
+ const x = [1, 0, -1, 0, 1, 0, -1, 0];
+ const X = rfft(x);
+ const n = fft(x).length;
+ const xBack = irfft(X, n);
+ for (let i = 0; i < x.length; i++) {
+ expect(xBack[i] ?? 0).toBeCloseTo(x[i] ?? 0, 8);
+ }
+ });
+
+ test("rfftFreq length matches rfft output", () => {
+ const x = new Array(16).fill(1) as number[];
+ const X = rfft(x);
+ const freqs = rfftFreq(fft(x).length, 1 / 100);
+ expect(freqs.length).toBe(X.length);
+ expect(freqs[0]).toBeCloseTo(0);
+ expect(freqs.at(-1)).toBeCloseTo(50); // Nyquist at 50 Hz for fs=100
+ });
+});
+
+// ─── fftFreq / fftshift / ifftshift ──────────────────────────────────────────
+
+describe("fftFreq", () => {
+ test("DC bin is 0", () => {
+ const f = fftFreq(8, 1);
+ expect(f[0]).toBe(0);
+ });
+
+ test("positive and negative frequencies", () => {
+ const f = fftFreq(8, 1);
+ expect(f[1]).toBeCloseTo(0.125);
+ expect(f[4]).toBeCloseTo(0.5);
+ expect(f[5]).toBeCloseTo(-0.375);
+ expect(f[7]).toBeCloseTo(-0.125);
+ });
+
+ test("sample spacing scales frequencies", () => {
+ const fs = 100;
+ const f = fftFreq(8, 1 / fs);
+ expect(f[1]).toBeCloseTo(fs / 8);
+ });
+});
+
+describe("fftshift / ifftshift", () => {
+ test("even length: round-trip", () => {
+ const x = [0, 1, 2, 3];
+ const shifted = fftshift(x);
+ expect(ifftshift(shifted)).toEqual(x);
+ });
+
+ test("odd length: fftshift matches numpy", () => {
+ const x = [0, 1, 2, 3, 4];
+ const shifted = fftshift(x);
+ expect(shifted).toEqual([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.at(-1) ?? 0;
+ expect(Math.abs(maxF - fs / 2)).toBeLessThan(fs / n + 1);
+ });
+});
+
+// ─── property-based ───────────────────────────────────────────────────────────
+
+describe("FFT properties (fast-check)", () => {
+ test("Parseval's theorem holds for all power-of-2 signals", () => {
+ fc.assert(
+ fc.property(
+ fc.array(fc.float({ min: -10, max: 10, noNaN: true }), { minLength: 8, maxLength: 8 }),
+ (x) => {
+ const X = fft(x);
+ const N = X.length;
+ const timePower = x.reduce((s, v) => s + v * v, 0);
+ const freqPower = X.reduce((s, c) => s + c.re * c.re + c.im * c.im, 0) / N;
+ return Math.abs(timePower - freqPower) < 1e-6 * (1 + timePower);
+ },
+ ),
+ );
+ });
+
+ test("fftshift round-trip for all lengths 1..20", () => {
+ for (let n = 1; n <= 20; n++) {
+ const x = Array.from({ length: n }, (_, i) => i);
+ const roundTrip = ifftshift(fftshift(x));
+ for (let i = 0; i < n; i++) {
+ expect(roundTrip[i]).toBe(x[i]);
+ }
+ }
+ });
+});