Create new variables with assign, use built-in math functions and spatial statistics inside lambda expressions, and do simple arithmetic directly on datasets.
Variable creation is done with assign, which works in a similar way to the method available in pandas, using lambda functions evaluated per grid cell, per time step. Let's say we have a dataset with a variable var and want to add 10 to it, calling the new variable new:
ds.assign(new=lambda x: x.var + 10)
The x after lambda represents the dataset; whatever follows the colon is the actual equation to evaluate, and x.var refers to var from the dataset. By default assign keeps the original variables — use drop to keep only the new one(s):
ds.assign(new=lambda x: x.var + 10, drop=True)
Each call to assign currently needs to be written on one line — splitting the arguments across multiple lines will not parse correctly.
assign uses kwargs, so drop can go anywhere in the call — the following are equivalent:
ds.assign(new=lambda x: x.var + 10, drop=True) ds.assign(drop=True, new=lambda x: x.var + 10)
Multiple variables can be created in one call — each can reference the ones created before it. Here, var2 is var1 plus 1, and diff (which should be 1 everywhere) is the difference between them:
ds.assign(var1=lambda x: x.var + 1, var2=lambda x: x.var1 + 1, diff=lambda x: x.var2 - x.var1)
The lambda is evaluated for every grid cell at every time step, so every part of the expression must evaluate to a number. A plain Python variable works fine as a constant:
k = 273.15 ds.assign(drop=True, sst_k=lambda x: x.sst + k)
But if that constant isn't actually a number, assign will throw an error:
k = "273.15" # a string, not a number ds.assign(drop=True, sst_k=lambda x: x.sst + k) # raises an error
Standard mathematical functions, with the same names as their numpy equivalents, are available inside the lambda: abs, floor, ceil, sqrt, exp, log10, sin, cos, tan, arcsin, arccos and arctan.
ds.assign(new=lambda x: ceil(x.old)) ds.assign(new=lambda x: log10(x.old + 1))
assign carries out its calculations at each time step, and you can access spatial statistics for each time step while generating new variables: spatial_mean, spatial_max, spatial_min, spatial_sum, vertical_mean, vertical_max, vertical_min, vertical_sum, zonal_mean, zonal_max, zonal_min and zonal_sum. See the full list of helper functions on the API reference.
For example, mapping how much warmer than the global mean each part of the ocean is:
ds.assign(temp_comp=lambda x: x.temperature - spatial_mean(x.temperature), drop=True)
Comparisons work too. Identify cells hotter than the global average:
ds.assign(temp_comp=lambda x: x.temperature > spatial_mean(x.temperature), drop=True)
Or take it further — cells more than 3 degrees warmer than the global average:
ds.assign(temp_comp=lambda x: x.temperature > (spatial_mean(x.temperature) + 3), drop=True)
Logical operators work in the standard Python way. To find cells with a value between 1 and 10:
ds.assign(one2ten=lambda x: x.var > 1 & x.var < 10, drop=True)
For simple operations like adding or subtracting a constant, or combining two datasets, use add, subtract, divide and multiply — or the equivalent Python operators directly on datasets:
ds.add(10) # or: ds + 10 ds1.add(ds2) # or: ds1 + ds2
These methods require the datasets to be structured compatibly — each dimension must either match exactly, or one dataset must have a single value for that dimension (which is then applied to every matching step in the other). See Simple arithmetic & comparisons for worked examples of subtracting a monthly or annual climatology from a longer time series.
Logical comparisons follow the same pattern — compare, or the Python operators directly:
ds.compare(">0") # or: ds > 0 ds.compare("==0") # or: ds == 0
