ber-equalization-studio 0.1.2__tar.gz → 0.1.3__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (26) hide show
  1. {ber_equalization_studio-0.1.2 → ber_equalization_studio-0.1.3}/PKG-INFO +2 -1
  2. {ber_equalization_studio-0.1.2 → ber_equalization_studio-0.1.3}/pyproject.toml +5 -1
  3. {ber_equalization_studio-0.1.2 → ber_equalization_studio-0.1.3}/src/ber_equalization_studio/__init__.py +11 -16
  4. ber_equalization_studio-0.1.3/src/ber_equalization_studio/__init__.pyi +20 -0
  5. {ber_equalization_studio-0.1.2 → ber_equalization_studio-0.1.3}/src/ber_equalization_studio/api.py +106 -13
  6. {ber_equalization_studio-0.1.2 → ber_equalization_studio-0.1.3}/src/ber_equalization_studio/config.py +31 -0
  7. {ber_equalization_studio-0.1.2 → ber_equalization_studio-0.1.3}/src/ber_equalization_studio/models.py +8 -0
  8. ber_equalization_studio-0.1.3/src/ber_equalization_studio/py.typed +1 -0
  9. {ber_equalization_studio-0.1.2 → ber_equalization_studio-0.1.3}/src/ber_equalization_studio/results.py +10 -0
  10. {ber_equalization_studio-0.1.2 → ber_equalization_studio-0.1.3}/src/ber_equalization_studio.egg-info/PKG-INFO +2 -1
  11. {ber_equalization_studio-0.1.2 → ber_equalization_studio-0.1.3}/src/ber_equalization_studio.egg-info/SOURCES.txt +2 -0
  12. {ber_equalization_studio-0.1.2 → ber_equalization_studio-0.1.3}/README.md +0 -0
  13. {ber_equalization_studio-0.1.2 → ber_equalization_studio-0.1.3}/setup.cfg +0 -0
  14. {ber_equalization_studio-0.1.2 → ber_equalization_studio-0.1.3}/src/ber_equalization_studio/_legacy_backend/__init__.py +0 -0
  15. {ber_equalization_studio-0.1.2 → ber_equalization_studio-0.1.3}/src/ber_equalization_studio/_legacy_backend/ber_equalization.py +0 -0
  16. {ber_equalization_studio-0.1.2 → ber_equalization_studio-0.1.3}/src/ber_equalization_studio/_legacy_backend/efficient_kan/__init__.py +0 -0
  17. {ber_equalization_studio-0.1.2 → ber_equalization_studio-0.1.3}/src/ber_equalization_studio/_legacy_backend/efficient_kan/kan.py +0 -0
  18. {ber_equalization_studio-0.1.2 → ber_equalization_studio-0.1.3}/src/ber_equalization_studio/cli.py +0 -0
  19. {ber_equalization_studio-0.1.2 → ber_equalization_studio-0.1.3}/src/ber_equalization_studio/data.py +0 -0
  20. {ber_equalization_studio-0.1.2 → ber_equalization_studio-0.1.3}/src/ber_equalization_studio/experiment.py +0 -0
  21. {ber_equalization_studio-0.1.2 → ber_equalization_studio-0.1.3}/src/ber_equalization_studio/legacy.py +0 -0
  22. {ber_equalization_studio-0.1.2 → ber_equalization_studio-0.1.3}/src/ber_equalization_studio/visualization.py +0 -0
  23. {ber_equalization_studio-0.1.2 → ber_equalization_studio-0.1.3}/src/ber_equalization_studio.egg-info/dependency_links.txt +0 -0
  24. {ber_equalization_studio-0.1.2 → ber_equalization_studio-0.1.3}/src/ber_equalization_studio.egg-info/entry_points.txt +0 -0
  25. {ber_equalization_studio-0.1.2 → ber_equalization_studio-0.1.3}/src/ber_equalization_studio.egg-info/requires.txt +0 -0
  26. {ber_equalization_studio-0.1.2 → ber_equalization_studio-0.1.3}/src/ber_equalization_studio.egg-info/top_level.txt +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: ber-equalization-studio
3
- Version: 0.1.2
3
+ Version: 0.1.3
4
4
  Summary: Research studio for BER equalization experiments in nonlinear optical links.
5
5
  Keywords: ber,equalization,optical-communications,photonics,pytorch
6
6
  Classifier: Development Status :: 3 - Alpha
@@ -9,6 +9,7 @@ Classifier: Programming Language :: Python :: 3
9
9
  Classifier: Programming Language :: Python :: 3.10
10
10
  Classifier: Programming Language :: Python :: 3.11
11
11
  Classifier: Programming Language :: Python :: 3.12
12
+ Classifier: Programming Language :: Python :: 3.13
12
13
  Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
13
14
  Classifier: Topic :: Scientific/Engineering :: Physics
14
15
  Requires-Python: >=3.10
@@ -1,6 +1,6 @@
1
1
  [project]
2
2
  name = "ber-equalization-studio"
3
- version = "0.1.2"
3
+ version = "0.1.3"
4
4
  description = "Research studio for BER equalization experiments in nonlinear optical links."
5
5
  readme = "README.md"
6
6
  requires-python = ">=3.10"
@@ -20,6 +20,7 @@ classifiers = [
20
20
  "Programming Language :: Python :: 3.10",
21
21
  "Programming Language :: Python :: 3.11",
22
22
  "Programming Language :: Python :: 3.12",
23
+ "Programming Language :: Python :: 3.13",
23
24
  "Topic :: Scientific/Engineering :: Artificial Intelligence",
24
25
  "Topic :: Scientific/Engineering :: Physics",
25
26
  ]
@@ -52,6 +53,9 @@ build-backend = "setuptools.build_meta"
52
53
  [tool.setuptools.packages.find]
53
54
  where = ["src"]
54
55
 
56
+ [tool.setuptools.package-data]
57
+ ber_equalization_studio = ["py.typed", "__init__.pyi"]
58
+
55
59
  [[tool.uv.index]]
56
60
  name = "pytorch-cu130"
57
61
  url = "https://download.pytorch.org/whl/cu130"
@@ -1,5 +1,13 @@
1
- """Convenient research API for BER equalization experiments."""
1
+ """Notebook-friendly API for BER equalization experiments.
2
2
 
3
+ The top-level package exports the objects that are most useful in notebooks:
4
+ `Studio` for running experiments, config dataclasses for reproducible setups,
5
+ model registry helpers, and result table helpers.
6
+ """
7
+
8
+ from typing import Any
9
+
10
+ from ber_equalization_studio.api import RunResult, Studio, StudyResult, run, sweep
3
11
  from ber_equalization_studio.config import (
4
12
  DataConfig,
5
13
  EvaluationConfig,
@@ -9,6 +17,7 @@ from ber_equalization_studio.config import (
9
17
  StudioConfig,
10
18
  TrainingConfig,
11
19
  )
20
+ from ber_equalization_studio.models import available_models, build_model
12
21
 
13
22
  __all__ = [
14
23
  "DataConfig",
@@ -32,27 +41,13 @@ __all__ = [
32
41
  ]
33
42
 
34
43
 
35
- def __getattr__(name):
44
+ def __getattr__(name: str) -> Any:
36
45
  if name in {"ExperimentRunner", "run_experiment"}:
37
46
  from ber_equalization_studio.experiment import ExperimentRunner, run_experiment
38
47
 
39
48
  return {"ExperimentRunner": ExperimentRunner, "run_experiment": run_experiment}[name]
40
- if name in {"available_models", "build_model"}:
41
- from ber_equalization_studio.models import available_models, build_model
42
-
43
- return {"available_models": available_models, "build_model": build_model}[name]
44
49
  if name in {"compare_results", "load_result_table"}:
45
50
  from ber_equalization_studio.results import compare_results, load_result_table
46
51
 
47
52
  return {"compare_results": compare_results, "load_result_table": load_result_table}[name]
48
- if name in {"RunResult", "Studio", "StudyResult", "run", "sweep"}:
49
- from ber_equalization_studio.api import RunResult, Studio, StudyResult, run, sweep
50
-
51
- return {
52
- "RunResult": RunResult,
53
- "Studio": Studio,
54
- "StudyResult": StudyResult,
55
- "run": run,
56
- "sweep": sweep,
57
- }[name]
58
53
  raise AttributeError(name)
@@ -0,0 +1,20 @@
1
+ from ber_equalization_studio.api import RunResult as RunResult
2
+ from ber_equalization_studio.api import Studio as Studio
3
+ from ber_equalization_studio.api import StudyResult as StudyResult
4
+ from ber_equalization_studio.api import run as run
5
+ from ber_equalization_studio.api import sweep as sweep
6
+ from ber_equalization_studio.config import DataConfig as DataConfig
7
+ from ber_equalization_studio.config import EvaluationConfig as EvaluationConfig
8
+ from ber_equalization_studio.config import ExperimentConfig as ExperimentConfig
9
+ from ber_equalization_studio.config import ModelConfig as ModelConfig
10
+ from ber_equalization_studio.config import OutputConfig as OutputConfig
11
+ from ber_equalization_studio.config import StudioConfig as StudioConfig
12
+ from ber_equalization_studio.config import TrainingConfig as TrainingConfig
13
+ from ber_equalization_studio.experiment import ExperimentRunner as ExperimentRunner
14
+ from ber_equalization_studio.experiment import run_experiment as run_experiment
15
+ from ber_equalization_studio.models import available_models as available_models
16
+ from ber_equalization_studio.models import build_model as build_model
17
+ from ber_equalization_studio.results import compare_results as compare_results
18
+ from ber_equalization_studio.results import load_result_table as load_result_table
19
+
20
+ __all__: list[str]
@@ -4,7 +4,7 @@ import itertools
4
4
  import re
5
5
  from dataclasses import fields, replace
6
6
  from pathlib import Path
7
- from typing import Any, Iterable
7
+ from typing import TYPE_CHECKING, Any, Iterable
8
8
 
9
9
  from ber_equalization_studio.config import (
10
10
  DataConfig,
@@ -16,6 +16,9 @@ from ber_equalization_studio.config import (
16
16
  TrainingConfig,
17
17
  )
18
18
 
19
+ if TYPE_CHECKING:
20
+ import pandas as pd
21
+
19
22
 
20
23
  _CONFIG_SECTIONS = {
21
24
  "data": DataConfig,
@@ -165,7 +168,12 @@ def _pd():
165
168
 
166
169
 
167
170
  class RunResult:
168
- """Notebook-friendly wrapper around a single experiment run."""
171
+ """Notebook-friendly wrapper around a completed experiment run.
172
+
173
+ Use `results` for the metrics table, `best()` to select the best model row,
174
+ `history()` for raw training curves, and the plotting helpers for inline
175
+ Plotly figures.
176
+ """
169
177
 
170
178
  def __init__(self, payload: dict[str, Any], config: StudioConfig):
171
179
  self.payload = payload
@@ -173,14 +181,20 @@ class RunResult:
173
181
 
174
182
  @property
175
183
  def run_dir(self) -> Path:
184
+ """Directory containing config, metrics, checkpoints, histories, and plots."""
185
+
176
186
  return Path(self.payload["run_dir"])
177
187
 
178
188
  @property
179
189
  def results_csv(self) -> Path:
190
+ """Path to the CSV table with one row per evaluated model."""
191
+
180
192
  return Path(self.payload["results_csv"])
181
193
 
182
194
  @property
183
- def results(self):
195
+ def results(self) -> "pd.DataFrame":
196
+ """Return the run metrics as a pandas DataFrame."""
197
+
184
198
  from ber_equalization_studio.results import flatten_metrics, load_result_table
185
199
 
186
200
  if self.results_csv.exists():
@@ -189,9 +203,23 @@ class RunResult:
189
203
 
190
204
  @property
191
205
  def histories(self) -> dict[str, dict[str, Any]]:
206
+ """Training histories keyed by canonical model name."""
207
+
192
208
  return self.payload.get("histories", {})
193
209
 
194
- def best(self, by: str = "equalized_ber", ascending: bool | None = None) -> pd.Series:
210
+ def best(self, by: str = "equalized_ber", ascending: bool | None = None) -> "pd.Series":
211
+ """Return the best result row sorted by a metric column.
212
+
213
+ Parameters
214
+ ----------
215
+ by:
216
+ Result column to rank by, for example `equalized_ber`,
217
+ `improvement_db`, or `samples_per_sec`.
218
+ ascending:
219
+ Sort direction. When omitted, lower is better for error metrics and
220
+ higher is better for improvement, accuracy, and throughput metrics.
221
+ """
222
+
195
223
  df = self.results
196
224
  if by not in df.columns:
197
225
  raise KeyError(f"Column not found in results: {by}")
@@ -199,19 +227,28 @@ class RunResult:
199
227
  ascending = by not in {"improvement_rel", "improvement_db", "accuracy", "samples_per_sec"}
200
228
  return df.sort_values(by, ascending=ascending).iloc[0]
201
229
 
202
- def compare(self, sort_by: str = "equalized_ber") -> pd.DataFrame:
230
+ def compare(self, sort_by: str = "equalized_ber") -> "pd.DataFrame":
231
+ """Return a publication-friendly comparison table for this run."""
232
+
203
233
  from ber_equalization_studio.results import compare_results
204
234
 
205
235
  return compare_results(self.results_csv, sort_by=sort_by)
206
236
 
207
237
  def history(self, model: str | None = None) -> dict[str, Any]:
238
+ """Return the saved training history for one model.
239
+
240
+ Pass `model` when the run evaluated more than one model.
241
+ """
242
+
208
243
  if model is None:
209
244
  if len(self.histories) != 1:
210
245
  raise ValueError("Pass model=... when a run contains more than one model.")
211
246
  return next(iter(self.histories.values()))
212
247
  return self.histories[model]
213
248
 
214
- def plot_history(self, model: str | None = None, metric: str = "ber"):
249
+ def plot_history(self, model: str | None = None, metric: str = "ber") -> Any:
250
+ """Build an inline training-history figure for `ber` or `loss`."""
251
+
215
252
  from ber_equalization_studio.visualization import history_figure
216
253
 
217
254
  history = self.history(model)
@@ -224,7 +261,9 @@ class RunResult:
224
261
  y: str = "equalized_ber",
225
262
  color: str | None = None,
226
263
  kind: str = "bar",
227
- ):
264
+ ) -> Any:
265
+ """Build an inline comparison figure from the result table."""
266
+
228
267
  from ber_equalization_studio.visualization import comparison_figure
229
268
 
230
269
  return comparison_figure(self.results, x=x, y=y, color=color, kind=kind)
@@ -234,14 +273,16 @@ class RunResult:
234
273
 
235
274
 
236
275
  class StudyResult:
237
- """Result object for a sequence of runs, usually produced by Studio.sweep()."""
276
+ """Result object for a sequence of runs produced by `Studio.sweep()`."""
238
277
 
239
278
  def __init__(self, runs: list[RunResult], parameters: list[dict[str, Any]]):
240
279
  self.runs = runs
241
280
  self.parameters = parameters
242
281
 
243
282
  @property
244
- def results(self):
283
+ def results(self) -> "pd.DataFrame":
284
+ """Concatenate all sweep run result tables and attach sweep parameters."""
285
+
245
286
  pd = _pd()
246
287
  frames: list[pd.DataFrame] = []
247
288
  for index, run in enumerate(self.runs):
@@ -254,7 +295,9 @@ class StudyResult:
254
295
  return pd.DataFrame()
255
296
  return pd.concat(frames, ignore_index=True)
256
297
 
257
- def best(self, by: str = "equalized_ber", ascending: bool | None = None) -> pd.Series:
298
+ def best(self, by: str = "equalized_ber", ascending: bool | None = None) -> "pd.Series":
299
+ """Return the best row across all runs in the sweep."""
300
+
258
301
  df = self.results
259
302
  if by not in df.columns:
260
303
  raise KeyError(f"Column not found in results: {by}")
@@ -267,7 +310,9 @@ class StudyResult:
267
310
  x: str = "trainable_params",
268
311
  y: str = "equalized_ber",
269
312
  color: str | None = "model_type",
270
- ):
313
+ ) -> Any:
314
+ """Build a scatter plot for accuracy/complexity tradeoff analysis."""
315
+
271
316
  from ber_equalization_studio.visualization import comparison_figure
272
317
 
273
318
  return comparison_figure(self.results, x=x, y=y, color=color, kind="scatter")
@@ -277,7 +322,14 @@ class StudyResult:
277
322
 
278
323
 
279
324
  class Studio:
280
- """High-level research interface for notebooks and scripts."""
325
+ """High-level research interface for notebooks and scripts.
326
+
327
+ `Studio` keeps a reusable base configuration and lets each call override
328
+ only the fields that matter for the current experiment. Short aliases such
329
+ as `data_dirs`, `out_dir`, `epochs`, `lr`, `context_k`, and `max_test_files`
330
+ are accepted in addition to dotted config keys like
331
+ `training.early_stopping_patience`.
332
+ """
281
333
 
282
334
  def __init__(
283
335
  self,
@@ -285,13 +337,30 @@ class Studio:
285
337
  legacy_path: str | Path | None = None,
286
338
  **defaults: Any,
287
339
  ):
340
+ """Create a studio with optional default configuration overrides.
341
+
342
+ Parameters
343
+ ----------
344
+ base_config:
345
+ Existing structured config to use as the starting point.
346
+ legacy_path:
347
+ Optional path to a custom `ber_equalization.py` backend.
348
+ **defaults:
349
+ Default config overrides. Examples: `data_dirs=["../test"]`,
350
+ `out_dir="notebook_runs"`, `device="cuda"`, `epochs=10`.
351
+ """
352
+
288
353
  self.base_config = _apply_updates(base_config or StudioConfig(), defaults)
289
354
  self.legacy_path = Path(legacy_path) if legacy_path is not None else None
290
355
 
291
356
  def config(self, **updates: Any) -> StudioConfig:
357
+ """Return a `StudioConfig` with temporary overrides applied."""
358
+
292
359
  return _apply_updates(self.base_config, updates)
293
360
 
294
- def models(self) -> pd.DataFrame:
361
+ def models(self) -> "pd.DataFrame":
362
+ """List available model names and short architecture descriptions."""
363
+
295
364
  from ber_equalization_studio.models import available_models
296
365
 
297
366
  rows = [{"model": name, "description": note} for name, note in available_models().items()]
@@ -303,6 +372,20 @@ class Studio:
303
372
  models: str | Iterable[str] | None = None,
304
373
  **updates: Any,
305
374
  ) -> RunResult:
375
+ """Run one experiment and return a notebook-friendly `RunResult`.
376
+
377
+ Parameters
378
+ ----------
379
+ name:
380
+ Experiment folder name inside `out_dir`.
381
+ models:
382
+ One model name, a comma-separated string, or an iterable of model
383
+ names. When omitted, `StudioConfig.experiment.models` is used.
384
+ **updates:
385
+ Per-run config overrides such as `epochs=5`, `lr=1e-3`,
386
+ `context_k=32`, or `evaluation.ber_scale_steps=20`.
387
+ """
388
+
306
389
  if models is not None:
307
390
  updates["models"] = models
308
391
  updates["experiment_name"] = name
@@ -319,6 +402,12 @@ class Studio:
319
402
  models: str | Iterable[str] | None = None,
320
403
  **updates: Any,
321
404
  ) -> StudyResult:
405
+ """Run a Cartesian sweep over config values.
406
+
407
+ `grid` maps config aliases or dotted keys to candidate values, for
408
+ example `{"context_k": [16, 32], "lr": [1e-3, 3e-4]}`.
409
+ """
410
+
322
411
  runs: list[RunResult] = []
323
412
  parameters = _grid_items(grid)
324
413
  for index, params in enumerate(parameters, start=1):
@@ -335,6 +424,8 @@ def run(
335
424
  legacy_path: str | Path | None = None,
336
425
  **updates: Any,
337
426
  ) -> RunResult:
427
+ """Convenience wrapper for `Studio(...).run(...)`."""
428
+
338
429
  return Studio(legacy_path=legacy_path).run(name=name, models=models, **updates)
339
430
 
340
431
 
@@ -345,4 +436,6 @@ def sweep(
345
436
  legacy_path: str | Path | None = None,
346
437
  **updates: Any,
347
438
  ) -> StudyResult:
439
+ """Convenience wrapper for `Studio(...).sweep(...)`."""
440
+
348
441
  return Studio(legacy_path=legacy_path).sweep(name=name, grid=grid, models=models, **updates)
@@ -20,6 +20,13 @@ def _default_device() -> str:
20
20
 
21
21
  @dataclass(slots=True)
22
22
  class DataConfig:
23
+ """Dataset and windowing settings for BER equalization.
24
+
25
+ `data_dirs` is searched for symbol CSV files, `context_k` controls the
26
+ number of neighboring symbols on each side of the center sample, and
27
+ `seq_len` is computed as `2 * context_k + 1`.
28
+ """
29
+
23
30
  data_dirs: list[Path] = field(
24
31
  default_factory=lambda: [Path("symbols_new"), Path("Symbols_1m_1ch_PR"), Path(".")]
25
32
  )
@@ -39,11 +46,19 @@ class DataConfig:
39
46
 
40
47
  @property
41
48
  def seq_len(self) -> int:
49
+ """Window length in symbols: `2 * context_k + 1`."""
50
+
42
51
  return 2 * self.context_k + 1
43
52
 
44
53
 
45
54
  @dataclass(slots=True)
46
55
  class ModelConfig:
56
+ """Architecture hyperparameters shared by the model registry.
57
+
58
+ The active architecture is selected by `name`. Other fields tune model
59
+ families such as MLP, temporal CNN/TCN, EfficientKAN, and FastKAN variants.
60
+ """
61
+
47
62
  name: str = "complex_fastkan"
48
63
  hidden_dim: int = 96
49
64
  dropout: float = 0.2
@@ -70,6 +85,8 @@ class ModelConfig:
70
85
 
71
86
  @dataclass(slots=True)
72
87
  class TrainingConfig:
88
+ """Optimizer, loss, scheduling, precision, and early-stopping settings."""
89
+
73
90
  epochs: int = 250
74
91
  learning_rate: float = 1e-3
75
92
  weight_decay: float = 0.0
@@ -97,6 +114,8 @@ class TrainingConfig:
97
114
 
98
115
  @dataclass(slots=True)
99
116
  class EvaluationConfig:
117
+ """Evaluation settings for BER, per-file metrics, and timing benchmarks."""
118
+
100
119
  eval_batch_size: int = 65536
101
120
  eval_test_during_training: bool = False
102
121
  test_ber_every: int = 10
@@ -114,6 +133,8 @@ class EvaluationConfig:
114
133
 
115
134
  @dataclass(slots=True)
116
135
  class OutputConfig:
136
+ """Output folder, checkpoint, history, and plot settings for a run."""
137
+
117
138
  out_dir: Path = Path("studio_outputs")
118
139
  experiment_name: str = "experiment"
119
140
  save_checkpoint: bool = True
@@ -126,11 +147,15 @@ class OutputConfig:
126
147
 
127
148
  @property
128
149
  def run_dir(self) -> Path:
150
+ """Concrete directory for the current experiment."""
151
+
129
152
  return self.out_dir / self.experiment_name
130
153
 
131
154
 
132
155
  @dataclass(slots=True)
133
156
  class ExperimentConfig:
157
+ """Experiment-level model list and metadata."""
158
+
134
159
  models: list[str] = field(default_factory=lambda: ["complex_fastkan"])
135
160
  tags: dict[str, str] = field(default_factory=dict)
136
161
  notes: str = ""
@@ -138,6 +163,8 @@ class ExperimentConfig:
138
163
 
139
164
  @dataclass(slots=True)
140
165
  class StudioConfig:
166
+ """Complete structured configuration for one studio run."""
167
+
141
168
  data: DataConfig = field(default_factory=DataConfig)
142
169
  model: ModelConfig = field(default_factory=ModelConfig)
143
170
  training: TrainingConfig = field(default_factory=TrainingConfig)
@@ -147,12 +174,16 @@ class StudioConfig:
147
174
  device: str = field(default_factory=_default_device)
148
175
 
149
176
  def to_dict(self) -> dict[str, Any]:
177
+ """Convert the config to JSON-serializable dictionaries."""
178
+
150
179
  raw = asdict(self)
151
180
  raw["data"]["data_dirs"] = [str(path) for path in self.data.data_dirs]
152
181
  raw["output"]["out_dir"] = str(self.output.out_dir)
153
182
  return raw
154
183
 
155
184
  def with_updates(self, **updates: Any) -> "StudioConfig":
185
+ """Return a copy with dotted-key updates such as `training.epochs=10`."""
186
+
156
187
  values = self.to_dict()
157
188
  for dotted_key, value in updates.items():
158
189
  section, key = dotted_key.split(".", 1)
@@ -45,11 +45,15 @@ MODEL_NOTES: dict[str, str] = {
45
45
 
46
46
 
47
47
  def canonical_model_name(name: str) -> str:
48
+ """Normalize aliases such as `kan` or `fastkan` to registry names."""
49
+
48
50
  normalized = name.strip().lower()
49
51
  return MODEL_ALIASES.get(normalized, normalized)
50
52
 
51
53
 
52
54
  def available_models() -> dict[str, str]:
55
+ """Return model registry entries as `{model_name: description}`."""
56
+
53
57
  return dict(sorted(MODEL_NOTES.items()))
54
58
 
55
59
 
@@ -75,10 +79,14 @@ def build_model(
75
79
 
76
80
 
77
81
  def count_parameters(model: "nn.Module") -> int:
82
+ """Count trainable parameters in a PyTorch module."""
83
+
78
84
  return sum(parameter.numel() for parameter in model.parameters() if parameter.requires_grad)
79
85
 
80
86
 
81
87
  def describe_model(name: str) -> dict[str, Any]:
88
+ """Return a normalized model name and human-readable description."""
89
+
82
90
  model_name = canonical_model_name(name)
83
91
  return {
84
92
  "name": model_name,
@@ -26,6 +26,8 @@ PREFERRED_COLUMNS = [
26
26
 
27
27
 
28
28
  def flatten_metrics(metrics: dict[str, Any]) -> dict[str, Any]:
29
+ """Convert nested or tensor-like metrics to CSV-friendly scalar values."""
30
+
29
31
  flat: dict[str, Any] = {}
30
32
  for key, value in metrics.items():
31
33
  if isinstance(value, (str, int, float, bool)) or value is None:
@@ -38,11 +40,15 @@ def flatten_metrics(metrics: dict[str, Any]) -> dict[str, Any]:
38
40
 
39
41
 
40
42
  def save_json(path: Path, payload: dict[str, Any]) -> None:
43
+ """Write a JSON file, creating parent directories as needed."""
44
+
41
45
  path.parent.mkdir(parents=True, exist_ok=True)
42
46
  path.write_text(json.dumps(payload, indent=2, ensure_ascii=True), encoding="utf-8")
43
47
 
44
48
 
45
49
  def load_result_table(path: Path) -> pd.DataFrame:
50
+ """Load one result CSV or concatenate result CSVs below a directory."""
51
+
46
52
  if path.is_dir():
47
53
  candidates = sorted(path.glob("**/results.csv"))
48
54
  if not candidates:
@@ -55,6 +61,8 @@ def load_result_table(path: Path) -> pd.DataFrame:
55
61
 
56
62
 
57
63
  def compare_results(path: Path, sort_by: str = "equalized_ber") -> pd.DataFrame:
64
+ """Load, sort, and reorder result columns for quick model comparison."""
65
+
58
66
  df = load_result_table(path)
59
67
  if sort_by in df.columns:
60
68
  ascending = sort_by not in {"improvement_rel", "improvement_db", "accuracy", "samples_per_sec"}
@@ -65,6 +73,8 @@ def compare_results(path: Path, sort_by: str = "equalized_ber") -> pd.DataFrame:
65
73
 
66
74
 
67
75
  def write_result_table(path: Path, rows: list[dict[str, Any]]) -> pd.DataFrame:
76
+ """Write result rows to CSV and return the ordered DataFrame."""
77
+
68
78
  path.parent.mkdir(parents=True, exist_ok=True)
69
79
  df = pd.DataFrame([flatten_metrics(row) for row in rows])
70
80
  ordered = [column for column in PREFERRED_COLUMNS if column in df.columns]
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: ber-equalization-studio
3
- Version: 0.1.2
3
+ Version: 0.1.3
4
4
  Summary: Research studio for BER equalization experiments in nonlinear optical links.
5
5
  Keywords: ber,equalization,optical-communications,photonics,pytorch
6
6
  Classifier: Development Status :: 3 - Alpha
@@ -9,6 +9,7 @@ Classifier: Programming Language :: Python :: 3
9
9
  Classifier: Programming Language :: Python :: 3.10
10
10
  Classifier: Programming Language :: Python :: 3.11
11
11
  Classifier: Programming Language :: Python :: 3.12
12
+ Classifier: Programming Language :: Python :: 3.13
12
13
  Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
13
14
  Classifier: Topic :: Scientific/Engineering :: Physics
14
15
  Requires-Python: >=3.10
@@ -1,6 +1,7 @@
1
1
  README.md
2
2
  pyproject.toml
3
3
  src/ber_equalization_studio/__init__.py
4
+ src/ber_equalization_studio/__init__.pyi
4
5
  src/ber_equalization_studio/api.py
5
6
  src/ber_equalization_studio/cli.py
6
7
  src/ber_equalization_studio/config.py
@@ -8,6 +9,7 @@ src/ber_equalization_studio/data.py
8
9
  src/ber_equalization_studio/experiment.py
9
10
  src/ber_equalization_studio/legacy.py
10
11
  src/ber_equalization_studio/models.py
12
+ src/ber_equalization_studio/py.typed
11
13
  src/ber_equalization_studio/results.py
12
14
  src/ber_equalization_studio/visualization.py
13
15
  src/ber_equalization_studio.egg-info/PKG-INFO