Xarray examples#

To better understand xr.Dataset, let us take a 2D example and a 4D example.

Steady-state Flow past a cylinder#

  • Input (coordinates): x, y

  • Output (data variables): u, v, p

import xarray as xr
import numpy as np

/home/hell/anaconda3/lib/python3.9/site-packages/pandas/core/arrays/masked.py:60: UserWarning: Pandas requires version '1.3.6' or newer of 'bottleneck' (version '1.3.5' currently installed).
  from pandas.core import (

# Define the spatial domain
nx, ny = 100, 100  # Number of grid points in x and y
x = np.linspace(-1, 2, nx)  # x-coordinates (e.g., -1 to 2)
y = np.linspace(-1, 1, ny)  # y-coordinates (e.g., -1 to 1)

# Initialize variables (u, v, p) as zeros or some initial guess, doesn't matter
u = np.zeros((nx, ny))  # Velocity in x-direction
v = np.zeros((nx, ny))  # Velocity in y-direction
p = np.zeros((nx, ny))  # Pressure

# Create a Dataset to store the variables
dataset = xr.Dataset(
    { # data_vars : mapping {var name: (dimension name, array-like)}
        "u": (["x", "y"], u),
        "v": (["x", "y"], v),
        "p": (["x", "y"], p),
    },
    coords={
        "x": x, "y": y # inputs coords
        },
    attrs={
        "description": "2D Navier-Stokes flow past a cylinder",
        "equation": "Incompressible, viscous Navier-Stokes",
    },
)

dataset

<xarray.Dataset> Size: 242kB
Dimensions:  (x: 100, y: 100)
Coordinates:
  * x        (x) float64 800B -1.0 -0.9697 -0.9394 -0.9091 ... 1.939 1.97 2.0
  * y        (y) float64 800B -1.0 -0.9798 -0.9596 -0.9394 ... 0.9596 0.9798 1.0
Data variables:
    u        (x, y) float64 80kB 0.0 0.0 0.0 0.0 0.0 0.0 ... 0.0 0.0 0.0 0.0 0.0
    v        (x, y) float64 80kB 0.0 0.0 0.0 0.0 0.0 0.0 ... 0.0 0.0 0.0 0.0 0.0
    p        (x, y) float64 80kB 0.0 0.0 0.0 0.0 0.0 0.0 ... 0.0 0.0 0.0 0.0 0.0
Attributes:
    description:  2D Navier-Stokes flow past a cylinder
    equation:     Incompressible, viscous Navier-Stokes
dataset.sel(x=0.5, y=0.0, method='nearest')  # Get values at a specific point

<xarray.Dataset> Size: 40B
Dimensions:  ()
Coordinates:
    x        float64 8B 0.5152
    y        float64 8B -0.0101
Data variables:
    u        float64 8B 0.0
    v        float64 8B 0.0
    p        float64 8B 0.0
Attributes:
    description:  2D Navier-Stokes flow past a cylinder
    equation:     Incompressible, viscous Navier-Stokes

Steady-state flow past a sphere with parametric viscosity#

  • Input (coordinates): x, y, z, mu

  • Output (data variables): u, v, w, p

# Define the spatial domain
nx, ny, nz = 50, 50, 50  # Grid points for x, y, z
x = np.linspace(-1, 1, nx)  # x-coordinates
y = np.linspace(-1, 1, ny)  # y-coordinates
z = np.linspace(-1, 1, nz)  # z-coordinates

# Define viscosity range
n_mu = 10  # Number of viscosity values
mu = np.linspace(0.1, 1.0, n_mu)  # Viscosity values

# Initialize outputs as zero arrays with the shape of the grid
shape = (nx, ny, nz, n_mu)  # Shape of the 4D space
u = np.zeros(shape)  # Velocity in x-direction
v = np.zeros(shape)  # Velocity in y-direction
w = np.zeros(shape)  # Velocity in z-direction
p = np.zeros(shape)  # Pressure

dataset = xr.Dataset(
    {
        "u": (["x", "y", "z", "mu"], u),  # u depends on (x, y, z, mu)
        "v": (["x", "y", "z", "mu"], v),  # v depends on (x, y, z, mu)
        "w": (["x", "y", "z", "mu"], w),  # w depends on (x, y, z, mu)
        "p": (["x", "y", "z", "mu"], p),  # p depends on (x, y, z, mu)
    },
    coords={
        "x": x,   # x-coordinates
        "y": y,   # y-coordinates
        "z": z,   # z-coordinates
        "mu": mu, # Viscosity values
    },
    attrs={
        "description": "3D flow past a sphere with parametric viscosity",
        "equation": "Navier-Stokes equations with variable viscosity",
    }
)

dataset

<xarray.Dataset> Size: 40MB
Dimensions:  (x: 50, y: 50, z: 50, mu: 10)
Coordinates:
  * x        (x) float64 400B -1.0 -0.9592 -0.9184 -0.8776 ... 0.9184 0.9592 1.0
  * y        (y) float64 400B -1.0 -0.9592 -0.9184 -0.8776 ... 0.9184 0.9592 1.0
  * z        (z) float64 400B -1.0 -0.9592 -0.9184 -0.8776 ... 0.9184 0.9592 1.0
  * mu       (mu) float64 80B 0.1 0.2 0.3 0.4 0.5 0.6 0.7 0.8 0.9 1.0
Data variables:
    u        (x, y, z, mu) float64 10MB 0.0 0.0 0.0 0.0 0.0 ... 0.0 0.0 0.0 0.0
    v        (x, y, z, mu) float64 10MB 0.0 0.0 0.0 0.0 0.0 ... 0.0 0.0 0.0 0.0
    w        (x, y, z, mu) float64 10MB 0.0 0.0 0.0 0.0 0.0 ... 0.0 0.0 0.0 0.0
    p        (x, y, z, mu) float64 10MB 0.0 0.0 0.0 0.0 0.0 ... 0.0 0.0 0.0 0.0
Attributes:
    description:  3D flow past a sphere with parametric viscosity
    equation:     Navier-Stokes equations with variable viscosity
dataset.sel(x=0.5, y=0.0, z=0.0, mu=0.2, method='nearest')  # Get values at a specific point

<xarray.Dataset> Size: 64B
Dimensions:  ()
Coordinates:
    x        float64 8B 0.5102
    y        float64 8B 0.02041
    z        float64 8B 0.02041
    mu       float64 8B 0.2
Data variables:
    u        float64 8B 0.0
    v        float64 8B 0.0
    w        float64 8B 0.0
    p        float64 8B 0.0
Attributes:
    description:  3D flow past a sphere with parametric viscosity
    equation:     Navier-Stokes equations with variable viscosity
dataset.sel(x=0.5, y=0.0, z=0.0, method='nearest')  # if we omit mu, it will return all mu values

<xarray.Dataset> Size: 424B
Dimensions:  (mu: 10)
Coordinates:
    x        float64 8B 0.5102
    y        float64 8B 0.02041
    z        float64 8B 0.02041
  * mu       (mu) float64 80B 0.1 0.2 0.3 0.4 0.5 0.6 0.7 0.8 0.9 1.0
Data variables:
    u        (mu) float64 80B 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0
    v        (mu) float64 80B 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0
    w        (mu) float64 80B 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0
    p        (mu) float64 80B 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0
Attributes:
    description:  3D flow past a sphere with parametric viscosity
    equation:     Navier-Stokes equations with variable viscosity

Transient flow past a sphere with parametric viscosity#

  • Input (coordinates): x, y, z, mu, t

  • Output (data variables): u, v, w, p

# Define time
n_t = 50  # Number of time steps
time = np.linspace(0, 10, n_t)  # Time values from 0 to 10 seconds

# new shape with time dimension
shape = (len(x), len(y), len(z), len(mu), len(time))
u = np.zeros(shape)  # Velocity in x-direction
v = np.zeros(shape)  # Velocity in y-direction
w = np.zeros(shape)  # Velocity in z-direction
p = np.zeros(shape)  # Pressure

dataset = xr.Dataset(
    {
        "u": (["x", "y", "z", "mu", "time"], u),
        "v": (["x", "y", "z", "mu", "time"], v),
        "w": (["x", "y", "z", "mu", "time"], w),
        "p": (["x", "y", "z", "mu", "time"], p),
    },
    coords={
        "x": x,       # x-coordinates
        "y": y,       # y-coordinates
        "z": z,       # z-coordinates
        "mu": mu,     # Viscosity values
        "time": time  # Time steps
    },
    attrs={
        "description": "Transient 3D flow past a sphere with parametric viscosity",
        "equation": "Time-dependent Navier-Stokes equations with variable viscosity",
    }
)

dataset

<xarray.Dataset> Size: 2GB
Dimensions:  (x: 50, y: 50, z: 50, mu: 10, time: 50)
Coordinates:
  * x        (x) float64 400B -1.0 -0.9592 -0.9184 -0.8776 ... 0.9184 0.9592 1.0
  * y        (y) float64 400B -1.0 -0.9592 -0.9184 -0.8776 ... 0.9184 0.9592 1.0
  * z        (z) float64 400B -1.0 -0.9592 -0.9184 -0.8776 ... 0.9184 0.9592 1.0
  * mu       (mu) float64 80B 0.1 0.2 0.3 0.4 0.5 0.6 0.7 0.8 0.9 1.0
  * time     (time) float64 400B 0.0 0.2041 0.4082 0.6122 ... 9.592 9.796 10.0
Data variables:
    u        (x, y, z, mu, time) float64 500MB 0.0 0.0 0.0 0.0 ... 0.0 0.0 0.0
    v        (x, y, z, mu, time) float64 500MB 0.0 0.0 0.0 0.0 ... 0.0 0.0 0.0
    w        (x, y, z, mu, time) float64 500MB 0.0 0.0 0.0 0.0 ... 0.0 0.0 0.0
    p        (x, y, z, mu, time) float64 500MB 0.0 0.0 0.0 0.0 ... 0.0 0.0 0.0
Attributes:
    description:  Transient 3D flow past a sphere with parametric viscosity
    equation:     Time-dependent Navier-Stokes equations with variable viscosity
dataset.sel(time = 10.0, method='nearest') # at a specific time

<xarray.Dataset> Size: 40MB
Dimensions:  (x: 50, y: 50, z: 50, mu: 10)
Coordinates:
  * x        (x) float64 400B -1.0 -0.9592 -0.9184 -0.8776 ... 0.9184 0.9592 1.0
  * y        (y) float64 400B -1.0 -0.9592 -0.9184 -0.8776 ... 0.9184 0.9592 1.0
  * z        (z) float64 400B -1.0 -0.9592 -0.9184 -0.8776 ... 0.9184 0.9592 1.0
  * mu       (mu) float64 80B 0.1 0.2 0.3 0.4 0.5 0.6 0.7 0.8 0.9 1.0
    time     float64 8B 10.0
Data variables:
    u        (x, y, z, mu) float64 10MB 0.0 0.0 0.0 0.0 0.0 ... 0.0 0.0 0.0 0.0
    v        (x, y, z, mu) float64 10MB 0.0 0.0 0.0 0.0 0.0 ... 0.0 0.0 0.0 0.0
    w        (x, y, z, mu) float64 10MB 0.0 0.0 0.0 0.0 0.0 ... 0.0 0.0 0.0 0.0
    p        (x, y, z, mu) float64 10MB 0.0 0.0 0.0 0.0 0.0 ... 0.0 0.0 0.0 0.0
Attributes:
    description:  Transient 3D flow past a sphere with parametric viscosity
    equation:     Time-dependent Navier-Stokes equations with variable viscosity

Summary#

The size of data variables or the output should be equal to the product of the size of all the coordinates or the input. This helps xarray to created a high dimensional meshgrid.