engine refactor

internal/datasets/gefs/source.go

@@ -0,0 +1,151 @@
+// Package gefs implements datasets.Source for NOAA GEFS (Global Ensemble
+// Forecast System) forecasts.
+//
+// Each ensemble member is treated as its own dataset, selected via
+// DatasetID.Subset.Members. The download skeleton (HTTP, idx parsing,
+// parallel blit) lives in internal/datasets/grib; this package only
+// supplies GEFS-specific URL templating and member resolution.
+package gefs
+
+import (
+	"context"
+	"fmt"
+	"net/http"
+	"time"
+
+	"go.uber.org/zap"
+
+	"predictor-refactored/internal/datasets"
+	"predictor-refactored/internal/datasets/grib"
+	"predictor-refactored/internal/weather"
+	wgfs "predictor-refactored/internal/weather/gfs"
+)
+
+// Source is the GEFS implementation of datasets.Source.
+type Source struct {
+	Variant  *wgfs.Variant
+	Parallel int
+	Client   *http.Client
+	Log      *zap.Logger
+}
+
+// NewSource returns a default Source over variant. If variant is nil,
+// GEFS 0.5° 3-hour is used.
+func NewSource(variant *wgfs.Variant, log *zap.Logger) *Source {
+	if variant == nil {
+		variant = wgfs.GEFS0p50_3h
+	}
+	return &Source{
+		Variant:  variant,
+		Parallel: 8,
+		Client:   &http.Client{Timeout: 2 * time.Minute},
+		Log:      log,
+	}
+}
+
+func (s *Source) ID() string { return s.Variant.ID }
+
+func (s *Source) downloader() *grib.Downloader {
+	return &grib.Downloader{
+		Variant:  s.Variant,
+		URLs:     s.url,
+		Parallel: s.Parallel,
+		Client:   s.Client,
+		Log:      s.Log,
+	}
+}
+
+// url generates the GEFS URL for (date, runHour, member, step, levelSet).
+func (s *Source) url(date string, runHour, member, step int, ls wgfs.LevelSet) string {
+	if ls == wgfs.LevelSetB {
+		return wgfs.GefsGribURLB(date, runHour, member, step, s.Variant.ResToken)
+	}
+	return wgfs.GefsGribURL(date, runHour, member, step, s.Variant.ResToken)
+}
+
+// LatestEpoch HEAD-checks the control member's final forecast hour.
+func (s *Source) LatestEpoch(ctx context.Context) (time.Time, error) {
+	now := time.Now().UTC()
+	hour := now.Hour() - (now.Hour() % 6)
+	current := time.Date(now.Year(), now.Month(), now.Day(), hour, 0, 0, 0, time.UTC)
+
+	client := s.Client
+	if client == nil {
+		client = &http.Client{Timeout: 2 * time.Minute}
+	}
+	log := s.Log
+	if log == nil {
+		log = zap.NewNop()
+	}
+
+	for range 8 {
+		date := current.Format("20060102")
+		url := wgfs.GefsGribURL(date, current.Hour(), 0, s.Variant.MaxHour, s.Variant.ResToken) + ".idx"
+		req, err := http.NewRequestWithContext(ctx, http.MethodHead, url, nil)
+		if err == nil {
+			resp, err := client.Do(req)
+			if err == nil {
+				resp.Body.Close()
+				if resp.StatusCode == http.StatusOK {
+					log.Info("latest GEFS run discovered",
+						zap.Time("run", current),
+						zap.String("verified_url", url))
+					return current, nil
+				}
+			}
+		}
+		current = current.Add(-6 * time.Hour)
+	}
+	return time.Time{}, fmt.Errorf("no recent GEFS run found")
+}
+
+// Coverage returns the extent of id.
+func (s *Source) Coverage(id datasets.DatasetID) datasets.Coverage {
+	v := s.Variant
+	cov := datasets.Coverage{
+		Region:    datasets.Region{MinLat: -90, MaxLat: 90, MinLng: 0, MaxLng: 360},
+		StartTime: id.Epoch,
+		EndTime:   id.Epoch.Add(time.Duration(v.MaxHour) * time.Hour),
+	}
+	if r := id.Subset.Region; r != nil {
+		cov.Region = *r
+	}
+	if h := id.Subset.HourRange; h != nil {
+		cov.StartTime = id.Epoch.Add(time.Duration(h.MinHour) * time.Hour)
+		cov.EndTime = id.Epoch.Add(time.Duration(h.MaxHour) * time.Hour)
+	}
+	return cov
+}
+
+// Open loads a stored GEFS dataset as a WindField.
+func (s *Source) Open(_ context.Context, id datasets.DatasetID, store datasets.Storage) (weather.WindField, error) {
+	if !store.Exists(id) {
+		return nil, fmt.Errorf("dataset %s not found", id.Filename())
+	}
+	file, err := wgfs.Open(store.Path(id), s.Variant, id.Epoch.UTC())
+	if err != nil {
+		return nil, err
+	}
+	return wgfs.NewWind(file), nil
+}
+
+// memberOf extracts the single member index encoded by id.Subset.Members.
+func memberOf(id datasets.DatasetID) (int, error) {
+	if len(id.Subset.Members) != 1 {
+		return 0, fmt.Errorf("gefs dataset id must specify exactly one member (got %v)", id.Subset.Members)
+	}
+	m := id.Subset.Members[0]
+	if m < 0 || m >= wgfs.GEFSMembers {
+		return 0, fmt.Errorf("gefs member %d out of range [0, %d)", m, wgfs.GEFSMembers)
+	}
+	return m, nil
+}
+
+// Download fetches one ensemble member's dataset.
+func (s *Source) Download(ctx context.Context, id datasets.DatasetID, store datasets.Storage, prog datasets.ProgressSink, throttle datasets.Throttle) error {
+	member, err := memberOf(id)
+	if err != nil {
+		return err
+	}
+	return s.downloader().Run(ctx, id, member, store, prog, throttle)
+}