3
\$\begingroup\$

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.

\$\endgroup\$
6
  • 1
    \$\begingroup\$ Did you try profiling your code to see where the bottleneck is? Something that can help with profiling (and readability) is splitting your code in functions \$\endgroup\$ Commented Jun 4, 2020 at 16:28
  • \$\begingroup\$ @MaartenFabré Yes I did and the problem is the call `zr = [dask.delayed(calculate_OSA).... This loops over m x n grid points and then for each grid points loops over 140 wavelengths using numpy vectorized. I am hoping that the structure of my code could be changed to improve speed. \$\endgroup\$ Commented Jun 4, 2020 at 16:40
  • 1
    \$\begingroup\$ From the np.vectorize documentation: The vectorize function is provided primarily for convenience, not for performance. The implementation is essentially a for loop.. This will not result in a speedup. I would look into numba. I've seen it give a large speedup \$\endgroup\$ Commented Jun 4, 2020 at 16:47
  • \$\begingroup\$ @MaartenFabré That was a good suggestion. I removed the use of Dask on calculate_OSAand instead inserted use of numba. That reduced time of execution from 2 min 47s to 1 min 35s. \$\endgroup\$ Commented Jun 4, 2020 at 18:32
  • 1
    \$\begingroup\$ Hey, I've manually removed the change in Rev 3, as your change made half of Reinderien's answer invalid. Please understand that invalidating answers goes against the Q&A format as new users that come to the post can be confused by the disconnect between your question and the answers. In the future you can post a follow up question to get more feedback. \$\endgroup\$
    – Peilonrayz
    Commented Jun 4, 2020 at 22:37

3 Answers 3

2
+50
\$\begingroup\$

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.

\$\endgroup\$
5
\$\begingroup\$

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.

\$\endgroup\$
1
  • \$\begingroup\$ Thank you! I had forgotten to add the continued call to calculate_spectral_and_broadband_OSA. I agree explicit type makes it more readable. \$\endgroup\$ Commented Jun 3, 2020 at 2:59
1
\$\begingroup\$

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.

\$\endgroup\$
3
  • \$\begingroup\$ Unfortunately, most of the functions vary for each calculation as they take input such as clouds, sun position, waves etc. as input which varies for each timestep \$\endgroup\$ Commented Jun 4, 2020 at 18:53
  • \$\begingroup\$ They may vary less than you think, based on measurement precision - you'd have to look at the dataset to be sure. Might be worth trying memoization anyway - shouldn't cost much time and could gain you a bit. calculate_alpha_dir, in particular, looks ripe - it often gets called with a constant 1.34 as its initial value, so any time µ repeats in the data, memoization will save a bit of time. \$\endgroup\$
    – pjz
    Commented Jun 4, 2020 at 19:08
  • 1
    \$\begingroup\$ Thanks for your help. I tested with the use of memoized and I am sure it would have a great improvement, but the code can not be combined with the use of numba. I have now tested with numba as suggested by @MaartenFabré and reduced the time of calculation to half of what it used to be. If you know of a way to mix these tools that would be great, but I assume I have to pick one approach. \$\endgroup\$ Commented Jun 4, 2020 at 20:59

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.