feat: refactor

cmd/compare_step0/main.go

@@ -0,0 +1,104 @@
+package main
+
+import (
+	"context"
+	"fmt"
+	"os"
+	"time"
+
+	"predictor-refactored/internal/dataset"
+	"predictor-refactored/internal/downloader"
+
+	"go.uber.org/zap"
+)
+
+// Downloads step 0 of a given run and writes a minimal dataset for comparison.
+// Usage: go run ./cmd/compare_step0 <run_YYYYMMDDHH> <output_path>
+func main() {
+	if len(os.Args) < 3 {
+		fmt.Fprintf(os.Stderr, "Usage: %s <run_YYYYMMDDHH> <output_path>\n", os.Args[0])
+		os.Exit(1)
+	}
+
+	runStr := os.Args[1]
+	outPath := os.Args[2]
+
+	run, err := time.Parse("2006010215", runStr)
+	if err != nil {
+		fmt.Fprintf(os.Stderr, "Invalid run time %q: %v\n", runStr, err)
+		os.Exit(1)
+	}
+
+	log, _ := zap.NewDevelopment()
+
+	// Create a full-size dataset (we only fill step 0)
+	fmt.Printf("Creating dataset at %s (%d bytes)...\n", outPath, dataset.DatasetSize)
+	ds, err := dataset.Create(outPath)
+	if err != nil {
+		fmt.Fprintf(os.Stderr, "Create dataset: %v\n", err)
+		os.Exit(1)
+	}
+	defer ds.Close()
+
+	cfg := &downloader.Config{
+		DataDir:  os.TempDir(),
+		Parallel: 4,
+	}
+	dl := downloader.NewDownloader(cfg, log)
+
+	ctx := context.Background()
+	date := run.Format("20060102")
+	runHour := run.Hour()
+
+	// Download and blit step 0 from pgrb2
+	fmt.Println("Downloading pgrb2 step 0...")
+	urlA := dataset.GribURL(date, runHour, 0)
+	if err := dl.DownloadAndBlit(ctx, ds, urlA, 0, dataset.LevelSetA); err != nil {
+		fmt.Fprintf(os.Stderr, "pgrb2: %v\n", err)
+		os.Exit(1)
+	}
+	fmt.Println("  done")
+
+	// Download and blit step 0 from pgrb2b
+	fmt.Println("Downloading pgrb2b step 0...")
+	urlB := dataset.GribURLB(date, runHour, 0)
+	if err := dl.DownloadAndBlit(ctx, ds, urlB, 0, dataset.LevelSetB); err != nil {
+		fmt.Fprintf(os.Stderr, "pgrb2b: %v\n", err)
+		os.Exit(1)
+	}
+	fmt.Println("  done")
+
+	if err := ds.Flush(); err != nil {
+		fmt.Fprintf(os.Stderr, "Flush: %v\n", err)
+		os.Exit(1)
+	}
+
+	// Spot-check: print same values as the Python script
+	fmt.Println("\n=== Go dataset values (spot check) ===")
+	type testPoint struct {
+		varName  string
+		varIdx   int
+		levelIdx int
+		lat, lon int
+	}
+
+	points := []testPoint{
+		{"HGT", 0, 0, 0, 0},       // HGT @ 1000mb, lat=-90, lon=0
+		{"HGT", 0, 0, 180, 0},     // HGT @ 1000mb, lat=0, lon=0
+		{"HGT", 0, 0, 360, 0},     // HGT @ 1000mb, lat=+90, lon=0
+		{"HGT", 0, 20, 180, 360},  // HGT @ 500mb, lat=0, lon=180
+		{"UGRD", 1, 0, 180, 0},    // UGRD @ 1000mb, lat=0, lon=0
+		{"VGRD", 2, 0, 180, 0},    // VGRD @ 1000mb, lat=0, lon=0
+		{"UGRD", 1, 20, 284, 0},   // UGRD @ 500mb, lat=52N, lon=0
+	}
+
+	for _, p := range points {
+		val := ds.Val(0, p.levelIdx, p.varIdx, p.lat, p.lon)
+		actualLat := -90.0 + float64(p.lat)*0.5
+		actualLon := float64(p.lon) * 0.5
+		fmt.Printf("  %-4s %4dmb lat=%+7.1f lon=%6.1f: %12.4f\n",
+			p.varName, dataset.Pressures[p.levelIdx], actualLat, actualLon, val)
+	}
+
+	fmt.Printf("\nDataset written to %s\n", outPath)
+}