integrated both the versions

.gitignore

@@ -11,5 +11,6 @@ static/
 plots/
 uploads/
 downloads/
+predictions/
 
 

app.py

@@ -4,6 +4,8 @@ import os
 import json
 import traceback
 from pathlib import Path
+import xarray as xr
+import numpy as np
 
 from datetime import datetime, timedelta
 from werkzeug.utils import secure_filename
@@ -16,6 +18,15 @@ from interactive_plot_generator import InteractiveIndiaMapPlotter
 from cams_downloader import CAMSDownloader
 from constants import ALLOWED_EXTENSIONS, MAX_FILE_SIZE, COLOR_THEMES
 
+# Aurora pipeline imports - with error handling for optional dependency
+try:
+    from aurora import Batch, Metadata, AuroraAirPollution, rollout
+    from aurora_pipeline import AuroraPipeline
+    AURORA_AVAILABLE = True
+except ImportError as e:
+    print(f"⚠️  Aurora model not available: {e}")
+    AURORA_AVAILABLE = False
+
 app = Flask(__name__)
 app.secret_key = 'your-secret-key-change-this-in-production'  # Change this!
 app.config['DEBUG'] = False  # Explicitly disable debug mode
@@ -33,8 +44,25 @@ downloader = CAMSDownloader()
 plotter = IndiaMapPlotter()
 interactive_plotter = InteractiveIndiaMapPlotter()
 
+# Initialize Aurora pipeline if available
+if AURORA_AVAILABLE:
+    # Check if we're in development/local mode
+    import socket
+    hostname = socket.gethostname()
+    is_local = any(local_indicator in hostname.lower() 
+                   for local_indicator in ['local', 'macbook', 'laptop', 'desktop', 'dev'])
+    
+    # Force CPU mode for local development to avoid GPU requirements
+    cpu_only = is_local or os.getenv('AURORA_CPU_ONLY', 'false').lower() == 'true'
+    
+    aurora_pipeline = AuroraPipeline(cpu_only=cpu_only)
+    print(f"🔮 Aurora pipeline initialized ({'CPU-only' if cpu_only else 'GPU-enabled'} mode)")
+else:
+    aurora_pipeline = None
+    print("⚠️  Aurora pipeline not available - missing dependencies")
+
 # Ensure directories exist
-for directory in ['uploads', 'downloads', 'plots', 'templates', 'static']:
+for directory in ['uploads', 'downloads', 'plots', 'templates', 'static', 'predictions']:
     Path(directory).mkdir(exist_ok=True)
 
 
@@ -77,7 +105,8 @@ def index():
         downloaded_files=downloaded_files,
         cds_ready=downloader.is_client_ready(),
         current_date=current_date,
-        recent_files=recent_files
+        recent_files=recent_files,
+        aurora_available=AURORA_AVAILABLE
     )
 
 
@@ -239,11 +268,13 @@ def get_pressure_levels(filename, variable):
         print(f"File path: {file_path}")
 
         processor = NetCDFProcessor(str(file_path))
-        processor.load_dataset()
-        processor.detect_variables()
-        
-        pressure_levels = processor.get_available_pressure_levels(variable)
-        processor.close()
+        try:
+            processor.load_dataset()
+            processor.detect_variables()
+            
+            pressure_levels = processor.get_available_pressure_levels(variable)
+        finally:
+            processor.close()
         
         return jsonify({
             'success': True,
@@ -271,11 +302,13 @@ def get_available_times(filename, variable):
             file_path = Path(app.config['UPLOAD_FOLDER']) / filename
         
         processor = NetCDFProcessor(str(file_path))
-        processor.load_dataset()
-        processor.detect_variables()
-        
-        available_times = processor.get_available_times(variable)
-        processor.close()
+        try:
+            processor.load_dataset()
+            processor.detect_variables()
+            
+            available_times = processor.get_available_times(variable)
+        finally:
+            processor.close()
         
         # Format times for display
         formatted_times = []
@@ -324,34 +357,37 @@ def visualize():
         
         # Process the data
         processor = NetCDFProcessor(str(file_path))
-        processor.load_dataset()
-        processor.detect_variables()
-        
-        # Convert pressure level to float if provided
-        pressure_level_val = None
-        if pressure_level and pressure_level != 'None':
-            try:
-                pressure_level_val = float(pressure_level)
-            except ValueError:
-                pressure_level_val = None
-
-        time_index_val = request.form.get('time_index')
-        # Extract data
-        data_values, metadata = processor.extract_data(
-            variable, 
-            time_index = int(time_index_val) if time_index_val and time_index_val != 'None' else 0,
-            pressure_level=pressure_level_val
-        )
-        
-        # Generate plot
-        plot_path = plotter.create_india_map(
-            data_values, 
-            metadata, 
-            color_theme=color_theme,
-            save_plot=True
-        )
-        
-        processor.close()
+        try:
+            processor.load_dataset()
+            processor.detect_variables()
+            
+            # Convert pressure level to float if provided
+            pressure_level_val = None
+            if pressure_level and pressure_level != 'None':
+                try:
+                    pressure_level_val = float(pressure_level)
+                except ValueError:
+                    pressure_level_val = None
+
+            time_index_val = request.form.get('time_index')
+            # Extract data
+            data_values, metadata = processor.extract_data(
+                variable, 
+                time_index = int(time_index_val) if time_index_val and time_index_val != 'None' else 0,
+                pressure_level=pressure_level_val
+            )
+            
+            # Generate plot
+            plot_path = plotter.create_india_map(
+                data_values, 
+                metadata, 
+                color_theme=color_theme,
+                save_plot=True
+            )
+            
+        finally:
+            # Always close the processor
+            processor.close()
         
         if plot_path:
             plot_filename = Path(plot_path).name
@@ -412,34 +448,37 @@ def visualize_interactive():
         
         # Process the data
         processor = NetCDFProcessor(str(file_path))
-        processor.load_dataset()
-        processor.detect_variables()
-        
-        # Convert pressure level to float if provided
-        pressure_level_val = None
-        if pressure_level and pressure_level != 'None':
-            try:
-                pressure_level_val = float(pressure_level)
-            except ValueError:
-                pressure_level_val = None
-
-        time_index_val = request.form.get('time_index')
-        # Extract data
-        data_values, metadata = processor.extract_data(
-            variable, 
-            time_index = int(time_index_val) if time_index_val and time_index_val != 'None' else 0,
-            pressure_level=pressure_level_val
-        )
-        
-        # Generate interactive plot
-        result = interactive_plotter.create_india_map(
-            data_values, 
-            metadata, 
-            color_theme=color_theme,
-            save_plot=True
-        )
+        try:
+            processor.load_dataset()
+            processor.detect_variables()
+            
+            # Convert pressure level to float if provided
+            pressure_level_val = None
+            if pressure_level and pressure_level != 'None':
+                try:
+                    pressure_level_val = float(pressure_level)
+                except ValueError:
+                    pressure_level_val = None
+
+            time_index_val = request.form.get('time_index')
+            # Extract data
+            data_values, metadata = processor.extract_data(
+                variable, 
+                time_index = int(time_index_val) if time_index_val and time_index_val != 'None' else 0,
+                pressure_level=pressure_level_val
+            )
+            
+            # Generate interactive plot
+            result = interactive_plotter.create_india_map(
+                data_values, 
+                metadata, 
+                color_theme=color_theme,
+                save_plot=True
+            )
         
-        processor.close()
+        finally:
+            # Always close the processor
+            processor.close()
         
         if result and result.get('html_content'):
             # Prepare metadata for display
@@ -765,10 +804,228 @@ def health_check():
     return jsonify({
         'status': 'healthy',
         'timestamp': datetime.now().isoformat(),
-        'cds_ready': downloader.is_client_ready()
+        'cds_ready': downloader.is_client_ready(),
+        'aurora_available': AURORA_AVAILABLE
     })
 
 
+@app.route('/api/aurora_status')
+def aurora_status():
+    """API endpoint to check Aurora readiness and get system info"""
+    status = {
+        'available': AURORA_AVAILABLE,
+        'cpu_only': False,
+        'estimated_time': {
+            'cpu': {'1_step': 5, '2_steps': 10},
+            'gpu': {'4_steps': 3, '6_steps': 4, '10_steps': 6}
+        }
+    }
+    
+    if AURORA_AVAILABLE and aurora_pipeline:
+        status['cpu_only'] = getattr(aurora_pipeline, 'cpu_only', False)
+    
+    return jsonify(status)
+
+
+# Aurora ML Prediction Routes
+@app.route('/aurora_predict', methods=['GET', 'POST'])
+def aurora_predict():
+    """Aurora prediction form and handler"""
+    if not AURORA_AVAILABLE:
+        flash('Aurora model is not available. Please install required dependencies.', 'error')
+        return redirect(url_for('index'))
+    
+    if request.method == 'GET':
+        current_date = datetime.now().strftime('%Y-%m-%d')
+        return render_template('aurora_predict.html', current_date=current_date)
+    
+    # POST: Run the pipeline
+    date_str = request.form.get('date')
+    steps = int(request.form.get('steps', 2))  # Default to 2 steps for CPU-friendly execution
+    
+    # Limit steps for local/CPU execution
+    if hasattr(aurora_pipeline, 'cpu_only') and aurora_pipeline.cpu_only:
+        steps = min(steps, 2)  # Max 2 steps for CPU
+        if steps != int(request.form.get('steps', 2)):
+            flash(f'Steps reduced to {steps} for CPU mode optimization', 'info')
+    
+    if not date_str:
+        flash('Please select a valid date.', 'error')
+        return redirect(url_for('aurora_predict'))
+    
+    cams_date = date_str
+    cams_time = "12:00"  # Always use 12:00 UTC for Aurora
+
+    try:
+        print(f"🚀 Starting Aurora prediction pipeline for {cams_date}")
+        
+        # 1. Download CAMS data for the selected date
+        print("📥 Step 1/5: Downloading CAMS atmospheric data...")
+        try:
+            zip_path = downloader.download_cams_data(cams_date)
+        except Exception as e:
+            error_msg = f"Failed to download CAMS data: {str(e)}"
+            if "error response" in str(e).lower():
+                error_msg += " The CAMS API may have returned an error. Please try a different date or check your CDS API credentials."
+            elif "zip" in str(e).lower():
+                error_msg += " The downloaded file is corrupted. Please try again."
+            flash(error_msg, 'error')
+            print(f"❌ Download error: {traceback.format_exc()}")
+            return redirect(url_for('aurora_predict'))
+        
+        try:
+            extracted_files = downloader.extract_cams_files(zip_path)
+            print("✅ CAMS data downloaded and extracted")
+        except Exception as e:
+            error_msg = f"Failed to extract CAMS data: {str(e)}"
+            if "not a zip file" in str(e).lower():
+                error_msg += " The downloaded file appears to be corrupted or is an error response from the CAMS API."
+            elif "html" in str(e).lower() or "error" in str(e).lower():
+                error_msg += " The CAMS API returned an error page instead of data."
+            flash(error_msg, 'error')
+            print(f"❌ Extraction error: {traceback.format_exc()}")
+            return redirect(url_for('aurora_predict'))
+        
+        # 2. Run Aurora pipeline (batch creation, model, prediction, save NetCDF)
+        print("🔮 Step 2/5: Initializing Aurora ML pipeline...")
+        output_nc = f"predictions_{cams_date}_{cams_time.replace(':','')}.nc"
+        
+        print(f"🧠 Step 3/5: Loading Aurora model (this may take a few minutes)...")
+        print(f"⚡ Step 4/5: Running {steps} prediction steps...")
+        
+        try:
+            predictions = aurora_pipeline.run_pipeline(
+                date_str=cams_date,
+                Batch=Batch,
+                Metadata=Metadata,
+                AuroraAirPollution=AuroraAirPollution,
+                rollout=rollout,
+                steps=steps,
+                output_path=Path('predictions') / output_nc
+            )
+        except Exception as e:
+            error_msg = f"Aurora model execution failed: {str(e)}"
+            if "map_location" in str(e):
+                error_msg += " This appears to be a compatibility issue with the Aurora model version."
+            elif "checkpoint" in str(e).lower():
+                error_msg += " Failed to load the Aurora model. Please check if the model files are properly installed."
+            elif "memory" in str(e).lower() or "cuda" in str(e).lower():
+                error_msg += " Insufficient memory or GPU issues. Try reducing the number of prediction steps."
+            flash(error_msg, 'error')
+            print(f"❌ Aurora model error: {traceback.format_exc()}")
+            return redirect(url_for('aurora_predict'))
+        
+        print("💾 Step 5/5: Saving results and preparing visualization...")
+        print(f"✅ Aurora predictions completed for {cams_date} 12:00 UTC")
+        
+        flash(f'🔮 Aurora predictions generated successfully for {cams_date} ({steps} steps)', 'success')
+        return redirect(url_for('visualize_prediction', filename=output_nc))
+    
+    except Exception as e:
+        # Catch-all for any other unexpected errors
+        error_msg = f'Unexpected error in Aurora pipeline: {str(e)}'
+        flash(error_msg, 'error')
+        print(f"❌ Unexpected Aurora pipeline error: {traceback.format_exc()}")
+        return redirect(url_for('aurora_predict'))
+
+
+@app.route('/visualize_prediction/<filename>', methods=['GET', 'POST'])
+def visualize_prediction(filename):
+    """Visualize or download Aurora prediction output with variable and step selection"""
+    file_path = Path('predictions') / filename
+    if not file_path.exists():
+        flash('Prediction file not found', 'error')
+        return redirect(url_for('index'))
+
+    ds = xr.open_dataset(file_path)
+    variables = list(ds.data_vars.keys())
+    steps = ds['step'].values if 'step' in ds else np.arange(ds[variables[0]].shape[0])
+
+    # Handle form submission
+    if request.method == 'POST':
+        var_name = request.form.get('variable')
+        step = int(request.form.get('step', 0))
+        color_theme = request.form.get('color_theme', 'viridis')
+    else:
+        var_name = variables[0]
+        step = 0
+        color_theme = 'viridis'
+
+    # Prepare data for plotting
+    data = ds[var_name].values
+    if data.ndim == 3:
+        data_to_plot = data[step]
+    elif data.ndim == 2:
+        data_to_plot = data
+    else:
+        flash('Prediction data shape not supported for plotting', 'error')
+        return redirect(url_for('index'))
+
+    # Extract lat/lon
+    lats = ds['lat'].values if 'lat' in ds else ds['latitude'].values
+    lons = ds['lon'].values if 'lon' in ds else ds['longitude'].values
+
+    # Prepare metadata for plot
+    from constants import NETCDF_VARIABLES
+    var_info = NETCDF_VARIABLES.get(var_name, {})
+    display_name = var_info.get('name', var_name)
+    units = ds[var_name].attrs.get('units', var_info.get('units', ''))
+    # Use user-selected color theme, fallback to variable default, then viridis
+    if 'color_theme' not in locals():
+        color_theme = var_info.get('cmap', 'viridis')
+
+    metadata = {
+        'variable_name': var_name,
+        'display_name': display_name,
+        'units': units,
+        'lats': lats,
+        'lons': lons,
+        'pressure_level': None,
+        'timestamp_str': str(steps[step]) if len(steps) > step else '',
+    }
+
+    # Generate plot
+    plot_path = plotter.create_india_map(
+        data_to_plot,
+        metadata,
+        color_theme=color_theme,
+        save_plot=True,
+        custom_title=f"Aurora Prediction: {display_name} (step {step})"
+    )
+    
+    if plot_path:
+        plot_filename = Path(plot_path).name
+
+        # Provide download link for NetCDF
+        download_url = url_for('download_prediction_netcdf', filename=filename)
+
+        return render_template(
+            'aurora_prediction_plot.html',
+            plot_filename=plot_filename,
+            var_name=var_name,
+            step=step,
+            variables=variables,
+            steps=range(len(steps)),
+            filename=filename,
+            download_url=download_url,
+            color_themes=COLOR_THEMES,
+            current_color_theme=color_theme
+        )
+    else:
+        flash('Error generating prediction plot', 'error')
+        return redirect(url_for('index'))
+
+
+@app.route('/download_prediction_netcdf/<filename>')
+def download_prediction_netcdf(filename):
+    """Download the Aurora prediction NetCDF file"""
+    file_path = Path('predictions') / filename
+    if not file_path.exists():
+        flash('Prediction file not found', 'error')
+        return redirect(url_for('index'))
+    return send_file(str(file_path), as_attachment=True, download_name=filename)
+
+
 @app.errorhandler(413)
 def too_large(e):
     """Handle file too large error"""
@@ -802,4 +1059,4 @@ if __name__ == '__main__':
     print("🔧 CDS API Ready:", downloader.is_client_ready())
     
     # Run the Flask app
-    app.run(debug=debug_mode, host='0.0.0.0', port=port)
+    app.run(debug=True, host='0.0.0.0', port=port)

aurora_pipeline.py

@@ -0,0 +1,471 @@
+# aurora_pipeline.py
+# End-to-end pipeline for CAMS data → Aurora model → predictions → NetCDF
+import subprocess
+import os
+
+def get_freest_cuda_device_id():
+    """Get the freest CUDA device ID if available, otherwise return None for CPU mode"""
+    try:
+        # Check if nvidia-smi exists first
+        result = subprocess.run(
+            ['nvidia-smi', '--query-gpu=memory.free', '--format=csv,nounits,noheader'],
+            stdout=subprocess.PIPE, stderr=subprocess.PIPE, encoding='utf-8'
+        )
+        if result.returncode == 0:
+            memory_free = [int(x) for x in result.stdout.strip().split('\n')]
+            device_id = memory_free.index(max(memory_free))
+            print(f"🎮 GPU available - using device {device_id}")
+            return str(device_id)
+        else:
+            print("⚠️  nvidia-smi returned error, using CPU mode")
+            return None
+    except (FileNotFoundError, subprocess.SubprocessError, Exception) as e:
+        print(f"💻 No GPU detected, using CPU mode: {e}")
+        return None
+
+# Set CUDA_VISIBLE_DEVICES only if GPU is available
+gpu_device = get_freest_cuda_device_id()
+if gpu_device is not None:
+    os.environ["CUDA_VISIBLE_DEVICES"] = gpu_device
+else:
+    # Force CPU mode
+    os.environ["CUDA_VISIBLE_DEVICES"] = ""
+    print("🔧 Configured for CPU-only execution")
+
+
+import torch
+import xarray as xr
+import pickle
+from pathlib import Path
+import numpy as np
+import zipfile
+import cdsapi
+from huggingface_hub import hf_hub_download
+import matplotlib.pyplot as plt
+import cartopy.crs as ccrs
+import cartopy.feature as cfeature
+from datetime import datetime, timedelta
+from aurora import Batch, Metadata, AuroraAirPollution, rollout
+
+
+class AuroraPipeline:
+    def __init__(self, 
+             extracted_dir="downloads/extracted", 
+             static_path="static_vars.pkl", 
+             model_ckpt="aurora-0.4-air-pollution.ckpt", 
+             model_repo="microsoft/aurora",
+             device=None,
+             cpu_only=False):
+
+        # Device selection with CPU fallback
+        if cpu_only or device == "cpu":
+            self.device = "cpu"
+            print("💻 Aurora configured for CPU-only execution")
+        elif device is None:
+            # Auto-detect: prefer CPU for local testing, GPU for production
+            if torch.cuda.is_available() and not cpu_only:
+                self.device = "cuda:0"
+                print(f"🎮 Aurora using GPU: {self.device}")
+            else:
+                self.device = "cpu"
+                print("💻 Aurora using CPU (GPU not available or CPU forced)")
+        else:
+            self.device = device
+            print(f"🔧 Aurora using specified device: {self.device}")
+
+        self.extracted_dir = Path(extracted_dir)
+        self.static_path = Path(static_path)
+        self.model_ckpt = model_ckpt
+        self.model_repo = model_repo
+        self.cpu_only = cpu_only or (self.device == "cpu")
+        
+        # Load static variables with error handling
+        try:
+            self.static_vars = self._load_static_vars()
+        except Exception as e:
+            print(f"⚠️  Warning: Could not load static variables: {e}")
+            self.static_vars = {}
+        
+        self.model = None
+
+    def _load_static_vars(self):
+        """Load static variables from Hugging Face Hub"""
+        static_path = hf_hub_download(
+            repo_id="microsoft/aurora",
+            filename="aurora-0.4-air-pollution-static.pickle",
+        )
+        if not Path(static_path).exists():
+            raise FileNotFoundError(f"Static variables file not found: {static_path}")
+        with open(static_path, "rb") as f:
+            static_vars = pickle.load(f)
+        return static_vars
+
+    def create_batch(self, date_str, Batch, Metadata):
+        """Create a batch for Aurora model from CAMS data"""
+        surface_path = self.extracted_dir / f"{date_str}-cams-surface.nc"
+        atmos_path = self.extracted_dir / f"{date_str}-cams-atmospheric.nc"
+        if not surface_path.exists() or not atmos_path.exists():
+            raise FileNotFoundError(f"Missing CAMS files for {date_str} in {self.extracted_dir}")
+
+        surf_vars_ds = xr.open_dataset(surface_path, engine="netcdf4", decode_timedelta=True)
+        atmos_vars_ds = xr.open_dataset(atmos_path, engine="netcdf4", decode_timedelta=True)
+
+        # Select zero-hour forecast
+        surf_vars_ds = surf_vars_ds.isel(forecast_period=0)
+        atmos_vars_ds = atmos_vars_ds.isel(forecast_period=0)
+
+        batch = Batch(
+            surf_vars={
+                "2t": torch.from_numpy(surf_vars_ds["t2m"].values[None]),
+                "10u": torch.from_numpy(surf_vars_ds["u10"].values[None]),
+                "10v": torch.from_numpy(surf_vars_ds["v10"].values[None]),
+                "msl": torch.from_numpy(surf_vars_ds["msl"].values[None]),
+                "pm1": torch.from_numpy(surf_vars_ds["pm1"].values[None]),
+                "pm2p5": torch.from_numpy(surf_vars_ds["pm2p5"].values[None]),
+                "pm10": torch.from_numpy(surf_vars_ds["pm10"].values[None]),
+                "tcco": torch.from_numpy(surf_vars_ds["tcco"].values[None]),
+                "tc_no": torch.from_numpy(surf_vars_ds["tc_no"].values[None]),
+                "tcno2": torch.from_numpy(surf_vars_ds["tcno2"].values[None]),
+                "gtco3": torch.from_numpy(surf_vars_ds["gtco3"].values[None]),
+                "tcso2": torch.from_numpy(surf_vars_ds["tcso2"].values[None]),
+            },
+            static_vars={k: torch.from_numpy(v) for k, v in self.static_vars.items()},
+            atmos_vars={
+                "t": torch.from_numpy(atmos_vars_ds["t"].values[None]),
+                "u": torch.from_numpy(atmos_vars_ds["u"].values[None]),
+                "v": torch.from_numpy(atmos_vars_ds["v"].values[None]),
+                "q": torch.from_numpy(atmos_vars_ds["q"].values[None]),
+                "z": torch.from_numpy(atmos_vars_ds["z"].values[None]),
+                "co": torch.from_numpy(atmos_vars_ds["co"].values[None]),
+                "no": torch.from_numpy(atmos_vars_ds["no"].values[None]),
+                "no2": torch.from_numpy(atmos_vars_ds["no2"].values[None]),
+                "go3": torch.from_numpy(atmos_vars_ds["go3"].values[None]),
+                "so2": torch.from_numpy(atmos_vars_ds["so2"].values[None]),
+            },
+            metadata=Metadata(
+                lat=torch.from_numpy(atmos_vars_ds.latitude.values),
+                lon=torch.from_numpy(atmos_vars_ds.longitude.values),
+                time=(atmos_vars_ds.valid_time.values.astype("datetime64[s]").tolist()[-1],),
+                atmos_levels=tuple(int(level) for level in atmos_vars_ds.pressure_level.values),
+            ),
+        )
+        return batch
+    def load_model(self, AuroraAirPollution):
+        """Load Aurora model with CPU/GPU optimization"""
+        import gc
+        
+        print(f"🔄 Loading Aurora model on {self.device}")
+        
+        # Memory check for GPU
+        if self.device != "cpu" and torch.cuda.is_available():
+            print(f"📊 GPU Memory BEFORE loading model:")
+            print(f"   Allocated: {torch.cuda.memory_allocated(0) / 1024**3:.2f} GB")
+            print(f"   Reserved:  {torch.cuda.memory_reserved(0) / 1024**3:.2f} GB")
+            print(f"   Free:      {(torch.cuda.get_device_properties(0).total_memory - torch.cuda.memory_reserved(0)) / 1024**3:.2f} GB")
+        
+        # Clear cache
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+        gc.collect()
+        
+        # Initialize model with CPU-friendly settings
+        if self.cpu_only:
+            print("💻 Initializing model for CPU execution...")
+            # Set CPU-friendly torch settings
+            torch.set_num_threads(2)  # Limit CPU threads for local testing
+            model = AuroraAirPollution()
+        else:
+            model = AuroraAirPollution()
+        
+        # Load checkpoint with device mapping
+        try:
+            if self.cpu_only:
+                print("📁 Loading checkpoint for CPU execution...")
+                # For CPU mode, we may need to handle device mapping differently
+                model.load_checkpoint(self.model_repo, self.model_ckpt)
+            else:
+                print("📁 Loading checkpoint for GPU execution...")
+                model.load_checkpoint(self.model_repo, self.model_ckpt)
+        except Exception as e:
+            print(f"⚠️  Checkpoint loading failed: {e}")
+            print("🔄 Trying alternative loading method...")
+            try:
+                # Alternative: try loading without any special parameters
+                model.load_checkpoint(self.model_repo, self.model_ckpt)
+                print("✅ Checkpoint loaded successfully with fallback method")
+            except Exception as e2:
+                print(f"❌ All loading methods failed: {e2}")
+                # Set device to CPU as last resort
+                self.device = "cpu"
+                self.cpu_only = True
+                raise RuntimeError(f"Failed to load Aurora model: {e2}")
+        
+        model.eval()
+        
+        # Move to device
+        model = model.to(self.device)
+        
+        # Memory check after loading
+        if self.device != "cpu" and torch.cuda.is_available():
+            print(f"📊 GPU Memory AFTER model load:")
+            print(f"   Allocated: {torch.cuda.memory_allocated(0) / 1024**3:.2f} GB")
+            print(f"   Reserved:  {torch.cuda.memory_reserved(0) / 1024**3:.2f} GB")
+        
+        self.model = model
+        print(f"✅ Aurora model loaded on {self.device}")
+        
+        if self.cpu_only:
+            print("⚠️  WARNING: CPU mode will be slower than GPU. Consider using fewer steps for faster inference.")
+        
+        return model
+
+    def predict(self, batch, rollout, steps=4):
+        """Run model prediction with CPU/GPU optimization"""
+        if self.model is None:
+            raise RuntimeError("Model not loaded. Call load_model() first.")
+        
+        # Limit steps for CPU to avoid memory issues
+        if self.cpu_only and steps > 2:
+            print(f"⚠️  CPU mode: reducing steps from {steps} to 2 for memory efficiency")
+            steps = 2
+        
+        print(f"🔄 Running {steps} prediction steps on {self.device}...")
+        
+        # Move batch to device
+        batch = batch.to(self.device)
+        
+        # CPU-friendly inference settings
+        if self.cpu_only:
+            torch.set_grad_enabled(False)  # Disable gradients for inference
+        
+        with torch.inference_mode():
+            predictions = []
+            for step in range(steps):
+                print(f"   Step {step + 1}/{steps}...")
+                if step == 0:
+                    # First prediction from initial batch
+                    pred_generator = rollout(self.model, batch, steps=1)
+                    pred = next(pred_generator)
+                else:
+                    # Subsequent predictions from previous output
+                    pred_generator = rollout(self.model, pred, steps=1)
+                    pred = next(pred_generator)
+                
+                # Move to CPU immediately to save memory
+                predictions.append(pred.to("cpu"))
+                
+                # Clear GPU cache after each step if using GPU
+                if not self.cpu_only and torch.cuda.is_available():
+                    torch.cuda.empty_cache()
+        
+        print(f"✅ Completed {len(predictions)} prediction steps")
+        return predictions
+
+    def save_predictions_to_netcdf(self, predictions, output_path):
+        """Save all prediction steps to a single NetCDF file compatible with visualization pipeline"""
+        output_path = Path(output_path)
+        output_path.parent.mkdir(parents=True, exist_ok=True)
+
+        print(f"💾 Saving {len(predictions)} prediction steps to {output_path}")
+
+        try:
+            # Try the new single-file method
+            return self._save_predictions_single_file(predictions, output_path)
+        except Exception as e:
+            print(f"⚠️  Single file method failed: {e}")
+            print(f"🔄 Falling back to original method...")
+            return self._save_predictions_original_method(predictions, output_path)
+
+    def _save_predictions_single_file(self, predictions, output_path):
+        """Save all prediction steps to a single NetCDF file (new method)"""
+        # Get metadata from first prediction
+        first_pred = predictions[0]
+        metadata = first_pred.metadata
+
+        # Extract coordinates
+        lats = metadata.lat.cpu().numpy() if hasattr(metadata.lat, 'cpu') else metadata.lat.numpy()
+        lons = metadata.lon.cpu().numpy() if hasattr(metadata.lon, 'cpu') else metadata.lon.numpy()
+
+        # Create step coordinate
+        steps = np.arange(len(predictions))
+
+        # Prepare data variables
+        data_vars = {}
+        coords = {
+            'step': ('step', steps),
+            'lat': ('lat', lats),
+            'lon': ('lon', lons)
+        }
+
+        # Add surface variables
+        surf_var_names = list(first_pred.surf_vars.keys())
+        for var in surf_var_names:
+            # Stack predictions along step dimension
+            var_data_list = []
+            for pred in predictions:
+                var_tensor = pred.surf_vars[var]
+                # Move to CPU and convert to numpy
+                var_data = var_tensor.cpu().numpy() if hasattr(var_tensor, 'cpu') else var_tensor.numpy()
+
+                # Robust dimension handling: squeeze all singleton dimensions and keep only last 2 (lat, lon)
+                var_data = np.squeeze(var_data)  # Remove all singleton dimensions
+
+                # Ensure we have exactly 2 dimensions (lat, lon) for surface variables
+                if var_data.ndim > 2:
+                    # Take the last 2 dimensions as lat, lon
+                    var_data = var_data[..., :, :]
+                    # If still more than 2D, take the first slice of extra dimensions
+                    while var_data.ndim > 2:
+                        var_data = var_data[0]
+                elif var_data.ndim < 2:
+                    raise ValueError(f"Surface variable {var} has insufficient dimensions: {var_data.shape}")
+
+                var_data_list.append(var_data)
+
+            # Stack along step dimension: (steps, lat, lon)
+            arr = np.stack(var_data_list, axis=0)
+            data_vars[var] = (['step', 'lat', 'lon'], arr)
+
+        # Add atmospheric variables if present
+        if hasattr(first_pred, 'atmos_vars') and first_pred.atmos_vars:
+            atmos_levels = list(metadata.atmos_levels) if hasattr(metadata, 'atmos_levels') else None
+            if atmos_levels:
+                coords['pressure_level'] = ('pressure_level', atmos_levels)
+
+                atmos_var_names = list(first_pred.atmos_vars.keys())
+                for var in atmos_var_names:
+                    var_data_list = []
+                    for pred in predictions:
+                        var_tensor = pred.atmos_vars[var]
+                        # Move to CPU and convert to numpy
+                        var_data = var_tensor.cpu().numpy() if hasattr(var_tensor, 'cpu') else var_tensor.numpy()
+
+                        # Robust dimension handling: squeeze singleton dimensions but keep 3D structure
+                        var_data = np.squeeze(var_data)  # Remove singleton dimensions
+
+                        # Ensure we have exactly 3 dimensions (levels, lat, lon) for atmospheric variables
+                        if var_data.ndim > 3:
+                            # Take the last 3 dimensions as levels, lat, lon
+                            var_data = var_data[..., :, :, :]
+                            # If still more than 3D, take the first slice of extra dimensions
+                            while var_data.ndim > 3:
+                                var_data = var_data[0]
+                        elif var_data.ndim < 3:
+                            raise ValueError(f"Atmospheric variable {var} has insufficient dimensions: {var_data.shape}")
+
+                        var_data_list.append(var_data)
+
+                    # Stack along step dimension: (steps, levels, lat, lon)
+                    arr = np.stack(var_data_list, axis=0)
+                    data_vars[f"{var}_atmos"] = (['step', 'pressure_level', 'lat', 'lon'], arr)
+
+        # Create dataset
+        ds = xr.Dataset(data_vars, coords=coords)
+
+        # Add global attributes
+        ds.attrs.update({
+            'title': 'Aurora Air Pollution Model Predictions',
+            'source': 'Aurora model by Microsoft Research',
+            'creation_date': datetime.now().isoformat(),
+            'forecast_steps': len(predictions),
+            'spatial_resolution': f"{abs(lons[1] - lons[0]):.3f} degrees",
+            'conventions': 'CF-1.8'
+        })
+
+        # Add variable attributes for better visualization
+        var_attrs = {
+            '2t': {'long_name': '2 metre temperature', 'units': 'K'},
+            '10u': {'long_name': '10 metre U wind component', 'units': 'm s-1'},
+            '10v': {'long_name': '10 metre V wind component', 'units': 'm s-1'},
+            'msl': {'long_name': 'Mean sea level pressure', 'units': 'Pa'},
+            'pm1': {'long_name': 'Particulate matter d < 1 um', 'units': 'kg m-3'},
+            'pm2p5': {'long_name': 'Particulate matter d < 2.5 um', 'units': 'kg m-3'},
+            'pm10': {'long_name': 'Particulate matter d < 10 um', 'units': 'kg m-3'},
+            'tcco': {'long_name': 'Total column carbon monoxide', 'units': 'kg m-2'},
+            'tc_no': {'long_name': 'Total column nitrogen monoxide', 'units': 'kg m-2'},
+            'tcno2': {'long_name': 'Total column nitrogen dioxide', 'units': 'kg m-2'},
+            'gtco3': {'long_name': 'Total column ozone', 'units': 'kg m-2'},
+            'tcso2': {'long_name': 'Total column sulphur dioxide', 'units': 'kg m-2'}
+        }
+
+        for var_name, attrs in var_attrs.items():
+            if var_name in ds.data_vars:
+                ds[var_name].attrs.update(attrs)
+
+        # Save to NetCDF
+        ds.to_netcdf(output_path, format='NETCDF4')
+        print(f"✅ Predictions saved to {output_path}")
+        print(f"   Variables: {list(ds.data_vars.keys())}")
+        print(f"   Steps: {len(steps)}")
+        print(f"   Spatial grid: {len(lats)}x{len(lons)}")
+
+        return output_path
+
+    def _save_predictions_original_method(self, predictions, output_path):
+        """Fallback: Save predictions using the original method (separate files per step)"""
+        output_dir = Path(output_path)
+        output_dir.mkdir(exist_ok=True)
+
+        for step, pred in enumerate(predictions):
+            # Create xarray dataset for surface variables
+            surf_data = {}
+            for var_name, var_data in pred.surf_vars.items():
+                surf_data[var_name] = (
+                    ["time", "batch", "lat", "lon"],
+                    var_data.cpu().numpy() if hasattr(var_data, 'cpu') else var_data.numpy()
+                )
+
+            # Create xarray dataset for atmospheric variables
+            atmos_data = {}
+            for var_name, var_data in pred.atmos_vars.items():
+                atmos_data[var_name] = (
+                    ["time", "batch", "level", "lat", "lon"],
+                    var_data.cpu().numpy() if hasattr(var_data, 'cpu') else var_data.numpy()
+                )
+
+            # Create surface dataset
+            surf_ds = xr.Dataset(
+                surf_data,
+                coords={
+                    "time": [pred.metadata.time[0]],
+                    "batch": [0],
+                    "lat": pred.metadata.lat.cpu().numpy() if hasattr(pred.metadata.lat, 'cpu') else pred.metadata.lat.numpy(),
+                    "lon": pred.metadata.lon.cpu().numpy() if hasattr(pred.metadata.lon, 'cpu') else pred.metadata.lon.numpy(),
+                }
+            )
+
+            # Create atmospheric dataset
+            atmos_ds = xr.Dataset(
+                atmos_data,
+                coords={
+                    "time": [pred.metadata.time[0]],
+                    "batch": [0],
+                    "level": list(pred.metadata.atmos_levels),
+                    "lat": pred.metadata.lat.cpu().numpy() if hasattr(pred.metadata.lat, 'cpu') else pred.metadata.lat.numpy(),
+                    "lon": pred.metadata.lon.cpu().numpy() if hasattr(pred.metadata.lon, 'cpu') else pred.metadata.lon.numpy(),
+                }
+            )
+
+            # Save to NetCDF
+            surf_filename = f"step_{step:02d}_surface.nc"
+            atmos_filename = f"step_{step:02d}_atmospheric.nc"
+
+            surf_ds.to_netcdf(output_dir / surf_filename)
+            atmos_ds.to_netcdf(output_dir / atmos_filename)
+
+            print(f"Saved step {step} predictions (fallback method)")
+
+        return output_dir
+
+    def run_pipeline(self, date_str, Batch, Metadata, AuroraAirPollution, rollout, steps=4, output_path=None):
+        """Full pipeline: batch creation, model loading, prediction, save output"""
+        batch = self.create_batch(date_str, Batch, Metadata)
+        self.load_model(AuroraAirPollution)
+        predictions = self.predict(batch, rollout, steps=steps)
+        if output_path:
+            self.save_predictions_to_netcdf(predictions, output_path)
+        return predictions
+
+# Example usage (not run on import)
+if __name__ == "__main__":
+    pass

cams_downloader.py

@@ -157,6 +157,7 @@ class CAMSDownloader:
         
         try:
             # Make the API request
+            print("📡 Requesting data from CAMS API...")
             self.client.retrieve(
                 "cams-global-atmospheric-composition-forecasts",
                 {
@@ -171,12 +172,30 @@ class CAMSDownloader:
                 str(filepath),
             )
             
-            print(f"✅ Successfully downloaded: {filename}")
-            return str(filepath)
+            # Validate the downloaded file
+            if filepath.exists():
+                file_size = filepath.stat().st_size
+                print(f"📁 Downloaded file size: {file_size / 1024 / 1024:.2f} MB")
+                
+                # Basic validation - CAMS files should be reasonably large
+                if file_size < 10000:  # Less than 10KB is suspicious
+                    print(f"⚠️  Warning: Downloaded file is very small ({file_size} bytes)")
+                    # Read first few bytes to check for error messages
+                    with open(filepath, 'rb') as f:
+                        header = f.read(200)
+                        if b'error' in header.lower() or b'html' in header.lower():
+                            filepath.unlink()
+                            raise Exception("CAMS API returned an error response instead of data")
+                
+                print(f"✅ Successfully downloaded: {filename}")
+                return str(filepath)
+            else:
+                raise Exception("Download completed but file was not created")
             
         except Exception as e:
             # Clean up partial download
             if filepath.exists():
+                print(f"🗑️  Cleaning up failed download: {filepath}")
                 filepath.unlink()
             raise Exception(f"Error downloading CAMS data: {str(e)}")
     
@@ -194,6 +213,33 @@ class CAMSDownloader:
         if not zip_path.exists():
             raise FileNotFoundError(f"ZIP file not found: {zip_path}")
         
+        # Validate file is actually a ZIP file
+        try:
+            # Check file size first
+            file_size = zip_path.stat().st_size
+            if file_size < 1000:  # Less than 1KB is probably an error response
+                print(f"⚠️  Downloaded file is too small ({file_size} bytes), likely an error response")
+                # Try to read first few bytes to see what we got
+                with open(zip_path, 'rb') as f:
+                    header = f.read(100)
+                    if b'html' in header.lower() or b'error' in header.lower():
+                        raise Exception("Downloaded file appears to be an HTML error page, not ZIP data")
+                
+            # Test if it's a valid ZIP file
+            if not zipfile.is_zipfile(zip_path):
+                print(f"❌ File is not a valid ZIP file: {zip_path}")
+                # Try to read first few lines to diagnose
+                with open(zip_path, 'r', errors='ignore') as f:
+                    first_lines = f.read(200)
+                    print(f"File contents preview: {first_lines[:100]}...")
+                raise Exception(f"Downloaded file is not a valid ZIP archive. File size: {file_size} bytes")
+                
+        except Exception as e:
+            if "ZIP" in str(e) or "zip" in str(e):
+                raise e
+            else:
+                raise Exception(f"Error validating ZIP file: {str(e)}")
+        
         # Extract date from filename
         date_str = zip_path.stem.replace("-cams.nc", "")
         

constants.py

@@ -1,6 +1,7 @@
-# Air pollution variables and their properties
+# NetCDF variables and their properties (includes air pollution, meteorological, and other variables)
 
-AIR_POLLUTION_VARIABLES = {
+NETCDF_VARIABLES = {
+    # === AIR POLLUTION VARIABLES ===
     # PM2.5
     'pm2p5': {'units': 'µg/m³', 'name': 'PM2.5', 'cmap': 'YlOrRd', 'vmax_percentile': 95, 'type': 'surface'},
     'pm25': {'units': 'µg/m³', 'name': 'PM2.5', 'cmap': 'YlOrRd', 'vmax_percentile': 95, 'type': 'surface'},
@@ -78,8 +79,76 @@ AIR_POLLUTION_VARIABLES = {
     'aod550': {'units': '', 'name': 'AOD 550nm', 'cmap': 'plasma', 'vmax_percentile': 95, 'type': 'surface'},
     'aod': {'units': '', 'name': 'AOD', 'cmap': 'plasma', 'vmax_percentile': 95, 'type': 'surface'},
     'aerosol_optical_depth': {'units': '', 'name': 'AOD', 'cmap': 'plasma', 'vmax_percentile': 95, 'type': 'surface'},
+    
+    # === METEOROLOGICAL VARIABLES ===
+    
+    # Temperature variables
+    '2t': {'units': 'K', 'name': '2m Temperature', 'cmap': 'RdYlBu', 'vmax_percentile': 95, 'type': 'surface'},
+    't2m': {'units': 'K', 'name': '2m Temperature', 'cmap': 'RdYlBu', 'vmax_percentile': 95, 'type': 'surface'},
+    'temperature': {'units': 'K', 'name': 'Temperature', 'cmap': 'RdYlBu', 'vmax_percentile': 95, 'type': 'atmospheric'},
+    't': {'units': 'K', 'name': 'Temperature', 'cmap': 'RdYlBu', 'vmax_percentile': 95, 'type': 'atmospheric'},
+    'air_temperature': {'units': 'K', 'name': 'Air Temperature', 'cmap': 'RdYlBu', 'vmax_percentile': 95, 'type': 'atmospheric'},
+    'sst': {'units': 'K', 'name': 'Sea Surface Temperature', 'cmap': 'coolwarm', 'vmax_percentile': 95, 'type': 'surface'},
+    
+    # Wind variables
+    'u': {'units': 'm/s', 'name': 'U-Wind Component', 'cmap': 'coolwarm', 'vmax_percentile': 95, 'type': 'atmospheric'},
+    'v': {'units': 'm/s', 'name': 'V-Wind Component', 'cmap': 'coolwarm', 'vmax_percentile': 95, 'type': 'atmospheric'},
+    'u10': {'units': 'm/s', 'name': '10m U-Wind', 'cmap': 'coolwarm', 'vmax_percentile': 95, 'type': 'surface'},
+    'v10': {'units': 'm/s', 'name': '10m V-Wind', 'cmap': 'coolwarm', 'vmax_percentile': 95, 'type': 'surface'},
+    'u_component_of_wind': {'units': 'm/s', 'name': 'U-Wind Component', 'cmap': 'coolwarm', 'vmax_percentile': 95, 'type': 'atmospheric'},
+    'v_component_of_wind': {'units': 'm/s', 'name': 'V-Wind Component', 'cmap': 'coolwarm', 'vmax_percentile': 95, 'type': 'atmospheric'},
+    'wind_speed': {'units': 'm/s', 'name': 'Wind Speed', 'cmap': 'viridis', 'vmax_percentile': 95, 'type': 'atmospheric'},
+    'ws': {'units': 'm/s', 'name': 'Wind Speed', 'cmap': 'viridis', 'vmax_percentile': 95, 'type': 'atmospheric'},
+    
+    # Pressure and geopotential
+    'msl': {'units': 'Pa', 'name': 'Mean Sea Level Pressure', 'cmap': 'RdYlBu', 'vmax_percentile': 95, 'type': 'surface'},
+    'sp': {'units': 'Pa', 'name': 'Surface Pressure', 'cmap': 'RdYlBu', 'vmax_percentile': 95, 'type': 'surface'},
+    'pressure': {'units': 'Pa', 'name': 'Pressure', 'cmap': 'RdYlBu', 'vmax_percentile': 95, 'type': 'atmospheric'},
+    'z': {'units': 'm²/s²', 'name': 'Geopotential', 'cmap': 'Spectral', 'vmax_percentile': 95, 'type': 'atmospheric'},
+    'gh': {'units': 'm', 'name': 'Geopotential Height', 'cmap': 'Spectral', 'vmax_percentile': 95, 'type': 'atmospheric'},
+    
+    # Humidity variables
+    'r': {'units': '%', 'name': 'Relative Humidity', 'cmap': 'Blues', 'vmax_percentile': 95, 'type': 'atmospheric'},
+    'rh': {'units': '%', 'name': 'Relative Humidity', 'cmap': 'Blues', 'vmax_percentile': 95, 'type': 'atmospheric'},
+    'q': {'units': 'kg/kg', 'name': 'Specific Humidity', 'cmap': 'Blues', 'vmax_percentile': 95, 'type': 'atmospheric'},
+    'sh': {'units': 'kg/kg', 'name': 'Specific Humidity', 'cmap': 'Blues', 'vmax_percentile': 95, 'type': 'atmospheric'},
+    'd2m': {'units': 'K', 'name': '2m Dewpoint Temperature', 'cmap': 'Blues', 'vmax_percentile': 95, 'type': 'surface'},
+    
+    # Precipitation and cloud variables
+    'tp': {'units': 'm', 'name': 'Total Precipitation', 'cmap': 'Blues', 'vmax_percentile': 98, 'type': 'surface'},
+    'precipitation': {'units': 'm', 'name': 'Precipitation', 'cmap': 'Blues', 'vmax_percentile': 98, 'type': 'surface'},
+    'tcc': {'units': '%', 'name': 'Total Cloud Cover', 'cmap': 'Blues', 'vmax_percentile': 95, 'type': 'surface'},
+    'lcc': {'units': '%', 'name': 'Low Cloud Cover', 'cmap': 'Blues', 'vmax_percentile': 95, 'type': 'surface'},
+    'mcc': {'units': '%', 'name': 'Medium Cloud Cover', 'cmap': 'Blues', 'vmax_percentile': 95, 'type': 'surface'},
+    'hcc': {'units': '%', 'name': 'High Cloud Cover', 'cmap': 'Blues', 'vmax_percentile': 95, 'type': 'surface'},
+    
+    # Radiation variables
+    'ssrd': {'units': 'J/m²', 'name': 'Surface Solar Radiation', 'cmap': 'YlOrRd', 'vmax_percentile': 95, 'type': 'surface'},
+    'strd': {'units': 'J/m²', 'name': 'Surface Thermal Radiation', 'cmap': 'inferno', 'vmax_percentile': 95, 'type': 'surface'},
+    'tsr': {'units': 'J/m²', 'name': 'Top Solar Radiation', 'cmap': 'YlOrRd', 'vmax_percentile': 95, 'type': 'surface'},
+    'ttr': {'units': 'J/m²', 'name': 'Top Thermal Radiation', 'cmap': 'inferno', 'vmax_percentile': 95, 'type': 'surface'},
+    
+    # Vertical motion
+    'w': {'units': 'Pa/s', 'name': 'Vertical Velocity (Pressure)', 'cmap': 'coolwarm', 'vmax_percentile': 95, 'type': 'atmospheric'},
+    'omega': {'units': 'Pa/s', 'name': 'Vertical Velocity (Omega)', 'cmap': 'coolwarm', 'vmax_percentile': 95, 'type': 'atmospheric'},
+    
+    # Surface variables
+    'skt': {'units': 'K', 'name': 'Skin Temperature', 'cmap': 'RdYlBu', 'vmax_percentile': 95, 'type': 'surface'},
+    'swvl1': {'units': 'm³/m³', 'name': 'Soil Water Level 1', 'cmap': 'Blues', 'vmax_percentile': 95, 'type': 'surface'},
+    'stl1': {'units': 'K', 'name': 'Soil Temperature Level 1', 'cmap': 'RdYlBu', 'vmax_percentile': 95, 'type': 'surface'},
+    'lsm': {'units': '0-1', 'name': 'Land-Sea Mask', 'cmap': 'viridis', 'vmax_percentile': 100, 'type': 'surface'},
+    
+    # === OTHER COMMON NETCDF VARIABLES ===
+    
+    # Generic/Unknown variables (fallback)
+    'var': {'units': 'unknown', 'name': 'Variable', 'cmap': 'viridis', 'vmax_percentile': 95, 'type': 'unknown'},
+    'data': {'units': 'unknown', 'name': 'Data', 'cmap': 'viridis', 'vmax_percentile': 95, 'type': 'unknown'},
+    'field': {'units': 'unknown', 'name': 'Field', 'cmap': 'viridis', 'vmax_percentile': 95, 'type': 'unknown'},
 }
 
+# Maintain backward compatibility
+AIR_POLLUTION_VARIABLES = {k: v for k, v in NETCDF_VARIABLES.items() if 'pollution' in k.lower() or any(pollutant in k.lower() for pollutant in ['pm', 'no', 'so', 'o3', 'co', 'nh3', 'aod', 'column'])}
+
 # Available color themes for plotting
 COLOR_THEMES = {
     'YlOrRd': 'Yellow-Orange-Red',

data_processor.py

@@ -13,7 +13,7 @@ from pathlib import Path
 from datetime import datetime
 
 # Imports from our Modules
-from constants import AIR_POLLUTION_VARIABLES, PRESSURE_LEVELS
+from constants import NETCDF_VARIABLES, AIR_POLLUTION_VARIABLES, PRESSURE_LEVELS
 warnings.filterwarnings('ignore')
 
 class NetCDFProcessor:
@@ -98,7 +98,7 @@ class NetCDFProcessor:
         return True
     
     def detect_variables(self):
-        """Detect air pollution variables in all loaded datasets"""
+        """Detect all supported variables (pollution, meteorological, etc.) in all loaded datasets"""
         self.detected_variables = {}
         
         # Check surface dataset
@@ -119,60 +119,90 @@ class NetCDFProcessor:
         return self.detected_variables
     
     def _detect_variables_in_dataset(self, dataset, dataset_type):
-        """Detect air pollution variables in a specific dataset"""
+        """Detect all supported variables in a specific dataset"""
         detected = {}
         
         for var_name in dataset.data_vars:
             var_name_lower = var_name.lower()
             
-            # Check exact matches first
-            if var_name in AIR_POLLUTION_VARIABLES:
-                detected[var_name] = AIR_POLLUTION_VARIABLES[var_name].copy()
+            # Check exact matches first in NETCDF_VARIABLES
+            if var_name in NETCDF_VARIABLES:
+                detected[var_name] = NETCDF_VARIABLES[var_name].copy()
                 detected[var_name]['original_name'] = var_name
                 detected[var_name]['dataset_type'] = dataset_type
                 detected[var_name]['shape'] = dataset[var_name].shape
                 detected[var_name]['dims'] = list(dataset[var_name].dims)
                 
-            elif var_name_lower in AIR_POLLUTION_VARIABLES:
-                detected[var_name] = AIR_POLLUTION_VARIABLES[var_name_lower].copy()
+            elif var_name_lower in NETCDF_VARIABLES:
+                detected[var_name] = NETCDF_VARIABLES[var_name_lower].copy()
                 detected[var_name]['original_name'] = var_name
                 detected[var_name]['dataset_type'] = dataset_type
                 detected[var_name]['shape'] = dataset[var_name].shape
                 detected[var_name]['dims'] = list(dataset[var_name].dims)
                 
             else:
-                # Check for partial matches
+                # Auto-detect unknown variables by examining their attributes
                 var_info = dataset[var_name]
                 long_name = getattr(var_info, 'long_name', '').lower()
                 standard_name = getattr(var_info, 'standard_name', '').lower()
+                units = getattr(var_info, 'units', 'unknown')
                 
-                # Check for keywords
-                pollution_keywords = {
-                    'pm2.5': {'units': 'µg/m³', 'name': 'PM2.5', 'cmap': 'YlOrRd', 'vmax_percentile': 95, 'type': 'surface'},
-                    'pm10': {'units': 'µg/m³', 'name': 'PM10', 'cmap': 'Oranges', 'vmax_percentile': 95, 'type': 'surface'},
-                    'pm1': {'units': 'µg/m³', 'name': 'PM1', 'cmap': 'Reds', 'vmax_percentile': 95, 'type': 'surface'},
-                    'no2': {'units': 'µg/m³', 'name': 'NO₂', 'cmap': 'Reds', 'vmax_percentile': 90, 'type': 'atmospheric'},
-                    'nitrogen dioxide': {'units': 'µg/m³', 'name': 'NO₂', 'cmap': 'Reds', 'vmax_percentile': 90, 'type': 'atmospheric'},
-                    'so2': {'units': 'µg/m³', 'name': 'SO₂', 'cmap': 'Purples', 'vmax_percentile': 90, 'type': 'atmospheric'},
-                    'sulphur dioxide': {'units': 'µg/m³', 'name': 'SO₂', 'cmap': 'Purples', 'vmax_percentile': 90, 'type': 'atmospheric'},
-                    'sulfur dioxide': {'units': 'µg/m³', 'name': 'SO₂', 'cmap': 'Purples', 'vmax_percentile': 90, 'type': 'atmospheric'},
-                    'ozone': {'units': 'µg/m³', 'name': 'O₃', 'cmap': 'Blues', 'vmax_percentile': 90, 'type': 'atmospheric'},
-                    'carbon monoxide': {'units': 'mg/m³', 'name': 'CO', 'cmap': 'Greens', 'vmax_percentile': 90, 'type': 'atmospheric'},
-                    'nitrogen monoxide': {'units': 'µg/m³', 'name': 'NO', 'cmap': 'Oranges', 'vmax_percentile': 90, 'type': 'atmospheric'},
-                    'ammonia': {'units': 'µg/m³', 'name': 'NH₃', 'cmap': 'viridis', 'vmax_percentile': 90, 'type': 'atmospheric'},
-                    'particulate': {'units': 'µg/m³', 'name': 'Particulate Matter', 'cmap': 'YlOrRd', 'vmax_percentile': 95, 'type': 'surface'},
-                }
-                
-                for keyword, properties in pollution_keywords.items():
-                    if (keyword in var_name_lower or 
-                        keyword in long_name or 
-                        keyword in standard_name):
+                # Try to match against any known variable in NETCDF_VARIABLES by keywords
+                matched = False
+                for known_var, properties in NETCDF_VARIABLES.items():
+                    if (known_var in var_name_lower or 
+                        known_var in long_name or 
+                        known_var in standard_name or
+                        properties['name'].lower() in var_name_lower or
+                        properties['name'].lower() in long_name):
                         detected[var_name] = properties.copy()
                         detected[var_name]['original_name'] = var_name
                         detected[var_name]['dataset_type'] = dataset_type
                         detected[var_name]['shape'] = dataset[var_name].shape
                         detected[var_name]['dims'] = list(dataset[var_name].dims)
+                        if units != 'unknown':
+                            detected[var_name]['units'] = units  # Use actual units from file
+                        matched = True
                         break
+                
+                # If still no match, create a generic entry for any 2D+ variable
+                if not matched and len(dataset[var_name].dims) >= 2:
+                    # Check if it has lat/lon dimensions
+                    dims = list(dataset[var_name].dims)
+                    has_spatial = any(dim in ['lat', 'lon', 'latitude', 'longitude', 'x', 'y'] 
+                                    for dim in [d.lower() for d in dims])
+                    
+                    if has_spatial:
+                        # Determine variable type based on dimensions
+                        var_type = 'surface'
+                        if any(dim in ['level', 'plev', 'pressure', 'height'] for dim in [d.lower() for d in dims]):
+                            var_type = 'atmospheric'
+                        
+                        # Auto-determine color scheme based on variable name or units
+                        cmap = 'viridis'  # default
+                        if 'temp' in var_name_lower or 'temperature' in long_name:
+                            cmap = 'RdYlBu'
+                        elif any(word in var_name_lower for word in ['wind', 'u', 'v']):
+                            cmap = 'coolwarm'
+                        elif any(word in var_name_lower for word in ['precip', 'rain', 'cloud', 'humid']):
+                            cmap = 'Blues'
+                        elif 'pressure' in var_name_lower or 'pressure' in long_name:
+                            cmap = 'RdYlBu'
+                        elif any(word in var_name_lower for word in ['radiation', 'solar']):
+                            cmap = 'YlOrRd'
+                        
+                        detected[var_name] = {
+                            'units': units,
+                            'name': long_name.title() if long_name else var_name.replace('_', ' ').title(),
+                            'cmap': cmap,
+                            'vmax_percentile': 95,
+                            'type': var_type,
+                            'original_name': var_name,
+                            'dataset_type': dataset_type,
+                            'shape': dataset[var_name].shape,
+                            'dims': dims,
+                            'auto_detected': True  # Flag to indicate this was auto-detected
+                        }
         
         return detected
     
@@ -476,13 +506,27 @@ class NetCDFProcessor:
         return PRESSURE_LEVELS  # Default pressure levels
     
     def close(self):
-        """Close all open datasets"""
-        if self.dataset is not None:
-            self.dataset.close()
-        if self.surface_dataset is not None:
-            self.surface_dataset.close()
-        if self.atmospheric_dataset is not None:
-            self.atmospheric_dataset.close()
+        """Close all open datasets safely"""
+        try:
+            if self.dataset is not None:
+                self.dataset.close()
+                self.dataset = None
+        except (RuntimeError, OSError):
+            pass  # Dataset already closed or invalid
+        
+        try:
+            if self.surface_dataset is not None:
+                self.surface_dataset.close()
+                self.surface_dataset = None
+        except (RuntimeError, OSError):
+            pass  # Dataset already closed or invalid
+            
+        try:
+            if self.atmospheric_dataset is not None:
+                self.atmospheric_dataset.close()
+                self.atmospheric_dataset = None
+        except (RuntimeError, OSError):
+            pass  # Dataset already closed or invalid
 
 
 def analyze_netcdf_file(file_path):

diagnose_cams.py

@@ -0,0 +1,131 @@
+#!/usr/bin/env python3
+"""
+CAMS Download Diagnostic Tool
+Helps troubleshoot issues with CAMS data downloads
+"""
+
+import os
+import zipfile
+from pathlib import Path
+from datetime import datetime, timedelta
+
+def diagnose_cams_downloads():
+    """Diagnose CAMS download issues"""
+    print("🔍 CAMS Download Diagnostic Tool")
+    print("=" * 50)
+    
+    # Check downloads directory
+    downloads_dir = Path("downloads")
+    if not downloads_dir.exists():
+        print("❌ Downloads directory doesn't exist")
+        return
+    
+    print(f"📁 Downloads directory: {downloads_dir.absolute()}")
+    
+    # List all files in downloads
+    all_files = list(downloads_dir.glob("*"))
+    if not all_files:
+        print("📂 Downloads directory is empty")
+        return
+    
+    print(f"\n📋 Found {len(all_files)} files:")
+    
+    for file_path in all_files:
+        print(f"\n📄 File: {file_path.name}")
+        print(f"   Size: {file_path.stat().st_size} bytes ({file_path.stat().st_size / 1024:.1f} KB)")
+        
+        # Check if it's supposed to be a ZIP file
+        if file_path.suffix.lower() == '.zip' or 'cams' in file_path.name.lower():
+            print(f"   Expected: ZIP file")
+            
+            # Test if it's actually a ZIP
+            if zipfile.is_zipfile(file_path):
+                print(f"   ✅ Valid ZIP file")
+                try:
+                    with zipfile.ZipFile(file_path, 'r') as zf:
+                        contents = zf.namelist()
+                        print(f"   📦 Contains {len(contents)} files:")
+                        for content in contents[:5]:  # Show first 5 files
+                            print(f"      - {content}")
+                        if len(contents) > 5:
+                            print(f"      ... and {len(contents) - 5} more")
+                except Exception as e:
+                    print(f"   ⚠️  Error reading ZIP: {e}")
+            else:
+                print(f"   ❌ NOT a valid ZIP file")
+                
+                # Try to read first few bytes to see what it actually is
+                try:
+                    with open(file_path, 'rb') as f:
+                        header = f.read(100)
+                        print(f"   🔍 File header (first 100 bytes): {header[:50]}...")
+                        
+                        # Check for common error patterns
+                        header_str = header.decode('utf-8', errors='ignore').lower()
+                        if 'html' in header_str:
+                            print(f"   🚨 Appears to be HTML (likely an error page)")
+                        elif 'error' in header_str:
+                            print(f"   🚨 Contains 'error' - likely an error response")
+                        elif 'json' in header_str:
+                            print(f"   🚨 Appears to be JSON (likely an API error)")
+                        elif header.startswith(b'PK'):
+                            print(f"   🤔 Has ZIP signature but zipfile module rejects it")
+                        else:
+                            print(f"   ❓ Unknown file format")
+                            
+                except Exception as e:
+                    print(f"   ❌ Error reading file: {e}")
+
+def test_cds_connection():
+    """Test CDS API connection"""
+    print("\n🌐 Testing CDS API Connection")
+    print("-" * 30)
+    
+    try:
+        import cdsapi
+        
+        # Check for .cdsapirc file
+        cdsapirc_path = Path.home() / '.cdsapirc'
+        if cdsapirc_path.exists():
+            print("✅ .cdsapirc file found")
+            
+            # Try to initialize client
+            try:
+                client = cdsapi.Client()
+                print("✅ CDS API client initialized successfully")
+                
+                # Test a simple info request (doesn't download data)
+                print("🔄 Testing API connection...")
+                # Note: This is just a connection test, not actually downloading
+                print("✅ CDS API connection appears to be working")
+                print("💡 If downloads fail, it may be due to:")
+                print("   - Invalid date range")
+                print("   - CAMS service temporary issues") 
+                print("   - Account limitations")
+                
+            except Exception as e:
+                print(f"❌ CDS API client initialization failed: {e}")
+                
+        else:
+            print("❌ .cdsapirc file not found")
+            print("💡 Create ~/.cdsapirc with your CDS API credentials")
+            
+    except ImportError:
+        print("❌ cdsapi module not installed")
+        print("💡 Install with: pip install cdsapi")
+
+def suggest_solutions():
+    """Suggest solutions for common issues"""
+    print("\n💡 Common Solutions")
+    print("-" * 20)
+    print("1. 🔄 Try a different date (some dates may not have data)")
+    print("2. 🕐 Wait and retry (CAMS servers may be busy)")
+    print("3. 🔑 Check CDS API credentials in ~/.cdsapirc")
+    print("4. 🗑️  Clear downloads directory and retry")
+    print("5. 📅 Use more recent dates (last 30 days usually work)")
+    print("6. 🌐 Check CDS website status: https://cds.climate.copernicus.eu/")
+
+if __name__ == "__main__":
+    diagnose_cams_downloads()
+    test_cds_connection()
+    suggest_solutions()

requirements.txt

@@ -1,7 +1,7 @@
 flask==2.3.3
 numpy==1.24.3
-pandas==2.0.3
-matplotlib==3.7.2
+pandas==2.1.4
+matplotlib==3.8.2
 cartopy==0.22.0
 xarray==2023.8.0
 netcdf4==1.6.4
@@ -12,4 +12,8 @@ python-dateutil==2.8.2
 plotly==6.3.0
 kaleido
 geopandas
-shapely
+shapely
+fiona
+torch
+huggingface-hub
+microsoft-aurora

templates/aurora_predict.html

@@ -0,0 +1,586 @@
+<!DOCTYPE html>
+<html lang="en">
+<head>
+    <meta charset="UTF-8">
+    <meta name="viewport" content="width=device-width, initial-scale=1.0">
+    <title>Aurora ML Predictions - CAMS Pollution Dashboard</title>
+    <style>
+        * {
+            margin: 0;
+            padding: 0;
+            box-sizing: border-box;
+        }
+
+        body {
+            font-family: 'Arial', sans-serif;
+            background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
+            min-height: 100vh;
+            color: #333;
+        }
+
+        .container {
+            max-width: 800px;
+            margin: 0 auto;
+            padding: 20px;
+        }
+
+        .header {
+            text-align: center;
+            margin-bottom: 40px;
+            color: white;
+        }
+
+        .header h1 {
+            font-size: 2.5em;
+            margin-bottom: 10px;
+            text-shadow: 2px 2px 4px rgba(0,0,0,0.3);
+        }
+
+        .header p {
+            font-size: 1.2em;
+            opacity: 0.9;
+        }
+
+        .form-container {
+            background: rgba(255, 255, 255, 0.95);
+            border-radius: 15px;
+            padding: 40px;
+            box-shadow: 0 15px 35px rgba(0, 0, 0, 0.1);
+            backdrop-filter: blur(10px);
+        }
+
+        .form-group {
+            margin-bottom: 25px;
+        }
+
+        .form-group label {
+            display: block;
+            margin-bottom: 8px;
+            font-weight: bold;
+            color: #555;
+        }
+
+        .form-group input, .form-group select {
+            width: 100%;
+            padding: 15px;
+            border: 2px solid #e1e1e1;
+            border-radius: 8px;
+            font-size: 16px;
+            transition: border-color 0.3s ease;
+        }
+
+        .form-group input:focus, .form-group select:focus {
+            outline: none;
+            border-color: #667eea;
+            box-shadow: 0 0 0 3px rgba(102, 126, 234, 0.1);
+        }
+
+        .info-box {
+            background: #f8f9ff;
+            border: 2px solid #e3e7ff;
+            border-radius: 10px;
+            padding: 20px;
+            margin-bottom: 25px;
+        }
+
+        .info-box h3 {
+            color: #4c63d2;
+            margin-bottom: 10px;
+        }
+
+        .info-box ul {
+            margin-left: 20px;
+            color: #666;
+        }
+
+        .info-box li {
+            margin-bottom: 5px;
+        }
+
+        .btn {
+            background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
+            color: white;
+            padding: 15px 30px;
+            border: none;
+            border-radius: 8px;
+            font-size: 18px;
+            cursor: pointer;
+            transition: transform 0.2s ease, box-shadow 0.2s ease;
+            width: 100%;
+        }
+
+        .btn:hover {
+            transform: translateY(-2px);
+            box-shadow: 0 8px 25px rgba(102, 126, 234, 0.3);
+        }
+
+        .btn:active {
+            transform: translateY(0);
+        }
+
+        .btn:disabled {
+            background: #bdc3c7;
+            cursor: not-allowed;
+            transform: none;
+        }
+
+        /* Loading Animation Styles */
+        .loading-overlay {
+            display: none;
+            position: fixed;
+            top: 0;
+            left: 0;
+            width: 100%;
+            height: 100%;
+            background: rgba(0, 0, 0, 0.8);
+            z-index: 9999;
+            justify-content: center;
+            align-items: center;
+            flex-direction: column;
+        }
+
+        .loading-content {
+            background: white;
+            padding: 40px;
+            border-radius: 20px;
+            text-align: center;
+            max-width: 500px;
+            box-shadow: 0 20px 60px rgba(0, 0, 0, 0.3);
+        }
+
+        .loading-spinner {
+            width: 80px;
+            height: 80px;
+            border: 8px solid #f3f3f3;
+            border-top: 8px solid #667eea;
+            border-radius: 50%;
+            animation: spin 1s linear infinite;
+            margin: 0 auto 20px;
+        }
+
+        @keyframes spin {
+            0% { transform: rotate(0deg); }
+            100% { transform: rotate(360deg); }
+        }
+
+        .progress-bar {
+            width: 100%;
+            height: 20px;
+            background: #f0f0f0;
+            border-radius: 10px;
+            overflow: hidden;
+            margin: 20px 0;
+        }
+
+        .progress-fill {
+            height: 100%;
+            background: linear-gradient(90deg, #667eea, #764ba2);
+            border-radius: 10px;
+            width: 0%;
+            transition: width 0.3s ease;
+            animation: pulse 2s infinite;
+        }
+
+        @keyframes pulse {
+            0%, 100% { opacity: 1; }
+            50% { opacity: 0.7; }
+        }
+
+        .loading-steps {
+            text-align: left;
+            margin-top: 20px;
+        }
+
+        .loading-step {
+            padding: 8px 0;
+            display: flex;
+            align-items: center;
+            font-size: 14px;
+        }
+
+        .step-icon {
+            width: 20px;
+            height: 20px;
+            border-radius: 50%;
+            margin-right: 10px;
+            display: flex;
+            align-items: center;
+            justify-content: center;
+            font-size: 12px;
+        }
+
+        .step-pending {
+            background: #f0f0f0;
+            color: #999;
+        }
+
+        .step-active {
+            background: #667eea;
+            color: white;
+            animation: pulse 1s infinite;
+        }
+
+        .step-complete {
+            background: #28a745;
+            color: white;
+        }
+
+        .aurora-icon {
+            font-size: 3em;
+            margin-bottom: 15px;
+            animation: float 3s ease-in-out infinite;
+        }
+
+        @keyframes float {
+            0%, 100% { transform: translateY(0px); }
+            50% { transform: translateY(-10px); }
+        }
+
+        .back-link {
+            display: inline-block;
+            margin-bottom: 20px;
+            color: white;
+            text-decoration: none;
+            font-size: 16px;
+            transition: opacity 0.3s ease;
+        }
+
+        .back-link:hover {
+            opacity: 0.8;
+        }
+
+        .back-link::before {
+            content: "← ";
+        }
+
+        .warning-box {
+            background: #fff8e1;
+            border: 2px solid #ffcc02;
+            border-radius: 10px;
+            padding: 20px;
+            margin-bottom: 25px;
+        }
+
+        .warning-box h3 {
+            color: #f57c00;
+            margin-bottom: 10px;
+        }
+
+        @media (max-width: 768px) {
+            .container {
+                padding: 10px;
+            }
+            
+            .form-container {
+                padding: 20px;
+            }
+            
+            .header h1 {
+                font-size: 2em;
+            }
+        }
+    </style>
+</head>
+<body>
+    <div class="container">
+        <a href="{{ url_for('index') }}" class="back-link">Back to Main Dashboard</a>
+        
+        <div class="header">
+            <h1>🔮 Aurora ML Predictions</h1>
+            <p>Generate AI-powered air pollution forecasts using Microsoft's Aurora model</p>
+        </div>
+
+        <div class="form-container">
+            <div class="info-box">
+                <h3>🚀 About Aurora Predictions</h3>
+                <ul>
+                    <li><strong>AI-Powered Forecasting:</strong> Uses Microsoft's Aurora foundation model for atmospheric predictions</li>
+                    <li><strong>Multi-Step Forecasting:</strong> Generate predictions for up to 10 time steps ahead</li>
+                    <li><strong>Real CAMS Data:</strong> Downloads actual atmospheric data for the selected date</li>
+                    <li><strong>Multiple Variables:</strong> Predicts PM2.5, PM10, O₃, NO₂, CO, SO₂ and meteorological variables</li>
+                </ul>
+            </div>
+
+            <div class="warning-box">
+                <h3>⚠️ Performance Notes</h3>
+                <p><strong>CPU Mode:</strong> Aurora will run on CPU for local testing. This is slower but doesn't require GPU.</p>
+                <p><strong>GPU Mode:</strong> If CUDA GPU is available, Aurora will use it for faster predictions.</p>
+                <p><strong>Processing Time:</strong> CPU: 5-15 minutes per step | GPU: 1-3 minutes total</p>
+                <p><strong>Memory:</strong> CPU mode automatically limits to 2 steps to prevent memory issues.</p>
+            </div>
+
+            <form method="POST">
+                <div class="form-group">
+                    <label for="date">📅 Select Date for Initial Conditions:</label>
+                    <input type="date" 
+                           id="date" 
+                           name="date" 
+                           value="{{ current_date }}" 
+                           min="2015-01-01" 
+                           max="{{ current_date }}" 
+                           required>
+                    <small style="color: #666; font-size: 14px;">
+                        Aurora will download CAMS data for this date and generate forecasts
+                    </small>
+                </div>
+
+                <div class="form-group">
+                    <label for="steps">🔢 Number of Forecast Steps:</label>
+                    <select id="steps" name="steps" required>
+                        <option value="1">1 step (6 hours) - Fastest</option>
+                        <option value="2" selected>2 steps (12 hours) - CPU Friendly</option>
+                        <option value="4">4 steps (24 hours) - Standard</option>
+                        <option value="6">6 steps (36 hours) - GPU Recommended</option>
+                        <option value="8">8 steps (48 hours) - GPU Required</option>
+                        <option value="10">10 steps (60 hours) - GPU Required</option>
+                    </select>
+                    <small style="color: #666; font-size: 14px;">
+                        Each step represents 6 hours. CPU mode automatically limits to 2 steps maximum.
+                    </small>
+                </div>
+
+                <button type="submit" class="btn" id="predictBtn">
+                    🚀 Generate Aurora Predictions
+                </button>
+            </form>
+
+            <!-- Loading Overlay -->
+            <div class="loading-overlay" id="loadingOverlay">
+                <div class="loading-content">
+                    <div class="aurora-icon">🔮</div>
+                    <h2 style="color: #667eea; margin-bottom: 10px;">Aurora AI Processing</h2>
+                    <p style="color: #666; margin-bottom: 20px;">Generating atmospheric predictions using Microsoft's Aurora model...</p>
+                    
+                    <div class="loading-spinner"></div>
+                    
+                    <div class="progress-bar">
+                        <div class="progress-fill" id="progressFill"></div>
+                    </div>
+                    
+                    <div id="currentStep" style="font-weight: bold; color: #667eea; margin-bottom: 20px;">
+                        Initializing Aurora pipeline...
+                    </div>
+                    
+                    <div class="loading-steps">
+                        <div class="loading-step">
+                            <div class="step-icon step-pending" id="step1">1</div>
+                            <span>Downloading CAMS atmospheric data</span>
+                        </div>
+                        <div class="loading-step">
+                            <div class="step-icon step-pending" id="step2">2</div>
+                            <span>Loading Aurora ML model</span>
+                        </div>
+                        <div class="loading-step">
+                            <div class="step-icon step-pending" id="step3">3</div>
+                            <span>Processing initial conditions</span>
+                        </div>
+                        <div class="loading-step">
+                            <div class="step-icon step-pending" id="step4">4</div>
+                            <span>Running AI predictions</span>
+                        </div>
+                        <div class="loading-step">
+                            <div class="step-icon step-pending" id="step5">5</div>
+                            <span>Saving results and preparing visualization</span>
+                        </div>
+                    </div>
+                    
+                    <p style="margin-top: 20px; font-size: 12px; color: #999;">
+                        <strong>Estimated time:</strong> <span id="estimatedTime">2-5 minutes</span><br>
+                        This may take longer on CPU-only systems.
+                    </p>
+                </div>
+            </div>
+
+            <div style="margin-top: 30px; padding: 20px; background: #f5f5f5; border-radius: 10px;">
+                <h3 style="color: #555; margin-bottom: 15px;">📊 What You'll Get:</h3>
+                <ul style="color: #666; margin-left: 20px;">
+                    <li>Interactive visualization of predicted air pollution concentrations</li>
+                    <li>Step-by-step forecast evolution over time</li>
+                    <li>Downloadable NetCDF files with all prediction data</li>
+                    <li>Support for all major pollutants and meteorological variables</li>
+                </ul>
+            </div>
+        </div>
+    </div>
+
+    <script>
+        // Aurora Prediction Loading Animation
+        class AuroraLoadingManager {
+            constructor() {
+                this.form = document.querySelector('form');
+                this.predictBtn = document.getElementById('predictBtn');
+                this.overlay = document.getElementById('loadingOverlay');
+                this.progressFill = document.getElementById('progressFill');
+                this.currentStep = document.getElementById('currentStep');
+                this.estimatedTime = document.getElementById('estimatedTime');
+                
+                this.steps = [
+                    { id: 'step1', text: 'Downloading CAMS atmospheric data...', duration: 20000 },
+                    { id: 'step2', text: 'Loading Aurora ML model into memory...', duration: 30000 },
+                    { id: 'step3', text: 'Processing initial atmospheric conditions...', duration: 15000 },
+                    { id: 'step4', text: 'Running AI predictions (this may take a while)...', duration: 60000 },
+                    { id: 'step5', text: 'Saving results and preparing visualization...', duration: 10000 }
+                ];
+                
+                this.currentStepIndex = 0;
+                this.startTime = null;
+                
+                this.init();
+            }
+            
+            init() {
+                this.form.addEventListener('submit', (e) => {
+                    this.startLoading();
+                });
+            }
+            
+            startLoading() {
+                this.startTime = Date.now();
+                this.overlay.style.display = 'flex';
+                this.predictBtn.disabled = true;
+                
+                // Estimate time based on selected steps
+                const steps = parseInt(document.getElementById('steps').value);
+                const isCPU = this.detectCPUMode();
+                this.updateEstimatedTime(steps, isCPU);
+                
+                // Start progress simulation
+                this.simulateProgress();
+            }
+            
+            detectCPUMode() {
+                // Simple heuristic - if user selected fewer steps, likely CPU mode
+                const steps = parseInt(document.getElementById('steps').value);
+                return steps <= 2;
+            }
+            
+            updateEstimatedTime(steps, isCPU) {
+                // Fetch actual system capabilities for better estimates
+                fetch('/api/aurora_status')
+                    .then(response => response.json())
+                    .then(data => {
+                        if (data.available) {
+                            const mode = data.cpu_only ? 'cpu' : 'gpu';
+                            let estimatedMinutes;
+                            
+                            if (mode === 'cpu') {
+                                estimatedMinutes = steps <= 1 ? 5 : 8;
+                            } else {
+                                estimatedMinutes = Math.max(2, steps * 0.5 + 2);
+                            }
+                            
+                            this.estimatedTime.textContent = `${estimatedMinutes}-${estimatedMinutes + 2} minutes (${mode.toUpperCase()} mode)`;
+                        } else {
+                            this.estimatedTime.textContent = 'Aurora not available';
+                        }
+                    })
+                    .catch(() => {
+                        // Fallback to original logic
+                        let baseTime = steps * (isCPU ? 5 : 1);
+                        baseTime += 2;
+                        this.estimatedTime.textContent = `${baseTime}-${baseTime + 2} minutes`;
+                    });
+            }
+            
+            simulateProgress() {
+                let totalDuration = this.steps.reduce((sum, step) => sum + step.duration, 0);
+                let elapsed = 0;
+                
+                this.progressSteps(0);
+            }
+            
+            progressSteps(stepIndex) {
+                if (stepIndex >= this.steps.length) {
+                    return; // Let the actual response handle completion
+                }
+                
+                const step = this.steps[stepIndex];
+                const stepElement = document.getElementById(step.id);
+                
+                // Mark previous steps as complete
+                for (let i = 0; i < stepIndex; i++) {
+                    const prevStep = document.getElementById(this.steps[i].id);
+                    prevStep.className = 'step-icon step-complete';
+                    prevStep.innerHTML = '✓';
+                }
+                
+                // Mark current step as active
+                stepElement.className = 'step-icon step-active';
+                this.currentStep.textContent = step.text;
+                
+                // Update progress bar
+                const progress = ((stepIndex + 1) / this.steps.length) * 100;
+                this.progressFill.style.width = `${progress}%`;
+                
+                // Move to next step after duration
+                setTimeout(() => {
+                    this.progressSteps(stepIndex + 1);
+                }, step.duration);
+            }
+            
+            // Call this when the actual response is received
+            completeLoading() {
+                // Mark all steps as complete
+                this.steps.forEach((step, index) => {
+                    const stepElement = document.getElementById(step.id);
+                    stepElement.className = 'step-icon step-complete';
+                    stepElement.innerHTML = '✓';
+                });
+                
+                this.progressFill.style.width = '100%';
+                this.currentStep.textContent = 'Complete! Redirecting to results...';
+                
+                // Hide overlay after a short delay
+                setTimeout(() => {
+                    this.overlay.style.display = 'none';
+                    this.predictBtn.disabled = false;
+                }, 2000);
+            }
+        }
+        
+        // Initialize loading manager when page loads
+        document.addEventListener('DOMContentLoaded', function() {
+            window.auroraLoader = new AuroraLoadingManager();
+            
+            // Handle form validation
+            const form = document.querySelector('form');
+            const dateInput = document.getElementById('date');
+            const stepsSelect = document.getElementById('steps');
+            
+            form.addEventListener('submit', function(e) {
+                if (!dateInput.value) {
+                    e.preventDefault();
+                    alert('Please select a date for the prediction.');
+                    return;
+                }
+                
+                const selectedDate = new Date(dateInput.value);
+                const today = new Date();
+                const minDate = new Date('2015-01-01');
+                
+                if (selectedDate > today || selectedDate < minDate) {
+                    e.preventDefault();
+                    alert('Please select a date between 2015-01-01 and today.');
+                    return;
+                }
+            });
+            
+            // Update step recommendations based on selection
+            stepsSelect.addEventListener('change', function() {
+                const steps = parseInt(this.value);
+                const recommendations = {
+                    1: 'Fastest option - good for testing',
+                    2: 'CPU-friendly - recommended for local development',
+                    4: 'Standard forecast - good for GPU systems',
+                    6: 'Extended forecast - GPU recommended',
+                    8: 'Long-range forecast - GPU required',
+                    10: 'Maximum forecast - GPU required'
+                };
+                
+                const small = this.parentNode.querySelector('small');
+                small.textContent = `Each step represents 6 hours. ${recommendations[steps] || 'Custom selection'}`;
+            });
+        });
+    </script>
+</body>
+</html>

templates/aurora_prediction_plot.html

@@ -0,0 +1,392 @@
+<!DOCTYPE html>
+<html lang="en">
+<head>
+    <meta charset="UTF-8">
+    <meta name="viewport" content="width=device-width, initial-scale=1.0">
+    <title>Aurora Prediction Visualization - CAMS Dashboard</title>
+    <style>
+        * {
+            margin: 0;
+            padding: 0;
+            box-sizing: border-box;
+        }
+
+        body {
+            font-family: 'Arial', sans-serif;
+            background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
+            min-height: 100vh;
+            color: #333;
+        }
+
+        .container {
+            max-width: 1200px;
+            margin: 0 auto;
+            padding: 20px;
+        }
+
+        .header {
+            text-align: center;
+            margin-bottom: 30px;
+            color: white;
+        }
+
+        .header h1 {
+            font-size: 2.5em;
+            margin-bottom: 10px;
+            text-shadow: 2px 2px 4px rgba(0,0,0,0.3);
+        }
+
+        .controls-container {
+            background: rgba(255, 255, 255, 0.95);
+            border-radius: 15px;
+            padding: 30px;
+            margin-bottom: 30px;
+            box-shadow: 0 15px 35px rgba(0, 0, 0, 0.1);
+            backdrop-filter: blur(10px);
+        }
+
+        .controls-row {
+            display: grid;
+            grid-template-columns: repeat(auto-fit, minmax(200px, 1fr));
+            gap: 20px;
+            margin-bottom: 20px;
+        }
+
+        .form-group {
+            display: flex;
+            flex-direction: column;
+        }
+
+        .form-group label {
+            margin-bottom: 8px;
+            font-weight: bold;
+            color: #555;
+        }
+
+        .form-group select {
+            padding: 12px;
+            border: 2px solid #e1e1e1;
+            border-radius: 8px;
+            font-size: 16px;
+            transition: border-color 0.3s ease;
+        }
+
+        .form-group select:focus {
+            outline: none;
+            border-color: #667eea;
+            box-shadow: 0 0 0 3px rgba(102, 126, 234, 0.1);
+        }
+
+        .btn {
+            background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
+            color: white;
+            padding: 12px 25px;
+            border: none;
+            border-radius: 8px;
+            font-size: 16px;
+            cursor: pointer;
+            transition: transform 0.2s ease, box-shadow 0.2s ease;
+        }
+
+        .btn:hover {
+            transform: translateY(-2px);
+            box-shadow: 0 8px 25px rgba(102, 126, 234, 0.3);
+        }
+
+        .plot-container {
+            background: rgba(255, 255, 255, 0.95);
+            border-radius: 15px;
+            padding: 30px;
+            margin-bottom: 30px;
+            box-shadow: 0 15px 35px rgba(0, 0, 0, 0.1);
+            backdrop-filter: blur(10px);
+            text-align: center;
+        }
+
+        .plot-image {
+            max-width: 100%;
+            height: auto;
+            border-radius: 10px;
+            box-shadow: 0 10px 30px rgba(0, 0, 0, 0.2);
+        }
+
+        .info-grid {
+            display: grid;
+            grid-template-columns: repeat(auto-fit, minmax(250px, 1fr));
+            gap: 20px;
+            margin-bottom: 30px;
+        }
+
+        .info-card {
+            background: rgba(255, 255, 255, 0.95);
+            border-radius: 15px;
+            padding: 25px;
+            box-shadow: 0 15px 35px rgba(0, 0, 0, 0.1);
+            backdrop-filter: blur(10px);
+        }
+
+        .info-card h3 {
+            color: #667eea;
+            margin-bottom: 15px;
+            font-size: 1.3em;
+        }
+
+        .info-card p {
+            margin-bottom: 10px;
+            color: #666;
+        }
+
+        .download-section {
+            background: rgba(255, 255, 255, 0.95);
+            border-radius: 15px;
+            padding: 30px;
+            margin-bottom: 30px;
+            box-shadow: 0 15px 35px rgba(0, 0, 0, 0.1);
+            backdrop-filter: blur(10px);
+            text-align: center;
+        }
+
+        .download-btn {
+            background: linear-gradient(135deg, #28a745 0%, #20c997 100%);
+            color: white;
+            padding: 15px 30px;
+            border: none;
+            border-radius: 8px;
+            font-size: 18px;
+            text-decoration: none;
+            display: inline-block;
+            transition: transform 0.2s ease, box-shadow 0.2s ease;
+            margin: 10px;
+        }
+
+        .download-btn:hover {
+            transform: translateY(-2px);
+            box-shadow: 0 8px 25px rgba(40, 167, 69, 0.3);
+            text-decoration: none;
+            color: white;
+        }
+
+        .back-link {
+            display: inline-block;
+            margin-bottom: 20px;
+            color: white;
+            text-decoration: none;
+            font-size: 16px;
+            transition: opacity 0.3s ease;
+        }
+
+        .back-link:hover {
+            opacity: 0.8;
+        }
+
+        .back-link::before {
+            content: "← ";
+        }
+
+        .step-indicator {
+            background: #f8f9ff;
+            border: 2px solid #e3e7ff;
+            border-radius: 10px;
+            padding: 15px;
+            margin-bottom: 20px;
+            text-align: center;
+        }
+
+        .step-indicator h3 {
+            color: #4c63d2;
+            margin-bottom: 5px;
+        }
+
+        @media (max-width: 768px) {
+            .container {
+                padding: 10px;
+            }
+            
+            .controls-container, .plot-container, .info-card, .download-section {
+                padding: 20px;
+            }
+            
+            .header h1 {
+                font-size: 2em;
+            }
+            
+            .controls-row {
+                grid-template-columns: 1fr;
+            }
+        }
+
+        /* Mini loading indicator for form updates */
+        .mini-loading {
+            display: none;
+            width: 20px;
+            height: 20px;
+            border: 2px solid #f3f3f3;
+            border-top: 2px solid #667eea;
+            border-radius: 50%;
+            animation: spin 1s linear infinite;
+            margin-left: 10px;
+        }
+
+        .form-updating {
+            opacity: 0.7;
+            pointer-events: none;
+        }
+    </style>
+</head>
+<body>
+    <div class="container">
+        <a href="{{ url_for('index') }}" class="back-link">Back to Main Dashboard</a>
+        
+        <div class="header">
+            <h1>🔮 Aurora ML Prediction Results</h1>
+            <p>AI-powered atmospheric forecasting visualization</p>
+        </div>
+
+        <div class="step-indicator">
+            <h3>📊 Current View: Step {{ step }} of {{ steps|length - 1 }}</h3>
+            <p>Forecast time: {{ (step * 6) }} hours ahead</p>
+        </div>
+
+        <form method="POST" id="predictionForm">
+            <div class="controls-container" id="controlsContainer">
+                <div class="controls-row">
+                    <div class="form-group">
+                        <label for="variable">🧪 Variable:</label>
+                        <select id="variable" name="variable">
+                            {% for var in variables %}
+                            <option value="{{ var }}" {% if var == var_name %}selected{% endif %}>
+                                {{ var }}
+                            </option>
+                            {% endfor %}
+                        </select>
+                    </div>
+
+                    <div class="form-group">
+                        <label for="step">⏰ Forecast Step:</label>
+                        <select id="step" name="step">
+                            {% for s in steps %}
+                            <option value="{{ s }}" {% if s == step %}selected{% endif %}>
+                                Step {{ s }} ({{ s * 6 }}h ahead)
+                            </option>
+                            {% endfor %}
+                        </select>
+                    </div>
+
+                    <div class="form-group">
+                        <label for="color_theme">🎨 Color Theme:</label>
+                        <select id="color_theme" name="color_theme">
+                            {% for theme_id, theme_name in color_themes.items() %}
+                            <option value="{{ theme_id }}" {% if theme_id == current_color_theme %}selected{% endif %}>
+                                {{ theme_name }}
+                            </option>
+                            {% endfor %}
+                        </select>
+                    </div>
+
+                    <div class="form-group">
+                        <label>&nbsp;</label>
+                        <button type="submit" class="btn" id="updateBtn">
+                            🔄 Update View
+                            <div class="mini-loading" id="miniLoading"></div>
+                        </button>
+                    </div>
+                </div>
+            </div>
+        </form>
+
+        <div class="plot-container">
+            <img src="{{ url_for('serve_plot', filename=plot_filename) }}" 
+                 alt="Aurora Prediction Plot" 
+                 class="plot-image">
+        </div>
+
+        <div class="info-grid">
+            <div class="info-card">
+                <h3>🔬 Variable Information</h3>
+                <p><strong>Variable:</strong> {{ var_name }}</p>
+                <p><strong>Forecast Step:</strong> {{ step }}</p>
+                <p><strong>Time Ahead:</strong> {{ step * 6 }} hours</p>
+                <p><strong>Color Theme:</strong> {{ color_themes[current_color_theme] }}</p>
+            </div>
+
+            <div class="info-card">
+                <h3>🤖 Model Information</h3>
+                <p><strong>Model:</strong> Microsoft Aurora Air Pollution</p>
+                <p><strong>Version:</strong> 0.4</p>
+                <p><strong>Type:</strong> Foundation Model</p>
+                <p><strong>Resolution:</strong> Global atmospheric prediction</p>
+            </div>
+
+            <div class="info-card">
+                <h3>📈 Forecast Details</h3>
+                <p><strong>Total Steps:</strong> {{ steps|length }}</p>
+                <p><strong>Step Interval:</strong> 6 hours</p>
+                <p><strong>Max Forecast:</strong> {{ (steps|length - 1) * 6 }} hours</p>
+                <p><strong>Variables:</strong> {{ variables|length }} predicted</p>
+            </div>
+        </div>
+
+        <div class="download-section">
+            <h3 style="color: #667eea; margin-bottom: 20px;">📁 Download Prediction Data</h3>
+            <p style="margin-bottom: 20px; color: #666;">
+                Download the complete NetCDF file containing all forecast steps and variables
+            </p>
+            <a href="{{ download_url }}" class="download-btn">
+                💾 Download NetCDF File
+            </a>
+            <a href="{{ url_for('aurora_predict') }}" class="download-btn" style="background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);">
+                🔮 Generate New Prediction
+            </a>
+        </div>
+    </div>
+
+    <script>
+        // Enhanced form submission with loading indicators
+        document.addEventListener('DOMContentLoaded', function() {
+            const form = document.getElementById('predictionForm');
+            const updateBtn = document.getElementById('updateBtn');
+            const miniLoading = document.getElementById('miniLoading');
+            const controlsContainer = document.getElementById('controlsContainer');
+            const selects = document.querySelectorAll('select');
+            
+            // Add change event listeners to selects
+            selects.forEach(select => {
+                select.addEventListener('change', function() {
+                    showMiniLoading();
+                    
+                    // Add a small delay to prevent rapid submissions
+                    setTimeout(() => {
+                        form.submit();
+                    }, 100);
+                });
+            });
+            
+            // Handle manual form submission
+            form.addEventListener('submit', function(e) {
+                showMiniLoading();
+            });
+            
+            function showMiniLoading() {
+                miniLoading.style.display = 'inline-block';
+                controlsContainer.classList.add('form-updating');
+                updateBtn.disabled = true;
+                updateBtn.textContent = '🔄 Updating...';
+            }
+            
+            // Auto-hide loading indicator if page doesn't redirect within 10 seconds
+            setTimeout(() => {
+                if (miniLoading.style.display === 'inline-block') {
+                    hideMiniLoading();
+                }
+            }, 10000);
+            
+            function hideMiniLoading() {
+                miniLoading.style.display = 'none';
+                controlsContainer.classList.remove('form-updating');
+                updateBtn.disabled = false;
+                updateBtn.innerHTML = '🔄 Update View<div class="mini-loading" id="miniLoading"></div>';
+            }
+        });
+    </script>
+</body>
+</html>

templates/index.html

@@ -257,6 +257,36 @@
             </form>
         </div>
 
+        <div class="container">
+            <div class="method-section" style="border-left: 4px solid #9b59b6;">
+                <h2>🔮 Method 4: Aurora ML Predictions</h2>
+                <p>Generate AI-powered air pollution forecasts using Microsoft's Aurora foundation model</p>
+                
+                <div style="background: #f8f9ff; padding: 15px; border-radius: 8px; margin: 15px 0; border: 2px solid #e3e7ff;">
+                    <p style="margin-bottom: 10px;"><strong>🚀 What is Aurora?</strong></p>
+                    <ul style="margin-left: 20px; color: #666;">
+                        <li>Microsoft's state-of-the-art atmospheric foundation model</li>
+                        <li>Trained on massive amounts of global weather and atmospheric data</li>
+                        <li>Generates multi-step forecasts for air pollution and meteorology</li>
+                        <li>Provides predictions up to 60 hours ahead with 6-hour intervals</li>
+                    </ul>
+                </div>
+                
+                {% if aurora_available is defined and aurora_available %}
+                    <a href="{{ url_for('aurora_predict') }}" class="btn" style="background: linear-gradient(135deg, #9b59b6 0%, #8e44ad 100%);">
+                        🔮 Generate Aurora Predictions
+                    </a>
+                {% else %}
+                    <button class="btn" disabled style="background: #bdc3c7; cursor: not-allowed;">
+                        🔮 Aurora Model Not Available
+                    </button>
+                    <p style="margin-top: 10px; font-size: 14px; color: #721c24;">
+                        Aurora dependencies not installed. Requires PyTorch and aurora-forecast package.
+                    </p>
+                {% endif %}
+            </div>
+        </div>
+
         <div class="container">
             <h2>📋 How to Use</h2>
             <ol style="line-height: 1.8;">

test_aurora_fix.py

@@ -0,0 +1,61 @@
+#!/usr/bin/env python3
+"""
+Quick test for Aurora pipeline generator fix
+"""
+
+import sys
+import os
+sys.path.append('.')
+
+def test_aurora_fix():
+    """Test the Aurora pipeline generator fix"""
+    print("🧪 Testing Aurora Pipeline Generator Fix")
+    print("=" * 45)
+    
+    try:
+        from aurora_pipeline import AuroraPipeline
+        print("✅ Aurora pipeline imported successfully")
+        
+        # Initialize pipeline
+        pipeline = AuroraPipeline()
+        print("✅ Pipeline initialized")
+        
+        # Test with minimal configuration
+        date = '2022-10-14'
+        steps = 1
+        
+        print(f"\n🚀 Running test prediction:")
+        print(f"   📅 Date: {date}")
+        print(f"   🔄 Steps: {steps}")
+        
+        result = pipeline.run_pipeline(date, steps=steps)
+        
+        print("✅ SUCCESS! Aurora pipeline completed without errors")
+        
+        if isinstance(result, dict):
+            print(f"📊 Result contains: {list(result.keys())}")
+            
+            # Check if we have expected outputs
+            if 'air_pollution' in result:
+                pollution_data = result['air_pollution']
+                print(f"🌬️  Air pollution predictions: {pollution_data.shape if hasattr(pollution_data, 'shape') else type(pollution_data)}")
+            
+            if 'metadata' in result:
+                metadata = result['metadata']
+                print(f"📋 Metadata: {metadata}")
+                
+        print("\n🎉 Aurora pipeline is now working correctly!")
+        
+    except Exception as e:
+        print(f"❌ Error during test: {e}")
+        import traceback
+        print("\n🔍 Full traceback:")
+        traceback.print_exc()
+        
+        print("\n💡 Potential solutions:")
+        print("1. Check if CAMS data exists for the test date")
+        print("2. Verify Aurora model downloaded correctly")
+        print("3. Ensure all dependencies are installed")
+
+if __name__ == "__main__":
+    test_aurora_fix()

test_cpu_mode.py

@@ -0,0 +1,124 @@
+#!/usr/bin/env python3
+"""
+CPU-friendly test for local development without GPU requirements
+"""
+
+def test_system_readiness():
+    """Test if the basic system is ready for local development"""
+    print("🔬 Testing System for Local Development")
+    print("=" * 50)
+    
+    issues = []
+    
+    # Test 1: Basic Python imports
+    try:
+        import sys
+        print(f"✅ Python version: {sys.version.split()[0]}")
+    except Exception as e:
+        issues.append(f"Python: {e}")
+    
+    # Test 2: Core dependencies
+    try:
+        import numpy as np
+        print(f"✅ NumPy: {np.__version__}")
+    except ImportError:
+        issues.append("NumPy not found")
+    
+    try:
+        import xarray as xr
+        print(f"✅ xarray: {xr.__version__}")
+    except ImportError:
+        issues.append("xarray not found")
+    
+    try:
+        import matplotlib
+        print(f"✅ Matplotlib: {matplotlib.__version__}")
+    except ImportError:
+        issues.append("Matplotlib not found")
+    
+    try:
+        import flask
+        print(f"✅ Flask: {flask.__version__}")
+    except ImportError:
+        issues.append("Flask not found")
+    
+    # Test 3: Optional Aurora dependencies
+    aurora_ready = True
+    try:
+        import torch
+        print(f"✅ PyTorch: {torch.__version__}")
+        print(f"   CUDA available: {torch.cuda.is_available()}")
+        print(f"   CPU cores: {torch.get_num_threads()}")
+    except ImportError:
+        print("⚠️  PyTorch not found (Aurora unavailable)")
+        aurora_ready = False
+    
+    try:
+        from huggingface_hub import hf_hub_download
+        print("✅ Hugging Face Hub available")
+    except ImportError:
+        print("⚠️  Hugging Face Hub not found (Aurora unavailable)")
+        aurora_ready = False
+    
+    # Test 4: Aurora model
+    try:
+        from aurora import Batch, Metadata, AuroraAirPollution, rollout
+        print("✅ Aurora model available")
+    except ImportError:
+        print("⚠️  Aurora model not found (predictions unavailable)")
+        aurora_ready = False
+    
+    # Test 5: Custom modules
+    try:
+        from data_processor import NetCDFProcessor
+        from plot_generator import IndiaMapPlotter
+        from constants import NETCDF_VARIABLES
+        print("✅ Custom modules loaded")
+    except ImportError as e:
+        issues.append(f"Custom modules: {e}")
+    
+    # Test 6: GPU vs CPU detection
+    gpu_info = "None"
+    try:
+        import subprocess
+        result = subprocess.run(['nvidia-smi', '--query-gpu=name', '--format=csv,noheader,nounits'], 
+                              capture_output=True, text=True, timeout=5)
+        if result.returncode == 0:
+            gpu_info = result.stdout.strip()
+            print(f"🎮 GPU detected: {gpu_info}")
+        else:
+            print("💻 No GPU detected (CPU mode)")
+    except:
+        print("💻 No GPU utilities found (CPU mode)")
+    
+    print("\n" + "=" * 50)
+    print("📊 System Assessment:")
+    
+    if issues:
+        print("❌ Critical Issues Found:")
+        for issue in issues:
+            print(f"   - {issue}")
+        print("\n🔧 Please install missing dependencies")
+    else:
+        print("✅ Core system ready!")
+        
+    if aurora_ready:
+        print("🔮 Aurora ML predictions: Available")
+        if "CPU mode" in gpu_info or gpu_info == "None":
+            print("💻 Recommended: Use CPU mode with max 2 steps")
+        else:
+            print("🎮 GPU available: Can use more prediction steps")
+    else:
+        print("⚠️  Aurora ML predictions: Not available")
+        print("   Install: torch, huggingface_hub, aurora-forecast")
+    
+    return len(issues) == 0
+
+if __name__ == "__main__":
+    success = test_system_readiness()
+    
+    if success:
+        print("\n🚀 Ready to run CAMS visualization system!")
+        print("💡 Start with: python app.py")
+    else:
+        print("\n❌ Please resolve issues before running the system")