Skip to content

Single column implementation with 93 vertical levels #36

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
ma595 wants to merge 4 commits into
base: datawave_ml
Choose a base branch
from
Open

Single column implementation with 93 vertical levels #36

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
bfabe21
Run inference on CPU only
ma595 Oct 13, 2025
3afaeac
Reconfigure single column model with 93 levels
ma595 Nov 3, 2025
a0b401a
Add long-run and variable output to README.md
ma595 Nov 14, 2025
c1e58ae
Correct erroneous user_nl_cam restart settings
ma595 Nov 14, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
46 changes: 46 additions & 0 deletions README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -149,6 +149,52 @@ The case can be run with `./case.submit` from the case directory.
> To leave all output in `$SCRATCH/case/` switch 'short term archiving' off by running `./xmlchange DOUT_S=FALSE` in the
> case directory to change `DOUT_S` from `TRUE` to `FALSE`.

#### Output useful variables

CAM has been configured to output tendencies and fluxes to file. These are named `UTGW_NL`, `VTGW_NL` i.e. tendencies in u and v directions. Similarly, fluxes are named `UFLUX_NL` and `VFLUX_NL`.

For instance, add the following to `user_nl_cam`:

```
nhtfrq=0,5
mfilt=1,60
avgflag_pertape = 'A','I'

fincl2 = 'U','V','T','UTGW_NL','VTGW_NL','UFLUX_NL','VFLUX_NL'
```

This outputs to two history files, 0 and 1, where `fincl1` controls fields in h0 and `fincl2` controls fields in h1. In h0, contains the default variables or fields.

We output to h1 a reduced set at a higher frequency. In this case, the selected variables are written every 5 timesteps, and then output to disk once 60 time samples have been buffered. See documentation on customising CAM output [here](https://ncar.github.io/CESM-Tutorial/notebooks/namelist/output/output_cam.html). In the above instance files will be output in the following format `<job-name>.cam.h1.1979-01-01-[00000,21600,43200...].nc`.

#### Long runs

Long runs can be achieved by following the documentation [here](https://ncar.github.io/CESM-Tutorial/notebooks/modifications/xml/run_length/restarting.html).

The below sequence of `xmlchange` commands sets up a long run that automatically resubmits until 5 years have elapsed. Of course, resubmission can be performed manually by setting `RESUBMIT` to 0 and toggling `CONTINUE_RUN` to `True` after the first run completes.

```
./xmlchange CONTINUE_RUN=FALSE # first run from initial conditions
./xmlchange RESUBMIT=9 # 10 segments total => 5 years (10 × 6 months)

./xmlchange STOP_OPTION='nmonths'
./xmlchange STOP_N=6 # each job = 6 months

./xmlchange REST_OPTION='nmonths'
./xmlchange REST_N=3 # restart every 3 months
```

And for long runs, sensible settings of output frequencies in `user_nl_cam` look like:
```
! Sensible output frequencies
nhtfrq=0,-336 ! Every two weeks
mfilt=1,1
avgflag_pertape = 'A','I'

fincl2 = 'U','V','T','UTGW_NL','VTGW_NL','UFLUX_NL','VFLUX_NL'
```


## NOTE: This is **unsupported** development code and is subject to the [CESM developer's agreement](http://www.cgd.ucar.edu/cseg/development-code.html).

### CAM Documentation - https://ncar.github.io/CAM/doc/build/html/index.html
Expand Down
128 changes: 14 additions & 114 deletions src/physics/cam/gw_nlgw.F90
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -57,50 +57,6 @@ module gw_nlgw
real(r8) :: uflux_mean, vflux_mean
real(r8) :: uflux_std, vflux_std

integer, parameter :: pver_interp = 137 ! number of levels in ERA5

real(r8), parameter :: era5_ak(137) = [ &
1.000000000e+00_r8, 2.550781250e+00_r8, 3.884765625e+00_r8, 5.746093750e+00_r8, 8.289062500e+00_r8, 1.167968750e+01_r8, 1.610937500e+01_r8, 2.179687500e+01_r8, &
2.898437500e+01_r8, 3.793750000e+01_r8, 4.890625000e+01_r8, 6.225000000e+01_r8, 7.818750000e+01_r8, 9.712500000e+01_r8, 1.194375000e+02_r8, 1.453750000e+02_r8, &
1.752500000e+02_r8, 2.096250000e+02_r8, 2.487500000e+02_r8, 2.930000000e+02_r8, 3.427500000e+02_r8, 3.982500000e+02_r8, 4.600000000e+02_r8, 5.285000000e+02_r8, &
6.040000000e+02_r8, 6.865000000e+02_r8, 7.770000000e+02_r8, 8.750000000e+02_r8, 9.820000000e+02_r8, 1.097000000e+03_r8, 1.221000000e+03_r8, 1.355000000e+03_r8, &
1.498000000e+03_r8, 1.651000000e+03_r8, 1.813000000e+03_r8, 1.987000000e+03_r8, 2.170000000e+03_r8, 2.366000000e+03_r8, 2.572000000e+03_r8, 2.788000000e+03_r8, &
3.018000000e+03_r8, 3.258000000e+03_r8, 3.512000000e+03_r8, 3.776000000e+03_r8, 4.054000000e+03_r8, 4.344000000e+03_r8, 4.644000000e+03_r8, 4.960000000e+03_r8, &
5.288000000e+03_r8, 5.624000000e+03_r8, 5.976000000e+03_r8, 6.340000000e+03_r8, 6.720000000e+03_r8, 7.112000000e+03_r8, 7.520000000e+03_r8, 7.944000000e+03_r8, &
8.384000000e+03_r8, 8.840000000e+03_r8, 9.320000000e+03_r8, 9.816000000e+03_r8, 1.032800000e+04_r8, 1.084800000e+04_r8, 1.139200000e+04_r8, 1.193600000e+04_r8, &
1.248800000e+04_r8, 1.304800000e+04_r8, 1.360000000e+04_r8, 1.416000000e+04_r8, 1.470400000e+04_r8, 1.524000000e+04_r8, 1.576800000e+04_r8, 1.628000000e+04_r8, &
1.676800000e+04_r8, 1.723200000e+04_r8, 1.768000000e+04_r8, 1.811200000e+04_r8, 1.849600000e+04_r8, 1.886400000e+04_r8, 1.918400000e+04_r8, 1.948800000e+04_r8, &
1.974400000e+04_r8, 1.995200000e+04_r8, 2.014400000e+04_r8, 2.027200000e+04_r8, 2.036800000e+04_r8, 2.043200000e+04_r8, 2.043200000e+04_r8, 2.040000000e+04_r8, &
2.030400000e+04_r8, 2.017600000e+04_r8, 1.998400000e+04_r8, 1.974400000e+04_r8, 1.945600000e+04_r8, 1.910400000e+04_r8, 1.870400000e+04_r8, 1.825600000e+04_r8, &
1.774400000e+04_r8, 1.718400000e+04_r8, 1.657600000e+04_r8, 1.592800000e+04_r8, 1.524800000e+04_r8, 1.453600000e+04_r8, 1.380000000e+04_r8, 1.304800000e+04_r8, &
1.228800000e+04_r8, 1.152000000e+04_r8, 1.075200000e+04_r8, 9.992000000e+03_r8, 9.248000000e+03_r8, 8.520000000e+03_r8, 7.816000000e+03_r8, 7.136000000e+03_r8, &
6.488000000e+03_r8, 5.868000000e+03_r8, 5.280000000e+03_r8, 4.724000000e+03_r8, 4.208000000e+03_r8, 3.722000000e+03_r8, 3.274000000e+03_r8, 2.858000000e+03_r8, &
2.476000000e+03_r8, 2.128000000e+03_r8, 1.810000000e+03_r8, 1.524000000e+03_r8, 1.265000000e+03_r8, 1.035000000e+03_r8, 8.310000000e+02_r8, 6.515000000e+02_r8, &
4.962500000e+02_r8, 3.635000000e+02_r8, 2.525000000e+02_r8, 1.622500000e+02_r8, 9.243750000e+01_r8, 4.281250000e+01_r8, 1.329687500e+01_r8, 1.878906250e+00_r8, &
0.000000000e+00_r8]

real(r8), parameter :: era5_bk(137) = [ &
0.0000000000e+00_r8, 0.0000000000e+00_r8, 0.0000000000e+00_r8, 0.0000000000e+00_r8, 0.0000000000e+00_r8, 0.0000000000e+00_r8, 0.0000000000e+00_r8, 0.0000000000e+00_r8, &
0.0000000000e+00_r8, 0.0000000000e+00_r8, 0.0000000000e+00_r8, 0.0000000000e+00_r8, 0.0000000000e+00_r8, 0.0000000000e+00_r8, 0.0000000000e+00_r8, 0.0000000000e+00_r8, &
0.0000000000e+00_r8, 0.0000000000e+00_r8, 0.0000000000e+00_r8, 0.0000000000e+00_r8, 0.0000000000e+00_r8, 0.0000000000e+00_r8, 0.0000000000e+00_r8, 0.0000000000e+00_r8, &
0.0000000000e+00_r8, 0.0000000000e+00_r8, 0.0000000000e+00_r8, 0.0000000000e+00_r8, 0.0000000000e+00_r8, 0.0000000000e+00_r8, 0.0000000000e+00_r8, 0.0000000000e+00_r8, &
0.0000000000e+00_r8, 0.0000000000e+00_r8, 0.0000000000e+00_r8, 0.0000000000e+00_r8, 0.0000000000e+00_r8, 0.0000000000e+00_r8, 0.0000000000e+00_r8, 0.0000000000e+00_r8, &
0.0000000000e+00_r8, 0.0000000000e+00_r8, 0.0000000000e+00_r8, 0.0000000000e+00_r8, 0.0000000000e+00_r8, 0.0000000000e+00_r8, 0.0000000000e+00_r8, 0.0000000000e+00_r8, &
0.0000000000e+00_r8, 0.0000000000e+00_r8, 0.0000000000e+00_r8, 0.0000000000e+00_r8, 0.0000000000e+00_r8, 0.0000000000e+00_r8, 2.9802322387e-06_r8, 1.5974044799e-05_r8, &
4.1007995605e-05_r8, 8.4996223449e-05_r8, 1.5604496002e-04_r8, 2.6893615722e-04_r8, 4.5108795166e-04_r8, 7.2622299194e-04_r8, 1.1224746704e-03_r8, 1.6717910766e-03_r8, &
2.4242401123e-03_r8, 3.4141540527e-03_r8, 4.6730041503e-03_r8, 6.2561035156e-03_r8, 8.1939697265e-03_r8, 1.0536193847e-02_r8, 1.3313293457e-02_r8, 1.6571044921e-02_r8, &
2.0339965820e-02_r8, 2.4658203125e-02_r8, 2.9571533203e-02_r8, 3.5095214843e-02_r8, 4.1290283203e-02_r8, 4.8156738281e-02_r8, 5.5755615234e-02_r8, 6.4086914062e-02_r8, &
7.3181152343e-02_r8, 8.3129882812e-02_r8, 9.3872070312e-02_r8, 1.0546875000e-01_r8, 1.1798095703e-01_r8, 1.3134765625e-01_r8, 1.4575195312e-01_r8, 1.6101074218e-01_r8, &
1.7724609375e-01_r8, 1.9458007812e-01_r8, 2.1289062500e-01_r8, 2.3229980468e-01_r8, 2.5268554687e-01_r8, 2.7441406250e-01_r8, 2.9687500000e-01_r8, 3.2080078125e-01_r8, &
3.4570312500e-01_r8, 3.7133789062e-01_r8, 3.9770507812e-01_r8, 4.2480468750e-01_r8, 4.5214843750e-01_r8, 4.7998046875e-01_r8, 5.0781250000e-01_r8, 5.3564453125e-01_r8, &
5.6298828125e-01_r8, 5.9033203125e-01_r8, 6.1669921875e-01_r8, 6.4306640625e-01_r8, 6.6796875000e-01_r8, 6.9287109375e-01_r8, 7.1630859375e-01_r8, 7.3876953125e-01_r8, &
7.6025390625e-01_r8, 7.8076171875e-01_r8, 8.0029296875e-01_r8, 8.1835937500e-01_r8, 8.3544921875e-01_r8, 8.5156250000e-01_r8, 8.6669921875e-01_r8, 8.8085937500e-01_r8, &
8.9355468750e-01_r8, 9.0576171875e-01_r8, 9.1699218750e-01_r8, 9.2675781250e-01_r8, 9.3652343750e-01_r8, 9.4482421875e-01_r8, 9.5263671875e-01_r8, 9.5996093750e-01_r8, &
9.6630859375e-01_r8, 9.7216796875e-01_r8, 9.7753906250e-01_r8, 9.8242187500e-01_r8, 9.8632812500e-01_r8, 9.9023437500e-01_r8, 9.9365234375e-01_r8, 9.9609375000e-01_r8, &
9.9902343750e-01_r8]



contains

!==========================================================================
Expand All @@ -118,10 +74,6 @@ subroutine gw_nlgw_dp_ml(state_in, ptend, lchnk)
integer :: ninputs = 1, noutputs = 1
integer, dimension(2) :: layout = [1 , 2]

integer :: device_id

device_id = mod(iam, 2)

ncol = state_in%ncol

allocate(lat(ncol))
Expand All @@ -140,8 +92,8 @@ subroutine gw_nlgw_dp_ml(state_in, ptend, lchnk)
allocate(utgw(ncol,pver))
allocate(vtgw(ncol,pver))

allocate(net_inputs(ncol, 4*pver_interp+3))
allocate(net_outputs(ncol, 2*pver_interp))
allocate(net_inputs(ncol, 4*pver+3))
allocate(net_outputs(ncol, 2*pver))

! dims = (ncol)
lat = state_in%lat(:ncol)
Expand All @@ -162,7 +114,7 @@ subroutine gw_nlgw_dp_ml(state_in, ptend, lchnk)
call construct_input()

! send all columns from this process
call torch_tensor_from_array(tensor_in(1), net_inputs, layout, torch_kCUDA, device_id)
call torch_tensor_from_array(tensor_in(1), net_inputs, layout, torch_kCPU)
call torch_tensor_from_array(tensor_out(1), net_outputs, layout, torch_kCPU)

! Run net forward on data
Expand Down Expand Up @@ -215,12 +167,9 @@ end subroutine gw_nlgw_dp_ml
subroutine gw_nlgw_dp_init(model_path)

character(len=*), intent(in) :: model_path ! Filepath to PyTorch Torchscript net
integer :: device_id

device_id = mod(iam, 2)

! Load the convective drag net from TorchScript file
call torch_model_load(nlgw_model, model_path, device_type=torch_kCUDA, device_index=device_id)
call torch_model_load(nlgw_model, model_path, device_type=torch_kCPU)
! read in normalisation weights
call read_norms()

Expand Down Expand Up @@ -291,82 +240,33 @@ end subroutine normalise_data

subroutine construct_input()

! temporary arrays for interpolated values
real(r8), dimension(:,:), allocatable :: &
u_interp, &! zonal wind (m/s)
v_interp, &! meridional wind (m/s)
omega_interp, &! vertical pressure velocity (Pa/s)
theta_interp, &! potential temperature (K)
pmid_interp ! midpoint pressure (Pa)

integer :: idx_beg, idx_end, i

allocate(pmid_interp(ncol,pver_interp))
allocate(u_interp(ncol,pver_interp))
allocate(v_interp(ncol,pver_interp))
allocate(theta_interp(ncol,pver_interp))
allocate(omega_interp(ncol,pver_interp))

do i = 1, ncol
pmid_interp(i,:) = era5_ak(:) + ps(i) * era5_bk(:)
call lininterp(u(i,:), pmid(i,:), pver, u_interp(i,:), pmid_interp(i,:), pver_interp)
call lininterp(v(i,:), pmid(i,:), pver, v_interp(i,:), pmid_interp(i,:), pver_interp)
call lininterp(theta(i,:), pmid(i,:), pver, theta_interp(i,:), pmid_interp(i,:), pver_interp)
call lininterp(omega(i,:), pmid(i,:), pver, omega_interp(i,:), pmid_interp(i,:), pver_interp)
end do

net_inputs(:,1) = lat
net_inputs(:,2) = lon
net_inputs(:,3) = phis

idx_end = 3 ! last index written to was phis at position 3
idx_beg = idx_end + 1
idx_end = idx_beg + pver_interp - 1
net_inputs(:,idx_beg:idx_end) = u_interp
idx_end = idx_beg + pver - 1
net_inputs(:,idx_beg:idx_end) = u
idx_beg = idx_end + 1
idx_end = idx_beg + pver_interp - 1
net_inputs(:,idx_beg:idx_end) = v_interp
idx_end = idx_beg + pver - 1
net_inputs(:,idx_beg:idx_end) = v
idx_beg = idx_end + 1
idx_end = idx_beg + pver_interp - 1
net_inputs(:,idx_beg:idx_end) = theta_interp
idx_end = idx_beg + pver - 1
net_inputs(:,idx_beg:idx_end) = theta
idx_beg = idx_end + 1
idx_end = idx_beg + pver_interp - 1
net_inputs(:,idx_beg:idx_end) = omega_interp
idx_end = idx_beg + pver - 1
net_inputs(:,idx_beg:idx_end) = omega

deallocate(pmid_interp)
deallocate(u_interp)
deallocate(v_interp)
deallocate(theta_interp)
deallocate(omega_interp)

end subroutine construct_input

subroutine extract_output()

! temporary arrays for interpolated values
real(r8), dimension(:,:), allocatable :: &
uflux_interp, &! zonal wind (m/s)
vflux_interp, &! meridional wind (m/s)
pmid_interp ! midpoint pressure (Pa)

integer :: idx_beg, idx_end, i

allocate(pmid_interp(ncol,pver_interp))
allocate(uflux_interp(ncol,pver_interp))
allocate(vflux_interp(ncol,pver_interp))

uflux_interp(:, :) = net_outputs(:,:pver_interp)
vflux_interp(:, :) = net_outputs(:,pver_interp+1:)

do i = 1, ncol
pmid_interp(i,:) = era5_ak(:) + ps(i) * era5_bk(:)
call lininterp(uflux_interp(i,:), pmid_interp(i,:), pver_interp, uflux(i,:), pmid(i,:), pver)
call lininterp(vflux_interp(i,:), pmid_interp(i,:), pver_interp, vflux(i,:), pmid(i,:), pver)
end do

deallocate(pmid_interp)
deallocate(uflux_interp)
deallocate(vflux_interp)
uflux(:, :) = net_outputs(:,:pver)
vflux(:, :) = net_outputs(:,pver+1:)

end subroutine extract_output

Expand Down