# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# + [markdown] deletable=true editable=true
# # Ocean heat transport in CMIP5 models
#
#

# + [markdown] deletable=true editable=true
# ## Read data

# + deletable=true editable=true
import matplotlib.pyplot as plt
import iris
import iris.plot as iplt
import iris.coord_categorisation
import cf_units
import numpy

# + deletable=true editable=true
# %matplotlib inline 

# + deletable=true editable=true
infile = '/g/data/ua6/DRSv2/CMIP5/NorESM1-M/rcp85/mon/ocean/r1i1p1/hfbasin/latest/hfbasin_Omon_NorESM1-M_rcp85_r1i1p1_200601-210012.nc'

# + deletable=true editable=true
cube = iris.load_cube(infile)

# + deletable=true editable=true
print(cube)

# + deletable=true editable=true
dim_coord_names = [coord.name() for coord in cube.dim_coords]
print(dim_coord_names)

# + deletable=true editable=true
cube.coord('latitude').points

# + deletable=true editable=true
aux_coord_names = [coord.name() for coord in cube.aux_coords]
print(aux_coord_names)

# + deletable=true editable=true
cube.coord('region')

# + deletable=true editable=true
global_cube = cube.extract(iris.Constraint(region='global_ocean'))


# + deletable=true editable=true
def convert_to_annual(cube):
    """Convert data to annual timescale.
    Args:
      cube (iris.cube.Cube)
      full_months(bool): only include years with data for all 12 months
    """

    iris.coord_categorisation.add_year(cube, 'time')
    iris.coord_categorisation.add_month(cube, 'time')

    cube = cube.aggregated_by(['year'], iris.analysis.MEAN)
    cube.remove_coord('year')
    cube.remove_coord('month')

    return cube


# + deletable=true editable=true
global_cube_annual = convert_to_annual(global_cube)

# + deletable=true editable=true
print(global_cube_annual)

# + deletable=true editable=true
iplt.plot(global_cube_annual[5, ::])
iplt.plot(global_cube_annual[20, ::])
plt.show()


# + [markdown] deletable=true editable=true
# So for any given year, the annual mean shows ocean heat transport away from the tropics.

# + [markdown] deletable=true editable=true
# ## Trends

# + deletable=true editable=true
def convert_to_seconds(time_axis):
    """Convert time axis units to seconds.
    
    Args:
      time_axis(iris.DimCoord)
    
    """

    old_units = str(time_axis.units)
    old_timestep = old_units.split(' ')[0]
    new_units = old_units.replace(old_timestep, 'seconds') 

    new_unit = cf_units.Unit(new_units, calendar=time_axis.units.calendar)  
    time_axis.convert_units(new_unit)

    return time_axis


def linear_trend(data, time_axis):
    """Calculate the linear trend.
    
    polyfit returns [a, b] corresponding to y = a + bx

    """    

    masked_flag = False

    if type(data) == numpy.ma.core.MaskedArray:
        if type(data.mask) == numpy.bool_:
            if data.mask:
                masked_flag = True
        elif data.mask[0]:
            masked_flag = True
            
    if masked_flag:
        return data.fill_value
    else:
        return numpy.polynomial.polynomial.polyfit(time_axis, data, 1)[-1]


def calc_trend(cube):
    """Calculate linear trend.
    Args:
      cube (iris.cube.Cube)
      running_mean(bool, optional): 
        A 12-month running mean can first be applied to the data
      yr (bool, optional):
        Change units from per second to per year
    """

    time_axis = cube.coord('time')
    time_axis = convert_to_seconds(time_axis)

    trend = numpy.ma.apply_along_axis(linear_trend, 0, cube.data, time_axis.points)
    trend = numpy.ma.masked_values(trend, cube.data.fill_value)

    return trend


# + deletable=true editable=true
trend_data = calc_trend(global_cube_annual)

trend_cube = global_cube_annual[0, ::].copy()
trend_cube.data = trend_data
trend_cube.remove_coord('time')

#trend_unit = ' yr-1'
#trend_cube.units = str(global_cube_annual.units) + trend_unit

# + deletable=true editable=true
iplt.plot(trend_cube)
plt.show()

# + [markdown] deletable=true editable=true
# So the trends in ocean heat transport suggest reduced transport in the RCP 8.5 simulation (i.e. the trend plot is almost the inverse of the climatology plot).

# + [markdown] deletable=true editable=true
# ## Convergence

# + deletable=true editable=true
print(global_cube_annual)

# + deletable=true editable=true
diffs_data = numpy.diff(global_cube_annual.data, axis=1)
lats = global_cube_annual.coord('latitude').points
diffs_lats = (lats[1:] + lats[:-1]) / 2.

# + deletable=true editable=true
print(diffs_data.shape)
print(len(diffs_lats))

# + deletable=true editable=true
plt.plot(diffs_lats, diffs_data[0, :])
plt.plot(lats, global_cube_annual[0, ::].data / 10.0)
plt.show()

# + [markdown] deletable=true editable=true
# ## Convergence trend

# + deletable=true editable=true
time_axis = global_cube_annual.coord('time')
time_axis = convert_to_seconds(time_axis)

diffs_trend = numpy.ma.apply_along_axis(linear_trend, 0, diffs_data, time_axis.points)
diffs_trend = numpy.ma.masked_values(diffs_trend, global_cube_annual.data.fill_value)

# + deletable=true editable=true
print(diffs_trend.shape)

# + deletable=true editable=true
plt.plot(diffs_lats, diffs_trend * -1)
plt.axhline(y=0)
plt.show()

# + deletable=true editable=true
plt.plot(diffs_lats, diffs_trend * -1, color='black')
plt.axhline(y=0)
plt.axvline(x=30)
plt.axvline(x=50)
plt.axvline(x=77)
plt.xlim(20, 90)
plt.show()

# + [markdown] deletable=true editable=true
# When I try and replicate the HTC curve in Figure 1 of [Nummelin et al (2016)](http://onlinelibrary.wiley.com/doi/10.1002/2016GL071333/abstract;jsessionid=7BC4C1DF16F35341AE3D3689F363955B.f02t01) (above) it looks different because I've plotted $W s^{-1}$, whereas they plot $W m^{-2}$. So I need to divide by area and mutliply by their delta $\delta t/2$ (i.e. ($60 \times 60 \times 24 \times 365.25 \times 95) / 2$). 
#
# FIXME: Probably easiest to re-do this NorESM1-M analysis with hfy.
#
# Once I've done that to confirm that I can reproduce their results, I may actually want to stick with $W s^{-1}$. Dividing by area inflates the high latitude regions, whereas I'm more interested in where the big heat transports are taking place.

# + deletable=true editable=true


StarcoderdataJupyterScriptsDedupFiltered