Increase speed of global calculations of light under climate change using Dask

Question

I am implementing an algorithm for estimating light at the ocean surface as a function of wind (waves, surface roughness), chlorophyll, and zenith angle. I want to do this using climate projections from CMIP6 as input for the period 1950-2100 on a monthly basis. I use Python and Jupyter notebook to read global values of clouds, chlorophyll, and wind from Google cloud available CMIP6 climate models.

Full code is here available as Jupyter notebook.

I use the Python library pvlib to calculate direct and diffuse light at the ocean surface as a function of time of year, geographic location, and clouds from CMIP6 models. I use the Seferian et al. 2018 approach to calculate the albedo on estimated light from chlorophyll and waves for the same time and place. The bottle-neck in my code seems to be estimating the effects of waves and chlorophyll on light albedo in the function def calculate_OSAwhich estimates the reflection spectrally at all wavelengths 200-4000nm at 10 nm intervals. I use numpy vectorized to loop over all wavelengths for a given geographic grid point and I use dask.delayed to loop over all gridpoints. Gridpoints are 180x360 for global coverage.

def calculate_OSA(µ_deg, uv, chl, wavelengths, refractive_indexes, alpha_chl, alpha_w, beta_w, alpha_wc, solar_energy):
    if (µ_deg<0 or µ_deg>180):
        µ_deg=0

        µ = np.cos(np.radians(µ_deg))

        # Solar zenith angle
        # wind is wind at 10 m height (m/s)
        σ = np.sqrt(0.003+0.00512*uv)

        # Vectorize the functions
        vec_calculate_direct_reflection=np.vectorize(calculate_direct_reflection)
        vec_calculate_diffuse_reflection=np.vectorize(calculate_diffuse_reflection)
        vec_calculate_direct_reflection_from_chl=np.vectorize(calculate_direct_reflection_from_chl)
        vec_calculate_diffuse_reflection_from_chl=np.vectorize(calculate_diffuse_reflection_from_chl)

        # Direct reflection
        alpha_direct = vec_calculate_direct_reflection(refractive_indexes,µ,σ)

        # Diffuse reflection
        alpha_diffuse = vec_calculate_diffuse_reflection(refractive_indexes,σ)

        # Reflection from chlorophyll and biological pigments
        alpha_direct_chl = vec_calculate_direct_reflection_from_chl(wavelengths, chl, alpha_chl, alpha_w, beta_w, σ, µ, alpha_direct)

        # Diffuse reflection interior of water from chlorophyll
        alpha_diffuse_chl = vec_calculate_diffuse_reflection_from_chl(wavelengths, chl, alpha_chl, alpha_w, beta_w, σ, alpha_direct)

        # OSA
        return 

The entire script is written as a Jupyer notebook found here although it uses one subroutine for reading CMIP6 data and one notebook for albedo calculations. I know the script is long and complex but the main function that I believe could be improved is def calculate_OSA and the main calculate_light function. In calculate_light I believe I could improve on how I use dask and perhaps chunking, and perhaps how vectorizing the main loop in calculate_light could speed things up.

Currently, it takes 2.27 minutes to run one timestep on a mac mini with 16GB of RAM.

%%time
def calculate_light(config_pices_obj):

    selected_time=0
    wavelengths, refractive_indexes, alpha_chl, alpha_w, beta_w, alpha_wc, solar_energy = albedo.setup_parameters()
    startdate=datetime.datetime.now()

    regional=True
    create_plots=True

    southern_limit_latitude=45
    for key in config_pices_obj.dset_dict.keys():

        var_name = key.split("_")[0]
        model_name = key.split("_")[3]

        if var_name=="uas":

            key_v="vas"+key[3:]
            key_chl="chl"+key[3:]
            key_clt="clt"+key[3:]
            key_sisnconc="sisnconc"+key[3:]
            key_sisnthick="sisnthick"+key[3:]
            key_siconc="siconc"+key[3:]
            key_sithick="sithick"+key[3:]

            var_name_v = key_v.split("_")[0]
            model_name_v = key_v.split("_")[3]

            print("=> model: {} variable name: {}".format(key, var_name))
            print("=> model: {} variable name: {}".format(key_v, var_name_v))

            if model_name_v==model_name:

                ds_uas=config_pices_obj.dset_dict[key].isel(time=selected_time) 
                ds_vas=config_pices_obj.dset_dict[key_v].isel(time=selected_time)
                ds_chl=config_pices_obj.dset_dict[key_chl].isel(time=selected_time)
                ds_clt=config_pices_obj.dset_dict[key_clt].isel(time=selected_time)
                ds_sisnconc=config_pices_obj.dset_dict[key_sisnconc].isel(time=selected_time)
                ds_sisnthick=config_pices_obj.dset_dict[key_sisnthick].isel(time=selected_time)
                ds_siconc=config_pices_obj.dset_dict[key_siconc].isel(time=selected_time)
                ds_sithick=config_pices_obj.dset_dict[key_sithick].isel(time=selected_time)

                if regional:
                    ds_uas=ds_uas.sel(y=slice(southern_limit_latitude,90))
                    ds_vas=ds_vas.sel(y=slice(southern_limit_latitude,90))
                    ds_chl=ds_chl.sel(y=slice(southern_limit_latitude,90))
                    ds_clt=ds_clt.sel(y=slice(southern_limit_latitude,90))
                    ds_sisnconc=ds_sisnconc.sel(y=slice(southern_limit_latitude,90))
                    ds_sisnthick=ds_sisnthick.sel(y=slice(southern_limit_latitude,90))
                    ds_siconc=ds_siconc.sel(y=slice(southern_limit_latitude,90))
                    ds_sithick=ds_sithick.sel(y=slice(southern_limit_latitude,90))

                # Regrid to cartesian grid:
                # For any Amon related variables (wind, clouds), the resolution from CMIP6 models is less than 
                # 1 degree longitude x latitude. To interpolate to a 1x1 degree grid we therefore first interpolate to a 
                # 2x2 degrees grid and then subsequently to a 1x1 degree grid.

                ds_out_amon = xe.util.grid_2d(-180,180,2,southern_limit_latitude,90,2) 
                ds_out = xe.util.grid_2d(-180,180,1,southern_limit_latitude,90,1) #grid_global(1, 1)

                dr_out_uas_amon=regrid_variable("uas",ds_uas,ds_out_amon,transpose=True).to_dataset()
                dr_out_uas=regrid_variable("uas",dr_out_uas_amon,ds_out,transpose=False)

                dr_out_vas_amon=regrid_variable("vas",ds_vas,ds_out_amon,transpose=True).to_dataset()
                dr_out_vas=regrid_variable("vas",dr_out_vas_amon,ds_out,transpose=False)

                dr_out_clt_amon=regrid_variable("clt",ds_clt,ds_out_amon,transpose=True).to_dataset()
                dr_out_clt=regrid_variable("clt",dr_out_clt_amon,ds_out,transpose=False)
                dr_out_chl=regrid_variable("chl",ds_chl,ds_out,transpose=False)

                dr_out_sisnconc=regrid_variable("sisnconc",ds_sisnconc,ds_out,transpose=False)
                dr_out_sisnthick=regrid_variable("sisnthick",ds_sisnthick,ds_out,transpose=False)
                dr_out_siconc=regrid_variable("siconc",ds_siconc,ds_out,transpose=False)
                dr_out_sithick=regrid_variable("sithick",ds_sithick,ds_out,transpose=False)

                # Calculate scalar wind and organize the data arrays to be used for  given timestep (month-year)
                wind=np.sqrt(dr_out_uas**2+dr_out_vas**2).values

                lat=dr_out_uas.lat.values
                lon=dr_out_uas.lon.values

                clt=dr_out_clt.values
                chl=dr_out_chl.values
                sisnconc=dr_out_sisnconc.values
                sisnthick=dr_out_sisnthick.values
                siconc=dr_out_siconc.values
                sithick=dr_out_sithick.values

                m=len(wind[:,0])
                n=len(wind[0,:])
                month=6

                all_direct=[]
                all_OSA=[]
                for hour_of_day in range(12,13,1):
                    print("Running for hour {}".format(hour_of_day))

                    calc_radiation = [dask.delayed(radiation)(clt[j,:],lat[j,0],month,hour_of_day) for j in range(m)]

                    # https://github.com/dask/dask/issues/5464   
                    rad = dask.compute(calc_radiation, scheduler='processes')
                    rads=np.asarray(rad).reshape((m, n, 3))

                    zr = [dask.delayed(calculate_OSA)(rads[i,j,2], wind[i,j], chl[i,j], wavelengths, refractive_indexes, 
                                                    alpha_chl, alpha_w, beta_w, alpha_wc, solar_energy) 
                                  for i in range(m) 
                                  for j in range(n)]

                    OSA = np.asarray(dask.compute(zr)).reshape((m, n, 2))
                    nlevels=np.arange(0.01,0.04,0.001)

                    irradiance_water = (rads[:,:,0]*OSA[:,:,0]+rads[:,:,1]*OSA[:,:,1])/(OSA[:,:,0]+OSA[:,:,1])

                    print("Time to finish {} with mean OSA {}".format(datetime.datetime.now()-startdate,
                          np.mean(irradiance_water)))

                    # Write to file
                    data_array=xr.DataArray(data=irradiance_water,dims={'lat':lat,'lon':lon})
                    if not os.path.exists("ncfiles"):
                        os.mkdir("ncfiles")
                    data_array.to_netcdf("ncfiles/irradiance.nc")

Since I need to run this script for several CMIP6 models for 3 socio-economic pathways (SSP). For each model and SSP I have to calculate monthly light values for 150 years, spectrally for 140 wavelengths, on a global scale of 1x1 degrees resolution. This is CPU and memory consuming and I wonder if there are ways of improving my vectorization or better approaches for utilizing Dask. It would be great if someone could point me in the right direction for how to improve speedup.

Answer 1

Readability

formatting

You have very long lines, and don't follow the PEP8 suggestions everywhere. The quickest way to solve both problems in one go is to use black. this can be integrated in most IDEs and in jupyterlab

type hints

In this I have to agree with Reinderein. Now it is not clear which parameters to your function are scalars, and which are arrays. That makes it difficult for other people (this includes you in a few months of not working with this code) to understand what happens. I have a rather strict mypy configuration

[mypy]
check_untyped_defs = true
disallow_any_generics = true
disallow_untyped_defs = true
ignore_missing_imports = true
no_implicit_optional = true
warn_redundant_casts = true
warn_return_any = true
warn_unused_ignores = true

but this has allowed me to remove some bugs that would have been hard to spot otherwise.

To type a notebook, I use jupytext to sync the notebook with a python file, open that python file in an IDE and run a battery of linters (pylama, pydocstyle, ..), code formatters (isort and black), type check (mypy), adapt the code to the suggestions. then I go back to the notebook, and run everything to make sure the changes did not affect the calculations' correctness.

This .py file can then also be more easily versioned.

speedup

Vectorise as much as possible. You can use numba to speed up some calculations.

As an outsider it is difficult to see what parameters to function tend to change, and which stay constant. memoization can cache intermediate results. arrays are not hashable, so you won't be able to use functools.lru_chache, but there are third party modules that can help, like joblib.Memory

rearrange

your calculate_light is too complex. It also mixes in system input (datetime.datetime.now()), calculations and sytem output (print and writing the file to disc)

logging

Instead of print, I would use the logging module. This allows you, or users of this code to later very easily switch off printing,, allows you to write it to a log file and inspect later, ...)

output

Doesn't data_array.to_netcdf("ncfiles/irradiance.nc") overwrite the results in every iterations.

Apart from that I have 2 problems with this. You hardcode the output path in this function. If ever you want the results somewhere else, this is difficult to do.

But I would not write the results in this method. I would yield the results, and let the caller of this method worry on what to do with them. If the results are intermediate, you don't need them afterwards, you can keep em in memory if you have enough RAM, or write them to a temporary directory

negative checks / continue

You have some checks like if var_name=="uas": and if model_name_v==model_name:. If you reverse those checks, you save a level of indentation

if var_name != "uas":
    continue
...

DRY

You have a lot of repetition. For example the key[3:] If you need to change this to the 4th number, you need to think about changing all these intermediate positions. Extract that into its own variable. This will also serve as extra documentation

General

Try to implement these changes already. If you do, the code will be a lot more readable and understandable for outsiders, so we can give better advice on how to speed up certain parts, then you can post a new question.

Answer 2

These will not impact performance, but are useful to address nonetheless:

Type hints

Some wild guesses here, but:

def calculate_OSA(
    µ_deg: float,
    uv: float,
    chl: float,
    wavelengths: ndarray,
    refractive_indexes: ndarray,
    alpha_chl: float,
    alpha_w: float,
    beta_w: float,
    alpha_wc: float,
    solar_energy: float,
):

That said, given the high number of parameters, it may be easier to make a @dataclass with typed members and either pass that as an argument or make a method on it.

No-op return

Your final return can be deleted. But it's suspicious that alpha_diffuse, alpha_direct_chl and alpha_diffuse_chl are unused. Looking at your Github, it seems that you forgot to copy the call to calculate_spectral_and_broadband_OSA here.

Answer 3

Looking at the jupyter notebook, I wonder if a bit of caching might help? How many of those datapoints are really unique? Something as simple as wrapping the often-called functions in a memoization decorator might help. Any of the calculate_ functions that take just floats are good candidates - I don't think memoizing anything that takes vectors would help.