<commit_before>"""
Filename:     calc_climate_index.py
Author:       <NAME>, <EMAIL>
Description:  Calculates the selected climate index

Input:        List of netCDF files to plot
Output:       Text file

"""

# Import general Python modules

import sys, os
import argparse
import numpy
from cdo import *
cdo = Cdo()
import pdb

# 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 netcdf_io as nio
except ImportError:
    raise ImportError('Must run this script from anywhere within the phd git repo')



def calc_monthly_climatology(base_timeseries, months):
    """Calcuate the monthly climatology.
    
    The base_timeseries must have a monthly
    timescale that begins in January.
    
    """

    assert (months[0] == 1), 'the base period must start in January'

    ntime_base = len(base_timeseries)
    monthly_climatology_mean = numpy.ma.zeros(12)
    monthly_climatology_std = numpy.ma.zeros(12)
    for i in range(0, 12):
        monthly_climatology_mean[i] = numpy.ma.mean(base_timeseries[i:ntime_base:12])
	monthly_climatology_std[i] = numpy.ma.std(base_timeseries[i:ntime_base:12])

    return monthly_climatology_mean, monthly_climatology_std


def calc_monthly_anomaly(complete_timeseries, base_timeseries, months):
    """Calculate monthly anomaly.""" 
    
    monthly_climatology_mean = calc_monthly_climatology(base_timeseries, months)[0]
    
    ntime_complete = len(complete_timeseries)
    monthly_anomaly = numpy.ma.zeros(ntime_complete)
    for i in range(0, ntime_complete):
	month_index = months[i]
	monthly_anomaly[i] = numpy.ma.subtract(complete_timeseries[i], 
	                                       monthly_climatology_mean[month_index-1])
    
    return monthly_anomaly 


def monthly_normalisation(complete_timeseries, base_timeseries, months):
    """Normalise the monthly timeseries: (x - mean) / stdev."""  
    
    # Calculate the monthly climatology #
    
    monthly_climatology_mean, monthly_climatology_std = calc_monthly_climatology(base_timeseries, months)

    # Normalise the entire timeseries #
    
    ntime_complete = len(complete_timeseries)
    monthly_normalised = numpy.ma.zeros(ntime_complete)
    for i in range(0, ntime_complete):
	month_index = months[i]
	monthly_normalised[i] = numpy.ma.divide((numpy.ma.subtract(complete_timeseries[i], 
	                        monthly_climatology_mean[month_index-1])), monthly_climatology_std[month_index-1])
    
    return monthly_normalised


def calc_reg_anomaly_timeseries(data_complete, data_base):
    """Calculate the monthly anomaly timeseries for a given region.

    Input arguments must be netcdf_io.InputData instances.

    """

    assert isinstance(data_complete, nio.InputData), \
    'input arguments must be nio.InputData instances'

    assert isinstance(data_base, nio.InputData), \
    'input arguments must be nio.InputData instances'

    ntime_complete, nlats, nlons = numpy.shape(data_complete.data)
    ntime_base, nlats, nlons = numpy.shape(data_base.data)

    data_complete_flat = numpy.ma.reshape(data_complete.data, (int(ntime_complete), int(nlats * nlons)))    # Flattens the spatial dimension
    data_base_flat = numpy.ma.reshape(data_base.data, (int(ntime_base), int(nlats * nlons)))

    # Calculate the anomalies #

    complete_timeseries = numpy.ma.mean(data_complete_flat, axis=1)
    base_timeseries = numpy.ma.mean(data_base_flat, axis=1)

    anomaly_timeseries = calc_monthly_anomaly(complete_timeseries, 
                         base_timeseries, data_complete.months())

    return anomaly_timeseries
    

def map_std(stds, data, timescale):
   """Takes the array of stardard deviations and map it
   to the size of the data array (matching the correct month/day"""
   
   assert timescale in ['daily', 'monthly']
   
   times = data.getTime().asComponentTime()
   dts = nio.get_datetime(times)
   
   if timescale == 'monthly':
       months = map(lambda x: x.month, dts)
       result = map(lambda x: stds[x-1], months)
   elif timescale == 'daily':
       day_of_year = map(nio.day_of_year_366, dts)
       result = map(lambda x: stds[x-1], day_of_year)
   
   return result


def calc_zw3(index, ifile, var_id, base_period):
    """Calculate an index of the SH ZW3 pattern
    
    Method as per Raphael (2004)
    
    This function uses cdo instead of CDAT because the
    cdutil library doesn't have routines for calculating 
    the daily climatology or stdev.
    
    The running mean should have been applied to the input 
    data beforehand. Raphael (2004) uses a 3-month running 
    mean.
    
    Input data that Raphael (2004) uses is the 500hPa 
    geopotential height, the sea level pressure or
    from 500hPa zonal anomalies which are constructed by 
    removing the zonal mean of the geopotential height from
    each grid point (preferred).
    
    """

    # Calulate the index
    #pdb.set_trace()    
    index = {}
    for region in ['zw31', 'zw32', 'zw33']: 
        south_lat, north_lat = nio.regions[region][0][0: 2]
        west_lon, east_lon = nio.regions[region][1][0: 2]

        indata_complete = nio.InputData(ifile, var_id, region=region, spatave=True) 
        time_axis = indata_complete.data.getTime().asComponentTime()
        timescale = nio.get_timescale(nio.get_datetime(time_axis[0:2]))
        
	assert timescale in ['daily', 'monthly']
	tscale_abbrev = 'day' if timescale == 'daily' else 'mon' 
	
	sub_operator_text = 'cdo y%ssub ' %(tscale_abbrev)
	sub_operator_func = eval(sub_operator_text.replace(' ', '.', 1))
	std_operator_text = 'cdo y%sstd ' %(tscale_abbrev)
	std_operator_func = eval(std_operator_text.replace(' ', '.', 1))
	avg_operator_text = ' -y%savg ' %(tscale_abbrev)
	
        selregion = "-sellonlatbox,%d,%d,%d,%d %s " %(west_lon, east_lon, 
	                                              south_lat, north_lat, 
						      ifile)
        fldmean = "-fldmean "+selregion

        anom_cmd_entry = fldmean + avg_operator_text + fldmean
        print sub_operator_text + anom_cmd_entry
        anom = sub_operator_func(input=anom_cmd_entry, returnArray=var_id)

        print std_operator_text + fldmean
        std = std_operator_func(input=fldmean, returnArray=var_id)
	std = map_std(std, indata_complete.data, timescale)
	
	index[region] = numpy.squeeze(anom) / numpy.squeeze(std)

    zw3_timeseries = (index['zw31'] + index['zw32'] + index['zw33']) / 3.0
 
    # Define the output attributes
    	
    hx = 'Ref: ZW3 index of Raphael (2004)'
    var_atts = {'id': 'zw3',
                'long_name': 'Zonal Wave 3 index Index',
                'units': '',
                'history': hx}

    return zw3_timeseries, var_atts, indata_complete.global_atts, indata_complete.data.getTime()
    
    
def calc_sam(index, ifile, var_id, base_period):
    """Calculate an index of the Southern Annular Mode.
<commit_msg>Made the entire ZW3 index process cdo<commit_after>"""
Filename:     calc_climate_index.py
Author:       <NAME>, <EMAIL>
Description:  Calculates the selected climate index

Input:        List of netCDF files to plot
Output:       Text file

"""

# Import general Python modules

import sys, os
import argparse
import numpy
from cdo import *
cdo = Cdo()
import pdb

# 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 netcdf_io as nio
except ImportError:
    raise ImportError('Must run this script from anywhere within the phd git repo')



def calc_monthly_climatology(base_timeseries, months):
    """Calcuate the monthly climatology.
    
    The base_timeseries must have a monthly
    timescale that begins in January.
    
    """

    assert (months[0] == 1), 'the base period must start in January'

    ntime_base = len(base_timeseries)
    monthly_climatology_mean = numpy.ma.zeros(12)
    monthly_climatology_std = numpy.ma.zeros(12)
    for i in range(0, 12):
        monthly_climatology_mean[i] = numpy.ma.mean(base_timeseries[i:ntime_base:12])
	monthly_climatology_std[i] = numpy.ma.std(base_timeseries[i:ntime_base:12])

    return monthly_climatology_mean, monthly_climatology_std


def calc_monthly_anomaly(complete_timeseries, base_timeseries, months):
    """Calculate monthly anomaly.""" 
    
    monthly_climatology_mean = calc_monthly_climatology(base_timeseries, months)[0]
    
    ntime_complete = len(complete_timeseries)
    monthly_anomaly = numpy.ma.zeros(ntime_complete)
    for i in range(0, ntime_complete):
	month_index = months[i]
	monthly_anomaly[i] = numpy.ma.subtract(complete_timeseries[i], 
	                                       monthly_climatology_mean[month_index-1])
    
    return monthly_anomaly 


def monthly_normalisation(complete_timeseries, base_timeseries, months):
    """Normalise the monthly timeseries: (x - mean) / stdev."""  
    
    # Calculate the monthly climatology #
    
    monthly_climatology_mean, monthly_climatology_std = calc_monthly_climatology(base_timeseries, months)

    # Normalise the entire timeseries #
    
    ntime_complete = len(complete_timeseries)
    monthly_normalised = numpy.ma.zeros(ntime_complete)
    for i in range(0, ntime_complete):
	month_index = months[i]
	monthly_normalised[i] = numpy.ma.divide((numpy.ma.subtract(complete_timeseries[i], 
	                        monthly_climatology_mean[month_index-1])), monthly_climatology_std[month_index-1])
    
    return monthly_normalised


def calc_reg_anomaly_timeseries(data_complete, data_base):
    """Calculate the monthly anomaly timeseries for a given region.

    Input arguments must be netcdf_io.InputData instances.

    """

    assert isinstance(data_complete, nio.InputData), \
    'input arguments must be nio.InputData instances'

    assert isinstance(data_base, nio.InputData), \
    'input arguments must be nio.InputData instances'

    ntime_complete, nlats, nlons = numpy.shape(data_complete.data)
    ntime_base, nlats, nlons = numpy.shape(data_base.data)

    data_complete_flat = numpy.ma.reshape(data_complete.data, (int(ntime_complete), int(nlats * nlons)))    # Flattens the spatial dimension
    data_base_flat = numpy.ma.reshape(data_base.data, (int(ntime_base), int(nlats * nlons)))

    # Calculate the anomalies #

    complete_timeseries = numpy.ma.mean(data_complete_flat, axis=1)
    base_timeseries = numpy.ma.mean(data_base_flat, axis=1)

    anomaly_timeseries = calc_monthly_anomaly(complete_timeseries, 
                         base_timeseries, data_complete.months())

    return anomaly_timeseries
    

def map_std(stds, data, timescale):
   """Takes the array of stardard deviations and map it
   to the size of the data array (matching the correct month/day"""
   
   assert timescale in ['daily', 'monthly']
   
   times = data.getTime().asComponentTime()
   dts = nio.get_datetime(times)
   
   if timescale == 'monthly':
       months = map(lambda x: x.month, dts)
       result = map(lambda x: stds[x-1], months)
   elif timescale == 'daily':
       day_of_year = map(nio.day_of_year_366, dts)
       result = map(lambda x: stds[x-1], day_of_year)
   
   return result


def calc_zw3(index, ifile, var_id, base_period):
    """Calculate an index of the SH ZW3 pattern
    
    Method as per Raphael (2004)
    
    This function uses cdo instead of CDAT because the
    cdutil library doesn't have routines for calculating 
    the daily climatology or stdev.
    
    The running mean should have been applied to the input 
    data beforehand. Raphael (2004) uses a 3-month running 
    mean.
    
    Input data that Raphael (2004) uses is the 500hPa 
    geopotential height, the sea level pressure or
    from 500hPa zonal anomalies which are constructed by 
    removing the zonal mean of the geopotential height from
    each grid point (preferred).
    
    """

    # Calulate the index
    #pdb.set_trace()    
    index = {}
    for region in ['zw31', 'zw32', 'zw33']: 
        south_lat, north_lat = nio.regions[region][0][0: 2]
        west_lon, east_lon = nio.regions[region][1][0: 2]

        indata_complete = nio.InputData(ifile, var_id, region=region, spatave=True) 
        time_axis = indata_complete.data.getTime().asComponentTime()
        timescale = nio.get_timescale(nio.get_datetime(time_axis[0:2]))
        
	assert timescale in ['daily', 'monthly']
	tscale_abbrev = 'day' if timescale == 'daily' else 'mon' 
	
	
	div_operator_text = 'cdo y%sdiv ' %(tscale_abbrev)
	div_operator_func = eval(div_operator_text.replace(' ', '.', 1))
	sub_operator_text = ' -y%ssub ' %(tscale_abbrev)
	avg_operator_text = ' -y%savg ' %(tscale_abbrev)
	std_operator_text = ' -y%sstd ' %(tscale_abbrev)
	
	selregion = "-sellonlatbox,%d,%d,%d,%d %s " %(west_lon, east_lon, 
	                                              south_lat, north_lat, 
						      ifile)
        fldmean = "-fldmean "+selregion
        anomaly = sub_operator_text + fldmean + avg_operator_text + fldmean
        std = std_operator_text + fldmean
	
	print div_operator_text + anomaly + std   #e.g. cdo ydaydiv anomaly std
        result = div_operator_func(input=anomaly + std, returnArray=var_id)
	index[region] = numpy.squeeze(result)

    zw3_timeseries = (index['zw31'] + index['zw32'] + index['zw33']) / 3.0
 
    # Define the output attributes
    	
    hx = 'Ref: ZW3 index of Raphael (2004)'
    var_atts = {'id': 'zw3',
                'long_name': 'Zonal Wave 3 index Index',
                'units': '',
                'history': hx}

    return zw3_timeseries, var_atts, indata_complete.global_atts, indata_complete.data.getTime()
    
    
def calc_sam(index, ifile, var_id, base_period):
    """Calculate an index of the Southern Annular Mode.

{'fraction_non_alphanumeric': 0.09847404449515117, 'fraction_numerical': 0.011551625784369652, 'mean_word_length': 3.3501861042183623, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 7, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 6, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2540432', 'n_tokens_mistral': 4854, 'n_tokens_neox': 4576, 'n_words': 1117}