feat(landscape): window executor Compute with worker pool and grid checks
Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
This commit is contained in:
167
internal/landscape/compute.go
Normal file
167
internal/landscape/compute.go
Normal file
@@ -0,0 +1,167 @@
|
||||
package landscape
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"path/filepath"
|
||||
"runtime"
|
||||
"sync"
|
||||
|
||||
"gis/internal/raster"
|
||||
|
||||
"golang.org/x/sync/errgroup"
|
||||
)
|
||||
|
||||
// Compute runs one node over its input rasters. See the package doc for the
|
||||
// concurrency model: windows are processed by a worker pool; each worker owns
|
||||
// its own Reader handles (GDAL datasets are not goroutine-safe) and writes are
|
||||
// serialised per output. Window order does not affect the result.
|
||||
func Compute(ctx context.Context, rio raster.IO, node Node, inputs map[string]string, outDir string, p Params) (map[string]string, error) {
|
||||
p = p.WithDefaults()
|
||||
paths := make([]string, len(node.Inputs))
|
||||
for i, g := range node.Inputs {
|
||||
path, ok := inputs[g]
|
||||
if !ok || path == "" {
|
||||
return nil, fmt.Errorf("node %s (%s): missing input %q", node.Code, node.Name, g)
|
||||
}
|
||||
paths[i] = path
|
||||
}
|
||||
|
||||
// Validate grids with one set of handles before creating anything.
|
||||
grid, err := checkGrids(ctx, rio, node, paths)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
outPaths := make(map[string]string, len(node.Outputs))
|
||||
writers := make([]raster.Writer, len(node.Outputs))
|
||||
for j, g := range node.Outputs {
|
||||
path := filepath.Join(outDir, g+".tif")
|
||||
w, err := rio.Create(ctx, path, grid, raster.DefaultWriterOptions())
|
||||
if err != nil {
|
||||
closeWriters(writers[:j])
|
||||
return nil, fmt.Errorf("node %s: create output %s: %w", node.Code, g, err)
|
||||
}
|
||||
writers[j] = w
|
||||
outPaths[g] = path
|
||||
}
|
||||
|
||||
windows := raster.Windows(grid.Width, grid.Height, p.WindowSize)
|
||||
jobs := make(chan raster.Window)
|
||||
var writeMu sync.Mutex
|
||||
eg, gctx := errgroup.WithContext(ctx)
|
||||
workers := min(runtime.GOMAXPROCS(0), len(windows))
|
||||
for w := 0; w < workers; w++ {
|
||||
eg.Go(func() error {
|
||||
readers, err := openAll(gctx, rio, paths)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer closeReaders(readers)
|
||||
size := p.WindowSize * p.WindowSize
|
||||
in := make([][]float64, len(paths))
|
||||
for i := range in {
|
||||
in[i] = make([]float64, size)
|
||||
}
|
||||
out := make([][]int16, len(writers))
|
||||
for j := range out {
|
||||
out[j] = make([]int16, size)
|
||||
}
|
||||
for win := range jobs {
|
||||
n := win.Width * win.Height
|
||||
inView := make([][]float64, len(in))
|
||||
for i, r := range readers {
|
||||
if err := r.ReadFloat64(win, in[i][:n]); err != nil {
|
||||
return fmt.Errorf("node %s: read %s window %+v: %w", node.Code, node.Inputs[i], win, err)
|
||||
}
|
||||
inView[i] = in[i][:n]
|
||||
}
|
||||
outView := make([][]int16, len(out))
|
||||
for j := range out {
|
||||
outView[j] = out[j][:n]
|
||||
}
|
||||
node.Kernel(inView, outView, p)
|
||||
writeMu.Lock()
|
||||
for j, wr := range writers {
|
||||
if err := wr.WriteInt16(win, outView[j]); err != nil {
|
||||
writeMu.Unlock()
|
||||
return fmt.Errorf("node %s: write %s window %+v: %w", node.Code, node.Outputs[j], win, err)
|
||||
}
|
||||
}
|
||||
writeMu.Unlock()
|
||||
}
|
||||
return nil
|
||||
})
|
||||
}
|
||||
eg.Go(func() error {
|
||||
defer close(jobs)
|
||||
for _, win := range windows {
|
||||
select {
|
||||
case jobs <- win:
|
||||
case <-gctx.Done():
|
||||
return gctx.Err()
|
||||
}
|
||||
}
|
||||
return nil
|
||||
})
|
||||
runErr := eg.Wait()
|
||||
if err := closeWriters(writers); err != nil && runErr == nil {
|
||||
runErr = err
|
||||
}
|
||||
if runErr != nil {
|
||||
return nil, runErr
|
||||
}
|
||||
return outPaths, nil
|
||||
}
|
||||
|
||||
// checkGrids opens every input once, verifies all grids match the first, and
|
||||
// returns that grid.
|
||||
func checkGrids(ctx context.Context, rio raster.IO, node Node, paths []string) (raster.Grid, error) {
|
||||
readers, err := openAll(ctx, rio, paths)
|
||||
if err != nil {
|
||||
return raster.Grid{}, err
|
||||
}
|
||||
defer closeReaders(readers)
|
||||
ref := readers[0].Grid()
|
||||
for i := 1; i < len(readers); i++ {
|
||||
if g := readers[i].Grid(); !rio.SameGrid(ref, g) {
|
||||
return raster.Grid{}, fmt.Errorf("node %s: input %s grid (%dx%d, %v) differs from %s (%dx%d, %v); inputs must be aligned",
|
||||
node.Code, node.Inputs[i], g.Width, g.Height, g.GeoTransform, node.Inputs[0], ref.Width, ref.Height, ref.GeoTransform)
|
||||
}
|
||||
}
|
||||
return ref, nil
|
||||
}
|
||||
|
||||
func openAll(ctx context.Context, rio raster.IO, paths []string) ([]raster.Reader, error) {
|
||||
readers := make([]raster.Reader, 0, len(paths))
|
||||
for _, p := range paths {
|
||||
r, err := rio.Open(ctx, p)
|
||||
if err != nil {
|
||||
closeReaders(readers)
|
||||
return nil, err
|
||||
}
|
||||
readers = append(readers, r)
|
||||
}
|
||||
return readers, nil
|
||||
}
|
||||
|
||||
func closeReaders(rs []raster.Reader) {
|
||||
for _, r := range rs {
|
||||
if r != nil {
|
||||
_ = r.Close()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func closeWriters(ws []raster.Writer) error {
|
||||
var first error
|
||||
for _, w := range ws {
|
||||
if w == nil {
|
||||
continue
|
||||
}
|
||||
if err := w.Close(); err != nil && first == nil {
|
||||
first = err
|
||||
}
|
||||
}
|
||||
return first
|
||||
}
|
||||
101
internal/landscape/compute_test.go
Normal file
101
internal/landscape/compute_test.go
Normal file
@@ -0,0 +1,101 @@
|
||||
package landscape
|
||||
|
||||
import (
|
||||
"context"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
|
||||
"gis/internal/raster"
|
||||
)
|
||||
|
||||
func grid(w, h int) raster.Grid {
|
||||
return raster.Grid{Width: w, Height: h, GeoTransform: [6]float64{0, 30, 0, 0, 0, -30}, ProjectionWKT: "EPSG:32642"}
|
||||
}
|
||||
|
||||
func constRaster(w, h int, v float64) []float64 {
|
||||
d := make([]float64, w*h)
|
||||
for i := range d {
|
||||
d[i] = v
|
||||
}
|
||||
return d
|
||||
}
|
||||
|
||||
func TestCompute_E_TilesWindowsAndWritesEveryPixel(t *testing.T) {
|
||||
fio := newFakeIO()
|
||||
g := grid(10, 7)
|
||||
z := constRaster(10, 7, 2)
|
||||
u := constRaster(10, 7, 4)
|
||||
z[0] = -9999 // pixel 0 invalid
|
||||
fio.put("z.tif", g, z)
|
||||
fio.put("u.tif", g, u)
|
||||
node, _ := NodeByCode("10")
|
||||
p := DefaultParams()
|
||||
p.WindowSize = 4 // forces partial windows (10x7 with 4x4 tiles)
|
||||
out, err := Compute(context.Background(), fio, node, map[string]string{OutZIntegrated: "z.tif", OutUIntegrated: "u.tif"}, "/out", p)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
path := out[OutE]
|
||||
if path != filepath.Join("/out", OutE+".tif") {
|
||||
t.Fatalf("unexpected output path %q", path)
|
||||
}
|
||||
res := fio.files[path]
|
||||
if res.grid != g {
|
||||
t.Fatal("output grid must copy the first input's grid")
|
||||
}
|
||||
for i, v := range res.data {
|
||||
want := 4.0
|
||||
if i == 0 {
|
||||
want = -9999
|
||||
}
|
||||
if v != want {
|
||||
t.Fatalf("pixel %d = %v want %v", i, v, want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestCompute_FailsOnGridMismatch(t *testing.T) {
|
||||
fio := newFakeIO()
|
||||
fio.put("z.tif", grid(10, 7), constRaster(10, 7, 2))
|
||||
fio.put("u.tif", grid(9, 7), constRaster(9, 7, 4))
|
||||
node, _ := NodeByCode("10")
|
||||
_, err := Compute(context.Background(), fio, node, map[string]string{OutZIntegrated: "z.tif", OutUIntegrated: "u.tif"}, "/out", DefaultParams())
|
||||
if err == nil {
|
||||
t.Fatal("expected grid mismatch error")
|
||||
}
|
||||
if _, created := fio.files[filepath.Join("/out", OutE+".tif")]; created {
|
||||
t.Fatal("no output must be created on mismatch")
|
||||
}
|
||||
}
|
||||
|
||||
func TestCompute_FailsOnMissingInput(t *testing.T) {
|
||||
fio := newFakeIO()
|
||||
fio.put("z.tif", grid(2, 2), constRaster(2, 2, 2))
|
||||
node, _ := NodeByCode("10")
|
||||
if _, err := Compute(context.Background(), fio, node, map[string]string{OutZIntegrated: "z.tif"}, "/out", DefaultParams()); err == nil {
|
||||
t.Fatal("expected error for missing input group")
|
||||
}
|
||||
}
|
||||
|
||||
func TestCompute_MultiOutputNode(t *testing.T) {
|
||||
fio := newFakeIO()
|
||||
g := grid(3, 1)
|
||||
paths := map[string]string{}
|
||||
for i, s := range []string{SlotPIrrigation, SlotPSnow, SlotPCrop, SlotPPasture, SlotPTourism} {
|
||||
paths[s] = s + ".tif"
|
||||
fio.put(paths[s], g, constRaster(3, 1, float64(i+1)))
|
||||
}
|
||||
paths[OutE] = "e.tif"
|
||||
fio.put("e.tif", g, constRaster(3, 1, 3))
|
||||
node, _ := NodeByCode("20")
|
||||
out, err := Compute(context.Background(), fio, node, paths, "/o", DefaultParams())
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(out) != 5 {
|
||||
t.Fatalf("want 5 outputs, got %d", len(out))
|
||||
}
|
||||
if fio.files[out[OutCTourism]].data[0] != 3 || fio.files[out[OutCIrrigation]].data[0] != 1 {
|
||||
t.Fatal("min(P, E) not applied per output")
|
||||
}
|
||||
}
|
||||
75
internal/landscape/fakeio_test.go
Normal file
75
internal/landscape/fakeio_test.go
Normal file
@@ -0,0 +1,75 @@
|
||||
package landscape
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"sync"
|
||||
|
||||
"gis/internal/raster"
|
||||
)
|
||||
|
||||
// fakeIO keeps rasters in memory keyed by path.
|
||||
type fakeIO struct {
|
||||
mu sync.Mutex
|
||||
files map[string]*fakeRaster
|
||||
}
|
||||
|
||||
type fakeRaster struct {
|
||||
grid raster.Grid
|
||||
data []float64 // row-major
|
||||
}
|
||||
|
||||
func newFakeIO() *fakeIO { return &fakeIO{files: map[string]*fakeRaster{}} }
|
||||
|
||||
func (f *fakeIO) put(path string, g raster.Grid, data []float64) {
|
||||
f.files[path] = &fakeRaster{grid: g, data: data}
|
||||
}
|
||||
|
||||
type fakeReader struct{ r *fakeRaster }
|
||||
|
||||
func (r *fakeReader) Grid() raster.Grid { return r.r.grid }
|
||||
func (r *fakeReader) NoData() (float64, bool) { return -9999, true }
|
||||
func (r *fakeReader) Close() error { return nil }
|
||||
func (r *fakeReader) ReadFloat64(w raster.Window, dst []float64) error {
|
||||
for y := 0; y < w.Height; y++ {
|
||||
copy(dst[y*w.Width:(y+1)*w.Width], r.r.data[(w.Y+y)*r.r.grid.Width+w.X:(w.Y+y)*r.r.grid.Width+w.X+w.Width])
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
type fakeWriter struct {
|
||||
f *fakeIO
|
||||
r *fakeRaster
|
||||
}
|
||||
|
||||
func (w *fakeWriter) WriteInt16(win raster.Window, src []int16) error {
|
||||
w.f.mu.Lock()
|
||||
defer w.f.mu.Unlock()
|
||||
for y := 0; y < win.Height; y++ {
|
||||
for x := 0; x < win.Width; x++ {
|
||||
w.r.data[(win.Y+y)*w.r.grid.Width+win.X+x] = float64(src[y*win.Width+x])
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
func (w *fakeWriter) Close() error { return nil }
|
||||
|
||||
func (f *fakeIO) Open(_ context.Context, path string) (raster.Reader, error) {
|
||||
f.mu.Lock()
|
||||
defer f.mu.Unlock()
|
||||
r, ok := f.files[path]
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("fake: %s not found", path)
|
||||
}
|
||||
return &fakeReader{r: r}, nil
|
||||
}
|
||||
|
||||
func (f *fakeIO) Create(_ context.Context, path string, g raster.Grid, _ raster.WriterOptions) (raster.Writer, error) {
|
||||
f.mu.Lock()
|
||||
defer f.mu.Unlock()
|
||||
r := &fakeRaster{grid: g, data: make([]float64, g.Width*g.Height)}
|
||||
f.files[path] = r
|
||||
return &fakeWriter{f: f, r: r}, nil
|
||||
}
|
||||
|
||||
func (f *fakeIO) SameGrid(a, b raster.Grid) bool { return a == b }
|
||||
Reference in New Issue
Block a user