lucazsh/RadNet

import json
import torch
import torch.nn as nn
import numpy as np

class RadNet(nn.Module):
    def __init__(self, input_size=2, hidden_size=64, num_layers=2, output_size=2):
        super(RadNet, self).__init__()
        self.rnn = nn.LSTM(input_size, hidden_size, num_layers, batch_first=True,
                           dropout=0.2)
        self.fc = nn.Linear(hidden_size, output_size)

    def forward(self, x):
        # x = [batch_size, sequence_length, input_size]
        out, _ = self.rnn(x)  # [batch_size, sequence_length, hidden_size]
        return self.fc(out[:, -1, :])  # [batch_size, output_size]

def main():
    path = "./RadNet.pth"  # path to the pre-trained RadNet checkpoint
    data_path = "test.json"  # path to the JSON file with test data

    # load the data from JSON
    with open(data_path, 'r') as f:
        data = json.load(f)

    # load the checkpoint to the preferred device, in our case cpu
    checkpoint = torch.load(path, weights_only=False, map_location=torch.device('cpu'))
    cfg = checkpoint["config"]
    mean = checkpoint["mean"]
    std = checkpoint["std"]

    model = RadNet(hidden_size=64, num_layers=2)
    model.load_state_dict(checkpoint["model_state_dict"])
    model.eval()

    tmp_data = []
    for d in data:
        solar_storm = d.get("solar_storm_score", 0.0)
        daily_radiation = d.get("sv_per_day_mSv", 0.0)
        tmp_data.append([solar_storm, daily_radiation])

    raw_data = np.array(tmp_data, dtype=float)

    print(f"RadNet was trained with a sequence of: {cfg['seq_length']} days.")
    print(f"Available data in the JSON dataset: {len(raw_data)} days.")

    user_seq_len = int(input(f"Enter the number of days: "))
    if len(raw_data) < user_seq_len:
        print(f"You do not have enough data in the JSON file for {user_seq_len} days.")
        return

    input_seq = raw_data[-user_seq_len:]
    input_norm = (input_seq - mean) / (std + 1e-6)
    input_tensor = torch.tensor(input_norm, dtype=torch.float32).unsqueeze(0)

    # predict using the model
    with torch.no_grad():
        pred_norm = model(input_tensor).numpy()[0]
        pred_values = (pred_norm * (std + 1e-6)) + mean

    if user_seq_len == 1:
        print(f"Prediction based on the last day:")
    else:
        print(f"Prediction based on the last {user_seq_len} days:")

    print(f"Solar storm score: {pred_values[0]:.6f}")
    print(f"Radiation (mSv): {pred_values[1]:.6f}")

if __name__ == "__main__":
    main()