flopscope 0.2.0__py3-none-any.whl

flopscope/_opt_einsum/_helpers.py

@@ -0,0 +1,164 @@
+"""Contains helper functions for opt_einsum testing scripts."""
+
+from collections.abc import Collection, Iterable, Sequence
+from typing import Any, overload
+
+# Inline type aliases (formerly from ._typing, deleted in Task 7+8).
+ArrayIndexType = frozenset  # frozenset[str]
+ArrayType = object  # Any
+
+__all__ = ["compute_size_by_dict", "find_contraction", "flop_count"]
+
+_valid_chars = "abcdefghijklmopqABC"
+_sizes = [2, 3, 4, 5, 4, 3, 2, 6, 5, 4, 3, 2, 5, 7, 4, 3, 2, 3, 4]
+_default_dim_dict = dict(zip(_valid_chars, _sizes, strict=False))
+
+
+@overload
+def compute_size_by_dict(indices: Iterable[int], idx_dict: list[int]) -> int: ...
+
+
+@overload
+def compute_size_by_dict(indices: Collection[str], idx_dict: dict[str, int]) -> int: ...
+
+
+def compute_size_by_dict(indices: Any, idx_dict: Any) -> int:
+    """Computes the product of the elements in indices based on the dictionary
+    idx_dict.
+
+    Parameters
+    ----------
+    indices : iterable
+        Indices to base the product on.
+    idx_dict : dictionary
+        Dictionary of index _sizes
+
+    Returns:
+    -------
+    ret : int
+        The resulting product.
+
+    Examples:
+    --------
+    >>> compute_size_by_dict('abbc', {'a': 2, 'b':3, 'c':5})
+    90
+
+    """
+    ret = 1
+    for i in indices:  # lgtm [py/iteration-string-and-sequence]
+        ret *= idx_dict[i]
+    return ret
+
+
+def find_contraction(
+    positions: Collection[int],
+    input_sets: list[ArrayIndexType],
+    output_set: ArrayIndexType,
+) -> tuple[frozenset[str], list[ArrayIndexType], ArrayIndexType, ArrayIndexType]:
+    """Finds the contraction for a given set of input and output sets.
+
+    Parameters
+    ----------
+    positions : iterable
+        Integer positions of terms used in the contraction.
+    input_sets : list
+        List of sets that represent the lhs side of the einsum subscript
+    output_set : set
+        Set that represents the rhs side of the overall einsum subscript
+
+    Returns:
+    -------
+    new_result : set
+        The indices of the resulting contraction
+    remaining : list
+        List of sets that have not been contracted, the new set is appended to
+        the end of this list
+    idx_removed : set
+        Indices removed from the entire contraction
+    idx_contraction : set
+        The indices used in the current contraction
+
+    Examples:
+    --------
+    # A simple dot product test case
+    >>> pos = (0, 1)
+    >>> isets = [set('ab'), set('bc')]
+    >>> oset = set('ac')
+    >>> find_contraction(pos, isets, oset)
+    ({'a', 'c'}, [{'a', 'c'}], {'b'}, {'a', 'b', 'c'})
+
+    # A more complex case with additional terms in the contraction
+    >>> pos = (0, 2)
+    >>> isets = [set('abd'), set('ac'), set('bdc')]
+    >>> oset = set('ac')
+    >>> find_contraction(pos, isets, oset)
+    ({'a', 'c'}, [{'a', 'c'}, {'a', 'c'}], {'b', 'd'}, {'a', 'b', 'c', 'd'})
+    """
+    remaining = list(input_sets)
+    inputs = (remaining.pop(i) for i in sorted(positions, reverse=True))
+    idx_contract = frozenset.union(*inputs)
+    idx_remain = output_set.union(*remaining)
+
+    new_result = idx_remain & idx_contract
+    idx_removed = idx_contract - new_result
+    remaining.append(new_result)
+
+    return new_result, remaining, idx_removed, idx_contract
+
+
+def flop_count(
+    idx_contraction: Collection[str],
+    inner: bool,
+    num_terms: int,
+    size_dictionary: dict[str, int],
+    *,
+    input_subscripts: Sequence[str] | None = None,
+    output_subscript: str | None = None,
+    input_shapes: Sequence[Sequence[int]] | None = None,
+) -> int:
+    """Per-step dense FLOP count (no declared symmetry).
+
+    Routes through compute_accumulation_cost so the result is consistent
+    with what ``einsum_accumulation_cost`` would return for an einsum with
+    no declared input symmetry. The α/M formula produces textbook FMA=2
+    scalar-op counts (multiplies and adds counted separately, with the
+    off-by-one accumulator-init credit applied).
+
+    Legacy fallback (when subscripts/shapes are not provided): the old
+    overall_size * op_factor formula with FMA=1 (inner does not add an extra
+    op).
+
+    Examples:
+    --------
+    >>> flop_count('abc', False, 1, {'a': 2, 'b':3, 'c':5})
+    30
+
+    >>> flop_count('abc', True, 2, {'a': 2, 'b':3, 'c':5})
+    30
+
+    """
+    if input_subscripts is None or input_shapes is None:
+        overall_size = compute_size_by_dict(idx_contraction, size_dictionary)
+        op_factor = max(1, num_terms - 1)
+        return overall_size * op_factor
+
+    from flopscope._accumulation._cost import compute_accumulation_cost
+
+    canonical = ",".join(input_subscripts) + "->" + (output_subscript or "")
+    per_op_syms = (None,) * len(input_subscripts)
+    cost = compute_accumulation_cost(
+        canonical_subscripts=canonical,
+        input_parts=tuple(input_subscripts),
+        output_subscript=output_subscript or "",
+        shapes=tuple(tuple(s) for s in input_shapes),
+        per_op_symmetries=per_op_syms,
+        identity_pattern=None,
+    )
+    return max(cost.total, 1)
+
+
+def has_array_interface(array: ArrayType) -> ArrayType:
+    if hasattr(array, "__array_interface__"):
+        return True
+    else:
+        return False

flopscope/_opt_einsum/_hsluv.py

@@ -0,0 +1,273 @@
+"""Minimal HSLuv helpers for terminal label palettes.
+
+Adapted from the MIT-licensed ``hsluv-python`` reference implementation:
+https://github.com/hsluv/hsluv-python
+
+The palette construction idea mirrors seaborn's ``husl_palette`` API, but
+this module does not vendor seaborn code or add a seaborn dependency.
+See NOTICE in this package for attribution details.
+"""
+
+from __future__ import annotations
+
+import math
+from functools import cache
+
+_M = (
+    (3.240969941904521, -1.537383177570093, -0.498610760293),
+    (-0.96924363628087, 1.87596750150772, 0.041555057407175),
+    (0.055630079696993, -0.20397695888897, 1.056971514242878),
+)
+_M_INV = (
+    (0.41239079926595, 0.35758433938387, 0.18048078840183),
+    (0.21263900587151, 0.71516867876775, 0.072192315360733),
+    (0.019330818715591, 0.11919477979462, 0.95053215224966),
+)
+_REF_Y = 1.0
+_REF_U = 0.19783000664283
+_REF_V = 0.46831999493879
+_KAPPA = 903.2962962
+_EPSILON = 0.0088564516
+
+# Frozen to preserve the current Rich label colors. Keep this in sync with
+# qualitative_hsluv_palette(64, lightness=0.5, saturation=0.9).
+PRECOMPUTED_QUALITATIVE_HSLUV_PALETTE_64 = (
+    "#E3245A",
+    "#2F70EF",
+    "#228921",
+    "#B92FED",
+    "#987022",
+    "#268389",
+    "#7B5EF0",
+    "#D529A6",
+    "#D14923",
+    "#5E8222",
+    "#238755",
+    "#2A7CBC",
+    "#E7232C",
+    "#9D4AF0",
+    "#5569F0",
+    "#DD2680",
+    "#C82CCA",
+    "#B75F22",
+    "#7C7B22",
+    "#408722",
+    "#24856F",
+    "#22893B",
+    "#2C76D5",
+    "#2880A3",
+    "#E62443",
+    "#DF3623",
+    "#6864F0",
+    "#D92893",
+    "#8D55F0",
+    "#AC3DF0",
+    "#426DF0",
+    "#E1256D",
+    "#CF2AB8",
+    "#C12EDC",
+    "#A86822",
+    "#C55522",
+    "#6D7E22",
+    "#4F8422",
+    "#318822",
+    "#8A7622",
+    "#2E73E2",
+    "#25847C",
+    "#278196",
+    "#228848",
+    "#238662",
+    "#22892E",
+    "#2B79C8",
+    "#297EAF",
+    "#E42B23",
+    "#D84023",
+    "#E5244F",
+    "#E72337",
+    "#A544F0",
+    "#845AF0",
+    "#C52DD3",
+    "#CC2BC1",
+    "#DC278A",
+    "#5F67F0",
+    "#7161F0",
+    "#D22AAF",
+    "#9550F0",
+    "#D7289D",
+    "#DF2676",
+    "#4B6BF0",
+)
+
+
+def _distance_line_from_origin(line: tuple[float, float]) -> float:
+    slope, intercept = line
+    return abs(intercept) / math.sqrt(slope**2 + 1)
+
+
+def _length_of_ray_until_intersect(theta: float, line: tuple[float, float]) -> float:
+    slope, intercept = line
+    return intercept / (math.sin(theta) - slope * math.cos(theta))
+
+
+def _get_bounds(lightness: float) -> list[tuple[float, float]]:
+    result = []
+    sub1 = ((lightness + 16) ** 3) / 1560896
+    sub2 = sub1 if sub1 > _EPSILON else lightness / _KAPPA
+
+    for channel in range(3):
+        m1, m2, m3 = _M[channel]
+        for t in range(2):
+            top1 = (284517 * m1 - 94839 * m3) * sub2
+            top2 = (
+                838422 * m3 + 769860 * m2 + 731718 * m1
+            ) * lightness * sub2 - 769860 * t * lightness
+            bottom = (632260 * m3 - 126452 * m2) * sub2 + 126452 * t
+            result.append((top1 / bottom, top2 / bottom))
+
+    return result
+
+
+def _max_chroma_for_lh(lightness: float, hue: float) -> float:
+    hue_radians = math.radians(hue)
+    lengths = [
+        _length_of_ray_until_intersect(hue_radians, bound)
+        for bound in _get_bounds(lightness)
+    ]
+    return min(length for length in lengths if length >= 0)
+
+
+def _dot_product(
+    left: tuple[float, float, float], right: tuple[float, float, float]
+) -> float:
+    return sum(
+        left_value * right_value
+        for left_value, right_value in zip(left, right, strict=False)
+    )
+
+
+def _from_linear(channel: float) -> float:
+    if channel <= 0.0031308:
+        return 12.92 * channel
+    return 1.055 * math.pow(channel, 5 / 12) - 0.055
+
+
+def _l_to_y(lightness: float) -> float:
+    if lightness <= 8:
+        return _REF_Y * lightness / _KAPPA
+    return _REF_Y * (((lightness + 16) / 116) ** 3)
+
+
+def _lch_to_luv(color: tuple[float, float, float]) -> tuple[float, float, float]:
+    lightness, chroma, hue = color
+    hue_radians = math.radians(hue)
+    return (
+        lightness,
+        math.cos(hue_radians) * chroma,
+        math.sin(hue_radians) * chroma,
+    )
+
+
+def _luv_to_xyz(color: tuple[float, float, float]) -> tuple[float, float, float]:
+    lightness, u_value, v_value = color
+    if lightness == 0:
+        return (0.0, 0.0, 0.0)
+
+    var_u = u_value / (13 * lightness) + _REF_U
+    var_v = v_value / (13 * lightness) + _REF_V
+    y_value = _l_to_y(lightness)
+    x_value = y_value * 9 * var_u / (4 * var_v)
+    z_value = y_value * (12 - 3 * var_u - 20 * var_v) / (4 * var_v)
+    return (x_value, y_value, z_value)
+
+
+def _xyz_to_rgb(color: tuple[float, float, float]) -> tuple[float, float, float]:
+    return tuple(_from_linear(_dot_product(row, color)) for row in _M)  # type: ignore[return-value]
+
+
+def _hsluv_to_lch(color: tuple[float, float, float]) -> tuple[float, float, float]:
+    hue, saturation, lightness = color
+    if lightness > 100 - 1e-7:
+        return (100.0, 0.0, hue)
+    if lightness < 1e-8:
+        return (0.0, 0.0, hue)
+    max_chroma = _max_chroma_for_lh(lightness, hue)
+    return (lightness, max_chroma * saturation / 100, hue)
+
+
+def hsluv_to_rgb(color: tuple[float, float, float]) -> tuple[float, float, float]:
+    return _xyz_to_rgb(_luv_to_xyz(_lch_to_luv(_hsluv_to_lch(color))))
+
+
+def rgb_to_hex(color: tuple[float, float, float]) -> str:
+    channels = []
+    for channel in color:
+        channel = min(max(channel, 0.0), 1.0)
+        channels.append(int(math.floor(channel * 255 + 0.5)))
+    return "#{:02X}{:02X}{:02X}".format(*channels)
+
+
+def hsluv_to_hex(color: tuple[float, float, float]) -> str:
+    return rgb_to_hex(hsluv_to_rgb(color))
+
+
+def _hex_to_rgb(color: str) -> tuple[int, int, int]:
+    return tuple(int(color[index : index + 2], 16) for index in (1, 3, 5))  # type: ignore[return-value]
+
+
+def rgb_distance_hex(left: str, right: str) -> float:
+    left_rgb = _hex_to_rgb(left)
+    right_rgb = _hex_to_rgb(right)
+    return math.dist(left_rgb, right_rgb)
+
+
+@cache
+def rich_label_palette(slot_count: int) -> tuple[str, ...]:
+    if slot_count <= 64:
+        return PRECOMPUTED_QUALITATIVE_HSLUV_PALETTE_64
+    return tuple(qualitative_hsluv_palette(slot_count, lightness=0.5, saturation=0.9))
+
+
+def qualitative_hsluv_palette(
+    n_colors: int,
+    *,
+    hue: float = 0.01,
+    saturation: float = 0.9,
+    lightness: float = 0.5,
+    hue_samples: int | None = None,
+) -> list[str]:
+    """Build a deterministic, high-separation qualitative HSLuv palette.
+
+    Parameters mirror seaborn's ``husl_palette`` convention, using normalized
+    floats in ``[0, 1]`` for hue, saturation, and lightness.
+    """
+    if n_colors <= 0:
+        return []
+
+    if hue_samples is None:
+        hue_samples = max(48, n_colors * 32)
+
+    start_degrees = (hue % 1.0) * 360.0
+    candidates = [
+        hsluv_to_hex(
+            (
+                (start_degrees + 360.0 * index / hue_samples) % 360.0,
+                saturation * 100.0,
+                lightness * 100.0,
+            )
+        )
+        for index in range(hue_samples)
+    ]
+    candidates = list(dict.fromkeys(candidates))
+
+    selected = [candidates[0]]
+    while len(selected) < min(n_colors, len(candidates)):
+        remaining = [candidate for candidate in candidates if candidate not in selected]
+        best = max(
+            remaining,
+            key=lambda candidate: min(
+                rgb_distance_hex(candidate, chosen) for chosen in selected
+            ),
+        )
+        selected.append(best)
+
+    return selected[:n_colors]