forked from gsn/predictor
104 lines
2.8 KiB
Go
104 lines
2.8 KiB
Go
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)
|
|
}
|