1st commit

.gitignore

@@ -0,0 +1,11 @@
+# python's temps
+**/.ipynb_checkpoints/
+**/__pycache__/
+
+# training outputs
+runs/
+# data files
+data/
+# typically weights and data
+*.pt
+

Readme.md

@@ -0,0 +1,176 @@
+# ShimNet
+ShimNet is a data-driven AI solution to improve high-resolution nuclear magnetic resonance (NMR) spectra
+distorted by the inhomogeneous magnetic field (less than optimal shimming). To use it, the experimental training data has to be collected (see **Data collection** below).
+Example data can also be downloaded (see below). 
+
+Paper: ...
+
+## Installation
+
+Python 3.9+
+
+GPU version (for training and inference)
+```
+pip install -r requirements.txt
+```
+
+CPU version (for inference, not recommended for training)
+```
+pip install -r requirements-cpu.txt --extra-index-url https://download.pytorch.org/whl/cpu
+```
+
+## Usage
+To correct spectra presented in the paper:
+1. download weights (model parameters):
+```
+python download_files.py
+```
+or directly from [Google Drive 700MHz](https://drive.google.com/uc?export=download&id=17fTNWl7YW6mPbbZWga0EfdoF_6S8fCke) and [Google Drive 600MHz](https://drive.google.com/uc?export=download&id=1_VxOpFGJcFsOa5DHOW2GJbP8RvHCmC1N) and place it in `weights` directory
+
+
+2. : run correction (e.g. `Azarone_20ul_700MHz.csv`):
+```
+python predict.py sample_data/Azarone_20ul_700MHz.csv -o output --config configs/shimnet_700.yaml --weights weights/shimnet_700MHz.pt
+```
+The output will be `output/Azarone_20ul_700MHz_processed.csv` file
+
+Multiple files may be processed using "*" syntax:
+```
+python predict.py sample_data/*700MHz.csv -o output --config configs/shimnet_700.yaml --weights weights/shimnet_700MH
+z.pt
+```
+
+For 600 MHz data use `--config configs/shimnet_600.yaml` and  `--weights weights/shimnet_600MHz.pt`, e.g.:
+
+```
+python predict.py sample_data/CresolRed_after_styrene_600MHz.csv -o output --config configs/shimnet_600.yaml --weights weights/shimnet_600MHz.pt
+```
+
+### input format
+
+The spectrum file for reconstruction should be in the format of two columns separated by a space and without the sign at the end of the line at the end of the file(example below):
+```csv
+-1.97134	0.0167137	
+-1.97085	-0.00778748	
+-1.97036	-0.0109595	
+-1.96988	0.00825978	
+-1.96939	0.0133886	
+```
+
+## Train on your data
+
+For the model to function properly, it should be trained on calibration data from the spectrometer used for the measurements. To train a model on data from your spectrometer, please follow the instructions below.
+
+### Training data collection
+
+Below we describe the training data collection for Agilent/Varian spectrometers. For machines of other vendors similar procedure can be implemented.
+To collect ShimNet training data use Python script (sweep_shims_lineshape_Z1Z2.py) from the calibration_loop folder to drive the spectrometer:
+ 1. Install TReNDS package ( trends.spektrino.com )
+ 2. Open VnmrJ and type: 'listenon'
+ 3. Put the lineshape sample (1% CHCl3 in deuterated acetone), set standard PROTON parameters, and set nt=1 (do not modify sw and at!)
+ 4. Shim the sample and collect the data. Save the optimally shimmed dataset
+ 5. Edit the sweep_shims_lineshape_Z1Z2.py script
+ 6. Put optimum z1 and z2 shim values as optiz1 and optiz2 below
+ 7. Define the calibration range as range_z1 and range_z2 (default is ok)
+ 8. Start the python script:
+   ```
+     python3 ./sweep_shims_lineshape_Z1Z2.py
+   ```
+   The spectrometer will start collecting spectra
+
+### SCRF extraction
+Shim Coil Response Functions (SCRF) should be extracted from the spectra with `extract_scrf_from_fids.py` script.
+```
+python extract_scrf_from_fids.py
+```
+
+The script uses hardcoded paths to the NMR signals (fid-s) in Agilent/Varian format: a directory with optimal measurement (`opti_fid_path` available) and a directory with calibration loop measurements (`data_dir`):
+```python
+# input
+data_dir = "../../sample_run/loop"
+opti_fid_path = "../../sample_run/opti.fid"
+
+```
+
+The output files are also hardcoded:
+```python
+# output
+spectra_file = "../../sample_run/total.npy"
+spectra_file_names = "../../sample_run/total.csv"
+opi_spectrum_file = "../../sample_run/opti.npy"
+responses_file = "../../sample_run/scrf_61.pt"
+```
+where only the `responses_file` is used in ShimNet training.
+
+If the measurements are stored in a format other than Varian, you may need to change this line:
+```python
+dic, data = ng.varian.read(varian_fid_path)
+```
+(see nmrglue package documentation for details)
+
+### Training
+
+1. Download multiplets database:
+    ```
+    python download_files.py --multiplets
+    ```
+2. Configure run:
+  - create a run directory, e.g. `runs/my_lab_spectrometer_2025`
+  - create a configuration file:
+    1. copy `configs/shimnet_template.py` to the run directory and rename it to `config.yaml`
+       ```bash
+       cp configs/shimnet_template.py runs/my_lab_spectrometer_2025/config.yaml
+       ```
+    2. edit the SCRF in path in the config file:
+       ```yaml
+         response_functions_files:
+         - path/to/srcf_file
+       ```
+       e.g.
+       ```yaml
+         response_functions_files:
+         - ../../sample_run/scrf_61.pt
+       ```
+    3. adjust spectrometer frequency step `frq_step` to match your data (spectrometer range in Hz divided by number of points in spectrum):
+        ```yaml
+        frq_step: 0.34059797
+        ```
+    4. adjust spectromer frequency in the metadata
+        ```yaml
+        metadata: # additional metadata, not used in the training process
+          spectrometer_frequency: 700.0 # MHz
+        ```
+3. Run training:
+    ```
+    python train.py runs/my_lab_spectrometer_2025
+    ```
+    Training results will appear in `runs/my_lab_spectrometer_2025` directory.
+    Model parameters are stored in `runs/my_lab_spectrometer_2025/model.pt` file
+4. Use trained model:
+
+    use `--config runs/my_lab_spectrometer_2025/config.yaml` and  `--weights runs/my_lab_spectrometer_2025/model.pt` flags, e.g.
+    ```
+    python predict.py my_sample1.csv -o my_output --config runs/my_lab_spectrometer_2025/config.yaml --weights runs/my_lab_spectrometer_2025/model.pt
+    ```
+
+
+
+## Repeat training on our data
+
+If you want to train the network using the calibration data from our paper, follow the procedure below.
+
+1. Download multiplets database and our SCRF file:
+    ```
+    python download_files.py --multiplets --SCRF --no-weights
+    ```
+2. Configure run
+    ```bash
+    mkdir -p runs/repeat_paper_training
+    cp configs/shimnet_700.yaml runs/repeat_paper_training/config.yaml
+    ```
+3. Run training:
+    ```
+    python train.py runs/repeat_paper_training
+    ```
+    Training results will appear in `runs/repeat_paper_training` directory.
+

calibration_loop/maclib/shimator_loop

@@ -0,0 +1,33 @@
+" macro shimnet loop "
+" arguments: optimal z1, optimal z2, step in z1 loop, step in z2 loop, half number of steps in z1, half number of steps in z2"
+
+$svfdir='/home/nmrbox/kkazimierczuk/shimnet/data/'
+
+$opti_z1=$1
+$opti_z2=$2
+$step_z1=$3
+$step_z2=$4
+$steps_z1=$5
+$steps_z2=$6
+$nn=''
+$bb=''
+$msg=''
+format((($steps_z1*2+1)*($steps_z2*2+1)*(d1+at)*nt)/3600,5,1):$msg
+$msg='Time: '+$msg+' h'
+banner($msg)
+$j=0
+repeat
+  $i=0
+  $z2=$opti_z2-$steps_z2*$step_z2+$j*$step_z2
+  format($z2,5,1):$bb
+   repeat 
+	$z1=$opti_z1-$steps_z1*$step_z1+$i*$step_z1
+	"su"
+	"go"
+	format($z1,5,1):$nn
+	$filepath=$svfdir + 'z1_'+$nn+'z2_'+$bb+'.fid'
+	svf($filepath,'force')
+	$i=$i+1	
+   until $z1>$opti_z1+$steps_z1*$step_z1-1
+   $j=$j+1
+until $z2>$opti_z2+$steps_z2*$step_z2-1

calibration_loop/sweep_shims_lineshape_Z1Z2.py

@@ -0,0 +1,39 @@
+from TReND.send2vnmr import *
+import shutil
+import time
+import os
+import numpy as np
+from subprocess import call
+
+#This is the script to collect ShimNet trainig data on Agilent spectrometers. Execution:
+# 1. Open VnmrJ and type: 'listenon'
+# 2. Put the lineshape sample, set standard PROTON parameters and set one scan (do not modify sw and at!)
+# 3. Shim the sample and collect the data. Save the optimally shimmed datatset
+# 4. Put optimum z1 and z2 shim values as optiz1 and optiz2 below
+# 5. Define the calibration range as range_z1 and range_z2 (default is ok)
+# 6. Start the python script: 
+# python3 ./sweep_shims_lineshape_Z1Z2.py
+# the spectrometer will start collecting spectra
+ 
+Classic_path = '/home/nmr700/shimnet6/lshp' 
+optiz1= 8868 #put optimum shim values here
+optiz2=-297
+
+range_z1=100 #put optimum shim ranges here
+range_z2=100 #put optimum shim ranges here
+
+z1_sweep=np.arange(optiz1-range_z1,optiz1+range_z1+1,2.0)
+z2_sweep=np.arange(optiz2-range_z2,optiz2+range_z2+1,2.0)
+
+for i in range(1,np.shape(z1_sweep)[0]+1, 1):
+    for j in range(1,np.shape(z2_sweep)[0]+1, 1):
+        wait_until_idle()
+    
+        Run_macro("sethw('z1',"+str(z1_sweep[i-1])+ ")")
+        Run_macro("sethw('z2',"+str(z2_sweep[j-1])+ ")")
+ 
+        go_if_idle()
+        wait_until_idle()
+        time.sleep(0.5)
+
+        Save_experiment(Classic_path + '_z1_'+ str(int(z1_sweep[i-1])) + '_z2_'+ str(int(z2_sweep[j-1]))+ '.fid')

configs/shimnet_600.yaml

@@ -0,0 +1,46 @@
+model:
+  name: ShimNetWithSCRF
+  kwargs:
+    rensponse_length: 81
+    resnponse_head_dims:
+    - 128
+training:
+- batch_size: 64
+  learning_rate: 0.001
+  max_iters: 1600000
+- batch_size: 512
+  learning_rate: 0.001
+  max_iters: 25600000
+- batch_size: 512
+  learning_rate: 0.0005
+  max_iters: 12800000
+losses_weights:
+  clean: 1.0
+  noised: 1.0
+  response: 1.0
+data:
+  response_functions_files:
+  # Paste path to your SCRF file here
+  # - Can be absolute path
+  # - Can be relative to repository root
+  - data/600hz-2-Lnative/scrf_81.pt
+  atom_groups_data_file: data/multiplets_10000_parsed.txt
+  response_function_stretch_min: 1.0
+  response_function_stretch_max: 1.0
+  response_function_noise: 0.0
+  multiplicity_j1_min: 0.0
+  multiplicity_j1_max: 15
+  multiplicity_j2_min: 0.0
+  multiplicity_j2_max: 15
+  number_of_signals_min: 2
+  number_of_signals_max: 5
+  thf_min: 0.5
+  thf_max: 2
+  relative_height_min: 0.5
+  relative_height_max: 4
+  frq_step: 0.30048
+logging:
+  step: 1000000
+  num_plots: 32
+metadata:
+  spectrometer_frequency: 600.0

configs/shimnet_700.yaml

@@ -0,0 +1,46 @@
+model:
+  name: ShimNetWithSCRF
+  kwargs:
+    rensponse_length: 61
+    resnponse_head_dims:
+    - 128
+training:
+- batch_size: 64
+  learning_rate: 0.001
+  max_iters: 1600000
+- batch_size: 512
+  learning_rate: 0.001
+  max_iters: 6400000
+- batch_size: 512
+  learning_rate: 0.0005
+  max_iters: 12800000
+- batch_size: 800
+  learning_rate: 0.0005
+  max_iters: 12800000
+losses_weights:
+  clean: 1.0
+  noised: 1.0
+  response: 1.0
+data:
+  response_functions_files:
+  - data/scrf_61_700MHz.pt
+  atom_groups_data_file: data/multiplets_10000_parsed.txt
+  response_function_stretch_min: 1.0
+  response_function_stretch_max: 1.0
+  response_function_noise: 0.0
+  multiplicity_j1_min: 0.0
+  multiplicity_j1_max: 15
+  multiplicity_j2_min: 0.0
+  multiplicity_j2_max: 15
+  number_of_signals_min: 2
+  number_of_signals_max: 5
+  thf_min: 0.5
+  thf_max: 2
+  relative_height_min: 0.5
+  relative_height_max: 4
+  frq_step: 0.34059797
+logging:
+  step: 1000000
+  num_plots: 32
+metadata:
+  spectrometer_frequency: 700.0

configs/shimnet_template.yaml

@@ -0,0 +1,50 @@
+model:
+  name: ShimNetWithSCRF
+  kwargs:
+    rensponse_length: 61
+    resnponse_head_dims:
+    - 128
+training:
+- batch_size: 64
+  learning_rate: 0.001
+  max_iters: 1600000
+- batch_size: 512
+  learning_rate: 0.001
+  max_iters: 6400000
+- batch_size: 512
+  learning_rate: 0.0005
+  max_iters: 12800000
+- batch_size: 800
+  learning_rate: 0.0005
+  max_iters: 12800000
+losses_weights:
+  clean: 1.0
+  noised: 1.0
+  response: 1.0
+data:
+  response_functions_files:
+  # Specify the path to your SCRF file/files here.
+  # - It can be an absolute path.
+  # - It can be relative to the repository root.
+  # Multiple files can be listed in the following rows (YAML list format).
+  - path/to/srcf_file
+  atom_groups_data_file: data/multiplets_10000_parsed.txt
+  response_function_stretch_min: 1.0 # realtive
+  response_function_stretch_max: 1.0 # realtive
+  response_function_noise: 0.0 # arbitrary hight units
+  multiplicity_j1_min: 0.0 # Hz
+  multiplicity_j1_max: 15 # Hz
+  multiplicity_j2_min: 0.0 # Hz
+  multiplicity_j2_max: 15 # Hz
+  number_of_signals_min: 2
+  number_of_signals_max: 5
+  thf_min: 0.5 # arbitrary hight units
+  thf_max: 2 # arbitrary hight units
+  relative_height_min: 0.5 # realtive
+  relative_height_max: 4 # realtive
+  frq_step: 0.34059797	# Hz per point
+logging:
+  step: 1000000
+  num_plots: 32
+metadata: # additional metadata, not used in the training process
+  spectrometer_frequency: 700.0 # MHz

download_files.py

@@ -0,0 +1,90 @@
+import urllib.request
+from pathlib import Path
+import argparse
+
+ALL_FILES_TO_DOWNLOAD = {
+    "weights": [{
+        "url": "https://drive.google.com/uc?export=download&id=17fTNWl7YW6mPbbZWga0EfdoF_6S8fCke",
+        "destination": "weights/shimnet_700MHz.pt"
+    },
+    {
+        "url": "https://drive.google.com/uc?export=download&id=1_VxOpFGJcFsOa5DHOW2GJbP8RvHCmC1N",
+        "destination": "weights/shimnet_600MHz.pt"
+    }],
+    "SCRF": [{
+        "url": "https://drive.google.com/uc?export=download&id=113al7A__yYALx_2hkESuzFIDU3feVtNY",
+        "destination": "data/scrf_61_700MHz.pt"
+    }],
+    "mupltiplets": [{
+        "url": "https://drive.google.com/uc?export=download&id=1QGvV-Au50ZxaP1vFsmR_auI299Dw-Wrt",
+        "destination": "data/multiplets_10000_parsed.txt"
+    }],
+    "development": []
+}
+
+def parse_args():
+    parser = argparse.ArgumentParser(
+        description='Download files: weighs (default), SCRF (optional), multiplet data (optional)',
+    )
+    parser.add_argument(
+        '--weights',
+        action='store_true',
+        default=True,
+        help='Download weights file (default behavior). Use --no-weights to opt out.',
+    )
+    parser.add_argument(
+        '--no-weights',
+        action='store_false',
+        dest='weights',
+        help='Do not download weights file.',
+    )
+    parser.add_argument('--SCRF', action='store_true', help='Download SCRF file') 
+    parser.add_argument('--multiplets', action='store_true', help='Download multiplets data file')
+    parser.add_argument('--development', action='store_true', help='Download development weights file')
+
+    parser.add_argument('--all', action='store_true', help='Download all available files')
+
+    args = parser.parse_args()
+    # Set all individual flags if --all is specified
+    if args.all:
+        args.weights = True
+        args.SCRF = True 
+        args.multiplets = True
+        args.development = True
+        
+    return args
+
+def download_file(url, target):
+    target = Path(target)
+    if target.exists():
+        response = input(f"File {target} already exists. Overwrite? (y/n): ")
+        if response.lower() != 'y':
+            print(f"Download of {target} cancelled")
+            return
+    target.parent.mkdir(parents=True, exist_ok=True)
+    try:
+        urllib.request.urlretrieve(url, target)
+        print(f"Downloaded {target}")
+    except Exception as e:
+        print(f"Failed to download file from {url}:\n {e}")
+
+
+if __name__ == "__main__":
+    args = parse_args()
+
+    main_dir = Path(__file__).parent
+    if args.weights:
+        for file_data in ALL_FILES_TO_DOWNLOAD["weights"]:
+            download_file(file_data["url"], main_dir / file_data["destination"])
+    
+    if args.SCRF:
+        for file_data in ALL_FILES_TO_DOWNLOAD["SCRF"]:
+            download_file(file_data["url"], main_dir / file_data["destination"])
+    
+    if args.multiplets:
+        for file_data in ALL_FILES_TO_DOWNLOAD["mupltiplets"]:
+            download_file(file_data["url"], main_dir / file_data["destination"])
+    
+    if args.development:
+        for file_data in ALL_FILES_TO_DOWNLOAD["development"]:
+            download_file(file_data["url"], main_dir / file_data["destination"])

extract_scrf_from_fids.py

@@ -0,0 +1,132 @@
+import os
+import shutil
+from pathlib import Path
+import matplotlib.pyplot as plt
+import numpy as np
+import pandas as pd
+import torch
+import nmrglue as ng
+from tqdm import tqdm
+
+# input
+data_dir = "../../../dane/600Hz_20241227/loop"
+opti_fid_path = "../../../dane/600Hz_20241227/opti"
+
+# output
+spectra_file = "data/600Hz_20241227-Lnative/total.npy"
+spectra_file_names = "data/600Hz_20241227-Lnative/total.csv"
+opi_spectrum_file = "data/600Hz_20241227-Lnative/opti.npy"
+responses_file = "data/600Hz_20241227-Lnative/scrf_61.pt"
+losses_file = "data/600Hz_20241227-Lnative/losses_scrf_61.pt"
+
+# create directories for output files
+for file_path in [spectra_file, spectra_file_names, opi_spectrum_file, responses_file]:
+    Path(file_path).parent.mkdir(parents=True, exist_ok=True)
+
+# settings: fid to spectrum
+ph0_correction =-190.49
+ph1_correction = 0
+autophase_fn="acme"
+target_length = None
+
+# settings: SCRF extraction
+calibration_peak_center = "auto"#12277 # the center of the calibration peak (peak of)
+calibration_window_halfwidth = 128 # the half width of the calibration window
+steps = 6000
+kernel_size = 61
+kernel_sqrt = False # True to allow only positive values
+
+# fid to spectra processing
+
+def fid_to_spectrum(varian_fid_path, ph0_correction, ph1_correction, autophase_fn, target_length=None, sin_pod=False):
+    dic, data = ng.varian.read(varian_fid_path)
+    data[0] *= 0.5
+    if sin_pod:
+        data = ng.proc_base.sp(data, end=0.98)
+
+    if target_length is not None:
+        if (pad_length := target_length - len(data)) > 0:
+            data = ng.proc_base.zf(data, pad_length)
+        else:
+            data = data[:target_length]
+
+    spec=ng.proc_base.fft(data)
+    spec = ng.process.proc_autophase.autops(spec, autophase_fn, p0=ph0_correction, p1=ph1_correction, disp=False)
+
+    return spec
+
+# process optimal measurement fid to spectrum
+opti_spectrum_full = fid_to_spectrum(opti_fid_path, ph0_correction, ph1_correction, autophase_fn, target_length=target_length)
+
+if calibration_peak_center == "auto":
+    calibration_peak_center = np.argmax(abs(opti_spectrum_full))
+fitting_range = (calibration_peak_center - calibration_window_halfwidth, calibration_peak_center+calibration_window_halfwidth+1)
+
+opti_spectrum = opti_spectrum_full[fitting_range[0]:fitting_range[1]]
+np.save(opi_spectrum_file, opti_spectrum)
+print(f"Optimal spectrum extracted to {opi_spectrum_file}")
+
+# process loop fids to spectra
+spec_list=[]
+spec_names=[]
+
+print("Extracting spectra from fids...")
+for fid_path in tqdm(list(Path(data_dir).rglob('*.fid'))):
+    spec = fid_to_spectrum(fid_path, ph0_correction, ph1_correction, autophase_fn, target_length=target_length)[fitting_range[0]:fitting_range[1]]
+       
+    spec_list.append(spec)
+    spec_names.append(fid_path.name)
+
+total = np.array(spec_list)
+np.save(spectra_file, total)
+pd.DataFrame(spec_names).to_csv(spectra_file_names, header=False)
+# total = np.load(spectra_file)
+print(f"Spectra extracted to {spectra_file}")
+
+# process SCRF extraction
+def fit_kernel(base, target, kernel_size, kernel_sqrt=True, steps=20000, verbose=False):
+
+    kernel = torch.ones((1,1,kernel_size), dtype=base.dtype)
+    if kernel_sqrt:
+        kernel /= torch.sqrt(torch.sum(kernel**2))
+    else:
+        kernel /= kernel_size
+    kernel.requires_grad = True
+
+    optimizer = torch.optim.Adam([kernel])
+
+    for epoch in range(steps):
+        if kernel_sqrt:
+            spe_est = torch.conv1d(base, kernel**2, padding='same')
+        else:
+            spe_est = torch.conv1d(base, kernel, padding='same')
+        loss = torch.mean(abs(target - spe_est)**2) #torch.nn.functional.mse_loss(spe_est, target)
+        loss.backward()
+        optimizer.step()
+        optimizer.zero_grad()
+
+        if verbose and (epoch+1) % 100 == 0:
+            print(epoch, loss.item())
+    if kernel_sqrt:
+        return kernel.detach()**2, loss.item()
+    else:
+        return kernel.detach(), loss.item()
+
+responses = torch.empty(len(total), 1, 1, 1, 1, kernel_size)
+losses = torch.empty(len(total))
+base = torch.tensor(opti_spectrum.real).unsqueeze(0)
+targets = torch.tensor(total.real)
+
+# normalization
+base /= base.sum()
+targets /= targets.sum(dim=(-1,), keepdim=True)
+
+print("\nExtracting SCRFs...")
+for i, target in tqdm(enumerate(targets), total=len(targets)):    
+    kernel, loss = fit_kernel(base, target.unsqueeze(0), kernel_size, kernel_sqrt=kernel_sqrt, steps=steps)
+    responses[i, 0, 0] = kernel
+    losses[i] = loss
+
+torch.save(responses, responses_file)
+torch.save(losses, losses_file)
+print(f"SCRFs extracted to {responses_file}, losses saved to {losses_file}")

predict.py

@@ -0,0 +1,89 @@
+import torch
+torch.set_grad_enabled(False)
+import numpy as np
+import argparse
+from pathlib import Path
+import sys, os
+from omegaconf import OmegaConf
+
+from src.models import ShimNetWithSCRF, Predictor
+
+# silent deprecation warnings
+# https://github.com/pytorch/pytorch/issues/97207#issuecomment-1494781560
+import warnings
+warnings.filterwarnings('ignore', category=UserWarning, message='TypedStorage is deprecated')
+
+class Defaults:
+    SCALE = 16.0
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("input_files", help="Input files", nargs="+") 
+    parser.add_argument("--config", help="config file .yaml")
+    parser.add_argument("--weights", help="model weights")
+    parser.add_argument("-o", "--output_dir", default=".", help="Output directory")
+    parser.add_argument("--input_spectrometer_frequency", default=None, type=float, help="spectrometer frequency in MHz (input sample collection frequency). Empty if the same as in the training data")
+    args = parser.parse_args()
+    return args
+    
+# functions
+def resample_input_spectrum(input_freqs, input_spectrum, Mhz_per_point):
+    """resample input spectrum to match the model's frequency range"""
+    freqs = np.arange(input_freqs.min(), input_freqs.max(), Mhz_per_point)
+    spectrum = np.interp(freqs, input_freqs, input_spectrum)
+    return freqs, spectrum
+
+def resample_output_spectrum(input_freqs, freqs, prediction):
+    """resample prediction to match the input spectrum's frequency range"""
+    prediction = np.interp(input_freqs, freqs, prediction)
+    return prediction
+
+def initialize_predictor(config, weights_file):
+    model = ShimNetWithSCRF(**config.model.kwargs)
+    predictor = Predictor(model, weights_file)
+    return predictor
+
+# run
+if __name__ == "__main__":
+    args = parse_args()
+    output_dir = Path(args.output_dir)
+    output_dir.mkdir(exist_ok=True, parents=True)
+
+    config = OmegaConf.load(args.config)
+    model_ppm_per_point = config.data.frq_step / config.metadata.spectrometer_frequency
+    predictor = initialize_predictor(config, args.weights)
+
+    for input_file in args.input_files:
+        print(f"processing {input_file} ...")
+
+        # load data
+        input_data = np.loadtxt(input_file)
+        input_freqs_input_ppm, input_spectrum = input_data[:,0], input_data[:,1]
+
+        # convert input frequencies to model's frequency - correct for zero filling ad spectrometer frequency
+        if args.input_spectrometer_frequency is not None:
+            input_freqs_model_ppm = input_freqs_input_ppm * args.input_spectrometer_frequency / config.metadata.spectrometer_frequency
+        else:
+            input_freqs_model_ppm = input_freqs_input_ppm
+        
+        freqs, spectrum = resample_input_spectrum(input_freqs_model_ppm, input_spectrum, model_ppm_per_point)
+        
+        spectrum = torch.tensor(spectrum).float()
+        # scale height of the spectrum
+        scaling_factor = Defaults.SCALE / spectrum.max()
+        spectrum *= scaling_factor
+
+        # correct spectrum
+        prediction = predictor(spectrum).numpy()
+
+        # rescale height
+        prediction /= scaling_factor
+
+        # resample the output to match the input spectrum
+        output_prediction = resample_output_spectrum(input_freqs_model_ppm, freqs, prediction)
+
+        # save result
+        output_file = output_dir / f"{Path(input_file).stem}_processed{Path(input_file).suffix}"
+
+        np.savetxt(output_file, np.column_stack((input_freqs_input_ppm, output_prediction)))
+        print(f"saved to {output_file}")

requirements-cpu.txt

@@ -0,0 +1,9 @@
+torch==2.4.1+cpu
+torchaudio==2.4.1+cpu
+nmrglue==0.11
+torchdata==0.9.0
+numpy==2.0.2
+matplotlib==3.9.3
+pandas==2.2.3
+tqdm==4.67.1
+hydra-core==1.3.2

requirements.txt

@@ -0,0 +1,9 @@
+torch==2.4.1
+torchaudio==2.4.1
+nmrglue==0.11
+torchdata==0.9.0
+numpy==2.0.2
+matplotlib==3.9.3
+pandas==2.2.3
+tqdm==4.67.1
+hydra-core==1.3.2

src/generators.py

@@ -0,0 +1,280 @@
+import numpy as np
+import torch
+import torchdata
+# from itertools import islice
+
+def random_value(min_value, max_value):
+    return (min_value + torch.rand(1) * (max_value - min_value)).item()
+
+def random_loguniform(min_value, max_value):
+    return (min_value * torch.exp(torch.rand(1) * (torch.log(torch.tensor(max_value)) - torch.log(torch.tensor(min_value))))).item()
+
+def calculate_theoretical_spectrum(peaks_parameters: dict, frq_frq:torch.Tensor):
+    # extract parameters
+    tff_lin = peaks_parameters["tff_lin"]
+    twf_lin = peaks_parameters["twf_lin"]
+    thf_lin = peaks_parameters["thf_lin"]
+    trf_lin = peaks_parameters["trf_lin"]
+
+    lwf_lin = twf_lin
+    lhf_lin = thf_lin * (1. - trf_lin)
+    gwf_lin = twf_lin
+    gdf_lin = gwf_lin / torch.tensor(2.).log().mul(2.).sqrt()
+    ghf_lin = thf_lin * trf_lin
+    # calculate Lorenz peaks contriubutions
+    lsf_linfrq = lwf_lin[:, None] ** 2 / (lwf_lin[:, None] ** 2 + (frq_frq - tff_lin[:, None]) ** 2) * lhf_lin[:, None]
+    # calculate Gaussian peaks contriubutions
+    gsf_linfrq = torch.exp(-(frq_frq - tff_lin[:, None]) ** 2 / gdf_lin[:, None] ** 2 / 2.) * ghf_lin[:, None]
+    tsf_linfrq = lsf_linfrq + gsf_linfrq
+    # sum peaks contriubutions
+    tsf_frq = tsf_linfrq.sum(0, keepdim = True)
+    return tsf_frq
+
+
+pascal_triangle = [(1,), (1,1), (1,2,1), (1,3,3,1), (1,4,6,4,1), (1,5,10,10,5,1), (1,6,15,20,15,6,1), (1,7, 21,35,35,21,7,1)]
+normalized_pascal_triangle = [torch.tensor(x)/sum(x) for x in pascal_triangle]
+
+def pascal_multiplicity(multiplicity):
+    intensities = normalized_pascal_triangle[multiplicity-1]
+    n_peaks = len(intensities)
+    shifts = torch.arange(n_peaks)-((n_peaks-1)/2)
+    return shifts, intensities
+
+def double_multiplicity(multiplicity1, multiplicity2, j1=1, j2=1):
+    shifts1, intensities1 = pascal_multiplicity(multiplicity1)
+    shifts2, intensities2 = pascal_multiplicity(multiplicity2)
+
+    shifts = (j1*shifts1.reshape(-1,1) + j2*shifts2.reshape(1,-1)).flatten()
+    intensities = (intensities1.reshape(-1,1) * intensities2.reshape(1,-1)).flatten()
+    return shifts, intensities
+    
+def generate_multiplet_parameters(multiplicity, tff_lin, thf_lin, twf_lin, trf_lin, j1, j2):
+    shifts, intensities = double_multiplicity(multiplicity[0], multiplicity[1], j1, j2)
+    n_peaks = len(shifts)
+
+    return {
+        "tff_lin": shifts + tff_lin,
+        "thf_lin": intensities * thf_lin,
+        "twf_lin": torch.full((n_peaks,), twf_lin),
+        "trf_lin": torch.full((n_peaks,), trf_lin),
+    }
+
+def value_to_index(values, table):
+    span = table[-1] - table[0]
+    indices = ((values - table[0])/span * (len(table)-1)) #.round().type(torch.int64)
+    return indices  
+    
+def generate_theoretical_spectrum(
+    number_of_signals_min, number_of_signals_max,
+    spectrum_width_min, spectrum_width_max,
+    relative_width_min, relative_width_max, 
+    tff_min, tff_max,
+    thf_min, thf_max,
+    trf_min, trf_max,
+    relative_height_min, relative_height_max,
+    multiplicity_j1_min, multiplicity_j1_max,
+    multiplicity_j2_min, multiplicity_j2_max,
+    atom_groups_data,
+    frq_frq
+):
+    number_of_signals = torch.randint(number_of_signals_min, number_of_signals_max+1, [])
+    atom_group_indices = torch.randint(0, len(atom_groups_data), [number_of_signals])
+    width_spectrum = random_loguniform(spectrum_width_min, spectrum_width_max)
+    height_spectrum = random_loguniform(thf_min, thf_max)
+    
+    peak_parameters_data = []
+    theoretical_spectrum = None
+    for atom_group_index in atom_group_indices:
+        relative_intensity, multiplicity1, multiplicity2 = atom_groups_data[atom_group_index]
+        position = random_value(tff_min, tff_max)
+        j1 = random_value(multiplicity_j1_min, multiplicity_j1_max)
+        j2 = random_value(multiplicity_j2_min, multiplicity_j2_max)
+        width = width_spectrum*random_loguniform(relative_width_min, relative_width_max)
+        height = height_spectrum*relative_intensity*random_loguniform(relative_height_min, relative_height_max)
+        gaussian_contribution = random_value(trf_min, trf_max)
+
+        peaks_parameters = generate_multiplet_parameters(multiplicity=(multiplicity1, multiplicity2), tff_lin=position, thf_lin=height, twf_lin= width, trf_lin= gaussian_contribution, j1=j1, j2=j2)
+        peaks_parameters["tff_relative"] = value_to_index(peaks_parameters["tff_lin"], frq_frq)
+        peak_parameters_data.append(peaks_parameters)
+        spectrum_contribution = calculate_theoretical_spectrum(peaks_parameters, frq_frq)
+        if theoretical_spectrum is None:
+            theoretical_spectrum = spectrum_contribution
+        else:
+            theoretical_spectrum += spectrum_contribution
+    return theoretical_spectrum, peak_parameters_data
+
+
+def theoretical_generator(
+    atom_groups_data,
+    pixels=2048, frq_step=11160.7142857 / 32768,
+    number_of_signals_min=1, number_of_signals_max=8,
+    spectrum_width_min=0.2, spectrum_width_max=1,
+    relative_width_min=1, relative_width_max=2,
+    relative_height_min=1, relative_height_max=1,
+    relative_frequency_min=-0.4, relative_frequency_max=0.4,
+    thf_min=1/16, thf_max=16,
+    trf_min=0, trf_max=1,
+    multiplicity_j1_min=0, multiplicity_j1_max=15,
+    multiplicity_j2_min=0, multiplicity_j2_max=15,
+    ):
+    tff_min = relative_frequency_min * pixels * frq_step
+    tff_max = relative_frequency_max * pixels * frq_step
+    frq_frq = torch.arange(-pixels // 2, pixels // 2) * frq_step
+    
+    while True:
+        yield generate_theoretical_spectrum(
+            number_of_signals_min=number_of_signals_min,
+            number_of_signals_max=number_of_signals_max,
+            spectrum_width_min=spectrum_width_min,
+            spectrum_width_max=spectrum_width_max,
+            relative_width_min=relative_width_min,
+            relative_width_max=relative_width_max,
+            relative_height_min=relative_height_min,
+            relative_height_max=relative_height_max,
+            tff_min=tff_min, tff_max=tff_max,
+            thf_min=thf_min, thf_max=thf_max,
+            trf_min=trf_min, trf_max=trf_max,
+            multiplicity_j1_min=multiplicity_j1_min,
+            multiplicity_j1_max=multiplicity_j1_max,
+            multiplicity_j2_min=multiplicity_j2_min,
+            multiplicity_j2_max=multiplicity_j2_max,
+            atom_groups_data=atom_groups_data,
+            frq_frq=frq_frq
+        )
+
+class ResponseLibrary:
+    def __init__(self, reponse_files, normalize=True):
+        self.data = [torch.load(f, map_location='cpu', weights_only=True).flatten(0,-4) for f in reponse_files]
+        if normalize:
+            self.data = [data/torch.sum(data, dim=(-1,), keepdim=True) for data in self.data]
+        lengths = [len(data) for data in self.data]
+        self.start_indices = torch.cumsum(torch.tensor([0] + lengths[:-1]), 0)
+        self.total_length = sum(lengths)
+    
+    def __getitem__(self, idx):
+        if idx >= self.total_length:
+            raise ValueError(f'index {idx} out of range')
+        tensor_index = torch.searchsorted(self.start_indices, idx, right=True) - 1
+        return self.data[tensor_index][idx - self.start_indices[tensor_index]]
+    
+    def __len__(self):
+        return self.total_length
+
+def generator(
+    theoretical_generator_params,
+    response_function_library,
+    response_function_stretch_min=0.5,
+    response_function_stretch_max=2.0,
+    response_function_noise=0.,
+    spectrum_noise_min=0.,
+    spectrum_noise_max=1/64,
+    include_spectrum_data=False,
+    include_peak_mask=False,
+    include_response_function=False,
+    flip_response_function=False
+    
+):
+    for theoretical_spectrum, theoretical_spectrum_data in theoretical_generator(**theoretical_generator_params):
+        # get response function
+        response_function = response_function_library[torch.randint(0, len(response_function_library), [1])][0]
+        # stretch response function
+        padding_size = (response_function.shape[-1] - 1)//2
+        padding_size = round(random_loguniform(response_function_stretch_min, response_function_stretch_max)*padding_size) #torch.randint(round(padding_size*response_function_stretch_min), round(padding_size*response_function_stretch_max), [1]).item()
+        response_function = torch.nn.functional.interpolate(response_function, size=2*padding_size+1, mode='linear')        
+        response_function /= response_function.sum() # normalize sum of response function to 1
+        # add noise to response function
+        response_function += torch.randn(response_function.shape) * response_function_noise
+        response_function /= response_function.sum() # normalize sum of response function to 1
+        if flip_response_function and (torch.rand(1).item() < 0.5):
+            response_function = response_function.flip(-1)
+        # disturbed spectrum
+        disturbed_spectrum = torch.nn.functional.conv1d(theoretical_spectrum, response_function, padding=padding_size)
+        # add noise
+        noised_spectrum = disturbed_spectrum + torch.randn(disturbed_spectrum.shape) * random_value(spectrum_noise_min, spectrum_noise_max)
+        
+        out = {
+            # 'response_function': response_function,
+            'theoretical_spectrum': theoretical_spectrum,
+            'disturbed_spectrum': disturbed_spectrum,
+            'noised_spectrum': noised_spectrum,
+        }
+        if include_response_function:
+            out['response_function'] = response_function
+        if include_spectrum_data:
+            out["theoretical_spectrum_data"] = theoretical_spectrum_data
+        if include_peak_mask:
+            all_peaks_rel = torch.cat([peak_data["tff_relative"] for peak_data in theoretical_spectrum_data])
+            peaks_indices = all_peaks_rel.round().type(torch.int64)
+            out["peaks_mask"] = torch.scatter(torch.zeros(out["theoretical_spectrum"].shape[1]), 0, peaks_indices, 1.).unsqueeze(0)
+
+        yield out
+
+        
+def collate_with_spectrum_data(batch):
+    tensor_keys = set(batch[0].keys())
+    tensor_keys.remove('theoretical_spectrum_data')
+    out = {k: torch.stack([item[k] for item in batch]) for k in tensor_keys}    
+    out["theoretical_spectrum_data"] = [item["theoretical_spectrum_data"] for item in batch]
+    return out
+
+def get_datapipe(
+    response_functions_files,
+    atom_groups_data_file=None,
+    batch_size=64,
+    pixels=2048, frq_step=11160.7142857 / 32768,
+    number_of_signals_min=1, number_of_signals_max=8,
+    spectrum_width_min=0.2, spectrum_width_max=1,
+    relative_width_min=1, relative_width_max=2,
+    relative_height_min=1, relative_height_max=1,
+    relative_frequency_min=-0.4, relative_frequency_max=0.4,
+    thf_min=1/16, thf_max=16,
+    trf_min=0, trf_max=1,
+    multiplicity_j1_min=0, multiplicity_j1_max=15,
+    multiplicity_j2_min=0, multiplicity_j2_max=15,
+    response_function_stretch_min=0.5,
+    response_function_stretch_max=2.0,
+    response_function_noise=0.,
+    spectrum_noise_min=0.,
+    spectrum_noise_max=1/64,
+    include_spectrum_data=False,
+    include_peak_mask=False,
+    include_response_function=False,
+    flip_response_function=False
+):
+    # singlets
+    if atom_groups_data_file is None:
+        atom_groups_data = np.ones((1,3), dtype=int)
+    else:
+        atom_groups_data = np.loadtxt(atom_groups_data_file, usecols=(1,2,3), dtype=int)
+    response_function_library = ResponseLibrary(response_functions_files)
+    g = generator(
+        theoretical_generator_params=dict(
+            atom_groups_data=atom_groups_data,
+            pixels=pixels, frq_step=frq_step,
+            number_of_signals_min=number_of_signals_min, number_of_signals_max=number_of_signals_max,
+            spectrum_width_min=spectrum_width_min, spectrum_width_max=spectrum_width_max,
+            relative_width_min=relative_width_min, relative_width_max=relative_width_max,
+            relative_height_min=relative_height_min, relative_height_max=relative_height_max,
+            relative_frequency_min=relative_frequency_min, relative_frequency_max=relative_frequency_max,
+            thf_min=thf_min, thf_max=thf_max,
+            trf_min=trf_min, trf_max=trf_max,
+            multiplicity_j1_min=multiplicity_j1_min, multiplicity_j1_max=multiplicity_j1_max,
+            multiplicity_j2_min=multiplicity_j2_min, multiplicity_j2_max=multiplicity_j2_max
+        ),
+        response_function_library=response_function_library,
+        response_function_stretch_min=response_function_stretch_min,
+        response_function_stretch_max=response_function_stretch_max,
+        response_function_noise=response_function_noise,
+        spectrum_noise_min=spectrum_noise_min,
+        spectrum_noise_max=spectrum_noise_max,
+        include_spectrum_data=include_spectrum_data,
+        include_peak_mask=include_peak_mask,
+        include_response_function=include_response_function,
+        flip_response_function=flip_response_function
+    )
+    
+    pipe = torchdata.datapipes.iter.IterableWrapper(g, deepcopy=False)
+    pipe = pipe.batch(batch_size)
+    pipe = pipe.collate(collate_fn=collate_with_spectrum_data if include_spectrum_data else None)
+    
+    return pipe

src/models.py

@@ -0,0 +1,105 @@
+import torch
+
+class ConvEncoder(torch.nn.Module):
+    def __init__(self, hidden_dim=64, output_dim=None, dropout=0, kernel_size=7):
+        super().__init__()
+        if output_dim is None:
+            output_dim = hidden_dim
+        self.conv4 = torch.nn.Conv1d(1, hidden_dim, kernel_size)
+        self.conv3 = torch.nn.Conv1d(hidden_dim, hidden_dim, kernel_size)
+        self.conv2 = torch.nn.Conv1d(hidden_dim, hidden_dim, kernel_size)
+        self.conv1 = torch.nn.Conv1d(hidden_dim, output_dim, kernel_size)
+        self.dropout = torch.nn.Dropout(dropout)
+
+    def forward(self, feature):                                        #(samples, 1, 2048)
+        feature = self.dropout(self.conv4(feature))                    #(samples, 64, 2042)
+        feature = feature.relu()
+        feature = self.dropout(self.conv3(feature))                    #(samples, 64, 2036)
+        feature = feature.relu()
+        feature = self.dropout(self.conv2(feature))                    #(samples, 64, 2030)
+        feature = feature.relu()
+        feature = self.dropout(self.conv1(feature))                    #(samples, 64, 2024)
+        return feature
+
+class ConvDecoder(torch.nn.Module):
+    def __init__(self, input_dim=None, hidden_dim=64, output_dim=None, dropout=0, kernel_size=7):
+        super().__init__()
+        if output_dim is None:
+            output_dim = hidden_dim
+        self.convTranspose1 = torch.nn.ConvTranspose1d(input_dim, hidden_dim, kernel_size)
+        self.convTranspose2 = torch.nn.ConvTranspose1d(hidden_dim, hidden_dim, kernel_size)
+        self.convTranspose3 = torch.nn.ConvTranspose1d(hidden_dim, hidden_dim, kernel_size)
+        self.convTranspose4 = torch.nn.ConvTranspose1d(hidden_dim, 1, kernel_size)
+    
+    def forward(self, feature):                                        #(samples, 1, 2048)
+        feature = self.convTranspose1(feature)                         #(samples,  64, 2030)
+        feature = feature.relu()
+        feature = self.convTranspose2(feature)                         #(samples,  64, 2036)
+        feature = feature.relu()
+        feature = self.convTranspose3(feature)                         #(samples,  64, 2042)
+        feature = feature.relu()
+        feature = self.convTranspose4(feature) 
+        return feature
+
+class ResponseHead(torch.nn.Module):
+    def __init__(self, input_dim, output_length, hidden_dims=[128]):
+        super().__init__()
+        response_head_dims = [input_dim]+hidden_dims + [output_length]
+        response_head_layers = [torch.nn.Linear(response_head_dims[0], response_head_dims[1])]
+        for dims_in, dims_out in zip(response_head_dims[1:-1], response_head_dims[2:]):
+            response_head_layers.extend([
+                torch.nn.GELU(),
+                torch.nn.Linear(dims_in, dims_out)
+            ])
+        self.response_head = torch.nn.Sequential(*response_head_layers)
+
+    def forward(self, feature):
+        return self.response_head(feature)
+
+class ShimNetWithSCRF(torch.nn.Module):
+    def __init__(self,
+        encoder_hidden_dims=64,
+        encoder_dropout=0,
+        bottleneck_dim=64,
+        rensponse_length=61,
+        resnponse_head_dims=[128],
+        decoder_hidden_dims=64
+        ):
+        super().__init__()
+        self.encoder = ConvEncoder(hidden_dim=encoder_hidden_dims, output_dim=bottleneck_dim, dropout=encoder_dropout)
+        self.query = torch.nn.Parameter(torch.empty(1, 1, bottleneck_dim))
+        torch.nn.init.xavier_normal_(self.query)
+
+        self.decoder = ConvDecoder(input_dim=2*bottleneck_dim, hidden_dim=decoder_hidden_dims)
+        
+        self.rensponse_length = rensponse_length
+        self.response_head = ResponseHead(bottleneck_dim, rensponse_length, resnponse_head_dims)
+        
+    def forward(self, feature):                                        #(samples,   1, 2048)
+        feature = self.encoder(feature)                                #(samples,  64, 2042)
+        energy = self.query @ feature                                  #(samples,   1, 2024)
+        weight = torch.nn.functional.softmax(energy, 2)                #(samples,   1, 2024)
+        global_features = feature @ weight.transpose(1, 2)                    #(samples,  64,    1)
+        
+        response = self.response_head(global_features.squeeze(-1))
+        
+        feature, global_features = torch.broadcast_tensors(feature, global_features) #(samples,  64, 2048)
+        feature = torch.cat([feature, global_features], 1)                    #(samples, 128, 2024)
+        denoised_spectrum = self.decoder(feature)                            #(samples, 1, 2048)
+        
+        return {
+            'denoised': denoised_spectrum,
+            'response': response,
+            'attention': weight.squeeze(1)
+        }
+
+class Predictor:
+    def __init__(self, model=None, weights_file=None):
+        self.model = model
+        if weights_file is not None:
+            self.model.load_state_dict(torch.load(weights_file, map_location='cpu', weights_only=True))
+
+    def __call__(self, nsf_frq):
+        with torch.no_grad():
+            msf_frq = self.model(nsf_frq[None, None])["denoised"]
+        return msf_frq[0, 0]

train.py

@@ -0,0 +1,141 @@
+import torch, torchaudio
+import numpy as np
+from pathlib import Path
+from omegaconf import OmegaConf
+from hydra.utils import instantiate
+import datetime
+import sys
+import matplotlib.pyplot as plt
+
+
+import matplotlib
+matplotlib.use('Agg')
+
+# silent deprecation_warning() from datapipes
+import warnings
+warnings.filterwarnings("ignore", category=UserWarning, module='torchdata')
+
+from src import models
+from src.generators import get_datapipe
+
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
+if len(sys.argv) < 2:
+    print("Please provide the run directory as an argument.")
+    sys.exit(1)
+run_dir = Path(sys.argv[1])
+
+config = OmegaConf.load(run_dir / "config.yaml")
+
+if (run_dir / "train.txt").is_file():
+    minimum = np.min(np.loadtxt(run_dir / "train.txt")[:,2])
+else:
+    minimum = float("inf")
+
+# initialization        
+model = instantiate({"_target_": f"__main__.models.{config.model.name}", **config.model.kwargs}).to(device)
+model_weights_file = run_dir / f'model.pt'
+optimizer = torch.optim.Adam(model.parameters())
+optimizer_weights_file = run_dir / f'optimizer.pt'
+
+def evaluate_model(stage=0, epoch=0):
+    plot_dir = run_dir / "plots" / f"{stage}_{epoch}"
+    plot_dir.mkdir(exist_ok=True, parents=True)
+
+    torch.save(model.state_dict(), plot_dir / "model.pt")
+    torch.save(optimizer.state_dict(), plot_dir / "optimizer.pt")
+    
+    num_plots = config.logging.num_plots
+    pipe = get_datapipe(
+            **config.data,
+            include_response_function=True,
+            batch_size=num_plots
+        )
+    batch = next(iter(pipe))
+
+    with torch.no_grad():
+        out = model(batch['noised_spectrum'].to(device))
+        noised_est = torchaudio.functional.convolve(out['denoised'], out['response'].flip(dims=(-1,)).unsqueeze(1), mode="same").cpu()
+
+        for i in range(num_plots):
+            plt.figure(figsize=(30,6))
+            plt.plot(batch['theoretical_spectrum'].cpu().numpy()[i,0])
+            plt.plot(out['denoised'].cpu().numpy()[i,0])
+            plt.savefig(plot_dir / f"{i:03d}_spectrum_clean.png")
+
+            plt.figure(figsize=(30,6))
+            plt.plot(batch['noised_spectrum'].cpu().numpy()[i,0])
+            plt.plot(noised_est.cpu().numpy()[i,0])
+            plt.savefig(plot_dir / f"{i:03d}_spectrum_noise.png")
+
+            plt.figure(figsize=(10,6))
+            plt.plot(batch['response_function'].cpu().numpy()[i,0,0])
+            plt.plot(out['response'].cpu().numpy()[i])
+            plt.savefig(plot_dir / f"{i:03d}_response.png")
+
+            if "attention" in out:
+                plt.figure(figsize=(10, 6))
+                plt.plot(out['attention'].cpu().numpy()[i])
+                plt.savefig(plot_dir / f"{i:03d}_attention.png")
+            
+            plt.close("all")
+
+for i_stage, training_stage in enumerate(config.training):
+    if model_weights_file.is_file():
+        model.load_state_dict(torch.load(model_weights_file, weights_only=True))
+
+    if optimizer_weights_file.is_file():
+        optimizer.load_state_dict(torch.load(optimizer_weights_file, weights_only=True))
+    optimizer.param_groups[0]['lr'] = training_stage.learning_rate
+    
+    pipe = get_datapipe(
+        **config.data,
+        include_response_function=True,
+        batch_size=training_stage.batch_size
+    )
+
+    losses_history = []
+    losses_history_limit = 64*100 // training_stage.batch_size
+    
+    last_evaluation = 0
+    for epoch, batch in pipe.enumerate():
+        
+        # logging
+        iters_done = epoch*training_stage.batch_size
+        if (iters_done - last_evaluation) > config.logging.step:
+            evaluate_model(i_stage, epoch)
+            last_evaluation = iters_done
+            
+        if  iters_done > training_stage.max_iters:
+            evaluate_model(i_stage, epoch)
+            break
+        
+        # run model
+        out = model(batch['noised_spectrum'].to(device))
+        # calculate losses
+        loss_response = torch.nn.functional.mse_loss(out['response'], batch['response_function'].squeeze(dim=(1,2)).to(device))
+        loss_clean = torch.nn.functional.mse_loss(out['denoised'], batch['theoretical_spectrum'].to(device))
+        noised_est = torchaudio.functional.convolve(out['denoised'], out['response'].flip(dims=(-1,)).unsqueeze(1), mode="same")
+        loss_noised = torch.nn.functional.mse_loss(noised_est, batch['noised_spectrum'].to(device))
+        loss = config.losses_weights.response*loss_response + config.losses_weights.clean*loss_clean + config.losses_weights.noised*loss_noised
+        
+        # logging
+        losses_history.append(loss_clean.item())
+        losses_history = losses_history[-losses_history_limit:]
+        loss_avg = sum(losses_history)/len(losses_history)
+        message = f"{epoch:7d} {loss:0.3e} {loss_avg:0.3e} {loss_clean:0.3e} {loss_response:0.3e} {loss_noised:0.3e}"
+        # message = '%7i %.3e %.3e %.3e' % (epoch, loss, regress, classify)
+        with open(run_dir / f'train.txt', 'a') as f:
+            f.write(message + '\n')
+        print(message, flush = True)
+        
+        # save best
+        if loss_avg < minimum:
+            minimum = loss_avg
+            torch.save(model.state_dict(), model_weights_file)
+            torch.save(optimizer.state_dict(),optimizer_weights_file)
+        
+        # update weights
+        loss.backward()
+        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.)
+        optimizer.step()
+        optimizer.zero_grad()