dataframe-toolkit 0.4.1__py3-none-any.whl

dftk/commands/scat_cmd.py

@@ -0,0 +1,278 @@
+"""
+dftk.commands.scat_cmd — dftk scat subcommand.
+
+Scatter plot with optional OLS / robust regression line overlay.
+
+Port of dfscat.py with improvements:
+  - --groupcol  colour-codes points by a column
+  - --subgraphcol  splits into a subplot grid
+  - --size / --fontsize  publication-quality figure presets
+  - --fit  overlays an OLS (or robust) regression line via fit_cmd.regress_it
+  - --sizecol  point size from a data column (was --area)
+  - --colorcol  continuous colour mapping from a data column (was --color)
+"""
+
+import argparse
+import sys
+
+import matplotlib.lines as mlines
+import numpy as np
+
+from dftk.commands.base import BaseCommand
+from dftk.common.io import check_cols, io
+from dftk.common.plot import (
+    add_figure_arguments,
+    add_group_arguments,
+    add_legend_arguments,
+    add_xy_arguments,
+    apply_labels,
+    apply_limits,
+    apply_style,
+    apply_ticks,
+    make_figure,
+    make_grouplabel,
+    save_or_show,
+    subgraph_groups,
+    subgraph_layout,
+)
+
+# ---------------------------------------------------------------------------
+# Fit overlay
+# ---------------------------------------------------------------------------
+
+
+def _fit_overlay(ax, x_series, y_series, args):
+    """Fit a regression line and overlay it on ax. Prints summary to stdout."""
+    import pandas as pd
+
+    from dftk.commands.fit_cmd import regress_it
+
+    df_fit = pd.DataFrame({"_x": x_series.values, "_y": y_series.values}).dropna()
+    fit_args = argparse.Namespace(
+        formula="_y ~ _x" if not args.nointercept else "_y ~ _x - 1",
+        weights=None,
+        robust=getattr(args, "robust", False),
+    )
+    res = regress_it(df_fit, fit_args)
+    print(res.summary(), file=sys.stderr)
+
+    xmin, xmax = ax.get_xlim()
+    X_plot = np.linspace(xmin, xmax, 200)
+
+    if args.nointercept:
+        slope = res.params["_x"]
+        Y_plot = slope * X_plot
+        label_text = f"{slope:+.3g} × {args.xcol}"
+    else:
+        intercept = res.params["Intercept"]
+        slope = res.params["_x"]
+        Y_plot = intercept + slope * X_plot
+        label_text = f"{intercept:.3g}{slope:+.3g} × {args.xcol}"
+
+    if getattr(args, "pvalue", False):
+        pval = res.pvalues.get("_x", res.pvalues.iloc[-1])
+        label_text += f"  -log10(p)={-np.log10(pval):.2g}"
+
+    ax.plot(X_plot, Y_plot, "r--", linewidth=1)
+    ax.legend(
+        handles=[mlines.Line2D([], [], color="red", linestyle="--", label=label_text)],
+        loc=0,
+    )
+
+
+# ---------------------------------------------------------------------------
+# Single-axes scatter (reused per subplot)
+# ---------------------------------------------------------------------------
+
+
+def _scatter_ax(ax, df, args, title_suffix: str = ""):
+    """Draw scatter on a single Axes, with optional --groupcol colouring."""
+    import matplotlib.pyplot as plt
+
+    xcol = args.xcol
+    ycol = args.ycol
+    legendloc = tuple(args.legendloc) if args.legendloc else "best"
+
+    scatter_kw = dict(marker=getattr(args, "marker", "o"), alpha=0.7)
+
+    if args.groupcol is not None:
+        for groupname, gdf in subgraph_groups(
+            df, args.groupcol, getattr(args, "groupcolorder", None)
+        ):
+            label = make_grouplabel(
+                groupname, args.groupcol, getattr(args, "groupcolformat", None)
+            )
+            kw = dict(scatter_kw)
+            if args.sizecol:
+                kw["s"] = gdf[args.sizecol]
+            ax.scatter(gdf[xcol], gdf[ycol], label=label, **kw)
+        ax.legend(loc=legendloc, title=args.legendtitle)
+    else:
+        kw = dict(scatter_kw)
+        if args.sizecol:
+            kw["s"] = df[args.sizecol]
+        if args.colorcol:
+            sc = ax.scatter(
+                df[xcol], df[ycol], c=df[args.colorcol], cmap="viridis", **kw
+            )
+            plt.colorbar(sc, ax=ax, label=args.colorcol)
+        else:
+            label = args.legend or ""
+            ax.scatter(df[xcol], df[ycol], label=label if label else None, **kw)
+            if label:
+                ax.legend(loc=legendloc, title=args.legendtitle)
+
+    if getattr(args, "fit", False) and args.groupcol is None:
+        _fit_overlay(ax, df[xcol], df[ycol], args)
+
+    apply_limits(ax, args)
+    apply_ticks(ax, args)
+    apply_labels(ax, args, default_xlabel=xcol or "", default_ylabel=ycol or "")
+    if title_suffix:
+        existing = ax.get_title()
+        ax.set_title(f"{existing}  [{title_suffix}]" if existing else title_suffix)
+
+    ax.grid(True, linewidth=0.4, alpha=0.5)
+
+
+# ---------------------------------------------------------------------------
+# Command
+# ---------------------------------------------------------------------------
+
+_EPILOG = """\
+EXAMPLES
+--------
+Basic scatter:
+
+  dftk scat data.tsv -x weight -y height
+
+Colour by group:
+
+  dftk scat data.tsv -x x -y y -g group
+
+Subplot grid by condition, colour by treatment:
+
+  dftk scat data.tsv -x x -y y --subgraphcol condition -g treatment
+
+Fit OLS line:
+
+  dftk scat data.tsv -x x -y y --fit
+
+Publication figure (Nature single column, PDF):
+
+  dftk scat data.tsv -x x -y y --size single --fontsize publication -f fig.pdf
+
+Bubble chart (point size from column):
+
+  dftk scat data.tsv -x x -y y --sizecol area_col
+
+Continuous colour mapping:
+
+  dftk scat data.tsv -x x -y y --colorcol score
+"""
+
+
+class ScatCommand(BaseCommand):
+    name = "scat"
+    help = "Scatter plot with optional regression line overlay."
+
+    def add_arguments(self, parser: argparse.ArgumentParser) -> None:
+        parser.formatter_class = argparse.RawDescriptionHelpFormatter
+        parser.epilog = _EPILOG
+
+        self.add_read_arguments(parser)
+        add_xy_arguments(parser)
+        add_figure_arguments(parser)
+        add_group_arguments(parser)
+        add_legend_arguments(parser)
+
+        s = parser.add_argument_group("scatter options")
+        s.add_argument(
+            "-m",
+            "--marker",
+            default="o",
+            metavar="MARKER",
+            help="Matplotlib marker style (default: o).",
+        )
+        s.add_argument(
+            "--sizecol",
+            default=None,
+            metavar="COL",
+            help="Column to use for point sizes (bubble chart).",
+        )
+        s.add_argument(
+            "--colorcol",
+            default=None,
+            metavar="COL",
+            help="Column for continuous colour mapping.",
+        )
+
+        f = parser.add_argument_group("fit options")
+        f.add_argument(
+            "--fit",
+            action="store_true",
+            help="Overlay an OLS regression line. Prints summary to stderr.",
+        )
+        f.add_argument(
+            "-r",
+            "--robust",
+            action="store_true",
+            help="Use robust regression (RLM) for --fit.",
+        )
+        f.add_argument(
+            "--nointercept", action="store_true", help="Suppress intercept in --fit."
+        )
+        f.add_argument(
+            "--pvalue",
+            action="store_true",
+            help="Add −log₁₀(p) to the fit legend label.",
+        )
+
+    def execute(self, args: argparse.Namespace) -> None:
+        if not args.xcol or not args.ycol:
+            raise ValueError("--xcol (-x) and --ycol (-y) are required.")
+
+        df = io.read(args)
+        check_cols(df, [args.xcol, args.ycol], "-x/-y")
+        check_cols(df, args.groupcol, "-g/--groupcol")
+        check_cols(df, args.subgraphcol, "--subgraphcol")
+        if args.sizecol:
+            check_cols(df, [args.sizecol], "--sizecol")
+        if args.colorcol:
+            check_cols(df, [args.colorcol], "--colorcol")
+
+        # Use non-interactive backend when saving to file
+        if args.file:
+            import matplotlib
+
+            matplotlib.use("Agg")
+
+        import matplotlib.pyplot as plt
+
+        if args.usetex:
+            plt.rc("text", usetex=True)
+
+        apply_style(args)
+
+        if args.subgraphcol is not None:
+            groups = subgraph_groups(
+                df, args.subgraphcol, getattr(args, "subgraphorder", None)
+            )
+            nrows, ncols = subgraph_layout(len(groups), args.ncols)
+            fig, axes = make_figure(args, nrows=nrows, ncols=ncols)
+            for idx, (groupname, gdf) in enumerate(groups):
+                ax = axes[idx // ncols][idx % ncols]
+                suffix = make_grouplabel(
+                    groupname, args.subgraphcol, getattr(args, "subgraphformat", None)
+                )
+                _scatter_ax(ax, gdf, args, title_suffix=suffix)
+            # Hide unused axes
+            for idx in range(len(groups), nrows * ncols):
+                axes[idx // ncols][idx % ncols].set_visible(False)
+            fig.tight_layout()
+        else:
+            fig, axes = make_figure(args)
+            _scatter_ax(axes[0][0], df, args)
+            fig.tight_layout()
+
+        save_or_show(fig, args)

dftk/commands/split_cmd.py

@@ -0,0 +1,130 @@
+"""
+dftk.commands.split_cmd — dftk split subcommand.
+
+Port of dfsplit.py: split a dataframe into one TSV file per group.
+"""
+
+import argparse
+import copy
+import os
+
+from dftk.commands.base import BaseCommand
+from dftk.common.io import check_cols, io
+
+
+def _make_filename(groupname, prefix: str, suffix: str) -> str:
+    tag = (
+        "_".join(str(x) for x in groupname)
+        if isinstance(groupname, tuple)
+        else str(groupname)
+    )
+    return f"{prefix}{tag}{suffix}"
+
+
+class _SafeDict(dict):
+    """Format-map dict that leaves unknown keys unreplaced."""
+
+    def __missing__(self, key):
+        return "{" + key + "}"
+
+
+def _split(df, args: argparse.Namespace) -> None:
+    for groupname, groupdf in df.groupby(args.groups):
+        if args.template is not None:
+            key = (groupname,) if not isinstance(groupname, tuple) else groupname
+            group_dict = dict(zip(args.groups, key, strict=False))
+            filename = args.template.format_map(_SafeDict(group_dict))
+        else:
+            filename = _make_filename(groupname, args.prefix, args.suffix)
+
+        if args.noclobber and os.path.exists(filename):
+            continue
+
+        file_args = copy.copy(args)
+        file_args.output = filename
+        io.printdf(groupdf, file_args)
+
+
+_EPILOG = """\
+FILENAME PATTERNS
+-----------------
+  By default output files are named:  <prefix><group_key><suffix>
+
+    dftk split data.tsv -g country -p out/ -s .tsv
+    # → out/Iceland.tsv, out/Norway.tsv, ...
+
+  With --template, use {col} placeholders for the group column values:
+
+    dftk split data.tsv -g year month --template data_{year}_{month}.tsv
+
+  Unknown placeholders are left as-is (safe format_map).
+
+OPTIONS
+-------
+  --noclobber   Skip files that already exist instead of overwriting.
+
+  Note: -o/--output is not used by split; output paths are always determined
+  by -p/--prefix, -s/--suffix, or --template.
+
+EXAMPLES
+--------
+  dftk split data.tsv -g group
+  dftk split data.tsv -g country year -p results/ -s .tsv
+  dftk split data.tsv -g category --template out/{category}.tsv
+"""
+
+
+class SplitCommand(BaseCommand):
+    @property
+    def name(self) -> str:
+        return "split"
+
+    @property
+    def help(self) -> str:
+        return "Split a dataframe into one file per group."
+
+    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("split options")
+        g.add_argument(
+            "-g",
+            "--groups",
+            nargs="+",
+            required=True,
+            metavar="COL",
+            help="Column(s) to group and split by.",
+        )
+        g.add_argument(
+            "-p",
+            "--prefix",
+            default="split-",
+            metavar="PREFIX",
+            help="Filename prefix (default: 'split-').",
+        )
+        g.add_argument(
+            "-s",
+            "--suffix",
+            default="",
+            metavar="SUFFIX",
+            help="Filename suffix / extension (default: none).",
+        )
+        g.add_argument(
+            "--template",
+            default=None,
+            metavar="TEMPLATE",
+            help="Filename template with {col} placeholders, e.g. out/{group}.tsv.",
+        )
+        g.add_argument(
+            "--noclobber",
+            action="store_true",
+            help="Skip output files that already exist.",
+        )
+
+    def execute(self, args: argparse.Namespace) -> None:
+        df = io.read(args)
+        check_cols(df, args.groups, "-g/--groups")
+        _split(df, args)

dftk/commands/stat_cmd.py

@@ -0,0 +1,289 @@
+"""
+dftk.commands.stat_cmd — the ``dftk stat`` subcommand.
+
+Computes descriptive statistics for one or more numeric columns, optionally
+within groups.  Output always includes a ``name`` column identifying which
+input column each row describes, giving a consistent schema regardless of
+how many columns are requested.
+
+Output columns
+--------------
+  [groupcol(s),] name, count, sum, mean, std, min, max, median,
+  cilo, cihi, sem, skew, kurt
+
+Bootstrap mode
+--------------
+With ``--bootstrap N`` the full stat computation is repeated N times on
+resampled data.  Each iteration appends a ``samplenum`` column and output
+blocks are streamed without repeating the header, making it easy to pipe
+into ``dftk pivot`` to aggregate across samples.
+
+Example
+-------
+    dftk stat data.tsv -c height weight -g sex
+    dftk stat data.tsv -c value -g group --bootstrap 1000 --randomseed 42 -o \\
+        | dftk pivot -i group -v value_mean -o \\
+        | dftk print
+"""
+
+import argparse
+import logging
+
+import numpy as np
+import pandas as pd
+
+from dftk.commands.base import BaseCommand
+from dftk.common.io import check_cols, io
+from dftk.common.seed import normalize_seed
+
+logger = logging.getLogger("dftk")
+
+_CI_METHODS = ["linear", "lower", "higher", "midpoint", "nearest"]
+
+# Canonical order of statistic columns in the output.
+_STAT_NAMES = [
+    "count",
+    "sum",
+    "mean",
+    "std",
+    "min",
+    "max",
+    "median",
+    "cilo",
+    "cihi",
+    "sem",
+    "skew",
+    "kurt",
+]
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+
+def _samplefraction(x: str) -> float:
+    """argparse type validator for --samplefrac."""
+    try:
+        x = float(x)
+    except ValueError:
+        raise argparse.ArgumentTypeError(f"{x!r} is not a valid float") from None
+    if not (0.0 <= x <= 1.0):
+        raise argparse.ArgumentTypeError(f"{x} is not in range [0.0, 1.0]")
+    return x
+
+
+def _ci(name: str, level: float, method: str):
+    """Return a named percentile aggregation function for use with pandas agg.
+
+    The function name is set explicitly so pandas uses it as the output
+    column name rather than the generic ``<lambda>`` or ``cifunction``.
+    """
+
+    def cifunction(a):
+        return np.percentile(a, q=level, method=method)
+
+    cifunction.__name__ = name
+    return cifunction
+
+
+# ---------------------------------------------------------------------------
+# Core computation
+# ---------------------------------------------------------------------------
+
+
+def _compute_stats(df: pd.DataFrame, args: argparse.Namespace) -> pd.DataFrame:
+    """Return a dataframe of descriptive statistics.
+
+    Always produces a ``name`` column identifying the source column, giving
+    a consistent schema regardless of how many columns are in ``args.cols``.
+    """
+    cols: list[str] = args.cols
+    groupcol: list[str] | None = args.groupcol
+
+    conf_lo = (100.0 - args.confidencelevel) / 2
+    conf_hi = 100.0 - conf_lo
+    cilo_fn = _ci("cilo", conf_lo, args.confidencemethod)
+    cihi_fn = _ci("cihi", conf_hi, args.confidencemethod)
+
+    agg_funcs = [
+        "count",
+        "sum",
+        "mean",
+        "std",
+        "min",
+        "max",
+        "median",
+        cilo_fn,
+        cihi_fn,
+        "sem",
+        "skew",
+        "kurt",
+    ]
+
+    if groupcol is None:
+        # Ungrouped — agg returns shape (n_stats, n_cols); transpose to
+        # (n_cols, n_stats) then reset the column index to get a 'name' row.
+        stats = (
+            df[cols].agg(agg_funcs).T.reset_index().rename(columns={"index": "name"})
+        )
+        return stats[["name"] + _STAT_NAMES]
+
+    else:
+        # Grouped — agg returns MultiIndex columns (col, stat).
+        # Reshape to long format so every (group, col) combination is a row.
+        name_col = "name"
+        if name_col in groupcol:
+            name_col = "name_mangled"
+            logger.warning(
+                'A group column is named "name"; '
+                'stat name column renamed to "name_mangled"'
+            )
+
+        stats = (
+            df.groupby(groupcol)[cols]
+            .agg(agg_funcs)
+            .rename(columns={"amin": "min", "amax": "max"})  # numpy compat
+        )
+
+        # swaplevel: (col, stat) → (stat, col)
+        # stack:     fold the col level into the row index
+        # reset_index + rename: flatten to a plain DataFrame
+        level_col = f"level_{len(groupcol)}"
+        stats = (
+            stats.swaplevel(0, 1, axis=1)
+            .stack(future_stack=True)
+            .reset_index()
+            .rename(columns={level_col: name_col})
+            .sort_values(by=[name_col] + groupcol)
+        )
+
+        out_cols = groupcol + [name_col] + _STAT_NAMES
+        return stats[out_cols].reset_index(drop=True)
+
+
+# ---------------------------------------------------------------------------
+# Command
+# ---------------------------------------------------------------------------
+
+
+class StatCommand(BaseCommand):
+    """Compute descriptive statistics for one or more columns."""
+
+    name = "stat"
+    help = "Compute descriptive statistics (mean, std, CI, …) for one or more columns"
+
+    def add_arguments(self, parser: argparse.ArgumentParser) -> None:
+        self.add_io_arguments(parser)
+
+        g = parser.add_argument_group("statistics")
+        g.add_argument(
+            "-c",
+            "--cols",
+            help="Column(s) to compute statistics on",
+            nargs="+",
+            required=True,
+            metavar="COL",
+        )
+        g.add_argument(
+            "-g",
+            "--groupcol",
+            help="Group-by column(s)",
+            nargs="+",
+            default=None,
+            metavar="COL",
+        )
+        g.add_argument(
+            "--confidencelevel",
+            help="Confidence level in percent (default: 95.0)",
+            type=float,
+            default=95.0,
+            metavar="PCT",
+        )
+        g.add_argument(
+            "--confidencemethod",
+            help="Percentile interpolation method (default: linear)",
+            choices=_CI_METHODS,
+            default="linear",
+        )
+
+        bs = parser.add_argument_group("bootstrap")
+        bs.add_argument(
+            "--bootstrap",
+            help="Number of bootstrap samples to draw",
+            type=int,
+            default=None,
+            metavar="N",
+        )
+        bs.add_argument(
+            "--randomseed",
+            help=(
+                "Random seed — integer, or a string which is MD5-hashed "
+                "to a reproducible integer"
+            ),
+            default=None,
+            metavar="SEED",
+        )
+        size_grp = bs.add_mutually_exclusive_group()
+        size_grp.add_argument(
+            "--samplesize",
+            help="Rows per bootstrap sample (default: same size as input)",
+            type=int,
+            default=None,
+            metavar="N",
+        )
+        size_grp.add_argument(
+            "--samplefrac",
+            help="Fraction of rows per bootstrap sample, 0.0–1.0 (default: 1.0)",
+            type=_samplefraction,
+            default=1.0,
+            metavar="FRAC",
+        )
+        bs.add_argument(
+            "--noreplace",
+            help="Sample without replacement (default: with replacement)",
+            dest="replace",
+            action="store_false",
+        )
+        parser.set_defaults(replace=True)
+
+    def execute(self, args: argparse.Namespace) -> None:
+        args.randomseed = normalize_seed(args.randomseed)
+        np.random.seed(args.randomseed)
+
+        df = io.read(args)
+        if df.empty:
+            raise ValueError("Input data is empty")
+
+        check_cols(df, args.cols, "-c/--cols")
+        check_cols(df, args.groupcol, "-g/--groupcol")
+
+        if args.bootstrap is None:
+            result = _compute_stats(df, args)
+            io.printdf(result, args)
+
+        else:
+            logger.info(
+                f"Bootstrapping: {args.bootstrap} samples, "
+                f"seed={args.randomseed}, replace={args.replace}"
+            )
+            sample_kwargs: dict = {"replace": args.replace}
+            if args.samplesize is not None:
+                sample_kwargs["n"] = args.samplesize
+            else:
+                sample_kwargs["frac"] = args.samplefrac
+
+            frames = []
+            for i in range(1, args.bootstrap + 1):
+                if args.groupcol is None:
+                    sample = df.sample(**sample_kwargs)
+                else:
+                    sample = df.groupby(args.groupcol, group_keys=False).apply(
+                        lambda x: x.sample(**sample_kwargs)
+                    )
+
+                result = _compute_stats(sample, args)
+                result["samplenum"] = i
+                frames.append(result)
+
+            io.printdf(pd.concat(frames, ignore_index=True), args)