Common Patterns
Now that we have introduced the core features of NDSL and highlighted features which enable more
complex coding structures, it is prudent to provide examples of how NDSL is implemented.
Below are a plethora of patterns which may be useful when developing code for weather and
climate models. These examples have been generalized, but should be broadly applicable to a
variety of situations. Each example has a "source" Fortran program, and a "ported" NDSL
implementation.
Writing to a Z_INTERFACE_DIM
Variable
If working with Z_INTERFACE_DIM
quantities, one will inevitably need to write to the "extra"
point. The easiest way to achieve this is by declaring your compute_dims
to operate on
the Z_INTERFACE_DIM
instead of the Z_DIM
, and then restrict your interval and offset Z_DIM
stencil reads accordingly:
Fortran Code
| program compute_height
implicit none
! setup domain
integer, parameter :: X_DIM = 10
integer, parameter :: Y_DIM = 10
integer, parameter :: Z_DIM = 10
! allocate scalars
integer :: i, j, k, p_crit
real :: seed, total
! allocate arrays
real, dimension(X_DIM, Y_DIM, Z_DIM) :: d_height
real, dimension(X_DIM, Y_DIM, Z_DIM+1) :: height
! initalize data
do i = 1, X_DIM
do j = 1, Y_DIM
do k = 1, Z_DIM
call RANDOM_NUMBER(seed)
! generate a random change in height for each layer in the range 125:175
d_height(i, j, k) = 125 + FLOOR(51*seed)
enddo
enddo
enddo
do i = 1, X_DIM
do j = 1, Y_DIM
do k = 2, Z_DIM+1
height(i, j, k) = height(i, j, k-1) + d_height(i, j, k-1)
enddo
enddo
enddo
PRINT *, height(1, 1, :)
end program compute_height
|
NDSL Code
| from ndsl.dsl.gt4py import (
computation,
interval,
PARALLEL,
)
from ndsl import StencilFactory
from ndsl.boilerplate import get_factories_single_tile
from ndsl.constants import X_DIM, Y_DIM, Z_DIM, Z_INTERFACE_DIM
from ndsl.dsl.typing import FloatField
import random
def compute_height(
height: FloatField,
d_height: FloatField,
):
with computation(PARALLEL), interval(1, None):
# d_height is offset down one because this field operates on the Z_DIM,
# while the stencil is computing on the Z_INTERFACE_DIM
height = height[0, 0, -1] + d_height[0, 0, -1]
class Code:
def __init__(self, stencil_factory: StencilFactory):
# build stencil
self.constructed_stencil = stencil_factory.from_dims_halo(
func=compute_height,
compute_dims=[X_DIM, Y_DIM, Z_INTERFACE_DIM],
)
def __call__(
self,
height: FloatField,
d_height: FloatField,
):
# execute stencil
self.constructed_stencil(
height,
d_height,
)
if __name__ == "__main__":
# setup domain and generate factories
domain = (10, 10, 10)
nhalo = 0
stencil_factory, quantity_factory = get_factories_single_tile(
domain[0],
domain[1],
domain[2],
nhalo,
backend="dace:cpu",
)
# initialize data
height = quantity_factory.zeros([X_DIM, Y_DIM, Z_INTERFACE_DIM], "n/a")
d_height = quantity_factory.zeros([X_DIM, Y_DIM, Z_DIM], "n/a")
for i in range(d_height.field.shape[0]):
for j in range(d_height.field.shape[1]):
for k in range(d_height.field.shape[2]):
d_height.field[i, j, k] = round(random.uniform(125, 175))
# build stencil
code = Code(stencil_factory)
# execute stencil
code(height, d_height)
print(height.field[0, 0, :])
|
K-Dimension Dependent Computations
In weather and climate modeling, it is often necessary to identify a specific level, then perform
one or more operations based on that level (e.g. identify LCL, compute convective parameters).
Fortran Code
| program average_below_level
implicit none
! setup domain
integer, parameter :: X_DIM = 10
integer, parameter :: Y_DIM = 10
integer, parameter :: Z_DIM = 10
! allocate scalars
integer :: i, j, k, p_crit
real :: seed, total
! allocate arrays
integer, dimension(X_DIM, Y_DIM, Z_DIM) :: temperature, pressure
integer, dimension(X_DIM, Y_DIM) :: desired_level
real, dimension(X_DIM, Y_DIM) :: average_temperature
! initalize data
p_crit = 500
do k = 1, Z_DIM
pressure(:, :, k) = 1000 - (k-1) * 100 ! pressure decreases by 100mb per level
do j = 1, Y_DIM
do i = 1, X_DIM
call RANDOM_NUMBER(seed)
! generate a random temperature in the range 20:30
temperature(i, j, k) = 20 + FLOOR(11*seed)
enddo
enddo
enddo
! determine k level where pressure meets a critical value
do i = 1, X_DIM
do j = 1, Y_DIM
do k = 1, Z_DIM
if (pressure(i, j, k) < p_crit) then
desired_level(i, j) = k
! in this simplified example desired level is the same everywhere, but one can
! imagine a case where this is spatially variable (e.g. determine LCL or LFC)
exit
endif
enddo
enddo
enddo
! calculate average based on desired_level
do i = 1, X_DIM
do j = 1, Y_DIM
total = 0
do k = 1, desired_level(i, j)
total = total + temperature(i, j, k)
enddo
average_temperature(i, j) = total/desired_level(i, j)
enddo
enddo
PRINT *, average_temperature
end program average_below_level
|
NDSL Code
| from ndsl.dsl.gt4py import (
FORWARD,
computation,
interval,
)
from gt4py.cartesian.gtscript import THIS_K
from ndsl import StencilFactory
from ndsl.boilerplate import get_factories_single_tile
from ndsl.constants import X_DIM, Y_DIM, Z_DIM
from ndsl.dsl.typing import FloatField, FloatFieldIJ, IntFieldIJ, Int, Float
import random
def average_layers(
temperature: FloatField,
pressure: FloatField,
average_temperature: FloatFieldIJ,
desired_level: IntFieldIJ,
):
from __externals__ import p_crit, k_end
with computation(FORWARD), interval(...):
# determine critical level
if pressure > p_crit:
desired_level = THIS_K
with computation(FORWARD), interval(...):
# sum all temperatures below critical level
if THIS_K <= desired_level:
average_temperature = average_temperature + temperature
# calculate the average on the final level
if THIS_K == k_end:
average_temperature = average_temperature / desired_level
class Code:
def __init__(self, stencil_factory: StencilFactory):
self.constructed_stencil = stencil_factory.from_dims_halo(
func=average_layers,
compute_dims=[X_DIM, Y_DIM, Z_DIM],
externals={
"p_crit": 500,
},
)
def __call__(
self,
temperature: FloatField,
pressure: FloatField,
average_temperature: FloatFieldIJ,
desired_level: IntFieldIJ,
):
self.constructed_stencil(
temperature,
pressure,
average_temperature,
desired_level,
)
if __name__ == "__main__":
# setup domain and generate factories
domain = (10, 10, 10)
nhalo = 0
stencil_factory, quantity_factory = get_factories_single_tile(
domain[0],
domain[1],
domain[2],
nhalo,
backend="dace:cpu",
)
# initialize data
temperature = quantity_factory.zeros([X_DIM, Y_DIM, Z_DIM], "n/a")
for i in range(temperature.field.shape[0]):
for j in range(temperature.field.shape[1]):
for k in range(temperature.field.shape[2]):
temperature.field[i, j, k] = round(random.uniform(20, 30))
pressure = quantity_factory.zeros([X_DIM, Y_DIM, Z_DIM], "n/a")
for k in range(pressure.field.shape[2]):
pressure.field[:, :, k] = 1000 - k * 100
average_temperature = quantity_factory.zeros([X_DIM, Y_DIM], "n/a")
desired_level = quantity_factory.zeros([X_DIM, Y_DIM], "n/a", dtype=Int)
# build stencil
code = Code(stencil_factory)
# execute stencil
code(temperature, pressure, average_temperature, desired_level)
print(average_temperature.field)
|
Global Tables
There may be situations where a table is needed for referencing throughout the execution of a
component/model. It is highly unlikely that these tables will conform to the grid system used by
the rest of the model in any way. Generating these tables, therefore, requires a somewhat
unconventional use of systems. The following example is an adaptation of code used to generate one
of the saturation vapor pressure tables in the GEOS model, and displays the flexibility of
NDSL systems:
Fortran Code
| program make_table
implicit none
integer, parameter :: size = 1000 ! dimension not related to the size of the encompassing model
real :: constant_one, constant_two
real, dimension(size) :: table
integer :: l
constant_one = 10
constant_two = 2
! construct a table based on temperature input
do l = 1, size
table(l) = log(0.1 * (l-1) / constant_one) + constant_two
enddo
print *, table
end program make_table
|
NDSL Code
| from ndsl.dsl.gt4py import (
computation,
interval,
log,
PARALLEL,
GlobalTable,
)
from ndsl import StencilFactory, QuantityFactory
from ndsl.boilerplate import get_factories_single_tile
from ndsl.constants import X_DIM, Y_DIM, Z_DIM
from ndsl.dsl.typing import FloatField, Float
from gt4py.cartesian.gtscript import THIS_K
# NOTE Ideally, these next two statements should be done elsewhere (perhaps
# in a constants and types file, respectively) and imported
# declare table size
table_size = 1000
# define table type
GlobalTable_local_type = GlobalTable[(Float, int(table_size))]
def _compute_table(table: FloatField):
from __externals__ import constant_one, constant_two
with computation(PARALLEL), interval(1, None):
table = log(0.1 * THIS_K / constant_one) + constant_two
def _use_table(out_field: FloatField, table: GlobalTable_local_type):
with computation(PARALLEL), interval(...):
out_field = 2 * table.A[150]
class ConstructTable:
def __init__(self, stencil_factory: StencilFactory, quantity_factory: QuantityFactory):
# build stencil
self.compute_table = stencil_factory.from_dims_halo(
func=_compute_table,
compute_dims=[X_DIM, Y_DIM, Z_DIM],
externals={
"constant_one": 10,
"constant_two": 2,
},
)
self._table = quantity_factory.zeros([X_DIM, Y_DIM, Z_DIM], "n/a")
def __call__(self):
# execute stencil
self.compute_table(
self._table,
)
# return a GlobalTable object
@property
def table(self):
return self._table.field[0, 0, :]
if __name__ == "__main__":
# consider the following model domain and factories are presnet
domain = (10, 10, 10)
nhalo = 0
stencil_factory, quantity_factory = get_factories_single_tile(
domain[0],
domain[1],
domain[2],
nhalo,
backend="dace:cpu",
)
# we must create a new set for this table calculation, because we require an off-grid dimension
domain_table = (1, 1, table_size)
nhalo_table = 0
stencil_factory_table, quantity_factory_table = get_factories_single_tile(
domain_table[0],
domain_table[1],
domain_table[2],
nhalo_table,
backend="debug",
)
# initialize data
temperature = quantity_factory_table.zeros([X_DIM, Y_DIM, Z_DIM], "n/a")
for l in range(temperature.field.shape[2]):
temperature.field[:, :, l] = 10 + 0.01 * (l + 1)
# build stencil
construct_table = ConstructTable(stencil_factory_table, quantity_factory_table)
# execute stencil
construct_table()
# construct stencil that will use the table
use_table = stencil_factory.from_dims_halo(
func=_use_table,
compute_dims=[X_DIM, Y_DIM, Z_DIM],
)
# # initalize array for calculation using table
out_field = quantity_factory.zeros([X_DIM, Y_DIM, Z_DIM], "n/a")
use_table(out_field, construct_table.table)
print(construct_table.table)
print("Done 🚀")
|
These tables must be subsequently referenced as a local variant of the GlobalTable
type,
informing the system that the object features a single off-grid axis, and may not conform
to the larger model grid specifications:
NDSL Code
This type is defined using the following pattern:
GlobalTable_local_type = GlobalTable[(<data type (Float/Int/Bool)>, int(table_size))]
And accurately typed and referenced within the stencil:
def stencil(data: FloatField, table: GlobalTable_local_type):
with computation(PARALLEL), interval(...):
data = table.A[desired_index]
Exit Statements and Internal Masks
Often there are times where it is necessary to control the execution of specific code paths
dynamically within code execution (e.g. computing precipitation if there is sufficient liquid
water present). In Fortran, this can be done easily with an exit
statement. In NDSL, the Python equivalent
break
statement is strictly forbidden, as there is no way to stop a computation early. Instead, the correct
way to control execution of different coordinates on the X/Y plane is by using two dimensional boolean fields,
which act as masks to turn on/off chunks of code on a per-point basis.
Fortran Code
| program conditional_calculation
implicit none
! setup domain
integer, parameter :: X_DIM = 5
integer, parameter :: Y_DIM = 5
integer, parameter :: Z_DIM = 3
! allocate scalars
integer :: i, j, k
real :: seed, critical_value
! allocate arrays
integer, dimension(X_DIM, Y_DIM, Z_DIM) :: in_data, out_data
logical, dimension(X_DIM, Y_DIM) :: mask
critical_value = 5
! initalize data
do i = 1, X_DIM
do j = 1, Y_DIM
do k = 1, Z_DIM
call RANDOM_NUMBER(seed)
! generate a random number between 0 and 10
in_data(i, j, k) = nint(10*seed)
enddo
enddo
enddo
do i = 1, X_DIM
do j = 1, Y_DIM
do k = 1, Z_DIM
! if data is above a critical value anywhere in the column
! perform the subsequenc calculation on the entire column
if (in_data(i, j, k) > critical_value) then
mask(i, j) = .true.
exit
endif
enddo
if (mask(i,j) .eqv. .true.) then
out_data(i,j,:) = 10*in_data
else
out_data(i,j,:) = -999
endif
enddo
enddo
PRINT *, "Input data: ", in_data
PRINT *, "Mask data: ", mask
PRINT *, "Output data: ", out_data
end program conditional_calculation
|
NDSL Code
| from ndsl.dsl.gt4py import PARALLEL, FORWARD, computation, interval, exp, function
from ndsl import StencilFactory
from ndsl.boilerplate import get_factories_single_tile
from ndsl.constants import X_DIM, Y_DIM, Z_DIM
from ndsl.dsl.typing import FloatField, BoolFieldIJ, Bool
import random
def check_value(data: FloatField, mask: BoolFieldIJ):
from __externals__ import critical_value
with computation(FORWARD), interval(...):
# if the data surpasses a critical value anywhere in the column, set the mask to true
if data > critical_value:
mask = True
def computation_stencil(data_in: FloatField, data_out: FloatField, mask: BoolFieldIJ):
with computation(PARALLEL), interval(...):
if mask == True:
data_out = 2 * exp(data_in)
else:
data_out = -999
class DoSomeMath:
def __init__(self, stencil_factory: StencilFactory):
self.check_value = stencil_factory.from_dims_halo(
func=check_value,
compute_dims=[X_DIM, Y_DIM, Z_DIM],
externals={
"critical_value": 5,
},
)
self.computation_stencil = stencil_factory.from_dims_halo(
func=computation_stencil,
compute_dims=[X_DIM, Y_DIM, Z_DIM],
)
def __call__(self, in_field: FloatField, out_field: FloatField, mask_field: BoolFieldIJ):
self.check_value(in_field, mask_field)
self.computation_stencil(in_field, out_field, mask_field)
if __name__ == "__main__":
domain = (5, 5, 3)
nhalo = 0
stencil_factory, quantity_factory = get_factories_single_tile(
domain[0],
domain[1],
domain[2],
nhalo,
)
do_some_math = DoSomeMath(stencil_factory)
in_field = quantity_factory.zeros([X_DIM, Y_DIM, Z_DIM], "n/a")
for i in range(in_field.field.shape[0]):
for j in range(in_field.field.shape[1]):
for k in range(in_field.field.shape[2]):
in_field.field[i, j, k] = round(random.uniform(0, 10), 2)
out_field = quantity_factory.zeros([X_DIM, Y_DIM, Z_DIM], "n/a")
mask_field = quantity_factory.zeros([X_DIM, Y_DIM], "n/a", dtype=Bool)
do_some_math(in_field, out_field, mask_field)
print(f"Input data: {in_field.field}")
print(f"Mask data: {mask_field.field}")
print(f"Output data: {out_field.field}")
|
Nested K Loops
Occasionally there may be situations when a nested K loop is required (such as
accumulating precipitation in a column). NDSL does not have the ability to nest computation/interval
statements; however, such a calculation can be performed using a while
loop and clever indexing within
a single computation/iteration statement:
Below is an example of a nested K loop from the lagrangian_contributions stencil. It is not
currently possible in NDSL to nest a with computation(PARALLEL)
within a
with computation(PARALLEL)
, however a while
loop can be used to create a nested K loop
(lines x-x).
Fortran Code
| program nested_k_loop
implicit none
! setup domain
integer, parameter :: X_DIM = 5
integer, parameter :: Y_DIM = 5
integer, parameter :: Z_DIM = 10
! allocate scalars
integer :: i, j, k, kk
real :: seed
! allocate arrays
integer, dimension(X_DIM, Y_DIM, Z_DIM) :: in_data, out_data
! initalize data
do i = 1, X_DIM
do j = 1, Y_DIM
do k = 1, Z_DIM
call RANDOM_NUMBER(seed)
! generate a random value between 0 and 10
in_data(i, j, k) = nint(10*seed)
enddo
enddo
enddo
out_data(:,:,:) = 0
do i = 1, X_DIM
do j = 1, Y_DIM
do k = 1, Z_DIM
do kk = 1, k
! sum all values at and below the current point
out_data(i,j,k) = out_data(i,j,k) + in_data(i,j,kk)
enddo
enddo
enddo
enddo
PRINT *, "Input data: ", in_data(1,1,:)
PRINT *, "Output data: ", out_data(1,1,:)
end program nested_k_loop
|
NDSL Code
| from ndsl.dsl.gt4py import (
computation,
interval,
PARALLEL,
)
from gt4py.cartesian.gtscript import THIS_K
from ndsl import StencilFactory
from ndsl.boilerplate import get_factories_single_tile
from ndsl.constants import X_DIM, Y_DIM, Z_DIM
from ndsl.dsl.typing import FloatField
import random
def nested_k_calculation(
in_data: FloatField,
out_data: FloatField,
):
with computation(PARALLEL), interval(...):
nested_k_index = 0
while nested_k_index <= THIS_K:
out_data = out_data + in_data.at(K=nested_k_index)
nested_k_index += 1
class Code:
def __init__(self, stencil_factory: StencilFactory):
# build stencil
self.constructed_stencil = stencil_factory.from_dims_halo(
func=nested_k_calculation,
compute_dims=[X_DIM, Y_DIM, Z_DIM],
)
def __call__(
self,
in_data: FloatField,
out_data: FloatField,
):
# execute stencil
self.constructed_stencil(
in_data,
out_data,
)
if __name__ == "__main__":
# setup domain and generate factories
domain = (5, 5, 10)
nhalo = 0
stencil_factory, quantity_factory = get_factories_single_tile(
domain[0],
domain[1],
domain[2],
nhalo,
backend="debug",
)
# initialize data
out_data = quantity_factory.zeros([X_DIM, Y_DIM, Z_DIM], "n/a")
in_data = quantity_factory.zeros([X_DIM, Y_DIM, Z_DIM], "n/a")
for i in range(in_data.field.shape[0]):
for j in range(in_data.field.shape[1]):
for k in range(in_data.field.shape[2]):
in_data.field[i, j, k] = round(random.uniform(0, 10))
# build stencil
code = Code(stencil_factory)
# execute stencil
code(in_data, out_data)
print(in_data.field[0, 0, :])
print(out_data.field[0, 0, :])
|
Nested K loops provide an excellent example of how thinking outside the box - and possessing a willingness
to re-approach traditional coding practices - allows for a much wider application of the software and reveals
that the rules put in place by NDSL are not necessarily as restrictive as they may seem.
Goto Statements
In Fortran, the goto
construct allows the user to jump to another portion of the code. NDSL is
incapable of of "jumping" from one portion of code to another; stencils are always executed linearly
and completely. Since goto
statement are tremendously flexile, there is no "standard" way of translating
a piece of code with a goto
statement into NDSL. Indeed, the presence of goto
statements often signals
a non-parallelizable code structure, which requires refactoring to be implemented in NDSL. During this
process, however, the aforementioned patterns - particularly the use of masks - may be extremely useful.