diff-stream 0.1.0.dev0__py3-none-any.whl

diff_stream/_core/twostage_kernels.py

@@ -0,0 +1,342 @@
+"""Compiled passes of the two-stage (Gardner) fit and its first-stage reduction.
+
+Each pass visits every unit and does a handful of small weighted sums over the
+unit's periods and the P design columns (the period dummies and covariates).
+In NumPy that is a few hundred array operations per batch of units, which the
+interpreter's lock serializes. The kernels do the same sums in one compiled
+call per batch.
+
+Units are processed in fixed-size chunks, and a chunk adds its units in order,
+so the sums, and with them the estimates, do not depend on the number of
+threads; the chunks' partial sums are added in order afterwards. The NumPy code
+in ``linear.py`` and ``gardner.py`` remains the reference; the tests compare
+the two.
+"""
+
+import numpy as np
+from numba import njit, prange
+
+ENABLED = True  # tests switch it off to compare with the array path
+CHUNK = 64  # units per chunk; fixed so results do not depend on threads
+
+
+@njit(cache=True, inline="always")
+def _design(x, i, t, T1, columns, out):
+    """Row (unit i, period t) of the design restricted to ``columns``:
+    period dummies for periods 1.. then covariates, with missing values as 0."""
+    for a in range(columns.shape[0]):
+        j = columns[a]
+        if j < T1:
+            out[t, a] = 1.0 if (t >= 1 and j == t - 1) else 0.0
+        else:
+            v = x[i, t, j - T1]
+            out[t, a] = v if np.isfinite(v) else 0.0
+
+
+@njit(parallel=True, cache=True)
+def reduce_moments(y, x, w, wc, columns, do_clean, Q, q, norms, Qc, qc, bad):
+    """Within-unit moments of all columns for the rows with weights ``w``, and
+    (``Qc``, ``qc``) for the rows with weights ``wc``; ``norms`` are the columns'
+    weighted sums of squares. ``bad[chunk]`` is 1 when a unit has no weight.
+
+    A unit's demeaned design is nonzero only in the columns where one of its
+    rows is, which with fixed-effect dummies is a few of the P (its own levels,
+    not all of them), so when it is under half of them the products visit those
+    columns alone. The other entries would add exact zeros."""
+    n, T = y.shape
+    P = columns.shape[0]
+    T1 = T - 1
+    chunks = (n + CHUNK - 1) // CHUNK
+    for c in prange(chunks):
+        X = np.empty((T, P))
+        xm = np.zeros(P)
+        yv = np.zeros(T)
+        act = np.empty(P, dtype=np.int64)
+        for i in range(c * CHUNK, min(n, (c + 1) * CHUNK)):
+            for t in range(T):
+                _design(x, i, t, T1, columns, X)
+                yv[t] = y[i, t] if np.isfinite(y[i, t]) else 0.0
+            nact = 0
+            for a in range(P):
+                for t in range(T):
+                    if X[t, a] != 0.0:
+                        act[nact] = a
+                        nact += 1
+                        break
+            if 2 * nact >= P:
+                # mostly dense: contiguous loops over every column are faster than
+                # indexed ones, and the skipped entries are exact zeros either way
+                nact = P
+                for a in range(P):
+                    act[a] = a
+            for pass_ in range(2 if do_clean else 1):
+                mass = 0.0
+                for t in range(T):
+                    mass += w[i, t] if pass_ == 0 else wc[i, t]
+                if mass <= 0.0:
+                    bad[c] = 1
+                    continue
+                for a in range(P):
+                    xm[a] = 0.0
+                ym = 0.0
+                for t in range(T):
+                    wt = w[i, t] if pass_ == 0 else wc[i, t]
+                    ym += wt * yv[t]
+                    for ia in range(nact):
+                        a = act[ia]
+                        xm[a] += wt * X[t, a]
+                ym /= mass
+                for ia in range(nact):
+                    a = act[ia]
+                    xm[a] /= mass
+                for t in range(T):
+                    wt = w[i, t] if pass_ == 0 else wc[i, t]
+                    if wt == 0.0:
+                        continue
+                    ry = yv[t] - ym
+                    for ia in range(nact):
+                        a = act[ia]
+                        xa = (X[t, a] - xm[a]) * wt
+                        if pass_ == 0:
+                            q[c, a] += xa * ry
+                            if nact == P:
+                                # every column is active: a contiguous loop vectorizes
+                                for b in range(P):
+                                    Q[c, a, b] += xa * (X[t, b] - xm[b])
+                            else:
+                                for ib in range(nact):
+                                    b = act[ib]
+                                    Q[c, a, b] += xa * (X[t, b] - xm[b])
+                        else:
+                            qc[c, a] += xa * ry
+                            if nact == P:
+                                for b in range(P):
+                                    Qc[c, a, b] += xa * (X[t, b] - xm[b])
+                            else:
+                                for ib in range(nact):
+                                    b = act[ib]
+                                    Qc[c, a, b] += xa * (X[t, b] - xm[b])
+            for t in range(T):
+                for ia in range(nact):
+                    a = act[ia]
+                    norms[c, a] += w[i, t] * X[t, a] * X[t, a]
+
+
+@njit(cache=True, inline="always")
+def _first_stage(g, y, x, w, i, T1, columns, theta, times, anticipation, X, xm, u1, score1, obs):
+    """One unit's first-stage residuals ``u1`` (clean-sample imputation), the
+    clean-weighted mean ``xm`` of its design and its first-stage score.
+    Returns 0, or 1 when the unit has no untreated observation."""
+    T = times.shape[0]
+    P = columns.shape[0]
+    for t in range(T):
+        _design(x, i, t, T1, columns, X)
+    mass = 0.0
+    for a in range(P):
+        xm[a] = 0.0
+    alpha = 0.0
+    for t in range(T):
+        obs[t] = np.isfinite(y[i, t])
+        use = obs[t] and (g[i] == 0 or times[t] < g[i] - anticipation)
+        cw = w[i, t] if use else 0.0
+        fit = 0.0
+        for a in range(P):
+            fit += X[t, a] * theta[a]
+        u1[t] = (y[i, t] - fit) if obs[t] else 0.0
+        mass += cw
+        alpha += cw * u1[t]
+        for a in range(P):
+            xm[a] += cw * X[t, a]
+    if mass <= 0.0:
+        return 1
+    alpha /= mass
+    for a in range(P):
+        xm[a] /= mass
+        score1[a] = 0.0
+    for t in range(T):
+        use = obs[t] and (g[i] == 0 or times[t] < g[i] - anticipation)
+        cw = w[i, t] if use else 0.0
+        u1[t] = (u1[t] - alpha) if obs[t] else 0.0
+        for a in range(P):
+            score1[a] += cw * u1[t] * (X[t, a] - xm[a])
+    return 0
+
+
+@njit(cache=True, inline="always")
+def _indicator(g, i, t, times, anticipation, horizons, m, dynamic):
+    """Second-stage indicator of coefficient m for unit i in period t (before
+    the observed-row condition)."""
+    if g[i] == 0:
+        return 0.0
+    event = times[t] - g[i]
+    if dynamic:
+        return 1.0 if event == horizons[m] else 0.0
+    return 1.0 if event >= -anticipation else 0.0
+
+
+@njit(parallel=True, cache=True)
+def second_stage_sums(
+    g, y, x, w, columns, theta, times, anticipation, horizons, dynamic, mass, rhs, C, bad
+):
+    """The second stage's sums: Z'w, Z'w u1 and Z'w (X - xm) (``gardner._fit``)."""
+    n, T = y.shape
+    P = columns.shape[0]
+    M = mass.shape[1]
+    T1 = T - 1
+    chunks = (n + CHUNK - 1) // CHUNK
+    for c in prange(chunks):
+        X = np.empty((T, P))
+        xm = np.zeros(P)
+        u1 = np.zeros(T)
+        s1 = np.zeros(P)
+        obs = np.zeros(T, dtype=np.bool_)
+        for i in range(c * CHUNK, min(n, (c + 1) * CHUNK)):
+            if _first_stage(
+                g, y, x, w, i, T1, columns, theta, times, anticipation, X, xm, u1, s1, obs
+            ):
+                bad[c] = 1
+                continue
+            for t in range(T):
+                if not obs[t]:
+                    continue
+                for m in range(M):
+                    z = _indicator(g, i, t, times, anticipation, horizons, m, dynamic)
+                    if z == 0.0:
+                        continue
+                    mass[c, m] += w[i, t]
+                    rhs[c, m] += w[i, t] * u1[t]
+                    for a in range(P):
+                        C[c, m, a] += w[i, t] * (X[t, a] - xm[a])
+
+
+@njit(parallel=True, cache=True)
+def cluster_scores(
+    g, y, x, w, columns, theta, times, anticipation, horizons, dynamic, beta, H, mass, out, bad
+):
+    """Each unit's second-stage score, corrected for the first stage and scaled
+    by the indicators' mass: ``out[unit, m]``."""
+    n, T = y.shape
+    P = columns.shape[0]
+    M = beta.shape[0]
+    T1 = T - 1
+    chunks = (n + CHUNK - 1) // CHUNK
+    for c in prange(chunks):
+        X = np.empty((T, P))
+        xm = np.zeros(P)
+        u1 = np.zeros(T)
+        s1 = np.zeros(P)
+        obs = np.zeros(T, dtype=np.bool_)
+        z = np.zeros(M)
+        for i in range(c * CHUNK, min(n, (c + 1) * CHUNK)):
+            if _first_stage(
+                g, y, x, w, i, T1, columns, theta, times, anticipation, X, xm, u1, s1, obs
+            ):
+                bad[c] = 1
+                continue
+            score2 = np.zeros(M)
+            for t in range(T):
+                fitted = 0.0
+                for m in range(M):
+                    z[m] = _indicator(g, i, t, times, anticipation, horizons, m, dynamic)
+                    if not obs[t]:
+                        z[m] = 0.0
+                    fitted += z[m] * beta[m]
+                u2 = u1[t] - fitted
+                for m in range(M):
+                    score2[m] += w[i, t] * u2 * z[m]
+            for m in range(M):
+                adjust = 0.0
+                for a in range(P):
+                    adjust += s1[a] * H[a, m]
+                out[i, m] = (score2[m] - adjust) / mass[m]
+
+
+def _arrays(raw):
+    g, y, x, w = raw
+    g = np.ascontiguousarray(g, dtype=np.int64)
+    y = np.ascontiguousarray(y, dtype=np.float64)
+    x = np.ascontiguousarray(x, dtype=np.float64)
+    w = np.ascontiguousarray(w, dtype=np.float64)
+    return g, y, x, w
+
+
+def _chunks(n):
+    return (n + CHUNK - 1) // CHUNK
+
+
+class TwoStageKernels:
+    """The compiled passes for a store: ``times`` are its period labels."""
+
+    def __init__(self, times, anticipation):
+        self.times = np.ascontiguousarray(times, dtype=np.int64)
+        self.anticipation = int(anticipation)
+
+    def reduce(self, batches, n_columns, clean):
+        """(Q, q, norms, Qc, qc) over the batches of (g, y, x, w), all columns kept."""
+        columns = np.arange(n_columns, dtype=np.int64)
+        P = n_columns
+        Q, norms, Qc, qc = np.zeros((P, P)), np.zeros(P), np.zeros((P, P)), np.zeros(P)
+        q = np.zeros(P)
+        for raw in batches():
+            g, y, x, w = _arrays(raw)
+            wc = w * clean(g) if clean is not None else w
+            k = _chunks(len(g))
+            parts = dict(
+                Q=np.zeros((k, P, P)),
+                q=np.zeros((k, P)),
+                norms=np.zeros((k, P)),
+                Qc=np.zeros((k, P, P)),
+                qc=np.zeros((k, P)),
+                bad=np.zeros(k, dtype=np.int8),
+            )
+            reduce_moments(
+                y, x, w, wc, columns, clean is not None, parts["Q"], parts["q"], parts["norms"], parts["Qc"],
+                parts["qc"], parts["bad"],
+            )  # fmt: skip
+            if parts["bad"].any():
+                raise ValueError("each group requires a positive-weight estimation observation")
+            Q += parts["Q"].sum(axis=0)
+            q += parts["q"].sum(axis=0)
+            norms += parts["norms"].sum(axis=0)
+            Qc += parts["Qc"].sum(axis=0)
+            qc += parts["qc"].sum(axis=0)
+        return Q, q, norms, Qc, qc
+
+    def second_stage(self, batches, keep, theta, horizons, dynamic):
+        """(mass, rhs, C): the sums of the second stage's normal equations."""
+        columns = np.ascontiguousarray(keep, dtype=np.int64)
+        M, P = len(horizons), len(columns)
+        mass, rhs, C = np.zeros(M), np.zeros(M), np.zeros((M, P))
+        hz = np.ascontiguousarray(horizons, dtype=np.int64)
+        for raw in batches():
+            g, y, x, w = _arrays(raw)
+            k = _chunks(len(g))
+            parts = (
+                np.zeros((k, M)),
+                np.zeros((k, M)),
+                np.zeros((k, M, P)),
+                np.zeros(k, dtype=np.int8),
+            )
+            second_stage_sums(
+                g, y, x, w, columns, theta, self.times, self.anticipation, hz, dynamic, *parts
+            )
+            if parts[3].any():
+                raise ValueError("requires at least one untreated observation per unit")
+            mass += parts[0].sum(axis=0)
+            rhs += parts[1].sum(axis=0)
+            C += parts[2].sum(axis=0)
+        return mass, rhs, C
+
+    def scores(self, raw, keep, theta, horizons, dynamic, beta, H, mass):
+        columns = np.ascontiguousarray(keep, dtype=np.int64)
+        g, y, x, w = _arrays(raw)
+        out = np.zeros((len(g), len(beta)))
+        bad = np.zeros(_chunks(len(g)), dtype=np.int8)
+        cluster_scores(
+            g, y, x, w, columns, theta, self.times, self.anticipation,
+            np.ascontiguousarray(horizons, dtype=np.int64), dynamic, beta,
+            np.ascontiguousarray(H), mass, out, bad,
+        )  # fmt: skip
+        if bad.any():
+            raise ValueError("requires at least one untreated observation per unit")
+        return out

diff_stream/_core/weights.py

@@ -0,0 +1,26 @@
+"""Sampling-weight validation and optional weighted batch adapters.
+
+No treatment or panel conventions live here. Consumer code owns restrictions
+such as weights being constant within a sampling unit.
+"""
+
+import numpy as np
+
+
+def validate_weights(values):
+    """Require finite, strictly positive numeric sampling weights."""
+    if not np.isfinite(values).all() or np.any(values <= 0):
+        raise ValueError("sampling weights must be finite and strictly positive")
+
+
+def weighted_system(parts):
+    """Accept (X, d, Y[, weights]); preserve the unweighted consumer protocol."""
+    X, d, Y = parts[:3]
+    w = parts[3] if len(parts) == 4 else np.ones(len(d))
+    return X, d, Y, w
+
+
+def weighted_linear(parts):
+    """Accept (X, y[, weights]) for a weighted least-squares stratum."""
+    X, y = parts[:2]
+    return X, y, parts[2] if len(parts) == 3 else np.ones(len(y))

diff_stream/_core/within.py

@@ -0,0 +1,81 @@
+"""Bounded weighted elimination of group intercepts and small Gram reduction.
+
+Arrays contain one batch of whole groups. No treatment, time or panel policy
+lives here; these helpers are candidates for hdfe_stream.core.
+"""
+
+import numpy as np
+from scipy.linalg import solve_triangular
+
+
+def center_groups(X, y, weights):
+    """Eliminate intercepts using row weights (zero means excluded internally)."""
+    mass = weights.sum(axis=1)
+    if np.any(mass <= 0):
+        raise ValueError("each group requires a positive-weight estimation observation")
+    xm = row_sums(weights, X) / mass[:, None]
+    ym = (weights * y).sum(axis=1) / mass
+    return X - xm[:, None, :], y - ym[:, None], xm, ym
+
+
+def moments(X, y, weights):
+    xc, yc, _, _ = center_groups(X, y, weights)
+    flat = xc.reshape(-1, xc.shape[2])
+    weighted = weights.reshape(-1, 1) * flat
+    return flat.T @ weighted, weighted.T @ yc.ravel()
+
+
+# Batched sums over the period axis of (units, periods, columns) arrays, as
+# matrix products: einsum would copy its operands and run single-threaded.
+
+
+def row_sums(weights, X):
+    """sum_t weights[b, t] * X[b, t, k] -> (units, columns)."""
+    return np.matmul(weights[:, None, :], X)[:, 0, :]
+
+
+def weighted_gram(X, weights):
+    """sum_{b,t} weights[b, t] * X[b, t, i] * X[b, t, j] -> (columns, columns)."""
+    flat = X.reshape(-1, X.shape[2])
+    return flat.T @ (weights.reshape(-1, 1) * flat)
+
+
+def weighted_squares(X, weights):
+    """sum_{b,t} weights[b, t] * X[b, t, k] ** 2 -> (columns,)."""
+    flat = X.reshape(-1, X.shape[2])
+    return (flat * flat).T @ weights.ravel()
+
+
+def independent_columns(gram, *, original_norm=None, tolerance=1e-10):
+    """Select an ordered basis after scaling columns to unit Gram diagonal.
+
+    Earlier columns have priority. Return aliases explicitly for the caller's
+    diagnostics; consumers must separately protect identified target columns.
+
+    A column is kept when its squared distance from the span of the columns kept
+    before it exceeds `tolerance`: an ordered Cholesky factorization that grows
+    one row per kept column (a triangular solve per column, O(K^3) in all).
+    """
+    diagonal = np.maximum(np.diag(gram), 0)
+    scale = np.sqrt(diagonal)
+    live = diagonal > 0 if original_norm is None else diagonal > original_norm * 1e-24
+    safe = np.where(live, scale, 1)
+    correlation = gram / safe[:, None] / safe[None, :]
+    candidates = np.flatnonzero(live)
+    factor = np.zeros((len(candidates), len(candidates)))  # lower triangular
+    kept = []
+    for j in candidates:
+        m = len(kept)
+        if m:
+            row = solve_triangular(
+                factor[:m, :m], correlation[j, kept], lower=True, check_finite=False
+            )
+            remaining = correlation[j, j] - row @ row
+        else:
+            row = np.zeros(0)
+            remaining = correlation[j, j]
+        if remaining > tolerance:
+            factor[m, :m] = row
+            factor[m, m] = np.sqrt(remaining)
+            kept.append(j)
+    return np.asarray(kept, dtype=int)

diff_stream/_docs.py

@@ -0,0 +1,100 @@
+"""Parameter descriptions shared by the public functions' docstrings."""
+
+DATA = """\
+    data : polars.DataFrame, polars.LazyFrame, str or Path
+        The panel, or a Parquet path or glob. Rows are read in streaming
+        batches; the data need not fit in memory.
+    panel : tuple of str
+        The unit and time columns, ``(unit, time)``. Periods are positive
+        integers, evenly spaced (every year, every second year, ...). Event
+        times, ``anticipation`` and ``pretrend_leads`` count periods; results
+        label cohorts and periods as in the data.
+    cohort, treatment : str
+        Exactly one. ``cohort``: each unit's first treated period, constant
+        within the unit; 0 or null marks a never-treated unit. ``treatment``: a
+        0/1 column, absorbing (once 1, a unit stays treated), from which each
+        unit's first treated period is taken."""
+
+ANTICIPATION = """\
+    anticipation : int, default 0
+        Periods before treatment in which units may already respond. Clean
+        (untreated) observations are those before ``cohort - anticipation``;
+        event times stay relative to actual treatment."""
+
+WEIGHTS = """\
+    weights : str, optional
+        A column of sampling weights: positive, finite, and constant within
+        each unit unless ``allow_time_varying_weights=True``. They enter the
+        estimates, aggregation and clustered standard errors.
+    allow_time_varying_weights : bool, default False
+        Accept weights that change over time within a unit."""
+
+INCOMPLETE_UNITS = """\
+    incomplete_units : {"error", "drop", "keep"}, default "error"
+        What to do with units missing some periods: refuse the data, drop
+        those units, or keep their observed rows (an unbalanced panel)."""
+
+LAST_TREATED = """\
+    last_treated_controls : {"if_needed", "never", "always"}, default "if_needed"
+        When to use the last-treated cohort as a comparison group, keeping
+        only periods before it is treated. "if_needed" does so only when there
+        are no never-treated units (with a ``ReferenceWindowWarning``);
+        "never" requires never-treated units; "always" uses it even when
+        never-treated units exist (they remain controls too)."""
+
+FIXED_EFFECT_SOLVER = """\
+    fixed_effect_solver : {"auto", "dense", "iterative"}, default "auto"
+        How additional fixed effects (after the bar in ``fml``) are
+        estimated. "dense" adds a dummy for each level to the regression,
+        which is exact and fastest with few levels, and limited to 256 columns in
+        all. "iterative" absorbs them by alternating projections (hdfe-stream),
+        with a time that does not grow with the number of levels and no limit
+        on it, such as county-by-year. "auto" uses dense for up to 40 columns
+        (levels and covariates) and iterative above that. Both give the same
+        estimates and standard errors."""
+
+VCOV = """\
+    vcov : {"CRV0", "CRV1"}, default "CRV0"
+        Standard errors clustered by unit. CRV1 multiplies the variance by
+        N/(N-1)."""
+
+RESOURCES = """\
+    n_threads : int, optional
+        Threads for numba and BLAS during the fit (stream-limits'
+        ``thread_scope``). Default: every core the job has. Polars' threads
+        are fixed when polars is imported: set ``STREAM_LIMITS_MAX_THREADS``
+        in the job script.
+    max_memory_gb, memory_reserve : float, optional
+        The memory limit and the share of it kept in reserve for allocations
+        the estimates do not model (stream-limits' ``MemoryBudget``; default:
+        the limits it detects, including a scheduler job's, and a 15%
+        reserve). Fits that would exceed the limit stop with ``MemoryError``
+        before allocating.
+    workdir : str or Path, optional
+        Where temporary files are written. Default: the system temporary
+        directory.
+    profile : bool, default False
+        Return a ``Profile`` (hdfe-stream's) of the time, memory and disk the
+        fit needs, without fitting.
+    batch_units : int, optional
+        Advanced: units per numerical batch.
+    preprocess_mb : float, optional
+        Advanced: the working memory for sorting and validating the data, in
+        MB. Default: a quarter of the memory left under ``max_memory_gb``, so a
+        panel that fits is sorted in one piece; larger panels are split into
+        whole-unit pieces on disk. A smaller value uses less memory and more
+        time."""
+
+
+def document(function):
+    """Fill the shared parameter blocks into a public function's docstring."""
+    blocks = {
+        name: value
+        for name, value in globals().items()
+        if name.isupper() and isinstance(value, str)
+    }
+    text = function.__doc__
+    for name, value in blocks.items():
+        text = text.replace("    {" + name + "}\n", value + "\n")
+    function.__doc__ = text
+    return function

diff_stream/_options.py

@@ -0,0 +1,62 @@
+"""Public option names and values, translated to the estimators' internal ones.
+
+The public functions spell options out (``incomplete_units="drop"``); the
+estimation code was written against shorter internal names
+(``panel_mode="complete"``). Translation happens once, here, so that the
+estimator functions and refit_bootstrap() accept exactly the same options.
+"""
+
+# public option name -> internal option name
+NAMES = {
+    "n_threads": "threads",
+    "max_memory_gb": "memory_gb",
+    "incomplete_units": "panel_mode",
+    "last_treated_controls": "reference_group",
+    "adjustment": "est_method",
+    "propensity_trim": "trim_level",
+    "multiplier_bootstrap_reps": "bootstrap_reps",
+    "multiplier_bootstrap_seed": "bootstrap_seed",
+    "multiplier_bootstrap_weights": "bootstrap_weights",
+    "effective_sample_warning": "ess_warn_threshold",
+    "min_effective_sample": "ess_suppress_threshold",
+    "covariate_balance": "balance",
+    "balance_distributions": "balance_ecdf",
+    "residual_averaging": "avgeffectsby",
+    "leave_out": "leaveout",
+    "pretrend_leads": "pretrends",
+    "event_study": "specification",
+    "treat_sample_as_fixed": "inference",
+    "zero_tolerance": "zero_tol",
+    "fixed_effect_solver": "fe_solver",
+}
+
+# public option name -> {public value: internal value}
+VALUES = {
+    "control_group": {"never_treated": "nevertreated", "not_yet_treated": "notyettreated"},
+    "incomplete_units": {"error": "balanced", "drop": "complete", "keep": "unbalanced"},
+    "last_treated_controls": {
+        "if_needed": "auto",
+        "never": "nevertreated",
+        "always": "lasttreated",
+    },
+    "adjustment": {"regression": "reg", "ipw": "ipw", "doubly_robust": "dr"},
+    "event_study": {True: "dynamic", False: "static"},
+    "treat_sample_as_fixed": {False: "unconditional", True: "conditional"},
+}
+
+
+def internal(options):
+    """Translate public keyword options to the internal estimator's keywords."""
+    out = {}
+    for name, value in options.items():
+        if name in VALUES:
+            choices = VALUES[name]
+            if isinstance(next(iter(choices)), bool):
+                if not isinstance(value, bool):
+                    raise ValueError(f"{name} must be True or False")
+            elif not isinstance(value, str) or value not in choices:
+                listed = ", ".join(repr(c) for c in choices)
+                raise ValueError(f"{name} must be one of {listed}; got {value!r}")
+            value = choices[value]
+        out[NAMES.get(name, name)] = value
+    return out