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 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
}