removed unused variables and using fstrings

learner.py

@@ -15,10 +15,10 @@ def __init__(self, args):
         '''
         self.alpha = args.alpha
         config = []
+        # kernel_size = 2 if block==args.num_down_conv-1 else 3
+        kernel_size = 3
         for block in range(args.num_down_conv):
 
-            # kernel_size = 2 if block==args.num_down_conv-1 else 3
-            kernel_size = 3
             out_channels = (2**block)*args.hidden_dim
             if(block == 0):
                 config += [
@@ -104,7 +104,7 @@ def __init__(self, args):
         self.lr_scheduler = None # will be initialized in waveynet_trainer.py
         self.residual_terms = None # to store the residual connect for addition later
 
-        for i, (name, param) in enumerate(self.config):
+        for name, param in self.config:
             if name == 'conv2d':
                 # [ch_out, ch_in, kernelsz, kernelsz]
                 w = nn.Parameter(torch.ones(*param[:4]), requires_grad=True)
@@ -144,12 +144,12 @@ def extra_repr(self):
         for name, param in self.config:
             if name == 'conv2d':
                 tmp = 'conv2d:(ch_in:%d, ch_out:%d, k:%dx%d, stride:%d, padding:%d)'\
-                      %(param[1], param[0], param[2], param[3], param[4], param[5],)
+                          %(param[1], param[0], param[2], param[3], param[4], param[5],)
                 info += tmp + '\n'
 
             elif name == 'conv2d_b':
                 tmp = 'conv2d:(ch_in:%d, ch_out:%d, k:%dx%d, stride:%d, padding:%d)'\
-                      %(param[1], param[0], param[2], param[3], param[4], param[5],)
+                          %(param[1], param[0], param[2], param[3], param[4], param[5],)
                 info += tmp + '\n'
 
             elif name == 'leakyrelu':
@@ -160,8 +160,8 @@ def extra_repr(self):
                 tmp = 'max_pool2d:(k:%s, stride:%s, padding:%d)'%(param[0], param[1], param[2])
                 info += tmp + '\n'
             elif name in ['flatten', 'tanh', 'relu', 'upsample', 'reshape', 'sigmoid',\
-                          'use_logits', 'bn', 'residual']:
-                tmp = name + ':' + str(tuple(param))
+                              'use_logits', 'bn', 'residual']:
+                tmp = f'{name}:{tuple(param)}'
                 info += tmp + '\n'
             else:
                 raise NotImplementedError
@@ -183,7 +183,14 @@ def forward(self, x, vars=None, bn_training=True):
 
         blocks = []
         for name, param in self.config:
-            if name == 'conv2d':
+            if name == 'bn':
+                w, b = vars[idx], vars[idx + 1]
+                running_mean, running_var = self.vars_bn[bn_idx], self.vars_bn[bn_idx+1]
+                x = F.batch_norm(x, running_mean, running_var, weight=w, bias=b, training=bn_training)
+                idx += 2
+                bn_idx += 2
+
+            elif name == 'conv2d':
                 # periodic padding
                 kernel_width = param[2]
                 if kernel_width % 2:
@@ -222,16 +229,16 @@ def forward(self, x, vars=None, bn_training=True):
                 x = F.conv2d(x_pad, w, b, stride=param[4], padding=0)
                 idx += 2
 
-            elif name == 'bn':
-                w, b = vars[idx], vars[idx + 1]
-                running_mean, running_var = self.vars_bn[bn_idx], self.vars_bn[bn_idx+1]
-                x = F.batch_norm(x, running_mean, running_var, weight=w, bias=b, training=bn_training)
-                idx += 2
-                bn_idx += 2
-
             elif name == 'leakyrelu':
                 x = F.leaky_relu(x, negative_slope=param[0], inplace=param[1])
 
+            elif name == 'max_pool2d':
+                blocks.append(x)
+                x = F.max_pool2d(x, param[0], stride=param[1], padding=param[2])
+
+            elif name == 'residual':
+                x = x + self.residual_terms
+
             elif name == 'upsample':
                 if first_upsample:
                     first_upsample = False
@@ -240,13 +247,6 @@ def forward(self, x, vars=None, bn_training=True):
                 x = F.interpolate(x, size=(shortcut.shape[2],shortcut.shape[3]), mode='nearest')
                 x = torch.cat([shortcut,x], dim=1) # batch, channels, h, w
 
-            elif name == 'residual':
-                x = x + self.residual_terms
-
-            elif name == 'max_pool2d':
-                blocks.append(x)
-                x = F.max_pool2d(x, param[0], stride=param[1], padding=param[2])
-
             else:
                 print(name)
                 raise NotImplementedError

simulation_dataset.py

@@ -21,30 +21,30 @@ def __init__(self, data_folder, local_data, total_sample_number = None, transfor
 
         print("Loading training data...")
         print("Note: The training and test datasets are very large, which can take "\
-               "on the order of hours to complete, depending on your internet download "\
-               "speed. If you have a slower internet download speed, consider downloading the "\
-               "data locally from http://metanet.stanford.edu/search/waveynet-study/ so "\
-               "that you do not have to re-download the data each time you run the code."\
-               "See the README file for more information.")
+                   "on the order of hours to complete, depending on your internet download "\
+                   "speed. If you have a slower internet download speed, consider downloading the "\
+                   "data locally from http://metanet.stanford.edu/search/waveynet-study/ so "\
+                   "that you do not have to re-download the data each time you run the code."\
+                   "See the README file for more information.")
         #load the training data, either locally or from Metanet, as specified by
         #the boolean args.local_data
-        if(local_data):
-            train_data = np.load(data_folder+'/train_ds.npz')
+        if local_data:
+            train_data = np.load(f'{data_folder}/train_ds.npz')
         else:
             response = requests.get('http://metanet.stanford.edu/static/search/waveynet/'\
-                                    'data/train_ds.npz')
+                                        'data/train_ds.npz')
             response.raise_for_status()
             train_data = np.load(io.BytesIO(response.content))
 
         print("Loading test data...")
 
         #load the test data, either locally or from Metanet, as specified by
         #the boolean args.local_data
-        if(local_data):
-            test_data = np.load(data_folder+'/test_ds.npz')
+        if local_data:
+            test_data = np.load(f'{data_folder}/test_ds.npz')
         else:
             response = requests.get('http://metanet.stanford.edu/static/search/waveynet/'\
-                                    'data/test_ds.npz')
+                                        'data/test_ds.npz')
             response.raise_for_status()
             test_data = np.load(io.BytesIO(response.content))
 
@@ -72,9 +72,9 @@ def __init__(self, data_folder, local_data, total_sample_number = None, transfor
 
         for dev in range(len(self.input_imgs)):
                 self.input_imgs[dev][np.where(self.input_imgs[dev]==self.input_imgs[dev,0,0,0])]=\
-                                     die_change[dev]
+                                         die_change[dev]
                 self.input_imgs[dev]=(self.input_imgs[dev]*(np.sqrt(self.refr_idx_vec[dev]) - \
-                                     n_air) + n_air)**2
+                                         n_air) + n_air)**2
 
         #Use only a subset of the dataset if specified to do so by args.total_sample_number
         if total_sample_number:

waveynet_trainer.py

@@ -147,8 +147,15 @@ def H_to_H(Hz_R: np.array, Hz_I: np.array, dL: float, omega: float, eps_grid: np
 
     FD_Ex = Hz_to_Ex(Hz_R, Hz_I, dL, omega, eps_grid, eps_0)
     FD_Ez = Hz_to_Ey(Hz_R, Hz_I, dL, omega, eps_grid, eps_0)
-    FD_H = E_to_Hz(FD_Ez[:, 0, :-1], FD_Ez[:, 1, :-1], FD_Ex[:, 0], FD_Ex[:, 1], dL, omega, mu_0)
-    return FD_H
+    return E_to_Hz(
+        FD_Ez[:, 0, :-1],
+        FD_Ez[:, 1, :-1],
+        FD_Ex[:, 0],
+        FD_Ex[:, 1],
+        dL,
+        omega,
+        mu_0,
+    )
 
 def regConstScheduler(epoch, args, last_epoch_data_loss, last_epoch_physical_loss):
     '''
@@ -216,14 +223,16 @@ def main(args):
 
     test_loader = DataLoader(test_ds, batch_size=args.batch_size, shuffle=True, num_workers=0)
 
-    train_mean = 0
-    test_mean = 0
+    train_mean = sum(
+        torch.mean(torch.abs(sample_batched["field"]))
+        for sample_batched in train_loader
+    )
+
+    test_mean = sum(
+        torch.mean(torch.abs(sample_batched["field"]))
+        for sample_batched in test_loader
+    )
 
-    # first get the mean-absolute-field value:
-    for sample_batched in train_loader:
-        train_mean += torch.mean(torch.abs(sample_batched["field"]))
-    for sample_batched in test_loader:
-        test_mean += torch.mean(torch.abs(sample_batched["field"]))
     train_mean /= len(train_loader)
     test_mean /= len(test_loader)