<commit_before>"""
Filename:     plot_zonal_ensemble.py
Author:       <NAME>, <EMAIL>
Description:  Plot zonal mean for an ensemble of models  

"""

# Import general Python modules

import sys, os, pdb
import argparse
import numpy
import iris
import iris.plot as iplt
import matplotlib.pyplot as plt
from matplotlib import gridspec
import seaborn

# Import my modules

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

modules_dir = os.path.join(repo_dir, 'modules')
sys.path.append(modules_dir)
try:
    import general_io as gio
    import timeseries
    import grids
    import convenient_universal as uconv
except ImportError:
    raise ImportError('Must run this script from anywhere within the ocean-analysis git repo')


# Define functions

def calc_trend_cube(cube):
    """Calculate trend and put into appropriate cube."""
    
    trend_array = timeseries.calc_trend(cube, per_yr=True)
    new_cube = cube[0,:].copy()
    new_cube.remove_coord('time')
    new_cube.data = trend_array
    
    return new_cube


def get_colors(infiles):
    """Define a color for each model"""

    models = []
    for infile in infiles:
        filename = infile.split('/')[-1]
        model = filename.split('_')[2]

        models.append(model) 
    
    nmodels = len(set(models))
    cm = plt.get_cmap('nipy_spectral')
    colors = [cm(1. * i / nmodels) for i in range(nmodels)]
    color_dict = {}
    count = 0
    for model in models:
        if not model in color_dict.keys():
            color_dict[model] = colors[count]
            count = count + 1

    return color_dict


def get_ylabel(cube, inargs):
    """get the y axis label"""

    ylabel = '$%s' %(str(cube.units))
    if inargs.perlat:
        ylabel = ylabel + ' \: lat^{-1}'
    if inargs.time_agg == 'trend':
        ylabel = ylabel + ' \: yr^{-1}'
    ylabel = ylabel + '$'

    return ylabel


def get_line_width(realization, model):
    """Get the line width"""

    if model == 'FGOALS-g2':
        lw = 2.0
    else:
        lw = 2.0 if realization == 1 else 0.5

    return lw


def main(inargs):
    """Run the program."""
    
    time_constraint = gio.get_time_constraint(inargs.time)
    fig, ax = plt.subplots()
    plt.axhline(y=0, color='0.5', linestyle='--')
    color_dict = get_colors(inargs.infiles)
    
    legend_models = []
    for infile in inargs.infiles:
        with iris.FUTURE.context(cell_datetime_objects=True):
            cube = iris.load_cube(infile, gio.check_iris_var(inargs.var) & time_constraint)
        
        if inargs.perlat:
            grid_spacing = grids.get_grid_spacing(cube) 
            cube.data = cube.data / grid_spacing
 
        if inargs.time_agg == 'trend':
            agg_cube = calc_trend_cube(cube)
            plot_name = 'linear trend'
        elif inargs.time_agg == 'climatology':
            agg_cube = cube.collapsed('time', iris.analysis.MEAN)
            plot_name = 'climatology'

        model = cube.attributes['model_id']
        if model not in legend_models:
            label = model
            legend_models.append(model)
        else:
            label = None

        realization = cube.attributes['realization']
        lw = get_line_width(realization, model)
        iplt.plot(agg_cube, label=label, color=color_dict[model], linewidth=lw)

    title = '%s, %s-%s' %(plot_name, inargs.time[0][0:4], inargs.time[1][0:4])
    plt.title(title)
    plt.xlim(-90, 90)
    ylabel = get_ylabel(cube, inargs)
    plt.ylabel(ylabel)
    plt.xlabel('latitude')

    box = ax.get_position()
    ax.set_position([box.x0, box.y0, box.width * 0.8, box.height])
    ax.legend(loc='center left', bbox_to_anchor=(1, 0.5))

    plt.savefig(inargs.outfile, bbox_inches='tight')
    gio.write_metadata(inargs.outfile, file_info={infile: cube.attributes['history']})


if __name__ == '__main__':

    extra_info =""" 

author:
    <NAME>, <EMAIL>

"""

    description = 'Plot zonal ensemble'
    parser = argparse.ArgumentParser(description=description,<commit_msg>Added groupby funcionality to account for physics<commit_after>"""
Filename:     plot_zonal_ensemble.py
Author:       <NAME>, <EMAIL>
Description:  Plot zonal mean for an ensemble of models  

"""

# Import general Python modules

import sys, os, pdb
import argparse
import numpy
import iris
import iris.plot as iplt
import matplotlib.pyplot as plt
from matplotlib import gridspec
import seaborn

# Import my modules

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

modules_dir = os.path.join(repo_dir, 'modules')
sys.path.append(modules_dir)
try:
    import general_io as gio
    import timeseries
    import grids
    import convenient_universal as uconv
except ImportError:
    raise ImportError('Must run this script from anywhere within the ocean-analysis git repo')


# Define functions

def calc_trend_cube(cube):
    """Calculate trend and put into appropriate cube."""
    
    trend_array = timeseries.calc_trend(cube, per_yr=True)
    new_cube = cube[0,:].copy()
    new_cube.remove_coord('time')
    new_cube.data = trend_array
    
    return new_cube


def get_colors(infiles):
    """Define a color for each model"""

    models = []
    for infile in infiles:
        filename = infile.split('/')[-1]
        model = filename.split('_')[2]

        models.append(model) 
    
    nmodels = len(set(models))
    cm = plt.get_cmap('nipy_spectral')
    colors = [cm(1. * i / nmodels) for i in range(nmodels)]
    color_dict = {}
    count = 0
    for model in models:
        if not model in color_dict.keys():
            color_dict[model] = colors[count]
            count = count + 1

    return color_dict


def get_ylabel(cube, inargs):
    """get the y axis label"""

    ylabel = '$%s' %(str(cube.units))
    if inargs.perlat:
        ylabel = ylabel + ' \: lat^{-1}'
    if inargs.time_agg == 'trend':
        ylabel = ylabel + ' \: yr^{-1}'
    ylabel = ylabel + '$'

    return ylabel


def get_line_width(realization, model):
    """Get the line width"""

    if model == 'FGOALS-g2':
        lw = 2.0
    else:
        lw = 2.0 if realization == 1 else 0.5

    return lw


def plot_individual():
    """ """

    if model not in legend_models:
        label = model
            legend_models.append(model)
        else:
            label = None

        lw = get_line_width(realization, model)
        iplt.plot(agg_cube, label=label, color=color_dict[model], linewidth=lw)

    title = '%s, %s-%s' %(plot_name, inargs.time[0][0:4], inargs.time[1][0:4])
    plt.title(title)
    plt.xlim(-90, 90)
    ylabel = get_ylabel(cube, inargs)
    plt.ylabel(ylabel)
    plt.xlabel('latitude')

    box = ax.get_position()
    ax.set_position([box.x0, box.y0, box.width * 0.8, box.height])
    ax.legend(loc='center left', bbox_to_anchor=(1, 0.5))


def group_runs(data_dict):
    """Create model and model/physics groups"""

    info = data_dict.keys()

    model_list = []
    for key, group in groupby(info, lambda x: x[0]):
        model_list.append(key)

    family_list = []
    for key, group in groupby(info, lambda x: x[0:2]):
        family_list.append(key)

    return family_list
    

def main(inargs):
    """Run the program."""
    
    time_constraint = gio.get_time_constraint(inargs.time)
    fig, ax = plt.subplots()
    plt.axhline(y=0, color='0.5', linestyle='--')
    color_dict = get_colors(inargs.infiles)
    
    data_dict = {}
    for infile in inargs.infiles:
        with iris.FUTURE.context(cell_datetime_objects=True):
            cube = iris.load_cube(infile, gio.check_iris_var(inargs.var) & time_constraint)
        
        if inargs.perlat:
            grid_spacing = grids.get_grid_spacing(cube) 
            cube.data = cube.data / grid_spacing
 
        if inargs.time_agg == 'trend':
            agg_cube = calc_trend_cube(cube)
            plot_name = 'linear trend'
        elif inargs.time_agg == 'climatology':
            agg_cube = cube.collapsed('time', iris.analysis.MEAN)
            plot_name = 'climatology'

        model = cube.attributes['model_id']
        realization = 'r' + str(cube.attributes['realization'])
        physics = 'p' + str(cube.attributes['physics_version'])

        data_dict[(model, physics, realization)] = agg_cube

    model_list, family_list = group_runs(data_dict)

    plot_individual(data_dict)

    plt.savefig(inargs.outfile, bbox_inches='tight')
    gio.write_metadata(inargs.outfile, file_info={infile: cube.attributes['history']})


if __name__ == '__main__':

    extra_info =""" 

author:
    <NAME>, <EMAIL>

"""

    description = 'Plot zonal ensemble'
    parser = argparse.ArgumentParser(description=description,
{'fraction_non_alphanumeric': 0.12444996050998533, 'fraction_numerical': 0.008687803226898341, 'mean_word_length': 2.799399914273468, 'pattern_counts': {'":': 0, '<': 11, '<?xml version=': 0, '>': 11, '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': '393634', 'n_tokens_mistral': 3131, 'n_tokens_neox': 2997, 'n_words': 696}