Newer
Older
Javier Vegas-Regidor
committed
import numpy as np
from earthdiagnostics.box import Box
from earthdiagnostics.diagnostic import Diagnostic
from earthdiagnostics.utils import Utils, TempFile
Javier Vegas-Regidor
committed
class HeatContentLayer(Diagnostic):
"""
Point-wise Ocean Heat Content in a specified ocean thickness (J/m-2)
:original author: Isabel Andreu Burillo
:contributor: Virginie Guemas <virginie.guemas@bsc.es>
:contributor: Javier Vegas-Regidor<javier.vegas@bsc.es>
:created: June 2012
:last modified: June 2016
:param data_manager: data management object
:type data_manager: DataManager
:param startdate: startdate
:type startdate: str
:param member: member number
:type member: int
:param chunk: chunk's number
:type chunk: int
:param box: box to use for the calculations
:type box: Box
Javier Vegas-Regidor
committed
"""
alias = 'ohclayer'
"Diagnostic alias for the configuration file"
Javier Vegas-Regidor
committed
def __init__(self, data_manager, startdate, member, chunk, box, weight, min_level, max_level):
Javier Vegas-Regidor
committed
Diagnostic.__init__(self, data_manager)
self.startdate = startdate
self.member = member
self.chunk = chunk
Javier Vegas-Regidor
committed
self.box = box
Javier Vegas-Regidor
committed
self.min_level = min_level
self.max_level = max_level
Javier Vegas-Regidor
committed
self.required_vars = ['so', 'mlotst']
self.generated_vars = ['scvertsum']
def __str__(self):
return 'Heat content layer Startdate: {0} Member: {1} Chunk: {2}'.format(self.startdate, self.member,
self.chunk)
@classmethod
def generate_jobs(cls, diags, options):
"""
Creates a job for each chunk to compute the diagnostic
:param diags: Diagnostics manager class
:type diags: Diags
:param options: minimum depth, maximum depth, basin=Global
Javier Vegas-Regidor
committed
:type options: list[str]
"""
num_options = len(options) - 1
if num_options < 2:
raise Exception('You must specify the minimum and maximum depth to use')
if num_options > 3:
raise Exception('You must specify between 2 and 3 parameters for the heat content layer diagnostic')
Javier Vegas-Regidor
committed
box = Box(True)
box.min_depth = int(options[1])
box.max_depth = int(options[2])
if len(options) > 3:
basin = Basins.parse(options[3])
else:
basin = Basins.Global
Javier Vegas-Regidor
committed
job_list = list()
handler = Utils.openCdf('mesh_zgr.nc')
mask = Utils.get_mask(basin)
Javier Vegas-Regidor
committed
if 'e3t' in handler.variables:
mask = handler.variables['e3t'][:] * mask
Javier Vegas-Regidor
committed
elif 'e3t_0' in handler.variables:
mask = handler.variables['e3t_0'][:] * mask
Javier Vegas-Regidor
committed
else:
raise Exception('e3t variable can not be found')
if 'gdepw' in handler.variables:
depth = handler.variables['gdepw'][:]
elif 'gdepw_0' in handler.variables:
depth = handler.variables['gdepw_0'][:]
raise Exception('gdepw variable can not be found')
while len(depth.shape) < 4:
depth = np.expand_dims(depth, -1)
def calculate_weight(array):
"""
Calculates the weight for each level for the given later
:param array:
:return:
"""
while array[level+1] <= box.min_depth:
array[level] = 0
if level == array.size - 1:
array[level] = 0
return array
if array[level] != box.min_depth:
weight_value = (array[level + 1] - box.min_depth) / (array[level + 1] - array[level])
array[level] = weight_value
if level == array.size - 1:
array[level] = 0
return array
while array[level + 1] <= box.max_depth:
array[level] = 1
if level == array.size - 1:
array[level] = 0
return array
if array[level] != box.max_depth:
weight_value = (box.max_depth - array[level]) / (array[level + 1] - array[level])
array[level] = weight_value
if level == array.size - 1:
array[level] = 0
return array
weight = mask * np.apply_along_axis(calculate_weight, 1, depth) * 1020 * 4000
# Now we will reduce to the levels with any weight != 0 to avoid loading too much data on memory
Javier Vegas-Regidor
committed
levels = weight.shape[1]
min_level = 0
while min_level < levels and not weight[:, min_level, :].any():
min_level += 1
max_level = min_level
while max_level < (levels - 1) and weight[:, max_level + 1, :].any():
max_level += 1
weight = weight[:, min_level:max_level, :]
for startdate, member, chunk in diags.exp_manager.get_chunk_list():
Javier Vegas-Regidor
committed
job_list.append(HeatContentLayer(diags.data_manager, startdate, member, chunk, box,
weight, min_level, max_level))
def compute(self):
"""
Runs the diagnostic
"""
nco = Utils.nco
results = TempFile.get()
Javier Vegas-Regidor
committed
thetao_file = self.data_manager.get_file('ocean', 'thetao', self.startdate, self.member, self.chunk)
handler = Utils.openCdf(thetao_file)
Javier Vegas-Regidor
committed
heatc_sl = np.sum(handler.variables['thetao'][:, self.min_level:self.max_level, :] * self.weight, 1)
handler.sync()
handler.renameVariable('thetao', 'heatc_sl')
Javier Vegas-Regidor
committed
handler.close()
nco.ncks(input=thetao_file, output=results, options='-O -v lon,lat,time')
Utils.rename_variables(results, {'x': 'i', 'y': 'j'}, False, True)
handler_results = Utils.openCdf(results)
handler_results.createVariable('ohc', float, ('time', 'j', 'i'), fill_value=1.e20)
handler_results.sync()
Javier Vegas-Regidor
committed
handler_results.variables['ohc'][:] = heatc_sl
handler_results.close()
Utils.setminmax(results, 'ohc')
self.data_manager.send_file(results, 'ocean', 'ohc', self.startdate, self.member, self.chunk, box=self.box)