Add Ai2's ACE as an emulator component by andrewdnolan · Pull Request #7957 · E3SM-Project/E3SM

andrewdnolan · 2025-12-18T21:15:29Z

This PR adds support for emulator components, starting with an emulator atmosphere based on Ai2's ACE.

andrewdnolan · 2025-12-18T21:35:36Z

+        vbot(i, j) = net_outputs(1, 32, i, j) ! V_0
+        tbot(i, j) = net_outputs(1,  8, i, j) ! T_0
+        ptem(i, j) = net_outputs(1,  8, i, j) ! T_0
+        shum(i, j) = 0.0_R8                   ! TBD


datm calculates shum (bottom atm level spec humidity) here:

E3SM/components/data_comps/datm/src/datm_comp_mod.F90

Lines 1106 to 1134 in cae5c62

!--- specific humidity ---

tbot = a2x%rAttr(ktbot,n)

pbot = a2x%rAttr(kpbot,n)

if (sshum > 0) then

e = datm_shr_esat(tbot,tbot)

qsat = (0.622_R8 * e)/(pbot - 0.378_R8 * e)

if (qsat < a2x%rAttr(kshum,n)) then

a2x%rAttr(kshum,n) = qsat

endif

else if (srh > 0) then

e = avstrm%rAttr(srh,n) * 0.01_R8 * datm_shr_esat(tbot,tbot)

qsat = (0.622_R8 * e)/(pbot - 0.378_R8 * e)

a2x%rAttr(kshum,n) = qsat

if(wiso_datm) then

! isotopic forcing

! For tracer specific humidity, lnd_import_mct expects a delta, so

! just keep the delta from the input file - TW

a2x%rAttr(kshum_16O,n) = avstrm%rAttr(srh_16O,n)

a2x%rAttr(kshum_18O,n) = avstrm%rAttr(srh_18O,n)

a2x%rAttr(kshum_HDO,n) = avstrm%rAttr(srh_HDO,n)

end if

else if (stdew > 0) then

if (tdewmax < 50.0_R8) avstrm%rAttr(stdew,n) = avstrm%rAttr(stdew,n) + tkFrz

e = datm_shr_esat(avstrm%rAttr(stdew,n),tbot)

qsat = (0.622_R8 * e)/(pbot - 0.378_R8 * e)

a2x%rAttr(kshum,n) = qsat

else

call shr_sys_abort(subname//'ERROR: cannot compute shum')

endif

I figure we can probably use the specific_total_water_0 field to calculate shum, but I'm not certain about the details here...

specific total water in ace is defined as qv+qc+qi+qr (basically all water), but I think specific humidity here is just qv (vapor). I think this is somewhat important, so we shouldn't use specific total water (except for prelim testing)

andrewdnolan · 2025-12-18T21:50:15Z

+        ubot(i, j) = net_outputs(1, 24, i, j) ! U_0
+        vbot(i, j) = net_outputs(1, 32, i, j) ! V_0
+        tbot(i, j) = net_outputs(1,  8, i, j) ! T_0
+        ptem(i, j) = net_outputs(1,  8, i, j) ! T_0


datm sets ptem equal to tbot, so we have followed suite here for now.
https://github.com/E3SM-Project/E3SM/blob/cae5c6237cfd379968960ebfef7802b19a2b59dd/components/data_comps/datm/src/datm_comp_mod.F90#L1034C36-L1034C41

we can also calculate the potential temperature, e.g., https://en.wikipedia.org/wiki/Potential_temperature

rljacob · 2025-12-18T22:53:53Z

    Where for the EAM specific compsets below the following is supported
    TIME = Time period (e.g. 2000, HIST, RCP8...)
-    ATM  = [EAM, EAMXX, SATM, SCREAM]
+    ATM  = [EAM, EAMXX, DATM, EATM, SATM, SCREAM]


What does "EATM" stand for?

Emulated atmosphere, following the c onvention of data, stub, etc. atmospheres --- this pr isn't ready for review yet, we are working on addressing a few items internally first, but wanted to start iterating publicly #open-dev

we are calling them "emulator_comps"

sort of parallel to data_comps

That's a good name for the directory holding them. It doesn't have to be 4 letters so how about "EMATM" to make it more obvious.

I would vote to keep the following convention:

comp purpose

cmp active, physics (flagship version)

ecmp active, emulated (fast version)

dcmp "data" or just reading files

scmp "stub" or no-op

xcmp "dead" (only in testing)

but I don't care what what each specific one is called; note that eatm will (in the longer term) represent several of the emulated atmosphere models, not a specific one

In other words, I think it is less than ideal to have "eam", "eamxx" and "scream" etc. as components (replacing atm). Those all should be just called atm as comps, and then once in the atm comp, one can decide if they want to do one eam, eamxx, scream, etc. --- it makes better sense to organize components that way imho

rljacob · 2025-12-18T22:55:40Z

@@ -0,0 +1,132 @@
+#!/usr/bin/env python3


Thank you for not copying one of the perl versions of buildnml

I assume you're joking, right? lol

see components/emulator_comps/eatm/bld/build-namelist, but ideally, we should get rid of this

just cut-and-paste!

rljacob · 2025-12-18T23:29:58Z

@@ -0,0 +1,589 @@
+#!/usr/bin/env perl


Is this a copy of the EAM build-namelist? Why didn't you just copy the datm one which is all python?

I was just cutting-and-pasting to get something working. Our first task was to get something working for Luke's scidac as quickly as possible, with the intent of cleaning it up later

The question still stands. Why not copy the datm one? You don't need the complexity of the EAM namelist.

I'm happy to replace it once we get something working

mahf708

@andrewdnolan, this is excellent! Thanks for the PR and the dilligent work with @jonbob on putting it together. I responded to all your questions. In short, they are nuanced. We have two pathways:

Take easiset shortcuts to demostrate things are working, without worrying about scientific validity, and only assess if ace starts crashing like crazy
Preemptively start assembling utils in a shared location that can do these calculations in a documented way

I'd defer to you to decide; I would obviously prefer 2 longer term, but happy to wait for it.

If you decide we should go for 2, let's briefly discuss how we want to do this. A mild personal preference is that we write these utils in C++ (not F90). Again, that's a longer term issue for me: I want us to move to C++ as soon as we are able to do so, and I'd rather see close to 0% code in F90. We can write these small utils/things in Python too.

Another important todo item (that can also wait, so no pressure to get it done now) is to start writing docs about what we are doing alongside PRs, i.e., in components/emulator_comps/docs; I'm happy to put these docs together once we are closer to integrating (unless you'd rather take the lead on this).

Let's keep iterating until we feel more confident. I have a few commits I'd like to push in the near future.

mahf708 · 2025-12-19T04:43:54Z

+        vbot(i, j) = net_outputs(1, 32, i, j) ! V_0
+        tbot(i, j) = net_outputs(1,  8, i, j) ! T_0
+        ptem(i, j) = net_outputs(1,  8, i, j) ! T_0
+        shum(i, j) = 0.0_R8                   ! TBD


specific total water in ace is defined as qv+qc+qi+qr (basically all water), but I think specific humidity here is just qv (vapor). I think this is somewhat important, so we shouldn't use specific total water (except for prelim testing)

mahf708 · 2025-12-19T04:45:49Z

+        ubot(i, j) = net_outputs(1, 24, i, j) ! U_0
+        vbot(i, j) = net_outputs(1, 32, i, j) ! V_0
+        tbot(i, j) = net_outputs(1,  8, i, j) ! T_0
+        ptem(i, j) = net_outputs(1,  8, i, j) ! T_0


we can also calculate the potential temperature, e.g., https://en.wikipedia.org/wiki/Potential_temperature

mahf708 · 2026-01-05T00:49:45Z

An alternative design of this, almost entirely in C++ is in: #7964

Switch to layer 7 (c.f. 0) for the "bottom" layer index

And load the ACE restart file as R4 array

stores eatm output at t_i-1 and t_i+1 emulator evaluation states

Needed b/c ACE provides data at 6-hour intervals, but coupling interval can be much smaller than that (e.g. 30-minutes). Unles we linearly interpolate the ACE prediction over the coupling interval the ACE predicted changes are too large and cause stability issue for `ice` model. Plus this ensure that the `net_output` array is synchronized with the coupler clock.

Fixes issue where ACE output had large transience even when forcing data was constant

to link link

No longer need to read the ace restart file. Also dropped all ACE specific netcdf i/o in favor the reatrt read/write functionality.

normalization, denormalization, and correction steps

andrewdnolan added the AI and emulators label Dec 18, 2025

andrewdnolan commented Dec 18, 2025

View reviewed changes

rljacob reviewed Dec 18, 2025

View reviewed changes

mahf708 reviewed Dec 19, 2025

View reviewed changes

jonbob and others added 22 commits February 23, 2026 08:32

Initial commit of files to support a new eatm component

4996248

Clean up to get building and running through init

9e930b7

Get rid of hardwired filename now that the namelist works

f81fafe

Add descriptors to eatm fields

9426c00

Read lat/lon info from SCRIP file

c24ffa3

Rename exported array using CPL naming convention

eade9e6

Add ACE module and allocate input/output arrays

1505539

Correct indicies for fields directly from ACE

63d619d

Switch to layer 7 (c.f. 0) for the "bottom" layer index

Add calcs for dens and shum following datm.

f628351

Add support for ftorch on pm-cpu

ad50b67

avoid clash of declarations

cd1313d

Add support for reading R4 data with PIO

4d0f07a

And load the ACE restart file as R4 array

Reorder loop order

4d00594

Add inference using FTorch

aa7ae34

Add derived type used for linear interpolation

cefa36d

stores eatm output at t_i-1 and t_i+1 emulator evaluation states

minor fixes for warnings and code clean up

8607d01

Fix update to i+1 array

a97f43f

Update gauss180x360 scrip, mapping, and domain files

dd7fc7f

Add normalize/denormalize functionality

695e4d2

WIP: lat correction and more debugging

63219b1

Fix normalization and clip snow/rain

de186b5

andrewdnolan added 3 commits March 2, 2026 12:21

Add scratch net input array for normalization

d888b5f

Fixes issue where ACE output had large transience even when forcing data was constant

FTorch clean up

b66a171

Update pytorch (2.8.0) and add pytorch nccl

32450e6

to link link

andrewdnolan force-pushed the jonbob/add-eatm branch from e4d9dcb to 32450e6 Compare March 3, 2026 18:14

andrewdnolan and others added 13 commits March 5, 2026 16:49

Manually calculate lndfrac from ocn and ice frac

3c4d873

Pass coupler fields to ACE

c1fb9ef

Add restart read/write capability for ACE

892bc1e

build with GPU enabled FTorch and run ACE on GPU

76d57e1

Add code to abort if the eatm mpi communicator has more than one pe

4a6924e

Call inference on init to populate netoutputs

f492b9b

No longer need to read the ace restart file. Also dropped all ACE specific netcdf i/o in favor the reatrt read/write functionality.

Add ability to read 2D fields with time dim

247c849

BFB restarts by writing both time levels

30aa303

Read time-varying SOLIN forcing from data file

2f8fb6f

Inference with a traced model that includes...

751978d

normalization, denormalization, and correction steps

Pass nextsw_cday to cpl so albedos are updated

93e9512

add inline calculation for solin and persist phis

7fce378

Pressure altitude calculation for zbot

dfcabf9

	!--- specific humidity ---
	tbot = a2x%rAttr(ktbot,n)
	pbot = a2x%rAttr(kpbot,n)
	if (sshum > 0) then
	e = datm_shr_esat(tbot,tbot)
	qsat = (0.622_R8 * e)/(pbot - 0.378_R8 * e)
	if (qsat < a2x%rAttr(kshum,n)) then
	a2x%rAttr(kshum,n) = qsat
	endif
	else if (srh > 0) then
	e = avstrm%rAttr(srh,n) * 0.01_R8 * datm_shr_esat(tbot,tbot)
	qsat = (0.622_R8 * e)/(pbot - 0.378_R8 * e)
	a2x%rAttr(kshum,n) = qsat
	if(wiso_datm) then
	! isotopic forcing
	! For tracer specific humidity, lnd_import_mct expects a delta, so
	! just keep the delta from the input file - TW
	a2x%rAttr(kshum_16O,n) = avstrm%rAttr(srh_16O,n)
	a2x%rAttr(kshum_18O,n) = avstrm%rAttr(srh_18O,n)
	a2x%rAttr(kshum_HDO,n) = avstrm%rAttr(srh_HDO,n)
	end if
	else if (stdew > 0) then
	if (tdewmax < 50.0_R8) avstrm%rAttr(stdew,n) = avstrm%rAttr(stdew,n) + tkFrz
	e = datm_shr_esat(avstrm%rAttr(stdew,n),tbot)
	qsat = (0.622_R8 * e)/(pbot - 0.378_R8 * e)
	a2x%rAttr(kshum,n) = qsat
	else
	call shr_sys_abort(subname//'ERROR: cannot compute shum')
	endif

comp	purpose
`cmp`	active, physics (flagship version)
e`cmp`	active, emulated (fast version)
d`cmp`	"data" or just reading files
s`cmp`	"stub" or no-op
x`cmp`	"dead" (only in testing)

Conversation

andrewdnolan commented Dec 18, 2025

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

mahf708 Dec 18, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

mahf708 Dec 18, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

mahf708 left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

mahf708 commented Jan 5, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

mahf708 Dec 18, 2025 •

edited

Loading

mahf708 Dec 18, 2025 •

edited

Loading

mahf708 left a comment •

edited

Loading