<commit_before>'''
Created on 14Sep.,2016

@author: Alex
'''
import re
import numpy as np
import netCDF4
import math
import scipy.ndimage
from osgeo.osr import SpatialReference, CoordinateTransformation


class NetCDF2DUtils(object):
    '''
    NetCDF2DUtils class to do various fiddly things with NetCDF geophysics files.
    '''
    # Assume WGS84 lat/lon if no CRS is provided
    DEFAULT_CRS = "GEOGCS[\"WGS 84\",DATUM[\"WGS_1984\",SPHEROID[\"WGS 84\",6378137,298.257223563,AUTHORITY[\"EPSG\",\"7030\"]],AUTHORITY[\"EPSG\",\"6326\"]],PRIMEM[\"Greenwich\",0,AUTHORITY[\"EPSG\",\"8901\"]],UNIT[\"degree\",0.0174532925199433,AUTHORITY[\"EPSG\",\"9122\"]],AUTHORITY[\"EPSG\",\"4326\"]]"
    HORIZONTAL_VARIABLE_NAMES = ['lon', 'Easting', 'x', 'longitude']
    DEFAULT_MAX_BYTES = 500000000 # Default to 500,000,000 bytes for NCI's OPeNDAP
    def __init__(self, netcdf_dataset):
        '''
        NetCDF2DUtils Constructor - wraps a NetCDF dataset
        '''
        self.netcdf_dataset = netcdf_dataset
        
#        assert len(self.netcdf_dataset.dimensions) == 2, 'NetCDF dataset must be 2D' # This is not valid
        
        # Find variable with "grid_mapping" attribute - assumed to be 2D data variable
        try:
            self.data_variable = [variable for variable in self.netcdf_dataset.variables.values() if hasattr(variable, 'grid_mapping')][0]
        except:
            raise Exception('Unable to determine data variable (must have "grid_mapping" attribute')
        
        # Boolean flag indicating YX array ordering
        #TODO: Find a nicer way of dealing with this
        self.YX_order = self.data_variable.dimensions[1] in  NetCDF2DUtils.HORIZONTAL_VARIABLE_NAMES
        
        # Two-element list of dimension varibles. 
        self.dimension_arrays = [self.netcdf_dataset.variables[dimension_name][:] for dimension_name in self.data_variable.dimensions]
        
        self.grid_mapping_variable = netcdf_dataset.variables[self.data_variable.grid_mapping]
        
        self.GeoTransform = [float(number) for number in self.grid_mapping_variable.GeoTransform.strip().split(' ')]
        
        self.pixel_size = [abs(self.GeoTransform[1]), abs(self.GeoTransform[5])]
        if self.YX_order:
            self.pixel_size.reverse()
        
        self.min_extent = tuple([min(self.dimension_arrays[dim_index]) - self.pixel_size[dim_index]/2.0 for dim_index in range(2)])
        self.max_extent = tuple([max(self.dimension_arrays[dim_index]) + self.pixel_size[dim_index]/2.0 for dim_index in range(2)])
        
    
    def get_coordinate_transformation(self, from_crs=None, to_crs=None):
        '''
        Use GDAL to obtain a CoordinateTransformation object to transform to/from native NetCDF CRS
        @parameter from_crs: WKT or "EPSG:nnnn" string from which to transform
        @parameter from_crs: WKT or "EPSG:nnnn" string to which to transform
        '''
        from_crs = from_crs or self.grid_mapping_variable.spatial_ref
        to_crs = to_crs or self.grid_mapping_variable.spatial_ref
        
        # Assume native coordinates if no crs given
        if from_crs == to_crs:
            return None
        
        def get_spatial_ref_from_crs(crs):
            spatial_ref = SpatialReference()
            # Check for EPSG then Well Known Text
            epsg_match = re.match('^EPSG:(\d+)$', from_crs)
            if epsg_match:
                spatial_ref.ImportFromEPSG(int(epsg_match.group(1)))
            else: # Assume valid WKT definition
                spatial_ref.ImportFromWkt(crs)    
            return spatial_ref       
        
        from_spatial_ref = get_spatial_ref_from_crs(from_crs)
        to_spatial_ref = get_spatial_ref_from_crs(to_crs)
        
        # This is probably redundant
        if from_spatial_ref.ExportToWkt() == to_spatial_ref.ExportToWkt():
            return None
                   
        return CoordinateTransformation(from_spatial_ref, to_spatial_ref)
        
    
        
    def get_native_coords(self, coordinates, crs=None):
        '''
        Convert coordinates from specified CRS to native NetCDF CRS
        @parameter coordinates: iterable collection of coordinate pairs or single coordinate pair
        @parameter crs: Coordinate Reference System for coordinates. None == native NetCDF CRS
        '''
        coord_trans = self.get_coordinate_transformation(crs) # Transform from specified CRS to native CRS
        
        if not coord_trans: # No transformation required
            return list(coordinates)
        
        try: # Multiple coordinates
            return [coordinate[0:2] for coordinate in coord_trans.TransformPoints(coordinates)]
        except TypeError: # Single coordinate
            return coord_trans.TransformPoint(*coordinates)[0:2]
        
    
    def get_indices_from_coords(self, coordinates, crs=None):
        '''
        Returns list of netCDF array indices corresponding to coordinates to support nearest neighbour queries
        @parameter coordinates: iterable collection of coordinate pairs or single coordinate pair
        @parameter crs: Coordinate Reference System for coordinates. None == native NetCDF CRS
        '''
        native_coordinates = self.get_native_coords(coordinates, crs)
        
        # Convert coordinates to same order as array
        if self.YX_order:
            try:
                for coord_index in range(len(native_coordinates)):
                    if native_coordinates[coord_index] is not None:
                        native_coordinates[coord_index] = list(native_coordinates[coord_index])
                        native_coordinates[coord_index].reverse()
            except TypeError:
                native_coordinates = list(native_coordinates)            
                native_coordinates.reverse()
        try: # Multiple coordinates
            indices = [[np.where(abs(self.dimension_arrays[dim_index] - coordinate[dim_index]) <= (self.pixel_size[dim_index] / 2.0))[0][0] for dim_index in range(2)]
                       if not ([True for dim_index in range(2) if coordinate[dim_index] < self.min_extent[dim_index] or coordinate[dim_index] > self.max_extent[dim_index]])
                       else None
                       for coordinate in native_coordinates]
        except TypeError: # Single coordinate pair
            indices = ([np.where(abs(self.dimension_arrays[dim_index] - native_coordinates[dim_index]) <= (self.pixel_size[dim_index] / 2.0))[0][0] for dim_index in range(2)]
                       if not [True for dim_index in range(2) if native_coordinates[dim_index] < self.min_extent[dim_index] or native_coordinates[dim_index] > self.max_extent[dim_index]]
                       else None)
            
        return indices
        
        
    def get_fractional_indices_from_coords(self, coordinates, crs=None):
        '''
        Returns list of fractional array indices corresponding to coordinates to support interpolation
        @parameter coordinates: iterable collection of coordinate pairs or single coordinate pair
        @parameter crs: Coordinate Reference System for coordinates. None == native NetCDF CRS
        '''
        native_coordinates = self.get_native_coords(coordinates, crs)
        
        self.pixel_size
        
        # Convert coordinates to same order as array
        if self.YX_order:
            try:
                for coord_index in range(len(native_coordinates)):
                    if native_coordinates[coord_index] is not None:
                        native_coordinates[coord_index] = list(native_coordinates[coord_index])
                        native_coordinates[coord_index].reverse()
            except:
                native_coordinates = list(native_coordinates)            
                native_coordinates.reverse()
        #TODO: Make sure this still works with Southwards-positive datasets
        try: # Multiple coordinates
            fractional_indices = [[(coordinate[dim_index] - min(self.dimension_arrays[dim_index])) / self.pixel_size[dim_index] for dim_index in range(2)]
                       if not ([True for dim_index in range(2) if coordinate[dim_index] < self.min_extent[dim_index] or coordinate[dim_index] > self.max_extent[dim_index]])
                       else None
                       for coordinate in native_coordinates]
        except: # Single coordinate pair
            fractional_indices = ([(native_coordinates[dim_index] - min(self.dimension_arrays[dim_index])) / self.pixel_size[dim_index] for dim_index in range(2)]
                       if not [True for dim_index in range(2) if native_coordinates[dim_index] < self.min_extent[dim_index] or native_coordinates[dim_index] > self.max_extent[dim_index]]
                       else None)
            
        return fractional_indices<commit_msg>get_native_coords renamed to transform_coords and made more general. get_utm_crs function added<commit_after>'''
Created on 14Sep.,2016

@author: Alex
'''
import re
import numpy as np
import netCDF4
import math
import scipy.ndimage
from osgeo.osr import SpatialReference, CoordinateTransformation


class NetCDF2DUtils(object):
    '''
    NetCDF2DUtils class to do various fiddly things with NetCDF geophysics files.
    '''
    # Assume WGS84 lat/lon if no CRS is provided
    DEFAULT_CRS = "GEOGCS[\"WGS 84\",DATUM[\"WGS_1984\",SPHEROID[\"WGS 84\",6378137,298.257223563,AUTHORITY[\"EPSG\",\"7030\"]],AUTHORITY[\"EPSG\",\"6326\"]],PRIMEM[\"Greenwich\",0,AUTHORITY[\"EPSG\",\"8901\"]],UNIT[\"degree\",0.0174532925199433,AUTHORITY[\"EPSG\",\"9122\"]],AUTHORITY[\"EPSG\",\"4326\"]]"
    HORIZONTAL_VARIABLE_NAMES = ['lon', 'Easting', 'x', 'longitude']
    DEFAULT_MAX_BYTES = 500000000 # Default to 500,000,000 bytes for NCI's OPeNDAP
    def __init__(self, netcdf_dataset):
        '''
        NetCDF2DUtils Constructor - wraps a NetCDF dataset
        '''
        self.netcdf_dataset = netcdf_dataset
        
#        assert len(self.netcdf_dataset.dimensions) == 2, 'NetCDF dataset must be 2D' # This is not valid
        
        # Find variable with "grid_mapping" attribute - assumed to be 2D data variable
        try:
            self.data_variable = [variable for variable in self.netcdf_dataset.variables.values() if hasattr(variable, 'grid_mapping')][0]
        except:
            raise Exception('Unable to determine data variable (must have "grid_mapping" attribute')
        
        # Boolean flag indicating YX array ordering
        #TODO: Find a nicer way of dealing with this
        self.YX_order = self.data_variable.dimensions[1] in  NetCDF2DUtils.HORIZONTAL_VARIABLE_NAMES
        
        # Two-element list of dimension varibles. 
        self.dimension_arrays = [self.netcdf_dataset.variables[dimension_name][:] for dimension_name in self.data_variable.dimensions]
        
        self.grid_mapping_variable = netcdf_dataset.variables[self.data_variable.grid_mapping]
        
        self.GeoTransform = [float(number) for number in self.grid_mapping_variable.GeoTransform.strip().split(' ')]
        
        self.pixel_size = [abs(self.GeoTransform[1]), abs(self.GeoTransform[5])]
        if self.YX_order:
            self.pixel_size.reverse()
        
        self.min_extent = tuple([min(self.dimension_arrays[dim_index]) - self.pixel_size[dim_index]/2.0 for dim_index in range(2)])
        self.max_extent = tuple([max(self.dimension_arrays[dim_index]) + self.pixel_size[dim_index]/2.0 for dim_index in range(2)])
        
    
    def get_coordinate_transformation(self, from_crs=None, to_crs=None):
        '''
        Use GDAL to obtain a CoordinateTransformation object to transform coordinates between CRSs or None if no transformation required.
        @parameter from_crs: WKT or "EPSG:nnnn" string from which to transform. Defaults to native NetCDF CRS
        @parameter to_crs: WKT or "EPSG:nnnn" string to which to transform. Defaults to native NetCDF CRS
        '''
        from_crs = from_crs or self.grid_mapping_variable.spatial_ref
        to_crs = to_crs or self.grid_mapping_variable.spatial_ref
        
        # Assume native coordinates if no crs given
        if from_crs == to_crs:
            return None
        
        def get_spatial_ref_from_crs(crs):
            spatial_ref = SpatialReference()
            # Check for EPSG then Well Known Text
            epsg_match = re.match('^EPSG:(\d+)$', crs)
            if epsg_match:
                spatial_ref.ImportFromEPSG(int(epsg_match.group(1)))
            else: # Assume valid WKT definition
                spatial_ref.ImportFromWkt(crs)    
            return spatial_ref       
        
        from_spatial_ref = get_spatial_ref_from_crs(from_crs)
        to_spatial_ref = get_spatial_ref_from_crs(to_crs)
        
        # This is probably redundant
        if from_spatial_ref.ExportToWkt() == to_spatial_ref.ExportToWkt():
            return None
                   
        return CoordinateTransformation(from_spatial_ref, to_spatial_ref)
        
    def get_utm_crs(self, coordinate, from_crs=None):
        '''
        Function to return CRS for UTM zone of specified coordinates.
        Used to transform coords to metres
        '''
        def utm_getZone(longitude):
            return (int(1+(longitude+180.0)/6.0))
        
        def utm_isNorthern(latitude):
            return (latitude >= 0.0)    
        
        latlon_coord_trans = self.get_coordinate_transformation(from_crs, 'EPSG:4326')
        latlon_coord = coordinate if latlon_coord_trans is None else latlon_coord_trans.TransformPoint(*coordinate)[0:2]
      
        # Set UTM coordinate reference system
        utm_spatial_ref = SpatialReference()
        utm_spatial_ref.SetWellKnownGeogCS('WGS84')
        utm_spatial_ref.SetUTM(utm_getZone(latlon_coord[0]), utm_isNorthern(latlon_coord[1]));
        
        return utm_spatial_ref.ExportToPrettyWkt()

    def transform_coords(self, coordinates, from_crs=None, to_crs=None):
        '''
        Convert coordinates between specified coordinate reference systems
        @parameter coordinates: iterable collection of coordinate pairs or single coordinate pair
        @parameter from_crs: WKT or "EPSG:nnnn" string from which to transform. Defaults to native NetCDF CRS
        @parameter to_crs: WKT or "EPSG:nnnn" string to which to transform. Defaults to native NetCDF CRS
        '''
        from_crs = from_crs or self.grid_mapping_variable.spatial_ref
        to_crs = to_crs or self.grid_mapping_variable.spatial_ref

        coord_trans = self.get_coordinate_transformation(from_crs, to_crs) # Transform from specified CRS to native CRS
        
        if not coord_trans: # No transformation required
            return list(coordinates)
        
        try: # Multiple coordinates
            return [coordinate[0:2] for coordinate in coord_trans.TransformPoints(coordinates)]
        except TypeError: # Single coordinate
            return coord_trans.TransformPoint(*coordinates)[0:2]
        
    
    def get_indices_from_coords(self, coordinates, crs=None):
        '''
        Returns list of netCDF array indices corresponding to coordinates to support nearest neighbour queries
        @parameter coordinates: iterable collection of coordinate pairs or single coordinate pair
        @parameter crs: Coordinate Reference System for coordinates. None == native NetCDF CRS
        '''
        native_coordinates = self.transform_coords(coordinates, crs)
        
        # Convert coordinates to same order as array
        if self.YX_order:
            try:
                for coord_index in range(len(native_coordinates)):
                    if native_coordinates[coord_index] is not None:
                        native_coordinates[coord_index] = list(native_coordinates[coord_index])
                        native_coordinates[coord_index].reverse()
            except TypeError:
                native_coordinates = list(native_coordinates)            
                native_coordinates.reverse()
        try: # Multiple coordinates
            indices = [[np.where(abs(self.dimension_arrays[dim_index] - coordinate[dim_index]) <= (self.pixel_size[dim_index] / 2.0))[0][0] for dim_index in range(2)]
                       if not ([True for dim_index in range(2) if coordinate[dim_index] < self.min_extent[dim_index] or coordinate[dim_index] > self.max_extent[dim_index]])
                       else None
                       for coordinate in native_coordinates]
        except TypeError: # Single coordinate pair
            indices = ([np.where(abs(self.dimension_arrays[dim_index] - native_coordinates[dim_index]) <= (self.pixel_size[dim_index] / 2.0))[0][0] for dim_index in range(2)]
                       if not [True for dim_index in range(2) if native_coordinates[dim_index] < self.min_extent[dim_index] or native_coordinates[dim_index] > self.max_extent[dim_index]]
                       else None)
            
        return indices
        
        
    def get_fractional_indices_from_coords(self, coordinates, crs=None):
        '''
        Returns list of fractional array indices corresponding to coordinates to support interpolation
        @parameter coordinates: iterable collection of coordinate pairs or single coordinate pair
        @parameter crs: Coordinate Reference System for coordinates. None == native NetCDF CRS
        '''
        native_coordinates = self.transform_coords(coordinates, crs)
        
        self.pixel_size
        
        # Convert coordinates to same order as array
        if self.YX_order:
            try:
                for coord_index in range(len(native_coordinates)):
                    if native_coordinates[coord_index] is not None:
                        native_coordinates[coord_index] = list(native_coordinates[coord_index])
                        native_coordinates[coord_index].reverse()
            except:
                native_coordinates = list(native_coordinates)            
                native_coordinates.reverse()
        #TODO: Make sure this still works with Southwards-positive datasets
        try: # Multiple coordinates
            fractional_indices = [[(coordinate[dim_index] - min(self.dimension_arrays[dim_index])) / self.pixel_size[dim_index] for dim_index in range(2)]
                       if not ([True for dim_index in range(2) if coordinate[dim_index] < self.min_extent[dim_index] or coordinate[dim_index] > self.max_extent[dim_index]])
                       else None
                       for coordinate in native_coordinates]
        except: # Single coordinate pair
            fractional_indices = ([(native_coordinates[dim_index] - min(self.dimension_arrays[dim_index])) / self.pixel_size[dim_index] for dim_index in range(2)]
                       if not [True for dim_index in range(2) if native_coordinates[dim_index] < self.min_extent[dim_index] or native_coordinates[dim_index] > self.max_extent[dim_index]]
                       else None)
            
        return fractional_indices
{'fraction_non_alphanumeric': 0.08264245985937664, 'fraction_numerical': 0.014639521460803861, 'mean_word_length': 2.6048798940798186, 'pattern_counts': {'":': 0, '<': 15, '<?xml version=': 0, '>': 12, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 2, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3847262', 'n_tokens_mistral': 5715, 'n_tokens_neox': 5401, 'n_words': 1319}