⚙️ How-To Spatial machine learning Raster & rasterio Spatial Machine Learning · step 9 of 14 Updated

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.

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

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

Valid cells extracted into a feature matrix, predicted, and scattered back into a full-size array with NaN elsewhere.
Extract with a boolean mask, scatter back with the same mask. Any other route risks a reordering.

Step-by-step solution

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.

feature_names = model.feature_names_in_          # scikit-learn stores them
stack = np.stack([layers[name] for name in feature_names])

Storing the feature names with the model and rebuilding the stack from them is the only reliable defence. A hand-maintained list in two files drifts.

2. Predict only the valid cells

Feeding NaN to most models raises; feeding a nodata fill produces confident nonsense.

valid = np.isfinite(stack).all(axis=0)
X = stack[:, valid].T
print(f"{valid.mean():.1%} 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()] = predictions
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. Using np.where indices for one and a boolean for the other is where scrambling creeps in.

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.

5. Write the applicability mask as a second band 

dst.write(prediction, 1)
dst.write(inside.astype("float32"), 2)

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.

A feature stack assembled in a different order from training, producing a plausible but wrong prediction map with no error.
A reordered feature stack produces a map, not an exception. Rebuild it from the stored feature names.

Code examples

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, chunk_rows=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

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=512):
    """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

Blocks make memory constant regardless of raster size. Note that any feature needing a neighbourhood — a moving-window mean — must be computed with an overlapping read, or the block edges are wrong.

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

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.

Explanation

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.

Mixing conventions breaks it. Using np.where for extraction and a boolean for the scatter, or flattening one in Fortran order, produces a map that is spatially structured and completely 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. Rebuilding the stack from it, and raising when a layer is missing, removes the whole failure mode.

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. A 27-cell 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 the largest neighbourhood, compute, and write only the unbuffered part.

Features sampled cell by cell — elevation, reflectance, distance rasters — 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.

Five checks verifying a prediction raster against direct prediction on sampled cells.
Because the inputs are spatially structured, a permuted assignment still looks like terrain.

Edge cases or notes

  • Rebuild the feature stack from the model's stored feature names.
  • Extract and scatter with the same boolean mask in the same order.
  • Predict only complete cells; NaN raises and fill values produce nonsense.
  • 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 uses a neighbourhood.
  • Build overviews so the map is usable in a viewer.

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?

Rebuild the stack from the model's feature_names_in_ and raise when a layer is missing. A reordered stack produces a plausible, wrong map with no error.

What do I do about NaN cells?

Predict only cells where every feature is finite, and write NaN elsewhere. Most models raise on NaN, and a fill value produces confident nonsense.

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.