Enhance export routines

pygmtsar/pygmtsar/Stack_export.py

@@ -8,6 +8,8 @@
 # Licensed under the BSD 3-Clause License (see LICENSE for details)
 # ----------------------------------------------------------------------------
 from .Stack_ps import Stack_ps
+from .tqdm_dask import tqdm_dask
+from .NCubeVTK import NCubeVTK
 
 class Stack_export(Stack_ps):
 
@@ -53,16 +55,16 @@ def as_geo(self, da):
             x_dim = 'x'
         return da.rio.write_crs(epsg).rio.set_spatial_dims(y_dim=y_dim, x_dim=x_dim)
 
-    def export_stack_geotiff(self, data, name, caption='Exporting WGS84 GeoTIFFs'):
+    def export_geotiff(self, data, name, caption='Exporting WGS84 GeoTIFF(s)'):
         """
-        Export single GeoTIFF velocity.tif:
-        sbas.export_stack(velocity_ps, 'velocity')
+        Export single GeoTIFF file "velocity.tif":
+        sbas.export_geotiff(velocity, 'velocity')
         
-        Export pair-based stack of GeoTIFFs (like corr_ps.2024-01-08_2024-01-20.tif, ...):
-        sbas.export_stack(corr_ps, 'corr_ps')
+        Export pair-based stack of GeoTIFF files like corr.2024-01-01_2024-01-02.tif, ...:
+        sbas.export_geotiff(corr, 'corr')
         
-        Export date-based stack of GeoTIFFs (like disp_ps.2023-12-27.tif, ...):
-        sbas.export_stack(disp_ps, 'disp_ps')
+        Export date-based stack of GeoTIFF files (ike disp.2024-01-01.tif, ...:
+        sbas.export_geotiff(disp, 'disp')
 
         The exported GeoTIFF files can be converted to KMZ for Google Earth Engine using GDAL tools:
         gdalwarp -of KMLSUPEROVERLAY -co FORMAT=PNG velocity.tif velocity.kmz
@@ -93,16 +95,10 @@ def export_stack_geotiff(self, data, name, caption='Exporting WGS84 GeoTIFFs'):
                 self.as_geo(self.ra2ll(grid) if not self.is_geo(grid) else grid).rio.to_raster(filename)
                 pbar.update(1)
 
-    def export_stack_geojson(self, data, name, caption='Exporting WGS84 GeoJSON', pivotal=True, digits=1):
+    def export_geojson(self, data, name, caption='Exporting WGS84 GeoJSON', pivotal=True, digits=1):
         """
-        Export single GeoJSON velocity.geojson:
-        sbas.export_stack_geojson(velocity_ps, 'velocity')
-    
-        Export pair-based stack of GeoJSONs (like corr_ps.2024-01-08_2024-01-20.geojson, ...):
-        sbas.export_stack_geojson(corr_ps, 'corr_ps')
-    
-        Export date-based stack of GeoJSONs (like disp_ps.2023-12-27.geojson, ...):
-        sbas.export_stack_geojson(disp_ps, 'disp_ps')
+        Export GeoJSON file "velocity.geojson":
+        sbas.export_geojson(velocity, 'velocity')
         """
         import xarray as xr
         import geopandas as gpd
@@ -162,3 +158,168 @@ def export_stack_geojson(self, data, name, caption='Exporting WGS84 GeoJSON', pi
                 f.write(']}')
         del grid
 
+    def export_csv(self, data, name, caption='Exporting WGS84 CSV', digits=1, delimiter=','):
+        """
+        Export CSV file "velocity.csv":
+        sbas.export_csv(velocity, 'velocity')
+        """
+        import xarray as xr
+        import numpy as np
+        from tqdm.auto import tqdm
+        # disable "distributed.utils_perf - WARNING - full garbage collections ..."
+        from dask.distributed import utils_perf
+        utils_perf.disable_gc_diagnosis()
+
+        assert isinstance(data, xr.DataArray), 'Argument data is not an xr.DataArray object'
+
+        # determine if data has a stack dimension and what it is
+        stackvar = data.dims[0] if len(data.dims) == 3 else None
+
+        # convert the data to geographic coordinates if necessary
+        if not self.is_geo(data):
+            grid = self.ra2ll(data)
+        else:
+            grid = data
+
+        # split to equal chunks and rest
+        lats_blocks = np.array_split(np.arange(grid.lat.size), np.arange(0, grid.lat.size, self.netcdf_chunksize//2)[1:])
+        lons_blocks = np.array_split(np.arange(grid.lon.size), np.arange(0, grid.lon.size, self.netcdf_chunksize//2)[1:])
+
+        # prepare the progress bar
+        with tqdm(desc=caption, total=len(lats_blocks)*len(lons_blocks)) as pbar:
+            with open(f'{name}.csv', 'w') as f:
+                # CSV header
+                f.write(delimiter.join(filter(None, [stackvar, 'lon', 'lat', data.name])) + '\n')
+                for lats_block in lats_blocks:
+                    for lons_block in lons_blocks:
+                        block = grid.isel(lat=lats_block, lon=lons_block).compute()
+                        block_val = block.round(digits).values
+                        block_lat = block.lat.round(6).values
+                        block_lon = block.lon.round(6).values
+                        if stackvar is not None:
+                            stackvals = block[stackvar]
+                            if np.issubdtype(stackvals.dtype, np.datetime64):
+                                stackvals = stackvals.dt.date.astype(str)                           
+                            stackvals, lats, lons = np.meshgrid(stackvals, block_lat, block_lon, indexing='ij')
+                            block_csv = np.column_stack((stackvals.ravel(), lons.ravel(), lats.ravel(), block_val.ravel()))
+                            del stackvals, lats, lons
+                        else:
+                            lats, lons = np.meshgrid(block_lat, block_lon, indexing='ij')
+                            block_csv = np.column_stack((lons.ravel(), lats.ravel(), block_val.astype(str).ravel()))
+                            del lats, lons
+                        del block, block_lat, block_lon
+                        block_csv = block_csv[np.isfinite(block_val.ravel())]
+                        del block_val
+                        if block_csv.size > 0:
+                            np.savetxt(f, block_csv, delimiter=delimiter, fmt='%s')
+                        del block_csv
+                        pbar.update(1)
+        del grid
+
+    def export_netcdf(self, data, name, caption='Exporting WGS84 NetCDF', engine='netcdf4', format='NETCDF3_64BIT'):
+        """
+        Export NetCDF file "velocity.nc":
+        sbas.export_netcdf(velocity, 'velocity')
+        """
+        import xarray as xr
+        import numpy as np
+        import dask
+        import os
+        # disable "distributed.utils_perf - WARNING - full garbage collections ..."
+        from dask.distributed import utils_perf
+        utils_perf.disable_gc_diagnosis()
+
+        assert isinstance(data, xr.DataArray), 'Argument data is not an xr.DataArray object'
+
+        # convert the data to geographic coordinates if necessary
+        if not self.is_geo(data):
+            grid = self.ra2ll(data)
+        else:
+            grid = data
+
+        filename = f'{name}.nc'
+        if os.path.exists(filename):
+            os.remove(filename)
+        encoding = {data.name: self._compression(grid.shape)}
+        delayed = grid.to_netcdf(filename, engine=engine, encoding=encoding, format=format, compute=False)
+        tqdm_dask(dask.persist(delayed), desc=caption)
+        del grid
+
+    def export_vtk(self, data, name, caption='Exporting WGS84 VTK(s)', topo='auto', image=None, band_mask=None, use_sealevel=False):
+        """
+        Export VTK file "velocity.vtk":
+        sbas.export_vtk(velocity, 'velocity')
+        """
+        import xarray as xr
+        import numpy as np
+        from vtk import vtkStructuredGridWriter, vtkStringArray
+        from tqdm.auto import tqdm
+        import os
+
+        assert isinstance(data, xr.DataArray), 'Argument data is not an xr.DataArray object'
+
+        # determine if data has a stack dimension and what it is
+        stackvar = data.dims[0] if len(data.dims) == 3 else None
+        #print ('stackvar', stackvar)
+        if stackvar is not None and np.issubdtype(data[stackvar].dtype, np.datetime64):
+            stackvals = data[stackvar].dt.date.astype(str).values
+        elif stackvar is not None:
+            stackvals = data[stackvar].astype(str).values
+        else:
+            stackvals = [None]
+        #print ('stackvals', stackvals)
+
+        filename = f'{name}.vtk'
+        if os.path.exists(filename):
+            os.remove(filename)
+
+        # convert the data to geographic coordinates if necessary
+        if not self.is_geo(data):
+            data_ll = self.ra2ll(data)
+        else:
+            data_ll = data
+        # define 2D grid for interpolation 
+        grid2d = data_ll.min(stackvar) if stackvar is not None else data_ll
+        #print ('grid2d', grid2d)
+
+        dss = []
+        if isinstance(topo, str) and topo == 'auto':
+            dem = self.get_dem()
+        elif topo is not None:
+            # convert topography to geographic coordinates if necessary
+            dem = self.ra2ll(topo) if not self.is_geo(topo) else topo
+        if topo is not None:
+            dem = dem.reindex_like(grid2d, method='nearest').where(np.isfinite(grid2d))
+            dss.append(dem.rename('z'))
+
+        if image is not None:
+            dss.append(image.reindex_like(grid2d, method='nearest').round().astype(np.uint8).rename('colors'))
+        #print ('dss', dss)
+
+        # prepare the progress bar
+        with tqdm(desc=caption, total=len(stackvals)) as pbar:
+            for stackidx, stackval in enumerate(stackvals):
+                #print (stackidx, stackval)
+                if stackval is not None:
+                    ds = xr.merge([*dss, data_ll.isel({stackvar: stackidx})], compat='override')\
+                        .rename({'lat': 'y', 'lon': 'x'})\
+                        .drop_vars(stackvar)
+                    filename = f'{name}.{stackidx}.vtk'
+                else:
+                    ds = xr.merge([*dss, data_ll], compat='override').rename({'lat': 'y', 'lon': 'x'})
+                    filename = f'{name}.vtk'
+                #print ('ds', ds)
+
+                vtk_grid = NCubeVTK.ImageOnTopography(ds)
+                if stackval is not None:
+                    metadata = vtkStringArray()
+                    metadata.SetName(stackvar)
+                    metadata.InsertNextValue(stackval)
+                    vtk_grid.GetFieldData().AddArray(metadata)
+
+                # convert to VTK structure and save to file
+                writer = vtkStructuredGridWriter()
+                writer.SetFileName(filename)
+                writer.SetInputData(vtk_grid)
+                writer.Write()
+                pbar.update(1)