Scalar Fields and Plotting

Scalar Fields and Plotting#

Note

In this tutorial, we use the xarray library. Ensure that you install it before running the code snippets.

The ScalarField class is the fundamental class for all scalar fields used in FRIDOM. Essentially, a scalar field is a wrapper around multidimensional arrays that can be stored in various backends, such as numpy, cupy, or jax.numpy.

In addition to the array itself, a scalar field stores a range of metadata, including the field’s name, units, coordinate information, and dimensions. If you’re familiar with xarray, you’ll notice that the scalar field class is very similar to DataArray in xarray. In fact, later in this tutorial, we will see how scalar field can be easily converted to xarray DataArrays. This is particularly useful if you want to leverage the extensive plotting capabilities of xarray.

In this tutorial, you will learn how to create scalar fields, manipulate them, plot them, and explore a variety of useful functions to work with them effectively.

Creating Scalar Fields#

To create a scalar field, you need to provide the ModelSettings object and the name of the field. Additionally, a range of optional arguments can be specified, such as the field’s unit and a longer, more descriptive name. For a full list of possible arguments that can be passed to the scalar field class, refer to the API documentation.

In the following example, we create a scalar field for temperature:

Creating a scalar field#

import fridom.shallowwater as sw

grid = sw.grid.cartesian.Grid(N=(256,256), L=(1,1), periodic_bounds=(True, True))
mset = sw.ModelSettings(grid=grid)
mset.setup()

temperature = sw.ScalarField(mset, name="temp", units="K", long_name="Temperature")

By default, a scalar field is initialized with zeros. Let’s now look at how we can access and manipulate the data within scalar fields.

Manipulating Scalar Fields#

Scalar fields can primarily be manipulated in three ways: through arithmetic operations, built-in methods of the scalar field class, and numpy functions.

Arithmetic Operations#

Let’s start with arithmetic operations between scalar fields and scalars:

Arithmetic operations with scalars#

import fridom.shallowwater as sw

grid = sw.grid.cartesian.Grid(N=(256,256), L=(1,1), periodic_bounds=(True, True))
mset = sw.ModelSettings(grid=grid)
mset.setup()

# Create a scalar field
u = sw.ScalarField(mset, name="u")

# We can access the data of the scalar field using the .arr attribute
print(u.arr)  # 0 everywhere

# Adding a scalar to the scalar field
u += 3.14  # or alternatively: u = u + 3.14

print(u.arr)  # 3.14 everywhere

In addition to addition, subtraction (-), multiplication (*), division (/), and exponentiation (**) with scalars are also supported.

We can also perform these arithmetic operations between two scalar fields. When performing arithmetic operations between two scalar fields, the attributes of the first scalar field are carried over to the result. Consider the following example:

Arithmetic operations with scalar fields#

import fridom.shallowwater as sw

grid = sw.grid.cartesian.Grid(N=(256,256), L=(1,1), periodic_bounds=(True, True))
mset = sw.ModelSettings(grid=grid)
mset.setup()

# Create two scalar fields
u = sw.ScalarField(mset, name="u", long_name="1st variable") + 1
v = sw.ScalarField(mset, name="v", long_name="2nd variable") + 2

# Adding two scalar fields
w = u + v  # 1 + 2 = 3

# Check the data of the new scalar field
print(w)

As you can see, the new scalar field w inherits the attributes of the first scalar field u.

Hint

If you want to prevent the attributes of a scalar field from being carried over, you can simply just change the array of the scalar field:

u = sw.ScalarField(mset, name="u")
v = sw.ScalarField(mset, name="v")

v.arr = (u * 2 + v).arr

In the output, you will notice several attributes we haven’t discussed yet. The attributes position, topo, and bc_types will be covered later in this tutorial. We now take a look at the is_spectral attribute, which indicates that Fourier transformations can be applied to scalar fields.

Scalar Field Methods#

Fourier transformations are one of the built-in methods that can be applied to scalar fields. However, Fourier transformations are not possible for all grids and boundary conditions. In our doubly-periodic case, they are applicable. Consider the following example:

Fourier transformation#

import fridom.shallowwater as sw

grid = sw.grid.cartesian.Grid(N=(256,256), L=(1,1), periodic_bounds=(True, True))
mset = sw.ModelSettings(grid=grid)
mset.setup()

u = sw.ScalarField(mset, name="u")
v = u.fft()   # Fourier transform to spectral space
w = v.ifft()  # Inverse Fourier transform back to physical space

# Check the dtype of the variables
print(u.arr.dtype)  # float64
print(v.arr.dtype)  # complex128
print(w.arr.dtype)  # float64

Note

By default, all arrays are stored in double precision. FRIDOM also has the capability to work with single precision arrays. For more details, see here.

A variety of other methods are available to facilitate working in parallel settings, such as synchronizing halo regions (see here) or computing global values like the maximum, minimum, sum, and integral.

Scalar field methods#

import fridom.shallowwater as sw

grid = sw.grid.cartesian.Grid(N=(256,256), L=(1,1), periodic_bounds=(True, True))
mset = sw.ModelSettings(grid=grid)
mset.setup()

u = sw.ScalarField(mset, name="u")

u_max = u.max()  # -> float: global maximum
u_min = u.min()  # -> float: global minimum
u_sum = u.sum()  # -> float: global sum
u_total = u.integrate()  # -> float: global integral

u = u.sync()  # -> ScalarField: synchronize halo regions

Apply Numpy Functions#

numpy functions cannot be directly applied to scalar fields. However, they can be used by applying them to the array inside the scalar field. In the following example, we initialize a scalar field using the sine function:

Applying numpy functions#

import fridom.shallowwater as sw

grid = sw.grid.cartesian.Grid(N=(256,256), L=(1,1), periodic_bounds=(True, True))
mset = sw.ModelSettings(grid=grid)
mset.setup()

u = sw.ScalarField(mset, name="u")

# Get the meshgrid of the scalar field
X, Y = u.get_mesh()

# Access the numpy-like module and apply the sin function
ncp = sw.config.ncp
u.arr = ncp.sin(2 * ncp.pi * X)  # sin(2*pi*x)

Note

Recall that the numpy-like module can be accessed via the ncp attribute of the config module.

In a similar manner, most numpy functions can be applied to scalar fields. However, an exception to this is random fields, as different backends handle them differently. Instead, random arrays can be generated from the grid:

Random fields#

import fridom.shallowwater as sw

grid = sw.grid.cartesian.Grid(N=(256,256), L=(1,1), periodic_bounds=(True, True))
mset = sw.ModelSettings(grid=grid)
mset.setup()

u = sw.ScalarField(mset, name="u")
X, Y = u.get_mesh()

# Generate a random field
u.arr = grid.create_random_array(seed=12345)

The seed parameter controls the random number generator.

If you have completed the previous tutorial, you may have noticed that here we created the meshgrid using the scalar field, while in the previous tutorial, we used the grid class. The reason for this is that scalar fields can be located at different positions on the grid. Which means that their exact coordinates depend on the position of the scalar field. We will explore this further in the next section.

Position of Scalar Fields#

Grids can be broadly categorized into two types: staggered and non-staggered grids. In non-staggered grids, all scalar fields are located at the center of the cells, while in staggered grids, scalar fields can be positioned at different locations. For non-staggered grids, this section is not particularly relevant. However, for staggered grids, it is important to know where the scalar fields are positioned. A scalar field can either be located on the cell faces (FACE) or at the cell centers (CENTER) in each direction.

../../_images/grid_positions.svg — The four possible grid positions for a 2D cartesian grid. The dots represents the position of the middle grid cell. F stands for the face, and C for the center.#

In FRIDOM, positions can be defined in two ways: either by using the cell center as defined in the grid as a starting point, or by manually specifying the position. In the following example, we create a position that is on the FACE in the x-direction and on the CENTER in the y-direction:

Using the grid

import fridom.shallowwater as sw

grid = sw.grid.cartesian.Grid(N=(256,256), L=(1,1), periodic_bounds=(True, True))
mset = sw.ModelSettings(grid=grid)
mset.setup()

# Shift the cell center to the face in the x-direction
position = grid.cell_center.shift(axis=0)
print(position)

Manual definition

import fridom.shallowwater as sw

grid = sw.grid.cartesian.Grid(N=(256,256), L=(1,1), periodic_bounds=(True, True))
mset = sw.ModelSettings(grid=grid)
mset.setup()

# Get the AxisPosition enum
AxisPosition = sw.grid.AxisPosition

# Create a position object
position = sw.grid.Position((AxisPosition.FACE, AxisPosition.CENTER))
print(position)

The position can be passed with the keyword argument position when creating a scalar field. By default, scalar fields are placed on the cell center.

Plotting Scalar Fields with Xarray#

The easiest way to plot scalar fields is to convert them into an xarray DataArray and then utilize the extensive plotting functions provided by xarray. There are essentially two ways to convert a scalar field into an xarray DataArray: you can either convert the entire scalar field using the .xr method or convert a slice of the scalar field using the .xrs method.

Converting everything

import fridom.shallowwater as sw

grid = sw.grid.cartesian.Grid(N=(256,256), L=(1,1), periodic_bounds=(True, True))
mset = sw.ModelSettings(grid=grid)
mset.setup()

ncp = sw.config.ncp
u = sw.ScalarField(mset, name="u", long_name="Velocity", units="m/s")
X, Y = u.get_mesh()
u.arr = ncp.sin(2 * ncp.pi * X)

# Convert the entire scalar field to a DataArray and plot it
u.xr.plot()

Converting a slice

import fridom.shallowwater as sw

grid = sw.grid.cartesian.Grid(N=(256,256), L=(1,1), periodic_bounds=(True, True))
mset = sw.ModelSettings(grid=grid)
mset.setup()

ncp = sw.config.ncp
u = sw.ScalarField(mset, name="u", long_name="Velocity", units="m/s")
X, Y = u.get_mesh()
u.arr = ncp.sin(2 * ncp.pi * X)

# Convert a slice of the scalar field to a DataArray and plot it
u.xrs[:, 0].plot()

Tip

If you are working with 3D scalar fields, you can plot a section of the field in a 2D plot by using the .xrs method with two slices:

u.xrs[:, :, 0].plot()  # plot the z=0 section

You may wonder why we have the .xrs method if you can achieve the same result using the .sel method in xarray. The reason lies in performance. For example, if the array lies on the GPU, it must first be copied to the CPU before being converted into an xarray DataArray. This process is obviously faster when only a part of the array needs to be converted.

Note

The xarray conversion is not compatible when running the framework in parallel.

Differentiation and Interpolation#

When modeling partial differential equations, one is often interested in the derivatives of scalar fields. Simple partial derivatives can be calculated using the diff method. The method takes as arguments the axis along which the derivative should be computed and the order of the derivative.

Differentiation#

import matplotlib.pyplot as plt
import fridom.shallowwater as sw

grid = sw.grid.cartesian.Grid(N=(256,256), L=(1,1), periodic_bounds=(True, True))
mset = sw.ModelSettings(grid=grid)
mset.setup()

ncp = sw.config.ncp
u = sw.ScalarField(mset, name="u")
X, Y = u.get_mesh()
u.arr = ncp.sin(2 * ncp.pi * X)

# Calculate the first derivative in the x-direction
du_dx = u.diff(axis=0, order=1)

u.xrs[:, 0].plot(label="u")
du_dx.xrs[:, 0].plot(label="du/dx")
plt.legend()

How these derivatives are calculated depends on the type of grid being used. For the Cartesian grid in this example, the derivative is calculated using forward and backward differences. Here, the position of the scalar field becomes important. If the field is located at the CENTER in a given direction, the first derivative is computed using a forward difference, and the resulting field is located at the cell FACE. Conversely, if the field is located on the FACE, a backward difference is used to compute the derivative.

Let’s review the positions of the fields from the above example:

print(f"u: {u.position}")
print(f"du_dx: {du_dx.position}")

In some cases, you might want to interpolate a field from one position to another. This can be achieved using the interpolate method, which takes as its argument the position to which the field should be interpolated.

Interpolation#

du_dx_center = du_dx.interpolate(u.position)

The interpolation method used depends on the grid and the underlying interpolation module. In this example, linear interpolation is performed. However, other interpolation methods of higher order are available as well.

Finally, let’s introduce a few useful differentiation functions, all of which are based on the diff method:

Useful differentiation functions#

# Calculate the gradient of a scalar field
grad = u.grad()  # returns a list of scalar fields

# Calculate the Laplacian of a scalar field
lap = u.laplacian()

Flat Axes#

So far, we have only considered scalar fields that have the same dimensions as the grid. However, it is also possible to create scalar fields that do not extend in certain directions. This can be useful, for instance, when applying 2D surface forcing in a 3D domain or when creating a vertical profile. To achieve this, you can use the topo argument, which is a list of booleans. If topo is True in a given direction, the scalar field has an extent in that direction. If topo is False, the scalar field has no extent in that direction.

Consider the following example, in which we create a 2D scalar field that has an extent in the x-direction but not in the y-direction:

Flat axes#

import fridom.shallowwater as sw

grid = sw.grid.cartesian.Grid(N=(256,256), L=(1,1), periodic_bounds=(True, True))
mset = sw.ModelSettings(grid=grid)
mset.setup()

u = sw.ScalarField(mset, name="u", topo=(True, False))

If you add this field to a scalar field that extends in all directions, the result will be a scalar field that extends in all directions. The scalar field without extent in a particular direction is assumed to be constant along that direction:

Adding flat axes#

import fridom.shallowwater as sw

grid = sw.grid.cartesian.Grid(N=(256,256), L=(1,1), periodic_bounds=(True, True))
mset = sw.ModelSettings(grid=grid)
mset.setup()

u = sw.ScalarField(mset, name="v")
v = sw.ScalarField(mset, name="u", topo=(True, False))

print(u.arr.shape)  # (258, 258)
print(v.arr.shape)  # (258, 1)

u += 1.0
v += 9.0
w = u + v  # w is now 10 everywhere

Note

The number of grid points in the extended directions is 258 and not 256 because the scalar field is extended by one grid pointp in each direction to account for the halo regions.

Boundary Conditions#

Warning

Boundary conditions are still under development and are likely to change in the future.

Boundary conditions are particularly important for spectral methods on non-periodic grids. The type of boundary condition can be set individually for each direction. The possible boundary conditions are:

NEUMANN: \(\partial_n u = 0\) (zero normal derivative at the boundary)
DIRICHLET: \(u = 0\) (zero value at the boundary)

By default, all boundaries are set to NEUMANN. They can be customized using the bc_types argument in the scalar field class:

Boundary conditions#

import fridom.shallowwater as sw

grid = sw.grid.cartesian.Grid(N=(256,256), L=(1,1), periodic_bounds=(True, True))
mset = sw.ModelSettings(grid=grid)
mset.setup()

NEUMANN = sw.grid.BCType.NEUMANN
DIRICHLET = sw.grid.BCType.DIRICHLET

# u should be zero at the x-boundaries and have zero normal derivative at the y-boundaries
u = sw.ScalarField(mset, name="u", bc_types=(DIRICHLET, NEUMANN))

Summary#

In this tutorial, we explored how to work with scalar fields in the FRIDOM framework. We began by creating a scalar field and learned how to manipulate its data using arithmetic operations, built-in methods, and numpy functions. We also discussed positioning of scalar fields on staggered and non-staggered grids, showing how to define a variable’s location on the grid and how this affects derivative computations.

We demonstrated how to utilize xarray to convert scalar fields for efficient plotting, either converting the entire variable or just a slice for better performance. We then covered calculating derivatives of scalar fields using the diff method and how the grid type and variable position impact these calculations. Interpolation between positions was shown using the interpolate method, and we provided additional useful differentiation functions like grad and laplacian.

The tutorial also discussed handling flat axes when a scalar field does not extend in certain grid dimensions, which is particularly useful for applying surface forcing or creating vertical profiles. Lastly, we briefly introduced boundary conditions, showing how they can be set for different grid directions.

With these tools and techniques, you now have a foundational understanding of how to create, manipulate, analyze, and visualize scalar fields in FRIDOM, equipping you to create own custom initial conditions, which we will cover in the next tutorial.

Scalar Fields and Plotting

Contents

Scalar Fields and Plotting#

Creating Scalar Fields#

Manipulating Scalar Fields#

Arithmetic Operations#

Scalar Field Methods#

Apply Numpy Functions#

Position of Scalar Fields#

Plotting Scalar Fields with Xarray#

Differentiation and Interpolation#

Flat Axes#

Boundary Conditions#

Summary#