From 4f4bcd027d30511d280afc7fecb95b5cfa6b63e0 Mon Sep 17 00:00:00 2001 From: DGoettlich Date: Fri, 22 May 2026 15:44:22 +0200 Subject: [PATCH 1/2] expose dtype option --- README.md | 2 ++ src/scpc/core.py | 6 +++++- src/scpc/types.py | 3 ++- src/scpc/utils/matrix.py | 31 ++++++++++++++++++++----------- src/scpc/utils/spatial.py | 37 ++++++++++++++++++++++--------------- tests/test_get_distmat.py | 8 ++++++++ tests/test_get_w.py | 8 ++++++++ tests/test_scpc.py | 20 ++++++++++++++++++++ tests/test_set_oms_wfin.py | 23 +++++++++++++++++++++++ 9 files changed, 110 insertions(+), 28 deletions(-) diff --git a/README.md b/README.md index dd05998..9488a95 100644 --- a/README.md +++ b/README.md @@ -98,6 +98,8 @@ The most important `scpc()` arguments in the workflow above are: - `avc`: upper bound on the average pairwise correlation - `uncond`: whether to skip the conditional adjustment - `cvs`: whether to store the extra critical values +- `dtype`: spatial array dtype, one of `"float16"`, `"float32"`, or `"float64"` + (default) ## Documentation diff --git a/src/scpc/core.py b/src/scpc/core.py index ee6695c..9e1089a 100644 --- a/src/scpc/core.py +++ b/src/scpc/core.py @@ -6,7 +6,7 @@ import numpy as np -from .types import DataFrameLike, ModelLike, SCPCResult +from .types import DType, DataFrameLike, ModelLike, SCPCResult from .utils.data import ( get_coef_names, get_conditional_projection_setup, @@ -49,6 +49,7 @@ def scpc( large_n_seed: int = 1, uncond: bool = False, cvs: bool = False, + dtype: DType = "float64", ) -> SCPCResult: """Run spatial correlation-robust inference. @@ -69,6 +70,7 @@ def scpc( large_n_seed: Seed for the large-n approximation branch. uncond: Whether to skip the conditional adjustment. cvs: Whether to return additional critical values. + dtype: Floating point dtype for memory-heavy spatial arrays. Returns: The fitted SCPC result object. @@ -79,6 +81,7 @@ def scpc( if avc <= 0.001 or avc >= 0.99: raise ValueError("Option avc() must be in (0.001, 0.99).") + assert dtype in {"float16", "float32", "float64"} method = validate_scpc_method(method) large_n_seed = validate_large_n_seed(large_n_seed) @@ -174,6 +177,7 @@ def scpc( latlong, method=method, large_n_seed=large_n_seed, + dtype=dtype, ) d = None if not spc.large_n: diff --git a/src/scpc/types.py b/src/scpc/types.py index 55a85a6..aee8bd6 100644 --- a/src/scpc/types.py +++ b/src/scpc/types.py @@ -2,7 +2,7 @@ from collections.abc import Sequence from dataclasses import dataclass -from typing import Any, TypeAlias, TypedDict +from typing import Any, Literal, TypeAlias, TypedDict import numpy as np import pandas as pd @@ -13,6 +13,7 @@ MatrixLike: TypeAlias = Any ModelLike: TypeAlias = Any DataFrameLike: TypeAlias = Any +DType: TypeAlias = Literal["float16", "float32", "float64"] SCPC_STATS_COLUMNS = ["Coef", "Std_Err", "t", "P>|t|", "2.5 %", "97.5 %"] SCPC_CV_COLUMNS = ["32%", "10%", "5%", "1%"] diff --git a/src/scpc/utils/matrix.py b/src/scpc/utils/matrix.py index 7656a4c..23c84e8 100644 --- a/src/scpc/utils/matrix.py +++ b/src/scpc/utils/matrix.py @@ -5,10 +5,10 @@ import numpy as np from scipy.sparse.linalg import eigsh -from ..types import ArrayLike, MatrixLike +from ..types import ArrayLike, DType, MatrixLike -def lvech(mat: MatrixLike) -> ArrayLike: +def lvech(mat: MatrixLike, dtype: DType = "float64") -> ArrayLike: """Collect the strict lower triangle of a matrix. This helper turns a square matrix into the vector of pairwise entries @@ -21,11 +21,11 @@ def lvech(mat: MatrixLike) -> ArrayLike: Returns: The strict lower-triangular entries as a vector-like object. """ - mat = np.asarray(mat, dtype=float) + mat = np.asarray(mat, dtype=dtype) return mat[np.tril_indices(mat.shape[0], k=-1)] -def demeanmat(mat: MatrixLike) -> MatrixLike: +def demeanmat(mat: MatrixLike, dtype: DType = "float64") -> MatrixLike: """Remove row and column averages from a matrix. This helper converts a matrix into its double-demeaned version. In SCPC, @@ -39,7 +39,7 @@ def demeanmat(mat: MatrixLike) -> MatrixLike: Returns: The double-demeaned matrix. """ - mat = np.asarray(mat, dtype=float) + mat = np.asarray(mat, dtype=dtype) # remove row means first, then column means from the row-demeaned matrix mat = mat - np.mean(mat, axis=1, keepdims=True) @@ -47,7 +47,12 @@ def demeanmat(mat: MatrixLike) -> MatrixLike: return mat -def get_w(distmat: MatrixLike, c0: float, qmax: int) -> MatrixLike: +def get_w( + distmat: MatrixLike, + c0: float, + qmax: int, + dtype: DType = "float64", +) -> MatrixLike: """Build the candidate spatial projection basis. This helper converts a distance matrix and kernel scale into the matrix of @@ -62,22 +67,26 @@ def get_w(distmat: MatrixLike, c0: float, qmax: int) -> MatrixLike: Returns: The candidate spatial projection matrix. """ - distmat = np.asarray(distmat, dtype=float) + distmat = np.asarray(distmat, dtype=dtype) n = distmat.shape[0] sig = np.exp(-c0 * distmat) - sig_d = demeanmat(sig) + sig_d = demeanmat(sig, dtype=dtype) + sig_d_eig = np.asarray(sig_d, dtype="float32" if dtype == "float16" else dtype) if qmax < n - 1: - eigvals, v = eigsh(sig_d, k=qmax, which="LM") + eigvals, v = eigsh(sig_d_eig, k=qmax, which="LM") order = np.argsort(np.abs(eigvals))[::-1] v = v[:, order] else: - eigvals, v = np.linalg.eigh(sig_d) + eigvals, v = np.linalg.eigh(sig_d_eig) order = np.argsort(eigvals)[::-1] v = v[:, order[:qmax]] # prepend the normalized constant direction, matching the R code - return np.column_stack((np.full(n, 1 / math.sqrt(n)), v)) + return np.column_stack((np.full(n, 1 / math.sqrt(n), dtype=dtype), v)).astype( + dtype, + copy=False, + ) def get_tau(y: ArrayLike, w: MatrixLike) -> float: diff --git a/src/scpc/utils/spatial.py b/src/scpc/utils/spatial.py index e76570b..eeee243 100644 --- a/src/scpc/utils/spatial.py +++ b/src/scpc/utils/spatial.py @@ -6,7 +6,7 @@ from scipy import stats from scipy.sparse.linalg import eigsh -from ..types import ArrayLike, MatrixLike, SpatialSetup +from ..types import ArrayLike, DType, MatrixLike, SpatialSetup from .matrix import demeanmat, get_w, lvech LARGE_N_THRESHOLD = 4500 @@ -36,11 +36,15 @@ def get_avc(c: float, dist: ArrayLike) -> float: Returns: The implied average pairwise correlation. """ - dist = np.asarray(dist, dtype=float) + dist = np.asarray(dist) return float(np.mean(np.exp(-c * dist))) -def get_distmat(s: MatrixLike, latlong: bool) -> MatrixLike: +def get_distmat( + s: MatrixLike, + latlong: bool, + dtype: DType = "float64", +) -> MatrixLike: """Compute the pairwise distance matrix from observation coordinates. This helper is the boundary between raw location data and the rest of the @@ -57,7 +61,7 @@ def get_distmat(s: MatrixLike, latlong: bool) -> MatrixLike: Raises: ValueError: Raised later for invalid coordinate shapes. """ - s = np.asarray(s, dtype=float) + s = np.asarray(s, dtype=dtype) if s.ndim == 1: s = s.reshape(-1, 1) @@ -83,7 +87,7 @@ def get_distmat(s: MatrixLike, latlong: bool) -> MatrixLike: else: d = np.sqrt(np.sum((s[:, None, :] - s[None, :, :]) ** 2, axis=2)) - return d + return d.astype(dtype, copy=False) def get_distvec(s1: MatrixLike, s2: MatrixLike, latlong: bool) -> ArrayLike: @@ -467,6 +471,7 @@ def set_final_w( oms: list[MatrixLike], w: MatrixLike, qmax: int, + dtype: DType = "float64", ) -> tuple[MatrixLike, float, int]: """Choose how many spatial principal components to keep. @@ -488,7 +493,7 @@ def set_final_w( Raises: ValueError: Raised later if the inputs are dimensionally incompatible. """ - w = np.asarray(w, dtype=float) + w = np.asarray(w, dtype=dtype) cvs = np.empty(qmax, dtype=float) lengths = np.empty(qmax, dtype=float) @@ -528,6 +533,7 @@ def get_oms( cmax: float, w: MatrixLike, cgridfac: float, + dtype: DType = "float64", ) -> list[MatrixLike]: """Build omega matrices over the spatial correlation grid. @@ -545,14 +551,14 @@ def get_oms( Returns: Omega matrices across the spatial correlation grid. """ - distmat = np.asarray(distmat, dtype=float) - w = np.asarray(w, dtype=float) + distmat = np.asarray(distmat, dtype=dtype) + w = np.asarray(w, dtype=dtype) nc = get_nc(c0, cmax, cgridfac) oms: list[MatrixLike] = [np.empty((0, 0)) for _ in range(nc)] c = c0 for i in range(nc): - oms[i] = w.T @ (np.exp(-c * distmat) @ w) + oms[i] = w.T @ (np.exp(-c * distmat).astype(dtype, copy=False) @ w) c = c * cgridfac return oms @@ -564,6 +570,7 @@ def set_oms_wfin( latlong: bool, method: str = "auto", large_n_seed: int = 1, + dtype: DType = "float64", ) -> SpatialSetup: """Build the unconditional spatial setup from coordinates and an AVC bound. @@ -580,7 +587,7 @@ def set_oms_wfin( Returns: The complete spatial setup used by the main inference routine. """ - coords = np.asarray(coords, dtype=float) + coords = np.asarray(coords, dtype=dtype) n = coords.shape[0] if method == "auto": @@ -606,12 +613,12 @@ def set_oms_wfin( qmax = min(qmax, n - 1) if method_actual == "exact": - distmat = get_distmat(coords, latlong) - distv = lvech(distmat) + distmat = get_distmat(coords, latlong, dtype=dtype) + distv = lvech(distmat, dtype=dtype) c0 = get_c0_from_avc(distv, avc0) cmax = get_c0_from_avc(distv, MINAVC) - w = get_w(distmat, c0, qmax) - oms = get_oms(distmat, c0, cmax, w, CGRIDFAC) + w = get_w(distmat, c0, qmax, dtype=dtype) + oms = get_oms(distmat, c0, cmax, w, CGRIDFAC, dtype=dtype) coords_use = coords perm = np.arange(n, dtype=int) random_t = None @@ -640,7 +647,7 @@ def set_oms_wfin( ) distmat = None - wfin, cvfin, q = set_final_w(oms, w, qmax) + wfin, cvfin, q = set_final_w(oms, w, qmax, dtype=dtype) if q < qmax or qmax == n - 1: break qmax = round(qmax + qmax / 2) diff --git a/tests/test_get_distmat.py b/tests/test_get_distmat.py index 24180e5..31247e1 100644 --- a/tests/test_get_distmat.py +++ b/tests/test_get_distmat.py @@ -22,6 +22,14 @@ def test_get_distmat_returns_euclidean_distances() -> None: npt.assert_allclose(get_distmat(coords, False), expected, atol=1e-12, rtol=0.0) +def test_get_distmat_preserves_requested_dtype() -> None: + coords = np.array([[0.0, 0.0], [3.0, 4.0]]) + + result = get_distmat(coords, False, dtype="float32") + + assert result.dtype == np.float32 + + @pytest.mark.skipif(R is None, reason="Rscript not installed") def test_python_r_parity_get_distmat() -> None: coords = np.array([[0.0, 0.0], [3.0, 4.0], [6.0, 8.0]]) diff --git a/tests/test_get_w.py b/tests/test_get_w.py index 9b398e0..014ab87 100644 --- a/tests/test_get_w.py +++ b/tests/test_get_w.py @@ -32,6 +32,14 @@ def test_get_w_builds_the_expected_two_point_basis() -> None: assert_columns_allclose_up_to_sign(get_w(distmat, math.log(2.0), 1), expected) +def test_get_w_preserves_requested_float16_dtype() -> None: + distmat = np.array([[0.0, 1.0], [1.0, 0.0]]) + + result = get_w(distmat, math.log(2.0), 1, dtype="float16") + + assert result.dtype == np.float16 + + @pytest.mark.skipif(R is None, reason="Rscript not installed") def test_python_r_parity_get_w() -> None: distmat = [[0.0, 1.0], [1.0, 0.0]] diff --git a/tests/test_scpc.py b/tests/test_scpc.py index 6259cbe..9e3cd75 100644 --- a/tests/test_scpc.py +++ b/tests/test_scpc.py @@ -46,6 +46,26 @@ def test_scpc_returns_a_result_with_the_expected_structure() -> None: assert result.coef_names == ["Intercept", "x"] +def test_scpc_rejects_invalid_dtype() -> None: + data = pd.DataFrame( + { + "y": [1.0, 1.8, 2.9, 3.7, 5.1], + "x": [0.0, 1.0, 2.0, 3.0, 4.0], + "coord_x": [0.0, 1.0, 0.5, 1.5, 2.0], + "coord_y": [0.0, 0.0, 1.0, 1.0, 1.5], + } + ) + model = smf.ols("y ~ x", data=data).fit() + + with pytest.raises(AssertionError): + scpc( + model, + data, + coords_euclidean=("coord_x", "coord_y"), + dtype="float128", # type: ignore # intentional invalid dtype + ) + + def test_scpc_stores_only_reported_coef_names_when_ncoef_is_set() -> None: data = pd.DataFrame( { diff --git a/tests/test_set_oms_wfin.py b/tests/test_set_oms_wfin.py index 2a93f25..325e353 100644 --- a/tests/test_set_oms_wfin.py +++ b/tests/test_set_oms_wfin.py @@ -53,6 +53,29 @@ def test_set_oms_wfin_returns_a_coherent_exact_spatial_setup() -> None: assert result.random_state is None +def test_set_oms_wfin_preserves_requested_exact_spatial_dtype() -> None: + coords = np.array( + [ + [0.0, 0.0], + [1.0, 0.0], + [0.5, 1.0], + ] + ) + + result = set_oms_wfin( + coords, + 0.1, + latlong=False, + method="exact", + large_n_seed=1, + dtype="float32", + ) + + assert result.distmat is not None + assert np.asarray(result.distmat).dtype == np.float32 + assert np.asarray(result.wfin).dtype == np.float32 + + def test_set_oms_wfin_returns_a_coherent_approximate_spatial_setup() -> None: coords = np.array([[0.0], [1.0], [2.0], [3.0]]) From f985dd759771a304683fb765cc84d8008b1a2c24 Mon Sep 17 00:00:00 2001 From: DGoettlich Date: Fri, 22 May 2026 19:59:50 +0200 Subject: [PATCH 2/2] preserve exact distance dtype --- src/scpc/core.py | 4 ++-- tests/test_scpc.py | 46 +++++++++++++++++++++++++++++++++++++++++++++- 2 files changed, 47 insertions(+), 3 deletions(-) diff --git a/src/scpc/core.py b/src/scpc/core.py index 9e1089a..c9e61de 100644 --- a/src/scpc/core.py +++ b/src/scpc/core.py @@ -183,7 +183,7 @@ def scpc( if not spc.large_n: if spc.distmat is None: raise ValueError("Internal error: exact SCPC setup is missing `distmat`.") - d = np.asarray(spc.distmat, dtype=float) + d = spc.distmat wfin = np.asarray(spc.wfin, dtype=float) cvfin = spc.cvfin @@ -293,7 +293,7 @@ def scpc( raise ValueError( "Internal error: exact SCPC conditional branch is missing `distmat`." ) - omsx = get_oms(d, spc.c0, spc.cmax, wx, 1.2) + omsx = get_oms(d, spc.c0, spc.cmax, wx, 1.2, dtype=dtype) p_c = max_rp(omsx, q, abs(tau_u) / math.sqrt(q))[0] cvx = get_cv(omsx, q, 0.05) p_final = max(p_u, p_c) diff --git a/tests/test_scpc.py b/tests/test_scpc.py index 9e3cd75..f22e267 100644 --- a/tests/test_scpc.py +++ b/tests/test_scpc.py @@ -7,7 +7,8 @@ import statsmodels.formula.api as smf from scpc import scpc -from scpc.types import SCPCResult +import scpc.core as scpc_core +from scpc.types import DType, SCPCResult from tests.config import ATOL, RTOL from tests.utils import R, execute_r_code @@ -66,6 +67,49 @@ def test_scpc_rejects_invalid_dtype() -> None: ) +def test_scpc_conditional_exact_keeps_requested_distance_dtype( + monkeypatch: pytest.MonkeyPatch, +) -> None: + data = pd.DataFrame( + { + "y": [1.0, 1.8, 2.9, 3.7, 5.1], + "x": [0.0, 1.0, 2.0, 3.0, 4.0], + "coord_x": [0.0, 1.0, 0.5, 1.5, 2.0], + "coord_y": [0.0, 0.0, 1.0, 1.0, 1.5], + } + ) + model = smf.ols("y ~ x", data=data).fit() + calls: list[tuple[np.dtype, DType]] = [] + original_get_oms = scpc_core.get_oms + + def recording_get_oms( + distmat, + c0, + cmax, + w, + cgridfac, + dtype: DType = "float64", + ): + calls.append((np.asarray(distmat).dtype, dtype)) + return original_get_oms(distmat, c0, cmax, w, cgridfac, dtype=dtype) + + monkeypatch.setattr(scpc_core, "get_oms", recording_get_oms) + + scpc( + model, + data, + coords_euclidean=("coord_x", "coord_y"), + avc=0.1, + method="exact", + uncond=False, + dtype="float32", + ) + + assert calls + assert all(dist_dtype == np.float32 for dist_dtype, _ in calls) + assert all(dtype == "float32" for _, dtype in calls) + + def test_scpc_stores_only_reported_coef_names_when_ncoef_is_set() -> None: data = pd.DataFrame( {