scripts/animate.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
import argparse
import gc
import glob
import logging
import multiprocessing
import os
import traceback

import netCDF4
import numpy as np
import PIL
from mpl_toolkits.basemap import Basemap

from libs.utils import setup_logging
from libs.utils import verbose as vprint

setup_logging()
log = logging.getLogger(__name__)
CONFIG = {}
V = 1
V_IGNORE = []  # Debug, Warning, Error
# print(os.getcwd())


def vis(files, layers=["SM1", "SM2", "SM3", "SM4", "SM5", "DD"], days="all"):
    """Save the visualisation of the soil moisture content for given layers,
    year files, and days.

    Parameters
    ----------
    files : list
        Iterable of the absolute or relative paths to the netcdf data ending with .nc.
    layers : list
        Soil layers to visualise. Default is ["SM1","SM2","SM3","SM4","SM5","DD"].
    days : list / str
        Iterable days of year range to visualise. Default is 'all'.

    Returns
    -------
        None
    """
    inputs = [(file, layer, days) for file in files for layer in layers]
    num_processes = multiprocessing.cpu_count()
    with multiprocessing.Pool(processes=max(1, int(num_processes / 2))) as pool:
        # pool.starmap(vis_process, [(f, l, days) for f in files for l in layers])
        # pool.apply_async(vis_process, args=inputs)
        for result in pool.imap_unordered(vis_process, inputs):
            print(f"Result file: {result}", flush=True)
    # pool.close()
    # pool.join()


def vis_process(input_nc_layer_days):
    """Save the visualisation of the soil moisture content for a given layer,
    year file, and days.

    Parameters
    ----------
    input_nc : str
        Absolute or relative path to the netcdf data ending with .nc.
    layer : str
        Soil layer to visualise.
    days : list / str
        Iterable days of year range to visualise.

    Returns
    -------
        None
    """
    import matplotlib
    import matplotlib.pyplot as plt

    input_nc, layer, days = input_nc_layer_days
    nc = netCDF4.Dataset(input_nc)
    # ref = nc.variables["spatial_ref"]
    # vprint(1, V, V_IGNORE, Debug=ref)
    lons = nc.variables["x"][:]
    lats = nc.variables["y"][:]

    mp = Basemap(
        projection="merc",
        llcrnrlon=lons.min() - 0.02,  # lower longitude
        llcrnrlat=lats.min() - 0.02,  # lower latitude
        urcrnrlon=lons.max() + 0.02,  # uppper longitude
        urcrnrlat=lats.max() + 0.02,  # uppper latitude
        resolution="i",
    )
    lon, lat = np.meshgrid(lons, lats)  # this converts coordinates into 2D arrray
    x, y = mp(lon, lat)  # mapping them together

    layer_data = nc.variables[layer][:]
    if days == "all":
        days = np.arange(0, layer_data.shape[0])

    # Generate daily images
    for i in days:
        day = i + 1
        matplotlib.use("Agg")
        plt.figure(figsize=(6, 8))  # figure size
        c_scheme = mp.pcolor(x, y, np.squeeze(layer_data[i, :, :]), cmap="jet")
        # mp.etopo()
        mp.shadedrelief()
        mp.drawcoastlines()
        mp.drawstates()
        mp.drawcountries()

        mp.colorbar(c_scheme, location="right", pad="10%")

        plt.title("Soil moisture content for day " + str(day) + " of the year")
        plt.clim(layer_data.min(), layer_data.max())

        plt.savefig(f".tmp/{layer}_{str(day)}.jpg")
        plt.clf()
    plt.close("all")

    # Generate gif animation
    image_frames = []  # creating a empty list to be appended later on
    for day in days:
        new_fram = PIL.Image.open(f".tmp/{layer}_{str(day + 1)}.jpg")
        image_frames.append(new_fram)
    o_name = input_nc.replace(".nc", f"_{layer}.gif")
    image_frames[0].save(
        o_name,
        format="GIF",
        append_images=image_frames[1:],
        save_all=True,
        duration=150,
        loop=0,
    )

    # Delete temporary files
    del (
        nc,
        layer_data,
        mp,
        image_frames,
        c_scheme,
        plt,
        new_fram,
        lon,
        lat,
        x,
        y,
        lons,
        lats,
        day,
        input_nc,
        layer,
        days,
        i,
    )
    gc.collect()
    return o_name


if __name__ == "__main__":
    # Load Configs
    parser = argparse.ArgumentParser(
        description="Download rainfall data from Google Earth Engine for a range of dates.",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    parser.add_argument(
        "-i",
        "--input",
        help="Absolute or relative path to the netcdf data ending with .nc. By dfault it is set to data.nc",
        default="data.nc",
    )
    parser.add_argument(
        "-l",
        "--layer",
        help="Soil layer to visualise. Default is all. Select between SM1 to SM5 or DD.",
        default="all",
    )
    parser.add_argument(
        "-d",
        "--days",
        help="Iterable days of year range to visualise. Default is all.",
        default="all",
    )

    args = parser.parse_args()

    try:
        # Get the files
        if os.path.isdir(args.input):
            files = glob.glob(os.path.join(args.input, "*.nc"))
        else:
            files = [args.input]

        # Get the layers
        if args.layer == "all":
            layers = ["SM1", "SM2", "SM3", "SM4", "SM5", "DD"]
        else:
            layers = str(args.layer).split(",")

        # Get the days
        if args.days == "all":
            days = "all"
        else:
            days = str(args.days).split(",")

        vis(files, layers, days)
    except Exception as e:
        vprint(
            0,
            V,
            V_IGNORE,
            Error="Failed to execute the main function:",
            ErrorMessage=e,
        )
        traceback.print_exc()
        raise e