cnnkan

tasks/audio.py

@@ -29,7 +29,6 @@ DESCRIPTION = "Conformer"
 ROUTE = "/audio"
 
 
-
 @router.post(ROUTE, tags=["Audio Task"],
              description=DESCRIPTION)
 async def evaluate_audio(request: AudioEvaluationRequest):
@@ -133,11 +132,11 @@ async def evaluate_audio(request: AudioEvaluationRequest):
     
     return results
 
-# if __name__ == "__main__":
-#     sample_request = AudioEvaluationRequest(
-#         dataset_name="rfcx/frugalai",  # Replace with actual dataset name
-#         test_size=0.2,  # Example values
-#         test_seed=42
-#     )
-# #
-#     asyncio.run(evaluate_audio(sample_request))
+if __name__ == "__main__":
+    sample_request = AudioEvaluationRequest(
+        dataset_name="rfcx/frugalai",  # Replace with actual dataset name
+        test_size=0.2,  # Example values
+        test_seed=42
+    )
+#
+    asyncio.run(evaluate_audio(sample_request))

tasks/inr_database/inr_database.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:13ead5ed23d1fa59062f4872fd784f609575fa7dd4876ea0ef562f9f817801c1
+size 50872350

tasks/models/frugal_2025-01-27/frugal_kan_2.pth

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:3f520cff8b9531981e16a8b009b6a55fb8ca98573fc4d3dc6806df60b07a49c2
-size 1710980
+oid sha256:58e353129b2750993441ea459485b150b2a45b39cbdb7e49bd1839809e4671e2
+size 1363844

tasks/models/frugal_2025-01-28/frugal_kan_2.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9591abf1e617bfedd8414f7c031de3394d7e1bccc64763f7060cef0ee13fab65
+size 1710980

tasks/models/frugal_2025-01-29/frugal_kan_2.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f11c67e65bf6acd1f7e2055dce4a92de0d4d1d88fb12c136eb70198c8ac6eab8
+size 1710980

tasks/run.py

@@ -1,14 +1,17 @@
 from torch.utils.data import DataLoader
 from .utils.data import FFTDataset, SplitDataset
 from datasets import load_dataset
-from .utils.train import Trainer
-from .utils.models import CNNKan, KanEncoder
+from .utils.train import Trainer, XGBoostTrainer
+from .utils.models import CNNKan, KanEncoder, CNNKanFeaturesEncoder
 from .utils.data_utils import *
 from huggingface_hub import login
 import yaml
 import datetime
 import json
 import numpy as np
+import pandas as pd
+import seaborn as sns
+import matplotlib.pyplot as plt
 from collections import OrderedDict
 
 # local_rank = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
@@ -19,9 +22,11 @@ data_args = Container(**yaml.safe_load(open(args_dir, 'r'))['Data'])
 exp_num = data_args.exp_num
 model_name = data_args.model_name
 model_args = Container(**yaml.safe_load(open(args_dir, 'r'))['CNNEncoder'])
+mlp_args = Container(**yaml.safe_load(open(args_dir, 'r'))['MLP'])
 model_args_f = Container(**yaml.safe_load(open(args_dir, 'r'))['CNNEncoder_f'])
 conformer_args = Container(**yaml.safe_load(open(args_dir, 'r'))['Conformer'])
 kan_args = Container(**yaml.safe_load(open(args_dir, 'r'))['KAN'])
+boost_args = Container(**yaml.safe_load(open(args_dir, 'r'))['XGBoost'])
 if not os.path.exists(f"{data_args.log_dir}/{datetime_dir}"):
     os.makedirs(f"{data_args.log_dir}/{datetime_dir}")
 
@@ -44,26 +49,62 @@ val_dl = DataLoader(val_ds,batch_size=data_args.batch_size, collate_fn=collate_f
 test_ds = FFTDataset(dataset["test"])
 test_dl = DataLoader(test_ds,batch_size=data_args.batch_size, collate_fn=collate_fn)
 
-# for i, batch in enumerate(train_dl):
-#     x, x_f, y = batch['audio']['array'], batch['audio']['fft'], batch['label']
-#     print(x.shape, x_f.shape, y.shape)
-#     if i > 10:
+# data = []
+#
+# # Iterate over the dataset
+# for i, batch in enumerate(train_ds):
+#     label = batch['label']
+#     features = batch['audio']['features']
+#
+#     # Flatten the nested dictionary structure
+#     feature_dict = {'label': label}
+#     for k, v in features.items():
+#         if isinstance(v, dict):
+#             for sub_k, sub_v in v.items():
+#                 feature_dict[f"{k}_{sub_k}"] = sub_v[0].item()  # Aggregate (e.g., mean)
+#         else:
+#             print(k, v.shape)  # Aggregate (e.g., mean)
+#
+#     data.append(feature_dict)
+#     print(i)
+#
+#     if i > 1000:  # Limit to 10 iterations
 #         break
+#
+# # Convert to DataFrame
+# df = pd.DataFrame(data)
+
+# Plot distributions colored by label
+# plt.figure()
+# for col in df.columns:
+#     if col != 'label':
+#         sns.kdeplot(df, x=col, hue='label', fill=True, alpha=0.5)
+#         plt.title(f'Distribution of {col}')
+#         plt.show()
+# exit()
+
+# trainer = XGBoostTrainer(boost_args.get_dict(), train_ds, val_ds, test_ds)
+# res = trainer.fit()
+# trainer.predict()
+# trainer.plot_results(res)
 # exit()
 
 # model = DualEncoder(model_args, model_args_f, conformer_args)
 # model = FasterKAN([18000,64,64,16,1])
 model = CNNKan(model_args, conformer_args, kan_args.get_dict())
+# model = CNNKanFeaturesEncoder(model_args, mlp_args, kan_args.get_dict())
 # model.kan.speed()
 # model = KanEncoder(kan_args.get_dict())
 model = model.to(local_rank)
-state_dict = torch.load(data_args.checkpoint_path, map_location=torch.device('cpu'))
-new_state_dict = OrderedDict()
-for key, value in state_dict.items():
-    if key.startswith('module.'):
-        key = key[7:]
-    new_state_dict[key] = value
-missing, unexpected = model.load_state_dict(new_state_dict)
+
+# state_dict = torch.load(data_args.checkpoint_path, map_location=torch.device('cpu'))
+# new_state_dict = OrderedDict()
+# for key, value in state_dict.items():
+#     if key.startswith('module.'):
+#         key = key[7:]
+#     new_state_dict[key] = value
+# missing, unexpected = model.load_state_dict(new_state_dict)
+
 # model = DDP(model, device_ids=[local_rank], output_device=local_rank)
 num_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
 print(f"Number of parameters: {num_params}")
@@ -92,7 +133,7 @@ fit_res = trainer.fit(num_epochs=100, device=local_rank,
 output_filename = f'{data_args.log_dir}/{datetime_dir}/{model_name}_frugal_{exp_num}.json'
 with open(output_filename, "w") as f:
     json.dump(fit_res, f, indent=2)
-preds, acc = trainer.predict(test_dl, local_rank)
+preds, tru, acc = trainer.predict(test_dl, local_rank)
 print(f"Accuracy: {acc}")
 
 

tasks/run_inr.py

@@ -0,0 +1,200 @@
+import torch.nn
+from torch.utils.data import DataLoader
+from utils.data import FFTDataset, SplitDataset, AudioINRDataset
+from datasets import load_dataset
+from utils.train import Trainer, INRTrainer
+from utils.models import MultiGraph, ImplicitEncoder
+from omegaconf import OmegaConf
+
+# from .utils.models import CNNKan, KanEncoder
+from utils.inr import INR
+from utils.data_utils import *
+from huggingface_hub import login
+import yaml
+import datetime
+import json
+import numpy as np
+from tqdm import tqdm
+import matplotlib.pyplot as plt
+from scipy.signal import savgol_filter as savgol
+from utils.kan import FasterKAN
+from utils.relational_transformer import RelationalTransformer
+from collections import OrderedDict
+def plot_results(dims, i, data, losses, pred_values):
+    data = savgol(data.cpu().detach().numpy(), window_length=250, polyorder=1)
+    pred_values = pred_values.transpose(-1, -2).unflatten(-1, data.shape[-2:]).squeeze(0).cpu().detach().numpy()
+    pred_values = (pred_values - np.min(pred_values)) / (np.max(pred_values) - np.min(pred_values))
+    data = (data - np.min(data)) / (np.max(data) - np.min(data))
+    plt.plot(data.squeeze())
+    plt.plot(pred_values.squeeze())
+    # axes[0].set_title('Original')
+    # axes[1].set_title('Reconstruction')
+    plt.show()
+    # plt.plot(np.arange(len(losses)), losses)
+    # plt.xlabel('Iteration')
+    # plt.ylabel('Reconstruction MSE Error')
+    # plt.show()
+
+# local_rank = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+current_date = datetime.date.today().strftime("%Y-%m-%d")
+datetime_dir = f"frugal_{current_date}"
+args_dir = 'utils/config.yaml'
+data_args = Container(**yaml.safe_load(open(args_dir, 'r'))['Data'])
+exp_num = data_args.exp_num
+model_name = data_args.model_name
+rt_args = Container(**yaml.safe_load(open(args_dir, 'r'))['RelationalTransformer'])
+cnn_args = Container(**yaml.safe_load(open(args_dir, 'r'))['CNNEncoder_f'])
+conformer_args = Container(**yaml.safe_load(open(args_dir, 'r'))['Conformer'])
+kan_args = Container(**yaml.safe_load(open(args_dir, 'r'))['KAN_INR'])
+inr_args = Container(**yaml.safe_load(open(args_dir, 'r'))['INR'])
+if not os.path.exists(f"{data_args.log_dir}/{datetime_dir}"):
+    os.makedirs(f"{data_args.log_dir}/{datetime_dir}")
+
+with open("../../logs/token.txt", "r") as f:
+    api_key = f.read()
+
+# local_rank, world_size, gpus_per_node = setup()
+local_rank = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+login(api_key)
+dataset = load_dataset("rfcx/frugalai", streaming=True)
+
+train_ds = SplitDataset(FFTDataset(dataset["train"]), is_train=True)
+
+train_dl = DataLoader(train_ds, batch_size=data_args.batch_size)
+
+val_ds = SplitDataset(FFTDataset(dataset["train"]), is_train=False)
+
+val_dl = DataLoader(val_ds, batch_size=data_args.batch_size)
+
+test_ds = AudioINRDataset(FFTDataset(dataset["test"]))
+test_dl = DataLoader(test_ds, batch_size=data_args.batch_size)
+
+# for i, batch in enumerate(train_ds):
+#     fft_phase, fft_mag, audio = batch['audio']['fft_phase'], batch['audio']['fft_mag'], batch['audio']['array']
+#     label = batch['label']
+#     fig, axes = plt.subplots(nrows=1, ncols=3)
+#     axes = axes.flatten()
+#     axes[0].plot(fft_phase)
+#     axes[1].plot(fft_mag)
+#     axes[2].plot(audio)
+#     fig.suptitle(label)
+#     plt.tight_layout()
+#     plt.show()
+#     if i > 20:
+#         break
+# model = DualEncoder(model_args, model_args_f, conformer_args)
+# model = FasterKAN([18000,64,64,16,1])
+# model = INR(in_features=1)
+# model.kan.speed()
+# model = KanEncoder(kan_args.get_dict())
+# model = model.to(local_rank)
+
+# state_dict = torch.load(data_args.checkpoint_path, map_location=torch.device('cpu'))
+# new_state_dict = OrderedDict()
+# for key, value in state_dict.items():
+#     if key.startswith('module.'):
+#         key = key[7:]
+#     new_state_dict[key] = value
+# missing, unexpected = model.load_state_dict(new_state_dict)
+
+# model = DDP(model, device_ids=[local_rank], output_device=local_rank)
+# num_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
+# print(f"Number of parameters: {num_params}")
+#
+# optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
+# total_steps = int(data_args.num_epochs) * 1000
+# scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer,
+#                                                        T_max=total_steps,
+#                                                        eta_min=float((5e-4) / 10))
+
+loss_fn = torch.nn.BCEWithLogitsLoss()
+inr_criterion = torch.nn.MSELoss()
+
+# for i, batch in enumerate(train_ds):
+#     coords, fft, audio = batch['audio']['coords'], batch['audio']['fft_mag'], batch['audio']['array']
+#     coords = coords.to(local_rank)
+#     fft = fft.to(local_rank)
+#     audio = audio.to(local_rank)
+#     values = torch.cat((audio.unsqueeze(-1), fft.unsqueeze(-1)), dim=-1)
+#     # model = INR(hidden_features=128, n_layers=3,
+#     #             in_features=1,
+#     #             out_features=1).to(local_rank)
+#     model = FasterKAN(**kan_args.get_dict()).to(local_rank)
+#     optimizer = torch.optim.Adam([{'params': model.parameters()}], lr=1e-3)
+#     pbar = tqdm(range(200))
+#     losses = []
+#     print(coords.shape)
+#     for t in pbar:
+#         optimizer.zero_grad()
+#         pred_values = model(coords.to(local_rank)).float()
+#         loss = inr_criterion(pred_values, values)
+#         loss.backward()
+#         optimizer.step()
+#         pbar.set_description(f'loss: {loss.item()}')
+#         losses.append(loss.item())
+#     state_dict = model.state_dict()
+#     torch.save(state_dict, 'test')
+#     # print(f'Sample {i+offset} label {label} saved in {inr_path}')
+#     plot_results(1, i, fft, losses, pred_values)
+# #
+# exit()
+
+
+# missing, unexpected = model.load_state_dict(torch.load(model_args.checkpoint_path))
+# print(f"Missing keys: {missing}")
+# print(f"Unexpected keys: {unexpected}")
+layer_layout = [inr_args.in_features] + [inr_args.hidden_features for _ in range(inr_args.n_layers)]  + [inr_args.out_features]
+
+graph_constructor = OmegaConf.create(
+        {
+            "_target_": "utils.graph_constructor.GraphConstructor",
+            "_recursive_": False,
+            "_convert_": "all",
+            "d_in": 1,
+            "d_edge_in": 1,
+            "zero_out_bias": False,
+            "zero_out_weights": False,
+            "sin_emb": True,
+            "sin_emb_dim": rt_args.d_node,
+            "use_pos_embed": False,
+            "input_layers": 1,
+            "inp_factor": 1,
+            "num_probe_features": 0,
+            "inr_model": None,
+            "stats": None,
+            "sparsify": False,
+            'sym_edges': False,
+        }
+    )
+
+rt_model = RelationalTransformer(layer_layout=layer_layout, graph_constructor=graph_constructor,
+                              **rt_args.get_dict()).to(local_rank)
+rt_model.proj_out= torch.nn.Identity()
+multi_graph = MultiGraph(rt_model, cnn_args)
+implicit_net = INR(**inr_args.get_dict())
+model = ImplicitEncoder(implicit_net, multi_graph).to(local_rank)
+num_parameters = sum(p.numel() for p in model.parameters() if p.requires_grad)
+print(f"Number of parameters: {num_parameters}")
+optimizer = torch.optim.Adam([{'params': model.parameters()}], lr=1e-3)
+trainer = Trainer(model=model, optimizer=optimizer,
+                  criterion=loss_fn, output_dim=1, scaler=None,
+                  scheduler=None, train_dataloader=train_dl,
+                  val_dataloader=val_dl, device=local_rank,
+                  exp_num=datetime_dir, log_path=data_args.log_dir,
+                  range_update=None,
+                  accumulation_step=1, max_iter=100,
+                  exp_name=f"frugal_kan_{exp_num}")
+fit_res = trainer.fit(num_epochs=100, device=local_rank,
+                      early_stopping=10, only_p=False, best='loss', conf=True)
+output_filename = f'{data_args.log_dir}/{datetime_dir}/{model_name}_frugal_{exp_num}.json'
+with open(output_filename, "w") as f:
+    json.dump(fit_res, f, indent=2)
+preds, acc = trainer.predict(test_dl, local_rank)
+print(f"Accuracy: {acc}")
+
+
+
+
+
+
+

tasks/tasks/models/frugal_2025-01-28/CNNEncoder_frugal_2.json

@@ -0,0 +1,420 @@
+{
+  "num_epochs": 100,
+  "train_loss": [
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0
+  ],
+  "val_loss": [
+    0.7067973613739014,
+    0.6959081888198853,
+    0.7170985341072083,
+    0.6722545623779297,
+    0.7166763544082642,
+    0.7339211702346802,
+    0.7830350399017334,
+    0.7616254091262817,
+    0.7465198040008545,
+    0.7474430799484253,
+    0.7112622261047363,
+    0.8131649494171143,
+    0.7090165019035339,
+    0.697528600692749,
+    0.7792525291442871,
+    0.700302243232727,
+    0.7315454483032227,
+    0.7067973613739014,
+    0.6959081888198853,
+    0.7170985341072083,
+    0.6722545623779297,
+    0.7166763544082642,
+    0.7339211702346802,
+    0.7830350399017334,
+    0.7616254091262817,
+    0.7465198040008545,
+    0.7474430799484253,
+    0.7112622261047363,
+    0.8131649494171143,
+    0.7090165019035339,
+    0.697528600692749,
+    0.7792525291442871,
+    0.700302243232727,
+    0.7315454483032227,
+    0.7067973613739014,
+    0.6959081888198853,
+    0.7170985341072083,
+    0.6722545623779297,
+    0.7166763544082642,
+    0.7339211702346802,
+    0.7830350399017334,
+    0.7616254091262817,
+    0.7465198040008545,
+    0.7474430799484253,
+    0.7112622261047363,
+    0.8131649494171143,
+    0.7090165019035339,
+    0.697528600692749,
+    0.7792525291442871,
+    0.700302243232727,
+    0.7315454483032227,
+    0.7067973613739014,
+    0.6959081888198853,
+    0.7170985341072083,
+    0.6722545623779297,
+    0.7166763544082642,
+    0.7339211702346802,
+    0.7830350399017334,
+    0.7616254091262817,
+    0.7465198040008545,
+    0.7474430799484253,
+    0.7112622261047363,
+    0.8131649494171143,
+    0.7090165019035339,
+    0.697528600692749,
+    0.7792525291442871,
+    0.700302243232727,
+    0.7315454483032227,
+    0.7067973613739014,
+    0.6959081888198853,
+    0.7170985341072083,
+    0.6722545623779297,
+    0.7166763544082642,
+    0.7339211702346802,
+    0.7830350399017334,
+    0.7616254091262817,
+    0.7465198040008545,
+    0.7474430799484253,
+    0.7112622261047363,
+    0.8131649494171143,
+    0.7090165019035339,
+    0.697528600692749,
+    0.7792525291442871,
+    0.700302243232727,
+    0.7315454483032227,
+    0.7067973613739014,
+    0.6959081888198853,
+    0.7170985341072083,
+    0.6722545623779297,
+    0.7166763544082642,
+    0.7339211702346802,
+    0.7830350399017334,
+    0.7616254091262817,
+    0.7465198040008545,
+    0.7474430799484253,
+    0.7112622261047363,
+    0.8131649494171143,
+    0.7090165019035339,
+    0.697528600692749,
+    0.7792525291442871,
+    0.700302243232727,
+    0.7315454483032227,
+    0.7067973613739014,
+    0.6959081888198853,
+    0.7170985341072083,
+    0.6722545623779297,
+    0.7166763544082642,
+    0.7339211702346802,
+    0.7830350399017334,
+    0.7616254091262817,
+    0.7465198040008545,
+    0.7474430799484253,
+    0.7112622261047363,
+    0.8131649494171143,
+    0.7090165019035339,
+    0.697528600692749,
+    0.7792525291442871,
+    0.700302243232727,
+    0.7315454483032227,
+    0.7067973613739014,
+    0.6959081888198853,
+    0.7170985341072083,
+    0.6722545623779297,
+    0.7166763544082642,
+    0.7339211702346802,
+    0.7830350399017334,
+    0.7616254091262817,
+    0.7465198040008545,
+    0.7474430799484253,
+    0.7112622261047363,
+    0.8131649494171143,
+    0.7090165019035339,
+    0.697528600692749,
+    0.7792525291442871,
+    0.700302243232727,
+    0.7315454483032227,
+    0.7067973613739014,
+    0.6959081888198853,
+    0.7170985341072083,
+    0.6722545623779297,
+    0.7166763544082642,
+    0.7339211702346802,
+    0.7830350399017334,
+    0.7616254091262817,
+    0.7465198040008545,
+    0.7474430799484253,
+    0.7112622261047363,
+    0.8131649494171143,
+    0.7090165019035339,
+    0.697528600692749,
+    0.7792525291442871,
+    0.700302243232727,
+    0.7315454483032227,
+    0.7067973613739014,
+    0.6959081888198853,
+    0.7170985341072083,
+    0.6722545623779297,
+    0.7166763544082642,
+    0.7339211702346802,
+    0.7830350399017334,
+    0.7616254091262817,
+    0.7465198040008545,
+    0.7474430799484253,
+    0.7112622261047363,
+    0.8131649494171143,
+    0.7090165019035339,
+    0.697528600692749,
+    0.7792525291442871,
+    0.700302243232727,
+    0.7315454483032227,
+    0.7067973613739014,
+    0.6959081888198853,
+    0.7170985341072083,
+    0.6722545623779297,
+    0.7166763544082642,
+    0.7339211702346802,
+    0.7830350399017334,
+    0.7616254091262817,
+    0.7465198040008545,
+    0.7474430799484253,
+    0.7112622261047363,
+    0.8131649494171143,
+    0.7090165019035339,
+    0.697528600692749,
+    0.7792525291442871,
+    0.700302243232727,
+    0.7315454483032227
+  ],
+  "train_acc": [
+    0.0,
+    0.0,
+    0.0,
+    0.0,
+    0.0,
+    0.0,
+    0.0,
+    0.0,
+    0.0,
+    0.0,
+    0.0
+  ],
+  "val_acc": [
+    0.0,
+    0.0,
+    0.0,
+    0.0,
+    0.0,
+    0.0,
+    0.0,
+    0.0,
+    0.0,
+    0.0,
+    0.0
+  ],
+  "lrs": [
+    0.001,
+    0.001,
+    0.001,
+    0.001,
+    0.001,
+    0.001,
+    0.001,
+    0.001,
+    0.001,
+    0.001,
+    0.001
+  ]
+}

tasks/tasks/models/frugal_2025-01-28/frugal_kan_2.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7b22195a96ca5dd9cf4d20bb1886c6d22327d42ea3fbe4b9d86b9a281a132024
+size 820491

tasks/utils/config.yaml

@@ -12,7 +12,7 @@ Data:
   max_days_lc: 270
   lc_freq: 0.0208
   create_umap: True
-  checkpoint_path: 'tasks/models/frugal_2025-01-27/frugal_kan_2.pth'
+  checkpoint_path: 'tasks/models/frugal_2025-01-29/frugal_kan_2.pth'
 
 CNNEncoder:
   # Model
@@ -27,9 +27,14 @@ CNNEncoder:
   load_checkpoint: False
   checkpoint_num: 1
   activation: "silu"
-  sine_w0: 1.0
+  sine_w0: 30.0
   avg_output: False
 
+MLP:
+  input_dim: 6
+  hidden_dims: [16,32]
+  dropout: 0.2
+
 KAN:
   layers_hidden: [1125,32,8,1]
   grid_min:  -1.2
@@ -37,9 +42,16 @@ KAN:
   num_grids:  8
   exponent:  2
 
+KAN_INR:
+  layers_hidden: [1,1024,128,128,1]
+  grid_min:  -1.2
+  grid_max:  1.2
+  num_grids:  8
+  exponent:  2
+
 CNNEncoder_f:
   # Model
-  in_channels: 1
+  in_channels: 32
   num_layers: 4
   stride: 1
   encoder_dims: [32,64,128]
@@ -64,10 +76,39 @@ Conformer:
   dropout_p: 0.2
   norm: "postnorm"
 
+RelationalTransformer:
+  d_node: 32
+  d_edge: 32
+  d_attn_hid: 16
+  d_node_hid: 16
+  d_edge_hid: 16
+  d_out_hid: 16
+  d_out: 1
+  n_layers: 4
+  n_heads: 4
+  dropout: 0.1
+
+
+INR:
+  in_features : 2
+  n_layers : 2
+  hidden_features : 64
+  out_features : 32
+
+XGBoost:
+  objective : 'binary:logistic'
+  eval_metric : 'logloss'
+  use_label_encoder : False
+  n_estimators : 500
+  learning_rate : 0.1
+  max_depth : 5
+  subsample : 0.8
+  colsample_bytree : 0.8
+  random_state : 42
 
 Optimization:
   # Optimization
   max_lr: 1e-5
   weight_decay: 5e-6
   warmup_pct: 0.3
-  steps_per_epoch: 3500
+  steps_per_epoch: 3500

tasks/utils/data.py

@@ -1,11 +1,38 @@
 import torch
 from torch.utils.data import IterableDataset
-from torch.fft import fft
+from torch.fft import fft, fftshift
 import torch.nn.functional as F
 from itertools import tee
 import random
 import torchaudio.transforms as T
+import hashlib
+from typing import NamedTuple, Tuple, Union
+from .transforms import compute_all_features
 
+from scipy.signal import savgol_filter as savgol
+
+
+class WeightsBatch(NamedTuple):
+    weights: Tuple
+    biases: Tuple
+    label: Union[torch.Tensor, int]
+
+    def _assert_same_len(self):
+        assert len(set([len(t) for t in self])) == 1
+
+    def as_dict(self):
+        return self._asdict()
+
+    def to(self, device):
+        """move batch to device"""
+        return self.__class__(
+            weights=tuple(w.to(device) for w in self.weights),
+            biases=tuple(w.to(device) for w in self.biases),
+            label=self.label.to(device),
+        )
+
+    def __len__(self):
+        return len(self.weights[0])
 
 class SplitDataset(IterableDataset):
     def __init__(self, dataset, is_train=True, train_ratio=0.8):
@@ -28,22 +55,101 @@ class FFTDataset(IterableDataset):
     def __init__(self, original_dataset, max_len=72000, orig_sample_rate=12000, target_sample_rate=3000):
         self.dataset = original_dataset
         self.resampler = T.Resample(orig_freq=orig_sample_rate, new_freq=target_sample_rate)
+        self.target_sample_rate = target_sample_rate
         self.max_len = max_len
-        
+
+
+    def normalize_audio(self, audio):
+        """Normalize audio to [0, 1] range"""
+        audio_min = audio.min()
+        audio_max = audio.max()
+        audio = (audio - audio_min) / (audio_max - audio_min)
+        return audio
+
+    def generate_unique_id(self, array):
+        # Convert the array to bytes
+        array_bytes = array.tobytes()
+        # Hash the bytes using SHA256
+        hash_object = hashlib.sha256(array_bytes)
+        # Return the hexadecimal representation of the hash
+        return hash_object.hexdigest()
+
     def __iter__(self):
         for item in self.dataset:
-            # Assuming your audio data is in item['audio']
-            # Modify this based on your actual data structure
-            audio_data = torch.tensor(item['audio']['array']).float()
-            # pad audio
-            # if len(audio_data) == 0:
-            #     continue
+            # audio_data = savgol(item['audio']['array'], 500, polyorder=1)
+            audio_data = item['audio']['array']
+            # item['id'] = self.generate_unique_id(audio_data)
+            audio_data = torch.tensor(audio_data).float()
+
             pad_len = self.max_len - len(audio_data)
             audio_data = F.pad(audio_data, (0, pad_len), mode='constant')
             audio_data = self.resampler(audio_data)
+
+            audio_data = self.normalize_audio(audio_data)
             fft_data = fft(audio_data)
-            
-            # Update the item with FFT data
-            item['audio']['fft'] = fft_data
-            item['audio']['array'] = audio_data
+            magnitude = torch.abs(fft_data)
+            phase = torch.angle(fft_data)
+            # features = compute_all_features(audio_data, sample_rate=self.target_sample_rate)
+            # features_arr = torch.tensor([v for _, v in features['frequency_domain'].items()])
+            magnitude_centered = fftshift(magnitude)
+            phase_centered = fftshift(phase)
+            # cwt = features['cwt_power']
+
+            # Optionally, remove the DC component
+            magnitude_centered[len(magnitude_centered) // 2] = 0  # Set DC component to zero
+
+            item['audio']['fft_mag'] = torch.nan_to_num(magnitude_centered, 0)
+            item['audio']['fft_phase'] = torch.nan_to_num(phase_centered, 0)
+            # item['audio']['cwt_mag'] = torch.nan_to_num(cwt, 0)
+            item['audio']['array'] = torch.nan_to_num(audio_data, 0)
+            # item['audio']['features'] = features
+            # item['audio']['features_arr'] = torch.nan_to_num(features_arr, 0)
+            yield item
+
+
+class AudioINRDataset(IterableDataset):
+    def __init__(self, original_dataset, max_len=18000, sample_size=1024, dim=1, normalize=True):
+        """
+        Convert audio data into coordinate-value pairs for INR training.
+
+        Args:
+            original_dataset: Original audio dataset
+            max_len: Maximum length of audio to process
+            batch_size: Number of points to sample per audio clip
+            normalize: Whether to normalize the audio values to [0, 1]
+        """
+        self.dataset = original_dataset
+        self.max_len = max_len
+        self.dim = dim
+        self.normalize = normalize
+        self.sample_size = sample_size
+
+    def get_coordinates(self, audio_len):
+        """Generate time coordinates"""
+        # Create normalized time coordinates in [0, 1]
+        coords = torch.linspace(0, 1, audio_len).unsqueeze(-1).expand(audio_len, self.dim)
+        return coords  # Shape: [audio_len, 1]
+
+    def sample_points(self, coords, values):
+        """Randomly sample points from the audio"""
+        if len(coords) > self.sample_size:
+            idx = torch.randperm(len(coords))[:self.sample_size]
+            coords = coords[idx]
+            values = values[idx]
+        return coords, values
+
+    def __iter__(self):
+        for item in self.dataset:
+            # Get audio data
+            audio_data = torch.tensor(item['audio']['array']).float()
+
+            # Generate coordinates
+            coords = self.get_coordinates(len(audio_data))
+
+            item['audio']['coords'] = coords
+
+            # Sample random points
+            # coords, values = self.sample_points(coords, audio_data)
+
+            # Create the INR training sample
             yield item

tasks/utils/data_utils.py

@@ -6,18 +6,25 @@ from torch.nn.utils.rnn import pad_sequence
 def collate_fn(batch):
     # Extract audio arrays and FFT data from the batch of dictionaries
     audio_arrays = [torch.tensor(item['audio']['array']) for item in batch]
-    fft_arrays = [torch.tensor(item['audio']['fft']) for item in batch]
+    fft_arrays = [torch.tensor(item['audio']['fft_mag']) for item in batch]
+    # cwt_arrays = [torch.tensor(item['audio']['cwt_mag']) for item in batch]
+    # features = [item['audio']['features'] for item in batch]
+    # features_arr = torch.stack([item['audio']['features_arr'] for item in batch])
     labels = [torch.tensor(item['label']) for item in batch]
     
     # Pad both sequences
     padded_audio = pad_sequence(audio_arrays, batch_first=True, padding_value=0)
     padded_fft = pad_sequence(fft_arrays, batch_first=True, padding_value=0)
+    # padded_features = pad_sequence(features_arr, batch_first=True, padding_value=0)
     
     # Return as dictionary with the same structure
     return {
         'audio': {
             'array': padded_audio,
-            'fft': padded_fft
+            'fft_mag': padded_fft,
+            # 'features': features,
+            # 'features_arr': features_arr,
+            # 'cwt_mag': padded_cwt,
         },
         'label': torch.stack(labels)
         

tasks/utils/graph_constructor.py

@@ -0,0 +1,214 @@
+import torch
+import torch.nn as nn
+from rff.layers import GaussianEncoding
+
+# from nn.probe_features import GraphProbeFeatures
+
+
+def sparsify_graph(edges, fraction=0.1):
+    abs_edges = torch.abs(edges)
+    flat_abs_tensor = abs_edges.flatten()
+    sorted_tensor, _ = torch.sort(flat_abs_tensor, descending=True)
+    num_elements = flat_abs_tensor.numel()
+    top_k = int(num_elements * fraction)
+    topk_values, topk_indices = torch.topk(flat_abs_tensor, top_k)
+    mask = torch.zeros_like(flat_abs_tensor, dtype=torch.bool)
+    mask[topk_indices] = True
+    mask = mask.view(edges.shape)
+    return mask
+
+def batch_to_graphs(
+    weights,
+    biases,
+    weights_mean=None,
+    weights_std=None,
+    biases_mean=None,
+    biases_std=None,
+    sparsify=False,
+    sym_edges=False
+):
+    device = weights[0].device
+    bsz = weights[0].shape[0]
+    num_nodes = weights[0].shape[1] + sum(w.shape[2] for w in weights)
+
+    node_features = torch.zeros(bsz, num_nodes, biases[0].shape[-1], device=device)
+    edge_features = torch.zeros(
+        bsz, num_nodes, num_nodes, weights[0].shape[-1], device=device
+    )
+
+    row_offset = 0
+    col_offset = weights[0].shape[1]  # no edge to input nodes
+
+    for i, w in enumerate(weights):
+        _, num_in, num_out, _ = w.shape
+        w_mean = weights_mean[i] if weights_mean is not None else 0
+        w_std = weights_std[i] if weights_std is not None else 1
+        w = (w - w_mean) / w_std
+        if sparsify:
+            w[~sparsify_graph(w)] = 0
+        edge_features[
+            :, row_offset : row_offset + num_in, col_offset : col_offset + num_out
+        ] = w
+        if sym_edges:
+            edge_features[
+            :, col_offset: col_offset + num_out, row_offset: row_offset + num_in
+            ] = torch.swapaxes(w, 1,2)
+        row_offset += num_in
+        col_offset += num_out
+
+    row_offset = weights[0].shape[1]  # no bias in input nodes
+    for i, b in enumerate(biases):
+        _, num_out, _ = b.shape
+        b_mean = biases_mean[i] if biases_mean is not None else 0
+        b_std = biases_std[i] if biases_std is not None else 1
+        node_features[:, row_offset : row_offset + num_out] = (b - b_mean) / b_std
+        row_offset += num_out
+
+    return node_features, edge_features
+
+
+class GraphConstructor(nn.Module):
+    def __init__(
+        self,
+        d_in,
+        d_edge_in,
+        d_node,
+        d_edge,
+        layer_layout,
+        rev_edge_features=False,
+        zero_out_bias=False,
+        zero_out_weights=False,
+        inp_factor=1,
+        input_layers=1,
+        sin_emb=False,
+        sin_emb_dim=128,
+        use_pos_embed=False,
+        num_probe_features=0,
+        inr_model=None,
+        stats=None,
+        sparsify=False,
+        sym_edges=False,
+    ):
+        super().__init__()
+        self.rev_edge_features = rev_edge_features
+        self.nodes_per_layer = layer_layout
+        self.zero_out_bias = zero_out_bias
+        self.zero_out_weights = zero_out_weights
+        self.use_pos_embed = use_pos_embed
+        self.stats = stats if stats is not None else {}
+        self._d_node = d_node
+        self._d_edge = d_edge
+        self.sparse = sparsify
+        self.sym_edges = sym_edges
+
+        self.pos_embed_layout = (
+            [1] * layer_layout[0] + layer_layout[1:-1] + [1] * layer_layout[-1]
+        )
+        self.pos_embed = nn.Parameter(torch.randn(len(self.pos_embed_layout), d_node))
+
+        if not self.zero_out_weights:
+            proj_weight = []
+            if sin_emb:
+                proj_weight.append(
+                    GaussianEncoding(
+                        sigma=inp_factor,
+                        input_size=d_edge_in
+                        + (2 * d_edge_in if rev_edge_features else 0),
+                        encoded_size=sin_emb_dim,
+                    )
+                )
+                proj_weight.append(nn.Linear(2 * sin_emb_dim, d_edge))
+            else:
+                proj_weight.append(
+                    nn.Linear(
+                        d_edge_in + (2 * d_edge_in if rev_edge_features else 0), d_edge
+                    )
+                )
+
+            for i in range(input_layers - 1):
+                proj_weight.append(nn.SiLU())
+                proj_weight.append(nn.Linear(d_edge, d_edge))
+
+            self.proj_weight = nn.Sequential(*proj_weight)
+        if not self.zero_out_bias:
+            proj_bias = []
+            if sin_emb:
+                proj_bias.append(
+                    GaussianEncoding(
+                        sigma=inp_factor,
+                        input_size=d_in,
+                        encoded_size=sin_emb_dim,
+                    )
+                )
+                proj_bias.append(nn.Linear(2 * sin_emb_dim, d_node))
+            else:
+                proj_bias.append(nn.Linear(d_in, d_node))
+
+            for i in range(input_layers - 1):
+                proj_bias.append(nn.SiLU())
+                proj_bias.append(nn.Linear(d_node, d_node))
+
+            self.proj_bias = nn.Sequential(*proj_bias)
+
+        self.proj_node_in = nn.Linear(d_node, d_node)
+        self.proj_edge_in = nn.Linear(d_edge, d_edge)
+
+        if num_probe_features > 0:
+            self.gpf = GraphProbeFeatures(
+                d_in=layer_layout[0],
+                num_inputs=num_probe_features,
+                inr_model=inr_model,
+                input_init=None,
+                proj_dim=d_node,
+            )
+        else:
+            self.gpf = None
+
+    def forward(self, inputs):
+        node_features, edge_features = batch_to_graphs(*inputs, **self.stats,
+                                                       )
+        mask = edge_features.sum(dim=-1, keepdim=True) != 0
+        if self.rev_edge_features:
+            rev_edge_features = edge_features.transpose(-2, -3)
+            edge_features = torch.cat(
+                [edge_features, rev_edge_features, edge_features + rev_edge_features],
+                dim=-1,
+            )
+            mask = mask | mask.transpose(-3, -2)
+
+        if self.zero_out_weights:
+            edge_features = torch.zeros(
+                (*edge_features.shape[:-1], self._d_edge),
+                device=edge_features.device,
+                dtype=edge_features.dtype,
+            )
+        else:
+            edge_features = self.proj_weight(edge_features)
+        if self.zero_out_bias:
+            # only zero out bias, not gpf
+            node_features = torch.zeros(
+                (*node_features.shape[:-1], self._d_node),
+                device=node_features.device,
+                dtype=node_features.dtype,
+            )
+        else:
+            node_features = self.proj_bias(node_features)
+
+        if self.gpf is not None:
+            probe_features = self.gpf(*inputs)
+            node_features = node_features + probe_features
+
+        node_features = self.proj_node_in(node_features)
+        edge_features = self.proj_edge_in(edge_features)
+
+        if self.use_pos_embed:
+            pos_embed = torch.cat(
+                [
+                    # repeat(self.pos_embed[i], "d -> 1 n d", n=n)
+                    self.pos_embed[i].unsqueeze(0).expand(1, n, -1)
+                    for i, n in enumerate(self.pos_embed_layout)
+                ],
+                dim=1,
+            )
+            node_features = node_features + pos_embed
+        return node_features, edge_features, mask

tasks/utils/inr.py

@@ -0,0 +1,147 @@
+import copy
+import math
+from typing import Optional
+
+import torch
+import torch.nn.functional as F
+from rff.layers import GaussianEncoding, PositionalEncoding
+from torch import nn
+from .kan.fasterkan import FasterKAN
+
+
+
+class Sine(nn.Module):
+    def __init__(self, w0=1.0):
+        super().__init__()
+        self.w0 = w0
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return torch.sin(self.w0 * x)
+
+
+def params_to_tensor(params):
+    return torch.cat([p.flatten() for p in params]), [p.shape for p in params]
+
+
+def tensor_to_params(tensor, shapes):
+    params = []
+    start = 0
+    for shape in shapes:
+        size = torch.prod(torch.tensor(shape)).item()
+        param = tensor[start : start + size].reshape(shape)
+        params.append(param)
+        start += size
+    return tuple(params)
+
+
+def wrap_func(func, shapes):
+    def wrapped_func(params, *args, **kwargs):
+        params = tensor_to_params(params, shapes)
+        return func(params, *args, **kwargs)
+
+    return wrapped_func
+
+
+class Siren(nn.Module):
+    def __init__(
+        self,
+        dim_in,
+        dim_out,
+        w0=30.0,
+        c=6.0,
+        is_first=False,
+        use_bias=True,
+        activation=None,
+    ):
+        super().__init__()
+        self.w0 = w0
+        self.c = c
+        self.dim_in = dim_in
+        self.dim_out = dim_out
+        self.is_first = is_first
+
+        weight = torch.zeros(dim_out, dim_in)
+        bias = torch.zeros(dim_out) if use_bias else None
+        self.init_(weight, bias, c=c, w0=w0)
+
+        self.weight = nn.Parameter(weight)
+        self.bias = nn.Parameter(bias) if use_bias else None
+        self.activation = Sine(w0) if activation is None else activation
+
+    def init_(self, weight: torch.Tensor, bias: torch.Tensor, c: float, w0: float):
+        dim = self.dim_in
+
+        w_std = (1 / dim) if self.is_first else (math.sqrt(c / dim) / w0)
+        weight.uniform_(-w_std, w_std)
+
+        if bias is not None:
+            # bias.uniform_(-w_std, w_std)
+            bias.zero_()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        out = F.linear(x, self.weight, self.bias)
+        out = self.activation(out)
+        return out
+
+
+class INR(nn.Module):
+    def __init__(
+        self,
+        in_features: int = 2,
+        n_layers: int = 3,
+        hidden_features: int = 32,
+        out_features: int = 1,
+        pe_features: Optional[int] = None,
+        fix_pe=True,
+    ):
+        super().__init__()
+
+        if pe_features is not None:
+            if fix_pe:
+                self.layers = [PositionalEncoding(sigma=10, m=pe_features)]
+                encoded_dim = in_features * pe_features * 2
+            else:
+                self.layers = [
+                    GaussianEncoding(
+                        sigma=10, input_size=in_features, encoded_size=pe_features
+                    )
+                ]
+                encoded_dim = pe_features * 2
+            self.layers.append(Siren(dim_in=encoded_dim, dim_out=hidden_features))
+        else:
+            self.layers = [Siren(dim_in=in_features, dim_out=hidden_features)]
+        for i in range(n_layers - 2):
+            self.layers.append(Siren(hidden_features, hidden_features))
+        self.layers.append(nn.Linear(hidden_features, out_features))
+        self.seq = nn.Sequential(*self.layers)
+        self.num_layers = len(self.layers)
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.seq(x) + 0.5
+
+
+class INRPerLayer(INR):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        nodes = [x]
+        for layer in self.seq:
+            nodes.append(layer(nodes[-1]))
+        nodes[-1] = nodes[-1] + 0.5
+        return nodes
+
+
+def make_functional(mod, disable_autograd_tracking=False):
+    params_dict = dict(mod.named_parameters())
+    params_names = params_dict.keys()
+    params_values = tuple(params_dict.values())
+
+    stateless_mod = copy.deepcopy(mod)
+    stateless_mod.to("meta")
+
+    def fmodel(new_params_values, *args, **kwargs):
+        new_params_dict = {
+            name: value for name, value in zip(params_names, new_params_values)
+        }
+        return torch.func.functional_call(stateless_mod, new_params_dict, args, kwargs)
+
+    if disable_autograd_tracking:
+        params_values = torch.utils._pytree.tree_map(torch.Tensor.detach, params_values)
+    return fmodel, params_values

tasks/utils/models.py

@@ -4,11 +4,56 @@ from .Modules.conformer import ConformerEncoder, ConformerDecoder
 from .Modules.mhsa_pro import RotaryEmbedding, ContinuousRotaryEmbedding
 from .kan.fasterkan import FasterKAN
 
+
+class Sine(nn.Module):
+    def __init__(self, w0=1.0):
+        super().__init__()
+        self.w0 = w0
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return torch.sin(self.w0 * x)
+
+
+class MLPEncoder(nn.Module):
+    def __init__(self, args):
+        """
+        Initialize an MLP with hidden layers, BatchNorm, and Dropout.
+
+        Args:
+            input_dim (int): Dimension of the input features.
+            hidden_dims (list of int): List of dimensions for hidden layers.
+            output_dim (int): Dimension of the output.
+            dropout (float): Dropout probability (default: 0.0).
+        """
+        super(MLPEncoder, self).__init__()
+
+        layers = []
+        prev_dim = args.input_dim
+
+        # Add hidden layers
+        for hidden_dim in args.hidden_dims:
+            layers.append(nn.Linear(prev_dim, hidden_dim))
+            layers.append(nn.BatchNorm1d(hidden_dim))
+            layers.append(nn.SiLU())
+            if args.dropout > 0.0:
+                layers.append(nn.Dropout(args.dropout))
+            prev_dim = hidden_dim
+        self.model = nn.Sequential(*layers)
+        self.output_dim = hidden_dim
+
+    def forward(self, x):
+        # if x.dim() == 2:
+        #     x = x.unsqueeze(-1)
+        x = self.model(x)
+        # x = x.mean(-1)
+        return x
 class ConvBlock(nn.Module):
   def __init__(self, args, num_layer) -> None:
     super().__init__()
     if args.activation == 'silu':
         self.activation = nn.SiLU()
+    elif args.activation == 'sine':
+        self.activation = Sine(w0=args.sine_w0)
     else:
         self.activation = nn.ReLU()
     in_channels = args.encoder_dims[num_layer-1] if num_layer < len(args.encoder_dims) else args.encoder_dims[-1]
@@ -31,6 +76,8 @@ class CNNEncoder(nn.Module):
         print("Using CNN encoder wit activation: ", args.activation, 'args avg_output: ', args.avg_output)
         if args.activation == 'silu':
             self.activation = nn.SiLU()
+        elif args.activation == 'sine':
+            self.activation = Sine(w0=args.sine_w0)
         else:
             self.activation = nn.ReLU()
         self.embedding = nn.Sequential(nn.Conv1d(in_channels = args.in_channels,
@@ -125,6 +172,21 @@ class CNNKan(nn.Module):
         x = x.mean(dim=1)
         return self.kan(x)
 
+class CNNKanFeaturesEncoder(nn.Module):
+    def __init__(self, args, mlp_args, kan_args):
+        super().__init__()
+        self.backbone = CNNEncoder(args)
+        self.mlp = MLPEncoder(mlp_args)
+        kan_args['layers_hidden'][0] += self.mlp.output_dim
+        self.kan = FasterKAN(**kan_args)
+
+    def forward(self, x: torch.Tensor, f: torch.Tensor) -> torch.Tensor:
+        x = self.backbone(x)
+        x = x.mean(dim=1)
+        f = self.mlp(f)
+        x_f = torch.cat([x, f], dim=-1)
+        return self.kan(x_f)
+
 class KanEncoder(nn.Module):
     def __init__(self, args):
         super().__init__()
@@ -138,3 +200,44 @@ class KanEncoder(nn.Module):
         out = torch.cat([x, f], dim=-1)
         return self.kan_out(out)
 
+
+class MultiGraph(nn.Module):
+    def __init__(self, graph_net, args):
+        super().__init__()
+        self.graph_net = graph_net
+        self.cnn = CNNEncoder(args)
+        total_output_dim = args.encoder_dims[-1]
+        self.projection = nn.Sequential(
+            nn.Linear(total_output_dim, total_output_dim // 2),
+            nn.BatchNorm1d(total_output_dim // 2),
+            nn.SiLU(),
+            nn.Linear(total_output_dim // 2, 1)
+        )
+
+    def forward(self, g: torch.Tensor, x:torch.Tensor) -> torch.Tensor:
+        # g_out = self.graph_net(g)
+        x_out = self.cnn(x)
+        # g_out = g_out.expand(x.shape[0], -1)
+        # features = torch.cat([g_out, x_out], dim=-1)
+        return self.projection(x_out)
+
+class ImplicitEncoder(nn.Module):
+    def __init__(self, transform_net, encoder_net):
+        super().__init__()
+        self.transform_net = transform_net
+        self.encoder_net = encoder_net
+
+    def get_weights_and_bises(self):
+        state_dict = self.transform_net.state_dict()
+        weights = tuple(
+            [v.permute(1, 0).unsqueeze(-1).unsqueeze(0) for w, v in state_dict.items() if "weight" in w]
+        )
+        biases = tuple([v.unsqueeze(-1).unsqueeze(0) for w, v in state_dict.items() if "bias" in w])
+        return weights, biases
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        transformed_x = self.transform_net(x.permute(0, 2, 1)).permute(0, 2, 1)
+        inputs = self.get_weights_and_bises()
+        outputs = self.encoder_net(inputs, transformed_x)
+        return outputs
+

tasks/utils/pooling.py

@@ -0,0 +1,199 @@
+import torch
+import torch.nn as nn
+from torch_geometric.nn.aggr import (
+    AttentionalAggregation,
+    GraphMultisetTransformer,
+    MaxAggregation,
+    MeanAggregation,
+    SetTransformerAggregation,
+)
+
+class CatAggregation(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.flatten = nn.Flatten(1, 2)
+
+    def forward(self, x, index=None):
+        return self.flatten(x)
+
+
+class HeterogeneousAggregator(nn.Module):
+    def __init__(
+        self,
+        input_dim,
+        hidden_dim,
+        output_dim,
+        pooling_method,
+        pooling_layer_idx,
+        input_channels,
+        num_classes,
+    ):
+        super().__init__()
+        self.pooling_method = pooling_method
+        self.pooling_layer_idx = pooling_layer_idx
+        self.input_channels = input_channels
+        self.num_classes = num_classes
+
+        if pooling_layer_idx == "all":
+            self._pool_layer_idx_fn = self.get_all_layer_indices
+        elif pooling_layer_idx == "last":
+            self._pool_layer_idx_fn = self.get_last_layer_indices
+        elif isinstance(pooling_layer_idx, int):
+            self._pool_layer_idx_fn = self.get_nth_layer_indices
+        else:
+            raise ValueError(f"Unknown pooling layer index {pooling_layer_idx}")
+
+        if pooling_method == "mean":
+            self.pool = MeanAggregation()
+        elif pooling_method == "max":
+            self.pool = MaxAggregation()
+        elif pooling_method == "cat":
+            self.pool = CatAggregation()
+        elif pooling_method == "attentional_aggregation":
+            self.pool = AttentionalAggregation(
+                gate_nn=nn.Sequential(
+                    nn.Linear(input_dim, hidden_dim),
+                    nn.SiLU(),
+                    nn.Linear(hidden_dim, 1),
+                ),
+                nn=nn.Sequential(
+                    nn.Linear(input_dim, hidden_dim),
+                    nn.SiLU(),
+                    nn.Linear(hidden_dim, output_dim),
+                ),
+            )
+        elif pooling_method == "set_transformer":
+            self.pool = SetTransformerAggregation(
+                input_dim, heads=8, num_encoder_blocks=4, num_decoder_blocks=4
+            )
+        elif pooling_method == "graph_multiset_transformer":
+            self.pool = GraphMultisetTransformer(input_dim, k=8, heads=8)
+        else:
+            raise ValueError(f"Unknown pooling method {pooling_method}")
+
+    def get_last_layer_indices(
+        self, x, layer_layouts, node_mask=None, return_dense=False
+    ):
+        batch_size = x.shape[0]
+        device = x.device
+
+        # NOTE: node_mask needs to exist in the heterogeneous case only
+        if node_mask is None:
+            node_mask = torch.ones_like(x[..., 0], dtype=torch.bool, device=device)
+
+        valid_layer_indices = (
+            torch.arange(node_mask.shape[1], device=device)[None, :] * node_mask
+        )
+        last_layer_indices = valid_layer_indices.topk(
+            k=self.num_classes, dim=1
+        ).values.fliplr()
+
+        if return_dense:
+            return torch.arange(batch_size, device=device)[:, None], last_layer_indices
+
+        batch_indices = torch.arange(batch_size, device=device).repeat_interleave(
+            self.num_classes
+        )
+        return batch_indices, last_layer_indices.flatten()
+
+    def get_nth_layer_indices(
+        self, x, layer_layouts, node_mask=None, return_dense=False
+    ):
+        batch_size = x.shape[0]
+        device = x.device
+
+        cum_layer_layout = [
+            torch.cumsum(torch.tensor([0] + layer_layout), dim=0)
+            for layer_layout in layer_layouts
+        ]
+
+        layer_sizes = torch.tensor(
+            [layer_layout[self.pooling_layer_idx] for layer_layout in layer_layouts],
+            dtype=torch.long,
+            device=device,
+        )
+        batch_indices = torch.arange(batch_size, device=device).repeat_interleave(
+            layer_sizes
+        )
+        layer_indices = torch.cat(
+            [
+                torch.arange(
+                    layout[self.pooling_layer_idx],
+                    layout[self.pooling_layer_idx + 1],
+                    device=device,
+                )
+                for layout in cum_layer_layout
+            ]
+        )
+        return batch_indices, layer_indices
+
+    def get_all_layer_indices(
+        self, x, layer_layouts, node_mask=None, return_dense=False
+    ):
+        """Imitate flattening with indexing"""
+        batch_size, num_nodes = x.shape[:2]
+        device = x.device
+        batch_indices = torch.arange(batch_size, device=device).repeat_interleave(
+            num_nodes
+        )
+        layer_indices = torch.arange(num_nodes, device=device).repeat(batch_size)
+        return batch_indices, layer_indices
+
+    def forward(self, x, layer_layouts, node_mask=None):
+        # NOTE: `cat` only works with `pooling_layer_idx == "last"`
+        return_dense = self.pooling_method == "cat" and self.pooling_layer_idx == "last"
+        batch_indices, layer_indices = self._pool_layer_idx_fn(
+            x, layer_layouts, node_mask=node_mask, return_dense=return_dense
+        )
+
+        flat_x = x[batch_indices, layer_indices]
+        return self.pool(flat_x, index=batch_indices)
+
+
+class HomogeneousAggregator(nn.Module):
+    def __init__(
+        self,
+        pooling_method,
+        pooling_layer_idx,
+        layer_layout,
+    ):
+        super().__init__()
+        self.pooling_method = pooling_method
+        self.pooling_layer_idx = pooling_layer_idx
+        self.layer_layout = layer_layout
+
+    def forward(self, node_features, edge_features):
+        if self.pooling_method == "mean" and self.pooling_layer_idx == "all":
+            graph_features = node_features.mean(dim=1)
+        elif self.pooling_method == "max" and self.pooling_layer_idx == "all":
+            graph_features = node_features.max(dim=1).values
+        elif self.pooling_method == "mean" and self.pooling_layer_idx == "last":
+            graph_features = node_features[:, -self.layer_layout[-1] :].mean(dim=1)
+        elif self.pooling_method == "cat" and self.pooling_layer_idx == "last":
+            graph_features = node_features[:, -self.layer_layout[-1] :].flatten(1, 2)
+        elif self.pooling_method == "mean" and isinstance(self.pooling_layer_idx, int):
+            graph_features = node_features[
+                :,
+                self.layer_idx[self.pooling_layer_idx] : self.layer_idx[
+                    self.pooling_layer_idx + 1
+                ],
+            ].mean(dim=1)
+        elif self.pooling_method == "cat_mean" and self.pooling_layer_idx == "all":
+            graph_features = torch.cat(
+                [
+                    node_features[:, self.layer_idx[i] : self.layer_idx[i + 1]].mean(
+                        dim=1
+                    )
+                    for i in range(len(self.layer_layout))
+                ],
+                dim=1,
+            )
+        elif self.pooling_method == "mean_edge" and self.pooling_layer_idx == "all":
+            graph_features = edge_features.mean(dim=(1, 2))
+        elif self.pooling_method == "max_edge" and self.pooling_layer_idx == "all":
+            graph_features = edge_features.flatten(1, 2).max(dim=1).values
+        elif self.pooling_method == "mean_edge" and self.pooling_layer_idx == "last":
+            graph_features = edge_features[:, :, -self.layer_layout[-1] :].mean(
+                dim=(1, 2)
+            )
+        return graph_features

tasks/utils/probe_features.py

@@ -0,0 +1,54 @@
+import hydra
+import torch
+import torch.nn as nn
+from einops.layers.torch import Rearrange
+
+from nn.inr import make_functional, params_to_tensor, wrap_func
+
+
+class GraphProbeFeatures(nn.Module):
+    def __init__(self, d_in, num_inputs, inr_model, input_init=None, proj_dim=None):
+        super().__init__()
+        inr = hydra.utils.instantiate(inr_model)
+        fmodel, params = make_functional(inr)
+
+        vparams, vshapes = params_to_tensor(params)
+        self.sirens = torch.vmap(wrap_func(fmodel, vshapes))
+
+        inputs = (
+            input_init
+            if input_init is not None
+            else 2 * torch.rand(1, num_inputs, d_in) - 1
+        )
+        self.inputs = nn.Parameter(inputs, requires_grad=input_init is None)
+
+        self.reshape_weights = Rearrange("b i o 1 -> b (o i)")
+        self.reshape_biases = Rearrange("b o 1 -> b o")
+
+        self.proj_dim = proj_dim
+        if proj_dim is not None:
+            self.proj = nn.ModuleList(
+                [
+                    nn.Sequential(
+                        nn.Linear(num_inputs, proj_dim),
+                        nn.LayerNorm(proj_dim),
+                    )
+                    for _ in range(inr.num_layers + 1)
+                ]
+            )
+
+    def forward(self, weights, biases):
+        weights = [self.reshape_weights(w) for w in weights]
+        biases = [self.reshape_biases(b) for b in biases]
+        params_flat = torch.cat(
+            [w_or_b for p in zip(weights, biases) for w_or_b in p], dim=-1
+        )
+
+        out = self.sirens(params_flat, self.inputs.expand(params_flat.shape[0], -1, -1))
+        if self.proj_dim is not None:
+            out = [proj(out[i].permute(0, 2, 1)) for i, proj in enumerate(self.proj)]
+            out = torch.cat(out, dim=1)
+            return out
+        else:
+            out = torch.cat(out, dim=-1)
+            return out.permute(0, 2, 1)

tasks/utils/relational_transformer.py

@@ -0,0 +1,361 @@
+import hydra
+import torch
+import torch.nn.functional as F
+from einops.layers.torch import Rearrange
+
+from utils.pooling import HomogeneousAggregator
+import torch.nn as nn
+
+
+class RelationalTransformer(nn.Module):
+    def __init__(
+        self,
+        d_node,
+        d_edge,
+        d_attn_hid,
+        d_node_hid,
+        d_edge_hid,
+        d_out_hid,
+        d_out,
+        n_layers,
+        n_heads,
+        layer_layout,
+        graph_constructor,
+        dropout=0.0,
+        node_update_type="rt",
+        disable_edge_updates=False,
+        use_cls_token=False,
+        pooling_method="cat",
+        pooling_layer_idx="last",
+        rev_edge_features=False,
+        modulate_v=True,
+        use_ln=True,
+        tfixit_init=False,
+    ):
+        super().__init__()
+        assert use_cls_token == (pooling_method == "cls_token")
+        self.pooling_method = pooling_method
+        self.pooling_layer_idx = pooling_layer_idx
+        self.rev_edge_features = rev_edge_features
+        self.nodes_per_layer = layer_layout
+        self.construct_graph = hydra.utils.instantiate(
+            graph_constructor,
+            d_node=d_node,
+            d_edge=d_edge,
+            layer_layout=layer_layout,
+            rev_edge_features=rev_edge_features,
+        )
+
+        self.use_cls_token = use_cls_token
+        if use_cls_token:
+            self.cls_token = nn.Parameter(torch.randn(d_node))
+
+        self.layers = nn.ModuleList(
+            [
+                torch.jit.script(
+                    RTLayer(
+                        d_node,
+                        d_edge,
+                        d_attn_hid,
+                        d_node_hid,
+                        d_edge_hid,
+                        n_heads,
+                        dropout,
+                        node_update_type=node_update_type,
+                        disable_edge_updates=(
+                            (disable_edge_updates or (i == n_layers - 1))
+                            and pooling_method != "mean_edge"
+                            and pooling_layer_idx != "all"
+                        ),
+                        modulate_v=modulate_v,
+                        use_ln=use_ln,
+                        tfixit_init=tfixit_init,
+                        n_layers=n_layers,
+                    )
+                )
+                for i in range(n_layers)
+            ]
+        )
+
+        if pooling_method != "cls_token":
+            self.pool = HomogeneousAggregator(
+                pooling_method,
+                pooling_layer_idx,
+                layer_layout,
+            )
+
+        self.num_graph_features = (
+            layer_layout[-1] * d_node
+            if pooling_method == "cat" and pooling_layer_idx == "last"
+            else d_edge if pooling_method in ("mean_edge", "max_edge") else d_node
+        )
+        self.proj_out = nn.Sequential(
+            nn.Linear(self.num_graph_features, d_out_hid),
+            nn.ReLU(),
+            # nn.Linear(d_out_hid, d_out_hid),
+            # nn.ReLU(),
+            nn.Linear(d_out_hid, d_out),
+        )
+
+        self.final_features = (None,None,None,None)
+
+    def forward(self, inputs):
+        attn_weights = None
+        node_features, edge_features, mask = self.construct_graph(inputs)
+        if self.use_cls_token:
+            node_features = torch.cat(
+                [
+                    # repeat(self.cls_token, "d -> b 1 d", b=node_features.size(0)),
+                    self.cls_token.unsqueeze(0).expand(node_features.size(0), 1, -1),
+                    node_features,
+                ],
+                dim=1,
+            )
+            edge_features = F.pad(edge_features, (0, 0, 1, 0, 1, 0), value=0)
+        for layer in self.layers:
+            node_features, edge_features, attn_weights = layer(node_features, edge_features, mask)
+
+        if self.pooling_method == "cls_token":
+            graph_features = node_features[:, 0]
+        else:
+            graph_features = self.pool(node_features, edge_features)
+        self.final_features = (graph_features, node_features, edge_features, attn_weights)
+        return self.proj_out(graph_features)
+
+
+class RTLayer(nn.Module):
+    def __init__(
+        self,
+        d_node,
+        d_edge,
+        d_attn_hid,
+        d_node_hid,
+        d_edge_hid,
+        n_heads,
+        dropout,
+        node_update_type="rt",
+        disable_edge_updates=False,
+        modulate_v=True,
+        use_ln=True,
+        tfixit_init=False,
+        n_layers=None,
+    ):
+        super().__init__()
+        self.node_update_type = node_update_type
+        self.disable_edge_updates = disable_edge_updates
+        self.use_ln = use_ln
+        self.n_layers = n_layers
+
+        self.self_attn = torch.jit.script(
+            RTAttention(
+                d_node,
+                d_edge,
+                d_attn_hid,
+                n_heads,
+                modulate_v=modulate_v,
+                use_ln=use_ln,
+            )
+        )
+        # self.self_attn = RTAttention(d_hid, d_hid, d_hid, n_heads)
+        self.lin0 = Linear(d_node, d_node)
+        self.dropout0 = nn.Dropout(dropout)
+        if use_ln:
+            self.node_ln0 = nn.LayerNorm(d_node)
+            self.node_ln1 = nn.LayerNorm(d_node)
+        else:
+            self.node_ln0 = nn.Identity()
+            self.node_ln1 = nn.Identity()
+
+        act_fn = nn.GELU
+
+        self.node_mlp = nn.Sequential(
+            Linear(d_node, d_node_hid, bias=False),
+            act_fn(),
+            Linear(d_node_hid, d_node),
+            nn.Dropout(dropout),
+        )
+
+        if not self.disable_edge_updates:
+            self.edge_updates = EdgeLayer(
+                d_node=d_node,
+                d_edge=d_edge,
+                d_edge_hid=d_edge_hid,
+                dropout=dropout,
+                act_fn=act_fn,
+                use_ln=use_ln,
+            )
+        else:
+            self.edge_updates = NoEdgeLayer()
+
+        if tfixit_init:
+            self.fixit_init()
+
+    def fixit_init(self):
+        temp_state_dict = self.state_dict()
+        n_layers = self.n_layers
+        for name, param in self.named_parameters():
+            if "weight" in name:
+                if name.split(".")[0] in ["node_mlp", "edge_mlp0", "edge_mlp1"]:
+                    temp_state_dict[name] = (0.67 * (n_layers) ** (-1.0 / 4.0)) * param
+                elif name.split(".")[0] in ["self_attn"]:
+                    temp_state_dict[name] = (0.67 * (n_layers) ** (-1.0 / 4.0)) * (
+                        param * (2**0.5)
+                    )
+
+        self.load_state_dict(temp_state_dict)
+
+    def node_updates(self, node_features, edge_features, mask):
+        out = self.self_attn(node_features, edge_features, mask)
+        attn_out, attn_weights = out
+        node_features = self.node_ln0(
+            node_features
+            + self.dropout0(
+                self.lin0(attn_out)
+            )
+        )
+        node_features = self.node_ln1(node_features + self.node_mlp(node_features))
+
+        return node_features, attn_weights
+
+    def forward(self, node_features, edge_features, mask):
+        node_features, attn_weights = self.node_updates(node_features, edge_features, mask)
+        edge_features = self.edge_updates(node_features, edge_features, mask)
+
+        return node_features, edge_features, attn_weights
+
+
+class EdgeLayer(nn.Module):
+    def __init__(
+        self,
+        *,
+        d_node,
+        d_edge,
+        d_edge_hid,
+        dropout,
+        act_fn,
+        use_ln=True,
+    ) -> None:
+        super().__init__()
+        self.edge_mlp0 = EdgeMLP(
+            d_edge=d_edge,
+            d_node=d_node,
+            d_edge_hid=d_edge_hid,
+            act_fn=act_fn,
+            dropout=dropout,
+        )
+        self.edge_mlp1 = nn.Sequential(
+            Linear(d_edge, d_edge_hid, bias=False),
+            act_fn(),
+            Linear(d_edge_hid, d_edge),
+            nn.Dropout(dropout),
+        )
+        if use_ln:
+            self.eln0 = nn.LayerNorm(d_edge)
+            self.eln1 = nn.LayerNorm(d_edge)
+        else:
+            self.eln0 = nn.Identity()
+            self.eln1 = nn.Identity()
+
+    def forward(self, node_features, edge_features, mask):
+        edge_features = self.eln0(
+            edge_features + self.edge_mlp0(node_features, edge_features)
+        )
+        edge_features = self.eln1(edge_features + self.edge_mlp1(edge_features))
+        return edge_features
+
+
+class NoEdgeLayer(nn.Module):
+    def forward(self, node_features, edge_features, mask):
+        return edge_features
+
+
+class EdgeMLP(nn.Module):
+    def __init__(self, *, d_node, d_edge, d_edge_hid, act_fn, dropout):
+        super().__init__()
+        self.reverse_edge = Rearrange("b n m d -> b m n d")
+        self.lin0_e = Linear(2 * d_edge, d_edge_hid)
+        self.lin0_s = Linear(d_node, d_edge_hid)
+        self.lin0_t = Linear(d_node, d_edge_hid)
+        self.act = act_fn()
+        self.lin1 = Linear(d_edge_hid, d_edge)
+        self.drop = nn.Dropout(dropout)
+
+    def forward(self, node_features, edge_features):
+        source_nodes = (
+            self.lin0_s(node_features)
+            .unsqueeze(-2)
+            .expand(-1, -1, node_features.size(-2), -1)
+        )
+        target_nodes = (
+            self.lin0_t(node_features)
+            .unsqueeze(-3)
+            .expand(-1, node_features.size(-2), -1, -1)
+        )
+
+        # reversed_edge_features = self.reverse_edge(edge_features)
+        edge_features = self.lin0_e(
+            torch.cat([edge_features, self.reverse_edge(edge_features)], dim=-1)
+        )
+        edge_features = edge_features + source_nodes + target_nodes
+        edge_features = self.act(edge_features)
+        edge_features = self.lin1(edge_features)
+        edge_features = self.drop(edge_features)
+
+        return edge_features
+
+
+class RTAttention(nn.Module):
+    def __init__(self, d_node, d_edge, d_hid, n_heads, modulate_v=None, use_ln=True):
+        super().__init__()
+        self.n_heads = n_heads
+        self.d_node = d_node
+        self.d_edge = d_edge
+        self.d_hid = d_hid
+        self.use_ln = use_ln
+        self.modulate_v = modulate_v
+        self.scale = 1 / (d_hid**0.5)
+        self.split_head_node = Rearrange("b n (h d) -> b h n d", h=n_heads)
+        self.split_head_edge = Rearrange("b n m (h d) -> b h n m d", h=n_heads)
+        self.cat_head_node = Rearrange("... h n d -> ... n (h d)", h=n_heads)
+
+        self.qkv_node = Linear(d_node, 3 * d_hid, bias=False)
+        self.edge_factor = 4 if modulate_v else 3
+        self.qkv_edge = Linear(d_edge, self.edge_factor * d_hid, bias=False)
+        self.proj_out = Linear(d_hid, d_node)
+
+    def forward(self, node_features, edge_features, mask):
+        qkv_node = self.qkv_node(node_features)
+        # qkv_node = rearrange(qkv_node, "b n (h d) -> b h n d", h=self.n_heads)
+        qkv_node = self.split_head_node(qkv_node)
+        q_node, k_node, v_node = torch.chunk(qkv_node, 3, dim=-1)
+
+        qkv_edge = self.qkv_edge(edge_features)
+        # qkv_edge = rearrange(qkv_edge, "b n m (h d) -> b h n m d", h=self.n_heads)
+        qkv_edge = self.split_head_edge(qkv_edge)
+        qkv_edge = torch.chunk(qkv_edge, self.edge_factor, dim=-1)
+        # q_edge, k_edge, v_edge, q_edge_b, k_edge_b, v_edge_b = torch.chunk(
+        #     qkv_edge, 6, dim=-1
+        # )
+        # qkv_edge = [item.masked_fill(mask.unsqueeze(1) == 0, 0) for item in qkv_edge]
+
+        q = q_node.unsqueeze(-2) + qkv_edge[0]  # + q_edge_b
+        k = k_node.unsqueeze(-3) + qkv_edge[1]  # + k_edge_b
+        if self.modulate_v:
+            v = v_node.unsqueeze(-3) * qkv_edge[3] + qkv_edge[2]
+        else:
+            v = v_node.unsqueeze(-3) + qkv_edge[2]
+        dots = self.scale * torch.einsum("b h i j d, b h i j d -> b h i j", q, k)
+        # dots.masked_fill_(mask.unsqueeze(1).squeeze(-1) == 0, -1e-9)
+
+        attn = F.softmax(dots, dim=-1)
+        out = torch.einsum("b h i j, b h i j d -> b h i d", attn, v)
+        out = self.cat_head_node(out)
+        return self.proj_out(out), attn
+
+
+def Linear(in_features, out_features, bias=True):
+    m = nn.Linear(in_features, out_features, bias)
+    nn.init.xavier_uniform_(m.weight)  # , gain=1 / math.sqrt(2))
+    if bias:
+        nn.init.constant_(m.bias, 0.0)
+    return m

tasks/utils/train.py

@@ -9,6 +9,14 @@ import glob
 from collections import OrderedDict
 from tqdm import tqdm
 import torch.distributed as dist
+import pandas as pd
+import xgboost as xgb
+from sklearn.metrics import accuracy_score, classification_report, roc_auc_score
+
+
+from torch.nn import ModuleList
+# from inr import INR
+# from kan import FasterKAN
 
 class Trainer(object):
     """
@@ -217,9 +225,12 @@ class Trainer(object):
         return train_loss, all_accs/total
     
     def train_batch(self, batch, batch_idx, device):
-        x, fft, y = batch['audio']['array'], batch['audio']['fft'], batch['label']
+        x, fft, y = batch['audio']['array'], batch['audio']['fft_mag'], batch['label']
+        # features = batch['audio']['features_arr'].to(device).float()
+        # cwt = batch['audio']['cwt_mag']
         x = x.to(device).float()
         fft = fft.to(device).float()
+        # cwt = cwt.to(device).float()
         y = y.to(device).float()
         x_fft = torch.cat((x.unsqueeze(dim=1), fft.unsqueeze(dim=1)), dim=1)
         y_pred = self.model(x_fft).squeeze()
@@ -255,7 +266,8 @@ class Trainer(object):
         return val_loss, all_accs/total
 
     def eval_batch(self, batch, batch_idx, device):
-        x, fft, y = batch['audio']['array'], batch['audio']['fft'], batch['label']
+        x, fft, y = batch['audio']['array'], batch['audio']['fft_mag'], batch['label']
+        # features = batch['audio']['features_arr'].to(device).float()
         x = x.to(device).float()
         fft = fft.to(device).float()
         x_fft = torch.cat((x.unsqueeze(dim=1), fft.unsqueeze(dim=1)), dim=1)
@@ -279,7 +291,7 @@ class Trainer(object):
         true_labels = []
         pbar = tqdm(test_dataloader)
         for i,batch in enumerate(pbar):
-            x, fft, y = batch['audio']['array'], batch['audio']['fft'], batch['label']
+            x, fft, y = batch['audio']['array'], batch['audio']['fft_mag'], batch['label']
             x = x.to(device).float()
             fft = fft.to(device).float()
             x_fft = torch.cat((x.unsqueeze(dim=1), fft.unsqueeze(dim=1)), dim=1)
@@ -297,4 +309,411 @@ class Trainer(object):
             pbar.set_description("acc: {:.4f}".format(acc))
             if i > self.max_iter:
                 break
-        return predictions, true_labels, all_accs/total
+        return predictions, true_labels, all_accs/total
+
+
+class INRDatabase:
+    """Database to store and manage INRs persistently."""
+
+    def __init__(self, save_dir='./inr_database'):
+        self.inrs = {}  # Maps sample_id -> INR
+        self.optimizers = {}  # Maps sample_id -> optimizer state
+        self.save_dir = save_dir
+        os.makedirs(save_dir, exist_ok=True)
+
+    def get_or_create_inr(self, sample_id, create_fn, device):
+        """Get existing INR or create new one if not exists."""
+        if sample_id not in self.inrs:
+            # Create new INR
+            inr = create_fn().to(device)
+            optimizer = torch.optim.Adam(inr.parameters())
+            self.inrs[sample_id] = inr
+            self.optimizers[sample_id] = optimizer
+        return self.inrs[sample_id], self.optimizers[sample_id]
+
+    def set_inr(self, sample_id, inr, optimizer):
+        self.inrs[sample_id] = inr
+        self.optimizers[sample_id] = optimizer
+
+    def save_state(self):
+        """Save all INRs and optimizer states to disk."""
+        state = {
+            'inrs': {
+                sample_id: inr.state_dict()
+                for sample_id, inr in self.inrs.items()
+            },
+            'optimizers': {
+                sample_id: opt.state_dict()
+                for sample_id, opt in self.optimizers.items()
+            }
+        }
+        torch.save(state, os.path.join(self.save_dir, 'inr_database.pt'))
+
+    def load_state(self, create_fn, device):
+        """Load INRs and optimizer states from disk."""
+        path = os.path.join(self.save_dir, 'inr_database.pt')
+        if os.path.exists(path):
+            state = torch.load(path, map_location=device)
+
+            # Restore INRs
+            for sample_id, inr_state in state['inrs'].items():
+                inr = create_fn().to(device)
+                inr.load_state_dict(inr_state)
+                self.inrs[sample_id] = inr
+
+            # Restore optimizers
+            for sample_id, opt_state in state['optimizers'].items():
+                optimizer = torch.optim.Adam(self.inrs[sample_id].parameters())
+                optimizer.load_state_dict(opt_state)
+                self.optimizers[sample_id] = optimizer
+
+
+class INRTrainer(Trainer):
+    def __init__(self, hidden_features=128, n_layers=3, in_features=1, out_features=1,
+                 num_steps=5000, lr=1e-3, inr_criterion=torch.nn.MSELoss(), save_dir='./inr_database', *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.hidden_features = hidden_features
+        self.n_layers = n_layers
+        self.in_features = in_features
+        self.out_features = out_features
+        self.num_steps = num_steps
+        self.lr = lr
+        self.inr_criterion = inr_criterion
+
+        # Initialize INR database
+        self.db = INRDatabase(save_dir)
+
+        # Load existing INRs if available
+        self.db.load_state(self.create_inr, self.device)
+
+    def create_inr(self):
+        """Factory function to create new INR instances."""
+        return INR(
+            hidden_features=self.hidden_features,
+            n_layers=self.n_layers,
+            in_features=self.in_features,
+            out_features=self.out_features
+        )
+
+    def create_kan(self):
+        return FasterKAN(layers_hidden=[self.in_features] + [self.hidden_features] * (self.n_layers) + [self.out_features],)
+
+    def get_sample_id(self, batch, idx):
+        """Extract unique identifier for a sample in the batch.
+        Override this method based on your data structure."""
+        # Example: if your batch contains unique IDs
+        if 'id' in batch:
+            return batch['id'][idx]
+        # Fallback: create hash from the sample data
+        sample_data = batch['audio']['array'][idx]
+        return hash(sample_data.cpu().numpy().tobytes())
+
+    def train_inr(self, optimizer, model, coords, values, num_iters=10, plot=False):
+        # pbar = tqdm(range(num_iters))
+        for _ in range(num_iters):
+            optimizer.zero_grad()
+            pred_values = model(coords.to(self.device)).float()
+            loss = self.inr_criterion(pred_values.squeeze(), values)
+            loss.backward()
+            optimizer.step()
+            # pbar.set_description(f'loss: {loss.item()}')
+        if plot:
+            plt.plot(values.cpu().detach().numpy())
+            plt.plot(pred_values.cpu().detach().numpy())
+            plt.title(loss.item())
+            plt.show()
+        return model, optimizer
+
+    def train_batch(self, batch, batch_idx, device):
+        """Train INRs for each sample in batch, persisting progress."""
+        coords = batch['audio']['coords'].to(device)  # [B, N, 1]
+        fft = batch['audio']['fft_mag'].to(device)  # [B, N]
+        audio = batch['audio']['array'].to(device)  # [B, N]
+        y = batch['label'].to(device).float()
+
+        batch_size = coords.shape[0]
+
+        values = audio
+
+        batch_losses = []
+        batch_optimizers = []
+        batch_inrs = []
+        batch_weights = tuple()
+        batch_biases = tuple()
+        # Training loop
+        # pbar = tqdm(range(self.num_steps), desc="Training INRs")
+        plot = batch_idx == 0
+        for i in range(batch_size):
+            sample_id = self.get_sample_id(batch, i)
+            inr, optimizer = self.db.get_or_create_inr(sample_id, self.create_inr, device)
+            inr, optimizer = self.train_inr(optimizer, inr, coords[i], values[i])
+            self.db.set_inr(sample_id, inr, optimizer)
+            # pred_values = inr(coords[i]).squeeze()
+            # batch_losses.append(self.inr_criterion(pred_values, values[i]))
+            # batch_optimizers.append(optimizer)
+            state_dict = inr.state_dict()
+            weights = tuple(
+                [v.permute(1, 0).unsqueeze(-1).unsqueeze(0).to(device) for w, v in state_dict.items() if "weight" in w]
+            )
+            biases = tuple([v.unsqueeze(-1).unsqueeze(0).to(device) for w, v in state_dict.items() if "bias" in w])
+            if not len(batch_weights):
+                batch_weights = weights
+            else:
+                batch_weights = tuple(
+                    [torch.cat((weights[i], batch_weights[i]), dim=0) for i in range(len(weights))]
+                )
+            if not len(batch_biases):
+                batch_biases = biases
+            else:
+                batch_biases = tuple(
+                    [torch.cat((biases[i], batch_biases[i]), dim=0) for i in range(len(biases))]
+                )
+        # loss_preds = torch.tensor([0])
+        # acc = 0
+        y_pred = self.model(inputs=(batch_weights, batch_biases)).squeeze()
+        loss_preds = self.criterion(y_pred, y)
+        self.optimizer.zero_grad()
+        loss_preds.backward()
+        self.optimizer.step()
+        # for i in range(batch_size):
+        #     batch_optimizers[i].zero_grad()
+        #     batch_losses[i] += loss_preds
+        #     batch_losses[i].backward()
+        #     batch_optimizers[i].step()
+
+
+        if batch_idx % 10 == 0:  # Adjust frequency as needed
+            self.db.save_state()
+
+        probs = torch.sigmoid(y_pred)
+        cls_pred = (probs > 0.5).float()
+        acc = (cls_pred == y).sum()
+
+
+        return loss_preds, acc, y
+
+    def eval_batch(self, batch, batch_idx, device):
+        """Evaluate INRs for each sample in batch."""
+        coords = batch['audio']['coords'].to(device)
+        fft = batch['audio']['fft_mag'].to(device)
+        audio = batch['audio']['array'].to(device)
+
+        batch_size = coords.shape[0]
+        # values = torch.cat((
+        #     audio.unsqueeze(-1),
+        #     fft.unsqueeze(-1)
+        # ), dim=-1)
+        values = audio
+        # Get INRs for each sample
+        batch_inrs = []
+        for i in range(batch_size):
+            sample_id = self.get_sample_id(batch, i)
+            inr, _ = self.db.get_or_create_inr(sample_id, self.create_inr, device)
+            batch_inrs.append(inr)
+
+        # Evaluate
+        with torch.no_grad():
+            all_preds = torch.stack([
+                inr(coords[i])
+                for i, inr in enumerate(batch_inrs)
+            ])
+
+            batch_losses = torch.stack([
+                self.criterion(all_preds[i].squeeze(), values[i])
+                for i in range(batch_size)
+            ])
+
+            avg_loss = batch_losses.mean().item()
+
+        acc = torch.zeros(self.output_dim, device=device)
+        y = values
+
+        return torch.tensor(avg_loss), acc, y
+
+
+def verify_parallel_gradient_isolation(trainer, batch_size=4, sequence_length=1000):
+    """
+    Verify that gradients remain isolated in parallel training.
+    """
+    device = trainer.device
+
+    # Create test data
+    coords = torch.linspace(0, 1, sequence_length).unsqueeze(-1)  # [N, 1]
+    coords = coords.unsqueeze(0).repeat(batch_size, 1, 1)  # [B, N, 1]
+
+    # Create synthetic signals
+    targets = torch.stack([
+        torch.sin(2 * torch.pi * (i + 1) * coords.squeeze(-1))
+        for i in range(batch_size)
+    ]).to(device)
+
+    # Create batch of INRs
+    inrs = trainer.create_batch_inrs()
+
+    # Store initial parameters
+    initial_params = [{name: param.clone().detach()
+                       for name, param in inr.named_parameters()}
+                      for inr in inrs]
+
+    # Create mock batch
+    batch = {
+        'audio': {
+            'coords': coords.to(device),
+            'fft_mag': targets.clone(),
+            'array': targets.clone()
+        }
+    }
+
+    # Run one training step
+    trainer.train_batch(batch, 0, device)
+
+    # Verify parameter changes
+    isolation_verified = True
+    for i, inr in enumerate(inrs):
+        params_changed = False
+        for name, param in inr.named_parameters():
+            if not torch.allclose(param, initial_params[i][name]):
+                params_changed = True
+                # Verify that the changes are unique to this INR
+                for j, other_inr in enumerate(inrs):
+                    if i != j:
+                        other_param = dict(other_inr.named_parameters())[name]
+                        if not torch.allclose(other_param, initial_params[j][name]):
+                            isolation_verified = False
+                            print(f"Warning: Parameter {name} of INR {j} changed when only INR {i} should have changed")
+
+    return isolation_verified
+
+class XGBoostTrainer():
+    def __init__(self, model_args, train_ds, val_ds, test_ds):
+        self.train_ds = train_ds
+        self.test_ds = test_ds
+        print("creating train dataframe...")
+        self.x_train, self.y_train = self.create_dataframe(train_ds, save_name='train')
+        print("creating validation dataframe...")
+        self.x_val, self.y_val = self.create_dataframe(val_ds, save_name='val')
+        print("creating test dataframe...")
+        self.x_test, self.y_test = self.create_dataframe(test_ds, save_name='test')
+
+        # Convert the data to DMatrix format
+        self.dtrain = xgb.DMatrix(self.x_train, label=self.y_train)
+        self.dval = xgb.DMatrix(self.x_val, label=self.y_val)
+        self.dtest = xgb.DMatrix(self.x_test, label=self.y_test)
+
+        # Model initialization
+        self.model_args = model_args
+        self.model = xgb.XGBClassifier(**model_args)
+
+    def create_dataframe(self, ds, save_name='train'):
+        try:
+            df = pd.read_csv(f"tasks/utils/dfs/{save_name}.csv")
+        except FileNotFoundError:
+            data = []
+
+            # Iterate over the dataset
+            pbar = tqdm(enumerate(ds))
+            for i, batch in pbar:
+                label = batch['label']
+                features = batch['audio']['features']
+
+                # Flatten the nested dictionary structure
+                feature_dict = {'label': label}
+                for k, v in features.items():
+                    if isinstance(v, dict):
+                        for sub_k, sub_v in v.items():
+                            feature_dict[f"{k}_{sub_k}"] = sub_v[0].item()  # Aggregate (e.g., mean)
+                data.append(feature_dict)
+            # Convert to DataFrame
+            df = pd.DataFrame(data)
+            print(os.getcwd())
+            df.to_csv(f"tasks/utils/dfs/{save_name}.csv", index=False)
+        X = df.drop(columns=['label'])
+        y = df['label']
+        return X, y
+
+    def fit(self):
+        # Train using the `train` method with early stopping
+        params = {
+            'objective': 'binary:logistic',
+            'eval_metric': 'logloss',
+            **self.model_args
+        }
+
+        evals_result = {}
+        num_boost_round = 1000  # Set a large number of boosting rounds
+
+        # Watchlist to monitor performance on train and validation data
+        watchlist = [(self.dtrain, 'train'), (self.dval, 'eval')]
+
+        # Train the model
+        self.model = xgb.train(
+            params,
+            self.dtrain,
+            num_boost_round=num_boost_round,
+            evals=watchlist,
+            early_stopping_rounds=10,  # Early stopping after 10 rounds with no improvement
+            evals_result=evals_result,
+            verbose_eval=True  # Show evaluation results for each iteration
+        )
+
+        return evals_result
+
+    def train_xgboost_in_batches(self, dataloader, eval_metric="logloss"):
+        evals_result = {}
+        for i, batch in enumerate(dataloader):
+            # Convert batch data to NumPy arrays
+            X_batch = torch.cat([batch[key].view(batch[key].size(0), -1) for key in batch if key != "label"],
+                                dim=1).numpy()
+            y_batch = batch["label"].numpy()
+
+            # Create DMatrix for XGBoost
+            dtrain = xgb.DMatrix(X_batch, label=y_batch)
+
+            # Use `train` with each batch
+            self.model = xgb.train(
+                params,
+                dtrain,
+                num_boost_round=1000,  # Use a large number of rounds
+                evals=[(self.dval, 'eval')],
+                eval_metric=eval_metric,
+                early_stopping_rounds=10,
+                evals_result=evals_result,
+                verbose_eval=False  # Avoid printing every iteration
+            )
+
+            # Optionally print progress
+            if i % 10 == 0:
+                print(f"Batch {i + 1}/{len(dataloader)} processed.")
+
+        return evals_result
+
+    def predict(self):
+        # Predict probabilities for class 1
+        y_prob = self.model.predict(self.dtest, output_margin=False)
+
+        # Convert probabilities to binary labels (0 or 1) using a threshold (e.g., 0.5)
+        y_pred = (y_prob >= 0.5).astype(int)
+
+        # Evaluate performance
+        accuracy = accuracy_score(self.y_test, y_pred)
+        roc_auc = roc_auc_score(self.y_test, y_prob)
+
+        print(f'Accuracy: {accuracy:.4f}')
+        print(f'ROC AUC Score: {roc_auc:.4f}')
+        print(classification_report(self.y_test, y_pred))
+
+    def plot_results(self, evals_result):
+        train_logloss = evals_result["train"]["logloss"]
+        val_logloss = evals_result["eval"]["logloss"]
+        iterations = range(1, len(train_logloss) + 1)
+
+        # Plot
+        plt.figure(figsize=(8, 5))
+        plt.plot(iterations, train_logloss, label="Train LogLoss", color="blue")
+        plt.plot(iterations, val_logloss, label="Validation LogLoss", color="red")
+        plt.xlabel("Boosting Round (Iteration)")
+        plt.ylabel("Log Loss")
+        plt.title("Training and Validation Log Loss over Iterations")
+        plt.legend()
+        plt.grid()
+        plt.show()

tasks/utils/transforms.py

@@ -0,0 +1,272 @@
+import numpy as np
+import librosa
+import torch
+import torch.nn as nn
+import pywt
+from scipy import signal
+
+
+
+def compute_cwt_power_spectrum(audio, sample_rate, num_freqs=128, f_min=20, f_max=None):
+    """
+    Compute the power spectrum of continuous wavelet transform using Morlet wavelet.
+
+    Parameters:
+        audio: torch.Tensor
+            Input audio signal
+        sample_rate: int
+            Sampling rate of the audio
+        num_freqs: int
+            Number of frequency bins for the CWT
+        f_min: float
+            Minimum frequency to analyze
+        f_max: float or None
+            Maximum frequency to analyze (defaults to Nyquist frequency)
+
+    Returns:
+        torch.Tensor: CWT power spectrum
+    """
+    # Convert to numpy
+    audio_np = audio.cpu().numpy()
+
+    # Set default f_max to Nyquist frequency if not specified
+    if f_max is None:
+        f_max = sample_rate // 2
+
+    # Generate frequency bins (logarithmically spaced)
+    frequencies = np.logspace(
+        np.log10(f_min),
+        np.log10(f_max),
+        num=num_freqs
+    )
+
+    # Compute the width of the wavelet (in samples)
+    widths = sample_rate / (2 * frequencies * np.pi)
+
+    # Compute CWT using Morlet wavelet
+    cwt = signal.cwt(
+        audio_np,
+        signal.morlet2,
+        widths,
+        w=5.0  # Width parameter of Morlet wavelet
+    )
+
+    # Compute power spectrum (magnitude squared)
+    power_spectrum = np.abs(cwt) ** 2
+
+    # Convert to torch tensor
+    power_spectrum_tensor = torch.FloatTensor(power_spectrum)
+
+    return power_spectrum_tensor
+
+def compute_wavelet_transform(audio, wavelet, decompos_level):
+    """Compute wavelet decomposition of the audio signal."""
+    # Convert to numpy and ensure 1D
+    audio_np = audio.cpu().numpy()
+
+    # Perform wavelet decomposition
+    coeffs = pywt.wavedec(audio_np, wavelet, level=decompos_level)
+
+    # Stack coefficients into a 2D array
+    # First, pad all coefficient arrays to the same length
+    max_len = max(len(c) for c in coeffs)
+    padded_coeffs = []
+    for coeff in coeffs:
+        pad_len = max_len - len(coeff)
+        if pad_len > 0:
+            padded_coeff = np.pad(coeff, (0, pad_len), mode='constant')
+        else:
+            padded_coeff = coeff
+        padded_coeffs.append(padded_coeff)
+
+    # Stack into 2D array where each row is a different scale
+    wavelet_features = np.stack(padded_coeffs)
+
+    # Convert to tensor
+    return torch.FloatTensor(wavelet_features)
+
+
+def compute_melspectrogram(audio, sample_rate):
+    mel_spec = librosa.feature.melspectrogram(
+        y=audio.cpu().numpy(),
+        sr=sample_rate,
+        n_mels=128
+    )
+    return torch.FloatTensor(librosa.power_to_db(mel_spec))
+
+
+def compute_mfcc(audio, sample_rate):
+    mfcc = librosa.feature.mfcc(
+        y=audio.cpu().numpy(),
+        sr=sample_rate,
+        n_mfcc=20
+    )
+    return torch.FloatTensor(mfcc)
+
+
+def compute_chroma(audio, sample_rate):
+    chroma = librosa.feature.chroma_stft(
+        y=audio.cpu().numpy(),
+        sr=sample_rate
+    )
+    return torch.FloatTensor(chroma)
+
+
+def compute_time_domain_features(audio, sample_rate, frame_length=2048, hop_length=128):
+    """
+    Compute time-domain features from audio signal.
+    Returns a dictionary of features.
+    """
+    # Convert to numpy
+    audio_np = audio.cpu().numpy()
+
+    # Initialize dictionary for features
+    features = {}
+
+    # 1. Zero Crossing Rate
+    zcr = librosa.feature.zero_crossing_rate(
+        y=audio_np,
+        frame_length=frame_length,
+        hop_length=hop_length
+    )
+    features['zcr'] = torch.Tensor([zcr.sum()])
+
+    # 2. Root Mean Square Energy
+    rms = librosa.feature.rms(
+        y=audio_np,
+        frame_length=frame_length,
+        hop_length=hop_length
+    )
+    features['rms_energy'] = torch.Tensor([rms.mean()])
+
+    # 3. Temporal Statistics
+    frames = librosa.util.frame(audio_np, frame_length=frame_length, hop_length=hop_length)
+    features['mean'] = torch.Tensor([np.mean(frames, axis=0).mean()])
+    features['std'] = torch.Tensor([np.std(frames, axis=0).mean()])
+    features['max'] = torch.Tensor([np.max(frames, axis=0).mean()])
+
+    # 4. Tempo and Beat Features
+    onset_env = librosa.onset.onset_strength(y=audio_np, sr=sample_rate)
+    tempo = librosa.beat.tempo(onset_envelope=onset_env, sr=sample_rate)
+    features['tempo'] = torch.Tensor(tempo)
+
+    # 5. Amplitude Envelope
+    envelope = np.abs(librosa.stft(audio_np, n_fft=frame_length, hop_length=hop_length))
+    features['envelope'] = torch.Tensor([np.mean(envelope, axis=0).mean()])
+
+    return features
+
+
+def compute_frequency_domain_features(audio, sample_rate, n_fft=2048, hop_length=512):
+    """
+    Compute frequency-domain features from audio signal.
+    Returns a dictionary of features.
+    """
+    # Convert to numpy
+    audio_np = audio.cpu().numpy()
+
+    # Initialize dictionary for features
+    features = {}
+
+    # 1. Spectral Centroid
+    try:
+        spectral_centroids = librosa.feature.spectral_centroid(
+            y=audio_np,
+            sr=sample_rate,
+            n_fft=n_fft,
+            hop_length=hop_length,
+
+        )
+        features['spectral_centroid'] = torch.FloatTensor([spectral_centroids.max()])
+    except Exception as e:
+        features['spectral_centroid'] = torch.FloatTensor([np.nan])
+
+    # 2. Spectral Rolloff
+    try:
+        spectral_rolloff = librosa.feature.spectral_rolloff(
+            y=audio_np,
+            sr=sample_rate,
+            n_fft=n_fft,
+            hop_length=hop_length,
+
+        )
+        features['spectral_rolloff'] = torch.FloatTensor([spectral_rolloff.max()])
+    except Exception as e:
+        features['spectral_rolloff'] = torch.FloatTensor([np.nan])
+
+    # 3. Spectral Bandwidth
+    try:
+        spectral_bandwidth = librosa.feature.spectral_bandwidth(
+            y=audio_np,
+            sr=sample_rate,
+            n_fft=n_fft,
+            hop_length=hop_length
+        )
+        features['spectral_bandwidth'] = torch.FloatTensor([spectral_bandwidth.max()])
+    except Exception as e:
+        features['spectral_bandwidth'] = torch.FloatTensor([np.nan])
+    # 4. Spectral Contrast
+    try:
+        spectral_contrast = librosa.feature.spectral_contrast(
+            y=audio_np,
+            sr=sample_rate,
+            n_fft=n_fft,
+            hop_length=hop_length,
+            fmin=20,  # Lower minimum frequency
+            n_bands=4,  # Reduce number of bands
+            quantile=0.02
+        )
+        features['spectral_contrast'] = torch.FloatTensor([spectral_contrast.mean()])
+    except Exception as e:
+        features['spectral_contrast'] = torch.FloatTensor([np.nan])
+
+    # 5. Spectral Flatness
+    try:
+        spectral_flatness = librosa.feature.spectral_flatness(
+            y=audio_np,
+            n_fft=n_fft,
+            hop_length=hop_length
+        )
+        features['spectral_flatness'] = torch.FloatTensor([spectral_flatness.max()])
+    except Exception as e:
+        features['spectral_flatness'] = torch.FloatTensor([np.nan])
+
+    # 6. Spectral Flux
+    try:
+        stft = np.abs(librosa.stft(audio_np, n_fft=n_fft, hop_length=hop_length))
+        spectral_flux = np.diff(stft, axis=1)
+        spectral_flux = np.pad(spectral_flux, ((0, 0), (1, 0)), mode='constant')
+        features['spectral_flux'] = torch.FloatTensor([np.std(spectral_flux)])
+    except Exception as e:
+        features['spectral_flux'] = torch.FloatTensor([np.nan])
+
+    return features
+
+
+def compute_all_features(audio, sample_rate, wavelet='db1', decompos_level=4):
+    """
+    Compute all available features and return them in a dictionary.
+    """
+    features = {}
+
+    # Basic transformations
+    # features['wavelet'] = compute_wavelet_transform(audio, wavelet, decompos_level)
+    # features['melspectrogram'] = compute_melspectrogram(audio, sample_rate)
+    # features['mfcc'] = compute_mfcc(audio, sample_rate)
+    # features['chroma'] = compute_chroma(audio, sample_rate)
+
+    # features['cwt_power'] = compute_cwt_power_spectrum(
+    #     audio,
+    #     sample_rate,
+    #     num_freqs=128,  # Same as mel bands for consistency
+    #     f_min=20,  # Standard lower frequency bound
+    #     f_max=sample_rate // 2  # Nyquist frequency
+    # )
+
+    # Time domain features
+    # features['time_domain'] = compute_time_domain_features(audio, sample_rate)
+
+    # Frequency domain features
+    features['frequency_domain'] = compute_frequency_domain_features(audio, sample_rate)
+
+    return features