Fix wave training script imports

README.md

@@ -31,6 +31,7 @@ Before diving into waveform vector exploitation with WAVE, ensure your environme
   ```bash
   python -m pip install -U -r requirements.txt
   ```
+
   - `numpy`
   - `scipy`
   - `torch`

gcp/wave_pytorch_gcp.py

@@ -4,10 +4,11 @@
 import os
 import time
 
+import numpy as np
 import scipy.io
 import torch
 
-from utils import *
+from utils.utils import normalize, splitdata, stdpt
 
 # set printoptions
 torch.set_printoptions(linewidth=320, precision=8)
@@ -41,11 +42,11 @@ def runexample(H, model, str, lr=0.001, amsgrad=False):
     mat = scipy.io.loadmat(pathd + data)
     x = mat["inputs"]  # inputs (nx512) [waveform1 waveform2]
     y = mat["outputs"][:, 1:2]  # outputs (nx4) [position(mm), time(ns), PE, E(MeV)]
-    nz, nx = x.shape
+    _nz, _nx = x.shape
     ny = y.shape[1]
 
     x, _, _ = normalize(x, 1)  # normalize each input row
-    y, ymu, ys = normalize(y, 0)  # normalize each output column
+    y, _ymu, ys = normalize(y, 0)  # normalize each output column
     x, y = torch.Tensor(x), torch.Tensor(y)
     x, y, xv, yv, xt, yt = splitdata(x, y, train=0.70, validate=0.15, test=0.15, shuffle=True)
     labels = ["train", "validate", "test"]
@@ -209,8 +210,7 @@ def forward(self, x):
 
 
 def tslr():  # TS learning rate
-    """Generate and save learning rate (LR) logs for time-series models with varying LRs using WAVE and TanH
-    activation.
+    """Generate and save learning rate (LR) logs for time-series models with varying LRs using WAVE and TanH activation.
     """
     tsv = np.logspace(-5, -2, 13)
     tsy = []
@@ -224,7 +224,7 @@ def tsams():  # TS AMSgrad
     tsv = [False, True]
     tsy = []
     for a in tsv:
-        tsy.extend(runexample(H, model=WAVE(H), str=f".TanhAMS{str(a)}", amsgrad=a) for _ in range(3))
+        tsy.extend(runexample(H, model=WAVE(H), str=f".TanhAMS{a!s}", amsgrad=a) for _ in range(3))
     scipy.io.savemat(f"{pathr}TS.AMSgrad.mat", dict(tsv=tsv, tsy=np.array(tsy)))
 
 
@@ -243,7 +243,13 @@ def tsshape():  # TS network shape
 
     # tsv = ['Tanh', 'LogSigmoid', 'Softsign', 'ELU']
     # tsv = np.logspace(-4, -2, 11)
-    tsv = [[512, 23, 1], [512, 64, 8, 1], [512, 108, 23, 5, 1], [512, 147, 42, 12, 3, 1], [512, 181, 64, 23, 8, 3, 1]]
+    tsv = [
+        [512, 23, 1],
+        [512, 64, 8, 1],
+        [512, 108, 23, 5, 1],
+        [512, 147, 42, 12, 3, 1],
+        [512, 181, 64, 23, 8, 3, 1],
+    ]
     H = tsv[0]
 
     class WAVE(torch.nn.Module):

train.py

@@ -3,11 +3,13 @@
 import argparse
 import os
 
+import numpy as np
 import scipy.io
+import torch
 import torch.nn as nn
 
-from utils.torch_utils import *
-from utils.utils import *
+from utils.torch_utils import init_seeds, select_device
+from utils.utils import model_info, normalize, patienceStopper, splitdata
 
 torch.backends.cudnn.benchmark = True  # unsuitable for multiscale
 
@@ -35,11 +37,11 @@ def train(H, model, str, lr=0.001):
     mat = scipy.io.loadmat(pathd + data)
     x = mat["inputs"][:]  # inputs (nx512) [waveform1 waveform2]
     y = mat["outputs"][:, 0:2]  # outputs (nx4) [position(mm), time(ns), PE, E(MeV)]
-    nz, nx = x.shape
+    _nz, _nx = x.shape
     ny = y.shape[1]
 
     x, _, _ = normalize(x, 1)  # normalize each input row
-    y, ymu, ys = normalize(y, 0)  # normalize each output column
+    y, _ymu, ys = normalize(y, 0)  # normalize each output column
     x, y = torch.Tensor(x), torch.Tensor(y)
     x, y, xv, yv, xt, yt = splitdata(x, y, train=0.70, validate=0.15, test=0.15, shuffle=False)
 
@@ -116,7 +118,10 @@ def train(H, model, str, lr=0.001):
             std[i] = r.std(0).cpu().numpy() * ys
         print(f"{loss[i]:.5f} {std[i, :]} {labels[i]}")
 
-    scipy.io.savemat(pathr + name + ".mat", dict(bestepoch=stopper.bestloss, loss=loss, std=std, L=L, name=name))
+    scipy.io.savemat(
+        pathr + name + ".mat",
+        dict(bestepoch=stopper.bestloss, loss=loss, std=std, L=L, name=name),
+    )
     # files.download(pathr + name + '.mat')
 
     return np.concatenate(([stopper.bestloss], np.array(loss), np.array(std.ravel())))
@@ -185,20 +190,37 @@ def __init__(self, n_out=2):
         super().__init__()
         n = 32
         self.layer1 = nn.Sequential(
-            nn.Conv2d(in_channels=2, out_channels=n, kernel_size=(1, 33), stride=(1, 2), padding=(0, 16), bias=False),
+            nn.Conv2d(
+                in_channels=2,
+                out_channels=n,
+                kernel_size=(1, 33),
+                stride=(1, 2),
+                padding=(0, 16),
+                bias=False,
+            ),
             nn.BatchNorm2d(n),
             nn.LeakyReLU(0.1),
         )
         self.layer2 = nn.Sequential(
             nn.Conv2d(
-                in_channels=n, out_channels=n * 2, kernel_size=(1, 17), stride=(1, 2), padding=(0, 8), bias=False
+                in_channels=n,
+                out_channels=n * 2,
+                kernel_size=(1, 17),
+                stride=(1, 2),
+                padding=(0, 8),
+                bias=False,
             ),
             nn.BatchNorm2d(n * 2),
             nn.LeakyReLU(0.1),
         )
         self.layer3 = nn.Sequential(
             nn.Conv2d(
-                in_channels=n * 2, out_channels=n * 4, kernel_size=(1, 9), stride=(1, 2), padding=(0, 4), bias=False
+                in_channels=n * 2,
+                out_channels=n * 4,
+                kernel_size=(1, 9),
+                stride=(1, 2),
+                padding=(0, 4),
+                bias=False,
             ),
             nn.BatchNorm2d(n * 4),
             nn.LeakyReLU(0.1),

train_tf.py

@@ -3,12 +3,13 @@
 import os
 import time
 
+import numpy as np
 import plotly.graph_objs as go
 import scipy.io
 import tensorflow as tf
 from plotly.offline import plot
 
-from utils.utils import *
+from utils.utils import normalize, splitdata, stdtf
 
 tf.enable_eager_execution()
 
@@ -38,7 +39,7 @@ def runexample(H, model, str):
     mat = scipy.io.loadmat(path + data)
     x = mat["inputs"]  # inputs (nx512) [waveform1 waveform2]
     y = mat["outputs"][:, 0:2]  # outputs (nx4) [position(mm), time(ns), PE, E(MeV)]
-    nz, nx = x.shape
+    _nz, _nx = x.shape
     ny = y.shape[1]
 
     if model is None:
@@ -56,7 +57,7 @@ def runexample(H, model, str):
         )
 
     x, _, _ = normalize(x, 1)  # normalize each input row
-    y, ymu, ys = normalize(y, 0)  # normalize each output column
+    y, _ymu, ys = normalize(y, 0)  # normalize each output column
     x, y, xv, yv, xt, yt = splitdata(x, y, train=0.70, validate=0.15, test=0.15, shuffle=False)
     labels = ["train", "validate", "test"]
 
@@ -104,7 +105,10 @@ def criteria(y_pred, y):  # MSE
                 ticb = time.time()
 
             # Apply the gradient to the model
-            optimizer.apply_gradients(zip(grads, model.variables), global_step=tf.train.get_or_create_global_step())
+            optimizer.apply_gradients(
+                zip(grads, model.variables),
+                global_step=tf.train.get_or_create_global_step(),
+            )
         else:
             print("WARNING: Validation loss still decreasing after %g epochs (train longer)." % (i + 1))
         # torch.save(best[2], path + 'models/' + name + '.pt')
@@ -122,12 +126,15 @@ def criteria(y_pred, y):  # MSE
     data = []
     for i, s in enumerate(labels):
         data.append(go.Scatter(x=np.arange(epochs), y=L[:, i], mode="markers+lines", name=s))
-    layout = go.Layout(xaxis=dict(type="linear", autorange=True), yaxis=dict(type="log", autorange=True))
+    layout = go.Layout(
+        xaxis=dict(type="linear", autorange=True),
+        yaxis=dict(type="log", autorange=True),
+    )
     # configure_plotly_browser_state()
     plot(go.Figure(data=data, layout=layout))
 
 
 if __name__ == "__main__":
     H = [128, 32, 8]
     for i in range(1):
-        runexample(H, None, f".{str(i)}")
+        runexample(H, None, f".{i!s}")

utils/utils.py

@@ -38,7 +38,14 @@ def splitdata(x, y, train=0.7, validate=0.15, test=0.15, shuffle=False):  # spli
     i = round(n * train)  # train
     j = round(n * validate) + i  # validate
     k = round(n * test) + j  # test
-    return x[:i], y[:i], x[i:j], y[i:j], x[j:k], y[j:k]  # xy train, xy validate, xy test
+    return (
+        x[:i],
+        y[:i],
+        x[i:j],
+        y[i:j],
+        x[j:k],
+        y[j:k],
+    )  # xy train, xy validate, xy test
 
 
 def stdpt(r, ys):  # MSE loss + standard deviation (pytorch)