flopscope 0.2.0__py3-none-any.whl

flopscope/_unwrap.py

@@ -0,0 +1,70 @@
+# src/flopscope/_unwrap.py
+"""Unwrap wrapper with FLOP counting."""
+
+from __future__ import annotations
+
+import numpy as _np
+from numpy.typing import ArrayLike
+
+from flopscope._budget import _call_numpy, _counted_wrapper
+from flopscope._docstrings import attach_docstring
+from flopscope._ndarray import FlopscopeArray, _to_base_ndarray
+from flopscope._validation import require_budget
+
+
+def unwrap_cost(shape: tuple[int, ...]) -> int:
+    """FLOP cost of phase unwrapping.
+
+    Parameters
+    ----------
+    shape : tuple of int
+        Input array shape.
+
+    Returns
+    -------
+    int
+        Estimated FLOP count: ``7 * numel(input)``.
+
+    Notes
+    -----
+    NumPy's unwrap (``numpy.lib._function_base_impl.unwrap``) does
+    approximately 6 full-array ufunc passes: ``diff``, ``mod``, ``add``,
+    ``subtract``, ``cumsum``, and ``where``. The ``where`` pass with two
+    real-array arguments charges 2*(N-1) elements rather than N-1,
+    bringing the effective total to ~7*(N-1). We charge a fixed
+    ``7 * numel`` to approximate the sum without tracking numpy's
+    internal implementation byte-by-byte. Issue #69.
+    """
+    numel = 1
+    for d in shape:
+        numel *= d
+    return max(7 * numel, 1)
+
+
+@_counted_wrapper
+def unwrap(
+    p: ArrayLike,
+    discont: float | None = None,
+    axis: int = -1,
+    *,
+    period: float = 6.283185307179586,
+) -> FlopscopeArray:
+    budget = require_budget()
+    if not isinstance(p, _np.ndarray):
+        p = _np.asarray(p)
+    cost = unwrap_cost(p.shape)
+    kwargs = {"axis": axis, "period": period}
+    if discont is not None:
+        kwargs["discont"] = discont
+    with budget.deduct("unwrap", flop_cost=cost, subscripts=None, shapes=(p.shape,)):
+        result = _call_numpy(_np.unwrap, _to_base_ndarray(p), **kwargs)
+    return result  # type: ignore[return-value]
+
+
+attach_docstring(unwrap, _np.unwrap, "counted_custom", "7 * numel(input) FLOPs")
+
+import sys as _sys  # noqa: E402
+
+from flopscope._ndarray import wrap_module_returns as _wrap_module_returns  # noqa: E402
+
+_wrap_module_returns(_sys.modules[__name__])

flopscope/_validation.py

@@ -0,0 +1,83 @@
+"""Input validation and NaN/Inf checking for flopscope operations."""
+
+from __future__ import annotations
+
+import warnings
+
+import numpy as np
+
+from flopscope._budget import get_active_budget
+from flopscope.errors import FlopscopeWarning
+
+
+def require_budget():
+    """Return the active budget, auto-activating the global default if needed."""
+    budget = get_active_budget()
+    if budget is not None:
+        return budget
+    from flopscope._budget import _get_global_default
+
+    return _get_global_default()
+
+
+def validate_ndarray(*arrays: object) -> None:
+    """Validate that all arguments are numpy ndarrays."""
+    for arr in arrays:
+        if not isinstance(arr, np.ndarray):
+            raise TypeError(f"Expected numpy.ndarray, got {type(arr).__name__}")
+
+
+def coerce_array(x):
+    """Convert input to ndarray if not already one, matching NumPy's behavior."""
+    if isinstance(x, np.ndarray):
+        return x
+    return np.asarray(x)
+
+
+def coerce_arrays(*arrays):
+    """Convert multiple inputs to ndarrays."""
+    return tuple(coerce_array(a) for a in arrays)
+
+
+def check_nan_inf(result: np.ndarray, op_name: str) -> None:
+    """Issue a warning if result contains NaN or Inf values.
+
+    Skips dtypes that don't support `np.isnan`/`np.isinf` (e.g. object,
+    integer, complex with object content) — these can never contain NaN
+    or Inf as ndarray values, so the check is a no-op.
+    """
+    if not isinstance(result, np.ndarray):
+        return
+    # np.isnan/np.isinf only support float and complex dtypes. For other
+    # dtypes (object, integer, bool, structured), there are no NaN/Inf
+    # values to detect, so skip the check.
+    if result.dtype.kind not in ("f", "c"):
+        return
+    # Strip flopscope subclasses so np.isnan / np.isinf do not re-dispatch
+    # through __array_ufunc__ and recurse into me.isnan / me.isfinite,
+    # which in turn would call check_nan_inf again.
+    plain = result.view(np.ndarray) if type(result) is not np.ndarray else result
+    nan_count = int(np.isnan(plain).sum())
+    inf_count = int(np.isinf(plain).sum())
+    if nan_count > 0 or inf_count > 0:
+        warnings.warn(
+            f"{op_name} produced {nan_count} NaN and {inf_count} Inf values "
+            f"in output of shape {result.shape}",
+            FlopscopeWarning,
+            stacklevel=3,
+        )
+
+
+def maybe_check_nan_inf(result: object, op_name: str) -> None:
+    """Run :func:`check_nan_inf` only if the global ``check_nan_inf`` setting is on.
+
+    Production scoring runs with the setting off (default) and pays no
+    per-op O(n) scan cost.  Debug callers opt in via
+    ``flopscope.configure(check_nan_inf=True)``.
+    """
+    from flopscope._config import get_setting
+
+    if not get_setting("check_nan_inf"):
+        return
+    if isinstance(result, np.ndarray):
+        check_nan_inf(result, op_name)

flopscope/_version_check.py

@@ -0,0 +1,46 @@
+"""NumPy version checking for flopscope."""
+
+from __future__ import annotations
+
+import warnings
+
+import numpy as _np
+
+from flopscope.errors import FlopscopeWarning
+
+
+def check_numpy_version(supported_range: str) -> None:
+    """Warn if the installed numpy version is outside the supported range.
+
+    Parameters
+    ----------
+    supported_range : str
+        A PEP 440 version specifier like ``">=2.0.0,<2.3.0"``.
+    """
+    installed = _np.__version__
+    installed_parts = tuple(int(x) for x in installed.split(".")[:2])
+
+    # Parse the supported range: ">=2.0.0,<2.3.0"
+    min_version = None
+    max_version = None
+    for spec in supported_range.split(","):
+        spec = spec.strip()
+        if spec.startswith(">="):
+            min_version = tuple(int(x) for x in spec[2:].split(".")[:2])
+        elif spec.startswith("<"):
+            max_version = tuple(int(x) for x in spec[1:].split(".")[:2])
+
+    in_range = True
+    if min_version and installed_parts < min_version:
+        in_range = False
+    if max_version and installed_parts >= max_version:
+        in_range = False
+
+    if not in_range:
+        warnings.warn(
+            f"flopscope supports numpy {supported_range} but numpy "
+            f"{installed} is installed. Some functions may be missing or "
+            f"behave differently.",
+            FlopscopeWarning,
+            stacklevel=3,
+        )

flopscope/_weights.py

@@ -0,0 +1,195 @@
+"""Configurable per-operation FLOP weights for flopscope.
+
+Flopscope supports official packaged weights, custom overrides via
+``FLOPSCOPE_WEIGHTS_FILE``, and complete disabling via ``FLOPSCOPE_DISABLE_WEIGHTS=1``.
+Packaged official weights are loaded automatically on import unless explicitly
+overridden or disabled.
+
+The JSON payload must have a ``"weights"`` key mapping
+``op_name -> float multiplier``.
+"""
+
+from __future__ import annotations
+
+import json
+import math
+import os
+import warnings
+from importlib import resources
+
+_ACTIVE_WEIGHTS: dict[str, float] = {}
+_WARNED_MESSAGES: set[str] = set()
+
+
+def _warn_once(message: str) -> None:
+    if message in _WARNED_MESSAGES:
+        return
+    warnings.warn(message, RuntimeWarning, stacklevel=2)
+    _WARNED_MESSAGES.add(message)
+
+
+def _weights_disabled() -> bool:
+    value = os.environ.get("FLOPSCOPE_DISABLE_WEIGHTS")
+    if value is None:
+        return False
+    return value.strip().lower() in {"1", "true", "yes", "on"}
+
+
+def _extract_weights(data: dict, *, source: str) -> dict[str, float]:
+    weights = data.get("weights")
+    if not isinstance(weights, dict):
+        raise ValueError(f"{source} is missing a valid 'weights' mapping")
+    validated: dict[str, float] = {}
+    for op_name, value in weights.items():
+        if not isinstance(op_name, str):
+            raise ValueError(f"{source} has a non-string operation name: {op_name!r}")
+        if isinstance(value, bool) or not isinstance(value, (int, float)):
+            raise ValueError(
+                f"{source} weight for {op_name!r} must be a non-negative finite number"
+            )
+        numeric_value = float(value)
+        if not math.isfinite(numeric_value) or numeric_value < 0:
+            raise ValueError(
+                f"{source} weight for {op_name!r} must be a non-negative finite number"
+            )
+        validated[op_name] = numeric_value
+    return validated
+
+
+def _read_weights_file(path: str, *, source: str) -> dict[str, float]:
+    with open(path, encoding="utf-8") as f:
+        data = json.load(f)
+    return _extract_weights(data, source=source)
+
+
+def _load_packaged_weights() -> dict[str, float] | None:
+    try:
+        resource = resources.files("flopscope").joinpath("data/default_weights.json")
+        with resource.open("r", encoding="utf-8") as f:
+            data = json.load(f)
+        return _extract_weights(data, source="Packaged official weights")
+    except Exception as exc:  # pragma: no cover - defensive fallback path
+        _warn_once(
+            "Flopscope could not load packaged official weights "
+            f"({exc}); falling back to unit weights (1.0 for all operations)."
+        )
+        return None
+
+
+# Modern-Generator → legacy-module-level rename map (issue #18 follow-up).
+# Methods whose names changed between the legacy API and the modern Generator API.
+# Everything else uses simple prefix-stripping (no entry needed).
+_RANDOM_METHOD_RENAMES: dict[str, str] = {
+    "integers": "randint",  # Generator.integers → random.randint
+    "random": "random_sample",  # Generator.random   → random.random_sample
+}
+
+
+def get_weight(op_name: str) -> float:
+    """Return the FLOP weight multiplier for an operation.
+
+    Parameters
+    ----------
+    op_name : str
+        Operation name as passed to ``BudgetContext.deduct()``,
+        e.g. ``"exp"``, ``"linalg.cholesky"``, ``"fft.fft"``.
+
+    Returns
+    -------
+    float
+        Multiplicative weight. Defaults to 1.0 if not configured.
+
+    Notes
+    -----
+    For method-level random ops (``random.Generator.<m>`` and
+    ``random.RandomState.<m>``), if no explicit weight is set, falls
+    through to the legacy module-level weight via name-stripping plus
+    ``_RANDOM_METHOD_RENAMES`` for the few renamed methods. This keeps
+    same-physics ops at the same FLOP cost without requiring 94 explicit
+    JSON entries (issue #18 follow-up).
+    """
+    # 1. Explicit entry — primary path.
+    if op_name in _ACTIVE_WEIGHTS:
+        return _ACTIVE_WEIGHTS[op_name]
+
+    # 2. Method-level alias fallback for random.Generator.* / random.RandomState.*
+    for prefix in ("random.Generator.", "random.RandomState."):
+        if op_name.startswith(prefix):
+            short = op_name[len(prefix) :]
+            legacy = _RANDOM_METHOD_RENAMES.get(short, short)
+            legacy_op = f"random.{legacy}"
+            if legacy_op in _ACTIVE_WEIGHTS:
+                return _ACTIVE_WEIGHTS[legacy_op]
+            break
+
+    # 3. Default fallback (unchanged).
+    return 1.0
+
+
+def load_weights(path: str | None = None, *, use_packaged_default: bool = True) -> None:
+    """Resolve and load active FLOP weights.
+
+    Parameters
+    ----------
+    path : str or None
+        Explicit path to a weights JSON file. If None, reads from the
+        ``FLOPSCOPE_WEIGHTS_FILE`` environment variable when present.
+    use_packaged_default : bool, optional
+        Whether to fall back to the packaged official weights when no valid
+        override is available.
+
+    Notes
+    -----
+    Resolution order:
+
+    1. If ``FLOPSCOPE_DISABLE_WEIGHTS=1``, all weights are disabled.
+    2. If an explicit path or ``FLOPSCOPE_WEIGHTS_FILE`` is provided and valid,
+       it is used.
+    3. If the override is unusable, a warning is emitted and Flopscope falls back
+       to the packaged official weights when enabled, otherwise to unit
+       weights.
+    4. If packaged weights are enabled but unavailable, a warning is emitted
+       and Flopscope falls back to unit weights.
+    """
+    _ACTIVE_WEIGHTS.clear()
+
+    if _weights_disabled():
+        return
+
+    override_path = (
+        path if path is not None else os.environ.get("FLOPSCOPE_WEIGHTS_FILE")
+    )
+
+    if override_path is not None:
+        try:
+            _ACTIVE_WEIGHTS.update(
+                _read_weights_file(
+                    override_path,
+                    source=f"Custom weights file '{override_path}'",
+                )
+            )
+            return
+        except Exception as exc:
+            fallback_target = (
+                "packaged official weights" if use_packaged_default else "unit weights"
+            )
+            _warn_once(
+                f"Flopscope could not load custom weights from '{override_path}' "
+                f"({exc}); falling back to {fallback_target}."
+            )
+
+    if not use_packaged_default:
+        return
+
+    packaged_weights = _load_packaged_weights()
+    if packaged_weights is not None:
+        _ACTIVE_WEIGHTS.update(packaged_weights)
+
+
+def reset_weights() -> None:
+    """Clear all loaded weights. For testing."""
+    _ACTIVE_WEIGHTS.clear()
+    _WARNED_MESSAGES.clear()
+
+
+load_weights()

flopscope/_window.py

@@ -0,0 +1,177 @@
+# src/flopscope/_window.py
+"""Window function wrappers with FLOP counting."""
+
+from __future__ import annotations
+
+import numpy as _np
+
+from flopscope._budget import _call_numpy, _counted_wrapper
+from flopscope._docstrings import attach_docstring
+from flopscope._ndarray import FlopscopeArray
+from flopscope._validation import require_budget
+
+
+def bartlett_cost(n: int) -> int:
+    """FLOP cost of Bartlett window generation.
+
+    Parameters
+    ----------
+    n : int
+        Window length.
+
+    Returns
+    -------
+    int
+        Estimated FLOP count: n.
+
+    Notes
+    -----
+    One linear evaluation per sample.
+    """
+    return max(n, 1)
+
+
+@_counted_wrapper
+def bartlett(M: int) -> FlopscopeArray:
+    budget = require_budget()
+    cost = bartlett_cost(M)
+    with budget.deduct("bartlett", flop_cost=cost, subscripts=None, shapes=((M,),)):
+        result = _call_numpy(_np.bartlett, M)
+    return result  # type: ignore[return-value]
+
+
+attach_docstring(bartlett, _np.bartlett, "counted_custom", "n FLOPs")
+
+
+def blackman_cost(n: int) -> int:
+    """FLOP cost of Blackman window generation.
+
+    Parameters
+    ----------
+    n : int
+        Window length.
+
+    Returns
+    -------
+    int
+        Estimated FLOP count: 3n.
+
+    Notes
+    -----
+    Three cosine terms per sample.
+    """
+    return max(3 * n, 1)
+
+
+@_counted_wrapper
+def blackman(M: int) -> FlopscopeArray:
+    budget = require_budget()
+    cost = blackman_cost(M)
+    with budget.deduct("blackman", flop_cost=cost, subscripts=None, shapes=((M,),)):
+        result = _call_numpy(_np.blackman, M)
+    return result  # type: ignore[return-value]
+
+
+attach_docstring(blackman, _np.blackman, "counted_custom", "3n FLOPs")
+
+
+def hamming_cost(n: int) -> int:
+    """FLOP cost of Hamming window generation.
+
+    Parameters
+    ----------
+    n : int
+        Window length.
+
+    Returns
+    -------
+    int
+        Estimated FLOP count: 2n (FMA=2 textbook: 1 multiply + 1 add per sample).
+
+    Notes
+    -----
+    Two ops per sample under FMA=2 convention (1 multiply + 1 add).
+    """
+    return max(2 * n, 1)
+
+
+@_counted_wrapper
+def hamming(M: int) -> FlopscopeArray:
+    budget = require_budget()
+    cost = hamming_cost(M)
+    with budget.deduct("hamming", flop_cost=cost, subscripts=None, shapes=((M,),)):
+        result = _call_numpy(_np.hamming, M)
+    return result  # type: ignore[return-value]
+
+
+attach_docstring(hamming, _np.hamming, "counted_custom", "2n FLOPs (FMA=2)")
+
+
+def hanning_cost(n: int) -> int:
+    """FLOP cost of Hanning window generation.
+
+    Parameters
+    ----------
+    n : int
+        Window length.
+
+    Returns
+    -------
+    int
+        Estimated FLOP count: 2n (FMA=2 textbook: 1 multiply + 1 add per sample).
+
+    Notes
+    -----
+    Two ops per sample under FMA=2 convention (1 multiply + 1 add).
+    """
+    return max(2 * n, 1)
+
+
+@_counted_wrapper
+def hanning(M: int) -> FlopscopeArray:
+    budget = require_budget()
+    cost = hanning_cost(M)
+    with budget.deduct("hanning", flop_cost=cost, subscripts=None, shapes=((M,),)):
+        result = _call_numpy(_np.hanning, M)
+    return result  # type: ignore[return-value]
+
+
+attach_docstring(hanning, _np.hanning, "counted_custom", "2n FLOPs (FMA=2)")
+
+
+def kaiser_cost(n: int) -> int:
+    """FLOP cost of Kaiser window generation.
+
+    Parameters
+    ----------
+    n : int
+        Window length.
+
+    Returns
+    -------
+    int
+        Estimated FLOP count: 3n.
+
+    Notes
+    -----
+    Bessel function evaluation per sample.
+    """
+    return max(3 * n, 1)
+
+
+@_counted_wrapper
+def kaiser(M: int, beta: float) -> FlopscopeArray:
+    budget = require_budget()
+    cost = kaiser_cost(M)
+    with budget.deduct("kaiser", flop_cost=cost, subscripts=None, shapes=((M,),)):
+        result = _call_numpy(_np.kaiser, M, beta)
+    return result  # type: ignore[return-value]
+
+
+attach_docstring(kaiser, _np.kaiser, "counted_custom", "3n FLOPs")
+
+import sys as _sys  # noqa: E402
+
+from flopscope._ndarray import wrap_module_returns as _wrap_module_returns  # noqa: E402
+
+_wrap_module_returns(_sys.modules[__name__])