ber-equalization-studio 0.1.0__py3-none-any.whl

ber_equalization_studio/visualization.py

@@ -0,0 +1,186 @@
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Any
+
+import matplotlib.pyplot as plt
+import pandas as pd
+
+
+def _try_plotly():
+    try:
+        import plotly.express as px
+        import plotly.graph_objects as go
+
+        return px, go
+    except Exception:
+        return None, None
+
+
+def history_figure(history: dict[str, list[Any]], model_name: str, metric: str = "ber"):
+    """Build a Plotly history figure for notebooks, returning None if Plotly is unavailable."""
+
+    _, go = _try_plotly()
+    if go is None:
+        return None
+
+    epochs = history.get("val_epochs") or list(range(1, len(history.get("train_loss", [])) + 1))
+    fig = go.Figure()
+    if metric == "loss":
+        if history.get("train_loss"):
+            fig.add_trace(go.Scatter(x=epochs, y=history["train_loss"], name="train loss", mode="lines"))
+        if history.get("val_loss"):
+            fig.add_trace(go.Scatter(x=epochs, y=history["val_loss"], name="val loss", mode="lines"))
+        fig.update_layout(
+            title=f"{model_name}: loss curves",
+            xaxis_title="Epoch",
+            yaxis_title="Loss",
+            template="plotly_white",
+        )
+        return fig
+
+    if history.get("val_ber"):
+        fig.add_trace(go.Scatter(x=epochs, y=history["val_ber"], name="val BER", mode="lines+markers"))
+    if history.get("test_ber"):
+        fig.add_trace(
+            go.Scatter(
+                x=history.get("test_epochs", []),
+                y=history["test_ber"],
+                name="test BER",
+                mode="markers",
+            )
+        )
+    fig.update_layout(
+        title=f"{model_name}: BER",
+        xaxis_title="Epoch",
+        yaxis_title="BER",
+        yaxis_type="log",
+        template="plotly_white",
+    )
+    return fig
+
+
+def comparison_figure(
+    df: pd.DataFrame,
+    x: str = "model_type",
+    y: str = "equalized_ber",
+    color: str | None = None,
+    kind: str = "bar",
+):
+    """Build a Plotly comparison figure for notebooks, returning None if Plotly is unavailable."""
+
+    px, _ = _try_plotly()
+    if px is None or x not in df.columns or y not in df.columns:
+        return None
+
+    if color is None:
+        color = "improvement_rel" if kind == "bar" and "improvement_rel" in df.columns else None
+    hover_cols = [column for column in ["trainable_params", "improvement_rel", "samples_per_sec"] if column in df]
+    plot_df = df.sort_values(y) if y in df.columns else df
+    common = {
+        "x": x,
+        "y": y,
+        "color": color if color in df.columns else None,
+        "hover_data": hover_cols,
+        "template": "plotly_white",
+    }
+    if kind == "scatter":
+        fig = px.scatter(plot_df, **common)
+    else:
+        fig = px.bar(plot_df, **common)
+    layout = {
+        "title": f"{y} by {x}",
+        "xaxis_title": x.replace("_", " ").title(),
+        "yaxis_title": y.replace("_", " ").title(),
+    }
+    if "ber" in y.lower():
+        layout["yaxis_type"] = "log"
+    fig.update_layout(
+        **layout,
+    )
+    return fig
+
+
+def plot_history(history: dict[str, list[Any]], model_name: str, out_dir: Path) -> list[Path]:
+    out_dir.mkdir(parents=True, exist_ok=True)
+    paths: list[Path] = []
+    px, go = _try_plotly()
+    epochs = history.get("val_epochs") or list(range(1, len(history.get("train_loss", [])) + 1))
+
+    if go is not None:
+        fig = history_figure(history, model_name, metric="loss")
+        path = out_dir / f"{model_name}_loss.html"
+        fig.write_html(path)
+        paths.append(path)
+
+        fig = history_figure(history, model_name, metric="ber")
+        path = out_dir / f"{model_name}_ber.html"
+        fig.write_html(path)
+        paths.append(path)
+        return paths
+
+    plt.figure(figsize=(9, 5))
+    if history.get("train_loss"):
+        plt.plot(epochs, history["train_loss"], label="train loss")
+    if history.get("val_loss"):
+        plt.plot(epochs, history["val_loss"], label="val loss")
+    plt.xlabel("Epoch")
+    plt.ylabel("Loss")
+    plt.legend()
+    plt.grid(alpha=0.25)
+    path = out_dir / f"{model_name}_loss.png"
+    plt.savefig(path, dpi=170, bbox_inches="tight")
+    plt.close()
+    paths.append(path)
+    return paths
+
+
+def plot_comparison(df: pd.DataFrame, out_dir: Path, filename: str = "model_comparison") -> list[Path]:
+    out_dir.mkdir(parents=True, exist_ok=True)
+    paths: list[Path] = []
+    px, _ = _try_plotly()
+    if px is not None and {"model_type", "equalized_ber"}.issubset(df.columns):
+        hover_cols = [column for column in ["trainable_params", "improvement_rel", "samples_per_sec"] if column in df]
+        fig = px.bar(
+            df.sort_values("equalized_ber"),
+            x="model_type",
+            y="equalized_ber",
+            color="improvement_rel" if "improvement_rel" in df.columns else None,
+            hover_data=hover_cols,
+            log_y=True,
+            template="plotly_white",
+            title="Equalized BER by model",
+        )
+        fig.update_layout(xaxis_title="Model", yaxis_title="BER")
+        path = out_dir / f"{filename}_ber.html"
+        fig.write_html(path)
+        paths.append(path)
+
+    if {"trainable_params", "equalized_ber"}.issubset(df.columns):
+        if px is not None:
+            fig = px.scatter(
+                df,
+                x="trainable_params",
+                y="equalized_ber",
+                color="model_type" if "model_type" in df.columns else None,
+                hover_name="model_type" if "model_type" in df.columns else None,
+                log_y=True,
+                template="plotly_white",
+                title="Complexity vs BER",
+            )
+            fig.update_layout(xaxis_title="Trainable parameters", yaxis_title="BER")
+            path = out_dir / f"{filename}_complexity.html"
+            fig.write_html(path)
+            paths.append(path)
+        else:
+            plt.figure(figsize=(8, 5))
+            plt.scatter(df["trainable_params"], df["equalized_ber"])
+            plt.yscale("log")
+            plt.xlabel("Trainable parameters")
+            plt.ylabel("BER")
+            plt.grid(alpha=0.25)
+            path = out_dir / f"{filename}_complexity.png"
+            plt.savefig(path, dpi=170, bbox_inches="tight")
+            plt.close()
+            paths.append(path)
+    return paths

ber_equalization_studio-0.1.0.dist-info/METADATA

@@ -0,0 +1,266 @@
+Metadata-Version: 2.4
+Name: ber-equalization-studio
+Version: 0.1.0
+Summary: Research studio for BER equalization experiments in nonlinear optical links.
+Keywords: ber,equalization,optical-communications,photonics,pytorch
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Science/Research
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: Topic :: Scientific/Engineering :: Physics
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+Requires-Dist: matplotlib>=3.8
+Requires-Dist: numpy>=1.24
+Requires-Dist: pandas>=2.0
+Requires-Dist: plotly>=5.18
+Requires-Dist: torch>=2.1
+Provides-Extra: kan
+Provides-Extra: notebook
+Requires-Dist: ipykernel>=6.29; extra == "notebook"
+Requires-Dist: ipython>=8.18; extra == "notebook"
+Requires-Dist: ipywidgets>=8.1; extra == "notebook"
+Requires-Dist: jupyterlab>=4.0; extra == "notebook"
+Requires-Dist: nbformat>=5.9; extra == "notebook"
+Provides-Extra: dev
+Requires-Dist: build>=1.2; extra == "dev"
+Requires-Dist: pytest>=8.0; extra == "dev"
+Requires-Dist: twine>=5.1; extra == "dev"
+Provides-Extra: publish
+Requires-Dist: build>=1.2; extra == "publish"
+Requires-Dist: twine>=5.1; extra == "publish"
+
+# BER Equalization Studio
+
+Research library for BER equalization experiments in nonlinear fiber-optic links.
+
+The studio wraps the existing, validated `BER_minimization_survey/ber_equalization.py`
+backend with a cleaner API:
+
+- structured dataclass configs instead of editing one large global `Config`;
+- model registry and aliases;
+- reusable experiment runner;
+- per-model run folders with checkpoints, metrics, histories, and plots;
+- result comparison helpers for paper tables;
+- Plotly HTML visualizations for BER/loss curves and model tradeoffs.
+
+## Layout
+
+```text
+ber-equalization-studio/
+  src/ber_equalization_studio/
+    config.py          # DataConfig, ModelConfig, TrainingConfig, StudioConfig
+    legacy.py          # adapter to the original ber_equalization.py backend
+    models.py          # model registry and constructors
+    data.py            # dataset preparation wrapper
+    experiment.py      # ExperimentRunner
+    api.py             # notebook-friendly Studio, RunResult, StudyResult
+    results.py         # result loading/comparison helpers
+    visualization.py   # BER/loss/complexity plots
+    cli.py             # ber-studio command
+  examples/
+    run_smoke.py
+```
+
+## Quick Start
+
+Install from PyPI:
+
+```powershell
+python -m pip install "ber-equalization-studio[notebook]"
+```
+
+For local development from the repository root:
+
+```powershell
+cd ber-equalization-studio
+python -m pip install -e ".[notebook,dev]"
+ber-studio models
+```
+
+Run a short smoke experiment:
+
+```powershell
+ber-studio run `
+  --name smoke_complex_fastkan `
+  --models complex_fastkan,mlp `
+  --data-dir ..\BER_minimization_survey\symbols_new `
+  --epochs 5 `
+  --max-test-files 1
+```
+
+If your symbol CSV files live directly under `BER_minimization_survey`, pass that
+folder instead. The backend searches for files named:
+
+```text
+Symbols_1m_1ch_PR_*.csv
+```
+
+The PyPI package includes the legacy training backend and the local EfficientKAN
+implementation used by the studio. You only need to pass data directories with symbol
+CSV files.
+
+## Notebook API
+
+The easiest way to work from Jupyter is the high-level `Studio` interface:
+
+```python
+from ber_equalization_studio import Studio
+
+studio = Studio(
+    data_dirs=["../test"],
+    out_dir="notebook_runs",
+    device="cuda",
+)
+
+studio.models()
+```
+
+Run one experiment:
+
+```python
+run = studio.run(
+    name="fastkan_smoke",
+    models=["complex_fastkan", "mlp"],
+    epochs=5,
+    lr=1e-3,
+    context_k=32,
+    max_test_files=1,
+)
+
+run.results
+```
+
+Useful result helpers:
+
+```python
+run.best()
+run.compare()
+run.history("complex_fastkan")
+run.plot_history("complex_fastkan")
+run.plot_comparison()
+run.run_dir
+```
+
+Sweep a small research grid:
+
+```python
+study = studio.sweep(
+    name="context_lr_sweep",
+    models=["complex_fastkan"],
+    grid={
+        "context_k": [16, 32, 64],
+        "lr": [1e-3, 3e-4],
+    },
+    epochs=50,
+    max_test_files=2,
+)
+
+study.results.sort_values("equalized_ber")
+study.best()
+study.plot_tradeoff(x="trainable_params", y="equalized_ber")
+```
+
+Short parameters such as `epochs`, `lr`, `context_k`, `max_test_files`, `data_dirs`,
+and `out_dir` are mapped to the structured config. For advanced settings, use dotted
+keys:
+
+```python
+run = studio.run(
+    name="custom_eval",
+    models="complex_fastkan",
+    **{
+        "evaluation.ber_scale_steps": 20,
+        "training.early_stopping_patience": 100,
+    },
+)
+```
+
+## Low-Level Python API
+
+```python
+from pathlib import Path
+
+from ber_equalization_studio import (
+    DataConfig,
+    ExperimentConfig,
+    ExperimentRunner,
+    OutputConfig,
+    StudioConfig,
+    TrainingConfig,
+)
+
+config = StudioConfig(
+    data=DataConfig(
+        data_dirs=[Path("../BER_minimization_survey/symbols_new")],
+        context_k=32,
+        max_test_files=1,
+    ),
+    training=TrainingConfig(epochs=10, learning_rate=1e-3),
+    output=OutputConfig(out_dir=Path("studio_outputs"), experiment_name="fastkan_baseline"),
+    experiment=ExperimentConfig(models=["complex_fastkan", "efficient_kan_baseline", "mlp"]),
+)
+
+result = ExperimentRunner(config).run()
+print(result["results_csv"])
+```
+
+## Available Models
+
+Use:
+
+```powershell
+ber-studio models
+```
+
+Important current models:
+
+- `complex_fastkan`: lightweight complex temporal encoder + RBF/FastKAN regression head.
+- `complex_fastkan_classifier`: same encoder + RBF/FastKAN 16-class detector.
+- `efficient_kan_baseline`: flat IQ window + B-spline EfficientKAN regression.
+- `kan_classifier`: flat IQ window + B-spline EfficientKAN classifier.
+- `cnn_kan`: temporal CNN + EfficientKAN head.
+- `mlp`, `cnn`, `tcn`, `lstm`, `transformer`: neural baselines.
+
+## Outputs
+
+Each run creates:
+
+```text
+studio_outputs/<experiment_name>/
+  config.json
+  dataset_summary.json
+  results.csv
+  model_comparison_ber.html
+  model_comparison_complexity.html
+  <model_name>/
+    metrics.json
+    history.json
+    final_state_dict.pt
+    <model_name>_loss.html
+    <model_name>_ber.html
+```
+
+Compare previous runs:
+
+```powershell
+ber-studio compare studio_outputs
+```
+
+## Research Workflow
+
+Recommended progression:
+
+1. In Jupyter, start with `Studio(...).models()` to choose candidate models.
+2. Run `complex_fastkan`, `efficient_kan_baseline`, and `mlp` as a short smoke test.
+3. Increase epochs and use the full test split for promising candidates.
+4. Use `run.results`, `study.results`, and saved `results.csv` files for paper tables.
+5. Sweep one family at a time with `studio.sweep(...)`, changing only a few parameters per study.
+
+The current implementation intentionally keeps the original backend intact. That means
+existing BER computation, Gray labels, file-level split protocol, normalization, KAN
+implementations, and training loop behavior are preserved while the new API makes
+future experiments much easier to compose.

ber_equalization_studio-0.1.0.dist-info/RECORD

@@ -0,0 +1,19 @@
+ber_equalization_studio/__init__.py,sha256=Q7ViLrj1-HbuMJQYjnyaVKrQR6lHBJ028spHjIb3FSU,1731
+ber_equalization_studio/api.py,sha256=go2sNKiQBP8z3MTAKeXyxPV6NcwKXo5h-6F18Abe740,12440
+ber_equalization_studio/cli.py,sha256=vl0wffZMBqfxyNuin3ZwiqKQswJ8I1esZ684d2qMfNQ,3777
+ber_equalization_studio/config.py,sha256=TanndlwrR3gpxwjDqzl1RKWKhyPsz4fkmHyYNDJRSK4,5142
+ber_equalization_studio/data.py,sha256=AlD5gyvPD_ZE7eyoosgJEHn7iTCM5Oyt6qN6zsIR3Mo,1168
+ber_equalization_studio/experiment.py,sha256=lEtBjujfockDuYMyatKlc2osVZgXtfxAs_obgi-3vFo,3592
+ber_equalization_studio/legacy.py,sha256=1JOmVvVuMSD0BAgILS6WQh6doqRO6jHq-ur7tcElIYE,6351
+ber_equalization_studio/models.py,sha256=Pmc_GD7aJZIT0t3K5cWKePXKMzGgh9ZGjcFuVZxQuhg,3368
+ber_equalization_studio/results.py,sha256=kQzj1m9HaoulJ9LtCJnB5Br6YdikT-BQCLOB5XdadEY,2417
+ber_equalization_studio/visualization.py,sha256=gnezSWiNBP1dIJ527P1pxL1AusWV1bRi7G9eVpjIniA,6288
+ber_equalization_studio/_legacy_backend/__init__.py,sha256=DEvlIvfNzTUPOUQC4VsL_rPtC_sm4Ly3Puush4QcILQ,61
+ber_equalization_studio/_legacy_backend/ber_equalization.py,sha256=9Od0JwoSHkPhARoYGNfW_ChH8FJqmNYz3UyvcGDGXFU,153794
+ber_equalization_studio/_legacy_backend/efficient_kan/__init__.py,sha256=xfwSQFMEZyj4IZzRwvziGYFBTVvtEPa1j9Gftzvx3_s,67
+ber_equalization_studio/_legacy_backend/efficient_kan/kan.py,sha256=6XZS5JDTm4ENYoJiTMwgpGQ4L9p1huoHebhuZAu2RQM,7929
+ber_equalization_studio-0.1.0.dist-info/METADATA,sha256=c8Ny3IUhMT0XYbdKhmbwOkqZmd9X32E9k3a6Qc2_Uic,7477
+ber_equalization_studio-0.1.0.dist-info/WHEEL,sha256=aeYiig01lYGDzBgS8HxWXOg3uV61G9ijOsup-k9o1sk,91
+ber_equalization_studio-0.1.0.dist-info/entry_points.txt,sha256=1YV0HoWUnDJ_X9WJsFP1qgP8b9CXst_WK9PtwJDT1PA,64
+ber_equalization_studio-0.1.0.dist-info/top_level.txt,sha256=MoXxAEwvBtpZGfQK3yYI1vhr2k8v9al3D1GK_3KUPlc,24
+ber_equalization_studio-0.1.0.dist-info/RECORD,,

ber_equalization_studio-0.1.0.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (82.0.1)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

ber_equalization_studio-0.1.0.dist-info/entry_points.txt

@@ -0,0 +1,2 @@
+[console_scripts]
+ber-studio = ber_equalization_studio.cli:main

ber_equalization_studio-0.1.0.dist-info/top_level.txt

@@ -0,0 +1 @@
+ber_equalization_studio