Bug fixes

.gitignore

@@ -172,3 +172,11 @@ cython_debug/
 
 # PyPI configuration file
 .pypirc
+
+.vscode/sftp.json
+
+.DS_Store
+
+.cursorignore
+
+wandb/

flaring/forecasting/inference/evaluation.py

@@ -838,7 +838,7 @@ class SolarFlareEvaluator:
 
             #plt.suptitle(f'Timestamp: {timestamp}', fontsize=14)
             #plt.tight_layout()
-            plt.savefig(save_path, dpi=500, bbox_inches='tight')
+            plt.savefig(save_path, dpi=500)
             plt.close()
 
             print(f"Worker {os.getpid()}: Completed {timestamp}")
@@ -941,10 +941,10 @@ class SolarFlareEvaluator:
 
 if __name__ == "__main__":
     # Example paths - replace with your actual paths
-    vit_csv = "/mnt/data/ML-Ready-mixed/ML-Ready-mixed/output/final_epoch_patch.csv"
+    vit_csv = "/mnt/data/ML-READY/output/final_epoch_patch.csv"
     baseline_results_csv = ""
-    aia_data = "/mnt/data/ML-Ready-mixed/ML-Ready-mixed/AIA/test/"
-    weights_directory = "/mnt/data/ML-Ready-mixed/ML-Ready-mixed/final_epoch_patch_weights_final"
+    aia_data = "/mnt/data/ML-READY/AIA/test/"
+    weights_directory = "/mnt/data/ML-READY/final_epoch_patch_weights_final"
 
     # Sample timestamps - Fixed the datetime generation
     start_time = datetime(2023, 8, 5, 20,30,00)
@@ -962,7 +962,7 @@ if __name__ == "__main__":
         baseline_csv_path=baseline_results_csv,
         aia_dir=aia_data,
         weight_path=weights_directory,
-        output_dir="/mnt/data/ML-Ready-mixed/ML-Ready-mixed/solar_flare_comparison_results/improving_patch"
+        output_dir="/mnt/data/ML-READY/solar_flare_comparison_results/improving_patch"
     )
 
     # Run complete evaluation with baseline comparison and uncertainties

flaring/forecasting/inference/inference.py

@@ -6,6 +6,13 @@ import torch
 import numpy as np
 from torch.utils.checkpoint import checkpoint
 from torch.utils.data import DataLoader
+
+from pathlib import Path
+
+# Add project root to Python path
+PROJECT_ROOT = Path(__file__).parent.parent.parent.parent.absolute()
+sys.path.insert(0, str(PROJECT_ROOT))
+
 from flaring.forecasting.data_loaders.SDOAIA_dataloader import AIA_GOESDataset
 import flaring.forecasting.models as models
 from flaring.forecasting.models.vision_transformer_custom import ViT

flaring/forecasting/inference/inference_on_patch.py

@@ -1,6 +1,12 @@
 import argparse
 import re
 import sys
+from pathlib import Path
+
+# Add project root to Python path
+PROJECT_ROOT = Path(__file__).parent.parent.parent.parent.absolute()
+sys.path.insert(0, str(PROJECT_ROOT))
+
 import pandas as pd
 import torch
 import numpy as np

flaring/forecasting/inference/inference_on_patch_config.yaml

@@ -1,4 +1,4 @@
-base_data_dir: "/mnt/data/ML-Ready-mixed/ML-Ready-mixed/"  # Change this line for different datasets
+base_data_dir: "/mnt/data/ML-READY/"  # Change this line for different datasets
 output_path: "${base_data_dir}/output/final_epoch_patch.csv"
 weight_path: "${base_data_dir}/final_epoch_patch_weights_final/"
 flux_path: "${base_data_dir}/patch_flux/"
@@ -25,7 +25,7 @@ data:
   sxr_dir:
     "${base_data_dir}/SXR"
   sxr_norm_path:
-    "/mnt/data/ML-Ready-mixed/ML-Ready-mixed/SXR/normalized_sxr.npy"
+    "/mnt/data/ML-READY/SXR/normalized_sxr.npy"
   checkpoint_path:
-    "/mnt/data/ML-Ready-mixed/ML-Ready-mixed/new-checkpoint/vit-16-higher-weight-lower-decay-epoch=119-val_total_loss=0.0419.ckpt"
+    "/mnt/data/ML-READY/new-checkpoint/vit-16-higher-weight-lower-decay-epoch=288-val_total_loss=0.0385.ckpt"
 

flaring/forecasting/models/efficientnet.py

@@ -1,45 +0,0 @@
-import torchvision
-from torch import nn
-from torch.nn import HuberLoss
-from irradiance.models.base_model import BaseModel
-
-
-class EfficientnetIrradiance(BaseModel):
-
-    def __init__(self, d_input, d_output, eve_norm, loss_func=HuberLoss(), model='efficientnet_b0', dp=0.75, lr=1e-4):
-        if model == 'efficientnet_b0':
-            model = torchvision.models.efficientnet_b0(pretrained=True)
-        elif model == 'efficientnet_b1': 
-            model = torchvision.models.efficientnet_b1(pretrained=True)
-        elif model == 'efficientnet_b2': 
-            model = torchvision.models.efficientnet_b2(pretrained=True)
-        elif model == 'efficientnet_b3':
-            model = torchvision.models.efficientnet_b3(pretrained=True)
-        elif model == 'efficientnet_b4': 
-            model = torchvision.models.efficientnet_b4(pretrained=True)
-        elif model == 'efficientnet_b5': 
-            model = torchvision.models.efficientnet_b5(pretrained=True)
-        elif model == 'efficientnet_b6': 
-            model = torchvision.models.efficientnet_b6(pretrained=True)
-        elif model == 'efficientnet_b7': 
-            model = torchvision.models.efficientnet_b7(pretrained=True)
-        conv1_out = model.features[0][0].out_channels
-        model.features[0][0] = nn.Conv2d(d_input, conv1_out, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
-
-        lin_in = model.classifier[1].in_features
-        # consider adding average pool of full image(s)
-        classifier = nn.Sequential(nn.Dropout(p=dp, inplace=True),
-                                   nn.Linear(in_features=lin_in, out_features=d_output, bias=True))
-        model.classifier = classifier
-        # set all dropouts to 0.75
-        # TODO: other dropout values?
-        for m in model.modules():
-            if m.__class__.__name__.startswith('Dropout'):
-                m.p = dp
-        model = model
-
-        super().__init__(model=model, eve_norm=eve_norm, loss_func=loss_func, lr=lr)
-    
-    def forward(self, x):
-        x = self.model(x)
-        return x

flaring/forecasting/models/kan_success.py

@@ -1,219 +0,0 @@
-# Copyright 2024 Li, Ziyao
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#      http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-from typing import *
-from torch.nn import HuberLoss
-from irradiance.models.base_model import BaseModel
-
-
-class SplineLinear(nn.Linear):
-    def __init__(self, in_features: int, out_features: int, init_scale: float = 0.1, **kw) -> None:
-        self.init_scale = init_scale
-        super().__init__(in_features, out_features, bias=False, **kw)
-
-    def reset_parameters(self) -> None:
-        nn.init.trunc_normal_(self.weight, mean=0, std=self.init_scale)
-
-
-class RadialBasisFunction(nn.Module):
-    def __init__(
-        self,
-        grid_min: float = -2.,
-        grid_max: float = 2.,
-        num_grids: int = 8,
-        denominator: float = None,  # larger denominators lead to smoother basis
-    ):
-        super().__init__()
-        self.grid_min = grid_min
-        self.grid_max = grid_max
-        self.num_grids = num_grids
-        grid = torch.linspace(grid_min, grid_max, num_grids)
-        self.grid = torch.nn.Parameter(grid, requires_grad=False)
-        self.denominator = denominator or (grid_max - grid_min) / (num_grids - 1)
-
-    def forward(self, x):
-        return torch.exp(-((x[..., None] - self.grid) / self.denominator) ** 2)
-
-class FastKANLayer(nn.Module):
-    def __init__(
-        self,
-        input_dim: int,
-        output_dim: int,
-        grid_min: float = -2.,
-        grid_max: float = 2.,
-        num_grids: int = 8,
-        use_base_update: bool = True,
-        use_layernorm: bool = True,
-        base_activation = F.silu,
-        spline_weight_init_scale: float = 0.1,
-    ) -> None:
-        super().__init__()
-        self.input_dim = input_dim
-        self.output_dim = output_dim
-        self.layernorm = None
-        if use_layernorm:
-            assert input_dim > 1, "Do not use layernorms on 1D inputs. Set `use_layernorm=False`."
-            self.layernorm = nn.LayerNorm(input_dim)
-        self.rbf = RadialBasisFunction(grid_min, grid_max, num_grids)
-        self.spline_linear = SplineLinear(input_dim * num_grids, output_dim, spline_weight_init_scale)
-        self.use_base_update = use_base_update
-        if use_base_update:
-            self.base_activation = base_activation
-            self.base_linear = nn.Linear(input_dim, output_dim)
-
-    def forward(self, x, use_layernorm=True):
-        if self.layernorm is not None and use_layernorm:
-            spline_basis = self.rbf(self.layernorm(x))
-        else:
-            spline_basis = self.rbf(x)
-        ret = self.spline_linear(spline_basis.view(*spline_basis.shape[:-2], -1))
-        if self.use_base_update:
-            base = self.base_linear(self.base_activation(x))
-            ret = ret + base
-        return ret
-
-    def plot_curve(
-        self,
-        input_index: int,
-        output_index: int,
-        num_pts: int = 1000,
-        num_extrapolate_bins: int = 2
-    ):
-        '''this function returns the learned curves in a FastKANLayer.
-        input_index: the selected index of the input, in [0, input_dim) .
-        output_index: the selected index of the output, in [0, output_dim) .
-        num_pts: num of points sampled for the curve.
-        num_extrapolate_bins (N_e): num of bins extrapolating from the given grids. The curve 
-            will be calculate in the range of [grid_min - h * N_e, grid_max + h * N_e].
-        '''
-        ng = self.rbf.num_grids
-        h = self.rbf.denominator
-        assert input_index < self.input_dim
-        assert output_index < self.output_dim
-        w = self.spline_linear.weight[
-            output_index, input_index * ng : (input_index + 1) * ng
-        ]   # num_grids,
-        x = torch.linspace(
-            self.rbf.grid_min - num_extrapolate_bins * h,
-            self.rbf.grid_max + num_extrapolate_bins * h,
-            num_pts
-        )   # num_pts, num_grids
-        with torch.no_grad():
-            y = (w * self.rbf(x.to(w.dtype))).sum(-1)
-        return x, y
-
-
-class FastKANIrradiance(BaseModel):
-    def __init__(
-        self,
-        eve_norm,
-        layers_hidden: List[int],
-        grid_min: float = -2.,
-        grid_max: float = 2.,
-        num_grids: int = 8,
-        use_base_update: bool = True,
-        base_activation = F.silu,
-        spline_weight_init_scale: float = 0.1,
-        loss_func = HuberLoss(),
-        lr=1e-4,
-        use_std=False   
-    ) -> None:
-        super().__init__(model=None, eve_norm=eve_norm, loss_func=loss_func, lr=lr)
-        self.use_std = use_std
-        if use_std:
-            layers_hidden[0] = layers_hidden[0]*2
-        self.layers = nn.ModuleList([
-            FastKANLayer(
-                in_dim, out_dim,
-                grid_min=grid_min,
-                grid_max=grid_max,
-                num_grids=num_grids,
-                use_base_update=use_base_update,
-                base_activation=base_activation,
-                spline_weight_init_scale=spline_weight_init_scale,
-            ) for in_dim, out_dim in zip(layers_hidden[:-1], layers_hidden[1:])
-        ])
-
-    def forward(self, x):
-        # Calculating mean and std of images to take them as input to 1D KAN
-        mean_irradiance = torch.torch.mean(x, dim=(2,3))
-        std_irradiance = torch.torch.std(x, dim=(2,3))
-        if self.use_std:
-            x = torch.cat((mean_irradiance, std_irradiance), dim=1)
-        else:
-            x = mean_irradiance
-        for layer in self.layers:
-            x = layer(x)
-        return x
-
-
-class AttentionWithFastKANTransform(nn.Module):
-    
-    def __init__(
-        self,
-        q_dim: int,
-        k_dim: int,
-        v_dim: int,
-        head_dim: int,
-        num_heads: int,
-        gating: bool = True,
-    ):
-        super(AttentionWithFastKANTransform, self).__init__()
-
-        self.num_heads = num_heads
-        total_dim = head_dim * self.num_heads
-        self.gating = gating
-        self.linear_q = FastKANLayer(q_dim, total_dim)
-        self.linear_k = FastKANLayer(k_dim, total_dim)
-        self.linear_v = FastKANLayer(v_dim, total_dim)
-        self.linear_o = FastKANLayer(total_dim, q_dim)
-        self.linear_g = None
-        if self.gating:
-            self.linear_g = FastKANLayer(q_dim, total_dim)
-        # precompute the 1/sqrt(head_dim)
-        self.norm = head_dim**-0.5
-
-    def forward(
-        self,
-        q: torch.Tensor,
-        k: torch.Tensor,
-        v: torch.Tensor,
-        bias: torch.Tensor = None,      # additive attention bias
-    ) -> torch.Tensor:         
-
-        wq = self.linear_q(q).view(*q.shape[:-1], 1, self.num_heads, -1) * self.norm     # *q1hc
-        wk = self.linear_k(k).view(*k.shape[:-2], 1, k.shape[-2], self.num_heads, -1)    # *1khc
-        att = (wq * wk).sum(-1).softmax(-2)     # *qkh
-        del wq, wk
-        if bias is not None:
-            att = att + bias[..., None]
-
-        wv = self.linear_v(v).view(*v.shape[:-2],1, v.shape[-2], self.num_heads, -1)     # *1khc
-        o = (att[..., None] * wv).sum(-3)        # *qhc
-        del att, wv
-
-        o = o.view(*o.shape[:-2], -1)           # *q(hc)
-
-        if self.linear_g is not None:
-            # gating, use raw query input
-            g = self.linear_g(q)
-            o = torch.sigmoid(g) * o
-
-        # merge heads
-        o = self.linear_o(o)
-        return o

flaring/forecasting/models/linear_and_hybrid.py

@@ -1,6 +1,11 @@
 import torch
 import torch.nn as nn
 from torch.nn import HuberLoss
+from pathlib import Path
+import sys
+# Add project root to Python path
+PROJECT_ROOT = Path(__file__).parent.parent.parent.parent.absolute()
+sys.path.insert(0, str(PROJECT_ROOT))
 from flaring.forecasting.models.base_model import BaseModel
 from torchvision.models import resnet18
 

flaring/forecasting/models/vision_transformer_custom.py

@@ -27,7 +27,7 @@ class ViT(pl.LightningModule):
         filtered_kwargs = dict(model_kwargs)
         filtered_kwargs.pop('lr', None)
         self.model = VisionTransformer(**filtered_kwargs)
-        self.adaptive_loss = SXRRegressionDynamicLoss(window_size=500)
+        self.adaptive_loss = SXRRegressionDynamicLoss(window_size=1500)
         self.sxr_norm = sxr_norm
 
     def forward(self, x, return_attention=True):
@@ -284,7 +284,8 @@ class SXRRegressionDynamicLoss:
         }
 
     def calculate_loss(self, preds_squeezed, sxr, sxr_un, preds_squeezed_un):
-        base_loss = F.huber_loss(preds_squeezed, sxr, delta=1.0, reduction='none')
+        #base_loss = F.huber_loss(preds_squeezed, sxr, delta=1.0, reduction='none')
+        base_loss = F.mse_loss(preds_squeezed, sxr, reduction='none')
         weights = self._get_adaptive_weights(sxr_un, preds_squeezed_un, base_loss)
         self._update_tracking(sxr_un, preds_squeezed_un, base_loss)
         weighted_loss = base_loss * weights
@@ -299,13 +300,13 @@ class SXRRegressionDynamicLoss:
             self.quiet_errors, max_multiplier=1.5, min_multiplier=0.5, sensitivity=2.0, sxrclass = 'quiet'
         )
         c_mult = self._get_performance_multiplier(
-            self.c_errors, max_multiplier=2.0, min_multiplier=0.7, sensitivity=2.5, sxrclass = 'c_class'
+            self.c_errors, max_multiplier=5.0, min_multiplier=0.5, sensitivity=2.5, sxrclass = 'c_class'
         )
         m_mult = self._get_performance_multiplier(
-            self.m_errors, max_multiplier=7.0, min_multiplier=0.8, sensitivity=3.0, sxrclass = 'm_class'
+            self.m_errors, max_multiplier=7.0, min_multiplier=0.5, sensitivity=3.0, sxrclass = 'm_class'
         )
         x_mult = self._get_performance_multiplier(
-            self.x_errors, max_multiplier=15.0, min_multiplier=0.9, sensitivity=4.0, sxrclass = 'x_class'
+            self.x_errors, max_multiplier=15.0, min_multiplier=0.5, sensitivity=4.0, sxrclass = 'x_class'
         )
 
         quiet_weight = self.base_weights['quiet'] * quiet_mult

flaring/forecasting/training/config.yaml

@@ -1,11 +1,11 @@
 
 #Base directories - change these to switch datasets
-base_data_dir: "/mnt/data/ML-Ready-mixed/ML-Ready-mixed"  # Change this line for different datasets
-base_checkpoint_dir: "/mnt/data/ML-Ready-mixed/ML-Ready-mixed"    # Change this line for different datasets
+base_data_dir: "/mnt/data/ML-READY"  # Change this line for different datasets
+base_checkpoint_dir: "/mnt/data/ML-READY"    # Change this line for different datasets
 wavelengths: [94, 131, 171, 193, 211, 304]  # AIA wavelengths in Angstroms
 # Model configuration
 selected_model: "ViT Patch"  # Options: "cnn", "vit",
-batch_size:    4
+batch_size:    80
 epochs:        500
 oversample: false
 balance_strategy: "upsample_minority"
@@ -21,9 +21,9 @@ vit_custom:
     embed_dim: 512
     num_channels: 6
     num_classes: 1
-    patch_size: 8
-    num_patches: 4096
-    hidden_dim: 512
+    patch_size: 16
+    num_patches: 1024
+    hidden_dim: 1024
     num_heads: 8
     num_layers: 6
     dropout: 0.1
@@ -48,5 +48,5 @@ wandb:
     - aia
     - sxr
     - regression
-  wb_name: vit-8-patch
+  wb_name: vit-16-MSE-deeper
   notes: Regression from AIA images (6 channels) to GOES SXR flux

flaring/forecasting/training/train.py

@@ -12,6 +12,12 @@ from pytorch_lightning import Trainer
 from pytorch_lightning.loggers import WandbLogger
 from pytorch_lightning.callbacks import ModelCheckpoint
 from torch.nn import MSELoss
+from pathlib import Path
+import sys
+# Add project root to Python path
+PROJECT_ROOT = Path(__file__).parent.parent.parent.parent.absolute()
+sys.path.insert(0, str(PROJECT_ROOT))
+
 from flaring.forecasting.data_loaders.SDOAIA_dataloader import AIA_GOESDataModule
 from flaring.forecasting.models.vision_transformer_custom import ViT
 from flaring.forecasting.models.linear_and_hybrid import LinearIrradianceModel, HybridIrradianceModel

flaring/vision_transformers/__init__.py

@@ -0,0 +1 @@
+# Vision Transformers module