agdc_flux_extractor.py

#!/g/data/v10/public/modules/dea-env/20181015/bin/python

import datacube as dc
from datacube.helpers import ga_pq_fuser
from datacube.storage import masking
import pyproj
import numpy as np
import xarray as xr
import datetime
import warnings
import argparse
import os


locations = """'Yanco',-34.9878,146.2908,2012,2017
'WombatStateForest',-37.4222,144.0944,2010,2017
'Whroo',-36.6732,145.0294,2011,2017
'Warra',-43.09502,146.65452,2015,2017
'WallabyCreek',-34.00206,140.58912,2005,2013
'Tumbarumba',-35.6566,148.1516,2002,2017
'TiTreeEast',-22.287,133.64,2012,2017
'SturtPlains',-17.1507,133.3502,2008,2017
'Samford',-27.38806,152.87778,2010,2017
'RobsonCreek',-17.11746943,145.6301375,2013,2017
'RiggsCreek',-36.656,145.576,2011,2017
'RedDirtMelonFarm',-14.563639,132.477567,2011,2013
'Otway',-38.525,142.81,2007,2010
'Loxton',-34.47035,140.65512,2008,2009
'Litchfield',-13.17904178,130.7945459,2015,2017
'HowardSprings',-12.4952,131.15005,2002,2017
'GreatWesternWoodlands',-30.1914,120.6541667,2013,2017
'Gingin',-31.375,115.65,2011,2017
'FoggDam',-12.5452,131.3072,2006,2008
'Emerald',-23.85872,148.4746,2011,2013
'DryRiver',-15.2588,132.3706,2008,2017
'DalyUncleared',-14.1592,131.3881,2008,2017
'DalyPasture',-17.1507,133.3502,2008,2013
'CumberlandPlain',-33.615278,150.723611,2012,2016
'CowBay',-16.238189,145.42715,2009,2016
'CapeTribulation',-16.103219,145.446922,2010,2017
'Calperum',-34.00206,140.58912,2010,2017
'AliceSpringsMulga',-22.283,133.249,2010,2017
'AdelaideRiver',-13.0769,131.1178,2007,2009"""



class BurnCube(dc.Datacube):
    def __init__(self):
        super(BurnCube, self).__init__(app='TreeMapping.getLandsatStack')

        self.band_names = ['red', 'green', 'blue', 'nir', 'swir1', 'swir2']
        self.dataset = None
        self.geomed = None
        self.dists = None
        self.outlrs = None


    def to_netcdf(self, path):
        """Saves input data to a netCDF4 on disk
            Args:
                path string: path to a file on disk
        """
        self.dataset.to_netcdf(path)

    def open_dataset(self, path):
        """Loads input data from a netCDF4 on disk
            Args:
                path string: path to a file on disk
        """
        self.dataset = xr.open_dataset(path)

    def _load_pq(self, x, y, res, period, n_landsat):
        query = {
            'time': period,
            'x': x,
            'y': y,
            'crs': 'EPSG:3577',
            'measurements': ['pixelquality'],
            'resolution': res,
        }

        pq_stack = []
        for n in n_landsat:
            pq_stack.append(self.load(product='ls{}_pq_albers'.format(n),
                                      group_by='solar_day', fuse_func=ga_pq_fuser,
                                      resampling='nearest', **query))

        pq_stack = xr.concat(pq_stack, dim='time').sortby('time')

        # Land/sea mask isn't used at the moment. Possible alternatives are WOFS and ITEM.
        #pq_stack['land'] = masking.make_mask(pq_stack.pixelquality, land_sea='land')
        # masking cloud, saturation and invalid data (contiguous)
        pq_stack['good_pixel'] = masking.make_mask(pq_stack.pixelquality, cloud_acca='no_cloud',
                                                   cloud_fmask='no_cloud', cloud_shadow_acca='no_cloud_shadow',
                                                   cloud_shadow_fmask='no_cloud_shadow',
                                                   blue_saturated=False,
                                                   green_saturated=False,
                                                   red_saturated=False,
                                                   nir_saturated=False,
                                                   swir1_saturated=False,
                                                   swir2_saturated=False,
                                                   contiguous=True)

        return pq_stack

    def _load_nbart(self, x, y, res, period, n_landsat):
        query = {
            'time': period,
            'x': x,
            'y': y,
            'crs': 'EPSG:3577',
            'measurements': self.band_names,
            'resolution': res,
        }

        nbart_stack = []
        for n in n_landsat:
            print(n)
            dss = self.find_datasets(product='ls{}_nbart_albers'.format(n), **query)
            if not dss:
                return None
            nbart_stack.append(self.load(product='ls{}_nbart_albers'.format(n),
                                         group_by='solar_day', datasets=dss, resampling='bilinear',
                                         **query))

        print(nbart_stack)
        nbart_stack = xr.concat(nbart_stack, dim='time').sortby('time')

        return nbart_stack

    def load_cube(self, x, y, res, period, n_landsat):
        """Loads the Landsat data for the selected region and sensors
            Note:
                This method loads the data into the self.dataset variable.
            Args:
                x list float: horizontal min max range
                y list float: vertical min max range
                res float: pixel resolution for the input data
                period list datetime: temporal range for input data
                n_landsat int: number of the Landsat mission
        """

        nbart_stack = self._load_nbart(x, y, res, period, n_landsat)
        if nbart_stack is None:
            return

        pq_stack = self._load_pq(x, y, res, period, n_landsat)
        pq_stack, nbart_stack = xr.align(pq_stack, nbart_stack, join='inner')
        mask = np.nanmean(pq_stack.good_pixel.values.reshape(pq_stack.good_pixel.shape[0], -1), axis=1) > .2

        self.dataset = nbart_stack.sel(time=mask).where(pq_stack.good_pixel.sel(time=mask), 0, drop=False) # keep data as integer
        del self.dataset['red'].attrs['spectral_definition']
        del self.dataset['red'].attrs['crs']
        del self.dataset['green'].attrs['spectral_definition']
        del self.dataset['green'].attrs['crs']
        del self.dataset['blue'].attrs['spectral_definition']
        del self.dataset['blue'].attrs['crs']
        del self.dataset['nir'].attrs['spectral_definition']
        del self.dataset['nir'].attrs['crs']
        del self.dataset['swir1'].attrs['spectral_definition']
        del self.dataset['swir1'].attrs['crs']
        del self.dataset['swir2'].attrs['spectral_definition']
        del self.dataset['swir2'].attrs['crs']
        del self.dataset.attrs['crs']
        del self.dataset['time'].attrs['units']

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="""Flux tower AGDC data extractor""")
    parser.add_argument('-dst', '--destination', required=True, type=str, help="Full path to destination.")
    args = parser.parse_args()

    for ls_n in [5,7,8]:
        for line in locations.splitlines():
            fields = line.split(',')
            station_name = fields[0].strip("\'")
            lat = float(fields[1])
            lon = float(fields[2])
            start_year = int(fields[3])
            end_year = int(fields[4])

            # convert to projected centre coordinates
            wgs84 = pyproj.Proj(init='epsg:4326')
            gda94 = pyproj.Proj(init='epsg:3577')
            easting, northing = pyproj.transform(wgs84, gda94, lon, lat)

            # define projected region extent
            x = (easting+3000, easting-3000)   # 3000m is half of the required window size
            y = (northing+3000, northing-3000)

            bc = BurnCube()
            bc.load_cube(x, y, (25, 25), ('{}-01-01'.format(start_year), '{}-12-31'.format(end_year)), [ls_n])

            if bc.dataset:
                bc.dataset.to_netcdf(os.path.join(args.destination, "{}_LS{}.nc".format(station_name, ls_n)))

To execute this on the VDI:

1.- Download and make this code executable
chmod u+x agdc_flux_extractor.py

2.- Load the DEA module
module use /g/data/v10/public/modules/modulefiles
module load dea

3.- Run specifying a path to the destination directory:
./agdc_flux_extractor.py -dst /g/data/xc0/user/Manolo