step one

internal/weather/gfs/file.go

@@ -0,0 +1,150 @@
+package gfs
+
+import (
+	"encoding/binary"
+	"fmt"
+	"math"
+	"os"
+	"time"
+
+	mmap "github.com/edsrzf/mmap-go"
+)
+
+// 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).
+type File struct {
+	mm       mmap.MMap
+	file     *os.File
+	writable bool
+	// Epoch is the forecast run time (UTC) the file represents.
+	Epoch time.Time
+}
+
+// Open opens an existing dataset file for reading.
+func Open(path string, epoch time.Time) (*File, error) {
+	f, err := os.Open(path)
+	if err != nil {
+		return nil, fmt.Errorf("open dataset: %w", err)
+	}
+	info, err := f.Stat()
+	if err != nil {
+		f.Close()
+		return nil, fmt.Errorf("stat dataset: %w", err)
+	}
+	if info.Size() != DatasetSize {
+		f.Close()
+		return nil, fmt.Errorf("dataset should be %d bytes (was %d)", 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
+}
+
+// Create creates a new dataset file of the canonical size, mmap'd read-write.
+func Create(path string) (*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 {
+		f.Close()
+		return nil, fmt.Errorf("truncate dataset: %w", err)
+	}
+	mm, err := mmap.MapRegion(f, int(DatasetSize), 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
+}
+
+// OpenWritable opens an existing dataset file for read-write access.
+// Used when resuming a partial download.
+func OpenWritable(path string) (*File, error) {
+	f, err := os.OpenFile(path, os.O_RDWR, 0o644)
+	if err != nil {
+		return nil, fmt.Errorf("open dataset rw: %w", err)
+	}
+	info, err := f.Stat()
+	if err != nil {
+		f.Close()
+		return nil, fmt.Errorf("stat dataset: %w", err)
+	}
+	if info.Size() != DatasetSize {
+		f.Close()
+		return nil, fmt.Errorf("dataset should be %d bytes (was %d)", DatasetSize, info.Size())
+	}
+	mm, err := mmap.MapRegion(f, int(DatasetSize), 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
+}
+
+// offset returns the byte offset of the [hour][level][variable][lat][lng] cell.
+func offset(hour, level, variable, lat, lng int) int64 {
+	idx := int64(hour)
+	idx = idx*int64(NumLevels) + int64(level)
+	idx = idx*int64(NumVariables) + int64(variable)
+	idx = idx*int64(NumLatitudes) + int64(lat)
+	idx = idx*int64(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)
+	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)
+	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.
+func (d *File) BlitGribData(hourIdx, levelIdx, varIdx int, gribData []float64) error {
+	expected := NumLatitudes * 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
+			d.SetVal(hourIdx, levelIdx, varIdx, lat, lng, float32(gribData[gribIdx]))
+		}
+	}
+	return nil
+}
+
+// Flush flushes the mmap to disk.
+func (d *File) Flush() error {
+	if d.mm != nil {
+		return d.mm.Flush()
+	}
+	return nil
+}
+
+// Close unmaps and closes the file.
+func (d *File) Close() error {
+	if d.mm != nil {
+		if err := d.mm.Unmap(); err != nil {
+			d.file.Close()
+			return fmt.Errorf("unmap: %w", err)
+		}
+		d.mm = nil
+	}
+	if d.file != nil {
+		err := d.file.Close()
+		d.file = nil
+		return err
+	}
+	return nil
+}