oscura 0.0.1__py3-none-any.whl → 0.1.1__py3-none-any.whl

oscura/loaders/sigrok.py

@@ -0,0 +1,321 @@
+"""Sigrok session file (.sr) loader.
+
+This module provides loading of sigrok session files containing
+logic analyzer captures. Sigrok sessions are ZIP archives containing
+metadata and binary signal data.
+
+
+Example:
+    >>> from oscura.loaders.sigrok import load_sigrok
+    >>> trace = load_sigrok("capture.sr")
+    >>> print(f"Sample rate: {trace.metadata.sample_rate} Hz")
+    >>> print(f"Channels: {len(trace.data)}")
+"""
+
+from __future__ import annotations
+
+import zipfile
+from pathlib import Path
+from typing import TYPE_CHECKING, Any
+
+import numpy as np
+from numpy.typing import NDArray
+
+from oscura.core.exceptions import FormatError, LoaderError
+from oscura.core.types import DigitalTrace, TraceMetadata
+
+if TYPE_CHECKING:
+    from os import PathLike
+
+
+def load_sigrok(
+    path: str | PathLike[str],
+    *,
+    channel: str | int | None = None,
+) -> DigitalTrace:
+    """Load a sigrok session file (.sr).
+
+    Sigrok session files are ZIP archives containing:
+    - metadata: JSON file with capture settings
+    - logic-1-*: Binary files with sample data
+
+    Args:
+        path: Path to the sigrok .sr session file.
+        channel: Optional channel name or index to load. If None,
+            loads the first channel or merges all channels.
+
+    Returns:
+        DigitalTrace containing the digital signal data and metadata.
+
+    Raises:
+        LoaderError: If the file cannot be loaded.
+        FormatError: If the file is not a valid sigrok session.
+
+    Example:
+        >>> trace = load_sigrok("capture.sr")
+        >>> print(f"Sample rate: {trace.metadata.sample_rate} Hz")
+        >>> print(f"Duration: {trace.duration:.6f} seconds")
+
+    References:
+        sigrok session file format specification
+    """
+    path = Path(path)
+
+    if not path.exists():
+        raise LoaderError(
+            "File not found",
+            file_path=str(path),
+        )
+
+    if not zipfile.is_zipfile(path):
+        raise FormatError(
+            "File is not a valid sigrok session (not a ZIP archive)",
+            file_path=str(path),
+            expected="ZIP archive",
+        )
+
+    try:
+        with zipfile.ZipFile(path, "r") as zf:
+            # Parse metadata
+            metadata_dict = _parse_metadata(zf, path)
+
+            # Get sample rate from metadata
+            sample_rate = metadata_dict.get("samplerate", 1_000_000)
+
+            # Get channel information
+            channels = metadata_dict.get("channels", [])
+            total_channels = metadata_dict.get("total probes", len(channels))
+
+            # Find and read logic data files
+            logic_files = [name for name in zf.namelist() if name.startswith("logic-1")]
+
+            if not logic_files:
+                raise FormatError(
+                    "No logic data found in sigrok session",
+                    file_path=str(path),
+                    expected="logic-1-* data files",
+                )
+
+            # Read and combine logic data
+            data = _read_logic_data(zf, logic_files, total_channels)
+
+            # Select specific channel if requested
+            if channel is not None:
+                if isinstance(channel, int):
+                    if channel < 0 or channel >= data.shape[0]:
+                        raise LoaderError(
+                            f"Channel index {channel} out of range",
+                            file_path=str(path),
+                            details=f"Available channels: 0-{data.shape[0] - 1}",
+                        )
+                    channel_data = data[channel]
+                    channel_name = channels[channel] if channel < len(channels) else f"D{channel}"
+                elif isinstance(channel, str):
+                    if channel in channels:
+                        idx = channels.index(channel)
+                        channel_data = data[idx]
+                        channel_name = channel
+                    else:
+                        raise LoaderError(
+                            f"Channel '{channel}' not found",
+                            file_path=str(path),
+                            details=f"Available channels: {channels}",
+                        )
+                else:
+                    channel_data = data[0]  # type: ignore[unreachable]
+                    channel_name = channels[0] if channels else "D0"
+            else:
+                # Default to first channel
+                channel_data = data[0] if data.ndim > 1 else data
+                channel_name = channels[0] if channels else "D0"
+
+            # Compute edges
+            edges = _compute_edges(channel_data, sample_rate)
+
+            # Build metadata
+            trace_metadata = TraceMetadata(
+                sample_rate=float(sample_rate),
+                source_file=str(path),
+                channel_name=channel_name,
+                trigger_info=metadata_dict.get("trigger", None),
+            )
+
+            return DigitalTrace(
+                data=channel_data,
+                metadata=trace_metadata,
+                edges=edges,
+            )
+
+    except zipfile.BadZipFile as e:
+        raise FormatError(
+            "Corrupted sigrok session file",
+            file_path=str(path),
+            expected="Valid ZIP archive",
+        ) from e
+    except Exception as e:
+        if isinstance(e, LoaderError | FormatError):
+            raise
+        raise LoaderError(
+            "Failed to load sigrok session",
+            file_path=str(path),
+            details=str(e),
+            fix_hint="Ensure the file is a valid sigrok session (.sr) file.",
+        ) from e
+
+
+def _parse_metadata(zf: zipfile.ZipFile, path: Path) -> dict[str, Any]:
+    """Parse sigrok session metadata.
+
+    Args:
+        zf: Open ZipFile object.
+        path: Path to the session file (for error messages).
+
+    Returns:
+        Dictionary of metadata values.
+    """
+    metadata: dict[str, Any] = {}
+
+    # Try to read metadata file (JSON format in newer versions)
+    if "metadata" in zf.namelist():
+        try:
+            with zf.open("metadata") as f:
+                content = f.read().decode("utf-8")
+                # Parse key=value format (sigrok classic format)
+                for line in content.strip().split("\n"):
+                    line = line.strip()
+                    if "=" in line:
+                        key, value = line.split("=", 1)
+                        key = key.strip()
+                        value = value.strip()
+                        # Try to convert numeric values
+                        try:
+                            if "." in value:
+                                metadata[key] = float(value)
+                            else:
+                                metadata[key] = int(value)
+                        except ValueError:
+                            metadata[key] = value
+        except Exception:
+            pass  # Use defaults if metadata parsing fails
+
+    # Extract channel names from probe entries
+    channels: list[str] = []
+    for key, value in metadata.items():
+        if key.startswith("probe"):
+            try:
+                idx = int(key.replace("probe", ""))
+                while len(channels) <= idx:
+                    channels.append(f"D{len(channels)}")
+                channels[idx] = value
+            except ValueError:
+                pass
+
+    if channels:
+        metadata["channels"] = channels
+
+    return metadata
+
+
+def _read_logic_data(
+    zf: zipfile.ZipFile,
+    logic_files: list[str],
+    total_channels: int,
+) -> NDArray[np.bool_]:
+    """Read and decode logic data from sigrok session.
+
+    Args:
+        zf: Open ZipFile object.
+        logic_files: List of logic data file names.
+        total_channels: Total number of digital channels.
+
+    Returns:
+        Boolean array of shape (channels, samples).
+    """
+    # Sort logic files to ensure correct order
+    logic_files = sorted(logic_files)
+
+    # Determine bytes per sample based on channel count
+    bytes_per_sample = (total_channels + 7) // 8
+
+    # Read all logic data
+    all_data = []
+    for logic_file in logic_files:
+        with zf.open(logic_file) as f:
+            raw_data = f.read()
+            all_data.append(raw_data)
+
+    # Combine data
+    combined = b"".join(all_data)
+
+    # Convert to numpy array
+    if bytes_per_sample == 1:
+        raw = np.frombuffer(combined, dtype=np.uint8)
+    elif bytes_per_sample == 2:
+        raw = np.frombuffer(combined, dtype=np.uint16)
+    elif bytes_per_sample <= 4:
+        # Pad to 4 bytes and read as uint32
+        padded = combined + b"\x00" * (len(combined) % 4)
+        raw = np.frombuffer(padded, dtype=np.uint32)
+    else:
+        # Handle larger sample widths
+        raw = np.frombuffer(combined, dtype=np.uint8)
+
+    # Extract individual channel bits
+    n_samples = len(raw)
+    channels_data = np.zeros((total_channels, n_samples), dtype=np.bool_)
+
+    for ch in range(total_channels):
+        if bytes_per_sample <= 4:
+            channels_data[ch] = (raw >> ch) & 1
+        else:
+            # For larger widths, calculate byte and bit position
+            byte_idx = ch // 8
+            bit_idx = ch % 8
+            byte_data = raw[byte_idx::bytes_per_sample]
+            channels_data[ch, : len(byte_data)] = (byte_data >> bit_idx) & 1
+
+    return channels_data
+
+
+def _compute_edges(
+    data: NDArray[np.bool_],
+    sample_rate: float,
+) -> list[tuple[float, bool]]:
+    """Compute edge timestamps from digital data.
+
+    Args:
+        data: Boolean array of digital samples.
+        sample_rate: Sample rate in Hz.
+
+    Returns:
+        List of (timestamp, is_rising) tuples.
+    """
+    edges: list[tuple[float, bool]] = []
+
+    if len(data) < 2:
+        return edges
+
+    # Find transitions
+    diff = np.diff(data.astype(np.int8))
+    rising_indices = np.where(diff == 1)[0]
+    falling_indices = np.where(diff == -1)[0]
+
+    time_per_sample = 1.0 / sample_rate
+
+    # Add rising edges
+    for idx in rising_indices:
+        timestamp = (idx + 1) * time_per_sample
+        edges.append((timestamp, True))
+
+    # Add falling edges
+    for idx in falling_indices:
+        timestamp = (idx + 1) * time_per_sample
+        edges.append((timestamp, False))
+
+    # Sort by timestamp
+    edges.sort(key=lambda x: x[0])
+
+    return edges
+
+
+__all__ = ["load_sigrok"]

oscura/loaders/tdms.py

@@ -0,0 +1,367 @@
+"""NI TDMS (Technical Data Management Streaming) file loader.
+
+This module provides loading of NI LabVIEW TDMS files using the
+npTDMS library when available.
+
+
+Example:
+    >>> from oscura.loaders.tdms import load_tdms
+    >>> trace = load_tdms("measurement.tdms")
+    >>> print(f"Sample rate: {trace.metadata.sample_rate} Hz")
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import TYPE_CHECKING, Any
+
+import numpy as np
+
+from oscura.core.exceptions import FormatError, LoaderError
+from oscura.core.types import TraceMetadata, WaveformTrace
+
+if TYPE_CHECKING:
+    from os import PathLike
+
+# Try to import npTDMS for TDMS support
+try:
+    from nptdms import TdmsFile
+
+    NPTDMS_AVAILABLE = True
+except ImportError:
+    NPTDMS_AVAILABLE = False
+
+
+def load_tdms(
+    path: str | PathLike[str],
+    *,
+    channel: str | int | None = None,
+    group: str | None = None,
+) -> WaveformTrace:
+    """Load an NI TDMS file.
+
+    TDMS files contain hierarchical data with groups and channels.
+    Each channel can have associated properties including sample rate.
+
+    Args:
+        path: Path to the TDMS file.
+        channel: Channel name or index to load. If None, loads the
+            first channel found.
+        group: Group name to select from. If None, uses the first group.
+
+    Returns:
+        WaveformTrace containing the channel data and metadata.
+
+    Raises:
+        Exception: If the file cannot be read or parsed.
+        LoaderError: If the file cannot be loaded or npTDMS is not installed.
+
+    Example:
+        >>> trace = load_tdms("measurement.tdms", group="Voltage", channel="CH1")
+        >>> print(f"Sample rate: {trace.metadata.sample_rate} Hz")
+        >>> print(f"Duration: {trace.duration:.6f} seconds")
+
+    References:
+        NI TDMS File Format: https://www.ni.com/en-us/support/documentation/
+    """
+    path = Path(path)
+
+    if not path.exists():
+        raise LoaderError(
+            "File not found",
+            file_path=str(path),
+        )
+
+    if not NPTDMS_AVAILABLE:
+        raise LoaderError(
+            "npTDMS library required for TDMS files",
+            file_path=str(path),
+            fix_hint="Install npTDMS: pip install npTDMS",
+        )
+
+    try:
+        return _load_with_nptdms(path, channel=channel, group=group)
+    except Exception as e:
+        if isinstance(e, LoaderError | FormatError):
+            raise
+        raise LoaderError(
+            "Failed to load TDMS file",
+            file_path=str(path),
+            details=str(e),
+            fix_hint="Ensure the file is a valid NI TDMS format.",
+        ) from e
+
+
+def _load_with_nptdms(
+    path: Path,
+    *,
+    channel: str | int | None = None,
+    group: str | None = None,
+) -> WaveformTrace:
+    """Load TDMS using npTDMS library.
+
+    Args:
+        path: Path to the TDMS file.
+        channel: Channel name or index.
+        group: Group name to select.
+
+    Returns:
+        WaveformTrace with channel data and metadata.
+
+    Raises:
+        FormatError: If file is not valid TDMS format or has no data.
+        LoaderError: If channel or group not found.
+    """
+    try:
+        tdms_file = TdmsFile.read(str(path))
+    except Exception as e:
+        raise FormatError(
+            "Failed to parse TDMS file",
+            file_path=str(path),
+            expected="Valid NI TDMS format",
+        ) from e
+
+    # Get available groups
+    groups = list(tdms_file.groups())
+
+    if not groups:
+        raise FormatError(
+            "No groups found in TDMS file",
+            file_path=str(path),
+        )
+
+    # Select group
+    if group is not None:
+        target_group = None
+        for g in groups:
+            if g.name == group:
+                target_group = g
+                break
+        if target_group is None:
+            available_groups = [g.name for g in groups]
+            raise LoaderError(
+                f"Group '{group}' not found",
+                file_path=str(path),
+                details=f"Available groups: {available_groups}",
+            )
+    else:
+        target_group = groups[0]
+
+    # Get channels in group
+    channels = list(target_group.channels())
+
+    if not channels:
+        raise FormatError(
+            f"No channels found in group '{target_group.name}'",
+            file_path=str(path),
+        )
+
+    # Select channel
+    if channel is not None:
+        if isinstance(channel, int):
+            if channel < 0 or channel >= len(channels):
+                raise LoaderError(
+                    f"Channel index {channel} out of range",
+                    file_path=str(path),
+                    details=f"Available channels: 0-{len(channels) - 1}",
+                )
+            target_channel = channels[channel]
+        elif isinstance(channel, str):
+            target_channel = None
+            for ch in channels:
+                if ch.name == channel:
+                    target_channel = ch
+                    break
+            if target_channel is None:
+                available_channels = [ch.name for ch in channels]
+                raise LoaderError(
+                    f"Channel '{channel}' not found",
+                    file_path=str(path),
+                    details=f"Available channels: {available_channels}",
+                )
+        else:
+            target_channel = channels[0]  # type: ignore[unreachable]
+    else:
+        target_channel = channels[0]
+
+    # Get channel data
+    data = target_channel.data
+    if data is None or len(data) == 0:
+        raise FormatError(
+            f"Channel '{target_channel.name}' has no data",
+            file_path=str(path),
+        )
+
+    # Convert to float64
+    data = np.asarray(data, dtype=np.float64)
+
+    # Extract sample rate from properties
+    sample_rate = _get_sample_rate(target_channel, target_group, tdms_file)
+
+    # Extract other metadata
+    vertical_scale = target_channel.properties.get("NI_Scale[0]_Linear_Slope")
+    vertical_offset = target_channel.properties.get("NI_Scale[0]_Linear_Y_Intercept")
+
+    # Get units if available
+    target_channel.properties.get("unit_string", None)
+
+    # Build metadata
+    metadata = TraceMetadata(
+        sample_rate=sample_rate,
+        vertical_scale=float(vertical_scale) if vertical_scale is not None else None,
+        vertical_offset=float(vertical_offset) if vertical_offset is not None else None,
+        source_file=str(path),
+        channel_name=target_channel.name,
+        trigger_info=_extract_tdms_properties(target_channel),
+    )
+
+    return WaveformTrace(data=data, metadata=metadata)
+
+
+def _get_sample_rate(
+    channel: Any,
+    group: Any,
+    tdms_file: Any,
+) -> float:
+    """Extract sample rate from TDMS channel properties.
+
+    Checks multiple common property names used by different NI software.
+
+    Args:
+        channel: TDMS channel object.
+        group: TDMS group object.
+        tdms_file: TDMS file object.
+
+    Returns:
+        Sample rate in Hz.
+    """
+    # Common property names for sample rate
+    sample_rate_keys = [
+        "wf_samples",  # DAQmx
+        "wf_increment",  # Waveform dt (inverse of sample rate)
+        "NI_RF_IQ_Rate",  # RF signal analyzer
+        "SamplingFrequency",  # SignalExpress
+        "dt",  # Delta time
+        "Fs",  # Sample rate
+        "SampleRate",
+        "sample_rate",
+    ]
+
+    # Check channel properties
+    for key in sample_rate_keys:
+        value = channel.properties.get(key)
+        if value is not None:
+            if key in ("wf_increment", "dt"):
+                # These are time intervals, invert for sample rate
+                if value > 0:
+                    return 1.0 / float(value)
+            else:
+                return float(value)
+
+    # Check group properties
+    for key in sample_rate_keys:
+        value = group.properties.get(key)
+        if value is not None:
+            if key in ("wf_increment", "dt"):
+                if value > 0:
+                    return 1.0 / float(value)
+            else:
+                return float(value)
+
+    # Check file properties
+    for key in sample_rate_keys:
+        value = tdms_file.properties.get(key)
+        if value is not None:
+            if key in ("wf_increment", "dt"):
+                if value > 0:
+                    return 1.0 / float(value)
+            else:
+                return float(value)
+
+    # Default sample rate if not found
+    return 1.0e6  # 1 MHz default
+
+
+def _extract_tdms_properties(channel: Any) -> dict[str, Any] | None:
+    """Extract relevant properties from TDMS channel.
+
+    Args:
+        channel: TDMS channel object.
+
+    Returns:
+        Dictionary of properties, or None if no useful properties found.
+    """
+    props: dict[str, Any] = {}
+
+    # Common useful properties
+    useful_keys = [
+        "unit_string",
+        "NI_ChannelName",
+        "wf_start_time",
+        "wf_start_offset",
+        "description",
+        "NI_Scale[0]_Linear_Slope",
+        "NI_Scale[0]_Linear_Y_Intercept",
+    ]
+
+    for key in useful_keys:
+        value = channel.properties.get(key)
+        if value is not None:
+            props[key] = value
+
+    return props if props else None
+
+
+def list_tdms_channels(
+    path: str | PathLike[str],
+) -> dict[str, list[str]]:
+    """List all groups and channels in a TDMS file.
+
+    Args:
+        path: Path to the TDMS file.
+
+    Returns:
+        Dictionary mapping group names to lists of channel names.
+
+    Raises:
+        LoaderError: If the file cannot be loaded.
+
+    Example:
+        >>> channels = list_tdms_channels("measurement.tdms")
+        >>> for group, chans in channels.items():
+        ...     print(f"Group '{group}': {chans}")
+    """
+    path = Path(path)
+
+    if not path.exists():
+        raise LoaderError(
+            "File not found",
+            file_path=str(path),
+        )
+
+    if not NPTDMS_AVAILABLE:
+        raise LoaderError(
+            "npTDMS library required for TDMS files",
+            file_path=str(path),
+            fix_hint="Install npTDMS: pip install npTDMS",
+        )
+
+    try:
+        tdms_file = TdmsFile.read(str(path))
+        result: dict[str, list[str]] = {}
+
+        for group in tdms_file.groups():
+            channel_names = [ch.name for ch in group.channels()]
+            result[group.name] = channel_names
+
+        return result
+
+    except Exception as e:
+        raise LoaderError(
+            "Failed to read TDMS file",
+            file_path=str(path),
+            details=str(e),
+        ) from e
+
+
+__all__ = ["list_tdms_channels", "load_tdms"]