<SYSTEM_TASK:>
Given the following text description, write Python code to implement the functionality described below step by step
<END_TASK>
<USER_TASK:>
Description:
 Heat budget validation
I'm doing an analysis of the CMIP5 single forcing experiments (historicalGHG and historicalAA) to explore the separate influence of anthropogenic aerosols (AAs) and greenhouse gases (GHGs) on historical ocean change.
To make sure that my heat budget calculations are correct, I'm first trying to reproduce the heat uptake values in Table 2 of Frolicher et al (2015). In particular, the following is my attempt to reproduce the heat uptake south of 30S for the GISS-E2-R model.
In summary, the steps I take to calculcate the cumulative oceanic heat uptake south of 30S between 1870 (1861-80) and 1995 (1986-2005) are as follows
Step2:  Step 1
Step5:  Step 2
Step7:  Step 3
Step9:  Step 4
Step10:  Step 5
Step14:  Step 6
Step15:  Final result
<ASSISTANT_TASK:>
Python Code:
import re
import glob
import numpy
import iris
import iris.coord_categorisation
from iris.experimental.equalise_cubes import equalise_attributes
import warnings
warnings.filterwarnings('ignore')
Explanation: Heat budget validation
I'm doing an analysis of the CMIP5 single forcing experiments (historicalGHG and historicalAA) to explore the separate influence of anthropogenic aerosols (AAs) and greenhouse gases (GHGs) on historical ocean change.
To make sure that my heat budget calculations are correct, I'm first trying to reproduce the heat uptake values in Table 2 of Frolicher et al (2015). In particular, the following is my attempt to reproduce the heat uptake south of 30S for the GISS-E2-R model.
In summary, the steps I take to calculcate the cumulative oceanic heat uptake south of 30S between 1870 (1861-80) and 1995 (1986-2005) are as follows:
1. Load the monthly timescale hfds (surface downward heat flux) data for the region south of 30S
2. Convert the units from $W m^{-2}$ to $J m^{-2}$ by multiplying by the number of seconds in each timestep
3. Convert to an annual timescale by summing the 12 monthly values for each year
4. Convert the units from $J m^{-2}$ to $J$ by multiplying each value by the corresponding grid-cell area
5. Calculate the spatial sum (i.e. collapse the latitude and longitude dimensions)
6. Calculate the 20-year sum for the time periods of interest
7. Calculate the final result
My result differs from Frolicher et al (2015) by a factor of 3. I haven't regridded the data to a $1 \times 1$ grid like they did, but I doubt that explains such as discrepancy. I'm stumped as to what the issue could be.
End of explanation
lat_constraint = iris.Constraint(latitude=lambda cell: cell <= -30)
def read_hfds_data(file_list, lat_constraint):
    Read in data for a given latitude constraint
    
    cube = iris.load(file_list, 'surface_downward_heat_flux_in_sea_water' & lat_constraint)
    
    equalise_attributes(cube)
    iris.util.unify_time_units(cube)
    cube = cube.concatenate_cube()
    
    return cube
hfds_control_files = glob.glob('/g/data/ua6/DRSv2/CMIP5/GISS-E2-R/piControl/mon/ocean/r1i1p1/hfds/latest/hfds_Omon_GISS-E2-R_piControl_r1i1p1_*.nc')
hfds_control_cube = read_hfds_data(hfds_control_files, lat_constraint)
hfds_control_cube
hfds_historical_files = glob.glob('/g/data/ua6/DRSv2/CMIP5/GISS-E2-R/historical/mon/ocean/r1i1p1/hfds/latest/hfds_Omon_GISS-E2-R_historical_r1i1p1_*.nc')
hfds_historical_cube = read_hfds_data(hfds_historical_files, lat_constraint)
hfds_historical_cube
Explanation: Step 1: Read data
The code below simply loads all data south of 30S.
End of explanation
def broadcast_array(array, axis_index, shape):
    Broadcast an array to a target shape.
    
    Args:
      array (numpy.ndarray)
      axis_index (int or tuple): Postion in the target shape that the 
        axis/axes of the array corresponds to
          e.g. if array corresponds to (depth, lat, lon) in (time, depth, lat, lon)
          then axis_index = [1, 3]
          e.g. if array corresponds to (lat) in (time, depth, lat, lon)
          then axis_index = 2
      shape (tuple): shape to broadcast to
      
    For a one dimensional array, make start_axis_index = end_axis_index
    
    
    if type(axis_index) in [float, int]:
        start_axis_index = end_axis_index = axis_index
    else:
        assert len(axis_index) == 2
        start_axis_index, end_axis_index = axis_index
    
    dim = start_axis_index - 1
    while dim >= 0:
        array = array[numpy.newaxis, ...]
        array = numpy.repeat(array, shape[dim], axis=0)
        dim = dim - 1
    
    dim = end_axis_index + 1
    while dim < len(shape):    
        array = array[..., numpy.newaxis]
        array = numpy.repeat(array, shape[dim], axis=-1)
        dim = dim + 1
    return array
def convert_to_joules(cube):
    Convert units to Joules
    assert 'W' in str(cube.units)
    assert 'days' in str(cube.coord('time').units)
    
    time_span_days = cube.coord('time').bounds[:, 1] - cube.coord('time').bounds[:, 0]
    time_span_seconds = time_span_days * 60 * 60 * 24
    
    cube.data = cube.data * broadcast_array(time_span_seconds, 0, cube.shape)
    cube.units = str(cube.units).replace('W', 'J')
    
    return cube
hfds_control_cube = convert_to_joules(hfds_control_cube) 
hfds_control_cube
hfds_historical_cube = convert_to_joules(hfds_historical_cube) 
hfds_historical_cube
Explanation: Step 2: Convert from W m-2 to J m-2
End of explanation
def annual_sum(cube):
    Calculate the annual sum.
    iris.coord_categorisation.add_year(cube, 'time')
    cube = cube.aggregated_by(['year'], iris.analysis.SUM)
    cube.remove_coord('year')
    return cube
hfds_control_cube = annual_sum(hfds_control_cube) 
hfds_control_cube
hfds_historical_cube = annual_sum(hfds_historical_cube) 
hfds_historical_cube
Explanation: Step 3: Calculate the annual sum
End of explanation
def multiply_by_area(cube, area_cube):
    Multiply each cell of cube by its area. 
    area_data = broadcast_array(area_cube.data, [1, 2], cube.shape)
    cube.data = cube.data * area_data
    units = str(cube.units)
    cube.units = units.replace('m-2', '')
    
    return cube
area_file = '/g/data/ua6/DRSv2/CMIP5/GISS-E2-R/piControl/fx/ocean/r0i0p0/areacello/latest/areacello_fx_GISS-E2-R_piControl_r0i0p0.nc'
area_cube = iris.load_cube(area_file, 'cell_area' & lat_constraint)
area_cube
hfds_control_cube = multiply_by_area(hfds_control_cube, area_cube) 
hfds_control_cube
hfds_historical_cube = multiply_by_area(hfds_historical_cube, area_cube) 
hfds_historical_cube
Explanation: Step 4: Convert from J m-2 to J (i.e. multiply by area)
End of explanation
hfds_control_cube = hfds_control_cube.collapsed(['latitude', 'longitude'], iris.analysis.SUM)
hfds_control_cube
hfds_historical_cube = hfds_historical_cube.collapsed(['latitude', 'longitude'], iris.analysis.SUM)
hfds_historical_cube
Explanation: Step 5: Calculate the spatial sum
End of explanation
def get_time_constraint(time_list):
    Get the time constraint used for subsetting an iris cube.
    
    start_date, end_date = time_list
    date_pattern = '([0-9]{4})-([0-9]{1,2})-([0-9]{1,2})'
    assert re.search(date_pattern, start_date)
    assert re.search(date_pattern, end_date)
    start_year, start_month, start_day = start_date.split('-') 
    end_year, end_month, end_day = end_date.split('-')
    time_constraint = iris.Constraint(time=lambda t: iris.time.PartialDateTime(year=int(start_year), month=int(start_month), day=int(start_day)) <= t.point <= iris.time.PartialDateTime(year=int(end_year), month=int(end_month), day=int(end_day)))
    return time_constraint
def get_control_time_constraint(control_cube, hist_cube, time_bounds):
    Define the time constraints for the control data.
    iris.coord_categorisation.add_year(control_cube, 'time')
    iris.coord_categorisation.add_year(hist_cube, 'time')
    branch_time = hist_cube.attributes['branch_time']
    
    index = 0
    for bounds in control_cube.coord('time').bounds:
        lower, upper = bounds
        if lower <= branch_time < upper:
            break
        else:
            index = index + 1
    branch_year = control_cube.coord('year').points[index]
    hist_start_year = hist_cube.coord('year').points[0]
    start_gap = int(time_bounds[0].split('-')[0]) - hist_start_year
    end_gap = int(time_bounds[1].split('-')[0]) - hist_start_year
    control_start_year = branch_year + start_gap
    control_end_year = branch_year + end_gap
    control_start_date = str(control_start_year).zfill(4)+'-01-01'
    control_end_date = str(control_end_year).zfill(4)+'-01-01'
    time_constraint = get_time_constraint([control_start_date, control_end_date])
    control_cube.remove_coord('year')
    hist_cube.remove_coord('year')
    return time_constraint
def temporal_sum(cube, time_constraint):
    Calculate temporal sum over a given time period.
    cube = cube.copy()
    temporal_subset = cube.extract(time_constraint)
    result = temporal_subset.collapsed('time', iris.analysis.SUM)
    return float(result.data)
period_1870 = ['1861-01-01', '1880-12-31']
period_1995 = ['1986-01-01', '2005-12-31']
hist_1870_constraint = get_time_constraint(period_1870)
hist_1995_constraint = get_time_constraint(period_1995)
hist_1870 = temporal_sum(hfds_historical_cube, hist_1870_constraint)
hist_1995 = temporal_sum(hfds_historical_cube, hist_1995_constraint)
print('historial, 1870: ', hist_1870, 'J')
print('historial, 1995: ', hist_1995, 'J')
control_1870_constraint = get_control_time_constraint(hfds_control_cube, hfds_historical_cube, period_1870)
control_1995_constraint = get_control_time_constraint(hfds_control_cube, hfds_historical_cube, period_1995)
control_1870 = temporal_sum(hfds_control_cube, control_1870_constraint)
control_1995 = temporal_sum(hfds_control_cube, control_1995_constraint)
print('control, 1870: ', control_1870, 'J')
print('control, 1995: ', control_1995, 'J')
Explanation: Step 6: Calculate the 20-year sum for the time periods of interest
End of explanation
change = (hist_1995 - hist_1870) - (control_1995 - control_1870)
print('Cumulative oceanic heat uptake south of 30S between 1870 (1861-80) and 1995 (1986-2005):', change)
Explanation: Final result
End of explanation
<END_TASK>