dataframe-toolkit 0.15.0__tar.gz → 0.16.0__tar.gz

{dataframe_toolkit-0.15.0 → dataframe_toolkit-0.16.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: dataframe-toolkit
-Version: 0.15.0
+Version: 0.16.0
 Summary: A command-line toolkit for DataFrame analysis and manipulation
 Keywords: dataframe,statistics,data analysis,CLI,TSV,pandas
 Author: Gunnar Pálsson
@@ -139,7 +139,7 @@ dftk dataset tips -o \
 | `wstat`    | Weighted descriptive statistics (wmean, wstd, weighted quantile CI)                                 |
 | `fit`      | OLS/robust/weighted regression via R-style formulas; tidy table, `--summary`, `--anova`             |
 | `test`     | P-values between column pairs: t-test, Mann-Whitney, Wilcoxon, KS, correlations, bootstrap; groups  |
-| `corr`     | Pairwise column correlations (Pearson/Spearman/Kendall) with optional BCa bootstrap CI              |
+| `corr`     | Pairwise column correlations (Pearson/Spearman/Kendall) with optional BCa bootstrap CI; globs (`-c '*:*'`) or a full matrix (`--all`) |
 | `describe` | Quick column-level summary (dtype, n, n_unique, n_null, sample values)                              |
 | `info`     | Per-column dtype, null counts, and memory usage; `--summary` for dataset-level totals               |
 | `randvar`  | Sample from a distribution and append as a new column (norm, alpha, beta, …)                        |
@@ -157,6 +157,8 @@ publication-quality output.
 | `line`  | Line plot; optional error bars (`--yerr`) or CI bands (`--yci lo,hi`) and fit overlay                                                                                               |
 | `hist`  | Histogram or KDE (`--kde`); normalisation, cumulative mode, mean±σ annotation (`--stats`)                                                                                           |
 | `chi`   | Corner-resolved chi-plot: whether tail dependence sits in the lower or upper corner                                                                                                 |
+| `heatmap` | Coloured grid of values from a long (`-x -y -v`) or wide (`-y`) table; correlograms from `corr --all`, `--annot`, `--lower` |
+| `pairs` | Scatter-plot matrix with KDEs (or histograms) on the diagonal; `-g` colours groups, `--upper corr` puts r in the upper triangle |
 
 ### Utilities
 
@@ -409,6 +411,34 @@ dftk eval segments.tsv -f 'gap = Begin - End.groupby(Chr).shift()'
 Without `groupby(Chr)` the first segment of a chromosome is compared with
 the last one of the previous chromosome.
 
+### Correlograms and heatmaps
+
+`corr --all` correlates every pair of numeric columns both ways round, plus the diagonal, which is the square table `heatmap` draws. Values of both signs get diverging colours centred on 0; `--vmin -1 --vmax 1` keeps the full scale for r, so a weak correlation doesn't look strong. `--annot` writes the values in the cells and `--lower` keeps the lower triangle:
+
+```bash
+dftk dataset iris -o | dftk corr --all -o \
+  | dftk heatmap -x col1 -y col2 -v correlation --annot --vmin -1 --vmax 1 -f iris_corr.png
+dftk dataset diamonds -o | dftk corr --all --method spearman -o \
+  | dftk heatmap -x col1 -y col2 -v correlation --annot --lower --square \
+      --vmin -1 --vmax 1 -f diamonds_corr.png
+```
+
+Missing values are dropped pair by pair, so `nobs` can differ between cells; `dftk print --dropna -o | dftk corr --all` uses the complete rows throughout. p-values are not adjusted for multiple testing.
+
+For the shape of each relationship, not just r, `pairs` draws every pair of columns as a scatter plot with each column's KDE on the diagonal; `-g` colours groups (each group's KDE scaled by its share of the rows) and `--upper corr` writes r in the upper triangle, with `-g` also r within each group, which can differ even in sign from r over all rows:
+
+```bash
+dftk dataset iris -o | dftk pairs -g species -f iris_pairs.png
+dftk dataset iris -o | dftk pairs -g species --upper corr --diag hist -f iris_pairs_r.png
+```
+
+`heatmap` draws any table with one row per cell (`-x -y -v`), or a wide one with `-y` naming the row labels and the other numeric columns as heatmap columns, as `pivot` makes it:
+
+```bash
+dftk dataset penguins -o | dftk pivot -i species -g island -v body_mass_g -f count -o \
+  | dftk heatmap -y species --annot --cmap Blues -f penguin_counts.png
+```
+
 ### Bootstrap confidence intervals
 
 Bootstrap mode repeats the full stat computation N times on resampled data

{dataframe_toolkit-0.15.0 → dataframe_toolkit-0.16.0}/README.md

@@ -109,7 +109,7 @@ dftk dataset tips -o \
 | `wstat`    | Weighted descriptive statistics (wmean, wstd, weighted quantile CI)                                 |
 | `fit`      | OLS/robust/weighted regression via R-style formulas; tidy table, `--summary`, `--anova`             |
 | `test`     | P-values between column pairs: t-test, Mann-Whitney, Wilcoxon, KS, correlations, bootstrap; groups  |
-| `corr`     | Pairwise column correlations (Pearson/Spearman/Kendall) with optional BCa bootstrap CI              |
+| `corr`     | Pairwise column correlations (Pearson/Spearman/Kendall) with optional BCa bootstrap CI; globs (`-c '*:*'`) or a full matrix (`--all`) |
 | `describe` | Quick column-level summary (dtype, n, n_unique, n_null, sample values)                              |
 | `info`     | Per-column dtype, null counts, and memory usage; `--summary` for dataset-level totals               |
 | `randvar`  | Sample from a distribution and append as a new column (norm, alpha, beta, …)                        |
@@ -127,6 +127,8 @@ publication-quality output.
 | `line`  | Line plot; optional error bars (`--yerr`) or CI bands (`--yci lo,hi`) and fit overlay                                                                                               |
 | `hist`  | Histogram or KDE (`--kde`); normalisation, cumulative mode, mean±σ annotation (`--stats`)                                                                                           |
 | `chi`   | Corner-resolved chi-plot: whether tail dependence sits in the lower or upper corner                                                                                                 |
+| `heatmap` | Coloured grid of values from a long (`-x -y -v`) or wide (`-y`) table; correlograms from `corr --all`, `--annot`, `--lower` |
+| `pairs` | Scatter-plot matrix with KDEs (or histograms) on the diagonal; `-g` colours groups, `--upper corr` puts r in the upper triangle |
 
 ### Utilities
 
@@ -379,6 +381,34 @@ dftk eval segments.tsv -f 'gap = Begin - End.groupby(Chr).shift()'
 Without `groupby(Chr)` the first segment of a chromosome is compared with
 the last one of the previous chromosome.
 
+### Correlograms and heatmaps
+
+`corr --all` correlates every pair of numeric columns both ways round, plus the diagonal, which is the square table `heatmap` draws. Values of both signs get diverging colours centred on 0; `--vmin -1 --vmax 1` keeps the full scale for r, so a weak correlation doesn't look strong. `--annot` writes the values in the cells and `--lower` keeps the lower triangle:
+
+```bash
+dftk dataset iris -o | dftk corr --all -o \
+  | dftk heatmap -x col1 -y col2 -v correlation --annot --vmin -1 --vmax 1 -f iris_corr.png
+dftk dataset diamonds -o | dftk corr --all --method spearman -o \
+  | dftk heatmap -x col1 -y col2 -v correlation --annot --lower --square \
+      --vmin -1 --vmax 1 -f diamonds_corr.png
+```
+
+Missing values are dropped pair by pair, so `nobs` can differ between cells; `dftk print --dropna -o | dftk corr --all` uses the complete rows throughout. p-values are not adjusted for multiple testing.
+
+For the shape of each relationship, not just r, `pairs` draws every pair of columns as a scatter plot with each column's KDE on the diagonal; `-g` colours groups (each group's KDE scaled by its share of the rows) and `--upper corr` writes r in the upper triangle, with `-g` also r within each group, which can differ even in sign from r over all rows:
+
+```bash
+dftk dataset iris -o | dftk pairs -g species -f iris_pairs.png
+dftk dataset iris -o | dftk pairs -g species --upper corr --diag hist -f iris_pairs_r.png
+```
+
+`heatmap` draws any table with one row per cell (`-x -y -v`), or a wide one with `-y` naming the row labels and the other numeric columns as heatmap columns, as `pivot` makes it:
+
+```bash
+dftk dataset penguins -o | dftk pivot -i species -g island -v body_mass_g -f count -o \
+  | dftk heatmap -y species --annot --cmap Blues -f penguin_counts.png
+```
+
 ### Bootstrap confidence intervals
 
 Bootstrap mode repeats the full stat computation N times on resampled data

{dataframe_toolkit-0.15.0 → dataframe_toolkit-0.16.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "dataframe-toolkit"
-version = "0.15.0"
+version = "0.16.0"
 description = "A command-line toolkit for DataFrame analysis and manipulation"
 readme = "README.md"
 keywords = [

{dataframe_toolkit-0.15.0 → dataframe_toolkit-0.16.0}/pyproject.toml.orig

@@ -1,6 +1,6 @@
 [project]
 name = "dataframe-toolkit"
-version = "0.15.0"
+version = "0.16.0"
 description = "A command-line toolkit for DataFrame analysis and manipulation"
 readme = "README.md"
 authors = [

{dataframe_toolkit-0.15.0 → dataframe_toolkit-0.16.0}/src/dftk/commands/__init__.py

@@ -62,6 +62,8 @@ COMMAND_GROUPS: list[tuple[str, list[tuple[str, str]]]] = [
             ("line", "line_cmd:LineCommand"),
             ("hist", "hist_cmd:HistCommand"),
             ("chi", "chi_cmd:ChiCommand"),
+            ("heatmap", "heatmap_cmd:HeatmapCommand"),
+            ("pairs", "pairs_cmd:PairsCommand"),
         ],
     ),
     (

{dataframe_toolkit-0.15.0 → dataframe_toolkit-0.16.0}/src/dftk/commands/corr_cmd.py

@@ -11,6 +11,8 @@ Bootstrap mode (--bootstrap N) produces a single summary row per pair with:
 """
 
 import argparse
+import fnmatch
+import itertools
 import logging
 
 import numpy as np
@@ -140,14 +142,82 @@ def _tail_stats(a: np.ndarray, b: np.ndarray, q_grid: np.ndarray) -> tuple:
     return tuple(out)
 
 
+_GLOB_CHARS = frozenset("*?[")
+
+
+def _numeric_columns(df: pd.DataFrame, groups: list[str]) -> list[str]:
+    """Numeric, non-boolean columns of *df* other than the group columns."""
+    return [
+        c
+        for c in df.columns
+        if c not in groups
+        and pd.api.types.is_numeric_dtype(df[c])
+        and not pd.api.types.is_bool_dtype(df[c])
+    ]
+
+
+def _is_glob(pattern: str, numeric: list[str]) -> bool:
+    """A glob pattern, unless it is the name of a column (e.g. x[1])."""
+    return bool(_GLOB_CHARS & set(pattern)) and pattern not in numeric
+
+
+def _match(pattern: str, numeric: list[str], option: str) -> list[str]:
+    """A column name as is, or the numeric columns a glob pattern matches."""
+    if not _is_glob(pattern, numeric):
+        return [pattern]
+    matched = fnmatch.filter(numeric, pattern)
+    if not matched:
+        raise ValueError(
+            f"{option}: {pattern!r} matches no numeric column (numeric columns: "
+            f"{numeric})"
+        )
+    return matched
+
+
+def _expand_pairs(cols: list[str], numeric: list[str]) -> list[tuple[str, str]]:
+    """-c COL1:COL2 items as pairs; a glob on either side expands.
+
+    A pair from a glob is left out when both sides are the same column or
+    when the pair already came up the other way round, so -c '*:*' gives
+    every pair once.  Pairs named in full are kept as given.
+    """
+    bad = [c for c in cols if len(c.split(":", maxsplit=1)) != 2]
+    if bad:
+        raise ValueError(f"-c pairs must be col1:col2, got: {bad}")
+    pairs: list[tuple[str, str]] = []
+    for item in cols:
+        left, right = item.split(":", maxsplit=1)
+        if not (_is_glob(left, numeric) or _is_glob(right, numeric)):
+            pairs.append((left, right))
+            continue
+        for a, b in itertools.product(
+            _match(left, numeric, "-c"), _match(right, numeric, "-c")
+        ):
+            if a != b and (a, b) not in pairs and (b, a) not in pairs:
+                pairs.append((a, b))
+    return pairs
+
+
+def _all_pairs(cols: list[str], numeric: list[str]) -> list[tuple[str, str]]:
+    """--all [COL ...]: every ordered pair, diagonal included (a square matrix)."""
+    names: list[str] = []
+    for pattern in cols or ["*"]:
+        names += [c for c in _match(pattern, numeric, "--all") if c not in names]
+    if len(names) < 2:
+        raise ValueError(f"--all needs at least two numeric columns, got {names}")
+    return list(itertools.product(names, names))
+
+
 def _corr(
-    df: pd.DataFrame, args: argparse.Namespace, rng: np.random.Generator
+    df: pd.DataFrame,
+    args: argparse.Namespace,
+    rng: np.random.Generator,
+    pairs: list[tuple[str, str]] | None = None,
 ) -> pd.DataFrame:
+    """Correlations for *pairs* (default: args.cols, globs expanded)."""
+    if pairs is None:
+        pairs = _expand_pairs(args.cols, _numeric_columns(df, args.groups))
     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"
     tail = getattr(args, "tail", False)
@@ -163,20 +233,49 @@ def _corr(
             header += [f"lambda_lower_q{q:g}", f"lambda_upper_q{q:g}"]
 
     rows = []
+    n_stats = len(header) - 2  # nobs onwards
 
     def _add_pairs(sub: pd.DataFrame, prefix: tuple = ()):
+        # Every statistic is symmetric in the two columns, so --all's mirrored
+        # pair (b, a) reuses the result for (a, b)
+        done: dict[frozenset, 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)
-            extra = ()
-            if args.bootstrap is not None:
-                extra += _bootstrap_stats(
-                    a, b, method_fn, args.bootstrap, rng, args.confidence / 100.0
-                )
-            if tail:
-                extra += _tail_stats(a, b, q_grid)
-            rows.append((*prefix, col1, col2, *cr, *extra))
+            key = frozenset((col1, col2))
+            if key not in done:
+                if col1 == col2:  # the diagonal of --all; sub[[c, c]] is 2-D
+                    a = b = sub[col1].dropna().values
+                else:
+                    pair_df = sub[[col1, col2]].dropna()
+                    a, b = pair_df[col1].values, pair_df[col2].values
+                if len(a) < 2:
+                    where = f" (group {', '.join(map(str, prefix))})" if prefix else ""
+                    logger.warning(
+                        f"corr: {col1}:{col2}{where} has {len(a)} complete "
+                        "row(s); statistics left empty"
+                    )
+                    stats = (len(a),) + (np.nan,) * (n_stats - 1)
+                elif col1 == col2:
+                    # r is 1 by definition; a p-value or interval of a column
+                    # with itself means nothing and would count as a finding,
+                    # and ties push the tail statistics above 1
+                    r = _compute_correlation(a, b, method_fn, False)[1]
+                    r = 1.0 if np.isfinite(r) else r  # NaN: a constant column
+                    stats = (len(a), r) + (np.nan,) * (n_stats - 2)
+                else:
+                    stats = _compute_correlation(a, b, method_fn, ci)
+                    if args.bootstrap is not None:
+                        stats += _bootstrap_stats(
+                            a,
+                            b,
+                            method_fn,
+                            args.bootstrap,
+                            rng,
+                            args.confidence / 100.0,
+                        )
+                    if tail:
+                        stats += _tail_stats(a, b, q_grid)
+                done[key] = stats
+            rows.append((*prefix, col1, col2, *done[key]))
 
     if args.groups:
         header = [*args.groups, *header]
@@ -197,6 +296,36 @@ COLUMN PAIRS (-c)
     dftk corr data.tsv -c x:y
     dftk dataset tips -o | dftk corr -c total_bill:tip tip:size
 
+  A glob on either side expands to the numeric columns it matches (quote
+  it). '*:*' gives every pair once; 'sepal*:petal*' every sepal column
+  with every petal column:
+
+    dftk dataset iris -o | dftk corr -c '*:*'
+    dftk dataset iris -o | dftk corr -c 'sepal*:petal*' --method spearman
+
+CORRELATION MATRIX (--all)
+--------------------------
+  --all gives every pair both ways round plus the diagonal, in column order:
+  the square table a correlogram needs.  Name columns (globs allowed) to
+  limit it; -g gives one matrix per group.  Each pair is computed once.  On
+  the diagonal r is 1 and the other statistics are empty.
+
+    dftk dataset iris -o | dftk corr --all
+    dftk dataset iris -o | dftk corr --all -o \\
+      | dftk heatmap -x col1 -y col2 -v correlation --annot --vmin -1 --vmax 1
+    dftk dataset diamonds -o | dftk corr --all carat depth price -g cut
+
+  Missing values are dropped pair by pair, so nobs can differ from cell to
+  cell, and the matrix need not be a valid correlation matrix.  For one set
+  of rows throughout, drop incomplete rows first:
+
+    dftk dataset penguins -o | dftk print --dropna -o | dftk corr --all
+
+  p-values are not adjusted for multiple testing, and --all lists each
+  p-value twice.  Adjust over one triangle:
+
+    dftk dataset iris -o | dftk corr --all -o | dftk query -q 'col1 < col2'
+
 METHODS (--method)
 ------------------
   pearson    Pearson r (default)
@@ -287,13 +416,25 @@ class CorrCommand(BaseCommand):
         self.add_io_arguments(parser)
 
         g = parser.add_argument_group("corr options")
-        g.add_argument(
+        which = g.add_mutually_exclusive_group(required=True)
+        which.add_argument(
             "-c",
             "--cols",
             nargs="+",
-            required=True,
             metavar="COL1:COL2",
-            help="Column pair(s) to correlate, e.g. x:y.",
+            help="Column pair(s) to correlate, e.g. x:y.  A glob on either "
+            "side expands to numeric columns: '*:*' is every pair once, "
+            "'sepal*:petal*' every sepal column with every petal column.",
+        )
+        which.add_argument(
+            "--all",
+            nargs="*",
+            default=None,
+            metavar="COL",
+            help="Correlation matrix: every pair of the named columns (globs "
+            "allowed; default: all numeric columns) both ways round, plus the "
+            "diagonal, in column order.  Pipe it to dftk heatmap for a "
+            "correlogram.",
         )
         g.add_argument(
             "-g",
@@ -355,10 +496,15 @@ class CorrCommand(BaseCommand):
         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)
+        numeric = _numeric_columns(df, args.groups)
+        if args.all is not None:
+            pairs = _all_pairs(args.all, numeric)
+            check_cols(df, sorted({c for p in pairs for c in p}), "--all")
+        else:
+            pairs = _expand_pairs(args.cols, numeric)
+            check_cols(df, [c for p in pairs for c in p], "-c/--cols")
+
+        result = _corr(df, args, rng, pairs)
         io.printdf(result, args)

dataframe_toolkit-0.16.0/src/dftk/commands/heatmap_cmd.py

@@ -0,0 +1,253 @@
+"""
+dftk.commands.heatmap_cmd — dftk heatmap subcommand.
+
+A coloured grid of values, e.g. a correlogram from ``dftk corr --all``.
+
+Two input shapes:
+
+  long  (-x COL -y COL -v COL)  one row per cell, as corr and stat give it;
+                                 rows and columns keep their first-seen order
+  wide  (-y COL)                 one row per heatmap row; every other numeric
+                                 column (or --select) is a heatmap column,
+                                 as pivot gives it or a matrix file holds it
+"""
+
+import argparse
+
+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.plot import (
+    add_figure_arguments,
+    apply_style,
+    make_figure,
+    save_or_show,
+)
+
+
+def _matrix(df: pd.DataFrame, args: argparse.Namespace) -> pd.DataFrame:
+    """The heatmap as a DataFrame: index = y labels, columns = x labels."""
+    if args.xcol is None and args.value is None:
+        if args.ycol is None:
+            raise ValueError(
+                "give -x, -y and -v (one row per cell), or -y alone (one row "
+                "per heatmap row, other numeric columns as heatmap columns)"
+            )
+        check_cols(df, [args.ycol], "-y/--ycol")
+        values = [
+            c
+            for c in df.columns
+            if c != args.ycol
+            and pd.api.types.is_numeric_dtype(df[c])
+            and not pd.api.types.is_bool_dtype(df[c])
+        ]
+        if not values:
+            raise ValueError(f"no numeric columns besides -y {args.ycol}")
+        if df[args.ycol].duplicated().any():
+            raise ValueError(f"-y {args.ycol} has repeated values: one row each")
+        return df.set_index(args.ycol)[values].astype(float)
+
+    if None in (args.xcol, args.ycol, args.value):
+        raise ValueError("long input needs all of -x, -y and -v")
+    check_cols(df, [args.xcol, args.ycol, args.value], "-x/-y/-v")
+    keys = [args.ycol, args.xcol]
+    if df.duplicated(keys).any():
+        raise ValueError(
+            f"several rows for the same {args.xcol} and {args.ycol}: summarise "
+            "first (dftk pivot), or filter (--prequery)"
+        )
+    xs = pd.unique(df[args.xcol])
+    ys = pd.unique(df[args.ycol])
+    matrix = df.pivot(index=args.ycol, columns=args.xcol, values=args.value)
+    return matrix.reindex(index=ys, columns=xs).astype(float)
+
+
+def _fmt(values: np.ndarray, fmt: str | None) -> str:
+    """--fmt, or g for whole numbers (counts) and .2f otherwise."""
+    if fmt:
+        return fmt
+    finite = values[np.isfinite(values)]
+    return "g" if finite.size and np.all(finite == np.round(finite)) else ".2f"
+
+
+def _heatmap(ax, fig, matrix: pd.DataFrame, args: argparse.Namespace) -> None:
+    values = matrix.to_numpy()
+    shown = np.ma.masked_invalid(values)
+    if args.lower:
+        if matrix.shape[0] != matrix.shape[1]:
+            raise ValueError("--lower needs a square heatmap (as corr --all gives)")
+        shown = np.ma.masked_where(np.triu(np.ones_like(values), k=1) > 0, shown)
+
+    vmin, vmax, center = args.vmin, args.vmax, args.center
+    finite = values[np.isfinite(values)]
+    if center is None and finite.size and finite.min() < 0 < finite.max():
+        center = 0.0  # values of both signs: diverging colours around 0
+    cmap = args.cmap or ("RdBu_r" if center is not None else "viridis")
+    if center is not None and finite.size:
+        # Symmetric limits so that center gets the middle colour
+        half = max(abs(finite.max() - center), abs(finite.min() - center))
+        vmin = center - half if vmin is None else vmin
+        vmax = center + half if vmax is None else vmax
+
+    mesh = ax.pcolormesh(shown, cmap=cmap, vmin=vmin, vmax=vmax)
+    fig.colorbar(mesh, ax=ax, label=args.value or "")
+
+    ax.set_xticks(np.arange(matrix.shape[1]) + 0.5)
+    ax.set_xticklabels([str(c) for c in matrix.columns], rotation=45, ha="right")
+    ax.set_yticks(np.arange(matrix.shape[0]) + 0.5)
+    ax.set_yticklabels([str(i) for i in matrix.index])
+    ax.invert_yaxis()  # first row at the top, as in the table
+    ax.set_xlabel(args.xlabel if args.xlabel is not None else (args.xcol or ""))
+    ax.set_ylabel(args.ylabel if args.ylabel is not None else (args.ycol or ""))
+    ax.set_title(args.title or "")
+    ax.set_aspect("equal" if args.square else "auto")
+
+    if args.annot:
+        fmt = _fmt(values, args.fmt)
+        for (i, j), v in np.ndenumerate(shown.filled(np.nan)):
+            if np.isfinite(v):
+                # Dark text on light cells, light text on dark cells
+                r, g, b, _ = mesh.cmap(mesh.norm(v))
+                light = 0.299 * r + 0.587 * g + 0.114 * b > 0.5
+                ax.text(
+                    j + 0.5,
+                    i + 0.5,
+                    format(v, fmt),
+                    ha="center",
+                    va="center",
+                    color="black" if light else "white",
+                )
+
+
+_EPILOG = """\
+INPUT
+-----
+Long, one row per cell (-x, -y, -v): rows and columns of the heatmap keep
+the order in which their labels first appear.
+
+  dftk heatmap data.tsv -x sample -y gene -v expression
+  dftk dataset tips -o | dftk pivot -i day time -v tip -f mean -o \\
+    | dftk heatmap -x time -y day -v mean_tip --annot
+
+Wide, one row per heatmap row (-y alone): every other numeric column is a
+heatmap column (--select picks and orders them):
+
+  dftk dataset tips -o | dftk pivot -i day -g time -v tip -f mean -o \\
+    | dftk heatmap -y day --annot
+
+CORRELOGRAM
+-----------
+corr --all gives every pair of columns both ways round, which is a square
+heatmap. Values of both signs get diverging colours centred on 0.  Give
+--vmin -1 --vmax 1 for correlations: otherwise the scale stretches to the
+largest |r|, and a weak correlation looks as strong as a strong one would:
+
+  dftk dataset iris -o | dftk corr --all -o \\
+    | dftk heatmap -x col1 -y col2 -v correlation --annot --vmin -1 --vmax 1
+
+  --lower keeps the lower triangle and the diagonal; --square makes the
+  cells square:
+
+  dftk dataset diamonds -o | dftk corr --all --method spearman -o \\
+    | dftk heatmap -x col1 -y col2 -v correlation --annot --lower --square \\
+        --vmin -1 --vmax 1
+
+COLOURS AND NUMBERS
+-------------------
+--cmap takes any matplotlib colour map (viridis, magma, RdBu_r, coolwarm,
+...). --annot writes each value in its cell, with --fmt as a Python format
+(.2f, .1%, g; default g for whole numbers, else .2f). Empty cells stay
+blank.
+
+  dftk dataset penguins -o | dftk pivot -i species -g island -v body_mass_g \\
+      -f count -o \\
+    | dftk heatmap -y species --annot --cmap Blues -f counts.png
+"""
+
+
+class HeatmapCommand(BaseCommand):
+    name = "heatmap"
+    help = "Coloured grid of values, e.g. a correlogram from corr --all."
+
+    def add_arguments(self, parser: argparse.ArgumentParser) -> None:
+        parser.formatter_class = argparse.RawDescriptionHelpFormatter
+        parser.epilog = _EPILOG
+
+        self.add_read_arguments(parser)
+
+        g = parser.add_argument_group("heatmap")
+        g.add_argument("-x", "--xcol", metavar="COL", help="Heatmap column labels.")
+        g.add_argument(
+            "-y",
+            "--ycol",
+            metavar="COL",
+            help="Heatmap row labels.  Alone, without -x and -v: the other "
+            "numeric columns are the heatmap columns.",
+        )
+        g.add_argument("-v", "--value", metavar="COL", help="Cell values.")
+        g.add_argument(
+            "--select",
+            nargs="+",
+            default=None,
+            metavar="COL",
+            help="Wide input: the value columns to show, in this order.",
+        )
+        g.add_argument("--xlabel", default=None, help="X-axis label (default: -x).")
+        g.add_argument("--ylabel", default=None, help="Y-axis label (default: -y).")
+        g.add_argument("--annot", action="store_true", help="Write values in cells.")
+        g.add_argument(
+            "--fmt",
+            default=None,
+            metavar="FMT",
+            help="Format of --annot values, e.g. .2f or .1%% (default: g for "
+            "whole numbers, else .2f).",
+        )
+        g.add_argument(
+            "--cmap",
+            default=None,
+            help="Matplotlib colour map (default: RdBu_r around --center, "
+            "else viridis).",
+        )
+        g.add_argument(
+            "--center",
+            type=float,
+            default=None,
+            help="Value that gets the middle colour (default: 0 when values "
+            "have both signs).",
+        )
+        g.add_argument("--vmin", type=float, default=None, help="Colour scale minimum.")
+        g.add_argument("--vmax", type=float, default=None, help="Colour scale maximum.")
+        g.add_argument(
+            "--lower",
+            action="store_true",
+            help="Show only the lower triangle and diagonal (square heatmaps).",
+        )
+        g.add_argument("--square", action="store_true", help="Square cells.")
+
+        add_figure_arguments(parser)
+
+    def execute(self, args: argparse.Namespace) -> None:
+        df = io.read(args)
+        if args.select:
+            check_cols(df, args.select, "--select")
+            keep = [args.ycol] if args.ycol else []
+            df = df[keep + [c for c in args.select if c not in keep]]
+        matrix = _matrix(df, args)
+
+        if args.file:
+            import matplotlib
+
+            matplotlib.use("Agg")
+
+        import matplotlib.pyplot as plt
+
+        if getattr(args, "usetex", False):
+            plt.rc("text", usetex=True)
+
+        apply_style(args)
+        fig, axes = make_figure(args)
+        _heatmap(axes[0][0], fig, matrix, args)
+        fig.tight_layout()
+        save_or_show(fig, args)

{dataframe_toolkit-0.15.0 → dataframe_toolkit-0.16.0}/src/dftk/commands/merge_cmd.py

@@ -29,6 +29,7 @@ import uuid
 import pandas as pd
 
 from dftk.commands.base import BaseCommand
+from dftk.common import docformats
 from dftk.common.io import check_cols, io
 
 logger = logging.getLogger(__name__)
@@ -113,6 +114,24 @@ def _add_cross_col(
     args.drop = drop
 
 
+def _carry_json_columns(
+    merged: pd.DataFrame,
+    left: pd.DataFrame,
+    right: pd.DataFrame,
+    suffixes: tuple[str, str],
+) -> None:
+    """Mark the JSON-text columns of both inputs in *merged* (see docformats)."""
+    names = []
+    for side, other, suffix in ((left, right, suffixes[0]), (right, left, suffixes[1])):
+        for col in docformats.json_columns(side):
+            name = col
+            if col in other.columns and f"{col}{suffix}" in merged.columns:
+                name = f"{col}{suffix}"
+            if name in merged.columns and name not in names:
+                names.append(name)
+    docformats.mark_json_columns(merged, names)
+
+
 def _gap_sides(args: argparse.Namespace) -> tuple[bool, bool]:
     """Return (left, right): which side's columns can get missing values.
 
@@ -227,6 +246,7 @@ def _do_merge(
     if "_merge" in df_merged.columns:
         df_merged = df_merged.drop(columns=["_merge"])
 
+    _carry_json_columns(df_merged, df_left, df_right, suffixes)
     return df_merged