oscura 0.5.0__py3-none-any.whl → 0.6.0__py3-none-any.whl

oscura/{math → utils/math}/arithmetic.py

@@ -6,7 +6,7 @@ and integration.
 
 
 Example:
-    >>> from oscura.math import add, differentiate
+    >>> from oscura.utils.math import add, differentiate
     >>> combined = add(trace1, trace2)
     >>> derivative = differentiate(trace)
 
@@ -716,48 +716,21 @@ class _SafeExpressionEvaluator(ast.NodeVisitor):
         raise AnalysisError(f"AST node type {node.__class__.__name__} not allowed")
 
 
-def math_expression(
-    expression: str,
-    traces: dict[str, WaveformTrace],
-    *,
-    channel_name: str | None = None,
-) -> WaveformTrace:
-    """Evaluate a mathematical expression on traces.
-
-    Evaluates an expression string using named traces as variables.
-    Supports standard mathematical operations and numpy functions.
+def _validate_trace_compatibility(
+    traces: dict[str, WaveformTrace], ref_trace: WaveformTrace
+) -> None:
+    """Validate all traces have same length and sample rate.
 
     Args:
-        expression: Math expression (e.g., "CH1 + CH2", "abs(CH1 - CH2)").
-        traces: Dictionary mapping variable names to traces.
-        channel_name: Name for the result trace (optional).
-
-    Returns:
-        Result WaveformTrace.
+        traces: Dictionary of traces to validate.
+        ref_trace: Reference trace for comparison.
 
     Raises:
-        AnalysisError: If expression is invalid or traces are incompatible.
-
-    Example:
-        >>> power = math_expression(
-        ...     "voltage * current",
-        ...     {"voltage": v_trace, "current": i_trace}
-        ... )
-
-    Security:
-        Uses AST-based safe evaluation (not eval()). Only whitelisted
-        operations are permitted: arithmetic, comparisons, and whitelisted
-        numpy functions. No arbitrary code execution is possible.
+        AnalysisError: If traces have incompatible dimensions.
     """
-    if not traces:
-        raise AnalysisError("No traces provided for expression evaluation")
-
-    # Get a reference trace for metadata
-    ref_trace = next(iter(traces.values()))
+    ref_len = len(ref_trace.data)
     sample_rate = ref_trace.metadata.sample_rate
 
-    # Validate all traces have same length and sample rate
-    ref_len = len(ref_trace.data)
     for name, trace in traces.items():
         if len(trace.data) != ref_len:
             raise AnalysisError(
@@ -771,8 +744,17 @@ def math_expression(
                 details={"expected": sample_rate, "got": trace.metadata.sample_rate},  # type: ignore[arg-type]
             )
 
-    # Create namespace with trace data and safe functions
-    safe_namespace = {
+
+def _build_safe_namespace(traces: dict[str, WaveformTrace]) -> dict[str, Any]:
+    """Build safe namespace with trace data and whitelisted functions.
+
+    Args:
+        traces: Dictionary of traces.
+
+    Returns:
+        Namespace dictionary with safe functions and trace data.
+    """
+    safe_namespace: dict[str, Any] = {
         "np": np,
         "abs": np.abs,
         "sqrt": np.sqrt,
@@ -789,28 +771,67 @@ def math_expression(
         "pi": np.pi,
     }
 
-    # Add trace data to namespace
     for name, trace in traces.items():
         safe_namespace[name] = trace.data.astype(np.float64)
 
-    # Use safe AST-based evaluator instead of eval()
-    evaluator = _SafeExpressionEvaluator(safe_namespace)
+    return safe_namespace
+
+
+def _evaluate_expression(expression: str, namespace: dict[str, Any]) -> Any:
+    """Evaluate expression using safe AST-based evaluator.
+
+    Args:
+        expression: Mathematical expression string.
+        namespace: Safe namespace with available functions and variables.
+
+    Returns:
+        Evaluated result.
+
+    Raises:
+        AnalysisError: If evaluation fails.
+    """
+    evaluator = _SafeExpressionEvaluator(namespace)
     try:
-        result = evaluator.eval(expression)
+        return evaluator.eval(expression)
     except AnalysisError:
-        raise  # Re-raise AnalysisError from evaluator
+        raise
     except Exception as e:
         raise AnalysisError(
             f"Failed to evaluate expression: {e}",
             details={"expression": expression},  # type: ignore[arg-type]
         ) from e
 
+
+def _ensure_array_result(result: Any, expected_len: int) -> NDArray[np.float64]:
+    """Ensure result is array of expected length.
+
+    Args:
+        result: Evaluation result.
+        expected_len: Expected array length.
+
+    Returns:
+        Result as float64 array.
+    """
     if not isinstance(result, np.ndarray):
-        # Scalar result - broadcast to array
-        result = np.full(ref_len, result, dtype=np.float64)
+        return np.full(expected_len, result, dtype=np.float64)
+    return result
 
-    new_metadata = TraceMetadata(
-        sample_rate=sample_rate,
+
+def _build_expression_metadata(
+    ref_trace: WaveformTrace, expression: str, channel_name: str | None
+) -> TraceMetadata:
+    """Build metadata for expression result trace.
+
+    Args:
+        ref_trace: Reference trace for metadata.
+        expression: Expression string (for default naming).
+        channel_name: Optional channel name override.
+
+    Returns:
+        Metadata for result trace.
+    """
+    return TraceMetadata(
+        sample_rate=ref_trace.metadata.sample_rate,
         vertical_scale=None,
         vertical_offset=None,
         acquisition_time=ref_trace.metadata.acquisition_time,
@@ -819,4 +840,49 @@ def math_expression(
         channel_name=channel_name or f"expr({expression[:20]})",
     )
 
-    return WaveformTrace(data=result.astype(np.float64), metadata=new_metadata)
+
+def math_expression(
+    expression: str,
+    traces: dict[str, WaveformTrace],
+    *,
+    channel_name: str | None = None,
+) -> WaveformTrace:
+    """Evaluate a mathematical expression on traces.
+
+    Evaluates an expression string using named traces as variables.
+    Supports standard mathematical operations and numpy functions.
+
+    Args:
+        expression: Math expression (e.g., "CH1 + CH2", "abs(CH1 - CH2)").
+        traces: Dictionary mapping variable names to traces.
+        channel_name: Name for the result trace (optional).
+
+    Returns:
+        Result WaveformTrace.
+
+    Raises:
+        AnalysisError: If expression is invalid or traces are incompatible.
+
+    Example:
+        >>> power = math_expression(
+        ...     "voltage * current",
+        ...     {"voltage": v_trace, "current": i_trace}
+        ... )
+
+    Security:
+        Uses AST-based safe evaluation (not eval()). Only whitelisted
+        operations are permitted: arithmetic, comparisons, and whitelisted
+        numpy functions. No arbitrary code execution is possible.
+    """
+    if not traces:
+        raise AnalysisError("No traces provided for expression evaluation")
+
+    ref_trace = next(iter(traces.values()))
+    _validate_trace_compatibility(traces, ref_trace)
+
+    safe_namespace = _build_safe_namespace(traces)
+    result = _evaluate_expression(expression, safe_namespace)
+    result = _ensure_array_result(result, len(ref_trace.data))
+
+    metadata = _build_expression_metadata(ref_trace, expression, channel_name)
+    return WaveformTrace(data=result.astype(np.float64), metadata=metadata)

oscura/{math → utils/math}/interpolation.py

@@ -5,7 +5,7 @@ functions for waveform data.
 
 
 Example:
-    >>> from oscura.math import resample, align_traces
+    >>> from oscura.utils.math import resample, align_traces
     >>> resampled = resample(trace, new_sample_rate=1e6)
     >>> aligned = align_traces(trace1, trace2)
 
@@ -16,7 +16,7 @@ References:
 from __future__ import annotations
 
 import warnings
-from typing import TYPE_CHECKING, Literal
+from typing import TYPE_CHECKING, Any, Literal
 
 import numpy as np
 from scipy import interpolate as sp_interp
@@ -45,13 +45,8 @@ def interpolate(
     Args:
         trace: Input trace.
         new_time: New time points in seconds.
-        method: Interpolation method:
-            - "linear": Linear interpolation (default)
-            - "cubic": Cubic spline interpolation
-            - "nearest": Nearest neighbor
-            - "zero": Zero-order hold (step function)
+        method: Interpolation method ("linear", "cubic", "nearest", "zero").
         fill_value: Value for points outside original range.
-            Can be a single value or (below, above) tuple.
         channel_name: Name for the result trace (optional).
 
     Returns:
@@ -73,55 +68,49 @@ def interpolate(
             analysis_type="interpolate",
         )
 
-    original_time = trace.time_vector
-    data = trace.data.astype(np.float64)
+    # Create interpolator and interpolate
+    interp_func = _create_interpolator(
+        trace.time_vector, trace.data.astype(np.float64), method, fill_value
+    )
+    result_data = interp_func(new_time)
 
-    # Create interpolator
-    if method == "linear":
-        interp_func = sp_interp.interp1d(
-            original_time,
-            data,
-            kind="linear",
-            bounds_error=False,
-            fill_value=fill_value,
-        )
-    elif method == "cubic":
-        interp_func = sp_interp.interp1d(
-            original_time,
-            data,
-            kind="cubic",
-            bounds_error=False,
-            fill_value=fill_value,
-        )
-    elif method == "nearest":
-        interp_func = sp_interp.interp1d(
-            original_time,
-            data,
-            kind="nearest",
-            bounds_error=False,
-            fill_value=fill_value,
-        )
-    elif method == "zero":
-        interp_func = sp_interp.interp1d(
-            original_time,
-            data,
-            kind="zero",
-            bounds_error=False,
-            fill_value=fill_value,
-        )
-    else:
+    # Build result trace
+    new_sample_rate = _calculate_new_sample_rate(new_time, trace.metadata.sample_rate)
+    new_metadata = _create_interpolated_metadata(trace, new_sample_rate, channel_name)
+
+    return WaveformTrace(data=result_data.astype(np.float64), metadata=new_metadata)
+
+
+def _create_interpolator(
+    original_time: NDArray[np.float64],
+    data: NDArray[np.float64],
+    method: str,
+    fill_value: float | tuple[float, float],
+) -> Any:
+    """Create scipy interpolation function."""
+    valid_methods = {"linear", "cubic", "nearest", "zero"}
+    if method not in valid_methods:
         raise ValueError(f"Unknown interpolation method: {method}")
 
-    # Interpolate
-    result_data = interp_func(new_time)
+    return sp_interp.interp1d(
+        original_time, data, kind=method, bounds_error=False, fill_value=fill_value
+    )
+
 
-    # Calculate new sample rate from time points
+def _calculate_new_sample_rate(new_time: NDArray[np.float64], original_sample_rate: float) -> float:
+    """Calculate new sample rate from time points."""
     if len(new_time) > 1:
-        new_sample_rate = 1.0 / np.mean(np.diff(new_time))
-    else:
-        new_sample_rate = trace.metadata.sample_rate
+        mean_diff: np.floating[Any] = np.mean(np.diff(new_time))
+        new_rate: float = float(1.0 / mean_diff)
+        return new_rate
+    return original_sample_rate
 
-    new_metadata = TraceMetadata(
+
+def _create_interpolated_metadata(
+    trace: WaveformTrace, new_sample_rate: float, channel_name: str | None
+) -> TraceMetadata:
+    """Create metadata for interpolated trace."""
+    return TraceMetadata(
         sample_rate=new_sample_rate,
         vertical_scale=trace.metadata.vertical_scale,
         vertical_offset=trace.metadata.vertical_offset,
@@ -131,7 +120,91 @@ def interpolate(
         channel_name=channel_name or f"{trace.metadata.channel_name or 'trace'}_interp",
     )
 
-    return WaveformTrace(data=result_data.astype(np.float64), metadata=new_metadata)
+
+def _calculate_target_params(
+    new_sample_rate: float | None,
+    num_samples: int | None,
+    original_rate: float,
+    original_samples: int,
+) -> tuple[float, int]:
+    """Calculate target sample rate and sample count for resampling."""
+    if new_sample_rate is not None:
+        target_rate = new_sample_rate
+        target_samples = round(original_samples * target_rate / original_rate)
+    else:
+        target_samples = num_samples  # type: ignore[assignment]
+        target_rate = original_rate * target_samples / original_samples
+    return target_rate, target_samples
+
+
+def _check_nyquist_violation(
+    data: NDArray[np.float64], original_rate: float, target_rate: float
+) -> None:
+    """Validate Nyquist criterion when downsampling and warn if violated."""
+    fft_data = np.fft.rfft(data)
+    fft_freqs = np.fft.rfftfreq(len(data), 1 / original_rate)
+    power = np.abs(fft_data) ** 2
+    power_threshold = 0.01 * np.max(power)
+    significant_freqs = fft_freqs[power > power_threshold]
+
+    if len(significant_freqs) > 0:
+        max_frequency = np.max(significant_freqs)
+        nyquist_required = 2 * max_frequency
+        if target_rate < nyquist_required:
+            warnings.warn(
+                f"Downsampling to {target_rate:.2e} Hz violates Nyquist criterion. "
+                f"Maximum signal frequency is ~{max_frequency:.2e} Hz, "
+                f"requiring ≥{nyquist_required:.2e} Hz sample rate. "
+                f"Aliasing may occur.",
+                UserWarning,
+                stacklevel=3,
+            )
+
+
+def _apply_anti_alias_filter(
+    data: NDArray[np.float64], target_rate: float, original_rate: float
+) -> NDArray[np.float64]:
+    """Apply lowpass anti-aliasing filter before downsampling."""
+    nyquist = target_rate / 2
+    cutoff = nyquist / original_rate * 2  # Normalized frequency
+    if cutoff < 1.0:
+        b, a = sp_signal.butter(8, min(cutoff * 0.9, 0.99), btype="low")
+        filtered: NDArray[np.float64] = np.asarray(sp_signal.filtfilt(b, a, data), dtype=np.float64)
+        return filtered
+    return data
+
+
+def _perform_resampling(
+    data: NDArray[np.float64],
+    method: Literal["fft", "polyphase", "interp"],
+    target_samples: int,
+    original_samples: int,
+    original_rate: float,
+    target_rate: float,
+) -> NDArray[np.float64]:
+    """Perform the actual resampling based on selected method."""
+    if method == "fft":
+        resampled: NDArray[np.float64] = np.asarray(
+            sp_signal.resample(data, target_samples), dtype=np.float64
+        )
+        return resampled
+    elif method == "polyphase":
+        from fractions import Fraction
+
+        ratio = Fraction(target_samples, original_samples).limit_denominator(1000)
+        up, down = ratio.numerator, ratio.denominator
+        result = sp_signal.resample_poly(data, up, down)
+        truncated: NDArray[np.float64] = np.asarray(result[:target_samples], dtype=np.float64)
+        return truncated
+    elif method == "interp":
+        old_time = np.arange(original_samples) / original_rate
+        new_time = np.arange(target_samples) / target_rate
+        interpolated: NDArray[np.float64] = np.asarray(
+            np.interp(new_time, old_time, data), dtype=np.float64
+        )
+        return interpolated
+    else:
+        raise ValueError(f"Unknown resampling method: {method}")
 
 
 def resample(
@@ -175,9 +248,9 @@ def resample(
     References:
         MEM-012 (downsampling for memory management)
     """
+    # Validate inputs
     if (new_sample_rate is None) == (num_samples is None):
         raise ValueError("Specify exactly one of new_sample_rate or num_samples")
-
     if len(trace.data) < 2:
         raise InsufficientDataError(
             "Need at least 2 samples for resampling",
@@ -186,73 +259,33 @@ def resample(
             analysis_type="resample",
         )
 
+    # Setup
     data = trace.data.astype(np.float64)
     original_rate = trace.metadata.sample_rate
     original_samples = len(data)
 
     # Calculate target parameters
-    if new_sample_rate is not None:
-        target_rate = new_sample_rate
-        target_samples = round(original_samples * target_rate / original_rate)
-    else:
-        target_samples = num_samples  # type: ignore[assignment]
-        target_rate = original_rate * target_samples / original_samples
+    target_rate, target_samples = _calculate_target_params(
+        new_sample_rate, num_samples, original_rate, original_samples
+    )
 
     if target_samples < 1:
         raise ValueError("Target number of samples must be at least 1")
 
     # REQ: API-019 - Validate Nyquist criterion when downsampling
     if target_rate < original_rate:
-        # Estimate maximum frequency using FFT
-        fft_data = np.fft.rfft(data)
-        fft_freqs = np.fft.rfftfreq(len(data), 1 / original_rate)
-        # Find frequency with 90% of max power as max frequency
-        power = np.abs(fft_data) ** 2
-        power_threshold = 0.01 * np.max(power)  # 1% of max power
-        significant_freqs = fft_freqs[power > power_threshold]
-        if len(significant_freqs) > 0:
-            max_frequency = np.max(significant_freqs)
-            nyquist_required = 2 * max_frequency
-            if target_rate < nyquist_required:
-                warnings.warn(
-                    f"Downsampling to {target_rate:.2e} Hz violates Nyquist criterion. "
-                    f"Maximum signal frequency is ~{max_frequency:.2e} Hz, "
-                    f"requiring ≥{nyquist_required:.2e} Hz sample rate. "
-                    f"Aliasing may occur.",
-                    UserWarning,
-                    stacklevel=2,
-                )
-
-    # Check if downsampling and apply anti-alias filter
+        _check_nyquist_violation(data, original_rate, target_rate)
+
+    # Apply anti-aliasing filter if downsampling
     if anti_alias and target_samples < original_samples:
-        # Lowpass filter at Nyquist of new rate
-        nyquist = target_rate / 2
-        cutoff = nyquist / original_rate * 2  # Normalized frequency
-        if cutoff < 1.0:
-            # Design lowpass filter
-            b, a = sp_signal.butter(8, min(cutoff * 0.9, 0.99), btype="low")
-            data = sp_signal.filtfilt(b, a, data)
-
-    # Resample
-    if method == "fft":
-        result_data = sp_signal.resample(data, target_samples)
-    elif method == "polyphase":
-        # Find rational approximation for polyphase resampling
-        from fractions import Fraction
+        data = _apply_anti_alias_filter(data, target_rate, original_rate)
 
-        ratio = Fraction(target_samples, original_samples).limit_denominator(1000)
-        up, down = ratio.numerator, ratio.denominator
-        result_data = sp_signal.resample_poly(data, up, down)
-        # Trim to exact length
-        result_data = result_data[:target_samples]
-    elif method == "interp":
-        # Simple interpolation
-        old_time = np.arange(original_samples) / original_rate
-        new_time = np.arange(target_samples) / target_rate
-        result_data = np.interp(new_time, old_time, data)
-    else:
-        raise ValueError(f"Unknown resampling method: {method}")
+    # Perform resampling
+    result_data = _perform_resampling(
+        data, method, target_samples, original_samples, original_rate, target_rate
+    )
 
+    # Build output trace
     new_metadata = TraceMetadata(
         sample_rate=target_rate,
         vertical_scale=trace.metadata.vertical_scale,