distributed runs of biosnicar using Dask delayed

biosnicar-distributed.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

# save this alongside snicar-driver.py and run in the terminal using
# `python src/biosnicar/biosnicar-distributed.py`

import numpy as np
from pathlib import Path
import dask

from adding_doubling_solver import adding_doubling_solver
from column_OPs import get_layer_OPs, mix_in_impurities
from display import display_out_data, plot_albedo
from setup_snicar import setup_snicar
from toon_rt_solver import toon_solver
from validate_inputs import validate_inputs
from dask.distributed import Client

if __name__ == "__main__":
    client = Client()

    BIOSNICAR_SRC_PATH = Path(__file__).resolve().parent

    # define input file
    input_file = BIOSNICAR_SRC_PATH.joinpath("inputs.yaml").as_posix()
    print(input_file)

    lyrList = [0, 1]
    densList = [400, 500, 600, 700, 800]
    reffList = [200, 400, 600, 800, 1000]
    zenList = [30, 40, 50, 60]
    bcList = [500, 1000, 2000]
    dzList = [
        [0.08, 0.1],
        [0.10, 0.15],
        [0.2, 0.5],
        [0.3, 0.5],
        [1, 10],
    ]

    ncols = (
        len(lyrList)
        * len(densList)
        * len(reffList)
        * len(zenList)
        * len(bcList)
        * len(dzList)
    )

    @dask.delayed
    def run_using_dask(input_file, lyr, dens, reff, zen, bc, dz):
        (
            ice,
            illumination,
            rt_config,
            model_config,
            plot_config,
            impurities,
        ) = setup_snicar(input_file)

        ice.dz = dz
        ice.nbr_lyr = 2
        ice.layer_type = [lyr] * len(ice.dz)
        ice.rho = [dens] * len(ice.dz)
        ice.rds = [reff] * len(ice.dz)
        illumination.solzen = zen
        illumination.calculate_irradiance()
        impurities[0].conc = [
            bc,
            bc,
        ]  # bc in all layers
        ice.calculate_refractive_index(input_file)
        illumination.calculate_irradiance()

        ssa_snw, g_snw, mac_snw = get_layer_OPs(ice, model_config)
        tau, ssa, g, L_snw = mix_in_impurities(
            ssa_snw,
            g_snw,
            mac_snw,
            ice,
            impurities,
            model_config,
        )
        outputs = adding_doubling_solver(
            tau, ssa, g, L_snw, ice, illumination, model_config
        )

        return outputs.albedo

    Out = []
    # now use the reduced LUT to call snicar and obtain best matching spectrum
    for lyr in lyrList:
        for dens in densList:
            for reff in reffList:
                for zen in zenList:
                    for bc in bcList:
                        for dz in dzList:
                            out = run_using_dask(
                                input_file, lyr, dens, reff, zen, bc, dz
                            )
                            Out.append(out)

    result = dask.compute(*Out, scheduler="processes")