update

Signed-off-by: Yiyu Ni <niyiyu@uw.edu>

.pre-commit-config.yaml

@@ -0,0 +1,11 @@
+repos:
+  - repo: https://github.com/PyCQA/isort
+    rev: 5.12.0
+    hooks:
+      - id: isort
+
+  - repo: https://github.com/psf/black
+    rev: 23.7.0
+    hooks:
+    - id: black
+

data/download.py

@@ -1,8 +1,9 @@
 # Download data from Cook Inlet DAS Experiment
 # Yiyu Ni (niyiyu@uw.edu)
 
-import pandas as pd
 import os
+
+import pandas as pd
 import wget
 
 # download data for TERRA?
@@ -34,8 +35,8 @@
 
 # download data for each event
 for idx, i in df.iterrows():
-    date = i['date']
-    eid = i['IRIS ID']
+    date = i["date"]
+    eid = i["IRIS ID"]
     print("=============================================")
     print(f"INFO: processing event {eid} on {date}.")
     if TERRA:

scripts/datasets.py

@@ -1,10 +1,12 @@
 import numpy as np
 import torch
 
+
 class DASDataset(torch.utils.data.Dataset):
-    '''
+    """
     torch dataset class for SHallow REcurrent Network.
-    '''
+    """
+
     def __init__(self, inputs, outputs):
         if isinstance(inputs, torch.Tensor):
             self.inputs = inputs
@@ -21,11 +23,13 @@ def __getitem__(self, index):
         Y = self.outputs[index, :]
 
         return X, Y
-
+
+
 class DASDataset2DOnTheFly(torch.utils.data.Dataset):
-    '''
+    """
     torch dataset class for Random Fourier Feature Network.
-    '''
+    """
+
     def __init__(self, B, T, X, outputs):
         self.B = torch.Tensor(B)
         self.T = torch.Tensor(T)
@@ -38,21 +42,25 @@ def __len__(self):
     def __getitem__(self, index):
         _t = self.T[index]
         _x = self.X[index]
-        _Bv = torch.matmul(torch.Tensor(torch.concat([_t, _x], axis = -1)), self.B.T)
-        X = torch.cat([torch.cos(2*torch.pi*_Bv), torch.sin(2*torch.pi*_Bv)], axis = -1)
+        _Bv = torch.matmul(torch.Tensor(torch.concat([_t, _x], axis=-1)), self.B.T)
+        X = torch.cat(
+            [torch.cos(2 * torch.pi * _Bv), torch.sin(2 * torch.pi * _Bv)], axis=-1
+        )
         Y = self.outputs[index, :]
 
         return X, Y
 
+
 class DASDataset2D(torch.utils.data.Dataset):
-    '''
+    """
     torch dataset class for SIREN model.
-    '''
+    """
+
     def __init__(self, T, X, outputs):
         self.T = torch.Tensor(T)
         self.X = torch.Tensor(X)
         self.outputs = torch.Tensor(outputs)
-        
+
     def __len__(self):
         return len(self.outputs)
 
@@ -62,4 +70,4 @@ def __getitem__(self, index):
         X = self.X[index]
         Y = self.outputs[index]
 
-        return torch.concat([T, X], axis=-1), Y
+        return torch.concat([T, X], axis=-1), Y

scripts/models.py

@@ -1,21 +1,25 @@
 import torch
 
+
 class SHRED(torch.nn.Module):
     def __init__(self, input_size, hidden_size, output_size, num_layers):
         super().__init__()
-        self.lstm = torch.nn.LSTM(input_size, hidden_size, num_layers, batch_first=True, dropout=0.1)
-        self.sdn1 = torch.nn.Linear(hidden_size, output_size//2)
-        self.sdn3 = torch.nn.Linear(output_size//2, output_size)
+        self.lstm = torch.nn.LSTM(
+            input_size, hidden_size, num_layers, batch_first=True, dropout=0.1
+        )
+        self.sdn1 = torch.nn.Linear(hidden_size, output_size // 2)
+        self.sdn3 = torch.nn.Linear(output_size // 2, output_size)
         self.relu = torch.nn.ReLU()
 
     def forward(self, x):
-        x = self.lstm(x)[1][0][-1] # should be -1
+        x = self.lstm(x)[1][0][-1]  # should be -1
         x = self.relu(self.sdn1(x))
         x = self.sdn3(x)
-        return 
-
+        return
+
+
 class RandomFourierFeature(torch.nn.Module):
-    def __init__(self, nfeature, n_layers, n_outputs = 1):
+    def __init__(self, nfeature, n_layers, n_outputs=1):
         super().__init__()
         self.n_layers = n_layers
         self.n_outputs = n_outputs
@@ -33,9 +37,10 @@ def forward(self, x):
         for i in range(self.n_layers):
             layer = getattr(self, f"ln{i+1}")
             x = self.relu(layer(x))
-        x = 2*self.sigmoid(self.outputs(x))-1
+        x = 2 * self.sigmoid(self.outputs(x)) - 1
         return x
 
+
 class SIREN(torch.nn.Module):
     def __init__(self, n_input, n_output, n_layers, n_units, omega):
         super().__init__()
@@ -62,6 +67,6 @@ def forward(self, x):
             layer = getattr(self, f"ln{i+1}")
             bias = getattr(self, f"ln{i+1}_bias")
             x = torch.sin(self.omega * layer(x) + bias)
-        x = 2*self.sigmoid(self.outputs(x) + self.outputs_bias) - 1
+        x = 2 * self.sigmoid(self.outputs(x) + self.outputs_bias) - 1
 
-        return x
+        return x

scripts/rffn/run.py

@@ -1,7 +1,8 @@
+import argparse
+import time
+
 import numpy as np
 from tqdm import tqdm
-import time
-import argparse
 
 parser = argparse.ArgumentParser(description="Fitting GCI DAS data")
 parser.add_argument("-f", "--file", required=True)
@@ -16,12 +17,14 @@
 gpu_id = args.gpu
 
 import os
-os.environ['CUDA_VISIBLE_DEVICES'] = str(gpu_id)
 
+os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu_id)
+
+import h5py
 import torch
-from torch.utils.data import DataLoader
 from torch.nn import MSELoss
-import h5py
+from torch.utils.data import DataLoader
+
 
 class DASDatasetOnTheFly(torch.utils.data.Dataset):
     def __init__(self, B, T, X, outputs):
@@ -36,13 +39,16 @@ def __len__(self):
     def __getitem__(self, index):
         _t = self.T[index]
         _x = self.X[index]
-        _Bv = torch.matmul(torch.Tensor(torch.concat([_t, _x], axis = -1)), self.B.T)
-        X = torch.cat([torch.cos(2*torch.pi*_Bv), torch.sin(2*torch.pi*_Bv)], axis = -1)
+        _Bv = torch.matmul(torch.Tensor(torch.concat([_t, _x], axis=-1)), self.B.T)
+        X = torch.cat(
+            [torch.cos(2 * torch.pi * _Bv), torch.sin(2 * torch.pi * _Bv)], axis=-1
+        )
         y = self.outputs[index, :]
         return X, y
-
+
+
 class RandomFourierFeature(torch.nn.Module):
-    def __init__(self, nfeature, n_layers, n_outputs = 1):
+    def __init__(self, nfeature, n_layers, n_outputs=1):
         super().__init__()
         self.n_layers = n_layers
         self.n_outputs = n_outputs
@@ -60,13 +66,16 @@ def forward(self, x):
         for i in range(self.n_layers):
             layer = getattr(self, f"ln{i+1}")
             x = self.relu(layer(x))
-        x = 2*self.sigmoid(self.outputs(x))-1
+        x = 2 * self.sigmoid(self.outputs(x)) - 1
         return x
-
-print(f"====REPORT: working on {filetime}, channel {channel}-{channel+1000}, GPU {gpu_id}====")
 
-df = h5py.File(filename, 'r')
-data = df['/Acquisition/Raw[0]/RawData'][:, channel:channel+1000].T
+
+print(
+    f"====REPORT: working on {filetime}, channel {channel}-{channel+1000}, GPU {gpu_id}===="
+)
+
+df = h5py.File(filename, "r")
+data = df["/Acquisition/Raw[0]/RawData"][:, channel : channel + 1000].T
 df.close()
 
 nc, ns = data.shape
@@ -77,18 +86,21 @@ def forward(self, x):
 data = torch.Tensor(data)
 
 in_x, in_t = torch.meshgrid(torch.arange(nc), torch.arange(ns))
-T = in_t/(ns-1); X = in_x/(nc-1)
+T = in_t / (ns - 1)
+X = in_x / (nc - 1)
 
-device = torch.device('cuda')
+device = torch.device("cuda")
 
 n_feature = 196
 
-B = torch.Tensor(np.random.normal(scale = 20, size = (n_feature, 2)))
+B = torch.Tensor(np.random.normal(scale=20, size=(n_feature, 2)))
 
-datasetonthefly = DASDatasetOnTheFly(B.to(device), 
-                                    T.reshape([-1, 1]).to(device), 
-                                    X.reshape([-1, 1]).to(device), 
-                                    data.reshape([-1, 1]).to(device))
+datasetonthefly = DASDatasetOnTheFly(
+    B.to(device),
+    T.reshape([-1, 1]).to(device),
+    X.reshape([-1, 1]).to(device),
+    data.reshape([-1, 1]).to(device),
+)
 
 data_loader = DataLoader(datasetonthefly, batch_size=4096, shuffle=True, num_workers=0)
 
@@ -98,31 +110,31 @@ def forward(self, x):
         i.weight.data.normal_(mean=0.0, std=0.1)
         i.bias.data.normal_(mean=0.0, std=0.1)
 
-model.to(device);
+model.to(device)
 
-optimizer = torch.optim.Adam(model.parameters(), lr = 3e-4)
+optimizer = torch.optim.Adam(model.parameters(), lr=3e-4)
 loss_fn = MSELoss()
 nepoch = 8
 train_loss_log = []
 t0 = time.time()
 for t in range(nepoch):
     model.train()
     train_loss = 0
-    
+
     for batch_id, batch in enumerate(data_loader):
         pred = model(batch[0].cuda())
         loss = loss_fn(pred, batch[1].cuda())
         optimizer.zero_grad()
         loss.backward()
         optimizer.step()
         train_loss += loss.item()
-    
+
     train_loss /= len(data_loader)
-    
+
     train_loss_log.append(train_loss)
-    print(f"{t}: train loss: %.6f"  % train_loss)
+    print(f"{t}: train loss: %.6f" % train_loss)
     if train_loss < 1e-4:
         break
 print(time.time() - t0)
 
-torch.save(model.state_dict(), f"./weights/rff_{filetime}_{channel}-{channel+1000}.pt")
+torch.save(model.state_dict(), f"./weights/rff_{filetime}_{channel}-{channel+1000}.pt")

scripts/rffn/submit.py

@@ -1,26 +1,31 @@
 from datetime import datetime
 
-
 from mpi4py import MPI
+
 comm = MPI.COMM_WORLD
 rank = comm.Get_rank()
 size = comm.Get_size()
 
-import os
 import glob
+import os
+
 import numpy as np
 
-with open("./log.log", 'a') as f:
+with open("./log.log", "a") as f:
     f.write(f"{rank}: ==============={datetime.now()}===========\n")
 
-flist = sorted(glob.glob("/home/niyiyu/Research/DAS-NIR/datasets/KKFL-S_FiberA_2p5Hz/decimator3_2023-07-16_*_UTC.h5"))
+flist = sorted(
+    glob.glob(
+        "/home/niyiyu/Research/DAS-NIR/datasets/KKFL-S_FiberA_2p5Hz/decimator3_2023-07-16_*_UTC.h5"
+    )
+)
 
 for f in flist:
-    for ic, c in enumerate(np.linspace(1000, 6000, 6).astype('int')):
+    for ic, c in enumerate(np.linspace(1000, 6000, 6).astype("int")):
         if ic % size == rank:
             gpuid = rank % 4
             command = f"/home/niyiyu/anaconda3/envs/dasnir/bin/python /home/niyiyu/Research/DAS-NIR/gci-summary/tests/rff/run.py -f {f} -c {c} -g 3"
             os.system(command)
 
-with open("./log.log", 'w') as f:
-    f.write(f"{rank}: ==============={datetime.now()}===========\n")
+with open("./log.log", "w") as f:
+    f.write(f"{rank}: ==============={datetime.now()}===========\n")