Initial implementation of pyscream module(s) for running eamxx from python #2851
Conversation
|
Here's an example of how it can be used. It still crashes, but gives an idea. Also, forget about units, I got tired and just wanted to get all fields in. #!/usr/bin/env python3
# These lines won't be needed once we package things correctly
import sys
sys.path.append('/home/lbertag/workdir/scream/scream-src/master/components/eamxx/ctest-build/full_debug/src/share/python')
sys.path.append('/home/lbertag/workdir/scream/scream-src/master/components/eamxx/ctest-build/full_debug/src/physics/p3/python')
from mpi4py import MPI
import pyscream as ps
import pyp3 as P3
from pyscream.units import m,s,kg,K,N,J,W,Pa,one
#######################################
def run ():
#########################################
# Create a comm, and print specs
comm = MPI.COMM_WORLD
print (f"hello from rank{comm.Get_rank()}/{comm.Get_size()}")
# Create a physics-like grid
ncols = 100
nlevs = 128
grid = ps.Grid("Physics",ncols,nlevs)
# Check that units operators overloads work
m2 = m*m
Wm2 = W / m2
print (f"Wm2: {Wm2.str()}")
# Create a bunch of (managed) fields
T_mid = grid.scalar3d_mid("T_mid",K)
qv = grid.scalar3d_mid("qv",kg/kg)
qc = grid.scalar3d_mid("qc",kg/kg)
qi = grid.scalar3d_mid("qi",kg/kg)
qr = grid.scalar3d_mid("qr",kg/kg)
qm = grid.scalar3d_mid("qm",kg/kg)
nc = grid.scalar3d_mid("nc",one/kg)
nr = grid.scalar3d_mid("nr",one/kg)
ni = grid.scalar3d_mid("ni",one/kg)
bm = grid.scalar3d_mid("bm",one/kg)
precip_l = grid.scalar2d("precip_liq_surf_mass",kg)
precip_i = grid.scalar2d("precip_ice_surf_mass",kg)
eff_r_qc = grid.scalar3d_mid("eff_radius_qc",m)
eff_r_qi = grid.scalar3d_mid("eff_radius_qi",m)
eff_r_qr = grid.scalar3d_mid("eff_radius_qr",m)
p_mid = grid.scalar3d_mid("p_mid",m)
p_dry_mid = grid.scalar3d_mid("p_dry_mid",m)
rho = grid.scalar3d_mid("pseudo_density",m)
rho_dry = grid.scalar3d_mid("pseudo_density_dry",m)
cldfrac_tot = grid.scalar3d_mid("cldfrac_tot",m)
qv_prev = grid.scalar3d_mid("qv_prev_micro_step",m)
T_prev = grid.scalar3d_mid("T_prev_micro_step",m)
rain_frac = grid.scalar3d_mid("rainfrac",m)
nc_nuceat_tend = grid.scalar3d_mid("nc_nuceat_tend",m)
nccn = grid.scalar3d_mid("nccn",m)
ni_activated = grid.scalar3d_mid("ni_activated",m)
inv_qc_relvar = grid.scalar3d_mid("inv_qc_relvar",m)
liq_ice_exchange = grid.scalar3d_mid("micro_liq_ice_exchange",m)
vap_liq_exchange = grid.scalar3d_mid("micro_vap_liq_exchange",m)
vap_ice_exchange = grid.scalar3d_mid("micro_vap_ice_exchange",m)
# Create an init p3
p3 = P3.P3(grid)
p3.set_fields([qv,T_mid,qc,qi,qr,qm,nc,nr,ni,bm,precip_l,precip_i,eff_r_qc,eff_r_qi,eff_r_qr,p_mid,p_dry_mid,rho,rho_dry,cldfrac_tot,qv_prev,rain_frac,nc_nuceat_tend,nccn,ni_activated,inv_qc_relvar,T_prev,liq_ice_exchange,vap_liq_exchange,vap_ice_exchange])
p3.initialize()
#######################################
def main ():
#######################################
# This level of indirection ensures all pybind structs are destroyed
# before we finalize eamxx (and hence kokkos)
ps.init()
run ()
ps.finalize()
####################################
if __name__ == "__main__":
main() |
|
quick thought - we should probably call the library pyeamxx or something like that since folks will be using it for low-res applications as well. Though pyeamxx sounds lame. py3sm? eampy? |
The main package (currently someone --- I don't know who!!! --- is calling it screaminpy) should definitely be changed to something more agreeable. The names of wrappers of components should match their nearest neighbor imho (so things like pyscream, as the python of scream_share, should remain as internal components) |
I thought about that too. |
I can switch the overall package (in eamxx/src/python) to pyeamxx later today. I don't think we need to change pyscream to that since if I understand your logic correctly it is just a wrapper for scream_share lib (which is an actual thing in our code right now, eamxx/src/share compiles to scream_share). I think it is important to think about this in terms of what faces the user: ultimately, the user will only interact with the package in eamxx/src/python which will have more python-centric tools that we will write (for example, hiding away all of the details of init, grid, etc. stuff unless the user wants to override them). In your code snippet above, I envision the user to do: from pyeamxx import pyscream as ps
from pyeamxx import pyp3 as P3
#######################################
def run ():
#########################################
# Create a physics-like grid
ncols = 100
nlevs = 128
with ps.Grid("Physics",ncols,nlevs) as psg:
# this with-as logic allows us to do a lot of init/de-init stuff
# by merely entering and existing this block, including MPI
# Create an init p3
p3 = P3.P3(psg)
# we will write an abstraction such that this
# P3.P3 call endows psg with all the fields needed
# init-ed to something reasonable (or via p3_input.yaml if available)
# users can access them as p3.T_mid, etc.
p3.initialize()
# we can add more convenient functionalities here based on user feedback
# exposing/hiding things as needed.
# NOTE: we will enable the user to override our hiding of things!
#######################################
def main ():
#######################################
run ()
####################################
if __name__ == "__main__":
main() |
|
Uhm, that may be a better approach. Do not let the user create fields, just access the ones they need. I'll work on that. It may also allow to trim down the pygrid file, since we don't need to expose those capabilities anymore. |
I don't think you need to change anything on the cpp end, we can do all the hiding/exposing in python 😉 |
There was a problem hiding this comment.
Two quick technical comments (for the record here, in case someone else wants to weigh in, but we will continue privately)
- pybind11 vs nanobind. So, nanobind is young and much more exciting. I couldn't get it to work. Did you test it at all? Since you're much more familiar with the build system, I think you will likely get it work. If we can get nanobind to work, we are better off just going with nanobind from the get-go. Also, for my understanding, how does cmake find the pybind11 submodule? Is there a logic in up dir to find all the stuff in externs?
- array api. I think you're using the pybind11/numpy headers to expose stuff. That's fine for now. But like pybind11 vs nanobind, if we are able to get the more general array api to work from the get-go, that will save us a lot of time later. I can look into this point myself to see where the state of the art is these days; I recently saw the generic array api being used with nanobind fwiw. My hope is that it will be simply transferrable with very minor edits.
|
It may limit us to numpy arrays exclusively, which may become burdensome down the line. Supporting a generic array api will enable us to exchange arbitrary arrays as long as they conform to the api. An example is: consider in the future we want to take this to device, numpy doesn't support gpus and won't, so we will need to figure out some workarounds. Another: general ML applications and frameworks with their "arrays" (tensors). Overall, I think the interoperability is robust and we can live with numpy.h if we have to, but if we don't have to, we should prefer the generic api (e.g., one offered by nanobind). side reading on the "python array api" https://data-apis.org/array-api/latest/purpose_and_scope.html (nanobind has something else, which seems operational: https://nanobind.readthedocs.io/en/latest/ndarray.html) |
we likely should set this up under eamxx/pippkg or something, tbd
* All other files are hpp files included by pyscream.cpp * Each file handle separate EAMxx concept * Remove hard-coding of pybind11_DIR. User must provide pybind11_ROOT cmake/env var
* Add header guards to pyxyz.hpp * Add PyGrid type * PyGrid can create managed/unmanaged fields * Add constructors to PyField (not all py-interoperable)
Start the pattern where folder blah has python subfolder if pybind hooks are supported
When needed, build a GridsManager on the fly, using the ad-hoc type SingleGridGM
* Fields are created by pyscream, no possibility to pass pre-built numpy array * Add generic impl of atm proc wrapper in PyAtmProc * Users can grab fields out of a py AP based on field name * Removed pointless stuff from py grid * Removed no longer needed wrappers for units
bartgol
force-pushed
the
bartgol/eamxx/pyscream
branch
from
0c81180 to
253774d
Compare
With PUBLIC, the compiled mod files magically appeared in the build folder of targets that link against tms. I'm not sure if this is the right fix, but it seems to work
|
@mahf708 I was puzzled an confused by the two modules (pyeamxx and libpyeamxx_ext), so I renamed the second to pyeamxx. My guess is that the two clash when you try to package it as a conda package? I'm not sure. We can discuss when you get back. Meanwhile, pyeamxx can now be used with any of the physics parametrizations. The only limitation is that the user must create a grid for which we have an IC file. We can work on adding some logic (in eamxx or pyeamxx) to allow reading an existing IC file (with N cols) into a grid with 1 col, perhaps maybe recycling IOP implementation: specify the lat/lon of the col, and reading in the closest col to that lat/lon. When Naser returns, we can iterate on the pyeamxx conda package impl, which I'm confident I broke. Meanwhile, pyeamxx is available as follows: $ cd components/eamxx
$ ./scripts/test-all-scream -t opt -m MACHINE --config-only
$ cd ctest-build/release/src/python
$ make -jAnd then in your python script, do something like import sys
sys.path.append('/path/to/eamxx/ctest-build/full_debug/src/python')
sys.path.append('/path/to/eamxx/scripts')
import mpi4py
mpi4py.rc.initialize = False # do not initialize MPI automatically
mpi4py.rc.finalize = False # do not finalize MPI automatically
from mpi4py import MPI
import pyeamxx
from pathlib import Path
from utils import ensure_yaml
ensure_yaml()
import yaml
def main():
# Get timestepping params
with open('input.yaml','r') as fd:
yaml_input = yaml.load(fd,Loader=yaml.SafeLoader)
nsteps = yaml_input['time_stepping']['number_of_steps']
dt = yaml_input['time_stepping']['time_step']
t0_str = yaml_input['time_stepping']['run_t0']
# Create the grid
ncols = 218
nlevs = 72
grid = pyeamxx.Grid("Physics",ncols,nlevs)
ic_file = Path('/home/lbertag/workdir/e3sm/e3sm-data/inputdata/atm/scream/init/screami_unit_tests_ne2np4L72_20220822.nc')
rrtmgp_params = pyeamxx.ParameterList(yaml_input['atmosphere_processes']['rrtmgp'],'rrtmgp')
rrtmgp = pyeamxx.AtmProc(rrtmgp_params,grid)
missing = rrtmgp.read_ic(str(ic_file))
print (f"WARNING! The following input fields were not found in the IC file, and must be manually initialized: {missing}")
rrtmgp.initialize(str(t0_str))
rrtmgp.setup_output("my_output_rrtmgp.yaml")
# Time looop
for n in range(0,nsteps):
rrtmgp.run(dt)
if __name__ == "__main__":
# This level of indirection ensures all pybind structs are destroyed
# before we finalize eamxx (and hence kokkos)
MPI.Init()
pyeamxx.init()
main ()
pyeamxx.finalize()
MPI.Finalize()The above assumes that @PeterCaldwell @hassanbeydoun I know you guys were interesting in running rrtmgp for multiple IC perturbations. If you want, you can start using the above example. You can access all input/output fields of the parametrization, and modify them at will. E.g.: def main():
...
for ens in range(0,num_ensembles):
# create rrtmgp as above
# f is a numpy array, of the proper dimension
f = rrtmgp.get_field('some_input')
data = f.get()
data[1,2,3] *= some_rand_number
f.sync_to_dev()
for n in range(0,nsteps):
rrtmgp.run(dt) |
bartgol
force-pushed
the
bartgol/eamxx/pyscream
branch
from
605d452 to
8d8e2af
Compare
|
Status Flag 'Pull Request AutoTester' - Testing Jenkins Projects: Pull Request Auto Testing STARTING (click to expand)Build InformationTest Name: SCREAM_PullRequest_Autotester_Mappy
Jenkins Parameters
Build InformationTest Name: SCREAM_PullRequest_Autotester_Weaver
Jenkins Parameters
Using Repos:
Pull Request Author: bartgol |
|
Status Flag 'Pull Request AutoTester' - Jenkins Testing: all Jobs PASSED Pull Request Auto Testing has PASSED (click to expand)Build InformationTest Name: SCREAM_PullRequest_Autotester_Mappy
Jenkins Parameters
Build InformationTest Name: SCREAM_PullRequest_Autotester_Weaver
Jenkins Parameters
|
|
Status Flag 'Pre-Merge Inspection' - - This Pull Request Requires Inspection... The code must be inspected by a member of the Team before Testing/Merging |
|
All Jobs Finished; status = PASSED, target_sha=d0bbf72587216ad2cec0d6e3a0ec3ff0fa902e96, However Inspection must be performed before merge can occur... |
|
Status Flag 'Pull Request AutoTester' - Testing Jenkins Projects: Pull Request Auto Testing STARTING (click to expand)Build InformationTest Name: SCREAM_PullRequest_Autotester_Mappy
Jenkins Parameters
Build InformationTest Name: SCREAM_PullRequest_Autotester_Weaver
Jenkins Parameters
Using Repos:
Pull Request Author: bartgol |
|
Status Flag 'Pull Request AutoTester' - Jenkins Testing: 1 or more Jobs FAILED Note: Testing will normally be attempted again in approx. 2 Hrs. If a change to the PR source branch occurs, the testing will be attempted again on next available autotester run. Pull Request Auto Testing has FAILED (click to expand)Build InformationTest Name: SCREAM_PullRequest_Autotester_Mappy
Jenkins Parameters
Build InformationTest Name: SCREAM_PullRequest_Autotester_Weaver
Jenkins Parameters
SCREAM_PullRequest_Autotester_Mappy # 5580 FAILED (click to see last 100 lines of console output)SCREAM_PullRequest_Autotester_Weaver # 5832 PASSED (click to see last 100 lines of console output) |
|
Status Flag 'Pull Request AutoTester' - User Requested Retest - Label AT: RETEST will be reset after testing. |
|
Status Flag 'Pull Request AutoTester' - Testing Jenkins Projects: Pull Request Auto Testing STARTING (click to expand)Build InformationTest Name: SCREAM_PullRequest_Autotester_Mappy
Jenkins Parameters
Build InformationTest Name: SCREAM_PullRequest_Autotester_Weaver
Jenkins Parameters
Using Repos:
Pull Request Author: bartgol |
|
Status Flag 'Pull Request AutoTester' - Jenkins Testing: all Jobs PASSED Pull Request Auto Testing has PASSED (click to expand)Build InformationTest Name: SCREAM_PullRequest_Autotester_Mappy
Jenkins Parameters
Build InformationTest Name: SCREAM_PullRequest_Autotester_Weaver
Jenkins Parameters
|
The desire is to run a single parametrization from a python script. This requires a few hooks, that this PR tries to put together. Work is still ongoing, and the PR may be rebased/reset/closed/reopen.
The idea is to create a separate python module for each scream library:
The wrappers do NOT expose all the methods of the wrapped class, nor they are built in the same way as the wrapped class. Instead, we expose the minimum necessary to be able to create/setup/use scream objects.
If you want to try building this, you need to
-c EAMXX_ENABLE_PYBIND=ON -c BUILD_SHARED_LIBS=ONLots of stuff to do/decide