03-DINCAE-tutorial.jl

# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .jl
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Julia 1.8.0
#     language: julia
#     name: julia-1.8
# ---

# # Tutorial on DINCAE
#
# This script/notebook reconstructs missing data in satellite data using a neural
# network architecture called convolutional auto-encoder described in the
# following articles:
#
# * Barth, A., Alvera-Azcárate, A., Licer, M., and Beckers, J.-M.: DINCAE 1.0: a convolutional neural network with error estimates to reconstruct sea surface temperature satellite observations, Geosci. Model Dev., 13, 1609–1622, https://doi.org/10.5194/gmd-13-1609-2020, 2020.
# * Barth, A., Alvera-Azcárate, A., Troupin, C., & Beckers, J.-M. (2022). DINCAE 2.0: multivariate convolutional neural network with error estimates to reconstruct sea surface temperature satellite and altimetry observations. Geoscientific Model Development, 15(5), 2183–2196. https://doi.org/10.5194/gmd-15-2183-2022
#
# The example here uses MODIS sea surface temperature from the Physical
# Oceanography Distributed Active Archive Center (PO.DAAC) JPL, NASA.
# More information is available at https://dx.doi.org/10.5067/MODST-1D4D9
#
# This notebook/script is indented to be run on a GPU with CUDA support (NVIDIA GPU)
# with a least 8 GB of RAM.

# Dependencies are installed with:
#
# ```julia
# using Pkg
# Pkg.add(["CUDA","Knet","NCDatasets","PyPlot"])
# Pkg.add(url="https://github.com/gher-uliege/DINCAE_utils.jl", rev="main")
# ```
#
# See also https://github.com/gher-ulg/DINCAE.jl#installation
#
# Load the necessary modules

using CUDA
using DINCAE
using DINCAE_utils
using Dates
using Knet
using NCDatasets
using PyPlot
using Printf
using Downloads: download

# ## Data download/processing
#
# First we define the domain and time interval.
#
# The example is tested here for a short time frame but for realistic
# applications one should use a much longer time-range (like 10, 20 years or more)
#
# Place the files modis_subset.nc and data-avg.nc (if you have them already)
# in your current working directory.
# They are downloaded is not present, but they are very large (400 MB).


# longitude range (east, west)
lon_range = [-33, -24]
# latitude range (south, north)
lat_range = [33.0, 40.0]

# time range (start, end)
time_range = [DateTime(2017,1,1), DateTime(2019,12,31)]

# local directory
localdir = "."
# create directory
mkpath(localdir)
# filename of the subset
fname_subset = joinpath(localdir,"modis_subset.nc")
# filename of the clean data
fname = joinpath(localdir,"modis_cleanup.nc")
# filename of the data with added clouds for cross-validation
fname_cv = joinpath(localdir,"modis_cleanup_add_clouds.nc")
varname = "sst"

# Results of DINCAE will be placed in a sub-directory under `localdir`

outdir = "."
mkpath(outdir)

# Download the MODIS subset:

if !isfile(fname_subset)
    download("https://dox.ulg.ac.be/index.php/s/a1Lk2sG4onZR5Gg/download",fname_subset)
end

# The MODIS subset was generated by the following code:
#
# ```julia
# for dt in time_range[1]:Day(1):time_range[end]
#     yyyy = Dates.format(dt,"yyyy")
#     mm = Dates.format(dt,"mm")
#     dd = Dates.format(dt,"dd")
#     doy = @sprintf("%03d",Dates.dayofyear(dt))
#     url = "https://podaac-opendap.jpl.nasa.gov:443/opendap/allData/modis/L3/terra/11um/v2019.0/4km/daily/$yyyy/$doy/TERRA_MODIS.$(yyyy)$(mm)$(dd).L3m.DAY.NSST.sst.4km.nc"
#     filename = joinpath(localdir,"TERRA_MODIS.$(yyyy)$(mm)$(dd).L3m.DAY.NSST.sst.4km.nc")
#     if isfile(filename)
#         continue
#     end
#     try
#         ds = NCDataset(url)
#         ds_subset = NCDatasets.@select(ds,$(lon_range[1]) <= lon <= $(lon_range[2]) && $(lat_range[1]) <= lat <= $(lat_range[2]))
#         write(filename,ds_subset)
#         push!(filenames,filename)
#         close(ds)
#         println("$dt downloaded")
#     catch e
#         showerror(stdout, e)
#         @warn("no data for $dt")
#     end
# end
#
# filenames = String[]
# time = DateTime[]
#
# for dt in time_range[1]:Day(1):time_range[end]
#     yyyy = Dates.format(dt,"yyyy")
#     mm = Dates.format(dt,"mm")
#     dd = Dates.format(dt,"dd")
#     filename = joinpath(localdir,"TERRA_MODIS.$(yyyy)$(mm)$(dd).L3m.DAY.NSST.sst.4km.nc")
#     if isfile(filename)
#         push!(filenames,filename)
#         push!(time,dt)
#     end
# end
#
# ds = NCDataset(filenames,aggdim="time",isnewdim=true,
#                constvars = ["lon","lat"]);
# ds2 = NCDataset(fname_subset,"c")
# write(ds2,ds)
# defVar(ds2,"time",time,("time",))
# close(ds2)
# close(ds)
# ```

ds = NCDataset(fname_subset,"r")
sst = ds["sst"][:,:,:];
qual = ds["qual_sst"][:,:,:];

@show extrema(skipmissing(qual))
@show extrema(skipmissing(sst))

# We ignore all data points with missing quality flags,
# quality indicator exceeding 3 and temperature
# larger than 40°C (not necessary for this particular dataset)

sst_t = copy(sst)
sst_t[(qual .> 3) .& .!ismissing.(qual)] .= missing
sst_t[(sst_t .> 40) .& .!ismissing.(sst_t)] .= missing

@info "number of missing observations: $(count(ismissing,sst_t))"
@info "number of valid observations: $(count(.!ismissing,sst_t))"

# Clean-up the data to them write to disk

isfile("modis_cleanup.nc") && rm("modis_cleanup.nc")
varname = "sst"
fname = joinpath(localdir,"modis_cleanup.nc")
ds2 = NCDataset(fname,"c")
write(ds2,ds,exclude = ["sst","qual"])
defVar(ds2,varname,sst_t,("lon","lat","time"))
close(ds2)

# Add a land-sea mask to the file grid points with less than 5% of
# valid data are considerd as land

DINCAE_utils.add_mask(fname,varname; minseafrac = 0.05)

# Choose cross-validation points by adding clouds to the cleanest
# images (copied from the cloudiest images). This function will generate
# a file `fname_cv`.

DINCAE_utils.addcvpoint(fname,varname; mincvfrac = 0.10);

# # Reconstruct missing data
#
#
# F is the floating point number type for neural network, here we use
# single precision

const F = Float32

# Test if CUDA is functional to use the GPU, otherwise the CPU is used.

if CUDA.functional()
    Atype = KnetArray{F}
else
    @warn "No supported GPU found. We will use the CPU which is very slow. Please check https://developer.nvidia.com/cuda-gpus"
    Atype = Array{F}
end

Knet.atype() = Atype

# Setting the paramenters of neural network
# see document of DINCAE.reconstruct for more information

epochs = 1000
batch_size = 32
enc_nfilter_internal = round.(Int,32 * 2 .^ (0:4))
clip_grad = 5.0
regularization_L2_beta = 0
ntime_win = 3
upsampling_method = :nearest
loss_weights_refine = (0.3,0.7)
save_epochs = 200:10:epochs


data = [
   (filename = fname_cv,
    varname = varname,
    obs_err_std = 1,
    jitter_std = 0.05,
    isoutput = true,
   )
]
data_test = data;
fnames_rec = [joinpath(outdir,"data-avg.nc")]
data_all = [data,data_test]

# Use these parameters for a quick test:

# epochs = 10
# save_epochs = epochs:epochs

# Start the training and reconstruction of the neural network
# This takes a long time. You can cancel this step and get the
# sample result.
#
# ```julia
# loss = DINCAE.reconstruct(
#     Atype,data_all,fnames_rec;
#     epochs = epochs,
#     batch_size = batch_size,
#     enc_nfilter_internal = enc_nfilter_internal,
#     clip_grad = clip_grad,
#     save_epochs = save_epochs,
#     upsampling_method = upsampling_method,
#     loss_weights_refine = loss_weights_refine,
#     ntime_win = ntime_win,
# )
# ```

# Download the result and load the loss

if !isfile(fnames_rec[1])
    download("https://dox.ulg.ac.be/index.php/s/91OnKoK0CUE3a8V/download",fnames_rec[1])
end
loss =
    NCDataset(fnames_rec[1]) do ds
        ds["losses"][:]
    end


# Plot the loss function

plot(loss)
ylim(extrema(loss[2:end]))
xlabel("epochs")
ylabel("loss");

# # Post process results
#
# Compute the RMS with the independent validation data

case = (
    fname_orig = fname,
    fname_cv = fname_cv,
    varname = varname,
)
fnameavg = joinpath(outdir,"data-avg.nc")
cvrms = DINCAE_utils.cvrms(case,fnameavg)
@info "Cross-validation RMS error is: $cvrms"

# Next we plot all time instances. The figures with be placed in the
# directory `figdir`

figdir = joinpath(outdir,"Fig")
DINCAE_utils.plotres(case,fnameavg, clim = nothing, figdir = figdir,
                     clim_quantile = (0.01,0.99),
                     which_plot = :cv)
@info "Figures are in $(figdir)"