[codex] expose dtype option

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
 

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,12 +177,13 @@ def scpc(
         latlong,
         method=method,
         large_n_seed=large_n_seed,
+        dtype=dtype,
     )
     d = None
     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
@@ -289,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)

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%"]

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,15 +39,20 @@ 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)
     mat = mat - np.mean(mat, axis=0, keepdims=True)
     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:

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)

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]])

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]]