densitty 0.8.2__py3-none-any.whl → 1.0.0__py3-none-any.whl

densitty/__init__.py

@@ -0,0 +1,9 @@
+"""In-terminal 2-D histogram/density/heatmap plotting"""
+
+from .detect import densityplot2d, grid_heatmap, histplot2d, plot
+from .detect import GRAYSCALE, FADE_IN, RAINBOW, REV_RAINBOW
+from .axis import Axis
+from .plotting import Plot
+from .colorbar import make_colorbar
+from .binning import histogram2d
+from .smoothing import smooth2d

densitty/ansi.py

@@ -8,7 +8,6 @@ from typing import Optional, Sequence
 
 from .util import nearest
 
-
 RESET = "\033[0m"
 
 

densitty/axis.py

@@ -4,15 +4,10 @@ import dataclasses
 from decimal import Decimal
 import itertools
 import math
-from typing import Optional
+from typing import Optional, Sequence
 
-from . import lineart
-from .util import FloatLike, ValueRange, pick_step_size
-
-MIN_X_TICKS_PER_LABEL = 4
-MIN_Y_TICKS_PER_LABEL = 2
-DEFAULT_X_COLS_PER_TICK = 4
-DEFAULT_Y_ROWS_PER_TICK = 2
+from . import lineart, util
+from .util import FloatLike, ValueRange
 
 
 @dataclasses.dataclass
@@ -31,16 +26,63 @@ x_border = {False: BorderChars(" ", " ", " "), True: BorderChars("╶", "─", "
 # Helper functions used by the Axis class below
 
 
-def add_label(line: list[str], label: str, ctr_pos: int):
+def pick_step_size(lower_bound) -> Decimal:
+    """Try to pick a step size that gives nice round values for step positions."""
+    if lower_bound <= 0:
+        raise ValueError("pick_step_size called with 0 or negative value")
+
+    _, frac, exp = util.sfrexp10(lower_bound)  # 0.1 <= frac < 1.0
+
+    # round up to an appropriate "round" value that starts with 1/2/5
+    if frac <= Decimal("0.2"):
+        out = Decimal((0, (2,), exp - 1))
+    elif frac <= Decimal("0.5"):
+        out = Decimal((0, (5,), exp - 1))
+    else:
+        out = Decimal((0, (1,), exp))
+
+    if out >= 10:
+        # this will be printed as 1E1 or such. Give it more significant figures so
+        # it is printed as an integer:
+        return out.quantize(1)
+    return out
+
+
+def add_x_label(line: list[str], label: str, ctr_pos: FloatLike):
     """Adds the label string to the output line, centered at specified position
     The output line is a list of single-character strings, to make this kind of thing
     straightforward"""
     width = len(label)
-    start_col = max(ctr_pos - width // 2, 0)
+    start_col = int(max(float(ctr_pos) + 0.49 - width / 2, 0))
     end_col = start_col + width
     line[start_col:end_col] = list(label)
 
 
+def add_x_tick(line: list[str], start_pos: float, left_margin: int, num_cols: int):
+    """Adds a (possibly fractional) tick to the output line, centered at specified position"""
+    base_idx = int(start_pos)
+    fractional_pos = start_pos - base_idx
+    margined_idx = base_idx + left_margin
+    if 0.25 <= fractional_pos <= 0.75:
+        # add a tick in the center of the appropriate bin
+        line[margined_idx] = lineart.merge_chars("│", line[margined_idx])
+        return
+    if fractional_pos < 0.25:
+        left_idx = max(margined_idx - 1, 0)
+    else:
+        left_idx = margined_idx
+
+    if base_idx == 0:
+        line[left_idx] = lineart.merge_chars("│", line[left_idx])
+    else:
+        line[left_idx] = "╱"
+
+    if base_idx >= num_cols - 1:
+        line[left_idx + 1] = lineart.merge_chars("│", line[left_idx + 1])
+    else:
+        line[left_idx + 1] = "╲"
+
+
 def gen_tick_values(value_range, tick_step):
     """Produce tick values in the specified range. Basically numpy.arange"""
 
@@ -50,29 +92,142 @@ def gen_tick_values(value_range, tick_step):
         tick += tick_step
 
 
-def gen_labels(
-    value_range: ValueRange, num_ticks, min_ticks_per_label, fmt, label_end_ticks=False
+def positions_to_labels(
+    printed_positions: Sequence[Decimal], ticked_positions: Sequence[FloatLike], fmt: str
+) -> dict[FloatLike, str]:
+    """Given positions, construct a label dict mapping position to string to be printed there.
+    The first positions get a printed value, the second set get a blank, for a bare 'tick'"""
+
+    tick_dict = {p: "" for p in ticked_positions}
+    printed_dict = {p: fmt.format(p) for p in util.sanitize_decimals(printed_positions)}
+    return tick_dict | printed_dict
+
+
+def calc_min_step(
+    value_range: ValueRange,
+    bin_width: FloatLike,
+    accomodate_values: bool,
+    tick_space: int,
+    fmt: str,
 ):
+    """Calculate minimum step size for tick placement.
+
+    When accomodate_values is True, considers the width of printed labels.
+    Otherwise, only considers tick spacing requirements.
+    """
+    if accomodate_values:
+        # find a representative printed label width to use for basic calculations
+        test_tick_step = pick_step_size((value_range.max - value_range.min) / 5)
+        test_values_printed = tuple(
+            fmt.format(value) for value in gen_tick_values(value_range, test_tick_step)
+        )
+        widths_in_bins = tuple(len(p) for p in test_values_printed)
+        # get the 3'd lowest width, or highest if there are less than 3 due to tick-step roundup:
+        example_width = sorted(widths_in_bins)[:3][-1]
+        min_printed_step = (example_width + 1) * bin_width
+        # If the printed labels are small (single-digit), the ticks themselves might be the
+        # limiting factor, especially if they are X-axis fractional ticks like "/\"
+        return max(min_printed_step, bin_width * (2 + tick_space * 2))
+    return bin_width * (2 + tick_space)
+
+
+def gen_position_subsets(positions: tuple, tick_step: Decimal) -> list[Sequence[Decimal]]:
+    """Generate candidate label position subsets based on tick step digit.
+
+    For tick steps starting with 1 or 2: generates every-5th subsets (5 variants).
+    For tick steps starting with 5: generates every-2nd subsets (2 variants).
+    """
+    step_digit = tick_step.as_tuple().digits[0]  # leading digit of tick step: 1, 2, or 5
+    # we want to pick different position subsets depending on whether we're advancing by
+    # 1eX, 2eX, or 5eX:
+    position_subsets = []
+    if step_digit in (1, 2):
+        # print on every fifth one, starting at 0..4
+        position_subsets += list(util.sanitize_decimals(positions[start::5]) for start in range(5))
+    elif step_digit == 5:
+        # print on every second one, starting at 0 or 1
+        position_subsets += list(util.sanitize_decimals(positions[start::2]) for start in range(2))
+    return position_subsets
+
+
+def label_ends_only(positions, tick_step, bin_width, accomodate_values, fmt):
+    """See if printing just the labels for the first and last ticks will fit"""
+
+    if not accomodate_values:
+        # For Y axis / we don't care about printed widths
+        return positions_to_labels((positions[0], positions[-1]), positions, fmt)
+
+    end_positions = (positions[0], positions[-1])
+    ends_printed = tuple(positions_to_labels(end_positions, [], fmt).values())
+
+    half_len_first = len(ends_printed[0]) // 2
+    half_len_last = (len(ends_printed[1]) + 1) // 2
+    space_available = math.floor(tick_step / bin_width)
+
+    if half_len_first + half_len_last <= space_available:
+        # We can fit values on the first and last ticks
+        return positions_to_labels(end_positions, positions, fmt)
+
+    # Not enough space to print two values, so just print the first
+    return positions_to_labels(end_positions[:1], positions, fmt)
+
+
+def find_fitting_subset(position_subsets, ticks_per_bin, accomodate_values, fmt):
+    """Find roundest label subset that fits within space constraints"""
+    for position_subset in util.roundness_ordered(position_subsets):
+        if not accomodate_values:  # it doesn't matter what the printed widths are, it will fit
+            return position_subset
+        # We're printing at most one value for every 2 ticks, so just make sure
+        # that the printed values will not run over the adjacent ticks' area
+        # Given the initial min_step logic, this will likely always be true
+        printed_widths = (
+            len(label) for label in positions_to_labels(position_subset, [], fmt).values()
+        )
+        allowed_width = ticks_per_bin * 2 - 2
+        if max(printed_widths) <= allowed_width:
+            return position_subset
+    return tuple()
+
+
+def gen_full_labels(value_range: ValueRange, num_bins, accomodate_values, tick_space, fmt):
     """Generate positions for labels (plain ticks & ticks with value)"""
-    tick_step, label_step = pick_step_size(value_range, num_ticks, min_ticks_per_label)
+    bin_width = (value_range.max - value_range.min) / num_bins
+
+    if bin_width <= 0:
+        # we don't have a sensible range for the axis values, so just have empty ticks
+        return {}
+
+    min_step = calc_min_step(value_range, bin_width, accomodate_values, tick_space, fmt)
+    cur_tick_step = pick_step_size(min_step)
+
+    while True:
+        positions = tuple(gen_tick_values(value_range, cur_tick_step))
+
+        if len(positions) == 0:
+            # No suitable ticks/labels.
+            if accomodate_values:
+                # Printed value labels won't fit, try without them
+                return gen_full_labels(value_range, num_bins, False, tick_space, "")
+            # Nothing fits
+            return {}
+        if len(positions) == 1:
+            return positions_to_labels(positions, [], fmt)
 
-    ticks = list(gen_tick_values(value_range, tick_step))
-    label_values = list(gen_tick_values(value_range, label_step))
-    if label_end_ticks or len(label_values) <= 2:
-        # ensure that first & last ticks have labels:
-        if label_values[0] != ticks[0]:
-            label_values = ticks[:1] + label_values
-        if label_values[-1] != ticks[-1]:
-            label_values += ticks[-1:]
+        position_subsets = gen_position_subsets(positions, cur_tick_step)
 
-    # sanity: if all but one ticks have labels, just label them all
-    if len(label_values) >= len(ticks) - 1:
-        label_values = ticks
+        # Check to see if all generated label subsets only have a single entry
+        if max(len(subset) for subset in position_subsets) == 1:
+            # Try to just label the ends:
+            return label_ends_only(positions, cur_tick_step, bin_width, accomodate_values, fmt)
 
-    ticks_only = {value: "" for value in ticks}
-    labeled_ticks = {value: fmt.format(value) for value in label_values}
+        best_subset = find_fitting_subset(
+            position_subsets, cur_tick_step / bin_width, accomodate_values, fmt
+        )
+
+        if best_subset:
+            return positions_to_labels(best_subset, positions, fmt)
 
-    return ticks_only | labeled_ticks
+        cur_tick_step = pick_step_size(float(cur_tick_step) * 1.01)
 
 
 def calc_edges(value_range, num_bins, values_are_edges):
@@ -88,7 +243,7 @@ def calc_edges(value_range, num_bins, values_are_edges):
     values_are_edges: bool
                  Indicates that value_range specifies outside edges rather than bin centers
     """
-    if values_are_edges:
+    if values_are_edges or num_bins == 1:
         bin_delta = (value_range.max - value_range.min) / num_bins
         first_bin_min = value_range.min
     else:
@@ -107,7 +262,9 @@ class Axis:
     """Options for axis generation."""
 
     value_range: ValueRange  # can also specify as a tuple of (min, max)
-    labels: Optional[dict[float, str]] = None  # map axis value to label (plus tick) at that value
+    labels: Optional[dict[FloatLike, str]] = (
+        None  # map axis value to label (plus tick) at that value
+    )
     label_fmt: str = "{}"  # format for generated labels
     border_line: bool = False  # embed ticks in a horizontal X-axis or vertical Y-axis line
     values_are_edges: bool = False  # N+1 values, indicating boundaries between pixels, not centers
@@ -116,7 +273,7 @@ class Axis:
     def __init__(
         self,
         value_range: ValueRange | tuple[FloatLike, FloatLike],
-        labels: Optional[dict[float, str]] = None,
+        labels: Optional[dict[FloatLike, str]] = None,
         label_fmt: str = "{}",
         border_line: bool = False,
         values_are_edges: bool = False,
@@ -137,10 +294,11 @@ class Axis:
     def _unjustified_y_axis(self, num_rows: int):
         """Returns the Y axis string for each line of the plot"""
         if self.labels is None:
-            labels = gen_labels(
+            labels = gen_full_labels(
                 self.value_range,
-                num_rows // DEFAULT_Y_ROWS_PER_TICK,
-                MIN_Y_TICKS_PER_LABEL,
+                num_rows,
+                False,
+                0,
                 self.label_fmt,
             )
         else:
@@ -154,10 +312,12 @@ class Axis:
             if label_values and row_min <= label_values[0] <= row_max:
                 label_str = labels[label_values[0]]
 
-                offset_frac = (label_values[0] - row_min) / (row_max - row_min)
-                if offset_frac < 0.25 and self.fractional_tick_pos:
+                offset_frac = (float(label_values[0]) - float(row_min)) / float(row_max - row_min)
+                # Try to avoid our cutoffs being exactly where roundoff errors can pop up and
+                # cause inconsistent behavior. So 0.249 rather than 0.25, 0.751 rather than 0.75:
+                if offset_frac < 0.249 and self.fractional_tick_pos:
                     tick_char = "▔"
-                elif offset_frac > 0.75 and self.fractional_tick_pos:
+                elif offset_frac > 0.751 and self.fractional_tick_pos:
                     tick_char = "▁"
                 else:
                     tick_char = "─"
@@ -215,14 +375,9 @@ class Axis:
         left_margin: int
                      chars to the left of leftmost data col. May have Labels/border-line.
         """
-
         if self.labels is None:
-            labels = gen_labels(
-                self.value_range,
-                num_cols // DEFAULT_X_COLS_PER_TICK,
-                MIN_X_TICKS_PER_LABEL,
-                self.label_fmt,
-            )
+            tick_space = 1 if self.fractional_tick_pos else 0
+            labels = gen_full_labels(self.value_range, num_cols, True, tick_space, self.label_fmt)
         else:
             labels = self.labels
 
@@ -242,24 +397,32 @@ class Axis:
         for col_idx, (col_min, col_max) in enumerate(bins):
             # use Decimal.next_plus to accomodate rounding error/truncation
             if label_values and col_min <= label_values[0] <= col_max.next_plus():
-                add_label(label_line, labels[label_values[0]], col_idx + left_margin)
-                tick_idx = left_margin + col_idx
-                offset_frac = (label_values[0] - col_min) / (col_max - col_min)
-                if self.fractional_tick_pos and offset_frac < 0.25:
-                    if col_idx == 0:
-                        tick_line[tick_idx - 1] = lineart.merge_chars("│", tick_line[tick_idx - 1])
-                    else:
-                        tick_line[tick_idx - 1] = "╱"
-                    tick_line[tick_idx] = "╲"
-                elif self.fractional_tick_pos and offset_frac > 0.75:
-                    tick_line[tick_idx] = "╱"
-                    if col_idx < num_cols - 1:
-                        tick_line[tick_idx + 1] = "╲"
-                    else:
-                        tick_line[tick_idx + 1] = lineart.merge_chars("│", tick_line[tick_idx + 1])
+                if self.fractional_tick_pos:
+                    offset_frac = (label_values[0] - col_min) / (col_max - col_min)
                 else:
-                    tick_line[tick_idx] = lineart.merge_chars("│", tick_line[tick_idx])
+                    offset_frac = 0.5  # not doing fractional tick positioning == center the tick
+
+                add_x_tick(tick_line, col_idx + offset_frac, left_margin, num_cols)
+                add_x_label(
+                    label_line, labels[label_values[0]], col_idx + left_margin + offset_frac
+                )
 
                 label_values = label_values[1:]  # pop that first label since we added it
 
         return "".join(tick_line), "".join(label_line)
+
+    def upscale(self, new_num_bins, multiplier):
+        """Adjust axis for an upscaled plot"""
+        if self.values_are_edges:
+            # upscaling doesn't move the axis edges. First edge of first bin is the same.
+            return
+        old_num_bins = new_num_bins // multiplier
+        old_bin_width = (self.value_range.max - self.value_range.min) / (old_num_bins - 1)
+        new_bin_width = old_bin_width / multiplier
+        # we used to have a value for the center of the first bin, but now that bin is
+        # 'multiplier' bins, and the value corresponds to the center of that _group_ of bins
+        # Find the offset between that group center, and the center of the new edge bin:
+        offset = new_bin_width * (multiplier - 1) / 2
+        self.value_range = util.make_value_range(
+            (self.value_range.min - offset, self.value_range.max + offset)
+        )