Fix and refactor change_pr_units (closes #107)

climix/util/__init__.py

@@ -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)
 

climix/util/change_pr_units.py

 # -*- 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