Structural summaries: Sequential.summary and Graph.summary¶

This notebook walks through the shape inference protocol in xrtoolz — the same idea Keras uses for Model.summary(), but for xrtoolz pipelines.

Three pieces:

1. Signature (pipekit's) — a tiny value object: ordered (dim_name, size | None) pairs plus a dtype. It carries what shape an op expects/produces without holding any data.
2. Operator.compute_output_signature(input_signature) — every operator can declare how it transforms shapes. The base default is shape-preserving; ops that rename, subset, reduce, regrid, etc. override it.
3. Sequential.summary(signature) and Graph.summary(...) — render an end-to-end table of input/output signatures per step, no real data needed.

Below: the basics, four worked pipelines on different tasks, multi-input metric graphs, nested Graph-in-Graph composition, and a couple of common error/edge cases.

Setup¶

Colab-friendly install + imports.

In [1]:

import subprocess
import sys


try:
    import google.colab  # noqa: F401

    IN_COLAB = True
except ImportError:
    IN_COLAB = False

if IN_COLAB:
    subprocess.run(
        [sys.executable, "-m", "pip", "install", "-q", "xrtoolz"],
        check=True,
    )

import subprocess import sys try: import google.colab # noqa: F401 IN_COLAB = True except ImportError: IN_COLAB = False if IN_COLAB: subprocess.run( [sys.executable, "-m", "pip", "install", "-q", "xrtoolz"], check=True, )

In [2]:

import warnings

warnings.filterwarnings("ignore", message=r".*IProgress.*")

import numpy as np
import xarray as xr

# Render xarray reprs as HTML and keep every section collapsed by default —
# readers expand only what they need; rendered notebooks stay tidy on the
# docs site.
_ = xr.set_options(
    display_style="html",
    display_expand_attrs=False,
    display_expand_coords=False,
    display_expand_data=False,
    display_expand_data_vars=False,
    display_expand_groups=False,
    display_expand_indexes=False,
)

from pipekit import Graph, Input, Sequential, Signature
from xrops.anomaly import CalculateClimatology
from xrops.cf import RenameCoords, ValidateCoords
from xrops.geom import SubsetBBox, SubsetTime
from xrops.utils import ReduceDims
from xrops.interpolate.operators import Coarsen, RegridLike, ResampleTime
from xrops.metrics.operators import RMSE


def sig(dtype: str | None = None, **sizes: int | None) -> Signature:
    """A Signature from dim sizes, in the order given."""
    return Signature(dims=tuple(sizes.items()), dtype=dtype)

import warnings warnings.filterwarnings("ignore", message=r".*IProgress.*") import numpy as np import xarray as xr # Render xarray reprs as HTML and keep every section collapsed by default — # readers expand only what they need; rendered notebooks stay tidy on the # docs site. _ = xr.set_options( display_style="html", display_expand_attrs=False, display_expand_coords=False, display_expand_data=False, display_expand_data_vars=False, display_expand_groups=False, display_expand_indexes=False, ) from pipekit import Graph, Input, Sequential, Signature from xrops.anomaly import CalculateClimatology from xrops.cf import RenameCoords, ValidateCoords from xrops.geom import SubsetBBox, SubsetTime from xrops.utils import ReduceDims from xrops.interpolate.operators import Coarsen, RegridLike, ResampleTime from xrops.metrics.operators import RMSE def sig(dtype: str | None = None, **sizes: int | None) -> Signature: """A Signature from dim sizes, in the order given.""" return Signature(dims=tuple(sizes.items()), dtype=dtype)

1. Signature basics¶

A Signature captures dim sizes and a dtype. None means "size not known symbolically" — you'll see this on operators where the output size depends on the data (irregular bbox crops, time-resamples, climatology grouping).

In [3]:

sig_year = sig(time=365, lat=181, lon=360, dtype="float32")
print(sig_year.format())

# Unknown sizes render as `?`:
print(sig(time=None, lat=4).format())

sig_year = sig(time=365, lat=181, lon=360, dtype="float32") print(sig_year.format()) # Unknown sizes render as `?`: print(sig(time=None, lat=4).format())

(time=365, lat=181, lon=360) float32
(time=?, lat=4) ?

Signature is a value object: it's frozen, and its dims are ordered (name, size) pairs, the carrier's axis order. So equality compares dims in order, plus the dtype label.

In [4]:

a = sig(time=12, lat=4, dtype="float32")
b = sig(lat=4, time=12, dtype="float32")  # the same sizes, another axis order

print(a == sig(time=12, lat=4, dtype="float32"), a == b)
print(a.dims)

a = sig(time=12, lat=4, dtype="float32") b = sig(lat=4, time=12, dtype="float32") # the same sizes, another axis order print(a == sig(time=12, lat=4, dtype="float32"), a == b) print(a.dims)

True False
(('time', 12), ('lat', 4))

Helpers replace_dims and drop_dims return new signatures (no in-place mutation), and matches tests the dims against a pattern:

In [5]:

print(sig_year.replace_dims({"lat": None}).format())  # a regrid the size of which we don't know
print(sig_year.drop_dims(("lat", "lon")).format())  # spatial reduction
print(sig_year.matches("time ... lon"), sig_year.matches("... y x"))  # dims patterns

print(sig_year.replace_dims({"lat": None}).format()) # a regrid the size of which we don't know print(sig_year.drop_dims(("lat", "lon")).format()) # spatial reduction print(sig_year.matches("time ... lon"), sig_year.matches("... y x")) # dims patterns

(time=365, lat=?, lon=360) float32
(time=365) float32
True False

2. Sequential pipeline — SSH preprocessing¶

The canonical "open a CMEMS / GLORYS Dataset → validate → crop to a region → resample time → spatially reduce" preprocessing chain. Each summary row shows what shape the step receives and what it emits.

In [6]:

ssh_pipe = Sequential([
    ValidateCoords(),
    SubsetBBox(lon_bnds=(-80, -50), lat_bnds=(30, 45)),  # Gulf Stream
    ResampleTime(freq="1D", method="mean"),
    ReduceDims(op="mean", dim=("lat", "lon")),
])

# Start with a daily 1°×1° year-long CMEMS-like cube:
input_signature = sig(time=365, lat=181, lon=360, dtype="float32")

print(ssh_pipe.summary(input_signature))

ssh_pipe = Sequential([ ValidateCoords(), SubsetBBox(lon_bnds=(-80, -50), lat_bnds=(30, 45)), # Gulf Stream ResampleTime(freq="1D", method="mean"), ReduceDims(op="mean", dim=("lat", "lon")), ]) # Start with a daily 1°×1° year-long CMEMS-like cube: input_signature = sig(time=365, lat=181, lon=360, dtype="float32") print(ssh_pipe.summary(input_signature))

Operator            Input sig                                        Output sig                                     
--------------------------------------------------------------------------------------------------------------------
[0] ValidateCoords  Signature((time=365, lat=181, lon=360) float32)  Signature((time=365, lat=181, lon=360) float32)
[1] SubsetBBox      Signature((time=365, lat=181, lon=360) float32)  Signature((time=365, lat=?, lon=?) float32)    
[2] ResampleTime    Signature((time=365, lat=?, lon=?) float32)      Signature((time=?, lat=?, lon=?) float32)      
[3] ReduceDims      Signature((time=?, lat=?, lon=?) float32)        Signature((time=?) float32)                    

Reading the table:

• ValidateCoords is shape-preserving — no override needed; the default kicks in.
• SubsetBBox knows it touches lat/lon but can't tell how many points land inside the bbox without the actual coords; both axes go to ?.
• ResampleTime similarly fuzzes time.
• ReduceDims(dim=("lat","lon")) drops the spatial axes; only time remains.

3. Sequential pipeline — climatology + anomaly setup¶

A different shape-changing pattern: CalculateClimatology swaps the time axis for a month (or dayofyear/year) axis whose length depends on the group-by — so it's ? symbolically.

In [7]:

clim_pipe = Sequential([
    ValidateCoords(),
    SubsetTime(time_min="2000-01-01", time_max="2010-12-31"),
    CalculateClimatology(freq="month"),
])

print(clim_pipe.summary(sig(time=4017, lat=181, lon=360, dtype="float32")))

clim_pipe = Sequential([ ValidateCoords(), SubsetTime(time_min="2000-01-01", time_max="2010-12-31"), CalculateClimatology(freq="month"), ]) print(clim_pipe.summary(sig(time=4017, lat=181, lon=360, dtype="float32")))

Operator                  Input sig                                         Output sig                                      
----------------------------------------------------------------------------------------------------------------------------
[0] ValidateCoords        Signature((time=4017, lat=181, lon=360) float32)  Signature((time=4017, lat=181, lon=360) float32)
[1] SubsetTime            Signature((time=4017, lat=181, lon=360) float32)  Signature((time=?, lat=181, lon=360) float32)   
[2] CalculateClimatology  Signature((time=?, lat=181, lon=360) float32)     Signature((month=?, lat=181, lon=360) float32)  

4. Sequential pipeline — coordinate reshape + regrid¶

This one mixes a CF rename, a coarsen, and a regrid onto another grid. RegridLike knows the target's exact size, so the output is fully concrete.

In [8]:

target_grid = xr.Dataset(coords={
    "lat": np.linspace(-90, 90, 91),  # 2° lat
    "lon": np.linspace(0, 360, 180, endpoint=False),  # 2° lon
})

regrid_pipe = Sequential([
    RenameCoords({"longitude": "lon", "latitude": "lat"}),
    Coarsen({"lat": 2, "lon": 2}, boundary="trim"),
    RegridLike(target_grid),
])

print(regrid_pipe.summary(sig(time=12, latitude=720, longitude=1440, dtype="float64")))

target_grid = xr.Dataset(coords={ "lat": np.linspace(-90, 90, 91), # 2° lat "lon": np.linspace(0, 360, 180, endpoint=False), # 2° lon }) regrid_pipe = Sequential([ RenameCoords({"longitude": "lon", "latitude": "lat"}), Coarsen({"lat": 2, "lon": 2}, boundary="trim"), RegridLike(target_grid), ]) print(regrid_pipe.summary(sig(time=12, latitude=720, longitude=1440, dtype="float64")))

Operator          Input sig                                                   Output sig                                     
-----------------------------------------------------------------------------------------------------------------------------
[0] RenameCoords  Signature((time=12, latitude=720, longitude=1440) float64)  Signature((time=12, lat=720, lon=1440) float64)
[1] Coarsen       Signature((time=12, lat=720, lon=1440) float64)             Signature((time=12, lat=360, lon=720) float64) 
[2] RegridLike    Signature((time=12, lat=360, lon=720) float64)              Signature((time=12, lat=91, lon=180) float64)  

Notice the rename flows through cleanly — the table tracks dim names, not just sizes.

5. Graph.summary — multi-input metric¶

Sequential is for linear chains. The functional Graph API handles DAGs, including ops with multiple inputs (the canonical case: pixel metrics like RMSE that take pred and ref). The shape protocol propagates through:

In [9]:

pred = Input("pred")
ref = Input("ref")
score = RMSE(variable="ssh", dims=("time", "lat"))(pred, ref)
metric_graph = Graph(inputs={"pred": pred, "ref": ref}, outputs={"score": score})

# Both inputs share the same shape; metric reduces (time, lat) and keeps lon:
shape = sig(time=365, lat=181, lon=360, dtype="float32")
print(metric_graph.summary({"pred": shape, "ref": shape}))

pred = Input("pred") ref = Input("ref") score = RMSE(variable="ssh", dims=("time", "lat"))(pred, ref) metric_graph = Graph(inputs={"pred": pred, "ref": ref}, outputs={"score": score}) # Both inputs share the same shape; metric reduces (time, lat) and keeps lon: shape = sig(time=365, lat=181, lon=360, dtype="float32") print(metric_graph.summary({"pred": shape, "ref": shape}))

Node                 Input sig                                                                   Output sig
-----------------------------------------------------------------------------------------------------------
[0] RMSE(pred, ref)  (time=365, lat=181, lon=360) float32, (time=365, lat=181, lon=360) float32  (lon=360) float64
→ score = [0]                                                                                    (lon=360) float64

output_signatures returns each output's signature — what running the graph on real data would give, minus the data:

In [10]:

print(metric_graph.output_signatures({"pred": shape, "ref": shape}))

print(metric_graph.output_signatures({"pred": shape, "ref": shape}))

{'score': Signature((lon=360) float64)}

What happens with mismatched shapes?¶

A node with several inputs gets its operator's signature only when the inputs agree; otherwise the graph reports it as unknown (?) rather than guessing. The pixel metric itself validates dim names and sizes, so asking it directly says what disagrees — before any xr.apply_ufunc broadcast error at runtime.

In [11]:

mismatched = sig(time=365, lat=181, lon=360, dtype="float32")
shorter = sig(time=90, lat=181, lon=360, dtype="float32")

print(metric_graph.summary({"pred": mismatched, "ref": shorter}))

try:
    RMSE(variable="ssh", dims=("time", "lat")).compute_output_signature((mismatched, shorter))
except ValueError as exc:
    print("caught:", exc)

mismatched = sig(time=365, lat=181, lon=360, dtype="float32") shorter = sig(time=90, lat=181, lon=360, dtype="float32") print(metric_graph.summary({"pred": mismatched, "ref": shorter})) try: RMSE(variable="ssh", dims=("time", "lat")).compute_output_signature((mismatched, shorter)) except ValueError as exc: print("caught:", exc)

Node                 Input sig                                                                  Output sig
----------------------------------------------------------------------------------------------------------
[0] RMSE(pred, ref)  (time=365, lat=181, lon=360) float32, (time=90, lat=181, lon=360) float32  ?
→ score = [0]                                                                                   ?
caught: Metric input signature sizes do not match: {'time': (365, 90)}.

6. Single-input Graph (the convenience form)¶

Graphs with one input/output accept a bare Signature (you don't have to wrap it in a dict). Same UX as Sequential.summary:

In [12]:

x = Input("ds")
y = SubsetBBox(lon_bnds=(0, 30), lat_bnds=(-30, 30))(x)
y = ResampleTime(freq="1MS")(y)
single_graph = Graph(inputs={"ds": x}, outputs={"out": y})

print(single_graph.summary(sig(time=365, lat=181, lon=360, dtype="float32")))

x = Input("ds") y = SubsetBBox(lon_bnds=(0, 30), lat_bnds=(-30, 30))(x) y = ResampleTime(freq="1MS")(y) single_graph = Graph(inputs={"ds": x}, outputs={"out": y}) print(single_graph.summary(sig(time=365, lat=181, lon=360, dtype="float32")))

Node                   Input sig                             Output sig
-----------------------------------------------------------------------
[0] SubsetBBox(ds)     (time=365, lat=181, lon=360) float32  (time=365, lat=?, lon=?) float32
[1] ResampleTime([0])  (time=365, lat=?, lon=?) float32      (time=?, lat=?, lon=?) float32
→ out = [1]                                                  (time=?, lat=?, lon=?) float32

7. Nested Graph-in-Graph composition¶

Because Graph is itself an Operator, you can drop a fully-built Graph into a larger Graph. A graph returns a dict of outputs rather than one carrier, so to the outer graph the inner graph's output shape is unknown (?), and so are the steps after it. Summarise the inner graph on its own for its shapes:

In [13]:

# Inner graph: time-subset + region crop on a single input.
inner_in = Input("ds_in")
cropped = SubsetTime(time_min="2010-01-01", time_max="2010-12-31")(inner_in)
cropped = SubsetBBox(lon_bnds=(-80, 0), lat_bnds=(20, 60))(cropped)
inner = Graph(inputs={"ds_in": inner_in}, outputs={"cropped": cropped})

# Outer graph: feed the input through the inner graph, then resample + reduce.
outer_in = Input("ds")
hooked = inner(outer_in)
resampled = ResampleTime(freq="1ME")(hooked)
reduced = ReduceDims(op="mean", dim=("lat", "lon"))(resampled)
outer = Graph(inputs={"ds": outer_in}, outputs={"global_mean": reduced})

year = sig(time=3650, lat=181, lon=360, dtype="float32")
print(outer.summary(year))
print()
print(inner.summary(year))

# Inner graph: time-subset + region crop on a single input. inner_in = Input("ds_in") cropped = SubsetTime(time_min="2010-01-01", time_max="2010-12-31")(inner_in) cropped = SubsetBBox(lon_bnds=(-80, 0), lat_bnds=(20, 60))(cropped) inner = Graph(inputs={"ds_in": inner_in}, outputs={"cropped": cropped}) # Outer graph: feed the input through the inner graph, then resample + reduce. outer_in = Input("ds") hooked = inner(outer_in) resampled = ResampleTime(freq="1ME")(hooked) reduced = ReduceDims(op="mean", dim=("lat", "lon"))(resampled) outer = Graph(inputs={"ds": outer_in}, outputs={"global_mean": reduced}) year = sig(time=3650, lat=181, lon=360, dtype="float32") print(outer.summary(year)) print() print(inner.summary(year))

Node                   Input sig                              Output sig
------------------------------------------------------------------------
[0] Graph(ds)          (time=3650, lat=181, lon=360) float32  ?
[1] ResampleTime([0])  ?                                      ?
[2] ReduceDims([1])    ?                                      ?
→ global_mean = [2]                                           ?

Node                   Input sig                              Output sig
------------------------------------------------------------------------
[0] SubsetTime(ds_in)  (time=3650, lat=181, lon=360) float32  (time=?, lat=181, lon=360) float32
[1] SubsetBBox([0])    (time=?, lat=181, lon=360) float32     (time=?, lat=?, lon=?) float32
→ cropped = [1]                                               (time=?, lat=?, lon=?) float32

8. Common mistakes the summary catches early¶

• Coarsen with boundary="exact" and a non-divisible factor. summary refuses to render a misleading shape and points at the offending dim:

In [14]:

try:
    Sequential([Coarsen({"lat": 7}, boundary="exact")]).summary(
        sig(time=12, lat=10, lon=8, dtype="float32"),
    )
except ValueError as exc:
    print("caught:", exc)

try: Sequential([Coarsen({"lat": 7}, boundary="exact")]).summary( sig(time=12, lat=10, lon=8, dtype="float32"), ) except ValueError as exc: print("caught:", exc)

caught: coarsen boundary='exact' requires 'lat' size 10 to be divisible by factor 7.

• Coarsen with an unknown boundary string. Constructor-time validation catches the typo before you ever call summary:

In [15]:

try:
    Coarsen({"lat": 2}, boundary="trip")
except ValueError as exc:
    print("caught:", exc)

try: Coarsen({"lat": 2}, boundary="trip") except ValueError as exc: print("caught:", exc)

caught: Coarsen boundary must be one of ('exact', 'trim', 'pad'), got 'trip'.

9. When to reach for which¶

• Sequential.describe() — structural tree only (operator names + configs). Doesn't need a Signature. Use when you want to show the plan of a pipeline.
• Sequential.summary(signature) — adds per-step input/output shape. Use when shape questions matter: "how many lat points after the regrid?", "where did the time axis go?".
• Graph.summary(...) — same idea for DAGs. Required when an operator has multiple inputs (metrics) or multiple outputs (multi-task pipelines).

All three are eager-free: no real data is touched, just signatures threaded through compute_output_signature overrides.