<commit_before>"""
Collection of commonly used classes, functions and global variables
for reading/writing variables to a netCDF file

To import:
module_dir = os.path.join(os.environ['HOME'], 'modules')
sys.path.insert(0, module_dir)

Included functions:
convert_units        -- Convert units
coordinate_pairs     -- Produce all lat/lon pairs for a given grid
dict_filter          -- Filter dictionary 
get_datetime         -- Return datetime instances for list of dates/times
hi_lo                -- Update highest and lowest value
list_kwargs          -- List keyword arguments of a function
match_dates          --
regrid_uniform       -- Regrid data to a uniform output grid
running_average      -- Calculate running average
scale_offset         -- Apply scaling and offset factors
single2list          -- Check if item is a list, then convert if not
split_dt             -- Split a getTime().asComponentTime() date/time into year, month and day parts
temporal_aggregation -- Create a temporal aggregate of 
                        the input data (i.e. raw, climatology or anomaly)
time_axis_check      -- Check whether 2 time axes are the same
write_netcdf         -- Write an output netCDF file
xy_axis_check        -- Check whether 2 lat or lon axes are the same

Included classes:
InputData            -- Extract and subset data

"""

## Import general Python modules ##

import os
import sys
import re
import copy

import inspect
import calendar
import datetime
from dateutil.parser import *

import numpy
from scipy import stats

import cdutil
import genutil
import cdms2
#if hasattr(cdms2, 'setNetcdfDeflateFlag'):
cdms2.setNetcdfShuffleFlag(0)
cdms2.setNetcdfDeflateFlag(0)
cdms2.setNetcdfDeflateLevelFlag(0)
import MV2
import regrid2


## Import my modules ##

cwd = os.getcwd()
repo_dir = '/'
for directory in cwd.split('/')[1:]:
    repo_dir = os.path.join(repo_dir, directory)
    if directory == 'phd':
        break

modules_dir = os.path.join(repo_dir, 'modules')
sys.path.append(modules_dir)

try:
    import general_io as gio
except ImportError:
    raise ImportError('Must run this script from anywhere within the phd git repo')

  
## Define regions ##

# The 3rd argument here works as follows:
# - First two letters: 'c' or 'o' for closed or open upper/loweer bounds
# - Third letter represents the search method: 
#   'b' bounds, 'n' node, 'e' extranode, 's' select

# [(minlat, maxlat), (minlon, maxlon)]

regions = {'aus': [(-45, -10, 'cc'), (110, 160, 'cc')],
           'ausnz': [(-50, 0, 'cc'), (100, 185, 'cc')],
           'emia': [(-10, 10, 'cc'), (165, 220, 'cc')],
           'emib': [(-15, 5, 'cc'), (250, 290, 'cc')],
           'emic': [(-10, 20, 'cc'), (125, 145, 'cc')],
           'eqpacific': [(-30, 30, 'cc'), (120, 280, 'cc')],
           'nino1': [(-10, -5, 'cc'), (270,280,'cc')],
           'nino2': [(-5, 0, 'cc'), (270, 280, 'cc')],
           'nino12': [(-10, 0, 'cc'), (270, 280, 'cc')],
           'nino3': [(-5, 5, 'cc'), (210, 270, 'cc')],
           'nino34': [(-5, 5, 'cc'), (190, 240, 'cc')],
           'nino4': [(-5, 5, 'cc'), (160, 210, 'cc')],
           'sh': [(-90, 0, 'cc'), (0, 360, 'co')],
	   'shextropics20': [(-90, -20, 'cc'), (0, 360, 'co')],
           'shextropics30': [(-90, -30, 'cc'), (0, 360, 'co')],
           'tropics': [(-30, 30, 'cc'), (0, 360, 'co')],
	   'glatt': [(20, 80, 'cc'), (-180, 180, 'co')],
           'nonpolar70': [(-70, 70, 'cc'), (0, 360, 'co')],
	   'nonpolar80': [(-80, 80, 'cc'), (0, 360, 'co')],
	   'sh-psa': [(-90, 0, 'cc'), (90, 450, 'co')],
	   'sh-psa-extra': [(-90, 30, 'cc'), (90, 450, 'co')],
	   'world-dateline': [(-90, 90, 'cc'), (0, 360, 'co')],
	   'world-dateline-duplicate360': [(-90, 90, 'cc'), (0, 360, 'cc')],
	   'world-greenwich': [(-90, 90, 'cc'), (-180, 180, 'co')],
	   'world-psa': [(-90, 90, 'cc'), (90, 450, 'co')],
	   'zw31': [(-50, -45, 'cc'), (45, 60, 'cc')],
           'zw32': [(-50, -45, 'cc'), (161, 171, 'cc')],
           'zw33': [(-50, -45, 'cc'), (279, 289, 'cc')],
           }


## Classes/functions ##

class InputData:
    """Extract and subset data."""

    def __init__(self, fname, var_id, convert=False, **kwargs):
	"""Extract desired data from an input file.
	
	Keyword arguments (with examples):
	
	SUBSETTORS
        latitude  -- (-30, 30)
        level     -- (1000.)
        longitude -- (120, 165)
        region    -- aus
        time      -- ('1979-01-01', '2000-12-31', 'MONTH/SEASON'), 
		     MONTH/SEASON (optional 3rd argument): 
                     'JAN', 'FEB', ..., 'DEC'
		     'DJF', ... 'SON'	

	MANIPULATORS
	agg       -- (quantity, season, lower_time_bound, upper_time_bound)
                     quantity: 'raw', 'climatology' or 'anomaly'
                     season: ANNUALCYCLE, SEASONALCYCLE, DJF, MJJASO etc
                     time_bounds (optional):
                     e.g. '1979-01-01', '1980-12-31'
        convert   -- True or False (for converting units)
        grid      -- (startLat,nlat,deltaLat,
	              startLon,nlon,deltaLon)
        runave    -- window size for the running average 

        The order of operations is as follows: subset data,
        temporal aggregation (agg), running average (runave),
        regrid (grid), convert units 
       		
        self.data has all the attributes and methods
	of a typical cdms2 variable. For instance:
        - self.data.getLatitude()
	- self.data.getLatitude()
	- self.data.getTime()
	- self.data.attributes (dictionary incl. _FillValue, units etc)       
	
        """

        infile = cdms2.open(fname)       	
        _infile_attribute_check(infile, var_id)
        kwargs['order'] = _define_order(infile, var_id)

        # Subset input data #

        if kwargs.has_key('region'):   
            try:
                kwargs['latitude'], kwargs['longitude'] = regions[kwargs['region']]
                self.minlat, self.maxlat = kwargs['latitude'][0:2]
                self.minlon, self.maxlon = kwargs['longitude'][0:2]
                self.region = kwargs['region']
	    except KeyError:
                print 'region not defined - using all spatial data...'    
	    del kwargs['region']

        #remove None values
        for key in kwargs:
            if not kwargs[key]:
	        del kwargs[key]

        subsettors = ['latitude', 'level', 'longitude', 'time']
        subset_kwargs = {}
        for key in kwargs.keys():
            if key in subsettors:
                subset_kwargs[key] = kwargs[key]

        data = _subset_data(infile, var_id, **subset_kwargs)        
       
        # Manipulate the subsetted data #  

	if kwargs.has_key('agg'):
            quantity = kwargs['agg'][0]
	    timescale = kwargs['agg'][1]
            times = [kwargs['agg'][2], kwargs['agg'][3]] if len(kwargs['agg']) > 2 else None
            data = temporal_aggregation(data, timescale, quantity, time_period=times)

        if kwargs.has_key('runave'):
	    window = kwargs['runave']
            data = running_average(data, window) if window > 1 else data

        if kwargs.has_key('grid'):
            data = regrid_uniform(data, kwargs['grid'])            

        if convert:
	    data = convert_units(data)

        # Set object attributes #
        
	if 'x' in data.getOrder():
            if not (data.getLongitude()[0] - (data.getLongitude()[-1] - 360)) > 0:
	        print 'WARNING: There are duplicate longitude values (can be problematic for some applications)'  
<commit_msg>Added spatial averaging to netcdf_io.py<commit_after>"""
Collection of commonly used classes, functions and global variables
for reading/writing variables to a netCDF file

To import:
module_dir = os.path.join(os.environ['HOME'], 'modules')
sys.path.insert(0, module_dir)

Included functions:
convert_units        -- Convert units
coordinate_pairs     -- Produce all lat/lon pairs for a given grid
dict_filter          -- Filter dictionary 
get_datetime         -- Return datetime instances for list of dates/times
hi_lo                -- Update highest and lowest value
list_kwargs          -- List keyword arguments of a function
match_dates          --
regrid_uniform       -- Regrid data to a uniform output grid
running_average      -- Calculate running average
scale_offset         -- Apply scaling and offset factors
single2list          -- Check if item is a list, then convert if not
split_dt             -- Split a getTime().asComponentTime() date/time into year, month and day parts
temporal_aggregation -- Create a temporal aggregate of 
                        the input data (i.e. raw, climatology or anomaly)
time_axis_check      -- Check whether 2 time axes are the same
write_netcdf         -- Write an output netCDF file
xy_axis_check        -- Check whether 2 lat or lon axes are the same

Included classes:
InputData            -- Extract and subset data

"""

## Import general Python modules ##

import os
import sys
import re
import copy

import inspect
import calendar
import datetime
from dateutil.parser import *

import numpy
from scipy import stats

import cdutil
import genutil
import cdms2
#if hasattr(cdms2, 'setNetcdfDeflateFlag'):
cdms2.setNetcdfShuffleFlag(0)
cdms2.setNetcdfDeflateFlag(0)
cdms2.setNetcdfDeflateLevelFlag(0)
import MV2
import regrid2


## Import my modules ##

cwd = os.getcwd()
repo_dir = '/'
for directory in cwd.split('/')[1:]:
    repo_dir = os.path.join(repo_dir, directory)
    if directory == 'phd':
        break

modules_dir = os.path.join(repo_dir, 'modules')
sys.path.append(modules_dir)

try:
    import general_io as gio
except ImportError:
    raise ImportError('Must run this script from anywhere within the phd git repo')

  
## Define regions ##

# The 3rd argument here works as follows:
# - First two letters: 'c' or 'o' for closed or open upper/loweer bounds
# - Third letter represents the search method: 
#   'b' bounds, 'n' node, 'e' extranode, 's' select

# [(minlat, maxlat), (minlon, maxlon)]

regions = {'aus': [(-45, -10, 'cc'), (110, 160, 'cc')],
           'ausnz': [(-50, 0, 'cc'), (100, 185, 'cc')],
           'emia': [(-10, 10, 'cc'), (165, 220, 'cc')],
           'emib': [(-15, 5, 'cc'), (250, 290, 'cc')],
           'emic': [(-10, 20, 'cc'), (125, 145, 'cc')],
           'eqpacific': [(-30, 30, 'cc'), (120, 280, 'cc')],
           'nino1': [(-10, -5, 'cc'), (270,280,'cc')],
           'nino2': [(-5, 0, 'cc'), (270, 280, 'cc')],
           'nino12': [(-10, 0, 'cc'), (270, 280, 'cc')],
           'nino3': [(-5, 5, 'cc'), (210, 270, 'cc')],
           'nino34': [(-5, 5, 'cc'), (190, 240, 'cc')],
           'nino4': [(-5, 5, 'cc'), (160, 210, 'cc')],
           'sh': [(-90, 0, 'cc'), (0, 360, 'co')],
	   'shextropics20': [(-90, -20, 'cc'), (0, 360, 'co')],
           'shextropics30': [(-90, -30, 'cc'), (0, 360, 'co')],
           'tropics': [(-30, 30, 'cc'), (0, 360, 'co')],
	   'glatt': [(20, 80, 'cc'), (-180, 180, 'co')],
           'nonpolar70': [(-70, 70, 'cc'), (0, 360, 'co')],
	   'nonpolar80': [(-80, 80, 'cc'), (0, 360, 'co')],
	   'sh-psa': [(-90, 0, 'cc'), (90, 450, 'co')],
	   'sh-psa-extra': [(-90, 30, 'cc'), (90, 450, 'co')],
	   'world-dateline': [(-90, 90, 'cc'), (0, 360, 'co')],
	   'world-dateline-duplicate360': [(-90, 90, 'cc'), (0, 360, 'cc')],
	   'world-greenwich': [(-90, 90, 'cc'), (-180, 180, 'co')],
	   'world-psa': [(-90, 90, 'cc'), (90, 450, 'co')],
	   'zw31': [(-50, -45, 'cc'), (45, 60, 'cc')],
           'zw32': [(-50, -45, 'cc'), (161, 171, 'cc')],
           'zw33': [(-50, -45, 'cc'), (279, 289, 'cc')],
           }


## Classes/functions ##

class InputData:
    """Extract and subset data."""

    def __init__(self, fname, var_id, convert=False, **kwargs):
	"""Extract desired data from an input file.
	
	Keyword arguments (with examples):
	
	SUBSETTORS
        latitude  -- (-30, 30)
        level     -- (1000.)
        longitude -- (120, 165)
        region    -- aus
        time      -- ('1979-01-01', '2000-12-31', 'MONTH/SEASON'), 
		     MONTH/SEASON (optional 3rd argument): 
                     'JAN', 'FEB', ..., 'DEC'
		     'DJF', ... 'SON'	

	MANIPULATORS
	agg       -- (quantity, season, lower_time_bound, upper_time_bound)
                     quantity: 'raw', 'climatology' or 'anomaly'
                     season: ANNUALCYCLE, SEASONALCYCLE, DJF, MJJASO etc
                     time_bounds (optional):
                     e.g. '1979-01-01', '1980-12-31'
        convert   -- True or False (for converting units)
        grid      -- (startLat,nlat,deltaLat,
	              startLon,nlon,deltaLon)
        runave    -- window size for the running average 
	spatave   -- True (for returning average over all spatial dimensions)

        The order of operations is as follows: subset data, 
	spatial averaging (spatave), temporal aggregation (agg), 
	running average (runave), regrid (grid), convert units 
       		
        self.data has all the attributes and methods
	of a typical cdms2 variable. For instance:
        - self.data.getLatitude()
	- self.data.getLatitude()
	- self.data.getTime()
	- self.data.attributes (dictionary incl. _FillValue, units etc)       
	
        """

        infile = cdms2.open(fname)       	
        _infile_attribute_check(infile, var_id)
        kwargs['order'] = _define_order(infile, var_id)

        # Subset input data #

        if kwargs.has_key('region'):   
            try:
                kwargs['latitude'], kwargs['longitude'] = regions[kwargs['region']]
                self.minlat, self.maxlat = kwargs['latitude'][0:2]
                self.minlon, self.maxlon = kwargs['longitude'][0:2]
                self.region = kwargs['region']
	    except KeyError:
                print 'region not defined - using all spatial data...'    
	    del kwargs['region']

        #remove None values
        for key in kwargs:
            if not kwargs[key]:
	        del kwargs[key]

        subsettors = ['latitude', 'level', 'longitude', 'time']
        subset_kwargs = {}
        for key in kwargs.keys():
            if key in subsettors:
                subset_kwargs[key] = kwargs[key]

        data = _subset_data(infile, var_id, **subset_kwargs)        
       
        # Manipulate the subsetted data #  

        if kwargs.has_key('spatave'):
            ave_axes = data.getOrder().translate(None, 't')
            data = cdutil.averager(data, axis=ave_axes, weights=['unweighted']*len(ave_axes))

	if kwargs.has_key('agg'):
            quantity = kwargs['agg'][0]
	    timescale = kwargs['agg'][1]
            times = [kwargs['agg'][2], kwargs['agg'][3]] if len(kwargs['agg']) > 2 else None
            data = temporal_aggregation(data, timescale, quantity, time_period=times)

        if kwargs.has_key('runave'):
	    window = kwargs['runave']
            data = running_average(data, window) if window > 1 else data

        if kwargs.has_key('grid'):
            data = regrid_uniform(data, kwargs['grid'])            

        if convert:
	    data = convert_units(data)

        # Set object attributes #
        
	if 'x' in data.getOrder():
            if not (data.getLongitude()[0] - (data.getLongitude()[-1] - 360)) > 0:
	        print 'WARNING: There are duplicate longitude values (can be problematic for some applications)'  

{'fraction_non_alphanumeric': 0.15876044476610304, 'fraction_numerical': 0.04618724916112902, 'mean_word_length': 2.374777975133215, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 9, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, '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': '2184289', 'n_tokens_mistral': 5526, 'n_tokens_neox': 4914, 'n_words': 1493}