predictor/internal/weather/gfs/file.go

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 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
	// Epoch is the forecast run time (UTC) the file represents.
	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, variant *Variant, 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() != variant.DatasetSize() {
		f.Close()
		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{variant: variant, mm: mm, file: f, writable: false, Epoch: epoch}, nil
}

// 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)
	}
	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(size), mmap.RDWR, 0, 0)
	if err != nil {
		f.Close()
		return nil, fmt.Errorf("mmap dataset: %w", err)
	}
	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, 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)
	}
	info, err := f.Stat()
	if err != nil {
		f.Close()
		return nil, fmt.Errorf("stat dataset: %w", err)
	}
	if info.Size() != variant.DatasetSize() {
		f.Close()
		return nil, fmt.Errorf("dataset should be %d bytes (was %d)", variant.DatasetSize(), info.Size())
	}
	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{variant: variant, mm: mm, file: f, writable: true}, nil
}

// offset returns the byte offset of the [hour][level][variable][lat][lng] cell.
func (d *File) offset(hour, level, variable, lat, lng int) int64 {
	v := d.variant
	idx := int64(hour)
	idx = idx*int64(v.NumLevels()) + int64(level)
	idx = idx*int64(NumVariables) + int64(variable)
	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 := d.offset(hour, level, variable, lat, lng)
	return math.Float32frombits(binary.LittleEndian.Uint32(d.mm[off : off+4]))
}

// ValByElem reads the float32 at a precomputed flat element index (not a byte
// offset). The wind sampler uses this to read the eight interpolation corners
// after computing their flat indices once via cube strides.
func (d *File) ValByElem(elem int64) float32 {
	off := elem * ElementSize
	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 := 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.
func (d *File) BlitGribData(hourIdx, levelIdx, varIdx int, gribData []float64) error {
	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)
	}
	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]))
		}
	}
	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
}