forked from gsn/predictor
140 lines
3.9 KiB
Go
140 lines
3.9 KiB
Go
package dataset
|
|
|
|
import (
|
|
"encoding/binary"
|
|
"fmt"
|
|
"math"
|
|
"os"
|
|
"time"
|
|
|
|
mmap "github.com/edsrzf/mmap-go"
|
|
)
|
|
|
|
// File represents an mmap-backed wind dataset file.
|
|
type File struct {
|
|
mm mmap.MMap
|
|
file *os.File
|
|
writable bool
|
|
DSTime time.Time // forecast run time (UTC)
|
|
}
|
|
|
|
// Open opens an existing dataset file for reading.
|
|
func Open(path string, dsTime 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, DSTime: dsTime}, nil
|
|
}
|
|
|
|
// Create creates a new dataset file for writing.
|
|
// The file is truncated to the correct size and 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
|
|
}
|
|
|
|
// offset computes the byte offset for element [hour][level][variable][lat][lon].
|
|
// Row-major C-order indexing matching the reference implementation:
|
|
// shape = (65, 47, 3, 361, 720)
|
|
func offset(hour, level, variable, lat, lon 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(lon)
|
|
return idx * int64(ElementSize)
|
|
}
|
|
|
|
// Val reads a float32 value from the dataset at [hour][level][variable][lat][lon].
|
|
func (d *File) Val(hour, level, variable, lat, lon int) float32 {
|
|
off := offset(hour, level, variable, lat, lon)
|
|
bits := binary.LittleEndian.Uint32(d.mm[off : off+4])
|
|
return math.Float32frombits(bits)
|
|
}
|
|
|
|
// SetVal writes a float32 value to the dataset at [hour][level][variable][lat][lon].
|
|
// Only valid on writable (created) datasets.
|
|
func (d *File) SetVal(hour, level, variable, lat, lon int, val float32) {
|
|
off := offset(hour, level, variable, lat, lon)
|
|
binary.LittleEndian.PutUint32(d.mm[off:off+4], math.Float32bits(val))
|
|
}
|
|
|
|
// BlitGribData copies decoded GRIB grid data into the dataset at the given position.
|
|
// gribData is 361*720 float64 values in GRIB scan order (north-to-south, west-to-east).
|
|
// This function flips the latitude so that dataset index 0 = -90 (south) and 360 = +90 (north).
|
|
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 := 0; lat < NumLatitudes; lat++ {
|
|
for lon := 0; lon < NumLongitudes; lon++ {
|
|
// GRIB scans north-to-south: row 0 = 90N, row 360 = 90S
|
|
// Dataset stores south-to-north: index 0 = -90 (90S), index 360 = +90 (90N)
|
|
gribIdx := (360-lat)*NumLongitudes + lon
|
|
val := float32(gribData[gribIdx])
|
|
d.SetVal(hourIdx, levelIdx, varIdx, lat, lon, val)
|
|
}
|
|
}
|
|
|
|
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 dataset 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
|
|
}
|