dataframe-toolkit 0.4.1__py3-none-any.whl

dftk/commands/corr_cmd.py

@@ -0,0 +1,281 @@
+"""
+dftk.commands.corr_cmd — dftk corr subcommand.
+
+Port of dfcorr.py: pairwise column correlations with optional grouping,
+analytical confidence intervals (Pearson), and BCa bootstrap inference.
+
+Bootstrap mode (--bootstrap N) produces a single summary row per pair with:
+  p_perm      — permutation p-value (fraction of |r_null| >= |r_obs|)
+  ci_boot_lo  — BCa 95% CI lower bound (all methods)
+  ci_boot_hi  — BCa 95% CI upper bound (all methods)
+"""
+
+import argparse
+import logging
+
+import numpy as np
+import pandas as pd
+import scipy.stats as ss
+
+from dftk.commands.base import BaseCommand
+from dftk.common.io import check_cols, io
+from dftk.common.seed import normalize_seed
+
+logger = logging.getLogger(__name__)
+
+METHODS = ["pearson", "spearman", "kendall"]
+_METHOD_FUNC = {
+    "pearson": ss.pearsonr,
+    "spearman": ss.spearmanr,
+    "kendall": ss.kendalltau,
+}
+
+
+def _compute_correlation(a: np.ndarray, b: np.ndarray, method_fn, ci: bool) -> tuple:
+    result = method_fn(a, b)
+    row = (len(a), result.statistic, result.pvalue)
+    if ci:
+        interval = result.confidence_interval()
+        row = (*row, interval.low, interval.high)
+    return row
+
+
+# BCa jackknife acceleration is estimated from leave-one-out resamples.
+# scipy.stats.bootstrap materialises an (n × n-1) array for this step,
+# which OOMs at large n.  We stream one leave-one-out at a time and cap
+# at _MAX_JACK observations so peak memory is O(min(n, _MAX_JACK)).
+_MAX_JACK = 500
+
+
+def _bootstrap_stats(
+    a: np.ndarray,
+    b: np.ndarray,
+    method_fn,
+    n_boot: int,
+    rng: np.random.Generator,
+    confidence: float = 0.95,
+) -> tuple[float, float, float]:
+    """Return (p_perm, ci_boot_lo, ci_boot_hi) via permutation test and BCa CI."""
+    n = len(a)
+    r_obs = method_fn(a, b).statistic
+
+    r_obs_abs = abs(r_obs)
+    n_extreme = 0
+    for _ in range(n_boot):
+        if abs(method_fn(a, rng.permutation(b)).statistic) >= r_obs_abs:
+            n_extreme += 1
+    p_perm = n_extreme / n_boot
+
+    theta_boot = np.empty(n_boot)
+    for i in range(n_boot):
+        idx = rng.integers(0, n, size=n)
+        theta_boot[i] = method_fn(a[idx], b[idx]).statistic
+
+    if n <= _MAX_JACK:
+        jack_indices = np.arange(n)
+    else:
+        jack_indices = rng.choice(n, size=_MAX_JACK, replace=False)
+
+    theta_jack = np.empty(len(jack_indices))
+    mask = np.ones(n, dtype=bool)
+    for j, i in enumerate(jack_indices):
+        mask[i] = False
+        theta_jack[j] = method_fn(a[mask], b[mask]).statistic
+        mask[i] = True
+
+    jmean = theta_jack.mean()
+    diff = jmean - theta_jack
+    denom = np.sum(diff**2)
+    a_hat = np.sum(diff**3) / (6.0 * denom**1.5) if denom > 1e-20 else 0.0
+
+    prop = np.clip((theta_boot < r_obs).mean(), 1e-10, 1 - 1e-10)
+    z0 = ss.norm.ppf(prop)
+
+    alpha = (1.0 - confidence) / 2.0
+    z_lo, z_hi = ss.norm.ppf(alpha), ss.norm.ppf(1.0 - alpha)
+
+    def _adj(z: float) -> float:
+        denom = 1.0 - a_hat * (z0 + z)
+        if abs(denom) < 1e-10:
+            return float("nan")
+        return ss.norm.cdf(z0 + (z0 + z) / denom)
+
+    a1, a2 = _adj(z_lo), _adj(z_hi)
+
+    if np.isnan(a1) or np.isnan(a2):
+        ci_lo = float(np.percentile(theta_boot, 100.0 * alpha))
+        ci_hi = float(np.percentile(theta_boot, 100.0 * (1.0 - alpha)))
+    else:
+        ci_lo = float(np.percentile(theta_boot, 100.0 * a1))
+        ci_hi = float(np.percentile(theta_boot, 100.0 * a2))
+
+    return p_perm, ci_lo, ci_hi
+
+
+def _corr(
+    df: pd.DataFrame, args: argparse.Namespace, rng: np.random.Generator
+) -> pd.DataFrame:
+    method_fn = _METHOD_FUNC[args.method]
+    pairs = [c.split(":", maxsplit=1) for c in args.cols]
+    bad = [c for c, p in zip(args.cols, pairs, strict=False) if len(p) != 2]
+    if bad:
+        raise ValueError(f"-c pairs must be col1:col2, got: {bad}")
+
+    ci = args.ci and args.method == "pearson"
+
+    header = ["col1", "col2", "nobs", "correlation", "pvalue"]
+    if ci:
+        header += ["cilo", "cihi"]
+    if args.bootstrap is not None:
+        header += ["p_perm", "ci_boot_lo", "ci_boot_hi"]
+
+    rows = []
+
+    def _add_pairs(sub: pd.DataFrame, prefix: tuple = ()):
+        for col1, col2 in pairs:
+            pair_df = sub[[col1, col2]].dropna()
+            a, b = pair_df[col1].values, pair_df[col2].values
+            cr = _compute_correlation(a, b, method_fn, ci)
+            if args.bootstrap is None:
+                rows.append((*prefix, col1, col2, *cr))
+            else:
+                boot = _bootstrap_stats(
+                    a, b, method_fn, args.bootstrap, rng, args.confidence / 100.0
+                )
+                rows.append((*prefix, col1, col2, *cr, *boot))
+
+    if args.groups:
+        header = [*args.groups, *header]
+        for group, gdf in df.groupby(args.groups):
+            prefix = group if isinstance(group, tuple) else (group,)
+            _add_pairs(gdf, prefix)
+    else:
+        _add_pairs(df)
+
+    return pd.DataFrame(rows, columns=header)
+
+
+_EPILOG = """\
+COLUMN PAIRS (-c)
+-----------------
+  Pairs are specified as col1:col2.  Multiple pairs may be given:
+
+    dftk corr data.tsv -c x:y
+    dftk corr data.tsv -c x:y a:b
+
+METHODS (--method)
+------------------
+  pearson    Pearson r (default)
+  spearman   Spearman rank correlation
+  kendall    Kendall tau
+
+CONFIDENCE INTERVAL (--ci)
+--------------------------
+  Adds cilo and cihi columns using the analytical Fisher z-transform formula.
+  Only available for Pearson; silently ignored for Spearman/Kendall.
+  Use --bootstrap to get non-parametric CIs for any method.
+
+BOOTSTRAP (--bootstrap N)
+--------------------------
+  Runs N bootstrap resamples and N permutation draws, adding three columns:
+
+    p_perm      Permutation p-value: fraction of |r_null| >= |r_obs| when b
+                is repeatedly shuffled to break the pairing with a.  This is
+                a non-parametric test of the null hypothesis r = 0.
+
+    ci_boot_lo  BCa (bias-corrected and accelerated) CI lower bound.
+    ci_boot_hi  BCa CI upper bound.
+
+  The CI level defaults to 95% and is controlled by --confidence.
+  BCa corrects for both bias (z0) and skewness (acceleration factor a_hat)
+  in the bootstrap distribution.  The acceleration factor is estimated via a
+  streaming jackknife: leave-one-out statistics are computed one at a time
+  (capped at 500 observations when n > 500) so peak memory is O(n) rather
+  than O(n^2).  Falls back to percentile CI if BCa is degenerate.
+
+  Use --randomseed for reproducibility.
+
+EXAMPLES
+--------
+  dftk corr data.tsv -c x:y
+  dftk corr data.tsv -c x:y -g group --method spearman
+  dftk corr data.tsv -c x:y --ci
+  dftk corr data.tsv -c x:y --method spearman --bootstrap 2000 --randomseed 42
+"""
+
+
+class CorrCommand(BaseCommand):
+    @property
+    def name(self) -> str:
+        return "corr"
+
+    @property
+    def help(self) -> str:
+        return "Pairwise column correlations (Pearson, Spearman, Kendall) with optional bootstrap."  # noqa: E501
+
+    def add_arguments(self, parser: argparse.ArgumentParser) -> None:
+        parser.formatter_class = argparse.RawDescriptionHelpFormatter
+        parser.epilog = _EPILOG
+
+        self.add_io_arguments(parser)
+
+        g = parser.add_argument_group("corr options")
+        g.add_argument(
+            "-c",
+            "--cols",
+            nargs="+",
+            required=True,
+            metavar="COL1:COL2",
+            help="Column pair(s) to correlate, e.g. x:y.",
+        )
+        g.add_argument(
+            "-g",
+            "--groups",
+            nargs="+",
+            default=[],
+            metavar="COL",
+            help="Group column(s); correlation is computed within each group.",
+        )
+        g.add_argument(
+            "--method",
+            choices=METHODS,
+            default="pearson",
+            help="Correlation method (default: pearson).",
+        )
+        g.add_argument(
+            "--ci",
+            action="store_true",
+            help="Include analytical 95%% CI columns cilo/cihi (Pearson only).",
+        )
+        g.add_argument(
+            "--bootstrap",
+            type=int,
+            default=None,
+            metavar="N",
+            help="Bootstrap N resamples; adds p_perm and BCa CI columns.",
+        )
+        g.add_argument(
+            "--confidence",
+            type=float,
+            default=95.0,
+            metavar="PCT",
+            help="CI level in percent for BCa bootstrap bounds (default: 95.0).",
+        )
+        g.add_argument(
+            "--randomseed",
+            default=None,
+            metavar="SEED",
+            help="Random seed for bootstrap/permutation (integer or string).",
+        )
+
+    def execute(self, args: argparse.Namespace) -> None:
+        seed = normalize_seed(args.randomseed) if args.randomseed is not None else None
+        rng = np.random.default_rng(seed)
+
+        all_cols = [c for pair in args.cols for c in pair.split(":", maxsplit=1)]
+        df = io.read(args)
+        check_cols(df, all_cols, "-c/--cols")
+        check_cols(df, args.groups, "-g/--groups")
+
+        result = _corr(df, args, rng)
+        io.printdf(result, args)

dftk/commands/dataset_cmd.py

@@ -0,0 +1,181 @@
+"""
+dftk.commands.dataset_cmd — the ``dftk dataset`` subcommand.
+
+Provides access to curated example datasets from seaborn, statsmodels,
+and (optionally) pydataset.  Datasets are written to stdout as TSV or
+into the dftk parquet pipe with -o.
+
+Examples
+--------
+    dftk dataset iris
+    dftk dataset iris --source seaborn
+    dftk dataset longley --source statsmodels
+    dftk dataset --list
+    dftk dataset --list --source seaborn
+    dftk dataset iris -o | dftk stat -c sepal_length -g species
+"""
+
+import argparse
+import sys
+
+import pandas as pd
+
+from dftk.commands.base import BaseCommand
+from dftk.common.io import io
+
+# ---------------------------------------------------------------------------
+# Per-source load/list helpers
+# ---------------------------------------------------------------------------
+
+
+def _load_seaborn(name: str) -> pd.DataFrame:
+    import seaborn as sns
+
+    return sns.load_dataset(name)
+
+
+def _list_seaborn() -> list[tuple[str, str, str]]:
+    import seaborn as sns
+
+    return [("seaborn", n, "") for n in sorted(sns.get_dataset_names())]
+
+
+def _load_statsmodels(name: str) -> pd.DataFrame:
+    import statsmodels.api as sm
+
+    mod = getattr(sm.datasets, name, None)
+    if mod is not None and hasattr(mod, "load_pandas"):
+        return mod.load_pandas().data
+    raise ValueError(f"statsmodels dataset {name!r} not found")
+
+
+def _list_statsmodels() -> list[tuple[str, str, str]]:
+    import statsmodels.api as sm
+
+    rows = []
+    for name in sorted(dir(sm.datasets)):
+        if name.startswith("_"):
+            continue
+        mod = getattr(sm.datasets, name, None)
+        if mod is None or not hasattr(mod, "load_pandas"):
+            continue
+        title = (getattr(mod, "TITLE", "") or "").strip()
+        desc = title or (getattr(mod, "DESCRSHORT", "") or "").strip()
+        rows.append(("statsmodels", name, desc))
+    return rows
+
+
+try:
+    import pydataset as _pydataset  # noqa: F401
+
+    _HAS_PYDATASET = True
+except ImportError:
+    _HAS_PYDATASET = False
+
+
+def _load_pydataset(name: str) -> pd.DataFrame:
+    import pydataset
+
+    return pydataset.data(name)
+
+
+def _list_pydataset() -> list[tuple[str, str, str]]:
+    import pydataset
+
+    index = pydataset.data()
+    return [
+        ("pydataset", str(row["item"]), str(row.get("title", "")))
+        for _, row in index.iterrows()
+    ]
+
+
+# Registry — built at import time so choices are accurate
+_SOURCES: dict[str, tuple] = {
+    "seaborn": (_load_seaborn, _list_seaborn),
+    "statsmodels": (_load_statsmodels, _list_statsmodels),
+}
+if _HAS_PYDATASET:
+    _SOURCES["pydataset"] = (_load_pydataset, _list_pydataset)
+
+
+# ---------------------------------------------------------------------------
+# Command
+# ---------------------------------------------------------------------------
+
+
+class DatasetCommand(BaseCommand):
+    name = "dataset"
+    help = "Load a curated example dataset (seaborn, statsmodels, pydataset)"
+
+    def add_arguments(self, parser: argparse.ArgumentParser) -> None:
+        io.parser_output(parser)
+
+        g = parser.add_argument_group("dataset options")
+        g.add_argument(
+            "NAME",
+            nargs="?",
+            default=None,
+            help=(
+                "Dataset name (e.g. 'iris', 'longley').  "
+                "Omit with --list to show all available datasets."
+            ),
+        )
+        g.add_argument(
+            "--source",
+            nargs="+",
+            default=None,
+            choices=list(_SOURCES),
+            metavar="SOURCE",
+            help=(
+                f"Source(s) to search: {', '.join(_SOURCES)}.  "
+                "Default: try all sources in order."
+            ),
+        )
+        g.add_argument(
+            "--list",
+            action="store_true",
+            dest="list_datasets",
+            help="List available datasets as TSV (source, name, description).",
+        )
+
+    def execute(self, args: argparse.Namespace) -> None:
+        sources = args.source or list(_SOURCES)
+
+        if args.list_datasets:
+            self._print_list(sources)
+            return
+
+        if args.NAME is None:
+            raise ValueError(
+                "Provide a dataset NAME, or use --list to browse available datasets."
+            )
+
+        df = self._load(args.NAME, sources)
+        io.printdf(df, args)
+
+    # ------------------------------------------------------------------ #
+    # Helpers                                                              #
+    # ------------------------------------------------------------------ #
+
+    def _load(self, name: str, sources: list[str]) -> pd.DataFrame:
+        errors: list[str] = []
+        for src in sources:
+            load_fn, _ = _SOURCES[src]
+            try:
+                return load_fn(name)
+            except Exception as exc:
+                errors.append(f"  {src}: {exc}")
+        raise ValueError(
+            f"Dataset {name!r} not found in: {', '.join(sources)}.\n"
+            + "\n".join(errors)
+        )
+
+    def _print_list(self, sources: list[str]) -> None:
+        print("source\tname\tdescription")
+        for src in sources:
+            _, list_fn = _SOURCES[src]
+            try:
+                for source, name, desc in list_fn():
+                    print(f"{source}\t{name}\t{desc}")
+            except Exception as exc:
+                print(f"Warning: could not list {src} datasets: {exc}", file=sys.stderr)