<commit_before>"""
Filename:     plot_interhemispheric_heat_difference_timeseries.py
Author:       <NAME>, <EMAIL>
Description:  Plot ensemble interhemispheric difference timeseries

"""

# Import general Python modules

import sys, os, pdb, re
import argparse
import numpy
import iris
from iris.experimental.equalise_cubes import equalise_attributes
import iris.plot as iplt
import matplotlib.pyplot as plt
from matplotlib import gridspec
import matplotlib as mpl
import seaborn

import cmdline_provenance as cmdprov

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


# Define functions

var_colors = {'ohc': 'blue', 'hfds': 'orange', 'rndt': 'red'}
exp_colors = {'historical-rcp85': 'black', 'historical': 'black', 'GHG-only': 'red',
              'AA-only': 'blue', '1pctCO2': 'orange'}

names = {'thetao': 'Sea Water Potential Temperature',
         'ohc': 'ocean heat content',
         'hfds': 'Downward Heat Flux at Sea Water Surface',
         'rndt': 'TOA Incoming Net Radiation'}

titles = {'rndt': 'netTOA', 'hfds': 'OHU', 'ohc': 'OHC', 'thetao': 'Average Ocean Temperature'}

plot_variables = {'thetao': ' Average Ocean Temperature',
                  'ohc': 'OHC',
                  'hfds': 'OHU',
                  'rndt': 'netTOA'}

scale_factors = {'CanESM2': -54567.5, 'CCSM4': -729817.5}

linestyles = {'historical-rcp85': 'solid', 'historical': 'solid', 'GHG-only': '--', 'AA-only': ':'}

grid_configs = {1: (1, 1), 2: (1, 2), 3: (1, 3), 4: (2, 2)} 

seaborn.set(style='ticks')

mpl.rcParams['axes.labelsize'] = 24
mpl.rcParams['axes.titlesize'] = 28
mpl.rcParams['xtick.labelsize'] = 24
mpl.rcParams['ytick.labelsize'] = 24
mpl.rcParams['legend.fontsize'] = 20


def equalise_time_axes(cube_list):
    """Make all the time axes the same."""

    iris.util.unify_time_units(cube_list)
    reference_cube = cube_list[0]
    new_cube_list = iris.cube.CubeList([])
    for cube in cube_list:
        assert len(cube.coord('time').points) == len(reference_cube.coord('time').points)
        cube.coord('time').points = reference_cube.coord('time').points
        cube.coord('time').bounds = reference_cube.coord('time').bounds
        cube.coord('time').units = reference_cube.coord('time').units
        cube.coord('time').attributes = reference_cube.coord('time').attributes
        new_cube_list.append(cube)
    
    return new_cube_list


def ensemble_aggregation(cube_list, operator):
    """Calculate the ensemble mean."""

    aggregators = {'mean': iris.analysis.MEAN, 'median': iris.analysis.MEDIAN}

    if len(cube_list) > 1:
        equalise_attributes(cube_list)
        equalise_time_axes(cube_list)
        ensemble_cube = cube_list.merge_cube()
        ensemble_agg = ensemble_cube.collapsed('ensemble_member', aggregators[operator])
        ensemble_spread = ensemble_cube.collapsed('ensemble_member', iris.analysis.PERCENTILE, percent=[25, 75])
    else:
        ensemble_agg = cube_list[0]
        ensemble_spread = None

    return ensemble_agg, ensemble_spread


def calc_anomaly(cube):
    """Calculate the anomaly."""
    
    anomaly = cube.copy()
    anomaly.data = anomaly.data - anomaly.data[0]
    
    return anomaly


def get_simulation_attributes(cube):
    """Get model, experiment and mip information."""

    model = cube.attributes['model_id']
    experiment = cube.attributes['experiment_id']
    physics = cube.attributes['physics_version']
    run = cube.attributes['realization']
    mip = 'r%si1p%s' %(run, physics)

    if experiment == 'historicalMisc':
        experiment = 'historicalAA'

    return model, experiment, mip


def calc_hemispheric_value(sh_file, nh_file, val_type, var, time_constraint, ensemble_number):
    """Calculate the interhemispheric difference timeseries."""

    agg = 'mean' if var =='thetao' else 'sum'

    nh_name = names[var] + ' nh ' + agg
    nh_cube = iris.load_cube(nh_file, nh_name & time_constraint)
    nh_attributes = get_simulation_attributes(nh_cube)
    nh_anomaly = calc_anomaly(nh_cube)
    nh_anomaly.data = nh_anomaly.data.astype(numpy.float32)

    sh_name = names[var] + ' sh ' + agg
    sh_cube = iris.load_cube(sh_file, sh_name & time_constraint)
    sh_attributes = get_simulation_attributes(sh_cube)
    sh_anomaly = calc_anomaly(sh_cube)
    sh_anomaly.data = sh_anomaly.data.astype(numpy.float32)

    assert nh_attributes == sh_attributes 

    metric = nh_cube.copy()
    if var == 'ohc-adjusted':
        globe_data = nh_cube.data + sh_cube.data
        globe_anomaly = sh_anomaly.data + nh_anomaly.data<commit_msg>Quick fix for var names<commit_after>"""
Filename:     plot_interhemispheric_heat_difference_timeseries.py
Author:       <NAME>, <EMAIL>
Description:  Plot ensemble interhemispheric difference timeseries

"""

# Import general Python modules

import sys, os, pdb, re
import argparse
import numpy
import iris
from iris.experimental.equalise_cubes import equalise_attributes
import iris.plot as iplt
import matplotlib.pyplot as plt
from matplotlib import gridspec
import matplotlib as mpl
import seaborn

import cmdline_provenance as cmdprov

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


# Define functions

var_colors = {'ohc': 'blue', 'hfds': 'orange', 'rndt': 'red'}
exp_colors = {'historical-rcp85': 'black', 'historical': 'black', 'GHG-only': 'red',
              'AA-only': 'blue', '1pctCO2': 'orange'}

names = {'thetao': 'Sea Water Potential Temperature',
         'ohc': 'ocean heat content',
         'hfds': 'Downward Heat Flux at Sea Water Surface',
         'rndt': 'TOA Incoming Net Radiation'}

titles = {'rndt': 'netTOA', 'hfds': 'OHU', 'ohc': 'OHC', 'thetao': 'Average Ocean Temperature'}

plot_variables = {'thetao': ' Average Ocean Temperature',
                  'ohc': 'OHC',
                  'hfds': 'OHU',
                  'rndt': 'netTOA'}

scale_factors = {'CanESM2': -54567.5, 'CCSM4': -729817.5}

linestyles = {'historical-rcp85': 'solid', 'historical': 'solid', 'GHG-only': '--', 'AA-only': ':'}

grid_configs = {1: (1, 1), 2: (1, 2), 3: (1, 3), 4: (2, 2)} 

seaborn.set(style='ticks')

mpl.rcParams['axes.labelsize'] = 24
mpl.rcParams['axes.titlesize'] = 28
mpl.rcParams['xtick.labelsize'] = 24
mpl.rcParams['ytick.labelsize'] = 24
mpl.rcParams['legend.fontsize'] = 20


def equalise_time_axes(cube_list):
    """Make all the time axes the same."""

    iris.util.unify_time_units(cube_list)
    reference_cube = cube_list[0]
    new_cube_list = iris.cube.CubeList([])
    for cube in cube_list:
        assert len(cube.coord('time').points) == len(reference_cube.coord('time').points)
        cube.coord('time').points = reference_cube.coord('time').points
        cube.coord('time').bounds = reference_cube.coord('time').bounds
        cube.coord('time').units = reference_cube.coord('time').units
        cube.coord('time').attributes = reference_cube.coord('time').attributes
        new_cube_list.append(cube)
    
    return new_cube_list


def ensemble_aggregation(cube_list, operator):
    """Calculate the ensemble mean."""

    aggregators = {'mean': iris.analysis.MEAN, 'median': iris.analysis.MEDIAN}

    if len(cube_list) > 1:
        equalise_attributes(cube_list)
        equalise_time_axes(cube_list)
        ensemble_cube = cube_list.merge_cube()
        ensemble_agg = ensemble_cube.collapsed('ensemble_member', aggregators[operator])
        ensemble_spread = ensemble_cube.collapsed('ensemble_member', iris.analysis.PERCENTILE, percent=[25, 75])
    else:
        ensemble_agg = cube_list[0]
        ensemble_spread = None

    return ensemble_agg, ensemble_spread


def calc_anomaly(cube):
    """Calculate the anomaly."""
    
    anomaly = cube.copy()
    anomaly.data = anomaly.data - anomaly.data[0]
    
    return anomaly


def get_simulation_attributes(cube):
    """Get model, experiment and mip information."""

    model = cube.attributes['model_id']
    experiment = cube.attributes['experiment_id']
    physics = cube.attributes['physics_version']
    run = cube.attributes['realization']
    mip = 'r%si1p%s' %(run, physics)

    if experiment == 'historicalMisc':
        experiment = 'historicalAA'

    return model, experiment, mip


def calc_hemispheric_value(sh_file, nh_file, val_type, var, time_constraint, ensemble_number):
    """Calculate the interhemispheric difference timeseries."""

    agg = 'mean' if var =='thetao' else 'sum'

    nh_name = names[var] + ' nh ' + agg
    nh_cube = iris.load_cube(nh_file, nh_name & time_constraint)
    nh_cube.var_name = nh_cube.var_name.replace('_', '-')
    nh_attributes = get_simulation_attributes(nh_cube)
    nh_anomaly = calc_anomaly(nh_cube)
    nh_anomaly.data = nh_anomaly.data.astype(numpy.float32)

    sh_name = names[var] + ' sh ' + agg
    sh_cube = iris.load_cube(sh_file, sh_name & time_constraint)
    sh_cube.var_name = sh_cube.var_name.replace('_', '-')
    sh_attributes = get_simulation_attributes(sh_cube)
    sh_anomaly = calc_anomaly(sh_cube)
    sh_anomaly.data = sh_anomaly.data.astype(numpy.float32)

    assert nh_attributes == sh_attributes 

    metric = nh_cube.copy()
    if var == 'ohc-adjusted':
        globe_data = nh_cube.data + sh_cube.data
        globe_anomaly = sh_anomaly.data + nh_anomaly.data
{'fraction_non_alphanumeric': 0.13411196133709222, 'fraction_numerical': 0.011478050745066451, 'mean_word_length': 3.843003412969283, 'pattern_counts': {'":': 0, '<': 7, '<?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': '1092329', 'n_tokens_mistral': 3653, 'n_tokens_neox': 3374, 'n_words': 733}