cortexforge 0.1.0__py3-none-any.whl

cortexforge/cli/forge.py

@@ -0,0 +1,58 @@
+"""CLI argument parser for CorteXForge forge."""
+
+from argparse import ArgumentParser, Namespace
+from pathlib import Path
+
+
+def build_parser() -> ArgumentParser:
+    """
+    Build and configure the cli argument parser for CorteXForge forge.
+
+    Returns:
+        ArgumentParser: Configured argument parser instance with
+    """
+    parser = ArgumentParser(prog="CorteXForge forge", description="Dataset Generator")
+    sub = parser.add_subparsers(
+        dest="role",
+        required=True,
+        help="Radio role for this node (rx=receiver, tx=transmitter)",
+    )
+
+    rx = sub.add_parser("rx", help="Receiver mode")
+    rx.add_argument(
+        "--frequency", type=int, required=True, help="Receiver center frequency (Hz)"
+    )
+    rx.add_argument(
+        "--sample-rate", type=int, required=True, help="Receiver sample rate (Sps)"
+    )
+    rx.add_argument("--gain", type=int, required=True, help="Receiver gain (dB)")
+    rx.add_argument(
+        "--duration", type=int, required=True, help="Capture duration (seconds)"
+    )
+    rx.add_argument("--timeline", type=Path, required=True, help="Path to timeline CSV")
+    rx.add_argument(
+        "--output-path",
+        type=Path,
+        required=True,
+        help=(
+            "Path to output directory for results "
+            "(e.g. /cortexlab/homes/<username>/out)"
+        ),
+    )
+
+    tx = sub.add_parser("tx", help="Transmitter")
+    tx.add_argument("--timeline", type=Path, required=True, help="Path to timeline CSV")
+    tx.add_argument(
+        "--record-node", type=str, required=True, help="Select the recorder node"
+    )
+    tx.add_argument(
+        "--frequency", type=int, required=True, help="Transmitter center frequency (Hz)"
+    )
+    tx.add_argument("--gain", type=int, required=True, help="Transmitter gain (dB)")
+    return parser
+
+
+def parse_args(argv: list[str] | None = None) -> Namespace:
+    """Parse command line arguments."""
+    parser = build_parser()
+    return parser.parse_args(argv)

cortexforge/cli/planner.py

@@ -0,0 +1,45 @@
+"""CLI argument parser for CorteXForge planner."""
+
+from argparse import ArgumentParser, Namespace
+from pathlib import Path
+
+
+def build_parser() -> ArgumentParser:
+    """
+    Build and configure the cli argument parser for CorteXForge planner.
+
+    Returns:
+        ArgumentParser: Configured argument parser instance with
+    """
+    parser = ArgumentParser(prog="CorteXForge planner", description="Dataset Generator")
+    parser.add_argument(
+        "--username", required=True, type=str, help="username on CorteXlab"
+    )
+    parser.add_argument(
+        "--duration", type=int, default=60, help="Experiment duration in seconds"
+    )
+    parser.add_argument(
+        "--rx-frequency", type=int, default=2450000000, help="Receiver frequency"
+    )
+    parser.add_argument("--rx-gain", type=int, default=10, help="Receiver gain")
+    parser.add_argument(
+        "--rx-sample-rate", type=int, default=250000, help="Receiver sample-rate"
+    )
+    parser.add_argument(
+        "--overlapping",
+        action="store_true",
+        help="Allow overlapping signals in timeline",
+    )
+    parser.add_argument(
+        "--nodes-path",
+        type=Path,
+        default="configs/nodes.yaml",
+        help="Path to nodes.yaml file",
+    )
+    return parser
+
+
+def parse_args(argv: list[str] | None = None) -> Namespace:
+    """Parse command line arguments."""
+    parser = build_parser()
+    return parser.parse_args(argv)

cortexforge/datasets/__init__.py

@@ -0,0 +1,17 @@
+from .api import (
+    describe_dataset,
+    list_datasets,
+    list_datasets_with_versions,
+    list_versions,
+    load_dataset,
+    load_manifest,
+)
+
+__all__ = [
+    "load_dataset",
+    "list_datasets",
+    "list_datasets_with_versions",
+    "describe_dataset",
+    "load_manifest",
+    "list_versions",
+]

cortexforge/datasets/api.py

@@ -0,0 +1,88 @@
+from pathlib import Path
+
+from .local import load_local_dataset
+from .manifest import load_manifest
+
+
+def load_dataset(
+    name: str,
+    version: str,
+    split: str,
+    root: str | Path,
+    verify: bool = True,
+):
+    dataset_dir = Path(root) / name / version
+    return load_local_dataset(dataset_dir=dataset_dir, split=split, verify=verify)
+
+
+def describe_dataset(
+    name: str,
+    version: str,
+    root: str | Path,
+):
+    dataset_dir = Path(root) / name / version
+    return load_manifest(dataset_dir / "manifest.json")
+
+
+def list_datasets(root: str | Path) -> list[str]:
+    root = Path(root)
+
+    if not root.exists():
+        return []
+
+    datasets = []
+
+    for dataset_dir in root.iterdir():
+        if not dataset_dir.is_dir():
+            continue
+
+        has_valid_version = any(
+            (version_dir / "manifest.json").exists()
+            for version_dir in dataset_dir.iterdir()
+            if version_dir.is_dir()
+        )
+
+        if has_valid_version:
+            datasets.append(dataset_dir.name)
+
+    return sorted(datasets)
+
+
+def list_versions(name: str, root: str | Path) -> list[str]:
+    root = Path(root)
+    dataset_dir = root / name
+
+    if not dataset_dir.exists() or not dataset_dir.is_dir():
+        return []
+
+    versions = [
+        version_dir.name
+        for version_dir in dataset_dir.iterdir()
+        if version_dir.is_dir() and (version_dir / "manifest.json").exists()
+    ]
+
+    return sorted(versions)
+
+
+def list_datasets_with_versions(root: str | Path) -> dict[str, list[str]]:
+    root = Path(root)
+
+    if not root.exists():
+        return {}
+
+    result = {}
+
+    for dataset_dir in root.iterdir():
+        if not dataset_dir.is_dir():
+            continue
+
+        versions = [
+            version_dir.name
+            for version_dir in dataset_dir.iterdir()
+            if (version_dir / "manifest.json").exists()
+        ]
+
+        if versions:
+            result[dataset_dir.name] = sorted(versions)
+
+    return result

cortexforge/datasets/hash.py

@@ -0,0 +1,11 @@
+import hashlib
+from pathlib import Path
+
+
+def _sha512_hex(path: Path, chunk_size: int = 1024 * 1024) -> str:
+    """Compute SHA-512 hash of a file and return it as a hex string."""
+    h = hashlib.sha512()
+    with path.open("rb") as f:
+        for chunk in iter(lambda: f.read(chunk_size), b""):
+            h.update(chunk)
+    return h.hexdigest()

cortexforge/datasets/local.py

@@ -0,0 +1,87 @@
+import json
+from pathlib import Path
+
+from .hash import _sha512_hex
+from .manifest import load_manifest
+from .types import LocalDataset
+
+
+def validate_sigmf_recording(meta_path: Path) -> None:
+    if not meta_path.exists():
+        raise FileNotFoundError(f"Missing SigMF metadata: {meta_path}")
+
+    # Load JSON
+    with meta_path.open("r", encoding="utf-8") as f:
+        meta = json.load(f)
+
+    if "global" not in meta:
+        raise ValueError(f"Invalid SigMF meta (missing 'global'): {meta_path}")
+
+    global_meta = meta["global"]
+
+    # ---- core:data_file ----
+    data_file = global_meta.get("core:data_file")
+    if not data_file:
+        raise ValueError(f"Missing 'core:data_file' in {meta_path}")
+
+    data_path = meta_path.parent / data_file
+
+    if not data_path.exists():
+        raise FileNotFoundError(
+            f"SigMF data file not found: {data_path} (referenced in {meta_path})"
+        )
+
+    # ---- SHA512 ----
+    expected_hash = global_meta.get("core:sha512")
+
+    if expected_hash:
+        actual_hash = _sha512_hex(data_path)
+
+        if actual_hash != expected_hash:
+            raise ValueError(
+                f"SHA512 mismatch for {data_path}\n"
+                f"Expected: {expected_hash}\n"
+                f"Got:      {actual_hash}"
+            )
+
+
+def _associated_data_path(meta_path: Path) -> Path:
+    if meta_path.name.endswith(".sigmf-meta"):
+        return meta_path.with_name(meta_path.name.replace(".sigmf-meta", ".sigmf-data"))
+    raise ValueError(f"Expected a .sigmf-meta file, got: {meta_path}")
+
+
+def load_local_dataset(
+    dataset_dir: str | Path,
+    split: str,
+    verify: bool = True,
+) -> LocalDataset:
+    dataset_dir = Path(dataset_dir)
+    manifest = load_manifest(dataset_dir / "manifest.json")
+
+    if split not in manifest.splits:
+        raise ValueError(f"Unknown split: {split!r}")
+
+    recordings = []
+    for rel_path in manifest.splits[split].recordings:
+        meta_path = dataset_dir / rel_path
+
+        if verify:
+            validate_sigmf_recording(meta_path)
+
+        if not meta_path.exists():
+            raise FileNotFoundError(f"Missing SigMF metadata file: {meta_path}")
+
+        data_path = _associated_data_path(meta_path)
+        if not data_path.exists():
+            raise FileNotFoundError(f"Missing SigMF data file: {data_path}")
+
+        recordings.append(meta_path)
+
+    return LocalDataset(
+        name=manifest.name,
+        version=manifest.version,
+        split=split,
+        root=dataset_dir,
+        recordings=recordings,
+    )

cortexforge/datasets/manifest.py

@@ -0,0 +1,68 @@
+import json
+from pathlib import Path
+
+from .types import DatasetManifest, GeneratorManifest, SplitManifest
+
+REQUIRED_TOP_LEVEL_FIELDS = {"name", "version", "format", "generator", "splits"}
+REQUIRED_GENERATOR_FIELDS = {"name", "version"}
+
+
+def validate_manifest(path: str | Path) -> None:
+    path = Path(path)
+
+    if not path.exists():
+        raise FileNotFoundError(f"Manifest not found: {path}")
+
+    with path.open("r", encoding="utf-8") as f:
+        raw = json.load(f)
+
+    missing = REQUIRED_TOP_LEVEL_FIELDS - raw.keys()
+    if missing:
+        raise ValueError(f"Missing required manifest fields: {sorted(missing)}")
+
+    generator = raw["generator"]
+    if not isinstance(generator, dict):
+        raise ValueError("'generator' must be a dictionary")
+
+    missing_generator = REQUIRED_GENERATOR_FIELDS - generator.keys()
+    if missing_generator:
+        raise ValueError(
+            f"Missing required generator fields: {sorted(missing_generator)}"
+        )
+
+    splits = raw["splits"]
+    if not isinstance(splits, dict) or not splits:
+        raise ValueError("'splits' must be a non-empty dictionary")
+
+    for split_name, split_data in splits.items():
+        if not isinstance(split_data, dict):
+            raise ValueError(f"Split '{split_name}' must be a dictionary")
+
+        if "recordings" not in split_data:
+            raise ValueError(f"Split '{split_name}' is missing 'recordings'")
+
+        if not isinstance(split_data["recordings"], list):
+            raise ValueError(f"Split '{split_name}' field 'recordings' must be a list")
+
+
+def load_manifest(path: str | Path) -> DatasetManifest:
+    path = Path(path)
+    validate_manifest(path)
+    with path.open("r", encoding="utf-8") as f:
+        raw = json.load(f)
+
+    splits = {
+        split_name: SplitManifest(recordings=split_data["recordings"])
+        for split_name, split_data in raw["splits"].items()
+    }
+
+    generator = GeneratorManifest(**raw["generator"])
+
+    return DatasetManifest(
+        name=raw["name"],
+        version=raw["version"],
+        description=raw.get("description", ""),
+        format=raw["format"],
+        generator=generator,
+        splits=splits,
+    )

cortexforge/datasets/types.py

@@ -0,0 +1,33 @@
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Dict, List
+
+
+@dataclass(slots=True)
+class SplitManifest:
+    recordings: List[str]
+
+
+@dataclass(slots=True)
+class GeneratorManifest:
+    name: str
+    version: str
+
+
+@dataclass(slots=True)
+class DatasetManifest:
+    name: str
+    version: str
+    description: str
+    format: str
+    generator: GeneratorManifest
+    splits: Dict[str, SplitManifest]
+
+
+@dataclass(slots=True)
+class LocalDataset:
+    name: str
+    version: str
+    split: str
+    root: Path
+    recordings: List[Path]

cortexforge/forge/main.py

@@ -0,0 +1,25 @@
+from cortexforge.cli.forge import parse_args
+from cortexforge.utils.logger import setup_logger
+
+logger = setup_logger()
+
+
+def main() -> None:
+    args = parse_args()
+    logger.info("Starting CorteXForge...")
+    logger.info(f"Args={args}")
+
+    if args.role == "rx":
+        from cortexforge.forge.radio.rx import main as rx_main
+
+        rx_main(args)
+    elif args.role == "tx":
+        from cortexforge.forge.radio.tx import main as tx_main
+
+        tx_main(args)
+    else:
+        raise ValueError(f"Unknown role: {args.role}")
+
+
+if __name__ == "__main__":
+    main()

cortexforge/forge/radio/rx.py

@@ -0,0 +1,100 @@
+from logging import getLogger
+from time import sleep
+from datetime import datetime, timezone
+from gnuradio import uhd
+
+from cortexforge.forge.radio.rx_recorder import RxRecorder
+from cortexforge.forge.utils.sigmf_writer import write_sigmf
+from cortexforge.forge.utils.compute_baseline import compute_baseline
+from cortexforge.forge.utils.load_timeline import load_timeline
+from cortexforge.forge.utils.sync_barrier.rx_barrier_server import RxBarrierServer
+from cortexforge.forge.utils.sync_barrier.sync_config import SyncConfig
+from cortexforge.forge.utils.uhd_time import arm_time_reset_next_pps
+from cortexforge.forge.utils.sigmf.sigmf_annotations import timeline_to_sigmf_annotations
+from cortexforge.forge.utils.node_identity import get_node_name
+
+
+logger = getLogger(__name__)
+
+
+def main(args):
+    out_dir = args.output_path
+    out_dir.mkdir(parents=True, exist_ok=True)
+
+    timeline = load_timeline(args.timeline)
+
+    radios = {str(ev["radio"]) for ev in timeline if ev.get("radio")}
+
+    logger.debug(f"Radios in timeline: {radios}, expecting {len(radios)} to synchronize")
+
+    raw_path = out_dir / "temp.cf32"
+
+    tb = RxRecorder(
+        usrp_args="",
+        freq=args.frequency,
+        rate=args.sample_rate,
+        gain=args.gain,
+        out_path=str(raw_path),
+    )
+    cfg = SyncConfig(
+        server_host=get_node_name(),
+        port_reg=5555,
+        port_pub=5556,
+        expected_tx=len(radios),
+    )
+    barrier = RxBarrierServer(cfg)
+    barrier.wait_for_all()
+
+    barrier.broadcast({"type": "GO"})
+    arm_time_reset_next_pps(tb.src)
+
+    capture_start_uhd = 1.0
+    if hasattr(tb.src, "set_start_time"):
+        tb.src.set_start_time(uhd.time_spec(capture_start_uhd))
+        logger.info("RX stream scheduled at UHD t=%.3f s", capture_start_uhd)
+        rx_uhd_t0 = capture_start_uhd
+    else:
+        rx_uhd_t0 = None
+        logger.warning("RX source has no set_start_time(); using runtime t0 estimate")
+
+    tb.start()
+    if rx_uhd_t0 is None:
+        rx_uhd_t0 = tb.src.get_time_now().get_real_secs()
+
+    while tb.src.get_time_now().get_real_secs() < rx_uhd_t0 + args.duration:
+        sleep(0.001)
+
+    tb.stop()
+    tb.wait()
+
+    stats = compute_baseline(
+        path=str(raw_path),
+        sample_rate=args.sample_rate
+    )
+
+    stamp = datetime.now(timezone.utc).strftime("%Y%m%d_%H%M%S")
+    base_path = out_dir / f"{stamp}"
+
+
+    logger.info(f"Recording stats: {stats}")
+    data_path, meta_path = write_sigmf(
+        base_path=str(base_path),
+        data_file=str(raw_path),
+        stat=stats,
+        sample_rate=args.sample_rate,
+        center_freq=args.frequency,
+        hardware=tb.src.get_usrp_info().get('mboard_id'),
+        author="Andrea Joly",
+        description="CorteXForge recording",
+        gain=args.gain,
+        annotations = timeline_to_sigmf_annotations(
+            events=timeline,
+            rx_sample_rate=args.sample_rate,
+            rx_uhd_t0=rx_uhd_t0,
+            tx_center_freq=args.frequency,
+            tx_gain=args.gain,
+            rx_data_path=str(raw_path),
+            baseline_stat=stats,
+        ),
+    )
+    logger.info(f"SigMF written: {data_path} and {meta_path}")

cortexforge/forge/radio/rx_recorder.py

@@ -0,0 +1,40 @@
+from gnuradio import blocks, gr, uhd
+from logging import getLogger
+
+logger = getLogger(__name__)
+
+
+class RxRecorder(gr.top_block):
+    """
+    GNU Radio top block for recording from a USRP device.
+    """
+
+    def __init__(
+        self, usrp_args: str, freq: float, rate: float, gain: float, out_path: str
+    ) -> None:
+        super().__init__("Rx Recorder")
+
+        self.rx_channel = 0
+
+        self.src = uhd.usrp_source(
+            usrp_args,
+            uhd.stream_args(cpu_format="fc32", channels=[self.rx_channel]),
+        )
+        self.src.set_clock_source("external", self.rx_channel)
+        self.src.set_time_source("external", self.rx_channel)
+        self.src.set_samp_rate(rate)
+        self.src.set_center_freq(freq, self.rx_channel)
+        self._try_set_rx_agc(False)
+        self.src.set_gain(gain, self.rx_channel)
+        self.src.set_antenna("TX/RX", self.rx_channel)
+        self.sink = blocks.file_sink(gr.sizeof_gr_complex, out_path, False)
+        self.sink.set_unbuffered(False)
+        self.connect(self.src, self.sink)
+
+
+    def _try_set_rx_agc(self, enable: bool = False) -> None:
+        try:
+            self.src.set_rx_agc(enable)
+        except RuntimeError as e:
+            logger.warning("RX AGC not supported by this radio; msg: %s", str(e))
+                            

cortexforge/forge/radio/tx.py

@@ -0,0 +1,77 @@
+import time
+from logging import getLogger
+
+from cortexforge.forge.radio.tx_burst import TxTimeline
+from cortexforge.forge.radio.waveforms import make_burst
+from cortexforge.forge.utils.load_timeline import load_timeline
+from cortexforge.forge.utils.node_identity import get_node_name
+from cortexforge.forge.utils.sync_barrier.sync_config import SyncConfig
+from cortexforge.forge.utils.sync_barrier.tx_barrier_client import TxBarrierClient
+from cortexforge.forge.utils.uhd_time import arm_time_reset_next_pps
+
+logger = getLogger(__name__)
+
+
+def main(args):
+    node_name = get_node_name()
+    timeline_all = load_timeline(args.timeline)
+    timeline = [ev for ev in timeline_all if ev.get("radio") == node_name]
+
+    if not timeline:
+        logger.warning("No TX events found for node %s.", node_name)
+        return
+
+    # Prepare I/Q burst
+    events_with_iq = []
+    for ev in timeline:
+        iq = make_burst(
+            modulation=ev["modulation"],
+            sample_rate=ev["sample_rate_sps"],
+            symbol_rate=ev["symbol_rate"],
+            duration_s=ev["duration_s"],
+            rolloff=ev["rolloff"],
+            amplitude=ev["amplitude"],
+        ).astype("complex64")
+
+        ev2 = dict(ev)
+        ev2["iq"] = iq
+        events_with_iq.append(ev2)
+
+        logger.info(
+            f"event start={ev['t_start_s']} dur={ev['duration_s']} "
+            f"modulation={ev['modulation']}"
+        )
+
+    tb = TxTimeline(
+        usrp_args="",
+        rate=events_with_iq[0]["sample_rate_sps"],
+        center_freq=args.frequency,
+        gain=args.gain,
+        events_with_iq=events_with_iq,
+    )
+
+    cfg = SyncConfig(server_host=args.record_node, port_reg=5555, port_pub=5556)
+    client = TxBarrierClient(cfg, node_name=node_name)
+    client.register()
+
+    msg = client.wait_broadcast()
+    assert msg["type"] == "GO"
+
+    # logger.info(f"UHD before reset: {tb.sink.get_time_now().get_real_secs():.6f} s")
+
+    arm_time_reset_next_pps(tb.sink)
+    # logger.info(f"UHD after reset: {tb.sink.get_time_now().get_real_secs():.6f} s")
+
+    tb.start()
+    # logger.info(f"UHD time at start: {tb.sink.get_time_now().get_real_secs():.6f} s")
+
+    # 5) Attendre jusqu’à la fin du dernier burst
+    last_ev = max(events_with_iq, key=lambda e: e["t_start_s"] + e["duration_s"])
+    t_end = float(last_ev["t_start_s"] + last_ev["duration_s"]) + 1.0
+
+    while tb.sink.get_time_now().get_real_secs() < t_end:
+        time.sleep(0.001)
+
+    tb.stop()
+    tb.wait()
+    logger.info("Timeline complete.")