It is possible to create 1d_profile plots and zonalmean plots

Author: FranziskaWinterstein

esmvaltool/diag_scripts/lifetime/lifetime.py

@@ -343,8 +343,7 @@ def __init__(self, config):
             'timeseries',
             'annual_cycle',
             'zonalmean',
-            '1d_profile'
-        ]
+            '1d_profile' ]
         for (plot_type, plot_options) in self.plots.items():
             if plot_type not in self.supported_plot_types:
                 raise ValueError(
@@ -353,7 +352,7 @@ def __init__(self, config):
             if plot_options is None:
                 self.plots[plot_type] = {}
 
-            # Default options for the different plot types
+            # Default options for the differ N MMNMJHJ JMHTHent plot types
             if plot_type == 'timeseries':
                 self.plots[plot_type].setdefault('annual_mean', False)
                 self.plots[plot_type].setdefault('annual_mean_kwargs', {})
@@ -371,7 +370,9 @@ def __init__(self, config):
 
             if plot_type == 'zonalmean':
                 self.plots[plot_type].setdefault(
-                    'cbar_label', '{short_name} [{units}]')
+                    'cbar_label',
+                    f"{chr(964)}({self._get_name('reactant').upper()})"
+                    ' [{units}]')
                 self.plots[plot_type].setdefault(
                     'cbar_label_bias', 'Δ{short_name} [{units}]')
                 self.plots[plot_type].setdefault(
@@ -569,20 +570,23 @@ def _calculate_coefficients(self):
                 z_coord = reaction.coord('air_pressure', dim_coords=True)
                 z_coord.attributes['positive'] = 'down'
                 z_coord.convert_units('Pa')
+                use_z_coord = 'air_pressure'
             elif reaction.coords('atmosphere_hybrid_sigma_pressure_coordinate',
                                  dim_coords=True):
                 self.z_coord = 'atmosphere_hybrid_sigma_pressure_coordinate'
                 z_coord = reaction.coord(
                     'atmosphere_hybrid_sigma_pressure_coordinate',
                     dim_coords=True)
                 z_coord.attributes['positive'] = 'down'
+                use_z_coord = 'model_level_number'
             else:
                 raise NotImplementedError(
                     "Lifetime calculation is not implemented"
                     " for the present type of vertical coordinate."
                 )
 
-            if not set(['TROP', 'STRA']).isdisjoint(self.cfg['regions']):
+            if (not set(['TROP', 'STRA']).isdisjoint(self.cfg['regions']) and
+                not set(['timeseries', 'annual_cycle']).isdisjoint(self.plots)):
 
                 # calculate climatological tropopause pressure (tp_clim)
                 # but only if no tropopause is given by data
@@ -600,16 +604,19 @@ def _calculate_coefficients(self):
                 #          - tp_i and model_level_number
                 #          - ptp and (derived) air_pressure
                 #          - tp_clim and (derived) air_pressure
-                use_z_coord = 'air_pressure'
                 if z_coord.name() == 'air_pressure':
                     tropopause = variables.get('ptp', tropopause)
                 elif (z_coord.name()
                       == 'atmosphere_hybrid_sigma_pressure_coordinate'):
                     if 'tp_i' in variables:
                         tropopause = variables['tp_i']
-                        use_z_coord = 'model_level_number'
                     else:
                         tropopause = variables.get('ptp', tropopause)
+                        # fall back to air_pressure
+                        use_z_coord = 'air_pressure'
+
+                input_data_dataset[dataset]['tropopause'] = tropopause
+
 
             weight = self._define_weight(variables)
 
@@ -620,7 +627,6 @@ def _calculate_coefficients(self):
 
             input_data_dataset[dataset]['z_coord'] = z_coord
             input_data_dataset[dataset]['use_z_coord'] = use_z_coord
-            input_data_dataset[dataset]['tropopause'] = tropopause
             input_data_dataset[dataset]['variables'] = variables
             input_data_dataset[dataset]['reaction'] = reaction
             input_data_dataset[dataset]['weight'] = weight
@@ -1030,8 +1036,8 @@ def create_timeseries_plot(self, region, input_data, base_datasets):
                     slice_dataset['tropopause'] = tp_slice
 
                     cube_slice = calculate_lifetime(slice_dataset,
-                                                          plot_type,
-                                                          region)
+                                                    plot_type,
+                                                    region)
                     # make it real to reduce memory demand later
                     cube_slice.data
                     cube_slices.append(cube_slice)
@@ -1164,7 +1170,7 @@ def create_annual_cycle_plot(self, region, input_data, base_datasets):
             list(base_datasets.values()))
         axes.set_title(f'{self.info["long_name"]} in region {region}')
         axes.set_xlabel('Month')
-        axes.set_ylabel(f"$\tau$({self._get_name('reactant').upper()})"
+        axes.set_ylabel(f"{chr(964)}({self._get_name('reactant').upper()})"
                         " [{self.info['units']}]")
         axes.set_xticks(range(1, 13), [str(m) for m in range(1, 13)])
         gridline_kwargs = self._get_gridline_kwargs(plot_type)
@@ -1329,8 +1335,8 @@ def create_1d_profile_plot(self, region, input_data, base_datasets):
         # Default plot appearance
         multi_dataset_facets = self._get_multi_dataset_facets(
             list(base_datasets.values()))
-        axes.set_title(f'{self.info["long_name"]} in region {region}')
-        axes.set_xlabel(f"$\tau$({self._get_name('reactant').upper()})"
+        axes.set_title(f'{self.info["long_name"]}')
+        axes.set_xlabel(f"{chr(964)}({self._get_name('reactant').upper()})"
                         f" [{self.info['units']}]")
         z_coord = cube.coord(axis='Z')
         axes.set_ylabel(f'{z_coord.long_name} [{z_coord.units}]')

esmvaltool/diag_scripts/lifetime/lifetime_func.py

@@ -46,11 +46,14 @@ def calculate_gridmassdry(press, hus, z_coord):
         pmin = da.nanmin(press)
 
     # surface air pressure as cube
-    pmax = press[:, 0, :, :].copy()
+    surface = press.coord(z_coord)[0].points
+    pmax = press.extract(iris.Constraint(coord_values={z: surface for z in [z_coord]})).copy()
     if 'surface_air_pressure' in press.coords():
         pmax.data = press.coord('surface_air_pressure').lazy_points()
     else:
-        pmax.data[:, :, :] = 101525.
+        pmax.data = da.full_like(
+            press.extract(iris.Constraint(coord_values={z: surface for z in [z_coord]})),
+            101525.)
 
     # grid area -> m2
     lat_lon_dims = sorted(
@@ -66,10 +69,10 @@ def calculate_gridmassdry(press, hus, z_coord):
     )
 
     # delta pressure levels
-    delta_p = dpres_plevel_4d(press,
-                              pmin,
-                              pmax,
-                              z_coord=z_coord)
+    delta_p = dpres_plevel(press,
+                           pmin,
+                           pmax,
+                           z_coord=z_coord)
 
     # gridmass of dry air
     gridmassdry = deepcopy(press)
@@ -188,6 +191,58 @@ def _number_density_dryair_by_grid(grmassdry, grvol):
 
     return rho
 
+def dpres_plevel(plev, pmin, pmax, z_coord='air_pressure'):
+    """
+    Call the appropriate pressure level calculation with
+    respect to the given coordinates.
+    """
+
+    cube_coords = [coord.name() for coord in plev.dim_coords]
+
+    if [z_coord, 'latitude', 'longitude'] == cube_coords:
+        dplev = dpres_plevel_3d(plev, pmin, pmax, z_coord)
+    elif ['time', z_coord, 'latitude', 'longitude'] == cube_coords:
+        dplev = dpres_plevel_4d(plev, pmin, pmax, z_coord)
+    else:
+        raise NotImplementedError(
+            "Pressurelevel calculation is not implemented"
+            " for the present coordinates.")
+
+    return dplev
+
+
+def dpres_plevel_3d(plev, pmin, pmax, z_coord='air_pressure'):
+    """Calculate delta pressure levels.
+
+    The delta pressure levels are based
+    on the given pressure as a
+    four dimensional cube.
+    """
+    axis = plev.coord_dims(z_coord)
+
+    # shifted vectors:
+    # - p_m1: shifted by one index down
+    #         (the value of the last index is the former first)
+    # - p_p1: shifted by one index up
+    #         (the value of the first index is the former last)
+    # both vector values are divided by two
+    p_p1 = da.roll(plev.lazy_data(), 1, axis=axis) / 2.
+    p_m1 = da.roll(plev.lazy_data(), -1, axis=axis) / 2.
+
+    # modify the first entry in p_p1
+    # note: p_p1[1, :, :] .eq. plev[0, :, :] / 2.
+    p_p1[0, :, :] = pmax.lazy_data() - p_p1[1, :, :]
+
+    # and the last entry in p_m1
+    # note: p_m1[-2, :, :] .eq. plev[-1, :, :] / 2.
+    p_m1[-1, :, :] = pmin - p_m1[-2, :, :]
+
+    # calculate difference
+    dplev = deepcopy(plev)
+    dplev.data = p_p1 - p_m1
+
+    return dplev
+
 
 def dpres_plevel_4d(plev, pmin, pmax, z_coord='air_pressure'):
     """Calculate delta pressure levels.
@@ -225,8 +280,12 @@ def dpres_plevel_4d(plev, pmin, pmax, z_coord='air_pressure'):
 def calculate_lifetime(dataset, plot_type, region):
     """Calculate the lifetime for the given plot_type and region."""
     # extract region from weights and reaction
-    reaction = extract_region(dataset, region, case='reaction')
-    weight = extract_region(dataset, region, case='weight')
+    if plot_type in ['timeseries', 'annual_cycle']:
+        reaction = extract_region(dataset, region, case='reaction')
+        weight = extract_region(dataset, region, case='weight')
+    else:
+        reaction = dataset['reaction']
+        weight = dataset['weight']
 
     # calculate nominator and denominator
     # and sum of nominator and denominator via plot_type dimensions