skxperiments 0.1.0.dev0__py3-none-any.whl

skxperiments/design/balance.py

@@ -0,0 +1,182 @@
+"""Covariate balance diagnostics for experimental designs.
+
+Provides check_balance, a standalone function that computes the
+standardized mean difference (SMD) between treatment and control
+groups for each covariate in an Assignment.
+"""
+
+import numpy as np
+import pandas as pd
+
+from skxperiments.core.assignment import BaseAssignment
+from skxperiments.core.exceptions import (
+    InsufficientDataError,
+    InvalidDesignError,
+)
+
+
+def check_balance(
+    assignment: BaseAssignment,
+    covariates: list[str] | None = None,
+) -> pd.DataFrame:
+    """Compute covariate balance between treatment and control groups.
+
+    For each covariate, returns the mean in each group, the pooled
+    standard deviation, and the standardized mean difference (SMD).
+
+    Parameters
+    ----------
+    assignment : BaseAssignment
+        Assignment object produced by a design. Must expose ``data_``
+        (DataFrame with treatment column attached) and ``treatment_col_``.
+    covariates : list of str or None, optional
+        Names of covariates to check. If None, all numeric columns in
+        ``assignment.data_`` except the treatment column are used, in
+        the order they appear in the DataFrame. Boolean columns count
+        as numeric. By default None.
+
+    Returns
+    -------
+    pd.DataFrame
+        DataFrame with one row per covariate and columns:
+        ``covariate``, ``mean_treated``, ``mean_control``,
+        ``std_pooled``, ``smd``. The index is a default RangeIndex.
+
+    Raises
+    ------
+    InvalidDesignError
+        If a name in ``covariates`` is not a column of ``assignment.data_``,
+        or if any selected covariate contains NaN values.
+    InsufficientDataError
+        If ``covariates`` is None and no numeric columns are available
+        after excluding the treatment column.
+
+    Notes
+    -----
+    The pooled standard deviation follows the convention common in
+    the SMD literature for randomized experiments (Austin 2009;
+    Stuart 2010):
+
+        std_pooled = sqrt((var_treated + var_control) / 2)
+
+    where each variance is computed with ``ddof=1``. When
+    ``std_pooled == 0`` (no within-group variation), the SMD is NaN
+    rather than raising an exception.
+
+    The function does not modify ``assignment.data_``.
+
+    References
+    ----------
+    Austin, P. C. (2009). Balance diagnostics for comparing the
+    distribution of baseline covariates between treatment groups in
+    propensity-score matched samples. Statistics in Medicine.
+
+    Stuart, E. A. (2010). Matching methods for causal inference:
+    A review and a look forward. Statistical Science.
+
+    Examples
+    --------
+    >>> import numpy as np
+    >>> import pandas as pd
+    >>> from skxperiments.core.assignment import CRDAssignment
+    >>> rng = np.random.default_rng(42)
+    >>> df = pd.DataFrame({
+    ...     "x1": rng.normal(size=100),
+    ...     "x2": rng.normal(size=100),
+    ...     "treatment": rng.integers(0, 2, size=100),
+    ... })
+    >>> assignment = CRDAssignment(
+    ...     data=df, treatment_col="treatment", design=None, seed=42
+    ... )
+    >>> result = check_balance(assignment)
+    >>> set(result.columns) == {
+    ...     "covariate", "mean_treated", "mean_control",
+    ...     "std_pooled", "smd",
+    ... }
+    True
+    """
+    data = assignment.data_
+    treatment_col = assignment.treatment_col_
+
+    # Resolve covariate list
+    if covariates is None:
+        numeric_cols = [
+            col
+            for col in data.columns
+            if col != treatment_col and pd.api.types.is_numeric_dtype(data[col])
+        ]
+        if len(numeric_cols) == 0:
+            raise InsufficientDataError(
+                context=(
+                    "check_balance with covariates=None "
+                    "(no numeric columns available after excluding "
+                    f"treatment column '{treatment_col}')"
+                ),
+                minimum=1,
+                received=0,
+            )
+        selected = numeric_cols
+    else:
+        missing = [c for c in covariates if c not in data.columns]
+        if missing:
+            raise InvalidDesignError(
+                f"Covariates not found in assignment.data_: {missing}. "
+                f"Available columns: {list(data.columns)}."
+            )
+        selected = list(covariates)
+
+    # Validate no NaN in any selected covariate
+    cols_with_nan = [c for c in selected if data[c].isna().any()]
+    if cols_with_nan:
+        raise InvalidDesignError(
+            f"Covariates contain NaN values: {cols_with_nan}. "
+            f"check_balance requires complete data; impute or drop NaN "
+            f"before calling."
+        )
+
+    # Compute group masks
+    treatment_values = data[treatment_col].values
+    treated_mask = treatment_values == 1
+    control_mask = treatment_values == 0
+
+    # Compute statistics per covariate
+    rows: list[dict[str, float | str]] = []
+    for cov in selected:
+        values = data[cov].astype(float).values
+        treated_vals = values[treated_mask]
+        control_vals = values[control_mask]
+
+        mean_t = float(np.mean(treated_vals))
+        mean_c = float(np.mean(control_vals))
+
+        var_t = float(np.var(treated_vals, ddof=1))
+        var_c = float(np.var(control_vals, ddof=1))
+
+        std_pooled = float(np.sqrt((var_t + var_c) / 2.0))
+
+        if std_pooled == 0.0:
+            smd: float = float("nan")
+        else:
+            smd = (mean_t - mean_c) / std_pooled
+
+        rows.append(
+            {
+                "covariate": cov,
+                "mean_treated": mean_t,
+                "mean_control": mean_c,
+                "std_pooled": std_pooled,
+                "smd": smd,
+            }
+        )
+
+    result = pd.DataFrame(
+        rows,
+        columns=[
+            "covariate",
+            "mean_treated",
+            "mean_control",
+            "std_pooled",
+            "smd",
+        ],
+    )
+    return result.reset_index(drop=True)

skxperiments/design/blocked_crd.py

@@ -0,0 +1,157 @@
+"""Blocked Completely Randomized Design.
+
+Randomizes treatment independently within each block, preserving the
+treatment proportion within every block. Useful when there are
+pre-experiment covariates that define meaningful subgroups (e.g.,
+geography, device type) where balance must be guaranteed.
+"""
+
+import numpy as np
+import pandas as pd
+
+from skxperiments.core.assignment import BlockedAssignment
+from skxperiments.core.base import BaseDesign
+from skxperiments.core.exceptions import (
+    InsufficientDataError,
+    InvalidDesignError,
+)
+
+
+class BlockedCRD(BaseDesign):
+    """Blocked Completely Randomized Design.
+
+    Treatment is randomized independently within each block defined by
+    ``block_col``. The treatment proportion ``p`` is applied uniformly
+    to all blocks. Within each block, the number of treated units is
+    ``round(p * n_block)``.
+
+    Parameters
+    ----------
+    block_col : str
+        Name of the column in the DataFrame that defines blocks.
+    p : float
+        Treatment proportion in (0, 1), applied uniformly across blocks.
+    seed : int or None, optional
+        Random seed for reproducibility, by default None.
+    treatment_col : str, optional
+        Name of the treatment column to be added to the output, by
+        default ``"treatment"``.
+
+    Examples
+    --------
+    >>> import pandas as pd
+    >>> df = pd.DataFrame({
+    ...     "x": range(8),
+    ...     "region": ["A", "A", "A", "A", "B", "B", "B", "B"],
+    ... })
+    >>> design = BlockedCRD(block_col="region", p=0.5, seed=42)
+    >>> assignment = design.randomize(df)
+    >>> assignment.block_sizes_
+    {'A': 4, 'B': 4}
+    """
+
+    def __init__(
+        self,
+        block_col: str,
+        p: float | None = None,
+        seed: int | None = None,
+        treatment_col: str = "treatment",
+    ) -> None:
+        if p is None:
+            raise InvalidDesignError(
+                "BlockedCRD requires a treatment proportion p; "
+                "received p=None."
+            )
+        if not (0.0 < p < 1.0):
+            raise InvalidDesignError(
+                f"Treatment proportion p must be in (0, 1), but received {p}."
+            )
+
+        self.block_col = block_col
+        self.p = p
+        self.seed = seed
+        self.treatment_col = treatment_col
+
+    def randomize(self, df: pd.DataFrame) -> BlockedAssignment:
+        """Perform blocked randomization and return a BlockedAssignment.
+
+        Parameters
+        ----------
+        df : pd.DataFrame
+            DataFrame containing the experimental units. Must contain
+            ``block_col`` and must not contain ``treatment_col``.
+
+        Returns
+        -------
+        BlockedAssignment
+            Assignment with treatment column added. Original DataFrame
+            is not modified.
+
+        Raises
+        ------
+        InvalidDesignError
+            If ``block_col`` is missing from ``df``, if ``treatment_col``
+            already exists in ``df``, or if rounding results in 0 or n
+            treated units in any block.
+        InsufficientDataError
+            If any block has fewer than 2 units.
+        """
+        if self.block_col not in df.columns:
+            raise InvalidDesignError(
+                f"Block column '{self.block_col}' not found in DataFrame. "
+                f"Available columns: {list(df.columns)}."
+            )
+        if self.treatment_col in df.columns:
+            raise InvalidDesignError(
+                f"Treatment column '{self.treatment_col}' already exists "
+                f"in DataFrame. Drop or rename it before calling randomize()."
+            )
+
+        df_out = df.copy()
+        treatment = np.zeros(len(df_out), dtype=int)
+
+        rng = np.random.default_rng(self.seed)
+
+        # Iterate blocks in stable order (sorted by label) for
+        # reproducibility independent of pandas grouping internals.
+        block_labels = sorted(df_out[self.block_col].unique(), key=lambda x: str(x))
+
+        block_sizes: dict = {}
+
+        for label in block_labels:
+            block_iloc = np.where(df_out[self.block_col].values == label)[0]
+            n_block = len(block_iloc)
+            block_sizes[label] = n_block
+
+            if n_block < 2:
+                raise InsufficientDataError(
+                    context=(
+                        f"BlockedCRD randomization for block '{label}'"
+                    ),
+                    minimum=2,
+                    received=n_block,
+                )
+
+            n_treated_block = int(round(self.p * n_block))
+
+            if n_treated_block == 0 or n_treated_block == n_block:
+                raise InvalidDesignError(
+                    f"Block '{label}' has size {n_block}; with p={self.p}, "
+                    f"rounding yields {n_treated_block} treated units. "
+                    f"Each block must have at least 1 treated and 1 control "
+                    f"unit. Increase block size or adjust p."
+                )
+
+            chosen = rng.choice(block_iloc, size=n_treated_block, replace=False)
+            treatment[chosen] = 1
+
+        df_out[self.treatment_col] = treatment
+
+        return BlockedAssignment(
+            data=df_out,
+            treatment_col=self.treatment_col,
+            design=self,
+            block_col=self.block_col,
+            block_sizes=block_sizes,
+            seed=self.seed,
+        )

skxperiments/design/crd.py

@@ -0,0 +1,162 @@
+"""Completely Randomized Design (CRD).
+
+Assigns units to treatment uniformly at random, with either a fixed
+absolute count of treated units or a fixed treatment proportion.
+"""
+
+import numpy as np
+import pandas as pd
+
+from skxperiments.core.assignment import CRDAssignment
+from skxperiments.core.base import BaseDesign
+from skxperiments.core.exceptions import (
+    InsufficientDataError,
+    InvalidDesignError,
+)
+
+
+class CRD(BaseDesign):
+    """Completely Randomized Design.
+
+    Treatment is assigned uniformly at random to a fixed number of
+    units. The user provides exactly one of ``n_treated`` (absolute
+    count) or ``p`` (proportion); rounding for ``p`` uses
+    ``round(p * n)``.
+
+    Parameters
+    ----------
+    n_treated : int or None, optional
+        Absolute number of units to assign to treatment. Mutually
+        exclusive with ``p``. By default None.
+    p : float or None, optional
+        Treatment proportion in (0, 1). Mutually exclusive with
+        ``n_treated``. By default None.
+    seed : int or None, optional
+        Random seed for reproducibility, by default None.
+    treatment_col : str, optional
+        Name of the treatment column to be added to the output, by
+        default ``"treatment"``.
+
+    Examples
+    --------
+    >>> import pandas as pd
+    >>> df = pd.DataFrame({"x": range(100)})
+    >>> design = CRD(p=0.5, seed=42)
+    >>> assignment = design.randomize(df)
+    >>> assignment.n_treated_
+    50
+    """
+
+    def __init__(
+        self,
+        n_treated: int | None = None,
+        p: float | None = None,
+        seed: int | None = None,
+        treatment_col: str = "treatment",
+    ) -> None:
+        # Mutual exclusivity: exactly one of n_treated or p.
+        if n_treated is None and p is None:
+            raise InvalidDesignError(
+                "CRD requires exactly one of n_treated or p; both "
+                "are None."
+            )
+        if n_treated is not None and p is not None:
+            raise InvalidDesignError(
+                "CRD requires exactly one of n_treated or p; both "
+                "were provided."
+            )
+
+        if n_treated is not None:
+            if not isinstance(n_treated, (int, np.integer)) or n_treated <= 0:
+                raise InvalidDesignError(
+                    f"n_treated must be a positive integer, but received "
+                    f"{n_treated!r}."
+                )
+
+        if p is not None:
+            if not isinstance(p, (int, float)) or not (0.0 < p < 1.0):
+                raise InvalidDesignError(
+                    f"p must be in (0, 1), but received {p!r}."
+                )
+
+        self.n_treated = n_treated
+        self.p = p
+        self.seed = seed
+        self.treatment_col = treatment_col
+
+    def randomize(self, df: pd.DataFrame) -> CRDAssignment:
+        """Perform complete randomization and return a CRDAssignment.
+
+        Parameters
+        ----------
+        df : pd.DataFrame
+            DataFrame with experimental units. Must not contain
+            ``treatment_col``.
+
+        Returns
+        -------
+        CRDAssignment
+            Assignment with the treatment column added. The original
+            DataFrame is not modified.
+
+        Raises
+        ------
+        InvalidDesignError
+            If ``treatment_col`` already exists in ``df``, or if the
+            resolved number of treated is 0 or N (no treatment
+            contrast possible).
+        InsufficientDataError
+            If ``len(df) < n_treated`` (when ``n_treated`` was given),
+            or if ``len(df) < 2`` (no contrast possible).
+        """
+        n_total = len(df)
+
+        if n_total < 2:
+            raise InsufficientDataError(
+                context="CRD randomization",
+                minimum=2,
+                received=n_total,
+            )
+
+        if self.treatment_col in df.columns:
+            raise InvalidDesignError(
+                f"Treatment column '{self.treatment_col}' already "
+                f"exists in DataFrame. Drop or rename it before "
+                f"calling randomize()."
+            )
+
+        # Resolve n_treated.
+        if self.n_treated is not None:
+            if n_total < self.n_treated:
+                raise InsufficientDataError(
+                    context="CRD randomization",
+                    minimum=self.n_treated,
+                    received=n_total,
+                )
+            n_treated_resolved = self.n_treated
+        else:
+            n_treated_resolved = int(round(self.p * n_total))
+
+        if n_treated_resolved <= 0 or n_treated_resolved >= n_total:
+            raise InvalidDesignError(
+                f"Resolved n_treated={n_treated_resolved} for N={n_total}; "
+                f"must be strictly between 0 and N. Adjust n_treated "
+                f"or p."
+            )
+
+        # Defensive copy. Build treatment vector.
+        df_out = df.copy()
+        rng = np.random.default_rng(self.seed)
+
+        treatment = np.zeros(n_total, dtype=int)
+        chosen = rng.choice(n_total, size=n_treated_resolved, replace=False)
+        treatment[chosen] = 1
+
+        df_out[self.treatment_col] = treatment
+
+        return CRDAssignment(
+            data=df_out,
+            treatment_col=self.treatment_col,
+            design=self,
+            seed=self.seed,
+        )

skxperiments/design/factorial.py

@@ -0,0 +1,174 @@
+"""2^K Factorial Design.
+
+Randomly assigns units to one of 2^K cells defined by the values of
+K binary factors, with all cells of equal size.
+"""
+
+import numpy as np
+import pandas as pd
+
+from skxperiments.core.assignment import FactorialAssignment
+from skxperiments.core.base import BaseDesign
+from skxperiments.core.exceptions import (
+    InsufficientDataError,
+    InvalidDesignError,
+)
+
+
+class FactorialDesign(BaseDesign):
+    """2^K Factorial Design with equal cell sizes.
+
+    Randomly assigns units to one of 2^K cells defined by K binary
+    factors. Each cell contains exactly ``n_per_cell`` units, so the
+    DataFrame must have ``n_per_cell * 2^K`` rows.
+
+    Parameters
+    ----------
+    factors : list of str
+        Names of the K factors. These will be added as columns to the
+        Assignment's data, alongside the synthetic ``"_cell"`` column.
+        Must be non-empty and contain no duplicates.
+    n_per_cell : int
+        Number of units per cell. All cells have equal size in this
+        version. Must be >= 1.
+    seed : int or None, optional
+        Random seed for reproducibility, by default None.
+
+    Notes
+    -----
+    Cell encoding convention (little-endian):
+
+        cell_index = sum(factor_value * 2**i
+                         for i, factor_value in enumerate(factors))
+
+    For K=2 with factors ``["A", "B"]``:
+        A=0, B=0 -> cell 0
+        A=1, B=0 -> cell 1
+        A=0, B=1 -> cell 2
+        A=1, B=1 -> cell 3
+
+    Examples
+    --------
+    >>> import pandas as pd
+    >>> df = pd.DataFrame({"x": range(8)})
+    >>> design = FactorialDesign(factors=["A", "B"], n_per_cell=2, seed=42)
+    >>> assignment = design.randomize(df)
+    >>> assignment.n_cells_
+    4
+    >>> assignment.cell_sizes_
+    {0: 2, 1: 2, 2: 2, 3: 2}
+    """
+
+    def __init__(
+        self,
+        factors: list[str],
+        n_per_cell: int,
+        seed: int | None = None,
+    ) -> None:
+        if not isinstance(factors, list) or len(factors) == 0:
+            raise InvalidDesignError(
+                "FactorialDesign requires a non-empty list of factors."
+            )
+
+        if len(set(factors)) != len(factors):
+            duplicates = [f for f in factors if factors.count(f) > 1]
+            raise InvalidDesignError(
+                f"FactorialDesign factors must be unique. "
+                f"Duplicates found: {sorted(set(duplicates))}."
+            )
+
+        if not isinstance(n_per_cell, int) or n_per_cell < 1:
+            raise InvalidDesignError(
+                f"n_per_cell must be a positive integer, "
+                f"received {n_per_cell!r}."
+            )
+
+        self.factors = factors
+        self.n_per_cell = n_per_cell
+        self.seed = seed
+
+    def randomize(self, df: pd.DataFrame) -> FactorialAssignment:
+        """Perform factorial randomization.
+
+        Parameters
+        ----------
+        df : pd.DataFrame
+            DataFrame with experimental units. Must have exactly
+            ``n_per_cell * 2^K`` rows and must not contain any column
+            named in ``factors`` or named ``"_cell"``.
+
+        Returns
+        -------
+        FactorialAssignment
+            Assignment with factor columns and ``"_cell"`` added.
+
+        Raises
+        ------
+        InvalidDesignError
+            If column-name collisions exist or if any cell has size 0.
+        InsufficientDataError
+            If ``len(df) != n_per_cell * 2^K``.
+        """
+        k = len(self.factors)
+        n_cells = 2**k
+        n_required = self.n_per_cell * n_cells
+
+        if len(df) != n_required:
+            raise InsufficientDataError(
+                context=(
+                    f"FactorialDesign with K={k} factors and "
+                    f"n_per_cell={self.n_per_cell} requires exactly "
+                    f"{n_required} units"
+                ),
+                minimum=n_required,
+                received=len(df),
+            )
+
+        # Detect column-name collisions
+        forbidden = set(self.factors) | {"_cell"}
+        collisions = sorted(forbidden & set(df.columns))
+        if collisions:
+            raise InvalidDesignError(
+                f"DataFrame already contains columns reserved for "
+                f"FactorialDesign output: {collisions}. Drop or rename "
+                f"them before calling randomize()."
+            )
+
+        df_out = df.copy()
+        rng = np.random.default_rng(self.seed)
+
+        # Shuffle iloc positions and assign sequentially to cells.
+        shuffled = rng.permutation(n_required)
+        cell_assignment = np.empty(n_required, dtype=int)
+        for cell_idx in range(n_cells):
+            start = cell_idx * self.n_per_cell
+            end = start + self.n_per_cell
+            cell_assignment[shuffled[start:end]] = cell_idx
+
+        df_out["_cell"] = cell_assignment
+
+        # Decode cell index back into binary factor values
+        # using little-endian convention.
+        for i, factor_name in enumerate(self.factors):
+            df_out[factor_name] = (cell_assignment >> i) & 1
+
+        # Build cell_sizes and validate non-empty cells (defensive;
+        # by construction every cell has n_per_cell >= 1).
+        cell_sizes: dict = {}
+        for cell_idx in range(n_cells):
+            size = int((cell_assignment == cell_idx).sum())
+            if size == 0:
+                raise InvalidDesignError(
+                    f"Cell {cell_idx} has zero units after randomization. "
+                    f"This should not happen with n_per_cell={self.n_per_cell}; "
+                    f"please report as a bug."
+                )
+            cell_sizes[cell_idx] = size
+
+        return FactorialAssignment(
+            data=df_out,
+            design=self,
+            factor_cols=self.factors,
+            cell_sizes=cell_sizes,
+            seed=self.seed,
+        )