Skip to content

Commit

Permalink
Merge to main
Browse files Browse the repository at this point in the history
  • Loading branch information
atb1995 committed Feb 13, 2025
2 parents 4ba7c2f + 22906e5 commit 154048e
Show file tree
Hide file tree
Showing 96 changed files with 5,191 additions and 1,674 deletions.
22 changes: 22 additions & 0 deletions .github/workflows/build.yml
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,13 @@ on:
# Scheduled build at 0330 UTC on Monday mornings to detect bitrot.
- cron: '30 3 * * 1'

concurrency:
# Cancels jobs running if new commits are pushed
group: >
${{ github.workflow }}-
${{ github.event.pull_request.number || github.ref }}
cancel-in-progress: true

jobs:
build:
name: "Build Gusto"
Expand Down Expand Up @@ -39,6 +46,10 @@ jobs:
firedrake-clean
export GUSTO_PARALLEL_LOG=FILE
export PYOP2_CFLAGS=-O0
python -m pip uninstall -y netCDF4
export HDF5_DIR=$PETSC_DIR/packages
export NETCDF4_DIR=$PETSC_DIR/packages
python -m pip install --no-binary netCDF4 --no-build-isolation netCDF4
python -m pytest \
-n 12 --dist worksteal \
--durations=100 \
Expand All @@ -53,6 +64,17 @@ jobs:
mkdir logs
cd /tmp/pytest-of-firedrake/pytest-0/
find . -name "*.log" -exec cp --parents {} /__w/gusto/gusto/logs/ \;
- name: Test serial netCDF
run: |
. /home/firedrake/firedrake/bin/activate
python -m pip uninstall -y netCDF4
python -m pip cache remove netCDF4
python -m pip install --only-binary netCDF4 netCDF4
firedrake-clean
export GUSTO_PARALLEL_LOG=FILE
export PYOP2_CFLAGS=-O0
python -m pytest -n 3 -v integration-tests/model/test_nc_outputting.py
timeout-minutes: 10
- name: Upload artifact
if: always()
uses: actions/upload-pages-artifact@v3
Expand Down
24 changes: 24 additions & 0 deletions docs/about_gusto.md
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,30 @@ cd firedrake/src/gusto
make test
```

#### Parallel output with netCDF

The [`netCDF4`](https://pypi.org/project/netCDF4/) package installed by Gusto does not support parallel usage.
This means that, when Gusto is run in parallel, distributed data structures must first be gathered onto rank 0 before they can be output.
This is *extremely inefficient* at high levels of parallelism.

To avoid this it is possible to build a parallel-aware version of `netCDF4`.
The steps to do this are as follows:

1. Activate the Firedrake virtual environment.
2. Uninstall the existing `netCDF4` package:
```
$ pip uninstall netCDF4
```
3. Set necessary environment variables (note that this assumes that PETSc was built by Firedrake):
```
$ export HDF5_DIR=$VIRTUAL_ENV/src/petsc/default
$ export NETCDF4_DIR=$HDF5_DIR
```
4. Install the parallel version of `netCDF4`:
```
$ pip install --no-build-isolation --no-binary netCDF4 netCDF4
```
### Getting Started
Once you have a working installation, the best way to get started with Gusto is to play with some of the examples in the `gusto/examples` directory.
Expand Down
Original file line number Diff line number Diff line change
@@ -1,68 +1,71 @@
"""
The hydrostatic 1 metre high mountain test case from Melvin et al, 2010:
``An inherently mass-conserving iterative semi-implicit semi-Lagrangian
discretization of the non-hydrostatic vertical-slice equations.'', QJRMS.
The Schär mountain test case of Schär et al, 2002:
``A new terrain-following vertical coordinate formulation for atmospheric
prediction models.'', MWR.
This test describes a wave over a mountain in a hydrostatic atmosphere.
This test describes a wave over a set of idealised mountains, testing how the
discretisation handles orography.
The setup used here uses the order 1 finite elements.
"""

from argparse import ArgumentParser, ArgumentDefaultsHelpFormatter

from firedrake import (
as_vector, VectorFunctionSpace, PeriodicIntervalMesh, ExtrudedMesh,
SpatialCoordinate, exp, pi, cos, Function, conditional, Mesh, Constant
SpatialCoordinate, exp, pi, cos, Function, Mesh, Constant
)
from gusto import (
Domain, IO, OutputParameters, SemiImplicitQuasiNewton, SSPRK3, DGUpwind,
TrapeziumRule, SUPGOptions, ZComponent, Perturbation,
CompressibleParameters, HydrostaticCompressibleEulerEquations,
CompressibleSolver, compressible_hydrostatic_balance, HydrostaticImbalance,
SpongeLayerParameters, MinKernel, MaxKernel, logger
Domain, CompressibleParameters, CompressibleSolver, logger,
OutputParameters, IO, SSPRK3, DGUpwind, SemiImplicitQuasiNewton,
compressible_hydrostatic_balance, SpongeLayerParameters, Exner, ZComponent,
Perturbation, SUPGOptions, TrapeziumRule, MaxKernel, MinKernel,
CompressibleEulerEquations, SubcyclingOptions, RungeKuttaFormulation
)

mountain_hydrostatic_defaults = {
'ncolumns': 200,
'nlayers': 120,
'dt': 5.0,
'tmax': 15000.,
'dumpfreq': 1500,
'dirname': 'mountain_hydrostatic'
schaer_mountain_defaults = {
'ncolumns': 100,
'nlayers': 50,
'dt': 8.0,
'tmax': 5*60*60., # 5 hours
'dumpfreq': 2250, # dump at end with default settings
'dirname': 'schaer_mountain'
}


def mountain_hydrostatic(
ncolumns=mountain_hydrostatic_defaults['ncolumns'],
nlayers=mountain_hydrostatic_defaults['nlayers'],
dt=mountain_hydrostatic_defaults['dt'],
tmax=mountain_hydrostatic_defaults['tmax'],
dumpfreq=mountain_hydrostatic_defaults['dumpfreq'],
dirname=mountain_hydrostatic_defaults['dirname']
def schaer_mountain(
ncolumns=schaer_mountain_defaults['ncolumns'],
nlayers=schaer_mountain_defaults['nlayers'],
dt=schaer_mountain_defaults['dt'],
tmax=schaer_mountain_defaults['tmax'],
dumpfreq=schaer_mountain_defaults['dumpfreq'],
dirname=schaer_mountain_defaults['dirname']
):

# ------------------------------------------------------------------------ #
# Parameters for test case
# ------------------------------------------------------------------------ #

domain_width = 240000. # width of domain in x direction, in m
domain_height = 50000. # height of model top, in m
a = 10000. # scale width of mountain, in m
hm = 1. # height of mountain, in m
zh = 5000. # height at which mesh is no longer distorted, in m
Tsurf = 250. # temperature of surface, in K
initial_wind = 20.0 # initial horizontal wind, in m/s
sponge_depth = 20000.0 # depth of sponge layer, in m
g = 9.80665 # acceleration due to gravity, in m/s^2
cp = 1004. # specific heat capacity at constant pressure
sponge_mu = 0.15 # parameter for strength of sponge layer, in J/kg/K
domain_width = 100000. # width of domain in x direction, in m
domain_height = 30000. # height of model top, in m
a = 5000. # scale width of mountain profile, in m
lamda = 4000. # scale width of individual mountains, in m
hm = 250. # height of mountain, in m
Tsurf = 288. # temperature of surface, in K
initial_wind = 10.0 # initial horizontal wind, in m/s
sponge_depth = 10000.0 # depth of sponge layer, in m
g = 9.810616 # acceleration due to gravity, in m/s^2
cp = 1004.5 # specific heat capacity at constant pressure
mu_dt = 1.2 # strength of sponge layer, no units
exner_surf = 1.0 # maximum value of Exner pressure at surface
max_iterations = 10 # maximum number of hydrostatic balance iterations
tolerance = 1e-7 # tolerance for hydrostatic balance iteration
max_iterations = 20 # maximum number of hydrostatic balance iterations
tolerance = 1e-8 # tolerance for hydrostatic balance iteration

# ------------------------------------------------------------------------ #
# Our settings for this set up
# ------------------------------------------------------------------------ #

spinup_steps = 5 # Not necessary but helps balance initial conditions
alpha = 0.51 # Necessary to absorb grid scale waves
element_order = 1
u_eqn_type = 'vector_invariant_form'

Expand All @@ -81,9 +84,9 @@ def mountain_hydrostatic(
# Describe the mountain
xc = domain_width/2.
x, z = SpatialCoordinate(ext_mesh)
zs = hm * a**2 / ((x - xc)**2 + a**2)
zs = hm * exp(-((x - xc)/a)**2) * (cos(pi*(x - xc)/lamda))**2
xexpr = as_vector(
[x, conditional(z < zh, z + cos(0.5 * pi * z / zh)**6 * zs, z)]
[x, z + ((domain_height - z) / domain_height) * zs]
)

# Make new mesh
Expand All @@ -95,64 +98,50 @@ def mountain_hydrostatic(
# Equation
parameters = CompressibleParameters(g=g, cp=cp)
sponge = SpongeLayerParameters(
H=domain_height, z_level=domain_height-sponge_depth, mubar=sponge_mu/dt
H=domain_height, z_level=domain_height-sponge_depth, mubar=mu_dt/dt
)
eqns = HydrostaticCompressibleEulerEquations(
eqns = CompressibleEulerEquations(
domain, parameters, sponge_options=sponge, u_transport_option=u_eqn_type
)

# I/O
output = OutputParameters(
dirname=dirname, dumpfreq=dumpfreq, dump_vtus=True, dump_nc=False
dirname=dirname, dumpfreq=dumpfreq, dump_vtus=False, dump_nc=True
)
diagnostic_fields = [
ZComponent('u'), HydrostaticImbalance(eqns),
Perturbation('theta'), Perturbation('rho')
Exner(parameters), ZComponent('u'), Perturbation('theta'),
Perturbation('rho')
]
io = IO(domain, output, diagnostic_fields=diagnostic_fields)

# Transport schemes
subcycling_opts = SubcyclingOptions(subcycle_by_courant=0.25)
theta_opts = SUPGOptions()
transported_fields = [
TrapeziumRule(domain, "u"),
SSPRK3(domain, "rho"),
SSPRK3(domain, "theta", options=theta_opts)
TrapeziumRule(domain, "u", subcycling_options=subcycling_opts),
SSPRK3(
domain, "rho", rk_formulation=RungeKuttaFormulation.predictor,
subcycling_options=subcycling_opts
),
SSPRK3(
domain, "theta", options=theta_opts,
subcycling_options=subcycling_opts
)
]
transport_methods = [
DGUpwind(eqns, "u"),
DGUpwind(eqns, "rho"),
DGUpwind(eqns, "rho", advective_then_flux=True),
DGUpwind(eqns, "theta", ibp=theta_opts.ibp)
]

# Linear solver
params = {'mat_type': 'matfree',
'ksp_type': 'preonly',
'pc_type': 'python',
'pc_python_type': 'firedrake.SCPC',
# Velocity mass operator is singular in the hydrostatic case.
# So for reconstruction, we eliminate rho into u
'pc_sc_eliminate_fields': '1, 0',
'condensed_field': {'ksp_type': 'fgmres',
'ksp_rtol': 1.0e-8,
'ksp_atol': 1.0e-8,
'ksp_max_it': 100,
'pc_type': 'gamg',
'pc_gamg_sym_graph': True,
'mg_levels': {'ksp_type': 'gmres',
'ksp_max_it': 5,
'pc_type': 'bjacobi',
'sub_pc_type': 'ilu'}}}

alpha = 0.51 # off-centering parameter
linear_solver = CompressibleSolver(
eqns, alpha, solver_parameters=params,
overwrite_solver_parameters=True
)
tau_values = {'rho': 1.0, 'theta': 1.0}
linear_solver = CompressibleSolver(eqns, alpha, tau_values=tau_values)

# Time stepper
stepper = SemiImplicitQuasiNewton(
eqns, io, transported_fields, transport_methods,
linear_solver=linear_solver, alpha=alpha
linear_solver=linear_solver, alpha=alpha, spinup_steps=spinup_steps
)

# ------------------------------------------------------------------------ #
Expand Down Expand Up @@ -189,7 +178,8 @@ def mountain_hydrostatic(
bottom_boundary = Constant(exner_surf, domain=mesh)
logger.info(f'Solving hydrostatic with bottom Exner of {exner_surf}')
compressible_hydrostatic_balance(
eqns, theta_b, rho_b, exner, top=False, exner_boundary=bottom_boundary
eqns, theta_b, rho_b, exner, top=False, exner_boundary=bottom_boundary,
solve_for_rho=True
)

# Solve hydrostatic balance again, but now use minimum value from first
Expand Down Expand Up @@ -243,7 +233,7 @@ def mountain_hydrostatic(

theta0.assign(theta_b)
rho0.assign(rho_b)
u0.project(as_vector([initial_wind, 0.0]), bcs=eqns.bcs['u'])
u0.project(as_vector([initial_wind, 0.0]))

stepper.set_reference_profiles([('rho', rho_b), ('theta', theta_b)])

Expand All @@ -268,38 +258,38 @@ def mountain_hydrostatic(
'--ncolumns',
help="The number of columns in the vertical slice mesh.",
type=int,
default=mountain_hydrostatic_defaults['ncolumns']
default=schaer_mountain_defaults['ncolumns']
)
parser.add_argument(
'--nlayers',
help="The number of layers for the mesh.",
type=int,
default=mountain_hydrostatic_defaults['nlayers']
default=schaer_mountain_defaults['nlayers']
)
parser.add_argument(
'--dt',
help="The time step in seconds.",
type=float,
default=mountain_hydrostatic_defaults['dt']
default=schaer_mountain_defaults['dt']
)
parser.add_argument(
"--tmax",
help="The end time for the simulation in seconds.",
type=float,
default=mountain_hydrostatic_defaults['tmax']
default=schaer_mountain_defaults['tmax']
)
parser.add_argument(
'--dumpfreq',
help="The frequency at which to dump field output.",
type=int,
default=mountain_hydrostatic_defaults['dumpfreq']
default=schaer_mountain_defaults['dumpfreq']
)
parser.add_argument(
'--dirname',
help="The name of the directory to write to.",
type=str,
default=mountain_hydrostatic_defaults['dirname']
default=schaer_mountain_defaults['dirname']
)
args, unknown = parser.parse_known_args()

mountain_hydrostatic(**vars(args))
schaer_mountain(**vars(args))
Loading

0 comments on commit 154048e

Please sign in to comment.