Commit cad7fd5f authored by Javier Vegas-Regidor's avatar Javier Vegas-Regidor
Browse files

Simplify basin treatment and bump diagonals version

parent 0f875683
......@@ -58,7 +58,7 @@ class Psi(Diagnostic):
)
def __str__(self):
return "PSI Startdate: {0} Member: {1} Chunk: {2}".format(
return "PSI Startdate: {0} Member: {1} Chunk: {2} Basins:{3}".format(
self.startdate,
self.member,
self.chunk,
......
......@@ -89,6 +89,9 @@ class RegionMean(Diagnostic):
self.lat_name = "lat"
self.lon_name = "lon"
self._length = max(len(basin.name) for basin in self.basins)
Log.debug(f'Max basin length: {self._length}')
def __eq__(self, other):
if self._different_type(other):
return False
......@@ -329,22 +332,23 @@ class RegionMean(Diagnostic):
def _save_result_2d(self, var, result, data):
final_name = "{1}{0}".format(var, self.variable)
temp = TempFile.get()
handler_source = Utils.open_cdf(self.variable_file.local_file)
handler_temp = Utils.open_cdf(temp, "w")
Utils.copy_variable(handler_source, handler_temp, "time", True, True)
handler_temp.createDimension("region", len(result))
handler_temp.createDimension("region_length", 50)
handler_temp.createDimension("region_length", self._length)
var_region = handler_temp.createVariable(
"region", "S1", ("region", "region_length")
)
var = handler_temp.createVariable(
"{1}{0}".format(var, self.variable), float, ("time", "region",),
"{1}{0}".format(var, self.variable), float, ("region", "time"),
)
var.units = "{0}".format(data.units)
for i, basin in enumerate(result):
var_region[i, ...] = netCDF4.stringtoarr(str(basin), 50)
var[..., i] = result[basin]
var_region[i, ...] = netCDF4.stringtoarr(str(basin), self._length)
var[i, ...] = result[basin]
handler_temp.close()
self.declared[final_name].set_local_file(temp, diagnostic=self)
......@@ -355,7 +359,7 @@ class RegionMean(Diagnostic):
handler_temp = Utils.open_cdf(temp, "w")
Utils.copy_variable(handler_source, handler_temp, "time", True, True)
handler_temp.createDimension("region", len(result))
handler_temp.createDimension("region_length", 50)
handler_temp.createDimension("region_length", self._length)
handler_temp.createDimension("lev", data.shape[1])
var_level = handler_temp.createVariable("lev", float, "lev")
var_level[...] = data.coord("depth").points
......@@ -376,7 +380,7 @@ class RegionMean(Diagnostic):
var.units = "{0}".format(data.units)
for i, basin in enumerate(result):
var_region[i, ...] = netCDF4.stringtoarr(str(basin), 50)
var_region[i, ...] = netCDF4.stringtoarr(str(basin), self._length)
var[..., i] = result[basin]
handler_temp.close()
self.declared[final_name].set_local_file(temp, diagnostic=self)
# coding=utf-8
"""Computes the seaice extent, area and volume"""
import os
import six
import iris
......@@ -83,7 +82,7 @@ class Sivolume(Diagnostic):
:return:
"""
options_available = (
DiagnosticBasinListOption("basins", [Basins().Global]),
DiagnosticBasinListOption("basins", Basins().Global.name),
)
options = cls.process_options(options, options_available)
......@@ -200,48 +199,9 @@ class Sivolume(Diagnostic):
def _save_file(self, data, var):
generated_file = self.generated[var]
temp = TempFile.get()
region = data.coord("region").points
data.remove_coord("region")
data.remove_coord('latitude')
data.remove_coord('longitude')
data.remove_coord(data.coord(var_name='i'))
data.remove_coord(data.coord(var_name='j'))
iris.save(data, temp, zlib=True)
if len(region) > 1:
Utils.rename_variable(temp, "dim0", "region", False)
handler = Utils.open_cdf(temp)
var = handler.createVariable("region2", str, ("region",))
var[...] = region
handler.close()
Utils.rename_variable(temp, "region2", "region", True)
else:
handler = Utils.open_cdf(temp)
if "region" not in handler.dimensions:
new_file = TempFile.get()
new_handler = Utils.open_cdf(new_file, "w")
new_handler.createDimension("region", 1)
for dimension in handler.dimensions:
Utils.copy_dimension(handler, new_handler, dimension)
for variable in handler.variables.keys():
if variable in (var, "region"):
continue
Utils.copy_variable(handler, new_handler, variable)
old_var = handler.variables[var]
new_var = new_handler.createVariable(
var,
old_var.dtype,
("region",) + old_var.dimensions,
zlib=True,
fill_value=1.0e20,
)
Utils.copy_attributes(new_var, old_var)
new_var[0, :] = old_var[:]
new_var = new_handler.createVariable(
"region", str, ("region",)
)
new_var[0] = region[0]
new_handler.close()
os.remove(temp)
temp = new_file
handler.close()
generated_file.set_local_file(temp)
......@@ -24,7 +24,7 @@ REQUIREMENTS = {
"cdo>=1.3.4",
"cfgrib",
"dask[array]",
"diagonals>=0.3.1",
"diagonals>=0.3.2",
"netCDF4",
"nco>=0.0.3",
"numba",
......
......@@ -40,5 +40,5 @@ class TestPsi(TestCase):
psi = Psi(self.data_manager, "20000101", 1, 1, {Basins().Global: None})
self.assertEqual(
str(psi),
"PSI Startdate: 20000101 Member: 1 Chunk: 1",
"PSI Startdate: 20000101 Member: 1 Chunk: 1 Basins:Global",
)
......@@ -56,7 +56,7 @@ class TestRegionMean(TestCase):
"var",
box,
True,
Basins().Global,
[Basins().Global],
"t",
Frequencies.monthly,
),
......@@ -72,7 +72,7 @@ class TestRegionMean(TestCase):
"var",
box,
True,
Basins().Global,
[Basins().Global],
"t",
Frequencies.monthly,
),
......@@ -93,7 +93,7 @@ class TestRegionMean(TestCase):
"var",
box,
True,
Basins().Global,
[Basins().Global],
"u",
Frequencies.monthly,
),
......@@ -109,7 +109,7 @@ class TestRegionMean(TestCase):
"var",
box,
True,
Basins().Global,
[Basins().Global],
"u",
Frequencies.monthly,
),
......@@ -130,7 +130,7 @@ class TestRegionMean(TestCase):
"var",
box,
True,
Basins().Global,
[Basins().Global],
"u",
Frequencies.monthly,
),
......@@ -146,7 +146,7 @@ class TestRegionMean(TestCase):
"var",
box,
True,
Basins().Global,
[Basins().Global],
"u",
Frequencies.monthly,
),
......@@ -172,7 +172,7 @@ class TestRegionMean(TestCase):
"var",
box,
True,
Basins().Global,
[Basins().Global],
"u",
Frequencies.monthly,
),
......@@ -188,7 +188,7 @@ class TestRegionMean(TestCase):
"var",
box,
True,
Basins().Global,
[Basins().Global],
"u",
Frequencies.monthly,
),
......@@ -219,7 +219,7 @@ class TestRegionMean(TestCase):
"var",
box,
False,
Basins().Global,
[Basins().Global],
"u",
Frequencies.monthly,
),
......@@ -235,7 +235,7 @@ class TestRegionMean(TestCase):
"var",
box,
False,
Basins().Global,
[Basins().Global],
"u",
Frequencies.monthly,
),
......@@ -267,7 +267,7 @@ class TestRegionMean(TestCase):
"var",
box,
False,
Basins().Global,
[Basins().Global],
"grid",
Frequencies.monthly,
),
......@@ -283,7 +283,7 @@ class TestRegionMean(TestCase):
"var",
box,
False,
Basins().Global,
[Basins().Global],
"grid",
Frequencies.monthly,
),
......@@ -316,7 +316,7 @@ class TestRegionMean(TestCase):
"var",
box,
False,
Basins().Global,
[Basins().Global],
"grid",
Frequencies.daily,
),
......@@ -332,7 +332,7 @@ class TestRegionMean(TestCase):
"var",
box,
False,
Basins().Global,
[Basins().Global],
"grid",
Frequencies.daily,
),
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment