lucazsh/RadNet

Introduction:

RadNet it is composed of 2 LSTM (Long-Short Term Memory) layers that helps the model to identify patterns from the dataset, that was trained on (lucazsh/RadNet), which will facilitate the prediction of the daily radiation amount as well as solar storms in LEO (Low Earth Orbit). This model was also featured in the National Space Society (NSS) contest under the project Aletheia.

The radiation data was extracted from the NASA OSDR EDA API (https://visualization.osdr.nasa.gov/eda/), from the RR's missions (RR-1, RR-3, RR-4, RR-6, RR-8, RR-9, RR-12, RR-17, and RR-19).

Installation

To install the dependencies, you need to make sure you have Python 3.10+ installed on your system. If not, you can download it from python.org.

Install dependencies:

pip install numpy json torch

Code:

Below is the Python code that runs the RadNet model and displays the daily radiation dose and solar storm:

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()