engine refactor

internal/weather/gfs/file.go

@@ -11,8 +11,10 @@ import (
 )
 
 // File is an mmap-backed wind dataset file. The layout is a flat C-order
-// row-major array of float32 values, shape (hour, level, variable, lat, lng).
+// row-major float32 array, shape (hour, level, variable, lat, lng), with
+// the per-axis sizes coming from Variant.
 type File struct {
+	variant  *Variant
 	mm       mmap.MMap
 	file     *os.File
 	writable bool
@@ -20,8 +22,11 @@ type File struct {
 	Epoch time.Time
 }
 
+// Variant returns the Variant the file was created with.
+func (d *File) Variant() *Variant { return d.variant }
+
 // Open opens an existing dataset file for reading.
-func Open(path string, epoch time.Time) (*File, error) {
+func Open(path string, variant *Variant, epoch time.Time) (*File, error) {
 	f, err := os.Open(path)
 	if err != nil {
 		return nil, fmt.Errorf("open dataset: %w", err)
@@ -31,39 +36,40 @@ func Open(path string, epoch time.Time) (*File, error) {
 		f.Close()
 		return nil, fmt.Errorf("stat dataset: %w", err)
 	}
-	if info.Size() != DatasetSize {
+	if info.Size() != variant.DatasetSize() {
 		f.Close()
-		return nil, fmt.Errorf("dataset should be %d bytes (was %d)", DatasetSize, info.Size())
+		return nil, fmt.Errorf("dataset should be %d bytes (was %d)", variant.DatasetSize(), info.Size())
 	}
 	mm, err := mmap.Map(f, mmap.RDONLY, 0)
 	if err != nil {
 		f.Close()
 		return nil, fmt.Errorf("mmap dataset: %w", err)
 	}
-	return &File{mm: mm, file: f, writable: false, Epoch: epoch}, nil
+	return &File{variant: variant, mm: mm, file: f, writable: false, Epoch: epoch}, nil
 }
 
-// Create creates a new dataset file of the canonical size, mmap'd read-write.
-func Create(path string) (*File, error) {
+// Create creates a new dataset file sized for variant, mmap'd read-write.
+func Create(path string, variant *Variant) (*File, error) {
 	f, err := os.Create(path)
 	if err != nil {
 		return nil, fmt.Errorf("create dataset: %w", err)
 	}
-	if err := f.Truncate(DatasetSize); err != nil {
+	size := variant.DatasetSize()
+	if err := f.Truncate(size); err != nil {
 		f.Close()
 		return nil, fmt.Errorf("truncate dataset: %w", err)
 	}
-	mm, err := mmap.MapRegion(f, int(DatasetSize), mmap.RDWR, 0, 0)
+	mm, err := mmap.MapRegion(f, int(size), mmap.RDWR, 0, 0)
 	if err != nil {
 		f.Close()
 		return nil, fmt.Errorf("mmap dataset: %w", err)
 	}
-	return &File{mm: mm, file: f, writable: true}, nil
+	return &File{variant: variant, mm: mm, file: f, writable: true}, nil
 }
 
-// OpenWritable opens an existing dataset file for read-write access.
-// Used when resuming a partial download.
-func OpenWritable(path string) (*File, error) {
+// OpenWritable opens an existing dataset file for read-write access. Used
+// when resuming a partial download.
+func OpenWritable(path string, variant *Variant) (*File, error) {
 	f, err := os.OpenFile(path, os.O_RDWR, 0o644)
 	if err != nil {
 		return nil, fmt.Errorf("open dataset rw: %w", err)
@@ -73,51 +79,55 @@ func OpenWritable(path string) (*File, error) {
 		f.Close()
 		return nil, fmt.Errorf("stat dataset: %w", err)
 	}
-	if info.Size() != DatasetSize {
+	if info.Size() != variant.DatasetSize() {
 		f.Close()
-		return nil, fmt.Errorf("dataset should be %d bytes (was %d)", DatasetSize, info.Size())
+		return nil, fmt.Errorf("dataset should be %d bytes (was %d)", variant.DatasetSize(), info.Size())
 	}
-	mm, err := mmap.MapRegion(f, int(DatasetSize), mmap.RDWR, 0, 0)
+	mm, err := mmap.MapRegion(f, int(info.Size()), mmap.RDWR, 0, 0)
 	if err != nil {
 		f.Close()
 		return nil, fmt.Errorf("mmap dataset: %w", err)
 	}
-	return &File{mm: mm, file: f, writable: true}, nil
+	return &File{variant: variant, mm: mm, file: f, writable: true}, nil
 }
 
 // offset returns the byte offset of the [hour][level][variable][lat][lng] cell.
-func offset(hour, level, variable, lat, lng int) int64 {
+func (d *File) offset(hour, level, variable, lat, lng int) int64 {
+	v := d.variant
 	idx := int64(hour)
-	idx = idx*int64(NumLevels) + int64(level)
+	idx = idx*int64(v.NumLevels()) + int64(level)
 	idx = idx*int64(NumVariables) + int64(variable)
-	idx = idx*int64(NumLatitudes) + int64(lat)
-	idx = idx*int64(NumLongitudes) + int64(lng)
+	idx = idx*int64(v.NumLatitudes()) + int64(lat)
+	idx = idx*int64(v.NumLongitudes()) + int64(lng)
 	return idx * int64(ElementSize)
 }
 
 // Val reads one cell as a float32.
 func (d *File) Val(hour, level, variable, lat, lng int) float32 {
-	off := offset(hour, level, variable, lat, lng)
+	off := d.offset(hour, level, variable, lat, lng)
 	return math.Float32frombits(binary.LittleEndian.Uint32(d.mm[off : off+4]))
 }
 
 // SetVal writes one cell. Only valid on writable files.
 func (d *File) SetVal(hour, level, variable, lat, lng int, val float32) {
-	off := offset(hour, level, variable, lat, lng)
+	off := d.offset(hour, level, variable, lat, lng)
 	binary.LittleEndian.PutUint32(d.mm[off:off+4], math.Float32bits(val))
 }
 
 // BlitGribData copies one decoded GRIB grid into the dataset, flipping the
 // latitude axis from GRIB's north-to-south scan order to our south-to-north
-// storage order. gribData must be 361*720 = 259920 float64 values.
+// storage order.
 func (d *File) BlitGribData(hourIdx, levelIdx, varIdx int, gribData []float64) error {
-	expected := NumLatitudes * NumLongitudes
+	v := d.variant
+	expected := v.NumLatitudes() * v.NumLongitudes()
 	if len(gribData) != expected {
 		return fmt.Errorf("grib data has %d values, expected %d", len(gribData), expected)
 	}
-	for lat := range NumLatitudes {
-		for lng := range NumLongitudes {
-			gribIdx := (360-lat)*NumLongitudes + lng
+	lats := v.NumLatitudes()
+	lngs := v.NumLongitudes()
+	for lat := range lats {
+		for lng := range lngs {
+			gribIdx := (lats-1-lat)*lngs + lng
 			d.SetVal(hourIdx, levelIdx, varIdx, lat, lng, float32(gribData[gribIdx]))
 		}
 	}