How to Turn Model Predictions Back into a Raster

Problem statement

A model produces a one-dimensional array of predictions. A map needs a georeferenced raster. The conversion is mechanical and has four ways to go wrong, all of them silent:

• the predictions get reshaped in the wrong order, producing a scrambled map
• the masked cells are not put back, so the array is the wrong length
• the georeferencing does not match the feature stack
• the applicability mask is not written, so nobody knows where the model was guessing

The last one is not a bug in the code; it is a bug in the product. A prediction raster without a mask invites every reader to trust the whole map equally.

The examples run on a real model: a random forest predicting Sentinel-2 NDVI across a 10 km window on Snowdon from five terrain features derived from the Copernicus DEM, trained on 3,000 cells of the western half.

Quick answer

Predict only the valid cells, then scatter them back into a full-size array:

import numpy as np
import rasterio


def predict_to_raster(model, feature_stack, valid, profile, out_path):
    """feature_stack: (bands, height, width); valid: (height, width) bool."""
    height, width = valid.shape
    X = feature_stack[:, valid].T                     # (n_valid, n_features)

    predicted = np.full(height * width, np.nan, dtype="float32")
    predicted[valid.ravel()] = model.predict(X)
    predicted = predicted.reshape(height, width)

    profile = dict(profile) | {"count": 1, "dtype": "float32",
                               "nodata": np.nan, "compress": "deflate"}
    with rasterio.open(out_path, "w", **profile) as dst:
        dst.write(predicted, 1)
    return predicted


valid = np.isfinite(stack).all(axis=0)
predicted = predict_to_raster(model_array, stack, valid, profile, "ndvi_predicted.tif")
with rasterio.open("ndvi_predicted.tif") as src:
    print(src.shape, src.crs.to_epsg(), src.transform == profile["transform"],
          f"{np.isfinite(src.read(1)).mean():.1%} of cells predicted")

(509, 543) 32630 True 100.0% of cells predicted

feature_stack[:, valid].T and predicted[valid.ravel()] use the same ordering, which is what keeps the round trip correct.

Same feature stack, valid cells, same mask, same profile

1. Build the feature stack exactly as at training time

The band order must match the training feature order, or the model receives elevation where it expects slope and produces a plausible, wrong map.

scikit-learn stores the feature names — but only when the model was fitted on a DataFrame. Fitted on a plain array, it has none:

print("fitted on a DataFrame:", list(model.feature_names_in_))
print("fitted on an array:   ", hasattr(model_array, "feature_names_in_"))

fitted on a DataFrame: ['elevation', 'slope', 'northness', 'eastness', 'tpi']
fitted on an array:    False

When the names are there, rebuild the stack from them rather than from a list maintained by hand:

feature_names = list(model.feature_names_in_)
stack = np.stack([layers[name] for name in feature_names])

With a DataFrame, scikit-learn also checks the order at prediction time — but only if you give it a DataFrame again. An array in the wrong order gets a warning at most:

import warnings
import pandas as pd

sample = pd.DataFrame(stack[:, valid][:, :5].T, columns=feature_names)
try:
    model.predict(sample[feature_names[::-1]])
except ValueError as err:
    print("DataFrame, reversed columns:", str(err).splitlines()[0])

with warnings.catch_warnings(record=True) as caught:
    warnings.simplefilter("always")
    model.predict(sample[feature_names[::-1]].to_numpy())
print("array, reversed columns:", [str(w.message) for w in caught])

DataFrame, reversed columns: The feature names should match those that were passed during fit.
array, reversed columns: ['X does not have valid feature names, but RandomForestRegressor was fitted with feature names']

The warning is the same one a correctly ordered array gets, so it tells you nothing about the order.

2. Predict only the valid cells

What happens to a NaN depends on the model. A linear model raises; scikit-learn's random forest, since version 1.4, accepts it and returns a prediction:

from sklearn.linear_model import LinearRegression

row = stack[:, valid][:, :1].T.copy()
row[0, 0] = np.nan                                   # elevation missing
print("random forest:", model_array.predict(row))
try:
    LinearRegression().fit(stack[:, valid][:, :500].T, ndvi[valid][:500]).predict(row)
except ValueError as err:
    print("linear model:", str(err).splitlines()[0])

random forest: [0.75292033]
linear model: Input X contains NaN.

A forest that sees NaN routes the cell down whichever branch its training data sent missing values to — and here the training data had none, so that is an arbitrary side of every split. It is a number, not an error, which is worse. Mask first:

valid = np.isfinite(stack).all(axis=0)
X = stack[:, valid].T
print(f"{valid.mean():.1%} of cells have every feature")

100.0% of cells have every feature

The valid fraction is worth printing. If one feature layer has large gaps, the prediction map inherits them, and it is better to know before writing the file.

3. Scatter the predictions back with the same mask

out = np.full(valid.size, np.nan, dtype="float32")
out[valid.ravel()] = model_array.predict(X)
out = out.reshape(valid.shape)

Both the extraction and the scatter use C-order flattening of the same boolean array, so the correspondence is exact. np.nonzero(valid) yields cells in the same C order, so index arrays from it are just as safe. What scrambles a map is a reshape that disagrees about the axes:

rows, cols = np.nonzero(valid)
print("np.nonzero order equals boolean order:", np.array_equal(stack[:, rows, cols], stack[:, valid]))

height, width = valid.shape
wrong_features = stack.T.reshape(-1, stack.shape[0])        # (width, height, bands) flattened
right_features = stack.reshape(stack.shape[0], -1).T
wrong = model_array.predict(np.nan_to_num(wrong_features)).reshape(height, width)
right = model_array.predict(np.nan_to_num(right_features)).reshape(height, width)
ok = target_ok
print(f"correlation with observed NDVI: right {np.corrcoef(right[ok], ndvi[ok])[0, 1]:.3f}, "
      f"axes swapped {np.corrcoef(wrong[ok], ndvi[ok])[0, 1]:.3f}")

np.nonzero order equals boolean order: True
correlation with observed NDVI: right 0.468, axes swapped -0.035

Transposing the stack before flattening feeds every cell's features to a different cell, and nothing complains, because the array sizes still agree.

4. Copy the georeferencing from the feature stack

The output must share the transform, CRS and shape of the layers the features came from. Copying the profile from one of them and overriding only count, dtype and nodata is the safe pattern — the quick answer's check compares the written transform with the source's.

5. Write the applicability mask as a second band

from scipy.spatial import cKDTree
from sklearn.preprocessing import StandardScaler

train = X_train.to_numpy()
scaler = StandardScaler().fit(train)
tree = cKDTree(scaler.transform(train))
threshold = np.quantile(tree.query(scaler.transform(train), k=2)[0][:, 1], 0.95)

def applicability(X):
    return tree.query(scaler.transform(X), k=1)[0] <= threshold

inside = applicability(X)
print(f"{inside.mean():.1%} of predicted cells are inside the area of applicability")

85.4% of predicted cells are inside the area of applicability

A prediction is only meaningful where the covariates resemble the training data. Shipping the mask alongside costs one band and is what stops the map being over-read.

Masked prediction, out-of-memory rasters and the round trip

The data used on this page

A 10 km window on Snowdon, 4.16–4.00° W, 53.03–53.12° N, on the 20 m grid of the least cloudy Sentinel-2 L2A scene of June–September 2025 from Earth Search. The target is that scene's NDVI; the features are elevation, slope, northness, eastness and a 300 m topographic position index from the Copernicus DEM GLO-30. layers maps each feature name to its grid, stack stacks them in training order, and profile is the grid's rasterio profile. model is a random forest fitted on a DataFrame of 3,000 western cells (X_train), model_array the same forest fitted on the plain array. Each layer is also written to a GeoTIFF for Example 2.

import numpy as np
import pandas as pd
import rasterio
import rasterio.merge
from pystac_client import Client
from rasterio.enums import Resampling
from rasterio.warp import reproject, transform_bounds
from scipy import ndimage
from sklearn.ensemble import RandomForestRegressor

AOI = (-4.16, 53.03, -4.00, 53.12)
items = Client.open("https://earth-search.aws.element84.com/v1").search(
    collections=["sentinel-2-l2a"], bbox=list(AOI), datetime="2025-06-01/2025-10-01",
    query={"eo:cloud_cover": {"lt": 60}}).items()
item = min(items, key=lambda i: i.properties["eo:cloud_cover"])

with rasterio.open(item.assets["scl"].href) as ds:                 # the 20 m grid
    win = rasterio.windows.from_bounds(*transform_bounds("EPSG:4326", ds.crs, *AOI),
                                       transform=ds.transform).round_offsets().round_lengths()
    scl, transform, crs = ds.read(1, window=win), ds.window_transform(win), ds.crs
    bounds = rasterio.windows.bounds(win, ds.transform)
shape = scl.shape
band = {}
for name in ("red", "nir"):
    with rasterio.open(item.assets[name].href) as ds:
        band[name] = ds.read(1, window=rasterio.windows.from_bounds(*bounds, transform=ds.transform),
                             out_shape=shape, resampling=Resampling.average).astype("float32") / 1e4
ndvi = (band["nir"] - band["red"]) / (band["nir"] + band["red"])

tiles = [rasterio.open(f"https://copernicus-dem-30m.s3.amazonaws.com/Copernicus_DSM_COG_10_N53_00_{t}_00_DEM/"
                       f"Copernicus_DSM_COG_10_N53_00_{t}_00_DEM.tif") for t in ("W005", "W004")]
mosaic, mosaic_transform = rasterio.merge.merge(tiles, bounds=(-4.18, 53.02, -3.98, 53.13))
elev = np.full(shape, np.nan, dtype="float32")
reproject(mosaic[0], elev, src_transform=mosaic_transform, src_crs=tiles[0].crs,
          dst_transform=transform, dst_crs=crs, resampling=Resampling.bilinear)
dz_dy, dz_dx = np.gradient(elev, -20, 20)
aspect = np.arctan2(-dz_dx, -dz_dy)
layers = {
    "elevation": elev,
    "slope": np.degrees(np.arctan(np.hypot(dz_dx, dz_dy))).astype("float32"),
    "northness": np.cos(aspect).astype("float32"),
    "eastness": np.sin(aspect).astype("float32"),
    "tpi": (elev - ndimage.uniform_filter(elev, size=15)).astype("float32"),
}
feature_names = list(layers)
stack = np.stack([layers[name] for name in feature_names])
target_ok = np.isin(scl, (2, 4, 5, 7)) & np.isfinite(ndvi)

west = np.tile(np.arange(shape[1]) < shape[1] // 2, (shape[0], 1))
rows, cols = np.nonzero(target_ok & west)
pick = np.random.default_rng(0).choice(len(rows), 3000, replace=False)
X_train = pd.DataFrame(stack[:, rows[pick], cols[pick]].T, columns=feature_names)
y_train = ndvi[rows[pick], cols[pick]]
model = RandomForestRegressor(n_estimators=200, min_samples_leaf=2, random_state=0, n_jobs=4).fit(X_train, y_train)
model_array = RandomForestRegressor(n_estimators=200, min_samples_leaf=2, random_state=0, n_jobs=4).fit(
    X_train.to_numpy(), y_train)

profile = {"driver": "GTiff", "height": shape[0], "width": shape[1], "count": 1, "dtype": "float32",
           "crs": crs, "transform": transform, "nodata": np.nan}
layer_paths = {}
for name, grid in layers.items():
    layer_paths[name] = f"{name}.tif"
    with rasterio.open(layer_paths[name], "w", **profile) as dst:
        dst.write(grid, 1)
print(f"{item.id}: {shape[1]} x {shape[0]} cells of 20 m, {len(X_train):,} training cells")

S2C_30UVD_20250712_0_L2A: 543 x 509 cells of 20 m, 3,000 training cells

Example 1 — prediction with the mask and the uncertainty

import numpy as np
import rasterio


def write_prediction(model, layers, feature_names, profile, out_path,
                     applicability=None, uncertainty=None):
    """Prediction, applicability mask and spread, as one multi-band raster."""
    missing = [n for n in feature_names if n not in layers]
    if missing:
        raise KeyError(f"feature layers missing: {missing}")

    stack = np.stack([layers[name].astype("float32") for name in feature_names])
    height, width = stack.shape[1:]
    valid = np.isfinite(stack).all(axis=0)
    print(f"  {valid.mean():.1%} of {height * width:,} cells have every feature")

    predicted = np.full(height * width, np.nan, dtype="float32")
    X = stack[:, valid].T
    predicted[valid.ravel()] = model.predict(X)
    bands = [predicted.reshape(height, width)]
    descriptions = ["prediction"]

    if applicability is not None:
        inside = np.full(height * width, np.nan, dtype="float32")
        inside[valid.ravel()] = applicability(X).astype("float32")
        bands.append(inside.reshape(height, width))
        descriptions.append("1 = inside the area of applicability")

    if uncertainty is not None:
        spread = np.full(height * width, np.nan, dtype="float32")
        spread[valid.ravel()] = uncertainty(X)
        bands.append(spread.reshape(height, width))
        descriptions.append("prediction spread across ensemble members")

    profile = dict(profile) | {
        "count": len(bands), "dtype": "float32", "nodata": np.nan,
        "compress": "deflate", "tiled": True, "blockxsize": 512,
        "blockysize": 512,
    }
    with rasterio.open(out_path, "w", **profile) as dst:
        for i, (band, description) in enumerate(zip(bands, descriptions), 1):
            dst.write(band, i)
            dst.set_band_description(i, description)
        dst.update_tags(features=",".join(feature_names),
                        model=type(model).__name__,
                        note="band 1 is supported only where band 2 is 1")
        dst.build_overviews([2, 4, 8, 16])

    print(f"  wrote {out_path} with {len(bands)} bands")
    return out_path


def tree_spread(X):
    return np.std([tree.predict(X) for tree in model_array.estimators_], axis=0)

write_prediction(model_array, layers, feature_names, profile, "ndvi_prediction.tif",
                 applicability=applicability, uncertainty=tree_spread)
with rasterio.open("ndvi_prediction.tif") as src:
    print("  ", src.descriptions)
    print("  ", src.tags())
    print("   overviews:", src.overviews(1))
    band1, band2, band3 = src.read()
    print(f"   mean spread inside the mask {np.nanmean(band3[band2 == 1]):.3f}, "
          f"outside {np.nanmean(band3[band2 == 0]):.3f}")

  100.0% of 276,387 cells have every feature
  wrote ndvi_prediction.tif with 3 bands
   ('prediction', '1 = inside the area of applicability', 'prediction spread across ensemble members')
   {'features': 'elevation,slope,northness,eastness,tpi', 'model': 'RandomForestRegressor', 'note': 'band 1 is supported only where band 2 is 1', 'AREA_OR_POINT': 'Area'}
   overviews: [2, 4, 8, 16]
   mean spread inside the mask 0.102, outside 0.144

Writing the feature list into the file's tags means a prediction raster carries its own provenance. Six months later, that is the difference between a reproducible product and a mystery.

Example 2 — predicting a raster larger than memory

import numpy as np
import rasterio
from rasterio.windows import Window


def predict_windowed(model, layer_paths, feature_names, out_path,
                     block=128):
    """Read, predict and write one block at a time."""
    with rasterio.open(layer_paths[feature_names[0]]) as src:
        profile = src.profile
        height, width = src.height, src.width

    profile.update(count=1, dtype="float32", nodata=np.nan,
                   compress="deflate", tiled=True,
                   blockxsize=block, blockysize=block)

    sources = {name: rasterio.open(path) for name, path in layer_paths.items()}
    try:
        with rasterio.open(out_path, "w", **profile) as dst:
            for row in range(0, height, block):
                for col in range(0, width, block):
                    window = Window(col, row,
                                    min(block, width - col),
                                    min(block, height - row))
                    stack = np.stack([
                        sources[name].read(1, window=window).astype("float32")
                        for name in feature_names])
                    valid = np.isfinite(stack).all(axis=0)

                    out = np.full(valid.shape, np.nan, dtype="float32")
                    if valid.any():
                        out[valid] = model.predict(stack[:, valid].T)
                    dst.write(out, 1, window=window)
            dst.build_overviews([2, 4, 8, 16])
    finally:
        for src in sources.values():
            src.close()

    print(f"  wrote {out_path} in {block}x{block} blocks")
    return out_path


predict_windowed(model_array, layer_paths, feature_names, "ndvi_windowed.tif")
with rasterio.open("ndvi_windowed.tif") as src:
    windowed = src.read(1)
print(f"  largest difference from the whole-raster prediction: {np.nanmax(np.abs(windowed - predicted)):.2e}")

  wrote ndvi_windowed.tif in 128x128 blocks
  largest difference from the whole-raster prediction: 0.00e+00

Blocks make memory constant regardless of raster size. It matches the whole-raster prediction exactly here because the TPI layer was computed on the whole grid before it was written; a feature computed from a neighbourhood inside the loop would need an overlapping read, or the block edges come out wrong — see prediction raster artefacts.

Example 3 — checking the round trip

import numpy as np


def verify_round_trip(model, stack, valid, predicted_raster, n=200, seed=0):
    """Predict a few cells directly and compare with the written raster."""
    rng = np.random.default_rng(seed)
    rows, cols = np.nonzero(valid)
    pick = rng.choice(len(rows), size=min(n, len(rows)), replace=False)

    X = stack[:, rows[pick], cols[pick]].T
    direct = model.predict(X)
    from_raster = predicted_raster[rows[pick], cols[pick]]

    difference = np.abs(direct - from_raster)
    print(f"  {len(pick)} sampled cells, max |difference| "
          f"{np.nanmax(difference):.3e}")
    if np.nanmax(difference) > 1e-4:
        print("  ! the raster does not match direct prediction — the reshape "
              "or the mask ordering is wrong")
    return difference


print("correct raster:");         verify_round_trip(model_array, stack, valid, predicted)
print("axes-swapped raster:");    verify_round_trip(model_array, stack, valid, wrong)

correct raster:
  200 sampled cells, max |difference| 2.959e-08
axes-swapped raster:
  200 sampled cells, max |difference| 7.455e-01
  ! the raster does not match direct prediction — the reshape or the mask ordering is wrong

This is the check that catches a scrambled map. A reordering produces a raster that looks like plausible output — smooth, in range, spatially structured — and only a direct comparison reveals it.

Reshape order and feature order

Why the reshape order matters

NumPy flattens in C order by default: the last axis varies fastest. stack[:, valid] returns columns in the order np.nonzero(valid) produces them, which is also C order.

So predicted[valid.ravel()] = values puts each prediction back where its features came from, exactly — and so does indexing with the arrays np.nonzero returns.

Mixing conventions breaks it. Flattening a transposed stack, or one array in Fortran order and the other in C order, produces a map that is spatially structured and wrong — because the features themselves are spatially structured, so a permuted assignment still looks like terrain.

Why the feature order is the most dangerous mismatch

A model given features in the wrong order produces predictions without complaint. Elevation values arrive where slope was expected, and the model applies the splits it learned for slope.

The output is in range and spatially smooth. Nothing raises, nothing looks obviously wrong, and the map is meaningless.

Scikit-learn's feature_names_in_ exists for this, when the model was fitted on a DataFrame. Rebuilding the stack from it, and raising when a layer is missing, removes the whole failure mode; predicting from a DataFrame lets scikit-learn check the order for you.

Why the mask has to be a band

A single-band prediction raster is used by people who did not build it. They will not read the documentation, and they cannot see which pixels the model was extrapolating.

A second band, with a description and a file tag explaining it, travels with the data through every copy and reprojection. It costs one band and it is the difference between a product and a liability.

The same argument applies to any per-pixel uncertainty the model can produce — the spread across ensemble members for a forest, the predictive variance for a Gaussian process.

Why windowed prediction needs care with neighbourhood features

Predicting block by block keeps memory constant, and it breaks any feature computed from a neighbourhood inside the loop. A moving mean at a block edge sees only the cells inside the block.

The fix is the same as for any tiled raster processing: read each block with a buffer larger than half the largest neighbourhood, compute, and write only the unbuffered part.

Features sampled cell by cell — elevation, reflectance, distance rasters, or a neighbourhood feature already written as its own raster — need no buffer, which is why building the feature stack as rasters first and predicting from them is simpler than computing features inside the prediction loop.

Stored feature names, buffered blocks and overviews

• Rebuild the feature stack from the model's stored feature names — which exist only if it was fitted on a DataFrame.
• Extract and scatter with the same boolean mask in the same order.
• Predict only complete cells. A linear model raises on NaN; a random forest in scikit-learn 1.4 or later returns a number for it.
• Copy the profile from a feature layer, overriding only count, dtype and nodata.
• Write the applicability mask as a band, with a description and a tag.
• Verify the round trip on a sample of cells.
• Buffer blocks if any feature is computed from a neighbourhood inside the loop.
• Build overviews so the map is usable in a viewer.

Internal links

• How to build spatial features for a machine learning model — building the stack this consumes
• Extrapolation in space explained: the area of applicability — the mask to write
• My prediction raster is striped, blocky or full of NoData — diagnosing a bad output
• How to evaluate a spatial model honestly — what to report with the map
• Spatial machine learning explained — the wider context
• How to write a cloud-optimised GeoTIFF in Python — the output format
• How to classify land cover from satellite imagery in Python — the classification case
• How to batch process rasters with rasterio in Python — the windowed pattern

FAQ

How do I turn model predictions into a raster?

Extract the valid cells with a boolean mask, predict them, scatter the results back with the same mask, reshape, and write with the feature stack's georeferencing.

Why is my prediction map scrambled?

The extraction and the scatter used different orderings. Use the same boolean mask for both, and verify by predicting a sample of cells directly.

How do I make sure the feature order is right?

Fit on a DataFrame so the model has feature_names_in_, rebuild the stack from it, and raise when a layer is missing. A reordered array produces a plausible, wrong map with no error; a reordered DataFrame is refused.

What do I do about NaN cells?

Predict only cells where every feature is finite, and write NaN elsewhere. A linear model raises on NaN, but scikit-learn's random forest (1.4 and later) quietly returns a prediction for it, and a fill value produces confident nonsense in any model.

Should I write an uncertainty band?

Yes, if the model can produce one, and always write the applicability mask. A single-band prediction invites readers to trust the whole map equally.

How do I predict a raster larger than memory?

Block by block. Buffer the reads if any feature uses a neighbourhood, and write only the unbuffered part.

How do I check the raster is correct?

Predict a sample of cells directly from the feature stack and compare against the written raster. A reordering is invisible any other way.