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

ber_equalization_studio/_legacy_backend/efficient_kan/__init__.py

@@ -0,0 +1,3 @@
+from .kan import KAN, KANLinear
+
+__all__ = ["KANLinear", "KAN"]

ber_equalization_studio/_legacy_backend/efficient_kan/kan.py

@@ -0,0 +1,218 @@
+import math
+
+import torch
+import torch.nn.functional as F
+
+
+class KANLinear(torch.nn.Module):
+    def __init__(
+        self,
+        in_features,
+        out_features,
+        grid_size=5,
+        spline_order=3,
+        scale_noise=0.1,
+        scale_base=1.0,
+        scale_spline=1.0,
+        enable_standalone_scale_spline=True,
+        base_activation=torch.nn.SiLU,
+        grid_eps=0.02,
+        grid_range=[-1, 1],
+    ):
+        super().__init__()
+        self.in_features = in_features
+        self.out_features = out_features
+        self.grid_size = grid_size
+        self.spline_order = spline_order
+
+        h = (grid_range[1] - grid_range[0]) / grid_size
+        grid = (
+            (torch.arange(-spline_order, grid_size + spline_order + 1) * h + grid_range[0])
+            .expand(in_features, -1)
+            .contiguous()
+        )
+        self.register_buffer("grid", grid)
+
+        self.base_weight = torch.nn.Parameter(torch.empty(out_features, in_features))
+        self.spline_weight = torch.nn.Parameter(
+            torch.empty(out_features, in_features, grid_size + spline_order)
+        )
+        if enable_standalone_scale_spline:
+            self.spline_scaler = torch.nn.Parameter(torch.empty(out_features, in_features))
+
+        self.scale_noise = scale_noise
+        self.scale_base = scale_base
+        self.scale_spline = scale_spline
+        self.enable_standalone_scale_spline = enable_standalone_scale_spline
+        self.base_activation = base_activation()
+        self.grid_eps = grid_eps
+
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        torch.nn.init.kaiming_uniform_(self.base_weight, a=math.sqrt(5) * self.scale_base)
+        with torch.no_grad():
+            noise = (
+                (torch.rand(self.grid_size + 1, self.in_features, self.out_features) - 0.5)
+                * self.scale_noise
+                / self.grid_size
+            )
+            self.spline_weight.data.copy_(
+                (self.scale_spline if not self.enable_standalone_scale_spline else 1.0)
+                * self.curve2coeff(self.grid.T[self.spline_order : -self.spline_order], noise)
+            )
+            if self.enable_standalone_scale_spline:
+                torch.nn.init.kaiming_uniform_(self.spline_scaler, a=math.sqrt(5) * self.scale_spline)
+
+    def b_splines(self, x: torch.Tensor):
+        assert x.dim() == 2 and x.size(1) == self.in_features
+
+        grid: torch.Tensor = self.grid
+        x = x.unsqueeze(-1)
+        bases = ((x >= grid[:, :-1]) & (x < grid[:, 1:])).to(x.dtype)
+        for k in range(1, self.spline_order + 1):
+            bases = (
+                (x - grid[:, : -(k + 1)])
+                / (grid[:, k:-1] - grid[:, : -(k + 1)])
+                * bases[:, :, :-1]
+            ) + (
+                (grid[:, k + 1 :] - x)
+                / (grid[:, k + 1 :] - grid[:, 1:-k])
+                * bases[:, :, 1:]
+            )
+
+        assert bases.size() == (
+            x.size(0),
+            self.in_features,
+            self.grid_size + self.spline_order,
+        )
+        return bases.contiguous()
+
+    def curve2coeff(self, x: torch.Tensor, y: torch.Tensor):
+        assert x.dim() == 2 and x.size(1) == self.in_features
+        assert y.size() == (x.size(0), self.in_features, self.out_features)
+
+        a = self.b_splines(x).transpose(0, 1)
+        b = y.transpose(0, 1)
+        solution = torch.linalg.lstsq(a, b).solution
+        result = solution.permute(2, 0, 1)
+
+        assert result.size() == (
+            self.out_features,
+            self.in_features,
+            self.grid_size + self.spline_order,
+        )
+        return result.contiguous()
+
+    @property
+    def scaled_spline_weight(self):
+        if not self.enable_standalone_scale_spline:
+            return self.spline_weight
+        return self.spline_weight * self.spline_scaler.unsqueeze(-1)
+
+    def forward(self, x: torch.Tensor):
+        assert x.size(-1) == self.in_features
+
+        original_shape = x.shape
+        x = x.reshape(-1, self.in_features)
+
+        base_output = F.linear(self.base_activation(x), self.base_weight)
+        spline_output = F.linear(
+            self.b_splines(x).view(x.size(0), -1),
+            self.scaled_spline_weight.view(self.out_features, -1),
+        )
+        output = base_output + spline_output
+        return output.reshape(*original_shape[:-1], self.out_features)
+
+    @torch.no_grad()
+    def update_grid(self, x: torch.Tensor, margin=0.01):
+        assert x.dim() == 2 and x.size(1) == self.in_features
+
+        batch = x.size(0)
+        splines = self.b_splines(x).permute(1, 0, 2)
+        orig_coeff = self.scaled_spline_weight.permute(1, 2, 0)
+        unreduced_spline_output = torch.bmm(splines, orig_coeff).permute(1, 0, 2)
+
+        x_sorted = torch.sort(x, dim=0)[0]
+        grid_adaptive = x_sorted[
+            torch.linspace(0, batch - 1, self.grid_size + 1, dtype=torch.int64, device=x.device)
+        ]
+
+        uniform_step = (x_sorted[-1] - x_sorted[0] + 2 * margin) / self.grid_size
+        grid_uniform = (
+            torch.arange(self.grid_size + 1, dtype=torch.float32, device=x.device).unsqueeze(1)
+            * uniform_step
+            + x_sorted[0]
+            - margin
+        )
+
+        grid = self.grid_eps * grid_uniform + (1 - self.grid_eps) * grid_adaptive
+        grid = torch.concatenate(
+            [
+                grid[:1] - uniform_step * torch.arange(self.spline_order, 0, -1, device=x.device).unsqueeze(1),
+                grid,
+                grid[-1:] + uniform_step * torch.arange(1, self.spline_order + 1, device=x.device).unsqueeze(1),
+            ],
+            dim=0,
+        )
+
+        self.grid.copy_(grid.T)
+        self.spline_weight.data.copy_(self.curve2coeff(x, unreduced_spline_output))
+
+    def regularization_loss(self, regularize_activation=1.0, regularize_entropy=1.0):
+        l1_fake = self.spline_weight.abs().mean(-1)
+        regularization_loss_activation = l1_fake.sum()
+        p = l1_fake / regularization_loss_activation
+        regularization_loss_entropy = -torch.sum(p * p.log())
+        return (
+            regularize_activation * regularization_loss_activation
+            + regularize_entropy * regularization_loss_entropy
+        )
+
+
+class KAN(torch.nn.Module):
+    def __init__(
+        self,
+        layers_hidden,
+        grid_size=5,
+        spline_order=3,
+        scale_noise=0.1,
+        scale_base=1.0,
+        scale_spline=1.0,
+        base_activation=torch.nn.SiLU,
+        grid_eps=0.02,
+        grid_range=[-1, 1],
+    ):
+        super().__init__()
+        self.grid_size = grid_size
+        self.spline_order = spline_order
+        self.layers = torch.nn.ModuleList()
+
+        for in_features, out_features in zip(layers_hidden, layers_hidden[1:]):
+            self.layers.append(
+                KANLinear(
+                    in_features,
+                    out_features,
+                    grid_size=grid_size,
+                    spline_order=spline_order,
+                    scale_noise=scale_noise,
+                    scale_base=scale_base,
+                    scale_spline=scale_spline,
+                    base_activation=base_activation,
+                    grid_eps=grid_eps,
+                    grid_range=grid_range,
+                )
+            )
+
+    def forward(self, x: torch.Tensor, update_grid=False):
+        for layer in self.layers:
+            if update_grid:
+                layer.update_grid(x)
+            x = layer(x)
+        return x
+
+    def regularization_loss(self, regularize_activation=1.0, regularize_entropy=1.0):
+        return sum(
+            layer.regularization_loss(regularize_activation, regularize_entropy)
+            for layer in self.layers
+        )

ber_equalization_studio/api.py

@@ -0,0 +1,348 @@
+from __future__ import annotations
+
+import itertools
+import re
+from dataclasses import fields, replace
+from pathlib import Path
+from typing import Any, Iterable
+
+from ber_equalization_studio.config import (
+    DataConfig,
+    EvaluationConfig,
+    ExperimentConfig,
+    ModelConfig,
+    OutputConfig,
+    StudioConfig,
+    TrainingConfig,
+)
+
+
+_CONFIG_SECTIONS = {
+    "data": DataConfig,
+    "model": ModelConfig,
+    "training": TrainingConfig,
+    "evaluation": EvaluationConfig,
+    "output": OutputConfig,
+    "experiment": ExperimentConfig,
+}
+
+_SHORTCUTS: dict[str, tuple[str, str]] = {
+    "data_dirs": ("data", "data_dirs"),
+    "context_k": ("data", "context_k"),
+    "max_files": ("data", "max_files"),
+    "max_test_files": ("data", "max_test_files"),
+    "train_portion": ("data", "train_portion"),
+    "val_portion_within_train": ("data", "val_portion_within_train"),
+    "randomize_file_split": ("data", "randomize_file_split"),
+    "split_seed": ("data", "split_seed"),
+    "input_dim": ("data", "input_dim"),
+    "power_normalize": ("data", "power_normalize"),
+    "model_name": ("model", "name"),
+    "model": ("model", "name"),
+    "hidden_dim": ("model", "hidden_dim"),
+    "dropout": ("model", "dropout"),
+    "mlp_layers": ("model", "mlp_layers"),
+    "fastkan_hidden_dim": ("model", "fastkan_hidden_dim"),
+    "fastkan_layers": ("model", "fastkan_layers"),
+    "fastkan_num_grids": ("model", "fastkan_num_grids"),
+    "fastkan_grids": ("model", "fastkan_num_grids"),
+    "efficient_kan_hidden_dim": ("model", "efficient_kan_hidden_dim"),
+    "efficient_kan_layers": ("model", "efficient_kan_layers"),
+    "efficient_kan_grid_size": ("model", "efficient_kan_grid_size"),
+    "efficient_kan_grid": ("model", "efficient_kan_grid_size"),
+    "epochs": ("training", "epochs"),
+    "lr": ("training", "learning_rate"),
+    "learning_rate": ("training", "learning_rate"),
+    "weight_decay": ("training", "weight_decay"),
+    "optimizer": ("training", "optimizer"),
+    "loss": ("training", "loss"),
+    "train_block_size": ("training", "train_block_size"),
+    "use_amp": ("training", "use_amp"),
+    "use_torch_compile": ("training", "use_torch_compile"),
+    "grad_clip_norm": ("training", "grad_clip_norm"),
+    "early_stopping": ("training", "early_stopping"),
+    "early_stopping_patience": ("training", "early_stopping_patience"),
+    "seed": ("training", "seed"),
+    "eval_batch_size": ("evaluation", "eval_batch_size"),
+    "eval_test_during_training": ("evaluation", "eval_test_during_training"),
+    "test_ber_every": ("evaluation", "test_ber_every"),
+    "compute_per_file_metrics": ("evaluation", "compute_per_file_metrics"),
+    "ber_scale_search": ("evaluation", "ber_scale_search"),
+    "out_dir": ("output", "out_dir"),
+    "experiment_name": ("output", "experiment_name"),
+    "save_checkpoint": ("output", "save_checkpoint"),
+    "save_history": ("output", "save_history"),
+    "save_plots": ("output", "save_plots"),
+    "plot_backend": ("output", "plot_backend"),
+    "models": ("experiment", "models"),
+    "tags": ("experiment", "tags"),
+    "notes": ("experiment", "notes"),
+}
+
+
+def _field_index() -> dict[str, tuple[str, str] | None]:
+    index: dict[str, tuple[str, str] | None] = {}
+    for section, section_type in _CONFIG_SECTIONS.items():
+        for field in fields(section_type):
+            if field.name in index:
+                index[field.name] = None
+            else:
+                index[field.name] = (section, field.name)
+    return index
+
+
+_FIELD_INDEX = _field_index()
+
+
+def _normalize_value(section: str, key: str, value: Any) -> Any:
+    if section == "data" and key == "data_dirs":
+        if isinstance(value, (str, Path)):
+            return [Path(value)]
+        return [Path(item) for item in value]
+    if section == "output" and key == "out_dir":
+        return Path(value)
+    if section == "experiment" and key == "models":
+        if isinstance(value, str):
+            return [item.strip() for item in value.split(",") if item.strip()]
+        return list(value)
+    return value
+
+
+def _locate_update(key: str) -> tuple[str, str]:
+    if "." in key:
+        section, field = key.split(".", 1)
+        if section not in _CONFIG_SECTIONS and section != "root":
+            raise KeyError(f"Unknown config section: {section}")
+        return section, field
+    if key == "device":
+        return "root", "device"
+    if key in _SHORTCUTS:
+        return _SHORTCUTS[key]
+    inferred = _FIELD_INDEX.get(key)
+    if inferred is None:
+        raise KeyError(f"Ambiguous or unknown config field: {key}. Use dotted form, for example 'training.{key}'.")
+    return inferred
+
+
+def _apply_updates(config: StudioConfig, updates: dict[str, Any]) -> StudioConfig:
+    section_updates: dict[str, dict[str, Any]] = {section: {} for section in _CONFIG_SECTIONS}
+    root_updates: dict[str, Any] = {}
+    for key, value in updates.items():
+        if value is None:
+            continue
+        section, field = _locate_update(key)
+        if section == "root":
+            root_updates[field] = value
+            continue
+        section_updates[section][field] = _normalize_value(section, field, value)
+
+    next_config = config
+    for section, values in section_updates.items():
+        if not values:
+            continue
+        current_section = getattr(next_config, section)
+        next_config = replace(next_config, **{section: replace(current_section, **values)})
+    if root_updates:
+        next_config = replace(next_config, **root_updates)
+    return next_config
+
+
+def _slug(value: str) -> str:
+    slug = re.sub(r"[^a-zA-Z0-9_.-]+", "_", value.strip())
+    return slug.strip("_") or "run"
+
+
+def _grid_items(grid: dict[str, Iterable[Any]]) -> list[dict[str, Any]]:
+    keys = list(grid)
+    values = [list(grid[key]) for key in keys]
+    return [dict(zip(keys, combination, strict=True)) for combination in itertools.product(*values)]
+
+
+def _pd():
+    import pandas as pd
+
+    return pd
+
+
+class RunResult:
+    """Notebook-friendly wrapper around a single experiment run."""
+
+    def __init__(self, payload: dict[str, Any], config: StudioConfig):
+        self.payload = payload
+        self.config = config
+
+    @property
+    def run_dir(self) -> Path:
+        return Path(self.payload["run_dir"])
+
+    @property
+    def results_csv(self) -> Path:
+        return Path(self.payload["results_csv"])
+
+    @property
+    def results(self):
+        from ber_equalization_studio.results import flatten_metrics, load_result_table
+
+        if self.results_csv.exists():
+            return load_result_table(self.results_csv)
+        return _pd().DataFrame([flatten_metrics(row) for row in self.payload.get("results", [])])
+
+    @property
+    def histories(self) -> dict[str, dict[str, Any]]:
+        return self.payload.get("histories", {})
+
+    def best(self, by: str = "equalized_ber", ascending: bool | None = None) -> pd.Series:
+        df = self.results
+        if by not in df.columns:
+            raise KeyError(f"Column not found in results: {by}")
+        if ascending is None:
+            ascending = by not in {"improvement_rel", "improvement_db", "accuracy", "samples_per_sec"}
+        return df.sort_values(by, ascending=ascending).iloc[0]
+
+    def compare(self, sort_by: str = "equalized_ber") -> pd.DataFrame:
+        from ber_equalization_studio.results import compare_results
+
+        return compare_results(self.results_csv, sort_by=sort_by)
+
+    def history(self, model: str | None = None) -> dict[str, Any]:
+        if model is None:
+            if len(self.histories) != 1:
+                raise ValueError("Pass model=... when a run contains more than one model.")
+            return next(iter(self.histories.values()))
+        return self.histories[model]
+
+    def plot_history(self, model: str | None = None, metric: str = "ber"):
+        from ber_equalization_studio.visualization import history_figure
+
+        history = self.history(model)
+        model_name = model or next(iter(self.histories))
+        return history_figure(history, model_name, metric=metric)
+
+    def plot_comparison(
+        self,
+        x: str = "model_type",
+        y: str = "equalized_ber",
+        color: str | None = None,
+        kind: str = "bar",
+    ):
+        from ber_equalization_studio.visualization import comparison_figure
+
+        return comparison_figure(self.results, x=x, y=y, color=color, kind=kind)
+
+    def _repr_html_(self) -> str:
+        return self.results._repr_html_()
+
+
+class StudyResult:
+    """Result object for a sequence of runs, usually produced by Studio.sweep()."""
+
+    def __init__(self, runs: list[RunResult], parameters: list[dict[str, Any]]):
+        self.runs = runs
+        self.parameters = parameters
+
+    @property
+    def results(self):
+        pd = _pd()
+        frames: list[pd.DataFrame] = []
+        for index, run in enumerate(self.runs):
+            df = run.results.copy()
+            df["run_dir"] = str(run.run_dir)
+            for key, value in self.parameters[index].items():
+                df[key] = value
+            frames.append(df)
+        if not frames:
+            return pd.DataFrame()
+        return pd.concat(frames, ignore_index=True)
+
+    def best(self, by: str = "equalized_ber", ascending: bool | None = None) -> pd.Series:
+        df = self.results
+        if by not in df.columns:
+            raise KeyError(f"Column not found in results: {by}")
+        if ascending is None:
+            ascending = by not in {"improvement_rel", "improvement_db", "accuracy", "samples_per_sec"}
+        return df.sort_values(by, ascending=ascending).iloc[0]
+
+    def plot_tradeoff(
+        self,
+        x: str = "trainable_params",
+        y: str = "equalized_ber",
+        color: str | None = "model_type",
+    ):
+        from ber_equalization_studio.visualization import comparison_figure
+
+        return comparison_figure(self.results, x=x, y=y, color=color, kind="scatter")
+
+    def _repr_html_(self) -> str:
+        return self.results._repr_html_()
+
+
+class Studio:
+    """High-level research interface for notebooks and scripts."""
+
+    def __init__(
+        self,
+        base_config: StudioConfig | None = None,
+        legacy_path: str | Path | None = None,
+        **defaults: Any,
+    ):
+        self.base_config = _apply_updates(base_config or StudioConfig(), defaults)
+        self.legacy_path = Path(legacy_path) if legacy_path is not None else None
+
+    def config(self, **updates: Any) -> StudioConfig:
+        return _apply_updates(self.base_config, updates)
+
+    def models(self) -> pd.DataFrame:
+        from ber_equalization_studio.models import available_models
+
+        rows = [{"model": name, "description": note} for name, note in available_models().items()]
+        return _pd().DataFrame(rows)
+
+    def run(
+        self,
+        name: str = "experiment",
+        models: str | Iterable[str] | None = None,
+        **updates: Any,
+    ) -> RunResult:
+        if models is not None:
+            updates["models"] = models
+        updates["experiment_name"] = name
+        config = self.config(**updates)
+        from ber_equalization_studio.experiment import ExperimentRunner
+
+        payload = ExperimentRunner(config, legacy_path=self.legacy_path).run()
+        return RunResult(payload, config)
+
+    def sweep(
+        self,
+        name: str,
+        grid: dict[str, Iterable[Any]],
+        models: str | Iterable[str] | None = None,
+        **updates: Any,
+    ) -> StudyResult:
+        runs: list[RunResult] = []
+        parameters = _grid_items(grid)
+        for index, params in enumerate(parameters, start=1):
+            parts = [f"{_slug(key)}-{_slug(str(value))}" for key, value in params.items()]
+            run_name = _slug("_".join([name, f"{index:03d}", *parts]))
+            run_updates = {**updates, **params}
+            runs.append(self.run(name=run_name, models=models, **run_updates))
+        return StudyResult(runs, parameters)
+
+
+def run(
+    name: str = "experiment",
+    models: str | Iterable[str] | None = None,
+    legacy_path: str | Path | None = None,
+    **updates: Any,
+) -> RunResult:
+    return Studio(legacy_path=legacy_path).run(name=name, models=models, **updates)
+
+
+def sweep(
+    name: str,
+    grid: dict[str, Iterable[Any]],
+    models: str | Iterable[str] | None = None,
+    legacy_path: str | Path | None = None,
+    **updates: Any,
+) -> StudyResult:
+    return Studio(legacy_path=legacy_path).sweep(name=name, grid=grid, models=models, **updates)

ber_equalization_studio/cli.py

@@ -0,0 +1,92 @@
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+
+from ber_equalization_studio.config import DataConfig, ExperimentConfig, ModelConfig, OutputConfig, StudioConfig, TrainingConfig
+
+
+def _split_models(value: str) -> list[str]:
+    return [item.strip() for item in value.split(",") if item.strip()]
+
+
+def build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(prog="ber-studio", description="BER equalization experiment studio")
+    sub = parser.add_subparsers(dest="command", required=True)
+
+    run = sub.add_parser("run", help="Train one or more equalizer models")
+    run.add_argument("--models", default="complex_fastkan", help="Comma-separated model names")
+    run.add_argument("--data-dir", action="append", type=Path, default=None, help="Directory with Symbols_1m_1ch_PR_*.csv")
+    run.add_argument("--out-dir", type=Path, default=Path("studio_outputs"))
+    run.add_argument("--name", default="experiment")
+    run.add_argument("--legacy-path", type=Path, default=None)
+    run.add_argument("--epochs", type=int, default=250)
+    run.add_argument("--lr", type=float, default=1e-3)
+    run.add_argument("--context-k", type=int, default=32)
+    run.add_argument("--max-files", type=int, default=64)
+    run.add_argument("--max-test-files", type=int, default=None)
+    run.add_argument("--device", default=None)
+    run.add_argument("--fastkan-hidden", type=int, default=96)
+    run.add_argument("--fastkan-grids", type=int, default=8)
+    run.add_argument("--efficient-kan-hidden", type=int, default=128)
+    run.add_argument("--efficient-kan-grid", type=int, default=8)
+    run.add_argument("--no-per-file", action="store_true")
+    run.add_argument("--no-plots", action="store_true")
+
+    sub.add_parser("models", help="List available model names")
+
+    compare = sub.add_parser("compare", help="Compare result CSV files below a run directory")
+    compare.add_argument("path", type=Path)
+    compare.add_argument("--sort-by", default="equalized_ber")
+
+    return parser
+
+
+def main(argv: list[str] | None = None) -> None:
+    parser = build_parser()
+    args = parser.parse_args(argv)
+
+    if args.command == "models":
+        from ber_equalization_studio.models import available_models
+
+        for name, note in available_models().items():
+            print(f"{name:32s} {note}")
+        return
+
+    if args.command == "compare":
+        from ber_equalization_studio.results import compare_results
+
+        df = compare_results(args.path, sort_by=args.sort_by)
+        print(df.to_string(index=False))
+        return
+
+    data_dirs = args.data_dir if args.data_dir else DataConfig().data_dirs
+    config = StudioConfig(
+        data=DataConfig(
+            data_dirs=data_dirs,
+            max_files=args.max_files,
+            context_k=args.context_k,
+            max_test_files=args.max_test_files,
+        ),
+        model=ModelConfig(
+            fastkan_hidden_dim=args.fastkan_hidden,
+            fastkan_num_grids=args.fastkan_grids,
+            efficient_kan_hidden_dim=args.efficient_kan_hidden,
+            efficient_kan_grid_size=args.efficient_kan_grid,
+        ),
+        training=TrainingConfig(epochs=args.epochs, learning_rate=args.lr),
+        output=OutputConfig(out_dir=args.out_dir, experiment_name=args.name, save_plots=not args.no_plots),
+        experiment=ExperimentConfig(models=_split_models(args.models)),
+        device=args.device or StudioConfig().device,
+    )
+    if args.no_per_file:
+        config.evaluation.compute_per_file_metrics = False
+
+    from ber_equalization_studio.experiment import ExperimentRunner
+
+    result = ExperimentRunner(config, legacy_path=args.legacy_path).run()
+    print(f"Saved results: {result['results_csv']}")
+
+
+if __name__ == "__main__":
+    main()