Spaces:
Sleeping
Sleeping
File size: 5,786 Bytes
6b83428 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 |
#!/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
|