refactor: remove unused Optuna visualization utility

feat: add utility functions for fiber dataset visualization and hyperparameter training;
housekeeping: rename dataset.py -> datasets.py

.gitignore

@@ -1,4 +1,5 @@
 src/**/*.ini
+.data
 
 # VSCode
 .vscode

LICENSE

@@ -10,6 +10,7 @@ use is covered by a right of the copyright holder of the Work).
 The Work is provided under the terms of this Licence when the Licensor (as
 defined below) has placed the following notice immediately following the
 copyright notice for the Work:
+
 ```raw
 Licensed under the EUPL
 ```

README.md

@@ -13,7 +13,7 @@ Full license text in LICENSE file
 
 # optical-regeneration
 
-## Notes on cloning:
+## Notes on cloning
 
 - `pypho` is added as a submodule -> `--recurse-submodules`
 - This repo has about 7.5GB of datasets in it. The `git lfs fetch` step will take a while.
@@ -29,4 +29,5 @@ git lfs checkout
 ```
 
 ## License
+
 This project is licensed under EUPL-1.2.

notes/cuda-12-4-install.md

@@ -1,6 +1,6 @@
 # CUDA 12.4 Install
 
-> https://unknowndba.blogspot.com/2024/04/cuda-getting-started-on-wsl.html (@\_freddenis\_, 2024)
+> <https://unknowndba.blogspot.com/2024/04/cuda-getting-started-on-wsl.html> (@\_freddenis\_, 2024)
 
 ```bash
 # get md5sums
@@ -49,10 +49,9 @@ make
 ./devicequery
 ```
 
-
 ### if the cuda-toolkit install fails with unmet dependencies
 
->https://askubuntu.com/a/1493087 (jspinella, 2023, CC BY-SA 4.0)
+><https://askubuntu.com/a/1493087> (jspinella, 2023, CC BY-SA 4.0)
 
 1. Open the *new* file for storing the sources list
 
@@ -71,4 +70,3 @@ make
     ```
 
 3. Save the file and run `sudo apt update` - now the install command for CUDA should work.
-

notes/links.md

@@ -1,4 +1,4 @@
 # useful links
 
-- (Optuna)[https://optuna.org] Hyperparameter optimization framework
+- [Optuna](https://optuna.org) Hyperparameter optimization framework
   `pip install optuna`

notes/pyenv-install.md

@@ -1,46 +1,42 @@
-# pyenv install
+# pyenv installation
 
-## install
+## pyenv
 
-nice to have:
+1. Install pyenv
 
-```bash
-sudo apt install python-is-python3
-```
+    ```bash
+    curl https://pyenv.run | bash
+    ```
 
-```bash
-curl https://pyenv.run | bash
-```
+2. setup zsh
 
-## setup zsh
+    add the following to `.zshrc`:
 
-add the following to `.zshrc`:
+    ```bash
+    export PYENV_ROOT="$HOME/.pyenv"
+    [[ -d $PYENV_ROOT/bin ]] && export PATH="$PYENV_ROOT/bin:$PATH"
+    eval "$(pyenv init -)"
+    ```
 
-```bash
-export PYENV_ROOT="$HOME/.pyenv"
-[[ -d $PYENV_ROOT/bin ]] && export PATH="$PYENV_ROOT/bin:$PATH"
-eval "$(pyenv init -)"
-```
+## python installation
 
-## pyenv install
+1. prerequisites
 
-prerequisites:
+    ```bash
+    sudo apt update
+    sudo apt install build-essential libssl-dev zlib1g-dev \
+    libbz2-dev libreadline-dev libsqlite3-dev curl git \
+    libncursesw5-dev xz-utils tk-dev libxml2-dev libxmlsec1-dev \
+    libffi-dev liblzma-dev python3-pip
+    ```
 
-```bash
-sudo apt update
-sudo apt install build-essential libssl-dev zlib1g-dev \
-libbz2-dev libreadline-dev libsqlite3-dev curl git \
-libncursesw5-dev xz-utils tk-dev libxml2-dev libxmlsec1-dev \
-libffi-dev liblzma-dev python3-pip
-```
+2. install
 
-install:
+    ```bash
+    # using python 3.12.7 as an example
+    pyenv install 3.12.7
 
-```bash
-# using python 3.12.7 as an example
-pyenv install 3.12.7
-
-# optional
-pyenv global 3.12.7
-pyenv versions
-```
+    # optional
+    pyenv global 3.12.7
+    pyenv versions
+    ```

notes/pytorch-install.md

@@ -8,7 +8,8 @@ source ./.venv/bin/activate
 ```
 
 ## install pytorch
-> https://pytorch.org/get-started/locally/
+
+> <https://pytorch.org/get-started/locally/>
 
 ```bash
 pip install torch torchvision torchaudio

src/single-core-data-gen/add_pypho.py

@@ -0,0 +1,37 @@
+
+# add_pypho.py
+# 
+# This file is part of the repo "optical-regeneration"
+# https://git.suuppl.dev/seppl/optical-regeneration.git
+# 
+# (c) Joseph Hopfmüller, 2024
+# Licensed under the EUPL
+# 
+# Full license text in LICENSE file
+
+###
+
+# copy this file into the directory where you want to use pypho
+
+import sys
+from pathlib import Path
+
+__log = []
+
+# add the dir above the one where this file lives
+__parent_dir = Path(__file__).parent
+
+# search for a dir containing ./pypho/pypho, then add the lower ./pypho
+while not (__parent_dir / "pypho" / "pypho").exists() and __parent_dir != Path("/"):
+    __parent_dir = __parent_dir.parent
+
+if __parent_dir != Path("/"):
+    sys.path.append(str(__parent_dir / "pypho"))
+    __log.append(f"Added '{__parent_dir/ "pypho"}' to 'PATH'")
+else:
+    __log.append('pypho not found')
+
+
+def show_log():
+    for entry in __log:
+        print(entry)

src/single-core-data-gen/generate_signal.py

@@ -19,9 +19,8 @@ import time
 from matplotlib import pyplot as plt  # noqa: F401
 import numpy as np
 
-import _path_fix  # noqa: F401
+import path_fix
 import pypho
-# import inspect
 
 default_config = f"""
 [glova]
@@ -498,6 +497,7 @@ def plot_eye_diagram(
 
 
 if __name__ == "__main__":
+    path_fix.show_log()
     config = get_config()
 
     length_ranges = [1000, 10000]

src/single-core-regen/_path_fix.py

@@ -1,9 +0,0 @@
-import sys
-from pathlib import Path
-
-# hack to add the parent directory to the path -> pypho doesn't have to be installed as package
-parent_dir = Path(__file__).parent
-while not (parent_dir / "pypho" / "pypho").exists() and parent_dir != Path("/"):
-    parent_dir = parent_dir.parent
-print(f"Adding '{parent_dir / "pypho"}' to 'sys.path' to enable import of '{parent_dir / 'pypho' / 'pypho'}'")
-sys.path.append(str(parent_dir / "pypho"))

src/single-core-regen/regen.py

@@ -0,0 +1,287 @@
+from dataclasses import dataclass
+from datetime import datetime
+import time
+import matplotlib.pyplot as plt
+
+import numpy as np
+import optuna
+import warnings
+
+import torch
+import torch.nn as nn
+import torch.functional as F
+import torch.optim as optim
+import torch.utils.data
+
+import multiprocessing
+
+from util.datasets import FiberRegenerationDataset
+import util
+
+
+# global settings
+@dataclass
+class GlobalSettings:
+    seed: int = 42
+
+
+# data settings
+@dataclass
+class DataSettings:
+    config_path: str = "data/*-128-16384-10000-0-0-17-0-PAM4-0.ini"
+    symbols_range: tuple = (1, 100)
+    data_size_range: tuple = (1, 20)
+    target_delay: float = 0
+    xy_delay_range: tuple = (0, 1)
+    drop_first: int = 1000
+    train_split: float = 0.8
+
+
+# pytorch settings
+@dataclass
+class PytorchSettings:
+    device: str = "cuda"
+    batchsize: int = 128
+    epochs: int = 100
+
+
+# model settings
+@dataclass
+class ModelSettings:
+    output_size: int = 2
+    n_layer_range: tuple = (1, 3)
+    n_units_range: tuple = (4, 128)
+    activation_func_range: tuple = ("ReLU",)
+
+
+@dataclass
+class OptimizerSettings:
+    optimizer_range: tuple = ("Adam", "RMSprop", "SGD")
+    lr_range: tuple = (1e-5, 1e-1)
+
+
+# optuna settings
+@dataclass
+class OptunaSettings:
+    n_trials: int = 128
+    n_threads: int = 16
+    timeout: int = 600
+    directions: tuple = ("maximize",)
+
+    limit_examples: bool = True
+    n_train_examples: int = PytorchSettings.batchsize * 30
+    n_valid_examples: int = PytorchSettings.batchsize * 10
+    storage: str = "sqlite:///optuna_single_core_regen.db"
+    study_name: str = (
+        f"single_core_regen_{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}"
+    )
+    metrics_names: tuple = ("accuracy",)
+
+class HyperTraining:
+    def __init__(self):
+        self.global_settings = GlobalSettings()
+        self.data_settings = DataSettings()
+        self.pytorch_settings = PytorchSettings()
+        self.model_settings = ModelSettings()
+        self.optimizer_settings = OptimizerSettings()
+        self.optuna_settings = OptunaSettings()
+
+        # set some extra settings to make the code more readable
+        self._extra_optuna_settings()
+    
+    def setup_study(self):
+        self.study = optuna.create_study(
+            study_name=self.optuna_settings.study_name,
+            storage=self.optuna_settings.storage,
+            load_if_exists=True,
+            direction=self.optuna_settings.direction,
+            directions=self.optuna_settings.directions,
+        )
+
+        with warnings.catch_warnings(action="ignore"):
+            self.study.set_metric_names(self.optuna_settings.metrics_names)
+        
+        self.n_threads = min(self.optuna_settings.n_trials, self.optuna_settings.n_threads)
+        self.processes = []
+        for _ in range(self.n_threads):
+            p = multiprocessing.Process(target=self._run_optimize)
+            self.processes.append(p)
+    
+    def run_study(self):
+        for p in self.processes:
+            p.start()
+        
+        for p in self.processes:
+            p.join()
+        
+        remaining_trials = self.optuna_settings.n_trials - self.optuna_settings.n_trials % self.optuna_settings.n_threads
+        if remaining_trials:
+            self._run_optimize(remaining_trials)
+    
+    def _run_optimize(self, n_trials):
+        self.study.optimize(self.objective, n_trials=n_trials, timeout=self.optuna_settings.timeout)
+
+    def eye(self, show=True):
+        util.plot.eye(self.data_settings.config_path, show=show)
+
+    def _extra_optuna_settings(self):
+        self.optuna_settings.multi_objective = len(self.optuna_settings.directions) > 1
+        if self.optuna_settings.multi_objective:
+            self.optuna_settings.direction = None
+        else:
+            self.optuna_settings.direction = self.optuna_settings.directions[0]
+            self.optuna_settings.directions = None
+
+        self.optuna_settings.n_train_examples = (
+            self.optuna_settings.n_train_examples
+            if self.optuna_settings.limit_examples
+            else float("inf")
+        )
+        self.optuna_settings.n_valid_examples = (
+            self.optuna_settings.n_valid_examples
+            if self.optuna_settings.limit_examples
+            else float("inf")
+        )
+
+    def define_model(self, trial: optuna.Trial):
+        n_layers = trial.suggest_int(
+            "model_n_layers", *self.model_settings.n_layer_range
+        )
+        layers = []
+
+        # REVIEW does that work?
+        in_features = trial.params["dataset_data_size"] * 2
+        for i in range(n_layers):
+            out_features = trial.suggest_int(
+                f"model_n_units_l{i}", *self.model_settings.n_units_range
+            )
+            activation_func = trial.suggest_categorical(
+                f"model_activation_func_l{i}", self.model_settings.activation_func_range
+            )
+
+            layers.append(nn.Linear(in_features, out_features))
+            layers.append(getattr(nn, activation_func))
+            in_features = out_features
+
+        layers.append(nn.Linear(in_features, self.model_settings.output_size))
+
+        return nn.Sequential(*layers)
+
+    def get_sliced_data(self, trial: optuna.Trial):
+        assert ModelSettings.input_size % 2 == 0, "input_dim must be even"
+        symbols = trial.suggest_float(
+            "dataset_symbols", *self.data_settings.symbols_range, log=True
+        )
+        xy_delay = trial.suggest_float(
+            "dataset_xy_delay", *self.data_settings.xy_delay_range
+        )
+        data_size = trial.suggest_int(
+            "dataset_data_size", *self.data_settings.data_size_range
+        )
+        # get dataset
+        dataset = FiberRegenerationDataset(
+            file_path=self.data_settings.config_path,
+            symbols=symbols,
+            data_size=data_size,  # two channels (x,y)
+            target_delay=self.data_settings.target_delay,
+            xy_delay=xy_delay,
+            drop_first=self.data_settings.drop_first,
+        )
+
+        dataset_size = len(dataset)
+        indices = list(range(dataset_size))
+        split = int(np.floor(self.data_settings.train_split * dataset_size))
+        np.random.seed(self.global_settings.seed)
+        np.random.shuffle(indices)
+        train_indices, valid_indices = indices[:split], indices[split:]
+
+        train_sampler = torch.utils.data.SubsetRandomSampler(train_indices)
+        valid_sampler = torch.utils.data.SubsetRandomSampler(valid_indices)
+
+        train_loader = torch.utils.data.DataLoader(
+            dataset, batch_size=self.pytorch_settings.batchsize, sampler=train_sampler
+        )
+        valid_loader = torch.utils.data.DataLoader(
+            dataset, batch_size=self.pytorch_settings.batchsize, sampler=valid_sampler
+        )
+
+        return train_loader, valid_loader
+
+
+    def train_model(self, model, optimizer, train_loader):
+        model.train()
+        for batch_idx, (data, target) in enumerate(train_loader):
+            if (batch_idx * train_loader.batchsize
+                >= self.optuna_settings.n_train_examples):
+                break
+            optimizer.zero_grad()
+            data, target = (
+                data.to(self.pytorch_settings.device),
+                target.to(self.pytorch_settings.device),
+            )
+            target_pred = model(data)
+            loss = F.mean_squared_error(target_pred, target)
+            loss.backward()
+            optimizer.step()
+
+
+    def eval_model(self, model, valid_loader):
+        model.eval()
+        correct = 0
+        with torch.no_grad():
+            for batch_idx, (data, target) in enumerate(valid_loader):
+                if (
+                    batch_idx * valid_loader.batchsize
+                    >= self.optuna_settings.n_valid_examples
+                ):
+                    break
+                data, target = (
+                    data.to(self.pytorch_settings.device),
+                    target.to(self.pytorch_settings.device),
+                )
+                target_pred = model(data)
+                pred = target_pred.argmax(dim=1, keepdim=True)
+                correct += pred.eq(target.view_as(pred)).sum().item()
+
+        accuracy = correct / len(valid_loader.dataset)
+        # num_params = sum(p.numel() for p in model.parameters())
+        return accuracy
+
+    def objective(self, trial: optuna.Trial):
+        model = self.define_model(trial).to(self.pytorch_settings.device)
+
+        optimizer_name = trial.suggest_categorical(
+            "optimizer", self.optimizer_settings.optimizer_range
+        )
+        lr = trial.suggest_float("lr", *self.optimizer_settings.lr_range, log=True)
+        optimizer = getattr(optim, optimizer_name)(model.parameters(), lr=lr)
+
+        train_loader, valid_loader = self.get_sliced_data(trial)
+
+        for epoch in range(self.pytorch_settings.epochs):
+            self.train_model(model, optimizer, train_loader)
+            accuracy = self.eval_model(model, valid_loader)
+
+            if len(self.optuna_settings.directions) == 1:
+                trial.report(accuracy, epoch)
+                if trial.should_prune():
+                    raise optuna.exceptions.TrialPruned()
+
+        return accuracy
+
+
+if __name__ == "__main__":
+    # plt.ion()
+    hyper_training = HyperTraining()
+    hyper_training.eye()
+
+    # hyper_training.setup_study()
+    # hyper_training.run_study()
+    for i in range(10):
+        #simulate some work
+        print(i)
+        time.sleep(0.2)
+
+    plt.show()
+    ...
+

src/single-core-regen/testing/learn_optuna.py

@@ -0,0 +1,194 @@
+from datetime import datetime
+from pathlib import Path
+
+import optuna
+import warnings
+from util.optuna_vis import show_figures
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.optim as optim
+import torch.utils.data
+
+from torchvision import datasets
+from torchvision import transforms
+import multiprocessing
+
+# from util.dataset import SlicedDataset
+
+DEVICE = torch.device("cuda")
+BATCHSIZE = 128
+CLASSES = 10
+DIR = Path(__file__).parent
+EPOCHS = 100
+N_TRAIN_EXAMPLES = BATCHSIZE * 30
+N_VALID_EXAMPLES = BATCHSIZE * 10
+
+n_trials = 128
+n_threads = 16
+
+
+def define_model(trial):
+    n_layers = trial.suggest_int("n_layers", 1, 3)
+    layers = []
+
+    in_features = 28 * 28
+    for i in range(n_layers):
+        out_features = trial.suggest_int(f"n_units_l{i}", 4, 128)
+        layers.append(nn.Linear(in_features, out_features))
+        layers.append(nn.ReLU())
+        p = trial.suggest_float(f"dropout_l{i}", 0.2, 0.5)
+        layers.append(nn.Dropout(p))
+
+        in_features = out_features
+
+    layers.append(nn.Linear(in_features, CLASSES))
+    layers.append(nn.LogSoftmax(dim=1))
+
+    return nn.Sequential(*layers)
+
+
+def get_mnist():
+    # Load FashionMNIST dataset.
+    train_loader = torch.utils.data.DataLoader(
+        datasets.FashionMNIST(
+            DIR / ".data", train=True, download=True, transform=transforms.ToTensor()
+        ),
+        batch_size=BATCHSIZE,
+        shuffle=True,
+    )
+    valid_loader = torch.utils.data.DataLoader(
+        datasets.FashionMNIST(
+            DIR / ".data", train=False, transform=transforms.ToTensor()
+        ),
+        batch_size=BATCHSIZE,
+        shuffle=True,
+    )
+
+    return train_loader, valid_loader
+
+
+def objective(trial):
+    model = define_model(trial).to(DEVICE)
+
+    optimizer_name = trial.suggest_categorical("optimizer", ["Adam", "RMSprop", "SGD"])
+    lr = trial.suggest_float("lr", 1e-5, 1e-1, log=True)
+    optimizer = getattr(optim, optimizer_name)(model.parameters(), lr=lr)
+
+    train_loader, valid_loader = get_mnist()
+
+    for epoch in range(EPOCHS):
+        train_model(model, optimizer, train_loader)
+        accuracy, num_params = eval_model(model, valid_loader)
+
+    return accuracy, num_params
+
+
+def eval_model(model, valid_loader):
+    model.eval()
+    correct = 0
+    with torch.no_grad():
+        for batch_idx, (data, target) in enumerate(valid_loader):
+            if batch_idx * BATCHSIZE >= N_VALID_EXAMPLES:
+                break
+
+            data, target = data.view(data.size(0), -1).to(DEVICE), target.to(DEVICE)
+            output = model(data)
+
+            pred = output.argmax(dim=1, keepdim=True)
+            correct += pred.eq(target.view_as(pred)).sum().item()
+
+    accuracy = correct / min(len(valid_loader.dataset), N_VALID_EXAMPLES)
+    num_params = sum(p.numel() for p in model.parameters())
+    return accuracy, num_params
+
+
+def train_model(model, optimizer, train_loader):
+    model.train()
+    for batch_idx, (data, target) in enumerate(train_loader):
+        if batch_idx * BATCHSIZE >= N_TRAIN_EXAMPLES:
+            break
+
+        data, target = data.view(data.size(0), -1).to(DEVICE), target.to(DEVICE)
+
+        optimizer.zero_grad()
+        output = model(data)
+        loss = F.nll_loss(output, target)
+        loss.backward()
+        optimizer.step()
+
+
+def run_optimize(n_trials, study):
+    study.optimize(objective, n_trials=n_trials, timeout=600)
+
+
+if __name__ == "__main__":
+    study_name = f"{datetime.now().strftime('%Y-%m-%d %H:%M:%S')} mnist example"
+    storage = "sqlite:///db.sqlite3"
+    directions = ["maximize", "minimize"]
+
+    study = optuna.create_study(
+        directions=directions,
+        storage=storage,
+        study_name=study_name,
+    )
+    
+    with warnings.catch_warnings(action="ignore"):
+        study.set_metric_names(["accuracy", "num params"])
+
+    n_threads = min(n_trials, n_threads)
+
+    processes = []
+    for _ in range(n_threads):
+        p = multiprocessing.Process(
+            target=run_optimize, args=(n_trials // n_threads, study)
+        )
+        p.start()
+        processes.append(p)
+
+    for p in processes:
+        p.join()
+
+    remaining_trials = n_trials - ((n_trials // n_threads) * n_threads)
+    if remaining_trials:
+        print(
+            f"\nRunning last {remaining_trials} trial{'s' if remaining_trials > 1 else ''}:"
+        )
+        run_optimize(directions, remaining_trials, study_name, storage)
+
+    print(f"Number of trials on the Pareto front: {len(study.best_trials)}")
+
+    trial_with_highest_accuracy = max(study.best_trials, key=lambda t: t.values[1])
+    print("Trial with highest accuracy: ")
+    print(f"\tnumber: {trial_with_highest_accuracy.number}")
+    print(f"\tparams: {trial_with_highest_accuracy.params}")
+    print(f"\tvalues: {trial_with_highest_accuracy.values}")
+
+    # for trial in trials:
+    #     print(f"Trial {trial.number}")
+    #     print(f"  Accuracy: {trial.values[0]}")
+    #     print(f"  n_params: {int(trial.values[1])}")
+    #     print( "  Params: ")
+    #     for key, value in trial.params.items():
+    #         print("    {}: {}".format(key, value))
+    #     print()
+
+    # print("  Value: ", trial.value)
+
+    # print("  Params: ")
+    # for key, value in trial.params.items():
+    #     print("    {}: {}".format(key, value))
+
+    figures = []
+    figures.append(
+        optuna.visualization.plot_pareto_front(
+            study, target_names=["accuracy", "num_params"]
+        )
+    )
+    figures.append(optuna.visualization.plot_timeline(study))
+
+    plt = show_figures(*figures)
+
+    print()
+    # plt.show()

src/single-core-regen/testing/sliced_dataset_test.py

@@ -0,0 +1,51 @@
+# move into dir single-core-regen before running
+
+from util.dataset import SlicedDataset
+from torch.utils.data import DataLoader
+from matplotlib import pyplot as plt
+import numpy as np
+
+def eye_dataset(dataset, no_symbols=None, offset=False, show=True):
+    if no_symbols is None:
+        no_symbols = len(dataset)
+    _, axs = plt.subplots(2,2, sharex=True, sharey=True)
+
+    xaxis = np.linspace(0,dataset.symbols_per_slice,dataset.samples_per_slice)
+    roll = dataset.samples_per_symbol//2 if offset else 0
+    for E_out, E_in in dataset[roll:dataset.samples_per_symbol*no_symbols+roll:dataset.samples_per_symbol]:
+        E_in_x, E_in_y, E_out_x, E_out_y = E_in[0], E_in[1], E_out[0], E_out[1]
+        axs[0,0].plot(xaxis, np.abs( E_in_x.numpy())**2, alpha=0.05, color='C0')
+        axs[1,0].plot(xaxis, np.abs( E_in_y.numpy())**2, alpha=0.05, color='C0')
+        axs[0,1].plot(xaxis, np.abs(E_out_x.numpy())**2, alpha=0.05, color='C0')
+        axs[1,1].plot(xaxis, np.abs(E_out_y.numpy())**2, alpha=0.05, color='C0')
+
+    if show:
+        plt.show()
+
+# def plt_dataloader(dataloader, show=True):
+#     _, axs = plt.subplots(2,2, sharex=True, sharey=True)
+
+#     E_outs, E_ins = next(iter(dataloader))
+#     for i, (E_out, E_in) in enumerate(zip(E_outs, E_ins)):
+#         xaxis = np.linspace(dataset.symbols_per_slice*i,dataset.symbols_per_slice+dataset.symbols_per_slice*i,dataset.samples_per_slice)
+#         E_in_x, E_in_y, E_out_x, E_out_y = E_in[0], E_in[1], E_out[0], E_out[1]
+#         axs[0,0].plot(xaxis, np.abs(E_in_x.numpy())**2)
+#         axs[1,0].plot(xaxis, np.abs(E_in_y.numpy())**2)
+#         axs[0,1].plot(xaxis, np.abs(E_out_x.numpy())**2)
+#         axs[1,1].plot(xaxis, np.abs(E_out_y.numpy())**2)
+
+#     if show:
+#         plt.show()
+
+if __name__ == "__main__":
+
+    dataset = SlicedDataset("data/20241115-175517-128-16384-10000-0-0-17-0-PAM4-0.ini", symbols=1, drop_first=100)
+    print(dataset[0][0].shape)
+
+    eye_dataset(dataset, 1000, offset=True, show=False)
+
+    train_loader = DataLoader(dataset, batch_size=10, shuffle=False)
+
+    # plt_dataloader(train_loader, show=False)
+
+    plt.show()

src/single-core-regen/testing/torch-import-test.py

@@ -0,0 +1,68 @@
+import torch
+import time
+
+def print_torch_env():
+    print("Torch version: ", torch.__version__)
+    print("CUDA available: ", torch.cuda.is_available())
+    print("CUDA version: ", torch.version.cuda)
+    print("CUDNN version: ", torch.backends.cudnn.version())
+    print("Device count: ", torch.cuda.device_count())
+    print("Current device: ", torch.cuda.current_device())
+    print("Device name: ", torch.cuda.get_device_name(0))
+    print("Device capability: ", torch.cuda.get_device_capability(0))
+    print("Device memory: ", torch.cuda.get_device_properties(0).total_memory)
+
+def measure_runtime(func):
+    """
+    Measure the runtime of a function.
+
+    :param func: Function to measure
+    :type func: function
+    :return: Wrapped function with runtime measurement
+    :rtype: function
+    """
+    def wrapper(*args, **kwargs):
+        start_time = time.time()
+        result = func(*args, **kwargs)
+        end_time = time.time()
+        print(f"Runtime: {end_time - start_time:.6f} seconds")
+        return result, end_time - start_time
+    return wrapper
+
+@measure_runtime
+def tensor_addition(a, b):
+    """
+    Perform tensor addition.
+
+    :param a: First tensor
+    :type a: torch.Tensor
+    :param b: Second tensor
+    :type b: torch.Tensor
+    :return: Sum of tensors
+    :rtype: torch.Tensor
+    """
+    return a + b
+
+def runtime_test():
+    x = torch.rand(2**18, 2**10)
+    y = torch.rand(2**18, 2**10)
+
+    print("Tensor addition on CPU")
+    _, cpu_time = tensor_addition(x, y)
+
+    print()
+    print("Tensor addition on GPU")
+    if not torch.cuda.is_available():
+        print("CUDA is not available")
+        return
+    
+    _, gpu_time = tensor_addition(x.cuda(), y.cuda())
+
+    print()
+    print(f"Speedup: {cpu_time / gpu_time *100:.2f}%")
+
+
+if __name__ == "__main__":
+    print_torch_env()
+    print()
+    runtime_test()

src/single-core-regen/util/__init__.py

@@ -0,0 +1,2 @@
+from .datasets import FiberRegenerationDataset  # noqa: F401
+from .plot import eye # noqa: F401

src/single-core-regen/util/datasets.py

@@ -0,0 +1,256 @@
+from pathlib import Path
+import torch
+from torch.utils.data import Dataset
+import numpy as np
+import configparser
+
+
+def load_data(config_path, skipfirst=0, num_symbols=None):
+    filepath = Path(config_path)
+    filepath = filepath.parent.glob(filepath.name)
+    config = configparser.ConfigParser()
+    config.read(filepath)
+    path_elements = (
+        config["data"]["dir"],
+        config["data"]["npy_dir"],
+        config["data"]["file"],
+    )
+    datapath = Path("/".join(path_elements).replace('"', ""))
+    sps = int(config["glova"]["sps"])
+
+    if num_symbols is None:
+        num_symbols = int(config["glova"]["nos"]) - skipfirst
+
+    data = np.load(datapath)[skipfirst * sps : num_symbols * sps + skipfirst * sps]
+    config["glova"]["nos"] = str(num_symbols)
+
+    return data, config
+
+
+def roll_along(arr, shifts, dim):
+    # https://stackoverflow.com/a/76920720
+    # (c) Mateen Ulhaq, 2023
+    # CC BY-SA 4.0
+    shifts = torch.tensor(shifts)
+    assert arr.ndim - 1 == shifts.ndim
+    dim %= arr.ndim
+    shape = (1,) * dim + (-1,) + (1,) * (arr.ndim - dim - 1)
+    dim_indices = torch.arange(arr.shape[dim]).reshape(shape)
+    indices = (dim_indices - shifts.unsqueeze(dim)) % arr.shape[dim]
+    return torch.gather(arr, dim, indices)
+
+
+class FiberRegenerationDataset(Dataset):
+    """
+    Dataset for fiber regeneration training.
+
+    The dataset is loaded from a configuration file, which must contain (at least) the following sections:
+    ```
+    [data]
+    dir = <data_dir>
+    npy_dir = <npy_dir>
+    file = <data_file>
+
+    [glova]
+    sps = <samples per symbol>
+    ```
+    The data is loaded from the file `<data_dir>/<npy_dir>/<data_file>` and is assumed to be in the following format:
+    ```
+    [ E_in_x,
+      E_in_y,
+      E_out_x,
+      E_out_y ]
+    ```
+
+    The dataset is sliced into slices, where each slice consists of a (fractional) number of symbols.
+    The target can be delayed relative to the input data by a (fractional) number of symbols.
+    The x and y channels can be delayed relative to each other by a (fractional) number of symbols.
+    """
+
+    def __init__(
+        self,
+        file_path: str | Path,
+        symbols: int | float,
+        *,
+        data_size: int = None,
+        target_delay: float | int = 0,
+        xy_delay: float | int = 0,
+        drop_first: float | int = 0,
+        **kwargs,
+    ):
+        """
+        Initialize the dataset.
+
+        :param file_path: Path to the data file. Can contain wildcards (*). The first
+        :type file_path: str | pathlib.Path
+        :param symbols: Number of symbols in each slice. Can be a float to specify a fraction of a symbol.
+        :type symbols: float | int
+        :param data_size: Number of samples in each slice. The data is reduced by taking equally spaced samples. If unset, each slice will contain symbols*samples_per_symbol samples.
+        :type data_size: int, optional
+        :param target_delay: Delay (in fractional symbols) between data and target. A positive delay means the target is delayed relative to the data. Default is 0.
+        :type target_delay: float | int, optional
+        :param xy_delay: Delay (in fractional symbols) between the x and y channels. A positive delay means the y channel is delayed relative to the x channel. Default is 0.
+        :type xy_delay: float | int, optional
+        :param drop_first: Number of (fractional) symbols to drop from the beginning
+        :type drop_first: float | int
+        """
+
+        # check types
+        assert isinstance(file_path, str), "file_path must be a string"
+        assert isinstance(symbols, (float, int)), (
+            "symbols must be a float or an integer"
+        )
+        assert data_size is None or isinstance(data_size, int), (
+            "output_len must be an integer"
+        )
+        assert isinstance(target_delay, (float, int)), (
+            "target_delay must be a float or an integer"
+        )
+        assert isinstance(xy_delay, (float, int)), (
+            "xy_delay must be a float or an integer"
+        )
+        assert isinstance(drop_first, int), "drop_first must be an integer"
+
+        # check values
+        assert symbols > 0, "symbols must be positive"
+        assert data_size is None or data_size > 0, "output_len must be positive or None"
+        assert drop_first >= 0, "drop_first must be non-negative"
+
+        faux = kwargs.pop("faux", False)
+
+        if faux:
+            data_raw = np.array(
+                [[i + 0.1j, i + 0.2j, i + 1.1j, i + 1.2j] for i in range(12800)],
+                dtype=np.complex128,
+            )
+            self.config = {
+                "data": {"dir": '"."', "npy_dir": '"."', "file": "faux"},
+                "glova": {"sps": 128},
+            }
+        else:
+            data_raw, self.config = load_data(file_path)
+
+        self.samples_per_symbol = int(self.config["glova"]["sps"])
+        self.samples_per_slice = int(symbols * self.samples_per_symbol)
+        self.symbols_per_slice = self.samples_per_slice / self.samples_per_symbol
+
+        self.data_size = data_size or self.samples_per_slice
+        self.target_delay = target_delay or 0
+        self.xy_delay = xy_delay or 0
+
+        ovrd_target_delay_samples = kwargs.pop("ovrd_target_delay_samples", None)
+        ovrd_xy_delay_samples = kwargs.pop("ovrd_xy_delay_samples", None)
+        ovrd_drop_first_samples = kwargs.pop("ovrd_drop_first_samples", None)
+
+        self.target_delay_samples = (
+            ovrd_target_delay_samples
+            if ovrd_target_delay_samples is not None
+            else int(self.target_delay * self.samples_per_symbol)
+        )
+        self.xy_delay_samples = (
+            ovrd_xy_delay_samples
+            if ovrd_xy_delay_samples is not None
+            else int(self.xy_delay * self.samples_per_symbol)
+        )
+        drop_first_samples = (
+            ovrd_drop_first_samples
+            if ovrd_drop_first_samples is not None
+            else int(drop_first * self.samples_per_symbol)
+        )
+
+        # drop samples from the beginning
+        data_raw = data_raw[drop_first_samples:]
+
+        # data layout
+        # [ [E_in_x0, E_in_y0, E_out_x0, E_out_y0],
+        #   [E_in_x1, E_in_y1, E_out_x1, E_out_y1],
+        #        ... 
+        #   [E_in_xN, E_in_yN, E_out_xN, E_out_yN] ]
+
+        data_raw = data_raw.transpose(0, 1)
+
+        # data layout
+        # [  E_in_x[0:N],  
+        #    E_in_y[0:N],  
+        #   E_out_x[0:N], 
+        #   E_out_y[0:N] ]
+
+        # shift x data by xy_delay_samples relative to the y data (example value: 3)
+        # [ E_in_x [0:N],         [ E_in_x [ 0:N  ],        [ E_in_x [3:N  ],
+        #   E_in_y [0:N],    ->     E_in_y [-3:N-3],    ->    E_in_y [0:N-3],
+        #   E_out_x[0:N],           E_out_x[ 0:N  ],          E_out_x[3:N  ],
+        #   E_out_y[0:N] ]          E_out_y[-3:N-3] ]         E_out_y[0:N-3] ]
+
+        if self.xy_delay_samples != 0:
+            data_raw = roll_along(
+                data_raw, [0, self.xy_delay_samples, 0, self.xy_delay_samples], dim=1
+            )
+        if self.xy_delay_samples > 0:
+            data_raw = data_raw[:, self.xy_delay_samples :]
+        elif self.xy_delay_samples < 0:
+            data_raw = data_raw[:, : self.xy_delay_samples]
+
+        # shift fiber input data (target) by target_delay_samples relative to the fiber output data (input)
+        # (example value: 5)
+        # [ E_in_x [0:N],         [ E_in_x [-5:N-5],        [ E_in_x [0:N-5],
+        #   E_in_y [0:N],    ->     E_in_y [-5:N-5],    ->    E_in_y [0:N-5],
+        #   E_out_x[0:N],           E_out_x[ 0:N  ],          E_out_x[5:N  ],
+        #   E_out_y[0:N] ]          E_out_y[ 0:N  ] ]         E_out_y[5:N  ] ]
+
+        if self.target_delay_samples != 0:
+            data_raw = roll_along(
+                data_raw,
+                [self.target_delay_samples, self.target_delay_samples, 0, 0],
+                dim=1,
+            )
+        if self.target_delay_samples > 0:
+            data_raw = data_raw[:, self.target_delay_samples :]
+        elif self.target_delay_samples < 0:
+            data_raw = data_raw[:, : self.target_delay_samples]
+
+        data_raw = data_raw.view(2, 2, -1)
+        # data layout
+        # [ [E_in_x,  E_in_y],
+        #   [E_out_x, E_out_y] ]
+
+        self.data = data_raw.unfold(dimension=-1, size=self.samples_per_slice, step=1)
+        self.data = self.data.movedim(-2, 0)
+        # -> [no_slices, 2, 2, samples_per_slice]
+
+        # data layout
+        # [
+        #   [ [E_in_x[0:N+0], E_in_y[0:N+0] ], [ E_out_x[0:N+0], E_out_y[0:N+0] ] ],
+        #   [ [E_in_x[1:N+1], E_in_y[1:N+1] ], [ E_out_x[1:N+1], E_out_y[1:N+1] ] ],
+        #   ...
+        # ] -> [no_slices, 2, 2, samples_per_slice]
+
+        ...
+
+    def __len__(self):
+        return self.data.shape[0]
+
+    def __getitem__(self, idx):
+        if isinstance(idx, slice):
+            return [self.__getitem__(i) for i in range(*idx.indices(len(self)))]
+        else:
+            data, target = self.data[idx, 1].squeeze(), self.data[idx, 0].squeeze()
+
+            # reduce by by taking self.output_dim equally spaced samples
+            data = data[:, : data.shape[1] // self.data_size * self.data_size]
+            data = data.view(data.shape[0], self.data_size, -1)
+            data = data[:, :, 0]
+
+            # target is corresponding to the latest data point -> try to regenerate that
+            target = target[:, : target.shape[1] // self.data_size * self.data_size]
+            target = target.view(target.shape[0], self.data_size, -1)
+            target = target[:, 0, 0]
+
+            data = data.transpose(0, 1).flatten().squeeze()
+            target = target.flatten().squeeze()
+
+            # data layout:
+            # [sample_x0, sample_y0, sample_x1, sample_y1, ...]
+            # target layout:
+            # [sample_x0, sample_y0]
+
+            return data, target

src/single-core-regen/util/plot.py

@@ -0,0 +1,34 @@
+import matplotlib.pyplot as plt
+import numpy as np
+from .datasets import load_data
+
+def eye(path, title=None, head=1000, skipfirst=1000, show=True):
+    """Plot an eye diagram for the data given by filepath.
+    
+    Args:
+        path (str): Path to the data description file.
+        title (str): Title of the plot.
+        head (int): Number of symbols to plot.
+        skipfirst (int): Number of symbols to skip.
+        show (bool): Whether to call plt.show().
+    """
+    data, config = load_data(path, skipfirst, head)
+    sps = int(config["glova"]["sps"])
+
+    xaxis = np.linspace(0, 2, 2*sps, endpoint=False)
+    fig, axs = plt.subplots(2, 2, figsize=(10, 10), sharex=True, sharey=True)
+    for i in range(head-1):
+        inx, iny, outx, outy = data[i*sps:(i+2)*sps].T
+        axs[0, 0].plot(xaxis, np.abs(inx)**2, color="C0", alpha=0.1)
+        axs[0, 1].plot(xaxis, np.abs(outx)**2, color="C0", alpha=0.1)
+        axs[1, 0].plot(xaxis, np.abs(iny)**2, color="C0", alpha=0.1)
+        axs[1, 1].plot(xaxis, np.abs(outy)**2, color="C0", alpha=0.1)
+    
+    axs[0, 0].set_title("Input x")
+    axs[0, 1].set_title("Output x")
+    axs[1, 0].set_title("Input y")
+    axs[1, 1].set_title("Output y")
+    fig.suptitle(title)
+
+    if show:
+        plt.show(block=False)