# Sinusoidal Frequency Predictor using PyTorch

# This script generates sinusoidal data, defines a neural network to predict frequencies,

# trains the model, and evaluates its performance using various metrics and visual comparisons.

import torch

import torch.nn as nn

import torch.optim as optim

import numpy as np

from torch.utils.data import Dataset, DataLoader

# Hyperparameters

NUM_SAMPLES = 1000  # Number of sinusoids to generate

SEQ_LENGTH = 100  # Number of points per sinusoid

FREQ_RANGE = (1, 10)  # Frequency range of sinusoids

BATCH_SIZE = 32

EPOCHS = 50

LEARNING_RATE = 0.01

# Generate sinusoidal data

def generate_sinusoid(frequency, length=SEQ_LENGTH, sample_rate=100):

    t = np.linspace(0, 1, sample_rate, endpoint=False)

    return np.sin(2 * np.pi * frequency * t)

class SinusoidDataset(Dataset):

    def __init__(self, num_samples=NUM_SAMPLES, freq_range=FREQ_RANGE):

        self.data = []

        self.labels = []

        for _ in range(num_samples):

            freq = np.random.uniform(freq_range[0], freq_range[1])

            signal = generate_sinusoid(freq)

            self.data.append(signal)

            self.labels.append(freq)

        self.data = torch.tensor(self.data, dtype=torch.float32)

        self.labels = torch.tensor(self.labels, dtype=torch.float32).view(-1, 1)

    def __len__(self):

        return len(self.data)

    def __getitem__(self, idx):

        return self.data[idx], self.labels[idx]

# Define the neural network

class FrequencyPredictor(nn.Module):

    def __init__(self):

        super(FrequencyPredictor, self).__init__()

        self.model = nn.Sequential(

            nn.Linear(SEQ_LENGTH, 64),

            nn.ReLU(),

            nn.Linear(64, 32),

            nn.ReLU(),

            nn.Linear(32, 1)

        )

    def forward(self, x):

        return self.model(x)

# Initialize dataset and dataloaders

dataset = SinusoidDataset()

dataloader = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True)

# Initialize model, loss function, and optimizer

model = FrequencyPredictor()

criterion = nn.MSELoss()

optimizer = optim.Adam(model.parameters(), lr=LEARNING_RATE)

# Training loop

for epoch in range(EPOCHS):

    epoch_loss = 0.0

    for signals, freqs in dataloader:

        optimizer.zero_grad()

        outputs = model(signals)

        loss = criterion(outputs, freqs)

        loss.backward()

        optimizer.step()

        epoch_loss += loss.item()

    print(f"Epoch {epoch+1}/{EPOCHS}, Loss: {epoch_loss/len(dataloader):.4f}")

print("Training complete!")

# Evaluation

import matplotlib.pyplot as plt

from sklearn.metrics import mean_absolute_error, r2_score

# Generate test data

test_frequencies = np.linspace(FREQ_RANGE[0], FREQ_RANGE[1], 20)

test_signals = np.array([generate_sinusoid(f) for f in test_frequencies])

test_signals_tensor = torch.tensor(test_signals, dtype=torch.float32)

# Predict frequencies with trained model

model.eval()

with torch.no_grad():

    predicted_frequencies = model(test_signals_tensor).numpy().flatten()

# Method 1: Compute Mean Absolute Error (MAE)

mae = mean_absolute_error(test_frequencies, predicted_frequencies)

print(f"Mean Absolute Error: {mae:.4f}")

# Method 2: Compute R-squared score (R²)

r2 = r2_score(test_frequencies, predicted_frequencies)

print(f"R-squared score: {r2:.4f}")

# Method 3: Visual Comparison of Predictions

plt.figure(figsize=(8, 6))

plt.plot(test_frequencies, test_frequencies, label="True Frequencies", linestyle='dashed')

plt.scatter(test_frequencies, predicted_frequencies, color='red', label="Predicted Frequencies")

plt.xlabel("True Frequency")

plt.ylabel("Predicted Frequency")

plt.legend()

plt.title("True vs Predicted Frequencies")

plt.show()

# Conclusion

# This script successfully trains a neural network to predict frequencies of generated 

# sinusoidal signals. Evaluation metrics such as MAE and R² score are computed, and the 

# results are visually compared. Additionally, we experimented by changing the epochs and 

# learning rate to observe different outcomes in the model's performance.