Victor Phan
733 lines
25 KiB
Python
Executable File
733 lines
25 KiB
Python
Executable File
"""
|
|
Functions for working with the Wetlands Insight Tool (WIT)
|
|
"""
|
|
|
|
# Import required packages
|
|
|
|
# Force GeoPandas to use Shapely instead of PyGEOS
|
|
# In a future release, GeoPandas will switch to using Shapely by default.
|
|
import os
|
|
os.environ['USE_PYGEOS'] = '0'
|
|
|
|
import warnings
|
|
import numpy as np
|
|
import pandas as pd
|
|
import geopandas as gpd
|
|
import seaborn as sns
|
|
import xarray as xr
|
|
import matplotlib.pyplot as plt
|
|
from skimage import exposure
|
|
import matplotlib.animation as animation
|
|
import matplotlib.patheffects as PathEffects
|
|
from mpl_toolkits.axes_grid1.inset_locator import inset_axes
|
|
from dask.distributed import progress
|
|
|
|
import datacube
|
|
from datacube.utils import masking
|
|
from datacube.utils import geometry
|
|
|
|
from deafrica_tools.bandindices import calculate_indices
|
|
from deafrica_tools.datahandling import load_ard, wofs_fuser
|
|
from deafrica_tools.spatial import xr_rasterize
|
|
from deafrica_tools.classification import HiddenPrints
|
|
|
|
|
|
def WIT_drill(
|
|
gdf,
|
|
time,
|
|
min_gooddata=0.85,
|
|
TCW_threshold=-0.035,
|
|
resample_frequency=None,
|
|
export_csv=None,
|
|
dask_chunks=None,
|
|
verbose=False,
|
|
verbose_progress=False,
|
|
):
|
|
"""
|
|
The Wetlands Insight Tool run onver an extent covered by a polygon.
|
|
This function loads FC, WOfS, and Landsat data, and calculates tasseled
|
|
cap wetness, in order to determine the dominant land cover class
|
|
within a polygon at each satellite observation.
|
|
|
|
The output is a pandas dataframe containing a timeseries of the relative
|
|
fractions of each class at each time-step. This forms the input to produce
|
|
a stacked line-plot.
|
|
|
|
Last modified: Oct 2021
|
|
|
|
Parameters
|
|
----------
|
|
gdf : geopandas.GeoDataFrame
|
|
The dataframe must only contain a single row,
|
|
containing the polygon you wish to interrograte.
|
|
time : tuple
|
|
a tuple containing the time range over which to run the WIT.
|
|
e.g. ('2015-01' , '2019-12')
|
|
min_gooddata : Float, optional
|
|
A number between 0 and 1 (e.g 0.8) indicating the minimum percentage
|
|
of good quality pixels required for a satellite observation to be loaded
|
|
and therefore included in the WIT plot. This number should, at a minimum,
|
|
be set to 0.80 to limit biases in the result if not resampling the time-series.
|
|
If resampling the data using the parameter `resample_frequency`, then
|
|
setting this number to 0 (or a low float number) is acceptable.
|
|
TCW_threshold : Int, optional
|
|
The tasseled cap wetness threshold, beyond which a pixel will be
|
|
considered 'wet'. Defaults to -0.035.
|
|
resample_frequency : str
|
|
Option for resampling time-series of input datasets. This option is useful
|
|
for either smoothing the WIT plot, or because the area of analysis is larger
|
|
than a scene width and therefore requires composites. Options include any
|
|
str accepted by `xarray.resample(time=)`. The resampling method used is .max()
|
|
export_csv : str, optional
|
|
To save the returned pandas dataframe as a .csv file, pass a
|
|
a location string (e.g. 'output/results.csv')
|
|
dask_chunks : dict, optional
|
|
To lazily load the datasets using dask, pass a dictionary containing
|
|
the dimensions over which to chunk e.g. {'time':-1, 'x':250, 'y':250}.
|
|
verbose: bool, optional
|
|
If true, print statements are putput detailing the progress of the tool.
|
|
verbose_progress: bool, optional
|
|
For use with Dask progress bar
|
|
|
|
Returns
|
|
-------
|
|
df : Pandas.Dataframe
|
|
A pandas dataframe containing the timeseries of relative fractions
|
|
of each land cover class (WOfs, FC, TCW)
|
|
|
|
"""
|
|
# add geom to dc query dict
|
|
if isinstance(gdf, datacube.utils.geometry._base.Geometry):
|
|
gdf = gpd.GeoDataFrame({'col1':['name'],'geometry':gdf.geom}, crs=gdf.crs)
|
|
geom = geometry.Geometry(geom=gdf.iloc[0].geometry, crs=gdf.crs)
|
|
query = {"geopolygon": geom, "time": time}
|
|
|
|
# Create a datacube instance
|
|
dc = datacube.Datacube(app="wetlands insight tool")
|
|
|
|
# load landsat 5,7,8 data
|
|
warnings.filterwarnings("ignore")
|
|
|
|
if verbose_progress:
|
|
print("Loading Landsat data")
|
|
ds_ls = load_ard(
|
|
dc=dc,
|
|
products=["ls8_sr", "ls7_sr", "ls5_sr"],
|
|
output_crs="epsg:6933",
|
|
min_gooddata=min_gooddata,
|
|
mask_filters=(['opening', 3], ['dilation', 3]),
|
|
measurements=["red", "green", "blue", "nir", "swir_1", "swir_2"],
|
|
dask_chunks=dask_chunks,
|
|
group_by="solar_day",
|
|
resolution=(-30, 30),
|
|
verbose=verbose,
|
|
**query,
|
|
)
|
|
|
|
# create polygon mask
|
|
mask = xr_rasterize(gdf.iloc[[0]], ds_ls)
|
|
ds_ls = ds_ls.where(mask)
|
|
|
|
# calculate tasselled cap wetness within masked AOI
|
|
if verbose:
|
|
print("calculating tasseled cap wetness index ")
|
|
|
|
with HiddenPrints(): #suppres the prints from this func
|
|
tcw = calculate_indices(
|
|
ds_ls, index=["TCW"], normalise=False, satellite_mission="ls", drop=True
|
|
)
|
|
|
|
if resample_frequency is not None:
|
|
if verbose:
|
|
print('Resampling TCW to '+ resample_frequency)
|
|
tcw = tcw.resample(time=resample_frequency).max()
|
|
|
|
tcw = tcw.TCW >= TCW_threshold
|
|
tcw = tcw.where(mask, 0)
|
|
tcw = tcw.persist()
|
|
|
|
if verbose:
|
|
print("Loading WOfS layers ")
|
|
|
|
wofls = dc.load(
|
|
product="wofs_ls",
|
|
like=ds_ls,
|
|
fuse_func=wofs_fuser,
|
|
dask_chunks=dask_chunks,
|
|
collection_category="T1",
|
|
)
|
|
|
|
# boolean of wet/dry
|
|
wofls_wet = masking.make_mask(wofls.water, wet=True)
|
|
|
|
if resample_frequency is not None:
|
|
if verbose:
|
|
print('Resampling WOfS to '+ resample_frequency)
|
|
wofls_wet = wofls_wet.resample(time=resample_frequency).max()
|
|
|
|
# mask sure wofs matches other datasets
|
|
wofls_wet = wofls_wet.where(wofls_wet.time == tcw.time)
|
|
|
|
# apply the polygon mask
|
|
wofls_wet = wofls_wet.where(mask)
|
|
|
|
# load Fractional cover
|
|
if verbose:
|
|
print("Loading fractional Cover")
|
|
|
|
# load fractional cover
|
|
fc_ds = dc.load(
|
|
product="fc_ls",
|
|
time=time,
|
|
dask_chunks=dask_chunks,
|
|
like=ds_ls,
|
|
measurements=["pv", "npv", "bs"],
|
|
collection_category="T1",
|
|
)
|
|
|
|
# use wofls mask to cloud mask FC
|
|
clear_and_dry = masking.make_mask(wofls, dry=True).water
|
|
fc_ds = fc_ds.where(clear_and_dry)
|
|
|
|
if resample_frequency is not None:
|
|
if verbose:
|
|
print('Resampling FC to '+ resample_frequency)
|
|
fc_ds = fc_ds.resample(time=resample_frequency).max()
|
|
|
|
# mask sure fc matches other datasets
|
|
fc_ds = fc_ds.where(fc_ds.time == tcw.time)
|
|
|
|
# mask with polygon
|
|
fc_ds = fc_ds.where(mask)
|
|
|
|
# mask with TC wetness
|
|
fc_ds_noTCW = fc_ds.where(tcw == False)
|
|
|
|
if verbose:
|
|
print("Generating classification")
|
|
|
|
# Cast the dataset to a dataarray
|
|
fc_ds_noTCW = fc_ds_noTCW.to_array(dim="variable", name="fc_ds_noTCW")
|
|
|
|
# turn FC array into integer only as nanargmax doesn't
|
|
# seem to handle floats the way we want it to
|
|
fc_int = fc_ds_noTCW.astype("int8")
|
|
|
|
# use nanargmax to get the index of the maximum value
|
|
BSPVNPV = fc_int.argmax(dim="variable")
|
|
|
|
#int dytype remocves NaNs so we need to create mask again
|
|
FC_mask = np.isfinite(fc_ds_noTCW).all(dim="variable")
|
|
BSPVNPV = BSPVNPV.where(FC_mask)
|
|
|
|
# Restack the Fractional cover dataset all together
|
|
# CAUTION:ARGMAX DEPENDS ON ORDER OF VARIABALES IN
|
|
# DATASET. NEED TO ADJUST BELOW DEPENDING ON ORDER OF FC VARIABLES
|
|
|
|
FC_dominant = xr.Dataset(
|
|
{
|
|
"bs": (BSPVNPV == 2).where(FC_mask),
|
|
"pv": (BSPVNPV == 0).where(FC_mask),
|
|
"npv": (BSPVNPV == 1).where(FC_mask),
|
|
}
|
|
)
|
|
|
|
# pixel counts
|
|
pixels = mask.sum(dim=["x", "y"])
|
|
|
|
|
|
if verbose_progress:
|
|
print("Computing wetness")
|
|
tcw_pixel_count = tcw.sum(dim=["x", "y"]).compute()
|
|
|
|
if verbose_progress:
|
|
print("Computing green veg, dry veg, and bare soil")
|
|
FC_count = FC_dominant.sum(dim=["x", "y"]).compute()
|
|
|
|
if verbose_progress:
|
|
print("Computing open water")
|
|
wofs_pixels = wofls_wet.sum(dim=["x", "y"]).compute()
|
|
|
|
# count percentages
|
|
wofs_area_percent = (wofs_pixels / pixels) * 100
|
|
tcw_area_percent = (tcw_pixel_count / pixels) * 100
|
|
tcw_less_wofs = tcw_area_percent - wofs_area_percent # wet not wofs
|
|
|
|
# Fractional cover pixel count method
|
|
# Get number of FC pixels, divide by total number of pixels per polygon
|
|
# Work out the number of nodata pixels in the data
|
|
BS_percent = (FC_count.bs / pixels) * 100
|
|
PV_percent = (FC_count.pv / pixels) * 100
|
|
NPV_percent = (FC_count.npv / pixels) * 100
|
|
NoData_count = ((
|
|
100 - wofs_area_percent - tcw_less_wofs - PV_percent - NPV_percent - BS_percent
|
|
) / 100) * pixels
|
|
|
|
# re-do percentages but now handling any no-data pixels within polygon
|
|
BS_percent = (FC_count.bs / (pixels - NoData_count)) * 100
|
|
PV_percent = (FC_count.pv / (pixels - NoData_count)) * 100
|
|
NPV_percent = (FC_count.npv / (pixels - NoData_count)) * 100
|
|
wofs_area_percent = (wofs_pixels / (pixels - NoData_count)) * 100
|
|
tcw_area_percent = (tcw_pixel_count / (pixels - NoData_count)) * 100
|
|
tcw_less_wofs = tcw_area_percent - wofs_area_percent
|
|
|
|
# Sometimes when we resample datastes, WOfS extent can be
|
|
# greater than the wetness extent, thus make negative values == zero
|
|
tcw_less_wofs = tcw_less_wofs.where(tcw_less_wofs>=0, 0)
|
|
|
|
# start setup of dataframe by adding only one dataset
|
|
df = pd.DataFrame(
|
|
data=wofs_area_percent.data,
|
|
index=wofs_area_percent.time.values,
|
|
columns=["wofs_area_percent"],
|
|
)
|
|
|
|
# add data into pandas dataframe for export
|
|
df["wet_percent"] = tcw_less_wofs.data
|
|
df["green_veg_percent"] = PV_percent.data
|
|
df["dry_veg_percent"] = NPV_percent.data
|
|
df["bare_soil_percent"] = BS_percent.data
|
|
|
|
# round numbers
|
|
df = df.round(2)
|
|
|
|
# save the csv of the output data used to create the stacked plot for the polygon drill
|
|
if export_csv:
|
|
if verbose:
|
|
print("exporting csv: " + export_csv)
|
|
df.to_csv(export_csv, index_label="Datetime")
|
|
|
|
return df
|
|
|
|
|
|
def animated_timeseries_WIT(
|
|
ds,
|
|
df,
|
|
output_path,
|
|
width_pixels=1000,
|
|
interval=200,
|
|
bands=["red", "green", "blue"],
|
|
percentile_stretch=(0.02, 0.98),
|
|
image_proc_func=None,
|
|
title=False,
|
|
show_date=True,
|
|
annotation_kwargs={},
|
|
onebandplot_cbar=True,
|
|
onebandplot_kwargs={},
|
|
shapefile_path=None,
|
|
shapefile_kwargs={},
|
|
pandasplot_kwargs={},
|
|
time_dim="time",
|
|
x_dim="x",
|
|
y_dim="y",
|
|
):
|
|
|
|
###############
|
|
# Setup steps #
|
|
###############
|
|
|
|
# Test if all dimensions exist in dataset
|
|
if time_dim in ds and x_dim in ds and y_dim in ds:
|
|
|
|
# Test if there is one or three bands, and that all exist in both datasets:
|
|
if ((len(bands) == 3) | (len(bands) == 1)) & all(
|
|
[(b in ds.data_vars) for b in bands]
|
|
):
|
|
|
|
# Import xarrays as lists of three band numpy arrays
|
|
imagelist, vmin, vmax = _ds_to_arrraylist(
|
|
ds,
|
|
bands=bands,
|
|
time_dim=time_dim,
|
|
x_dim=x_dim,
|
|
y_dim=y_dim,
|
|
percentile_stretch=percentile_stretch,
|
|
image_proc_func=image_proc_func,
|
|
)
|
|
|
|
# Get time, x and y dimensions of dataset and calculate width vs height of plot
|
|
timesteps = len(ds[time_dim])
|
|
width = len(ds[x_dim])
|
|
height = len(ds[y_dim])
|
|
width_ratio = float(width) / float(height)
|
|
height = 10.0 / width_ratio
|
|
|
|
# If title is supplied as a string, multiply out to a list with one string per timestep.
|
|
# Otherwise, use supplied list for plot titles.
|
|
if isinstance(title, str) or isinstance(title, bool):
|
|
title_list = [title] * timesteps
|
|
else:
|
|
title_list = title
|
|
|
|
# Set up annotation parameters that plt.imshow plotting for single band array images.
|
|
# The nested dict structure sets default values which can be overwritten/customised by the
|
|
# manually specified `onebandplot_kwargs`
|
|
onebandplot_kwargs = dict(
|
|
{
|
|
"cmap": "Greys",
|
|
"interpolation": "bilinear",
|
|
"vmin": vmin,
|
|
"vmax": vmax,
|
|
"tick_colour": "black",
|
|
"tick_fontsize": 11,
|
|
},
|
|
**onebandplot_kwargs,
|
|
)
|
|
|
|
# Use pop to remove the two special tick kwargs from the onebandplot_kwargs dict, and save individually
|
|
onebandplot_tick_colour = onebandplot_kwargs.pop("tick_colour")
|
|
onebandplot_tick_fontsize = onebandplot_kwargs.pop("tick_fontsize")
|
|
|
|
# Set up annotation parameters that control font etc. The nested dict structure sets default
|
|
# values which can be overwritten/customised by the manually specified `annotation_kwargs`
|
|
annotation_kwargs = dict(
|
|
{
|
|
"xy": (1, 1),
|
|
"xycoords": "axes fraction",
|
|
"xytext": (-5, -5),
|
|
"textcoords": "offset points",
|
|
"horizontalalignment": "right",
|
|
"verticalalignment": "top",
|
|
"fontsize": 15,
|
|
"color": "white",
|
|
"path_effects": [
|
|
PathEffects.withStroke(linewidth=3, foreground="black")
|
|
],
|
|
},
|
|
**annotation_kwargs,
|
|
)
|
|
|
|
# Define default plotting parameters for the overlaying shapefile(s). The nested dict structure sets
|
|
# default values which can be overwritten/customised by the manually specified `shapefile_kwargs`
|
|
shapefile_kwargs = dict(
|
|
{"linewidth": 2, "edgecolor": "black", "facecolor": "#00000000"},
|
|
**shapefile_kwargs,
|
|
)
|
|
|
|
# Define default plotting parameters for the right-hand line plot. The nested dict structure sets
|
|
# default values which can be overwritten/customised by the manually specified `pandasplot_kwargs`
|
|
pandasplot_kwargs = dict({}, **pandasplot_kwargs)
|
|
|
|
###################
|
|
# Initialise plot #
|
|
###################
|
|
|
|
# Set up figure
|
|
fig, (ax1, ax2) = plt.subplots(
|
|
ncols=2, gridspec_kw={"width_ratios": [1, 2]}
|
|
)
|
|
fig.subplots_adjust(left=0, bottom=0, right=1, top=1, wspace=0.2, hspace=0)
|
|
fig.set_size_inches(10.0, height * 0.5, forward=True)
|
|
ax1.axis("off")
|
|
ax2.margins(x=0.01)
|
|
ax2.xaxis.label.set_visible(False)
|
|
|
|
# Initialise axesimage objects to be updated during animation, setting extent from dims
|
|
extents = [
|
|
float(ds[x_dim].min()),
|
|
float(ds[x_dim].max()),
|
|
float(ds[y_dim].min()),
|
|
float(ds[y_dim].max()),
|
|
]
|
|
im = ax1.imshow(imagelist[0], extent=extents, **onebandplot_kwargs)
|
|
|
|
# Initialise right panel and set y axis limits
|
|
# set up color palette
|
|
pal = [
|
|
sns.xkcd_rgb["cobalt blue"],
|
|
sns.xkcd_rgb["neon blue"],
|
|
sns.xkcd_rgb["grass"],
|
|
sns.xkcd_rgb["beige"],
|
|
sns.xkcd_rgb["brown"],
|
|
]
|
|
|
|
# make a stacked area plot
|
|
ax2.stackplot(
|
|
df.index,
|
|
df.wofs_area_percent,
|
|
df.wet_percent,
|
|
df.green_veg_percent,
|
|
df.dry_veg_percent,
|
|
df.bare_soil_percent,
|
|
labels=["open water", "wet", "green veg", "dry veg", "bare soil"],
|
|
colors=pal,
|
|
alpha=0.6,
|
|
**pandasplot_kwargs,
|
|
)
|
|
|
|
ax2.legend(loc="lower left", framealpha=0.6)
|
|
|
|
df1 = pd.DataFrame(
|
|
{
|
|
"wofs_area_percent": df.wofs_area_percent,
|
|
"wet_percent": df.wofs_area_percent + df.wet_percent,
|
|
"green_veg_percent": df.wofs_area_percent
|
|
+ df.wet_percent
|
|
+ df.green_veg_percent,
|
|
"dry_veg_percent": df.wofs_area_percent
|
|
+ df.wet_percent
|
|
+ df.green_veg_percent
|
|
+ df.dry_veg_percent,
|
|
"bare_soil_percent": df.dry_veg_percent
|
|
+ df.green_veg_percent
|
|
+ df.wofs_area_percent
|
|
+ df.wet_percent
|
|
+ df.bare_soil_percent,
|
|
}
|
|
)
|
|
df1 = df1.set_index(df.index)
|
|
|
|
line_test = df1.plot(
|
|
ax=ax2, legend=False, color="black", **pandasplot_kwargs
|
|
)
|
|
|
|
# set axis limits to the min and max
|
|
ax2.set(xlim=(df.index[0], df.index[-1]), ylim=(0, 100))
|
|
|
|
# add a legend and a tight plot box
|
|
|
|
ax2.set_title("Fractional Cover, Wetness, and Water")
|
|
|
|
# Initialise annotation objects to be updated during animation
|
|
t = ax1.annotate("", **annotation_kwargs)
|
|
|
|
#########################
|
|
# Add optional overlays #
|
|
#########################
|
|
|
|
# Optionally add shapefile overlay(s) from either string path or list of string paths
|
|
if isinstance(shapefile_path, str):
|
|
|
|
shapefile = gpd.read_file(shapefile_path)
|
|
shapefile.plot(**shapefile_kwargs, ax=ax1)
|
|
|
|
elif isinstance(shapefile_path, list):
|
|
|
|
# Iterate through list of string paths
|
|
for shapefile in shapefile_path:
|
|
|
|
shapefile = gpd.read_file(shapefile)
|
|
shapefile.plot(**shapefile_kwargs, ax=ax1)
|
|
|
|
# After adding shapefile, fix extents of plot
|
|
ax1.set_xlim(extents[0], extents[1])
|
|
ax1.set_ylim(extents[2], extents[3])
|
|
|
|
# Optionally add colourbar for one band images
|
|
if (len(bands) == 1) & onebandplot_cbar:
|
|
_add_colourbar(
|
|
ax1,
|
|
im,
|
|
tick_fontsize=onebandplot_tick_fontsize,
|
|
tick_colour=onebandplot_tick_colour,
|
|
vmin=onebandplot_kwargs["vmin"],
|
|
vmax=onebandplot_kwargs["vmax"],
|
|
)
|
|
|
|
########################################
|
|
# Create function to update each frame #
|
|
########################################
|
|
|
|
# Function to update figure
|
|
|
|
def update_figure(frame_i):
|
|
|
|
####################
|
|
# Plot image panel #
|
|
####################
|
|
|
|
# If possible, extract dates from time dimension
|
|
try:
|
|
|
|
# Get human-readable date info (e.g. "16 May 1990")
|
|
ts = ds[time_dim][{time_dim: frame_i}].dt
|
|
year = ts.year.item()
|
|
month = ts.month.item()
|
|
day = ts.day.item()
|
|
date_string = "{} {} {}".format(
|
|
day, calendar.month_abbr[month], year
|
|
)
|
|
|
|
except:
|
|
|
|
date_string = ds[time_dim][{time_dim: frame_i}].values.item()
|
|
|
|
# Create annotation string based on title and date specifications:
|
|
title = title_list[frame_i]
|
|
if title and show_date:
|
|
title_date = "{}\n{}".format(date_string, title)
|
|
elif title and not show_date:
|
|
title_date = "{}".format(title)
|
|
elif show_date and not title:
|
|
title_date = "{}".format(date_string)
|
|
else:
|
|
title_date = ""
|
|
|
|
# Update left panel with annotation and image
|
|
im.set_array(imagelist[frame_i])
|
|
t.set_text(title_date)
|
|
|
|
########################
|
|
# Plot linegraph panel #
|
|
########################
|
|
|
|
# Create list of artists to return
|
|
artist_list = [im, t]
|
|
|
|
# Update right panel with temporal line subset, adding each new line into artist_list
|
|
for i, line in enumerate(line_test.lines):
|
|
|
|
# Clip line data to current time, and get x and y values
|
|
y = df1[
|
|
df1.index
|
|
<= datetime(year=year, month=month, day=day, hour=23, minute=59)
|
|
].iloc[:, i]
|
|
x = df1[
|
|
df1.index
|
|
<= datetime(year=year, month=month, day=day, hour=23, minute=59)
|
|
].index
|
|
|
|
# Plot lines after stripping NaNs (this produces continuous, unbroken lines)
|
|
line.set_data(x[y.notnull()], y[y.notnull()])
|
|
artist_list.extend([line])
|
|
|
|
# Return the artists set
|
|
return artist_list
|
|
|
|
# Nicely space subplots
|
|
fig.tight_layout()
|
|
|
|
##############################
|
|
# Generate and run animation #
|
|
##############################
|
|
|
|
# Generate animation
|
|
ani = animation.FuncAnimation(
|
|
fig=fig,
|
|
func=update_figure,
|
|
frames=timesteps,
|
|
interval=interval,
|
|
blit=True,
|
|
)
|
|
|
|
# Export as either MP4 or GIF
|
|
if output_path[-3:] == "mp4":
|
|
print(" Exporting animation to {}".format(output_path))
|
|
ani.save(output_path, dpi=width_pixels / 10.0)
|
|
|
|
elif output_path[-3:] == "wmv":
|
|
print(" Exporting animation to {}".format(output_path))
|
|
ani.save(
|
|
output_path,
|
|
dpi=width_pixels / 10.0,
|
|
writer=animation.FFMpegFileWriter(
|
|
fps=1000 / interval, bitrate=4000, codec="wmv2"
|
|
),
|
|
)
|
|
|
|
elif output_path[-3:] == "gif":
|
|
print(" Exporting animation to {}".format(output_path))
|
|
ani.save(output_path, dpi=width_pixels / 10.0, writer="imagemagick")
|
|
|
|
else:
|
|
print(" Output file type must be either .mp4, .wmv or .gif")
|
|
|
|
else:
|
|
print(
|
|
"Please select either one or three bands that all exist in the input dataset"
|
|
)
|
|
|
|
else:
|
|
print(
|
|
"At least one x, y or time dimension does not exist in the input dataset. Please use the `time_dim`,"
|
|
"`x_dim` or `y_dim` parameters to override the default dimension names used for plotting"
|
|
)
|
|
|
|
|
|
# Define function to convert xarray dataset to list of one or three band numpy arrays
|
|
|
|
|
|
def _ds_to_arrraylist(
|
|
ds, bands, time_dim, x_dim, y_dim, percentile_stretch, image_proc_func=None
|
|
):
|
|
"""
|
|
Converts an xarray dataset to a list of numpy arrays for plt.imshow plotting
|
|
"""
|
|
|
|
# Compute percents
|
|
p_low, p_high = ds[bands].to_array().quantile(percentile_stretch).values
|
|
|
|
array_list = []
|
|
for i, timestep in enumerate(ds[time_dim]):
|
|
|
|
# Select single timestep from the data array
|
|
ds_i = ds[{time_dim: i}]
|
|
|
|
# Get shape of array
|
|
x = len(ds[x_dim])
|
|
y = len(ds[y_dim])
|
|
|
|
if len(bands) == 1:
|
|
|
|
# Create new one band array
|
|
img_toshow = exposure.rescale_intensity(
|
|
ds_i[bands[0]].values, in_range=(p_low, p_high), out_range="image"
|
|
)
|
|
|
|
else:
|
|
|
|
# Create new three band array
|
|
rawimg = np.zeros((y, x, 3), dtype=np.float32)
|
|
|
|
# Add xarray bands into three dimensional numpy array
|
|
for band, colour in enumerate(bands):
|
|
|
|
rawimg[:, :, band] = ds_i[colour].values
|
|
|
|
# Stretch contrast using percentile values
|
|
img_toshow = exposure.rescale_intensity(
|
|
rawimg, in_range=(p_low, p_high), out_range=(0, 1.0)
|
|
)
|
|
|
|
# Optionally image processing
|
|
if image_proc_func:
|
|
|
|
img_toshow = image_proc_func(img_toshow).clip(0, 1)
|
|
|
|
array_list.append(img_toshow)
|
|
|
|
return array_list, p_low, p_high
|
|
|
|
|
|
def _add_colourbar(
|
|
ax, im, vmin, vmax, cmap="Greys", tick_fontsize=15, tick_colour="black"
|
|
):
|
|
"""
|
|
Add a nicely formatted colourbar to an animation panel
|
|
"""
|
|
|
|
# Add colourbar
|
|
axins2 = inset_axes(ax, width="97%", height="4%", loc=8, borderpad=1)
|
|
plt.gcf().colorbar(
|
|
im, cax=axins2, orientation="horizontal", ticks=np.linspace(vmin, vmax, 3)
|
|
)
|
|
axins2.xaxis.set_ticks_position("top")
|
|
axins2.tick_params(axis="x", colors=tick_colour, labelsize=tick_fontsize)
|
|
|
|
# Justify left and right labels to edge of plot
|
|
axins2.get_xticklabels()[0].set_horizontalalignment("left")
|
|
axins2.get_xticklabels()[-1].set_horizontalalignment("right")
|
|
labels = [item.get_text() for item in axins2.get_xticklabels()]
|
|
labels[0] = " " + labels[0]
|
|
labels[-1] = labels[-1] + " "
|
|
|
|
|
|
if __name__ == "__main__":
|
|
# print that we are running the testing
|
|
print("Testing..")
|
|
# import doctest to test our module for documentation
|
|
import doctest
|
|
|
|
doctest.testmod()
|
|
print("Testing done")
|