Commit fb8bde17 authored by Klaus Zimmermann's avatar Klaus Zimmermann
Browse files

Fix and refactor change_pr_units (closes #107)

parent ea637c43
......@@ -6,9 +6,7 @@ from .cube_diffs import cube_diffs
def change_units(cube, new_units, new_standard_name=None):
if is_precipitation(new_standard_name):
change_pr_units(cube,
new_unit=new_units,
new_stname=new_standard_name)
change_pr_units(cube, new_units, new_standard_name)
else:
cube.convert_units(new_units)
......
# -*- coding: utf-8 -*-
PRECIPITATION_STANDARD_NAMES = [
'lwe_precipitation_rate',
'lwe_thickness_of_precipitation_amount',
'precipitation_amount',
'precipitation_flux',
]
from collections import namedtuple, OrderedDict
from cf_units import Unit
import six
DENSITY_WATER = Unit('1000 kg m-3')
PrecipQty = namedtuple('PrecipQty', ['units', 'is_rate_based', 'is_mass_based'])
PRECIPITATION_INFO = OrderedDict([
('lwe_precipitation_rate', PrecipQty(Unit('m s-1'), True, False)),
('lwe_thickness_of_precipitation_amount', PrecipQty(Unit('m'), False, False)),
('thickness_of_rainfall_amount', PrecipQty(Unit('m'), False, False)),
('precipitation_amount', PrecipQty(Unit('kg m-2'), False, True)),
('precipitation_flux', PrecipQty(Unit('kg m-2 s-1'), True, True)),
])
def is_precipitation(standard_name):
return standard_name in PRECIPITATION_STANDARD_NAMES
return standard_name in PRECIPITATION_INFO.keys()
def get_standard_name_for_units(units):
for standard_name, qty in PRECIPITATION_INFO.items():
if units.is_convertible(qty.units):
return standard_name
raise ValueError(f'change_pr_units: Cannot handle new units {units}.')
def ensure_units_and_standard_name(units, standard_name):
if isinstance(units, six.string_types):
units = Unit(units)
if units is None and standard_name is None:
raise ValueError('No new units or new standard name is given.')
if standard_name is None:
standard_name = get_standard_name_for_units(units)
try:
units_dict = PRECIPITATION_INFO[standard_name]
except KeyError:
raise ValueError(f'Unknown standard_name {standard_name}.')
if units is None:
units = units_dict.units
if not units.is_convertible(units_dict.units):
raise ValueError(f'Cannot handle new units {units}. '
f'Not fitting standard_name, {standard_name}.')
return units, standard_name
def get_rate_conversion(old_qty, new_qty, integration_time):
if old_qty.is_rate_based and not new_qty.is_rate_based:
return integration_time
if not old_qty.is_rate_based and new_qty.is_rate_based:
return integration_time.invert()
return 1.
def change_pr_units(cube, new_unit=None, new_stname=None):
def get_mass_conversion(old_qty, new_qty, density):
if old_qty.is_mass_based and not new_qty.is_mass_based:
return density.invert()
if not old_qty.is_mass_based and new_qty.is_mass_based:
return density
return 1.
def change_pr_units(cube, new_units=None, new_standard_name=None,
integration_time=Unit('1 day'),
density=DENSITY_WATER):
"""Converts precipitation units considering standard names.
Converts precipitation units and cube.data to requested units.
......@@ -22,9 +79,9 @@ def change_pr_units(cube, new_unit=None, new_stname=None):
----------
cube: iris.cube.Cube
cube with precipitation data
new_unit: cf_units.Unit
unit to convert to
new_stname: str
new_units: cf_units.Unit
units to convert to
new_standard_name: str
standard name to convert to
Returns
......@@ -33,95 +90,21 @@ def change_pr_units(cube, new_unit=None, new_stname=None):
cube with new units and standard name
"""
from cf_units import Unit
# logging.debug('change_pr_units')
rho = Unit('1000 kg m-3')
T = Unit('1 s')
u = Unit('1')
stin_dict2 = {
'lwe_precipitation_rate': {
'unit': 'm s-1',
'lwe_precipitation_rate': u,
'lwe_thickness_of_precipitation_amount': u/T,
'precipitation_amount': u/T/rho,
'precipitation_flux': u/rho},
'lwe_thickness_of_precipitation_amount': {
'unit': 'm',
'lwe_precipitation_rate': T,
'lwe_thickness_of_precipitation_amount': u,
'precipitation_amount': u/rho,
'precipitation_flux': T/rho},
'precipitation_amount': {
'unit': 'kg m-2',
'lwe_precipitation_rate': rho*T,
'lwe_thickness_of_precipitation_amount': rho,
'precipitation_amount': u,
'precipitation_flux': T},
'precipitation_flux': {
'unit': 'kg m-2 s-1',
'lwe_precipitation_rate': rho,
'lwe_thickness_of_precipitation_amount': rho/T,
'precipitation_amount': u/T,
'precipitation_flux': u}
}
if new_unit is None and new_stname is None:
# logging.debug('change_pr_units: no new unit or new standard name '
# 'is given. Nothing happens.')
return
elif new_stname is None:
try:
new_unit2 = Unit(new_unit)
for stn in stin_dict2.keys():
unit_dict = stin_dict2[stn]
if new_unit2.is_convertible(Unit(unit_dict['unit'])):
new_stname = stn
break
else:
raise ValueError
except ValueError:
# logging.warn('change_pr_units: Cannot handle new '
# 'unit %s.', new_unit)
raise
else:
try:
unit_dict = stin_dict2[new_stname]
if new_unit is None:
new_unit = unit_dict['unit']
try:
new_unit2 = Unit(new_unit)
except ValueError:
# logging.warn('change_pr_units: Cannot handle new unit %s.',
# new_unit)
raise
if not new_unit2.is_convertible(Unit(unit_dict['unit'])):
# logging.error('change_pr_units: Cannot handle new unit %s. '
# 'Not fitting standard_name, %s',
# new_unit, new_stname)
raise ValueError
except KeyError:
# logging.error('change_pr_units: Illegal standard_name %s.',
# new_stname)
raise
old_stname = cube.standard_name
try:
unit_dict_old = stin_dict2[old_stname]
except KeyError:
# logging.error('change_pr_units: The cube must have a valid '
# 'precipitation standard name. %s.',
# old_stname)
raise
old_unit = str(cube.units)
try:
old_unit2 = Unit(old_unit)
old_canon = Unit(unit_dict_old['unit'])
if not old_unit2.is_convertible(old_canon):
raise ValueError
except ValueError:
# logging.error('The cube must have a valid precipitation unit %s.',
# old_unit)
raise
c = unit_dict_old[new_stname]
conv = (old_unit2/Unit(c)).convert(1., new_unit2)
cube.data = cube.core_data() * float(conv)
cube.units = new_unit2
cube.standard_name = new_stname
new_units, new_standard_name = ensure_units_and_standard_name(
new_units,
new_standard_name)
old_units, old_standard_name = ensure_units_and_standard_name(
cube.units,
cube.standard_name)
old_qty = PRECIPITATION_INFO[old_standard_name]
new_qty = PRECIPITATION_INFO[new_standard_name]
conv = (get_rate_conversion(old_qty, new_qty, integration_time)
* get_mass_conversion(old_qty, new_qty, density))
conv_factor = (old_units*conv).convert(1., new_units)
cube.data = cube.core_data() * float(conv_factor)
cube.units = new_units
cube.standard_name = new_standard_name
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