Bugfix for #14

Pass the aspect ratio from the model params object to the diagnostic
each time a new such object is provided to the diagnostic.

qgs/basis/base.py

@@ -94,9 +94,11 @@ def subs_functions(self, extra_subs=None):
         sf = list()
 
         for f in self.functions:
-            ff = f.subs(self.substitutions)
             if extra_subs is not None:
-                ff = ff.subs(extra_subs)
+                ff = f.subs(extra_subs)
+            else:
+                ff = f
+            ff = ff.subs(self.substitutions)
             sf.append(ff)
 
         return sf

qgs/diagnostics/base.py

@@ -28,12 +28,17 @@
 from IPython.display import HTML, display
 from ipywidgets import interactive
 
+from sympy import symbols
+
 from qgs.diagnostics.misc import *
 
 # TODO: - need to introduce an oro_basis specific to the orography !!
 #       - no orography plot when oro_basis is none
 
 
+_n = symbols('n', real=True, nonnegative=True)  # only this works with sympy substitution, don't ask me why...
+
+
 class Diagnostic(ABC):
     """General base class to create diagnostics.
 
@@ -65,6 +70,7 @@ def __init__(self, model_params, dimensional):
         self._diagnostic_data = None
         self._diagnostic_data_dimensional = None
         self._time = None
+        self._subs = list()
 
         self.dimensional = dimensional
 
@@ -114,6 +120,7 @@ def set_params(self, model_params, kwargs=None):
         """
 
         self._set_params(model_params)
+        self._subs = [(_n, model_params.scale_params.n)]
         if kwargs is not None:
             self._configure(**kwargs)
 

qgs/diagnostics/streamfunctions.py

@@ -97,7 +97,7 @@ def _configure(self, delta_x=None, delta_y=None):
         self._compute_grid(delta_x, delta_y)
         basis = self._model_params.atmospheric_basis
 
-        self._grid_basis = create_grid_basis(basis, self._X, self._Y)
+        self._grid_basis = create_grid_basis(basis, self._X, self._Y, self._subs)
 
 
 class LowerLayerAtmosphericStreamfunctionDiagnostic(AtmosphericStreamfunctionDiagnostic):
@@ -353,7 +353,7 @@ def _configure(self, model_params=None, delta_x=None, delta_y=None):
         self._compute_grid(delta_x, delta_y)
         basis = self._model_params.oceanic_basis
 
-        self._grid_basis = create_grid_basis(basis, self._X, self._Y)
+        self._grid_basis = create_grid_basis(basis, self._X, self._Y, self._subs)
 
 
 class OceanicLayerStreamfunctionDiagnostic(OceanicStreamfunctionDiagnostic):
@@ -396,7 +396,7 @@ def __init__(self, model_params, delta_x=None, delta_y=None, dimensional=True, c
         self._fields_average = list()
         basis = self._model_params.oceanic_basis
 
-        for func in basis.num_functions():
+        for func in basis.num_functions(self._subs):
             average = dblquad(func, 0, np.pi, 0, 2*np.pi/model_params.scale_params.n)
             self._fields_average.append(average[0])
 

qgs/diagnostics/temperatures.py

@@ -95,7 +95,7 @@ def _configure(self, model_params=None, delta_x=None, delta_y=None):
         self._compute_grid(delta_x, delta_y)
         basis = self._model_params.atmospheric_basis
 
-        self._grid_basis = create_grid_basis(basis, self._X, self._Y)
+        self._grid_basis = create_grid_basis(basis, self._X, self._Y, self._subs)
 
 
 class MiddleAtmosphericTemperatureAnomalyDiagnostic(AtmosphericTemperatureDiagnostic):
@@ -285,7 +285,7 @@ def _configure(self, model_params=None, delta_x=None, delta_y=None):
         self._compute_grid(delta_x, delta_y)
         basis = self._model_params.oceanic_basis
 
-        self._grid_basis = create_grid_basis(basis, self._X, self._Y)
+        self._grid_basis = create_grid_basis(basis, self._X, self._Y, self._subs)
 
 
 class OceanicLayerTemperatureAnomalyDiagnostic(OceanicTemperatureDiagnostic):
@@ -473,7 +473,7 @@ def _configure(self, model_params=None, delta_x=None, delta_y=None):
         self._compute_grid(delta_x, delta_y)
         basis = self._model_params.ground_basis
 
-        self._grid_basis = create_grid_basis(basis, self._X, self._Y)
+        self._grid_basis = create_grid_basis(basis, self._X, self._Y, self._subs)
 
     def _get_diagnostic(self, dimensional):
 

qgs/diagnostics/util.py

@@ -9,7 +9,7 @@
 import numpy as np
 
 
-def create_grid_basis(basis, X, Y):
+def create_grid_basis(basis, X, Y, extra_subs=None):
     """Create an array from a basis with the basis functions evaluated on a grid.
 
 
@@ -21,6 +21,10 @@ def create_grid_basis(basis, X, Y):
         The zonal x-coordinates of the grid points.
     Y: ~numpy.ndarray
         The meridional y-coordinates of the grid points.
+    extra_subs: list(tuple), optional
+        List of 2-tuples containing extra substitutions to be made with the functions providing
+        the grib basis, before transforming them into python callable.
+        The 2-tuples contain first a `Sympy`_  expression and then the value to substitute.
 
     Returns
     -------
@@ -31,7 +35,7 @@ def create_grid_basis(basis, X, Y):
 
     grid_basis = list()
 
-    for func in basis.num_functions():
+    for func in basis.num_functions(extra_subs):
         grid = func(X, Y)
         if isinstance(grid, (int, float)):
             grid = np.ones_like(X) * grid

qgs/diagnostics/variables.py

@@ -721,7 +721,7 @@ def _compute_functions_value(self):
 
         basis = self._model_params.atmospheric_basis
 
-        funcs_list = basis.num_functions()
+        funcs_list = basis.num_functions(self._subs)
 
         for point in self._point1:
             self._func_points1.append(list())

qgs/diagnostics/wind.py

@@ -105,11 +105,11 @@ def _configure(self, delta_x=None, delta_y=None):
 
         if self.type == "V":
             dx_basis = self._model_params.atmospheric_basis.x_derivative
-            self._grid_basis = create_grid_basis(dx_basis, self._X, self._Y)
+            self._grid_basis = create_grid_basis(dx_basis, self._X, self._Y, self._subs)
 
         elif self.type == "U":
             dy_basis = self._model_params.atmospheric_basis.y_derivative
-            self._grid_basis = create_grid_basis(dy_basis, self._X, self._Y)
+            self._grid_basis = create_grid_basis(dy_basis, self._X, self._Y, self._subs)
 
         elif self.type is None:
             warnings.warn("AtmosphericWindDiagnostic: Basis type note specified." +